vibeproxy/internal/streams/tcp.go

// Package streams implements raw TCP and UDP proxying, independent of the
// HTTP reverse proxy, for arbitrary protocols (game servers, VPNs, etc).
package streams

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"net"
	"strconv"
	"sync"
	"time"

	"github.com/owen/vibeproxy/internal/metrics"
	"github.com/owen/vibeproxy/internal/models"
)

// TCPProxy listens for TCP connections and forwards each to the
// configured destination, optionally prefixing the outbound connection
// with a PROXY protocol v1 header so the backend can see the real client
// address.
type TCPProxy struct {
	model    models.Stream
	listener net.Listener
	log      *slog.Logger
	wg       sync.WaitGroup
}

func NewTCPProxy(model models.Stream, log *slog.Logger) *TCPProxy {
	return &TCPProxy{model: model, log: log.With("stream", streamLabel(model))}
}

func streamLabel(m models.Stream) string {
	return fmt.Sprintf("tcp:%s:%d", m.ListenIP, m.ListenPort)
}

func (p *TCPProxy) Start(ctx context.Context) error {
	addr := net.JoinHostPort(p.model.ListenIP, strconv.Itoa(p.model.ListenPort))
	lc := net.ListenConfig{}
	ln, err := lc.Listen(ctx, "tcp", addr)
	if err != nil {
		return fmt.Errorf("listen %s: %w", addr, err)
	}
	p.listener = ln
	p.log.Info("tcp stream listening", "addr", addr, "destination", fmt.Sprintf("%s:%d", p.model.DestinationHost, p.model.DestinationPort))

	go func() {
		<-ctx.Done()
		ln.Close()
	}()

	p.wg.Add(1)
	go p.acceptLoop(ctx)
	return nil
}

func (p *TCPProxy) acceptLoop(ctx context.Context) {
	defer p.wg.Done()
	for {
		conn, err := p.listener.Accept()
		if err != nil {
			select {
			case <-ctx.Done():
				return
			default:
				p.log.Warn("accept error", "err", err)
				return
			}
		}
		p.wg.Add(1)
		go func() {
			defer p.wg.Done()
			p.handleConn(ctx, conn)
		}()
	}
}

func (p *TCPProxy) handleConn(ctx context.Context, client net.Conn) {
	defer client.Close()
	metrics.ActiveConnections.WithLabelValues("tcp").Inc()
	defer metrics.ActiveConnections.WithLabelValues("tcp").Dec()
	metrics.TCPConnectionsTotal.WithLabelValues(streamLabel(p.model)).Inc()

	dest := net.JoinHostPort(p.model.DestinationHost, strconv.Itoa(p.model.DestinationPort))
	connectTimeout := time.Duration(p.model.ConnectTimeoutMS) * time.Millisecond
	if connectTimeout <= 0 {
		connectTimeout = 5 * time.Second
	}
	d := net.Dialer{Timeout: connectTimeout}
	backend, err := d.DialContext(ctx, "tcp", dest)
	if err != nil {
		p.log.Warn("backend dial failed", "dest", dest, "err", err)
		return
	}
	defer backend.Close()

	if p.model.ProxyProtocol {
		if err := writeProxyProtocolV1(backend, client.RemoteAddr(), backend.LocalAddr()); err != nil {
			p.log.Warn("failed to write PROXY protocol header", "err", err)
			return
		}
	}

	idleTimeout := time.Duration(p.model.IdleTimeoutMS) * time.Millisecond
	if idleTimeout <= 0 {
		idleTimeout = 5 * time.Minute
	}

	var wg sync.WaitGroup
	wg.Add(2)
	pipe := func(dst, src net.Conn, direction string) {
		defer wg.Done()
		buf := make([]byte, 32*1024)
		for {
			src.SetReadDeadline(time.Now().Add(idleTimeout))
			n, err := src.Read(buf)
			if n > 0 {
				metrics.BytesTransferred.WithLabelValues(direction, "tcp").Add(float64(n))
				if _, werr := dst.Write(buf[:n]); werr != nil {
					return
				}
			}
			if err != nil {
				if err != io.EOF {
					if ne, ok := dst.(interface{ CloseWrite() error }); ok {
						ne.CloseWrite()
					}
				}
				return
			}
		}
	}
	go pipe(backend, client, "in")
	go pipe(client, backend, "out")
	wg.Wait()
}

func (p *TCPProxy) Stop() {
	if p.listener != nil {
		p.listener.Close()
	}
	p.wg.Wait()
}

// writeProxyProtocolV1 writes a human-readable PROXY protocol v1 header
// (as used by nginx/HAProxy) so the backend can recover the real client
// address even though the connection now originates from vibeproxy.
func writeProxyProtocolV1(w io.Writer, clientAddr, proxyAddr net.Addr) error {
	ct, ok1 := clientAddr.(*net.TCPAddr)
	pt, ok2 := proxyAddr.(*net.TCPAddr)
	if !ok1 || !ok2 {
		_, err := fmt.Fprintf(w, "PROXY UNKNOWN\r\n")
		return err
	}
	family := "TCP4"
	if ct.IP.To4() == nil {
		family = "TCP6"
	}
	_, err := fmt.Fprintf(w, "PROXY %s %s %s %d %d\r\n", family, ct.IP.String(), pt.IP.String(), ct.Port, pt.Port)
	return err
}