fixed handling of Transfer-Encoding header during streaming; wrote tests for streaming support

1 files changed, 22 insertions, 16 deletions

diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index f3d5f666..2c3c8f97 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -632,24 +632,21 @@ class HTTPResponse(HTTPMessage):
         return 'HTTP/%s.%s %s %s' % \
                (self.httpversion[0], self.httpversion[1], self.code, self.msg)
 
-    def _assemble_headers(self):
+    def _assemble_headers(self, preserve_transfer_encoding=False):
         headers = self.headers.copy()
-        utils.del_all(
-            headers,
-            [
-                'Proxy-Connection',
-                'Transfer-Encoding'
-            ]
-        )
+        utils.del_all(headers,['Proxy-Connection'])
+        if not preserve_transfer_encoding:
+            utils.del_all(headers,['Transfer-Encoding'])
+
         if self.content:
             headers["Content-Length"] = [str(len(self.content))]
-        elif 'Transfer-Encoding' in self.headers:  # add content-length for chuncked transfer-encoding with no content
+        elif not preserve_transfer_encoding and 'Transfer-Encoding' in self.headers:  # add content-length for chuncked transfer-encoding with no content
             headers["Content-Length"] = ["0"]
 
         return str(headers)
 
-    def _assemble_head(self):
-        return '%s\r\n%s\r\n' % (self._assemble_first_line(), self._assemble_headers())
+    def _assemble_head(self, preserve_transfer_encoding=False):
+        return '%s\r\n%s\r\n' % (self._assemble_first_line(), self._assemble_headers(preserve_transfer_encoding=preserve_transfer_encoding))
 
     def _assemble(self):
         """
@@ -928,7 +925,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
                     # flow.stream_expecting_body now contains the state info of whether or not 
                     # body still remains to be read
 
-                # call the appropriate script hook - this is an opportunity for 
+                # call the appropriate script hook - this is an opportunity for an inline script to set flow.stream = True
                 responseheaders_reply = self.c.channel.ask("responseheaders", flow.response)
                 # hm - do we need to do something with responseheaders_reply??
 
@@ -944,6 +941,8 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
             if response_reply is None or response_reply == KILL:
                 return False
 
+            disconnected_while_streaming = False
+
             if not flow.stream or not flow.stream_expecting_body:
                 # if not streaming or there is no body to be read, we'll already have the body, just send it
                 self.c.client_conn.send(flow.response._assemble())
@@ -952,15 +951,17 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
                 # if streaming, we still need to read the body and stream its bits back to the client
 
                 # start with head
-                h = flow.response._assemble_head()
+                h = flow.response._assemble_head(preserve_transfer_encoding=True)
                 self.c.client_conn.send(h)
 
                 # if chunked then we send back each chunk
                 if http.has_chunked_encoding(flow.response.headers):
                     while 1:
                         content = http.read_next_chunk(self.c.server_conn.rfile, flow.response.headers, False)
-                        if not http.write_chunk(self.c.client_conn.rfile, content):
+                        if not http.write_chunk(self.c.client_conn.wfile, content):
                             break
+                        self.c.client_conn.wfile.flush()
+                    self.c.client_conn.wfile.flush()
 
                 else: # not chunked, we send back 4k at a time
                     clen = http.expected_http_body_size(flow.response.headers, False)
@@ -969,17 +970,21 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
                     blocksize = 4096
                     while 1:
                         bytes_to_read = min(blocksize, clen - rcount)
+                        if bytes_to_read == 0:
+                            break
                         content = self.c.server_conn.rfile.read(bytes_to_read)
                         if content == "": # check for EOF
+                            disconnected_while_streaming = True
                             break
                         rcount += len(content)
-                        self.c.client_conn.rfile.write(content)
+                        self.c.client_conn.wfile.write(content)
+                        self.c.client_conn.wfile.flush()
                         if rcount >= clen: # check for having read up to clen
                             break
 
             flow.timestamp_end = utils.timestamp()
 
-            if (http.connection_close(flow.request.httpversion, flow.request.headers) or
+            if (disconnected_while_streaming or http.connection_close(flow.request.httpversion, flow.request.headers) or
                     http.connection_close(flow.response.httpversion, flow.response.headers)):
                 return False
 
@@ -989,6 +994,7 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
             # If the user has changed the target server on this connection,
             # restore the original target server
             self.restore_server()
+
             return True
         except (HttpAuthenticationError, http.HttpError, proxy.ProxyError, tcp.NetLibError), e:
             self.handle_error(e, flow)