You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
97 lines
2.9 KiB
97 lines
2.9 KiB
#!/usr/bin/env python3 |
|
"""TCP test client for tcp_proxy integration test. |
|
|
|
Connects to host:port with SO_MARK=1, sends N bytes, receives echo, verifies. |
|
""" |
|
|
|
import os |
|
import socket |
|
import sys |
|
import time |
|
|
|
SO_MARK = 36 |
|
|
|
|
|
def main(): |
|
import argparse |
|
|
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("--host", required=True) |
|
parser.add_argument("--port", type=int, required=True) |
|
parser.add_argument("--size", type=int, required=True) |
|
parser.add_argument("--verify", action="store_true") |
|
parser.add_argument("--half-close", action="store_true") |
|
parser.add_argument("--sleep-before", type=float, default=0) |
|
parser.add_argument("--timeout", type=float, default=30.0) |
|
parser.add_argument("--mark", type=int, default=1, help="SO_MARK (0=off)") |
|
parser.add_argument("--name", default="test") |
|
args = parser.parse_args() |
|
|
|
payload = os.urandom(args.size) |
|
t0 = time.monotonic() |
|
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
|
if args.mark != 0: |
|
s.setsockopt(socket.SOL_SOCKET, SO_MARK, args.mark) |
|
s.settimeout(args.timeout) |
|
|
|
try: |
|
s.connect((args.host, args.port)) |
|
except OSError as e: |
|
print(f"[FAIL] {args.name}: connect failed: {e}", flush=True) |
|
s.close() |
|
sys.exit(1) |
|
|
|
if args.sleep_before > 0: |
|
time.sleep(args.sleep_before) |
|
|
|
try: |
|
s.sendall(payload) |
|
except OSError as e: |
|
print(f"[FAIL] {args.name}: sendall failed: {e}", flush=True) |
|
s.close() |
|
sys.exit(1) |
|
|
|
if args.half_close: |
|
s.shutdown(socket.SHUT_WR) |
|
|
|
if not args.verify and not args.half_close: |
|
elapsed = time.monotonic() - t0 |
|
print(f"[PASS] {args.name}: {args.size} bytes sent in {elapsed:.3f}s " |
|
f"({args.size / 1e6 / elapsed:.2f} MB/s)", flush=True) |
|
s.close() |
|
sys.exit(0) |
|
|
|
received = bytearray() |
|
expected = args.size |
|
try: |
|
while len(received) < expected: |
|
chunk = s.recv(min(65536, expected - len(received))) |
|
if not chunk: |
|
break |
|
received += chunk |
|
except OSError as e: |
|
print(f"[FAIL] {args.name}: recv error after {len(received)} bytes: {e}", flush=True) |
|
s.close() |
|
sys.exit(1) |
|
|
|
s.close() |
|
elapsed = time.monotonic() - t0 |
|
|
|
if len(received) != expected: |
|
print(f"[FAIL] {args.name}: {len(received)}/{expected} bytes in {elapsed:.3f}s", flush=True) |
|
sys.exit(1) |
|
|
|
if payload != received: |
|
for i in range(min(len(payload), len(received))): |
|
if payload[i] != received[i]: |
|
print(f"[FAIL] {args.name}: data mismatch at byte {i}: " |
|
f"sent={payload[i]:02x} recv={received[i]:02x}", flush=True) |
|
sys.exit(1) |
|
|
|
throughput = (expected * 2) / 1e6 / elapsed |
|
print(f"[PASS] {args.name}: {expected} bytes echoed in {elapsed:.3f}s ({throughput:.2f} MB/s)", flush=True) |
|
|
|
|
|
if __name__ == "__main__": |
|
main()
|
|
|