v0.32.0

scripts/test_stream_tcp.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+"""TCP Stream black-box test: use against a running RustProxy and mock backends.
+
+The 'selftest' mode validates this harness only; it does NOT run RustProxy.
+All traffic uses loopback by default. Do not point it at a production service.
+"""
+from __future__ import annotations
+
+import argparse
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import secrets
+import socket
+import socketserver
+import sys
+import threading
+import time
+
+MAX_LINE = 1024
+
+
+class TCPServer(socketserver.ThreadingTCPServer):
+    allow_reuse_address = True
+    daemon_threads = True
+    request_queue_size = 128
+
+
+class TCPBackend(socketserver.StreamRequestHandler):
+    def handle(self) -> None:
+        self.request.settimeout(10)
+        backend_id = self.server.backend_id.encode('ascii')
+        try:
+            while True:
+                line = self.rfile.readline(MAX_LINE + 1)
+                if not line or len(line) > MAX_LINE or not line.endswith(b'\n'):
+                    return
+                self.wfile.write(backend_id + b'|' + line)
+                self.wfile.flush()
+        except (OSError, TimeoutError):
+            return
+
+
+def run_server(host: str, port: int, backend_id: str) -> None:
+    with TCPServer((host, port), TCPBackend) as server:
+        server.backend_id = backend_id
+        print(f'TCP mock {backend_id} listening on {server.server_address}', flush=True)
+        try:
+            server.serve_forever(poll_interval=0.1)
+        except KeyboardInterrupt:
+            pass
+
+
+def check_one(index: int, host: str, port: int, messages: int, timeout: float, hold_ms: float, nonce: str) -> str:
+    backend = None
+    with socket.create_connection((host, port), timeout=timeout) as connection:
+        connection.settimeout(timeout)
+        with connection.makefile('rb') as reader:
+            for seq in range(messages):
+                payload = f'rp-tcp:{nonce}:{index}:{seq}\n'.encode('ascii')
+                connection.sendall(payload)
+                response = reader.readline(MAX_LINE + 1)
+                try:
+                    observed, echoed = response.split(b'|', 1)
+                except ValueError as exc:
+                    raise AssertionError(f'session {index}: missing mock backend response: {response!r}') from exc
+                assert echoed == payload, f'session {index}: response mismatch: {response!r}'
+                assert observed, f'session {index}: empty backend ID'
+                if backend is None:
+                    backend = observed
+                assert observed == backend, f'session {index}: TCP backend changed during connection'
+                if hold_ms and seq != messages - 1:
+                    time.sleep(hold_ms / 1000)
+    return backend.decode('ascii')
+
+
+def run_client(host: str, port: int, sessions: int, messages: int, parallel: int, timeout: float,
+               hold_ms: float, min_backends: int) -> Counter:
+    nonce = secrets.token_hex(8)
+    started = time.monotonic()
+    counts: Counter[str] = Counter()
+    with ThreadPoolExecutor(max_workers=parallel) as pool:
+        futures = [pool.submit(check_one, i, host, port, messages, timeout, hold_ms, nonce)
+                   for i in range(sessions)]
+        for future in as_completed(futures):
+            counts[future.result()] += 1
+    assert len(counts) >= min_backends, f'observed {len(counts)} backend(s), required {min_backends}: {dict(counts)}'
+    print(f'TCP PASS: {sessions} sessions, {sessions * messages} messages, '
+          f'{len(counts)} backend(s), {time.monotonic() - started:.2f}s; assignments={dict(sorted(counts.items()))}')
+    return counts
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    sub = parser.add_subparsers(dest='mode', required=True)
+    server = sub.add_parser('server', help='run a local TCP mock backend')
+    server.add_argument('--host', default='127.0.0.1')
+    server.add_argument('--port', type=int, required=True)
+    server.add_argument('--backend-id', required=True, help='unique ASCII marker (e.g. node-a)')
+    client = sub.add_parser('client', help='test forwarding via an already-running RustProxy listener')
+    client.add_argument('--host', default='127.0.0.1', help='RustProxy TCP listener host')
+    client.add_argument('--port', type=int, required=True, help='RustProxy TCP listener port')
+    client.add_argument('--sessions', type=int, default=24)
+    client.add_argument('--messages', type=int, default=3)
+    client.add_argument('--parallel', type=int, default=8)
+    client.add_argument('--timeout', type=float, default=4.0)
+    client.add_argument('--hold-ms', type=float, default=20.0, help='pause between messages to keep TCP sessions alive')
+    client.add_argument('--min-backends', type=int, default=1, help='minimum distinct backend IDs observed')
+    sub.add_parser('selftest', help='test mock peer + harness directly, WITHOUT RustProxy')
+    args = parser.parse_args()
+    if args.mode == 'server':
+        if not args.backend_id.isascii() or not args.backend_id or '|' in args.backend_id or '\n' in args.backend_id:
+            parser.error('--backend-id must be nonempty ASCII without | or newline')
+        run_server(args.host, args.port, args.backend_id)
+    elif args.mode == 'client':
+        if min(args.port, args.sessions, args.messages, args.parallel, args.min_backends) < 1 or args.port > 65535 or args.timeout <= 0 or args.hold_ms < 0:
+            parser.error('invalid port, count, timeout or hold delay')
+        run_client(args.host, args.port, args.sessions, args.messages, args.parallel,
+                   args.timeout, args.hold_ms, args.min_backends)
+    else:
+        with TCPServer(('127.0.0.1', 0), TCPBackend) as server:
+            server.backend_id = 'local-mock'
+            thread = threading.Thread(target=server.serve_forever, daemon=True)
+            thread.start()
+            try:
+                run_client('127.0.0.1', server.server_address[1], 8, 3, 4, 2.0, 5, 1)
+                print('SELFTEST PASS: tested mock TCP peer, not RustProxy')
+            finally:
+                server.shutdown()
+                thread.join(timeout=2)
+
+
+if __name__ == '__main__':
+    try:
+        main()
+    except (OSError, AssertionError, UnicodeError) as exc:
+        print(f'TCP FAIL: {exc}', file=sys.stderr)
+        sys.exit(1)