websocket: inject messages via flow

2 files changed, 55 insertions, 0 deletions

diff --git a/test/mitmproxy/proxy/protocol/test_websocket.py b/test/mitmproxy/proxy/protocol/test_websocket.py
index 1f4e2bca..0b26ed29 100644
--- a/test/mitmproxy/proxy/protocol/test_websocket.py
+++ b/test/mitmproxy/proxy/protocol/test_websocket.py
@@ -467,3 +467,46 @@ class TestExtension(_WebSocketTest):
         assert self.master.state.flows[1].messages[3].type == websockets.OPCODE.BINARY
         assert self.master.state.flows[1].messages[4].content == b'\xde\xad\xbe\xef'
         assert self.master.state.flows[1].messages[4].type == websockets.OPCODE.BINARY
+
+
+class TestInjectMessageClient(_WebSocketTest):
+
+    @classmethod
+    def handle_websockets(cls, rfile, wfile):
+        pass
+
+    def test_inject_message_client(self):
+        class Inject:
+            def websocket_start(self, flow):
+                flow.inject_message(flow.client_conn, 'This is an injected message!')
+
+        self.proxy.set_addons(Inject())
+        self.setup_connection()
+
+        frame = websockets.Frame.from_file(self.client.rfile)
+        assert frame.header.opcode == websockets.OPCODE.TEXT
+        assert frame.payload == b'This is an injected message!'
+
+
+class TestInjectMessageServer(_WebSocketTest):
+
+    @classmethod
+    def handle_websockets(cls, rfile, wfile):
+        frame = websockets.Frame.from_file(rfile)
+        assert frame.header.opcode == websockets.OPCODE.TEXT
+        success = frame.payload == b'This is an injected message!'
+
+        wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=str(success).encode())))
+        wfile.flush()
+
+    def test_inject_message_server(self):
+        class Inject:
+            def websocket_start(self, flow):
+                flow.inject_message(flow.server_conn, 'This is an injected message!')
+
+        self.proxy.set_addons(Inject())
+        self.setup_connection()
+
+        frame = websockets.Frame.from_file(self.client.rfile)
+        assert frame.header.opcode == websockets.OPCODE.TEXT
+        assert frame.payload == b'True'
diff --git a/test/mitmproxy/test_websocket.py b/test/mitmproxy/test_websocket.py
index fcacec36..bd4bb4c9 100644
--- a/test/mitmproxy/test_websocket.py
+++ b/test/mitmproxy/test_websocket.py
@@ -92,3 +92,15 @@ class TestWebSocketFlow:
         assert not f.messages[-1].killed
         f.messages[-1].kill()
         assert f.messages[-1].killed
+
+    def test_inject_message(self):
+        f = tflow.twebsocketflow()
+
+        with pytest.raises(ValueError):
+            f.inject_message(None, 'foobar')
+
+        f.inject_message(f.client_conn, 'foobar')
+        assert f._inject_messages_client.qsize() == 1
+
+        f.inject_message(f.server_conn, 'foobar')
+        assert f._inject_messages_client.qsize() == 1