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

fix #4452 (#4454)

Browse Source
This commit is contained in:
Maximilian Hils 2021-02-16 17:38:08 +01:00 committed by GitHub
parent aebc40c408
commit c8eca9a396
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 36 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

59
mitmproxy/proxy/layers/http/_http2.py
View File

@ -64,11 +64,29 @@ class Http2Connection(HttpConnection):
def is_closed(self, stream_id: int) -> bool: def is_closed(self, stream_id: int) -> bool:
"""Check if a non-idle stream is closed""" """Check if a non-idle stream is closed"""
stream = self.h2_conn.streams.get(stream_id, None) stream = self.h2_conn.streams.get(stream_id, None)
if stream is not None: if (
return stream.closed stream is not None
and
stream.state_machine.state is not h2.stream.StreamState.CLOSED
):
return False
else: else:
return True return True
def is_open_for_us(self, stream_id: int) -> bool:
"""Check if we can write to a non-idle stream."""
stream = self.h2_conn.streams.get(stream_id, None)
if (
stream is not None
and
stream.state_machine.state is not h2.stream.StreamState.HALF_CLOSED_LOCAL
and
stream.state_machine.state is not h2.stream.StreamState.CLOSED
):
return True
else:
return False
def _handle_event(self, event: Event) -> CommandGenerator[None]: def _handle_event(self, event: Event) -> CommandGenerator[None]:
if isinstance(event, Start): if isinstance(event, Start):
self.h2_conn.initiate_connection() self.h2_conn.initiate_connection()
@ -77,24 +95,18 @@ class Http2Connection(HttpConnection):
elif isinstance(event, HttpEvent): elif isinstance(event, HttpEvent):
if isinstance(event, self.SendData): if isinstance(event, self.SendData):
assert isinstance(event, (RequestData, ResponseData)) assert isinstance(event, (RequestData, ResponseData))
self.h2_conn.send_data(event.stream_id, event.data) if self.is_open_for_us(event.stream_id):
self.h2_conn.send_data(event.stream_id, event.data)
elif isinstance(event, self.SendEndOfMessage): elif isinstance(event, self.SendEndOfMessage):
stream = self.h2_conn.streams.get(event.stream_id) if self.is_open_for_us(event.stream_id):
if stream.state_machine.state not in (h2.stream.StreamState.HALF_CLOSED_LOCAL,
h2.stream.StreamState.CLOSED):
self.h2_conn.end_stream(event.stream_id) self.h2_conn.end_stream(event.stream_id)
if self.is_closed(event.stream_id):
self.streams.pop(event.stream_id, None)
elif isinstance(event, self.SendProtocolError): elif isinstance(event, self.SendProtocolError):
assert isinstance(event, (RequestProtocolError, ResponseProtocolError)) assert isinstance(event, (RequestProtocolError, ResponseProtocolError))
stream = self.h2_conn.streams.get(event.stream_id) if not self.is_closed(event.stream_id):
if stream.state_machine.state is not h2.stream.StreamState.CLOSED:
code = { code = {
status_codes.CLIENT_CLOSED_REQUEST: h2.errors.ErrorCodes.CANCEL, status_codes.CLIENT_CLOSED_REQUEST: h2.errors.ErrorCodes.CANCEL,
}.get(event.code, h2.errors.ErrorCodes.INTERNAL_ERROR) }.get(event.code, h2.errors.ErrorCodes.INTERNAL_ERROR)
self.h2_conn.reset_stream(event.stream_id, code) self.h2_conn.reset_stream(event.stream_id, code)
if self.is_closed(event.stream_id):
self.streams.pop(event.stream_id, None)
else: else:
raise AssertionError(f"Unexpected event: {event}") raise AssertionError(f"Unexpected event: {event}")
data_to_send = self.h2_conn.data_to_send() data_to_send = self.h2_conn.data_to_send()
@ -250,18 +262,19 @@ class Http2Server(Http2Connection):
def _handle_event(self, event: Event) -> CommandGenerator[None]: def _handle_event(self, event: Event) -> CommandGenerator[None]:
if isinstance(event, ResponseHeaders): if isinstance(event, ResponseHeaders):
headers = [ if self.is_open_for_us(event.stream_id):
(b":status", b"%d" % event.response.status_code), headers = [
*event.response.headers.fields (b":status", b"%d" % event.response.status_code),
] *event.response.headers.fields
if not event.response.is_http2: ]
headers = normalize_h1_headers(headers, False) if not event.response.is_http2:
headers = normalize_h1_headers(headers, False)
self.h2_conn.send_headers( self.h2_conn.send_headers(
event.stream_id, event.stream_id,
headers, headers,
) )
yield SendData(self.conn, self.h2_conn.data_to_send()) yield SendData(self.conn, self.h2_conn.data_to_send())
else: else:
yield from super()._handle_event(event) yield from super()._handle_event(event)