Merge pull request #1624 from cortesi/streambodies

Response body streaming to addon

11 files changed, 76 insertions, 53 deletions

diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py
index 4bd2e56d..455f8bd8 100644
--- a/mitmproxy/builtins/__init__.py
+++ b/mitmproxy/builtins/__init__.py
@@ -1,13 +1,14 @@
 from mitmproxy.builtins import anticache
 from mitmproxy.builtins import anticomp
+from mitmproxy.builtins import clientplayback
 from mitmproxy.builtins import filestreamer
-from mitmproxy.builtins import stickyauth
-from mitmproxy.builtins import stickycookie
-from mitmproxy.builtins import script
 from mitmproxy.builtins import replace
+from mitmproxy.builtins import script
 from mitmproxy.builtins import setheaders
 from mitmproxy.builtins import serverplayback
-from mitmproxy.builtins import clientplayback
+from mitmproxy.builtins import stickyauth
+from mitmproxy.builtins import stickycookie
+from mitmproxy.builtins import streambodies
 
 
 def default_addons():
@@ -18,6 +19,7 @@ def default_addons():
         stickycookie.StickyCookie(),
         script.ScriptLoader(),
         filestreamer.FileStreamer(),
+        streambodies.StreamBodies(),
         replace.Replace(),
         setheaders.SetHeaders(),
         serverplayback.ServerPlayback(),
diff --git a/mitmproxy/builtins/streambodies.py b/mitmproxy/builtins/streambodies.py
new file mode 100644
index 00000000..b3e5d6b2
--- /dev/null
+++ b/mitmproxy/builtins/streambodies.py
@@ -0,0 +1,34 @@
+from netlib.http import http1
+from netlib import exceptions
+from mitmproxy import ctx
+
+
+class StreamBodies:
+    def __init__(self):
+        self.max_size = None
+
+    def configure(self, options, updated):
+        self.max_size = options.stream_large_bodies
+
+    def run(self, f, is_request):
+        if self.max_size:
+            r = f.request if is_request else f.response
+            try:
+                expected_size = http1.expected_http_body_size(
+                    f.request, f.response if not is_request else None
+                )
+            except exceptions.HTTPException:
+                f.reply.kill()
+                return
+            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)
+
+    # FIXME! Request streaming doesn't work at the moment.
+    def requestheaders(self, f):
+        self.run(f, True)
+
+    def responseheaders(self, f):
+        self.run(f, False)
diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py
index 148f23bc..1a413990 100644
--- a/mitmproxy/console/master.py
+++ b/mitmproxy/console/master.py
@@ -236,8 +236,6 @@ class ConsoleMaster(flow.FlowMaster):
         if options.filter:
             self.set_view_filter(options.filter)
 
-        self.set_stream_large_bodies(options.stream_large_bodies)
-
         self.palette = options.palette
         self.palette_transparent = options.palette_transparent
 
diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py
index e281af3b..e7944a9e 100644
--- a/mitmproxy/console/statusbar.py
+++ b/mitmproxy/console/statusbar.py
@@ -196,10 +196,10 @@ class StatusBar(urwid.WidgetWrap):
             opts.append("no-upstream-cert")
         if self.master.state.follow_focus:
             opts.append("following")
-        if self.master.stream_large_bodies:
+        if self.master.options.stream_large_bodies:
             opts.append(
                 "stream:%s" % human.pretty_size(
-                    self.master.stream_large_bodies.max_size
+                    self.master.options.stream_large_bodies
                 )
             )
 
diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py
index 35f6e1d3..d3f894fa 100644
--- a/mitmproxy/dump.py
+++ b/mitmproxy/dump.py
@@ -40,7 +40,6 @@ class DumpMaster(flow.FlowMaster):
         self.addons.add(dumper.Dumper())
         # This line is just for type hinting
         self.options = self.options  # type: Options
