Add TCP support to DNS server and client

README.md

@@ -37,12 +37,8 @@ $adapter = new Native('0.0.0.0', 5300);
 $zone = new Zone(
     name: 'example.test',
     records: [
-        new Record(name: 'example.test', type: Record::TYPE_A, rdata: '127.0.0.1', ttl: 60),
+The server listens on UDP and TCP port `5300` (RFC 5966) and answers queries for `example.test` from the in-memory zone. Implement the [`Utopia\DNS\Resolver`](src/DNS/Resolver.php) interface to serve records from databases, APIs, or other stores.
         new Record(name: 'www.example.test', type: Record::TYPE_CNAME, rdata: 'example.test', ttl: 60),
-        new Record(name: 'example.test', type: Record::TYPE_TXT, rdata: '"demo record"', ttl: 60),
-    ],
-    soa: new Record(
-        name: 'example.test',
         type: Record::TYPE_SOA,
         rdata: 'ns1.example.test hostmaster.example.test 1 7200 1800 1209600 3600',
         ttl: 60

composer.json

@@ -8,7 +8,7 @@
     "scripts": {
         "lint": "./vendor/bin/pint --test --config pint.json",
         "format": "./vendor/bin/pint --config pint.json",
-        "check": "./vendor/bin/phpstan analyse --level 8 -c phpstan.neon src tests",
+        "check": "./vendor/bin/phpstan analyse --level max -c phpstan.neon src tests",
         "test": "./vendor/bin/phpunit --configuration phpunit.xml"
     },
     "authors": [

docker-compose.yml

@@ -15,5 +15,6 @@ services:
       - dns
     ports:
       - '5300:5300/udp'
+      - '5300:5300/tcp'
 networks:
   dns:

src/DNS/Adapter.php

@@ -15,8 +15,8 @@ abstract public function onWorkerStart(callable $callback): void;
     /**
      * Packet handler
      *
-     * @param callable(string $buffer, string $ip, int $port): string $callback
-     * @phpstan-param callable(string $buffer, string $ip, int $port):string $callback
+     * @param callable(string $buffer, string $ip, int $port, ?int $maxResponseSize): string $callback
+     * @phpstan-param callable(string $buffer, string $ip, int $port, ?int $maxResponseSize):string $callback
      */
     abstract public function onPacket(callable $callback): void;
 

src/DNS/Adapter/Native.php

@@ -8,31 +8,46 @@
 
 class Native extends Adapter
 {
-    protected Socket $server;
+    protected Socket $udpServer;
 
-    /** @var callable(string $buffer, string $ip, int $port): string */
+    protected ?Socket $tcpServer = null;
+
+    /** @var array<int, Socket> */
+    protected array $tcpClients = [];
+
+    /** @var array<int, string> */
+    protected array $tcpBuffers = [];
+
+    /** @var callable(string $buffer, string $ip, int $port, ?int $maxResponseSize): string */
     protected mixed $onPacket;
 
     /** @var list<callable(int $workerId): void> */
     protected array $onWorkerStart = [];
 
-    protected string $host;
-    protected int $port;
-
-    /**
-     * @param string $host
-     * @param int $port
-     */
-    public function __construct(string $host = '0.0.0.0', int $port = 8053)
-    {
-        $this->host = $host;
-        $this->port = $port;
+    public function __construct(
+        protected string $host = '0.0.0.0',
+        protected int $port = 8053,
+        protected bool $enableTcp = true,
+        protected int $maxTcpClients = 100,
+        protected int $maxTcpBufferSize = 16384,
+        protected int $maxTcpFrameSize = 8192
+    ) {
 
         $server = \socket_create(AF_INET, SOCK_DGRAM, SOL_UDP);
         if (!$server) {
             throw new Exception('Could not start server.');
         }
-        $this->server = $server;
+        $this->udpServer = $server;
+
+        if ($this->enableTcp) {
+            $tcp = \socket_create(AF_INET, SOCK_STREAM, SOL_TCP);
+            if (!$tcp) {
+                throw new Exception('Could not start TCP server.');
+            }
+
+            socket_set_option($tcp, SOL_SOCKET, SO_REUSEADDR, 1);
+            $this->tcpServer = $tcp;
+        }
     }
 
     /**
@@ -48,7 +63,7 @@ public function onWorkerStart(callable $callback): void
 
     /**
      * @param callable $callback
-     * @phpstan-param callable(string $buffer, string $ip, int $port):string $callback
+     * @phpstan-param callable(string $buffer, string $ip, int $port, ?int $maxResponseSize):string $callback
      */
     public function onPacket(callable $callback): void
     {
@@ -60,27 +75,93 @@ public function onPacket(callable $callback): void
      */
     public function start(): void
     {
-        if (socket_bind($this->server, $this->host, $this->port) == false) {
+        if (socket_bind($this->udpServer, $this->host, $this->port) == false) {
             throw new Exception('Could not bind server to a server.');
         }
 
+        if ($this->tcpServer) {
+            if (socket_bind($this->tcpServer, $this->host, $this->port) == false) {
+                throw new Exception('Could not bind TCP server.');
+            }
+
+            if (socket_listen($this->tcpServer, 128) == false) {
+                throw new Exception('Could not listen on TCP server.');
+            }
+
+            socket_set_nonblock($this->tcpServer);
+        }
+
         foreach ($this->onWorkerStart as $callback) {
             \call_user_func($callback, 0);
         }
 
         /** @phpstan-ignore-next-line */
         while (1) {
-            $buf = '';
-            $ip = '';
-            $port = null;
-            $len = socket_recvfrom($this->server, $buf, 1024 * 4, 0, $ip, $port);
+            $readSockets = [$this->udpServer];
 
-            if ($len > 0) {
-                $answer = call_user_func($this->onPacket, $buf, $ip, $port);
+            if ($this->tcpServer) {
+                $readSockets[] = $this->tcpServer;
+            }
 
-                if (socket_sendto($this->server, $answer, strlen($answer), 0, $ip, $port) === false) {
-                    printf('Error in socket\n');
+            foreach ($this->tcpClients as $client) {
+                $readSockets[] = $client;
+            }
+
+            $write = [];
+            $except = [];
+
+            $changed = socket_select($readSockets, $write, $except, null);
+
+            if ($changed === false) {
+                continue;
+            }
+
+            foreach ($readSockets as $socket) {
+                if ($socket === $this->udpServer) {
+                    $buf = '';
+                    $ip = '';
+                    $port = 0;
+                    $len = socket_recvfrom($this->udpServer, $buf, 1024 * 4, 0, $ip, $port);
+
+                    if ($len > 0 && is_string($buf) && is_string($ip) && is_int($port)) {
+                        $answer = call_user_func($this->onPacket, $buf, $ip, $port, 512);
+
+                        if ($answer !== '') {
+                            socket_sendto($this->udpServer, $answer, strlen($answer), 0, $ip, $port);
+                        }
+                    }
+
+                    continue;
                 }
+
+                if ($this->tcpServer !== null && $socket === $this->tcpServer) {
+                    $client = @socket_accept($this->tcpServer);
+
+                    if ($client instanceof Socket) {
+                        if (count($this->tcpClients) >= $this->maxTcpClients) {
+                            @socket_close($client);
+                            continue;
+                        }
+
+                        if (@socket_set_nonblock($client) === false) {
+                            @socket_close($client);
+                            continue;
+                        }
+
+                        socket_set_option($client, SOL_SOCKET, SO_KEEPALIVE, 1);
+                        socket_set_option($client, SOL_SOCKET, SO_RCVTIMEO, ['sec' => 5, 'usec' => 0]);
+                        socket_set_option($client, SOL_SOCKET, SO_SNDTIMEO, ['sec' => 5, 'usec' => 0]);
+
+                        $id = spl_object_id($client);
+                        $this->tcpClients[$id] = $client;
+                        $this->tcpBuffers[$id] = '';
+                    }
+
+                    continue;
+                }
+
+                // Remaining readable sockets are TCP clients.
+                $this->handleTcpClient($socket);
             }
         }
     }
@@ -94,4 +175,97 @@ public function getName(): string
     {
         return 'native';
     }
+
+    protected function handleTcpClient(Socket $client): void
+    {
+        $clientId = spl_object_id($client);
+
+        $chunk = @socket_read($client, 8192, PHP_BINARY_READ);
+
+        if ($chunk === '' || $chunk === false) {
+            $error = socket_last_error($client);
+
+            if ($chunk === '' || !in_array($error, [SOCKET_EAGAIN, SOCKET_EWOULDBLOCK], true)) {
+                $this->closeTcpClient($client);
+            }
+
+            return;
+        }
+
+        $currentBufferSize = strlen($this->tcpBuffers[$clientId] ?? '');
+        $chunkSize = strlen($chunk);
+
+        if ($currentBufferSize + $chunkSize > $this->maxTcpBufferSize) {
+            printf("TCP buffer size limit exceeded for client %d\n", $clientId);
+            $this->closeTcpClient($client);
+            return;
+        }
+
+        $this->tcpBuffers[$clientId] = ($this->tcpBuffers[$clientId] ?? '') . $chunk;
+
+        while (strlen($this->tcpBuffers[$clientId]) >= 2) {
+            $unpacked = unpack('n', substr($this->tcpBuffers[$clientId], 0, 2));
+            $payloadLength = (is_array($unpacked) && array_key_exists(1, $unpacked) && is_int($unpacked[1])) ? $unpacked[1] : 0;
+
+            if ($payloadLength > $this->maxTcpFrameSize) {
+                printf("Invalid TCP frame size %d for client %d\n", $payloadLength, $clientId);
+                $this->closeTcpClient($client);
+                return;
+            }
+
+            if (strlen($this->tcpBuffers[$clientId]) < ($payloadLength + 2)) {
+                return;
+            }
+
+            $message = substr($this->tcpBuffers[$clientId], 2, $payloadLength);
+            $this->tcpBuffers[$clientId] = substr($this->tcpBuffers[$clientId], $payloadLength + 2);
+
+            $ip = '';
+            $port = 0;
+            socket_getpeername($client, $ip, $port);
+
+            if (is_string($ip) && is_int($port)) {
+                $answer = call_user_func($this->onPacket, $message, $ip, $port, null);
+
+                if ($answer !== '') {
+                    $this->sendTcpResponse($client, $answer);
+                }
+            }
+        }
+    }
+
+    protected function sendTcpResponse(Socket $client, string $payload): void
+    {
+        $frame = pack('n', strlen($payload)) . $payload;
+        $total = strlen($frame);
+        $sent = 0;
+
+        while ($sent < $total) {
+            $written = @socket_write($client, substr($frame, $sent));
+
+            if ($written === false) {
+                $error = socket_last_error($client);
+
+                if (in_array($error, [SOCKET_EAGAIN, SOCKET_EWOULDBLOCK], true)) {
+                    socket_clear_error($client);
+                    usleep(1000);
+                    continue;
+                }
+
+                $this->closeTcpClient($client);
+                return;
+            }
+
+            $sent += $written;
+        }
+    }
+
+    protected function closeTcpClient(Socket $client): void
+    {
+        $id = spl_object_id($client);
+
+        unset($this->tcpClients[$id], $this->tcpBuffers[$id]);
+
+        @socket_close($client);
+    }
 }