138 lines
6.2 KiB
Python
Executable File
138 lines
6.2 KiB
Python
Executable File
#!/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)
|