fastopen.py

fastopen.py

#!/usr/bin/env python3

"""
fastopen.py
Copyright (C) 2023 Storj Labs, Inc.

This utility allows you to create test TCP_FASTOPEN connections.

You can run it as a TCP_FASTOPEN server, where it will print to stdout
everything it receives on the first connection.

  ./fastopen.py server :5996

You can run it as a TCP_FASTOPEN client, where it will send the string
"complete" over its socket and close it.

  ./fastopen client 127.0.0.1:5996

Or you can have it run both locally for a quick test:

  ./fastopen both

To enable TCP_FASTOPEN server-side support, you may need to run:

  sysctl -w net.ipv4.tcp_fastopen=3

Make sure to watch kernel TCP statistics:

  netstat -s | grep TCPFastOpen

"""

import sys, socket, threading

if sys.platform in ("cygwin", "win32"):
  TCP_FASTOPEN = 15
else:
  TCP_FASTOPEN = 0x17
TCP_FASTOPEN_QUEUE = 32
MSG_FASTOPEN = 0x20000000


def server(addr, barrier=None):
    host, port = addr.rsplit(":", 1)
    if host.strip() == "":
        host = "0.0.0.0"
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    if sys.platform in ("cygwin", "win32"):
      sock.setsockopt(socket.SOL_TCP, TCP_FASTOPEN, 1)
    else:
      sock.setsockopt(socket.SOL_TCP, TCP_FASTOPEN, TCP_FASTOPEN_QUEUE)
    sock.bind((host, int(port)))
    sock.listen(TCP_FASTOPEN_QUEUE)
    if barrier is not None:
        barrier.release()
    conn, addr = sock.accept()
    while True:
        data = conn.recv(1024)
        if not data:
            conn.close()
            break
        print("%r" % data)
    sock.close()


def client(addr):
    host, port = addr.rsplit(":", 1)
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.setsockopt(socket.SOL_TCP, TCP_FASTOPEN, 1)
    sock.sendto(b"complete", MSG_FASTOPEN, (host, int(port)))
    sock.close()


def both(addr):
    barrier = threading.Lock()
    barrier.acquire()
    t1 = threading.Thread(target=server, args=(addr, barrier))
    t1.start()
    barrier.acquire()
    barrier.release()
    t2 = threading.Thread(target=client, args=(addr,))
    t2.start()
    t2.join()
    t1.join()


def main():
    dispatch = {"client": client, "server": server, "both": both}
    command = "both"
    addr = "127.0.0.1:5996"
    if len(sys.argv) > 1:
        command = sys.argv[1]
        if len(sys.argv) not in (2, 3) or sys.argv[1].lower() not in dispatch:
            print("usage:")
            print("  %s server [:port]" % sys.argv[0])
            print("  %s client [host:port]" % sys.argv[0])
            print("  %s [both]" % sys.argv[0])
            sys.exit(1)
        if len(sys.argv) == 3:
            addr = sys.argv[2]
    dispatch[command](addr)


if __name__ == "__main__":
    main()