Roll tsproxy to the latest version

telemetry/third_party/tsproxy/tsproxy.py

 #!/usr/bin/python
 """
 Copyright 2016 Google Inc. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 
     http://www.apache.org/licenses/LICENSE-2.0
 
@@ skipping 139 matching lines @@
       self.available_bytes = .0
     self.last_tick = now
 
     return processed_messages
 
 
 ########################################################################################################################
 #   Threaded DNS resolver
 ########################################################################################################################
 class AsyncDNS(threading.Thread):
-  def __init__(self, client_id, hostname, port, result_pipe):
+  def __init__(self, client_id, hostname, port, is_localhost, result_pipe):
     threading.Thread.__init__(self)
     self.hostname = hostname
     self.port = port
     self.client_id = client_id
+    self.is_localhost = is_localhost
     self.result_pipe = result_pipe
 
   def run(self):
     global lock, background_activity_count
     try:
       logging.debug('[{0:d}] AsyncDNS - calling getaddrinfo for {1}:{2:d}'.format(self.client_id, self.hostname, self.port))
       addresses = socket.getaddrinfo(self.hostname, self.port)
       logging.info('[{0:d}] Resolving {1}:{2:d} Completed'.format(self.client_id, self.hostname, self.port))
     except:
       addresses = ()
       logging.info('[{0:d}] Resolving {1}:{2:d} Failed'.format(self.client_id, self.hostname, self.port))
-    message = {'message': 'resolved', 'connection': self.client_id, 'addresses': addresses}
+    message = {'message': 'resolved', 'connection': self.client_id, 'addresses': addresses, 'localhost': self.is_localhost}
     self.result_pipe.SendMessage(message, False)
     lock.acquire()
     if background_activity_count > 0:
       background_activity_count -= 1
     lock.release()
     # open and close a local socket which will interrupt the long polling loop to process the message
     s = socket.socket()
     s.connect((server.ipaddr, server.port))
     s.close()
 
@@ skipping 13 matching lines @@
     asyncore.dispatcher.__init__(self)
     self.client_id = client_id
     self.state = self.STATE_IDLE
     self.buffer = ''
     self.addr = None
     self.dns_thread = None
     self.hostname = None
     self.port = None
     self.needs_config = True
     self.needs_close = False
-    self.is_localhost = False
     self.did_resolve = False
 
   def SendMessage(self, type, message):
     message['message'] = type
     message['connection'] = self.client_id
     in_pipe.SendMessage(message)
 
   def handle_message(self, message):
     if message['message'] == 'data' and 'data' in message and len(message['data']):
       self.buffer += message['data']
@@ skipping 64 matching lines @@
             logging.debug('[{0:d}] TCP <= {1:d} byte(s)'.format(self.client_id, len(data)))
             self.SendMessage('data', {'data': data})
         else:
           return
     except:
       pass
 
   def HandleResolve(self, message):
     global in_pipe,  map_localhost, lock, background_activity_count
     self.did_resolve = True
+    is_localhost = False
     if 'hostname' in message:
       self.hostname = message['hostname']
     self.port = 0
     if 'port' in message:
       self.port = message['port']
     logging.info('[{0:d}] Resolving {1}:{2:d}'.format(self.client_id, self.hostname, self.port))
     if self.hostname == 'localhost':
       self.hostname = '127.0.0.1'
     if self.hostname == '127.0.0.1':
       logging.info('[{0:d}] Connection to localhost detected'.format(self.client_id))
-      self.is_localhost = True
-    if (dest_addresses is not None) and (not self.is_localhost or map_localhost):
+      is_localhost = True
+    if (dest_addresses is not None) and (not is_localhost or map_localhost):
       logging.info('[{0:d}] Resolving {1}:{2:d} to mapped address {3}'.format(self.client_id, self.hostname, self.port, dest_addresses))
-      self.SendMessage('resolved', {'addresses': dest_addresses})
+      self.SendMessage('resolved', {'addresses': dest_addresses, 'localhost': False})
     else:
       lock.acquire()
       background_activity_count += 1
       lock.release()
       self.state = self.STATE_RESOLVING
-      self.dns_thread = AsyncDNS(self.client_id, self.hostname, self.port, in_pipe)
+      self.dns_thread = AsyncDNS(self.client_id, self.hostname, self.port, is_localhost, in_pipe)
       self.dns_thread.start()
 
   def HandleConnect(self, message):
     global map_localhost
     if 'addresses' in message and len(message['addresses']):
       self.state = self.STATE_CONNECTING
-      if not self.did_resolve and message['addresses'][0] == '127.0.0.1':
+      is_localhost = False
+      if 'localhost' in message:
+        is_localhost = message['localhost']
+      elif not self.did_resolve and message['addresses'][0] == '127.0.0.1':
         logging.info('[{0:d}] Connection to localhost detected'.format(self.client_id))
-        self.is_localhost = True
-      if (dest_addresses is not None) and (not self.is_localhost or map_localhost):
+        is_localhost = True
+      if (dest_addresses is not None) and (not is_localhost or map_localhost):
         self.addr = dest_addresses[0]
       else:
         self.addr = message['addresses'][0]
       self.create_socket(self.addr[0], socket.SOCK_STREAM)
       addr = self.addr[4][0]
-      if not self.is_localhost or map_localhost:
+      if not is_localhost or map_localhost:
         port = GetDestPort(message['port'])
       else:
         port = message['port']
       logging.info('[{0:d}] Connecting to {1}:{2:d}'.format(self.client_id, addr, port))
       self.connect((addr, port))
 
 
 ########################################################################################################################
 #   Socks5 Server
 ########################################################################################################################
@@ skipping 136 matching lines @@
                 for i in range(16):
                   self.ip += '{0:02x}'.format(ord(data[4 + i]))
                   if i % 2 and i < 15:
                     self.ip += ':'
               if port_offset and connections[self.client_id]['server'] is not None:
                 self.port = 256 * ord(data[port_offset]) + ord(data[port_offset + 1])
                 if self.port:
                   if self.ip is None and self.hostname is not None:
                     if self.hostname in dns_cache:
                       self.state = self.STATE_CONNECTING
-                      self.addresses = dns_cache[self.hostname]
-                      self.SendMessage('connect', {'addresses': self.addresses, 'port': self.port})
+                      cache_entry = dns_cache[self.hostname]
+                      self.addresses = cache_entry['addresses']
+                      self.SendMessage('connect', {'addresses': self.addresses, 'port': self.port, 'localhost': cache_entry['localhost']})
                     else:
                       self.state = self.STATE_RESOLVING
                       self.SendMessage('resolve', {'hostname': self.hostname, 'port': self.port})
                   elif self.ip is not None:
                     self.state = self.STATE_CONNECTING
                     logging.debug('[{0:d}] Socks Connect - calling getaddrinfo for {1}:{2:d}'.format(self.client_id, self.ip, self.port))
                     self.addresses = socket.getaddrinfo(self.ip, self.port)
                     self.SendMessage('connect', {'addresses': self.addresses, 'port': self.port})
         else:
           return
@@ skipping 14 matching lines @@
           del connections[self.client_id]
     except:
       pass
 
   def HandleResolved(self, message):
     global dns_cache
     if self.state == self.STATE_RESOLVING:
       if 'addresses' in message and len(message['addresses']):
         self.state = self.STATE_CONNECTING
         self.addresses = message['addresses']
-        dns_cache[self.hostname] = self.addresses
+        dns_cache[self.hostname] = {'addresses': self.addresses, 'localhost': message['localhost']}
         logging.debug('[{0:d}] Resolved {1}, Connecting'.format(self.client_id, self.hostname))
-        self.SendMessage('connect', {'addresses': self.addresses, 'port': self.port})
+        self.SendMessage('connect', {'addresses': self.addresses, 'port': self.port, 'localhost': message['localhost']})
       else:
         # Send host unreachable error
         self.state = self.STATE_ERROR
         self.buffer += chr(0x05) + chr(0x04) + self.requested_address
         self.handle_write()
 
   def HandleConnected(self, message):
     if 'success' in message and self.state == self.STATE_CONNECTING:
       response = chr(0x05)
       if message['success']:
@@ skipping 231 matching lines @@
     if src == '*':
       port_mappings['default'] = int(dest)
       logging.debug("Default port mapped to port {0}".format(dest))
     else:
       logging.debug("Port {0} mapped to port {1}".format(src, dest))
       port_mappings[src] = int(dest)
 
 
 if '__main__' == __name__:
   main()