vee1e/mitmproxy
1
0
Fork 0
mirror of https://github.com/vee1e/mitmproxy.git synced 2026-09-04 03:37:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

remove websocket_error event, fixes #4674

Browse source
This commit is contained in:
Alexander Prinzhorn 2021-07-12 16:35:25 +02:00
parent a51dc10d8b
commit eef6522a5e
12 changed files with 40 additions and 73 deletions
Download patch file Download diff file Expand all files Collapse all files

1
docs/scripts/api-events.py
View file

@ -133,7 +133,6 @@ with outfile.open("w") as f, contextlib.redirect_stdout(f):
websocket.WebsocketStartHook,
websocket.WebsocketMessageHook,
websocket.WebsocketEndHook,
websocket.WebsocketErrorHook,
]
)

6
examples/contrib/jsondump.py
View file

@ -247,11 +247,5 @@ class JSONDumper:
"""
self.queue.put(flow.get_state())
def websocket_error(self, flow):
"""
Dump websocket errors.
"""
self.queue.put(flow.get_state())
addons = [JSONDumper()] # pylint: disable=invalid-name

30
mitmproxy/addons/dumper.py
View file

@ -10,10 +10,12 @@ from mitmproxy import ctx
from mitmproxy import exceptions
from mitmproxy import flowfilter
from mitmproxy import http
from mitmproxy import flow
from mitmproxy.tcp import TCPFlow, TCPMessage
from mitmproxy.utils import human
from mitmproxy.utils import strutils
from mitmproxy.websocket import WebSocketMessage
from mitmproxy.websocket import WebSocketMessage, WebSocketData
from wsproto.frame_protocol import CloseReason
def indent(n: int, text: str) -> str:
@ -277,12 +279,6 @@ class Dumper:
if self.match(f):
self.echo_flow(f)
def websocket_error(self, f: http.HTTPFlow):
self.echo_error(
f"Error in WebSocket connection to {human.format_address(f.server_conn.address)}: {f.error}",
fg="red"
)
def websocket_message(self, f: http.HTTPFlow):
assert f.websocket is not None # satisfy type checker
if self.match(f):
@ -300,8 +296,24 @@ class Dumper:
def websocket_end(self, f: http.HTTPFlow):
assert f.websocket is not None # satisfy type checker
if self.match(f):
c = 'client' if f.websocket.closed_by_client else 'server'
self.echo(f"WebSocket connection closed by {c}: {f.websocket.close_code} {f.websocket.close_reason}")
if f.websocket.close_code in {1000, 1001, 1005}:
c = 'client' if f.websocket.closed_by_client else 'server'
self.echo(f"WebSocket connection closed by {c}: {f.websocket.close_code} {f.websocket.close_reason}")
else:
error = flow.Error(f"WebSocket Error: {self.format_websocket_error(f.websocket)}")
self.echo_error(
f"Error in WebSocket connection to {human.format_address(f.server_conn.address)}: {error}",
fg="red"
)
def format_websocket_error(self, websocket: WebSocketData) -> str:
try:
ret = CloseReason(websocket.close_code).name
except ValueError:
ret = f"UNKNOWN_ERROR={websocket.close_code}"
if websocket.close_reason:
ret += f" (reason: {websocket.close_reason})"
return ret
def tcp_error(self, f):
if self.match(f):

5
mitmproxy/addons/save.py
View file

@ -94,15 +94,12 @@ class Save:
self.stream.add(flow)
self.active_flows.discard(flow)
def websocket_error(self, flow: http.HTTPFlow):
self.websocket_end(flow)
def request(self, flow: http.HTTPFlow):
if self.stream:
self.active_flows.add(flow)
def response(self, flow: http.HTTPFlow):
# websocket flows will receive either websocket_end or websocket_error,
# websocket flows will receive a websocket_end,
# we don't want to persist them here already
if self.stream and flow.websocket is None:
self.stream.add(flow)

5
mitmproxy/eventsequence.py
View file

@ -24,10 +24,7 @@ def _iterate_http(f: http.HTTPFlow) -> TEventGenerator:
for m in message_queue:
f.websocket.messages.append(m)
yield layers.websocket.WebsocketMessageHook(f)
if f.error:
yield layers.websocket.WebsocketErrorHook(f)
else:
yield layers.websocket.WebsocketEndHook(f)
yield layers.websocket.WebsocketEndHook(f)
elif f.error:
yield layers.http.HttpErrorHook(f)

31
mitmproxy/proxy/layers/websocket.py
View file

@ -6,14 +6,14 @@ import wsproto
import wsproto.extensions
import wsproto.frame_protocol
import wsproto.utilities
from mitmproxy import connection, flow, http, websocket
from mitmproxy import connection, http, websocket
from mitmproxy.proxy import commands, events, layer
from mitmproxy.proxy.commands import StartHook
from mitmproxy.proxy.context import Context
from mitmproxy.proxy.events import MessageInjected
from mitmproxy.proxy.utils import expect
from wsproto import ConnectionState
from wsproto.frame_protocol import CloseReason, Opcode
from wsproto.frame_protocol import Opcode
@dataclass
@ -39,21 +39,12 @@ class WebsocketMessageHook(StartHook):
class WebsocketEndHook(StartHook):
"""
A WebSocket connection has ended.
You can check `flow.websocket.close_code` to determine why it ended.
"""
flow: http.HTTPFlow
@dataclass
class WebsocketErrorHook(StartHook):
"""
A WebSocket connection has had an error.
Every WebSocket flow will receive either a websocket_error or a websocket_end event, but not both.
"""
flow: http.HTTPFlow
class WebSocketMessageInjected(MessageInjected[websocket.WebSocketMessage]):
"""
The user has injected a custom WebSocket message.
@ -200,11 +191,7 @@ class WebsocketLayer(layer.Layer):
# response == original event, so no need to differentiate here.
yield ws.send2(ws_event)
yield commands.CloseConnection(ws.conn)
if ws_event.code in {1000, 1001, 1005}:
yield WebsocketEndHook(self.flow)
else:
self.flow.error = flow.Error(f"WebSocket Error: {format_close_event(ws_event)}")
yield WebsocketErrorHook(self.flow)
yield WebsocketEndHook(self.flow)
self._handle_event = self.done
else: # pragma: no cover
raise AssertionError(f"Unexpected WebSocket event: {ws_event}")
@ -214,16 +201,6 @@ class WebsocketLayer(layer.Layer):
yield from ()
def format_close_event(event: wsproto.events.CloseConnection) -> str:
try:
ret = CloseReason(event.code).name
except ValueError:
ret = f"UNKNOWN_ERROR={event.code}"
if event.reason:
ret += f" (reason: {event.reason})"
return ret
class Fragmentizer:
"""
Theory (RFC 6455):

6
mitmproxy/test/tflow.py
View file

@ -75,7 +75,11 @@ def twebsocketflow(messages=True, err=None) -> http.HTTPFlow:
websocket.WebSocketMessage(Opcode.TEXT, False, b"it's me", 946681205),
]
if err is True:
flow.error = terr()
# ABNORMAL_CLOSURE
flow.websocket.close_code = 1006
else:
# NORMAL_CLOSURE
flow.websocket.close_code = 1000
flow.reply = controller.DummyReply()
return flow

2
test/mitmproxy/addons/test_dumper.py
View file

@ -234,7 +234,7 @@ def test_websocket():
assert "WebSocket connection closed by" in sio.getvalue()
f = tflow.twebsocketflow(err=True)
d.websocket_error(f)
d.websocket_end(f)
assert "Error in WebSocket" in sio_err.getvalue()

2
test/mitmproxy/addons/test_save.py
View file

@ -60,7 +60,7 @@ def test_websocket(tmpdir):
f = tflow.twebsocketflow()
sa.request(f)
sa.websocket_error(f)
sa.websocket_end(f)
tctx.configure(sa, save_stream_file=None)
assert len(rd(p)) == 2

8
test/mitmproxy/proxy/layers/test_websocket.py
View file

@ -188,7 +188,7 @@ def test_protocol_error(ws_testdata):
<< CloseConnection(tctx.server)
<< SendData(tctx.client, b"\x88/\x03\xeaexpected CONTINUATION, got <Opcode.BINARY: 2>")
<< CloseConnection(tctx.client)
<< websocket.WebsocketErrorHook(flow)
<< websocket.WebsocketEndHook(flow)
>> reply()
)
@ -245,11 +245,10 @@ def test_close_disconnect(ws_testdata):
<< CloseConnection(tctx.server)
<< SendData(tctx.client, b"\x88\x02\x03\xe8")
<< CloseConnection(tctx.client)
<< websocket.WebsocketErrorHook(flow)
<< websocket.WebsocketEndHook(flow)
>> reply()
>> ConnectionClosed(tctx.client)
)
assert "ABNORMAL_CLOSURE" in flow.error.msg
def test_close_error(ws_testdata):
@ -263,10 +262,9 @@ def test_close_error(ws_testdata):
<< CloseConnection(tctx.server)
<< SendData(tctx.client, b"\x88\x02\x0f\xa0")
<< CloseConnection(tctx.client)
<< websocket.WebsocketErrorHook(flow)
<< websocket.WebsocketEndHook(flow)
>> reply()
)
assert "UNKNOWN_ERROR=4000" in flow.error.msg
def test_deflate(ws_testdata):

10
test/mitmproxy/test_eventsequence.py
View file

@ -23,9 +23,8 @@ def test_http_flow(resp, err):
assert isinstance(next(i), layers.http.HttpErrorHook)
@pytest.mark.parametrize("err", [False, True])
def test_websocket_flow(err):
f = tflow.twebsocketflow(err=err)
def test_websocket_flow():
f = tflow.twebsocketflow()
i = eventsequence.iterate(f)
assert isinstance(next(i), layers.http.HttpRequestHeadersHook)
@ -41,10 +40,7 @@ def test_websocket_flow(err):
assert len(f.websocket.messages) == 2
assert isinstance(next(i), layers.websocket.WebsocketMessageHook)
assert len(f.websocket.messages) == 3
if err:
assert isinstance(next(i), layers.websocket.WebsocketErrorHook)
else:
assert isinstance(next(i), layers.websocket.WebsocketEndHook)
assert isinstance(next(i), layers.websocket.WebsocketEndHook)
@pytest.mark.parametrize("err", [False, True])

7
test/mitmproxy/test_flowfilter.py
View file

@ -461,9 +461,6 @@ class TestMatchingWebSocketFlow:
def flow(self) -> http.HTTPFlow:
return tflow.twebsocketflow()
def err(self) -> http.HTTPFlow:
return tflow.twebsocketflow(err=True)
def q(self, q, o):
return flowfilter.parse(q)(o)
@ -484,10 +481,6 @@ class TestMatchingWebSocketFlow:
f = tflow.tflow(resp=True)
assert not self.q("~websocket", f)
def test_ferr(self):
e = self.err()
assert self.q("~e", e)
def test_domain(self):
q = self.flow()
assert self.q("~d example.com", q)