From 5ee192b75896b7527372943655c393c447f9caaf Mon Sep 17 00:00:00 2001 From: Thomas Kriechbaumer Date: Tue, 29 Nov 2016 21:00:42 +0100 Subject: websocket: fix empty frame with fin=1 --- mitmproxy/proxy/protocol/websocket.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py index 15d9a288..c94763e0 100644 --- a/mitmproxy/proxy/protocol/websocket.py +++ b/mitmproxy/proxy/protocol/websocket.py @@ -82,7 +82,15 @@ class WebSocketLayer(base.Layer): mask=(False if is_server else 1), masking_key=(b'' if is_server else os.urandom(4))) for i in chunks ] - frms[-1].header.fin = 1 + + if len(frms) > 0: + frms[-1].header.fin = True + else: + frms.append(websockets.Frame( + fin=True, + opcode=frame.header.opcode, + mask=(False if is_server else 1), + masking_key=(b'' if is_server else os.urandom(4)))) for frm in frms: other_conn.send(bytes(frm)) -- cgit v1.2.3 From 6e15e766c579b179e3ae6e06197670f5d573bb40 Mon Sep 17 00:00:00 2001 From: Thomas Kriechbaumer Date: Tue, 29 Nov 2016 22:01:41 +0100 Subject: websocket: fix close handshake and re-chunking --- mitmproxy/proxy/protocol/websocket.py | 47 ++++++++++++++++++++++++----------- 1 file changed, 33 insertions(+), 14 deletions(-) diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py index c94763e0..31d734fd 100644 --- a/mitmproxy/proxy/protocol/websocket.py +++ b/mitmproxy/proxy/protocol/websocket.py @@ -59,28 +59,41 @@ class WebSocketLayer(base.Layer): fb.append(frame) if frame.header.fin: + payload = b''.join(f.payload for f in fb) + original_chunk_sizes = [len(f.payload) for f in fb] + fb.clear() + if frame.header.opcode == websockets.OPCODE.TEXT: t = WebSocketTextMessage else: t = WebSocketBinaryMessage - payload = b''.join(f.payload for f in fb) - fb.clear() - websocket_message = t(self.flow, not is_server, payload) + length = len(websocket_message.content) self.flow.messages.append(websocket_message) self.channel.ask("websocket_message", self.flow) - # chunk payload into multiple 10kB frames, and send them - payload = websocket_message.content - chunk_size = 10240 # 10kB - chunks = range(0, len(payload), chunk_size) + def get_chunk(payload): + if len(payload) == length: + # message has the same length, we can reuse the same sizes + pos = 0 + for s in original_chunk_sizes: + yield payload[pos:pos + s] + pos += s + else: + # just re-chunk everything into 10kB frames + chunk_size = 10240 + chunks = range(0, len(payload), chunk_size) + for i in chunks: + yield payload[i:i + chunk_size] + frms = [ websockets.Frame( - payload=payload[i:i + chunk_size], + payload=chunk, opcode=frame.header.opcode, mask=(False if is_server else 1), - masking_key=(b'' if is_server else os.urandom(4))) for i in chunks + masking_key=(b'' if is_server else os.urandom(4))) + for chunk in get_chunk(websocket_message.content) ] if len(frms) > 0: @@ -113,7 +126,7 @@ class WebSocketLayer(base.Layer): other_conn.send(bytes(frame)) - # close the connection + # initiate close handshake return False def _handle_unknown_frame(self, frame, source_conn, other_conn, is_server): @@ -134,10 +147,11 @@ class WebSocketLayer(base.Layer): client = self.client_conn.connection server = self.server_conn.connection conns = [client, server] + close_received = False try: while not self.channel.should_exit.is_set(): - r = tcp.ssl_read_select(conns, 1) + r = tcp.ssl_read_select(conns, 0.5) for conn in r: source_conn = self.client_conn if conn == client else self.server_conn other_conn = self.server_conn if conn == client else self.client_conn @@ -145,10 +159,15 @@ class WebSocketLayer(base.Layer): frame = websockets.Frame.from_file(source_conn.rfile) - if not self._handle_frame(frame, source_conn, other_conn, is_server): - return + cont = self._handle_frame(frame, source_conn, other_conn, is_server) + if not cont: + if close_received: + return + else: + close_received = True except (socket.error, exceptions.TcpException, SSL.Error) as e: - self.flow.error = flow.Error("WebSocket connection closed unexpectedly: {}".format(repr(e))) + s = 'server' if is_server else 'client' + self.flow.error = flow.Error("WebSocket connection closed unexpectedly by {}: {}".format(s, repr(e))) self.channel.tell("websocket_error", self.flow) finally: self.channel.tell("websocket_end", self.flow) -- cgit v1.2.3 From bd8ae910d23821cc88437f486cf9eb3f2bef9470 Mon Sep 17 00:00:00 2001 From: Thomas Kriechbaumer Date: Tue, 29 Nov 2016 22:13:59 +0100 Subject: websocket: fix message type on chunking --- mitmproxy/proxy/protocol/websocket.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py index 31d734fd..ec1c6ebc 100644 --- a/mitmproxy/proxy/protocol/websocket.py +++ b/mitmproxy/proxy/protocol/websocket.py @@ -61,9 +61,10 @@ class WebSocketLayer(base.Layer): if frame.header.fin: payload = b''.join(f.payload for f in fb) original_chunk_sizes = [len(f.payload) for f in fb] + message_type = fb[0].header.opcode fb.clear() - if frame.header.opcode == websockets.OPCODE.TEXT: + if message_type == websockets.OPCODE.TEXT: t = WebSocketTextMessage else: t = WebSocketBinaryMessage @@ -101,10 +102,12 @@ class WebSocketLayer(base.Layer): else: frms.append(websockets.Frame( fin=True, - opcode=frame.header.opcode, + opcode=websockets.OPCODE.CONTINUE, mask=(False if is_server else 1), masking_key=(b'' if is_server else os.urandom(4)))) + frms[0].header.opcode = message_type + for frm in frms: other_conn.send(bytes(frm)) -- cgit v1.2.3 From 073a286098f09775f821db2e668fd2d1a771bb74 Mon Sep 17 00:00:00 2001 From: Thomas Kriechbaumer Date: Tue, 29 Nov 2016 22:17:33 +0100 Subject: websocket: reduce connection timeout --- mitmproxy/proxy/protocol/websocket.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py index ec1c6ebc..1b859dc7 100644 --- a/mitmproxy/proxy/protocol/websocket.py +++ b/mitmproxy/proxy/protocol/websocket.py @@ -154,7 +154,7 @@ class WebSocketLayer(base.Layer): try: while not self.channel.should_exit.is_set(): - r = tcp.ssl_read_select(conns, 0.5) + r = tcp.ssl_read_select(conns, 0.1) for conn in r: source_conn = self.client_conn if conn == client else self.server_conn other_conn = self.server_conn if conn == client else self.client_conn -- cgit v1.2.3 From ea6de424a344e5a4fcbaf8c54703331e20079268 Mon Sep 17 00:00:00 2001 From: Thomas Kriechbaumer Date: Tue, 29 Nov 2016 22:32:55 +0100 Subject: websocket: carry over per-message compression bit --- mitmproxy/proxy/protocol/websocket.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mitmproxy/proxy/protocol/websocket.py b/mitmproxy/proxy/protocol/websocket.py index 1b859dc7..d0b12540 100644 --- a/mitmproxy/proxy/protocol/websocket.py +++ b/mitmproxy/proxy/protocol/websocket.py @@ -62,6 +62,7 @@ class WebSocketLayer(base.Layer): payload = b''.join(f.payload for f in fb) original_chunk_sizes = [len(f.payload) for f in fb] message_type = fb[0].header.opcode + compressed_message = fb[0].header.rsv1 fb.clear() if message_type == websockets.OPCODE.TEXT: @@ -107,6 +108,7 @@ class WebSocketLayer(base.Layer): masking_key=(b'' if is_server else os.urandom(4)))) frms[0].header.opcode = message_type + frms[0].header.rsv1 = compressed_message for frm in frms: other_conn.send(bytes(frm)) -- cgit v1.2.3 From d2216801ddb1443bc7365290ecce2c045046c57e Mon Sep 17 00:00:00 2001 From: Thomas Kriechbaumer Date: Tue, 29 Nov 2016 23:01:55 +0100 Subject: websocket: make flowfilter work --- mitmproxy/flowfilter.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/mitmproxy/flowfilter.py b/mitmproxy/flowfilter.py index 3a44efd1..304670db 100644 --- a/mitmproxy/flowfilter.py +++ b/mitmproxy/flowfilter.py @@ -37,6 +37,7 @@ import sys import functools from mitmproxy import http +from mitmproxy import websocket from mitmproxy import tcp from mitmproxy import flow @@ -99,6 +100,14 @@ class FHTTP(_Action): return True +class FWebSocket(_Action): + code = "websocket" + help = "Match WebSocket flows" + + @only(websocket.WebSocketFlow) + def __call__(self, f): + return True + class FTCP(_Action): code = "tcp" help = "Match TCP flows" @@ -245,7 +254,7 @@ class FBod(_Rex): help = "Body" flags = re.DOTALL - @only(http.HTTPFlow, tcp.TCPFlow) + @only(http.HTTPFlow, websocket.WebSocketFlow, tcp.TCPFlow) def __call__(self, f): if isinstance(f, http.HTTPFlow): if f.request and f.request.raw_content: @@ -254,7 +263,7 @@ class FBod(_Rex): if f.response and f.response.raw_content: if self.re.search(f.response.get_content(strict=False)): return True - elif isinstance(f, tcp.TCPFlow): + elif isinstance(f, websocket.WebSocketFlow) or isinstance(f, tcp.TCPFlow): for msg in f.messages: if self.re.search(msg.content): return True @@ -266,13 +275,13 @@ class FBodRequest(_Rex): help = "Request body" flags = re.DOTALL - @only(http.HTTPFlow, tcp.TCPFlow) + @only(http.HTTPFlow, websocket.WebSocketFlow, tcp.TCPFlow) def __call__(self, f): if isinstance(f, http.HTTPFlow): if f.request and f.request.raw_content: if self.re.search(f.request.get_content(strict=False)): return True - elif isinstance(f, tcp.TCPFlow): + elif isinstance(f, websocket.WebSocketFlow) or isinstance(f, tcp.TCPFlow): for msg in f.messages: if msg.from_client and self.re.search(msg.content): return True @@ -283,13 +292,13 @@ class FBodResponse(_Rex): help = "Response body" flags = re.DOTALL - @only(http.HTTPFlow, tcp.TCPFlow) + @only(http.HTTPFlow, websocket.WebSocketFlow, tcp.TCPFlow) def __call__(self, f): if isinstance(f, http.HTTPFlow): if f.response and f.response.raw_content: if self.re.search(f.response.get_content(strict=False)): return True - elif isinstance(f, tcp.TCPFlow): + elif isinstance(f, websocket.WebSocketFlow) or isinstance(f, tcp.TCPFlow): for msg in f.messages: if not msg.from_client and self.re.search(msg.content): return True -- cgit v1.2.3 From cc6aa1f54201aa0634da8a4fa4fafca36609a9c3 Mon Sep 17 00:00:00 2001 From: Thomas Kriechbaumer Date: Tue, 29 Nov 2016 23:47:19 +0100 Subject: websocket: update close handshake tests --- test/mitmproxy/protocol/test_websocket.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/mitmproxy/protocol/test_websocket.py b/test/mitmproxy/protocol/test_websocket.py index e1c3e49a..e42250e0 100644 --- a/test/mitmproxy/protocol/test_websocket.py +++ b/test/mitmproxy/protocol/test_websocket.py @@ -276,6 +276,7 @@ class TestClose(_WebSocketTest): def handle_websockets(cls, rfile, wfile): frame = websockets.Frame.from_file(rfile) wfile.write(bytes(frame)) + wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) wfile.flush() with pytest.raises(exceptions.TcpDisconnect): @@ -287,6 +288,7 @@ class TestClose(_WebSocketTest): client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) client.wfile.flush() + websockets.Frame.from_file(client.rfile) with pytest.raises(exceptions.TcpDisconnect): websockets.Frame.from_file(client.rfile) @@ -296,6 +298,7 @@ class TestClose(_WebSocketTest): client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42'))) client.wfile.flush() + websockets.Frame.from_file(client.rfile) with pytest.raises(exceptions.TcpDisconnect): websockets.Frame.from_file(client.rfile) @@ -305,6 +308,7 @@ class TestClose(_WebSocketTest): client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42foobar'))) client.wfile.flush() + websockets.Frame.from_file(client.rfile) with pytest.raises(exceptions.TcpDisconnect): websockets.Frame.from_file(client.rfile) -- cgit v1.2.3