Merge remote-tracking branch 'sakurai/feature/issue319'

3 files changed, 28 insertions, 1 deletions

diff --git a/doc-src/features/responsestreaming.html b/doc-src/features/responsestreaming.html
index 47fafef7..20785269 100644
--- a/doc-src/features/responsestreaming.html
+++ b/doc-src/features/responsestreaming.html
@@ -40,6 +40,9 @@ 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>
 
diff --git a/examples/stream_modify.py b/examples/stream_modify.py
new file mode 100644
index 00000000..a28d95c7
--- /dev/null
+++ b/examples/stream_modify.py
@@ -0,0 +1,19 @@
+"""
+This inline script won't work with --stream SIZE command line option.
+
+That's because flow.response.stream will be overwritten to True if the
+command line option exists.
+"""
+
+def modify(chunks):
+    """
+    chunks is a generator that can be used to iterate over all chunks.
+    Each chunk is a (prefix, content, suffix) tuple. 
+    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
+
+def responseheaders(ctx, flow):
+    flow.response.stream = modify
+    flow.response.stream_large_bodies = 1024 # = 1KB
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 046d0b42..51fd503f 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -1332,7 +1332,12 @@ class HTTPHandler(ProtocolHandler):
             # incrementally:
             h = flow.response._assemble_head(preserve_transfer_encoding=True)
             self.c.client_conn.send(h)
-            for chunk in http.read_http_body_chunked(self.c.server_conn.rfile,
+            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):