From f9f0c6cdc8fd6586566c9c9f908cc360db274be6 Mon Sep 17 00:00:00 2001
From: Niklas Keller <me@kelunik.com>
Date: Thu, 29 Jun 2017 12:36:09 +0200
Subject: [PATCH] Fix queries over TCP, add tests

---
 composer.json              |  5 ++++
 lib/Internal/TcpSocket.php |  2 +-
 test/SocketTest.php        | 31 +++++++++++++++++++++++
 test/TcpSocketTest.php     | 52 ++++++++++++++++++++++++++++++++++++++
 test/UdpSocketTest.php     | 18 +++++++++++++
 5 files changed, 107 insertions(+), 1 deletion(-)
 create mode 100644 test/SocketTest.php
 create mode 100644 test/TcpSocketTest.php
 create mode 100644 test/UdpSocketTest.php

diff --git a/composer.json b/composer.json
index 9c57cf4..b23496d 100644
--- a/composer.json
+++ b/composer.json
@@ -56,5 +56,10 @@
         "files": [
             "lib/functions.php"
         ]
+    },
+    "autoload-dev": {
+        "psr-4": {
+            "Amp\\Dns\\Test\\": "test"
+        }
     }
 }
diff --git a/lib/Internal/TcpSocket.php b/lib/Internal/TcpSocket.php
index dadf186..9a0e41d 100644
--- a/lib/Internal/TcpSocket.php
+++ b/lib/Internal/TcpSocket.php
@@ -94,7 +94,7 @@ class TcpSocket extends Socket {
         if ($this->queue->isEmpty()) {
             return call(function () {
                 do {
-                    $chunk = $this->read();
+                    $chunk = yield $this->read();
 
                     if ($chunk === null) {
                         $this->isAlive = false;
diff --git a/test/SocketTest.php b/test/SocketTest.php
new file mode 100644
index 0000000..07f907a
--- /dev/null
+++ b/test/SocketTest.php
@@ -0,0 +1,31 @@
+<?php
+
+namespace Amp\Dns\Test;
+
+use Amp\Dns;
+use Amp\Loop;
+use Amp\PHPUnit\TestCase;
+use Amp\Promise;
+use LibDNS\Messages\Message;
+use LibDNS\Messages\MessageTypes;
+use LibDNS\Records\QuestionFactory;
+
+abstract class SocketTest extends TestCase {
+    abstract protected function connect(): Promise;
+
+    public function testAsk() {
+        Loop::run(function () {
+            $question = (new QuestionFactory)->create(Dns\Record::A);
+            $question->setName("google.com");
+
+            /** @var Dns\Internal\Socket $socket */
+            $socket = yield $this->connect();
+
+            /** @var Message $result */
+            $result = yield $socket->ask($question, 5000);
+
+            $this->assertInstanceOf(Message::class, $result);
+            $this->assertSame(MessageTypes::RESPONSE, $result->getType());
+        });
+    }
+}
diff --git a/test/TcpSocketTest.php b/test/TcpSocketTest.php
new file mode 100644
index 0000000..41c3583
--- /dev/null
+++ b/test/TcpSocketTest.php
@@ -0,0 +1,52 @@
+<?php
+
+namespace Amp\Dns\Test;
+
+use Amp\Delayed;
+use Amp\Dns;
+use Amp\Loop;
+use Amp\Promise;
+use LibDNS\Messages\Message;
+use LibDNS\Messages\MessageTypes;
+use LibDNS\Records\QuestionFactory;
+use function Amp\Promise\wait;
+
+class TcpSocketTest extends SocketTest {
+    protected function connect(): Promise {
+        return Dns\Internal\TcpSocket::connect("tcp://8.8.8.8:53");
+    }
+
+    public function testTimeout() {
+        $this->expectException(Dns\TimeoutException::class);
+        wait(Dns\Internal\TcpSocket::connect("tcp://8.8.8.8:53", 0));
+    }
+
+    public function testInvalidUri() {
+        $this->expectException(Dns\ResolutionException::class);
+        wait(Dns\Internal\TcpSocket::connect("tcp://8.8.8.8"));
+    }
+
+    public function testAfterConnectionTimedOut() {
+        Loop::run(function () {
+            $question = (new QuestionFactory)->create(Dns\Record::A);
+            $question->setName("google.com");
+
+            /** @var Dns\Internal\Socket $socket */
+            $socket = yield $this->connect();
+
+            /** @var Message $result */
+            $result = yield $socket->ask($question, 3000);
+
+            $this->assertInstanceOf(Message::class, $result);
+            $this->assertSame(MessageTypes::RESPONSE, $result->getType());
+
+            // Google's DNS times out really fast
+            yield new Delayed(3000);
+
+            $this->expectException(Dns\ResolutionException::class);
+            $this->expectExceptionMessage("Reading from the server failed");
+
+            yield $socket->ask($question, 3000);
+        });
+    }
+}
diff --git a/test/UdpSocketTest.php b/test/UdpSocketTest.php
new file mode 100644
index 0000000..cdaa772
--- /dev/null
+++ b/test/UdpSocketTest.php
@@ -0,0 +1,18 @@
+<?php
+
+namespace Amp\Dns\Test;
+
+use Amp\Dns;
+use Amp\Promise;
+use function Amp\Promise\wait;
+
+class UdpSocketTest extends SocketTest {
+    protected function connect(): Promise {
+        return Dns\Internal\UdpSocket::connect("udp://8.8.8.8:53");
+    }
+
+    public function testInvalidUri() {
+        $this->expectException(Dns\ResolutionException::class);
+        wait(Dns\Internal\UdpSocket::connect("udp://8.8.8.8"));
+    }
+}