fix #319

doc-src/features/responsestreaming.html

@@ -40,10 +40,6 @@ Responses that should be tagged for streaming by setting their respective .strea
 
 $!example("examples/stream.py")!$
 
-In addition, if the .stream attribute is callable, .stream will work as a hook in chunk data processing.
-
-$!example("examples/stream_modify.py")!$
-
 <h2>Implementation Details</h2>
 
 When response streaming is enabled, portions of the code which would have otherwise performed changes
@@ -52,6 +48,11 @@ on the response body will see an empty response body instead (<code>libmproxy.pr
 Streamed responses are usually sent in chunks of 4096 bytes. If the response is sent with a <code>Transfer-Encoding:
     chunked</code> header, the response will be streamed one chunk at a time.
 
+<h2>Modifying streamed data</h2>
+If the <code>.stream</code> attribute is callable, .stream will work as a hook in chunk data processing.
+
+$!example("examples/stream_modify.py")!$
+
 ### See Also
 
 - [Ignore Domains](@!urlTo("passthrough.html")!@)

examples/stream_modify.py

@@ -1,9 +1,12 @@
 """
-This inline script won't work with --stream SIZE command line option.
+This inline script modifies a streamed response.
+If you do not need streaming, see the modify_response_body example.
+Be aware that content replacement isn't trivial:
+    - If the transfer encoding isn't chunked, you cannot simply change the content length.
+    - If you want to replace all occurences of "foobar", make sure to catch the cases
+      where one chunk ends with [...]foo" and the next starts with "bar[...].
+"""
 
-That's because flow.response.stream will be overwritten to True if the
-command line option exists.
-"""
 
 def modify(chunks):
     """
@@ -12,8 +15,8 @@ def modify(chunks):
     For example, in the case of chunked transfer encoding: ("3\r\n","foo","\r\n")
     """
     for prefix, content, suffix in chunks:
-        yield prefix, content.replace("foo","bar"), suffix
+        yield prefix, content.replace("foo", "bar"), suffix
 
-def responseheaders(ctx, flow):
-    flow.response.stream = modify
-    flow.response.stream_large_bodies = 1024 # = 1KB
+
+def responseheaders(context, flow):
+    flow.response.stream = modify

libmproxy/flow.py

@@ -165,7 +165,7 @@ class StreamLargeBodies(object):
             r.headers, is_request, flow.request.method, code
         )
         if not (0 <= expected_size <= self.max_size):
-            r.stream = True
+            r.stream = r.stream or True  # r.stream may already be a callable, which we want to preserve.
 
 
 class ClientPlaybackState:

libmproxy/protocol/http.py

@@ -1332,15 +1332,19 @@ class HTTPHandler(ProtocolHandler):
             # incrementally:
             h = flow.response._assemble_head(preserve_transfer_encoding=True)
             self.c.client_conn.send(h)
-            for chunk in callable(flow.response.stream) and \
-                            flow.response.stream(http.read_http_body_chunked(self.c.server_conn.rfile,
-                                                     flow.response.headers,
-                                                     self.c.config.body_size_limit, flow.request.method,
-                                                     flow.response.code, False, 4096)) or \
-                            http.read_http_body_chunked(self.c.server_conn.rfile,
-                                                     flow.response.headers,
-                                                     self.c.config.body_size_limit, flow.request.method,
-                                                     flow.response.code, False, 4096):
+
+            chunks = http.read_http_body_chunked(
+                self.c.server_conn.rfile,
+                flow.response.headers,
+                self.c.config.body_size_limit,
+                flow.request.method,
+                flow.response.code,
+                False,
+                4096
+            )
+            if callable(flow.response.stream):
+                chunks = flow.response.stream(chunks)
+            for chunk in chunks:
                 for part in chunk:
                     self.c.client_conn.wfile.write(part)
                 self.c.client_conn.wfile.flush()

test/scripts/stream_modify.py

@@ -0,0 +1,7 @@
+def modify(chunks):
+    for prefix, content, suffix in chunks:
+        yield prefix, content.replace("foo", "bar"), suffix
+
+
+def responseheaders(context, flow):
+    flow.response.stream = modify

test/test_server.py

@@ -266,6 +266,12 @@ class TestHTTP(tservers.HTTPProxTest, CommonMixin, AppMixin):
         assert self.master.state.view[-1].response.content == CONTENT_MISSING
         self.master.set_stream_large_bodies(None)
 
+    def test_stream_modify(self):
+        self.master.load_script(tutils.test_data.path("scripts/stream_modify.py"))
+        d = self.pathod('200:b"foo"')
+        assert d.content == "bar"
+        self.master.unload_scripts()
+
 class TestHTTPAuth(tservers.HTTPProxTest):
     authenticator = http_auth.BasicProxyAuth(http_auth.PassManSingleUser("test", "test"), "realm")
     def test_auth(self):