diff --git a/mitmproxy/proxy/layers/http/__init__.py b/mitmproxy/proxy/layers/http/__init__.py index 1671f15d5..cae90c0d0 100644 --- a/mitmproxy/proxy/layers/http/__init__.py +++ b/mitmproxy/proxy/layers/http/__init__.py @@ -17,10 +17,10 @@ from mitmproxy.proxy.utils import expect from mitmproxy.utils import human from mitmproxy.websocket import WebSocketData from ._base import HttpCommand, HttpConnection, ReceiveHttp, StreamId -from ._events import HttpEvent, RequestData, RequestEndOfMessage, RequestHeaders, RequestProtocolError, ResponseData, \ +from ._events import HttpEvent, RequestData, RequestEndOfMessage, RequestHeaders, RequestProtocolError, ResponseData, ResponseTrailers, \ ResponseEndOfMessage, ResponseHeaders, ResponseProtocolError from ._hooks import HttpConnectHook, HttpErrorHook, HttpRequestHeadersHook, HttpRequestHook, HttpResponseHeadersHook, \ - HttpResponseHook + HttpResponseHook, HttpResponseTrailersHook from ._http1 import Http1Client, Http1Server from ._http2 import Http2Client, Http2Server from ...context import Context @@ -297,10 +297,13 @@ class HttpStream(layer.Layer): elif isinstance(event, ResponseEndOfMessage): yield from self.send_response(already_streamed=True) - @expect(ResponseData, ResponseEndOfMessage) + @expect(ResponseData, ResponseTrailers, ResponseEndOfMessage) def state_consume_response_body(self, event: events.Event) -> layer.CommandGenerator[None]: if isinstance(event, ResponseData): self.response_body_buf += event.data + elif isinstance(event, ResponseTrailers): + self.flow.response.trailers = event.trailers + yield HttpResponseTrailersHook(self.flow) elif isinstance(event, ResponseEndOfMessage): assert self.flow.response self.flow.response.data.content = self.response_body_buf @@ -336,6 +339,8 @@ class HttpStream(layer.Layer): yield SendHttp(ResponseHeaders(self.stream_id, self.flow.response, not content), self.context.client) if content: yield SendHttp(ResponseData(self.stream_id, content), self.context.client) + if self.flow.response.trailers: + yield SendHttp(ResponseTrailers(self.stream_id, self.flow.response.trailers, end_stream=True), self.context.client) yield SendHttp(ResponseEndOfMessage(self.stream_id), self.context.client) diff --git a/mitmproxy/proxy/layers/http/_events.py b/mitmproxy/proxy/layers/http/_events.py index f138b4ba2..71c6d6666 100644 --- a/mitmproxy/proxy/layers/http/_events.py +++ b/mitmproxy/proxy/layers/http/_events.py @@ -45,6 +45,12 @@ class ResponseData(HttpEvent): self.data = data +@dataclass +class ResponseTrailers(HttpEvent): + trailers: http.Headers + end_stream: bool = True + + @dataclass class RequestEndOfMessage(HttpEvent): def __init__(self, stream_id: int): @@ -86,6 +92,7 @@ __all__ = [ "RequestEndOfMessage", "ResponseHeaders", "ResponseData", + "ResponseTrailers", "ResponseEndOfMessage", "RequestProtocolError", "ResponseProtocolError", diff --git a/mitmproxy/proxy/layers/http/_hooks.py b/mitmproxy/proxy/layers/http/_hooks.py index 34019bd95..cf33626df 100644 --- a/mitmproxy/proxy/layers/http/_hooks.py +++ b/mitmproxy/proxy/layers/http/_hooks.py @@ -44,6 +44,13 @@ class HttpResponseHook(commands.StartHook): name = "response" flow: http.HTTPFlow +@dataclass +class HttpResponseTrailersHook(commands.StartHook): + """ + The HTTP response trailers has been read. + """ + name = "responsetrailers" + flow: http.HTTPFlow @dataclass class HttpErrorHook(commands.StartHook): diff --git a/mitmproxy/proxy/layers/http/_http2.py b/mitmproxy/proxy/layers/http/_http2.py index 08c10122d..efd752207 100644 --- a/mitmproxy/proxy/layers/http/_http2.py +++ b/mitmproxy/proxy/layers/http/_http2.py @@ -17,7 +17,7 @@ from mitmproxy.connection import Connection from mitmproxy.net.http import status_codes, url from mitmproxy.utils import human from . import RequestData, RequestEndOfMessage, RequestHeaders, RequestProtocolError, ResponseData, \ - ResponseEndOfMessage, ResponseHeaders, ResponseProtocolError + ResponseEndOfMessage, ResponseHeaders, ResponseTrailers, ResponseProtocolError from ._base import HttpConnection, HttpEvent, ReceiveHttp, format_error from ._http_h2 import BufferedH2Connection, H2ConnectionLogger from ...commands import CloseConnection, Log, SendData @@ -97,6 +97,13 @@ class Http2Connection(HttpConnection): assert isinstance(event, (RequestData, ResponseData)) if self.is_open_for_us(event.stream_id): self.h2_conn.send_data(event.stream_id, event.data) + elif isinstance(event, ResponseTrailers): + if self.is_open_for_us(event.stream_id): + trailers = [ + *event.trailers.fields + ] + r = event.trailers.fields + self.h2_conn.send_headers(event.stream_id, trailers, event.end_stream) elif isinstance(event, self.SendEndOfMessage): if self.is_open_for_us(event.stream_id): self.h2_conn.end_stream(event.stream_id) @@ -216,7 +223,7 @@ class Http2Connection(HttpConnection): elif isinstance(event, h2.events.PingAckReceived): pass elif isinstance(event, h2.events.TrailersReceived): - yield Log("Received HTTP/2 trailers, which are currently unimplemented and silently discarded", "error") + yield Log("Received HTTP/2 request trailers, which are currently unimplemented and silently discarded", "error") elif isinstance(event, h2.events.PushedStreamReceived): yield Log("Received HTTP/2 push promise, even though we signalled no support.", "error") elif isinstance(event, h2.events.UnknownFrameReceived): @@ -448,6 +455,9 @@ class Http2Client(Http2Connection): self.streams[event.stream_id] = StreamState.HEADERS_RECEIVED yield ReceiveHttp(ResponseHeaders(event.stream_id, response, bool(event.stream_ended))) return False + elif isinstance(event, h2.events.TrailersReceived): + pseudo_trailers, trailers = split_pseudo_headers(event.headers) + yield ReceiveHttp(ResponseTrailers(event.stream_id, trailers, bool(event.stream_ended))) elif isinstance(event, h2.events.RequestReceived): yield from self.protocol_error(f"HTTP/2 protocol error: received request from server") return True diff --git a/test/mitmproxy/proxy/layers/http/test_http2.py b/test/mitmproxy/proxy/layers/http/test_http2.py index 7e87eab29..95dfc0ba3 100644 --- a/test/mitmproxy/proxy/layers/http/test_http2.py +++ b/test/mitmproxy/proxy/layers/http/test_http2.py @@ -30,6 +30,10 @@ example_response_headers = ( (b':status', b'200'), ) +example_response_trailers = ( + (b'my-trailer-b', b'0'), + (b'my-trailer-b', b'0') +) def decode_frames(data: bytes) -> List[hyperframe.frame.Frame]: # swallow preamble @@ -103,6 +107,50 @@ def test_simple(tctx): assert flow().response.text == "Hello, World!" +def test_trailers(tctx): + playbook, cff = start_h2_client(tctx) + flow = Placeholder(HTTPFlow) + server = Placeholder(Server) + initial = Placeholder(bytes) + assert ( + playbook + >> DataReceived(tctx.client, + cff.build_headers_frame(example_request_headers, flags=["END_STREAM"]).serialize()) + << http.HttpRequestHeadersHook(flow) + >> reply() + << http.HttpRequestHook(flow) + >> reply() + << OpenConnection(server) + >> reply(None, side_effect=make_h2) + << SendData(server, initial) + ) + frames = decode_frames(initial()) + assert [type(x) for x in frames] == [ + hyperframe.frame.SettingsFrame, + hyperframe.frame.HeadersFrame, + ] + sff = FrameFactory() + assert ( + playbook + # a conforming h2 server would send settings first, we disregard this for now. + >> DataReceived(server, sff.build_headers_frame(example_response_headers).serialize()) + << http.HttpResponseHeadersHook(flow) + >> reply() + >> DataReceived(server, sff.build_data_frame(b"Hello, World!").serialize()) + >> DataReceived(server, sff.build_headers_frame(example_response_trailers, flags=["END_STREAM"]).serialize()) + << http.HttpResponseTrailersHook(flow) + >> reply() + << http.HttpResponseHook(flow) + >> reply() + << SendData(tctx.client, + cff.build_headers_frame(example_response_headers).serialize() + + cff.build_data_frame(b"Hello, World!").serialize() + + cff.build_headers_frame(example_response_trailers, flags=["END_STREAM"]).serialize()) + ) + assert flow().request.url == "http://example.com/" + assert flow().response.text == "Hello, World!" + + def test_upstream_error(tctx): playbook, cff = start_h2_client(tctx) flow = Placeholder(HTTPFlow)