Substitute HTTP Host header when proxying and fix HTTP socket reading.

vaurien/protocols/http.py

@@ -4,11 +4,9 @@
 from vaurien.util import chunked
 
 
-RE_LEN = re.compile('Content-Length: (\d+)', re.M | re.I)
 RE_KEEPALIVE = re.compile('Connection: Keep-Alive')
-RE_MEMCACHE_COMMAND = re.compile('(.*)\r\n')
+HOST_REPLACE = re.compile(r'\r\nHost: .+\r\n')
 
-EOH = '\r\n\r\n'
 CRLF = '\r\n'
 
 
@@ -22,6 +20,7 @@ def _handle(self, source, dest, to_backend):
 
         # Getting the HTTP query
         data = self._get_data(source)
+        data = HOST_REPLACE.sub('\r\nHost: %s\r\n' % self.proxy.backend, data)
 
         if not data:
             self._abort_handling(to_backend, dest)
@@ -31,39 +30,18 @@ def _handle(self, source, dest, to_backend):
         dest.sendall(data)
 
         # Receiving the response
-        buffer = self._get_data(dest, buffer_size)
+        data = self._get_data(dest, buffer_size)
+        source.sendall(data)
 
-        source.sendall(buffer)
-
-        # Reading the HTTP Headers
-        while EOH not in buffer:
-            data = self._get_data(dest, buffer_size)
-            buffer += data
-            source.sendall(data)
+        buffer = data
+        while data:
+          data = self._get_data(dest, buffer_size)
+          buffer += data
+          source.sendall(data)
 
         # keep alive header ?
         keep_alive = RE_KEEPALIVE.search(buffer) is not None
 
-        # content-length header - to see if we need to suck more
-        # data.
-        match = RE_LEN.search(buffer)
-        if match:
-            resp_len = int(match.group(1))
-            left_to_read = resp_len - len(buffer)
-            if left_to_read > 0:
-                for chunk in chunked(left_to_read, buffer_size):
-                    data = self._get_data(dest, chunk)
-                    buffer += data
-                    source.sendall(data)
-        else:
-            # embarrassing...
-            # just sucking until recv() returns ''
-            while True:
-                data = self._get_data(dest, buffer_size)
-                if data == '':
-                    break
-                source.sendall(data)
-
         # do we close the client ?
         if not keep_alive and not self.option('keep_alive'):
             source.close()

vaurien/proxy.py

@@ -25,6 +25,7 @@ def __init__(self, proxy, backend, protocol='tcp', behaviors=None,
 
         logger.info('Starting the Chaos TCP Server')
         parsed_proxy = parse_address(proxy)
+        self.backend = backend
         dest = parse_address(backend)
         backlog = cfg.get('backlog', 8192)
         StreamServer.__init__(self, parsed_proxy, backlog=backlog, **kwargs)
@@ -58,6 +59,7 @@ def __init__(self, proxy, backend, protocol='tcp', behaviors=None,
         args = settings['args']
         self.handler.update_settings(extract_settings(args, 'protocol',
                                                       self.protocol))
+        self.handler.proxy = self
 
         logger.info('Options:')
         logger.info('* proxies from %s to %s' % (proxy, backend))