[proxy] Reverse-proxy middleware framework, chain, and request plumbing

The per-target middleware chain (slots, dispatcher, mutation gate, metadata merger), body capture, access-log terminal sink, and the proxy wiring that builds + runs chains for synthesized agent-network services.
2026-09-26 16:49:08 +02:00 · 2026-06-27 00:43:35 +02:00
parent 6613d194ef
commit d4d158a8f3
33 changed files with 4790 additions and 52 deletions
@@ -55,6 +55,8 @@ import (
 	"github.com/netbirdio/netbird/proxy/internal/health"
 	"github.com/netbirdio/netbird/proxy/internal/k8s"
 	proxymetrics "github.com/netbirdio/netbird/proxy/internal/metrics"
+	"github.com/netbirdio/netbird/proxy/internal/middleware"
+	mwbuiltin "github.com/netbirdio/netbird/proxy/internal/middleware/builtin"
 	"github.com/netbirdio/netbird/proxy/internal/netutil"
 	"github.com/netbirdio/netbird/proxy/internal/proxy"
 	"github.com/netbirdio/netbird/proxy/internal/restrict"
@@ -77,29 +79,36 @@ type portRouter struct {

 type Server struct {
 	ctx               context.Context
-	mgmtClient        proto.ProxyServiceClient
-	proxy             *proxy.ReverseProxy
-	netbird           *roundtrip.NetBird
-	acme              *acme.Manager
+	mgmtClient    proto.ProxyServiceClient
+	proxy         *proxy.ReverseProxy
+	netbird       *roundtrip.NetBird
+	acme          *acme.Manager
 	staticCertWatcher *certwatch.Watcher
-	auth              *auth.Middleware
-	http              *http.Server
-	https             *http.Server
-	debug             *http.Server
-	healthServer      *health.Server
-	healthChecker     *health.Checker
-	meter             *proxymetrics.Metrics
-	accessLog         *accesslog.Logger
-	mainRouter        *nbtcp.Router
-	mainPort          uint16
-	udpMu             sync.Mutex
-	udpRelays         map[types.ServiceID]*udprelay.Relay
-	udpRelayWg        sync.WaitGroup
-	portMu            sync.RWMutex
-	portRouters       map[uint16]*portRouter
-	svcPorts          map[types.ServiceID][]uint16
-	lastMappings      map[types.ServiceID]*proto.ProxyMapping
-	portRouterWg      sync.WaitGroup
+	auth          *auth.Middleware
+	http          *http.Server
+	https         *http.Server
+	debug         *http.Server
+	healthServer  *health.Server
+	healthChecker *health.Checker
+	meter         *proxymetrics.Metrics
+	accessLog     *accesslog.Logger
+	// middlewareManager drives per-target middleware dispatch. Always
+	// constructed during boot; an empty registry produces empty chains and
+	// the reverse-proxy stays on the no-capture fast path.
+	middlewareManager *middleware.Manager
+	// middlewareRegistry is the source of registered middleware factories.
+	// Concrete middlewares register themselves through init().
+	middlewareRegistry *middleware.Registry
+	mainRouter    *nbtcp.Router
+	mainPort      uint16
+	udpMu         sync.Mutex
+	udpRelays     map[types.ServiceID]*udprelay.Relay
+	udpRelayWg    sync.WaitGroup
+	portMu        sync.RWMutex
+	portRouters   map[uint16]*portRouter
+	svcPorts      map[types.ServiceID][]uint16
+	lastMappings  map[types.ServiceID]*proto.ProxyMapping
+	portRouterWg  sync.WaitGroup

 	// hijackTracker tracks hijacked connections (e.g. WebSocket upgrades)
 	// so they can be closed during graceful shutdown, since http.Server.Shutdown
@@ -236,8 +245,20 @@ type Server struct {
 	// in processMappings before the receive loop reconnects to resync.
 	// Zero uses defaultMappingBatchWatchdog.
 	MappingBatchWatchdog time.Duration
+	// MiddlewareDataDir is the base directory the middleware system uses to
+	// resolve file-backed configuration (e.g. the cost_meter pricing table).
+	// Empty means any middleware that requires a file fails at configure time.
+	MiddlewareDataDir string
+	// MiddlewareCaptureBudgetBytes overrides the proxy-wide in-flight capture
+	// budget passed to middleware.NewManager. Zero or negative values fall
+	// back to defaultMiddlewareCaptureBudgetBytes (256 MiB).
+	MiddlewareCaptureBudgetBytes int64
 }

+// defaultMiddlewareCaptureBudgetBytes is the proxy-wide in-flight capture cap
+// passed to middleware.NewManager when MiddlewareCaptureBudgetBytes is unset.
+const defaultMiddlewareCaptureBudgetBytes = 256 << 20
+
 // clampIdleTimeout returns d capped to MaxSessionIdleTimeout when configured.
 func (s *Server) clampIdleTimeout(d time.Duration) time.Duration {
 	if s.MaxSessionIdleTimeout > 0 && d > s.MaxSessionIdleTimeout {
@@ -343,6 +364,15 @@ func (s *Server) Start(ctx context.Context) error {
 		return err
 	}

+	// Management client must be initialised BEFORE the middleware manager —
+	// initMiddlewareManager passes s.mgmtClient into the builtin FactoryContext
+	// that the limit-check / limit-record middlewares pull from. Reversed
+	// order would silently disable enforcement (mgmt=nil → allow-without-
+	// attribution + no-record).
+	if err := s.initMiddlewareManager(ctx); err != nil {
+		return fmt.Errorf("init middleware manager: %w", err)
+	}
+
 	runCtx, runCancel := context.WithCancel(ctx)
 	s.runCancel = runCancel

@@ -562,7 +592,11 @@ func (s *Server) initNetBirdClient() {
 // proxy host's resolver instead of the tunnel's DNS.
 func (s *Server) initReverseProxy() {
 	upstreamRT := roundtrip.NewMultiTransport(s.netbird, s.Logger)
-	s.proxy = proxy.NewReverseProxy(s.meter.RoundTripper(upstreamRT), s.ForwardedProto, s.TrustedProxies, s.Logger)
+	var rpOpts []proxy.Option
+	if s.middlewareManager != nil {
+		rpOpts = append(rpOpts, proxy.WithMiddlewareManager(s.middlewareManager))
+	}
+	s.proxy = proxy.NewReverseProxy(s.meter.RoundTripper(upstreamRT), s.ForwardedProto, s.TrustedProxies, s.Logger, rpOpts...)
 }

 // initGeoLookup configures the GeoLite2 lookup used for country-based
@@ -2047,9 +2081,94 @@ func (s *Server) updateMapping(ctx context.Context, mapping *proto.ProxyMapping)
 	m := s.protoToMapping(ctx, mapping)
 	s.proxy.AddMapping(m)
 	s.meter.AddMapping(m)
+	s.rebuildMiddlewareChains(svcID, m)
 	return nil
 }

+// initMiddlewareManager wires the middleware subsystem at boot. It configures
+// the per-process FactoryContext concrete middlewares consult, installs the
+// live-service check, and binds the resolver to the registry concrete
+// middlewares register themselves into via init().
+func (s *Server) initMiddlewareManager(ctx context.Context) error {
+	if s.meter == nil {
+		return fmt.Errorf("middleware manager requires metrics bundle")
+	}
+	otelMeter := s.meter.Meter()
+	mwbuiltin.Configure(ctx, s.MiddlewareDataDir, otelMeter, s.Logger, s.mgmtClient)
+
+	mwMetrics, err := middleware.NewMetrics(otelMeter)
+	if err != nil {
+		return fmt.Errorf("init middleware metrics: %w", err)
+	}
+	budgetBytes := s.MiddlewareCaptureBudgetBytes
+	if budgetBytes <= 0 {
+		budgetBytes = defaultMiddlewareCaptureBudgetBytes
+	}
+
+	registry := mwbuiltin.DefaultRegistry()
+	mgr := middleware.NewManager(budgetBytes, mwMetrics, s.Logger)
+	mgr.SetResolver(middleware.NewResolver(registry))
+	mgr.SetLiveServiceCheck(s.isLiveService)
+
+	s.middlewareRegistry = registry
+	s.middlewareManager = mgr
+	ids := registry.IDs()
+	s.Logger.Infof("middleware system enabled: %d built-in middlewares registered %v, capture budget %d bytes",
+		len(ids), ids, budgetBytes)
+	return nil
+}
+
+// rebuildMiddlewareChains converts m into per-path bindings and calls
+// Manager.Rebuild. Short-circuits when the middleware manager is unset.
+func (s *Server) rebuildMiddlewareChains(svcID types.ServiceID, m proxy.Mapping) {
+	if s.middlewareManager == nil {
+		return
+	}
+	bindings := buildMiddlewareBindings(svcID, m)
+	if err := s.middlewareManager.Rebuild(string(svcID), bindings); err != nil {
+		s.Logger.WithError(err).WithField("service_id", svcID).Error("failed to rebuild middleware chains")
+	}
+}
+
+// isLiveService reports whether svcID is currently present in the live
+// mapping cache. Used by the middleware manager to confirm a chain is still
+// referenced before rebuilding it from cached bindings.
+func (s *Server) isLiveService(svcID string) bool {
+	s.portMu.RLock()
+	defer s.portMu.RUnlock()
+	_, ok := s.lastMappings[types.ServiceID(svcID)]
+	return ok
+}
+
+// invalidateMiddlewareChains drops every middleware chain registered for svcID.
+func (s *Server) invalidateMiddlewareChains(svcID types.ServiceID) {
+	if s.middlewareManager == nil {
+		return
+	}
+	s.middlewareManager.Invalidate(string(svcID))
+}
+
+// buildMiddlewareBindings converts the path targets of m into the per-path
+// binding list the middleware manager's Rebuild expects. Targets without any
+// middleware specs are skipped.
+func buildMiddlewareBindings(svcID types.ServiceID, m proxy.Mapping) []middleware.PathTargetBinding {
+	if len(m.Paths) == 0 {
+		return nil
+	}
+	bindings := make([]middleware.PathTargetBinding, 0, len(m.Paths))
+	for pathID, pt := range m.Paths {
+		if pt == nil || len(pt.Middlewares) == 0 {
+			continue
+		}
+		bindings = append(bindings, middleware.PathTargetBinding{
+			ServiceID: string(svcID),
+			PathID:    pathID,
+			Specs:     pt.Middlewares,
+		})
+	}
+	return bindings
+}
+
 // removeMapping tears down routes/relays and the NetBird peer for a service.
 // Uses the stored mapping state when available to ensure all previously
 // configured routes are cleaned up.
@@ -2085,6 +2204,8 @@ func (s *Server) cleanupMappingRoutes(mapping *proto.ProxyMapping) {
 	svcID := types.ServiceID(mapping.GetId())
 	host := mapping.GetDomain()

+	s.invalidateMiddlewareChains(svcID)
+
 	// HTTP/TLS cleanup (only relevant when a domain is set).
 	if host != "" {
 		d := domain.Domain(host)
@@ -2192,6 +2313,12 @@ func (s *Server) protoToMapping(ctx context.Context, mapping *proto.ProxyMapping
 				pt.RequestTimeout = d.AsDuration()
 			}
 			pt.DirectUpstream = opts.GetDirectUpstream()
+			// Agent-network middleware specs + capture config + flag ride on
+			// the same per-target options.
+			pt.CaptureConfig = translateMiddlewareCaptureConfig(mapping.GetId(), opts)
+			pt.Middlewares = translateMiddlewareConfigs(ctx, mapping.GetId(), opts.GetMiddlewares(), s.middlewareRegistry)
+			pt.AgentNetwork = opts.GetAgentNetwork()
+			pt.DisableAccessLog = opts.GetDisableAccessLog()
 		}
 		pt.RequestTimeout = s.clampDialTimeout(pt.RequestTimeout)
 		paths[pathMapping.GetPath()] = pt