Properly filter to only return IPv4 addresses for type A queries

Author: clue

src/Query/HostsFileExecutor.php

@@ -28,13 +28,15 @@ public function __construct(HostsFile $hosts, ExecutorInterface $fallback)
     public function query($nameserver, Query $query)
     {
         if ($query->class === Message::CLASS_IN && $query->type === Message::TYPE_A) {
-            $ips = $this->hosts->getIpsForHost($query->name);
-            if ($ips) {
-                $records = array();
-                foreach ($ips as $ip) {
+            $records = array();
+            foreach ($this->hosts->getIpsForHost($query->name) as $ip) {
+                // ensure this is an IPv4 address
+                if (strpos($ip, ':') === false) {
                     $records[] = new Record($query->name, $query->type, $query->class, 0, $ip);
                 }
+            }
 
+            if ($records) {
                 return Promise\resolve(
                     Message::createResponseWithAnswersForQuery($query, $records)
                 );

tests/Config/HostsFileTest.php

@@ -77,4 +77,11 @@ public function testMergesEntriesOverMultipleLines()
 
         $this->assertEquals(array('127.0.0.1', '127.0.0.2', '127.0.0.3', '127.0.0.4'), $hosts->getIpsForHost('localhost'));
     }
+
+    public function testMergesIpv4AndIpv6EntriesOverMultipleLines()
+    {
+        $hosts = new HostsFile("127.0.0.1 localhost\n::1 localhost");
+
+        $this->assertEquals(array('127.0.0.1', '::1'), $hosts->getIpsForHost('localhost'));
+    }
 }

tests/Query/HostsFileExecutorTest.php

@@ -43,4 +43,12 @@ public function testReturnsResponseMessageIfIpsWereFound()
 
         $ret = $this->executor->query('8.8.8.8', new Query('google.com', Message::TYPE_A, Message::CLASS_IN, 0));
     }
+
+    public function testFallsBackIfNoIpv4Matches()
+    {
+        $this->hosts->expects($this->once())->method('getIpsForHost')->willReturn(array('::1'));
+        $this->fallback->expects($this->once())->method('query');
+
+        $ret = $this->executor->query('8.8.8.8', new Query('google.com', Message::TYPE_A, Message::CLASS_IN, 0));
+    }
 }