diff options
| -rw-r--r-- | doc-src/features/responsestreaming.html | 9 | ||||
| -rw-r--r-- | examples/stream_modify.py | 19 | ||||
| -rw-r--r-- | libmproxy/flow.py | 2 | ||||
| -rw-r--r-- | libmproxy/protocol/http.py | 22 | ||||
| -rw-r--r-- | test/scripts/stream_modify.py | 7 | ||||
| -rw-r--r-- | test/test_server.py | 6 |
6 files changed, 43 insertions, 22 deletions
diff --git a/doc-src/features/responsestreaming.html b/doc-src/features/responsestreaming.html index 20785269..6511e913 100644 --- a/doc-src/features/responsestreaming.html +++ b/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")!@) diff --git a/examples/stream_modify.py b/examples/stream_modify.py index a28d95c7..56d26e6d 100644 --- a/examples/stream_modify.py +++ b/examples/stream_modify.py @@ -1,10 +1,13 @@ """ -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. +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[...]. """ + def modify(chunks): """ chunks is a generator that can be used to iterate over all 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
\ No newline at end of file diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 14497964..43580109 100644 --- a/libmproxy/flow.py +++ b/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: diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py index 51fd503f..49310ec3 100644 --- a/libmproxy/protocol/http.py +++ b/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() diff --git a/test/scripts/stream_modify.py b/test/scripts/stream_modify.py new file mode 100644 index 00000000..9a98a7ee --- /dev/null +++ b/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
\ No newline at end of file diff --git a/test/test_server.py b/test/test_server.py index e387293f..26770f29 100644 --- a/test/test_server.py +++ b/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): |