246 lines
7.5 KiB
Python
246 lines
7.5 KiB
Python
"""Per-bottle TCP forwarder for smolmachines sidecar VMs.
|
|
|
|
smolVM currently publishes guest ports on host loopback as
|
|
``HOST:GUEST`` port pairs, without an address-scoped bind. The agent
|
|
VM's TSI allowlist is IP-only, so bot-bottle exposes a second, stricter
|
|
surface: listeners bound only to the bottle's allocated host address,
|
|
forwarding to the raw smolVM-published loopback ports.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import selectors
|
|
import signal
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
from dataclasses import asdict, dataclass
|
|
from typing import Any, Iterable, Sequence
|
|
|
|
from ...log import die, warn
|
|
|
|
|
|
_BUFFER_SIZE = 64 * 1024
|
|
_LOOPBACK_TARGETS = {"127.0.0.1", "::1"}
|
|
_FORBIDDEN_LISTEN_HOSTS = {"", "0.0.0.0", "::", "127.0.0.1", "::1"}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ForwardSpec:
|
|
label: str
|
|
listen_host: str
|
|
listen_port: int
|
|
target_host: str
|
|
target_port: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ForwarderHandle:
|
|
process: Any
|
|
forwards: tuple[ForwardSpec, ...]
|
|
|
|
|
|
def start_forwarder(specs: Sequence[ForwardSpec]) -> ForwarderHandle:
|
|
"""Start one helper process for a bottle and wait until all
|
|
listeners are bound. ``listen_port`` may be 0; the returned specs
|
|
contain the kernel-assigned ports printed by the helper."""
|
|
if not specs:
|
|
return ForwarderHandle(process=_CompletedProcess(), forwards=())
|
|
_validate_specs(specs)
|
|
argv = [
|
|
sys.executable,
|
|
"-c",
|
|
"from bot_bottle.backend.smolmachines.port_forward import main; "
|
|
"raise SystemExit(main())",
|
|
"--spec-json",
|
|
json.dumps([asdict(s) for s in specs], separators=(",", ":")),
|
|
]
|
|
proc = subprocess.Popen(
|
|
argv,
|
|
stdin=subprocess.DEVNULL,
|
|
stdout=subprocess.PIPE,
|
|
stderr=None,
|
|
text=True,
|
|
env=os.environ,
|
|
)
|
|
assert proc.stdout is not None
|
|
line = proc.stdout.readline()
|
|
proc.stdout.close()
|
|
if not line:
|
|
proc.terminate()
|
|
proc.wait(timeout=5)
|
|
die("smolmachines forwarder failed before reporting readiness")
|
|
try:
|
|
payload = json.loads(line)
|
|
bound = tuple(ForwardSpec(**item) for item in payload["forwards"])
|
|
except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
|
proc.terminate()
|
|
proc.wait(timeout=5)
|
|
die(f"smolmachines forwarder returned invalid readiness payload: {exc}")
|
|
return ForwarderHandle(process=proc, forwards=bound)
|
|
|
|
|
|
def stop_forwarder(handle: ForwarderHandle) -> None:
|
|
proc = handle.process
|
|
if proc.poll() is not None:
|
|
return
|
|
proc.terminate()
|
|
try:
|
|
proc.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
proc.wait(timeout=5)
|
|
|
|
|
|
def _validate_specs(specs: Iterable[ForwardSpec]) -> None:
|
|
for spec in specs:
|
|
if spec.listen_host in _FORBIDDEN_LISTEN_HOSTS:
|
|
die(f"refusing unsafe smolmachines forwarder bind: {spec.listen_host!r}")
|
|
if spec.target_host not in _LOOPBACK_TARGETS:
|
|
die(f"refusing non-loopback smolmachines forwarder target: {spec.target_host!r}")
|
|
if not 0 <= spec.listen_port <= 65535:
|
|
die(f"invalid smolmachines forwarder listen port: {spec.listen_port}")
|
|
if not 1 <= spec.target_port <= 65535:
|
|
die(f"invalid smolmachines forwarder target port: {spec.target_port}")
|
|
|
|
|
|
class _CompletedProcess:
|
|
"""Popen-like no-op used when a bottle has no forwards."""
|
|
|
|
def poll(self) -> int:
|
|
return 0
|
|
|
|
def terminate(self) -> None:
|
|
return None
|
|
|
|
def kill(self) -> None:
|
|
return None
|
|
|
|
def wait(self, timeout: float | None = None) -> int:
|
|
del timeout
|
|
return 0
|
|
|
|
|
|
class _Forwarder:
|
|
def __init__(self, specs: Sequence[ForwardSpec]):
|
|
_validate_specs(specs)
|
|
self.specs = tuple(specs)
|
|
self.selector = selectors.DefaultSelector()
|
|
self.stop_event = threading.Event()
|
|
self.listeners: list[tuple[socket.socket, ForwardSpec]] = []
|
|
self.bound: list[ForwardSpec] = []
|
|
|
|
def start(self) -> None:
|
|
for spec in self.specs:
|
|
listener = socket.create_server(
|
|
(spec.listen_host, spec.listen_port),
|
|
reuse_port=False,
|
|
backlog=64,
|
|
)
|
|
listener.setblocking(False)
|
|
host, port = listener.getsockname()[:2]
|
|
bound = ForwardSpec(
|
|
label=spec.label,
|
|
listen_host=str(host),
|
|
listen_port=int(port),
|
|
target_host=spec.target_host,
|
|
target_port=spec.target_port,
|
|
)
|
|
self.listeners.append((listener, bound))
|
|
self.bound.append(bound)
|
|
self.selector.register(listener, selectors.EVENT_READ, bound)
|
|
|
|
def serve(self) -> None:
|
|
while not self.stop_event.is_set():
|
|
try:
|
|
events = self.selector.select(timeout=0.25)
|
|
except OSError:
|
|
return
|
|
for key, _ in events:
|
|
listener = key.fileobj
|
|
if not isinstance(listener, socket.socket):
|
|
continue
|
|
spec = key.data
|
|
try:
|
|
client, _ = listener.accept()
|
|
except OSError:
|
|
continue
|
|
threading.Thread(
|
|
target=self._handle_client,
|
|
args=(client, spec),
|
|
daemon=True,
|
|
).start()
|
|
|
|
def stop(self) -> None:
|
|
self.stop_event.set()
|
|
for listener, _ in self.listeners:
|
|
try:
|
|
self.selector.unregister(listener)
|
|
except Exception: # noqa: BLE001 - best effort shutdown
|
|
pass
|
|
listener.close()
|
|
self.selector.close()
|
|
|
|
def _handle_client(self, client: socket.socket, spec: ForwardSpec) -> None:
|
|
with client:
|
|
try:
|
|
target = socket.create_connection((spec.target_host, spec.target_port))
|
|
except OSError as exc:
|
|
warn(f"smolmachines forwarder {spec.label}: target connect failed: {exc}")
|
|
return
|
|
with target:
|
|
left = threading.Thread(target=_pipe, args=(client, target), daemon=True)
|
|
right = threading.Thread(target=_pipe, args=(target, client), daemon=True)
|
|
left.start()
|
|
right.start()
|
|
left.join()
|
|
right.join()
|
|
|
|
|
|
def _pipe(src: socket.socket, dst: socket.socket) -> None:
|
|
while True:
|
|
try:
|
|
data = src.recv(_BUFFER_SIZE)
|
|
except OSError:
|
|
break
|
|
if not data:
|
|
break
|
|
try:
|
|
dst.sendall(data)
|
|
except OSError:
|
|
break
|
|
try:
|
|
dst.shutdown(socket.SHUT_WR)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--spec-json", required=True)
|
|
ns = parser.parse_args(argv)
|
|
specs = tuple(ForwardSpec(**item) for item in json.loads(ns.spec_json))
|
|
forwarder = _Forwarder(specs)
|
|
|
|
def request_stop(signum: int, _frame: object) -> None:
|
|
del signum
|
|
forwarder.stop()
|
|
|
|
signal.signal(signal.SIGTERM, request_stop)
|
|
signal.signal(signal.SIGINT, request_stop)
|
|
forwarder.start()
|
|
print(json.dumps({"forwards": [asdict(s) for s in forwarder.bound]}), flush=True)
|
|
try:
|
|
forwarder.serve()
|
|
finally:
|
|
forwarder.stop()
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|