diff --git a/framework/lrucache/lrucache.go b/framework/lrucache/lrucache.go index 75583cd7c33..e728c9077a5 100644 --- a/framework/lrucache/lrucache.go +++ b/framework/lrucache/lrucache.go @@ -37,6 +37,8 @@ import ( "container/list" "context" "fmt" + "strconv" + "strings" "sync" ) @@ -397,3 +399,56 @@ func (c *Cache[V]) Len() int { defer c.mu.Unlock() return c.order.Len() } + +// EncodeKey builds a collision-free composite cache key from parts, for +// callers whose key is a tuple of caller-controlled strings (e.g. +// (auth mode, identity, mcp client ID) — identity is frequently a +// caller-asserted string with no charset restriction). A naive +// separator-joined key lets one part's content forge a boundary — e.g. +// join("\x00", "a\x00b", "c") and join("\x00", "a", "b\x00c") would build +// the identical string, aliasing two distinct tuples onto one cache entry +// and skewing any eviction predicate that parses the key back apart. +// Length-prefixing each part makes that forgery impossible regardless of +// what bytes a part contains. Pair with DecodeKey to parse it back. +func EncodeKey(parts ...string) string { + var b strings.Builder + for _, p := range parts { + b.WriteString(strconv.Itoa(len(p))) + b.WriteByte(':') + b.WriteString(p) + } + return b.String() +} + +// DecodeKey is EncodeKey's inverse: it parses exactly n length-prefixed +// parts out of key, in the order EncodeKey wrote them. ok is false for a +// key that isn't in the length-prefixed form EncodeKey builds for n parts +// (impossible for keys a well-behaved caller produces with EncodeKey). +func DecodeKey(key string, n int) (parts []string, ok bool) { + if n < 0 { + return nil, false + } + rest := key + parts = make([]string, 0, n) + for range n { + i := strings.IndexByte(rest, ':') + if i < 0 { + return nil, false + } + lengthText := rest[:i] + length, err := strconv.Atoi(lengthText) + if err != nil || length < 0 || strconv.Itoa(length) != lengthText { + return nil, false + } + rest = rest[i+1:] + if length > len(rest) { + return nil, false + } + parts = append(parts, rest[:length]) + rest = rest[length:] + } + if rest != "" { + return nil, false + } + return parts, true +} diff --git a/framework/lrucache/lrucache_test.go b/framework/lrucache/lrucache_test.go index 2965b0bd62a..778dcfa652c 100644 --- a/framework/lrucache/lrucache_test.go +++ b/framework/lrucache/lrucache_test.go @@ -590,3 +590,54 @@ func TestNew_PanicsOnNonPositiveCapacity(t *testing.T) { assert.Panics(t, func() { New[string](0) }) assert.Panics(t, func() { New[string](-1) }) } + +func TestEncodeDecodeKey_RoundTrip(t *testing.T) { + parts := []string{"user", "alice", "client-1"} + key := EncodeKey(parts...) + got, ok := DecodeKey(key, len(parts)) + require.True(t, ok) + assert.Equal(t, parts, got) +} + +// TestEncodeDecodeKey_NoCollisionOnEmbeddedDelimiter pins the reason this +// codec exists over a plain separator join: a value that happens to +// contain the separator (or, here, digits and a colon shaped like a length +// prefix) must not let one tuple's key collide with a different tuple's. +func TestEncodeDecodeKey_NoCollisionOnEmbeddedDelimiter(t *testing.T) { + keyA := EncodeKey("user", "a\x00b", "c") + keyB := EncodeKey("user", "a", "b\x00c") + assert.NotEqual(t, keyA, keyB, "distinct tuples must not alias to the same key") + + gotA, ok := DecodeKey(keyA, 3) + require.True(t, ok) + assert.Equal(t, []string{"user", "a\x00b", "c"}, gotA) + + gotB, ok := DecodeKey(keyB, 3) + require.True(t, ok) + assert.Equal(t, []string{"user", "a", "b\x00c"}, gotB) +} + +func TestDecodeKey_RejectsMalformedInput(t *testing.T) { + tests := []struct { + name string + key string + n int + }{ + {"empty string", "", 3}, + {"no colon", "abc", 1}, + {"non-numeric length", "x:abc", 1}, + {"negative length", "-1:a", 1}, + {"length exceeds remaining bytes", "10:ab", 1}, + {"trailing garbage after all parts", "1:a1:btrailing", 2}, + {"too few parts", "4:user", 2}, + {"negative part count", "1:a", -1}, + {"non-canonical length prefix with leading plus", "+1:a", 1}, + {"non-canonical length prefix with leading zero", "01:a", 1}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + _, ok := DecodeKey(tc.key, tc.n) + assert.False(t, ok) + }) + } +} diff --git a/framework/mcp_headers/credentialcache.go b/framework/mcp_headers/credentialcache.go new file mode 100644 index 00000000000..23ab8d19c5c --- /dev/null +++ b/framework/mcp_headers/credentialcache.go @@ -0,0 +1,165 @@ +package mcp_headers + +import ( + "context" + + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/lrucache" +) + +// defaultCredentialCacheCapacity bounds the per-user header credential cache. +// Session-mode identities are caller-asserted strings, so the cache must stay +// bounded no matter what identities callers present; least-recently-used +// entries are dropped once the capacity is reached. +const defaultCredentialCacheCapacity = 4096 + +// cachedHeaderCredential is the value cached per (auth mode, identity, +// mcp client) binding: the row ID (for targeted eviction) and the parsed +// credential. Header credentials carry no expiry and have no refresh +// machinery, so there is no expiry-as-miss logic; explicit eviction is the +// only way a cached entry stops being served. The credential pointer is +// private to the cache: the provider hands callers a deep copy so a caller +// mutating the returned Headers map can never corrupt the cached value. +type cachedHeaderCredential struct { + credentialID string + credential *schemas.MCPHeadersUserCredential +} + +// headerCredentialCache adapts lrucache.Cache to per-user MCP header +// credential lookups: it owns the binding-key scheme, registers each entry +// under its credential row ID for targeted eviction (upsert, delete, +// revoke), and provides the scoped bulk evictions the credential lifecycle +// needs (by MCP client, virtual key, and user). No validator is installed: +// header credentials never expire, so explicit eviction is the only +// invalidation. +type headerCredentialCache struct { + cache *lrucache.Cache[cachedHeaderCredential] +} + +// headerCredentialCacheKey builds the cache key for a (mode, identity, +// mcp client) binding. identity is a caller-asserted string with no +// charset restriction, so this goes through lrucache.EncodeKey rather than +// a plain separator join — see its doc comment for why a naive join lets +// an identity value forge a component boundary and alias two distinct +// bindings onto one cache entry. Admin-mode bindings carry an empty +// identity component by design. +func headerCredentialCacheKey(mode schemas.MCPAuthMode, identity, mcpClientID string) string { + return lrucache.EncodeKey(string(mode), identity, mcpClientID) +} + +// splitHeaderCredentialCacheKey is headerCredentialCacheKey's inverse, for +// the scoped eviction predicates. +func splitHeaderCredentialCacheKey(key string) (mode, identity, clientID string, ok bool) { + parts, ok := lrucache.DecodeKey(key, 3) + if !ok { + return "", "", "", false + } + return parts[0], parts[1], parts[2], true +} + +func newHeaderCredentialCache(capacity int) *headerCredentialCache { + if capacity <= 0 { + capacity = defaultCredentialCacheCapacity + } + return &headerCredentialCache{cache: lrucache.New[cachedHeaderCredential](capacity)} +} + +// Get returns the cached value for key. Header credentials have no expiry, +// so a hit is served as-is until an eviction removes it. +func (c *headerCredentialCache) Get(key string) (cachedHeaderCredential, bool) { + if c == nil { + return cachedHeaderCredential{}, false + } + return c.cache.Get(key) +} + +// Fill runs fill for key with single-flight deduplication: concurrent +// callers for the same key wait for one leader and share its result and +// error, so a tool-call burst for one identity performs a single database +// read instead of a stampede. A successful result is cached under both the +// binding key and its credential row ID; errors are propagated but never +// cached. +func (c *headerCredentialCache) Fill(ctx context.Context, key string, fill func() (cachedHeaderCredential, error)) (cachedHeaderCredential, error) { + if c == nil { + return fill() + } + return c.cache.Fill(ctx, key, func() (cachedHeaderCredential, string, error) { + value, err := fill() + return value, value.credentialID, err + }) +} + +// Evict removes the entry for an exact binding key, if present. +func (c *headerCredentialCache) Evict(key string) { + if c == nil { + return + } + c.cache.Evict(key) +} + +// EvictByCredentialID removes the entry holding the given credential row ID, +// if any. +func (c *headerCredentialCache) EvictByCredentialID(credentialID string) { + if c == nil { + return + } + c.cache.EvictByIndex(credentialID) +} + +// EvictByMCPClient removes every cached entry bound to the given MCP client, +// across all auth modes and identities, including the admin-mode binding +// (whose identity component is empty). Used when a client-level change +// invalidates its credential rows as a set, such as a header schema change +// or client deletion. A linear sweep is fine here: these are rare admin +// operations and the cache is bounded. +func (c *headerCredentialCache) EvictByMCPClient(mcpClientID string) { + if c == nil || mcpClientID == "" { + return + } + c.cache.EvictWhere(func(key string) bool { + _, _, clientID, ok := splitHeaderCredentialCacheKey(key) + return ok && clientID == mcpClientID + }) +} + +// EvictByVirtualKey removes every cached vk-mode entry bound to the given +// virtual key, across all MCP clients. Used when a virtual key change +// orphans or deletes its credential rows as a set. +func (c *headerCredentialCache) EvictByVirtualKey(virtualKeyID string) { + if c == nil || virtualKeyID == "" { + return + } + c.cache.EvictWhere(func(key string) bool { + mode, identity, _, ok := splitHeaderCredentialCacheKey(key) + return ok && mode == string(schemas.MCPAuthModeVK) && identity == virtualKeyID + }) +} + +// EvictByUser removes every cached user-mode entry bound to the given user, +// across all MCP clients. Used when a user-level change orphans or deletes +// the user's credential rows as a set. +func (c *headerCredentialCache) EvictByUser(userID string) { + if c == nil || userID == "" { + return + } + c.cache.EvictWhere(func(key string) bool { + mode, identity, _, ok := splitHeaderCredentialCacheKey(key) + return ok && mode == string(schemas.MCPAuthModeUser) && identity == userID + }) +} + +// Flush drops every cached entry. +func (c *headerCredentialCache) Flush() { + if c == nil { + return + } + c.cache.Flush() +} + +// Len reports the number of cached entries. +func (c *headerCredentialCache) Len() int { + if c == nil { + return 0 + } + return c.cache.Len() +} diff --git a/framework/mcp_headers/credentialcache_test.go b/framework/mcp_headers/credentialcache_test.go new file mode 100644 index 00000000000..2cc461e636b --- /dev/null +++ b/framework/mcp_headers/credentialcache_test.go @@ -0,0 +1,630 @@ +package mcp_headers + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + bifrost "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/configstore/tables" +) + +func testCredential(id string, headers map[string]string) cachedHeaderCredential { + return cachedHeaderCredential{ + credentialID: id, + credential: &schemas.MCPHeadersUserCredential{ + ID: id, + Headers: headers, + Status: schemas.MCPHeadersUserCredentialStatusActive, + }, + } +} + +func fillWithCredential(v cachedHeaderCredential) func() (cachedHeaderCredential, error) { + return func() (cachedHeaderCredential, error) { return v, nil } +} + +func TestHeaderCredentialCache_HitAndMiss(t *testing.T) { + c := newHeaderCredentialCache(4) + + _, ok := c.Get("missing") + assert.False(t, ok, "empty cache must miss") + + v, err := c.Fill(context.Background(), "k1", fillWithCredential(testCredential("c1", map[string]string{"X-Api-Key": "v1"}))) + require.NoError(t, err) + assert.Equal(t, "c1", v.credentialID) + + got, ok := c.Get("k1") + require.True(t, ok, "filled key must hit") + assert.Equal(t, "c1", got.credentialID) + assert.Equal(t, "v1", got.credential.Headers["X-Api-Key"]) + + _, ok = c.Get("k2") + assert.False(t, ok, "unrelated key must miss") +} + +func TestHeaderCredentialCache_CapacityEviction(t *testing.T) { + c := newHeaderCredentialCache(2) + + for i := 1; i <= 3; i++ { + _, err := c.Fill(context.Background(), + fmt.Sprintf("k%d", i), + fillWithCredential(testCredential(fmt.Sprintf("c%d", i), nil)), + ) + require.NoError(t, err) + } + + assert.Equal(t, 2, c.Len(), "cache must stay at capacity") + _, ok := c.Get("k1") + assert.False(t, ok, "least recently used entry must be evicted") + _, ok = c.Get("k2") + assert.True(t, ok) + _, ok = c.Get("k3") + assert.True(t, ok) + + // The evicted entry's credential-ID index mapping must be gone: evicting + // by its row ID must not disturb the surviving entries. + c.EvictByCredentialID("c1") + assert.Equal(t, 2, c.Len()) +} + +func TestHeaderCredentialCache_EvictExactKey(t *testing.T) { + c := newHeaderCredentialCache(4) + _, err := c.Fill(context.Background(), "k1", fillWithCredential(testCredential("c1", nil))) + require.NoError(t, err) + + c.Evict("k1") + _, ok := c.Get("k1") + assert.False(t, ok) + assert.Equal(t, 0, c.Len()) + + // Evicting an absent key is a no-op, not a panic. + c.Evict("never-set") +} + +func TestHeaderCredentialCache_EvictByCredentialID(t *testing.T) { + c := newHeaderCredentialCache(4) + _, err := c.Fill(context.Background(), "k1", fillWithCredential(testCredential("c1", nil))) + require.NoError(t, err) + _, err = c.Fill(context.Background(), "k2", fillWithCredential(testCredential("c2", nil))) + require.NoError(t, err) + + c.EvictByCredentialID("c1") + _, ok := c.Get("k1") + assert.False(t, ok, "entry holding the evicted credential ID must be gone") + _, ok = c.Get("k2") + assert.True(t, ok, "unrelated entry must survive") + + // Unknown credential IDs are a no-op. + c.EvictByCredentialID("unknown") + assert.Equal(t, 1, c.Len()) +} + +func TestHeaderCredentialCache_EvictByMCPClient(t *testing.T) { + c := newHeaderCredentialCache(8) + keyA1 := headerCredentialCacheKey(schemas.MCPAuthModeUser, "u1", "client-a") + keyA2 := headerCredentialCacheKey(schemas.MCPAuthModeVK, "vk1", "client-a") + // Admin-mode entries carry an empty identity component and must be + // swept with the rest of the client's entries. + keyA3 := headerCredentialCacheKey(schemas.MCPAuthModeAdmin, "", "client-a") + keyB := headerCredentialCacheKey(schemas.MCPAuthModeUser, "u1", "client-b") + for i, k := range []string{keyA1, keyA2, keyA3, keyB} { + _, err := c.Fill(context.Background(), k, fillWithCredential(testCredential(fmt.Sprintf("c%d", i), nil))) + require.NoError(t, err) + } + + c.EvictByMCPClient("client-a") + _, ok := c.Get(keyA1) + assert.False(t, ok, "user-mode entry for the client must be gone") + _, ok = c.Get(keyA2) + assert.False(t, ok, "vk-mode entry for the client must be gone") + _, ok = c.Get(keyA3) + assert.False(t, ok, "admin-mode entry for the client must be gone") + _, ok = c.Get(keyB) + assert.True(t, ok, "entry for another client must survive") + assert.Equal(t, 1, c.Len()) + + // Empty client ID and unknown client IDs are no-ops. + c.EvictByMCPClient("") + c.EvictByMCPClient("client-x") + assert.Equal(t, 1, c.Len()) +} + +func TestHeaderCredentialCache_EvictByVirtualKey(t *testing.T) { + c := newHeaderCredentialCache(8) + keyVK1A := headerCredentialCacheKey(schemas.MCPAuthModeVK, "vk1", "client-a") + keyVK1B := headerCredentialCacheKey(schemas.MCPAuthModeVK, "vk1", "client-b") + keyVK2 := headerCredentialCacheKey(schemas.MCPAuthModeVK, "vk2", "client-a") + // A user-mode identity that happens to equal the VK ID must survive: + // the eviction is scoped to vk-mode entries only. + keyUser := headerCredentialCacheKey(schemas.MCPAuthModeUser, "vk1", "client-a") + for i, k := range []string{keyVK1A, keyVK1B, keyVK2, keyUser} { + _, err := c.Fill(context.Background(), k, fillWithCredential(testCredential(fmt.Sprintf("vc%d", i), nil))) + require.NoError(t, err) + } + + c.EvictByVirtualKey("vk1") + _, ok := c.Get(keyVK1A) + assert.False(t, ok) + _, ok = c.Get(keyVK1B) + assert.False(t, ok, "the VK's entries must be evicted across clients") + _, ok = c.Get(keyVK2) + assert.True(t, ok, "another VK's entry must survive") + _, ok = c.Get(keyUser) + assert.True(t, ok, "a user-mode identity equal to the VK ID must survive") + + // Empty and unknown VK IDs are no-ops. + c.EvictByVirtualKey("") + c.EvictByVirtualKey("vk-x") + assert.Equal(t, 2, c.Len()) +} + +func TestHeaderCredentialCache_EvictByUser(t *testing.T) { + c := newHeaderCredentialCache(8) + keyU1A := headerCredentialCacheKey(schemas.MCPAuthModeUser, "u1", "client-a") + keyU1B := headerCredentialCacheKey(schemas.MCPAuthModeUser, "u1", "client-b") + keyU2 := headerCredentialCacheKey(schemas.MCPAuthModeUser, "u2", "client-a") + // A vk-mode identity that happens to equal the user ID must survive: + // the eviction is scoped to user-mode entries only. + keyVK := headerCredentialCacheKey(schemas.MCPAuthModeVK, "u1", "client-a") + for i, k := range []string{keyU1A, keyU1B, keyU2, keyVK} { + _, err := c.Fill(context.Background(), k, fillWithCredential(testCredential(fmt.Sprintf("uc%d", i), nil))) + require.NoError(t, err) + } + + c.EvictByUser("u1") + _, ok := c.Get(keyU1A) + assert.False(t, ok) + _, ok = c.Get(keyU1B) + assert.False(t, ok, "the user's entries must be evicted across clients") + _, ok = c.Get(keyU2) + assert.True(t, ok, "another user's entry must survive") + _, ok = c.Get(keyVK) + assert.True(t, ok, "a vk-mode identity equal to the user ID must survive") + + // Empty and unknown user IDs are no-ops. + c.EvictByUser("") + c.EvictByUser("u-x") + assert.Equal(t, 2, c.Len()) +} + +func TestHeaderCredentialCache_Flush(t *testing.T) { + c := newHeaderCredentialCache(4) + for i := 1; i <= 3; i++ { + _, err := c.Fill(context.Background(), fmt.Sprintf("k%d", i), fillWithCredential(testCredential(fmt.Sprintf("c%d", i), nil))) + require.NoError(t, err) + } + + c.Flush() + assert.Equal(t, 0, c.Len()) + for i := 1; i <= 3; i++ { + _, ok := c.Get(fmt.Sprintf("k%d", i)) + assert.False(t, ok) + } +} + +func TestHeaderCredentialCache_InflightDedup(t *testing.T) { + c := newHeaderCredentialCache(4) + + var calls atomic.Int64 + release := make(chan struct{}) + started := make(chan struct{}) + + fill := func() (cachedHeaderCredential, error) { + calls.Add(1) + close(started) + <-release + return testCredential("c1", map[string]string{"X-Api-Key": "shared"}), nil + } + + var wg sync.WaitGroup + results := make([]cachedHeaderCredential, 2) + fillErrs := make([]error, 2) + wg.Add(1) + go func() { + defer wg.Done() + v, err := c.Fill(context.Background(), "k1", fill) + results[0], fillErrs[0] = v, err + }() + <-started + + wg.Add(1) + go func() { + defer wg.Done() + // This fill func must never run: the leader is already in flight. + v, err := c.Fill(context.Background(), "k1", func() (cachedHeaderCredential, error) { + calls.Add(1) + return testCredential("c-other", nil), nil + }) + results[1], fillErrs[1] = v, err + }() + + // Give the second goroutine a moment to register as a waiter, then let + // the leader finish. + time.Sleep(20 * time.Millisecond) + close(release) + wg.Wait() + + // require, not assert, calls t.FailNow(), which per the testing + // package's own contract must only be called from the goroutine running + // the test — hence collecting errors above and asserting them here. + require.NoError(t, fillErrs[0]) + require.NoError(t, fillErrs[1]) + + assert.Equal(t, int64(1), calls.Load(), "concurrent fills for one key must run the handler once") + assert.Equal(t, "c1", results[0].credentialID) + assert.Equal(t, "c1", results[1].credentialID, "waiter must share the leader's result") +} + +func TestHeaderCredentialCache_ErrorSharedNotCached(t *testing.T) { + c := newHeaderCredentialCache(4) + + var calls atomic.Int64 + release := make(chan struct{}) + started := make(chan struct{}) + fillErr := errors.New("db unavailable") + + var wg sync.WaitGroup + errs := make([]error, 2) + wg.Add(1) + go func() { + defer wg.Done() + _, err := c.Fill(context.Background(), "k1", func() (cachedHeaderCredential, error) { + calls.Add(1) + close(started) + <-release + return cachedHeaderCredential{}, fillErr + }) + errs[0] = err + }() + <-started + + wg.Add(1) + go func() { + defer wg.Done() + _, err := c.Fill(context.Background(), "k1", func() (cachedHeaderCredential, error) { + calls.Add(1) + return cachedHeaderCredential{}, fillErr + }) + errs[1] = err + }() + + time.Sleep(20 * time.Millisecond) + close(release) + wg.Wait() + + assert.Equal(t, int64(1), calls.Load(), "waiter must share the leader's error, not re-run the fill") + assert.ErrorIs(t, errs[0], fillErr) + assert.ErrorIs(t, errs[1], fillErr, "error must be shared with waiters") + assert.Equal(t, 0, c.Len(), "errors must never be cached") + + // A later fill runs the handler again: the failure was not cached. + _, err := c.Fill(context.Background(), "k1", func() (cachedHeaderCredential, error) { + calls.Add(1) + return testCredential("c1", nil), nil + }) + require.NoError(t, err) + assert.Equal(t, int64(2), calls.Load()) + got, ok := c.Get("k1") + require.True(t, ok) + assert.Equal(t, "c1", got.credentialID) +} + +func TestHeaderCredentialCache_GenerationGuardDiscardsStaleFill(t *testing.T) { + c := newHeaderCredentialCache(4) + + inFill := make(chan struct{}) + release := make(chan struct{}) + + done := make(chan struct{}) + go func() { + defer close(done) + v, err := c.Fill(context.Background(), "k1", func() (cachedHeaderCredential, error) { + close(inFill) + <-release + return testCredential("c-stale", nil), nil + }) + // The caller still receives the value it read; only the cache + // install is discarded. assert, not require: require.NoError calls + // t.FailNow(), which per the testing package's own contract must + // only be called from the goroutine running the test. + assert.NoError(t, err) + assert.Equal(t, "c-stale", v.credentialID) + }() + + <-inFill + // An eviction lands while the fill is mid-read: whatever the fill read + // is now suspect and must not be installed. + c.Evict("k1") + close(release) + <-done + + _, ok := c.Get("k1") + assert.False(t, ok, "a fill that raced an eviction must not install its result") + assert.Equal(t, 0, c.Len()) +} + +func TestHeaderCredentialCache_UpsertRebindsCredentialIDIndex(t *testing.T) { + c := newHeaderCredentialCache(4) + _, err := c.Fill(context.Background(), "k1", fillWithCredential(testCredential("c-old", nil))) + require.NoError(t, err) + + // Same key, new backing row (the binding resubmitted onto a fresh row): + // the index must follow the new row ID and drop the old one. + c.Evict("k1") + _, err = c.Fill(context.Background(), "k1", fillWithCredential(testCredential("c-new", nil))) + require.NoError(t, err) + + c.EvictByCredentialID("c-old") + _, ok := c.Get("k1") + assert.True(t, ok, "stale credential-ID mapping must not evict the rebound entry") + + c.EvictByCredentialID("c-new") + _, ok = c.Get("k1") + assert.False(t, ok, "current credential-ID mapping must evict the entry") +} + +func TestHeaderCredentialCache_NilSafety(t *testing.T) { + var c *headerCredentialCache + _, ok := c.Get("k") + assert.False(t, ok) + c.Evict("k") + c.EvictByCredentialID("c") + c.EvictByMCPClient("client") + c.EvictByVirtualKey("vk") + c.Flush() + assert.Equal(t, 0, c.Len()) + v, err := c.Fill(context.Background(), "k", fillWithCredential(testCredential("c-pass", nil))) + require.NoError(t, err) + assert.Equal(t, "c-pass", v.credentialID) +} + +// ---------- Integration through GetCredentialByMode ---------- + +// UpsertMCPPerUserHeaderCredential mirrors the real store's binding-keyed +// upsert semantics on the test double: an existing row for the same MCP +// client is reused and the passed row's ID is rewritten to it, exactly the +// contract the provider's post-upsert eviction relies on. +func (s *testConfigStore) UpsertMCPPerUserHeaderCredential(_ context.Context, cred *tables.TableMCPPerUserHeaderCredential) error { + s.mu.Lock() + defer s.mu.Unlock() + if existing, ok := s.credentials[cred.MCPClientID]; ok { + cred.ID = existing.ID + cred.CreatedAt = existing.CreatedAt + } else if cred.ID == "" { + cred.ID = uuid.NewString() + } + s.credentials[cred.MCPClientID] = bifrost.Ptr(*cred) + return nil +} + +func (s *testConfigStore) DeleteMCPPerUserHeaderCredential(_ context.Context, id string) error { + s.mu.Lock() + defer s.mu.Unlock() + for clientID, row := range s.credentials { + if row.ID == id { + delete(s.credentials, clientID) + } + } + return nil +} + +// seedHeaderCredential inserts a session-mode credential row into the double. +func seedHeaderCredential(store *testConfigStore, id, mcpClientID, status, headersJSON string) { + now := time.Now() + store.mu.Lock() + defer store.mu.Unlock() + store.credentials[mcpClientID] = &tables.TableMCPPerUserHeaderCredential{ + ID: id, + MCPClientID: mcpClientID, + AuthMode: "session", + SessionID: "sess-1", + Status: status, + HeadersJSON: headersJSON, + CreatedAt: now, + UpdatedAt: now, + } +} + +func lookupCalls(store *testConfigStore) int { + store.mu.Lock() + defer store.mu.Unlock() + return store.credentialLookupCalls +} + +func TestGetCredentialByMode_AdminIdentityNormalizedToOneCacheEntry(t *testing.T) { + store := newTestConfigStore() + now := time.Now() + store.mu.Lock() + store.credentials["mcp-1"] = &tables.TableMCPPerUserHeaderCredential{ + ID: "cred-admin", MCPClientID: "mcp-1", AuthMode: "admin", + Status: "active", HeadersJSON: `{"X-Api-Key":"v1"}`, + CreatedAt: now, UpdatedAt: now, + } + store.mu.Unlock() + provider := newTestProvider(store) + ctx := context.Background() + + // The store lookup ignores identity for admin mode, so lookups with and + // without a stray identity must share one cache entry; otherwise the + // byID index could only evict one of the aliases. + _, err := provider.GetCredentialByMode(ctx, schemas.MCPAuthModeAdmin, "stray", "mcp-1") + require.NoError(t, err) + _, err = provider.GetCredentialByMode(ctx, schemas.MCPAuthModeAdmin, "", "mcp-1") + require.NoError(t, err) + assert.Equal(t, 1, lookupCalls(store), "admin lookups must share one cache entry regardless of identity") + + provider.EvictCredentialByID("cred-admin") + _, err = provider.GetCredentialByMode(ctx, schemas.MCPAuthModeAdmin, "stray", "mcp-1") + require.NoError(t, err) + assert.Equal(t, 2, lookupCalls(store), "byID eviction must invalidate the admin entry for every identity spelling") +} + +func TestGetCredentialByMode_SecondCallServedFromCache(t *testing.T) { + store := newTestConfigStore() + seedHeaderCredential(store, "cred-1", "mcp-1", "active", `{"X-Api-Key":"v1"}`) + provider := newTestProvider(store) + ctx := context.Background() + + cred, err := provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + assert.Equal(t, "v1", cred.Headers["X-Api-Key"]) + assert.Equal(t, 1, lookupCalls(store)) + + cred, err = provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + assert.Equal(t, "v1", cred.Headers["X-Api-Key"]) + assert.Equal(t, 1, lookupCalls(store), "second call must be served from cache, not the store") +} + +func TestGetCredentialByMode_CachedHitReturnsIsolatedCopy(t *testing.T) { + store := newTestConfigStore() + seedHeaderCredential(store, "cred-1", "mcp-1", "active", `{"X-Api-Key":"v1"}`) + provider := newTestProvider(store) + ctx := context.Background() + + first, err := provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + // A caller scribbling on its copy must never leak into the cache. + first.Headers["X-Api-Key"] = "mutated" + first.Status = schemas.MCPHeadersUserCredentialStatusOrphaned + + second, err := provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + assert.Equal(t, "v1", second.Headers["X-Api-Key"], "cached value must be isolated from caller mutation") + assert.Equal(t, schemas.MCPHeadersUserCredentialStatusActive, second.Status) +} + +func TestGetCredentialByMode_UpsertEvictsSoNextReadSeesNewValues(t *testing.T) { + store := newTestConfigStore() + seedHeaderCredential(store, "cred-1", "mcp-1", "active", `{"X-Api-Key":"old"}`) + provider := newTestProvider(store) + ctx := context.Background() + + cred, err := provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + require.Equal(t, "old", cred.Headers["X-Api-Key"]) + + sessionID := "sess-1" + require.NoError(t, provider.UpsertCredential(ctx, &schemas.MCPHeadersUserCredential{ + MCPClientID: "mcp-1", + AuthMode: schemas.MCPAuthModeSession, + SessionID: &sessionID, + Headers: map[string]string{"X-Api-Key": "new"}, + Status: schemas.MCPHeadersUserCredentialStatusActive, + })) + + cred, err = provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + assert.Equal(t, "new", cred.Headers["X-Api-Key"], "upsert must evict so the next read serves the new values") + assert.Equal(t, 2, lookupCalls(store), "post-upsert read must come from the store") +} + +func TestGetCredentialByMode_DeleteEvictsSoNextReadIsNotFound(t *testing.T) { + store := newTestConfigStore() + seedHeaderCredential(store, "cred-1", "mcp-1", "active", `{"X-Api-Key":"v1"}`) + provider := newTestProvider(store) + ctx := context.Background() + + _, err := provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + + require.NoError(t, provider.DeleteCredential(ctx, "cred-1")) + + _, err = provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.Error(t, err) + assert.ErrorIs(t, err, schemas.ErrHeadersCredentialNotFound, "post-delete lookup must see the delete, not the cached credential") +} + +func TestGetCredentialByMode_EvictByIDAfterDirectStoreDelete(t *testing.T) { + store := newTestConfigStore() + seedHeaderCredential(store, "cred-1", "mcp-1", "active", `{"X-Api-Key":"v1"}`) + provider := newTestProvider(store) + ctx := context.Background() + + _, err := provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + + // Delete the row directly through the store (as the sessions revoke + // handler does) and evict by ID (as its cache callback does). + store.mu.Lock() + delete(store.credentials, "mcp-1") + store.mu.Unlock() + provider.EvictCredentialByID("cred-1") + + _, err = provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.Error(t, err) + assert.ErrorIs(t, err, schemas.ErrHeadersCredentialNotFound) +} + +func TestGetCredentialByMode_NeedsUpdateRowIsCachedWithStatus(t *testing.T) { + store := newTestConfigStore() + seedHeaderCredential(store, "cred-1", "mcp-1", "needs_update", `{"X-Api-Key":"stale"}`) + provider := newTestProvider(store) + ctx := context.Background() + + cred, err := provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + assert.Equal(t, schemas.MCPHeadersUserCredentialStatusNeedsUpdate, cred.Status) + + cred, err = provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + assert.Equal(t, schemas.MCPHeadersUserCredentialStatusNeedsUpdate, cred.Status, "the cached copy must carry the row's status through") + assert.Equal(t, "stale", cred.Headers["X-Api-Key"]) + assert.Equal(t, 1, lookupCalls(store), "needs_update rows are cacheable like active rows") +} + +func TestGetCredentialByMode_NegativeNotCached(t *testing.T) { + store := newTestConfigStore() + provider := newTestProvider(store) + ctx := context.Background() + + _, err := provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.Error(t, err) + assert.ErrorIs(t, err, schemas.ErrHeadersCredentialNotFound) + + // The user completes a submission flow: a fresh row appears. No eviction + // happens (there is nothing to evict) and the very next call must see it. + seedHeaderCredential(store, "cred-1", "mcp-1", "active", `{"X-Api-Key":"fresh"}`) + + cred, err := provider.GetCredentialByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + assert.Equal(t, "fresh", cred.Headers["X-Api-Key"], "a failed lookup must never be cached") +} + +func TestGetCredentialByMode_EvictByMCPClientCoversAdminMode(t *testing.T) { + store := newTestConfigStore() + now := time.Now() + store.mu.Lock() + store.credentials["mcp-1"] = &tables.TableMCPPerUserHeaderCredential{ + ID: "cred-admin", + MCPClientID: "mcp-1", + AuthMode: "admin", + Status: "active", + HeadersJSON: `{"X-Api-Key":"admin-v1"}`, + CreatedAt: now, + UpdatedAt: now, + } + store.mu.Unlock() + provider := newTestProvider(store) + ctx := context.Background() + + _, err := provider.GetCredentialByMode(ctx, schemas.MCPAuthModeAdmin, "", "mcp-1") + require.NoError(t, err) + require.Equal(t, 1, lookupCalls(store)) + + provider.EvictCredentialsByMCPClient("mcp-1") + + _, err = provider.GetCredentialByMode(ctx, schemas.MCPAuthModeAdmin, "", "mcp-1") + require.NoError(t, err) + assert.Equal(t, 2, lookupCalls(store), "client-scoped eviction must also drop the admin-mode binding") +} diff --git a/framework/mcp_headers/main.go b/framework/mcp_headers/main.go index c23ce8dba7c..e36e689a628 100644 --- a/framework/mcp_headers/main.go +++ b/framework/mcp_headers/main.go @@ -13,6 +13,7 @@ import ( "context" "errors" "fmt" + "maps" "strings" "sync/atomic" "time" @@ -44,6 +45,15 @@ type Provider struct { // once at startup and read lock-free on the request path. Mirrors // oauth2.OAuth2Provider.tempTokens exactly. tempTokens atomic.Pointer[temptoken.Service] + + // credentials caches per-user header credential lookups keyed by + // (auth mode, identity, mcp client). It owns its own locking; the + // provider itself is lock-free, so the cache's mutexes only ever guard + // map and list surgery. Header credentials have no expiry and no + // refresh machinery, so cached entries never self-heal: write paths in + // this file evict inline, and handler-level database writes that bypass + // the provider evict through the exported eviction methods. + credentials *headerCredentialCache } // NewProvider constructs a configstore-backed MCPHeadersProvider. Mirrors @@ -53,7 +63,11 @@ func NewProvider(configStore configstore.ConfigStore, logger schemas.Logger) *Pr if logger == nil { logger = bifrost.NewDefaultLogger(schemas.LogLevelInfo) } - return &Provider{configStore: configStore, logger: logger} + return &Provider{ + configStore: configStore, + logger: logger, + credentials: newHeaderCredentialCache(defaultCredentialCacheCapacity), + } } // SetTempTokenService installs the temp-token service used by @@ -92,6 +106,15 @@ func (p *Provider) mcpTempTokenAuthEnabled(ctx context.Context) bool { // absent so callers can switch on the sentinel. mode can also be // MCPAuthModeAdmin: the retained bootstrap credential has no per-caller // identity, so identity is allowed empty in that case alone. +// +// Reads are served from an in-memory cache when possible. Both 'active' and +// 'needs_update' rows are cached with their status carried through, since +// the store lookup deliberately returns both and callers decide usability +// themselves. Header credentials have no expiry, so a cached entry is +// served until an eviction removes it. Failed lookups (including not-found) +// are never cached, so a caller completing a submission flow becomes +// visible on the very next call. Every hit returns a fresh deep copy so a +// caller mutating the result can never corrupt the cached value. func (p *Provider) GetCredentialByMode(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (*schemas.MCPHeadersUserCredential, error) { if p.configStore == nil { return nil, schemas.ErrHeadersCredentialProviderNotAvailable @@ -102,18 +125,70 @@ func (p *Provider) GetCredentialByMode(ctx context.Context, mode schemas.MCPAuth if mode != schemas.MCPAuthModeAdmin && strings.TrimSpace(identity) == "" { return nil, schemas.ErrHeadersCredentialNotFound } + if mode == schemas.MCPAuthModeAdmin { + // The store lookup ignores identity for admin mode (the retained + // bootstrap credential is keyed by client alone), so normalize it + // before building the cache key. A non-empty identity would alias + // the same row under multiple keys, and the byID index can only + // evict one of them. + identity = "" + } + key := headerCredentialCacheKey(mode, identity, mcpClientID) + if cached, ok := p.credentials.Get(key); ok { + return copyHeaderCredential(cached.credential), nil + } + cached, err := p.credentials.Fill(ctx, key, func() (cachedHeaderCredential, error) { + return p.loadCredentialByMode(ctx, mode, identity, mcpClientID) + }) + if err != nil { + return nil, err + } + return copyHeaderCredential(cached.credential), nil +} + +// loadCredentialByMode is GetCredentialByMode's cache-miss path: the full +// database lookup plus row-to-schema conversion. +func (p *Provider) loadCredentialByMode(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (cachedHeaderCredential, error) { row, err := p.configStore.GetMCPPerUserHeaderCredentialByMode(ctx, mode, identity, mcpClientID) if err != nil { - return nil, fmt.Errorf("load mcp per-user header credential: %w", err) + return cachedHeaderCredential{}, fmt.Errorf("load mcp per-user header credential: %w", err) } if row == nil { - return nil, schemas.ErrHeadersCredentialNotFound + return cachedHeaderCredential{}, schemas.ErrHeadersCredentialNotFound } cred, err := rowToCredential(row) if err != nil { - return nil, err + return cachedHeaderCredential{}, err } - return cred, nil + return cachedHeaderCredential{credentialID: cred.ID, credential: cred}, nil +} + +// copyHeaderCredential returns a deep copy of a cached credential: fresh +// Headers map and fresh identity pointers, so no two callers (and never the +// cache itself) share mutable state. +func copyHeaderCredential(cred *schemas.MCPHeadersUserCredential) *schemas.MCPHeadersUserCredential { + if cred == nil { + return nil + } + out := *cred + if cred.Headers != nil { + headers := make(map[string]string, len(cred.Headers)) + maps.Copy(headers, cred.Headers) + out.Headers = headers + } + if cred.UserID != nil { + v := *cred.UserID + out.UserID = &v + } + if cred.VirtualKeyID != nil { + v := *cred.VirtualKeyID + out.VirtualKeyID = &v + } + if cred.SessionID != nil { + v := *cred.SessionID + out.SessionID = &v + } + return &out } // UpsertCredential persists the caller-supplied credential. The caller is @@ -141,6 +216,11 @@ func (p *Provider) UpsertCredential(ctx context.Context, cred *schemas.MCPHeader cred.ID = row.ID cred.CreatedAt = row.CreatedAt cred.UpdatedAt = row.UpdatedAt + // UpsertMCPPerUserHeaderCredential upserts by binding and rewrites + // row.ID to the reused row's ID when one existed, so this evicts exactly + // the row that was written; the next lookup for this binding reads the + // new values. A fresh row has nothing cached and the evict is a no-op. + p.EvictCredentialByID(row.ID) return nil } @@ -155,9 +235,81 @@ func (p *Provider) DeleteCredential(ctx context.Context, id string) error { if err := p.configStore.DeleteMCPPerUserHeaderCredential(ctx, id); err != nil { return fmt.Errorf("delete mcp per-user header credential: %w", err) } + p.EvictCredentialByID(id) return nil } +// EvictCredential drops the cached credential for one (mode, identity, +// mcp client) binding. Side-effect only and safe to call when the cache is +// absent; the next lookup for the binding reads the database. +func (p *Provider) EvictCredential(mode schemas.MCPAuthMode, identity, mcpClientID string) { + if p == nil { + return + } + if mode == schemas.MCPAuthModeAdmin { + // Admin-mode entries are always cached under an empty identity (see + // GetCredentialByMode's normalization); without this, an eviction + // call carrying a non-empty identity would compute a key that was + // never installed, leaving the stale admin credential cached. + identity = "" + } + p.credentials.Evict(headerCredentialCacheKey(mode, identity, mcpClientID)) +} + +// EvictCredentialByID drops the cached credential backed by the given row +// ID, if any binding currently holds it. Side-effect only and safe to call +// when the cache is absent or the ID is not cached. +func (p *Provider) EvictCredentialByID(id string) { + if p == nil { + return + } + p.credentials.EvictByCredentialID(id) +} + +// EvictCredentialsByMCPClient drops every cached credential bound to the +// given MCP client, across all auth modes and identities. Used after +// client-level mutations that invalidate its credential rows as a set, such +// as a needs_update schema flip, access reconciliation, or client deletion. +// Side-effect only and safe to call when the cache is absent. +func (p *Provider) EvictCredentialsByMCPClient(mcpClientID string) { + if p == nil { + return + } + p.credentials.EvictByMCPClient(mcpClientID) +} + +// EvictCredentialsByVirtualKey drops every cached vk-mode credential bound +// to the given virtual key, across all MCP clients. Used after virtual key +// mutations that orphan or delete its credential rows as a set. Side-effect +// only and safe to call when the cache is absent. +func (p *Provider) EvictCredentialsByVirtualKey(virtualKeyID string) { + if p == nil { + return + } + p.credentials.EvictByVirtualKey(virtualKeyID) +} + +// EvictCredentialsByUser drops every cached user-mode credential bound to +// the given user, across all MCP clients. Used after user-level mutations +// that orphan or delete the user's credential rows as a set. Side-effect +// only and safe to call when the cache is absent. +func (p *Provider) EvictCredentialsByUser(userID string) { + if p == nil { + return + } + p.credentials.EvictByUser(userID) +} + +// FlushCredentialCache drops every cached per-user header credential. The +// coarse fallback for mutations whose blast radius cannot be scoped to one +// client or virtual key. Side-effect only. +func (p *Provider) FlushCredentialCache() { + if p == nil { + return + } + p.credentials.Flush() +} + // InitiateUserSubmissionFlow creates a pending mcp_per_user_header_flows // row keyed by (mode, identity, mcp_client_id), mints a // mcp_headers_auth temp-token bound to the new row's ID, and returns the diff --git a/framework/mcp_headers/sweep.go b/framework/mcp_headers/sweep.go index a198d335fc4..13112ae893e 100644 --- a/framework/mcp_headers/sweep.go +++ b/framework/mcp_headers/sweep.go @@ -26,14 +26,14 @@ import ( // orphanRetention disables the orphan sweep entirely; the expired-flow sweep // always runs because flow rows have no semantic value past their expiry. type CredentialSweepWorker struct { - provider *Provider - orphanSweepEvery time.Duration - orphanRetention time.Duration - expiredFlowEvery time.Duration - stopCh chan struct{} - stopOnce sync.Once - cancel context.CancelFunc - logger schemas.Logger + provider *Provider + orphanSweepEvery time.Duration + orphanRetention time.Duration + expiredFlowEvery time.Duration + stopCh chan struct{} + stopOnce sync.Once + cancel context.CancelFunc + logger schemas.Logger } // NewCredentialSweepWorker creates a sweep worker with sensible defaults. @@ -111,6 +111,11 @@ func (w *CredentialSweepWorker) sweepOrphanedCredentials(ctx context.Context) { if w.orphanRetention <= 0 { return } + // No cache eviction here: the provider's credential cache can never hold + // an orphaned row (the ByMode lookup filters them at SQL), and the + // reconcile paths that flip rows to 'orphaned' already evict when the + // flip happens. Deleting the rows outright therefore cannot invalidate + // any cached entry. n, err := w.provider.configStore.DeleteOrphanedMCPPerUserHeaderCredentials(ctx, w.orphanRetention) if err != nil { if w.logger != nil { diff --git a/framework/oauth2/usertokencache.go b/framework/oauth2/usertokencache.go index d3f383ab8a0..ae8160beadd 100644 --- a/framework/oauth2/usertokencache.go +++ b/framework/oauth2/usertokencache.go @@ -2,9 +2,6 @@ package oauth2 import ( "context" - "fmt" - "strconv" - "strings" "time" "github.com/maximhq/bifrost/core/schemas" @@ -44,41 +41,20 @@ type userTokenCache struct { } // userTokenCacheKey builds the cache key for a (mode, identity, mcp client) -// binding. Length-prefixed rather than NUL-separated: identity is a -// caller-asserted string (see the doc comment above), so a NUL-separated -// scheme would let an identity value containing "\x00" forge a component -// boundary — e.g. (mode, "a\x00b", "c") and (mode, "a", "b\x00c") would build -// the identical key, letting one binding's cached access token be served to -// another. Prefixing each component with its own byte length makes that -// forgery impossible regardless of what bytes a component contains. +// binding. identity is a caller-asserted string (see the doc comment +// above) with no charset restriction, so this goes through +// lrucache.EncodeKey rather than a plain separator join — see its doc +// comment for why a naive join lets an identity value forge a component +// boundary and alias two distinct bindings onto one cache entry. func userTokenCacheKey(mode schemas.MCPAuthMode, identity, mcpClientID string) string { - return fmt.Sprintf("%d:%s%d:%s%d:%s", len(mode), mode, len(identity), identity, len(mcpClientID), mcpClientID) + return lrucache.EncodeKey(string(mode), identity, mcpClientID) } // splitUserTokenCacheKey is userTokenCacheKey's inverse, for the scoped -// eviction predicates. ok is false for a key that isn't in the -// length-prefixed form userTokenCacheKey builds (impossible for keys this -// package produces). +// eviction predicates. func splitUserTokenCacheKey(key string) (mode, identity, clientID string, ok bool) { - rest := key - parts := make([]string, 0, 3) - for range 3 { - i := strings.IndexByte(rest, ':') - if i < 0 { - return "", "", "", false - } - n, err := strconv.Atoi(rest[:i]) - if err != nil || n < 0 { - return "", "", "", false - } - rest = rest[i+1:] - if n > len(rest) { - return "", "", "", false - } - parts = append(parts, rest[:n]) - rest = rest[n:] - } - if rest != "" { + parts, ok := lrucache.DecodeKey(key, 3) + if !ok { return "", "", "", false } return parts[0], parts[1], parts[2], true diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index a85a9e1a1df..c3d14623108 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -2098,8 +2098,9 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { // VK's owner (user-keyed creds) against the new effective allowlist // (explicit rows ∪ MCPs with AllowOnAllVirtualKeys=true). OSS no-ops. // Must run before ReloadVirtualKey: the reload also evicts this VK's - // cached OAuth tokens, and an eviction that lands before these writes - // could be refilled from the pre-reconcile rows and then never dropped. + // cached OAuth tokens and header credentials, and an eviction that lands + // before these writes could be refilled from the pre-reconcile rows and + // then never dropped. if req.MCPConfigs != nil && h.configStore != nil { if err := h.configStore.ReconcileOauthAfterVKChange(ctx, vk.ID); err != nil { logger.Error("reconcile OAuth credentials after VK %s update failed: %v", vk.ID, err) @@ -2244,8 +2245,8 @@ func (h *GovernanceHandler) deleteVirtualKey(ctx *fasthttp.RequestCtx) { return } // Removing key from in-memory store. RemoveVirtualKey also evicts the - // VK's cached OAuth access tokens internally, covering the token rows - // the database delete above cascaded over. + // VK's cached OAuth access tokens and header credentials internally, + // covering the rows the database delete above cascaded over. err = h.governanceManager.RemoveVirtualKey(ctx, vk.ID) if err != nil { // But we ignore this error because its not diff --git a/transports/bifrost-http/handlers/mcp.go b/transports/bifrost-http/handlers/mcp.go index 952404b06c7..bc7d8b5f17b 100644 --- a/transports/bifrost-http/handlers/mcp.go +++ b/transports/bifrost-http/handlers/mcp.go @@ -60,11 +60,12 @@ type MCPHandler struct { mcpManager MCPManager governanceManager GovernanceManager oauthHandler *OAuthHandler - // mcpOauthTokenCacheManager invalidates cached OAuth access tokens after - // mutations that rewrite or delete token rows through the configstore - // (credential rotation, access reconciliation). Always wired by the - // server; see the interface doc for the non-nil requirement. - mcpOauthTokenCacheManager MCPOauthTokenCacheManager + // mcpCredentialCacheManager invalidates cached per-user credentials + // (OAuth access tokens and header credentials) after mutations that + // rewrite or delete their rows through the configstore (credential + // rotation, needs_update schema flips, access reconciliation). Always + // wired by the server; see the interface doc for the non-nil requirement. + mcpCredentialCacheManager MCPCredentialCacheManager } // NewMCPHandler creates a new MCP handler instance @@ -74,7 +75,7 @@ func NewMCPHandler( client *bifrost.Bifrost, store *lib.Config, oauthHandler *OAuthHandler, - mcpOauthTokenCacheManager MCPOauthTokenCacheManager, + mcpCredentialCacheManager MCPCredentialCacheManager, ) *MCPHandler { return &MCPHandler{ client: client, @@ -82,7 +83,7 @@ func NewMCPHandler( mcpManager: mcpManager, governanceManager: governanceManager, oauthHandler: oauthHandler, - mcpOauthTokenCacheManager: mcpOauthTokenCacheManager, + mcpCredentialCacheManager: mcpCredentialCacheManager, } } @@ -1996,7 +1997,7 @@ func (h *MCPHandler) updateMCPClient(ctx *fasthttp.RequestCtx) { return } if rotated { - h.mcpOauthTokenCacheManager.EvictOauthTokenCacheByMCPClient(ctx, id) + h.mcpCredentialCacheManager.EvictOauthTokenCacheByMCPClient(ctx, id) // Rotation just cascaded every token bound to this oauth_config to // needs_reauth in the DB, but the in-memory client above only had // its ExecutionConfig replaced — its live connection, if any, is @@ -2030,6 +2031,10 @@ func (h *MCPHandler) updateMCPClient(ctx *fasthttp.RequestCtx) { h.store.ConfigStore != nil { if err := h.store.ConfigStore.MarkMCPPerUserHeaderCredentialsNeedsUpdate(ctx, existingConfig.ID); err != nil { logger.Error(fmt.Sprintf("failed to flip per-user header credentials to needs_update for client %s: %v", existingConfig.ID, err)) + } else { + // Cached copies still carry the pre-flip status and values; drop + // them so the next lookup reads the needs_update rows. + h.mcpCredentialCacheManager.EvictMCPHeaderCredentialCacheByMCPClient(ctx, existingConfig.ID) } } @@ -2165,9 +2170,10 @@ func (h *MCPHandler) updateMCPClient(ctx *fasthttp.RequestCtx) { if err := h.store.ConfigStore.ReconcileMCPHeadersAfterMCPChange(ctx, id); err != nil { logger.Error(fmt.Sprintf("reconcile per-user-headers credentials after MCP %s update failed: %v", id, err)) } - // Reconciliation may have orphaned or reactivated token rows; - // cached copies no longer reflect the database. - h.mcpOauthTokenCacheManager.EvictOauthTokenCacheByMCPClient(ctx, id) + // Reconciliation may have orphaned or reactivated token and + // credential rows; cached copies no longer reflect the database. + h.mcpCredentialCacheManager.EvictOauthTokenCacheByMCPClient(ctx, id) + h.mcpCredentialCacheManager.EvictMCPHeaderCredentialCacheByMCPClient(ctx, id) } } @@ -2197,8 +2203,8 @@ func (h *MCPHandler) deleteMCPClient(ctx *fasthttp.RequestCtx) { } } // RemoveMCPClient also evicts the client's cached OAuth access tokens - // internally, covering the token rows the database delete above - // cascaded over. + // and header credentials internally, covering the rows the database + // delete above cascaded over. if err := h.mcpManager.RemoveMCPClient(ctx, id); err != nil { SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("failed to remove MCP client: %v", err)) return diff --git a/transports/bifrost-http/handlers/mcpsessions.go b/transports/bifrost-http/handlers/mcpsessions.go index 3af954a637e..4b7c07badf2 100644 --- a/transports/bifrost-http/handlers/mcpsessions.go +++ b/transports/bifrost-http/handlers/mcpsessions.go @@ -23,29 +23,33 @@ import ( "github.com/valyala/fasthttp" ) -// MCPOauthTokenCacheManager invalidates cached per-user MCP OAuth access -// tokens after a database write that bypasses the OAuth provider's own write -// paths. The base server implementation evicts from local memory; a -// clustered deployment overrides it to also notify peers so their caches -// stay current. Handlers call it directly after successful writes, so the -// server must always wire it (it is nil-safe internally when no provider is -// configured, but the interface value itself must be non-nil). -type MCPOauthTokenCacheManager interface { +// MCPCredentialCacheManager invalidates cached per-user MCP credentials +// (OAuth access tokens and header credentials) after a database write that +// bypasses the owning provider's own write paths. The base server +// implementation evicts from local memory; a clustered deployment overrides +// it to also notify peers so their caches stay current. Handlers call it +// directly after successful writes, so the server must always wire it (it is +// nil-safe internally when no provider is configured, but the interface +// value itself must be non-nil). +type MCPCredentialCacheManager interface { EvictOauthTokenCacheByID(ctx context.Context, tokenID string) EvictOauthTokenCacheByMCPClient(ctx context.Context, mcpClientID string) + EvictMCPHeaderCredentialCacheByID(ctx context.Context, credentialID string) + EvictMCPHeaderCredentialCacheByMCPClient(ctx context.Context, mcpClientID string) } // MCPSessionsHandler serves the sessions tab API. type MCPSessionsHandler struct { store *lib.Config - // mcpOauthTokenCacheManager invalidates cached OAuth access tokens after - // this handler writes token rows directly through the configstore. - mcpOauthTokenCacheManager MCPOauthTokenCacheManager + // mcpCredentialCacheManager invalidates cached per-user credentials after + // this handler writes token or credential rows directly through the + // configstore. + mcpCredentialCacheManager MCPCredentialCacheManager } // NewMCPSessionsHandler creates the handler. -func NewMCPSessionsHandler(store *lib.Config, mcpOauthTokenCacheManager MCPOauthTokenCacheManager) *MCPSessionsHandler { - return &MCPSessionsHandler{store: store, mcpOauthTokenCacheManager: mcpOauthTokenCacheManager} +func NewMCPSessionsHandler(store *lib.Config, mcpCredentialCacheManager MCPCredentialCacheManager) *MCPSessionsHandler { + return &MCPSessionsHandler{store: store, mcpCredentialCacheManager: mcpCredentialCacheManager} } // RegisterRoutes registers the sessions tab routes. @@ -638,6 +642,7 @@ func (h *MCPSessionsHandler) revoke(ctx *fasthttp.RequestCtx) { SendError(ctx, fasthttp.StatusInternalServerError, "Failed to delete MCP session") return } + h.mcpCredentialCacheManager.EvictMCPHeaderCredentialCacheByID(ctx, headerCred.ID) logger.Debug("[mcp/sessions] revoked header credential: id=%s mcp_client=%s mode=%s", rowID, headerCred.MCPClientID, headerCred.AuthMode) ctx.SetStatusCode(fasthttp.StatusNoContent) return @@ -682,7 +687,7 @@ func (h *MCPSessionsHandler) revoke(ctx *fasthttp.RequestCtx) { SendError(ctx, fasthttp.StatusInternalServerError, "Failed to delete MCP session") return } - h.mcpOauthTokenCacheManager.EvictOauthTokenCacheByID(ctx, tok.ID) + h.mcpCredentialCacheManager.EvictOauthTokenCacheByID(ctx, tok.ID) logger.Debug("[mcp/sessions] revoked: token=%s mcp_client=%s mode=%s", rowID, tok.MCPClientID, tok.AuthMode) ctx.SetStatusCode(fasthttp.StatusNoContent) } diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index 2c09b417106..5f0cd9113df 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -149,6 +149,15 @@ type ServerCallbacks interface { // mutation that invalidates its token rows as a set (credential // rotation, access reconciliation, client deletion). EvictOauthTokenCacheByMCPClient(ctx context.Context, mcpClientID string) + // EvictMCPHeaderCredentialCacheByID drops the cached per-user MCP header + // credential backed by the given credential row ID after a database + // write that bypassed the headers provider's own write paths. + EvictMCPHeaderCredentialCacheByID(ctx context.Context, credentialID string) + // EvictMCPHeaderCredentialCacheByMCPClient drops every cached per-user + // MCP header credential bound to the given MCP client after a + // client-level mutation that invalidates its credential rows as a set + // (needs_update schema flip, access reconciliation, client deletion). + EvictMCPHeaderCredentialCacheByMCPClient(ctx context.Context, mcpClientID string) } // LogRedactionMappingResolverProvider is implemented by servers that can attach reveal data to log-detail responses. @@ -333,6 +342,7 @@ func (s *BifrostHTTPServer) RemoveMCPClient(ctx context.Context, id string) erro logger.Warn("failed to sync MCP servers after removing client: %v", err) } s.Config.OAuthProvider.EvictUserTokensByMCPClient(id) + s.Config.MCPHeadersProvider.EvictCredentialsByMCPClient(id) return nil } @@ -493,6 +503,7 @@ func (s *BifrostHTTPServer) ReloadVirtualKey(ctx context.Context, id string) (*t } s.MCPServerHandler.SyncVKMCPServer(virtualKey) s.Config.OAuthProvider.EvictUserTokensByVirtualKey(id) + s.Config.MCPHeadersProvider.EvictCredentialsByVirtualKey(id) return virtualKey, nil } @@ -511,11 +522,14 @@ func (s *BifrostHTTPServer) RemoveVirtualKey(ctx context.Context, id string) err if preloadedVk == nil { // This could be broadcast message from other server, so we will just clean up in-memory store governancePlugin.GetGovernanceStore().DeleteVirtualKeyInMemory(ctx, id) + s.Config.OAuthProvider.EvictUserTokensByVirtualKey(id) + s.Config.MCPHeadersProvider.EvictCredentialsByVirtualKey(id) return nil } governancePlugin.GetGovernanceStore().DeleteVirtualKeyInMemory(ctx, id) s.MCPServerHandler.DeleteVKMCPServer(preloadedVk.Value.GetValue()) s.Config.OAuthProvider.EvictUserTokensByVirtualKey(id) + s.Config.MCPHeadersProvider.EvictCredentialsByVirtualKey(id) return nil } @@ -999,6 +1013,48 @@ func (s *BifrostHTTPServer) FlushOauthTokenCache(ctx context.Context) { s.Config.OAuthProvider.FlushUserTokenCache() } +// EvictMCPHeaderCredentialCacheByID drops the cached per-user MCP header +// credential backed by the given credential row ID from the in-memory cache +// after a database write. A clustered deployment overrides this to also +// notify peers. +func (s *BifrostHTTPServer) EvictMCPHeaderCredentialCacheByID(ctx context.Context, credentialID string) { + if s.Config == nil || s.Config.MCPHeadersProvider == nil { + return + } + s.Config.MCPHeadersProvider.EvictCredentialByID(credentialID) +} + +// EvictMCPHeaderCredentialCacheByMCPClient drops every cached per-user MCP +// header credential bound to the given MCP client from the in-memory cache. +// A clustered deployment overrides this to also notify peers. +func (s *BifrostHTTPServer) EvictMCPHeaderCredentialCacheByMCPClient(ctx context.Context, mcpClientID string) { + if s.Config == nil || s.Config.MCPHeadersProvider == nil { + return + } + s.Config.MCPHeadersProvider.EvictCredentialsByMCPClient(mcpClientID) +} + +// EvictMCPHeaderCredentialCacheByVirtualKey drops every cached vk-mode MCP +// header credential bound to the given virtual key from the in-memory cache. +// A clustered deployment overrides this to also notify peers. +func (s *BifrostHTTPServer) EvictMCPHeaderCredentialCacheByVirtualKey(ctx context.Context, virtualKeyID string) { + if s.Config == nil || s.Config.MCPHeadersProvider == nil { + return + } + s.Config.MCPHeadersProvider.EvictCredentialsByVirtualKey(virtualKeyID) +} + +// FlushMCPHeaderCredentialCache drops every cached per-user MCP header +// credential from the in-memory cache. The coarse fallback for mutations +// whose blast radius cannot be scoped to one client or virtual key. A +// clustered deployment overrides this to also notify peers. +func (s *BifrostHTTPServer) FlushMCPHeaderCredentialCache(ctx context.Context) { + if s.Config == nil || s.Config.MCPHeadersProvider == nil { + return + } + s.Config.MCPHeadersProvider.FlushCredentialCache() +} + // ReloadClientConfigFromConfigStore reloads the client config from config store func (s *BifrostHTTPServer) ReloadClientConfigFromConfigStore(ctx context.Context) error { if s.Config == nil || s.Config.ConfigStore == nil {