diff --git a/authbridge/authlib/pipeline/action.go b/authbridge/authlib/pipeline/action.go new file mode 100644 index 000000000..8a88a0adf --- /dev/null +++ b/authbridge/authlib/pipeline/action.go @@ -0,0 +1,17 @@ +package pipeline + +// ActionType represents the result of a plugin's processing. +type ActionType int + +const ( + Continue ActionType = iota + Reject +) + +// Action is returned by a plugin to indicate whether processing should continue +// or the request should be rejected. +type Action struct { + Type ActionType + Status int // HTTP status code (only for Reject) + Reason string // human-readable reason (only for Reject) +} diff --git a/authbridge/authlib/pipeline/context.go b/authbridge/authlib/pipeline/context.go new file mode 100644 index 000000000..5c08da972 --- /dev/null +++ b/authbridge/authlib/pipeline/context.go @@ -0,0 +1,41 @@ +package pipeline + +import ( + "net/http" + + "github.com/kagenti/kagenti-extensions/authbridge/authlib/routing" + "github.com/kagenti/kagenti-extensions/authbridge/authlib/validation" +) + +// Direction indicates whether a request is inbound (caller → this agent) or +// outbound (this agent → target service). +type Direction int + +const ( + Inbound Direction = iota + Outbound +) + +// Context is the shared state passed through the plugin pipeline. +// Plugins read and mutate fields directly — there is no separate mutation API. +type Context struct { + Direction Direction + Method string + Host string + Path string + Headers http.Header + Body []byte // nil unless at least one plugin declares BodyAccess: true + + Agent *AgentIdentity + Claims *validation.Claims // nil before jwt-validation runs + Route *routing.ResolvedRoute + + Extensions Extensions +} + +// AgentIdentity carries the agent's own workload identity. +type AgentIdentity struct { + ClientID string + WorkloadID string + TrustDomain string +} diff --git a/authbridge/authlib/pipeline/extensions.go b/authbridge/authlib/pipeline/extensions.go new file mode 100644 index 000000000..627cc2e90 --- /dev/null +++ b/authbridge/authlib/pipeline/extensions.go @@ -0,0 +1,104 @@ +package pipeline + +import "time" + +// Extensions holds typed extension slots for plugin-to-plugin communication. +// Each slot is populated by a specific plugin and consumed by downstream plugins. +type Extensions struct { + MCP *MCPExtension + A2A *A2AExtension + Security *SecurityExtension + Delegation *DelegationExtension + Custom map[string]any +} + +// MCPExtension carries parsed MCP JSON-RPC metadata. +// Exactly one of Tool, Resource, or Prompt is populated per request. +type MCPExtension struct { + Method string // JSON-RPC method: "tools/call", "resources/read", "prompts/get" + RPCID any // JSON-RPC id for request-response correlation + + Tool *MCPToolMetadata + Resource *MCPResourceMetadata + Prompt *MCPPromptMetadata +} + +// MCPToolMetadata is populated for tools/call requests. +type MCPToolMetadata struct { + Name string + Args map[string]any +} + +// MCPResourceMetadata is populated for resources/read requests. +type MCPResourceMetadata struct { + URI string +} + +// MCPPromptMetadata is populated for prompts/get requests. +type MCPPromptMetadata struct { + Name string + Args map[string]string +} + +// A2AExtension carries parsed A2A protocol metadata. +type A2AExtension struct { + TaskID string + Method string // "tasks/send", "tasks/get", etc. + Parts []A2APart +} + +// A2APart represents a message part in an A2A request. +type A2APart struct { + Type string // "text", "file", "data" + Content string +} + +// SecurityExtension carries guardrail output. +// Caller identity is already in ctx.Agent and ctx.Claims — this slot is only +// for downstream signals from content-inspection plugins. +type SecurityExtension struct { + Labels []string + Blocked bool + BlockReason string +} + +// DelegationExtension tracks the token delegation chain across hops. +// The chain is append-only and unexported to prevent forgery or truncation. +type DelegationExtension struct { + chain []DelegationHop + Origin string // original caller's subject ID + Actor string // current actor's subject ID +} + +// Chain returns a copy of the delegation chain. The copy prevents callers from +// mutating the backing slice (truncation, reordering, forgery). +func (d *DelegationExtension) Chain() []DelegationHop { + out := make([]DelegationHop, len(d.chain)) + copy(out, d.chain) + return out +} + +// Depth returns the number of hops in the delegation chain. +func (d *DelegationExtension) Depth() int { + return len(d.chain) +} + +// DelegationHop represents one hop in the delegation chain. +type DelegationHop struct { + SubjectID string + Scopes []string + Timestamp time.Time +} + +// AppendHop adds a hop to the delegation chain. This is the only way to extend +// the chain — direct mutation is prevented by the unexported slice. +// +// AppendHop is not safe for concurrent use. The pipeline guarantees sequential +// invocation. +func (d *DelegationExtension) AppendHop(hop DelegationHop) { + d.chain = append(d.chain, hop) + if d.Origin == "" { + d.Origin = hop.SubjectID + } + d.Actor = hop.SubjectID +} diff --git a/authbridge/authlib/pipeline/pipeline.go b/authbridge/authlib/pipeline/pipeline.go new file mode 100644 index 000000000..7aeba1aa9 --- /dev/null +++ b/authbridge/authlib/pipeline/pipeline.go @@ -0,0 +1,109 @@ +package pipeline + +import ( + "context" + "fmt" +) + +// Pipeline holds an ordered list of plugins and runs them sequentially. +type Pipeline struct { + plugins []Plugin +} + +// defaultSlots lists the built-in extension slot names. +var defaultSlots = map[string]bool{ + "mcp": true, + "a2a": true, + "security": true, + "delegation": true, + "custom": true, +} + +// Option configures pipeline construction. +type Option func(*options) + +type options struct { + extraSlots []string +} + +// WithSlots registers additional valid extension slot names beyond the built-in set. +// Use this when a bridge plugin (e.g., CPEX) produces extensions not in the default set. +func WithSlots(slots ...string) Option { + return func(o *options) { + o.extraSlots = append(o.extraSlots, slots...) + } +} + +// New creates a Pipeline from the given plugins after validating capability wiring. +// Returns an error if any plugin declares a read on a slot that no earlier plugin writes. +func New(plugins []Plugin, opts ...Option) (*Pipeline, error) { + var o options + for _, opt := range opts { + opt(&o) + } + validSlots := make(map[string]bool, len(defaultSlots)+len(o.extraSlots)) + for k, v := range defaultSlots { + validSlots[k] = v + } + for _, s := range o.extraSlots { + validSlots[s] = true + } + if err := validateCapabilities(plugins, validSlots); err != nil { + return nil, err + } + return &Pipeline{plugins: plugins}, nil +} + +// Run executes the request phase of the pipeline sequentially. +// If any plugin returns Reject, the pipeline stops and returns that action. +func (p *Pipeline) Run(ctx context.Context, pctx *Context) Action { + for _, plugin := range p.plugins { + if ctx.Err() != nil { + return Action{Type: Reject, Status: 499, Reason: "request cancelled"} + } + action := plugin.OnRequest(ctx, pctx) + if action.Type == Reject { + return action + } + } + return Action{Type: Continue} +} + +// RunResponse executes the response phase in reverse order. +// The last plugin in the chain sees the response first. +func (p *Pipeline) RunResponse(ctx context.Context, pctx *Context) Action { + for i := len(p.plugins) - 1; i >= 0; i-- { + if ctx.Err() != nil { + return Action{Type: Reject, Status: 499, Reason: "request cancelled"} + } + action := p.plugins[i].OnResponse(ctx, pctx) + if action.Type == Reject { + return action + } + } + return Action{Type: Continue} +} + +// 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 { + written := make(map[string]bool) + for _, plugin := range plugins { + caps := plugin.Capabilities() + for _, slot := range caps.Reads { + if !validSlots[slot] { + return fmt.Errorf("plugin %q declares read on unknown slot %q", plugin.Name(), slot) + } + if !written[slot] { + return fmt.Errorf("plugin %q reads slot %q but no earlier plugin writes it", plugin.Name(), slot) + } + } + for _, slot := range caps.Writes { + if !validSlots[slot] { + return fmt.Errorf("plugin %q declares write on unknown slot %q", plugin.Name(), slot) + } + written[slot] = true + } + } + return nil +} diff --git a/authbridge/authlib/pipeline/pipeline_test.go b/authbridge/authlib/pipeline/pipeline_test.go new file mode 100644 index 000000000..150a99dd6 --- /dev/null +++ b/authbridge/authlib/pipeline/pipeline_test.go @@ -0,0 +1,363 @@ +package pipeline + +import ( + "context" + "testing" +) + +// stubPlugin is a minimal Plugin implementation for testing. +type stubPlugin struct { + name string + caps PluginCapabilities + onReq func(ctx context.Context, pctx *Context) Action + onResp func(ctx context.Context, pctx *Context) Action +} + +func (s *stubPlugin) Name() string { return s.name } +func (s *stubPlugin) Capabilities() PluginCapabilities { return s.caps } +func (s *stubPlugin) OnRequest(ctx context.Context, pctx *Context) Action { + if s.onReq != nil { + return s.onReq(ctx, pctx) + } + return Action{Type: Continue} +} +func (s *stubPlugin) OnResponse(ctx context.Context, pctx *Context) Action { + if s.onResp != nil { + return s.onResp(ctx, pctx) + } + return Action{Type: Continue} +} + +func TestPipelineRun_EmptyPipeline(t *testing.T) { + p, err := New(nil) + if err != nil { + t.Fatalf("New(nil) returned error: %v", err) + } + pctx := &Context{} + action := p.Run(context.Background(), pctx) + if action.Type != Continue { + t.Errorf("empty pipeline returned %v, want Continue", action.Type) + } +} + +func TestPipelineRun_Sequential(t *testing.T) { + var order []string + p1 := &stubPlugin{ + name: "first", + onReq: func(_ context.Context, pctx *Context) Action { + order = append(order, "first") + pctx.Extensions.Custom = map[string]any{"key": "value"} + return Action{Type: Continue} + }, + } + p2 := &stubPlugin{ + name: "second", + caps: PluginCapabilities{Reads: []string{"custom"}}, + onReq: func(_ context.Context, pctx *Context) Action { + order = append(order, "second") + if pctx.Extensions.Custom["key"] != "value" { + t.Error("second plugin did not see mutation from first") + } + return Action{Type: Continue} + }, + } + p1.caps = PluginCapabilities{Writes: []string{"custom"}} + + pipe, err := New([]Plugin{p1, p2}) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + pctx := &Context{} + action := pipe.Run(context.Background(), pctx) + if action.Type != Continue { + t.Errorf("got %v, want Continue", action.Type) + } + if len(order) != 2 || order[0] != "first" || order[1] != "second" { + t.Errorf("execution order = %v, want [first second]", order) + } +} + +func TestPipelineRun_Reject(t *testing.T) { + called := false + p1 := &stubPlugin{ + name: "rejecter", + onReq: func(_ context.Context, _ *Context) Action { + return Action{Type: Reject, Status: 403, Reason: "forbidden"} + }, + } + p2 := &stubPlugin{ + name: "never-called", + onReq: func(_ context.Context, _ *Context) Action { + called = true + return Action{Type: Continue} + }, + } + pipe, err := New([]Plugin{p1, p2}) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + pctx := &Context{} + action := pipe.Run(context.Background(), pctx) + 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.Reason != "forbidden" { + t.Errorf("reason = %q, want %q", action.Reason, "forbidden") + } + if called { + t.Error("second plugin was called after first rejected") + } +} + +func TestPipelineRunResponse_ReverseOrder(t *testing.T) { + var order []string + p1 := &stubPlugin{ + name: "first", + onResp: func(_ context.Context, _ *Context) Action { + order = append(order, "first") + return Action{Type: Continue} + }, + } + p2 := &stubPlugin{ + name: "second", + onResp: func(_ context.Context, _ *Context) Action { + order = append(order, "second") + return Action{Type: Continue} + }, + } + pipe, err := New([]Plugin{p1, p2}) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + pctx := &Context{} + action := pipe.RunResponse(context.Background(), pctx) + if action.Type != Continue { + t.Errorf("got %v, want Continue", action.Type) + } + if len(order) != 2 || order[0] != "second" || order[1] != "first" { + t.Errorf("response order = %v, want [second first]", order) + } +} + +func TestPipelineRunResponse_Reject(t *testing.T) { + called := false + p1 := &stubPlugin{ + name: "first", + onResp: func(_ context.Context, _ *Context) Action { + called = true + return Action{Type: Continue} + }, + } + p2 := &stubPlugin{ + name: "rejecter", + onResp: func(_ context.Context, _ *Context) Action { + return Action{Type: Reject, Status: 500, Reason: "response blocked"} + }, + } + pipe, err := New([]Plugin{p1, p2}) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + pctx := &Context{} + action := pipe.RunResponse(context.Background(), pctx) + if action.Type != Reject { + t.Errorf("got %v, want Reject", action.Type) + } + if called { + t.Error("first plugin OnResponse was called after second rejected (reverse order)") + } +} + +func TestNew_ValidCapabilities(t *testing.T) { + plugins := []Plugin{ + &stubPlugin{ + name: "writer", + caps: PluginCapabilities{Writes: []string{"mcp"}}, + }, + &stubPlugin{ + name: "reader", + caps: PluginCapabilities{Reads: []string{"mcp"}}, + }, + } + _, err := New(plugins) + if err != nil { + t.Errorf("expected no error, got: %v", err) + } +} + +func TestNew_InvalidCapabilities_ReadBeforeWrite(t *testing.T) { + plugins := []Plugin{ + &stubPlugin{ + name: "reader", + caps: PluginCapabilities{Reads: []string{"mcp"}}, + }, + &stubPlugin{ + name: "writer", + caps: PluginCapabilities{Writes: []string{"mcp"}}, + }, + } + _, err := New(plugins) + if err == nil { + t.Fatal("expected error for read-before-write, got nil") + } +} + +func TestNew_InvalidCapabilities_UnknownSlot(t *testing.T) { + plugins := []Plugin{ + &stubPlugin{ + name: "bad-reader", + caps: PluginCapabilities{Reads: []string{"nonexistent"}}, + }, + } + _, err := New(plugins) + if err == nil { + t.Fatal("expected error for unknown slot, got nil") + } +} + +func TestNew_MultipleWriters(t *testing.T) { + plugins := []Plugin{ + &stubPlugin{ + name: "writer-1", + caps: PluginCapabilities{Writes: []string{"security"}}, + }, + &stubPlugin{ + name: "writer-2", + caps: PluginCapabilities{Writes: []string{"security"}}, + }, + &stubPlugin{ + name: "reader", + caps: PluginCapabilities{Reads: []string{"security"}}, + }, + } + _, err := New(plugins) + if err != nil { + t.Errorf("multiple writers should be valid, got: %v", err) + } +} + +func TestNew_NoCapabilities(t *testing.T) { + plugins := []Plugin{ + &stubPlugin{name: "simple"}, + } + _, err := New(plugins) + if err != nil { + t.Errorf("plugin with no capabilities should be valid, got: %v", err) + } +} + +func TestNew_CustomSlot(t *testing.T) { + plugins := []Plugin{ + &stubPlugin{ + name: "custom-writer", + caps: PluginCapabilities{Writes: []string{"custom"}}, + }, + &stubPlugin{ + name: "custom-reader", + caps: PluginCapabilities{Reads: []string{"custom"}}, + }, + } + _, err := New(plugins) + if err != nil { + t.Errorf("custom slot should be valid, got: %v", err) + } +} + +func TestNew_WithSlots(t *testing.T) { + plugins := []Plugin{ + &stubPlugin{ + name: "cpex-bridge", + caps: PluginCapabilities{Writes: []string{"cpex.completion"}}, + }, + &stubPlugin{ + name: "consumer", + caps: PluginCapabilities{Reads: []string{"cpex.completion"}}, + }, + } + // Without WithSlots, this should fail + _, err := New(plugins) + if err == nil { + t.Fatal("expected error for unregistered slot without WithSlots") + } + + // With WithSlots, this should succeed + _, err = New(plugins, WithSlots("cpex.completion")) + if err != nil { + t.Errorf("expected no error with WithSlots, got: %v", err) + } +} + +func TestDelegationExtension_AppendHop(t *testing.T) { + d := &DelegationExtension{} + + d.AppendHop(DelegationHop{SubjectID: "alice", Scopes: []string{"read", "write"}}) + if d.Depth() != 1 { + t.Errorf("Depth() = %d, want 1", d.Depth()) + } + if d.Origin != "alice" { + t.Errorf("origin = %q, want %q", d.Origin, "alice") + } + if d.Actor != "alice" { + t.Errorf("actor = %q, want %q", d.Actor, "alice") + } + + d.AppendHop(DelegationHop{SubjectID: "bob", Scopes: []string{"read"}}) + if d.Depth() != 2 { + t.Errorf("Depth() = %d, want 2", d.Depth()) + } + if d.Origin != "alice" { + t.Errorf("origin should stay %q, got %q", "alice", d.Origin) + } + if d.Actor != "bob" { + t.Errorf("actor = %q, want %q", d.Actor, "bob") + } + chain := d.Chain() + if len(chain) != 2 { + t.Errorf("Chain() length = %d, want 2", len(chain)) + } +} + +func TestDelegationExtension_ChainIsCopy(t *testing.T) { + d := &DelegationExtension{} + d.AppendHop(DelegationHop{SubjectID: "alice"}) + + chain := d.Chain() + chain[0].SubjectID = "tampered" + + original := d.Chain() + if original[0].SubjectID != "alice" { + t.Errorf("Chain() returned reference to backing slice, mutation leaked") + } +} + +func TestPipelineRun_ContextCancellation(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + called := false + p1 := &stubPlugin{ + name: "should-not-run", + onReq: func(_ context.Context, _ *Context) Action { + called = true + return Action{Type: Continue} + }, + } + pipe, err := New([]Plugin{p1}) + if err != nil { + t.Fatalf("New returned error: %v", err) + } + pctx := &Context{} + action := pipe.Run(ctx, pctx) + 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) + } + if called { + t.Error("plugin was called despite cancelled context") + } +} diff --git a/authbridge/authlib/pipeline/plugin.go b/authbridge/authlib/pipeline/plugin.go new file mode 100644 index 000000000..b2686db6e --- /dev/null +++ b/authbridge/authlib/pipeline/plugin.go @@ -0,0 +1,20 @@ +package pipeline + +import "context" + +// Plugin is the interface that all pipeline extensions implement. +type Plugin interface { + Name() string + Capabilities() PluginCapabilities + OnRequest(ctx context.Context, pctx *Context) Action + OnResponse(ctx context.Context, pctx *Context) Action +} + +// PluginCapabilities declares what extension slots a plugin reads and writes. +// The pipeline validates at startup that all reads are satisfied by an earlier +// plugin's writes. +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 +}