feat: add ConsentCodeStore and OIDCStateStore.PeekInfo

Add ConsentCodeStore to api-gateway for managing one-time consent codes
with TTL, capacity limits, concurrent-safe consumption, and background
eviction. Extend OIDCFlowState with RequestedScopes field and add
non-consuming PeekInfo method to OIDCStateStore for the consent-info
endpoint.

Author: bjcoombs

services/api-gateway/consent_code_store.go

@@ -0,0 +1,129 @@
+package gateway
+
+import (
+	"crypto/rand"
+	"encoding/base64"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+)
+
+const (
+	// consentCodeTTL is how long a consent code remains valid.
+	consentCodeTTL = 2 * time.Minute
+	// consentCodeEvictInterval is how often the store sweeps expired entries.
+	consentCodeEvictInterval = 1 * time.Minute
+	// consentCodeMaxEntries caps entries to prevent memory exhaustion.
+	consentCodeMaxEntries = 10_000
+	// consentCodeBytes is the number of random bytes in a generated consent code.
+	consentCodeBytes = 32
+)
+
+var errConsentCodeStoreFull = errors.New("consent code store is full")
+
+// ConsentCodeEntry holds the state stored alongside a consent code.
+type ConsentCodeEntry struct {
+	Email          string
+	TenantID       string   // UUID from JWT x-tenant-id claim
+	TenantSlug     string   // subdomain slug for cross-validation
+	MCPState       string   // key into OIDCStateStore
+	ClientID       string   // OAuth client_id
+	ApprovedScopes []string // e.g., ["mcp:default"]
+	CreatedAt      time.Time
+}
+
+// ConsentCodeStore is a thread-safe in-memory store for consent codes.
+// Each code can be consumed exactly once and expires after consentCodeTTL.
+type ConsentCodeStore struct {
+	mu        sync.Mutex
+	entries   map[string]ConsentCodeEntry
+	stop      chan struct{}
+	closeOnce sync.Once
+}
+
+// NewConsentCodeStore creates an empty ConsentCodeStore and starts the
+// background eviction goroutine. Call [ConsentCodeStore.Close] to stop it.
+func NewConsentCodeStore() *ConsentCodeStore {
+	s := &ConsentCodeStore{
+		entries: make(map[string]ConsentCodeEntry),
+		stop:    make(chan struct{}),
+	}
+	go s.evictLoop()
+	return s
+}
+
+// Close stops the background eviction goroutine. Safe to call multiple times.
+func (s *ConsentCodeStore) Close() {
+	s.closeOnce.Do(func() { close(s.stop) })
+}
+
+func (s *ConsentCodeStore) evictLoop() {
+	ticker := time.NewTicker(consentCodeEvictInterval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			s.evictExpired()
+		case <-s.stop:
+			return
+		}
+	}
+}
+
+func (s *ConsentCodeStore) evictExpired() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for code, entry := range s.entries {
+		if time.Since(entry.CreatedAt) > consentCodeTTL {
+			delete(s.entries, code)
+		}
+	}
+}
+
+// Store saves a consent code entry and returns the generated code.
+// Returns errConsentCodeStoreFull if the store has reached its capacity limit.
+func (s *ConsentCodeStore) Store(entry ConsentCodeEntry) (string, error) {
+	code, err := generateConsentCode()
+	if err != nil {
+		return "", err
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if len(s.entries) >= consentCodeMaxEntries {
+		return "", errConsentCodeStoreFull
+	}
+	s.entries[code] = entry
+	return code, nil
+}
+
+// Consume atomically retrieves and deletes a consent code.
+// Returns (entry, true) if the code exists and has not expired.
+// Returns (zero, false) if the code is unknown or expired.
+func (s *ConsentCodeStore) Consume(code string) (ConsentCodeEntry, bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	entry, ok := s.entries[code]
+	if !ok {
+		return ConsentCodeEntry{}, false
+	}
+
+	// Always delete (one-time use), even if expired.
+	delete(s.entries, code)
+
+	if time.Since(entry.CreatedAt) > consentCodeTTL {
+		return ConsentCodeEntry{}, false
+	}
+
+	return entry, true
+}
+
+// generateConsentCode returns a cryptographically random, URL-safe consent code.
+func generateConsentCode() (string, error) {
+	b := make([]byte, consentCodeBytes)
+	if _, err := rand.Read(b); err != nil {
+		return "", fmt.Errorf("generate consent code: %w", err)
+	}
+	return base64.RawURLEncoding.EncodeToString(b), nil
+}

services/api-gateway/consent_code_store_test.go

@@ -0,0 +1,175 @@
+package gateway
+
+import (
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func newTestConsentCodeStore(t *testing.T) *ConsentCodeStore {
+	t.Helper()
+	s := NewConsentCodeStore()
+	t.Cleanup(s.Close)
+	return s
+}
+
+func TestConsentCodeStore_StoreAndConsume(t *testing.T) {
+	s := newTestConsentCodeStore(t)
+
+	entry := ConsentCodeEntry{
+		Email:          "alice@example.com",
+		TenantID:       "tid-123",
+		TenantSlug:     "acme",
+		MCPState:       "state-abc",
+		ClientID:       "client-1",
+		ApprovedScopes: []string{"mcp:default"},
+		CreatedAt:      time.Now(),
+	}
+
+	code, err := s.Store(entry)
+	require.NoError(t, err)
+	assert.NotEmpty(t, code)
+
+	// First consume succeeds.
+	got, ok := s.Consume(code)
+	assert.True(t, ok)
+	assert.Equal(t, entry.Email, got.Email)
+	assert.Equal(t, entry.TenantID, got.TenantID)
+	assert.Equal(t, entry.TenantSlug, got.TenantSlug)
+	assert.Equal(t, entry.MCPState, got.MCPState)
+	assert.Equal(t, entry.ClientID, got.ClientID)
+	assert.Equal(t, entry.ApprovedScopes, got.ApprovedScopes)
+
+	// Second consume fails (one-time use).
+	_, ok = s.Consume(code)
+	assert.False(t, ok)
+}
+
+func TestConsentCodeStore_ConsumeExpired(t *testing.T) {
+	s := newTestConsentCodeStore(t)
+
+	entry := ConsentCodeEntry{
+		Email:     "expired@example.com",
+		CreatedAt: time.Now().Add(-consentCodeTTL - time.Second),
+	}
+
+	code, err := s.Store(entry)
+	require.NoError(t, err)
+
+	_, ok := s.Consume(code)
+	assert.False(t, ok, "expired entry should not be consumable")
+
+	// Entry should have been deleted despite being expired.
+	_, ok = s.Consume(code)
+	assert.False(t, ok, "expired entry should be cleaned up after first consume attempt")
+}
+
+func TestConsentCodeStore_ConcurrentConsume(t *testing.T) {
+	s := newTestConsentCodeStore(t)
+
+	entry := ConsentCodeEntry{
+		Email:     "concurrent@example.com",
+		CreatedAt: time.Now(),
+	}
+
+	code, err := s.Store(entry)
+	require.NoError(t, err)
+
+	const goroutines = 50
+	var (
+		wg        sync.WaitGroup
+		successes int32
+		mu        sync.Mutex
+	)
+
+	wg.Add(goroutines)
+	for i := 0; i < goroutines; i++ {
+		go func() {
+			defer wg.Done()
+			_, ok := s.Consume(code)
+			if ok {
+				mu.Lock()
+				successes++
+				mu.Unlock()
+			}
+		}()
+	}
+
+	wg.Wait()
+	assert.Equal(t, int32(1), successes, "exactly one goroutine should consume the code")
+}
+
+func TestConsentCodeStore_CapacityLimit(t *testing.T) {
+	s := newTestConsentCodeStore(t)
+
+	// Fill to capacity.
+	for i := 0; i < consentCodeMaxEntries; i++ {
+		_, err := s.Store(ConsentCodeEntry{
+			Email:     "fill@example.com",
+			CreatedAt: time.Now(),
+		})
+		require.NoError(t, err)
+	}
+
+	// Next store should fail.
+	_, err := s.Store(ConsentCodeEntry{
+		Email:     "overflow@example.com",
+		CreatedAt: time.Now(),
+	})
+	assert.ErrorIs(t, err, errConsentCodeStoreFull)
+}
+
+func TestConsentCodeStore_EvictionLoop(t *testing.T) {
+	// Create store with manual control over eviction (don't use the background loop).
+	s := &ConsentCodeStore{
+		entries: make(map[string]ConsentCodeEntry),
+		stop:    make(chan struct{}),
+	}
+	t.Cleanup(func() { s.closeOnce.Do(func() { close(s.stop) }) })
+
+	// Store an expired entry.
+	s.mu.Lock()
+	s.entries["expired-code"] = ConsentCodeEntry{
+		Email:     "old@example.com",
+		CreatedAt: time.Now().Add(-consentCodeTTL - time.Second),
+	}
+	// Store a valid entry.
+	s.entries["valid-code"] = ConsentCodeEntry{
+		Email:     "new@example.com",
+		CreatedAt: time.Now(),
+	}
+	s.mu.Unlock()
+
+	// Run eviction.
+	s.evictExpired()
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	assert.NotContains(t, s.entries, "expired-code", "expired entry should be evicted")
+	assert.Contains(t, s.entries, "valid-code", "valid entry should remain")
+}
+
+func TestConsentCodeStore_ConsumeNotFound(t *testing.T) {
+	s := newTestConsentCodeStore(t)
+
+	_, ok := s.Consume("nonexistent-code")
+	assert.False(t, ok)
+}
+
+func TestConsentCodeStore_UniqueCodeGeneration(t *testing.T) {
+	s := newTestConsentCodeStore(t)
+
+	codes := make(map[string]struct{})
+	for i := 0; i < 100; i++ {
+		code, err := s.Store(ConsentCodeEntry{
+			Email:     "unique@example.com",
+			CreatedAt: time.Now(),
+		})
+		require.NoError(t, err)
+		assert.NotContains(t, codes, code, "generated codes should be unique")
+		codes[code] = struct{}{}
+	}
+}

services/mcp-server/internal/auth/oidc.go

@@ -110,6 +110,8 @@ type OIDCFlowState struct {
 	DexCodeVerifier string
 	// TenantSlug extracted from the request subdomain.
 	TenantSlug string
+	// RequestedScopes are the OAuth scopes requested by the MCP client.
+	RequestedScopes []string
 	// IssuedAt is when this state was created.
 	IssuedAt time.Time
 }
@@ -191,6 +193,22 @@ func (s *OIDCStateStore) Consume(key string) (OIDCFlowState, bool) {
 	return entry, true
 }
 
+// PeekInfo returns selected fields from an OIDC flow state entry without
+// consuming it. Expired entries are cleaned up and reported as not found.
+func (s *OIDCStateStore) PeekInfo(key string) (clientID, redirectURI string, scopes []string, ok bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	entry, exists := s.entries[key]
+	if !exists {
+		return "", "", nil, false
+	}
+	if time.Since(entry.IssuedAt) > oidcStateTTL {
+		delete(s.entries, key)
+		return "", "", nil, false
+	}
+	return entry.MCPClientID, entry.MCPRedirectURI, entry.RequestedScopes, true
+}
+
 // TenantSlugResolver resolves a tenant slug (e.g., "acme") to its canonical
 // UUID. This ensures the x-tenant-id JWT claim contains a UUID consistent
 // with BFF-issued tokens, not the raw slug string.

services/mcp-server/internal/auth/oidc_test.go

@@ -1242,3 +1242,52 @@ func TestHandleCallback_NoResolver_FallsBackToSlug(t *testing.T) {
 
 	assert.Equal(t, "acme", claims.TenantID, "without resolver, JWT should contain raw slug")
 }
+
+func TestOIDCStateStore_PeekInfo(t *testing.T) {
+	s := newTestOIDCStateStore(t)
+
+	key, err := s.Store(auth.OIDCFlowState{
+		MCPClientID:     "client-1",
+		MCPRedirectURI:  "https://example.com/callback",
+		RequestedScopes: []string{"mcp:default", "mcp:admin"},
+		IssuedAt:        time.Now(),
+	})
+	require.NoError(t, err)
+
+	clientID, redirectURI, scopes, ok := s.PeekInfo(key)
+	assert.True(t, ok)
+	assert.Equal(t, "client-1", clientID)
+	assert.Equal(t, "https://example.com/callback", redirectURI)
+	assert.Equal(t, []string{"mcp:default", "mcp:admin"}, scopes)
+
+	// PeekInfo is non-consuming - a second call should also succeed.
+	_, _, _, ok = s.PeekInfo(key)
+	assert.True(t, ok, "PeekInfo should not consume the entry")
+
+	// Consume should still work after PeekInfo.
+	entry, ok := s.Consume(key)
+	assert.True(t, ok)
+	assert.Equal(t, "client-1", entry.MCPClientID)
+}
+
+func TestOIDCStateStore_PeekInfo_Expired(t *testing.T) {
+	s := newTestOIDCStateStore(t)
+
+	// Store a valid entry, then verify PeekInfo works for non-existent keys
+	// (expired entries are cleaned up internally by the store).
+	clientID, redirectURI, scopes, ok := s.PeekInfo("nonexistent-key")
+	assert.False(t, ok)
+	assert.Empty(t, clientID)
+	assert.Empty(t, redirectURI)
+	assert.Nil(t, scopes)
+}
+
+func TestOIDCStateStore_PeekInfo_NotFound(t *testing.T) {
+	s := newTestOIDCStateStore(t)
+
+	clientID, redirectURI, scopes, ok := s.PeekInfo("missing-key")
+	assert.False(t, ok)
+	assert.Empty(t, clientID)
+	assert.Empty(t, redirectURI)
+	assert.Nil(t, scopes)
+}