From 10ee19c13879bddafeb9217dd5e4e29af954baac Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Mon, 15 Mar 2021 11:38:21 +0100 Subject: [PATCH] fix race condition --- mitmproxy/proxy/layers/tcp.py | 2 +- mitmproxy/proxy/layers/websocket.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/mitmproxy/proxy/layers/tcp.py b/mitmproxy/proxy/layers/tcp.py index bf87e37d4..c6c75c583 100644 --- a/mitmproxy/proxy/layers/tcp.py +++ b/mitmproxy/proxy/layers/tcp.py @@ -130,6 +130,6 @@ class TCPLayer(layer.Layer): else: raise AssertionError(f"Unexpected event: {event}") - @expect(events.DataReceived, events.ConnectionClosed) + @expect(events.DataReceived, events.ConnectionClosed, TcpMessageInjected) def done(self, _) -> layer.CommandGenerator[None]: yield from () diff --git a/mitmproxy/proxy/layers/websocket.py b/mitmproxy/proxy/layers/websocket.py index f55ab07d6..ca15ed9ff 100644 --- a/mitmproxy/proxy/layers/websocket.py +++ b/mitmproxy/proxy/layers/websocket.py @@ -207,7 +207,7 @@ class WebsocketLayer(layer.Layer): else: # pragma: no cover raise AssertionError(f"Unexpected WebSocket event: {ws_event}") - @expect(events.DataReceived, events.ConnectionClosed) + @expect(events.DataReceived, events.ConnectionClosed, WebSocketMessageInjected) def done(self, _) -> layer.CommandGenerator[None]: yield from ()