mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-27 00:59:07 +02:00
[management,proxy] Agent network: per-account LLM gateway (policy, metering, multi-provider) (#6555)
* [agent-network] Shared proto, OpenAPI schema, and generated types * [agent-network] Management: store, manager, synthesizer, policy engine, provider catalog, HTTP/gRPC API Adds the account-scoped agent-network module: provider/policy/budget CRUD and store, the reverse-proxy service synthesizer, policy selection + limit enforcement, the provider catalog (incl. Vertex AI and AWS Bedrock entries), and the management HTTP + proxy gRPC surfaces. * [management] Fix agent-network proxy-peer fan-out on affected-peer recompute The affected-peers resolver loaded only persisted reverse-proxy services, but agent-network services are synthesized on demand and never persisted. As a result the embedded proxy peer was never folded into the affected set when a client's group changed, so the proxy received no network-map update for a newly authorised client and rejected its handshake until a full resync (restart). loadProxyServices now merges the synthesized agent-network services (injected via a registration hook to avoid an import cycle), so proxy peers learn newly authorised clients immediately. * [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] LLM parsers, pricing, and builtin middlewares (OpenAI, Anthropic, Vertex AI, AWS Bedrock) Request/response parsers and SSE/event-stream metering, the embedded pricing table, and the builtin middleware set: request parser, router, policy limit-check/record, cost meter, guardrail, identity inject, response parser. Includes the path-routed providers — Google Vertex AI (keyfile:: service-account OAuth minting) and AWS Bedrock (bearer auth, invoke/converse/streaming, optional /bedrock prefix) — plus the Models allowlist and unmeterable-publisher deny. * [proxy] IPv6 in-place apply and TCP accept-loop hardening on netstack listeners * [agent-network] End-to-end test suite, module docs, and deployment preset * [agent-network] Fix codespell typos and exclude false positives - labelgen word pool: vermillion -> vermilion, racoon -> raccoon. - codespell ignore list: add flate (Go compress/flate package), recordin (a test-local identifier), and unparseable (a valid alternative spelling used consistently across identifiers + a metadata-value constant). * [management] Set LastSeen on injected proxy peer in realstack test (MySQL strict-mode) The injected embedded proxy peer had a PeerStatus with a zero LastSeen, which serializes to '0000-00-00' and is rejected by MySQL in strict mode (SQLite tolerates it). Set LastSeen to a valid time so SaveAccount succeeds on both engines. * [agent-network] Remove e2e shell-script suite from this branch The end-to-end shell scripts under scripts/e2e/ are maintained in a separate testing suite and are not part of this change set. * [agent-network] Polish module docs: remove internal review scaffolding, fix links, verify diagrams Strip PR-review framing, commit references, absolute paths, and stale internal references from the agent-network module docs; fix broken relative links; verify all diagrams against the current architecture. Remove the internal AI-reviewer prompt file. * [management] Refine session expiration handling to support 3-state encoding for SSO deadlines * [agent-network] Relocate agentnetwork package to internals/modules Move management/server/agentnetwork (and its catalog/, labelgen/, types/ subpackages) to management/internals/modules/agentnetwork, alongside the reverse-proxy module, and rewrite all importers. Pure relocation: package names, the synthesizer + affectedpeers registration hook, and store access (shared store.Store) are unchanged, so no import cycle is introduced (affectedpeers still depends only on the agentnetwork/types leaf). * [agent-network] Co-locate HTTP handlers in the module (RegisterEndpoints) Move the agent-network HTTP handlers from server/http/handlers/agentnetwork into the module at internals/modules/agentnetwork/handlers (package handlers) and rename the entrypoint AddEndpoints -> RegisterEndpoints, matching the reverse-proxy module convention. Wiring in http/handler.go updated accordingly.
This commit is contained in:
Maycon Santos
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
Reference in New Issue
Block a user