From 441ef79bd0f3c0d3155ad27d70a542ba5e11aa88 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Wed, 17 Jun 2026 16:37:41 +0000 Subject: [PATCH] feat(coderd/httpmw): add AI Gateway key auth middleware and proto version Add ExtractAIGatewayKeyAuthenticated, which authenticates standalone AI Gateway replicas via the X-AI-Governance-Gateway-Key header, mirroring the external provisioner daemon key flow. Also add the aibridged proto CurrentVersion (1.0) for API compatibility negotiation, and the X-AI-Governance-Gateway-Key header constant. Part of AIGOV-308. Generated with Coder Agents. --- coderd/aibridged/proto/version.go | 19 +++++ coderd/httpmw/aigatewaykey.go | 97 ++++++++++++++++++++++++ coderd/httpmw/aigatewaykey_test.go | 117 +++++++++++++++++++++++++++++ codersdk/client.go | 4 + site/src/api/typesGenerated.ts | 7 ++ 5 files changed, 244 insertions(+) create mode 100644 coderd/aibridged/proto/version.go create mode 100644 coderd/httpmw/aigatewaykey.go create mode 100644 coderd/httpmw/aigatewaykey_test.go diff --git a/coderd/aibridged/proto/version.go b/coderd/aibridged/proto/version.go new file mode 100644 index 0000000000000..914189515d06b --- /dev/null +++ b/coderd/aibridged/proto/version.go @@ -0,0 +1,19 @@ +package proto + +import "github.com/coder/coder/v2/apiversion" + +// Version history: +// +// API v1.0: +// - Initial version. Serves the Recorder, MCPConfigurator, and Authorizer +// services to embedded and standalone AI Gateway daemons. +const ( + CurrentMajor = 1 + CurrentMinor = 0 +) + +// CurrentVersion is the current aibridged API version. +// Breaking changes to the aibridged API **MUST** increment CurrentMajor above. +// Non-breaking changes to the aibridged API **MUST** increment CurrentMinor +// above. +var CurrentVersion = apiversion.New(CurrentMajor, CurrentMinor) diff --git a/coderd/httpmw/aigatewaykey.go b/coderd/httpmw/aigatewaykey.go new file mode 100644 index 0000000000000..7184ddf404188 --- /dev/null +++ b/coderd/httpmw/aigatewaykey.go @@ -0,0 +1,97 @@ +package httpmw + +import ( + "context" + "crypto/subtle" + "net/http" + + "github.com/coder/coder/v2/coderd/apikey" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/httpapi" + "github.com/coder/coder/v2/codersdk" +) + +type aiGatewayKeyContextKey struct{} + +// AIGatewayKeyAuthOptional returns the AI Gateway key that authenticated the +// request, if any. The key is used by the /serve handler to record liveness +// against the authenticating key. +func AIGatewayKeyAuthOptional(r *http.Request) (database.AIGatewayKey, bool) { + key, ok := r.Context().Value(aiGatewayKeyContextKey{}).(database.AIGatewayKey) + return key, ok +} + +// ExtractAIGatewayKeyConfig configures ExtractAIGatewayKeyAuthenticated. +type ExtractAIGatewayKeyConfig struct { + DB database.Store + // Optional, when true, allows the request to proceed unauthenticated. The + // next handler can detect authentication via AIGatewayKeyAuthOptional. + Optional bool +} + +// ExtractAIGatewayKeyAuthenticated authenticates a request as a standalone AI +// Gateway replica using the X-AI-Governance-Gateway-Key header. The header +// value is hashed and compared against the ai_gateway_keys table, mirroring the +// external provisioner daemon key flow in +// ExtractProvisionerDaemonAuthenticated. +// +// One key may authenticate many replicas at once; keys are not unique per +// connection. A deleted key fails the next reconnect because the row no longer +// exists, but does not disconnect sessions already established. +func ExtractAIGatewayKeyAuthenticated(opts ExtractAIGatewayKeyConfig) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ctx := r.Context() + + handleOptional := func(code int, response codersdk.Response) { + if opts.Optional { + next.ServeHTTP(w, r) + return + } + httpapi.Write(ctx, w, code, response) + } + + key := r.Header.Get(codersdk.AIGatewayKeyHeader) + if key == "" { + handleOptional(http.StatusUnauthorized, codersdk.Response{ + Message: "AI Gateway key required.", + }) + return + } + + hashedKey := apikey.HashSecret(key) + // nolint:gocritic // System must look up the AI Gateway key to authenticate the request. + gatewayKey, err := opts.DB.GetAIGatewayKeyByHashedSecret(dbauthz.AsSystemRestricted(ctx), hashedKey) + if err != nil { + if httpapi.Is404Error(err) { + handleOptional(http.StatusUnauthorized, codersdk.Response{ + Message: "AI Gateway key invalid.", + }) + return + } + handleOptional(http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to look up AI Gateway key.", + Detail: err.Error(), + }) + return + } + + // Defense in depth: the lookup already matches on hashed_secret, but + // confirm equality in constant time to avoid relying solely on the + // query. + if subtle.ConstantTimeCompare(gatewayKey.HashedSecret, hashedKey) != 1 { + handleOptional(http.StatusUnauthorized, codersdk.Response{ + Message: "AI Gateway key invalid.", + }) + return + } + + ctx = context.WithValue(ctx, aiGatewayKeyContextKey{}, gatewayKey) + // nolint:gocritic // Authenticating as an AI Gateway replica, which + // acts as the AI Bridge daemon. + ctx = dbauthz.AsAIBridged(ctx) + next.ServeHTTP(w, r.WithContext(ctx)) + }) + } +} diff --git a/coderd/httpmw/aigatewaykey_test.go b/coderd/httpmw/aigatewaykey_test.go new file mode 100644 index 0000000000000..46788392efdb6 --- /dev/null +++ b/coderd/httpmw/aigatewaykey_test.go @@ -0,0 +1,117 @@ +package httpmw_test + +import ( + "database/sql" + "net/http" + "net/http/httptest" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/apikey" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbmock" + "github.com/coder/coder/v2/coderd/httpmw" + "github.com/coder/coder/v2/codersdk" +) + +func TestExtractAIGatewayKeyAuthenticated(t *testing.T) { + t.Parallel() + + const secret = "this-is-a-test-gateway-key" + + t.Run("MissingHeader", func(t *testing.T) { + t.Parallel() + db := dbmock.NewMockStore(gomock.NewController(t)) + rw := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodGet, "/", nil) + newGatewayHandler(db, false).ServeHTTP(rw, r) + require.Equal(t, http.StatusUnauthorized, rw.Code) + }) + + t.Run("InvalidKey", func(t *testing.T) { + t.Parallel() + db := dbmock.NewMockStore(gomock.NewController(t)) + db.EXPECT().GetAIGatewayKeyByHashedSecret(gomock.Any(), apikey.HashSecret(secret)). + Return(database.AIGatewayKey{}, sql.ErrNoRows) + rw := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodGet, "/", nil) + r.Header.Set(codersdk.AIGatewayKeyHeader, secret) + newGatewayHandler(db, false).ServeHTTP(rw, r) + require.Equal(t, http.StatusUnauthorized, rw.Code) + }) + + t.Run("LookupError", func(t *testing.T) { + t.Parallel() + db := dbmock.NewMockStore(gomock.NewController(t)) + db.EXPECT().GetAIGatewayKeyByHashedSecret(gomock.Any(), apikey.HashSecret(secret)). + Return(database.AIGatewayKey{}, xerrors.New("boom")) + rw := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodGet, "/", nil) + r.Header.Set(codersdk.AIGatewayKeyHeader, secret) + newGatewayHandler(db, false).ServeHTTP(rw, r) + require.Equal(t, http.StatusInternalServerError, rw.Code) + }) + + t.Run("Success", func(t *testing.T) { + t.Parallel() + db := dbmock.NewMockStore(gomock.NewController(t)) + key := database.AIGatewayKey{ + ID: uuid.New(), + Name: "test-key", + HashedSecret: apikey.HashSecret(secret), + } + db.EXPECT().GetAIGatewayKeyByHashedSecret(gomock.Any(), apikey.HashSecret(secret)). + Return(key, nil) + + var ( + gotKey database.AIGatewayKey + gotOK bool + ) + handler := httpmw.ExtractAIGatewayKeyAuthenticated(httpmw.ExtractAIGatewayKeyConfig{DB: db})( + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotKey, gotOK = httpmw.AIGatewayKeyAuthOptional(r) + w.WriteHeader(http.StatusOK) + }), + ) + + rw := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodGet, "/", nil) + r.Header.Set(codersdk.AIGatewayKeyHeader, secret) + handler.ServeHTTP(rw, r) + + require.Equal(t, http.StatusOK, rw.Code) + require.True(t, gotOK) + require.Equal(t, key.ID, gotKey.ID) + }) + + t.Run("OptionalPassesThrough", func(t *testing.T) { + t.Parallel() + db := dbmock.NewMockStore(gomock.NewController(t)) + var called bool + handler := httpmw.ExtractAIGatewayKeyAuthenticated(httpmw.ExtractAIGatewayKeyConfig{DB: db, Optional: true})( + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, ok := httpmw.AIGatewayKeyAuthOptional(r) + require.False(t, ok) + called = true + w.WriteHeader(http.StatusOK) + }), + ) + rw := httptest.NewRecorder() + r := httptest.NewRequest(http.MethodGet, "/", nil) + handler.ServeHTTP(rw, r) + require.Equal(t, http.StatusOK, rw.Code) + require.True(t, called) + }) +} + +func newGatewayHandler(db database.Store, optional bool) http.Handler { + return httpmw.ExtractAIGatewayKeyAuthenticated(httpmw.ExtractAIGatewayKeyConfig{DB: db, Optional: optional})( + http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + }), + ) +} diff --git a/codersdk/client.go b/codersdk/client.go index b01b5e4fb3a8f..0747edfaa3e24 100644 --- a/codersdk/client.go +++ b/codersdk/client.go @@ -96,6 +96,10 @@ const ( // ProvisionerDaemonKey contains the authentication key for an external provisioner daemon ProvisionerDaemonKey = "Coder-Provisioner-Daemon-Key" + // AIGatewayKeyHeader contains the authentication key for a standalone AI + // Gateway replica connecting to coderd. + AIGatewayKeyHeader = "X-AI-Governance-Gateway-Key" + // BuildVersionHeader contains build information of Coder. BuildVersionHeader = "X-Coder-Build-Version" diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 5ef2e86e3f519..ec8a74176569e 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -253,6 +253,13 @@ export interface AIGatewayKey { readonly last_used_at?: string; } +// From codersdk/client.go +/** + * AIGatewayKeyHeader contains the authentication key for a standalone AI + * Gateway replica connecting to coderd. + */ +export const AIGatewayKeyHeader = "X-AI-Governance-Gateway-Key"; + // From codersdk/aiproviders.go /** * AIProvider represents an AI provider configuration row as returned