Browse Source
- Fix queue_entry_free before queue_dgram_free in all proxy files (UAF bug causing segfault) - Add PROXY/RP prefixed diagnostic logs for transparent proxy flow tracking - Add tcp_proxy integration test: real utun instances, ETCP, remote_proxy, DNAT routing - 7 test scenarios: basic_1mb, concurrent_5, concurrent_20, half_close, idle, conn_refused, stress (10 threads x 20 iterations) - stress_client.py with --verbose per-thread phase timing (connect/send/recv)congestion
Evgeny
4 months ago
13 changed files with 673 additions and 56 deletions
@ -0,0 +1,29 @@
|
||||
[global] |
||||
my_node_id=0xAAAA000000000002 |
||||
my_private_key=4813d31d28b7e9829247f488c6be7672f2bdf61b2508333128e386d1759afed2 |
||||
my_public_key=c594f33c91f3a2222795c2c110c527bf214ad1009197ce14556cb13df3c461b3c373bed8f205a8dd1fc0c364f90bf471d7c6f5db49564c33e4235d268569ac71 |
||||
tun_ip=10.200.20.1/24 |
||||
tun_ifname=tun_test_cli |
||||
debug_level=error |
||||
|
||||
[server: s1] |
||||
addr=127.0.0.1:15002 |
||||
type=public |
||||
|
||||
[client: c1] |
||||
keepalive=1 |
||||
link=s1:127.0.0.1:15001 |
||||
peer_public_key=1c55e4ccae7c4470707759086738b10681bf88b81f198cc2ab54a647d1556e17c65e6b1833e0c771e5a39382c03067c388915a4c732191bc130480f20f8e00b9 |
||||
|
||||
[tcp_proxy] |
||||
enabled=yes |
||||
tun_name=tun_test_proxy |
||||
tun_ip=10.200.30.1 |
||||
via_node=0xAAAA000000000001 |
||||
|
||||
[debug] |
||||
connection=info |
||||
socket=info |
||||
general=info |
||||
traffic=info |
||||
tun=info |
||||
@ -0,0 +1,76 @@
|
||||
#!/usr/bin/env python3 |
||||
"""TCP echo server for tcp_proxy integration test. |
||||
|
||||
Accepts connections, reads all data, sends back. |
||||
Supports --read-timeout to close after idle (for half-close tests). |
||||
""" |
||||
|
||||
import socket |
||||
import sys |
||||
import threading |
||||
import signal |
||||
|
||||
running = True |
||||
|
||||
|
||||
def handle(conn: socket.socket, read_timeout: float): |
||||
try: |
||||
while True: |
||||
if read_timeout > 0: |
||||
conn.settimeout(read_timeout) |
||||
try: |
||||
data = conn.recv(65536) |
||||
except socket.timeout: |
||||
break |
||||
if not data: |
||||
break |
||||
conn.sendall(data) |
||||
except OSError: |
||||
pass |
||||
finally: |
||||
conn.close() |
||||
|
||||
|
||||
def main(): |
||||
if len(sys.argv) < 3: |
||||
print(f"Usage: {sys.argv[0]} <host> <port> [--read-timeout SEC]", file=sys.stderr) |
||||
sys.exit(1) |
||||
|
||||
host = sys.argv[1] |
||||
port = int(sys.argv[2]) |
||||
read_timeout = 0.0 |
||||
for i in range(3, len(sys.argv)): |
||||
if sys.argv[i] == "--read-timeout" and i + 1 < len(sys.argv): |
||||
read_timeout = float(sys.argv[i + 1]) |
||||
|
||||
global running |
||||
|
||||
def _sigterm(sig, frame): |
||||
global running |
||||
running = False |
||||
|
||||
signal.signal(signal.SIGTERM, _sigterm) |
||||
signal.signal(signal.SIGINT, _sigterm) |
||||
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
||||
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) |
||||
s.bind((host, port)) |
||||
s.listen(128) |
||||
print(f"ECHO: listening on {host}:{port}" + (f" timeout={read_timeout}s" if read_timeout > 0 else ""), flush=True) |
||||
|
||||
while running: |
||||
try: |
||||
s.settimeout(1.0) |
||||
conn, addr = s.accept() |
||||
except socket.timeout: |
||||
continue |
||||
except OSError: |
||||
break |
||||
t = threading.Thread(target=handle, args=(conn, read_timeout), daemon=True) |
||||
t.start() |
||||
|
||||
s.close() |
||||
|
||||
|
||||
if __name__ == "__main__": |
||||
main() |
||||
@ -0,0 +1,24 @@
|
||||
[global] |
||||
my_node_id=0xAAAA000000000001 |
||||
my_private_key=67b705a92b41bcaae105af2d6a17743faa7b26ccebba8b3b9b0af05e9cd1d5fb |
||||
my_public_key=1c55e4ccae7c4470707759086738b10681bf88b81f198cc2ab54a647d1556e17c65e6b1833e0c771e5a39382c03067c388915a4c732191bc130480f20f8e00b9 |
||||
tun_ip=10.200.10.1/24 |
||||
tun_ifname=tun_test_exit |
||||
debug_level=error |
||||
|
||||
[server: s1] |
||||
addr=127.0.0.1:15001 |
||||
type=public |
||||
|
||||
[allowed_keys] |
||||
allow_all=1 |
||||
|
||||
[remote_proxy] |
||||
enabled=yes |
||||
|
||||
[debug] |
||||
connection=info |
||||
socket=info |
||||
general=info |
||||
traffic=info |
||||
tun=info |
||||
@ -0,0 +1,263 @@
|
||||
#!/bin/bash |
||||
# run_test.sh — интеграционный тест tcp_proxy (transparent proxy) |
||||
# |
||||
# Трафик: |
||||
# tcp_client.py (SO_MARK=1, dest=10.200.100.N) → table 100 → tun_test_proxy |
||||
# → tcp_proxy(lwIP) → ETCP → exit(remote_proxy) |
||||
# → connect(10.200.100.N) → iptables DNAT (!mark=1) → 127.0.0.1 → echo → обратно |
||||
# |
||||
# Запуск: sudo ./run_test.sh |
||||
|
||||
set -euo pipefail |
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" |
||||
UTUN_BIN="$SCRIPT_DIR/../../utun" |
||||
LOG_DIR="$SCRIPT_DIR/log" |
||||
CLIENT_PY="$SCRIPT_DIR/tcp_client.py" |
||||
ECHO_PY="$SCRIPT_DIR/echo_server.py" |
||||
STRESS_PY="$SCRIPT_DIR/stress_client.py" |
||||
|
||||
ECHO_PORT=19090 |
||||
HC_PORT=19091 |
||||
REFUSED_PORT=19099 |
||||
N_IP=20 |
||||
|
||||
PASS=0; FAIL=0 |
||||
|
||||
# -------- setup / cleanup -------- |
||||
|
||||
setup_net() { |
||||
echo "=== Setting up networking ===" |
||||
sysctl -w net.ipv4.conf.all.rp_filter=0 |
||||
|
||||
# policy routing: marked packets → TUN (priority 199) |
||||
ip rule add fwmark 1 priority 199 table 100 2>/dev/null || true |
||||
local gw |
||||
gw=$(ip route show default | awk '/via/ {print $3; exit}') |
||||
[ -n "$gw" ] && ip route replace default via "$gw" table 100 2>/dev/null || true |
||||
|
||||
# DNAT: unmarked packets to fake IPs → real localhost |
||||
for i in $(seq 1 $N_IP); do |
||||
iptables -t nat -C OUTPUT -d "10.200.100.${i}" -p tcp -m mark ! --mark 1 -j DNAT --to-destination 127.0.0.1:"$ECHO_PORT" 2>/dev/null \ |
||||
|| iptables -t nat -A OUTPUT -d "10.200.100.${i}" -p tcp -m mark ! --mark 1 -j DNAT --to-destination 127.0.0.1:"$ECHO_PORT" |
||||
done |
||||
echo " Done" |
||||
} |
||||
|
||||
setup_tun_route() { |
||||
for i in $(seq 1 60); do |
||||
if [ -e "/proc/sys/net/ipv4/conf/tun_test_proxy/rp_filter" ]; then |
||||
sysctl -w net.ipv4.conf.tun_test_proxy.rp_filter=0 |
||||
ip route replace 10.200.100.0/24 dev tun_test_proxy table 100 |
||||
ip route replace 10.200.100.0/24 dev tun_test_proxy # main table anti-martian |
||||
echo " tun route ready (${i}x0.3s)" |
||||
return 0 |
||||
fi |
||||
sleep 0.3 |
||||
done |
||||
echo " WARN: tun_test_proxy not found after 18s" |
||||
} |
||||
|
||||
cleanup_net() { |
||||
echo "=== Cleaning up networking ===" |
||||
for i in $(seq 1 $N_IP); do |
||||
iptables -t nat -D OUTPUT -d "10.200.100.${i}" -p tcp -m mark ! --mark 1 -j DNAT --to-destination 127.0.0.1:"$ECHO_PORT" 2>/dev/null || true |
||||
done |
||||
ip rule del fwmark 1 priority 199 table 100 2>/dev/null || true |
||||
ip route del 10.200.100.0/24 dev tun_test_proxy table 100 2>/dev/null || true |
||||
ip route del 10.200.100.0/24 dev tun_test_proxy 2>/dev/null || true |
||||
ip route flush table 100 2>/dev/null || true |
||||
sysctl -w net.ipv4.conf.all.rp_filter=1 |
||||
sysctl -w net.ipv4.conf.tun_test_proxy.rp_filter=1 2>/dev/null || true |
||||
echo " Done" |
||||
} |
||||
|
||||
cleanup() { |
||||
echo ""; echo "=== Cleanup ===" |
||||
kill $EXIT_PID 2>/dev/null || true |
||||
kill $CLIENT_PID 2>/dev/null || true |
||||
kill $ECHO_PID 2>/dev/null || true |
||||
wait $EXIT_PID 2>/dev/null || true |
||||
wait $CLIENT_PID 2>/dev/null || true |
||||
wait $ECHO_PID 2>/dev/null || true |
||||
sleep 0.5; cleanup_net |
||||
} |
||||
|
||||
wait_for_etcp() { |
||||
for i in $(seq 1 30); do |
||||
grep -q "Connection established" "$LOG_DIR/exit_utun.log" 2>/dev/null && return 0 |
||||
sleep 0.5 |
||||
done |
||||
return 1 |
||||
} |
||||
|
||||
# -------- test runners -------- |
||||
|
||||
run_test() { |
||||
local name=$1; shift |
||||
local t_start log rc elapsed |
||||
t_start=$(date +%s%3N) |
||||
echo -n " $name ... " |
||||
log="$LOG_DIR/tcp_client_${name}.log" |
||||
if python3 "$CLIENT_PY" "$@" --name "$name" >"$log" 2>&1; then |
||||
rc=0 |
||||
else |
||||
rc=$? |
||||
fi |
||||
elapsed=$(($(date +%s%3N) - t_start)) |
||||
if [ $rc -eq 0 ]; then |
||||
echo "PASS (${elapsed}ms)" |
||||
tail -1 "$log" |
||||
return 0 |
||||
else |
||||
echo "FAIL (rc=$rc, ${elapsed}ms)" |
||||
tail -3 "$log" |
||||
return 1 |
||||
fi |
||||
} |
||||
|
||||
run_test_neg() { |
||||
local name=$1; shift |
||||
local t_start log rc elapsed |
||||
t_start=$(date +%s%3N) |
||||
echo -n " $name ... " |
||||
log="$LOG_DIR/tcp_client_${name}.log" |
||||
if ! python3 "$CLIENT_PY" "$@" --name "$name" >"$log" 2>&1; then |
||||
rc=0 |
||||
else |
||||
rc=1 |
||||
fi |
||||
elapsed=$(($(date +%s%3N) - t_start)) |
||||
if [ $rc -eq 0 ]; then |
||||
echo "PASS (expected error, ${elapsed}ms)"; return 0 |
||||
else |
||||
echo "FAIL (expected error, got success, ${elapsed}ms)"; return 1 |
||||
fi |
||||
} |
||||
|
||||
run_concurrent() { |
||||
local name=$1 count=$2 size=$3; shift 3 |
||||
local t_start pids ok elapsed |
||||
t_start=$(date +%s%3N) |
||||
pids=() |
||||
for i in $(seq 1 $count); do |
||||
python3 "$CLIENT_PY" \ |
||||
--host "10.200.100.${i}" --port "$ECHO_PORT" \ |
||||
--size "$size" --verify --timeout 15 \ |
||||
--name "${name}_${i}" \ |
||||
>"$LOG_DIR/tcp_client_${name}_${i}.log" 2>&1 & |
||||
pids+=($!) |
||||
done |
||||
ok=0 |
||||
for pid in "${pids[@]}"; do wait "$pid"; [ $? -eq 0 ] && ((++ok)); done |
||||
elapsed=$(($(date +%s%3N) - t_start)) |
||||
echo -n " $name ($count) ... " |
||||
if [ "$ok" -eq "$count" ]; then |
||||
echo "PASS ($ok/$count, ${elapsed}ms)"; return 0 |
||||
else |
||||
echo "FAIL ($ok/$count, ${elapsed}ms)" |
||||
local fail='' |
||||
for i in $(seq 1 $count); do |
||||
grep -q "\[PASS\]" "$LOG_DIR/tcp_client_${name}_${i}.log" 2>/dev/null || fail+=" ${i}" |
||||
done |
||||
echo " failed:$fail"; return 1 |
||||
fi |
||||
} |
||||
|
||||
run_stress() { |
||||
local t_start elapsed |
||||
t_start=$(date +%s%3N) |
||||
echo -n " stress_sessions ... " |
||||
local log="$LOG_DIR/stress.log" |
||||
if python3 "$STRESS_PY" \ |
||||
--host 10.200.100.1 --port-base "$ECHO_PORT" --ports 20 \ |
||||
--threads 10 --iters 20 --min-size 1024 --max-size 65536 --timeout 60 --verbose \ |
||||
>"$log" 2>&1; then |
||||
elapsed=$(($(date +%s%3N) - t_start)) |
||||
echo "PASS (${elapsed}ms)" |
||||
head -1 "$log"; return 0 |
||||
else |
||||
elapsed=$(($(date +%s%3N) - t_start)) |
||||
echo "FAIL (${elapsed}ms)" |
||||
tail -5 "$log"; return 1 |
||||
fi |
||||
} |
||||
|
||||
# -------- main -------- |
||||
|
||||
if [ "$(id -u)" -ne 0 ]; then echo "ERROR: must be root"; exit 1; fi |
||||
[ -x "$UTUN_BIN" ] || { echo "ERROR: utun not found. Build first."; exit 1; } |
||||
|
||||
echo "=== tcp_proxy full integration test ===" |
||||
|
||||
mkdir -p "$LOG_DIR"; rm -f "$LOG_DIR"/*.log |
||||
setup_net |
||||
|
||||
echo "Starting echo server on 127.0.0.1:$ECHO_PORT ..." |
||||
python3 "$ECHO_PY" 127.0.0.1 "$ECHO_PORT" >"$LOG_DIR/echo.log" 2>&1 & |
||||
ECHO_PID=$!; sleep 0.3 |
||||
|
||||
echo "Starting utun exit ..." |
||||
"$UTUN_BIN" -c "$SCRIPT_DIR/exit.conf" -f -l "$LOG_DIR/exit_utun.log" >"$LOG_DIR/exit_stdout.log" 2>&1 & |
||||
EXIT_PID=$! |
||||
|
||||
echo "Starting utun client ..." |
||||
"$UTUN_BIN" -c "$SCRIPT_DIR/client.conf" -f -l "$LOG_DIR/client_utun.log" >"$LOG_DIR/client_stdout.log" 2>&1 & |
||||
CLIENT_PID=$! |
||||
|
||||
trap cleanup EXIT |
||||
|
||||
setup_tun_route |
||||
|
||||
echo "Waiting for ETCP connection ..." |
||||
wait_for_etcp || { echo "ERROR: ETCP timeout"; tail -20 "$LOG_DIR/client_utun.log"; exit 1; } |
||||
echo "ETCP ready"; sleep 1 |
||||
|
||||
echo ""; echo "=== Running tests ===" |
||||
|
||||
# 1: stress first (catch resource leaks) |
||||
run_stress && ((++PASS)) || ((++FAIL)) |
||||
sleep 2 |
||||
|
||||
# 2: basic_1mb |
||||
run_test basic_1mb --host 10.200.100.1 --port "$ECHO_PORT" --size 1048576 --verify --timeout 10 \ |
||||
&& ((++PASS)) || ((++FAIL)) |
||||
sleep 1 |
||||
|
||||
# 3: concurrent_5 |
||||
run_concurrent concurrent_5 5 204800 && ((++PASS)) || ((++FAIL)) |
||||
sleep 1 |
||||
|
||||
# 4: concurrent_20 |
||||
run_concurrent concurrent_20 20 51200 && ((++PASS)) || ((++FAIL)) |
||||
sleep 1 |
||||
|
||||
# 5: half_close |
||||
echo " half_close: one-shot echo on $HC_PORT ..." |
||||
python3 "$ECHO_PY" 127.0.0.1 "$HC_PORT" --read-timeout 2 >"$LOG_DIR/echo_hc.log" 2>&1 & |
||||
HC_PID=$!; sleep 0.3 |
||||
|
||||
run_test half_close --host 10.200.100.1 --port "$HC_PORT" --size 65536 --half-close --timeout 8 \ |
||||
&& ((++PASS)) || ((++FAIL)) |
||||
|
||||
kill $HC_PID 2>/dev/null || true; wait $HC_PID 2>/dev/null || true |
||||
sleep 1 |
||||
|
||||
# 6: idle |
||||
run_test idle --host 10.200.100.1 --port "$ECHO_PORT" --size 32768 --verify --sleep-before 1 --timeout 5 \ |
||||
&& ((++PASS)) || ((++FAIL)) |
||||
sleep 1 |
||||
|
||||
# 7: conn_refused |
||||
run_test_neg conn_refused --host 10.200.100.99 --port "$REFUSED_PORT" --size 1 --verify --timeout 5 \ |
||||
&& ((++PASS)) || ((++FAIL)) |
||||
sleep 1 |
||||
|
||||
echo ""; echo "==========================================" |
||||
echo "Results: $PASS passed, $FAIL failed ($((PASS + FAIL)) total)" |
||||
echo "Logs: $LOG_DIR/" |
||||
echo "==========================================" |
||||
|
||||
trap - EXIT; cleanup |
||||
[ "$FAIL" -gt 0 ] && exit 1 |
||||
exit 0 |
||||
@ -0,0 +1,122 @@
|
||||
#!/usr/bin/env python3 |
||||
"""Stress test: parallel threads, sequential echo requests, data verification.""" |
||||
|
||||
import os |
||||
import socket |
||||
import sys |
||||
import threading |
||||
import random |
||||
import time |
||||
|
||||
SO_MARK = 36 |
||||
|
||||
|
||||
def worker(thread_id, host, ports, iters, min_size, max_size, timeout, verbose, results_raw): |
||||
ok = 0 |
||||
times = [] |
||||
for i in range(iters): |
||||
size = random.randint(min_size, max_size) |
||||
port = ports[i % len(ports)] |
||||
payload = os.urandom(size) |
||||
|
||||
t0 = time.monotonic() |
||||
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
||||
s.setsockopt(socket.SOL_SOCKET, SO_MARK, 1) |
||||
s.settimeout(timeout) |
||||
try: |
||||
t1 = time.monotonic() |
||||
s.connect((host, port)) |
||||
t2 = time.monotonic() |
||||
s.sendall(payload) |
||||
t3 = time.monotonic() |
||||
|
||||
received = bytearray() |
||||
while len(received) < size: |
||||
chunk = s.recv(min(65536, size - len(received))) |
||||
if not chunk: |
||||
break |
||||
received += chunk |
||||
t4 = time.monotonic() |
||||
|
||||
if len(received) != size or received != payload: |
||||
raise RuntimeError(f"data mismatch") |
||||
ok += 1 |
||||
times.append((size, int((t2-t1)*1000), int((t3-t2)*1000), int((t4-t3)*1000))) |
||||
except Exception as e: |
||||
results_raw[thread_id] = (False, f"FAIL iter={i} size={size}: {e} " |
||||
f"connect={int((t2-t1)*1000)}ms send={int((t3-t2)*1000)}ms recv={int((t4-t3)*1000)}ms") |
||||
try: |
||||
s.close() |
||||
except Exception: |
||||
pass |
||||
return |
||||
finally: |
||||
if s: |
||||
try: |
||||
s.close() |
||||
except Exception: |
||||
pass |
||||
|
||||
agg_connect = sum(t[1] for t in times) |
||||
agg_send = sum(t[2] for t in times) |
||||
agg_recv = sum(t[3] for t in times) |
||||
total_bytes = sum(t[0] for t in times) |
||||
results_raw[thread_id] = (True, f"PASS {ok}/{iters} c={agg_connect}ms s={agg_send}ms r={agg_recv}ms total={total_bytes}B") |
||||
|
||||
if verbose: |
||||
for i, (sz, cm, sm, rm) in enumerate(times): |
||||
print(f" TH{thread_id:02d} iter={i:02d} size={sz:6d} c={cm:4d}ms s={sm:4d}ms r={rm:4d}ms", flush=True) |
||||
|
||||
|
||||
def main(): |
||||
import argparse |
||||
|
||||
parser = argparse.ArgumentParser() |
||||
parser.add_argument("--host", required=True) |
||||
parser.add_argument("--port-base", type=int, required=True) |
||||
parser.add_argument("--ports", type=int, default=20) |
||||
parser.add_argument("--threads", type=int, default=10) |
||||
parser.add_argument("--iters", type=int, default=20) |
||||
parser.add_argument("--min-size", type=int, default=1024) |
||||
parser.add_argument("--max-size", type=int, default=1048576) |
||||
parser.add_argument("--timeout", type=float, default=60) |
||||
parser.add_argument("--verbose", action="store_true") |
||||
args = parser.parse_args() |
||||
|
||||
ports = [args.port_base + i for i in range(args.ports)] |
||||
random.seed(os.urandom(8)) |
||||
|
||||
threads = [] |
||||
results_raw = [None] * args.threads |
||||
|
||||
t0 = time.monotonic() |
||||
|
||||
for t in range(args.threads): |
||||
th = threading.Thread( |
||||
target=worker, |
||||
args=(t, args.host, ports, args.iters, args.min_size, args.max_size, args.timeout, args.verbose, results_raw), |
||||
daemon=True, |
||||
) |
||||
th.start() |
||||
threads.append(th) |
||||
|
||||
for th in threads: |
||||
th.join() |
||||
|
||||
elapsed = time.monotonic() - t0 |
||||
|
||||
passed = sum(1 for r in results_raw if r and r[0]) |
||||
total_iters = args.threads * args.iters |
||||
|
||||
print(f"[{'PASS' if passed == args.threads else 'FAIL'}] stress: " |
||||
f"{passed}/{args.threads} threads ({total_iters} sessions) in {elapsed:.1f}s", flush=True) |
||||
|
||||
for i, r in enumerate(results_raw): |
||||
if r: |
||||
print(f" thread {i}: {r[1]}", flush=True) |
||||
if passed < args.threads: |
||||
sys.exit(1) |
||||
|
||||
|
||||
if __name__ == "__main__": |
||||
main() |
||||
@ -0,0 +1,97 @@
|
||||
#!/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() |
||||
Loading…
Reference in new issue