-        self.set_stream_large_bodies(options.stream_large_bodies)
 
         if not self.options.no_server and server:
             self.add_log(
diff --git a/mitmproxy/flow/__init__.py b/mitmproxy/flow/__init__.py
index 9c1cb477..9d243f49 100644
--- a/mitmproxy/flow/__init__.py
+++ b/mitmproxy/flow/__init__.py
@@ -2,12 +2,12 @@ from mitmproxy.flow import export, modules
 from mitmproxy.flow.io import FlowWriter, FilteredFlowWriter, FlowReader, read_flows_from_paths
 from mitmproxy.flow.master import FlowMaster
 from mitmproxy.flow.modules import (
-    AppRegistry, StreamLargeBodies
+    AppRegistry
 )
 from mitmproxy.flow.state import State, DummyState, FlowView
 
 __all__ = [
     "export", "modules",
     "FlowWriter", "FilteredFlowWriter", "FlowReader", "read_flows_from_paths",
-    "FlowMaster", "AppRegistry", "StreamLargeBodies", "DummyState", "State", "FlowView",
+    "FlowMaster", "AppRegistry", "DummyState", "State", "FlowView",
 ]
diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py
index 14ba39e2..02e5246f 100644
--- a/mitmproxy/flow/master.py
+++ b/mitmproxy/flow/master.py
@@ -3,7 +3,6 @@ import sys
 
 from typing import Optional  # noqa
 
-import netlib.exceptions
 from netlib import http
 from mitmproxy import controller
 from mitmproxy import exceptions
@@ -53,18 +52,11 @@ class FlowMaster(controller.Master):
         if server:
             self.add_server(server)
         self.state = state
-        self.stream_large_bodies = None  # type: Optional[modules.StreamLargeBodies]
         self.apps = modules.AppRegistry()
 
     def start_app(self, host, port):
         self.apps.add(app.mapp, host, port)
 
-    def set_stream_large_bodies(self, max_size):
-        if max_size is not None:
-            self.stream_large_bodies = modules.StreamLargeBodies(max_size)
-        else:
-            self.stream_large_bodies = False
-
     def duplicate_flow(self, f):
         """
             Duplicate flow, and insert it into state without triggering any of
@@ -238,12 +230,7 @@ class FlowMaster(controller.Master):
 
     @controller.handler
     def responseheaders(self, f):
-        try:
-            if self.stream_large_bodies:
-                self.stream_large_bodies.run(f, False)
-        except netlib.exceptions.HttpException:
-            f.reply.kill()
-            return
+        pass
 
     @controller.handler
     def response(self, f):
diff --git a/mitmproxy/flow/modules.py b/mitmproxy/flow/modules.py
index ec86ba9f..4bc4d08b 100644
--- a/mitmproxy/flow/modules.py
+++ b/mitmproxy/flow/modules.py
@@ -1,6 +1,5 @@
 from netlib import wsgi
 from netlib import version
-from netlib.http import http1
 
 
 class AppRegistry:
@@ -28,17 +27,3 @@ class AppRegistry:
         if "host" in request.headers:
             host = request.headers["host"]
             return self.apps.get((host, request.port), None)
-
-
-class StreamLargeBodies:
-    def __init__(self, max_size):
-        self.max_size = max_size
-
-    def run(self, flow, is_request):
-        r = flow.request if is_request else flow.response
-        expected_size = http1.expected_http_body_size(
-            flow.request, flow.response if not is_request else None
-        )
-        if 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
diff --git a/mitmproxy/models/http.py b/mitmproxy/models/http.py
index 1d7f9eb8..91263b95 100644
--- a/mitmproxy/models/http.py
+++ b/mitmproxy/models/http.py
@@ -53,6 +53,7 @@ class HTTPRequest(http.Request):
 
         # Is this request replayed?
         self.is_replay = is_replay
+        self.stream = None
 
     def get_state(self):
         state = super().get_state()
@@ -125,7 +126,7 @@ class HTTPResponse(http.Response):
 
         # Is this request replayed?
         self.is_replay = is_replay
-        self.stream = False
+        self.stream = None
 
     @classmethod
     def wrap(self, response):
diff --git a/test/mitmproxy/builtins/test_streambodies.py b/test/mitmproxy/builtins/test_streambodies.py
new file mode 100644
index 00000000..0e8a82f2
--- /dev/null
+++ b/test/mitmproxy/builtins/test_streambodies.py
@@ -0,0 +1,29 @@
+from .. import tutils, mastertest
+from mitmproxy.flow import state
+from mitmproxy.flow import master
+from mitmproxy import options
+
+from mitmproxy.builtins import streambodies
+
+
+class TestStreamBodies(mastertest.MasterTest):
+    def test_simple(self):
+        s = state.DummyState()
+        o = options.Options(stream_large_bodies = 10)
+        m = master.FlowMaster(o, None, s)
+        sa = streambodies.StreamBodies()
+        m.addons.add(sa)
+
+        f = tutils.tflow()
+        f.request.content = b""
+        f.request.headers["Content-Length"] = "1024"
+        assert not f.request.stream
+        m.requestheaders(f)
+        assert f.request.stream
+
+        f = tutils.tflow(resp=True)
+        f.response.content = b""
+        f.response.headers["Content-Length"] = "1024"
+        assert not f.response.stream
+        m.responseheaders(f)
+        assert f.response.stream
diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py
index a175d193..e9db210e 100644
--- a/test/mitmproxy/test_server.py
+++ b/test/mitmproxy/test_server.py
@@ -288,18 +288,6 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin, AppMixin):
             resp = p.request("get:'http://foo':h':foo'='bar'")
         assert resp.status_code == 400
 
-    def test_stream(self):
-        self.master.set_stream_large_bodies(1024 * 2)
-
-        self.pathod("200:b@1k")
-        assert not self.master.state.view[-1].response.stream
-        assert len(self.master.state.view[-1].response.content) == 1024 * 1
-
-        self.pathod("200:b@3k")
-        assert self.master.state.view[-1].response.stream
-        assert self.master.state.view[-1].response.content is None
-        self.master.set_stream_large_bodies(None)
-
     def test_stream_modify(self):
         s = script.Script(
             tutils.test_data.path("data/addonscripts/stream_modify.py")