From 840e95f8f37b545de797bc3e69206cc43f93eb62 Mon Sep 17 00:00:00 2001
From: Robert Goldsmith <rgoldsmith@names.co.uk>
Date: Mon, 11 Sep 2017 17:54:21 +0100
Subject: [PATCH] Add the ability to limit address resolution by type

Allows restricting to A or AAAA only via the ClientConnectContext.

Relates to #35, but doesn't fully fix it. Further work needs to be done
to avoid IPv6 automatically when IPv6 support is not available.
---
 lib/ClientConnectContext.php | 17 +++++++++++++++++
 lib/functions.php            |  2 +-
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/lib/ClientConnectContext.php b/lib/ClientConnectContext.php
index 4a01e64..39243cf 100644
--- a/lib/ClientConnectContext.php
+++ b/lib/ClientConnectContext.php
@@ -2,12 +2,14 @@
 
 namespace Amp\Socket;
 
+use Amp\Dns\Record;
 use function Amp\Socket\Internal\normalizeBindToOption;
 
 final class ClientConnectContext {
     private $bindTo = null;
     private $connectTimeout = 10000;
     private $maxAttempts = 2;
+    private $typeRestriction = null;
 
     public function withBindTo(string $bindTo = null): self {
         $bindTo = normalizeBindToOption($bindTo);
@@ -52,6 +54,21 @@ public function getMaxAttempts(): int {
         return $this->maxAttempts;
     }
 
+    public function withDnsTypeRestriction(int $type = null): self {
+        if ($type !== null && $type !== Record::AAAA && $type !== Record::A) {
+            throw new \Error("Invalid resolver type restriction");
+        }
+
+        $clone = clone $this;
+        $clone->typeRestriction = $type;
+
+        return $clone;
+    }
+
+    public function getDnsTypeRestriction() {
+        return $this->typeRestriction;
+    }
+
     public function toStreamContextArray(): array {
         $options = [];
 
diff --git a/lib/functions.php b/lib/functions.php
index 50e17e8..a6232b4 100644
--- a/lib/functions.php
+++ b/lib/functions.php
@@ -76,7 +76,7 @@ function connect(string $uri, ClientConnectContext $socketContext = null, Cancel
             $uris = [$uri];
         } else {
             // Host is not an IP address, so resolve the domain name.
-            $records = yield Dns\resolve($host);
+            $records = yield Dns\resolve($host, $socketContext->getDnsTypeRestriction());
             foreach ($records as $record) {
                 /** @var Dns\Record $record */
                 if ($record->getType() === Dns\Record::AAAA) {