vibeproxy/internal/streams/udp.go

package streams

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

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

// UDPProxy is a connectionless proxy that tracks per-client "sessions" (a
// backend UDP socket per source address) so replies are routed back to the
// correct client, with idle sessions reaped after the configured timeout.
type UDPProxy struct {
	model models.Stream
	log   *slog.Logger
	conn  *net.UDPConn

	mu       sync.Mutex
	sessions map[string]*udpSession
}

type udpSession struct {
	clientAddr *net.UDPAddr
	backend    *net.UDPConn
	lastActive time.Time
}

func NewUDPProxy(model models.Stream, log *slog.Logger) *UDPProxy {
	return &UDPProxy{model: model, log: log.With("stream", udpLabel(model)), sessions: map[string]*udpSession{}}
}

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

func (p *UDPProxy) Start(ctx context.Context) error {
	addr := &net.UDPAddr{IP: net.ParseIP(p.model.ListenIP), Port: p.model.ListenPort}
	if addr.IP == nil {
		addr.IP = net.IPv4zero
	}
	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		return fmt.Errorf("listen udp %s:%d: %w", p.model.ListenIP, p.model.ListenPort, err)
	}
	p.conn = conn
	p.log.Info("udp stream listening", "addr", addr.String(), "destination", fmt.Sprintf("%s:%d", p.model.DestinationHost, p.model.DestinationPort))

	go func() {
		<-ctx.Done()
		conn.Close()
	}()
	go p.reaper(ctx)
	go p.readLoop(ctx)
	return nil
}

func (p *UDPProxy) idleTimeout() time.Duration {
	d := time.Duration(p.model.IdleTimeoutMS) * time.Millisecond
	if d <= 0 {
		return 2 * time.Minute
	}
	return d
}

func (p *UDPProxy) readLoop(ctx context.Context) {
	buf := make([]byte, 65535)
	dest := fmt.Sprintf("%s:%d", p.model.DestinationHost, p.model.DestinationPort)
	destAddr, err := net.ResolveUDPAddr("udp", dest)
	if err != nil {
		p.log.Error("failed to resolve destination", "dest", dest, "err", err)
		return
	}
	for {
		n, clientAddr, err := p.conn.ReadFromUDP(buf)
		if err != nil {
			select {
			case <-ctx.Done():
				return
			default:
				return
			}
		}
		metrics.UDPPacketsTotal.WithLabelValues(udpLabel(p.model), "in").Inc()
		metrics.BytesTransferred.WithLabelValues("in", "udp").Add(float64(n))

		sess := p.getOrCreateSession(ctx, clientAddr, destAddr)
		if sess == nil {
			continue
		}
		if _, err := sess.backend.Write(buf[:n]); err != nil {
			p.log.Warn("write to backend failed", "err", err)
		}
	}
}

func (p *UDPProxy) getOrCreateSession(ctx context.Context, clientAddr, destAddr *net.UDPAddr) *udpSession {
	key := clientAddr.String()
	p.mu.Lock()
	sess, ok := p.sessions[key]
	if ok {
		sess.lastActive = time.Now()
		p.mu.Unlock()
		return sess
	}
	p.mu.Unlock()

	backendConn, err := net.DialUDP("udp", nil, destAddr)
	if err != nil {
		p.log.Warn("failed to dial backend", "err", err)
		return nil
	}
	sess = &udpSession{clientAddr: clientAddr, backend: backendConn, lastActive: time.Now()}

	p.mu.Lock()
	p.sessions[key] = sess
	p.mu.Unlock()

	metrics.ActiveConnections.WithLabelValues("udp").Inc()

	go p.backendReadLoop(ctx, key, sess)
	return sess
}

func (p *UDPProxy) backendReadLoop(ctx context.Context, key string, sess *udpSession) {
	buf := make([]byte, 65535)
	for {
		sess.backend.SetReadDeadline(time.Now().Add(p.idleTimeout()))
		n, err := sess.backend.Read(buf)
		if err != nil {
			p.closeSession(key)
			return
		}
		sess.lastActive = time.Now()
		metrics.UDPPacketsTotal.WithLabelValues(udpLabel(p.model), "out").Inc()
		metrics.BytesTransferred.WithLabelValues("out", "udp").Add(float64(n))
		if _, err := p.conn.WriteToUDP(buf[:n], sess.clientAddr); err != nil {
			return
		}
	}
}

func (p *UDPProxy) closeSession(key string) {
	p.mu.Lock()
	sess, ok := p.sessions[key]
	if ok {
		delete(p.sessions, key)
	}
	p.mu.Unlock()
	if ok {
		sess.backend.Close()
		metrics.ActiveConnections.WithLabelValues("udp").Dec()
	}
}

func (p *UDPProxy) reaper(ctx context.Context) {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			cutoff := time.Now().Add(-p.idleTimeout())
			p.mu.Lock()
			for key, sess := range p.sessions {
				if sess.lastActive.Before(cutoff) {
					sess.backend.Close()
					delete(p.sessions, key)
					metrics.ActiveConnections.WithLabelValues("udp").Dec()
				}
			}
			p.mu.Unlock()
		}
	}
}

func (p *UDPProxy) Stop() {
	if p.conn != nil {
		p.conn.Close()
	}
	p.mu.Lock()
	for key, sess := range p.sessions {
		sess.backend.Close()
		delete(p.sessions, key)
	}
	p.mu.Unlock()
}