diff --git a/authbridge/authlib/README.md b/authbridge/authlib/README.md index 6b9a085a5..d8c16ca3c 100644 --- a/authbridge/authlib/README.md +++ b/authbridge/authlib/README.md @@ -14,6 +14,11 @@ A pure Go library providing reusable building blocks for JWT validation, OAuth 2 | `routing/` | Host-to-audience router with glob pattern matching | | `auth/` | Composition layer: `HandleInbound` + `HandleOutbound` | | `config/` | YAML config, mode presets, startup validation, URL derivation | +| `pipeline/` | Plugin pipeline — see [pipeline/README.md](./pipeline/README.md) for the full interface spec and bridge-plugin integration guide | +| `session/` | In-memory session store + `SessionSummary` aggregation, backing the `:9094` API | +| `sessionapi/` | HTTP API (`/v1/sessions`, `/v1/events` SSE, `/v1/pipeline`) exposing the session store | +| `plugins/` | Built-in protocol parsers: `a2a-parser`, `mcp-parser`, `inference-parser`, `jwt-validation` | +| `observe/` | OTEL + metrics helpers | ## Usage diff --git a/authbridge/authlib/pipeline/README.md b/authbridge/authlib/pipeline/README.md new file mode 100644 index 000000000..1b2f26a3b --- /dev/null +++ b/authbridge/authlib/pipeline/README.md @@ -0,0 +1,575 @@ +# pipeline — Plugin Pipeline Specification + +The `pipeline` package defines AuthBridge's plugin contract: how plugins are written, how they communicate through shared state, how they compose into ordered chains, and how those chains run inside each of the three listener modes (ext_proc, ext_authz, forward/reverse proxy). + +This document is the reference for AuthBridge's plugin contract. It covers the interface plugins implement, the shared state they communicate through, how the pipeline composes them, and how the listener renders their decisions. + +**Audience:** +- Go developers adding plugins to AuthBridge's native chain. +- Anyone debugging the plugin flow via `abctl` or the `:9094` session API. + +**Scope:** +- The Go surface in `authbridge/authlib/pipeline/` and `authbridge/authlib/session/`. +- The observability contract carried by `SessionEvent` on the `:9094` API. +- What the pipeline *does* and *does not* own at the boundary with the listener. + +--- + +## 1. Mental model + +AuthBridge intercepts HTTP traffic in two directions and runs a **separate plugin chain** for each. Each chain has two **phases** — request (headers/body going to the upstream) and response (headers/body coming back). + +``` + Inbound (caller → this agent) + ┌────────────────────────────────────────────────────┐ + │ Request phase → jwt-validation │ + │ → a2a-parser │ + │ → session-recorder (implicit) │ + │ Response phase ← a2a-parser OnResponse │ + │ ← jwt-validation OnResponse │ + └────────────────────────────────────────────────────┘ + + Outbound (this agent → target service) + ┌────────────────────────────────────────────────────┐ + │ Request phase → route-resolver │ + │ → token-exchange │ + │ → mcp-parser / inference-parser │ + │ Response phase ← mcp-parser / inference-parser │ + │ ← token-exchange OnResponse │ + └────────────────────────────────────────────────────┘ +``` + +**Key properties:** +- Plugins execute **sequentially** within a phase. +- Response phase runs plugins in **reverse order** (last plugin sees the response first — LIFO, matches middleware conventions). +- Inbound and outbound are **separate `Pipeline` instances**. A plugin that cares about both directions is registered on both. +- All state shared between plugins within one request/response cycle lives on `*pipeline.Context` (`pctx`). +- Cross-request state (per-session telemetry) lives in the `session.Store`, accessed read-only via `pctx.Session`. + +--- + +## 2. The `Plugin` interface + +```go +type Plugin interface { + Name() string + Capabilities() PluginCapabilities + OnRequest(ctx context.Context, pctx *Context) Action + OnResponse(ctx context.Context, pctx *Context) Action +} +``` + +### `Name() string` +A stable identifier. Used for logs, metrics, `GetState`/`SetState` keys (by convention), and pipeline introspection (`GET /v1/pipeline`). + +### `Capabilities() PluginCapabilities` + +```go +type PluginCapabilities struct { + Reads []string // extension slot names this plugin reads + Writes []string // extension slot names this plugin writes + BodyAccess bool // whether this plugin needs request/response body buffered +} +``` + +Declared once per plugin instance. `pipeline.New` validates that every `Read` is satisfied by an earlier plugin's `Write` — a plugin that depends on `mcp` being populated cannot be registered before `mcp-parser`. A mis-ordered registration fails fast at startup with: + +``` +plugin "guardrail" reads slot "mcp" but no earlier plugin writes it +``` + +`BodyAccess: true` on *any* plugin in a chain causes `Pipeline.NeedsBody()` to return true, which the **listener** uses to negotiate Envoy's `ProcessingMode` (BUFFERED vs HEADERS-only). Without this, the gRPC ext_proc server never asks for the body and parsers see `pctx.Body == nil`. + +### `OnRequest(ctx, pctx) Action` +Called when a request is entering the pipeline. Plugins typically read request headers / body, mutate one or more extension slots, and return `Continue` or `Reject`. + +### `OnResponse(ctx, pctx) Action` +Called after the upstream returns. `pctx.StatusCode`, `pctx.ResponseHeaders`, and `pctx.ResponseBody` are populated. Plugins typically enrich the telemetry extensions with response-side data (completion text, token usage, error code) or apply guardrails on the response content. + +Plugins that only care about the request set `OnResponse` to a no-op (`return Action{Type: Continue}`); same for response-only plugins on `OnRequest`. + +--- + +## 3. `pipeline.Context` — the shared state + +The entire surface a plugin sees: + +```go +type Context struct { + Direction Direction // Inbound | Outbound + Method string // HTTP method + Host string // :authority / Host + Path string // :path + Headers http.Header + Body []byte // nil unless a plugin declared BodyAccess: true + StartedAt time.Time // listener wall-clock at request entry + + Agent *AgentIdentity // this workload's SPIFFE / Keycloak identity + Claims *validation.Claims // inbound caller's JWT claims after jwt-validation + Route *routing.ResolvedRoute // outbound: resolved audience / token scopes + Session *SessionView // read-only view of the session bucket + + // Response-phase fields (populated by listener before RunResponse) + StatusCode int + ResponseHeaders http.Header + ResponseBody []byte + + Extensions Extensions +} +``` + +**Ownership rules:** +- Plugins **read** any field they declared in `Capabilities.Reads`. +- Plugins **write** fields they declared in `Capabilities.Writes`. By convention each extension slot has exactly one writer (the parser plugin). +- `Claims` is populated by `jwt-validation` and is read-only afterward. +- `Agent`, `Route`, `Session` are populated by the listener before `Run`. Plugins treat them as read-only. +- `ResponseBody` appears between `Run` and `RunResponse` — plugins must not read it in `OnRequest`. + +**Lifetime:** one `*Context` per HTTP transaction. Not reused across requests. Single-threaded — the pipeline guarantees sequential invocation of plugins within a phase, so plugins don't need internal locking for pctx reads/writes. + +--- + +## 4. `Extensions` — typed plugin-to-plugin communication + +```go +type Extensions struct { + MCP *MCPExtension + A2A *A2AExtension + Security *SecurityExtension + Delegation *DelegationExtension + Inference *InferenceExtension + Custom map[string]any +} +``` + +Two categories: + +### Named slots (telemetry-worthy) +MCP, A2A, Security, Delegation, Inference. These are: +- Part of the **published schema** carried on `SessionEvent` to `:9094` / `abctl`. +- Consumable by multiple downstream plugins. +- Added to the core struct only when the data has a public contract. + +Adding a named slot is an authlib-core change: edit `Extensions`, add a `SessionEventWire` field, update `snapshotXXX` helpers in the listener, and add filtering rules in `abctl`. + +### `Custom map[string]any` + `GetState[T]`/`SetState[T]` (plugin-private) +For state that's internal to a single plugin or to a bridge's sub-pipeline: + +```go +// Plugin's private state type: +type rlState struct { + TokensAtStart int + Decision string +} + +// In OnRequest: +pipeline.SetState(pctx, "rate-limiter", &rlState{TokensAtStart: 100}) + +// In OnResponse: +s := pipeline.GetState[rlState](pctx, "rate-limiter") +if s != nil { /* use s */ } +``` + +Convention: **key = plugin's Name()** so collisions across plugins don't happen. Storage is lazy (`Custom` is nil-initialized until first write). + +`GetState[T]` type-asserts and returns `nil` on mismatch instead of panicking — a plugin whose type evolves across versions degrades gracefully. + +### Built-in extension shapes + +All at `authbridge/authlib/pipeline/extensions.go`: + +```go +type MCPExtension struct { + Method string // JSON-RPC method, e.g. "tools/call" + RPCID any // JSON-RPC id (could be int or string) + Params map[string]any // request params + Result map[string]any // response result (mutually exclusive with Err) + Err *MCPError +} + +type A2AExtension struct { + Method string + RPCID any + SessionID string // contextId from the client, or server-assigned on first turn + MessageID string + TaskID string + Role string // "user" | "agent" + Parts []A2APart + FinalStatus string // response: "completed" | "failed" | "canceled" + Artifact string // response: assembled artifact text + ErrorMessage string // response: failure reason +} + +type InferenceExtension struct { + // Request side: + Model string + Messages []InferenceMessage + Temperature *float64 + MaxTokens *int + TopP *float64 + Stream bool + Tools []InferenceTool // full definition incl. parameters schema + ToolChoice any + // Response side: + Completion string + FinishReason string + PromptTokens int + CompletionTokens int + TotalTokens int + ToolCalls []InferenceToolCall +} + +type SecurityExtension struct { + Labels []string // classifier / guardrail output + Blocked bool + BlockReason string +} + +type DelegationExtension struct { + Origin string // original caller subject + Actor string // current actor subject + // chain is append-only via AppendHop; reads via Chain() +} +``` + +Mutability: **always assigned, never mutated in place** after the parser sets the slot. This guarantees that `snapshotXXX` in the listener (shallow-copy for event recording) stays correct even when OnResponse enriches the struct — the response snapshot is taken from the now-enriched pointer, but any earlier request-phase snapshot was taken of a frozen copy. + +--- + +## 5. `Action` — control flow + +```go +type Action struct { + Type ActionType // Continue | Reject + Violation *Violation // populated iff Type == Reject +} + +type Violation struct { + // Structured machine-readable error: + Code string // machine-readable, e.g. "auth.missing-token" + Reason string // short human message + Description string // longer explanation; optional + Details map[string]any // plugin-arbitrary structured context; optional + + // HTTP rendering hints — all optional; defaults from Code: + Status int // when 0, StatusFromCode(Code) is used + Body []byte // when nil, synthesized JSON + BodyType string // Content-Type for Body; defaults to application/json + Headers http.Header // merged into the response (e.g. WWW-Authenticate, Retry-After) + + // Framework-populated from Plugin.Name(); plugins leave it empty: + PluginName string +} +``` + +Returning `Reject` from `OnRequest` halts the request pipeline; from `OnResponse` halts the response pipeline. The listener calls `Violation.Render()` to produce `(status, headers, body)` and emits that as the HTTP response. The default body when `Body` is nil: + +```json +{ + "error": "auth.missing-token", + "message": "Bearer token required", + "description": "No Authorization header present", + "plugin": "jwt-validation", + "details": { "realm": "kagenti" } +} +``` + +Helper constructors cover the common cases so the reject site stays one line: + +```go +pipeline.Deny("auth.invalid-token", "token expired") +pipeline.DenyStatus(451, "policy.forbidden", "unavailable for legal reasons") +pipeline.DenyWithDetails("policy.rate-limited", "quota hit", map[string]any{ + "remaining": 0, "window": "1h", +}) +pipeline.Challenge("kagenti", "Authorization required") // 401 + WWW-Authenticate +pipeline.RateLimited(30*time.Second, "", "slow down") // 429 + Retry-After +``` + +The `Code` → HTTP-status mapping for well-known codes lives at `codeToStatus` in `action.go`; unknown codes default to 500. Plugins that need a non-default status set `Violation.Status` explicitly or use `DenyStatus`. + +There is no "soft error" channel today — a plugin that wants to fail open logs and returns `Continue`. A future iteration may add a per-plugin `on_error` policy. + +--- + +## 6. `Pipeline` — composition and execution + +```go +func New(plugins []Plugin, opts ...Option) (*Pipeline, error) +func (p *Pipeline) Run(ctx context.Context, pctx *Context) Action // request phase +func (p *Pipeline) RunResponse(ctx context.Context, pctx *Context) Action // response phase (reverse) +func (p *Pipeline) Start(ctx context.Context) error // invoke Init on Initializer plugins +func (p *Pipeline) Stop(ctx context.Context) // invoke Shutdown on Shutdowner plugins +func (p *Pipeline) Plugins() []Plugin // defensive copy +func (p *Pipeline) NeedsBody() bool // OR over all plugins' BodyAccess +``` + +`New` validates capability wiring at startup: every `Read` must be satisfied by some earlier plugin's `Write`. + +### Plugin lifecycle (`Start` / `Stop`) + +Plugins that need one-time setup (load a model, warm a cache, register metrics, spawn a background goroutine) implement the optional `Initializer` interface: + +```go +type Initializer interface { + Init(ctx context.Context) error +} +``` + +Plugins that need graceful cleanup (flush audit events, close a connection, cancel a goroutine) implement `Shutdowner`: + +```go +type Shutdowner interface { + Shutdown(ctx context.Context) error +} +``` + +Both are **optional** via Go's type-assertion idiom — a plugin that doesn't need them simply doesn't implement them, and the pipeline skips it. Existing plugins (jwt-validation, a2a-parser, mcp-parser, inference-parser, token-exchange) don't implement these; they keep working unchanged. + +The host (e.g. `cmd/authbridge/main.go`) drives the lifecycle: + +```go +// After pipeline.New, before listeners accept traffic: +initCtx, cancel := context.WithTimeout(context.Background(), 60*time.Second) +defer cancel() +if err := inboundPipeline.Start(initCtx); err != nil { + log.Fatalf("inbound pipeline Start: %v", err) // fail-fast on bad plugin init +} +if err := outboundPipeline.Start(initCtx); err != nil { + log.Fatalf("outbound pipeline Start: %v", err) +} + +// ... serve traffic ... + +// After listeners have drained on SIGTERM: +outboundPipeline.Stop(shutdownCtx) // reverse order within each pipeline +inboundPipeline.Stop(shutdownCtx) +``` + +Semantics: +- `Start` — Init runs **in declaration order**, fails fast on the first error. The returned error names the offending plugin. No Shutdown is invoked on plugins whose Init already ran successfully — the intent is hard-fail on startup, not unwind. +- `Stop` — Shutdown runs **in reverse declaration order (LIFO)** so a plugin that depends on an earlier plugin's resources can still use them while cleaning up. Best-effort: errors from one Shutdown are logged but do not stop the sequence. Bounded by the caller's ctx deadline. + +A minimal Init/Shutdown plugin example — a rate-limiter that refreshes its quota store in the background: + +```go +type RateLimiter struct { + store *quotaStore + cancel context.CancelFunc +} + +func (p *RateLimiter) Name() string { return "rate-limiter" } +func (p *RateLimiter) Capabilities() pipeline.PluginCapabilities { return pipeline.PluginCapabilities{} } + +func (p *RateLimiter) Init(ctx context.Context) error { + p.store = newQuotaStore() + bg, cancel := context.WithCancel(context.Background()) + p.cancel = cancel + go p.store.refreshLoop(bg, 10*time.Second) // lives until Shutdown + return nil +} + +func (p *RateLimiter) Shutdown(ctx context.Context) error { + p.cancel() // stop the refresh loop + return p.store.flush(ctx) // best-effort write-back of pending counters +} + +func (p *RateLimiter) OnRequest(ctx context.Context, pctx *pipeline.Context) pipeline.Action { + if !p.store.allow(pctx) { + return pipeline.RateLimited(30*time.Second, "", "quota exceeded") + } + return pipeline.Action{Type: pipeline.Continue} +} + +func (p *RateLimiter) OnResponse(context.Context, *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +``` + +### Extension slots known to the validator + +Built-in: `mcp`, `a2a`, `security`, `delegation`, `inference`, `custom`. + +**For plugins that write new slot names:** use the `WithSlots` option: + +```go +pipeline, err := pipeline.New(plugins, + pipeline.WithSlots("provenance", "risk-score")) +``` + +This tells the validator those slot names are legal, so a downstream plugin can `Capabilities.Reads = []string{"provenance"}` without being rejected as "unknown slot". + +### Execution order +- Request phase: `plugins[0].OnRequest → plugins[1].OnRequest → …` +- Response phase: `plugins[N-1].OnResponse → plugins[N-2].OnResponse → …` +- A `Reject` from any plugin halts its phase immediately. +- `ctx.Err() != nil` between plugins also halts with `Reject{Status: 499}`. + +### Concurrency model +Always sequential. No priority / mode / fire-and-forget semantics yet. This is the 80% case for auth-and-parse pipelines; richer modes would require an executor layer above the current loop. + +--- + +## 7. `Session` + `SessionEvent` — the observability side-channel + +The pipeline itself is **in-band** (plugins alter request handling). Alongside it runs an **out-of-band** observability layer: the listener snapshots `pctx` into a `SessionEvent` after each phase and appends it to a per-session bucket in the `session.Store`. This store is what powers the `:9094` HTTP API and `abctl`. + +```go +type SessionEvent struct { + SessionID string // bucket the event landed in + At time.Time + Direction Direction // inbound | outbound + Phase SessionPhase // request | response + A2A *A2AExtension // snapshot of pctx.Extensions.A2A + MCP *MCPExtension + Inference *InferenceExtension + Identity *EventIdentity // Subject, ClientID, AgentID, Scopes + StatusCode int // response phase only + Error *EventError // populated on 4xx/5xx + Host string // :authority + TargetAudience string // outbound: resolved OAuth audience + Duration time.Duration // response: wall-clock since request entry +} +``` + +**Plugins do not touch `SessionEvent` directly.** The listener records events; plugins only read `pctx.Session` (a `*SessionView`) when they want to correlate the current request with prior ones in the same conversation — e.g. a rate-limiter that counts a session's inference events. + +Wire format (`SessionEvent.MarshalJSON`) translates enums to strings and `Duration` to `DurationMs`. Round-trip stable — `json.Marshal(e) → json.Unmarshal → json.Marshal` is byte-identical. Tested at `pipeline/session_test.go:TestSessionEvent_JSONRoundTrip`. + +--- + +## 8. Boundary: pipeline vs listener + +The pipeline **does not own**: + +| Concern | Owner | Why | +|---|---|---| +| HTTP wire protocol (ext_proc gRPC, ext_authz, reverse/forward proxy) | `cmd/authbridge/listener/` | Each mode speaks a different wire; pipeline stays protocol-free | +| Body buffering negotiation (`ProcessingMode: BUFFERED`) | Listener reads `Pipeline.NeedsBody()` | Only listener can respond to the ext_proc handshake | +| JWT issuance, client registration, Keycloak admin calls | Outside the pipeline (agent sidecars / kagenti-operator) | Async concerns happening before/after any request flow | +| Session store writes (`Store.Append`) | Listener, called after each phase | Plugins see only the read-only `SessionView` | +| SSE streaming of events to abctl | `authlib/sessionapi` | Observability API, not a plugin concern | + +The pipeline **does own**: +- The `Plugin` interface contract. +- `pipeline.Context` structure and invariants. +- Validation of capability wiring at construction. +- Sequential dispatch and reject-short-circuit semantics. +- Typed extension slots and `GetState`/`SetState` helpers. +- The session-event *shape* (the listener uses it but doesn't define it). + +--- + +## 9. Writing a native plugin — a worked example + +A minimal outbound plugin that stamps an extra header on any request routed to GitHub: + +```go +package myplugins + +import ( + "context" + "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" +) + +type GitHubStamper struct{} + +func (GitHubStamper) Name() string { return "github-stamper" } + +func (GitHubStamper) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{ + // We don't need body buffering, and don't depend on other plugins' data. + } +} + +func (GitHubStamper) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { + if pctx.Route != nil && pctx.Route.Audience == "github-tool" { + pctx.Headers.Set("x-from-authbridge", "1") + } + return pipeline.Action{Type: pipeline.Continue} +} + +func (GitHubStamper) OnResponse(_ context.Context, _ *pipeline.Context) pipeline.Action { + return pipeline.Action{Type: pipeline.Continue} +} +``` + +Registered in the outbound pipeline alongside the built-in plugins at startup in `cmd/authbridge/main.go`. + +A plugin that reads the caller's SPIFFE ID from inbound claims and records a per-session counter via `GetState`/`SetState`: + +```go +type sessionCounter struct{ N int } + +func (p *Counter) Name() string { return "turn-counter" } + +func (p *Counter) Capabilities() pipeline.PluginCapabilities { + return pipeline.PluginCapabilities{Reads: []string{"a2a"}} +} + +func (p *Counter) OnRequest(_ context.Context, pctx *pipeline.Context) pipeline.Action { + if pctx.Direction != pipeline.Inbound || pctx.Extensions.A2A == nil { + return pipeline.Action{Type: pipeline.Continue} + } + state := pipeline.GetState[sessionCounter](pctx, p.Name()) + if state == nil { + state = &sessionCounter{} + pipeline.SetState(pctx, p.Name(), state) + } + state.N++ + if state.N > 10 { + return pipeline.RateLimited(30*time.Second, + "policy.rate-limited", "per-session turn limit exceeded") + } + return pipeline.Action{Type: pipeline.Continue} +} +``` + +Plugins express rejections through a structured `Violation` carrying a +machine-readable `Code`, a short `Reason`, an optional longer +`Description`, a `Details` map for plugin-arbitrary context, and HTTP- +rendering hints (`Status`, `Body`, `BodyType`, `Headers`). Default JSON +body synthesis covers the 95% case — set the hints only when overriding. +Helper constructors (`Deny`, `DenyStatus`, `DenyWithDetails`, +`Challenge`, `RateLimited`) make the common cases one-liners. See +`action.go` for the full surface and `action_test.go` for worked +examples. + +--- + +## 10. Open questions + +- **Priority / on-error policies.** Plugins don't declare these today. If fail-open / fail-closed behavior becomes important to express per plugin, it would be added to `PluginCapabilities` (or a sibling metadata struct) and interpreted by `Pipeline`. +- **Body mutation semantics.** Today plugins generally don't rewrite `pctx.Body` or `pctx.ResponseBody`. If a plugin needs to modify the payload, we'd need a clear contract about whether downstream plugins see the modified or original bytes. +- **Execution modes.** The pipeline is sequential-only. Concurrent or fire-and-forget modes would require an executor layer; no concrete use case yet. + +--- + +## 11. Versioning + +The plugin interface is **not** semver-stable yet (AuthBridge is pre-1.0). Changes since the initial release: +- Added `BodyAccess` to `PluginCapabilities`. +- Added `WithSlots` to `New` for bridge-plugin slot registration. +- Added `GetState[T]` / `SetState[T]` generic helpers. +- Extended `A2AExtension` with response-side fields (TaskID, FinalStatus, Artifact, ErrorMessage). +- Extended `InferenceExtension` with structured tools + tool calls + TopP / ToolChoice. +- Added `SessionEvent.MarshalJSON`/`UnmarshalJSON` round-trip contract. +- **Breaking**: replaced `Action.Status`/`Action.Reason` with `Action.Violation` (see §5). Migration: use `Deny()`, `DenyStatus()`, `Challenge()`, `RateLimited()` helpers. +- Added optional `Initializer` / `Shutdowner` interfaces + `Pipeline.Start` / `Pipeline.Stop` (see §6). Existing plugins are unaffected because the interfaces are opt-in via type-assertion. + +Breaking changes will be announced in `authbridge/CHANGELOG.md` (TBD) before a 1.0 tag. + +--- + +## 12. Cross-references + +- `pipeline.go` — `Pipeline` type, `New`, `Run`, `RunResponse`, `Start`, `Stop`, `Plugins`, `NeedsBody`. +- `plugin.go` — `Plugin` interface, `PluginCapabilities`, `Initializer`, `Shutdowner`. +- `action.go` — `Action`, `ActionType`, `Violation`, helper constructors (`Deny`, `DenyStatus`, `DenyWithDetails`, `Challenge`, `RateLimited`), `StatusFromCode`. +- `context.go` — `Context`, `Direction`, `AgentIdentity`. +- `extensions.go` — named extension types + `GetState`/`SetState`. +- `session.go` — `SessionEvent`, `SessionView`, `SessionPhase`, marshalers. +- `authlib/session/` — `Store`, `SessionSummary`, ring buffer, TTL / max-events caps. +- `authlib/sessionapi/` — HTTP API (`/v1/sessions`, `/v1/events`, `/v1/pipeline`) surfacing all of the above. +- `cmd/authbridge/listener/extproc/` — reference usage for all three phases. +- `cmd/abctl/` — TUI consumer of the session API, useful as a reference integrator. diff --git a/authbridge/authlib/pipeline/action.go b/authbridge/authlib/pipeline/action.go index 8a88a0adf..b5079d31a 100644 --- a/authbridge/authlib/pipeline/action.go +++ b/authbridge/authlib/pipeline/action.go @@ -1,17 +1,196 @@ package pipeline +import ( + "encoding/json" + "fmt" + "net/http" + "time" +) + // ActionType represents the result of a plugin's processing. type ActionType int const ( + // Continue the pipeline to the next plugin. Continue ActionType = iota + // Reject the request; the listener synthesizes an HTTP response from + // the accompanying Violation. Reject ) -// Action is returned by a plugin to indicate whether processing should continue -// or the request should be rejected. +// Action is returned by a plugin to indicate how the pipeline should proceed. +// For Continue, Violation is nil. For Reject, Violation is populated (either +// by the plugin itself or by a helper constructor such as Deny / Challenge). type Action struct { - Type ActionType - Status int // HTTP status code (only for Reject) - Reason string // human-readable reason (only for Reject) + Type ActionType + Violation *Violation +} + +// Violation is the structured denial a plugin returns when it rejects a +// request or response. The shape is intentionally close to CPEX's +// PluginViolation so bridges and cross-framework tooling can reason about +// authbridge and CPEX denials uniformly: Code is machine-readable, Reason +// is a short human message, Description is a longer explanation, Details +// is plugin-arbitrary structured context. +// +// HTTP-rendering hints (Status, Body, BodyType, Headers) are optional +// overrides — when empty, Render() synthesizes sensible defaults from Code, +// Reason, Description, Details, and PluginName. +type Violation struct { + // Structured fields — the authoritative description of the denial. + // Plugins that only care about semantics (not wire shape) set just + // Code and Reason; listeners and clients can inspect the rest. + Code string // machine-readable error code, e.g. "auth.missing-token" + Reason string // short human-readable message + Description string // longer explanation; optional + Details map[string]any // structured context; optional + + // HTTP rendering hints — all optional. When empty, Render() picks + // defaults. Plugins set these only when they need to override. + Status int // HTTP status; if 0, StatusFromCode(Code) is used + Body []byte // response body; if nil, synthesized JSON + BodyType string // Content-Type for Body; defaults to application/json + Headers http.Header // merged into response headers + + // PluginName is populated by the pipeline framework from Plugin.Name() + // when the Reject is returned. Plugins should leave this empty. + PluginName string +} + +// codeToStatus maps well-known violation codes to HTTP status. Plugins +// introducing new codes should either populate Violation.Status explicitly +// or the host can extend this table (future: make this configurable). +var codeToStatus = map[string]int{ + "auth.missing-token": http.StatusUnauthorized, + "auth.invalid-token": http.StatusUnauthorized, + "auth.audience-mismatch": http.StatusUnauthorized, + "auth.unauthorized": http.StatusUnauthorized, + "policy.forbidden": http.StatusForbidden, + "policy.rate-limited": http.StatusTooManyRequests, + "policy.content-blocked": http.StatusForbidden, + "upstream.unreachable": http.StatusServiceUnavailable, + "upstream.token-exchange-failed": http.StatusServiceUnavailable, + "upstream.timeout": http.StatusGatewayTimeout, + "pipeline.cancelled": 499, // client-closed request (nginx convention) +} + +// StatusFromCode returns the HTTP status a Violation maps to when its +// Status field is unset. Falls back to 500 for unknown codes. +func StatusFromCode(code string) int { + if s, ok := codeToStatus[code]; ok { + return s + } + return http.StatusInternalServerError +} + +// Render produces the HTTP response triple (status, headers, body) for a +// Violation. Called by each listener's reject path so the translation +// from structured Violation to wire bytes lives in one place. +// +// Defaults when fields are unset: +// - Status: StatusFromCode(Code), or 500 for unknown codes +// - Body: JSON {"error":Code,"message":Reason, plus description/plugin/details when non-empty} +// - BodyType: application/json +// - Content-Type header: set from BodyType if not already present in Headers +// +// Render does not mutate the receiver — it's safe to call multiple times. +// When Headers is nil, a fresh map is allocated. +func (v *Violation) Render() (int, http.Header, []byte) { + if v == nil { + return http.StatusInternalServerError, http.Header{}, []byte(`{"error":"internal","message":"reject without violation"}`) + } + status := v.Status + if status == 0 { + status = StatusFromCode(v.Code) + } + body := v.Body + bodyType := v.BodyType + if body == nil { + payload := map[string]any{ + "error": v.Code, + "message": v.Reason, + } + if v.Description != "" { + payload["description"] = v.Description + } + if v.PluginName != "" { + payload["plugin"] = v.PluginName + } + if len(v.Details) > 0 { + payload["details"] = v.Details + } + body, _ = json.Marshal(payload) + bodyType = "application/json" + } + headers := v.Headers + if headers == nil { + headers = http.Header{} + } else { + // Clone so the caller's headers are not mutated. + cloned := make(http.Header, len(headers)) + for k, vs := range headers { + cloned[k] = append([]string(nil), vs...) + } + headers = cloned + } + if bodyType != "" && headers.Get("Content-Type") == "" { + headers.Set("Content-Type", bodyType) + } + return status, headers, body +} + +// ---------------------------------------------------------------------------- +// Helper constructors — the 95% cases. Use these over manually constructing +// Action{Type: Reject, Violation: &Violation{...}} for readability. +// ---------------------------------------------------------------------------- + +// Deny returns a Reject action with a code and short reason. HTTP status +// is derived from the code via StatusFromCode. +func Deny(code, reason string) Action { + return Action{Type: Reject, Violation: &Violation{Code: code, Reason: reason}} +} + +// DenyStatus overrides the HTTP status inferred from Code. Use when the +// code-to-status default doesn't match the caller's intent (e.g. a policy +// plugin that wants a 451 Unavailable For Legal Reasons). +func DenyStatus(status int, code, reason string) Action { + return Action{Type: Reject, Violation: &Violation{ + Code: code, Reason: reason, Status: status, + }} +} + +// DenyWithDetails attaches plugin-arbitrary structured context. The +// details map surfaces as a "details" object in the default JSON body +// and is available to clients parsing the error programmatically. +func DenyWithDetails(code, reason string, details map[string]any) Action { + return Action{Type: Reject, Violation: &Violation{ + Code: code, Reason: reason, Details: details, + }} +} + +// Challenge returns a 401 with a Bearer challenge so clients know to +// present credentials. The realm appears in the WWW-Authenticate header +// per RFC 6750. +func Challenge(realm, reason string) Action { + h := http.Header{} + h.Set("WWW-Authenticate", fmt.Sprintf("Bearer realm=%q", realm)) + return Action{Type: Reject, Violation: &Violation{ + Code: "auth.missing-token", + Reason: reason, + Headers: h, + }} +} + +// RateLimited returns a 429 with a Retry-After header expressed in +// seconds. Code defaults to "policy.rate-limited" when empty so a plugin +// can call RateLimited(30*time.Second, "", "slow down") without ceremony. +func RateLimited(retryAfter time.Duration, code, reason string) Action { + if code == "" { + code = "policy.rate-limited" + } + h := http.Header{} + h.Set("Retry-After", fmt.Sprintf("%d", int(retryAfter.Seconds()))) + return Action{Type: Reject, Violation: &Violation{ + Code: code, Reason: reason, Headers: h, + }} } diff --git a/authbridge/authlib/pipeline/action_test.go b/authbridge/authlib/pipeline/action_test.go new file mode 100644 index 000000000..7dd79bdcf --- /dev/null +++ b/authbridge/authlib/pipeline/action_test.go @@ -0,0 +1,219 @@ +package pipeline + +import ( + "encoding/json" + "net/http" + "strings" + "testing" + "time" +) + +func TestStatusFromCode(t *testing.T) { + cases := map[string]int{ + "auth.missing-token": 401, + "auth.invalid-token": 401, + "policy.forbidden": 403, + "policy.rate-limited": 429, + "upstream.unreachable": 503, + "upstream.token-exchange-failed": 503, + "upstream.timeout": 504, + "pipeline.cancelled": 499, + "anything.unknown": 500, + "": 500, + } + for code, want := range cases { + if got := StatusFromCode(code); got != want { + t.Errorf("StatusFromCode(%q) = %d, want %d", code, got, want) + } + } +} + +// Default body synthesis covers the standard fields. This is the contract +// external clients rely on when a plugin uses Deny() without overriding Body. +func TestViolation_Render_DefaultBody(t *testing.T) { + v := &Violation{ + Code: "auth.missing-token", + Reason: "Bearer required", + Description: "No Authorization header present", + Details: map[string]any{"realm": "kagenti"}, + PluginName: "jwt-validation", + } + status, headers, body := v.Render() + + if status != 401 { + t.Errorf("status = %d, want 401", status) + } + if ct := headers.Get("Content-Type"); ct != "application/json" { + t.Errorf("Content-Type = %q, want application/json", ct) + } + + var got map[string]any + if err := json.Unmarshal(body, &got); err != nil { + t.Fatalf("body not valid JSON: %v\n body: %s", err, body) + } + if got["error"] != "auth.missing-token" { + t.Errorf("error = %v, want auth.missing-token", got["error"]) + } + if got["message"] != "Bearer required" { + t.Errorf("message = %v, want Bearer required", got["message"]) + } + if got["description"] != "No Authorization header present" { + t.Errorf("description = %v", got["description"]) + } + if got["plugin"] != "jwt-validation" { + t.Errorf("plugin = %v", got["plugin"]) + } + details, ok := got["details"].(map[string]any) + if !ok { + t.Fatalf("details not a map: %v", got["details"]) + } + if details["realm"] != "kagenti" { + t.Errorf("details.realm = %v", details["realm"]) + } +} + +// A plugin that supplies its own body bytes must have them written +// verbatim, with BodyType used for Content-Type. +func TestViolation_Render_CustomBody(t *testing.T) { + v := &Violation{ + Code: "policy.content-blocked", + Reason: "blocked", + Body: []byte("blocked by policy: rule-42"), + BodyType: "text/plain", + } + status, headers, body := v.Render() + if status != 403 { + t.Errorf("status = %d, want 403", status) + } + if string(body) != "blocked by policy: rule-42" { + t.Errorf("body = %q, want verbatim", body) + } + if ct := headers.Get("Content-Type"); ct != "text/plain" { + t.Errorf("Content-Type = %q, want text/plain", ct) + } +} + +// Optional fields (description, plugin, details) must be omitted from the +// default body when empty so the wire payload stays clean. +func TestViolation_Render_OmitsEmptyOptionals(t *testing.T) { + v := &Violation{Code: "auth.unauthorized", Reason: "go away"} + _, _, body := v.Render() + s := string(body) + for _, field := range []string{"description", "plugin", "details"} { + if strings.Contains(s, `"`+field+`":`) { + t.Errorf("expected %q omitted from %s", field, s) + } + } +} + +// Headers passed in by the plugin (e.g. WWW-Authenticate) must be +// preserved. Render must clone so the caller's map isn't mutated with +// Content-Type additions. +func TestViolation_Render_PreservesAndClonesHeaders(t *testing.T) { + // Use http.Header.Set so the key is canonicalized by the stdlib — this + // mirrors what a real plugin would do via the Challenge() helper. + origHeaders := http.Header{} + origHeaders.Set("WWW-Authenticate", `Bearer realm="kagenti"`) + v := &Violation{ + Code: "auth.missing-token", + Reason: "bearer required", + Headers: origHeaders, + } + _, h, _ := v.Render() + if got := h.Get("WWW-Authenticate"); got != `Bearer realm="kagenti"` { + t.Errorf("WWW-Authenticate = %q", got) + } + if h.Get("Content-Type") == "" { + t.Error("Content-Type not auto-set from synthesized body") + } + // Original must not have Content-Type appended. + if origHeaders.Get("Content-Type") != "" { + t.Error("Render mutated caller's headers map") + } +} + +// Explicit Status overrides the code-to-status lookup. +func TestViolation_Render_ExplicitStatusWins(t *testing.T) { + v := &Violation{Code: "auth.missing-token", Reason: "", Status: 418} + status, _, _ := v.Render() + if status != 418 { + t.Errorf("status = %d, want explicit 418", status) + } +} + +// Render must be safe on a nil receiver — used in error paths where a +// plugin returned Reject with no violation populated. +func TestViolation_Render_NilReceiver(t *testing.T) { + var v *Violation + status, h, body := v.Render() + if status != 500 { + t.Errorf("status = %d, want 500", status) + } + if h == nil { + t.Error("headers nil on nil receiver") + } + if len(body) == 0 { + t.Error("body empty on nil receiver") + } +} + +// ---------------------------------------------------------------------------- +// Helper constructors +// ---------------------------------------------------------------------------- + +func TestDeny(t *testing.T) { + a := Deny("auth.invalid-token", "expired") + if a.Type != Reject || a.Violation == nil { + t.Fatalf("not a Reject: %+v", a) + } + if a.Violation.Code != "auth.invalid-token" || a.Violation.Reason != "expired" { + t.Errorf("unexpected violation: %+v", a.Violation) + } + status, _, _ := a.Violation.Render() + if status != 401 { + t.Errorf("derived status = %d", status) + } +} + +func TestDenyStatus(t *testing.T) { + a := DenyStatus(451, "policy.forbidden", "unavailable for legal reasons") + status, _, _ := a.Violation.Render() + if status != 451 { + t.Errorf("status = %d, want 451", status) + } +} + +func TestDenyWithDetails(t *testing.T) { + a := DenyWithDetails("policy.rate-limited", "quota hit", map[string]any{ + "quota": 100, "remaining": 0, + }) + _, _, body := a.Violation.Render() + if !strings.Contains(string(body), `"details"`) { + t.Errorf("details missing from body: %s", body) + } +} + +func TestChallenge(t *testing.T) { + a := Challenge("kagenti", "Authorization required") + status, h, _ := a.Violation.Render() + if status != 401 { + t.Errorf("status = %d", status) + } + if got := h.Get("WWW-Authenticate"); got != `Bearer realm="kagenti"` { + t.Errorf("WWW-Authenticate = %q", got) + } +} + +func TestRateLimited(t *testing.T) { + a := RateLimited(30*time.Second, "", "slow down") + status, h, _ := a.Violation.Render() + if status != 429 { + t.Errorf("status = %d", status) + } + if got := h.Get("Retry-After"); got != "30" { + t.Errorf("Retry-After = %q, want 30", got) + } + if a.Violation.Code != "policy.rate-limited" { + t.Errorf("default code = %q", a.Violation.Code) + } +} diff --git a/authbridge/authlib/pipeline/pipeline.go b/authbridge/authlib/pipeline/pipeline.go index eb5137fff..fc164df21 100644 --- a/authbridge/authlib/pipeline/pipeline.go +++ b/authbridge/authlib/pipeline/pipeline.go @@ -57,17 +57,18 @@ func New(plugins []Plugin, opts ...Option) (*Pipeline, error) { } // Run executes the request phase of the pipeline sequentially. -// If any plugin returns Reject, the pipeline stops and returns that action. +// If any plugin returns Reject, the pipeline stops and returns that action +// with Violation.PluginName populated. func (p *Pipeline) Run(ctx context.Context, pctx *Context) Action { for _, plugin := range p.plugins { if ctx.Err() != nil { slog.Info("pipeline: request cancelled", "plugin", plugin.Name()) - return Action{Type: Reject, Status: 499, Reason: "request cancelled"} + return Deny("pipeline.cancelled", "request cancelled") } action := plugin.OnRequest(ctx, pctx) if action.Type == Reject { - slog.Info("pipeline: plugin rejected request", - "plugin", plugin.Name(), "status", action.Status, "reason", action.Reason) + stampPluginName(&action, plugin.Name()) + logReject(plugin.Name(), action, "pipeline: plugin rejected request") return action } slog.Debug("pipeline: plugin completed", "plugin", plugin.Name()) @@ -81,18 +82,42 @@ func (p *Pipeline) RunResponse(ctx context.Context, pctx *Context) Action { for i := len(p.plugins) - 1; i >= 0; i-- { if ctx.Err() != nil { slog.Info("pipeline: response cancelled", "plugin", p.plugins[i].Name()) - return Action{Type: Reject, Status: 499, Reason: "request cancelled"} + return Deny("pipeline.cancelled", "request cancelled") } action := p.plugins[i].OnResponse(ctx, pctx) if action.Type == Reject { - slog.Info("pipeline: plugin rejected response", - "plugin", p.plugins[i].Name(), "status", action.Status, "reason", action.Reason) + stampPluginName(&action, p.plugins[i].Name()) + logReject(p.plugins[i].Name(), action, "pipeline: plugin rejected response") return action } } return Action{Type: Continue} } +// stampPluginName annotates a reject action with the plugin that produced +// it, so listeners and clients can attribute the denial without the +// plugin remembering to set it. +func stampPluginName(action *Action, name string) { + if action.Violation == nil { + action.Violation = &Violation{Code: "plugin.unspecified", Reason: "plugin rejected without violation"} + } + if action.Violation.PluginName == "" { + action.Violation.PluginName = name + } +} + +// logReject emits a structured log for a rejected request/response, with +// the violation's code and reason. Keeps the two identical log statements +// in Run/RunResponse in one place. +func logReject(pluginName string, action Action, msg string) { + status, _, _ := action.Violation.Render() + slog.Info(msg, + "plugin", pluginName, + "status", status, + "code", action.Violation.Code, + "reason", action.Violation.Reason) +} + // Plugins returns a copy of the pipeline's plugin list in execution order. // The copy prevents callers from mutating the backing slice; individual // Plugin values are interface types and can be inspected freely. @@ -115,6 +140,51 @@ func (p *Pipeline) NeedsBody() bool { return false } +// Start invokes Init on every plugin that implements the Initializer +// interface, in declaration order. Returns the first error encountered; +// on error, later plugins are not initialized. Plugins without Init are +// silently skipped. +// +// Callers should invoke Start after Pipeline construction (pipeline.New) +// and before the listener accepts traffic. Safe to call at most once per +// Pipeline — plugins may assume Init runs exactly once per process. +func (p *Pipeline) Start(ctx context.Context) error { + for _, plugin := range p.plugins { + init, ok := plugin.(Initializer) + if !ok { + continue + } + slog.Debug("pipeline: initializing plugin", "plugin", plugin.Name()) + if err := init.Init(ctx); err != nil { + return fmt.Errorf("plugin %q Init: %w", plugin.Name(), err) + } + } + return nil +} + +// Stop invokes Shutdown on every plugin that implements the Shutdowner +// interface, in reverse declaration order (LIFO). Errors are logged but +// do not stop the sequence — every Shutdowner is given a chance to flush. +// The caller-supplied ctx carries the shutdown deadline; plugins are +// expected to respect it. +// +// Callers should invoke Stop after the listener has drained / stopped +// accepting new requests so in-flight work is allowed to complete first. +// Safe to call at most once per Pipeline. +func (p *Pipeline) Stop(ctx context.Context) { + for i := len(p.plugins) - 1; i >= 0; i-- { + sh, ok := p.plugins[i].(Shutdowner) + if !ok { + continue + } + slog.Debug("pipeline: shutting down plugin", "plugin", p.plugins[i].Name()) + if err := sh.Shutdown(ctx); err != nil { + slog.Warn("pipeline: plugin Shutdown returned error", + "plugin", p.plugins[i].Name(), "error", err) + } + } +} + // validateCapabilities checks that every slot a plugin reads has been written // by an earlier plugin in the chain. func validateCapabilities(plugins []Plugin, validSlots map[string]bool) error { diff --git a/authbridge/authlib/pipeline/pipeline_test.go b/authbridge/authlib/pipeline/pipeline_test.go index 3ba9654c0..92ea88542 100644 --- a/authbridge/authlib/pipeline/pipeline_test.go +++ b/authbridge/authlib/pipeline/pipeline_test.go @@ -2,6 +2,8 @@ package pipeline import ( "context" + "errors" + "strings" "testing" ) @@ -82,7 +84,7 @@ func TestPipelineRun_Reject(t *testing.T) { p1 := &stubPlugin{ name: "rejecter", onReq: func(_ context.Context, _ *Context) Action { - return Action{Type: Reject, Status: 403, Reason: "forbidden"} + return Deny("policy.forbidden", "forbidden") }, } p2 := &stubPlugin{ @@ -101,11 +103,21 @@ func TestPipelineRun_Reject(t *testing.T) { if action.Type != Reject { t.Errorf("got %v, want Reject", action.Type) } - if action.Status != 403 { - t.Errorf("status = %d, want 403", action.Status) + if action.Violation == nil { + t.Fatal("violation not populated") } - if action.Reason != "forbidden" { - t.Errorf("reason = %q, want %q", action.Reason, "forbidden") + if action.Violation.Code != "policy.forbidden" { + t.Errorf("code = %q, want policy.forbidden", action.Violation.Code) + } + if action.Violation.Reason != "forbidden" { + t.Errorf("reason = %q, want forbidden", action.Violation.Reason) + } + if action.Violation.PluginName != "rejecter" { + t.Errorf("PluginName = %q, want rejecter (auto-stamped by pipeline)", action.Violation.PluginName) + } + status, _, _ := action.Violation.Render() + if status != 403 { + t.Errorf("status = %d, want 403", status) } if called { t.Error("second plugin was called after first rejected") @@ -154,7 +166,7 @@ func TestPipelineRunResponse_Reject(t *testing.T) { p2 := &stubPlugin{ name: "rejecter", onResp: func(_ context.Context, _ *Context) Action { - return Action{Type: Reject, Status: 500, Reason: "response blocked"} + return DenyStatus(500, "policy.content-blocked", "response blocked") }, } pipe, err := New([]Plugin{p1, p2}) @@ -354,8 +366,12 @@ func TestPipelineRun_ContextCancellation(t *testing.T) { if action.Type != Reject { t.Errorf("got %v, want Reject for cancelled context", action.Type) } - if action.Status != 499 { - t.Errorf("status = %d, want 499", action.Status) + status, _, _ := action.Violation.Render() + if status != 499 { + t.Errorf("status = %d, want 499", status) + } + if action.Violation.Code != "pipeline.cancelled" { + t.Errorf("code = %q, want pipeline.cancelled", action.Violation.Code) } if called { t.Error("plugin was called despite cancelled context") @@ -386,3 +402,159 @@ func TestPipeline_Plugins_ReturnsCopy(t *testing.T) { t.Errorf("Plugins() returned aliased slice; backing data was mutated") } } + +// ---------------------------------------------------------------------------- +// Lifecycle (Initializer / Shutdowner) +// ---------------------------------------------------------------------------- + +// lifecyclePlugin implements Initializer and Shutdowner so tests can +// verify the order and error handling of Pipeline.Start / Pipeline.Stop. +type lifecyclePlugin struct { + stubPlugin + initErr error + shutdownErr error + onInit func(ctx context.Context) // runs before returning initErr + onShutdown func(ctx context.Context) +} + +func (l *lifecyclePlugin) Init(ctx context.Context) error { + if l.onInit != nil { + l.onInit(ctx) + } + return l.initErr +} +func (l *lifecyclePlugin) Shutdown(ctx context.Context) error { + if l.onShutdown != nil { + l.onShutdown(ctx) + } + return l.shutdownErr +} + +// Init is called in declaration order, once per plugin. Plugins without +// Init are silently skipped. +func TestPipelineStart_InvokesInitInOrder(t *testing.T) { + var order []string + a := &lifecyclePlugin{ + stubPlugin: stubPlugin{name: "a"}, + onInit: func(context.Context) { order = append(order, "a") }, + } + // b has no Init — should be skipped, not cause a panic. + b := &stubPlugin{name: "b"} + c := &lifecyclePlugin{ + stubPlugin: stubPlugin{name: "c"}, + onInit: func(context.Context) { order = append(order, "c") }, + } + + p, err := New([]Plugin{a, b, c}) + if err != nil { + t.Fatalf("New: %v", err) + } + if err := p.Start(context.Background()); err != nil { + t.Fatalf("Start: %v", err) + } + if len(order) != 2 || order[0] != "a" || order[1] != "c" { + t.Errorf("Init order = %v, want [a c]", order) + } +} + +// An Init error halts the sequence and the subsequent plugin's Init is +// never called. The returned error names the offending plugin so +// operators can debug a failed startup without reading the plugin's +// wrapped error message. +func TestPipelineStart_StopsOnInitError(t *testing.T) { + aInitCalled := false + cInitCalled := false + a := &lifecyclePlugin{ + stubPlugin: stubPlugin{name: "a"}, + onInit: func(context.Context) { aInitCalled = true }, + } + b := &lifecyclePlugin{ + stubPlugin: stubPlugin{name: "b"}, + initErr: errors.New("boom"), + } + c := &lifecyclePlugin{ + stubPlugin: stubPlugin{name: "c"}, + onInit: func(context.Context) { cInitCalled = true }, + } + + p, err := New([]Plugin{a, b, c}) + if err != nil { + t.Fatalf("New: %v", err) + } + err = p.Start(context.Background()) + if err == nil { + t.Fatal("expected Start to return an error") + } + if !aInitCalled { + t.Error("plugin a's Init should have been called before b failed") + } + if cInitCalled { + t.Error("plugin c's Init should NOT have been called after b failed") + } + if !strings.Contains(err.Error(), `plugin "b"`) { + t.Errorf("error = %q; expected it to name plugin b", err) + } +} + +// Shutdown is called in reverse declaration order (LIFO), so plugins +// that depend on earlier plugins' resources can still use them while +// cleaning up. +func TestPipelineStop_InvokesShutdownInReverseOrder(t *testing.T) { + var order []string + a := &lifecyclePlugin{ + stubPlugin: stubPlugin{name: "a"}, + onShutdown: func(context.Context) { order = append(order, "a") }, + } + b := &stubPlugin{name: "b"} // no Shutdown, skipped + c := &lifecyclePlugin{ + stubPlugin: stubPlugin{name: "c"}, + onShutdown: func(context.Context) { order = append(order, "c") }, + } + + p, err := New([]Plugin{a, b, c}) + if err != nil { + t.Fatalf("New: %v", err) + } + p.Stop(context.Background()) + if len(order) != 2 || order[0] != "c" || order[1] != "a" { + t.Errorf("Shutdown order = %v, want [c a] (reverse)", order) + } +} + +// Shutdown is best-effort: an error from one plugin must not prevent +// the rest from being shut down. Otherwise a buggy plugin could strand +// in-flight work in a later plugin. +func TestPipelineStop_ContinuesOnShutdownError(t *testing.T) { + var calls []string + a := &lifecyclePlugin{ + stubPlugin: stubPlugin{name: "a"}, + onShutdown: func(context.Context) { calls = append(calls, "a") }, + } + b := &lifecyclePlugin{ + stubPlugin: stubPlugin{name: "b"}, + shutdownErr: errors.New("b failed"), + onShutdown: func(context.Context) { calls = append(calls, "b") }, + } + + p, err := New([]Plugin{a, b}) + if err != nil { + t.Fatalf("New: %v", err) + } + p.Stop(context.Background()) // must not panic / abort + // Reverse order: b runs first (and errors), a must still run. + if len(calls) != 2 || calls[0] != "b" || calls[1] != "a" { + t.Errorf("calls = %v, want [b a] — a must run even after b errored", calls) + } +} + +// Shutdown with no Shutdowner-implementing plugins is a no-op (not a +// panic). Important for default pipelines that don't declare +// background state. +func TestPipelineStop_NoShutdownersIsNoop(t *testing.T) { + a := &stubPlugin{name: "a"} + p, err := New([]Plugin{a}) + if err != nil { + t.Fatalf("New: %v", err) + } + p.Stop(context.Background()) // must not panic +} diff --git a/authbridge/authlib/pipeline/plugin.go b/authbridge/authlib/pipeline/plugin.go index b2686db6e..ca8146ef4 100644 --- a/authbridge/authlib/pipeline/plugin.go +++ b/authbridge/authlib/pipeline/plugin.go @@ -18,3 +18,35 @@ type PluginCapabilities struct { Writes []string // extension slot names this plugin writes BodyAccess bool // whether this plugin needs request/response body buffered } + +// Initializer is an optional interface a plugin may implement when it +// needs to run work once before the pipeline starts serving traffic. +// Typical uses: load a model, warm a cache, open a database connection, +// register Prometheus metrics, spawn a background goroutine. Init is +// called by Pipeline.Start exactly once, in plugin declaration order. +// If any plugin's Init returns an error the pipeline fails fast — +// Pipeline.Start returns the error without calling Init on later +// plugins (nothing to unwind: earlier plugins succeeded). +// +// Plugins that don't need initialization simply don't implement this +// interface; the pipeline skips them. Keeping it optional preserves +// backward compatibility with every existing plugin. +type Initializer interface { + Init(ctx context.Context) error +} + +// Shutdowner is an optional interface a plugin may implement when it +// needs to release resources on graceful shutdown. Typical uses: flush +// in-flight audit events, close a DB connection, cancel a background +// goroutine it spawned in Init. Shutdown is called by Pipeline.Stop +// exactly once, in reverse declaration order (LIFO — symmetric with +// OnResponse dispatch) so a plugin that depends on an earlier plugin's +// resources can still use them while shutting down. +// +// Shutdown is best-effort: errors are logged but do not prevent other +// plugins from shutting down. The caller-supplied ctx carries a +// shutdown deadline; plugins must respect it and return rather than +// block indefinitely. +type Shutdowner interface { + Shutdown(ctx context.Context) error +} diff --git a/authbridge/authlib/plugins/jwtvalidation.go b/authbridge/authlib/plugins/jwtvalidation.go index 64565603a..51eb72f3c 100644 --- a/authbridge/authlib/plugins/jwtvalidation.go +++ b/authbridge/authlib/plugins/jwtvalidation.go @@ -48,7 +48,16 @@ func (p *JWTValidation) OnRequest(ctx context.Context, pctx *pipeline.Context) p result := p.auth.HandleInbound(ctx, authHeader, path, audience) if result.Action == auth.ActionDeny { - return pipeline.Action{Type: pipeline.Reject, Status: result.DenyStatus, Reason: result.DenyReason} + // result.DenyReason carries the specific failure (missing header, + // audience mismatch, expired, etc.). Pick a code whose default + // HTTP status matches what auth returned, so the fallback body is + // meaningful even before auth.HandleInbound grows a structured + // code of its own. + code := "auth.unauthorized" + if result.DenyStatus == 503 { + code = "upstream.unreachable" + } + return pipeline.DenyStatus(result.DenyStatus, code, result.DenyReason) } pctx.Claims = result.Claims return pipeline.Action{Type: pipeline.Continue} diff --git a/authbridge/authlib/plugins/plugins_test.go b/authbridge/authlib/plugins/plugins_test.go index fa0d8b158..d35d79c83 100644 --- a/authbridge/authlib/plugins/plugins_test.go +++ b/authbridge/authlib/plugins/plugins_test.go @@ -66,8 +66,8 @@ func TestJWTValidation_InvalidToken(t *testing.T) { if action.Type != pipeline.Reject { t.Fatalf("got %v, want Reject", action.Type) } - if action.Status != http.StatusUnauthorized { - t.Errorf("status = %d, want 401", action.Status) + status, _, _ := action.Violation.Render(); if status != http.StatusUnauthorized { + t.Errorf("status = %d, want 401", status) } } @@ -87,8 +87,8 @@ func TestJWTValidation_MissingHeader(t *testing.T) { if action.Type != pipeline.Reject { t.Fatalf("got %v, want Reject", action.Type) } - if action.Status != http.StatusUnauthorized { - t.Errorf("status = %d, want 401", action.Status) + status, _, _ := action.Violation.Render(); if status != http.StatusUnauthorized { + t.Errorf("status = %d, want 401", status) } } @@ -201,8 +201,8 @@ func TestTokenExchange_ExchangeFailure(t *testing.T) { if action.Type != pipeline.Reject { t.Fatalf("got %v, want Reject", action.Type) } - if action.Status != http.StatusServiceUnavailable { - t.Errorf("status = %d, want 503", action.Status) + status, _, _ := action.Violation.Render(); if status != http.StatusServiceUnavailable { + t.Errorf("status = %d, want 503", status) } } @@ -223,8 +223,8 @@ func TestTokenExchange_NoToken_Deny(t *testing.T) { if action.Type != pipeline.Reject { t.Fatalf("got %v, want Reject", action.Type) } - if action.Status != http.StatusUnauthorized { - t.Errorf("status = %d, want 401", action.Status) + status, _, _ := action.Violation.Render(); if status != http.StatusUnauthorized { + t.Errorf("status = %d, want 401", status) } } diff --git a/authbridge/authlib/plugins/tokenexchange.go b/authbridge/authlib/plugins/tokenexchange.go index 9904f338c..70778fe78 100644 --- a/authbridge/authlib/plugins/tokenexchange.go +++ b/authbridge/authlib/plugins/tokenexchange.go @@ -2,6 +2,7 @@ package plugins import ( "context" + "net/http" "github.com/kagenti/kagenti-extensions/authbridge/authlib/auth" "github.com/kagenti/kagenti-extensions/authbridge/authlib/pipeline" @@ -30,7 +31,15 @@ func (p *TokenExchange) OnRequest(ctx context.Context, pctx *pipeline.Context) p result := p.auth.HandleOutbound(ctx, authHeader, host) switch result.Action { case auth.ActionDeny: - return pipeline.Action{Type: pipeline.Reject, Status: result.DenyStatus, Reason: result.DenyReason} + // Outbound denials almost always come from failed token exchange + // at the IdP (upstream unreachable, bad credentials, audience + // refused). The auth layer returns the HTTP status it wants to + // expose; pick the closest well-known code for the body. + code := "upstream.token-exchange-failed" + if result.DenyStatus == http.StatusForbidden { + code = "policy.forbidden" + } + return pipeline.DenyStatus(result.DenyStatus, code, result.DenyReason) case auth.ActionReplaceToken: pctx.Headers.Set("Authorization", "Bearer "+result.Token) } diff --git a/authbridge/cmd/authbridge/listener/extauthz/server.go b/authbridge/cmd/authbridge/listener/extauthz/server.go index 773c1925e..a1da6d531 100644 --- a/authbridge/cmd/authbridge/listener/extauthz/server.go +++ b/authbridge/cmd/authbridge/listener/extauthz/server.go @@ -5,7 +5,6 @@ package extauthz import ( "context" - "encoding/json" "net/http" "strings" @@ -30,7 +29,8 @@ type Server struct { func (s *Server) Check(ctx context.Context, req *authv3.CheckRequest) (*authv3.CheckResponse, error) { httpReq := req.GetAttributes().GetRequest().GetHttp() if httpReq == nil { - return denied(codes.InvalidArgument, 400, "missing HTTP request attributes"), nil + return deniedFromAction(codes.InvalidArgument, + pipeline.DenyStatus(400, "auth.invalid-request", "missing HTTP request attributes")), nil } headers := httpReq.GetHeaders() @@ -49,7 +49,7 @@ func (s *Server) Check(ctx context.Context, req *authv3.CheckRequest) (*authv3.C } inAction := s.InboundPipeline.Run(ctx, inPctx) if inAction.Type == pipeline.Reject { - return denied(codes.Unauthenticated, inAction.Status, inAction.Reason), nil + return deniedFromAction(codes.Unauthenticated, inAction), nil } // Outbound exchange via pipeline @@ -62,7 +62,7 @@ func (s *Server) Check(ctx context.Context, req *authv3.CheckRequest) (*authv3.C originalAuth := outPctx.Headers.Get("Authorization") outAction := s.OutboundPipeline.Run(ctx, outPctx) if outAction.Type == pipeline.Reject { - return denied(codes.PermissionDenied, outAction.Status, outAction.Reason), nil + return deniedFromAction(codes.PermissionDenied, outAction), nil } newAuth := outPctx.Headers.Get("Authorization") @@ -87,27 +87,31 @@ func extractBearer(authHeader string) string { return "" } -func denied(code codes.Code, httpStatus int, msg string) *authv3.CheckResponse { - body, _ := json.Marshal(map[string]string{"error": msg}) +// deniedFromAction renders a pipeline Reject as an ext_authz CheckResponse +// preserving the plugin's status, headers, and body. The flat +// {"error":reason} body of the old denied() is gone — plugins can now +// supply structured bodies, Content-Type overrides, WWW-Authenticate +// challenges, Retry-After, etc., via action.Violation. +func deniedFromAction(code codes.Code, action pipeline.Action) *authv3.CheckResponse { + status, headers, body := action.Violation.Render() + setHeaders := make([]*corev3.HeaderValueOption, 0, len(headers)) + for k, vs := range headers { + for _, v := range vs { + setHeaders = append(setHeaders, &corev3.HeaderValueOption{ + Header: &corev3.HeaderValue{Key: k, Value: v}, + }) + } + } return &authv3.CheckResponse{ Status: &rpcstatus.Status{ Code: int32(code), - Message: msg, + Message: action.Violation.Reason, }, HttpResponse: &authv3.CheckResponse_DeniedResponse{ DeniedResponse: &authv3.DeniedHttpResponse{ - Status: &typev3.HttpStatus{ - Code: typev3.StatusCode(httpStatus), - }, - Body: string(body), - Headers: []*corev3.HeaderValueOption{ - { - Header: &corev3.HeaderValue{ - Key: "Content-Type", - Value: "application/json", - }, - }, - }, + Status: &typev3.HttpStatus{Code: typev3.StatusCode(status)}, + Body: string(body), + Headers: setHeaders, }, }, } diff --git a/authbridge/cmd/authbridge/listener/extproc/server.go b/authbridge/cmd/authbridge/listener/extproc/server.go index acea9d26f..08a50c623 100644 --- a/authbridge/cmd/authbridge/listener/extproc/server.go +++ b/authbridge/cmd/authbridge/listener/extproc/server.go @@ -130,8 +130,7 @@ func (s *Server) handleInbound(stream extprocv3.ExternalProcessor_ProcessServer, action := s.InboundPipeline.Run(ctx, pctx) if action.Type == pipeline.Reject { - return denyResponse(typev3.StatusCode(action.Status), - jsonError("unauthorized", action.Reason)), nil + return rejectFromAction(action), nil } s.recordInboundSession(pctx) @@ -150,8 +149,7 @@ func (s *Server) handleInboundBody(stream extprocv3.ExternalProcessor_ProcessSer action := s.InboundPipeline.Run(ctx, pctx) if action.Type == pipeline.Reject { - return denyResponse(typev3.StatusCode(action.Status), - jsonError("unauthorized", action.Reason)), nil + return rejectFromAction(action), nil } s.recordInboundSession(pctx) @@ -398,8 +396,7 @@ func (s *Server) handleOutbound(stream extprocv3.ExternalProcessor_ProcessServer originalAuth := pctx.Headers.Get("Authorization") action := s.OutboundPipeline.Run(ctx, pctx) if action.Type == pipeline.Reject { - return denyResponse(typev3.StatusCode_ServiceUnavailable, - jsonError("token_acquisition_failed", action.Reason)), nil + return rejectFromAction(action), nil } s.recordOutboundSession(pctx) @@ -434,8 +431,7 @@ func (s *Server) handleOutboundBody(stream extprocv3.ExternalProcessor_ProcessSe originalAuth := pctx.Headers.Get("Authorization") action := s.OutboundPipeline.Run(ctx, pctx) if action.Type == pipeline.Reject { - return denyResponse(typev3.StatusCode_ServiceUnavailable, - jsonError("token_acquisition_failed", action.Reason)), nil + return rejectFromAction(action), nil } s.recordOutboundSession(pctx) @@ -478,8 +474,7 @@ func (s *Server) handleResponseHeaders(ctx context.Context, headers *corev3.Head action := p.RunResponse(ctx, pctx) if action.Type == pipeline.Reject { - return denyResponse(typev3.StatusCode(action.Status), - jsonError("response_blocked", action.Reason)) + return rejectFromAction(action) } // No body phase will run; record the response event here. A2A responses @@ -516,8 +511,7 @@ func (s *Server) handleResponseBody(ctx context.Context, body []byte, pctx *pipe action := p.RunResponse(ctx, pctx) if action.Type == pipeline.Reject { - return denyResponse(typev3.StatusCode(action.Status), - jsonError("response_blocked", action.Reason)) + return rejectFromAction(action) } // The server's response may carry the server-assigned A2A contextId. If @@ -671,14 +665,29 @@ func replaceTokenResponse(token string) *extprocv3.ProcessingResponse { } } -func denyResponse(code typev3.StatusCode, body string) *extprocv3.ProcessingResponse { +// rejectFromAction turns a pipeline Reject into an Envoy ImmediateResponse, +// preserving the plugin's status/headers/body. Replaces the old +// denyResponse helper which hardcoded {"error":...,"message":...} at each +// call site. +func rejectFromAction(action pipeline.Action) *extprocv3.ProcessingResponse { + status, headers, body := action.Violation.Render() + immediate := &extprocv3.ImmediateResponse{ + Status: &typev3.HttpStatus{Code: typev3.StatusCode(status)}, + Body: body, + } + if len(headers) > 0 { + setHeaders := make([]*corev3.HeaderValueOption, 0, len(headers)) + for k, vs := range headers { + for _, v := range vs { + setHeaders = append(setHeaders, &corev3.HeaderValueOption{ + Header: &corev3.HeaderValue{Key: k, RawValue: []byte(v)}, + }) + } + } + immediate.Headers = &extprocv3.HeaderMutation{SetHeaders: setHeaders} + } return &extprocv3.ProcessingResponse{ - Response: &extprocv3.ProcessingResponse_ImmediateResponse{ - ImmediateResponse: &extprocv3.ImmediateResponse{ - Status: &typev3.HttpStatus{Code: code}, - Body: []byte(body), - }, - }, + Response: &extprocv3.ProcessingResponse_ImmediateResponse{ImmediateResponse: immediate}, } } @@ -694,11 +703,6 @@ func immediateResponse(httpStatus int, reason string) *extprocv3.ProcessingRespo } } -func jsonError(errorCode, message string) string { - b, _ := json.Marshal(map[string]string{"error": errorCode, "message": message}) - return string(b) -} - func requestHasBody(headers *corev3.HeaderMap) bool { method := getHeader(headers, ":method") if method == "GET" || method == "HEAD" || method == "OPTIONS" || method == "DELETE" { diff --git a/authbridge/cmd/authbridge/listener/extproc/server_test.go b/authbridge/cmd/authbridge/listener/extproc/server_test.go index 4b7aeb512..3ea6257dc 100644 --- a/authbridge/cmd/authbridge/listener/extproc/server_test.go +++ b/authbridge/cmd/authbridge/listener/extproc/server_test.go @@ -318,8 +318,13 @@ func TestExtProc_Outbound_Deny(t *testing.T) { if ir == nil { t.Fatal("expected ImmediateResponse (deny)") } - if ir.Status.Code != 503 { - t.Errorf("status = %d, want 503", ir.Status.Code) + // Pre-Violation, the listener used to hardcode 503 here regardless of + // what the plugin said. That's now fixed — we pass through whatever + // status the plugin (via auth.HandleOutbound) chose. NoTokenPolicyDeny + // returns 401 because "no auth credential present" is a caller problem, + // not an upstream-unavailable one. + if ir.Status.Code != 401 { + t.Errorf("status = %d, want 401", ir.Status.Code) } } diff --git a/authbridge/cmd/authbridge/listener/forwardproxy/server.go b/authbridge/cmd/authbridge/listener/forwardproxy/server.go index e150741c7..314270458 100644 --- a/authbridge/cmd/authbridge/listener/forwardproxy/server.go +++ b/authbridge/cmd/authbridge/listener/forwardproxy/server.go @@ -5,7 +5,6 @@ package forwardproxy import ( "bytes" - "encoding/json" "fmt" "io" "log/slog" @@ -81,10 +80,7 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) { action := s.OutboundPipeline.Run(r.Context(), pctx) if action.Type == pipeline.Reject { - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(action.Status) - body, _ := json.Marshal(map[string]string{"error": action.Reason}) - w.Write(body) + writeRejection(w, action) return } @@ -153,10 +149,7 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) { respAction := s.OutboundPipeline.RunResponse(r.Context(), pctx) if respAction.Type == pipeline.Reject { - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(respAction.Status) - body, _ := json.Marshal(map[string]string{"error": respAction.Reason}) - w.Write(body) + writeRejection(w, respAction) return } @@ -201,3 +194,18 @@ func extractBearer(authHeader string) string { } return "" } + +// writeRejection renders a pipeline Reject to the http.ResponseWriter. +// Preserves the plugin's status, headers, and body — the old flat +// {"error":reason} is gone; plugins now control the wire shape via +// action.Violation. +func writeRejection(w http.ResponseWriter, action pipeline.Action) { + status, headers, body := action.Violation.Render() + for k, vs := range headers { + for _, v := range vs { + w.Header().Add(k, v) + } + } + w.WriteHeader(status) + _, _ = w.Write(body) +} diff --git a/authbridge/cmd/authbridge/listener/reverseproxy/server.go b/authbridge/cmd/authbridge/listener/reverseproxy/server.go index f2e578a4f..e71296920 100644 --- a/authbridge/cmd/authbridge/listener/reverseproxy/server.go +++ b/authbridge/cmd/authbridge/listener/reverseproxy/server.go @@ -6,7 +6,6 @@ package reverseproxy import ( "bytes" "context" - "encoding/json" "fmt" "io" "log/slog" @@ -23,12 +22,20 @@ const maxBodySize = 1 << 20 // 1MB — matches Envoy's default per_stream_buffer type pctxKey struct{} +// responseRejectedError carries a pipeline Reject from the roundTripper +// back to the error handler, where it's rendered into the +// http.ResponseWriter. The embedded action keeps Violation.Render() and +// helper constructors available at the render site. type responseRejectedError struct { - status int - reason string + action pipeline.Action } -func (e *responseRejectedError) Error() string { return e.reason } +func (e *responseRejectedError) Error() string { + if e.action.Violation != nil { + return e.action.Violation.Reason + } + return "response rejected" +} // Server is an HTTP reverse proxy with inbound JWT validation. type Server struct { @@ -84,10 +91,7 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) { action := s.InboundPipeline.Run(r.Context(), pctx) if action.Type == pipeline.Reject { - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(action.Status) - body, _ := json.Marshal(map[string]string{"error": action.Reason}) - w.Write(body) + writeRejection(w, action) return } @@ -135,7 +139,7 @@ func (s *Server) modifyResponse(resp *http.Response) error { action := s.InboundPipeline.RunResponse(resp.Request.Context(), pctx) if action.Type == pipeline.Reject { - return &responseRejectedError{status: action.Status, reason: action.Reason} + return &responseRejectedError{action: action} } // If a plugin mutated ResponseBody, use the mutated version. @@ -149,10 +153,21 @@ func (s *Server) modifyResponse(resp *http.Response) error { func (s *Server) errorHandler(w http.ResponseWriter, _ *http.Request, err error) { if rErr, ok := err.(*responseRejectedError); ok { - w.Header().Set("Content-Type", "application/json") - w.WriteHeader(rErr.status) - json.NewEncoder(w).Encode(map[string]string{"error": rErr.reason}) + writeRejection(w, rErr.action) return } http.Error(w, `{"error":"bad gateway"}`, http.StatusBadGateway) } + +// writeRejection renders a pipeline Reject to the http.ResponseWriter, +// preserving the plugin's status, headers, and body. +func writeRejection(w http.ResponseWriter, action pipeline.Action) { + status, headers, body := action.Violation.Render() + for k, vs := range headers { + for _, v := range vs { + w.Header().Add(k, v) + } + } + w.WriteHeader(status) + _, _ = w.Write(body) +} diff --git a/authbridge/cmd/authbridge/main.go b/authbridge/cmd/authbridge/main.go index 72ed760cf..b071ec029 100644 --- a/authbridge/cmd/authbridge/main.go +++ b/authbridge/cmd/authbridge/main.go @@ -123,6 +123,21 @@ func main() { } } + // Invoke Init on any plugin implementing pipeline.Initializer. Done + // before listeners start so a plugin that fails to initialize takes + // down the process cleanly rather than serving traffic in a broken + // state. Init may do network I/O (register metrics, load a model, + // fetch a remote config); we give it the process context so an + // abort during init (Ctrl-C) cancels cooperatively. + initCtx, initCancel := context.WithTimeout(context.Background(), 60*time.Second) + defer initCancel() + if err := inboundPipeline.Start(initCtx); err != nil { + log.Fatalf("inbound pipeline Start: %v", err) + } + if err := outboundPipeline.Start(initCtx); err != nil { + log.Fatalf("outbound pipeline Start: %v", err) + } + // Build session store if enabled (nil when disabled — zero overhead). // Defaults to on; set session.enabled: false in runtime config to opt out. var sessions *session.Store @@ -244,6 +259,14 @@ func main() { if sessionAPISrv != nil { sessionAPISrv.Shutdown(shutdownCtx) } + + // Invoke plugin Shutdown hooks after listeners have stopped accepting + // traffic so in-flight work is allowed to complete before plugins + // flush state. Best-effort — individual errors are logged but do + // not abort the sequence. Bounded by shutdownCtx's deadline. + outboundPipeline.Stop(shutdownCtx) + inboundPipeline.Stop(shutdownCtx) + if sessions != nil { sessions.Close() } diff --git a/authbridge/demos/weather-agent/demo-ui.md b/authbridge/demos/weather-agent/demo-ui.md index 8c746eba9..8792e3968 100644 --- a/authbridge/demos/weather-agent/demo-ui.md +++ b/authbridge/demos/weather-agent/demo-ui.md @@ -11,7 +11,9 @@ that doesn't require token exchange. For a more advanced demo showing outbound token exchange and scope-based access control, see the [GitHub Issue Agent demo](../github-issue/demo.md). For the same weather images with **token exchange and AuthBridge on the tool** (plus a CI-style verify script), -see [Weather Agent — Advanced](demo-ui-advanced.md). +see [Weather Agent — Advanced](demo-ui-advanced.md). To observe the +plugin pipeline in real time while chatting with the agent, see +[Weather Agent with `abctl`](demo-with-abctl.md). ## What This Demo Shows diff --git a/authbridge/demos/weather-agent/demo-with-abctl.md b/authbridge/demos/weather-agent/demo-with-abctl.md new file mode 100644 index 000000000..fd6d17c6a --- /dev/null +++ b/authbridge/demos/weather-agent/demo-with-abctl.md @@ -0,0 +1,279 @@ +# Weather Agent Walkthrough with `abctl` + +This demo extends the standard [Weather Agent demo](./demo-ui.md) with a live view of AuthBridge's plugin pipeline using **`abctl`**, the terminal UI that reads AuthBridge's session API. You'll send chat messages from the Kagenti UI, then watch the request flow through inbound JWT validation → protocol parsers → outbound MCP calls → LLM inference → response — all with token counts, caller identity, request/response pairing, and request-scoped pipeline state visible in real time. + +## Prerequisites + +Before starting, complete these in order: + +1. **Install Kagenti** (operator + Keycloak + UI + SPIFFE/SPIRE): [Kagenti Installation Guide](https://github.com/kagenti/kagenti/blob/main/docs/install.md). Works on Kind (local) or OpenShift. +2. **Deploy the Weather Agent + Tool** through the Kagenti UI: [Weather Agent Demo with AuthBridge](./demo-ui.md). At the end of that demo you'll have: + - `weather-service` agent running in the `team1` namespace (with AuthBridge sidecars). + - `weather-tool-mcp` MCP tool running in `team1`. + - Keycloak configured with the agent's `agent-team1-weather-service-aud` scope. + - The Kagenti UI reachable at its route, with a working chat tab for the weather agent. + +Verify by sending one chat message and getting a weather response. Once that works, you're ready for this walkthrough. + +3. **Enable the protocol parser plugins** on the weather-service's AuthBridge. The default pipeline runs only `jwt-validation` (inbound) and `token-exchange` (outbound), which leaves abctl's Events pane without the `a2a` / `mcp` / `inf` protocol labels and without token counts. Patch the agent's runtime config to add the parsers, then restart the pod: + + ```sh + kubectl patch configmap authbridge-runtime-config -n team1 --type merge -p ' + data: + config.yaml: | + pipeline: + inbound: + plugins: + - jwt-validation + - a2a-parser + outbound: + plugins: + - token-exchange + - mcp-parser + - inference-parser + ' + kubectl rollout restart deployment/weather-service -n team1 + ``` + + Merging only the `pipeline:` key into the existing YAML is brittle across operator versions. If the patch above doesn't take effect, `kubectl edit configmap authbridge-runtime-config -n team1` and add the `pipeline:` section to the existing `config.yaml` by hand. See [mcp-parser demo](../mcp-parser/README.md) for the full config format and rationale. + +## 1. Build `abctl` + +`abctl` lives in the `kagenti-extensions` repo. Build it once: + +```sh +git clone https://github.com/kagenti/kagenti-extensions.git +cd kagenti-extensions/authbridge/cmd/abctl +go build . +``` + +Produces a single ~10 MB binary at `./abctl`. See [the abctl README](../../cmd/abctl/README.md) for full flags and keybindings. + +## 2. Port-forward the agent's session API + +AuthBridge exposes session telemetry on port `9094` of the agent pod. Forward it to your laptop: + +```sh +POD=$(kubectl get pod -n team1 -l app.kubernetes.io/name=weather-service \ + -o jsonpath='{.items[0].metadata.name}') +kubectl port-forward -n team1 "$POD" 9094:9094 & +``` + +Quick sanity check: + +```sh +curl -s http://localhost:9094/v1/sessions | python3 -m json.tool +``` + +Should return `{"sessions": [...]}` — possibly empty on a fresh pod, populated once you've sent a chat message. + +## 3. Launch `abctl` + +```sh +./abctl +``` + +Opens the terminal UI pointed at `http://localhost:9094`. The initial view is the **Sessions** pane: + +``` +╭─ abctl · http://localhost:9094 · [Sessions] Pipeline ─────────────────╮ +│ ID UPDATED EVENTS TOKENS ACTIVE │ +│ (no sessions yet) │ +│ │ +│ ● connected 0.0 ev/s drops: 0 │ +│ [↑↓] nav [↵] drill [tab] pipeline [/] filter [p] pause [?] help [q] │ +╰────────────────────────────────────────────────────────────────────────────────╯ +``` + +## 4. Send a chat message from the Kagenti UI + +In a browser: + +1. Open the Kagenti UI (route from `kubectl get route kagenti-ui -n kagenti-system` or port-forward). +2. Navigate to the Weather Service agent's chat. +3. Type a question: *"What's the weather in New York?"* +4. Wait for the agent to reply. + +`abctl` updates in real time. A new session bucket appears with ~23 events and a token count: + +``` + ID UPDATED EVENTS TOKENS ACTIVE +▸ 4647e888-db99-4739-926e-8bcceeb237c6 just now 23 462 ● +``` + +- **ID** — the A2A `contextId` assigned to this conversation. +- **UPDATED** — relative time since the last event in the bucket. +- **EVENTS** — number of pipeline events (inbound + outbound × request + response). +- **TOKENS** — total tokens consumed across all inference calls in this conversation. +- **ACTIVE** (●) — the most recently updated session. + +## 5. Drill into the session's events + +Select the row with `↑`/`↓` (or `j`/`k`) and press `Enter` to open the **Events** pane: + +``` +╭─ abctl · 4647e888-db99-4739-926e-8bcceeb237c6 ─────────────────────────╮ +│ ┌─ IDENTITY ─────────────────────────────────────────────────────────┐ │ +│ │ subject alice │ │ +│ │ client kagenti │ │ +│ │ scopes openid, email, profile +1 more │ │ +│ └─────────────────────────────────────────────────────────────────────┘ │ +│ │ +│ EVENTS (23) │ +│ TIME DIR PHASE PROTO METHOD STATUS DURATION TOKENS HOST │ +│ 12:24:39.77 in req a2a message/stream │ +│ 12:24:39.82 out req mcp tools/list │ +│ 12:24:39.85 out └resp mcp tools/list 200 31ms … │ +│ 12:24:39.90 out req mcp tools/call │ +│ 12:24:40.05 out └resp mcp tools/call 200 148ms … │ +│ 12:24:40.10 out req inf gpt-4 │ +│ 12:24:41.65 out └resp inf gpt-4 200 1550ms 185 … │ +│ … │ +│ 12:24:43.40 in └resp a2a message/stream 200 4256ms │ +│ │ +│ ● connected focus=4647e888 23 events │ +│ [↑↓] nav [↵] detail [esc] back [/] filter [p] pause [q] quit │ +╰──────────────────────────────────────────────────────────────────────────╯ +``` + +What to notice: + +- **IDENTITY banner** at the top of the pane — the caller subject (`alice` from Keycloak), the OAuth client (`kagenti`, the UI), and the scopes their token carried. This is per-session, extracted from the inbound JWT. +- **Event rows**: + - `DIR` — `in` = inbound (caller → agent), `out` = outbound (agent → tool or LLM). + - `PHASE` — `req` for request, `└resp` for the matching response (the `└` visually connects a response row back to its paired request). + - `PROTO` — `a2a` (user-facing protocol), `mcp` (tool calls), `inf` (LLM inference). + - `METHOD` — the JSON-RPC method or model name. + - `STATUS` / `DURATION` — populated on response rows. + - `TOKENS` — populated only on inference response rows; shows the total tokens for that one LLM call. + +The flow you're looking at: inbound A2A request → outbound MCP `tools/list` (discover what the weather tool offers) → LLM inference (plan) → outbound MCP `tools/call` (execute weather query) → LLM inference (synthesize answer) → inbound A2A response. + +## 6. Inspect an individual event + +Select any row and press `Enter` for the **Detail** pane. It shows the event as pretty-printed JSON, filtered to the fields relevant to that event's phase: + +**A request event** (you selected an `inf req` row): + +```json +{ + "at": "2026-05-06T12:24:40.10Z", + "direction": "outbound", + "phase": "request", + "sessionId": "4647e888-…", + "host": "litellm-prod.apps.…", + "inference": { + "model": "gpt-4", + "messages": [ + {"role": "system", "content": "You are a weather assistant."}, + {"role": "user", "content": "What's the weather in New York?"} + ], + "temperature": 0.7, + "stream": false, + "tools": [ + { + "name": "get_weather", + "description": "Fetch current weather for a city", + "parameters": { "type": "object", "properties": { … } } + } + ] + } +} +``` + +**A response event** (same call's `└resp` row): + +```json +{ + "at": "2026-05-06T12:24:41.65Z", + "direction": "outbound", + "phase": "response", + "statusCode": 200, + "durationMs": 1550, + "inference": { + "model": "gpt-4", + "completion": "The weather in New York is currently 72°F and sunny.", + "finishReason": "stop", + "promptTokens": 167, + "completionTokens": 18, + "totalTokens": 185 + } +} +``` + +Note the detail view **strictly separates request and response fields** — a request row shows only request-side data, a response row shows only response-side data. The `identity` field is filtered out here because it's already shown in the banner above the table. + +Press `y` (yank) to write the full, unfiltered wire-format JSON to `/tmp/abctl-event--.json` (mode `0600`). Useful for sharing with teammates or diffing across runs. + +Press `Esc` to go back to the events pane. + +## 7. View the plugin pipeline composition + +From any top-level pane, press `Tab` to switch between **Sessions** and **Pipeline**: + +``` +╭─ abctl · http://localhost:9094 · Sessions [Pipeline] ─────────────────────────╮ +│ # DIRECTION PLUGIN WRITES BODY EVENTS │ +│ 1 inbound jwt-validation no │ +│ 2 inbound a2a-parser a2a yes 2 │ +│ ── (app) ── │ +│ 1 outbound token-exchange no │ +│ 2 outbound mcp-parser mcp yes 10 │ +│ 3 outbound inference-parser inference yes 4 │ +│ │ +│ [↑↓] nav [↵] plugin detail [tab] sessions [q] quit │ +╰────────────────────────────────────────────────────────────────────────────────╯ +``` + +This view shows the plugin chain for both inbound and outbound directions, separated by a `── (app) ──` divider that represents the agent process sitting between them. Columns: + +- **#** — execution order within the direction (restarts at 1 for outbound). +- **DIRECTION** — which chain the plugin lives in. +- **PLUGIN** — name (as returned by `Plugin.Name()`). +- **WRITES** — extension slots the plugin populates on `pctx.Extensions`. +- **BODY** — `yes` if the plugin buffers bodies (drives Envoy `ProcessingMode: BUFFERED`), `no` otherwise. +- **EVENTS** — number of events in the currently-selected session whose data was populated by this plugin (empty when zero). Use this to see *which* plugins were active in this conversation. + +Session recording itself isn't a plugin — the listener (ext_proc / ext_authz / proxy) snapshots `pctx` after each phase and appends to `session.Store` directly, so it doesn't appear in this table. + +Select any plugin and press `Enter` for a plugin-detail pane with its declared reads / writes / body-access flag. + +## 8. Follow a conversation live + +Back in the Sessions or Events pane, everything is **streamed via SSE** from the `/v1/events` endpoint. Send another chat message and watch: + +- The session's `EVENTS` count tick up in real time. +- The `TOKENS` column grows by ~200 per inference call. +- New rows appear in the events pane if you've drilled in; auto-follow keeps you at the bottom unless you've scrolled up. + +The bottom-footer rate indicator (`3.2 ev/s`) is a live smoothed events-per-second gauge — useful to tell if traffic is flowing or something upstream is stuck. + +## 9. Filter and search + +Press `/` in Sessions or Events to open a substring filter. Filters apply to: + +- **Sessions pane**: session ID substring. +- **Events pane**: matches across `host`, `method`, `proto`, `A2A parts content`, `LLM completion`, `MCP error message`, caller `subject` / `clientId`. + +`Esc` to cancel, `Enter` to commit. Press `/` again and clear with `Esc` to remove the filter. + +## 10. Pause / resume + +Press `p` to pause stream rendering. The SSE connection stays open (events don't get dropped), but the UI stops refreshing — useful when you want to read a burst of events without chasing them. Press `p` again to resume. + +## Troubleshooting + +| Symptom | Likely cause | +|---|---| +| Footer shows `✗ failed: …` | `/v1/sessions` fetch failed — check the port-forward is still up and the pod is Ready. | +| Sessions appear and disappear | The pod restarted (session store is in-memory). | +| `TOKENS` shows `—` on a session | No inference events in that bucket yet, or the server is older than the token-aggregation change (client-side fallback sums from streamed events). | +| Multiple session buckets for one conversation | The UI backend isn't forwarding the A2A `contextId` on subsequent turns. See [kagenti#1481](https://github.com/kagenti/kagenti/pull/1481) for the fix. | +| Events missing identity | `jwt-validation` didn't run on this code path (e.g. bypass path) or the caller didn't send a Bearer token. | +| IDENTITY banner reports `subject —` | The JWT had no `sub` claim (common for Keycloak public-client tokens). Field is still extracted from `azp` / scopes. | + +## What to explore next + +- **GitHub-issue demo** — same idea with **outbound token exchange** and scope-based access control: [demo-ui.md](../github-issue/demo-ui.md). In `abctl` you'll see an additional outbound `exchange` plugin activity on the pipeline pane, and different `targetAudience` values on outbound events. +- **Advanced weather demo** — adds **AuthBridge on the tool side** so you can inspect a second set of inbound pipeline events when the agent calls the MCP tool: [demo-ui-advanced.md](./demo-ui-advanced.md). +- **The plugin pipeline spec** — if you want to understand the data structures (`pctx`, `Extensions`, `SessionEvent`, `GetState`/`SetState`), or integrate a custom plugin or sub-pipeline engine: [pipeline/README.md](../../authlib/pipeline/README.md).