From a8c9cee1c3e8b68b59461834eee08460aa8212c4 Mon Sep 17 00:00:00 2001
From: Niklas Keller <me@kelunik.com>
Date: Tue, 10 Jan 2023 23:29:50 +0100
Subject: [PATCH] Ignore malformed UDP packets

Relates to #102.
---
 src/Internal/Socket.php    | 22 +++++++++++++++++-----
 src/Internal/UdpSocket.php | 18 +++++++++++++-----
 test/UdpSocketTest.php     | 21 +++++++++++++++++++++
 3 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/src/Internal/Socket.php b/src/Internal/Socket.php
index 35efd93..1f14620 100755
--- a/src/Internal/Socket.php
+++ b/src/Internal/Socket.php
@@ -28,6 +28,8 @@ abstract class Socket
 
     private const MAX_CONCURRENT_REQUESTS = 500;
 
+    protected int $invalidPacketsReceived = 0;
+
     abstract public static function connect(string $uri): self;
 
     private readonly ReadableResourceStream $input;
@@ -159,8 +161,10 @@ abstract class Socket
         /** @var DeferredFuture<\Closure():Message> $deferred */
         $deferred = new DeferredFuture;
 
+        $invalidPacketsReceived = &$this->invalidPacketsReceived;
+
         /** @psalm-suppress InaccessibleProperty $this->pending is an ArrayObject */
-        $this->pending[$id] = new class($this->pending, $id, $deferred, $question, $timeout) {
+        $this->pending[$id] = new class($this->pending, $id, $deferred, $question, $timeout, $invalidPacketsReceived) {
             private readonly string $callbackId;
 
             public ?DeferredFuture $deferred;
@@ -171,15 +175,23 @@ abstract class Socket
                 DeferredFuture $deferred,
                 public readonly Question $question,
                 float $timeout,
+                int &$invalidPacketsReceived
             ) {
                 $this->deferred = $deferred;
 
                 $this->callbackId = EventLoop::unreference(EventLoop::delay(
                     $timeout,
-                    weakClosure(function () use ($id, $pending, $timeout): void {
-                        $this->deferred?->complete(static fn () => throw new DnsTimeoutException(
-                            "Didn't receive a response within {$timeout} seconds."
-                        ));
+                    weakClosure(function () use ($id, $pending, $timeout, &$invalidPacketsReceived): void {
+                        if ($invalidPacketsReceived > 0) {
+                            $this->deferred?->complete(static fn () => throw new DnsTimeoutException(
+                                "Didn't receive a response within {$timeout} seconds, but received {$invalidPacketsReceived} invalid packets on this socket"
+                            ));
+                        } else {
+                            $this->deferred?->complete(static fn () => throw new DnsTimeoutException(
+                                "Didn't receive a response within {$timeout} seconds."
+                            ));
+                        }
+
                         $this->deferred = null;
 
                         unset($pending[$id]);
diff --git a/src/Internal/UdpSocket.php b/src/Internal/UdpSocket.php
index 4ebdd5c..a449f1b 100644
--- a/src/Internal/UdpSocket.php
+++ b/src/Internal/UdpSocket.php
@@ -47,12 +47,20 @@ final class UdpSocket extends Socket
 
     protected function receive(): Message
     {
-        $data = $this->read();
+        while (true) {
+            $data = $this->read();
 
-        if ($data === null) {
-            throw new DnsException("Reading from the server failed");
+            if ($data === null) {
+                throw new DnsException("Reading from the server failed");
+            }
+
+            try {
+                return $this->decoder->decode($data);
+            } catch (\Exception) {
+                $this->invalidPacketsReceived++;
+
+                continue;
+            }
         }
-
-        return $this->decoder->decode($data);
     }
 }
diff --git a/test/UdpSocketTest.php b/test/UdpSocketTest.php
index 5e388c2..4398d1a 100644
--- a/test/UdpSocketTest.php
+++ b/test/UdpSocketTest.php
@@ -3,6 +3,8 @@
 namespace Amp\Dns\Test;
 
 use Amp\Dns;
+use LibDNS\Records\QuestionFactory;
+use Revolt\EventLoop;
 
 class UdpSocketTest extends SocketTest
 {
@@ -12,6 +14,25 @@ class UdpSocketTest extends SocketTest
         Dns\Internal\UdpSocket::connect("udp://8.8.8.8");
     }
 
+    public function testMalformedResponse(): void
+    {
+        $server = \stream_socket_server('udp://127.0.0.1:0', $errno, $errstr, \STREAM_SERVER_BIND);
+        EventLoop::onReadable($server, function () use ($server) {
+            \stream_socket_recvfrom($server, 512, 0, $addr);
+            \stream_socket_sendto($server, 'invalid', 0, $addr);
+        });
+
+        $question = (new QuestionFactory)->create(Dns\DnsRecord::A);
+        $question->setName("google.com");
+
+        $socket = Dns\Internal\UdpSocket::connect('udp://' . \stream_socket_get_name($server, false));
+
+        $this->expectException(Dns\DnsTimeoutException::class);
+        $this->expectErrorMessage("Didn't receive a response within 1 seconds, but received 1 invalid packets on this socket");
+
+        $socket->ask($question, 1);
+    }
+
     protected function connect(): Dns\Internal\Socket
     {
         return Dns\Internal\UdpSocket::connect("udp://8.8.8.8:53");