diff --git a/aibridge/bridge.go b/aibridge/bridge.go index b556b20d68..541a0bbf1b 100644 --- a/aibridge/bridge.go +++ b/aibridge/bridge.go @@ -7,11 +7,13 @@ import ( "net/http" "net/url" "regexp" + "strconv" "strings" "sync" "sync/atomic" "time" + "github.com/google/uuid" "github.com/hashicorp/go-multierror" "github.com/sony/gobreaker/v2" "go.opentelemetry.io/otel/codes" @@ -27,6 +29,7 @@ import ( "github.com/coder/coder/v2/aibridge/provider" "github.com/coder/coder/v2/aibridge/recorder" "github.com/coder/coder/v2/aibridge/tracing" + agplaibridge "github.com/coder/coder/v2/coderd/aibridge" "github.com/coder/quartz" ) @@ -245,6 +248,18 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC client := GuessClient(r) sessionID := GuessSessionID(client, r) + // Read and validate Agent Firewall correlation headers. The + // values are captured here and recorded below; the headers + // themselves are stripped from the upstream request by + // PrepareClientHeaders. Fail closed: reject the request if the + // headers are partial or malformed. + agentFirewallSessionID, agentFirewallSeqNumber, err := extractAgentFirewallHeaders(r) + if err != nil { + logger.Warn(ctx, "rejecting request with invalid agent firewall headers", slog.Error(err)) + http.Error(w, "invalid agent firewall headers", http.StatusBadRequest) + return + } + interceptor, err := p.CreateInterceptor(w, r.WithContext(ctx), tracer) if err != nil { span.SetStatus(codes.Error, fmt.Sprintf("failed to create interceptor: %v", err)) @@ -293,18 +308,20 @@ func newInterceptionProcessor(p provider.Provider, cbs *circuitbreaker.ProviderC interceptor.Setup(logger, asyncRecorder, mcpProxy) if err := rec.RecordInterception(ctx, &recorder.InterceptionRecord{ - ID: interceptor.ID().String(), - InitiatorID: actor.ID, - Metadata: actor.Metadata, - Model: interceptor.Model(), - Provider: p.Type(), - ProviderName: p.Name(), - UserAgent: r.UserAgent(), - Client: string(client), - ClientSessionID: sessionID, - CorrelatingToolCallID: interceptor.CorrelatingToolCallID(), - CredentialKind: string(cred.Kind()), - CredentialHint: cred.Hint(), + ID: interceptor.ID().String(), + InitiatorID: actor.ID, + Metadata: actor.Metadata, + Model: interceptor.Model(), + Provider: p.Type(), + ProviderName: p.Name(), + UserAgent: r.UserAgent(), + Client: string(client), + ClientSessionID: sessionID, + CorrelatingToolCallID: interceptor.CorrelatingToolCallID(), + AgentFirewallSessionID: agentFirewallSessionID, + AgentFirewallSequenceNumber: agentFirewallSeqNumber, + CredentialKind: string(cred.Kind()), + CredentialHint: cred.Hint(), }); err != nil { span.SetStatus(codes.Error, fmt.Sprintf("failed to record interception: %v", err)) logger.Warn(ctx, "failed to record interception", slog.Error(err)) @@ -461,3 +478,45 @@ func mergeContexts(base, other context.Context) context.Context { }() return ctx } + +// extractAgentFirewallHeaders reads and parses the Agent Firewall +// correlation headers from the request. Both headers must be present +// together with a valid UUID session ID and a non-negative int32 +// sequence number, or both must be absent. Partial or malformed headers +// return an error so the caller can reject the request (fail closed). +func extractAgentFirewallHeaders(r *http.Request) (sessionID *string, seqNumber *int32, err error) { + rawSessionID := r.Header.Get(agplaibridge.HeaderAgentFirewallSessionID) + rawSeqNumber := r.Header.Get(agplaibridge.HeaderAgentFirewallSequenceNumber) + + hasSessionID := rawSessionID != "" + hasSeqNumber := rawSeqNumber != "" + + switch { + case !hasSessionID && !hasSeqNumber: + // Neither header present; request did not traverse Agent Firewall. + return nil, nil, nil + case hasSessionID && !hasSeqNumber: + return nil, nil, xerrors.Errorf("agent firewall session ID header present without sequence number") + case !hasSessionID && hasSeqNumber: + return nil, nil, xerrors.Errorf("agent firewall sequence number header present without session ID") + } + + // Both headers present; validate the session ID is a UUID. Storing an + // invalid value would silently drop the firewall correlation to NULL + // downstream, so reject it here instead. + if _, parseErr := uuid.Parse(rawSessionID); parseErr != nil { + return nil, nil, xerrors.Errorf("invalid agent firewall session ID %q: %w", rawSessionID, parseErr) + } + + // Parse the sequence number. + n, err := strconv.ParseInt(rawSeqNumber, 10, 32) + if err != nil { + return nil, nil, xerrors.Errorf("invalid agent firewall sequence number %q: %w", rawSeqNumber, err) + } + if n < 0 { + return nil, nil, xerrors.Errorf("invalid agent firewall sequence number %q: must be non-negative", rawSeqNumber) + } + + n32 := int32(n) + return &rawSessionID, &n32, nil +} diff --git a/aibridge/bridge_internal_test.go b/aibridge/bridge_internal_test.go new file mode 100644 index 0000000000..561f758de1 --- /dev/null +++ b/aibridge/bridge_internal_test.go @@ -0,0 +1,131 @@ +package aibridge + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + agplaibridge "github.com/coder/coder/v2/coderd/aibridge" +) + +func TestExtractAgentFirewallHeaders(t *testing.T) { + t.Parallel() + + const validSessionID = "e5f6a7b8-1234-5678-9abc-def012345678" + + ptr := func(s string) *string { return &s } + + cases := []struct { + name string + // sessionID and seqNumber set the corresponding headers when + // non-nil. A nil value leaves the header unset. + sessionID *string + seqNumber *string + + wantErr bool + errContains string + wantSession *string + wantSeq *int32 + }{ + { + name: "both headers present", + sessionID: ptr(validSessionID), + seqNumber: ptr("42"), + wantSession: ptr(validSessionID), + wantSeq: int32Ptr(42), + }, + { + name: "no headers present", + }, + { + name: "only session ID returns error", + sessionID: ptr(validSessionID), + wantErr: true, + errContains: "without sequence number", + }, + { + name: "only sequence number returns error", + seqNumber: ptr("7"), + wantErr: true, + errContains: "without session ID", + }, + { + name: "sequence number zero", + sessionID: ptr(validSessionID), + seqNumber: ptr("0"), + wantSession: ptr(validSessionID), + wantSeq: int32Ptr(0), + }, + { + name: "invalid session ID returns error", + sessionID: ptr("not-a-uuid"), + seqNumber: ptr("42"), + wantErr: true, + errContains: "invalid agent firewall session ID", + }, + { + name: "invalid sequence number returns error", + sessionID: ptr(validSessionID), + seqNumber: ptr("not-a-number"), + wantErr: true, + errContains: "invalid agent firewall sequence number", + }, + { + name: "negative sequence number returns error", + sessionID: ptr(validSessionID), + seqNumber: ptr("-1"), + wantErr: true, + errContains: "must be non-negative", + }, + { + name: "sequence number exceeding int32 range returns error", + sessionID: ptr(validSessionID), + seqNumber: ptr("2147483648"), // max int32 + 1 + wantErr: true, + errContains: "invalid agent firewall sequence number", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, "/", nil) + require.NoError(t, err) + if tc.sessionID != nil { + req.Header.Set(agplaibridge.HeaderAgentFirewallSessionID, *tc.sessionID) + } + if tc.seqNumber != nil { + req.Header.Set(agplaibridge.HeaderAgentFirewallSequenceNumber, *tc.seqNumber) + } + + sessionID, seqNumber, extractErr := extractAgentFirewallHeaders(req) + + if tc.wantErr { + require.Error(t, extractErr) + assert.Contains(t, extractErr.Error(), tc.errContains) + assert.Nil(t, sessionID) + assert.Nil(t, seqNumber) + return + } + + require.NoError(t, extractErr) + if tc.wantSession == nil { + assert.Nil(t, sessionID) + } else { + require.NotNil(t, sessionID) + assert.Equal(t, *tc.wantSession, *sessionID) + } + if tc.wantSeq == nil { + assert.Nil(t, seqNumber) + } else { + require.NotNil(t, seqNumber) + assert.Equal(t, *tc.wantSeq, *seqNumber) + } + }) + } +} + +func int32Ptr(n int32) *int32 { return &n } diff --git a/aibridge/intercept/client_headers.go b/aibridge/intercept/client_headers.go index f7af0590e8..60a49523f7 100644 --- a/aibridge/intercept/client_headers.go +++ b/aibridge/intercept/client_headers.go @@ -47,6 +47,15 @@ var proxyHeaders = []string{ "Forwarded", } +// agentFirewallHeaders carry Agent Firewall correlation data used by +// AI Gateway for session correlation. AI Gateway records the values +// from the incoming request and strips the headers here so they are +// never forwarded to upstream LLM providers. +var agentFirewallHeaders = []string{ + "X-Coder-Agent-Firewall-Session-Id", + "X-Coder-Agent-Firewall-Sequence-Number", +} + // PrepareClientHeaders returns a copy of the client headers with hop-by-hop, // transport, auth, and proxy headers removed. func PrepareClientHeaders(clientHeaders http.Header) http.Header { @@ -63,6 +72,9 @@ func PrepareClientHeaders(clientHeaders http.Header) http.Header { for _, h := range proxyHeaders { prepared.Del(h) } + for _, h := range agentFirewallHeaders { + prepared.Del(h) + } return prepared } diff --git a/aibridge/intercept/client_headers_test.go b/aibridge/intercept/client_headers_test.go index d16d175d1d..c4ac270437 100644 --- a/aibridge/intercept/client_headers_test.go +++ b/aibridge/intercept/client_headers_test.go @@ -121,6 +121,22 @@ func TestPrepareClientHeaders(t *testing.T) { require.Equal(t, originalCopy, input) }) + + t.Run("agent firewall headers are removed", func(t *testing.T) { + t.Parallel() + + input := http.Header{ + "X-Coder-Agent-Firewall-Session-Id": {"e5f6a7b8-1234-5678-9abc-def012345678"}, + "X-Coder-Agent-Firewall-Sequence-Number": {"42"}, + "X-Custom": {"preserved"}, + } + + result := intercept.PrepareClientHeaders(input) + + assert.Empty(t, result.Get("X-Coder-Agent-Firewall-Session-Id")) + assert.Empty(t, result.Get("X-Coder-Agent-Firewall-Sequence-Number")) + assert.Equal(t, "preserved", result.Get("X-Custom")) + }) } func TestBuildUpstreamHeaders(t *testing.T) { diff --git a/aibridge/internal/integrationtest/agent_firewall_internal_test.go b/aibridge/internal/integrationtest/agent_firewall_internal_test.go new file mode 100644 index 0000000000..c63e04deed --- /dev/null +++ b/aibridge/internal/integrationtest/agent_firewall_internal_test.go @@ -0,0 +1,105 @@ +package integrationtest + +import ( + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tidwall/sjson" + + "github.com/coder/coder/v2/aibridge/config" + "github.com/coder/coder/v2/aibridge/fixtures" + "github.com/coder/coder/v2/aibridge/internal/testutil" + agplaibridge "github.com/coder/coder/v2/coderd/aibridge" +) + +func TestAgentFirewallHeaders(t *testing.T) { + t.Parallel() + + t.Run("valid headers are recorded and stripped", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong) + t.Cleanup(cancel) + + fix := fixtures.Parse(t, fixtures.OaiChatSimple) + upstream := testutil.NewMockUpstream(ctx, t, testutil.NewFixtureResponse(fix)) + + bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, withProvider(config.ProviderOpenAI)) + + reqBody, err := sjson.SetBytes(fix.Request(), "stream", false) + require.NoError(t, err) + + agentFirewallSessionID := "e5f6a7b8-1234-5678-9abc-def012345678" + agentFirewallSequenceNumber := int32(42) + resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIChatCompletions, reqBody, http.Header{ + agplaibridge.HeaderAgentFirewallSessionID: {agentFirewallSessionID}, + agplaibridge.HeaderAgentFirewallSequenceNumber: {fmt.Sprintf("%d", agentFirewallSequenceNumber)}, + }) + require.NoError(t, err) + defer resp.Body.Close() + require.Equal(t, http.StatusOK, resp.StatusCode) + + // Read the full response body so that AI Gateway can record the interception. + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + + // Verify firewall headers were recorded in the interception. + interceptions := bridgeServer.Recorder.RecordedInterceptions() + require.Len(t, interceptions, 1) + require.NotNil(t, interceptions[0].AgentFirewallSessionID) + assert.Equal(t, agentFirewallSessionID, *interceptions[0].AgentFirewallSessionID) + require.NotNil(t, interceptions[0].AgentFirewallSequenceNumber) + assert.Equal(t, agentFirewallSequenceNumber, *interceptions[0].AgentFirewallSequenceNumber) + + // Verify firewall headers were stripped before reaching upstream. + received := upstream.ReceivedRequests() + require.Len(t, received, 1) + assert.Empty(t, received[0].Header.Get(agplaibridge.HeaderAgentFirewallSessionID)) + assert.Empty(t, received[0].Header.Get(agplaibridge.HeaderAgentFirewallSequenceNumber)) + + bridgeServer.Recorder.VerifyAllInterceptionsEnded(t) + }) + + t.Run("invalid headers are rejected before reaching upstream", func(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(t.Context(), testutil.WaitLong) + t.Cleanup(cancel) + + fix := fixtures.Parse(t, fixtures.OaiChatSimple) + // Use a plain upstream that fails the test if called, since the + // request must be rejected before reaching the provider. + upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Error("upstream should not have been called") + })) + t.Cleanup(upstream.Close) + + bridgeServer := newBridgeTestServer(ctx, t, upstream.URL, withProvider(config.ProviderOpenAI)) + + reqBody, err := sjson.SetBytes(fix.Request(), "stream", false) + require.NoError(t, err) + + // Session ID without a sequence number is malformed; the rest of + // the validation matrix itself is covered by the unit tests for + // extractAgentFirewallHeaders. + resp, err := bridgeServer.makeRequest(t, http.MethodPost, pathOpenAIChatCompletions, reqBody, http.Header{ + agplaibridge.HeaderAgentFirewallSessionID: {"e5f6a7b8-1234-5678-9abc-def012345678"}, + }) + require.NoError(t, err) + defer resp.Body.Close() + assert.Equal(t, http.StatusBadRequest, resp.StatusCode) + + _, err = io.ReadAll(resp.Body) + require.NoError(t, err) + + // The request must fail closed: no interception recorded. + interceptions := bridgeServer.Recorder.RecordedInterceptions() + assert.Empty(t, interceptions) + }) +} diff --git a/aibridge/recorder/types.go b/aibridge/recorder/types.go index faa5713900..09dead839b 100644 --- a/aibridge/recorder/types.go +++ b/aibridge/recorder/types.go @@ -39,6 +39,14 @@ type InterceptionRecord struct { Client string UserAgent string CorrelatingToolCallID *string + // AgentFirewallSessionID is the UUID of the Agent Firewall session + // that produced this request. Nil when the request did not pass + // through Agent Firewall. + AgentFirewallSessionID *string + // AgentFirewallSequenceNumber is the monotonically increasing + // sequence number assigned by Agent Firewall. Nil when the request + // did not pass through Agent Firewall. + AgentFirewallSequenceNumber *int32 // CredentialKind is always set: either BYOK or centralized. CredentialKind string // CredentialHint is only set for BYOK, where the key is known diff --git a/coderd/aibridge/aibridge.go b/coderd/aibridge/aibridge.go index 6fff76e368..a4307257b8 100644 --- a/coderd/aibridge/aibridge.go +++ b/coderd/aibridge/aibridge.go @@ -20,6 +20,16 @@ const HeaderCoderToken = "X-Coder-AI-Governance-Token" //nolint:gosec // This is // request forwarded to aibridged for cross-service log correlation. const HeaderCoderRequestID = "X-Coder-AI-Governance-Request-Id" +// HeaderAgentFirewallSessionID is injected by Agent Firewall on requests +// routed through it. It carries the firewall session UUID so that AI +// Gateway can correlate interceptions with firewall audit events. +const HeaderAgentFirewallSessionID = "X-Coder-Agent-Firewall-Session-Id" + +// HeaderAgentFirewallSequenceNumber is injected alongside the session ID +// by Agent Firewall. It carries a monotonically increasing sequence +// number that orders network requests within a single firewall session. +const HeaderAgentFirewallSequenceNumber = "X-Coder-Agent-Firewall-Sequence-Number" + // Copilot provider. const ( ProviderCopilotBusiness = "copilot-business" diff --git a/coderd/aibridged/translator.go b/coderd/aibridged/translator.go index 6d251df0fe..fa3553f954 100644 --- a/coderd/aibridged/translator.go +++ b/coderd/aibridged/translator.go @@ -25,20 +25,22 @@ type recorderTranslation struct { func (t *recorderTranslation) RecordInterception(ctx context.Context, req *aibridge.InterceptionRecord) error { _, err := t.client.RecordInterception(ctx, &proto.RecordInterceptionRequest{ - Id: req.ID, - ApiKeyId: t.apiKeyID, - InitiatorId: req.InitiatorID, - Provider: req.Provider, - ProviderName: req.ProviderName, - Model: req.Model, - UserAgent: req.UserAgent, - Client: req.Client, - ClientSessionId: req.ClientSessionID, - Metadata: marshalForProto(req.Metadata), - StartedAt: timestamppb.New(req.StartedAt), - CorrelatingToolCallId: req.CorrelatingToolCallID, - CredentialKind: req.CredentialKind, - CredentialHint: req.CredentialHint, + Id: req.ID, + ApiKeyId: t.apiKeyID, + InitiatorId: req.InitiatorID, + Provider: req.Provider, + ProviderName: req.ProviderName, + Model: req.Model, + UserAgent: req.UserAgent, + Client: req.Client, + ClientSessionId: req.ClientSessionID, + Metadata: marshalForProto(req.Metadata), + StartedAt: timestamppb.New(req.StartedAt), + CorrelatingToolCallId: req.CorrelatingToolCallID, + CredentialKind: req.CredentialKind, + CredentialHint: req.CredentialHint, + AgentFirewallSessionId: req.AgentFirewallSessionID, + AgentFirewallSequenceNumber: req.AgentFirewallSequenceNumber, }) return err }