HTTP Keep-Alive Reconstruction

When writing reverse proxies, understanding if a request is part of a flow can be really useful. As the ServeHTTP has no concept of sessions, we can use the ConnState and ConnContext functions to keep track of Keep-Alive connections and propagate a unique ID to ServeHTTP.

proxy := httputil.NewSingleHostReverseProxy(remote)
proxy.Transport = &Transport{http.DefaultTransport}

mux := http.NewServeMux()
mux.HandleFunc("/", proxy.ServeHTTP)

// Build a map of all currently active sockets, this will be useful to reconstruct Keep-Alive connections
var sockets map[net.Conn]string = make(map[net.Conn]string)

server := http.Server{
    Addr:    addr,
    Handler: mux,
    ConnState: func(c net.Conn, cs http.ConnState) {
        switch cs {
        case http.StateClosed:
            // Cleanup connection when it's closed
            delete(sockets, c)
        }
    },
    ConnContext: func(ctx context.Context, c net.Conn) context.Context {
        // Assign a new UUID to each new connection
        id, ok := sockets[c]
        if !ok {
            id = uuid.New().String()
        }

        ctx = context.WithValue(ctx, "session", id)

        return ctx
    },
}

if err := server.ListenAndServe(); err != nil {
    panic(err)
}