Skip to content

Commit a13bc7d

Browse files
AchoArnoldCopilot
AchoArnold
and
Copilot
committed
feat(mcp): add API key tools
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9ff1e38a-b018-4cf7-a5e9-5044a2efd03c
1 parent 86afa17 commit a13bc7d

8 files changed

Lines changed: 1147 additions & 21 deletions

File tree

‎mcp/internal/auth/middleware.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,22 @@ func PrincipalFromContext(ctx context.Context) (Principal, bool) {
8787
return principal, ok
8888
}
8989

90+
// ClientIDFromContext returns the OAuth client ID carried by the MCP access
91+
// token that mcpauth.RequireBearerToken (configured with a Verifier's
92+
// VerifyMCPToken) has already validated for the current request, or false
93+
// if ctx carries no verified token. Tools use this to bind sensitive
94+
// confirmation state (see the rotate_user_api_key tool) to the exact OAuth
95+
// client that requested the operation, not just the authenticated user.
96+
func ClientIDFromContext(ctx context.Context) (string, bool) {
97+
info := mcpauth.TokenInfoFromContext(ctx)
98+
if info == nil {
99+
return "", false
100+
}
101+
102+
clientID, ok := info.Extra[tokenInfoClientIDKey].(string)
103+
return clientID, ok
104+
}
105+
90106
// RequireScope returns the Principal carried by ctx's already-validated MCP
91107
// access token, or an error if ctx carries no verified token or the token's
92108
// scopes do not include scope. It never calls the httpSMS API and never

‎mcp/internal/auth/middleware_test.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,12 @@ func newMiddlewareTestServer(t *testing.T, keys *auth.KeySet, requiredScopes []s
3636
principal, ok := auth.PrincipalFromContext(r.Context())
3737
require.True(t, ok, "auth.PrincipalFromContext must find the principal the middleware stored")
3838

39+
clientID, ok := auth.ClientIDFromContext(r.Context())
40+
require.True(t, ok, "auth.ClientIDFromContext must find the client ID the middleware stored")
41+
3942
w.Header().Set("X-Test-User-ID", info.UserID)
4043
w.Header().Set("X-Test-Principal-Email", principal.Email)
44+
w.Header().Set("X-Test-Client-ID", clientID)
4145
w.WriteHeader(http.StatusOK)
4246
})
4347

@@ -145,6 +149,7 @@ func TestRequireBearerTokenAcceptsValidTokenAndStoresTokenInfo(t *testing.T) {
145149
require.Equal(t, http.StatusOK, resp.StatusCode)
146150
assert.Equal(t, testFirebaseUserID, resp.Header.Get("X-Test-User-ID"))
147151
assert.Equal(t, testUserEmail, resp.Header.Get("X-Test-Principal-Email"))
152+
assert.Equal(t, "client", resp.Header.Get("X-Test-Client-ID"))
148153
}
149154

150155
func doBearerRequest(t *testing.T, url string, token string) *http.Response {
@@ -164,6 +169,11 @@ func TestPrincipalFromContextReturnsFalseWithoutToken(t *testing.T) {
164169
assert.False(t, ok)
165170
}
166171

172+
func TestClientIDFromContextReturnsFalseWithoutToken(t *testing.T) {
173+
_, ok := auth.ClientIDFromContext(t.Context())
174+
assert.False(t, ok)
175+
}
176+
167177
func TestRequireScopeReturnsErrorWithoutToken(t *testing.T) {
168178
_, err := auth.RequireScope(t.Context(), auth.ScopePhonesRead)
169179
require.Error(t, err)

‎mcp/internal/oauth/store.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -326,6 +326,21 @@ func (s *RedisStore) ConsumeConfirmation(ctx context.Context, handle string) (Co
326326
return record, err
327327
}
328328

329+
// confirmationHandleBytes is the amount of crypto/rand entropy (see
330+
// newRandomToken) encoded into a rotation confirmation handle.
331+
const confirmationHandleBytes = 32
332+
333+
// NewConfirmationHandle returns a new cryptographically random, one-time
334+
// confirmation handle for the primary-API-key-rotation confirmation flow
335+
// (see Confirmation, PutConfirmation, and ConsumeConfirmation). Callers
336+
// store it with PutConfirmation and hand it to the client -- as
337+
// mcp.CallToolResult.RequestState for MRTR-capable clients, or as plain
338+
// tool output text for legacy clients that must echo it back explicitly --
339+
// and later redeem it exactly once with ConsumeConfirmation.
340+
func NewConfirmationHandle() (string, error) {
341+
return newRandomToken(confirmationHandleBytes)
342+
}
343+
329344
// hashedKey returns the namespaced Redis key for publicValue under prefix:
330345
// prefix followed by the hex-encoded SHA-256 hash of publicValue. The raw
331346
// value is never used as key material.

‎mcp/internal/oauth/store_test.go‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,31 @@ func TestRedisStoreGetDynamicClientNotFound(t *testing.T) {
267267
require.ErrorIs(t, err, oauth.ErrNotFound)
268268
}
269269

270+
// TestNewConfirmationHandleIsRandomAndURLSafe asserts NewConfirmationHandle
271+
// returns a fresh, non-empty, URL-safe value on every call (never a fixed
272+
// or predictable value), and that the handle it returns actually works
273+
// end-to-end with PutConfirmation/ConsumeConfirmation.
274+
func TestNewConfirmationHandleIsRandomAndURLSafe(t *testing.T) {
275+
first, err := oauth.NewConfirmationHandle()
276+
require.NoError(t, err)
277+
assert.NotEmpty(t, first)
278+
assert.NotRegexp(t, `[^A-Za-z0-9_-]`, first, "confirmation handle must be URL-safe base64")
279+
280+
second, err := oauth.NewConfirmationHandle()
281+
require.NoError(t, err)
282+
assert.NotEqual(t, first, second, "two generated handles must never collide")
283+
284+
store, _ := newTestStore(t)
285+
ctx := context.Background()
286+
287+
confirmation := oauth.Confirmation{Handle: first, UserID: "firebase-uid", ClientID: "client-id", Operation: "rotate_user_api_key"}
288+
require.NoError(t, store.PutConfirmation(ctx, confirmation, time.Minute))
289+
290+
got, err := store.ConsumeConfirmation(ctx, first)
291+
require.NoError(t, err)
292+
assert.Equal(t, confirmation.UserID, got.UserID)
293+
}
294+
270295
func TestRedisStoreConsumeConfirmationIsOneTimeUse(t *testing.T) {
271296
store, _ := newTestStore(t)
272297
ctx := context.Background()

0 commit comments

Comments
 (0)