Week 5 Mini Project: TCP Proxy with Traffic Inspector

Architecture

Client → [TCP Proxy :8080] → [Upstream :9090]
                 ↓
          [Traffic Tap: tap.log]
          [Metrics: /debug/vars]

Full Implementation

// proxy/main.go
package main

import (
    "context"
    "expvar"
    "fmt"
    "io"
    "net"
    "net/http"
    _ "net/http/pprof"
    "os"
    "os/signal"
    "sync"
    "sync/atomic"
    "syscall"
    "time"
)

var (
    connAccepted  = expvar.NewInt("connections_accepted")
    connActive    = expvar.NewInt("connections_active")
    connRejected  = expvar.NewInt("connections_rejected")
    bytesInTotal  = expvar.NewInt("bytes_in_total")
    bytesOutTotal = expvar.NewInt("bytes_out_total")
)

const (
    maxConnections = 500
    idleTimeout    = 30 * time.Second
    upstreamAddr   = "127.0.0.1:9090"
    listenAddr     = ":8080"
)

type Proxy struct {
    listener net.Listener
    pool     *Pool
    tap      io.Writer
    active   atomic.Int64
    wg       sync.WaitGroup
    shutdown atomic.Bool
}

func NewProxy(listener net.Listener, upstreamAddr string, tap io.Writer) *Proxy {
    return &Proxy{
        listener: listener,
        pool:     NewPool(20, 30*time.Second),
        tap:      tap,
    }
}

func (p *Proxy) Serve(ctx context.Context) {
    for {
        conn, err := p.listener.Accept()
        if err != nil {
            select {
            case <-ctx.Done():
                return
            default:
                continue
            }
        }

        if p.active.Load() >= maxConnections {
            connRejected.Add(1)
            conn.Close()
            continue
        }

        connAccepted.Add(1)
        connActive.Add(1)
        p.active.Add(1)
        p.wg.Add(1)
        go p.handleConn(conn)
    }
}

func (p *Proxy) handleConn(client net.Conn) {
    defer client.Close()
    defer connActive.Add(-1)
    defer p.active.Add(-1)
    defer p.wg.Done()

    // Get upstream connection from pool
    upstream, err := p.pool.Get(upstreamAddr)
    if err != nil {
        fmt.Fprintf(os.Stderr, "upstream connect: %v\n", err)
        return
    }

    remote := client.RemoteAddr().String()
    var bytesIn, bytesOut int64
    var innerWg sync.WaitGroup
    innerWg.Add(2)

    // client → upstream
    go func() {
        defer innerWg.Done()
        tapped := io.TeeReader(client,
            &hexLogger{w: p.tap, dir: fmt.Sprintf("IN  [%s]", remote)})
        n, _ := io.Copy(upstream, tapped)
        atomic.AddInt64(&bytesIn, n)
        bytesInTotal.Add(n)
        if tc, ok := upstream.(*net.TCPConn); ok {
            tc.CloseWrite()
        }
    }()

    // upstream → client
    go func() {
        defer innerWg.Done()
        tapped := io.TeeReader(upstream,
            &hexLogger{w: p.tap, dir: fmt.Sprintf("OUT [%s]", remote)})
        n, _ := io.Copy(client, tapped)
        atomic.AddInt64(&bytesOut, n)
        bytesOutTotal.Add(n)
        if tc, ok := client.(*net.TCPConn); ok {
            tc.CloseWrite()
        }
    }()

    innerWg.Wait()
    p.pool.Put(upstreamAddr, upstream)
}

func (p *Proxy) Drain() {
    p.wg.Wait()
}

type hexLogger struct {
    w   io.Writer
    dir string
}

func (h *hexLogger) Write(p []byte) (int, error) {
    preview := p
    if len(preview) > 64 {
        preview = preview[:64]
    }
    fmt.Fprintf(h.w, "%s %d bytes: %q\n", h.dir, len(p), preview)
    return len(p), nil
}

func main() {
    // Traffic tap
    tap, _ := os.Create("tap.log")
    defer tap.Close()

    ln, err := net.Listen("tcp", listenAddr)
    if err != nil {
        fmt.Println("listen:", err)
        return
    }

    proxy := NewProxy(ln, upstreamAddr, tap)

    ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT)
    defer stop()

    // Observability server
    go http.ListenAndServe(":9091", nil) // /debug/vars, /debug/pprof

    fmt.Printf("TCP proxy %s → %s\n", listenAddr, upstreamAddr)
    fmt.Println("Metrics: http://localhost:9091/debug/vars")

    go proxy.Serve(ctx)
    <-ctx.Done()

    fmt.Println("shutting down...")
    ln.Close()
    proxy.Drain()
    fmt.Println("shutdown complete")
}

Engineer Takeaway

A TCP proxy is the foundation of load balancers (Week 7), WAFs (Week 10), and security gateways (Week 12). Cloudflare Spectrum, AWS NLB, and HAProxy in TCP mode all start with this: accept → bidirectional io.Copy → return to pool. The traffic tap is how compliance teams capture data for forensics without changing the application.