136 lines
6.3 KiB
Python
Executable File
136 lines
6.3 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""UDP Stream black-box test: sends multiple packets from each stable client socket.
|
|
|
|
Start separately labeled mock backends and target an existing RustProxy UDP
|
|
listener. 'selftest' only validates this Python harness, NOT RustProxy.
|
|
By default, the server responds PONG to PING for an explicitly configured probe.
|
|
"""
|
|
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_DATAGRAM = 1024
|
|
|
|
|
|
class UDPServer(socketserver.UDPServer):
|
|
allow_reuse_address = True
|
|
|
|
|
|
class UDPBackend(socketserver.BaseRequestHandler):
|
|
def handle(self) -> None:
|
|
data, connection = self.request
|
|
if data == b'PING':
|
|
reply = b'PONG'
|
|
elif data and len(data) <= MAX_DATAGRAM:
|
|
reply = self.server.backend_id.encode('ascii') + b'|' + data
|
|
else:
|
|
return
|
|
connection.sendto(reply, self.client_address)
|
|
|
|
|
|
def run_server(host: str, port: int, backend_id: str) -> None:
|
|
with UDPServer((host, port), UDPBackend) as server:
|
|
server.backend_id = backend_id
|
|
print(f'UDP mock {backend_id} listening on {server.server_address}; health probe PING -> PONG', 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, pause_ms: float, nonce: str) -> str:
|
|
info = socket.getaddrinfo(host, port, type=socket.SOCK_DGRAM)[0]
|
|
family, sock_type, proto, _, address = info
|
|
backend = None
|
|
# One socket = one stable client identity for the entire UDP Stream session.
|
|
with socket.socket(family, sock_type, proto) as connection:
|
|
connection.settimeout(timeout)
|
|
for seq in range(messages):
|
|
payload = f'rp-udp:{nonce}:{index}:{seq}'.encode('ascii')
|
|
connection.sendto(payload, address)
|
|
response, _ = connection.recvfrom(MAX_DATAGRAM + 128)
|
|
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}: UDP backend changed for the same client socket'
|
|
if pause_ms and seq != messages - 1:
|
|
time.sleep(pause_ms / 1000)
|
|
return backend.decode('ascii')
|
|
|
|
|
|
def run_client(host: str, port: int, sessions: int, messages: int, parallel: int, timeout: float,
|
|
pause_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, pause_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'UDP PASS: {sessions} client sockets, {sessions * messages} datagrams, '
|
|
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 local UDP echo and optional PING/PONG probe responder')
|
|
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 UDP listener')
|
|
client.add_argument('--host', default='127.0.0.1', help='RustProxy UDP listener host')
|
|
client.add_argument('--port', type=int, required=True, help='RustProxy UDP 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('--pause-ms', type=float, default=20.0, help='pause between datagrams on each stable client socket')
|
|
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.pause_ms < 0:
|
|
parser.error('invalid port, count, timeout or pause delay')
|
|
run_client(args.host, args.port, args.sessions, args.messages, args.parallel,
|
|
args.timeout, args.pause_ms, args.min_backends)
|
|
else:
|
|
with UDPServer(('127.0.0.1', 0), UDPBackend) 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 UDP peer, not RustProxy')
|
|
finally:
|
|
server.shutdown()
|
|
thread.join(timeout=2)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
try:
|
|
main()
|
|
except (OSError, AssertionError, UnicodeError) as exc:
|
|
print(f'UDP FAIL: {exc}', file=sys.stderr)
|
|
sys.exit(1)
|