[sans-io] add tunnel tests

mitmproxy/proxy2/tunnel.py

@@ -24,6 +24,11 @@ class TunnelLayer(layer.Layer):
     """The 'inner' connection which provides data I/O"""
     tunnel_state: TunnelState = TunnelState.INACTIVE
     command_to_reply_to: Optional[commands.OpenConnection] = None
+    _event_queue: List[events.Event]
+    """
+    If the connection already exists when we receive the start event,
+    we buffer commands until we have established the tunnel.
+    """
 
     def __init__(
             self,
@@ -35,6 +40,7 @@ class TunnelLayer(layer.Layer):
         self.tunnel_connection = tunnel_connection
         self.conn = conn
         self.child_layer = layer.NextLayer(self.context)
+        self._event_queue = []
 
     def __repr__(self):
         return f"{type(self).__name__}({self.tunnel_state.name.lower()})"
@@ -55,13 +61,12 @@ class TunnelLayer(layer.Layer):
                     if done:
                         if self.conn != self.tunnel_connection:
                             self.conn.state = context.ConnectionState.OPEN
-                        self.tunnel_state = TunnelState.OPEN
                     if err:
-                        self.tunnel_state = TunnelState.CLOSED
+                        if self.conn != self.tunnel_connection:
+                            self.conn.state = context.ConnectionState.CLOSED
                         yield from self.on_handshake_error(err)
-                    if (done or err) and self.command_to_reply_to:
+                    if done or err:
-                        yield from self.event_to_child(events.OpenConnectionReply(self.command_to_reply_to, err))
+                        yield from self._handshake_finished(err)
-                        self.command_to_reply_to = None
                 else:
                     yield from self.receive_data(event.data)
             elif isinstance(event, events.ConnectionClosed):
@@ -69,13 +74,32 @@ class TunnelLayer(layer.Layer):
                 if self.tunnel_state is TunnelState.OPEN:
                     yield from self.receive_close()
                 elif self.tunnel_state is TunnelState.ESTABLISHING:
-                    yield from self.on_handshake_error("connection closed without notice")
+                    err = "connection closed without notice"
-            else:
+                    yield from self.on_handshake_error(err)
+                    yield from self._handshake_finished(err)
+                self.tunnel_state = TunnelState.CLOSED
+            else:  # pragma: no cover
                 raise AssertionError(f"Unexpected event: {event}")
         else:
             yield from self.event_to_child(event)
 
+    def _handshake_finished(self, err: Optional[str]):
+        if err:
+            self.tunnel_state = TunnelState.CLOSED
+        else:
+            self.tunnel_state = TunnelState.OPEN
+        if self.command_to_reply_to:
+            yield from self.event_to_child(events.OpenConnectionReply(self.command_to_reply_to, err))
+            self.command_to_reply_to = None
+        else:
+            for evt in self._event_queue:
+                yield from self.event_to_child(evt)
+            self._event_queue.clear()
+
     def event_to_child(self, event: events.Event) -> layer.CommandGenerator[None]:
+        if self.tunnel_state is TunnelState.ESTABLISHING and not self.command_to_reply_to:
+            self._event_queue.append(event)
+            return
         for command in self.child_layer.handle_event(event):
             if isinstance(command, commands.ConnectionCommand) and command.connection == self.conn:
                 if isinstance(command, commands.SendData):
@@ -89,14 +113,15 @@ class TunnelLayer(layer.Layer):
                     yield from self.send_close(command.half_close)
                 elif isinstance(command, commands.OpenConnection):
                     # create our own OpenConnection command object that blocks here.
+                    self.command_to_reply_to = command
                     self.tunnel_state = TunnelState.ESTABLISHING
                     err = yield commands.OpenConnection(self.tunnel_connection)
                     if err:
                         yield from self.event_to_child(events.OpenConnectionReply(command, err))
+                        self.tunnel_state = TunnelState.CLOSED
                     else:
-                        self.command_to_reply_to = command
                         yield from self.start_handshake()
-                else:
+                else:  # pragma: no cover
                     raise AssertionError(f"Unexpected command: {command}")
             else:
                 yield command
@@ -105,6 +130,7 @@ class TunnelLayer(layer.Layer):
         yield from self._handle_event(events.DataReceived(self.tunnel_connection, b""))
 
     def receive_handshake_data(self, data: bytes) -> layer.CommandGenerator[Tuple[bool, Optional[str]]]:
+        """returns a (done, err) tuple"""
         yield from ()
         return True, None
 
@@ -142,14 +168,3 @@ class LayerStack:
         self._stack.append(other)
         return self
 
-
-class OpenConnectionStub(layer.Layer):
-    done = False
-    err = None
-
-    def _handle_event(self, event: events.Event) -> layer.CommandGenerator[None]:
-        if isinstance(event, events.Start):
-            self.err = yield commands.OpenConnection(self.context.server)
-            self.done = not self.err
-        else:
-            self.err = event

test/mitmproxy/proxy2/test_tunnel.py

@@ -0,0 +1,262 @@
+from typing import Tuple, Optional
+
+import pytest
+
+from mitmproxy.proxy2 import tunnel, layer
+from mitmproxy.proxy2.commands import SendData, Log, CloseConnection, OpenConnection
+from mitmproxy.proxy2.context import Context, Server, ConnectionState
+from mitmproxy.proxy2.events import Event, DataReceived, Start, ConnectionClosed
+from test.mitmproxy.proxy2.tutils import Playbook, reply
+
+
+class TChildLayer(layer.Layer):
+    child_layer: Optional[layer.Layer] = None
+
+    def _handle_event(self, event: Event) -> layer.CommandGenerator[None]:
+        if isinstance(event, Start):
+            yield Log(f"Got start. Server state: {self.context.server.state.name}")
+        elif isinstance(event, DataReceived) and event.data == b"client-hello":
+            yield SendData(self.context.client, b"client-hello-reply")
+        elif isinstance(event, DataReceived) and event.data == b"server-hello":
+            yield SendData(self.context.server, b"server-hello-reply")
+        elif isinstance(event, DataReceived) and event.data == b"open":
+            err = yield OpenConnection(self.context.server)
+            yield Log(f"Opened: {err=}. Server state: {self.context.server.state.name}")
+        elif isinstance(event, DataReceived) and event.data == b"half-close":
+            err = yield CloseConnection(event.connection, half_close=True)
+        elif isinstance(event, ConnectionClosed):
+            yield Log(f"Got {event.connection.__class__.__name__.lower()} close.")
+            yield CloseConnection(event.connection)
+        else:
+            raise AssertionError
+
+
+class TTunnelLayer(tunnel.TunnelLayer):
+    def start_handshake(self) -> layer.CommandGenerator[None]:
+        yield SendData(self.tunnel_connection, b"handshake-hello")
+
+    def receive_handshake_data(self, data: bytes) -> layer.CommandGenerator[Tuple[bool, Optional[str]]]:
+        yield SendData(self.tunnel_connection, data)
+        if data == b"handshake-success":
+            return True, None
+        else:
+            return False, "handshake error"
+
+    def send_data(self, data: bytes) -> layer.CommandGenerator[None]:
+        yield SendData(self.tunnel_connection, b"tunneled-" + data)
+
+    def receive_data(self, data: bytes) -> layer.CommandGenerator[None]:
+        yield from self.event_to_child(
+            DataReceived(self.conn, data.replace(b"tunneled-", b""))
+        )
+
+
+@pytest.mark.parametrize("success", ["success", "fail"])
+def test_tunnel_handshake_start(tctx: Context, success):
+    server = Server(("proxy", 1234))
+    server.state = ConnectionState.OPEN
+
+    tl = TTunnelLayer(tctx, server, tctx.server)
+    tl.child_layer = TChildLayer(tctx)
+    assert repr(tl)
+
+    playbook = Playbook(tl, logs=True)
+    (
+            playbook
+            << SendData(server, b"handshake-hello")
+            >> DataReceived(tctx.client, b"client-hello")
+            >> DataReceived(server, b"handshake-" + success.encode())
+            << SendData(server, b"handshake-" + success.encode())
+    )
+    if success == "success":
+        playbook << Log("Got start. Server state: OPEN")
+    else:
+        playbook << CloseConnection(server)
+        playbook << Log("Got start. Server state: CLOSED")
+
+    playbook << SendData(tctx.client, b"client-hello-reply")
+    if success == "success":
+        playbook >> DataReceived(server, b"tunneled-server-hello")
+        playbook << SendData(server, b"tunneled-server-hello-reply")
+
+    assert playbook
+
+
+@pytest.mark.parametrize("success", ["success", "fail"])
+def test_tunnel_handshake_command(tctx: Context, success):
+    server = Server(("proxy", 1234))
+
+    tl = TTunnelLayer(tctx, server, tctx.server)
+    tl.child_layer = TChildLayer(tctx)
+
+    playbook = Playbook(tl, logs=True)
+    (
+            playbook
+            << Log("Got start. Server state: CLOSED")
+            >> DataReceived(tctx.client, b"client-hello")
+            << SendData(tctx.client, b"client-hello-reply")
+            >> DataReceived(tctx.client, b"open")
+            << OpenConnection(server)
+            >> reply(None)
+            << SendData(server, b"handshake-hello")
+            >> DataReceived(server, b"handshake-" + success.encode())
+            << SendData(server, b"handshake-" + success.encode())
+    )
+    if success == "success":
+        assert (
+                playbook
+                << Log(f"Opened: err=None. Server state: OPEN")
+                >> DataReceived(server, b"tunneled-server-hello")
+                << SendData(server, b"tunneled-server-hello-reply")
+                >> ConnectionClosed(tctx.client)
+                << Log("Got client close.")
+                << CloseConnection(tctx.client)
+        )
+        assert tl.tunnel_state is tunnel.TunnelState.OPEN
+        assert (
+                playbook
+                >> ConnectionClosed(server)
+                << Log("Got server close.")
+                << CloseConnection(server)
+        )
+        assert tl.tunnel_state is tunnel.TunnelState.CLOSED
+    else:
+        assert (
+                playbook
+                << CloseConnection(server)
+                << Log("Opened: err='handshake error'. Server state: CLOSED")
+        )
+        assert tl.tunnel_state is tunnel.TunnelState.CLOSED
+
+
+def test_tunnel_default_impls(tctx: Context):
+    tctx.server.state = ConnectionState.OPEN
+    tl = tunnel.TunnelLayer(tctx, tctx.server, Server(None))
+    tl.child_layer = TChildLayer(tctx)
+    playbook = Playbook(tl, logs=True)
+    assert (
+            playbook
+            << Log("Got start. Server state: OPEN")
+            >> DataReceived(tctx.server, b"server-hello")
+            << SendData(tctx.server, b"server-hello-reply")
+    )
+    assert tl.tunnel_state is tunnel.TunnelState.OPEN
+    assert (
+            playbook
+            >> ConnectionClosed(tctx.server)
+            << Log("Got server close.")
+            << CloseConnection(tctx.server)
+    )
+    assert tl.tunnel_state is tunnel.TunnelState.CLOSED
+
+    assert (
+            playbook
+            >> DataReceived(tctx.client, b"open")
+            << OpenConnection(tctx.server)
+            >> reply(None)
+            << Log("Opened: err=None. Server state: OPEN")
+            >> DataReceived(tctx.server, b"half-close")
+            << CloseConnection(tctx.server, half_close=True)
+    )
+
+
+def test_tunnel_openconnection_error(tctx: Context):
+    server = Server(("proxy", 1234))
+
+    tl = TTunnelLayer(tctx, server, tctx.server)
+    tl.child_layer = TChildLayer(tctx)
+
+    playbook = Playbook(tl, logs=True)
+    assert (
+            playbook
+            << Log("Got start. Server state: CLOSED")
+            >> DataReceived(tctx.client, b"open")
+            << OpenConnection(server)
+    )
+    assert tl.tunnel_state is tunnel.TunnelState.ESTABLISHING
+    assert (
+            playbook
+            >> reply("IPoAC packet dropped.")
+            << Log("Opened: err='IPoAC packet dropped.'. Server state: CLOSED")
+    )
+    assert tl.tunnel_state is tunnel.TunnelState.CLOSED
+
+
+@pytest.mark.parametrize("disconnect", ["client", "server"])
+def test_disconnect_during_handshake_start(tctx: Context, disconnect):
+    server = Server(("proxy", 1234))
+    server.state = ConnectionState.OPEN
+
+    tl = TTunnelLayer(tctx, server, tctx.server)
+    tl.child_layer = TChildLayer(tctx)
+
+    playbook = Playbook(tl, logs=True)
+
+    assert (
+            playbook
+            << SendData(server, b"handshake-hello")
+    )
+    if disconnect == "client":
+        assert (
+                playbook
+                >> ConnectionClosed(tctx.client)
+                >> ConnectionClosed(server)  # proxyserver will cancel all other connections as well.
+                << CloseConnection(server)
+                << Log("Got start. Server state: CLOSED")
+                << Log("Got client close.")
+                << CloseConnection(tctx.client)
+        )
+    else:
+        assert (
+                playbook
+                >> ConnectionClosed(server)
+                << CloseConnection(server)
+                << Log("Got start. Server state: CLOSED")
+        )
+
+
+@pytest.mark.parametrize("disconnect", ["client", "server"])
+def test_disconnect_during_handshake_command(tctx: Context, disconnect):
+    server = Server(("proxy", 1234))
+
+    tl = TTunnelLayer(tctx, server, tctx.server)
+    tl.child_layer = TChildLayer(tctx)
+
+    playbook = Playbook(tl, logs=True)
+    assert (
+            playbook
+            << Log("Got start. Server state: CLOSED")
+            >> DataReceived(tctx.client, b"client-hello")
+            << SendData(tctx.client, b"client-hello-reply")
+            >> DataReceived(tctx.client, b"open")
+            << OpenConnection(server)
+            >> reply(None)
+            << SendData(server, b"handshake-hello")
+    )
+    if disconnect == "client":
+        assert (
+                playbook
+                >> ConnectionClosed(tctx.client)
+                >> ConnectionClosed(server)  # proxyserver will cancel all other connections as well.
+                << CloseConnection(server)
+                << Log("Opened: err='connection closed without notice'. Server state: CLOSED")
+                << Log("Got client close.")
+                << CloseConnection(tctx.client)
+        )
+    else:
+        assert (
+                playbook
+                >> ConnectionClosed(server)
+                << CloseConnection(server)
+                << Log("Opened: err='connection closed without notice'. Server state: CLOSED")
+        )
+
+
+def test_layer_stack(tctx):
+    stack = tunnel.LayerStack()
+    a = TChildLayer(tctx)
+    b = TChildLayer(tctx)
+    stack /= a
+    stack /= b
+    assert stack[0] == a
+    assert a.child_layer is b