Add localhost to default resolver

Author: jayvdb

dns_cache/__init__.py

@@ -7,9 +7,15 @@
 
 from .expiration import _NO_EXPIRY as NO_EXPIRY
 from .expiration import FIVE_MINS, MinExpirationCache, NoExpirationCache
+from .persistence import _LayeredCache
 from .pickle import PickableCache
 from .resolver import AggressiveCachingResolver, ExceptionCachingResolver
 
+try:
+    from .hosts import HostsCache
+except ImportError:
+    HostsCache = None
+
 __version__ = "0.2.0"
 
 
@@ -49,6 +55,9 @@ def override_system_resolver(
             else:
                 cache = MinExpirationCache(min_ttl=min_ttl)
 
+        if HostsCache:
+            cache = _LayeredCache(HostsCache(filename=None), cache)
+
     if not resolver:
         resolver = Resolver(configure=False)
         try:

tests/test_hosts.py

@@ -1,9 +1,67 @@
+import socket
 import unittest
 
+from dns.name import from_text
+from dns.rdataclass import IN
+from dns.rdatatype import A
+from dns.resolver import restore_system_resolver
+
+from dns_cache import override_system_resolver
+from dns_cache.block import dnspython_resolver_socket_block
 from dns_cache.hosts import loads
+from dns_cache.persistence import _LayeredCache
+from dns_cache.resolver import DNSPYTHON_2
+
+from .test_upstream import orig_gethostbyname
 
 
 class TestHostsSerializer(unittest.TestCase):
     def test_loads(self):
         data = loads()
         assert data
+
+
+class TestHostsCache(unittest.TestCase):
+
+    def test_hit_localhost(self):
+        name = "localhost"
+        assert socket.gethostbyname == orig_gethostbyname
+
+        try:
+            socket.gethostbyname(name)
+        except Exception as e:
+            raise unittest.SkipTest("gethostbyname: {}".format(e))
+
+        resolver = override_system_resolver()
+        assert isinstance(resolver.cache, _LayeredCache)
+
+        if DNSPYTHON_2:
+            query = resolver.resolve
+        else:
+            query = resolver.query
+
+        q1 = query(name)
+
+        assert len(resolver.cache._read_only_cache.data) >= 1
+        # The layering does a put, which pushes localhost into cache2
+        # TODO this needs to be blocked
+        assert len(resolver.cache._writable_cache.data) == 1
+
+        assert q1
+
+        name = from_text(name)
+
+        assert (name, A, IN) in resolver.cache.data
+        assert resolver.cache.get((name, A, IN))
+
+        with dnspython_resolver_socket_block():
+            q2 = query(name)
+
+        assert q2 is q1
+
+        with dnspython_resolver_socket_block():
+            ip = socket.gethostbyname(name)
+            assert ip == "127.0.0.1"
+
+        restore_system_resolver()
+        assert socket.gethostbyname == orig_gethostbyname

tests/test_persistence.py

@@ -3,6 +3,7 @@
 import os.path
 import pickle
 import shutil
+import socket
 import sys
 import unittest
 
@@ -26,6 +27,8 @@
     DNSPYTHON_2,
     dnspython_resolver_socket_block,
     get_test_resolver,
+    orig_gethostbyname,
+    restore_system_resolver,
 )
 
 try:
@@ -68,6 +71,10 @@ class _TestPersistentCacheBase(object):
     cache_cls = None
     kwargs = {}
 
+    def tearDown(self):
+        restore_system_resolver()
+        assert socket.gethostbyname == orig_gethostbyname
+
     def is_jsonpickle(self):
         serializer = self.kwargs.get("serializer", None)
         if hasattr(serializer, "startswith") and serializer.startswith(

tests/test_upstream.py

@@ -23,6 +23,7 @@
     Resolver,
     _getaddrinfo,
     override_system_resolver,
+    restore_system_resolver,
 )
 
 import pubdns
@@ -45,6 +46,7 @@
 PY2 = sys.version_info < (3, 0)
 
 pd = pubdns.PubDNS()
+orig_gethostbyname = socket.gethostbyname
 
 
 def compare_response(a, b):
@@ -144,6 +146,10 @@ class _TestCacheBase(object):
     cache_cls = Cache
     expiration = 60 * 5
 
+    def tearDown(self):
+        restore_system_resolver()
+        assert socket.gethostbyname == orig_gethostbyname
+
     def get_test_resolver(self, nameserver=None):
         resolver = get_test_resolver(self.resolver_cls, nameserver)
         resolver.cache = self.cache_cls()
@@ -571,6 +577,30 @@ def test_no_answer(self, expected_extra=0):
 
         return resolver
 
+    def test_miss_localhost(self):
+        name = "localhost"
+
+        assert socket.gethostbyname == orig_gethostbyname
+
+        try:
+            socket.gethostbyname(name)
+        except Exception as e:
+            raise unittest.SkipTest("gethostbyname: {}".format(e))
+
+        resolver = self.get_test_resolver()
+
+        if DNSPYTHON_2:
+            query = resolver.resolve
+        else:
+            query = resolver.query
+
+        with self.assertRaises(_SocketBlockedError):
+            with dnspython_resolver_socket_block():
+                query(name)
+
+        with self.assertRaises((NXDOMAIN, NoAnswer)):
+            query(name)
+
 
 @expand
 class TestCache(_TestCacheBase, unittest.TestCase):