[management,proxy,client] Add L4 capabilities (TLS/TCP/UDP) (#5530)

2026-04-18 16:26:38 +00:00 · 2026-03-14 01:36:44 +08:00
parent fe9b844511
commit 3e6baea405
90 changed files with 9611 additions and 1397 deletions
--- a/proxy/internal/tcp/router.go
+++ b/proxy/internal/tcp/router.go
@@ -0,0 +1,570 @@
+package tcp
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"net/netip"
+	"slices"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+
+	"github.com/netbirdio/netbird/proxy/internal/accesslog"
+	"github.com/netbirdio/netbird/proxy/internal/types"
+)
+
+// defaultDialTimeout is the fallback dial timeout when no per-route
+// timeout is configured.
+const defaultDialTimeout = 30 * time.Second
+
+// SNIHost is a typed key for SNI hostname lookups.
+type SNIHost string
+
+// RouteType specifies how a connection should be handled.
+type RouteType int
+
+const (
+	// RouteHTTP routes the connection through the HTTP reverse proxy.
+	RouteHTTP RouteType = iota
+	// RouteTCP relays the connection directly to the backend (TLS passthrough).
+	RouteTCP
+)
+
+const (
+	// sniPeekTimeout is the deadline for reading the TLS ClientHello.
+	sniPeekTimeout = 5 * time.Second
+	// DefaultDrainTimeout is the default grace period for in-flight relay
+	// connections to finish during shutdown.
+	DefaultDrainTimeout = 30 * time.Second
+	// DefaultMaxRelayConns is the default cap on concurrent TCP relay connections per router.
+	DefaultMaxRelayConns = 4096
+	// httpChannelBuffer is the capacity of the channel feeding HTTP connections.
+	httpChannelBuffer = 4096
+)
+
+// DialResolver returns a DialContextFunc for the given account.
+type DialResolver func(accountID types.AccountID) (types.DialContextFunc, error)
+
+// Route describes where a connection for a given SNI should be sent.
+type Route struct {
+	Type      RouteType
+	AccountID types.AccountID
+	ServiceID types.ServiceID
+	// Domain is the service's configured domain, used for access log entries.
+	Domain string
+	// Protocol is the frontend protocol (tcp, tls), used for access log entries.
+	Protocol accesslog.Protocol
+	// Target is the backend address for TCP relay (e.g. "10.0.0.5:5432").
+	Target string
+	// ProxyProtocol enables sending a PROXY protocol v2 header to the backend.
+	ProxyProtocol bool
+	// DialTimeout overrides the default dial timeout for this route.
+	// Zero uses defaultDialTimeout.
+	DialTimeout time.Duration
+}
+
+// l4Logger sends layer-4 access log entries to the management server.
+type l4Logger interface {
+	LogL4(entry accesslog.L4Entry)
+}
+
+// RelayObserver receives callbacks for TCP relay lifecycle events.
+// All methods must be safe for concurrent use.
+type RelayObserver interface {
+	TCPRelayStarted(accountID types.AccountID)
+	TCPRelayEnded(accountID types.AccountID, duration time.Duration, srcToDst, dstToSrc int64)
+	TCPRelayDialError(accountID types.AccountID)
+	TCPRelayRejected(accountID types.AccountID)
+}
+
+// Router accepts raw TCP connections on a shared listener, peeks at
+// the TLS ClientHello to extract the SNI, and routes the connection
+// to either the HTTP reverse proxy or a direct TCP relay.
+type Router struct {
+	logger *log.Logger
+	// httpCh is immutable after construction: set only in NewRouter, nil in NewPortRouter.
+	httpCh       chan net.Conn
+	httpListener *chanListener
+	mu           sync.RWMutex
+	routes       map[SNIHost][]Route
+	fallback     *Route
+	draining     bool
+	dialResolve  DialResolver
+	activeConns  sync.WaitGroup
+	activeRelays sync.WaitGroup
+	relaySem     chan struct{}
+	drainDone    chan struct{}
+	observer     RelayObserver
+	accessLog    l4Logger
+	// svcCtxs tracks a context per service ID. All relay goroutines for a
+	// service derive from its context; canceling it kills them immediately.
+	svcCtxs    map[types.ServiceID]context.Context
+	svcCancels map[types.ServiceID]context.CancelFunc
+}
+
+// NewRouter creates a new SNI-based connection router.
+func NewRouter(logger *log.Logger, dialResolve DialResolver, addr net.Addr) *Router {
+	httpCh := make(chan net.Conn, httpChannelBuffer)
+	return &Router{
+		logger:       logger,
+		httpCh:       httpCh,
+		httpListener: newChanListener(httpCh, addr),
+		routes:       make(map[SNIHost][]Route),
+		dialResolve:  dialResolve,
+		relaySem:     make(chan struct{}, DefaultMaxRelayConns),
+		svcCtxs:      make(map[types.ServiceID]context.Context),
+		svcCancels:   make(map[types.ServiceID]context.CancelFunc),
+	}
+}
+
+// NewPortRouter creates a Router for a dedicated port without an HTTP
+// channel. Connections that don't match any SNI route fall through to
+// the fallback relay (if set) or are closed.
+func NewPortRouter(logger *log.Logger, dialResolve DialResolver) *Router {
+	return &Router{
+		logger:      logger,
+		routes:      make(map[SNIHost][]Route),
+		dialResolve: dialResolve,
+		relaySem:    make(chan struct{}, DefaultMaxRelayConns),
+		svcCtxs:     make(map[types.ServiceID]context.Context),
+		svcCancels:  make(map[types.ServiceID]context.CancelFunc),
+	}
+}
+
+// HTTPListener returns a net.Listener that yields connections routed
+// to the HTTP handler. Use this with http.Server.ServeTLS.
+func (r *Router) HTTPListener() net.Listener {
+	return r.httpListener
+}
+
+// AddRoute registers an SNI route. Multiple routes for the same host are
+// stored and resolved by priority at lookup time (HTTP > TCP).
+// Empty host is ignored to prevent conflicts with ECH/ESNI fallback.
+func (r *Router) AddRoute(host SNIHost, route Route) {
+	if host == "" {
+		return
+	}
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	routes := r.routes[host]
+	for i, existing := range routes {
+		if existing.ServiceID == route.ServiceID {
+			r.cancelServiceLocked(route.ServiceID)
+			routes[i] = route
+			return
+		}
+	}
+	r.routes[host] = append(routes, route)
+}
+
+// RemoveRoute removes the route for the given host and service ID.
+// Active relay connections for the service are closed immediately.
+// If other routes remain for the host, they are preserved.
+func (r *Router) RemoveRoute(host SNIHost, svcID types.ServiceID) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	r.routes[host] = slices.DeleteFunc(r.routes[host], func(route Route) bool {
+		return route.ServiceID == svcID
+	})
+	if len(r.routes[host]) == 0 {
+		delete(r.routes, host)
+	}
+	r.cancelServiceLocked(svcID)
+}
+
+// SetFallback registers a catch-all route for connections that don't
+// match any SNI route. On a port router this handles plain TCP relay;
+// on the main router it takes priority over the HTTP channel.
+func (r *Router) SetFallback(route Route) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.fallback = &route
+}
+
+// RemoveFallback clears the catch-all fallback route and closes any
+// active relay connections for the given service.
+func (r *Router) RemoveFallback(svcID types.ServiceID) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.fallback = nil
+	r.cancelServiceLocked(svcID)
+}
+
+// SetObserver sets the relay lifecycle observer. Must be called before Serve.
+func (r *Router) SetObserver(obs RelayObserver) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.observer = obs
+}
+
+// SetAccessLogger sets the L4 access logger. Must be called before Serve.
+func (r *Router) SetAccessLogger(l l4Logger) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.accessLog = l
+}
+
+// getObserver returns the current relay observer under the read lock.
+func (r *Router) getObserver() RelayObserver {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	return r.observer
+}
+
+// IsEmpty returns true when the router has no SNI routes and no fallback.
+func (r *Router) IsEmpty() bool {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	return len(r.routes) == 0 && r.fallback == nil
+}
+
+// Serve accepts connections from ln and routes them based on SNI.
+// It blocks until ctx is canceled or ln is closed, then drains
+// active relay connections up to DefaultDrainTimeout.
+func (r *Router) Serve(ctx context.Context, ln net.Listener) error {
+	done := make(chan struct{})
+	defer close(done)
+
+	go func() {
+		select {
+		case <-ctx.Done():
+			_ = ln.Close()
+			if r.httpListener != nil {
+				r.httpListener.Close()
+			}
+		case <-done:
+		}
+	}()
+
+	for {
+		conn, err := ln.Accept()
+		if err != nil {
+			if ctx.Err() != nil || errors.Is(err, net.ErrClosed) {
+				if ok := r.Drain(DefaultDrainTimeout); !ok {
+					r.logger.Warn("timed out waiting for connections to drain")
+				}
+				return nil
+			}
+			r.logger.Debugf("SNI router accept: %v", err)
+			continue
+		}
+		r.activeConns.Add(1)
+		go func() {
+			defer r.activeConns.Done()
+			r.handleConn(ctx, conn)
+		}()
+	}
+}
+
+// handleConn peeks at the TLS ClientHello and routes the connection.
+func (r *Router) handleConn(ctx context.Context, conn net.Conn) {
+	// Fast path: when no SNI routes and no HTTP channel exist (pure TCP
+	// fallback port), skip the TLS peek entirely to avoid read errors on
+	// non-TLS connections and reduce latency.
+	if r.isFallbackOnly() {
+		r.handleUnmatched(ctx, conn)
+		return
+	}
+
+	if err := conn.SetReadDeadline(time.Now().Add(sniPeekTimeout)); err != nil {
+		r.logger.Debugf("set SNI peek deadline: %v", err)
+		_ = conn.Close()
+		return
+	}
+
+	sni, wrapped, err := PeekClientHello(conn)
+	if err != nil {
+		r.logger.Debugf("SNI peek: %v", err)
+		if wrapped != nil {
+			r.handleUnmatched(ctx, wrapped)
+		} else {
+			_ = conn.Close()
+		}
+		return
+	}
+
+	if err := wrapped.SetReadDeadline(time.Time{}); err != nil {
+		r.logger.Debugf("clear SNI peek deadline: %v", err)
+		_ = wrapped.Close()
+		return
+	}
+
+	host := SNIHost(sni)
+	route, ok := r.lookupRoute(host)
+	if !ok {
+		r.handleUnmatched(ctx, wrapped)
+		return
+	}
+
+	if route.Type == RouteHTTP {
+		r.sendToHTTP(wrapped)
+		return
+	}
+
+	if err := r.relayTCP(ctx, wrapped, host, route); err != nil {
+		r.logger.WithFields(log.Fields{
+			"sni":        host,
+			"service_id": route.ServiceID,
+			"target":     route.Target,
+		}).Warnf("TCP relay: %v", err)
+		_ = wrapped.Close()
+	}
+}
+
+// isFallbackOnly returns true when the router has no SNI routes and no HTTP
+// channel, meaning all connections should go directly to the fallback relay.
+func (r *Router) isFallbackOnly() bool {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	return len(r.routes) == 0 && r.httpCh == nil
+}
+
+// handleUnmatched routes a connection that didn't match any SNI route.
+// This includes ECH/ESNI connections where the cleartext SNI is empty.
+// It tries the fallback relay first, then the HTTP channel, and closes
+// the connection if neither is available.
+func (r *Router) handleUnmatched(ctx context.Context, conn net.Conn) {
+	r.mu.RLock()
+	fb := r.fallback
+	r.mu.RUnlock()
+
+	if fb != nil {
+		if err := r.relayTCP(ctx, conn, SNIHost("fallback"), *fb); err != nil {
+			r.logger.WithFields(log.Fields{
+				"service_id": fb.ServiceID,
+				"target":     fb.Target,
+			}).Warnf("TCP relay (fallback): %v", err)
+			_ = conn.Close()
+		}
+		return
+	}
+	r.sendToHTTP(conn)
+}
+
+// lookupRoute returns the highest-priority route for the given SNI host.
+// HTTP routes take precedence over TCP routes.
+func (r *Router) lookupRoute(host SNIHost) (Route, bool) {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+	routes, ok := r.routes[host]
+	if !ok || len(routes) == 0 {
+		return Route{}, false
+	}
+	best := routes[0]
+	for _, route := range routes[1:] {
+		if route.Type < best.Type {
+			best = route
+		}
+	}
+	return best, true
+}
+
+// sendToHTTP feeds the connection to the HTTP handler via the channel.
+// If no HTTP channel is configured (port router), the router is
+// draining, or the channel is full, the connection is closed.
+func (r *Router) sendToHTTP(conn net.Conn) {
+	if r.httpCh == nil {
+		_ = conn.Close()
+		return
+	}
+
+	r.mu.RLock()
+	draining := r.draining
+	r.mu.RUnlock()
+
+	if draining {
+		_ = conn.Close()
+		return
+	}
+
+	select {
+	case r.httpCh <- conn:
+	default:
+		r.logger.Warnf("HTTP channel full, dropping connection from %s", conn.RemoteAddr())
+		_ = conn.Close()
+	}
+}
+
+// Drain prevents new relay connections from starting and waits for all
+// in-flight connection handlers and active relays to finish, up to the
+// given timeout. Returns true if all completed, false on timeout.
+func (r *Router) Drain(timeout time.Duration) bool {
+	r.mu.Lock()
+	r.draining = true
+	if r.drainDone == nil {
+		done := make(chan struct{})
+		go func() {
+			r.activeConns.Wait()
+			r.activeRelays.Wait()
+			close(done)
+		}()
+		r.drainDone = done
+	}
+	done := r.drainDone
+	r.mu.Unlock()
+
+	select {
+	case <-done:
+		return true
+	case <-time.After(timeout):
+		return false
+	}
+}
+
+// cancelServiceLocked cancels and removes the context for the given service,
+// closing all its active relay connections. Must be called with mu held.
+func (r *Router) cancelServiceLocked(svcID types.ServiceID) {
+	if cancel, ok := r.svcCancels[svcID]; ok {
+		cancel()
+		delete(r.svcCtxs, svcID)
+		delete(r.svcCancels, svcID)
+	}
+}
+
+// relayTCP sets up and runs a bidirectional TCP relay.
+// The caller owns conn and must close it if this method returns an error.
+// On success (nil error), both conn and backend are closed by the relay.
+func (r *Router) relayTCP(ctx context.Context, conn net.Conn, sni SNIHost, route Route) error {
+	svcCtx, err := r.acquireRelay(ctx, route)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		<-r.relaySem
+		r.activeRelays.Done()
+	}()
+
+	backend, err := r.dialBackend(svcCtx, route)
+	if err != nil {
+		obs := r.getObserver()
+		if obs != nil {
+			obs.TCPRelayDialError(route.AccountID)
+		}
+		return err
+	}
+
+	if route.ProxyProtocol {
+		if err := writeProxyProtoV2(conn, backend); err != nil {
+			_ = backend.Close()
+			return fmt.Errorf("write PROXY protocol header: %w", err)
+		}
+	}
+
+	obs := r.getObserver()
+	if obs != nil {
+		obs.TCPRelayStarted(route.AccountID)
+	}
+
+	entry := r.logger.WithFields(log.Fields{
+		"sni":        sni,
+		"service_id": route.ServiceID,
+		"target":     route.Target,
+	})
+	entry.Debug("TCP relay started")
+
+	start := time.Now()
+	s2d, d2s := Relay(svcCtx, entry, conn, backend, DefaultIdleTimeout)
+	elapsed := time.Since(start)
+
+	if obs != nil {
+		obs.TCPRelayEnded(route.AccountID, elapsed, s2d, d2s)
+	}
+	entry.Debugf("TCP relay ended (client→backend: %d bytes, backend→client: %d bytes)", s2d, d2s)
+
+	r.logL4Entry(route, conn, elapsed, s2d, d2s)
+	return nil
+}
+
+// acquireRelay checks draining state, increments activeRelays, and acquires
+// a semaphore slot. Returns the per-service context on success.
+// The caller must release the semaphore and call activeRelays.Done() when done.
+func (r *Router) acquireRelay(ctx context.Context, route Route) (context.Context, error) {
+	r.mu.Lock()
+	if r.draining {
+		r.mu.Unlock()
+		return nil, errors.New("router is draining")
+	}
+	r.activeRelays.Add(1)
+	svcCtx := r.getOrCreateServiceCtxLocked(ctx, route.ServiceID)
+	r.mu.Unlock()
+
+	select {
+	case r.relaySem <- struct{}{}:
+		return svcCtx, nil
+	default:
+		r.activeRelays.Done()
+		obs := r.getObserver()
+		if obs != nil {
+			obs.TCPRelayRejected(route.AccountID)
+		}
+		return nil, errors.New("TCP relay connection limit reached")
+	}
+}
+
+// dialBackend resolves the dialer for the route's account and dials the backend.
+func (r *Router) dialBackend(svcCtx context.Context, route Route) (net.Conn, error) {
+	dialFn, err := r.dialResolve(route.AccountID)
+	if err != nil {
+		return nil, fmt.Errorf("resolve dialer: %w", err)
+	}
+
+	dialTimeout := route.DialTimeout
+	if dialTimeout <= 0 {
+		dialTimeout = defaultDialTimeout
+	}
+	dialCtx, dialCancel := context.WithTimeout(svcCtx, dialTimeout)
+	backend, err := dialFn(dialCtx, "tcp", route.Target)
+	dialCancel()
+	if err != nil {
+		return nil, fmt.Errorf("dial backend %s: %w", route.Target, err)
+	}
+	return backend, nil
+}
+
+// logL4Entry sends a TCP relay access log entry if an access logger is configured.
+func (r *Router) logL4Entry(route Route, conn net.Conn, duration time.Duration, bytesUp, bytesDown int64) {
+	r.mu.RLock()
+	al := r.accessLog
+	r.mu.RUnlock()
+
+	if al == nil {
+		return
+	}
+
+	var sourceIP netip.Addr
+	if remote := conn.RemoteAddr(); remote != nil {
+		if ap, err := netip.ParseAddrPort(remote.String()); err == nil {
+			sourceIP = ap.Addr().Unmap()
+		}
+	}
+
+	al.LogL4(accesslog.L4Entry{
+		AccountID:     route.AccountID,
+		ServiceID:     route.ServiceID,
+		Protocol:      route.Protocol,
+		Host:          route.Domain,
+		SourceIP:      sourceIP,
+		DurationMs:    duration.Milliseconds(),
+		BytesUpload:   bytesUp,
+		BytesDownload: bytesDown,
+	})
+}
+
+// getOrCreateServiceCtxLocked returns the context for a service, creating one
+// if it doesn't exist yet. The context is a child of the server context.
+// Must be called with mu held.
+func (r *Router) getOrCreateServiceCtxLocked(parent context.Context, svcID types.ServiceID) context.Context {
+	if ctx, ok := r.svcCtxs[svcID]; ok {
+		return ctx
+	}
+	ctx, cancel := context.WithCancel(parent)
+	r.svcCtxs[svcID] = ctx
+	r.svcCancels[svcID] = cancel
+	return ctx
+}