From b64005dfeac10228eab24e950cacc1bd05f7deaf Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Thu, 29 Jul 2021 14:51:19 +0200
Subject: [PATCH] make it possible to return multiple chunks from streaming,
 refs #4706

---
 mitmproxy/http.py                       |  2 +-
 mitmproxy/proxy/layers/http/__init__.py | 53 +++++++++++++++----------
 2 files changed, 32 insertions(+), 23 deletions(-)

diff --git a/mitmproxy/http.py b/mitmproxy/http.py
index f48fc1e1a..d2e385db8 100644
--- a/mitmproxy/http.py
+++ b/mitmproxy/http.py
@@ -250,7 +250,7 @@ class Message(serializable.Serializable):
         self.data.set_state(state)
 
     data: MessageData
-    stream: Union[Callable[[bytes], bytes], bool] = False
+    stream: Union[Callable[[bytes], Union[Iterable[bytes], bytes]], bool] = False
     """
     If `True`, the message body will not be buffered on the proxy
     but immediately streamed to the destination instead.
diff --git a/mitmproxy/proxy/layers/http/__init__.py b/mitmproxy/proxy/layers/http/__init__.py
index 05017f2cd..9e16b4d3a 100644
--- a/mitmproxy/proxy/layers/http/__init__.py
+++ b/mitmproxy/proxy/layers/http/__init__.py
@@ -231,31 +231,32 @@ class HttpStream(layer.Layer):
     def state_stream_request_body(self, event: Union[RequestData, RequestEndOfMessage]) -> layer.CommandGenerator[None]:
         if isinstance(event, RequestData):
             if callable(self.flow.request.stream):
-                event.data = self.flow.request.stream(event.data)
+                chunks = self.flow.request.stream(event.data)
+                if isinstance(chunks, bytes):
+                    chunks = [chunks]
+            else:
+                chunks = [event.data]
+            for chunk in chunks:
+                yield SendHttp(RequestData(self.stream_id, chunk), self.context.server)
         elif isinstance(event, RequestTrailers):
-            self.flow.request.trailers = event.trailers
             # we don't do anything further here, we wait for RequestEndOfMessage first to trigger the request hook.
-            return
+            self.flow.request.trailers = event.trailers
         elif isinstance(event, RequestEndOfMessage):
+            if callable(self.flow.request.stream):
+                chunks = self.flow.request.stream(b"")
+                if isinstance(chunks, bytes):
+                    chunks = [chunks]
+                for chunk in chunks:
+                    yield SendHttp(RequestData(self.stream_id, chunk), self.context.server)
+
             self.flow.request.timestamp_end = time.time()
             yield HttpRequestHook(self.flow)
             self.client_state = self.state_done
 
-        # edge case found while fuzzing:
-        # we may arrive here after a hook unpaused the stream,
-        # but the server may have sent us a RST_STREAM in the meantime.
-        # We need to 1) check the server state and 2) peek into the event queue to
-        # see if this is the case.
-        if self.server_state == self.state_errored:
-            return
-        for evt in self._paused_event_queue:
-            if isinstance(evt, ResponseProtocolError):
-                return
-        if self.flow.request.trailers:
-            # we've delayed sending trailers until after `request` has been triggered.
-            assert isinstance(event, RequestEndOfMessage)
-            yield SendHttp(RequestTrailers(self.stream_id, self.flow.request.trailers), self.context.server)
-        yield SendHttp(event, self.context.server)
+            if self.flow.request.trailers:
+                # we've delayed sending trailers until after `request` has been triggered.
+                yield SendHttp(RequestTrailers(self.stream_id, self.flow.request.trailers), self.context.server)
+            yield SendHttp(event, self.context.server)
 
     @expect(RequestData, RequestTrailers, RequestEndOfMessage)
     def state_consume_request_body(self, event: events.Event) -> layer.CommandGenerator[None]:
@@ -322,15 +323,23 @@ class HttpStream(layer.Layer):
         assert self.flow.response
         if isinstance(event, ResponseData):
             if callable(self.flow.response.stream):
-                data = self.flow.response.stream(event.data)
+                chunks = self.flow.response.stream(event.data)
+                if isinstance(chunks, bytes):
+                    chunks = [chunks]
             else:
-                data = event.data
-            yield SendHttp(ResponseData(self.stream_id, data), self.context.client)
+                chunks = [event.data]
+            for chunk in chunks:
+                yield SendHttp(ResponseData(self.stream_id, chunk), self.context.client)
         elif isinstance(event, ResponseTrailers):
-            assert self.flow.response
             self.flow.response.trailers = event.trailers
             # will be sent in send_response() after the response hook.
         elif isinstance(event, ResponseEndOfMessage):
+            if callable(self.flow.response.stream):
+                chunks = self.flow.response.stream(b"")
+                if isinstance(chunks, bytes):
+                    chunks = [chunks]
+                for chunk in chunks:
+                    yield SendHttp(ResponseData(self.stream_id, chunk), self.context.client)
             yield from self.send_response(already_streamed=True)
 
     @expect(ResponseData, ResponseTrailers, ResponseEndOfMessage)