fastopen.py

1 - 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()

2 - fastopen.go

/*
	fastopen.go

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
*/
package main

import (
	"fmt"
	"io"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"syscall"

	"github.com/dsnet/try"
	"golang.org/x/net/context"
	"golang.org/x/sync/errgroup"
)

const (
	TCP_FASTOPEN       = 0x17 // should be 0x0f if platform is win32 or cygwin
	TCP_FASTOPEN_QUEUE = 32
	MSG_FASTOPEN       = 0x20000000
)

func server(addr string, barrier *sync.Mutex) {
	l := try.E1((&net.ListenConfig{
		Control: func(network, address string, c syscall.RawConn) error {
			return c.Control(func(fd uintptr) {
				// on platforms win32 and cygwin, value should be 1.
				try.E(syscall.SetsockoptInt(int(fd), syscall.SOL_TCP, TCP_FASTOPEN, TCP_FASTOPEN_QUEUE))
			})
		},
	}).Listen(context.Background(), "tcp", addr))
	defer l.Close()
	if barrier != nil {
		barrier.Unlock()
	}
	for {
		conn := try.E1(l.Accept())
		data := try.E1(io.ReadAll(conn))
		try.E(conn.Close())
		try.E1(fmt.Printf("%q\n", string(data)))
		if barrier != nil { return }
	}
}

func client(addr string) {
	host, portStr := try.E2(net.SplitHostPort(addr))
	port := try.E1(strconv.ParseInt(portStr, 10, 64))
	ip := net.ParseIP(host).To4()
	if len(ip) != 4{
		panic(fmt.Sprintf("not ipv4"))
	}
	var ipBytes [4]byte
	copy(ipBytes[:], ip[:4])
	fd := try.E1(syscall.Socket(syscall.AF_INET, syscall.SOCK_STREAM, syscall.IPPROTO_TCP))
	try.E(syscall.SetsockoptInt(fd, syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 1))
	try.E(syscall.SetsockoptInt(fd, syscall.SOL_TCP, TCP_FASTOPEN, 1))
	try.E(syscall.Sendto(fd, []byte("complete"), MSG_FASTOPEN, &syscall.SockaddrInet4{Addr: ipBytes, Port: int(port)}))
	try.E(syscall.Close(fd))
}

func both(addr string) {
	barrier := &sync.Mutex{}
	barrier.Lock()
	var eg errgroup.Group
	eg.Go(func() error {
		server(addr, barrier)
		return nil
	})
	barrier.Lock()
	barrier.Unlock()
	eg.Go(func() error {
		client(addr)
		return nil
	})
	try.E(eg.Wait())
}

func main() {
	dispatch := map[string]func(addr string){
		"client": client,
		"both":   both,
		"server": func(addr string) { server(addr, nil) },
	}
	command := "both"
	addr := "127.0.0.1:5996"
	if len(os.Args) > 1 {
		command = os.Args[1]
		if _, found := dispatch[strings.ToLower(os.Args[1])]; (len(os.Args) != 2 && len(os.Args) != 3) || !found {
			try.E1(fmt.Printf("usage:\n"))
			try.E1(fmt.Printf("  %s server [:port]\n", os.Args[0]))
			try.E1(fmt.Printf("  %s client [host:port]\n", os.Args[0]))
			try.E1(fmt.Printf("  %s both\n", os.Args[0]))
			os.Exit(1)
		}
		if len(os.Args) == 3 {
			addr = os.Args[2]
		}
	}
	dispatch[command](addr)
}