Merge pull request #87 from andrey-malets/fix-dnsproxy

Fix dnsproxy

dnsproxy.py

@@ -34,11 +34,16 @@ class DnsProxyException(Exception):
   pass
 
 
+DEFAULT_DNS_PORT = 53
+
+
 class RealDnsLookup(object):
-  def __init__(self, name_servers):
-    if '127.0.0.1' in name_servers:
+  def __init__(self, name_servers, dns_forwarding, proxy_host, proxy_port):
+    if (proxy_host in name_servers and proxy_port == DEFAULT_DNS_PORT and
+        dns_forwarding):
       raise DnsProxyException(
-          'Invalid nameserver: 127.0.0.1 (causes an infinte loop)')
+          'Invalid nameserver: %s (causes an infinte loop)'.format(
+              proxy_host))
     self.resolver = resolver.get_default_resolver()
     self.resolver.nameservers = name_servers
     self.dns_cache_lock = threading.Lock()

httpclient_test.py

@@ -182,16 +182,17 @@ def test_ssl_get_connection_with_proxy_connects_to_proxy_ip(self):
   def test_ssl_get_connection_with_proxy_tunnels_to_host(self):
     """HTTPS (SSL) connection with proxy tunnels to target host."""
     self.set_https_proxy(host='proxy.com', port=8443)
-    connection = self.fetch._get_connection('example.com', None, is_ssl=True)
+    connection = self.fetch._get_connection('example.com', 9443, is_ssl=True)
     self.assertEqual('example.com', connection._tunnel_host)  # host name
-    self.assertEqual(None, connection._tunnel_port)  # host port
+    self.assertEqual(9443, connection._tunnel_port)  # host port
 
 
 class ActualNetworkFetchTest(test_utils.RealNetworkFetchTest):
 
   def testFetchNonSSLRequest(self):
     real_dns_lookup = dnsproxy.RealDnsLookup(
-        name_servers=[platformsettings.get_original_primary_nameserver()])
+        name_servers=[platformsettings.get_original_primary_nameserver()],
+        dns_forwarding=False, proxy_host='127.0.0.1', proxy_port=5353)
     fetch = httpclient.RealHttpFetch(real_dns_lookup)
     request = httparchive.ArchivedHttpRequest(
         command='GET', host='google.com', full_path='/search?q=dogs',
@@ -201,7 +202,8 @@ def testFetchNonSSLRequest(self):
 
   def testFetchSSLRequest(self):
     real_dns_lookup = dnsproxy.RealDnsLookup(
-        name_servers=[platformsettings.get_original_primary_nameserver()])
+        name_servers=[platformsettings.get_original_primary_nameserver()],
+        dns_forwarding=False, proxy_host='127.0.0.1', proxy_port=5353)
     fetch = httpclient.RealHttpFetch(real_dns_lookup)
     request = httparchive.ArchivedHttpRequest(
         command='GET', host='google.com', full_path='/search?q=dogs',

replay.py

@@ -315,8 +315,6 @@ def replay(options, replay_filename):
   if options.server:
     AddDnsForward(server_manager, options.server)
   else:
-    real_dns_lookup = dnsproxy.RealDnsLookup(
-        name_servers=[platformsettings.get_original_primary_nameserver()])
     if options.record:
       httparchive.HttpArchive.AssertWritable(replay_filename)
       if options.append and os.path.exists(replay_filename):
@@ -329,7 +327,6 @@ def replay(options, replay_filename):
       http_archive = httparchive.HttpArchive.Load(replay_filename)
       logging.info('Loaded %d responses from %s',
                    len(http_archive), replay_filename)
-    server_manager.AppendRecordCallback(real_dns_lookup.ClearCache)
     server_manager.AppendRecordCallback(http_archive.clear)
 
     ipfw_dns_host = None
@@ -340,6 +337,13 @@ def replay(options, replay_filename):
         ipfw_dns_host = platformsettings.get_server_ip_address(
             options.server_mode)
 
+    real_dns_lookup = dnsproxy.RealDnsLookup(
+        name_servers=[platformsettings.get_original_primary_nameserver()],
+        dns_forwarding=options.dns_forwarding,
+        proxy_host=ipfw_dns_host,
+        proxy_port=options.dns_port)
+    server_manager.AppendRecordCallback(real_dns_lookup.ClearCache)
+
     if options.dns_forwarding:
       if not options.server_mode and ipfw_dns_host == '127.0.0.1':
         AddDnsForward(server_manager, ipfw_dns_host)