feat(coderd): accept delegated API key ID from in-process aibridge callers (#25625)

Allows an `api_key_id` to be passed from a trusted in-memory transport (currently: `chatd`) to `aibridged` for use in authenticating LLM requests. This value can _only_ be passed via context, and all users of the in-memory transport _must_ provide it. It can be used in conjunction with BYOK headers. --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-09-24 15:04:27 +08:00 · 2026-05-25 11:08:07 +02:00
parent 814386dda7
commit eddd4a8c2f
10 changed files with 660 additions and 113 deletions
@@ -38,12 +38,13 @@ var (
 	// matching.
 	// TODO: return these errors to the client in a more structured/comparable
 	//       way.
-	ErrInvalidKey  = xerrors.New("invalid key")
-	ErrUnknownKey  = xerrors.New("unknown key")
-	ErrExpired     = xerrors.New("expired")
-	ErrUnknownUser = xerrors.New("unknown user")
-	ErrDeletedUser = xerrors.New("deleted user")
-	ErrSystemUser  = xerrors.New("system user")
+	ErrInvalidKey    = xerrors.New("invalid key")
+	ErrUnknownKey    = xerrors.New("unknown key")
+	ErrExpired       = xerrors.New("expired")
+	ErrUnknownUser   = xerrors.New("unknown user")
+	ErrDeletedUser   = xerrors.New("deleted user")
+	ErrSystemUser    = xerrors.New("system user")
+	ErrAmbiguousAuth = xerrors.New("both key and key_id set; exactly one required")

 	ErrNoExternalAuthLinkFound = xerrors.New("no external auth link found")
 )
@@ -550,6 +551,15 @@ externalAuthLoop:

 // IsAuthorized validates a given Coder API key and returns the user ID to which it belongs (if valid).
 //
+// SECURITY: when in.KeyId is set (the "delegated" path), this method trusts the
+// caller's claim of identity and skips the key-secret check. This is safe only
+// because the DRPCServer is reachable solely via the in-process
+// [aibridged.MemTransportPipe]; the handler itself cannot tell whether it was
+// invoked over the in-memory pipe or a network socket. If this RPC is ever
+// exposed over a network boundary, any caller who knows a valid 10-char key ID
+// (which is not secret) could authenticate as the key's owner without the
+// secret. Do not bind this DRPCServer to a network listener.
+//
 // NOTE: this should really be using the code from [httpmw.ExtractAPIKey]. That function not only validates the key
 // but handles many other cases like updating last used, expiry, etc. This code does not currently use it for
 // a few reasons:
@@ -565,10 +575,26 @@ func (s *Server) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest
 	//nolint:gocritic // AIBridged has specific authz rules.
 	ctx = dbauthz.AsAIBridged(ctx)

-	// Key matches expected format.
-	keyID, keySecret, err := httpmw.SplitAPIToken(in.GetKey())
-	if err != nil {
-		return nil, ErrInvalidKey
+	var (
+		keyID     string
+		keySecret string
+		// delegated requests skip the secret check: the caller never
+		// has the secret. Trust is established at the in-process
+		// transport boundary, not in this RPC.
+		delegated bool
+	)
+	switch {
+	case in.GetKey() != "" && in.GetKeyId() != "":
+		return nil, ErrAmbiguousAuth
+	case in.GetKeyId() != "":
+		keyID = in.GetKeyId()
+		delegated = true
+	default:
+		var err error
+		keyID, keySecret, err = httpmw.SplitAPIToken(in.GetKey())
+		if err != nil {
+			return nil, ErrInvalidKey
+		}
 	}

 	// Key exists.
@@ -584,8 +610,8 @@ func (s *Server) IsAuthorized(ctx context.Context, in *proto.IsAuthorizedRequest
 		return nil, ErrExpired
 	}

-	// Key secret matches.
-	if !apikey.ValidateHash(key.HashedSecret, keySecret) {
+	// Key secret matches (skipped for delegated callers).
+	if !delegated && !apikey.ValidateHash(key.HashedSecret, keySecret) {
 		return nil, ErrInvalidKey
 	}

@@ -199,6 +199,148 @@ func TestAuthorization(t *testing.T) {
 	}
 }

+// When IsAuthorizedRequest carries KeyId instead of Key, the server skips
+// the secret check and validates only that the key exists, is unexpired, and
+// belongs to a non-deleted non-system user. This is the path used by
+// in-process delegated callers (e.g., chatd) that hold only the key ID.
+func TestAuthorization_Delegated(t *testing.T) {
+	t.Parallel()
+
+	cases := []struct {
+		name        string
+		mocksFn     func(db *dbmock.MockStore, apiKey database.APIKey, user database.User)
+		bothFields  bool
+		expectedErr error
+	}{
+		{
+			name: "valid",
+			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
+				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
+				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
+			},
+		},
+		{
+			name:        "unknown key",
+			expectedErr: aibridgedserver.ErrUnknownKey,
+			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, _ database.User) {
+				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(database.APIKey{}, sql.ErrNoRows)
+			},
+		},
+		{
+			name:        "expired",
+			expectedErr: aibridgedserver.ErrExpired,
+			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, _ database.User) {
+				apiKey.ExpiresAt = dbtime.Now().Add(-time.Hour)
+				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
+			},
+		},
+		{
+			// Sending both Key and KeyId is an API misuse and must be
+			// rejected to avoid ambiguity about which path was taken.
+			name:        "both fields set",
+			bothFields:  true,
+			expectedErr: aibridgedserver.ErrAmbiguousAuth,
+		},
+		{
+			// A bogus secret has no effect on the delegated path because
+			// the secret is never checked. This is the load-bearing
+			// security property: trust is established out-of-band, not in
+			// this RPC.
+			name: "secret hash mismatch is ignored",
+			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
+				apiKey.HashedSecret = []byte("not-the-real-hash")
+				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
+				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(user, nil)
+			},
+		},
+		{
+			// The delegated path must still reject keys whose owner has
+			// been deleted; trust at the transport boundary does not
+			// extend to bypassing user-status checks.
+			name:        "deleted user",
+			expectedErr: aibridgedserver.ErrDeletedUser,
+			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
+				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
+				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(database.User{ID: user.ID, Deleted: true}, nil)
+			},
+		},
+		{
+			// Likewise, a system user must never be authenticated through
+			// the delegated path.
+			name:        "system user",
+			expectedErr: aibridgedserver.ErrSystemUser,
+			mocksFn: func(db *dbmock.MockStore, apiKey database.APIKey, user database.User) {
+				db.EXPECT().GetAPIKeyByID(gomock.Any(), apiKey.ID).Times(1).Return(apiKey, nil)
+				db.EXPECT().GetUserByID(gomock.Any(), user.ID).Times(1).Return(database.User{ID: user.ID, IsSystem: true}, nil)
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctrl := gomock.NewController(t)
+			db := dbmock.NewMockStore(ctrl)
+			logger := testutil.Logger(t)
+
+			now := dbtime.Now()
+			user := database.User{
+				ID:         uuid.New(),
+				Email:      "test@coder.com",
+				Username:   "test",
+				Name:       "Test User",
+				CreatedAt:  now,
+				UpdatedAt:  now,
+				RBACRoles:  []string{},
+				LoginType:  database.LoginTypePassword,
+				Status:     database.UserStatusActive,
+				LastSeenAt: now,
+			}
+			keyID, _ := cryptorand.String(10)
+			_, keySecretHashed, _ := apikey.GenerateSecret(22)
+			apiKey := database.APIKey{
+				ID:              keyID,
+				LifetimeSeconds: 86400,
+				HashedSecret:    keySecretHashed,
+				UserID:          user.ID,
+				LastUsed:        now,
+				ExpiresAt:       now.Add(time.Hour),
+				CreatedAt:       now,
+				UpdatedAt:       now,
+				LoginType:       database.LoginTypePassword,
+				Scopes:          []database.APIKeyScope{database.ApiKeyScopeCoderAll},
+			}
+
+			if tc.mocksFn != nil {
+				tc.mocksFn(db, apiKey, user)
+			}
+
+			srv, err := aibridgedserver.NewServer(t.Context(), db, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{})
+			require.NoError(t, err)
+			require.NotNil(t, srv)
+
+			req := &proto.IsAuthorizedRequest{KeyId: keyID}
+			if tc.bothFields {
+				req.Key = "anything-anything"
+			}
+
+			resp, err := srv.IsAuthorized(t.Context(), req)
+			if tc.expectedErr != nil {
+				require.Error(t, err)
+				require.ErrorIs(t, err, tc.expectedErr)
+				return
+			}
+			require.NoError(t, err)
+			require.Equal(t, &proto.IsAuthorizedResponse{
+				OwnerId:  user.ID.String(),
+				ApiKeyId: keyID,
+				Username: user.Username,
+			}, resp)
+		})
+	}
+}
+
 func TestGetMCPServerConfigs(t *testing.T) {
 	t.Parallel()