mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 00:01:36 +00:00
[sans-io] add tunnel tests
This commit is contained in:
Maximilian Hils
parent
f0bdf887fc
commit
8f49cde627
@ -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:
|
||||
yield from self.event_to_child(events.OpenConnectionReply(self.command_to_reply_to, err))
|
||||
self.command_to_reply_to = None
|
||||
if done or err:
|
||||
yield from self._handshake_finished(err)
|
||||
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")
|
||||
else:
|
||||
err = "connection closed without notice"
|
||||
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
|
||||
|
||||
@ -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
|
||||
Loading…
Reference in New Issue
Block a user