326 lines
10 KiB
Python
326 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import socket
|
|
import subprocess
|
|
import threading
|
|
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from bot_bottle.backend.smolmachines.port_forward import (
|
|
ForwardSpec,
|
|
ForwarderHandle,
|
|
_CompletedProcess,
|
|
_Forwarder,
|
|
_pipe,
|
|
main,
|
|
start_forwarder,
|
|
stop_forwarder,
|
|
)
|
|
|
|
|
|
def _echo_server() -> tuple[socket.socket, int]:
|
|
listener = socket.create_server(("127.0.0.1", 0))
|
|
port = int(listener.getsockname()[1])
|
|
|
|
def run() -> None:
|
|
with listener:
|
|
conn, _ = listener.accept()
|
|
with conn:
|
|
while True:
|
|
data = conn.recv(65536)
|
|
if not data:
|
|
return
|
|
conn.sendall(data)
|
|
|
|
threading.Thread(target=run, daemon=True).start()
|
|
return listener, port
|
|
|
|
|
|
class TestForwarderValidation(unittest.TestCase):
|
|
def test_empty_specs_returns_noop_handle(self):
|
|
handle = start_forwarder(())
|
|
|
|
self.assertEqual((), handle.forwards)
|
|
self.assertEqual(0, handle.process.poll())
|
|
stop_forwarder(handle)
|
|
|
|
def test_rejects_generic_localhost_listener(self):
|
|
with self.assertRaises(SystemExit):
|
|
start_forwarder((
|
|
ForwardSpec(
|
|
label="egress",
|
|
listen_host="127.0.0.1",
|
|
listen_port=0,
|
|
target_host="127.0.0.1",
|
|
target_port=9099,
|
|
),
|
|
))
|
|
|
|
def test_rejects_wildcard_listener(self):
|
|
with self.assertRaises(SystemExit):
|
|
start_forwarder((
|
|
ForwardSpec(
|
|
label="egress",
|
|
listen_host="0.0.0.0",
|
|
listen_port=0,
|
|
target_host="127.0.0.1",
|
|
target_port=9099,
|
|
),
|
|
))
|
|
|
|
def test_rejects_non_loopback_target(self):
|
|
with self.assertRaises(SystemExit):
|
|
start_forwarder((
|
|
ForwardSpec(
|
|
label="egress",
|
|
listen_host="127.0.0.16",
|
|
listen_port=0,
|
|
target_host="192.0.2.10",
|
|
target_port=9099,
|
|
),
|
|
))
|
|
|
|
def test_rejects_invalid_ports(self):
|
|
with self.assertRaises(SystemExit):
|
|
start_forwarder((
|
|
ForwardSpec("bad-listen", "127.0.0.2", 70000, "127.0.0.1", 9099),
|
|
))
|
|
with self.assertRaises(SystemExit):
|
|
start_forwarder((
|
|
ForwardSpec("bad-target", "127.0.0.2", 0, "127.0.0.1", 0),
|
|
))
|
|
|
|
def test_completed_process_methods_are_noops(self):
|
|
proc = _CompletedProcess()
|
|
|
|
proc.terminate()
|
|
proc.kill()
|
|
|
|
self.assertEqual(0, proc.poll())
|
|
self.assertEqual(0, proc.wait(timeout=1))
|
|
|
|
|
|
class _FakeStdout:
|
|
def __init__(self, line: str):
|
|
self.line = line
|
|
self.closed = False
|
|
|
|
def readline(self) -> str:
|
|
return self.line
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class TestForwarderProcess(unittest.TestCase):
|
|
def test_start_parses_ready_payload(self):
|
|
payload = {
|
|
"forwards": [{
|
|
"label": "egress",
|
|
"listen_host": "127.0.0.2",
|
|
"listen_port": 61000,
|
|
"target_host": "127.0.0.1",
|
|
"target_port": 55000,
|
|
}],
|
|
}
|
|
proc = MagicMock()
|
|
proc.stdout = _FakeStdout(json.dumps(payload) + "\n")
|
|
|
|
with patch(
|
|
"bot_bottle.backend.smolmachines.port_forward.subprocess.Popen",
|
|
return_value=proc,
|
|
) as popen:
|
|
handle = start_forwarder((
|
|
ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 55000),
|
|
))
|
|
|
|
self.assertIs(proc, handle.process)
|
|
self.assertEqual(61000, handle.forwards[0].listen_port)
|
|
self.assertTrue(proc.stdout.closed)
|
|
argv = popen.call_args.args[0]
|
|
self.assertIn("--spec-json", argv)
|
|
|
|
def test_start_dies_when_helper_exits_before_ready(self):
|
|
proc = MagicMock()
|
|
proc.stdout = _FakeStdout("")
|
|
|
|
with patch(
|
|
"bot_bottle.backend.smolmachines.port_forward.subprocess.Popen",
|
|
return_value=proc,
|
|
):
|
|
with self.assertRaises(SystemExit):
|
|
start_forwarder((
|
|
ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 55000),
|
|
))
|
|
|
|
proc.terminate.assert_called_once()
|
|
proc.wait.assert_called_once_with(timeout=5)
|
|
|
|
def test_start_dies_on_invalid_ready_payload(self):
|
|
proc = MagicMock()
|
|
proc.stdout = _FakeStdout('{"wrong": []}\n')
|
|
|
|
with patch(
|
|
"bot_bottle.backend.smolmachines.port_forward.subprocess.Popen",
|
|
return_value=proc,
|
|
):
|
|
with self.assertRaises(SystemExit):
|
|
start_forwarder((
|
|
ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 55000),
|
|
))
|
|
|
|
proc.terminate.assert_called_once()
|
|
proc.wait.assert_called_once_with(timeout=5)
|
|
|
|
def test_stop_kills_process_after_timeout(self):
|
|
proc = MagicMock()
|
|
proc.poll.return_value = None
|
|
proc.wait.side_effect = [subprocess.TimeoutExpired("cmd", 5), 0]
|
|
|
|
stop_forwarder(ForwarderHandle(process=proc, forwards=()))
|
|
|
|
proc.terminate.assert_called_once()
|
|
proc.kill.assert_called_once()
|
|
self.assertEqual(2, proc.wait.call_count)
|
|
|
|
|
|
class TestForwarding(unittest.TestCase):
|
|
def test_forwards_bytes_both_directions(self):
|
|
try:
|
|
probe = socket.create_server(("127.0.0.2", 0))
|
|
except OSError:
|
|
self.skipTest("127.0.0.2 is not bindable on this host")
|
|
else:
|
|
probe.close()
|
|
_listener, target_port = _echo_server()
|
|
handle = start_forwarder((
|
|
ForwardSpec(
|
|
label="egress",
|
|
listen_host="127.0.0.2",
|
|
listen_port=0,
|
|
target_host="127.0.0.1",
|
|
target_port=target_port,
|
|
),
|
|
))
|
|
try:
|
|
bound = handle.forwards[0]
|
|
with socket.create_connection((bound.listen_host, bound.listen_port)) as sock:
|
|
sock.sendall(b"hello")
|
|
self.assertEqual(b"hello", sock.recv(5))
|
|
finally:
|
|
stop_forwarder(handle)
|
|
|
|
|
|
class TestForwarderInProcess(unittest.TestCase):
|
|
def test_forwarder_serves_and_stops_in_process(self):
|
|
_listener, target_port = _echo_server()
|
|
forwarder = _Forwarder((
|
|
ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", target_port),
|
|
))
|
|
forwarder.start()
|
|
thread = threading.Thread(target=forwarder.serve, daemon=True)
|
|
thread.start()
|
|
try:
|
|
bound = forwarder.bound[0]
|
|
with socket.create_connection((bound.listen_host, bound.listen_port)) as sock:
|
|
sock.sendall(b"ping")
|
|
self.assertEqual(b"ping", sock.recv(4))
|
|
finally:
|
|
forwarder.stop()
|
|
thread.join(timeout=2)
|
|
|
|
self.assertFalse(thread.is_alive())
|
|
|
|
def test_forwarder_logs_target_connect_failure(self):
|
|
forwarder = _Forwarder((
|
|
ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 1),
|
|
))
|
|
forwarder.start()
|
|
thread = threading.Thread(target=forwarder.serve, daemon=True)
|
|
thread.start()
|
|
try:
|
|
bound = forwarder.bound[0]
|
|
with patch("bot_bottle.backend.smolmachines.port_forward.warn") as warn:
|
|
with socket.create_connection((bound.listen_host, bound.listen_port)):
|
|
pass
|
|
for _ in range(100):
|
|
if warn.called:
|
|
break
|
|
threading.Event().wait(0.01)
|
|
self.assertTrue(warn.called)
|
|
finally:
|
|
forwarder.stop()
|
|
thread.join(timeout=2)
|
|
|
|
def test_pipe_stops_when_recv_fails(self):
|
|
src, dst = socket.socketpair()
|
|
src.close()
|
|
try:
|
|
_pipe(src, dst)
|
|
finally:
|
|
dst.close()
|
|
|
|
def test_pipe_stops_when_send_fails(self):
|
|
src, src_peer = socket.socketpair()
|
|
_dst_peer, dst = socket.socketpair()
|
|
dst.close()
|
|
try:
|
|
src_peer.sendall(b"hello")
|
|
src_peer.shutdown(socket.SHUT_WR)
|
|
_pipe(src, dst)
|
|
finally:
|
|
src.close()
|
|
src_peer.close()
|
|
_dst_peer.close()
|
|
|
|
|
|
class TestMain(unittest.TestCase):
|
|
def test_main_starts_serves_and_stops_forwarder(self):
|
|
calls: list[str] = []
|
|
handlers = []
|
|
|
|
class FakeForwarder:
|
|
def __init__(self, specs): # type: ignore[no-untyped-def]
|
|
self.specs = specs
|
|
self.bound = specs
|
|
|
|
def start(self) -> None:
|
|
calls.append("start")
|
|
|
|
def serve(self) -> None:
|
|
calls.append("serve")
|
|
|
|
def stop(self) -> None:
|
|
calls.append("stop")
|
|
|
|
def record_signal(_signum, handler): # type: ignore[no-untyped-def]
|
|
handlers.append(handler)
|
|
|
|
spec = ForwardSpec("egress", "127.0.0.2", 0, "127.0.0.1", 9099)
|
|
with patch(
|
|
"bot_bottle.backend.smolmachines.port_forward._Forwarder",
|
|
FakeForwarder,
|
|
), patch(
|
|
"bot_bottle.backend.smolmachines.port_forward.signal.signal",
|
|
side_effect=record_signal,
|
|
), patch("builtins.print") as print_mock:
|
|
self.assertEqual(
|
|
0,
|
|
main(("--spec-json", json.dumps([{
|
|
"label": spec.label,
|
|
"listen_host": spec.listen_host,
|
|
"listen_port": spec.listen_port,
|
|
"target_host": spec.target_host,
|
|
"target_port": spec.target_port,
|
|
}]))),
|
|
)
|
|
handlers[0](15, None)
|
|
|
|
self.assertEqual(["start", "serve", "stop", "stop"], calls)
|
|
self.assertIn("forwards", print_mock.call_args.args[0])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|