svc-proxy/internal/bridge/manager.go

// Package bridge owns the UDP data plane: one Valve per backend, each Valve
// owns a public listener socket and a pool of per-client bridges that copy
// datagrams to the backend's voice address and back.
package bridge

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

	"git.timemachine.center/timemachine/svc-proxy/internal/pgsync"
)

// Manager is the top-level coordinator. Implements pgsync.Applier so the
// pgsync goroutine can hand it desired/undesired routes; Manager turns those
// into open/close calls on a Valve registry keyed by public port.
type Manager struct {
	ctx           context.Context
	bindHost      string
	bridgeIdleTTL time.Duration

	mu     sync.Mutex
	valves map[int]*Valve // key: public UDP port
}

func NewManager(ctx context.Context, bindHost string, idleTTL time.Duration) *Manager {
	return &Manager{
		ctx:           ctx,
		bindHost:      bindHost,
		bridgeIdleTTL: idleTTL,
		valves:        map[int]*Valve{},
	}
}

// Apply satisfies pgsync.Applier. Open first (so a backend-address change
// can flip-cleanly while the new listener takes over the new port), then
// close.
func (m *Manager) Apply(add []pgsync.Route, del []pgsync.Route) {
	m.mu.Lock()
	defer m.mu.Unlock()

	for _, r := range add {
		if existing, ok := m.valves[r.Port]; ok {
			// Same port, different backend — close, then re-open.
			existing.Close()
			delete(m.valves, r.Port)
		}
		v, err := openValve(m.ctx, m.bindHost, r, m.bridgeIdleTTL)
		if err != nil {
			slog.Error("valve open failed", "port", r.Port, "addr", r.Address, "name", r.Name, "err", err)
			continue
		}
		m.valves[r.Port] = v
		slog.Info("valve open", "port", r.Port, "addr", r.Address, "name", r.Name)
	}

	for _, r := range del {
		v, ok := m.valves[r.Port]
		if !ok {
			continue
		}
		v.Close()
		delete(m.valves, r.Port)
		slog.Info("valve close", "port", r.Port, "name", r.Name)
	}
}

// Shutdown closes every active valve. Safe to call once; idempotent for
// per-valve Close.
func (m *Manager) Shutdown() {
	m.mu.Lock()
	defer m.mu.Unlock()
	for port, v := range m.valves {
		v.Close()
		delete(m.valves, port)
	}
}

// Valve owns one public UDP listener and the per-client bridges hanging off
// it. Each bridge is a goroutine that copies datagrams from one ephemeral
// upstream socket back to the original client. The public socket itself is
// the egress for backend → client.
type Valve struct {
	route   pgsync.Route
	backend *net.UDPAddr
	pub     *net.UDPConn // 0.0.0.0:<route.Port>

	idleTTL time.Duration

	ctx    context.Context
	cancel context.CancelFunc

	mu      sync.Mutex
	bridges map[string]*clientBridge // key: client.RemoteAddr().String()
}

func openValve(parent context.Context, bindHost string, r pgsync.Route, idleTTL time.Duration) (*Valve, error) {
	backend, err := net.ResolveUDPAddr("udp", r.Address)
	if err != nil {
		return nil, fmt.Errorf("resolve backend %q: %w", r.Address, err)
	}
	pubAddr := &net.UDPAddr{IP: net.ParseIP(bindHost), Port: r.Port}
	if pubAddr.IP == nil {
		return nil, fmt.Errorf("bind host %q not an IP", bindHost)
	}
	pub, err := net.ListenUDP("udp", pubAddr)
	if err != nil {
		return nil, fmt.Errorf("bind %s: %w", pubAddr, err)
	}
	ctx, cancel := context.WithCancel(parent)
	v := &Valve{
		route:   r,
		backend: backend,
		pub:     pub,
		idleTTL: idleTTL,
		ctx:     ctx,
		cancel:  cancel,
		bridges: map[string]*clientBridge{},
	}
	go v.readLoop()
	go v.evictIdle()
	return v, nil
}

