MulticastDns.php

<?php

namespace Kelunik\LightD;

use Amp\Deferred;
use Amp\Delayed;
use Amp\Dns\DnsException;
use Amp\Dns\Record;
use Amp\Loop;
use Amp\Promise;
use LibDNS\Decoder\Decoder;
use LibDNS\Decoder\DecodingContextFactory;
use LibDNS\Encoder\Encoder;
use LibDNS\Encoder\EncodingContextFactory;
use LibDNS\Messages\MessageFactory;
use LibDNS\Messages\MessageTypes;
use LibDNS\Packets\PacketFactory;
use LibDNS\Records\QuestionFactory;
use LibDNS\Records\RDataBuilder;
use LibDNS\Records\RDataFactory;
use LibDNS\Records\ResourceBuilder;
use LibDNS\Records\ResourceFactory;
use LibDNS\Records\TypeDefinitions\FieldDefinitionFactory;
use LibDNS\Records\TypeDefinitions\TypeDefinitionFactory;
use LibDNS\Records\TypeDefinitions\TypeDefinitionManager;
use LibDNS\Records\Types\TypeBuilder;
use LibDNS\Records\Types\TypeFactory;
use function Amp\call;

final class MulticastDns
{
    public function query(string $name): Promise
    {
        return call(static function () use ($name) {
            for ($attempt = 1; $attempt < 3; $attempt++) {
                $context = \stream_context_create([
                    'socket' => [
                        'so_reuseport' => true,
                        'so_reuseaddr' => true,
                        'bindto' => '0.0.0.0:5353',
                    ],
                ]);

                $stream = \stream_socket_server('udp://0.0.0.0:5353', $errno, $error, \STREAM_SERVER_BIND, $context);
                $socket = \socket_import_stream($stream);

                if (!$socket) {
                    throw new \RuntimeException('Failed to access socket');
                }

                if (!\socket_set_option($socket, \IPPROTO_IP, \MCAST_JOIN_GROUP,
                    ['group' => '224.0.0.251', 'interface' => 0])) {
                    throw new \RuntimeException('Failed to join multicast group');
                }

                if (!\socket_set_option($socket, \IPPROTO_IP, \IP_MULTICAST_TTL, 255)) {
                    throw new \RuntimeException('Failed to set multicast TTL');
                }

                if (!\socket_set_option($socket, \IPPROTO_IP, \IP_MULTICAST_LOOP, 1)) {
                    throw new \RuntimeException('Failed to set multicast loop');
                }

                $question = (new QuestionFactory)->create(Record::PTR);
                $question->setName($name);

                $request = (new MessageFactory)->create(MessageTypes::QUERY);
                $request->getQuestionRecords()->add($question);
                $request->isRecursionDesired(false);
                $request->setID(0);

                $data = (new Encoder(new PacketFactory, new EncodingContextFactory))->encode($request);
                \stream_socket_sendto($stream, $data, null, '224.0.0.251:5353');

                $timeout = new Delayed(3000, false);
                $timeout->unreference();

                $devices = [];

                read:
                $deferred = new Deferred;

                $watcher = Loop::onReadable($stream, static function ($watcher) use ($deferred) {
                    Loop::cancel($watcher);
                    $deferred->resolve(true);
                });

                if (!yield Promise\first([$deferred->promise(), $timeout])) {
                    Loop::cancel($watcher);

                    if ($devices) {
                        return $devices;
                    }

                    continue;
                }

                $data = \stream_socket_recvfrom($stream, 8192, null, $addr);
                if ($data === null || $data === false) {
                    continue;
                }

                $typeBuilder = new TypeBuilder(new TypeFactory);

                $resourceBuilder = new ResourceBuilder(
                    new ResourceFactory,
                    new RDataBuilder(new RDataFactory, $typeBuilder),
                    new TypeDefinitionManager(new TypeDefinitionFactory, new FieldDefinitionFactory)
                );

                $response = (new Decoder(
                    new PacketFactory(),
                    new MessageFactory(),
                    new QuestionFactory(),
                    $resourceBuilder,
                    $typeBuilder,
                    new DecodingContextFactory
                ))->decode($data);

                if ($response->getType() !== MessageTypes::RESPONSE) {
                    goto read; // loop back or a question
                }

                $answers = $response->getAnswerRecords();
                $pointers = $answers->getRecordsByName($name);

                if (!$pointers) {
                    goto read; // response to another query
                }

                $records = [];

                foreach ($answers as $answer) {
                    $records[] = $answer;
                }

                foreach ($response->getAdditionalRecords() as $answer) {
                    $records[] = $answer;
                }

                $hosts = [];

                foreach ($records as $answer) {
                    \assert($answer instanceof \LibDNS\Records\Resource);

                    if ($answer->getType() === Record::A) {
                        $hosts[] = $answer->getData()->getFieldByName('address')->getValue();
                    }

                    if ($answer->getType() === Record::AAAA) {
                        $hosts[] = '[' . $answer->getData()->getFieldByName('address')->getValue() . ']';
                    }
                }

                if (empty($hosts)) {
                    throw new DnsException('Failed to find A or AAAA record for ' . $name);
                }

                // automatically prefers IPv4, because IPv6 starts with '[' and IPv4 starts with digits
                \sort($hosts);

                $devices[] = $hosts;

                goto read;
            }

            throw new DnsException('Query for ' . $name . ' timed out');
        });
    }
}