Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 00:01:36 +00:00
Code Issues Projects Releases Wiki Activity

[sans-io] http: emit error hook when killed

Browse Source
This commit is contained in:
Maximilian Hils 2020-12-07 14:19:24 +01:00
parent 7087539c4d
commit 4f0cbec308
3 changed files with 39 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
mitmproxy/proxy2/layers/http/__init__.py
View File

@ -183,7 +183,7 @@ class HttpStream(layer.Layer):
self.flow.request.host_header = self.context.server.address[0]
yield HttpRequestHeadersHook(self.flow)
if (yield from self.check_killed()):
if (yield from self.check_killed(True)):
return
if self.flow.request.headers.get("expect", "").lower() == "100-continue":
@ -227,13 +227,13 @@ class HttpStream(layer.Layer):
self.flow.request.data.content = self.request_body_buf
self.request_body_buf = b""
yield HttpRequestHook(self.flow)
if (yield from self.check_killed()):
if (yield from self.check_killed(True)):
return
elif self.flow.response:
# response was set by an inline script.
# we now need to emulate the responseheaders hook.
yield HttpResponseHeadersHook(self.flow)
if (yield from self.check_killed()):
if (yield from self.check_killed(True)):
return
yield from self.send_response()
else:
@ -253,7 +253,7 @@ class HttpStream(layer.Layer):
def state_wait_for_response_headers(self, event: ResponseHeaders) -> layer.CommandGenerator[None]:
self.flow.response = event.response
yield HttpResponseHeadersHook(self.flow)
if (yield from self.check_killed()):
if (yield from self.check_killed(True)):
return
elif self.flow.response.stream:
yield SendHttp(event, self.context.client)
@ -286,23 +286,9 @@ class HttpStream(layer.Layer):
yield from self.send_response()
self.server_state = self.state_done
def check_killed(self) -> layer.CommandGenerator[bool]:
killed_by_us = (
self.flow.error and self.flow.error.msg == flow.Error.KILLED_MESSAGE
)
killed_by_remote = (
self.context.client.state is not ConnectionState.OPEN
)
if killed_by_us or killed_by_remote:
if self.context.client.state & ConnectionState.CAN_WRITE:
yield commands.CloseConnection(self.context.client)
self._handle_event = self.state_errored
return True
return False
def send_response(self):
yield HttpResponseHook(self.flow)
if (yield from self.check_killed()):
if (yield from self.check_killed(False)):
return
has_content = bool(self.flow.response.raw_content)
yield SendHttp(ResponseHeaders(self.stream_id, self.flow.response, not has_content), self.context.client)
@ -310,6 +296,25 @@ class HttpStream(layer.Layer):
yield SendHttp(ResponseData(self.stream_id, self.flow.response.raw_content), self.context.client)
yield SendHttp(ResponseEndOfMessage(self.stream_id), self.context.client)
def check_killed(self, emit_error_hook: bool) -> layer.CommandGenerator[bool]:
killed_by_us = (
self.flow.error and self.flow.error.msg == flow.Error.KILLED_MESSAGE
)
killed_by_remote = (
self.context.client.state is not ConnectionState.OPEN
)
if killed_by_remote:
if not self.flow.error:
self.flow.error = flow.Error("Client disconnected.")
if killed_by_us or killed_by_remote:
if emit_error_hook:
yield HttpErrorHook(self.flow)
if self.context.client.state & ConnectionState.CAN_WRITE:
yield commands.CloseConnection(self.context.client)
self._handle_event = self.state_errored
return True
return False
def handle_protocol_error(
self,
event: typing.Union[RequestProtocolError, ResponseProtocolError]
@ -331,7 +336,7 @@ class HttpStream(layer.Layer):
self.flow.error = flow.Error(event.message)
yield HttpErrorHook(self.flow)
if (yield from self.check_killed()):
if (yield from self.check_killed(False)):
return
if isinstance(event, ResponseProtocolError):
@ -353,7 +358,7 @@ class HttpStream(layer.Layer):
def handle_connect(self) -> layer.CommandGenerator[None]:
yield HttpConnectHook(self.flow)
if (yield from self.check_killed()):
if (yield from self.check_killed(False)):
return
self.context.server.address = (self.flow.request.host, self.flow.request.port)

2
mitmproxy/proxy2/layers/http/_hooks.py
View File

@ -44,6 +44,8 @@ class HttpErrorHook(commands.Hook):
An HTTP error has occurred, e.g. invalid server responses, or
interrupted connections. This is distinct from a valid server HTTP
error response, which is simply a response with an HTTP error code.
Every flow will receive either an error or an response event, but not both.
"""
name = "error"
flow: http.HTTPFlow

18
test/mitmproxy/proxy2/layers/http/test_http.py
View File

@ -806,19 +806,23 @@ def test_kill_flow(tctx, when):
flow = Placeholder(HTTPFlow)
def kill(flow: HTTPFlow):
# Can't use flow.kill() here because that currently still depends on a reply object.
flow.error = Error(Error.KILLED_MESSAGE)
def assert_kill():
assert (playbook
>> reply(side_effect=kill)
<< CloseConnection(tctx.client))
def assert_kill(err_hook: bool = True):
playbook >> reply(side_effect=kill)
if err_hook:
playbook << http.HttpErrorHook(flow)
playbook >> reply()
playbook << CloseConnection(tctx.client)
assert playbook
playbook = Playbook(http.HttpLayer(tctx, HTTPMode.regular))
assert (playbook
>> DataReceived(tctx.client, b"CONNECT example.com:80 HTTP/1.1\r\n\r\n")
<< http.HttpConnectHook(connect_flow))
if when == "http_connect":
return assert_kill()
return assert_kill(False)
assert (playbook
>> reply()
<< SendData(tctx.client, b'HTTP/1.1 200 Connection established\r\n\r\n')
@ -853,14 +857,14 @@ def test_kill_flow(tctx, when):
>> reply()
>> DataReceived(server, b"!")
<< http.HttpResponseHook(flow))
return assert_kill()
return assert_kill(False)
elif when == "error":
assert (playbook
>> reply()
>> ConnectionClosed(server)
<< CloseConnection(server)
<< http.HttpErrorHook(flow))
return assert_kill()
return assert_kill(False)
else:
raise AssertionError