// readLoop runs forever copying packets from the public socket to per-client
// upstream sockets. The reverse direction (backend → client) is per-bridge
// goroutines on the upstream sockets writing back to v.pub.
func (v *Valve) readLoop() {
	buf := make([]byte, 2048) // SVC max datagram body
	for {
		n, src, err := v.pub.ReadFromUDP(buf)
		if err != nil {
			if v.ctx.Err() != nil || errors.Is(err, net.ErrClosed) {
				return
			}
			slog.Warn("valve read error", "port", v.route.Port, "err", err)
			continue
		}
		v.mu.Lock()
		b, ok := v.bridges[src.String()]
		if !ok {
			b, err = v.openBridge(src)
			if err != nil {
				v.mu.Unlock()
				slog.Warn("bridge open failed", "port", v.route.Port, "src", src, "err", err)
				continue
			}
			v.bridges[src.String()] = b
			slog.Debug("bridge open", "port", v.route.Port, "client", src.String())
		}
		v.mu.Unlock()
		b.touch()
		if _, err := b.upstream.Write(buf[:n]); err != nil {
			if v.ctx.Err() == nil {
				slog.Warn("bridge forward failed", "port", v.route.Port, "err", err)
			}
			continue
		}
		b.counters.bytesUp.Add(uint64(n))
	}
}

func (v *Valve) openBridge(src *net.UDPAddr) (*clientBridge, error) {
	up, err := net.DialUDP("udp", nil, v.backend)
	if err != nil {
		return nil, fmt.Errorf("dial backend: %w", err)
	}
	now := time.Now()
	b := &clientBridge{
		client:   src,
		upstream: up,
		valve:    v,
		openedAt: now,
	}
	b.lastSeen = now
	go b.readBackend()
	return b, nil
}

func (v *Valve) evictIdle() {
	t := time.NewTicker(15 * time.Second)
	defer t.Stop()
	for {
		select {
		case <-v.ctx.Done():
			return
		case <-t.C:
			cutoff := time.Now().Add(-v.idleTTL)
			v.mu.Lock()
			for k, b := range v.bridges {
				if b.lastUseBefore(cutoff) {
					slog.Debug("bridge idle evict", "port", v.route.Port, "client", k)
					b.close()
					delete(v.bridges, k)
				}
			}
			v.mu.Unlock()
		}
	}
}

func (v *Valve) Close() {
	v.cancel()
	v.pub.Close()
	v.mu.Lock()
	for k, b := range v.bridges {
		b.close()
		delete(v.bridges, k)
	}
	v.mu.Unlock()
}

type clientBridge struct {
	client   *net.UDPAddr
	upstream *net.UDPConn
	valve    *Valve

	counters counters // atomic — hot path

	mu       sync.Mutex
	lastSeen time.Time
	openedAt time.Time
}

func (b *clientBridge) touch() {
	b.mu.Lock()
	b.lastSeen = time.Now()
	b.mu.Unlock()
}

func (b *clientBridge) lastUseBefore(t time.Time) bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.lastSeen.Before(t)
}

func (b *clientBridge) close() {
	_ = b.upstream.Close()
}

// readBackend pumps datagrams from the backend back to the client via the
// public socket. Exits when the upstream socket is closed.
func (b *clientBridge) readBackend() {
	buf := make([]byte, 2048)
	for {
		n, err := b.upstream.Read(buf)
		if err != nil {
			return
		}
		b.touch()
		if _, err := b.valve.pub.WriteToUDP(buf[:n], b.client); err != nil {
			if b.valve.ctx.Err() == nil {
				slog.Warn("bridge reverse failed", "port", b.valve.route.Port, "err", err)
			}
			return
		}
		b.counters.bytesDown.Add(uint64(n))
	}
}