request streaming for HTTP/1

4 files changed, 24 insertions, 11 deletions

diff --git a/mitmproxy/addons/streambodies.py b/mitmproxy/addons/streambodies.py
index 181f0337..16c5978d 100644
--- a/mitmproxy/addons/streambodies.py
+++ b/mitmproxy/addons/streambodies.py
@@ -28,8 +28,7 @@ class StreamBodies:
             if expected_size and not r.raw_content and not (0 <= expected_size <= self.max_size):
                 # r.stream may already be a callable, which we want to preserve.
                 r.stream = r.stream or True
-                # FIXME: make message generic when we add rquest streaming
-                ctx.log.info("Streaming response from %s" % f.request.host)
+                ctx.log.info("Streaming {} {}".format("response from" if not is_request else "request to", f.request.host))
 
     # FIXME! Request streaming doesn't work at the moment.
     def requestheaders(self, f):
diff --git a/mitmproxy/proxy/protocol/http.py b/mitmproxy/proxy/protocol/http.py
index 45870830..9f120767 100644
--- a/mitmproxy/proxy/protocol/http.py
+++ b/mitmproxy/proxy/protocol/http.py
@@ -273,7 +273,10 @@ class HttpLayer(base.Layer):
                 self.send_response(http.expect_continue_response)
                 request.headers.pop("expect")
 
-            request.data.content = b"".join(self.read_request_body(request))
+            if f.request.stream:
+                f.request.data.content = None
+            else:
+                f.request.data.content = b"".join(self.read_request_body(request))
             request.timestamp_end = time.time()
         except exceptions.HttpException as e:
             # We optimistically guess there might be an HTTP client on the
@@ -327,7 +330,13 @@ class HttpLayer(base.Layer):
                 )
 
                 def get_response():
-                    self.send_request(f.request)
+                    if f.request.stream:
+                        self.send_request_headers(f.request)
+                        chunks = self.read_request_body(f.request)
+                        self.send_request_body(f.request, chunks)
+                    else:
+                        self.send_request(f.request)
+
                     f.response = self.read_response_headers()
 
                 try:
diff --git a/mitmproxy/proxy/protocol/http1.py b/mitmproxy/proxy/protocol/http1.py
index cafc2682..84cd6324 100644
--- a/mitmproxy/proxy/protocol/http1.py
+++ b/mitmproxy/proxy/protocol/http1.py
@@ -22,6 +22,16 @@ class Http1Layer(httpbase._HttpTransmissionLayer):
             self.config.options._processed.get("body_size_limit")
         )
 
+    def send_request_headers(self, request):
+        headers = http1.assemble_request_head(request)
+        self.server_conn.wfile.write(headers)
+        self.server_conn.wfile.flush()
+
+    def send_request_body(self, request, chunks):
+        for chunk in http1.assemble_body(request.headers, chunks):
+            self.server_conn.wfile.write(chunk)
+            self.server_conn.wfile.flush()
+
     def send_request(self, request):
         self.server_conn.wfile.write(http1.assemble_request(request))
         self.server_conn.wfile.flush()
diff --git a/test/mitmproxy/proxy/protocol/test_http1.py b/test/mitmproxy/proxy/protocol/test_http1.py
index 1eff8666..4cca370c 100644
--- a/test/mitmproxy/proxy/protocol/test_http1.py
+++ b/test/mitmproxy/proxy/protocol/test_http1.py
@@ -1,7 +1,6 @@
 from unittest import mock
 import pytest
 
-from mitmproxy import exceptions
 from mitmproxy.test import tflow
 from mitmproxy.net.http import http1
 from mitmproxy.net.tcp import TCPClient
@@ -108,9 +107,5 @@ class TestStreaming(tservers.HTTPProxyTest):
                 r = p.request("post:'%s/p/200:b@10000'" % self.server.urlbase)
                 assert len(r.content) == 10000
 
-                if streaming:
-                    with pytest.raises(exceptions.HttpReadDisconnect):  # as the assertion in assert_write fails
-                        # request with 10000 bytes
-                        p.request("post:'%s/p/200':b@10000" % self.server.urlbase)
-                else:
-                    assert p.request("post:'%s/p/200':b@10000" % self.server.urlbase)
+                # request with 10000 bytes
+                assert p.request("post:'%s/p/200':b@10000" % self.server.urlbase)