From d4d158a8f35de7ad69fb6f40ae96aa94733c6fcd Mon Sep 17 00:00:00 2001 From: mlsmaycon Date: Sat, 27 Jun 2026 00:43:35 +0200 Subject: [PATCH] [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. --- proxy/internal/accesslog/logger.go | 50 +++ proxy/internal/accesslog/middleware.go | 14 +- proxy/internal/accesslog/middleware_test.go | 185 ++++++++ proxy/internal/auth/middleware_test.go | 119 ++++- proxy/internal/metrics/metrics.go | 9 + proxy/internal/middleware/bodypolicy.go | 63 +++ proxy/internal/middleware/bodytap/request.go | 344 +++++++++++++++ proxy/internal/middleware/bodytap/response.go | 189 ++++++++ .../middleware/bodytap/routing_scan_test.go | 86 ++++ proxy/internal/middleware/chain.go | 320 ++++++++++++++ proxy/internal/middleware/chain_test.go | 370 ++++++++++++++++ proxy/internal/middleware/decision.go | 81 ++++ proxy/internal/middleware/dispatcher.go | 189 ++++++++ proxy/internal/middleware/headerpolicy.go | 99 +++++ proxy/internal/middleware/keys.go | 86 ++++ proxy/internal/middleware/manager.go | 412 ++++++++++++++++++ proxy/internal/middleware/metadata.go | 99 +++++ proxy/internal/middleware/metrics.go | 171 ++++++++ proxy/internal/middleware/middleware.go | 47 ++ proxy/internal/middleware/redaction.go | 79 ++++ proxy/internal/middleware/registry.go | 121 +++++ proxy/internal/middleware/spec.go | 44 ++ proxy/internal/middleware/types.go | 253 +++++++++++ .../agent_network_chain_realstack_test.go | 321 ++++++++++++++ proxy/internal/proxy/context.go | 43 +- proxy/internal/proxy/reverseproxy.go | 359 ++++++++++++++- proxy/internal/proxy/reverseproxy_test.go | 43 ++ proxy/internal/proxy/servicemapping.go | 16 + proxy/internal/proxy/strip_prefix_test.go | 30 ++ proxy/middleware_register.go | 16 + proxy/middleware_translate.go | 165 +++++++ proxy/middleware_translate_test.go | 246 +++++++++++ proxy/server.go | 173 +++++++- 33 files changed, 4790 insertions(+), 52 deletions(-) create mode 100644 proxy/internal/accesslog/middleware_test.go create mode 100644 proxy/internal/middleware/bodypolicy.go create mode 100644 proxy/internal/middleware/bodytap/request.go create mode 100644 proxy/internal/middleware/bodytap/response.go create mode 100644 proxy/internal/middleware/bodytap/routing_scan_test.go create mode 100644 proxy/internal/middleware/chain.go create mode 100644 proxy/internal/middleware/chain_test.go create mode 100644 proxy/internal/middleware/decision.go create mode 100644 proxy/internal/middleware/dispatcher.go create mode 100644 proxy/internal/middleware/headerpolicy.go create mode 100644 proxy/internal/middleware/keys.go create mode 100644 proxy/internal/middleware/manager.go create mode 100644 proxy/internal/middleware/metadata.go create mode 100644 proxy/internal/middleware/metrics.go create mode 100644 proxy/internal/middleware/middleware.go create mode 100644 proxy/internal/middleware/redaction.go create mode 100644 proxy/internal/middleware/registry.go create mode 100644 proxy/internal/middleware/spec.go create mode 100644 proxy/internal/middleware/types.go create mode 100644 proxy/internal/proxy/agent_network_chain_realstack_test.go create mode 100644 proxy/internal/proxy/strip_prefix_test.go create mode 100644 proxy/middleware_register.go create mode 100644 proxy/middleware_translate.go create mode 100644 proxy/middleware_translate_test.go diff --git a/proxy/internal/accesslog/logger.go b/proxy/internal/accesslog/logger.go index 3283f61db..db868b4e0 100644 --- a/proxy/internal/accesslog/logger.go +++ b/proxy/internal/accesslog/logger.go @@ -128,6 +128,7 @@ type logEntry struct { BytesDownload int64 Protocol Protocol Metadata map[string]string + AgentNetwork bool } // Protocol identifies the transport protocol of an access log entry. @@ -214,6 +215,54 @@ func (l *Logger) allowDenyLog(serviceID types.ServiceID, reason string) bool { return false } +// usageMetadataKeys is the allowlist of metadata retained on a stripped, +// usage-only agent-network entry. Mirrors the llm.* / cost.* keys in +// proxy/internal/middleware/keys.go — only the dimensions management needs to +// record a usage row (provider / model / tokens / cost / groups). +var usageMetadataKeys = map[string]struct{}{ + "llm.provider": {}, + "llm.model": {}, + "llm.resolved_provider_id": {}, + "llm.input_tokens": {}, + "llm.output_tokens": {}, + "llm.total_tokens": {}, + "cost.usd_total": {}, + "llm.authorising_groups": {}, +} + +// stripAgentNetworkEntryForUsage returns the entry reduced to what's needed to +// record usage/cost: it drops request detail (host / path / source IP) and any +// prompt capture, keeping the LLM usage metadata plus the caller identity +// (user / auth mechanism) needed for attribution. Shipped when an +// agent-network account has log collection disabled but usage must still be +// collected. logEntry is passed by value, so mutating it here is safe; Metadata +// is replaced with a fresh map rather than mutated in place. +func stripAgentNetworkEntryForUsage(entry logEntry) logEntry { + entry.Host = "" + entry.Path = "" + entry.SourceIP = netip.Addr{} + // Drop the rest of the per-request telemetry too — a usage-only entry + // must carry the LLM usage metadata and caller identity, nothing that + // describes the individual request. + entry.Method = "" + entry.ResponseCode = 0 + entry.DurationMs = 0 + entry.BytesUpload = 0 + entry.BytesDownload = 0 + entry.Protocol = "" + + if len(entry.Metadata) > 0 { + stripped := make(map[string]string, len(usageMetadataKeys)) + for k := range usageMetadataKeys { + if v, ok := entry.Metadata[k]; ok { + stripped[k] = v + } + } + entry.Metadata = stripped + } + return entry +} + func (l *Logger) log(entry logEntry) { // Fire off the log request in a separate routine. // This increases the possibility of losing a log message @@ -264,6 +313,7 @@ func (l *Logger) log(entry logEntry) { BytesDownload: entry.BytesDownload, Protocol: string(entry.Protocol), Metadata: entry.Metadata, + AgentNetwork: entry.AgentNetwork, }, }); err != nil { l.logger.WithFields(log.Fields{ diff --git a/proxy/internal/accesslog/middleware.go b/proxy/internal/accesslog/middleware.go index 5a0684c19..9c644418e 100644 --- a/proxy/internal/accesslog/middleware.go +++ b/proxy/internal/accesslog/middleware.go @@ -83,11 +83,23 @@ func (l *Logger) Middleware(next http.Handler) http.Handler { BytesDownload: bytesDownload, Protocol: ProtocolHTTP, Metadata: capturedData.GetMetadata(), + AgentNetwork: capturedData.GetAgentNetwork(), } l.logger.Debugf("response: request_id=%s method=%s host=%s path=%s status=%d duration=%dms source=%s origin=%s service=%s account=%s", requestID, r.Method, host, r.URL.Path, sw.status, duration.Milliseconds(), sourceIp, capturedData.GetOrigin(), capturedData.GetServiceID(), capturedData.GetAccountID()) - l.log(entry) + // Emit the access log unless the matched target opted out + // (agent-network synth targets do this when the account's + // EnableLogCollection toggle is off). For agent-network entries we + // still ship a stripped, usage-only record even when suppressed, so + // usage/cost is collected regardless of the log-collection toggle; + // request detail and prompt capture are dropped before sending. + switch { + case !capturedData.GetSuppressAccessLog(): + l.log(entry) + case entry.AgentNetwork: + l.log(stripAgentNetworkEntryForUsage(entry)) + } // Track usage for cost monitoring (upload + download) by domain l.trackUsage(host, bytesUpload+bytesDownload) diff --git a/proxy/internal/accesslog/middleware_test.go b/proxy/internal/accesslog/middleware_test.go new file mode 100644 index 000000000..cf91957a8 --- /dev/null +++ b/proxy/internal/accesslog/middleware_test.go @@ -0,0 +1,185 @@ +package accesslog + +import ( + "context" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + + "github.com/netbirdio/netbird/proxy/internal/proxy" + "github.com/netbirdio/netbird/shared/management/proto" +) + +// recorderClient is a minimal stub for the access-log gRPCClient interface. It +// counts SendAccessLog invocations and signals on every call so tests can +// deterministically wait for the goroutine inside Logger.log without sleeping. +type recorderClient struct { + mu sync.Mutex + calls int64 + lastEntry *proto.AccessLog + called chan struct{} +} + +func newRecorderClient() *recorderClient { + return &recorderClient{called: make(chan struct{}, 16)} +} + +func (r *recorderClient) SendAccessLog(_ context.Context, in *proto.SendAccessLogRequest, _ ...grpc.CallOption) (*proto.SendAccessLogResponse, error) { + r.mu.Lock() + r.calls++ + r.lastEntry = in.GetLog() + r.mu.Unlock() + select { + case r.called <- struct{}{}: + default: + } + return &proto.SendAccessLogResponse{}, nil +} + +func (r *recorderClient) callCount() int64 { + r.mu.Lock() + defer r.mu.Unlock() + return r.calls +} + +// newTestLogger builds a Logger backed by the supplied recorderClient. It is +// the same constructor production uses, just with a stub gRPC client — no +// mocks, no interface re-implementations. +func newTestLogger(t *testing.T, client *recorderClient) *Logger { + t.Helper() + logger := NewLogger(client, nil, nil) + t.Cleanup(logger.Close) + return logger +} + +// TestMiddleware_SuppressAccessLog_SkipsLogSink asserts the suppression gate. +// When the inner handler stamps SuppressAccessLog=true on CapturedData (mirrors +// what reverseproxy does when the matched target's DisableAccessLog flag is +// set), the middleware must NOT invoke the access-log sink. Bandwidth telemetry +// (trackUsage) keeps running — it's the call to SendAccessLog that we gate. +func TestMiddleware_SuppressAccessLog_SkipsLogSink(t *testing.T) { + client := newRecorderClient() + l := newTestLogger(t, client) + + inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + cd := proxy.CapturedDataFromContext(r.Context()) + require.NotNil(t, cd, "middleware must inject CapturedData into the request context") + cd.SetSuppressAccessLog(true) + w.WriteHeader(http.StatusOK) + }) + + srv := httptest.NewServer(l.Middleware(inner)) + t.Cleanup(srv.Close) + + resp, err := http.Get(srv.URL + "/agent-network/v1/chat/completions") + require.NoError(t, err, "GET against suppressed target must succeed") + require.NoError(t, resp.Body.Close()) + require.Equal(t, http.StatusOK, resp.StatusCode, "inner handler must run normally") + + // Give the goroutine fence a beat (Logger.log dispatches in a goroutine). + // The negative assertion needs a small window: if a send is going to + // happen, it happens promptly. + select { + case <-client.called: + t.Fatalf("access-log sink must not be invoked when SuppressAccessLog=true (got %d call(s))", client.callCount()) + case <-time.After(150 * time.Millisecond): + } + + assert.Equal(t, int64(0), client.callCount(), + "SendAccessLog must not be called for suppressed requests") +} + +// TestMiddleware_SuppressAccessLog_DefaultEmitsLog is the regression sanity: +// when nothing sets SuppressAccessLog (the universal default for every +// non-agent-network target), the middleware MUST still emit the access-log +// entry. This is the guarantee that wires-through to the EnableLogCollection +// gate without breaking anyone who isn't opted in. +func TestMiddleware_SuppressAccessLog_DefaultEmitsLog(t *testing.T) { + client := newRecorderClient() + l := newTestLogger(t, client) + + var innerRan atomic.Bool + inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + innerRan.Store(true) + // Intentionally DO NOT touch SuppressAccessLog — mirrors every + // non-agent-network target. + w.WriteHeader(http.StatusOK) + }) + + srv := httptest.NewServer(l.Middleware(inner)) + t.Cleanup(srv.Close) + + resp, err := http.Get(srv.URL + "/service/healthz") + require.NoError(t, err, "GET against default target must succeed") + require.NoError(t, resp.Body.Close()) + require.Equal(t, http.StatusOK, resp.StatusCode, "inner handler must run normally") + require.True(t, innerRan.Load(), "inner handler must have run") + + select { + case <-client.called: + case <-time.After(2 * time.Second): + t.Fatalf("SendAccessLog must be invoked for non-suppressed requests, none observed (calls=%d)", client.callCount()) + } + + assert.Equal(t, int64(1), client.callCount(), + "non-suppressed request must produce exactly one access-log send") +} + +// TestMiddleware_SuppressAccessLog_PreservesUsageTracking proves the gate is +// surgical: with SuppressAccessLog=true the access-log send is skipped, but +// the per-domain usage tracker still records the bytes transferred. This is +// the cost-monitoring guarantee called out in the gate's comment. +func TestMiddleware_SuppressAccessLog_PreservesUsageTracking(t *testing.T) { + client := newRecorderClient() + l := newTestLogger(t, client) + + payload := []byte("ok") + inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + cd := proxy.CapturedDataFromContext(r.Context()) + require.NotNil(t, cd, "middleware must inject CapturedData") + cd.SetSuppressAccessLog(true) + w.WriteHeader(http.StatusOK) + _, _ = w.Write(payload) + }) + + srv := httptest.NewServer(l.Middleware(inner)) + t.Cleanup(srv.Close) + + resp, err := http.Get(srv.URL + "/agent-network/v1/chat/completions") + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + + // Allow trackUsage to land — it runs synchronously after l.log(entry) is + // (would have been) called. + time.Sleep(50 * time.Millisecond) + + l.usageMux.Lock() + usage, present := l.domainUsage[hostNoPort(srv.URL)] + l.usageMux.Unlock() + require.True(t, present, "domain usage must be tracked even when the access-log is suppressed") + assert.Greater(t, usage.bytesTransferred, int64(0), "bytesTransferred must include the response payload") + assert.Equal(t, int64(0), client.callCount(), + "SendAccessLog must remain suppressed across the response write") +} + +// hostNoPort extracts the host name from an httptest server URL. The +// middleware strips the port before keying domain usage, so the test mirrors +// that to look the entry up. +func hostNoPort(url string) string { + // httptest URLs are always "http://127.0.0.1:PORT". + const prefix = "http://" + host := url[len(prefix):] + for i := 0; i < len(host); i++ { + if host[i] == ':' || host[i] == '/' { + return host[:i] + } + } + return host +} diff --git a/proxy/internal/auth/middleware_test.go b/proxy/internal/auth/middleware_test.go index c0ec5c94c..6608c2b22 100644 --- a/proxy/internal/auth/middleware_test.go +++ b/proxy/internal/auth/middleware_test.go @@ -297,6 +297,109 @@ func TestProtect_SessionCookieGroupsPropagate(t *testing.T) { assert.Equal(t, groups, capturedData.GetUserGroups(), "CapturedData groups must be retained after handler completes") } +// stubTunnelValidator implements SessionValidator for the tunnel-peer +// path. ValidateTunnelPeer returns a fixed response so tests can assert +// how the proxy maps it onto CapturedData, and records whether the +// fast-path actually reached management. +type stubTunnelValidator struct { + called bool + resp *proto.ValidateTunnelPeerResponse +} + +func (s *stubTunnelValidator) ValidateSession(context.Context, *proto.ValidateSessionRequest, ...grpc.CallOption) (*proto.ValidateSessionResponse, error) { + return nil, errors.New("not used in this test") +} + +func (s *stubTunnelValidator) ValidateTunnelPeer(context.Context, *proto.ValidateTunnelPeerRequest, ...grpc.CallOption) (*proto.ValidateTunnelPeerResponse, error) { + s.called = true + return s.resp, nil +} + +// TestProtect_PrivateService_TunnelPeerGroupsPropagate locks the agent-network +// auth path end-to-end at the proxy edge: a Private service must route through +// ValidateTunnelPeer and lift the returned peer_group_ids onto CapturedData so +// the llm_router group-authorisation pass can see them. Regression guard for +// the failure that surfaces downstream as llm_policy.no_authorised_provider — +// i.e. a synthesised service that reaches the proxy without private=true (so +// this path is skipped) leaves UserGroups empty and every request is denied. +func TestProtect_PrivateService_TunnelPeerGroupsPropagate(t *testing.T) { + groups := []string{"grp-admins", "grp-users"} + names := []string{"Admins", "Users"} + validator := &stubTunnelValidator{resp: &proto.ValidateTunnelPeerResponse{ + Valid: true, + UserId: "user-1", + UserEmail: "user@example.com", + SessionToken: "tunnel-session-token", + PeerGroupIds: groups, + PeerGroupNames: names, + }} + mw := NewMiddleware(log.StandardLogger(), validator, nil) + kp := generateTestKeyPair(t) + + // Private service: no operator schemes — auth gates solely on the tunnel peer. + require.NoError(t, mw.AddDomain("agent.example.com", nil, kp.PublicKey, time.Hour, "acct-1", "svc-1", nil, true)) + + cd := proxy.NewCapturedData("") + cd.SetClientIP(netip.MustParseAddr("100.90.1.14")) // CGNAT tunnel source + + var seenGroups []string + var seenUser string + handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c := proxy.CapturedDataFromContext(r.Context()) + require.NotNil(t, c, "captured data must be present in request context") + seenGroups = c.GetUserGroups() + seenUser = c.GetUserID() + w.WriteHeader(http.StatusOK) + })) + + lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) { + return PeerIdentity{}, true + }) + req := httptest.NewRequest(http.MethodPost, "http://agent.example.com/v1/chat/completions", nil) + req.RemoteAddr = "100.90.1.14:5000" + req = req.WithContext(WithTunnelLookup(proxy.WithCapturedData(req.Context(), cd), lookup)) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code, "private service must authorise a tunnel peer the validator accepts") + assert.Equal(t, groups, seenGroups, "ValidateTunnelPeer peer_group_ids must reach CapturedData.UserGroups for llm_router authorisation") + assert.Equal(t, "user-1", seenUser, "tunnel-peer principal must reach CapturedData") + assert.Equal(t, groups, cd.GetUserGroups(), "groups must persist on CapturedData after the handler returns") +} + +// TestProtect_PrivateService_TunnelPeerDenied verifies the deny path: when +// ValidateTunnelPeer rejects the peer, a Private service 403s and never reaches +// the upstream handler (no fall-through to unauthenticated pass-through). +func TestProtect_PrivateService_TunnelPeerDenied(t *testing.T) { + validator := &stubTunnelValidator{resp: &proto.ValidateTunnelPeerResponse{ + Valid: false, + DeniedReason: "not_in_group", + }} + mw := NewMiddleware(log.StandardLogger(), validator, nil) + kp := generateTestKeyPair(t) + require.NoError(t, mw.AddDomain("agent.example.com", nil, kp.PublicKey, time.Hour, "acct-1", "svc-1", nil, true)) + + cd := proxy.NewCapturedData("") + cd.SetClientIP(netip.MustParseAddr("100.90.1.14")) + + reached := false + handler := mw.Protect(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + reached = true + w.WriteHeader(http.StatusOK) + })) + lookup := TunnelLookupFunc(func(_ netip.Addr) (PeerIdentity, bool) { + return PeerIdentity{}, true + }) + req := httptest.NewRequest(http.MethodPost, "http://agent.example.com/v1/chat/completions", nil) + req.RemoteAddr = "100.90.1.14:5000" + req = req.WithContext(WithTunnelLookup(proxy.WithCapturedData(req.Context(), cd), lookup)) + rec := httptest.NewRecorder() + handler.ServeHTTP(rec, req) + + assert.Equal(t, http.StatusForbidden, rec.Code, "private service must 403 when the tunnel peer is rejected") + assert.False(t, reached, "denied private request must not reach the upstream handler") +} + func TestProtect_ExpiredSessionCookieIsRejected(t *testing.T) { mw := NewMiddleware(log.StandardLogger(), nil, nil) kp := generateTestKeyPair(t) @@ -1228,22 +1331,6 @@ func TestProtect_NonOIDCSchemes_PlainHTTP_NotBlocked(t *testing.T) { assert.Equal(t, http.StatusUnauthorized, rec.Code, "PIN-only domain should serve the login page on plain HTTP") } -// stubTunnelValidator records ValidateTunnelPeer calls so a test can -// assert whether the fast-path reached management. -type stubTunnelValidator struct { - called bool - resp *proto.ValidateTunnelPeerResponse -} - -func (s *stubTunnelValidator) ValidateSession(context.Context, *proto.ValidateSessionRequest, ...grpc.CallOption) (*proto.ValidateSessionResponse, error) { - return nil, errors.New("not used in this test") -} - -func (s *stubTunnelValidator) ValidateTunnelPeer(context.Context, *proto.ValidateTunnelPeerRequest, ...grpc.CallOption) (*proto.ValidateTunnelPeerResponse, error) { - s.called = true - return s.resp, nil -} - // TestProtect_TunnelPeerFastPath_RequiresInboundMarker guards the // anti-spoof gate: a request with an RFC1918 source IP arriving on the // public listener (no TunnelLookupFromContext attached) must not be diff --git a/proxy/internal/metrics/metrics.go b/proxy/internal/metrics/metrics.go index 41a6b0dd4..5fd23d934 100644 --- a/proxy/internal/metrics/metrics.go +++ b/proxy/internal/metrics/metrics.go @@ -17,6 +17,7 @@ import ( // Metrics collects OpenTelemetry metrics for the proxy. type Metrics struct { ctx context.Context + meter metric.Meter requestsTotal metric.Int64Counter activeRequests metric.Int64UpDownCounter configuredDomains metric.Int64UpDownCounter @@ -49,10 +50,18 @@ type Metrics struct { mappingPaths map[string]int } +// Meter returns the OpenTelemetry meter the bundle was built with, so other +// subsystems (e.g. the middleware manager) register instruments on the same +// meter. +func (m *Metrics) Meter() metric.Meter { + return m.meter +} + // New creates a Metrics instance using the given OpenTelemetry meter. func New(ctx context.Context, meter metric.Meter) (*Metrics, error) { m := &Metrics{ ctx: ctx, + meter: meter, mappingPaths: make(map[string]int), } diff --git a/proxy/internal/middleware/bodypolicy.go b/proxy/internal/middleware/bodypolicy.go new file mode 100644 index 000000000..f31486fc8 --- /dev/null +++ b/proxy/internal/middleware/bodypolicy.go @@ -0,0 +1,63 @@ +package middleware + +import ( + "bytes" + "errors" + "fmt" + "io" + "net/http" + "strconv" + "strings" +) + +// ErrExpectContinue is returned when a middleware attempts to replace +// the body of a request that advertised Expect: 100-continue. +var ErrExpectContinue = errors.New("body replace rejected: request has Expect: 100-continue") + +// ErrOriginalNotDrained is returned when the original body was not +// fully consumed before replacement. This prevents the backend from +// seeing a mix of original bytes and the replacement. +var ErrOriginalNotDrained = errors.New("body replace rejected: original body not drained") + +// ErrContentLengthMismatch is returned when the client-advertised +// Content-Length disagrees with the number of bytes actually read from +// the body (short-read). +var ErrContentLengthMismatch = errors.New("body replace rejected: content-length mismatch (short read)") + +// ValidateBodyReplace runs the smuggling-prevention rules before a +// body replacement is applied. Callers must pass originalDrained=true +// once they have read r.Body to EOF. +func ValidateBodyReplace(r *http.Request, newBody []byte, originalDrained bool) error { + if r == nil { + return errors.New("body replace rejected: nil request") + } + if strings.EqualFold(r.Header.Get("Expect"), "100-continue") { + return ErrExpectContinue + } + if !originalDrained { + return ErrOriginalNotDrained + } + if cl := r.Header.Get("Content-Length"); cl != "" && r.ContentLength > 0 { + parsed, err := strconv.ParseInt(cl, 10, 64) + if err == nil && parsed != r.ContentLength { + return fmt.Errorf("%w: header=%d actual=%d", ErrContentLengthMismatch, parsed, r.ContentLength) + } + } + return nil +} + +// ApplyBodyReplace swaps r.Body for a reader over newBody, recomputes +// Content-Length, and strips Transfer-Encoding and Trailer so no stale +// framing reaches the backend. +func ApplyBodyReplace(r *http.Request, newBody []byte) { + if r == nil { + return + } + r.Body = io.NopCloser(bytes.NewReader(newBody)) + r.ContentLength = int64(len(newBody)) + r.Header.Set("Content-Length", strconv.Itoa(len(newBody))) + r.Header.Del("Transfer-Encoding") + r.Header.Del("Trailer") + r.TransferEncoding = nil + r.Trailer = nil +} diff --git a/proxy/internal/middleware/bodytap/request.go b/proxy/internal/middleware/bodytap/request.go new file mode 100644 index 000000000..826883a96 --- /dev/null +++ b/proxy/internal/middleware/bodytap/request.go @@ -0,0 +1,344 @@ +// Package bodytap owns the framework-side body capture used by the +// middleware chain. Request capture buffers up to N bytes of the +// request body for middleware inspection while replaying the original +// stream to the upstream. Response capture tees up to N bytes off the +// streaming response while every byte continues to flow to the client +// untouched. +// +// The package is the single owner of body access — middlewares never +// read req.Body or hijack the response writer. All inspection happens +// against the buffer surfaced by the tap, so streaming remains +// transparent to the client even when middlewares need access to the +// payload. +package bodytap + +import ( + "bytes" + "encoding/json" + "errors" + "io" + "net/http" + "strconv" + "strings" + "sync" +) + +// MaxRoutingScanBytes bounds how far ScanRoutingFields will read into a +// request body to recover routing fields when the normal capture is +// bypassed for size. Sized to comfortably hold a 1M-token context +// request (whose `model` field a client may place after a multi-MB +// `messages` array) while still capping pathological inputs. +const MaxRoutingScanBytes int64 = 32 << 20 + +// Request bypass reasons emitted as the `mw.capture.bypass_reason` +// metadata key by the chain when a request body is not surfaced. +const ( + BypassUpgradeHeader = "upgrade_header" + BypassConnectionUpgrd = "connection_upgrade" + BypassContentType = "content_type_not_allowed" + BypassBudget = "capture_budget_exhausted" + BypassNoConfig = "no_capture_config" + BypassNoMiddlewares = "no_middlewares" + BypassCapZero = "cap_zero" + BypassContentLengthCap = "content_length_over_cap" +) + +// DefaultCaptureBudgetBytes is the default global capture-budget size. +const DefaultCaptureBudgetBytes int64 = 256 << 20 + +// Config holds per-target body capture limits after clamp validation. +// A zero MaxRequestBytes / MaxResponseBytes disables capture in that +// direction. +type Config struct { + MaxRequestBytes int64 + MaxResponseBytes int64 + ContentTypes []string +} + +// Budget is the global token-bucket semaphore shared across all +// in-flight captures so a single misbehaving target cannot exhaust the +// proxy. +type Budget interface { + Acquire(n int64) bool + Release(n int64) +} + +// NewBudget returns a Budget with the given total byte cap. A zero or +// negative total disables the budget check. +func NewBudget(total int64) Budget { + return &budget{total: total} +} + +type budget struct { + mu sync.Mutex + used int64 + total int64 +} + +func (b *budget) Acquire(n int64) bool { + if n <= 0 { + return true + } + b.mu.Lock() + defer b.mu.Unlock() + if b.total <= 0 { + return true + } + if b.used+n > b.total { + return false + } + b.used += n + return true +} + +func (b *budget) Release(n int64) { + if n <= 0 { + return + } + b.mu.Lock() + defer b.mu.Unlock() + if b.total <= 0 { + return + } + b.used -= n + if b.used < 0 { + b.used = 0 + } +} + +// CaptureRequest reads up to cfg.MaxRequestBytes from r.Body into a +// buffer suitable for middleware inspection, replacing r.Body with a +// replay reader so the upstream still sees the original bytes. When +// bypass != "" no body is read and r.Body is left untouched. The +// returned release function must be invoked once the request is fully +// processed; it returns the acquired budget tokens to the shared pool. +// release is always non-nil and is safe to defer immediately after the +// call. +func CaptureRequest(r *http.Request, cfg *Config, b Budget) (body []byte, truncated bool, originalSize int64, bypass string, release func(), err error) { + release = func() {} + if r == nil { + return nil, false, 0, BypassNoConfig, release, nil + } + if cfg == nil { + return nil, false, 0, BypassNoConfig, release, nil + } + if cfg.MaxRequestBytes <= 0 { + return nil, false, 0, BypassCapZero, release, nil + } + if r.Header.Get("Upgrade") != "" { + return nil, false, 0, BypassUpgradeHeader, release, nil + } + if strings.EqualFold(r.Header.Get("Connection"), "upgrade") { + return nil, false, 0, BypassConnectionUpgrd, release, nil + } + if !contentTypeAllowed(r.Header.Get("Content-Type"), cfg.ContentTypes) { + return nil, false, 0, BypassContentType, release, nil + } + + originalSize = parseContentLength(r.Header.Get("Content-Length")) + if originalSize > cfg.MaxRequestBytes { + return nil, true, originalSize, BypassContentLengthCap, release, nil + } + + limit := cfg.MaxRequestBytes + if b != nil && !b.Acquire(limit) { + return nil, false, originalSize, BypassBudget, release, nil + } + if b != nil { + var released sync.Once + release = func() { + released.Do(func() { b.Release(limit) }) + } + } + + if r.Body == nil || r.Body == http.NoBody { + release() + release = func() {} + return nil, false, originalSize, "", release, nil + } + + limited := io.LimitReader(r.Body, limit+1) + buf, readErr := io.ReadAll(limited) + if readErr != nil && !errors.Is(readErr, io.EOF) { + release() + release = func() {} + return nil, false, originalSize, "", release, readErr + } + + truncated = int64(len(buf)) > limit + if truncated { + replay := append([]byte(nil), buf...) + viewable := buf[:limit] + r.Body = &replayReadCloser{replay: bytes.NewReader(replay), tail: r.Body} + return viewable, true, originalSize, "", release, nil + } + _ = r.Body.Close() + r.Body = io.NopCloser(bytes.NewReader(buf)) + if originalSize <= 0 { + originalSize = int64(len(buf)) + } + return buf, false, originalSize, "", release, nil +} + +// replayReadCloser replays the captured prefix and then forwards the +// remaining bytes from the original body so the upstream sees the +// full request stream even when capture truncates. +type replayReadCloser struct { + replay *bytes.Reader + tail io.ReadCloser + drained bool +} + +func (r *replayReadCloser) Read(p []byte) (int, error) { + if !r.drained { + n, err := r.replay.Read(p) + if n > 0 { + return n, nil + } + if errors.Is(err, io.EOF) { + r.drained = true + } else if err != nil { + return 0, err + } + } + return r.tail.Read(p) +} + +func (r *replayReadCloser) Close() error { + return r.tail.Close() +} + +// ScanRoutingFields recovers the LLM routing fields ("model" and +// "stream") from a request whose normal capture was bypassed or +// truncated for size. It reads up to maxScan bytes of r.Body to locate +// the top-level keys — clients (e.g. Claude Code) may place `model` +// after a multi-MB `messages` array — then restores r.Body so the +// upstream still receives the full, untouched stream. Only the small +// routing fields are extracted; the prompt is never buffered for +// capture, keeping memory bounded. Returns ok=false when the body isn't +// a JSON object, the model field isn't found within maxScan, or on a +// read error. +func ScanRoutingFields(r *http.Request, maxScan int64) (model string, stream bool, ok bool) { + if r == nil || r.Body == nil || r.Body == http.NoBody || maxScan <= 0 { + return "", false, false + } + limited := io.LimitReader(r.Body, maxScan+1) + buf, readErr := io.ReadAll(limited) + if readErr != nil && !errors.Is(readErr, io.EOF) { + // Mid-stream read error (e.g. client disconnect): restore the bytes + // read so far plus the untouched tail and abort, rather than + // forwarding only the partial prefix as if it were the whole body. + r.Body = &replayReadCloser{replay: bytes.NewReader(append([]byte(nil), buf...)), tail: r.Body} + return "", false, false + } + if int64(len(buf)) > maxScan { + // Body exceeds the scan ceiling: restore the read prefix plus the + // untouched tail so the upstream still gets every byte. + r.Body = &replayReadCloser{replay: bytes.NewReader(append([]byte(nil), buf...)), tail: r.Body} + } else { + _ = r.Body.Close() + r.Body = io.NopCloser(bytes.NewReader(buf)) + } + return scanTopLevelModelStream(buf) +} + +// scanTopLevelModelStream walks the top level of a JSON object via a +// streaming token reader, extracting the "model" string and "stream" +// bool without materialising large values (each non-target value is +// skipped as a RawMessage). Tolerant of truncation: returns whatever was +// found before a malformed/short tail. +func scanTopLevelModelStream(body []byte) (model string, stream bool, ok bool) { + dec := json.NewDecoder(bytes.NewReader(body)) + tok, err := dec.Token() + if err != nil { + return "", false, false + } + if d, isDelim := tok.(json.Delim); !isDelim || d != '{' { + return "", false, false + } + for dec.More() { + keyTok, err := dec.Token() + if err != nil { + return model, stream, ok + } + key, _ := keyTok.(string) + switch key { + case "model": + var v string + if dec.Decode(&v) == nil { + model, ok = v, true + } + case "stream": + var v bool + if dec.Decode(&v) == nil { + stream = v + } + default: + // Skip the value by walking tokens instead of decoding it into + // a json.RawMessage — a multi-MB messages array would otherwise + // be materialised in full just to be discarded. + if err := skipValue(dec); err != nil { + return model, stream, ok + } + } + } + return model, stream, ok +} + +// skipValue consumes one JSON value from dec without materialising it. +// Scalars are a single token; objects/arrays are walked to their matching +// close delimiter so nested structures are skipped in bounded memory. +func skipValue(dec *json.Decoder) error { + tok, err := dec.Token() + if err != nil { + return err + } + d, isDelim := tok.(json.Delim) + if !isDelim || (d != '{' && d != '[') { + return nil + } + depth := 1 + for depth > 0 { + tok, err := dec.Token() + if err != nil { + return err + } + if d, ok := tok.(json.Delim); ok { + switch d { + case '{', '[': + depth++ + case '}', ']': + depth-- + } + } + } + return nil +} + +func contentTypeAllowed(ct string, allowed []string) bool { + if len(allowed) == 0 { + return false + } + media := ct + if idx := strings.Index(ct, ";"); idx >= 0 { + media = ct[:idx] + } + media = strings.TrimSpace(strings.ToLower(media)) + for _, a := range allowed { + if strings.EqualFold(strings.TrimSpace(a), media) { + return true + } + } + return false +} + +func parseContentLength(v string) int64 { + if v == "" { + return 0 + } + parsed, err := strconv.ParseInt(v, 10, 64) + if err != nil || parsed < 0 { + return 0 + } + return parsed +} diff --git a/proxy/internal/middleware/bodytap/response.go b/proxy/internal/middleware/bodytap/response.go new file mode 100644 index 000000000..c23e35b34 --- /dev/null +++ b/proxy/internal/middleware/bodytap/response.go @@ -0,0 +1,189 @@ +package bodytap + +import ( + "bytes" + "net/http" + "sync" + + "github.com/netbirdio/netbird/proxy/internal/responsewriter" +) + +// CapturingResponseWriter wraps an http.ResponseWriter, forwards bytes +// immediately to the client, and tees a bounded copy into an internal +// buffer for middleware inspection. Streaming-aware in the sense that +// every byte the upstream emits flows to the client without queuing +// — the tee just sees a bounded prefix. SSE-aware parsing happens in +// the response middleware against the buffered prefix; this writer +// makes no attempt to demux event boundaries. +// +// Flusher and Hijacker are preserved via responsewriter.PassthroughWriter. +type CapturingResponseWriter struct { + *responsewriter.PassthroughWriter + mu sync.Mutex + buf bytes.Buffer + cap int64 + status int + statusSet bool + written int64 + truncated bool + stopped bool + releaseBuf func() + released sync.Once + bypassed bool + bypassReas string + acquiredCap int64 +} + +// NewCapturingResponseWriter returns a writer that tees up to maxBytes +// into a capped buffer while forwarding bytes to the underlying writer +// immediately. When budget is non-nil the writer pre-acquires maxBytes +// from it and the returned wrapper must be released by calling +// Release() once the response is fully forwarded. If the budget cannot +// be acquired the writer falls back to forwarding the response +// unmodified, exposes Bypassed()=true with reason BypassBudget, and +// releases nothing. +func NewCapturingResponseWriter(w http.ResponseWriter, maxBytes int64, b Budget) *CapturingResponseWriter { + cw := &CapturingResponseWriter{ + PassthroughWriter: responsewriter.New(w), + cap: maxBytes, + status: http.StatusOK, + releaseBuf: func() {}, + } + if maxBytes <= 0 { + // Capture disabled: mark stopped so Write never tees and never + // flags truncation (a zero cap means "don't capture", not + // "captured nothing"). + cw.stopped = true + return cw + } + if b == nil { + return cw + } + if !b.Acquire(maxBytes) { + cw.bypassed = true + cw.bypassReas = BypassBudget + cw.cap = 0 + cw.stopped = true + return cw + } + cw.acquiredCap = maxBytes + cw.releaseBuf = func() { b.Release(maxBytes) } + return cw +} + +// Release returns the response capture budget acquired at construction +// back to the shared pool. Idempotent. Safe to call from a defer +// immediately after construction even when the writer ended up +// bypassing the budget. +func (c *CapturingResponseWriter) Release() { + if c == nil { + return + } + c.released.Do(func() { + if c.releaseBuf != nil { + c.releaseBuf() + } + }) +} + +// Bypassed reports whether the writer fell through to a no-tee +// passthrough because the response capture budget could not be +// acquired. +func (c *CapturingResponseWriter) Bypassed() bool { + if c == nil { + return false + } + c.mu.Lock() + defer c.mu.Unlock() + return c.bypassed +} + +// BypassReason returns the bypass code recorded by the budget check. +// Empty when capture proceeded normally. +func (c *CapturingResponseWriter) BypassReason() string { + if c == nil { + return "" + } + c.mu.Lock() + defer c.mu.Unlock() + return c.bypassReas +} + +// WriteHeader records the status code and forwards it to the underlying +// writer. Only the first call commits the status — matching HTTP semantics, +// where superfluous WriteHeader calls (and any call after the body has +// started) are ignored — so Status() reflects the code actually sent. +func (c *CapturingResponseWriter) WriteHeader(status int) { + c.mu.Lock() + if c.statusSet { + c.mu.Unlock() + return + } + c.status = status + c.statusSet = true + c.mu.Unlock() + c.PassthroughWriter.WriteHeader(status) +} + +// Write forwards p to the underlying writer unmodified and copies up +// to the remaining buffer capacity into the tee buffer. +func (c *CapturingResponseWriter) Write(p []byte) (int, error) { + n, err := c.PassthroughWriter.Write(p) + if n > 0 { + c.mu.Lock() + // The first byte commits the status (implicit 200 if WriteHeader was + // never called); a later WriteHeader must not change Status(). + c.statusSet = true + c.written += int64(n) + if !c.stopped { + remaining := c.cap - int64(c.buf.Len()) + if remaining <= 0 { + c.truncated = true + c.stopped = true + } else { + take := int64(n) + if take > remaining { + take = remaining + c.truncated = true + c.stopped = true + } + c.buf.Write(p[:take]) + } + } + c.mu.Unlock() + } + return n, err +} + +// Status returns the captured status code (defaults to 200 when +// WriteHeader has not been called). +func (c *CapturingResponseWriter) Status() int { + c.mu.Lock() + defer c.mu.Unlock() + return c.status +} + +// Body returns a copy of the buffered response prefix. +func (c *CapturingResponseWriter) Body() []byte { + c.mu.Lock() + defer c.mu.Unlock() + out := make([]byte, c.buf.Len()) + copy(out, c.buf.Bytes()) + return out +} + +// Truncated reports whether the buffered prefix stopped short of the +// full response stream. +func (c *CapturingResponseWriter) Truncated() bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.truncated +} + +// BytesWritten returns the total number of bytes forwarded to the +// underlying writer. +func (c *CapturingResponseWriter) BytesWritten() int64 { + c.mu.Lock() + defer c.mu.Unlock() + return c.written +} diff --git a/proxy/internal/middleware/bodytap/routing_scan_test.go b/proxy/internal/middleware/bodytap/routing_scan_test.go new file mode 100644 index 000000000..1748c989e --- /dev/null +++ b/proxy/internal/middleware/bodytap/routing_scan_test.go @@ -0,0 +1,86 @@ +package bodytap + +import ( + "fmt" + "io" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// makeBigAnthropicBody builds a request body shaped like Claude Code's: +// a multi-MB "messages" array with the routing fields (model, stream) +// placed AFTER it, which is the ordering that defeats a prefix-only +// capture. +func makeBigAnthropicBody(t *testing.T, model string, stream bool, messagesBytes int) string { + t.Helper() + filler := strings.Repeat("x", messagesBytes) + return fmt.Sprintf( + `{"max_tokens":64000,"messages":[{"role":"user","content":%q}],"model":%q,"stream":%t}`, + filler, model, stream, + ) +} + +func TestScanRoutingFields_ModelAfterLargeMessages(t *testing.T) { + body := makeBigAnthropicBody(t, "claude-opus-4-8", true, 3<<20) // 3 MiB messages + req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body)) + + model, stream, ok := ScanRoutingFields(req, MaxRoutingScanBytes) + require.True(t, ok, "model must be recovered even when it follows a multi-MB messages array") + assert.Equal(t, "claude-opus-4-8", model, "model field must be extracted") + assert.True(t, stream, "stream field must be extracted") + + // Body must be fully restored for the upstream. + got, err := io.ReadAll(req.Body) + require.NoError(t, err) + assert.Equal(t, body, string(got), "the full request body must be replayed to upstream after scanning") +} + +func TestScanRoutingFields_SmallBody(t *testing.T) { + body := `{"model":"claude-opus-4-8","stream":false,"messages":[]}` + req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body)) + + model, stream, ok := ScanRoutingFields(req, MaxRoutingScanBytes) + require.True(t, ok) + assert.Equal(t, "claude-opus-4-8", model) + assert.False(t, stream) + + got, _ := io.ReadAll(req.Body) + assert.Equal(t, body, string(got), "small bodies must also be restored intact") +} + +func TestScanRoutingFields_NoModel(t *testing.T) { + body := `{"stream":true,"messages":[]}` + req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body)) + + _, _, ok := ScanRoutingFields(req, MaxRoutingScanBytes) + assert.False(t, ok, "ok must be false when no model field is present") + + got, _ := io.ReadAll(req.Body) + assert.Equal(t, body, string(got), "body must be restored even when model is absent") +} + +func TestScanRoutingFields_NotJSON(t *testing.T) { + body := "this is not json at all" + req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body)) + + _, _, ok := ScanRoutingFields(req, MaxRoutingScanBytes) + assert.False(t, ok, "ok must be false for a non-JSON body") +} + +func TestScanRoutingFields_ModelBeyondScanCeiling(t *testing.T) { + // model sits after 4 MiB of messages but the scan ceiling is 1 MiB: + // model can't be recovered, yet the full body must still replay. + body := makeBigAnthropicBody(t, "claude-opus-4-8", true, 4<<20) + req := httptest.NewRequest("POST", "https://x/v1/messages", strings.NewReader(body)) + + _, _, ok := ScanRoutingFields(req, 1<<20) + assert.False(t, ok, "model beyond the scan ceiling is not recoverable") + + got, err := io.ReadAll(req.Body) + require.NoError(t, err) + assert.Equal(t, body, string(got), "the full body must still replay to upstream even when the scan gives up") +} diff --git a/proxy/internal/middleware/chain.go b/proxy/internal/middleware/chain.go new file mode 100644 index 000000000..45d32cdb0 --- /dev/null +++ b/proxy/internal/middleware/chain.go @@ -0,0 +1,320 @@ +package middleware + +import ( + "context" + "net/http" + "sync" +) + +// boundMiddleware pairs a validated spec with the resolved middleware +// instance the chain will invoke. +type boundMiddleware struct { + spec Spec + mw Middleware +} + +// Chain is the ordered set of middlewares that run for a specific +// target. Chains are immutable once built; Manager produces a new +// Chain on every Rebuild. +// +// Ordering: middlewares are kept in registration order. RunRequest +// iterates the SlotOnRequest middlewares in order; RunResponse +// iterates the SlotOnResponse middlewares in reverse order +// (middleware-style LIFO so the last to see the request is the first +// to see the response); RunTerminal iterates the SlotTerminal +// middlewares in registration order, after every on_response slot has +// emitted, so the metadata bag they observe is complete. +// +// Close drains in-flight invocations and tears down each middleware. +// Callers swapping a chain via Manager invoke Close on the old chain +// after the swap so live requests finish on the previous instance. +type Chain struct { + targetID string + all []boundMiddleware + onRequest []int + onResponse []int + terminal []int + dispatcher *Dispatcher + inflight sync.WaitGroup +} + +// NewChain assembles a Chain from the bound middlewares. The slice +// order is the registration order; the chain captures index slices +// per slot so iteration does not re-scan the slot field per call. +func NewChain(targetID string, bound []boundMiddleware, d *Dispatcher) *Chain { + c := &Chain{ + targetID: targetID, + all: bound, + dispatcher: d, + } + for i, bm := range bound { + switch bm.spec.Slot { + case SlotOnRequest: + c.onRequest = append(c.onRequest, i) + case SlotOnResponse: + c.onResponse = append(c.onResponse, i) + case SlotTerminal: + c.terminal = append(c.terminal, i) + } + } + return c +} + +// Close waits for outstanding invocations against this chain to +// finish (bounded by ctx) and releases the middleware instances bound +// to it. Safe to call once the chain has been removed from the +// routing snapshot. Subsequent Run* calls are still safe (return +// without invoking) but Close itself is one-shot. +func (c *Chain) Close(ctx context.Context) error { + if c == nil { + return nil + } + if ctx == nil { + ctx = context.Background() + } + done := make(chan struct{}) + go func() { + c.inflight.Wait() + close(done) + }() + select { + case <-done: + case <-ctx.Done(): + // Drain timed out: requests may still be running against these + // middleware instances, so tearing them down now risks a + // use-after-close. Leave them (a bounded leak) and surface the + // timeout; the runaway backstop in the Manager already alerts. + return ctx.Err() + } + for _, bm := range c.all { + if bm.mw == nil { + continue + } + if err := bm.mw.Close(); err != nil { + c.dispatcher.logger.Debugf("middleware %s close: %v", bm.spec.ID, err) + } + } + return nil +} + +// Empty reports whether the chain has no middlewares. +func (c *Chain) Empty() bool { + return c == nil || len(c.all) == 0 +} + +// TargetID returns the key used to find this chain. +func (c *Chain) TargetID() string { + if c == nil { + return "" + } + return c.targetID +} + +// IDs returns the ordered list of middleware IDs bound to this chain. +func (c *Chain) IDs() []string { + if c == nil { + return nil + } + out := make([]string, len(c.all)) + for i, bm := range c.all { + out[i] = bm.spec.ID + } + return out +} + +// RunRequest iterates the on_request slot in registration order. Deny +// short-circuits the remaining middlewares and returns the deny +// output. The caller owns applying mutations to the real request and +// merging the metadata returned in `merged` into the captured-data +// bag passed to subsequent slots. +// +// Each middleware sees the metadata emitted by earlier middlewares in +// the same slot — this is how llm_guardrail reads +// llm.request_prompt_raw from llm_request_parser without a side +// channel, and how cost_meter reads tokens emitted by +// llm_response_parser on the response leg. +// +// If any middleware emits a non-nil Mutations.RewriteUpstream while +// satisfying the mutation gates (CanMutate && MutationsSupported), the +// latest such value is returned to the caller. Last-write-wins so the +// last middleware in the slot can override an earlier rewrite. +func (c *Chain) RunRequest(ctx context.Context, r *http.Request, in *Input, acc *Accumulator) (denied *Output, merged []KV, rewrite *UpstreamRewrite, err error) { + if c.Empty() || len(c.onRequest) == 0 { + return nil, nil, nil, nil + } + c.inflight.Add(1) + defer c.inflight.Done() + running := append([]KV(nil), in.Metadata...) + for _, idx := range c.onRequest { + bm := c.all[idx] + call := cloneInputFor(in, SlotOnRequest) + call.Metadata = append([]KV(nil), running...) + out, invErr := c.dispatcher.Invoke(ctx, bm.spec, bm.mw, call) + if invErr != nil && out == nil { + continue + } + if out == nil { + continue + } + + accepted, rejected := acc.Emit(bm.spec.ID, bm.spec.MetadataKeys, out.Metadata) + for _, rej := range rejected { + c.dispatcher.metrics.IncMetadataRejected(ctx, bm.spec.ID, rej.Reason) + } + merged = append(merged, accepted...) + running = append(running, accepted...) + + if out.Decision == DecisionDeny { + c.dispatcher.metrics.IncRequest(ctx, bm.spec.ID, c.targetID, "deny") + return out, merged, rewrite, nil + } + c.dispatcher.metrics.IncRequest(ctx, bm.spec.ID, c.targetID, "allow") + + if rw := mutationRewrite(bm.spec, out.Mutations); rw != nil { + rewrite = rw + } + if r != nil && bm.spec.CanMutate && out.Mutations != nil { + applyMutations(ctx, c.dispatcher, bm.spec, r, out.Mutations) + } + } + return nil, merged, rewrite, nil +} + +// RunResponse iterates the on_response slot in reverse registration +// order, matching the middleware "last in, first out" convention so +// the last middleware to see the request is the first to see the +// response. Middlewares cannot deny; they emit metadata. +// +// As with RunRequest, each middleware sees the metadata emitted by +// earlier middlewares in this slot — accumulated in the order the +// middlewares run (LIFO of registration). cost_meter relies on this +// to read llm.input_tokens / llm.output_tokens that +// llm_response_parser emitted just before it. +func (c *Chain) RunResponse(ctx context.Context, in *Input, acc *Accumulator) (merged []KV) { + if c.Empty() || len(c.onResponse) == 0 { + return nil + } + c.inflight.Add(1) + defer c.inflight.Done() + running := append([]KV(nil), in.Metadata...) + for i := len(c.onResponse) - 1; i >= 0; i-- { + bm := c.all[c.onResponse[i]] + call := cloneInputFor(in, SlotOnResponse) + call.Metadata = append([]KV(nil), running...) + out, _ := c.dispatcher.Invoke(ctx, bm.spec, bm.mw, call) + if out == nil { + continue + } + accepted, rejected := acc.Emit(bm.spec.ID, bm.spec.MetadataKeys, out.Metadata) + for _, rej := range rejected { + c.dispatcher.metrics.IncMetadataRejected(ctx, bm.spec.ID, rej.Reason) + } + merged = append(merged, accepted...) + running = append(running, accepted...) + c.dispatcher.metrics.IncRequest(ctx, bm.spec.ID, c.targetID, "passthrough") + } + return merged +} + +// RunTerminal iterates the terminal slot in registration order, after +// every on_response middleware has emitted. Terminal middlewares +// observe the full metadata bag carried in `in.Metadata` plus any +// emissions from terminal middlewares that ran before them; they +// cannot deny and cannot mutate. +func (c *Chain) RunTerminal(ctx context.Context, in *Input, acc *Accumulator) (merged []KV) { + if c.Empty() || len(c.terminal) == 0 { + return nil + } + c.inflight.Add(1) + defer c.inflight.Done() + running := append([]KV(nil), in.Metadata...) + for _, idx := range c.terminal { + bm := c.all[idx] + call := cloneInputFor(in, SlotTerminal) + call.Metadata = append([]KV(nil), running...) + out, _ := c.dispatcher.Invoke(ctx, bm.spec, bm.mw, call) + if out == nil { + continue + } + accepted, rejected := acc.Emit(bm.spec.ID, bm.spec.MetadataKeys, out.Metadata) + for _, rej := range rejected { + c.dispatcher.metrics.IncMetadataRejected(ctx, bm.spec.ID, rej.Reason) + } + merged = append(merged, accepted...) + running = append(running, accepted...) + c.dispatcher.metrics.IncRequest(ctx, bm.spec.ID, c.targetID, "terminal") + } + return merged +} + +// mutationRewrite returns the upstream rewrite carried in m when the +// spec's mutation gates allow it. The rewrite itself is not applied +// here; the caller (reverse proxy) decides whether to honour it. +func mutationRewrite(spec Spec, m *Mutations) *UpstreamRewrite { + if m == nil || m.RewriteUpstream == nil { + return nil + } + if !spec.CanMutate || !spec.MutationsSupported { + return nil + } + return m.RewriteUpstream +} + +func applyMutations(ctx context.Context, d *Dispatcher, spec Spec, r *http.Request, m *Mutations) { + if m == nil { + return + } + add, remove, blocked := FilterHeaderMutations(m) + for _, h := range blocked { + d.metrics.IncHeaderMutationBlocked(ctx, spec.ID, h) + } + for _, name := range remove { + r.Header.Del(name) + } + for _, kv := range add { + r.Header.Add(kv.Key, kv.Value) + } + if len(m.BodyReplace) == 0 { + return + } + if err := ValidateBodyReplace(r, m.BodyReplace, true); err != nil { + d.logger.Warnf("middleware %s body replace rejected: %v", spec.ID, err) + return + } + ApplyBodyReplace(r, m.BodyReplace) +} + +// cloneInputFor deep-copies the mutation-prone fields of Input so +// each middleware receives an isolated view. +func cloneInputFor(in *Input, slot Slot) *Input { + if in == nil { + return nil + } + out := *in + out.Slot = slot + out.Headers = cloneKVs(in.Headers) + out.RespHeaders = cloneKVs(in.RespHeaders) + out.Metadata = cloneKVs(in.Metadata) + if len(in.UserGroups) > 0 { + out.UserGroups = append([]string(nil), in.UserGroups...) + } + if len(in.UserGroupNames) > 0 { + out.UserGroupNames = append([]string(nil), in.UserGroupNames...) + } + if len(in.Body) > 0 { + out.Body = append([]byte(nil), in.Body...) + } + if len(in.RespBody) > 0 { + out.RespBody = append([]byte(nil), in.RespBody...) + } + return &out +} + +func cloneKVs(in []KV) []KV { + if len(in) == 0 { + return nil + } + out := make([]KV, len(in)) + copy(out, in) + return out +} diff --git a/proxy/internal/middleware/chain_test.go b/proxy/internal/middleware/chain_test.go new file mode 100644 index 000000000..929ccee08 --- /dev/null +++ b/proxy/internal/middleware/chain_test.go @@ -0,0 +1,370 @@ +package middleware + +import ( + "context" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fakeMiddleware is a minimal Middleware for chain composition tests. +// It records the metadata the dispatcher hands to it and emits a +// caller-supplied Output. Tests use the recorded snapshot to assert +// that earlier-in-slot emissions are visible to later middlewares. +type fakeMiddleware struct { + id string + slot Slot + keys []string + emit []KV + decision Decision + mutationsSupported bool + canMutate bool + mutations *Mutations + + // seen captures the in.Metadata snapshot the dispatcher passed to + // Invoke, so tests can assert ordering and visibility. + seen []KV +} + +func (f *fakeMiddleware) ID() string { return f.id } +func (f *fakeMiddleware) Version() string { return "test" } +func (f *fakeMiddleware) Slot() Slot { return f.slot } +func (f *fakeMiddleware) AcceptedContentTypes() []string { return nil } +func (f *fakeMiddleware) MetadataKeys() []string { return f.keys } +func (f *fakeMiddleware) MutationsSupported() bool { return f.mutationsSupported } +func (f *fakeMiddleware) Close() error { return nil } + +func (f *fakeMiddleware) Invoke(_ context.Context, in *Input) (*Output, error) { + f.seen = append([]KV(nil), in.Metadata...) + out := &Output{Decision: f.decision, Metadata: append([]KV(nil), f.emit...)} + if f.mutations != nil { + m := *f.mutations + out.Mutations = &m + } + return out, nil +} + +// chainFor builds a Chain over the given middlewares with a noop +// dispatcher. +func chainFor(t *testing.T, mws ...*fakeMiddleware) *Chain { + t.Helper() + bound := make([]boundMiddleware, len(mws)) + for i, mw := range mws { + bound[i] = boundMiddleware{ + spec: Spec{ + ID: mw.id, + Slot: mw.slot, + Enabled: true, + MetadataKeys: mw.keys, + CanMutate: mw.canMutate, + MutationsSupported: mw.mutationsSupported, + }, + mw: mw, + } + } + disp := NewDispatcher(nil, nil) + return NewChain("t-1", bound, disp) +} + +// TestChain_RunRequest_ThreadsMetadataAcrossMiddlewares locks that +// each on_request middleware sees metadata emitted by earlier +// middlewares in the same slot. Regression cover for the original +// chain.go where every iteration cloned from the same source `in` and +// later middlewares (e.g. llm_guardrail) couldn't read what the first +// (e.g. llm_request_parser) had just emitted. +func TestChain_RunRequest_ThreadsMetadataAcrossMiddlewares(t *testing.T) { + first := &fakeMiddleware{ + id: "first", + slot: SlotOnRequest, + keys: []string{"foo.k"}, + emit: []KV{{Key: "foo.k", Value: "v"}}, + } + second := &fakeMiddleware{ + id: "second", + slot: SlotOnRequest, + keys: []string{"bar.k"}, + emit: []KV{{Key: "bar.k", Value: "z"}}, + } + c := chainFor(t, first, second) + acc := NewAccumulator(0) + + denied, merged, rewrite, err := c.RunRequest(context.Background(), nil, &Input{}, acc) + require.NoError(t, err) + assert.Nil(t, denied, "no deny without DecisionDeny") + assert.Nil(t, rewrite, "no rewrite without Mutations.RewriteUpstream") + + require.Len(t, second.seen, 1, "the second middleware must observe one prior emission") + assert.Equal(t, "foo.k", second.seen[0].Key, "second middleware must see the first middleware's key") + assert.Equal(t, "v", second.seen[0].Value, "second middleware must see the first middleware's value") + + require.Len(t, merged, 2, "merged slice contains both middleware emissions") +} + +// TestChain_RunResponse_ThreadsMetadataAcrossMiddlewares does the +// same for the response slot. The response slot iterates in reverse +// registration order, so the middleware registered LAST runs first. +// This test asserts that a middleware running later (in reverse +// order) sees the metadata emitted by the one that ran before it. +func TestChain_RunResponse_ThreadsMetadataAcrossMiddlewares(t *testing.T) { + // Registration order: [outer, inner]. + // Reverse iteration runs inner first, outer second. + // outer must see inner's emission. + outer := &fakeMiddleware{ + id: "outer", + slot: SlotOnResponse, + keys: []string{"outer.k"}, + emit: []KV{{Key: "outer.k", Value: "o"}}, + } + inner := &fakeMiddleware{ + id: "inner", + slot: SlotOnResponse, + keys: []string{"inner.k"}, + emit: []KV{{Key: "inner.k", Value: "i"}}, + } + c := chainFor(t, outer, inner) + acc := NewAccumulator(0) + + merged := c.RunResponse(context.Background(), &Input{}, acc) + + require.Len(t, outer.seen, 1, "outer must observe inner's emission") + assert.Equal(t, "inner.k", outer.seen[0].Key) + require.Len(t, merged, 2, "merged slice contains both response emissions") +} + +// TestChain_RunResponse_CostMeterScenario simulates the synth-service +// chain shape (response_parser registered AFTER cost_meter so reverse +// iter runs response_parser first). The cost_meter analogue must see +// the tokens response_parser just emitted — this is the exact +// regression that produced cost.skipped=missing_tokens in the live +// access logs. +func TestChain_RunResponse_CostMeterScenario(t *testing.T) { + // Synthesizer registers cost_meter first, response_parser second. + costMeter := &fakeMiddleware{ + id: "cost_meter", + slot: SlotOnResponse, + keys: []string{"cost.usd_total", "cost.skipped"}, + } + respParser := &fakeMiddleware{ + id: "llm_response_parser", + slot: SlotOnResponse, + keys: []string{"llm.input_tokens", "llm.output_tokens"}, + emit: []KV{ + {Key: "llm.input_tokens", Value: "13"}, + {Key: "llm.output_tokens", Value: "259"}, + }, + } + c := chainFor(t, costMeter, respParser) + acc := NewAccumulator(0) + + _ = c.RunResponse(context.Background(), &Input{}, acc) + + require.Len(t, costMeter.seen, 2, "cost_meter must observe both token keys emitted by response_parser") + keys := []string{costMeter.seen[0].Key, costMeter.seen[1].Key} + assert.ElementsMatch(t, []string{"llm.input_tokens", "llm.output_tokens"}, keys, + "cost_meter must see the exact keys response_parser emitted") + values := []string{costMeter.seen[0].Value, costMeter.seen[1].Value} + assert.ElementsMatch(t, []string{"13", "259"}, values, "cost_meter must see the exact token counts") + for _, kv := range costMeter.seen { + _, err := strconv.Atoi(kv.Value) + assert.NoError(t, err, "values handed to cost_meter must be numeric (regression for missing_tokens)") + } +} + +// TestChain_RunResponse_DetachedContextStillRecords guards the metering +// fix in reverseproxy.go. The response/terminal phase runs after the body +// is forwarded, so a streaming client has usually disconnected by then, +// cancelling its request context. The dispatcher derives each middleware's +// context from the one passed here and short-circuits to fail-mode the +// instant it's Done, which silently drops token/cost metering. The reverse +// proxy now detaches that phase with context.WithoutCancel; this proves a +// context detached from an already-cancelled parent still lets a response +// middleware emit. (The cancelled-parent direction is intentionally not +// asserted: the dispatcher's select over ctx.Done vs the result channel is +// racy when both are ready, which is exactly why the bug was intermittent.) +func TestChain_RunResponse_DetachedContextStillRecords(t *testing.T) { + resp := &fakeMiddleware{ + id: "recorder", + slot: SlotOnResponse, + keys: []string{"llm.input_tokens"}, + emit: []KV{{Key: "llm.input_tokens", Value: "42"}}, + decision: DecisionPassthrough, + } + c := chainFor(t, resp) + + clientCtx, cancel := context.WithCancel(context.Background()) + cancel() // client disconnected after the stream completed + require.Error(t, clientCtx.Err(), "client context must be cancelled for the test to be meaningful") + + detached := context.WithoutCancel(clientCtx) + require.NoError(t, detached.Err(), "detached context must not inherit the client's cancellation") + + acc := NewAccumulator(MaxRequestMetadataBytes) + merged := c.RunResponse(detached, &Input{Slot: SlotOnResponse}, acc) + + var got string + for _, kv := range merged { + if kv.Key == "llm.input_tokens" { + got = kv.Value + } + } + assert.Equal(t, "42", got, "response middleware must still emit token metadata under the detached context") +} + +// TestChain_RunRequest_LatestRewriteWins asserts that when two +// on_request middlewares both emit an UpstreamRewrite, the chain +// returns the value from the later middleware. +func TestChain_RunRequest_LatestRewriteWins(t *testing.T) { + first := &fakeMiddleware{ + id: "first", + slot: SlotOnRequest, + mutationsSupported: true, + canMutate: true, + mutations: &Mutations{RewriteUpstream: &UpstreamRewrite{Scheme: "https", Host: "first.test"}}, + } + second := &fakeMiddleware{ + id: "second", + slot: SlotOnRequest, + mutationsSupported: true, + canMutate: true, + mutations: &Mutations{RewriteUpstream: &UpstreamRewrite{Scheme: "https", Host: "second.test"}}, + } + c := chainFor(t, first, second) + acc := NewAccumulator(0) + + denied, _, rewrite, err := c.RunRequest(context.Background(), nil, &Input{}, acc) + require.NoError(t, err) + assert.Nil(t, denied, "neither middleware denies") + require.NotNil(t, rewrite, "chain must surface the rewrite emitted by the on_request slot") + assert.Equal(t, "https", rewrite.Scheme, "rewrite scheme must come from the later middleware") + assert.Equal(t, "second.test", rewrite.Host, "rewrite host must come from the later middleware (last-write-wins)") +} + +// TestChain_RunRequest_NoRewrite_NilReturn asserts the chain returns a +// nil rewrite when no middleware emits one. +func TestChain_RunRequest_NoRewrite_NilReturn(t *testing.T) { + first := &fakeMiddleware{id: "first", slot: SlotOnRequest} + second := &fakeMiddleware{id: "second", slot: SlotOnRequest} + c := chainFor(t, first, second) + acc := NewAccumulator(0) + + denied, _, rewrite, err := c.RunRequest(context.Background(), nil, &Input{}, acc) + require.NoError(t, err) + assert.Nil(t, denied, "neither middleware denies") + assert.Nil(t, rewrite, "chain must return nil rewrite when no middleware emits one") +} + +// TestChain_ApplyMutations_RewriteGatedOnCanMutate asserts that a +// middleware emitting an UpstreamRewrite with CanMutate=false has its +// rewrite filtered out by the chain. The dispatcher's filterOutput +// already clears Mutations when the gates fail; the chain's defensive +// gate inside mutationRewrite mirrors that contract so a stale +// Mutations field cannot leak through. +func TestChain_ApplyMutations_RewriteGatedOnCanMutate(t *testing.T) { + mw := &fakeMiddleware{ + id: "first", + slot: SlotOnRequest, + mutationsSupported: true, + canMutate: false, + mutations: &Mutations{RewriteUpstream: &UpstreamRewrite{Scheme: "https", Host: "denied.test"}}, + } + c := chainFor(t, mw) + acc := NewAccumulator(0) + + denied, _, rewrite, err := c.RunRequest(context.Background(), nil, &Input{}, acc) + require.NoError(t, err) + assert.Nil(t, denied, "middleware does not deny") + assert.Nil(t, rewrite, "rewrite must be filtered when CanMutate=false") +} + +// TestChain_RunRequest_PropagatesUserGroups asserts the chain forwards +// Input.UserGroups verbatim through cloneInputFor so policy-aware +// middlewares (e.g. llm_policy_check) can authorise without an extra +// management round-trip. +func TestChain_RunRequest_PropagatesUserGroups(t *testing.T) { + groupCapture := &userGroupCaptureMiddleware{ + id: "group-capture", + slot: SlotOnRequest, + } + c := chainFor(t, groupCapture.fake()) + groupCapture.bind(c) + acc := NewAccumulator(0) + + in := &Input{UserGroups: []string{"g1"}} + denied, _, _, err := c.RunRequest(context.Background(), nil, in, acc) + require.NoError(t, err) + assert.Nil(t, denied, "no deny without DecisionDeny") + + require.Len(t, groupCapture.seenGroups, 1, "middleware must observe the caller's UserGroups") + assert.Equal(t, "g1", groupCapture.seenGroups[0], "UserGroups must reach the middleware verbatim") +} + +// userGroupCaptureMiddleware is a fakeMiddleware variant that records +// Input.UserGroups during Invoke. It exists so the cloneInputFor +// behaviour for the new field can be asserted without leaking into +// every other chain test. +type userGroupCaptureMiddleware struct { + id string + slot Slot + seenGroups []string + fakeMW *fakeMiddleware +} + +func (u *userGroupCaptureMiddleware) fake() *fakeMiddleware { + u.fakeMW = &fakeMiddleware{id: u.id, slot: u.slot} + return u.fakeMW +} + +func (u *userGroupCaptureMiddleware) bind(c *Chain) { + for i, bm := range c.all { + if bm.spec.ID != u.id { + continue + } + c.all[i].mw = userGroupRecorder{ + fakeMiddleware: u.fakeMW, + parent: u, + } + } +} + +type userGroupRecorder struct { + *fakeMiddleware + parent *userGroupCaptureMiddleware +} + +func (r userGroupRecorder) Invoke(ctx context.Context, in *Input) (*Output, error) { + r.parent.seenGroups = append([]string(nil), in.UserGroups...) + return r.fakeMiddleware.Invoke(ctx, in) +} + +// TestChain_RunTerminal_SeesAccumulatedMetadata locks that terminal +// middlewares observe the full bag (the caller-supplied in.Metadata +// plus any prior terminal emissions). +func TestChain_RunTerminal_SeesAccumulatedMetadata(t *testing.T) { + first := &fakeMiddleware{ + id: "term-1", + slot: SlotTerminal, + keys: []string{"term.first"}, + emit: []KV{{Key: "term.first", Value: "1"}}, + } + second := &fakeMiddleware{ + id: "term-2", + slot: SlotTerminal, + keys: []string{"term.second"}, + } + c := chainFor(t, first, second) + acc := NewAccumulator(0) + + in := &Input{Metadata: []KV{{Key: "ext.k", Value: "ext"}}} + merged := c.RunTerminal(context.Background(), in, acc) + + require.Len(t, second.seen, 2, "second terminal must see ext bag + first terminal's emission") + got := map[string]string{} + for _, kv := range second.seen { + got[kv.Key] = kv.Value + } + assert.Equal(t, "ext", got["ext.k"], "external bag carries through") + assert.Equal(t, "1", got["term.first"], "first terminal's emission visible to second terminal") + assert.Len(t, merged, 1, "only first terminal emitted; second emitted nothing") +} diff --git a/proxy/internal/middleware/decision.go b/proxy/internal/middleware/decision.go new file mode 100644 index 000000000..0970bdea4 --- /dev/null +++ b/proxy/internal/middleware/decision.go @@ -0,0 +1,81 @@ +package middleware + +import ( + "encoding/json" + "net/http" + "regexp" +) + +var codeRegex = regexp.MustCompile(`^[a-z][a-z0-9._-]{0,63}$`) + +// denyResponse is the on-wire shape rendered by RenderDenyResponse. +// Keeping this as a typed struct ensures we never leak +// middleware-supplied bytes outside known fields. +type denyResponse struct { + Code string `json:"code"` + Message string `json:"message,omitempty"` + Details map[string]string `json:"details,omitempty"` + Middleware string `json:"middleware,omitempty"` +} + +// RenderDenyResponse writes a structured JSON deny body. Status is +// clamped to [400, 499] excluding 401 (to avoid conflicts with the +// proxy's auth flow). All middleware-supplied strings are redacted and +// truncated. On any validation failure the function writes a generic +// 403. +func RenderDenyResponse(w http.ResponseWriter, middlewareID string, reason *DenyReason, defaultStatus int) { + status := clampDenyStatus(defaultStatus) + + if reason == nil || !codeRegex.MatchString(reason.Code) { + writeGenericDeny(w, middlewareID, status) + return + } + + resp := denyResponse{ + Code: reason.Code, + Message: truncate(Scan(reason.Message), 256), + Middleware: truncate(Scan(middlewareID), 64), + } + if n := len(reason.Details); n > 0 { + resp.Details = make(map[string]string, min(n, 8)) + for k, v := range reason.Details { + if len(resp.Details) >= 8 { + break + } + safeKey := truncate(Scan(k), 64) + if safeKey == "" { + continue + } + resp.Details[safeKey] = truncate(Scan(v), 256) + } + } + + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(status) + if err := json.NewEncoder(w).Encode(resp); err != nil { + return + } +} + +func writeGenericDeny(w http.ResponseWriter, middlewareID string, status int) { + w.Header().Set("Content-Type", "application/json; charset=utf-8") + w.WriteHeader(status) + _ = json.NewEncoder(w).Encode(denyResponse{Code: "middleware.error", Middleware: truncate(Scan(middlewareID), 64)}) +} + +func clampDenyStatus(s int) int { + if s < 400 || s >= 500 { + return http.StatusForbidden + } + if s == http.StatusUnauthorized { + return http.StatusForbidden + } + return s +} + +func truncate(s string, n int) string { + if len(s) <= n { + return s + } + return s[:n] +} diff --git a/proxy/internal/middleware/dispatcher.go b/proxy/internal/middleware/dispatcher.go new file mode 100644 index 000000000..316604651 --- /dev/null +++ b/proxy/internal/middleware/dispatcher.go @@ -0,0 +1,189 @@ +package middleware + +import ( + "context" + "errors" + "fmt" + "reflect" + "runtime" + "time" + + log "github.com/sirupsen/logrus" +) + +// Dispatcher reliability kinds reported via +// proxy.middleware.errors_total{kind=...}. +const ( + ErrorKindPanic = "panic" + ErrorKindTimeout = "timeout" + ErrorKindInvokeError = "invoke_error" +) + +// Dispatcher drives a single middleware invocation with panic +// recovery, deadline, and output filtering. Safe for concurrent use. +type Dispatcher struct { + metrics *Metrics + logger *log.Logger +} + +// NewDispatcher returns a dispatcher that emits on the provided +// metrics bundle and logger. A nil metrics bundle falls back to a noop +// instrument set; a nil logger falls back to the standard logger. +func NewDispatcher(metrics *Metrics, logger *log.Logger) *Dispatcher { + if metrics == nil { + metrics, _ = NewMetrics(nil) + } + if logger == nil { + logger = log.StandardLogger() + } + return &Dispatcher{metrics: metrics, logger: logger} +} + +// Invoke runs a single middleware under the reliability wrappers: +// deadline, panic recovery (type + truncated stack only), fail-mode, +// metric emission, and output filtering. The returned output is always +// safe to apply. +func (d *Dispatcher) Invoke(ctx context.Context, spec Spec, mw Middleware, in *Input) (*Output, error) { + if mw == nil { + return nil, fmt.Errorf("middleware %s: instance unavailable", spec.ID) + } + + timeout := clampTimeout(spec.Timeout) + callCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + d.metrics.IncInvocation(ctx, spec.ID) + start := time.Now() + + type result struct { + out *Output + err error + } + ch := make(chan result, 1) + + go func() { + defer func() { + if r := recover(); r != nil { + stack := make([]byte, 4<<10) + n := runtime.Stack(stack, false) + requestID := "" + if in != nil { + requestID = in.RequestID + } + d.logger.Warnf("middleware %s panic: request_id=%s type=%s stack=%s", + spec.ID, requestID, reflect.TypeOf(r).String(), stack[:n]) + ch <- result{err: panicError{msg: fmt.Sprintf("middleware %s panic: %s", spec.ID, reflect.TypeOf(r).String())}} + } + }() + out, err := mw.Invoke(callCtx, in) + ch <- result{out: out, err: err} + }() + + var ( + out *Output + invErr error + kind string + ) + + select { + case <-callCtx.Done(): + invErr = callCtx.Err() + kind = ErrorKindTimeout + case res := <-ch: + out = res.out + invErr = res.err + if invErr != nil { + kind = d.classifyError(invErr) + } + } + + d.metrics.ObserveDuration(ctx, spec.ID, time.Since(start).Milliseconds()) + + if invErr != nil { + d.metrics.IncError(ctx, spec.ID, kind) + return d.failMode(spec, kind), invErr + } + + return d.filterOutput(spec, out), nil +} + +func (d *Dispatcher) classifyError(err error) string { + if err == nil { + return "" + } + if errors.Is(err, context.DeadlineExceeded) { + return ErrorKindTimeout + } + var pe panicError + if errors.As(err, &pe) { + return ErrorKindPanic + } + return ErrorKindInvokeError +} + +// panicError marks an error as coming from the recover branch so the +// classifier can tag it without string inspection. +type panicError struct{ msg string } + +func (p panicError) Error() string { return p.msg } + +// failMode converts an error into a synthesised output per the +// middleware's fail-mode. An mw..error_kind metadata entry is +// attached so operators can alert on error rate even when the +// decision is fail-open. Slot constraints still apply: response and +// terminal slots clamp deny back to passthrough in filterOutput. +func (d *Dispatcher) failMode(spec Spec, kind string) *Output { + meta := []KV{{Key: fmt.Sprintf(KeyFrameworkErrorKindFmt, spec.ID), Value: kind}} + if spec.FailMode == FailClosed && spec.Slot == SlotOnRequest { + return &Output{ + Decision: DecisionDeny, + DenyStatus: 500, + DenyReason: &DenyReason{Code: "middleware.error"}, + Metadata: meta, + } + } + return &Output{Decision: DecisionAllow, Metadata: meta} +} + +// filterOutput applies the output-filter pipeline (slot-aware decision +// clamp, mutations gate) so downstream consumers never see +// middleware-supplied values that violate the contract. Metadata is +// passed through; the Accumulator is the single owner of allowlist + +// caps + redaction (called by Chain). +func (d *Dispatcher) filterOutput(spec Spec, out *Output) *Output { + if out == nil { + return &Output{Decision: DecisionAllow} + } + if spec.Slot != SlotOnRequest && out.Decision == DecisionDeny { + out.Decision = DecisionPassthrough + out.DenyStatus = 0 + out.DenyReason = nil + } + if out.Decision == DecisionDeny { + if out.DenyStatus == 0 { + out.DenyStatus = 403 + } else { + out.DenyStatus = clampDenyStatus(out.DenyStatus) + } + } + if !spec.CanMutate || !spec.MutationsSupported { + out.Mutations = nil + } + if spec.Slot == SlotTerminal { + out.Mutations = nil + } + return out +} + +func clampTimeout(d time.Duration) time.Duration { + if d <= 0 { + return DefaultTimeout + } + if d < MinTimeout { + return MinTimeout + } + if d > MaxTimeout { + return MaxTimeout + } + return d +} diff --git a/proxy/internal/middleware/headerpolicy.go b/proxy/internal/middleware/headerpolicy.go new file mode 100644 index 000000000..d041ad1e1 --- /dev/null +++ b/proxy/internal/middleware/headerpolicy.go @@ -0,0 +1,99 @@ +package middleware + +import "strings" + +var denyHeaders = []string{ + "Authorization", + "Connection", + "Cookie", + "Set-Cookie", + "Forwarded", + "Keep-Alive", + "Proxy-Authorization", + "Proxy-Authenticate", + "Proxy-Connection", + "TE", + "Upgrade", + "Via", + "X-Real-IP", + "X-Request-ID", + "Host", + "Content-Length", + "Transfer-Encoding", + "Trailer", +} + +var denyHeaderPrefixes = []string{ + "X-Authenticated-", + "X-Forwarded-", + "X-Remote-", + "X-NetBird-", +} + +// IsHeaderMutable reports whether a middleware is allowed to mutate +// the named header. The check is case-insensitive and honours both +// exact matches and the compiled-in prefix denylist. +func IsHeaderMutable(name string) bool { + if name == "" { + return false + } + if !isHeaderFieldName(name) { + return false + } + for _, d := range denyHeaders { + if strings.EqualFold(d, name) { + return false + } + } + for _, p := range denyHeaderPrefixes { + if len(name) >= len(p) && strings.EqualFold(name[:len(p)], p) { + return false + } + } + return true +} + +// isHeaderFieldName reports whether name is a valid RFC 7230 header +// field-name (a non-empty token of tchar octets). Rejects names with +// spaces, control characters, or separators that could enable header +// injection or smuggling when applied to the outbound request. +func isHeaderFieldName(name string) bool { + for i := 0; i < len(name); i++ { + c := name[i] + if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') { + continue + } + switch c { + case '!', '#', '$', '%', '&', '\'', '*', '+', '-', '.', '^', '_', '`', '|', '~': + continue + default: + return false + } + } + return true +} + +// FilterHeaderMutations returns the subsets of HeadersAdd and +// HeadersRemove that are safe to apply, plus the list of blocked +// header names so the dispatcher can increment the blocked-header +// metric. +func FilterHeaderMutations(m *Mutations) (filteredAdd []KV, filteredRemove []string, blocked []string) { + if m == nil { + return nil, nil, nil + } + for _, kv := range m.HeadersAdd { + if IsHeaderMutable(kv.Key) { + filteredAdd = append(filteredAdd, kv) + continue + } + blocked = append(blocked, kv.Key) + } + for _, name := range m.HeadersRemove { + if IsHeaderMutable(name) { + filteredRemove = append(filteredRemove, name) + continue + } + blocked = append(blocked, name) + } + return filteredAdd, filteredRemove, blocked +} diff --git a/proxy/internal/middleware/keys.go b/proxy/internal/middleware/keys.go new file mode 100644 index 000000000..9c584ad82 --- /dev/null +++ b/proxy/internal/middleware/keys.go @@ -0,0 +1,86 @@ +package middleware + +// Metadata key namespace constants shared across the built-in +// middlewares. Each domain owns a prefix; middlewares declare their +// per-key allowlist drawn from these constants. Agents implementing +// the G2 middlewares import this file so the dashboard's expanded-row +// viewer and the access-log writer see a stable key surface. +// +// Key shape rules (enforced by the metadata accumulator): +// - Lowercase ASCII letters, digits, dot, underscore, hyphen. +// - At least one dot separating namespace from leaf. +// - Max length: MaxMetadataKeyBytes. +const ( + // LLM request-side metadata (emitted by llm_request_parser). + KeyLLMProvider = "llm.provider" + KeyLLMModel = "llm.model" + KeyLLMStream = "llm.stream" + KeyLLMRequestPromptRaw = "llm.request_prompt_raw" + KeyLLMCaptureTruncated = "llm.capture_truncated" + // KeyLLMSessionID groups requests of the same conversation / coding + // session, read from the per-provider session marker in the request + // body. Empty for clients that don't send one. + KeyLLMSessionID = "llm.session_id" + + // LLM response-side metadata (emitted by llm_response_parser). + //nolint:gosec // metadata key name, not a credential + KeyLLMInputTokens = "llm.input_tokens" + //nolint:gosec // metadata key name, not a credential + KeyLLMOutputTokens = "llm.output_tokens" + //nolint:gosec // metadata key name, not a credential + KeyLLMTotalTokens = "llm.total_tokens" + // LLM cached-input bucket. For OpenAI it's the SUBSET of input + // tokens that hit the prompt cache (prompt_tokens_details. + // cached_tokens) — billed at the cached_input_per_1k rate when + // configured. For Anthropic it's cache_read_input_tokens, which + // is ADDITIVE to llm.input_tokens — billed at cache_read_per_1k. + // cost_meter switches formula on llm.provider. + //nolint:gosec // metadata key name, not a credential + KeyLLMCachedInputTokens = "llm.cached_input_tokens" + // LLM cache-creation bucket (Anthropic only). ADDITIVE to + // llm.input_tokens; billed at cache_creation_per_1k. + //nolint:gosec // metadata key name, not a credential + KeyLLMCacheCreationTokens = "llm.cache_creation_tokens" + KeyLLMResponseCompletion = "llm.response_completion" + + // Guardrail outcomes (emitted by llm_guardrail). The guardrail + // also re-emits llm.request_prompt as a redacted variant of the + // raw prompt and drops llm.request_prompt_raw from the bag. + KeyLLMRequestPrompt = "llm.request_prompt" + KeyLLMPolicyDecision = "llm_policy.decision" + KeyLLMPolicyReason = "llm_policy.reason" + + // LLM router routing decision (emitted by llm_router). The router + // stamps the resolved provider id so downstream middlewares and + // the access-log emitter can attribute the request without + // re-parsing the body. + KeyLLMResolvedProviderID = "llm.resolved_provider_id" + + // LLM authorising groups for this request (emitted by llm_router + // on the allow path). Carries the comma-separated intersection of + // the caller's UserGroups with the resolved route's + // AllowedGroupIDs — i.e. the groups that actually authorise this + // specific request, NOT every group the peer happens to be in. + // Identity-stamping middlewares use this for per-request tag + // attribution so unrelated group memberships don't leak into + // downstream gateways' spend logs. + KeyLLMAuthorisingGroups = "llm.authorising_groups" + + // LLM policy attribution (emitted by llm_limit_check on the allow + // path). Names the policy that paid for this request and the + // dimension counters the post-flight llm_limit_record middleware + // must tick. Empty when no applicable policy has any caps + // configured (catch-all-allow attribution). + KeyLLMSelectedPolicyID = "llm.selected_policy_id" + KeyLLMAttributionGroupID = "llm.attribution_group_id" + KeyLLMAttributionWindowS = "llm.attribution_window_seconds" + + // Cost metering (emitted by cost_meter). + KeyCostUSDTotal = "cost.usd_total" + KeyCostSkipped = "cost.skipped" + + // Framework-emitted error markers. Use the mw..* prefix to + // distinguish framework-injected entries from middleware-emitted + // metadata. + KeyFrameworkErrorKindFmt = "mw.%s.error_kind" +) diff --git a/proxy/internal/middleware/manager.go b/proxy/internal/middleware/manager.go new file mode 100644 index 000000000..9b22edeff --- /dev/null +++ b/proxy/internal/middleware/manager.go @@ -0,0 +1,412 @@ +package middleware + +import ( + "context" + "fmt" + "strings" + "sync" + "sync/atomic" + "time" + + log "github.com/sirupsen/logrus" + + "github.com/netbirdio/netbird/proxy/internal/middleware/bodytap" +) + +// chainCloseTimeout bounds how long closeChainsAsync waits for an +// individual chain to drain before forcing teardown. Set to 2x +// MaxTimeout so a middleware blocked on the dispatcher's per-Invoke +// deadline always wins; anything running longer is a runaway and gets +// force-closed. +const chainCloseTimeout = 2 * MaxTimeout + +// PathTargetBinding is the minimal per-path binding the server passes +// to Rebuild. It carries the stable keys Manager uses for snapshot +// lookups plus the validated middleware spec list for that path. +type PathTargetBinding struct { + ServiceID string + PathID string + Specs []Spec +} + +// LiveServiceCheck reports whether the given service ID is still +// present in the proxy's live mapping cache. The Manager calls it +// during InvalidateMiddleware so a chain whose service has been +// removed since the last Rebuild is not resurrected from the binding +// cache, closing the auth-revocation race. +type LiveServiceCheck func(serviceID string) bool + +// chainTable holds the immutable per-target chain snapshot. It is +// cloned into a new instance on every Rebuild and swapped in via +// atomic.Pointer. The reverse index byMiddleware lets +// InvalidateMiddleware find the chain keys that reference a given +// middleware without scanning the whole table. +type chainTable struct { + byTarget map[string]*Chain + byMiddleware map[string]map[string]struct{} +} + +func newChainTable() *chainTable { + return &chainTable{ + byTarget: make(map[string]*Chain), + byMiddleware: make(map[string]map[string]struct{}), + } +} + +func (c *chainTable) clone() *chainTable { + out := newChainTable() + for k, v := range c.byTarget { + out.byTarget[k] = v + } + for id, keys := range c.byMiddleware { + set := make(map[string]struct{}, len(keys)) + for k := range keys { + set[k] = struct{}{} + } + out.byMiddleware[id] = set + } + return out +} + +func (c *chainTable) addChain(key string, ch *Chain) { + c.byTarget[key] = ch + if ch == nil { + return + } + for _, bm := range ch.all { + set, ok := c.byMiddleware[bm.spec.ID] + if !ok { + set = make(map[string]struct{}) + c.byMiddleware[bm.spec.ID] = set + } + set[key] = struct{}{} + } +} + +func (c *chainTable) removeChain(key string) (*Chain, []string) { + ch, ok := c.byTarget[key] + if !ok { + return nil, nil + } + delete(c.byTarget, key) + if ch == nil { + return nil, nil + } + ids := make([]string, 0, len(ch.all)) + for _, bm := range ch.all { + ids = append(ids, bm.spec.ID) + set, ok := c.byMiddleware[bm.spec.ID] + if !ok { + continue + } + delete(set, key) + if len(set) == 0 { + delete(c.byMiddleware, bm.spec.ID) + } + } + return ch, ids +} + +// Manager owns the per-target middleware chains, the global capture +// budget, and the shared dispatcher. Readers (ChainFor) are lock-free; +// writers (Rebuild, Invalidate*) serialise on writeMu so two +// concurrent mapping updates do not lose writes. +type Manager struct { + writeMu sync.Mutex + chains atomic.Pointer[chainTable] + budget bodytap.Budget + metrics *Metrics + logger *log.Logger + dispatcher *Dispatcher + resolver *Resolver + lastBindings map[string]PathTargetBinding + liveServiceCheck atomic.Pointer[LiveServiceCheck] +} + +// NewManager constructs a Manager with the given capture budget size. +// A zero or negative budget falls back to bodytap.DefaultCaptureBudgetBytes. +func NewManager(budgetBytes int64, metrics *Metrics, logger *log.Logger) *Manager { + if metrics == nil { + metrics, _ = NewMetrics(nil) + } + if logger == nil { + logger = log.StandardLogger() + } + if budgetBytes <= 0 { + budgetBytes = bodytap.DefaultCaptureBudgetBytes + } + m := &Manager{ + budget: bodytap.NewBudget(budgetBytes), + metrics: metrics, + logger: logger, + dispatcher: NewDispatcher(metrics, logger), + lastBindings: make(map[string]PathTargetBinding), + } + m.chains.Store(newChainTable()) + return m +} + +// SetResolver installs the resolver used by Rebuild. Safe to call +// once at boot before any Rebuild; not safe to swap concurrently. +func (m *Manager) SetResolver(r *Resolver) { + m.resolver = r +} + +// SetLiveServiceCheck installs a callback the Manager uses to confirm +// a service ID still maps to a live mapping before resurrecting its +// chain from the binding cache during InvalidateMiddleware. A nil fn +// disables the check. +func (m *Manager) SetLiveServiceCheck(fn LiveServiceCheck) { + if fn == nil { + m.liveServiceCheck.Store(nil) + return + } + m.liveServiceCheck.Store(&fn) +} + +// Budget returns the shared capture budget. +func (m *Manager) Budget() bodytap.Budget { + return m.budget +} + +// Metrics returns the shared metrics bundle. +func (m *Manager) Metrics() *Metrics { + return m.metrics +} + +// Dispatcher returns the shared dispatcher (primarily for testing). +func (m *Manager) Dispatcher() *Dispatcher { + return m.dispatcher +} + +// Rebuild replaces every chain keyed by serviceID with the provided +// bindings. Entries for other services are preserved. Replaced chains +// are closed asynchronously after the atomic swap so in-flight +// requests against the previous chain finish before middleware +// resources are released. +func (m *Manager) Rebuild(serviceID string, bindings []PathTargetBinding) error { + m.writeMu.Lock() + defer m.writeMu.Unlock() + + cur := m.chains.Load() + next := cur.clone() + + prefix := serviceID + "|" + var retired []*Chain + for k := range cur.byTarget { + if !strings.HasPrefix(k, prefix) { + continue + } + ch, _ := next.removeChain(k) + if ch != nil { + retired = append(retired, ch) + } + delete(m.lastBindings, k) + } + + for _, b := range bindings { + if b.ServiceID != serviceID { + return fmt.Errorf("binding service %q does not match rebuild service %q", b.ServiceID, serviceID) + } + key := chainKey(b.ServiceID, b.PathID) + m.lastBindings[key] = cloneBinding(b) + chain := m.buildChain(b) + if chain == nil || chain.Empty() { + delete(m.lastBindings, key) + continue + } + next.addChain(key, chain) + } + + m.chains.Store(next) + m.closeChainsAsync(retired) + return nil +} + +// Invalidate drops every chain for the given service ID. +func (m *Manager) Invalidate(serviceID string) { + m.writeMu.Lock() + defer m.writeMu.Unlock() + cur := m.chains.Load() + next := cur.clone() + prefix := serviceID + "|" + var retired []*Chain + for k := range cur.byTarget { + if !strings.HasPrefix(k, prefix) { + continue + } + ch, _ := next.removeChain(k) + if ch != nil { + retired = append(retired, ch) + } + delete(m.lastBindings, k) + } + for k := range m.lastBindings { + if strings.HasPrefix(k, prefix) { + delete(m.lastBindings, k) + } + } + m.chains.Store(next) + m.closeChainsAsync(retired) +} + +// InvalidateMiddleware rebuilds only the chains that reference id. +func (m *Manager) InvalidateMiddleware(id string) { + if id == "" { + return + } + m.writeMu.Lock() + defer m.writeMu.Unlock() + + cur := m.chains.Load() + keys, ok := cur.byMiddleware[id] + if !ok || len(keys) == 0 { + return + } + + affected := make([]string, 0, len(keys)) + for k := range keys { + affected = append(affected, k) + } + + next := cur.clone() + var retired []*Chain + check := m.loadLiveServiceCheck() + for _, k := range affected { + ch, _ := next.removeChain(k) + if ch != nil { + retired = append(retired, ch) + } + b, ok := m.lastBindings[k] + if !ok { + delete(m.lastBindings, k) + continue + } + if check != nil && !check(b.ServiceID) { + m.logger.Debugf("middleware %s: skipping rebuild for %s; service no longer live", id, k) + delete(m.lastBindings, k) + continue + } + chain := m.buildChain(b) + if chain == nil || chain.Empty() { + delete(m.lastBindings, k) + continue + } + next.addChain(k, chain) + } + + m.chains.Store(next) + m.closeChainsAsync(retired) +} + +func (m *Manager) loadLiveServiceCheck() LiveServiceCheck { + p := m.liveServiceCheck.Load() + if p == nil { + return nil + } + return *p +} + +// InvalidateAll drops every chain. +func (m *Manager) InvalidateAll() { + m.writeMu.Lock() + defer m.writeMu.Unlock() + cur := m.chains.Load() + retired := make([]*Chain, 0, len(cur.byTarget)) + for _, c := range cur.byTarget { + retired = append(retired, c) + } + m.chains.Store(newChainTable()) + for k := range m.lastBindings { + delete(m.lastBindings, k) + } + m.closeChainsAsync(retired) +} + +func (m *Manager) closeChainsAsync(retired []*Chain) { + if len(retired) == 0 { + return + } + chains := make([]*Chain, len(retired)) + copy(chains, retired) + go func() { + for _, c := range chains { + ctx, cancel := context.WithTimeout(context.Background(), chainCloseTimeout) + start := time.Now() + if err := c.Close(ctx); err != nil { + if m.metrics != nil { + m.metrics.IncError(context.Background(), c.TargetID(), "chain_close_timeout") + } + m.logger.Warnf("middleware chain %s close exceeded %s after %s: %v", + c.TargetID(), chainCloseTimeout, time.Since(start), err) + } + cancel() + } + }() +} + +// ChainFor returns the chain for serviceID/pathID or nil if none is +// registered. Lock-free. +func (m *Manager) ChainFor(serviceID, pathID string) *Chain { + tbl := m.chains.Load() + if tbl == nil { + return nil + } + c, ok := tbl.byTarget[chainKey(serviceID, pathID)] + if !ok { + return nil + } + return c +} + +// buildChain resolves each enabled spec and returns the assembled +// chain. Returns a nil chain when no middlewares are bound; resolver +// errors per middleware are logged and counted but do not abort the +// chain. +func (m *Manager) buildChain(b PathTargetBinding) *Chain { + if len(b.Specs) == 0 || m.resolver == nil { + return nil + } + + bound := make([]boundMiddleware, 0, len(b.Specs)) + for _, spec := range b.Specs { + if !spec.Enabled { + continue + } + mw, merged, err := m.resolver.Resolve(spec) + if err != nil { + m.logger.Warnf("middleware %s resolve on target %s/%s: %v", spec.ID, b.ServiceID, b.PathID, err) + m.metrics.IncError(context.Background(), spec.ID, "resolve_error") + continue + } + if mw == nil { + continue + } + bound = append(bound, boundMiddleware{spec: merged, mw: mw}) + } + if len(bound) == 0 { + return nil + } + return NewChain(chainKey(b.ServiceID, b.PathID), bound, m.dispatcher) +} + +// cloneBinding returns a deep copy of b suitable for caching across +// mapping updates. +func cloneBinding(b PathTargetBinding) PathTargetBinding { + out := PathTargetBinding{ + ServiceID: b.ServiceID, + PathID: b.PathID, + } + if len(b.Specs) == 0 { + return out + } + out.Specs = make([]Spec, len(b.Specs)) + for i, s := range b.Specs { + out.Specs[i] = s.Clone() + } + return out +} + +func chainKey(serviceID, pathID string) string { + return serviceID + "|" + pathID +} diff --git a/proxy/internal/middleware/metadata.go b/proxy/internal/middleware/metadata.go new file mode 100644 index 000000000..576c379ec --- /dev/null +++ b/proxy/internal/middleware/metadata.go @@ -0,0 +1,99 @@ +package middleware + +import "regexp" + +// keyRegex constrains metadata keys to the cross-domain shape +// described in keys.go. At least one dot, lowercase ASCII / digits / +// dot / underscore / hyphen only, length within MaxMetadataKeyBytes. +var keyRegex = regexp.MustCompile(`^[a-z][a-z0-9_-]*(\.[a-z0-9][a-z0-9_-]*)+$`) + +// MetadataRejection describes a single rejected key/value so the +// dispatcher can emit per-reason counter increments. +type MetadataRejection struct { + Key string + Reason string +} + +// Rejection reasons reported by Accumulator.Emit. +const ( + MetadataReasonBadKey = "bad_key" + MetadataReasonNotAllowlisted = "not_allowlisted" + MetadataReasonKeyTooLong = "key_too_long" + MetadataReasonValueTooLong = "value_too_long" + MetadataReasonMiddlewareCap = "middleware_cap" + MetadataReasonRequestCap = "request_cap" +) + +// Accumulator enforces per-middleware and per-request metadata caps. +// Not safe for concurrent use; callers hold one inside a single chain +// execution. +type Accumulator struct { + perMiddlewareUsed map[string]int + totalUsed int + maxPerRequest int +} + +// NewAccumulator returns an accumulator configured for the per-request +// total cap. A maxPerRequest of zero means use MaxRequestMetadataBytes. +func NewAccumulator(maxPerRequest int) *Accumulator { + if maxPerRequest <= 0 { + maxPerRequest = MaxRequestMetadataBytes + } + return &Accumulator{ + perMiddlewareUsed: make(map[string]int), + maxPerRequest: maxPerRequest, + } +} + +// Emit validates the candidate metadata against the middleware's +// allowlist and the global caps, redacts each accepted value, and +// returns the accepted entries plus any rejections for metric emission. +func (a *Accumulator) Emit(middlewareID string, allow []string, out []KV) ([]KV, []MetadataRejection) { + if len(out) == 0 { + return nil, nil + } + allowSet := make(map[string]struct{}, len(allow)) + for _, k := range allow { + allowSet[k] = struct{}{} + } + + accepted := make([]KV, 0, len(out)) + var rejected []MetadataRejection + + for _, kv := range out { + if len(kv.Key) == 0 || len(kv.Key) > MaxMetadataKeyBytes { + rejected = append(rejected, MetadataRejection{Key: kv.Key, Reason: MetadataReasonKeyTooLong}) + continue + } + if !keyRegex.MatchString(kv.Key) { + rejected = append(rejected, MetadataRejection{Key: kv.Key, Reason: MetadataReasonBadKey}) + continue + } + if _, ok := allowSet[kv.Key]; !ok { + rejected = append(rejected, MetadataRejection{Key: kv.Key, Reason: MetadataReasonNotAllowlisted}) + continue + } + if len(kv.Value) > MaxMetadataValueBytes { + rejected = append(rejected, MetadataRejection{Key: kv.Key, Reason: MetadataReasonValueTooLong}) + continue + } + + redacted := Scan(kv.Value) + cost := len(kv.Key) + len(redacted) + + if a.perMiddlewareUsed[middlewareID]+cost > MaxMiddlewareMetadataBytes { + rejected = append(rejected, MetadataRejection{Key: kv.Key, Reason: MetadataReasonMiddlewareCap}) + continue + } + if a.totalUsed+cost > a.maxPerRequest { + rejected = append(rejected, MetadataRejection{Key: kv.Key, Reason: MetadataReasonRequestCap}) + continue + } + + a.perMiddlewareUsed[middlewareID] += cost + a.totalUsed += cost + accepted = append(accepted, KV{Key: kv.Key, Value: redacted}) + } + + return accepted, rejected +} diff --git a/proxy/internal/middleware/metrics.go b/proxy/internal/middleware/metrics.go new file mode 100644 index 000000000..73745a86b --- /dev/null +++ b/proxy/internal/middleware/metrics.go @@ -0,0 +1,171 @@ +package middleware + +import ( + "context" + + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + "go.opentelemetry.io/otel/metric/noop" +) + +// Metrics is the bundle of OTel instruments emitted by the middleware +// dispatcher. The constructor falls back to a noop meter when given +// nil so tests can skip metrics wiring entirely. +type Metrics struct { + requestsTotal metric.Int64Counter + durationMs metric.Int64Histogram + invocationsTotal metric.Int64Counter + errorsTotal metric.Int64Counter + metadataRejectedTotal metric.Int64Counter + headerMutationBlocked metric.Int64Counter + captureBypassTotal metric.Int64Counter +} + +// NewMetrics registers the proxy.middleware.* instruments on the +// given meter. A nil meter is treated as the global no-op provider. +func NewMetrics(meter metric.Meter) (*Metrics, error) { + if meter == nil { + meter = noop.NewMeterProvider().Meter("proxy.middleware.noop") + } + + m := &Metrics{} + var err error + + m.requestsTotal, err = meter.Int64Counter( + "proxy.middleware.requests_total", + metric.WithUnit("1"), + metric.WithDescription("Middleware invocations grouped by outcome"), + ) + if err != nil { + return nil, err + } + + m.durationMs, err = meter.Int64Histogram( + "proxy.middleware.duration_ms", + metric.WithUnit("milliseconds"), + metric.WithDescription("Middleware Invoke latency"), + ) + if err != nil { + return nil, err + } + + m.invocationsTotal, err = meter.Int64Counter( + "proxy.middleware.invocations_total", + metric.WithUnit("1"), + metric.WithDescription("Middleware Invoke heartbeat counter"), + ) + if err != nil { + return nil, err + } + + m.errorsTotal, err = meter.Int64Counter( + "proxy.middleware.errors_total", + metric.WithUnit("1"), + metric.WithDescription("Middleware errors grouped by kind"), + ) + if err != nil { + return nil, err + } + + m.metadataRejectedTotal, err = meter.Int64Counter( + "proxy.middleware.metadata_rejected_total", + metric.WithUnit("1"), + metric.WithDescription("Middleware metadata entries rejected by the allowlist/caps"), + ) + if err != nil { + return nil, err + } + + m.headerMutationBlocked, err = meter.Int64Counter( + "proxy.middleware.header_mutation_blocked_total", + metric.WithUnit("1"), + metric.WithDescription("Middleware header mutations dropped by the denylist"), + ) + if err != nil { + return nil, err + } + + m.captureBypassTotal, err = meter.Int64Counter( + "proxy.middleware.capture_bypass_total", + metric.WithUnit("1"), + metric.WithDescription("Capture bypasses grouped by reason"), + ) + if err != nil { + return nil, err + } + + return m, nil +} + +// IncRequest increments proxy.middleware.requests_total with the +// middleware, target, and outcome labels. +func (m *Metrics) IncRequest(ctx context.Context, middlewareID, targetID, outcome string) { + if m == nil { + return + } + m.requestsTotal.Add(ctx, 1, metric.WithAttributes( + attribute.String("middleware", middlewareID), + attribute.String("target_id", targetID), + attribute.String("outcome", outcome), + )) +} + +// ObserveDuration records the middleware Invoke latency in milliseconds. +func (m *Metrics) ObserveDuration(ctx context.Context, middlewareID string, ms int64) { + if m == nil { + return + } + m.durationMs.Record(ctx, ms, metric.WithAttributes(attribute.String("middleware", middlewareID))) +} + +// IncInvocation increments the heartbeat counter regardless of outcome. +func (m *Metrics) IncInvocation(ctx context.Context, middlewareID string) { + if m == nil { + return + } + m.invocationsTotal.Add(ctx, 1, metric.WithAttributes(attribute.String("middleware", middlewareID))) +} + +// IncError increments the error counter with the given failure kind label. +func (m *Metrics) IncError(ctx context.Context, middlewareID, kind string) { + if m == nil { + return + } + m.errorsTotal.Add(ctx, 1, metric.WithAttributes( + attribute.String("middleware", middlewareID), + attribute.String("kind", kind), + )) +} + +// IncMetadataRejected increments the rejected-metadata counter for a reason. +func (m *Metrics) IncMetadataRejected(ctx context.Context, middlewareID, reason string) { + if m == nil { + return + } + m.metadataRejectedTotal.Add(ctx, 1, metric.WithAttributes( + attribute.String("middleware", middlewareID), + attribute.String("reason", reason), + )) +} + +// IncHeaderMutationBlocked increments the blocked-header counter. +func (m *Metrics) IncHeaderMutationBlocked(ctx context.Context, middlewareID, header string) { + if m == nil { + return + } + m.headerMutationBlocked.Add(ctx, 1, metric.WithAttributes( + attribute.String("middleware", middlewareID), + attribute.String("header", header), + )) +} + +// IncCaptureBypass increments the capture-bypass counter for a reason. +func (m *Metrics) IncCaptureBypass(ctx context.Context, targetID, reason string) { + if m == nil { + return + } + m.captureBypassTotal.Add(ctx, 1, metric.WithAttributes( + attribute.String("target_id", targetID), + attribute.String("reason", reason), + )) +} diff --git a/proxy/internal/middleware/middleware.go b/proxy/internal/middleware/middleware.go new file mode 100644 index 000000000..16d398d2c --- /dev/null +++ b/proxy/internal/middleware/middleware.go @@ -0,0 +1,47 @@ +package middleware + +import "context" + +// Middleware is the surface exposed by each concrete implementation. +// The Manager invokes it through the Dispatcher, passing a cloned +// Input. Each middleware lives in exactly one Slot. +// +// Close releases any resources owned by the middleware instance +// (background goroutines, file handles). It is invoked when the chain +// holding the middleware is replaced or torn down. Implementations +// must be idempotent and safe to call after construction even when +// Invoke was never called. +type Middleware interface { + ID() string + Version() string + Slot() Slot + + // AcceptedContentTypes lists the request/response content types + // the middleware needs the body for. Empty slice means the + // middleware does not inspect the body. + AcceptedContentTypes() []string + + // MetadataKeys is the closed set of metadata keys this middleware + // may emit. The accumulator drops anything outside this allowlist. + MetadataKeys() []string + + // MutationsSupported reports whether the middleware may emit + // header / body mutations. A spec with CanMutate=true is honoured + // only when the implementation also supports mutations. + MutationsSupported() bool + + Invoke(ctx context.Context, in *Input) (*Output, error) + + Close() error +} + +// Factory builds a configured Middleware instance from raw config +// bytes shipped on the wire. Each registered middleware ID has a +// single factory in the registry. Factory.New returns an error when +// the config is malformed or violates a per-middleware invariant; the +// chain build path logs the error, increments the resolve_error metric, +// and skips the middleware. +type Factory interface { + ID() string + New(rawConfig []byte) (Middleware, error) +} diff --git a/proxy/internal/middleware/redaction.go b/proxy/internal/middleware/redaction.go new file mode 100644 index 000000000..ebbe90c61 --- /dev/null +++ b/proxy/internal/middleware/redaction.go @@ -0,0 +1,79 @@ +package middleware + +import ( + "regexp" + "strings" +) + +// Redaction scope: Scan handles the narrow, high-signal set of +// secrets we are comfortable masking with a regex. The intent is +// "make accidental leaks impossible to miss at a glance", not "be a +// DLP product". Contributors adding more patterns should weigh false +// positives carefully — a metadata value that over-redacts benign +// strings is strictly worse than one that misses a rare format. +var ( + jwtRegex = regexp.MustCompile(`eyJ[A-Za-z0-9_-]{5,}\.[A-Za-z0-9_-]{5,}\.[A-Za-z0-9_-]{5,}`) + pemRegex = regexp.MustCompile(`-----BEGIN [A-Z ]+-----[\s\S]*?-----END [A-Z ]+-----`) + awsKeyRegex = regexp.MustCompile(`AKIA[0-9A-Z]{16}`) + bearerRegex = regexp.MustCompile(`(?i)\b(?:bearer|token|api[_-]?key|authorization)[\s:=]+([A-Za-z0-9_\-\.]{40,})`) + ccCandidateRgx = regexp.MustCompile(`\b(?:\d[ -]?){13,19}\b`) +) + +// Scan redacts high-signal secret patterns from value. Matches are +// replaced with `[REDACTED:]`. Non-matching input is returned +// unchanged. +func Scan(value string) string { + if value == "" { + return value + } + result := value + result = pemRegex.ReplaceAllString(result, "[REDACTED:pem]") + result = jwtRegex.ReplaceAllString(result, "[REDACTED:jwt]") + result = awsKeyRegex.ReplaceAllString(result, "[REDACTED:aws_key]") + result = bearerRegex.ReplaceAllStringFunc(result, func(match string) string { + sub := bearerRegex.FindStringSubmatch(match) + if len(sub) < 2 { + return "[REDACTED:bearer]" + } + return strings.Replace(match, sub[1], "[REDACTED:bearer]", 1) + }) + result = ccCandidateRgx.ReplaceAllStringFunc(result, func(match string) string { + digits := stripNonDigits(match) + if len(digits) < 13 || len(digits) > 19 { + return match + } + if !luhn(digits) { + return match + } + return "[REDACTED:cc]" + }) + return result +} + +func stripNonDigits(s string) string { + var b strings.Builder + b.Grow(len(s)) + for _, r := range s { + if r >= '0' && r <= '9' { + b.WriteRune(r) + } + } + return b.String() +} + +func luhn(digits string) bool { + sum := 0 + alt := false + for i := len(digits) - 1; i >= 0; i-- { + n := int(digits[i] - '0') + if alt { + n *= 2 + if n > 9 { + n -= 9 + } + } + sum += n + alt = !alt + } + return sum%10 == 0 +} diff --git a/proxy/internal/middleware/registry.go b/proxy/internal/middleware/registry.go new file mode 100644 index 000000000..c46552162 --- /dev/null +++ b/proxy/internal/middleware/registry.go @@ -0,0 +1,121 @@ +package middleware + +import ( + "fmt" + "sync" +) + +// Registry maps middleware IDs to their factories. The proxy installs +// a single Registry at boot; concrete middlewares register themselves +// from init() functions inside their own packages so the boot wiring +// only needs an anonymous import. +// +// Registry is safe for concurrent reads after boot. Register / Unregister +// take the write lock; Get and IDs take the read lock. +type Registry struct { + mu sync.RWMutex + factories map[string]Factory +} + +// NewRegistry returns an empty registry. +func NewRegistry() *Registry { + return &Registry{factories: make(map[string]Factory)} +} + +// Register installs the factory under its ID. Returns an error when an +// ID is already registered — collisions are programmer errors and must +// be visible at boot rather than silently last-write-wins. +func (r *Registry) Register(f Factory) error { + if f == nil { + return fmt.Errorf("middleware registry: nil factory") + } + id := f.ID() + if id == "" { + return fmt.Errorf("middleware registry: factory has empty id") + } + r.mu.Lock() + defer r.mu.Unlock() + if _, exists := r.factories[id]; exists { + return fmt.Errorf("middleware registry: %q already registered", id) + } + r.factories[id] = f + return nil +} + +// MustRegister panics on error. Intended for init() registration so +// duplicate IDs surface at startup. +func (r *Registry) MustRegister(f Factory) { + if err := r.Register(f); err != nil { + panic(err) + } +} + +// Get returns the factory for id, or nil when no factory is +// registered. +func (r *Registry) Get(id string) Factory { + r.mu.RLock() + defer r.mu.RUnlock() + return r.factories[id] +} + +// IDs returns the registered IDs in unspecified order. Used by the +// management translator to reject specs that reference unknown IDs at +// apply time. +func (r *Registry) IDs() []string { + r.mu.RLock() + defer r.mu.RUnlock() + out := make([]string, 0, len(r.factories)) + for id := range r.factories { + out = append(out, id) + } + return out +} + +// IsKnown reports whether id has a registered factory. +func (r *Registry) IsKnown(id string) bool { + return r.Get(id) != nil +} + +// Resolver wraps a Registry and produces a configured Middleware +// instance from a Spec. The Manager uses this during chain build. +type Resolver struct { + registry *Registry +} + +// NewResolver returns a resolver backed by the registry. +func NewResolver(registry *Registry) *Resolver { + if registry == nil { + registry = NewRegistry() + } + return &Resolver{registry: registry} +} + +// Resolve builds a Middleware instance and merges runtime-only fields +// (version, accepted content types, metadata key allowlist, mutation +// support) onto the spec. +// +// Return semantics: +// - (mw, mergedSpec, nil): instance built, include in chain. +// - (nil, spec, nil): id not registered; silently skip. +// - (nil, spec, err): factory rejected the config (logged + counted +// by Manager, other middlewares still bind). +func (r *Resolver) Resolve(spec Spec) (Middleware, Spec, error) { + f := r.registry.Get(spec.ID) + if f == nil { + return nil, spec, nil + } + mw, err := f.New(spec.RawConfig) + if err != nil { + return nil, spec, fmt.Errorf("middleware %s factory: %w", spec.ID, err) + } + if mw.Slot() != spec.Slot { + _ = mw.Close() + return nil, spec, fmt.Errorf("middleware %s slot mismatch: spec=%d impl=%d", spec.ID, spec.Slot, mw.Slot()) + } + merged := spec + merged.Version = mw.Version() + merged.MetadataKeys = append([]string(nil), mw.MetadataKeys()...) + merged.AcceptedContentTypes = append([]string(nil), mw.AcceptedContentTypes()...) + merged.MutationsSupported = mw.MutationsSupported() + return mw, merged, nil +} diff --git a/proxy/internal/middleware/spec.go b/proxy/internal/middleware/spec.go new file mode 100644 index 000000000..a154ceca2 --- /dev/null +++ b/proxy/internal/middleware/spec.go @@ -0,0 +1,44 @@ +package middleware + +import "time" + +// Spec is the apply-time, validated representation of a per-target +// middleware configuration merged with the runtime-only fields +// compiled into the middleware implementation. +// +// The wire shape is RawConfig (JSON bytes) instead of the older +// params map[string]string. Each middleware unmarshals RawConfig into +// its own typed config struct, surfacing structural validation errors +// at construction rather than per-invocation lookups. +type Spec struct { + ID string + Slot Slot + Version string + Enabled bool + FailMode FailMode + Timeout time.Duration + RawConfig []byte + CanMutate bool + + // Runtime-only fields populated from the registered middleware at + // chain build time; not sourced from proto. + MetadataKeys []string + AcceptedContentTypes []string + MutationsSupported bool +} + +// Clone returns a deep copy of the spec safe to cache across mapping +// updates. +func (s Spec) Clone() Spec { + out := s + if len(s.RawConfig) > 0 { + out.RawConfig = append([]byte(nil), s.RawConfig...) + } + if len(s.MetadataKeys) > 0 { + out.MetadataKeys = append([]string(nil), s.MetadataKeys...) + } + if len(s.AcceptedContentTypes) > 0 { + out.AcceptedContentTypes = append([]string(nil), s.AcceptedContentTypes...) + } + return out +} diff --git a/proxy/internal/middleware/types.go b/proxy/internal/middleware/types.go new file mode 100644 index 000000000..1b49e6159 --- /dev/null +++ b/proxy/internal/middleware/types.go @@ -0,0 +1,253 @@ +// Package middleware defines the per-target middleware chain that runs +// inside the reverse proxy hot path. It is the only chain wired into +// the request path. +// +// Concepts: +// - Slot: the position a middleware occupies in the chain. A +// middleware lives in exactly one slot — separate concerns become +// separate middlewares. +// - Decision: the on_request slot can DENY; on_response and terminal +// slots can only PASSTHROUGH. The dispatcher clamps decisions that +// violate this contract. +// - Metadata: the only side-channel between middlewares. Each +// middleware declares an allowlist of keys it may emit; the merger +// enforces caps and namespace rules. +package middleware + +import "time" + +// Slot identifies where in the request lifecycle a middleware runs. +// A middleware declares a single slot. Splitting per-purpose work +// (request parsing vs response parsing vs cost metering) into separate +// slot-keyed middlewares is the explicit architectural choice for the +// agent-network use case; no middleware participates in more than one +// slot. +type Slot int + +const ( + // SlotOnRequest runs before the upstream call. Middlewares in this + // slot may DENY the request, mutate headers/body (when permitted), + // and emit metadata derived from the request envelope. + SlotOnRequest Slot = 1 + // SlotOnResponse runs after the upstream returns. Middlewares in + // this slot observe the response, emit metadata, and may mutate + // response headers when permitted. They cannot DENY. + SlotOnResponse Slot = 2 + // SlotTerminal runs after every SlotOnResponse middleware has + // emitted. Terminal middlewares observe the full metadata bag and + // ship it to external sinks (access log, metrics export). They + // cannot DENY and cannot mutate the response. + SlotTerminal Slot = 3 +) + +// FailMode controls how the dispatcher reacts when a middleware +// returns an error, times out, or panics. Observer middlewares default +// to FailOpen; policy middlewares should default to FailClosed. +type FailMode int + +const ( + // FailOpen allows the request to proceed when a middleware fails. + FailOpen FailMode = 0 + // FailClosed denies the request when a middleware fails. Only + // meaningful for SlotOnRequest middlewares. + FailClosed FailMode = 1 +) + +// Decision captures the outcome of a middleware invocation as observed +// by the dispatcher. Response-phase middlewares always return +// DecisionPassthrough; the dispatcher clamps any other value. +type Decision int + +const ( + // DecisionAllow lets the request proceed. + DecisionAllow Decision = 0 + // DecisionDeny stops the chain and returns a rendered deny + // response. Only honoured in SlotOnRequest. + DecisionDeny Decision = 1 + // DecisionPassthrough is the response-phase neutral outcome. + DecisionPassthrough Decision = 2 +) + +// Resource limits enforced by the proxy at config apply time and by +// the dispatcher at runtime. Per-target values supplied by management +// are clamped to these bounds. +const ( + // MaxBodyCapBytes is the proxy-wide upper bound for per-direction + // body capture. Sized to hold a full LLM streaming response (token + // usage rides the trailing SSE event, so the captured prefix must + // reach the end of the stream); a single response is bounded by the + // model's max output tokens, so this is a real ceiling, not a + // treadmill. Request capture stays well under this — oversized + // requests use the tolerant routing scan instead of buffering. + MaxBodyCapBytes int64 = 8 << 20 + // MinTimeout is the proxy-wide lower bound for per-middleware + // Invoke timeouts. + MinTimeout = 10 * time.Millisecond + // MaxTimeout is the proxy-wide upper bound for per-middleware + // Invoke timeouts. + MaxTimeout = 5 * time.Second + // DefaultTimeout is used when the per-target timeout is zero or + // unset. + DefaultTimeout = 500 * time.Millisecond + + // MaxMiddlewareMetadataBytes is the per-middleware metadata total + // cap. + MaxMiddlewareMetadataBytes = 16 << 10 + // MaxRequestMetadataBytes is the per-request metadata total cap + // across all middlewares in the chain. Earlier middlewares win + // when the budget is exhausted. + MaxRequestMetadataBytes = 32 << 10 + // MaxMetadataKeyBytes is the maximum length of a metadata key. + MaxMetadataKeyBytes = 96 + // MaxMetadataValueBytes is the maximum length of a metadata value. + MaxMetadataValueBytes = 4 << 10 + // MaxMiddlewaresPerChain caps the number of middleware entries + // accepted per chain at the proxy translator and the management + // REST API. Mirrors the chain invocation cap so a misconfigured + // mapping cannot push the chain clone cost beyond a known bound. + MaxMiddlewaresPerChain = 16 +) + +// KV is the canonical header/metadata representation used across the +// middleware boundary. We use a slice of KV instead of http.Header +// because it preserves key order, is cheap to deep-copy per +// invocation, and is directly representable in a future protobuf +// envelope. +type KV struct { + Key string + Value string +} + +// Input is the immutable envelope handed to each middleware. The +// dispatcher deep-copies Headers, Body, Metadata, RespHeaders, and +// RespBody before each invocation so middlewares cannot mutate the +// shared in-flight copies; mutations must flow through Output.Mutations. +type Input struct { + Slot Slot + RequestID string + TargetID string + Method string + URL string + Headers []KV + Body []byte + BodyTruncated bool + OriginalBodySize int64 + + Status int + RespHeaders []KV + RespBody []byte + RespBodyTruncated bool + OriginalRespSize int64 + + ServiceID string + AccountID string + UserID string + // UserEmail is the calling user's email address when the auth path + // resolves a user record. Empty for non-OIDC schemes (PIN/Password/ + // Header) and for legacy session JWTs minted before the email claim + // was introduced. Identity-stamping middlewares (e.g. + // llm_identity_inject) prefer this over UserID for upstream gateways + // that key budgets / attribution on a human-readable identifier. + UserEmail string + AuthMethod string + SourceIP string + // UserGroups captures the calling peer's group memberships at + // request time, surfaced from the proxy's auth flow so policy-aware + // middlewares can authorise without an extra management round-trip. + UserGroups []string + // UserGroupNames carries the human-readable display names paired + // positionally with UserGroups (UserGroupNames[i] is the name of + // UserGroups[i]). Identity-stamping middlewares prefer names for + // upstream tags so attribution dashboards stay readable. Slice may + // be shorter than UserGroups for tokens minted before names were + // resolvable; consumers should fall back to ids for missing + // positions. + UserGroupNames []string + Metadata []KV + + // AgentNetwork is true when the target is a synthesised + // agent-network service. Carried on the input so the access-log + // terminal middleware can stamp the proto field without re-deriving + // from the service ID. + AgentNetwork bool +} + +// DenyReason is the structured payload a middleware returns alongside +// a DecisionDeny. The proxy renders it through a fixed JSON template +// so middlewares cannot emit arbitrary bytes to the wire. +type DenyReason struct { + Code string + Message string + Details map[string]string +} + +// Output is the value each middleware returns to the dispatcher. The +// dispatcher applies the output filter (clamp, mutations gate) before +// any side effect reaches the shared request. +type Output struct { + Decision Decision + DenyStatus int + DenyReason *DenyReason + Metadata []KV + Mutations *Mutations +} + +// Mutations describes the deltas a middleware wants applied to the +// in-flight request. The dispatcher filters HeadersAdd/HeadersRemove +// through the compiled-in denylist and runs BodyReplace through the +// body policy before anything is applied. RewriteUpstream redirects +// the outbound target (scheme + host) for the request; the chain +// returns the latest non-nil rewrite to the reverse proxy. +type Mutations struct { + HeadersAdd []KV + HeadersRemove []string + BodyReplace []byte + RewriteUpstream *UpstreamRewrite +} + +// UpstreamRewrite redirects the request's outbound target. Only +// scheme+host are honoured; path, query, and body are untouched. The +// reverse proxy reads the rewrite (when non-nil) instead of the +// PathTarget URL configured by the synth, so a single shared synth +// service can fan out to many upstreams selected per request. +// +// AuthHeader and StripHeaders carry the upstream auth substitution +// the router needs. They bypass the framework's HeadersAdd / +// HeadersRemove denylist (which blocks Authorization, Cookie, etc. +// from middleware mutation) on the grounds that the proxy itself is +// the entity rewriting auth here, not an arbitrary middleware. The +// reverse proxy applies them directly to the upstream request after +// the chain's regular mutation phase, so a malicious or misconfigured +// middleware can still emit RewriteUpstream but only the proxy's +// trusted upstream-build path actually unpacks AuthHeader. +type UpstreamRewrite struct { + Scheme string + Host string + // Path, when non-empty, replaces the path component of the + // proxy's effective upstream URL. The rewrite path is then joined + // with the agent's request path by httputil.ProxyRequest.SetURL — + // e.g. rewrite Path="/v1/{account}/{gateway}/compat" + agent + // request "/chat/completions" → outbound + // "/v1/{account}/{gateway}/compat/chat/completions". Used by + // llm_router to honor the operator-configured upstream path on + // gateways like Cloudflare AI Gateway whose URL contains + // account / gateway segments that the agent's app doesn't know + // about. Empty Path leaves the original target's path + // untouched (the historical behavior). + Path string + // StripPathPrefix, when non-empty, is removed from the front of the agent's + // request path before it is joined onto the upstream URL. Used for + // gateway-namespace prefixes (e.g. a client addressing Bedrock as + // "/bedrock/model/{id}/invoke") that must not reach the real upstream, whose + // native path is "/model/{id}/invoke". Empty leaves the request path intact. + StripPathPrefix string + AuthHeader *AuthHeader + StripHeaders []string +} + +// AuthHeader is a single name/value pair the proxy injects on the +// upstream request after stripping the client's auth headers. +type AuthHeader struct { + Name string + Value string +} diff --git a/proxy/internal/proxy/agent_network_chain_realstack_test.go b/proxy/internal/proxy/agent_network_chain_realstack_test.go new file mode 100644 index 000000000..7b7075256 --- /dev/null +++ b/proxy/internal/proxy/agent_network_chain_realstack_test.go @@ -0,0 +1,321 @@ +package proxy_test + +import ( + "context" + "net" + "net/http" + "net/http/httptest" + "net/url" + "runtime" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/test/bufconn" + + rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service" + mgmtgrpc "github.com/netbirdio/netbird/management/internals/shared/grpc" + "github.com/netbirdio/netbird/management/server/agentnetwork" + agentNetworkTypes "github.com/netbirdio/netbird/management/server/agentnetwork/types" + "github.com/netbirdio/netbird/management/server/store" + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/bodytap" + mwbuiltin "github.com/netbirdio/netbird/proxy/internal/middleware/builtin" + // Side-effect imports register every builtin middleware factory. + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/cost_meter" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_guardrail" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_identity_inject" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_check" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_record" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_request_parser" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_response_parser" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_router" + "github.com/netbirdio/netbird/proxy/internal/proxy" + nbproxytypes "github.com/netbirdio/netbird/proxy/internal/types" + "github.com/netbirdio/netbird/shared/management/proto" + + log "github.com/sirupsen/logrus" +) + +// TestReverseProxy_AgentNetworkRequest_FullChain is the self-contained Go +// replacement for the bash 50 + 51 legs. It drives a real agent-network +// request through proxy.ReverseProxy.ServeHTTP with the actual middleware +// chain the synthesizer produces, against an in-process management gRPC and a +// httptest fake upstream — no tilt, no docker, no real LLM provider, no +// WireGuard tunnel. The test guarantees: +// +// 1. The reverse proxy's response-leg input construction copies UserGroups +// onto respInput so llm_limit_record sends a non-empty group_ids field +// on RecordLLMUsage. This is the exact bug class that motivated the +// reverseproxy.go fix — its regression would land the request OK but +// leave consumption at zero, defeating any group-targeted budget rule. +// 2. With settings.RedactPii=true the parsers ship redacted text on both +// llm.request_prompt_raw and llm.response_completion — proving the +// end-to-end wiring (synth → proto → spec → parser config) carries the +// toggle through to runtime emission. +// 3. The full chain (request + response + recorder) runs against a real +// management stack and the consumption row for the bound group dim +// increments. +// +// If any of those three guarantees regresses, this single test fails. +func TestReverseProxy_AgentNetworkRequest_FullChain(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("sqlite store not supported on Windows") + } + + const ( + testAccountID = "acct-fullchain-1" + testAdminUser = "user-admin-1" + adminGroupID = "grp-admins" + providerID = "prov-openai-test" + cluster = "test.proxy.local" + subdomain = "fullchain" + ) + testLogger := log.New() + testLogger.SetLevel(log.PanicLevel) // keep test output clean + + ctx := context.Background() + + // ---- 1. Fake upstream that returns OpenAI-shaped JSON with PII in the + // completion. The reverse proxy's chain will redact this when the synth + // stamps redact_pii=true on the response parser config. + completion := "Sample record: Alice Johnson alice.johnson@example.com SSN 123-45-6789 phone (202) 555-0147 also Bob 202/555/0108" + upstreamBody := []byte(`{"id":"x","model":"gpt-5.4","choices":[{"message":{"role":"assistant","content":"` + completion + `"}}],"usage":{"prompt_tokens":12,"completion_tokens":40,"total_tokens":52}}`) + var upstreamHits atomic.Int64 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamHits.Add(1) + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(upstreamBody) + })) + t.Cleanup(upstream.Close) + upstreamHost := strings.TrimPrefix(upstream.URL, "http://") + + // ---- 2. In-process management gRPC server (bufconn) backed by a real + // sqlite store + real agentnetwork.Manager. The proxy's middlewares talk + // to this client. + st, cleanup, err := store.NewTestStoreFromSQL(ctx, "", t.TempDir()) + require.NoError(t, err, "real sqlite test store must come up") + t.Cleanup(cleanup) + + anMgr := agentnetwork.NewManager(st, nil, nil, nil) + server := &mgmtgrpc.ProxyServiceServer{} + server.SetAgentNetworkLimitsService(anMgr) + + lis := bufconn.Listen(1024 * 1024) + srv := grpc.NewServer() + proto.RegisterProxyServiceServer(srv, server) + go func() { _ = srv.Serve(lis) }() + t.Cleanup(srv.Stop) + + conn, err := grpc.NewClient("passthrough:///bufnet", + grpc.WithContextDialer(func(_ context.Context, _ string) (net.Conn, error) { return lis.Dial() }), + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.Close() }) + mgmtClient := proto.NewProxyServiceClient(conn) + + // ---- 3. Seed account state: settings (redact + capture on), provider + // whose upstream URL points at our fake server, policy (catch-all-allow + // over the Admins group → window=0 path), and a generous budget rule + // targeting Admins so the curl succeeds and we can prove the counter + // increments on the response leg. + require.NoError(t, st.SaveAgentNetworkSettings(ctx, &agentNetworkTypes.Settings{ + AccountID: testAccountID, + Cluster: cluster, + Subdomain: subdomain, + EnablePromptCollection: true, + EnableLogCollection: true, + RedactPii: true, + })) + require.NoError(t, st.SaveAgentNetworkProvider(ctx, &agentNetworkTypes.Provider{ + ID: providerID, + AccountID: testAccountID, + ProviderID: "openai_api", + Name: "openai-fullchain-test", + UpstreamURL: upstream.URL, // router rewrites to this + APIKey: "sk-test", + Enabled: true, + Models: []agentNetworkTypes.ProviderModel{{ID: "gpt-5.4"}}, + SessionPrivateKey: "priv", + SessionPublicKey: "pub", + })) + require.NoError(t, st.SaveAgentNetworkPolicy(ctx, &agentNetworkTypes.Policy{ + ID: "ainpol-fullchain", + AccountID: testAccountID, + Name: "admins-openai", + Enabled: true, + SourceGroups: []string{adminGroupID}, + DestinationProviderIDs: []string{providerID}, + // No token / budget caps → effectiveWindowSeconds=0 → exercises the + // catch-all-allow path that the GC-2 record-on-window=0 fix targets. + })) + require.NoError(t, st.SaveAgentNetworkBudgetRule(ctx, &agentNetworkTypes.AccountBudgetRule{ + ID: "ainbud-admins-fullchain", + AccountID: testAccountID, + Name: "admins-monthly", + Enabled: true, + TargetGroups: []string{adminGroupID}, + Limits: agentNetworkTypes.PolicyLimits{ + TokenLimit: agentNetworkTypes.PolicyTokenLimit{Enabled: true, GroupCap: 1_000_000, UserCap: 1_000_000, WindowSeconds: 60}, + }, + })) + + // ---- 4. Synth the service. This produces the exact middleware chain + // configuration the production reconcile path ships to the proxy. + services, err := agentnetwork.SynthesizeServices(ctx, st, testAccountID) + require.NoError(t, err) + require.Len(t, services, 1, "exactly one synth service expected") + synthSvc := services[0] + require.NotEmpty(t, synthSvc.Targets, "synth target must exist") + + // ---- 5. Wire the middleware framework — same registry the proxy uses + // in production, configured with our bufconn-backed management client. + mwbuiltin.Configure(ctx, t.TempDir(), nil, testLogger, mgmtClient) + registry := mwbuiltin.DefaultRegistry() + mwMetrics, err := middleware.NewMetrics(nil) + require.NoError(t, err) + mwMgr := middleware.NewManager(0, mwMetrics, testLogger) + mwMgr.SetResolver(middleware.NewResolver(registry)) + + // Convert the synth's rpservice.MiddlewareConfig list into proxy + // middleware.Spec values. Mirrors the proto→Spec translation server.go + // does at runtime; kept inline here so the test isn't coupled to the + // proxy server's private translateMiddlewareConfig helper. + specs := make([]middleware.Spec, 0, len(synthSvc.Targets[0].Options.Middlewares)) + for _, mw := range synthSvc.Targets[0].Options.Middlewares { + var slot middleware.Slot + switch mw.Slot { + case rpservice.MiddlewareSlotOnRequest: + slot = middleware.SlotOnRequest + case rpservice.MiddlewareSlotOnResponse: + slot = middleware.SlotOnResponse + case rpservice.MiddlewareSlotTerminal: + slot = middleware.SlotTerminal + default: + t.Fatalf("unknown middleware slot %q on %s", mw.Slot, mw.ID) + } + specs = append(specs, middleware.Spec{ + ID: mw.ID, + Slot: slot, + Enabled: mw.Enabled, + FailMode: middleware.FailOpen, + Timeout: middleware.DefaultTimeout, + RawConfig: append([]byte(nil), mw.ConfigJSON...), + CanMutate: mw.CanMutate, + }) + } + + serviceIDStr := synthSvc.ID + require.NoError(t, mwMgr.Rebuild(serviceIDStr, []middleware.PathTargetBinding{{ + ServiceID: serviceIDStr, + PathID: "/", + Specs: specs, + }})) + + // ---- 6. Build the reverse proxy, with a mapping whose target URL goes + // straight to the fake upstream (the router middleware rewriting upstream + // from the synth's noop placeholder isn't needed when we own the mapping + // in-process — point the target at the fake URL directly so the body + // arrives at the upstream the synth would have routed to). + upstreamURL, err := url.Parse(upstream.URL) + require.NoError(t, err) + + rp := proxy.NewReverseProxy(http.DefaultTransport, "auto", nil, testLogger, proxy.WithMiddlewareManager(mwMgr)) + rp.AddMapping(proxy.Mapping{ + ID: nbproxytypes.ServiceID(serviceIDStr), + AccountID: nbproxytypes.AccountID(testAccountID), + Host: synthSvc.Domain, + Paths: map[string]*proxy.PathTarget{ + "/": { + URL: upstreamURL, + DirectUpstream: true, + AgentNetwork: true, + Middlewares: specs, + CaptureConfig: &bodytap.Config{ + MaxRequestBytes: 1 << 20, + MaxResponseBytes: 1 << 20, + ContentTypes: []string{"application/json", "text/event-stream"}, + }, + }, + }, + }) + + // ---- 7. Send a request with the auth-stamped CapturedData (mimicking + // what the tunnel-peer auth middleware does at the edge of the proxy). + reqBody := `{"model":"gpt-5.4","client_metadata":{"session_id":"sess-fullchain-1"},"messages":[{"role":"user","content":"contact alice.johnson@example.com SSN 987-65-4321 phone (202)555-0156"}]}` + req := httptest.NewRequest("POST", "https://"+synthSvc.Domain+"/v1/chat/completions", strings.NewReader(reqBody)) + req.Host = synthSvc.Domain + req.Header.Set("Content-Type", "application/json") + + cd := proxy.NewCapturedData("test-request-1") + cd.SetServiceID(nbproxytypes.ServiceID(serviceIDStr)) + cd.SetAccountID(nbproxytypes.AccountID(testAccountID)) + cd.SetUserID(testAdminUser) + cd.SetUserGroups([]string{adminGroupID}) + cd.SetAuthMethod("tunnel_peer") + req = req.WithContext(proxy.WithCapturedData(req.Context(), cd)) + + w := httptest.NewRecorder() + rp.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code, "upstream call must succeed end-to-end; body=%s", w.Body.String()) + assert.GreaterOrEqual(t, upstreamHits.Load(), int64(1), "fake upstream must have been hit") + + // ---- 8. Assertions — the three guarantees this test exists for. + + // 8a. The reverseproxy.go respInput construction carried UserGroups + // into the response-leg middleware chain, so llm_limit_record sent a + // non-empty group_ids on RecordLLMUsage. Verifying via the management + // store directly bypasses the manager's permission gate (which is nil + // in this test) — we want to confirm the row landed, not who saw it. + require.Eventually(t, func() bool { + rows, lerr := st.ListAgentNetworkConsumption(ctx, store.LockingStrengthNone, testAccountID) + if lerr != nil { + return false + } + for _, r := range rows { + if r.DimensionKind == agentNetworkTypes.DimensionGroup && + r.DimensionID == adminGroupID && + r.WindowSeconds == 60 && + r.TokensInput+r.TokensOutput > 0 { + return true + } + } + return false + }, 5*time.Second, 50*time.Millisecond, + "Admins group consumption row must increment via the response leg — if this fails the proxy's respInput dropped UserGroups again or the parser/recorder wiring is broken") + + // 8b. Both the captured prompt and the captured completion are + // redacted — proves the synth threads redact_pii=true into BOTH parser + // configs and the parsers honour it at emission time. + md := cd.GetMetadata() + promptRaw := md["llm.request_prompt_raw"] + completionMeta := md["llm.response_completion"] + + // 8a-bis. The session id from client_metadata.session_id flows through + // the request parser into the captured metadata, so the access-log / + // usage rows can group this request with the rest of its conversation. + assert.Equal(t, "sess-fullchain-1", md["llm.session_id"], + "session id must be extracted from client_metadata.session_id and carried through the chain") + + assert.NotEmpty(t, promptRaw, "llm.request_prompt_raw must be present in captured metadata") + assert.Contains(t, promptRaw, "[REDACTED:", "captured raw prompt must carry redaction markers") + assert.NotContains(t, promptRaw, "alice.johnson@example.com", "raw email must NOT survive in prompt_raw") + assert.NotContains(t, promptRaw, "987-65-4321", "raw SSN must NOT survive in prompt_raw") + assert.NotContains(t, promptRaw, "(202)555-0156", "raw paren-no-space phone must NOT survive in prompt_raw") + + assert.NotEmpty(t, completionMeta, "llm.response_completion must be present in captured metadata") + assert.Contains(t, completionMeta, "[REDACTED:", "captured completion must carry redaction markers") + assert.NotContains(t, completionMeta, "alice.johnson@example.com", "raw email must NOT survive in completion") + assert.NotContains(t, completionMeta, "123-45-6789", "raw SSN must NOT survive in completion") + assert.NotContains(t, completionMeta, "(202) 555-0147", "raw paren+space phone must NOT survive in completion") + assert.NotContains(t, completionMeta, "202/555/0108", "raw slash phone must NOT survive in completion") + + _ = upstreamHost // kept for future header-inspection assertions if needed +} diff --git a/proxy/internal/proxy/context.go b/proxy/internal/proxy/context.go index e05ec78aa..09bb9a8d1 100644 --- a/proxy/internal/proxy/context.go +++ b/proxy/internal/proxy/context.go @@ -58,9 +58,11 @@ type CapturedData struct { // the JWT's group_names claim or from ValidateSession/Tunnel // responses. Slice may be shorter than userGroups for tokens minted // before names were resolvable. - userGroupNames []string - authMethod string - metadata map[string]string + userGroupNames []string + authMethod string + metadata map[string]string + agentNetwork bool + suppressAccessLog bool } // NewCapturedData creates a CapturedData with the given request ID. @@ -178,6 +180,41 @@ func (c *CapturedData) SetUserGroups(groups []string) { c.userGroups = append(c.userGroups[:0], groups...) } +// SetAgentNetwork records whether the request hit a synthesised +// agent-network target. The terminal access-log middleware stamps the +// flag onto the proto so management can distinguish synthetic traffic. +func (c *CapturedData) SetAgentNetwork(b bool) { + c.mu.Lock() + defer c.mu.Unlock() + c.agentNetwork = b +} + +// GetAgentNetwork reports whether the request matched a synthesised +// agent-network target. +func (c *CapturedData) GetAgentNetwork() bool { + c.mu.RLock() + defer c.mu.RUnlock() + return c.agentNetwork +} + +// SetSuppressAccessLog records whether the per-request access-log emission +// must be skipped for this request. Stamped from the matched target's +// DisableAccessLog flag so the access-log middleware can short-circuit +// log delivery for opted-out agent-network targets. +func (c *CapturedData) SetSuppressAccessLog(b bool) { + c.mu.Lock() + defer c.mu.Unlock() + c.suppressAccessLog = b +} + +// GetSuppressAccessLog reports whether access-log emission has been +// suppressed for this request. +func (c *CapturedData) GetSuppressAccessLog() bool { + c.mu.RLock() + defer c.mu.RUnlock() + return c.suppressAccessLog +} + // GetUserGroups returns a copy of the authenticated user's group // memberships. func (c *CapturedData) GetUserGroups() []string { diff --git a/proxy/internal/proxy/reverseproxy.go b/proxy/internal/proxy/reverseproxy.go index da0bf6552..c64d0fa26 100644 --- a/proxy/internal/proxy/reverseproxy.go +++ b/proxy/internal/proxy/reverseproxy.go @@ -2,6 +2,7 @@ package proxy import ( "context" + "encoding/json" "errors" "fmt" "net" @@ -11,10 +12,13 @@ import ( "net/url" "strings" "sync" + "time" log "github.com/sirupsen/logrus" "github.com/netbirdio/netbird/proxy/auth" + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/bodytap" "github.com/netbirdio/netbird/proxy/internal/roundtrip" "github.com/netbirdio/netbird/proxy/internal/types" "github.com/netbirdio/netbird/proxy/web" @@ -32,6 +36,25 @@ type ReverseProxy struct { mappingsMux sync.RWMutex mappings map[string]Mapping logger *log.Logger + // middlewareManager, when non-nil, drives per-target middleware + // dispatch. A nil manager (or an empty chain for the resolved + // target) keeps the reverse-proxy hot path on the no-capture fast + // path with no middleware overhead. + middlewareManager *middleware.Manager +} + +// Option configures optional ReverseProxy behavior. Options exist so the core +// constructor signature stays stable across additive features. +type Option func(*ReverseProxy) + +// WithMiddlewareManager attaches a middleware manager to the reverse +// proxy. When the manager is nil or returns an empty chain for the +// target, the request follows the fast path with no middleware +// overhead. +func WithMiddlewareManager(m *middleware.Manager) Option { + return func(p *ReverseProxy) { + p.middlewareManager = m + } } // NewReverseProxy configures a new NetBird ReverseProxy. @@ -40,17 +63,21 @@ type ReverseProxy struct { // between requested URLs and targets. // The internal mappings can be modified using the AddMapping // and RemoveMapping functions. -func NewReverseProxy(transport http.RoundTripper, forwardedProto string, trustedProxies []netip.Prefix, logger *log.Logger) *ReverseProxy { +func NewReverseProxy(transport http.RoundTripper, forwardedProto string, trustedProxies []netip.Prefix, logger *log.Logger, opts ...Option) *ReverseProxy { if logger == nil { logger = log.StandardLogger() } - return &ReverseProxy{ + p := &ReverseProxy{ transport: transport, forwardedProto: forwardedProto, trustedProxies: trustedProxies, mappings: make(map[string]Mapping), logger: logger, } + for _, opt := range opts { + opt(p) + } + return p } func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { @@ -89,9 +116,12 @@ func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { // Populate captured data if it exists (allows middleware to read after handler completes). // This solves the problem of passing data UP the middleware chain: we put a mutable struct // pointer in the context, and mutate the struct here so outer middleware can read it. - if capturedData := CapturedDataFromContext(ctx); capturedData != nil { + capturedData := CapturedDataFromContext(ctx) + if capturedData != nil { capturedData.SetServiceID(result.serviceID) capturedData.SetAccountID(result.accountID) + capturedData.SetAgentNetwork(result.target != nil && result.target.AgentNetwork) + capturedData.SetSuppressAccessLog(result.target != nil && result.target.DisableAccessLog) } pt := result.target @@ -99,28 +129,331 @@ func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) { if pt.SkipTLSVerify { ctx = roundtrip.WithSkipTLSVerify(ctx) } - if pt.RequestTimeout > 0 { - ctx = types.WithDialTimeout(ctx, pt.RequestTimeout) - } if pt.DirectUpstream { ctx = roundtrip.WithDirectUpstream(ctx) } + if pt.RequestTimeout > 0 { + ctx = types.WithDialTimeout(ctx, pt.RequestTimeout) + } rewriteMatchedPath := result.matchedPath if pt.PathRewrite == PathRewritePreserve { rewriteMatchedPath = "" } + chain := p.resolveChain(result) + if chain == nil || chain.Empty() { + rp := &httputil.ReverseProxy{ + Rewrite: p.rewriteFunc(pt.URL, rewriteMatchedPath, result.passHostHeader, pt.PathRewrite, pt.CustomHeaders, result.stripAuthHeaders), + Transport: p.transport, + FlushInterval: -1, + ErrorHandler: p.proxyErrorHandler, + } + if result.rewriteRedirects { + rp.ModifyResponse = p.rewriteLocationFunc(pt.URL, rewriteMatchedPath, r) //nolint:bodyclose + } + rp.ServeHTTP(w, r.WithContext(ctx)) + return + } + + middlewareIDs := chain.IDs() + p.logger.Debugf("middleware chain matched: service=%s path=%s middlewares=%v", result.serviceID, result.matchedPath, middlewareIDs) + + capturedBody, truncated, originalSize, bypass, releaseBudget, captureErr := bodytap.CaptureRequest(r, pt.CaptureConfig, p.middlewareManager.Budget()) + defer releaseBudget() + if captureErr != nil { + p.logger.Debugf("middleware request body capture error: %v", captureErr) + } + if bypass != "" { + if capturedData != nil { + capturedData.SetMetadata("mw.capture.bypass_reason", bypass) + } + p.middlewareManager.Metrics().IncCaptureBypass(ctx, string(result.serviceID), bypass) + } + + // Routing recovery for oversized agent-network requests: when the body + // exceeded the capture cap (bypassed or truncated), the captured copy + // can't be parsed for the model, so llm_router would deny with + // model_not_routable. Scan the full stream for just the routing fields + // and hand the request parser a minimal stub so routing succeeds; the + // prompt stays uncaptured and the upstream still gets the full body. + if pt.AgentNetwork && (truncated || capturedBody == nil) { + if model, stream, ok := bodytap.ScanRoutingFields(r, bodytap.MaxRoutingScanBytes); ok { + capturedBody = buildRoutingStub(model, stream) + truncated = false + p.logger.Debugf("agent-network routing recovery: extracted model=%s stream=%t from oversized request body (service=%s)", model, stream, result.serviceID) + } + } + + acc := middleware.NewAccumulator(middleware.MaxRequestMetadataBytes) + reqInput := buildRequestInput(r, result, capturedData, capturedBody, truncated, originalSize) + + denyOutput, requestMeta, upstreamRewrite, _ := chain.RunRequest(ctx, r, reqInput, acc) + if capturedData != nil { + for _, kv := range requestMeta { + capturedData.SetMetadata(kv.Key, kv.Value) + } + } + if denyOutput != nil { + middlewareID := "middleware" + if denyOutput.DenyReason != nil && denyOutput.DenyReason.Code != "" { + middlewareID = denyOutput.DenyReason.Code + } + // Policy/budget/routing/guardrail denials are expected runtime outcomes + // and can be high-volume under misconfigured or hostile clients; keep + // per-request detail at Debug and rely on metrics/access logs at scale. + p.logger.Debugf("middleware chain denied request: service=%s path=%s middlewares=%v reason=%s status=%d", + result.serviceID, result.matchedPath, middlewareIDs, middlewareID, denyOutput.DenyStatus) + middleware.RenderDenyResponse(w, middlewareID, denyOutput.DenyReason, denyOutput.DenyStatus) + return + } + + respWriter := http.ResponseWriter(w) + var capturingWriter *bodytap.CapturingResponseWriter + if pt.CaptureConfig != nil && pt.CaptureConfig.MaxResponseBytes > 0 { + capturingWriter = bodytap.NewCapturingResponseWriter(w, pt.CaptureConfig.MaxResponseBytes, p.middlewareManager.Budget()) + defer capturingWriter.Release() + if capturingWriter.Bypassed() { + if capturedData != nil { + capturedData.SetMetadata("mw.capture.bypass_reason", capturingWriter.BypassReason()) + } + p.middlewareManager.Metrics().IncCaptureBypass(ctx, string(result.serviceID), capturingWriter.BypassReason()) + capturingWriter = nil + } else { + respWriter = capturingWriter + } + } + + defer func() { + if capturingWriter == nil { + return + } + respInput := &middleware.Input{ + Slot: middleware.SlotOnResponse, + RequestID: reqInput.RequestID, + TargetID: reqInput.TargetID, + Method: reqInput.Method, + URL: reqInput.URL, + Headers: reqInput.Headers, + Status: capturingWriter.Status(), + RespHeaders: headerToKV(w.Header()), + RespBody: capturingWriter.Body(), + RespBodyTruncated: capturingWriter.Truncated(), + OriginalRespSize: capturingWriter.BytesWritten(), + ServiceID: reqInput.ServiceID, + AccountID: reqInput.AccountID, + UserID: reqInput.UserID, + // UserEmail / UserGroups / UserGroupNames must flow into the + // response leg too — llm_limit_record needs UserGroups to send + // group_ids on RecordLLMUsage so management's account-budget + // fan-out can match group-targeted rules; identity-stamping and + // any future response-side authorisation also depend on these. + UserEmail: reqInput.UserEmail, + UserGroups: reqInput.UserGroups, + UserGroupNames: reqInput.UserGroupNames, + AuthMethod: reqInput.AuthMethod, + SourceIP: reqInput.SourceIP, + Metadata: requestMeta, + AgentNetwork: reqInput.AgentNetwork, + } + // The response/terminal phase runs after the body is forwarded, so + // a streaming client (e.g. Codex) has usually disconnected by now, + // cancelling r.Context(). These middlewares only observe and record + // (token/cost metering, usage recording) and must still complete — + // otherwise the dispatcher short-circuits each to fail-mode and the + // usage is silently lost. Detach from client cancellation, keep ctx + // values, and bound the work. + obsCtx, obsCancel := context.WithTimeout(context.WithoutCancel(ctx), observabilityPhaseTimeout) + defer obsCancel() + + respMeta := chain.RunResponse(obsCtx, respInput, acc) + if capturedData != nil { + for _, kv := range respMeta { + capturedData.SetMetadata(kv.Key, kv.Value) + } + } + + // Terminal slot sees the merged metadata bag from request and + // response phases. + mergedMeta := append(append([]middleware.KV(nil), requestMeta...), respMeta...) + termInput := *respInput + termInput.Slot = middleware.SlotTerminal + termInput.Metadata = mergedMeta + termMeta := chain.RunTerminal(obsCtx, &termInput, acc) + if capturedData != nil { + for _, kv := range termMeta { + capturedData.SetMetadata(kv.Key, kv.Value) + } + } + + p.logger.Debugf("middleware chain ran: service=%s path=%s middlewares=%v status=%d req_meta=%d resp_meta=%d term_meta=%d", + result.serviceID, result.matchedPath, middlewareIDs, capturingWriter.Status(), len(requestMeta), len(respMeta), len(termMeta)) + }() + + effectiveURL := applyUpstreamRewrite(pt.URL, upstreamRewrite) + if upstreamRewrite != nil { + r.Host = effectiveURL.Host + applyUpstreamHeaders(r, upstreamRewrite) + stripUpstreamPathPrefix(r, upstreamRewrite.StripPathPrefix) + } + rp := &httputil.ReverseProxy{ - Rewrite: p.rewriteFunc(pt.URL, rewriteMatchedPath, result.passHostHeader, pt.PathRewrite, pt.CustomHeaders, result.stripAuthHeaders), + Rewrite: p.rewriteFunc(effectiveURL, rewriteMatchedPath, result.passHostHeader, pt.PathRewrite, pt.CustomHeaders, result.stripAuthHeaders), Transport: p.transport, FlushInterval: -1, ErrorHandler: p.proxyErrorHandler, } if result.rewriteRedirects { - rp.ModifyResponse = p.rewriteLocationFunc(pt.URL, rewriteMatchedPath, r) //nolint:bodyclose + rp.ModifyResponse = p.rewriteLocationFunc(effectiveURL, rewriteMatchedPath, r) //nolint:bodyclose } - rp.ServeHTTP(w, r.WithContext(ctx)) + rp.ServeHTTP(respWriter, r.WithContext(ctx)) +} + +// buildRoutingStub returns a minimal JSON request body carrying only the +// model and stream fields. It feeds the LLM request parser when the real +// body was too large to capture: the parser emits llm.model / llm.stream +// so llm_router can route, while ExtractPrompt on the stub yields nothing +// — no prompt is captured for oversized requests. +func buildRoutingStub(model string, stream bool) []byte { + b, err := json.Marshal(map[string]any{"model": model, "stream": stream}) + if err != nil { + return nil + } + return b +} + +// applyUpstreamRewrite returns the effective upstream URL after +// applying a middleware-emitted rewrite. When rewrite is nil or +// incomplete, the original target is returned unchanged. The original +// URL is never mutated; a clone is returned when a rewrite applies. +// +// Rewrite Path semantics: when non-empty, replaces the cloned URL's +// path entirely. httputil.ProxyRequest.SetURL then joins target.Path +// with the agent's request path, so an operator-configured upstream +// path like "/v1/{account}/{gateway}/compat" gets prepended to +// "/chat/completions" yielding the full Cloudflare-shaped path. +// Empty rewrite.Path preserves the original target's path (the +// historical, non-agent-network behavior). +func applyUpstreamRewrite(orig *url.URL, rewrite *middleware.UpstreamRewrite) *url.URL { + if rewrite == nil || orig == nil { + return orig + } + if rewrite.Scheme == "" || rewrite.Host == "" { + return orig + } + cloned := *orig + cloned.Scheme = rewrite.Scheme + cloned.Host = rewrite.Host + if rewrite.Path != "" { + cloned.Path = rewrite.Path + cloned.RawPath = "" + } + return &cloned +} + +// stripUpstreamPathPrefix removes a gateway-namespace prefix (e.g. "/bedrock") +// from the request path before it is forwarded, so the upstream receives its +// native path. The chain has already run by this point, so metering/logging +// keep the original client path; only the outbound path is rewritten. RawPath +// is cleared so the escaped form is recomputed from the trimmed Path. +func stripUpstreamPathPrefix(r *http.Request, prefix string) { + if r == nil || r.URL == nil || prefix == "" { + return + } + if !strings.HasPrefix(r.URL.Path, prefix+"/") && r.URL.Path != prefix { + return + } + r.URL.Path = strings.TrimPrefix(r.URL.Path, prefix) + if r.URL.Path == "" { + r.URL.Path = "/" + } + r.URL.RawPath = "" +} + +// applyUpstreamHeaders strips the headers the rewrite asks for and +// injects the resolved auth header on the in-flight request. It is +// the proxy-trusted counterpart to chain.applyMutations: regular +// middleware HeadersAdd/HeadersRemove pass through the framework +// denylist (which blocks Authorization, Cookie, etc.), but the +// router middleware needs to replace Authorization on the upstream +// request as a first-class operation. AuthHeader/StripHeaders ride +// on UpstreamRewrite so only the proxy's upstream-build path +// unpacks them — middlewares can't smuggle these in via the +// regular mutation surface. +func applyUpstreamHeaders(r *http.Request, rewrite *middleware.UpstreamRewrite) { + if r == nil || rewrite == nil { + return + } + for _, name := range rewrite.StripHeaders { + if name == "" { + continue + } + r.Header.Del(name) + } + if rewrite.AuthHeader != nil && rewrite.AuthHeader.Name != "" { + r.Header.Set(rewrite.AuthHeader.Name, rewrite.AuthHeader.Value) + } +} + +// resolveChain returns the middleware chain registered for the +// resolved target, or nil when middleware is disabled for the proxy +// or the target. +func (p *ReverseProxy) resolveChain(result targetResult) *middleware.Chain { + if p.middlewareManager == nil { + return nil + } + return p.middlewareManager.ChainFor(string(result.serviceID), result.matchedPath) +} + +// buildRequestInput gathers the per-request fields the middleware +// chain needs. Body and captured metadata are passed in; the rest are +// copied from the request and CapturedData. +func buildRequestInput(r *http.Request, result targetResult, cd *CapturedData, body []byte, truncated bool, originalSize int64) *middleware.Input { + in := &middleware.Input{ + Slot: middleware.SlotOnRequest, + TargetID: result.matchedPath, + Method: r.Method, + URL: r.URL.String(), + Headers: headerToKV(r.Header), + Body: body, + BodyTruncated: truncated, + OriginalBodySize: originalSize, + ServiceID: string(result.serviceID), + AccountID: string(result.accountID), + AgentNetwork: result.target != nil && result.target.AgentNetwork, + } + if cd != nil { + in.RequestID = cd.GetRequestID() + in.UserID = cd.GetUserID() + in.UserEmail = cd.GetUserEmail() + in.UserGroups = cd.GetUserGroups() + in.UserGroupNames = cd.GetUserGroupNames() + in.AuthMethod = cd.GetAuthMethod() + if ip := cd.GetClientIP(); ip.IsValid() { + in.SourceIP = ip.String() + } + } + return in +} + +// headerToKV flattens an http.Header into the KV slice shape expected +// by the middleware envelope, preserving value order under the same +// key. +func headerToKV(h http.Header) []middleware.KV { + if len(h) == 0 { + return nil + } + total := 0 + for _, v := range h { + total += len(v) + } + out := make([]middleware.KV, 0, total) + for k, vs := range h { + for _, v := range vs { + out = append(out, middleware.KV{Key: k, Value: v}) + } + } + return out } // isSelfTargetLoop reports whether an overlay-origin request is about to @@ -486,6 +819,14 @@ const ( // comma or any non-printable byte are dropped at stamp time so the // list is unambiguously splittable by consumers. headerNetBirdGroups = "X-NetBird-Groups" + + // observabilityPhaseTimeout bounds the detached response/terminal + // metering phase. It runs after the client connection (and its context) + // may be gone, so it can't borrow the request deadline; this ceiling + // keeps a slow management round-trip (RecordLLMUsage) from pinning the + // handler goroutine indefinitely while still allowing each middleware + // its own per-invoke timeout. + observabilityPhaseTimeout = 30 * time.Second ) // isHeaderValueSafe reports whether v is a valid RFC 7230 field-value: diff --git a/proxy/internal/proxy/reverseproxy_test.go b/proxy/internal/proxy/reverseproxy_test.go index a8244fa56..9bd427056 100644 --- a/proxy/internal/proxy/reverseproxy_test.go +++ b/proxy/internal/proxy/reverseproxy_test.go @@ -19,6 +19,7 @@ import ( "github.com/stretchr/testify/require" "github.com/netbirdio/netbird/proxy/auth" + "github.com/netbirdio/netbird/proxy/internal/middleware" "github.com/netbirdio/netbird/proxy/internal/roundtrip" "github.com/netbirdio/netbird/proxy/internal/types" "github.com/netbirdio/netbird/proxy/web" @@ -1407,3 +1408,45 @@ func TestStampNetBirdIdentity_CapturedDataPresentButEmpty(t *testing.T) { assert.Empty(t, pr.Out.Header.Get(headerNetBirdGroups), "X-NetBird-Groups must be stripped when CapturedData has no groups") } + +// TestBuildRequestInput_PropagatesIdentityAndGroups locks the final wiring link +// between auth and the middleware chain: CapturedData identity (user, groups, +// auth method, client IP) and the target's AgentNetwork flag must land on the +// middleware Input the chain runs against. If UserGroups stops flowing here, +// llm_router denies every request with no_authorised_provider. +func TestBuildRequestInput_PropagatesIdentityAndGroups(t *testing.T) { + cd := NewCapturedData("req-123") + cd.SetUserID("user-1") + cd.SetUserEmail("user@example.com") + cd.SetUserGroups([]string{"grp-admins", "grp-users"}) + cd.SetUserGroupNames([]string{"Admins", "Users"}) + cd.SetAuthMethod("oidc") + cd.SetClientIP(netip.MustParseAddr("100.90.1.14")) + + r := httptest.NewRequest(http.MethodPost, "http://agent.example.com/v1/chat/completions", nil) + r.Header.Set("Content-Type", "application/json") + + result := targetResult{ + target: &PathTarget{AgentNetwork: true}, + matchedPath: "/", + serviceID: types.ServiceID("svc-1"), + accountID: types.AccountID("acct-1"), + } + + body := []byte(`{"model":"gpt-5.4"}`) + in := buildRequestInput(r, result, cd, body, false, int64(len(body))) + + require.NotNil(t, in, "buildRequestInput must return an envelope") + assert.Equal(t, middleware.SlotOnRequest, in.Slot, "request input runs in the on-request slot") + assert.Equal(t, "svc-1", in.ServiceID, "service id must propagate") + assert.Equal(t, "acct-1", in.AccountID, "account id must propagate") + assert.Equal(t, "user-1", in.UserID, "user id must propagate") + assert.Equal(t, "user@example.com", in.UserEmail, "user email must propagate") + assert.Equal(t, []string{"grp-admins", "grp-users"}, in.UserGroups, + "CapturedData groups MUST reach the middleware Input — llm_router authorises against this") + assert.Equal(t, []string{"Admins", "Users"}, in.UserGroupNames, "group names must propagate") + assert.Equal(t, "oidc", in.AuthMethod, "auth method must propagate") + assert.Equal(t, "100.90.1.14", in.SourceIP, "client IP must propagate") + assert.True(t, in.AgentNetwork, "agent-network target flag must reach the Input") + assert.Equal(t, body, in.Body, "captured body must reach the Input") +} diff --git a/proxy/internal/proxy/servicemapping.go b/proxy/internal/proxy/servicemapping.go index 46b4d2e8d..64fccc42a 100644 --- a/proxy/internal/proxy/servicemapping.go +++ b/proxy/internal/proxy/servicemapping.go @@ -8,6 +8,8 @@ import ( "strings" "time" + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/bodytap" "github.com/netbirdio/netbird/proxy/internal/types" ) @@ -32,6 +34,20 @@ type PathTarget struct { // over the embedded NetBird WireGuard client when forwarding requests // to this target. Default false → embedded client (existing behaviour). DirectUpstream bool + // Middlewares is the validated per-target middleware chain. Nil or empty + // for non-agent-network targets, keeping them on the no-middleware fast path. + Middlewares []middleware.Spec + // CaptureConfig holds the per-target body-capture limits used by the + // middleware chain. Nil for targets without body-inspecting middlewares. + CaptureConfig *bodytap.Config + // AgentNetwork marks this target as a synthesised agent-network target so + // the proxy can tag access-log entries and gate agent-network behaviour. + AgentNetwork bool + // DisableAccessLog suppresses the per-request access-log emission for this + // target. Defaults false so non-agent-network targets continue to log + // unchanged. The agent-network synthesizer sets this true only when the + // account's EnableLogCollection toggle is off. + DisableAccessLog bool } // Mapping describes how a domain is routed by the HTTP reverse proxy. diff --git a/proxy/internal/proxy/strip_prefix_test.go b/proxy/internal/proxy/strip_prefix_test.go new file mode 100644 index 000000000..4ff364f6a --- /dev/null +++ b/proxy/internal/proxy/strip_prefix_test.go @@ -0,0 +1,30 @@ +package proxy + +import ( + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestStripUpstreamPathPrefix(t *testing.T) { + cases := []struct { + name string + path string + prefix string + want string + }{ + {"strips matching namespace prefix", "/bedrock/model/x/invoke", "/bedrock", "/model/x/invoke"}, + {"no-op when prefix absent", "/model/x/invoke", "/bedrock", "/model/x/invoke"}, + {"no-op on empty prefix", "/bedrock/model/x/invoke", "", "/bedrock/model/x/invoke"}, + {"no-op on non-segment match", "/bedrockfoo/model/x", "/bedrock", "/bedrockfoo/model/x"}, + {"bare prefix collapses to root", "/bedrock", "/bedrock", "/"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + r := httptest.NewRequest("POST", tc.path, nil) + stripUpstreamPathPrefix(r, tc.prefix) + assert.Equal(t, tc.want, r.URL.Path, "stripped path for %q", tc.path) + }) + } +} diff --git a/proxy/middleware_register.go b/proxy/middleware_register.go new file mode 100644 index 000000000..736ee04c1 --- /dev/null +++ b/proxy/middleware_register.go @@ -0,0 +1,16 @@ +package proxy + +// Anonymous imports trigger init() in each built-in middleware +// sub-package so they self-register into mwbuiltin.DefaultRegistry() +// before initMiddlewareManager builds the resolver. Add a new line +// here when introducing another built-in middleware. +import ( + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/cost_meter" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_guardrail" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_identity_inject" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_check" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_limit_record" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_request_parser" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_response_parser" + _ "github.com/netbirdio/netbird/proxy/internal/middleware/builtin/llm_router" +) diff --git a/proxy/middleware_translate.go b/proxy/middleware_translate.go new file mode 100644 index 000000000..c5d9fe016 --- /dev/null +++ b/proxy/middleware_translate.go @@ -0,0 +1,165 @@ +package proxy + +import ( + "context" + "time" + + log "github.com/sirupsen/logrus" + + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/proxy/internal/middleware/bodytap" + "github.com/netbirdio/netbird/shared/management/proto" +) + +// translateMiddlewareCaptureConfig builds the per-target capture +// limits used by the middleware chain. Returns nil when the options +// are nil or no capture field is set. Negative caps are normalised to +// zero; oversized caps are clamped to middleware.MaxBodyCapBytes. +func translateMiddlewareCaptureConfig(targetID string, opts *proto.PathTargetOptions) *bodytap.Config { + if opts == nil { + return nil + } + reqCap := clampMiddlewareCaptureBytes(targetID, "request", opts.GetCaptureMaxRequestBytes()) + respCap := clampMiddlewareCaptureBytes(targetID, "response", opts.GetCaptureMaxResponseBytes()) + types := opts.GetCaptureContentTypes() + if reqCap == 0 && respCap == 0 && len(types) == 0 { + return nil + } + return &bodytap.Config{ + MaxRequestBytes: reqCap, + MaxResponseBytes: respCap, + ContentTypes: types, + } +} + +func clampMiddlewareCaptureBytes(targetID, direction string, v int64) int64 { + if v < 0 { + log.Debugf("target %s %s capture cap %d clamped to 0", targetID, direction, v) + return 0 + } + if v > middleware.MaxBodyCapBytes { + log.Debugf("target %s %s capture cap %d clamped to %d", targetID, direction, v, middleware.MaxBodyCapBytes) + return middleware.MaxBodyCapBytes + } + return v +} + +// translateMiddlewareConfigs converts the proto MiddlewareConfig list +// into validated middleware.Spec values. The list is truncated to +// middleware.MaxMiddlewaresPerChain when the caller exceeds the cap. +// Entries with empty IDs, unknown IDs (when registry is non-nil), or +// unspecified slots are skipped with a warn log. Timeouts are clamped +// to [MinTimeout, MaxTimeout] and zero substitutes for DefaultTimeout. +// Returns nil when the resulting slice is empty so callers can leave +// PathTarget.Middlewares unset. +func translateMiddlewareConfigs( + ctx context.Context, + targetID string, + in []*proto.MiddlewareConfig, + registry *middleware.Registry, +) []middleware.Spec { + _ = ctx + if len(in) == 0 { + return nil + } + if len(in) > middleware.MaxMiddlewaresPerChain { + log.Warnf("middleware list for target %q truncated: %d entries exceeds cap of %d", + targetID, len(in), middleware.MaxMiddlewaresPerChain) + in = in[:middleware.MaxMiddlewaresPerChain] + } + + out := make([]middleware.Spec, 0, len(in)) + for _, cfg := range in { + spec, ok := translateMiddlewareConfig(targetID, cfg, registry) + if !ok { + continue + } + out = append(out, spec) + } + if len(out) == 0 { + return nil + } + return out +} + +// translateMiddlewareConfig validates and converts a single +// MiddlewareConfig. The second return value is false when the entry +// must be dropped from the chain. +func translateMiddlewareConfig(targetID string, cfg *proto.MiddlewareConfig, registry *middleware.Registry) (middleware.Spec, bool) { + if cfg == nil { + return middleware.Spec{}, false + } + id := cfg.GetId() + if id == "" { + log.Warnf("middleware config for target %q dropped: empty middleware id", targetID) + return middleware.Spec{}, false + } + if registry != nil && !registry.IsKnown(id) { + log.Warnf("unknown middleware %q configured for target %s; dropping", id, targetID) + return middleware.Spec{}, false + } + slot, ok := protoToMiddlewareSlot(cfg.GetSlot()) + if !ok { + log.Warnf("middleware %q on target %q dropped: slot is unspecified", id, targetID) + return middleware.Spec{}, false + } + + var rawConfig []byte + if src := cfg.GetConfigJson(); len(src) > 0 { + rawConfig = append([]byte(nil), src...) + } + + return middleware.Spec{ + ID: id, + Slot: slot, + Enabled: cfg.GetEnabled(), + FailMode: protoToMiddlewareFailMode(cfg.GetFailMode()), + Timeout: clampMiddlewareTimeout(id, cfg.GetTimeout().AsDuration()), + RawConfig: rawConfig, + CanMutate: cfg.GetCanMutate(), + }, true +} + +// protoToMiddlewareSlot maps the proto slot enum onto the internal +// middleware.Slot. Returns ok=false for the UNSPECIFIED value so the +// translator can drop the entry. +func protoToMiddlewareSlot(s proto.MiddlewareSlot) (middleware.Slot, bool) { + switch s { + case proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST: + return middleware.SlotOnRequest, true + case proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_RESPONSE: + return middleware.SlotOnResponse, true + case proto.MiddlewareSlot_MIDDLEWARE_SLOT_TERMINAL: + return middleware.SlotTerminal, true + default: + return 0, false + } +} + +// protoToMiddlewareFailMode maps the proto FailMode enum onto the +// internal middleware.FailMode, defaulting to FailOpen for any value +// other than FAIL_CLOSED. +func protoToMiddlewareFailMode(m proto.MiddlewareConfig_FailMode) middleware.FailMode { + if m == proto.MiddlewareConfig_FAIL_CLOSED { + return middleware.FailClosed + } + return middleware.FailOpen +} + +// clampMiddlewareTimeout enforces the proxy-wide [MinTimeout, MaxTimeout] +// bounds and substitutes DefaultTimeout for zero inputs. A warn is logged +// only on an actual clamp, not when filling the default. +func clampMiddlewareTimeout(id string, d time.Duration) time.Duration { + if d <= 0 { + return middleware.DefaultTimeout + } + if d < middleware.MinTimeout { + log.Debugf("middleware %s timeout %s clamped to %s", id, d, middleware.MinTimeout) + return middleware.MinTimeout + } + if d > middleware.MaxTimeout { + log.Debugf("middleware %s timeout %s clamped to %s", id, d, middleware.MaxTimeout) + return middleware.MaxTimeout + } + return d +} diff --git a/proxy/middleware_translate_test.go b/proxy/middleware_translate_test.go new file mode 100644 index 000000000..1a956090c --- /dev/null +++ b/proxy/middleware_translate_test.go @@ -0,0 +1,246 @@ +package proxy + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/types/known/durationpb" + + "github.com/netbirdio/netbird/proxy/internal/middleware" + "github.com/netbirdio/netbird/shared/management/proto" +) + +// stubFactory builds a stub Middleware so the registry's IsKnown check +// passes for the configured id. The translator never invokes the +// middleware, so the methods only need to satisfy the interface. +type stubFactory struct { + id string + slot middleware.Slot +} + +func (f stubFactory) ID() string { return f.id } +func (f stubFactory) New(_ []byte) (middleware.Middleware, error) { + return stubMiddleware(f), nil +} + +type stubMiddleware struct { + id string + slot middleware.Slot +} + +func (m stubMiddleware) ID() string { return m.id } +func (m stubMiddleware) Version() string { return "test" } +func (m stubMiddleware) Slot() middleware.Slot { return m.slot } +func (m stubMiddleware) AcceptedContentTypes() []string { return nil } +func (m stubMiddleware) MetadataKeys() []string { return nil } +func (m stubMiddleware) MutationsSupported() bool { return false } +func (m stubMiddleware) Close() error { return nil } +func (m stubMiddleware) Invoke(context.Context, *middleware.Input) (*middleware.Output, error) { + panic("stubMiddleware.Invoke must not be called in translator tests") +} + +// newTestRegistry returns a fresh registry pre-populated with the given +// middleware ids in the matching slot. +func newTestRegistry(t *testing.T, entries map[string]middleware.Slot) *middleware.Registry { + t.Helper() + r := middleware.NewRegistry() + for id, slot := range entries { + require.NoError(t, r.Register(stubFactory{id: id, slot: slot}), "stub registration must succeed") + } + return r +} + +func TestTranslateMiddlewareConfigs_EmptyInput(t *testing.T) { + assert.Nil(t, translateMiddlewareConfigs(context.Background(), "target-a", nil, nil), + "nil input should translate to nil") + assert.Nil(t, translateMiddlewareConfigs(context.Background(), "target-a", []*proto.MiddlewareConfig{}, nil), + "empty input should translate to nil") +} + +func TestTranslateMiddlewareConfigs_KnownIDs(t *testing.T) { + registry := newTestRegistry(t, map[string]middleware.Slot{ + "llm_request_parser": middleware.SlotOnRequest, + "llm_response_parser": middleware.SlotOnResponse, + }) + in := []*proto.MiddlewareConfig{ + { + Id: "llm_request_parser", + Enabled: true, + Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, + ConfigJson: []byte(`{"foo":"bar"}`), + FailMode: proto.MiddlewareConfig_FAIL_OPEN, + Timeout: durationpb.New(250 * time.Millisecond), + CanMutate: true, + }, + { + Id: "llm_response_parser", + Enabled: false, + Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_RESPONSE, + ConfigJson: nil, + FailMode: proto.MiddlewareConfig_FAIL_CLOSED, + Timeout: durationpb.New(50 * time.Millisecond), + }, + } + + out := translateMiddlewareConfigs(context.Background(), "target-a", in, registry) + require.Len(t, out, 2, "two known middlewares should produce two specs") + + assert.Equal(t, "llm_request_parser", out[0].ID, "first id should match") + assert.Equal(t, middleware.SlotOnRequest, out[0].Slot, "first slot should be on_request") + assert.True(t, out[0].Enabled, "first spec should be enabled") + assert.Equal(t, middleware.FailOpen, out[0].FailMode, "first spec should be fail-open") + assert.Equal(t, 250*time.Millisecond, out[0].Timeout, "first spec timeout should pass through") + assert.True(t, out[0].CanMutate, "first spec should permit mutations") + assert.Equal(t, []byte(`{"foo":"bar"}`), out[0].RawConfig, "first spec raw config should match") + + assert.Equal(t, "llm_response_parser", out[1].ID, "second id should match") + assert.Equal(t, middleware.SlotOnResponse, out[1].Slot, "second slot should be on_response") + assert.False(t, out[1].Enabled, "second spec should be disabled") + assert.Equal(t, middleware.FailClosed, out[1].FailMode, "second spec should be fail-closed") + assert.Equal(t, 50*time.Millisecond, out[1].Timeout, "second spec timeout should pass through") + assert.Nil(t, out[1].RawConfig, "second spec raw config should be nil") +} + +func TestTranslateMiddlewareConfigs_UnknownIDSkipped(t *testing.T) { + registry := newTestRegistry(t, map[string]middleware.Slot{ + "llm_request_parser": middleware.SlotOnRequest, + }) + in := []*proto.MiddlewareConfig{ + {Id: "llm_request_parser", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST}, + {Id: "not_registered", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST}, + } + out := translateMiddlewareConfigs(context.Background(), "target-unknown", in, registry) + require.Len(t, out, 1, "unknown id must be skipped") + assert.Equal(t, "llm_request_parser", out[0].ID, "remaining entry should be the known one") +} + +func TestTranslateMiddlewareConfigs_NilRegistrySkipsValidation(t *testing.T) { + in := []*proto.MiddlewareConfig{ + {Id: "anything_goes", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST}, + } + out := translateMiddlewareConfigs(context.Background(), "target-nilreg", in, nil) + require.Len(t, out, 1, "nil registry must accept any non-empty id") + assert.Equal(t, "anything_goes", out[0].ID, "id should pass through unchecked") +} + +func TestTranslateMiddlewareConfigs_TimeoutClamps(t *testing.T) { + registry := newTestRegistry(t, map[string]middleware.Slot{ + "llm_request_parser": middleware.SlotOnRequest, + }) + in := []*proto.MiddlewareConfig{ + {Id: "llm_request_parser", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, Timeout: nil}, + {Id: "llm_request_parser", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, Timeout: durationpb.New(time.Microsecond)}, + {Id: "llm_request_parser", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, Timeout: durationpb.New(time.Hour)}, + } + out := translateMiddlewareConfigs(context.Background(), "target-clamp", in, registry) + require.Len(t, out, 3, "clamping must keep all three entries") + assert.Equal(t, middleware.DefaultTimeout, out[0].Timeout, "zero timeout should default") + assert.Equal(t, middleware.MinTimeout, out[1].Timeout, "below-min timeout should clamp up") + assert.Equal(t, middleware.MaxTimeout, out[2].Timeout, "above-max timeout should clamp down") +} + +func TestTranslateMiddlewareConfigs_FailModeMapping(t *testing.T) { + registry := newTestRegistry(t, map[string]middleware.Slot{ + "llm_request_parser": middleware.SlotOnRequest, + }) + in := []*proto.MiddlewareConfig{ + {Id: "llm_request_parser", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST}, + {Id: "llm_request_parser", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, FailMode: proto.MiddlewareConfig_FAIL_CLOSED}, + } + out := translateMiddlewareConfigs(context.Background(), "target-failmode", in, registry) + require.Len(t, out, 2, "both entries should translate") + assert.Equal(t, middleware.FailOpen, out[0].FailMode, "default fail mode should be open") + assert.Equal(t, middleware.FailClosed, out[1].FailMode, "explicit fail closed should map") +} + +func TestTranslateMiddlewareConfigs_SlotMapping(t *testing.T) { + registry := newTestRegistry(t, map[string]middleware.Slot{ + "req": middleware.SlotOnRequest, + "resp": middleware.SlotOnResponse, + "term": middleware.SlotTerminal, + }) + in := []*proto.MiddlewareConfig{ + {Id: "req", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST}, + {Id: "resp", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_RESPONSE}, + {Id: "term", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_TERMINAL}, + {Id: "req", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_UNSPECIFIED}, + } + out := translateMiddlewareConfigs(context.Background(), "target-slot", in, registry) + require.Len(t, out, 3, "unspecified slot entry must be skipped") + assert.Equal(t, middleware.SlotOnRequest, out[0].Slot, "on_request slot mapping") + assert.Equal(t, middleware.SlotOnResponse, out[1].Slot, "on_response slot mapping") + assert.Equal(t, middleware.SlotTerminal, out[2].Slot, "terminal slot mapping") +} + +func TestTranslateMiddlewareConfigs_EmptyIDSkipped(t *testing.T) { + registry := newTestRegistry(t, map[string]middleware.Slot{ + "llm_request_parser": middleware.SlotOnRequest, + }) + in := []*proto.MiddlewareConfig{ + {Id: "", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST}, + {Id: "llm_request_parser", Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST}, + } + out := translateMiddlewareConfigs(context.Background(), "target-empty-id", in, registry) + require.Len(t, out, 1, "empty id must be dropped") + assert.Equal(t, "llm_request_parser", out[0].ID, "remaining entry should be valid") +} + +// TestTranslateMiddlewareConfigs_TruncatesAboveCap proves the translator +// truncates lists that exceed MaxMiddlewaresPerChain rather than dropping +// the whole slice, matching the documented G3 behaviour. +func TestTranslateMiddlewareConfigs_TruncatesAboveCap(t *testing.T) { + registry := newTestRegistry(t, map[string]middleware.Slot{ + "llm_request_parser": middleware.SlotOnRequest, + }) + overCap := middleware.MaxMiddlewaresPerChain + 1 + in := make([]*proto.MiddlewareConfig, 0, overCap) + for i := 0; i < overCap; i++ { + in = append(in, &proto.MiddlewareConfig{ + Id: "llm_request_parser", + Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, + }) + } + out := translateMiddlewareConfigs(context.Background(), "target-truncate", in, registry) + assert.Len(t, out, middleware.MaxMiddlewaresPerChain, "over-cap input must be truncated to MaxMiddlewaresPerChain") +} + +func TestTranslateMiddlewareConfigs_AllowsListAtCap(t *testing.T) { + registry := newTestRegistry(t, map[string]middleware.Slot{ + "llm_request_parser": middleware.SlotOnRequest, + }) + in := make([]*proto.MiddlewareConfig, 0, middleware.MaxMiddlewaresPerChain) + for i := 0; i < middleware.MaxMiddlewaresPerChain; i++ { + in = append(in, &proto.MiddlewareConfig{ + Id: "llm_request_parser", + Slot: proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, + }) + } + out := translateMiddlewareConfigs(context.Background(), "target-cap", in, registry) + assert.Len(t, out, middleware.MaxMiddlewaresPerChain, "list at the cap boundary must translate fully") +} + +func TestProtoToMiddlewareSlot(t *testing.T) { + cases := []struct { + name string + in proto.MiddlewareSlot + want middleware.Slot + wantOk bool + }{ + {"unspecified", proto.MiddlewareSlot_MIDDLEWARE_SLOT_UNSPECIFIED, 0, false}, + {"on_request", proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_REQUEST, middleware.SlotOnRequest, true}, + {"on_response", proto.MiddlewareSlot_MIDDLEWARE_SLOT_ON_RESPONSE, middleware.SlotOnResponse, true}, + {"terminal", proto.MiddlewareSlot_MIDDLEWARE_SLOT_TERMINAL, middleware.SlotTerminal, true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, ok := protoToMiddlewareSlot(tc.in) + assert.Equal(t, tc.wantOk, ok, "ok flag for %s", tc.name) + if tc.wantOk { + assert.Equal(t, tc.want, got, "slot mapping for %s", tc.name) + } + }) + } +} diff --git a/proxy/server.go b/proxy/server.go index 1d8a2451b..f28d580bd 100644 --- a/proxy/server.go +++ b/proxy/server.go @@ -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