latency_proxy.py raw

   1  #!/usr/bin/env python3
   2  """TCP latency proxy — adds configurable delay + jitter to all traffic.
   3  
   4  Sits between clients and a TCP server (e.g. a WebSocket relay), forwarding
   5  bytes in both directions with artificial delay. WebSocket-transparent: it
   6  operates on raw TCP, so HTTP upgrade and WebSocket frames pass through as-is.
   7  
   8  Usage:
   9      python3 test/latency_proxy.py --listen 127.0.0.1:3335 --target 127.0.0.1:3334 --latency 100 --jitter 30
  10  
  11      This proxies port 3335 → 3334 with 100ms base delay + 0-30ms random jitter
  12      on each forwarded chunk (both directions).
  13  """
  14  
  15  import argparse
  16  import random
  17  import select
  18  import socket
  19  import sys
  20  import threading
  21  import time
  22  
  23  
  24  def forward(src, dst, label, latency_s, jitter_s, stats):
  25      """Forward data from src to dst with delay."""
  26      try:
  27          while True:
  28              data = src.recv(65536)
  29              if not data:
  30                  break
  31              delay = latency_s + random.uniform(0, jitter_s)
  32              if delay > 0:
  33                  time.sleep(delay)
  34              dst.sendall(data)
  35              stats[label] += len(data)
  36      except (ConnectionResetError, BrokenPipeError, OSError):
  37          pass
  38      finally:
  39          try:
  40              dst.shutdown(socket.SHUT_WR)
  41          except OSError:
  42              pass
  43  
  44  
  45  def handle_conn(client_sock, target_addr, latency_s, jitter_s):
  46      """Handle one proxied connection."""
  47      try:
  48          server_sock = socket.create_connection(target_addr, timeout=5)
  49          server_sock.settimeout(None)  # back to blocking after connect
  50      except Exception as e:
  51          print(f"  proxy: connect to {target_addr} failed: {e}", file=sys.stderr)
  52          client_sock.close()
  53          return
  54  
  55      stats = {"c→s": 0, "s→c": 0}
  56      t1 = threading.Thread(target=forward, args=(client_sock, server_sock, "c→s", latency_s, jitter_s, stats), daemon=True)
  57      t2 = threading.Thread(target=forward, args=(server_sock, client_sock, "s→c", latency_s, jitter_s, stats), daemon=True)
  58      t1.start()
  59      t2.start()
  60      t1.join()
  61      t2.join()
  62      client_sock.close()
  63      server_sock.close()
  64  
  65  
  66  def run_proxy(listen_addr, target_addr, latency_ms, jitter_ms):
  67      """Run the latency proxy."""
  68      latency_s = latency_ms / 1000.0
  69      jitter_s = jitter_ms / 1000.0
  70  
  71      lhost, lport = listen_addr.rsplit(":", 1)
  72      thost, tport = target_addr.rsplit(":", 1)
  73      target = (thost, int(tport))
  74  
  75      sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  76      sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  77      sock.bind((lhost, int(lport)))
  78      sock.listen(64)
  79      print(f"latency proxy: {listen_addr} → {target_addr} (latency={latency_ms}ms jitter={jitter_ms}ms)")
  80  
  81      try:
  82          while True:
  83              client, addr = sock.accept()
  84              threading.Thread(target=handle_conn, args=(client, target, latency_s, jitter_s), daemon=True).start()
  85      except KeyboardInterrupt:
  86          pass
  87      finally:
  88          sock.close()
  89  
  90  
  91  def main():
  92      p = argparse.ArgumentParser(description="TCP latency proxy")
  93      p.add_argument("--listen", default="127.0.0.1:3335", help="listen address (default: 127.0.0.1:3335)")
  94      p.add_argument("--target", default="127.0.0.1:3334", help="target address (default: 127.0.0.1:3334)")
  95      p.add_argument("--latency", type=int, default=100, help="base latency in ms (default: 100)")
  96      p.add_argument("--jitter", type=int, default=30, help="random jitter in ms (default: 30)")
  97      args = p.parse_args()
  98      run_proxy(args.listen, args.target, args.latency, args.jitter)
  99  
 100  
 101  if __name__ == "__main__":
 102      main()
 103