mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-08 16:01:29 +02:00
[management,proxy] Reduce agent-network cognitive complexity (#6566)
Address the SonarCloud quality-gate findings in new agent-network code by extracting focused helpers. No behavior change. - synthesizer.go: split buildIdentityInjectConfigJSON into per-shape rule builders; extract mergeGuardrail from mergeGuardrails to cut nesting depth. - llm_identity_inject: extract injectionEmitsAnything validation predicate from New. - llm_response_parser/streaming.go: extract applyOpenAIStreamUsage and applyAnthropicStreamUsage (via a named anthropicStreamUsage type) and simplify the OpenAI scanner loop. - reverseproxy.go: decompose ServeHTTP into serveRouteError, buildTargetContext, serveDirect, serveWithChain, captureRequestForChain, serveDeny, newResponseWriter, observeResponse, and forwardUpstream, preserving the defer ordering so response observation still reads the captured writer before it is released.
This commit is contained in:
@@ -40,8 +40,8 @@ const Version = "1.0.0"
|
||||
// Middleware stamps NetBird identity onto upstream requests for the
|
||||
// configured set of resolved providers.
|
||||
type Middleware struct {
|
||||
cfg Config
|
||||
byID map[string]ProviderInjection
|
||||
cfg Config
|
||||
byID map[string]ProviderInjection
|
||||
}
|
||||
|
||||
// New constructs a Middleware from the supplied configuration. A nil
|
||||
@@ -49,45 +49,44 @@ type Middleware struct {
|
||||
func New(cfg Config) *Middleware {
|
||||
byID := make(map[string]ProviderInjection, len(cfg.Providers))
|
||||
for _, p := range cfg.Providers {
|
||||
if p.ProviderID == "" {
|
||||
if p.ProviderID == "" || !injectionEmitsAnything(p) {
|
||||
continue
|
||||
}
|
||||
// Drop entries that wouldn't inject anything — keeps the
|
||||
// runtime check tight. Also drop entries that set both
|
||||
// shapes (configuration error; refuse to guess which wins).
|
||||
// Extras alone are enough to keep the rule alive even if
|
||||
// neither identity shape is set.
|
||||
hasExtras := false
|
||||
for _, e := range p.ExtraHeaders {
|
||||
if e.Name != "" && e.Value != "" {
|
||||
hasExtras = true
|
||||
break
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case p.HeaderPair != nil && p.JSONMetadata != nil:
|
||||
continue
|
||||
case p.HeaderPair != nil:
|
||||
if p.HeaderPair.EndUserIDHeader == "" && p.HeaderPair.TagsHeader == "" && !p.HeaderPair.TagsInBody && !p.HeaderPair.EndUserIDInBody && !hasExtras {
|
||||
continue
|
||||
}
|
||||
case p.JSONMetadata != nil:
|
||||
if p.JSONMetadata.Header == "" {
|
||||
continue
|
||||
}
|
||||
if p.JSONMetadata.UserKey == "" && p.JSONMetadata.GroupsKey == "" && !hasExtras {
|
||||
continue
|
||||
}
|
||||
default:
|
||||
if !hasExtras {
|
||||
continue
|
||||
}
|
||||
}
|
||||
byID[p.ProviderID] = p
|
||||
}
|
||||
return &Middleware{cfg: cfg, byID: byID}
|
||||
}
|
||||
|
||||
// injectionEmitsAnything reports whether a provider injection rule would
|
||||
// stamp anything at runtime. Rules that set both identity shapes are a
|
||||
// configuration error (we refuse to guess which wins), and rules that
|
||||
// resolve to no headers are dropped to keep the runtime check tight.
|
||||
// Non-empty extras alone keep a rule alive even when neither identity
|
||||
// shape is set.
|
||||
func injectionEmitsAnything(p ProviderInjection) bool {
|
||||
hasExtras := false
|
||||
for _, e := range p.ExtraHeaders {
|
||||
if e.Name != "" && e.Value != "" {
|
||||
hasExtras = true
|
||||
break
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case p.HeaderPair != nil && p.JSONMetadata != nil:
|
||||
return false
|
||||
case p.HeaderPair != nil:
|
||||
return p.HeaderPair.EndUserIDHeader != "" || p.HeaderPair.TagsHeader != "" ||
|
||||
p.HeaderPair.TagsInBody || p.HeaderPair.EndUserIDInBody || hasExtras
|
||||
case p.JSONMetadata != nil:
|
||||
if p.JSONMetadata.Header == "" {
|
||||
return false
|
||||
}
|
||||
return p.JSONMetadata.UserKey != "" || p.JSONMetadata.GroupsKey != "" || hasExtras
|
||||
default:
|
||||
return hasExtras
|
||||
}
|
||||
}
|
||||
|
||||
// ID returns the registry identifier.
|
||||
func (m *Middleware) ID() string { return ID }
|
||||
|
||||
|
||||
@@ -84,15 +84,12 @@ func accumulateOpenAIStream(body []byte) (llm.Usage, string) {
|
||||
for {
|
||||
ev, err := scanner.Next()
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
break
|
||||
}
|
||||
if ev.Data == "" || ev.Data == openAIDoneSentinel {
|
||||
if ev.Data == openAIDoneSentinel {
|
||||
break
|
||||
}
|
||||
if ev.Data == openAIDoneSentinel {
|
||||
break
|
||||
}
|
||||
if ev.Data == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -113,26 +110,35 @@ func accumulateOpenAIStream(body []byte) (llm.Usage, string) {
|
||||
if u == nil && chunk.Response != nil {
|
||||
u = chunk.Response.Usage
|
||||
}
|
||||
if u != nil {
|
||||
usage.InputTokens = pickInt64(u.InputTokens, u.PromptTokens)
|
||||
usage.OutputTokens = pickInt64(u.OutputTokens, u.CompletionTokens)
|
||||
usage.TotalTokens = derefInt64(u.TotalTokens)
|
||||
if u.InputTokensDetails != nil {
|
||||
if v := derefInt64(u.InputTokensDetails.CachedTokens); v > 0 {
|
||||
usage.CachedInputTokens = v
|
||||
}
|
||||
}
|
||||
if usage.CachedInputTokens == 0 && u.PromptTokensDetails != nil {
|
||||
usage.CachedInputTokens = derefInt64(u.PromptTokensDetails.CachedTokens)
|
||||
}
|
||||
if usage.TotalTokens == 0 && (usage.InputTokens > 0 || usage.OutputTokens > 0) {
|
||||
usage.TotalTokens = usage.InputTokens + usage.OutputTokens
|
||||
}
|
||||
}
|
||||
applyOpenAIStreamUsage(u, &usage)
|
||||
}
|
||||
return usage, completion.String()
|
||||
}
|
||||
|
||||
// applyOpenAIStreamUsage lifts the token counts off a final-frame usage
|
||||
// block into the running usage, normalising the chat.completions
|
||||
// (prompt_/completion_) and Responses-API (input_/output_) names and
|
||||
// backfilling total tokens when the provider omits them.
|
||||
func applyOpenAIStreamUsage(u *openAIStreamUsage, usage *llm.Usage) {
|
||||
if u == nil {
|
||||
return
|
||||
}
|
||||
usage.InputTokens = pickInt64(u.InputTokens, u.PromptTokens)
|
||||
usage.OutputTokens = pickInt64(u.OutputTokens, u.CompletionTokens)
|
||||
usage.TotalTokens = derefInt64(u.TotalTokens)
|
||||
if u.InputTokensDetails != nil {
|
||||
if v := derefInt64(u.InputTokensDetails.CachedTokens); v > 0 {
|
||||
usage.CachedInputTokens = v
|
||||
}
|
||||
}
|
||||
if usage.CachedInputTokens == 0 && u.PromptTokensDetails != nil {
|
||||
usage.CachedInputTokens = derefInt64(u.PromptTokensDetails.CachedTokens)
|
||||
}
|
||||
if usage.TotalTokens == 0 && (usage.InputTokens > 0 || usage.OutputTokens > 0) {
|
||||
usage.TotalTokens = usage.InputTokens + usage.OutputTokens
|
||||
}
|
||||
}
|
||||
|
||||
// decodeJSONString unmarshals a JSON-encoded string value, returning
|
||||
// ok=false when the raw message is empty or not a string.
|
||||
func decodeJSONString(raw json.RawMessage) (string, bool) {
|
||||
@@ -149,26 +155,23 @@ func decodeJSONString(raw json.RawMessage) (string, bool) {
|
||||
// anthropicStreamEvent captures the union of Messages-API stream event
|
||||
// payloads we care about. Each named event on the wire fills only its
|
||||
// shape's fields; unknown keys are ignored.
|
||||
type anthropicStreamUsage struct {
|
||||
InputTokens *int64 `json:"input_tokens"`
|
||||
OutputTokens *int64 `json:"output_tokens"`
|
||||
CacheReadInputTokens *int64 `json:"cache_read_input_tokens"`
|
||||
CacheCreationInputTokens *int64 `json:"cache_creation_input_tokens"`
|
||||
}
|
||||
|
||||
type anthropicStreamEvent struct {
|
||||
Type string `json:"type"`
|
||||
Message *struct {
|
||||
Usage *struct {
|
||||
InputTokens *int64 `json:"input_tokens"`
|
||||
OutputTokens *int64 `json:"output_tokens"`
|
||||
CacheReadInputTokens *int64 `json:"cache_read_input_tokens"`
|
||||
CacheCreationInputTokens *int64 `json:"cache_creation_input_tokens"`
|
||||
} `json:"usage"`
|
||||
Usage *anthropicStreamUsage `json:"usage"`
|
||||
} `json:"message"`
|
||||
Delta *struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
} `json:"delta"`
|
||||
Usage *struct {
|
||||
InputTokens *int64 `json:"input_tokens"`
|
||||
OutputTokens *int64 `json:"output_tokens"`
|
||||
CacheReadInputTokens *int64 `json:"cache_read_input_tokens"`
|
||||
CacheCreationInputTokens *int64 `json:"cache_creation_input_tokens"`
|
||||
} `json:"usage"`
|
||||
Usage *anthropicStreamUsage `json:"usage"`
|
||||
}
|
||||
|
||||
// accumulateAnthropicStream tracks input_tokens from message_start,
|
||||
@@ -216,44 +219,42 @@ func accumulateAnthropicStream(body []byte) (llm.Usage, string) {
|
||||
func applyAnthropicStreamEvent(eventType string, payload anthropicStreamEvent, usage *llm.Usage, completion *strings.Builder) {
|
||||
switch eventType {
|
||||
case "message_start":
|
||||
if payload.Message != nil && payload.Message.Usage != nil {
|
||||
if v := derefInt64(payload.Message.Usage.InputTokens); v > 0 {
|
||||
usage.InputTokens = v
|
||||
}
|
||||
if v := derefInt64(payload.Message.Usage.OutputTokens); v > 0 {
|
||||
usage.OutputTokens = v
|
||||
}
|
||||
if v := derefInt64(payload.Message.Usage.CacheReadInputTokens); v > 0 {
|
||||
usage.CachedInputTokens = v
|
||||
}
|
||||
if v := derefInt64(payload.Message.Usage.CacheCreationInputTokens); v > 0 {
|
||||
usage.CacheCreationTokens = v
|
||||
}
|
||||
if payload.Message != nil {
|
||||
applyAnthropicStreamUsage(payload.Message.Usage, usage)
|
||||
}
|
||||
case "content_block_delta":
|
||||
if payload.Delta != nil && payload.Delta.Type == "text_delta" {
|
||||
completion.WriteString(payload.Delta.Text)
|
||||
}
|
||||
case "message_delta":
|
||||
if payload.Usage != nil {
|
||||
if v := derefInt64(payload.Usage.InputTokens); v > 0 {
|
||||
usage.InputTokens = v
|
||||
}
|
||||
if v := derefInt64(payload.Usage.OutputTokens); v > 0 {
|
||||
usage.OutputTokens = v
|
||||
}
|
||||
if v := derefInt64(payload.Usage.CacheReadInputTokens); v > 0 {
|
||||
usage.CachedInputTokens = v
|
||||
}
|
||||
if v := derefInt64(payload.Usage.CacheCreationInputTokens); v > 0 {
|
||||
usage.CacheCreationTokens = v
|
||||
}
|
||||
}
|
||||
applyAnthropicStreamUsage(payload.Usage, usage)
|
||||
case "message_stop":
|
||||
// No-op; Anthropic does not emit usage here.
|
||||
}
|
||||
}
|
||||
|
||||
// applyAnthropicStreamUsage folds a non-nil Anthropic usage block into the
|
||||
// running totals. Each field overwrites only when present and positive, so
|
||||
// message_delta's post-completion counts supersede the message_start seed
|
||||
// without zeroing dimensions a later event omits.
|
||||
func applyAnthropicStreamUsage(u *anthropicStreamUsage, usage *llm.Usage) {
|
||||
if u == nil {
|
||||
return
|
||||
}
|
||||
if v := derefInt64(u.InputTokens); v > 0 {
|
||||
usage.InputTokens = v
|
||||
}
|
||||
if v := derefInt64(u.OutputTokens); v > 0 {
|
||||
usage.OutputTokens = v
|
||||
}
|
||||
if v := derefInt64(u.CacheReadInputTokens); v > 0 {
|
||||
usage.CachedInputTokens = v
|
||||
}
|
||||
if v := derefInt64(u.CacheCreationInputTokens); v > 0 {
|
||||
usage.CacheCreationTokens = v
|
||||
}
|
||||
}
|
||||
|
||||
func pickInt64(preferred, fallback *int64) int64 {
|
||||
if preferred != nil {
|
||||
return *preferred
|
||||
|
||||
@@ -83,13 +83,8 @@ func NewReverseProxy(transport http.RoundTripper, forwardedProto string, trusted
|
||||
func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
result, exists := p.findTargetForRequest(r)
|
||||
if !exists {
|
||||
if cd := CapturedDataFromContext(r.Context()); cd != nil {
|
||||
cd.SetOrigin(OriginNoRoute)
|
||||
}
|
||||
requestID := getRequestID(r)
|
||||
web.ServeErrorPage(w, r, http.StatusNotFound, "Service Not Found",
|
||||
"The requested service could not be found. Please check the URL, try refreshing, or check if the peer is running. If that doesn't work, see our documentation for help.",
|
||||
requestID, web.ErrorStatus{Proxy: true, Destination: false})
|
||||
p.serveRouteError(w, r, http.StatusNotFound, "Service Not Found",
|
||||
"The requested service could not be found. Please check the URL, try refreshing, or check if the peer is running. If that doesn't work, see our documentation for help.")
|
||||
return
|
||||
}
|
||||
|
||||
@@ -99,19 +94,13 @@ func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
// with 421 (Misdirected Request) so the caller sees an explicit
|
||||
// error instead of silently doubling tunnel traffic.
|
||||
if p.isSelfTargetLoop(r, result.target.URL) {
|
||||
if cd := CapturedDataFromContext(r.Context()); cd != nil {
|
||||
cd.SetOrigin(OriginNoRoute)
|
||||
}
|
||||
requestID := getRequestID(r)
|
||||
web.ServeErrorPage(w, r, http.StatusMisdirectedRequest, "Loop Detected",
|
||||
"This peer is the target of the requested service. Reach the backend directly instead of dialing the public service URL from the same machine.",
|
||||
requestID, web.ErrorStatus{Proxy: true, Destination: false})
|
||||
p.serveRouteError(w, r, http.StatusMisdirectedRequest, "Loop Detected",
|
||||
"This peer is the target of the requested service. Reach the backend directly instead of dialing the public service URL from the same machine.")
|
||||
return
|
||||
}
|
||||
|
||||
ctx := r.Context()
|
||||
// Set the account ID in the context for the roundtripper to use.
|
||||
ctx = roundtrip.WithAccountID(ctx, result.accountID)
|
||||
pt := result.target
|
||||
ctx := p.buildTargetContext(r.Context(), result)
|
||||
|
||||
// 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
|
||||
@@ -124,8 +113,34 @@ func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
capturedData.SetSuppressAccessLog(result.target != nil && result.target.DisableAccessLog)
|
||||
}
|
||||
|
||||
pt := result.target
|
||||
rewriteMatchedPath := result.matchedPath
|
||||
if pt.PathRewrite == PathRewritePreserve {
|
||||
rewriteMatchedPath = ""
|
||||
}
|
||||
|
||||
chain := p.resolveChain(result)
|
||||
if chain == nil || chain.Empty() {
|
||||
p.serveDirect(w, r, ctx, result, rewriteMatchedPath)
|
||||
return
|
||||
}
|
||||
p.serveWithChain(w, r, ctx, result, chain, rewriteMatchedPath, capturedData)
|
||||
}
|
||||
|
||||
// serveRouteError marks the request as un-routed on any captured-data
|
||||
// context and renders the proxy error page.
|
||||
func (p *ReverseProxy) serveRouteError(w http.ResponseWriter, r *http.Request, status int, title, message string) {
|
||||
if cd := CapturedDataFromContext(r.Context()); cd != nil {
|
||||
cd.SetOrigin(OriginNoRoute)
|
||||
}
|
||||
web.ServeErrorPage(w, r, status, title, message, getRequestID(r),
|
||||
web.ErrorStatus{Proxy: true, Destination: false})
|
||||
}
|
||||
|
||||
// buildTargetContext layers the per-target roundtrip flags (account id,
|
||||
// TLS-verify skip, direct upstream, dial timeout) onto the request context.
|
||||
func (p *ReverseProxy) buildTargetContext(ctx context.Context, result targetResult) context.Context {
|
||||
pt := result.target
|
||||
ctx = roundtrip.WithAccountID(ctx, result.accountID)
|
||||
if pt.SkipTLSVerify {
|
||||
ctx = roundtrip.WithSkipTLSVerify(ctx)
|
||||
}
|
||||
@@ -135,32 +150,66 @@ func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
if pt.RequestTimeout > 0 {
|
||||
ctx = types.WithDialTimeout(ctx, pt.RequestTimeout)
|
||||
}
|
||||
return ctx
|
||||
}
|
||||
|
||||
rewriteMatchedPath := result.matchedPath
|
||||
if pt.PathRewrite == PathRewritePreserve {
|
||||
rewriteMatchedPath = ""
|
||||
// serveDirect forwards the request without a middleware chain — the common
|
||||
// path for plain reverse-proxy targets.
|
||||
func (p *ReverseProxy) serveDirect(w http.ResponseWriter, r *http.Request, ctx context.Context, result targetResult, rewriteMatchedPath string) {
|
||||
pt := result.target
|
||||
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,
|
||||
}
|
||||
|
||||
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
|
||||
if result.rewriteRedirects {
|
||||
rp.ModifyResponse = p.rewriteLocationFunc(pt.URL, rewriteMatchedPath, r) //nolint:bodyclose
|
||||
}
|
||||
rp.ServeHTTP(w, r.WithContext(ctx))
|
||||
}
|
||||
|
||||
// serveWithChain runs the per-target middleware chain around the upstream
|
||||
// request: request-leg capture and authorisation, then (on allow) the
|
||||
// upstream forward with response/terminal observation deferred so it reads
|
||||
// the captured response before the writer is released.
|
||||
func (p *ReverseProxy) serveWithChain(w http.ResponseWriter, r *http.Request, ctx context.Context, result targetResult, chain *middleware.Chain, rewriteMatchedPath string, capturedData *CapturedData) {
|
||||
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())
|
||||
capturedBody, truncated, originalSize, releaseBudget := p.captureRequestForChain(ctx, r, result, capturedData)
|
||||
defer releaseBudget()
|
||||
|
||||
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 {
|
||||
p.serveDeny(w, denyOutput, result, middlewareIDs)
|
||||
return
|
||||
}
|
||||
|
||||
respWriter, capturingWriter := p.newResponseWriter(ctx, w, result, capturedData)
|
||||
if capturingWriter != nil {
|
||||
defer capturingWriter.Release()
|
||||
defer p.observeResponse(ctx, chain, acc, reqInput, requestMeta, capturingWriter, w, capturedData, result, middlewareIDs)
|
||||
}
|
||||
|
||||
p.forwardUpstream(respWriter, r, ctx, result, rewriteMatchedPath, upstreamRewrite)
|
||||
}
|
||||
|
||||
// captureRequestForChain copies the request body for inspection by the
|
||||
// chain, records any capture bypass, and applies agent-network routing
|
||||
// recovery for oversized bodies. The returned release frees the capture
|
||||
// budget and must be deferred by the caller.
|
||||
func (p *ReverseProxy) captureRequestForChain(ctx context.Context, r *http.Request, result targetResult, capturedData *CapturedData) ([]byte, bool, int64, func()) {
|
||||
pt := result.target
|
||||
capturedBody, truncated, originalSize, bypass, releaseBudget, captureErr := bodytap.CaptureRequest(r, pt.CaptureConfig, p.middlewareManager.Budget())
|
||||
if captureErr != nil {
|
||||
p.logger.Debugf("middleware request body capture error: %v", captureErr)
|
||||
}
|
||||
@@ -184,112 +233,114 @@ func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
p.logger.Debugf("agent-network routing recovery: extracted model=%s stream=%t from oversized request body (service=%s)", model, stream, result.serviceID)
|
||||
}
|
||||
}
|
||||
return capturedBody, truncated, originalSize, releaseBudget
|
||||
}
|
||||
|
||||
acc := middleware.NewAccumulator(middleware.MaxRequestMetadataBytes)
|
||||
reqInput := buildRequestInput(r, result, capturedData, capturedBody, truncated, originalSize)
|
||||
// serveDeny renders the chain's deny response. Policy/budget/routing/guardrail
|
||||
// denials are expected runtime outcomes and can be high-volume under
|
||||
// misconfigured or hostile clients; per-request detail stays at Debug and
|
||||
// metrics/access logs carry the signal at scale.
|
||||
func (p *ReverseProxy) serveDeny(w http.ResponseWriter, denyOutput *middleware.Output, result targetResult, middlewareIDs []string) {
|
||||
middlewareID := "middleware"
|
||||
if denyOutput.DenyReason != nil && denyOutput.DenyReason.Code != "" {
|
||||
middlewareID = denyOutput.DenyReason.Code
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
||||
denyOutput, requestMeta, upstreamRewrite, _ := chain.RunRequest(ctx, r, reqInput, acc)
|
||||
// newResponseWriter returns the writer the upstream forward should use. When
|
||||
// response capture is enabled and not bypassed it wraps w in a capturing
|
||||
// writer (also returned so the caller can release it and feed the response
|
||||
// leg); otherwise the capturing writer is nil and w is used directly.
|
||||
func (p *ReverseProxy) newResponseWriter(ctx context.Context, w http.ResponseWriter, result targetResult, capturedData *CapturedData) (http.ResponseWriter, *bodytap.CapturingResponseWriter) {
|
||||
pt := result.target
|
||||
if pt.CaptureConfig == nil || pt.CaptureConfig.MaxResponseBytes <= 0 {
|
||||
return w, nil
|
||||
}
|
||||
capturingWriter := bodytap.NewCapturingResponseWriter(w, pt.CaptureConfig.MaxResponseBytes, p.middlewareManager.Budget())
|
||||
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.Release()
|
||||
return w, nil
|
||||
}
|
||||
return capturingWriter, capturingWriter
|
||||
}
|
||||
|
||||
// observeResponse runs the response and terminal middleware slots after the
|
||||
// body has been forwarded. It is deferred by serveWithChain so it reads the
|
||||
// captured response before the writer is released.
|
||||
func (p *ReverseProxy) observeResponse(ctx context.Context, chain *middleware.Chain, acc *middleware.Accumulator, reqInput *middleware.Input, requestMeta []middleware.KV, capturingWriter *bodytap.CapturingResponseWriter, w http.ResponseWriter, capturedData *CapturedData, result targetResult, middlewareIDs []string) {
|
||||
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 requestMeta {
|
||||
for _, kv := range respMeta {
|
||||
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
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
|
||||
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))
|
||||
}()
|
||||
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))
|
||||
}
|
||||
|
||||
// forwardUpstream applies any middleware-emitted upstream rewrite and proxies
|
||||
// the request to the effective upstream URL.
|
||||
func (p *ReverseProxy) forwardUpstream(respWriter http.ResponseWriter, r *http.Request, ctx context.Context, result targetResult, rewriteMatchedPath string, upstreamRewrite *middleware.UpstreamRewrite) {
|
||||
pt := result.target
|
||||
effectiveURL := applyUpstreamRewrite(pt.URL, upstreamRewrite)
|
||||
if upstreamRewrite != nil {
|
||||
r.Host = effectiveURL.Host
|
||||
|
||||
Reference in New Issue
Block a user