| 1 | """Rootless byte-counting TCP proxy for benchmark scaffolding.
|
| 2 |
|
| 3 | The proxy binds to localhost, forwards to a configured upstream, and records
|
| 4 | bytes in both directions plus connection counts. It needs no packet capture,
|
| 5 | pf/iptables, or privileged ports.
|
| 6 | """
|
| 7 |
|
| 8 | from __future__ import annotations
|
| 9 |
|
| 10 | import socket
|
| 11 | import threading
|
| 12 | import time
|
| 13 | from dataclasses import dataclass
|
| 14 | from typing import Optional
|
| 15 |
|
| 16 |
|
| 17 | @dataclass(frozen=True)
|
| 18 | class ByteProxyStats:
|
| 19 | connections: int
|
| 20 | client_to_upstream_bytes: int
|
| 21 | upstream_to_client_bytes: int
|
| 22 | connection_errors: int
|
| 23 | first_upstream_to_client_ms: float | None
|
| 24 | first_payload_bytes_server_to_client: int | None
|
| 25 |
|
| 26 | @property
|
| 27 | def total_bytes(self) -> int:
|
| 28 | return self.client_to_upstream_bytes + self.upstream_to_client_bytes
|
| 29 |
|
| 30 |
|
| 31 | class CountingProxy:
|
| 32 | def __init__(
|
| 33 | self,
|
| 34 | upstream_host: str,
|
| 35 | upstream_port: int,
|
| 36 | *,
|
| 37 | listen_host: str = "127.0.0.1",
|
| 38 | listen_port: int = 0,
|
| 39 | connect_timeout_s: float = 5.0,
|
| 40 | ) -> None:
|
| 41 | self.upstream_host = upstream_host
|
| 42 | self.upstream_port = upstream_port
|
| 43 | self.listen_host = listen_host
|
| 44 | self.listen_port = listen_port
|
| 45 | self.connect_timeout_s = connect_timeout_s
|
| 46 |
|
| 47 | self._lock = threading.Lock()
|
| 48 | self._stop = threading.Event()
|
| 49 | self._listener: Optional[socket.socket] = None
|
| 50 | self._accept_thread: Optional[threading.Thread] = None
|
| 51 | self._connection_threads: list[threading.Thread] = []
|
| 52 | self._active_sockets: set[socket.socket] = set()
|
| 53 |
|
| 54 | self._connections = 0
|
| 55 | self._client_to_upstream_bytes = 0
|
| 56 | self._upstream_to_client_bytes = 0
|
| 57 | self._connection_errors = 0
|
| 58 | self._count_start = time.monotonic()
|
| 59 | self._first_upstream_to_client_ms: float | None = None
|
| 60 | self._first_payload_bytes_server_to_client: int | None = None
|
| 61 |
|
| 62 | @property
|
| 63 | def address(self) -> tuple[str, int]:
|
| 64 | if self._listener is None:
|
| 65 | raise RuntimeError("proxy is not started")
|
| 66 | host, port = self._listener.getsockname()[:2]
|
| 67 | return str(host), int(port)
|
| 68 |
|
| 69 | def start(self) -> "CountingProxy":
|
| 70 | if self._listener is not None:
|
| 71 | return self
|
| 72 | listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
| 73 | listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
| 74 | listener.bind((self.listen_host, self.listen_port))
|
| 75 | listener.listen()
|
| 76 | listener.settimeout(0.2)
|
| 77 | self._listener = listener
|
| 78 | self._accept_thread = threading.Thread(target=self._accept_loop, name="oakbench-byteproxy-accept", daemon=True)
|
| 79 | self._accept_thread.start()
|
| 80 | return self
|
| 81 |
|
| 82 | def stop(self) -> None:
|
| 83 | self._stop.set()
|
| 84 | listener = self._listener
|
| 85 | if listener is not None:
|
| 86 | try:
|
| 87 | listener.close()
|
| 88 | except OSError:
|
| 89 | pass
|
| 90 | with self._lock:
|
| 91 | sockets = list(self._active_sockets)
|
| 92 | for sock in sockets:
|
| 93 | _close_socket(sock)
|
| 94 | if self._accept_thread is not None:
|
| 95 | self._accept_thread.join(timeout=1.0)
|
| 96 | for thread in list(self._connection_threads):
|
| 97 | thread.join(timeout=1.0)
|
| 98 | self._listener = None
|
| 99 |
|
| 100 | def snapshot(self) -> ByteProxyStats:
|
| 101 | with self._lock:
|
| 102 | return ByteProxyStats(
|
| 103 | connections=self._connections,
|
| 104 | client_to_upstream_bytes=self._client_to_upstream_bytes,
|
| 105 | upstream_to_client_bytes=self._upstream_to_client_bytes,
|
| 106 | connection_errors=self._connection_errors,
|
| 107 | first_upstream_to_client_ms=self._first_upstream_to_client_ms,
|
| 108 | first_payload_bytes_server_to_client=self._first_payload_bytes_server_to_client,
|
| 109 | )
|
| 110 |
|
| 111 | def reset_counts(self) -> None:
|
| 112 | with self._lock:
|
| 113 | self._connections = 0
|
| 114 | self._client_to_upstream_bytes = 0
|
| 115 | self._upstream_to_client_bytes = 0
|
| 116 | self._connection_errors = 0
|
| 117 | self._count_start = time.monotonic()
|
| 118 | self._first_upstream_to_client_ms = None
|
| 119 | self._first_payload_bytes_server_to_client = None
|
| 120 |
|
| 121 | def __enter__(self) -> "CountingProxy":
|
| 122 | return self.start()
|
| 123 |
|
| 124 | def __exit__(self, exc_type, exc, tb) -> None: # type: ignore[no-untyped-def]
|
| 125 | self.stop()
|
| 126 |
|
| 127 | def _accept_loop(self) -> None:
|
| 128 | assert self._listener is not None
|
| 129 | while not self._stop.is_set():
|
| 130 | try:
|
| 131 | client, _addr = self._listener.accept()
|
| 132 | except socket.timeout:
|
| 133 | continue
|
| 134 | except OSError:
|
| 135 | break
|
| 136 | thread = threading.Thread(target=self._handle_client, args=(client,), daemon=True)
|
| 137 | with self._lock:
|
| 138 | self._connection_threads.append(thread)
|
| 139 | self._active_sockets.add(client)
|
| 140 | thread.start()
|
| 141 |
|
| 142 | def _handle_client(self, client: socket.socket) -> None:
|
| 143 | upstream: Optional[socket.socket] = None
|
| 144 | try:
|
| 145 | upstream = socket.create_connection(
|
| 146 | (self.upstream_host, self.upstream_port),
|
| 147 | timeout=self.connect_timeout_s,
|
| 148 | )
|
| 149 | with self._lock:
|
| 150 | self._connections += 1
|
| 151 | self._active_sockets.add(upstream)
|
| 152 | t1 = threading.Thread(target=self._pump, args=(client, upstream, "client_to_upstream"), daemon=True)
|
| 153 | t2 = threading.Thread(target=self._pump, args=(upstream, client, "upstream_to_client"), daemon=True)
|
| 154 | t1.start()
|
| 155 | t2.start()
|
| 156 | t1.join()
|
| 157 | t2.join()
|
| 158 | except OSError:
|
| 159 | with self._lock:
|
| 160 | self._connection_errors += 1
|
| 161 | finally:
|
| 162 | _close_socket(client)
|
| 163 | if upstream is not None:
|
| 164 | _close_socket(upstream)
|
| 165 | with self._lock:
|
| 166 | self._active_sockets.discard(client)
|
| 167 | if upstream is not None:
|
| 168 | self._active_sockets.discard(upstream)
|
| 169 |
|
| 170 | def _pump(self, source: socket.socket, dest: socket.socket, direction: str) -> None:
|
| 171 | try:
|
| 172 | while not self._stop.is_set():
|
| 173 | chunk = source.recv(64 * 1024)
|
| 174 | if not chunk:
|
| 175 | break
|
| 176 | with self._lock:
|
| 177 | if direction == "client_to_upstream":
|
| 178 | self._client_to_upstream_bytes += len(chunk)
|
| 179 | else:
|
| 180 | if self._first_upstream_to_client_ms is None:
|
| 181 | self._first_upstream_to_client_ms = (time.monotonic() - self._count_start) * 1000
|
| 182 | self._first_payload_bytes_server_to_client = len(chunk)
|
| 183 | self._upstream_to_client_bytes += len(chunk)
|
| 184 | dest.sendall(chunk)
|
| 185 | except OSError:
|
| 186 | pass
|
| 187 | finally:
|
| 188 | try:
|
| 189 | dest.shutdown(socket.SHUT_WR)
|
| 190 | except OSError:
|
| 191 | pass
|
| 192 |
|
| 193 |
|
| 194 | def _close_socket(sock: socket.socket) -> None:
|
| 195 | try:
|
| 196 | sock.shutdown(socket.SHUT_RDWR)
|
| 197 | except OSError:
|
| 198 | pass
|
| 199 | try:
|
| 200 | sock.close()
|
| 201 | except OSError:
|
| 202 | pass
|