mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-18 00:15:05 +00:00
feat(mcp): MCP OAuth 2.1 client for tool servers (#1196)
* feat(mcp): MCP OAuth 2.1 client — full implementation with tests
Implements a complete MCP OAuth 2.1 authorization flow for tool servers that
require user-delegated access, covering all layers from DB to UI.
- discovery.go: RFC 9728 protected-resource → RFC 8414 AS metadata → OIDC
fallback chain with 5-min in-memory cache and InvalidateCache()
- dcr.go: RFC 7591 Dynamic Client Registration with response size guard
- flow.go: PKCE (S256) authorization code flow — StartFlow(), ExchangeCode(),
ClientCredentials(), auto-cleanup of expired flows; carries AS issuer through
PendingFlow for status display
- refresher.go: OAuthTokenProvider with in-memory token cache, automatic refresh
on expiry, per-user vs global slot isolation, InvalidateCache/InvalidateServer
- migrations/000074 + SQLite schema: mcp_oauth_tokens with AES-256-GCM encrypted
access/refresh tokens, partial unique index for global vs per-user rows,
ON DELETE CASCADE from mcp_servers
- store.MCPOAuthTokenStore: Upsert, Get/GetUser, Delete/DeleteUser, and
DeleteServerOAuthTokens (purge all rows for a server)
- PostgreSQL + SQLite implementations
- POST /v1/mcp/oauth/start — discovery + optional DCR + PKCE redirect URL;
client_credentials completes server-side (no redirect) and returns completed=true
- GET /v1/mcp/oauth/callback — exchange code, persist token, publish WS event;
payload built via json.Marshal (no reflected XSS via error_description)
- GET /v1/mcp/oauth/status/{id}, DELETE /v1/mcp/oauth/token/{id} — admin-gated
- POST /v1/mcp/oauth/discover/{id} — on-demand discovery probe
- All outbound calls go through the SSRF-safe client with pinned IPs
- pkg/protocol/mcp_events.go: EventMCPOAuthComplete routed only to the initiating
user (admins in-tenant included); fail-closed across tenants
- getUserMCPTools() injects Authorization: Bearer from OAuthTokenProvider; on a
401 for OAuth servers it purges the cached token so the next turn re-resolves
- handleUpdateServer purges all OAuth tokens (global + per-user), drops the
refresher cache, and evicts the pool when a server's URL or OAuth config
(client_id / endpoints / grant_type / scope / auth_type) changes — so the
status UI and agent never use a token minted for the old resource/AS
- MCPOAuthDialog (WS-driven), unified user-credentials dialog, OAuth settings
fields; handles the no-redirect client_credentials completion
- internal/mcp/oauth/*_test.go: discovery cache, PKCE, DCR, refresher
- internal/http/mcp_oauth_test.go + mcp_update_oauth_purge_test.go: routes, auth
gating, WS event, purge-on-URL/OAuth-config-change
- tests/integration: store + encryption + tenant isolation, E2E start→callback,
DeleteServerOAuthTokens
- internal/gateway/event_filter_test.go, internal/agent/loop_mcp_user_test.go
* fix(mcp): return 400 on OAuth callback with code but missing state
The callback handler rendered a 200 HTML page whenever code or state was
absent. An auth code WITH a missing state is a malformed / CSRF-risk
callback (state is the CSRF token), so reject that case with HTTP 400.
A bare hit with neither code nor state (user opening the URL directly),
provider errors, and exchange failures keep their 200 HTML popup page.
Adds a status code parameter to writeCallbackHTML. Fixes the
TestOAuthCallbackMissingState integration regression while keeping
TestHandleCallbackMissingCodeAndState (no params -> 200) green.
* fix(mcp): scope-based OAuth auth + honor manual OAuth endpoints
Addresses the two MCP/OAuth security-review findings.
Finding 1 — authorization. mcp_oauth_tokens is tenant-scoped, but
start/status/revoke were gated only by requireAuth(RoleAdmin), an RBAC
role check, not tenant membership, so a RoleAdmin caller could act on a
tenant they don't administer. A blanket requireTenantAdmin would have
broken per-user self-service, which the UI exposes (the per-user
MCPUserCredentialsDialog shows an "Authorize" button to regular users for
their own credentials). Instead mirror the existing per-user MCP
credentials model (resolveTargetUserID in mcp_user_credentials.go):
- start/status/revoke accept any authenticated user; each handler calls
authorizeOAuthScope.
- a caller may manage their OWN per-user token (self-service); the
global/server token (user_id="") and other users' tokens require
tenant-admin (owner bypass), so a RoleAdmin that is not a tenant admin
is rejected.
- discover stays admin-only (it only previews AS metadata for a server).
Add a TenantStore dependency. Tests cover self-service, on-behalf-of-
another (403), and global-by-non-tenant-admin (403).
Finding 2 — honor manual OAuth config end-to-end. The UI sent use_dcr /
auth_endpoint / token_endpoint and the update path fingerprinted them for
purge, but handleStart always discovered + DCR'd and ignored them. Now:
- use_dcr=false (a *bool, so legacy/absent stays discover+DCR) skips
discovery/registration and uses the operator endpoints, SSRF-validated.
- token_endpoint is always required; auth_endpoint only for auth-code
grants — client_credentials needs no authorization URL, matching the UI
which hides that field for that grant.
- the refresher already refreshes against the stored token_endpoint and
the callback persists it, so manual-mode tokens refresh correctly.
- oauthFingerprint includes use_dcr (nil normalized to true) so toggling
DCR mode purges stale tokens.
- the web form only serializes manual endpoints when use_dcr is off.
Audited all MCP dialogs (form, global OAuth, per-user credentials, grants,
tools): OAuth dialogs handle completed/auth_url identically and read
config from stored server settings; runtime connect uses the stored token
via the refresher (no re-discovery).
Tests: manual auth-code + client_credentials endpoints, missing/SSRF
endpoints, and the full self/global/on-behalf authorization matrix.
This commit is contained in:
1 parent
21fb1f18ee
commit
0ae55991bb
65 files changed
+7512
-132
No files matched your search
+41
-1
@@ -41,9 +41,11 @@ import (
|
||||
httpapi "github.com/nextlevelbuilder/goclaw/internal/http"
|
||||
kg "github.com/nextlevelbuilder/goclaw/internal/knowledgegraph"
|
||||
mcpbridge "github.com/nextlevelbuilder/goclaw/internal/mcp"
|
||||
mcpoauth "github.com/nextlevelbuilder/goclaw/internal/mcp/oauth"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/media"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/scheduler"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/skills"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
||||
@@ -366,10 +368,16 @@ func runGateway() {
|
||||
server.SetAgentStore(pgStores.Agents)
|
||||
}
|
||||
|
||||
// Build OAuth token refresher before wireExtras so the resolver can inject tokens.
|
||||
var mcpOAuthRefresher mcpbridge.OAuthTokenProvider
|
||||
if pgStores != nil && pgStores.MCPOAuthTokens != nil {
|
||||
mcpOAuthRefresher = mcpoauth.NewRefresher(pgStores.MCPOAuthTokens, security.NewSafeClient(15*time.Second))
|
||||
}
|
||||
|
||||
var mcpPool *mcpbridge.Pool
|
||||
var mediaStore *media.Store
|
||||
var postTurn tools.PostTurnProcessor
|
||||
contextFileInterceptor, mcpPool, mediaStore, postTurn = wireExtras(pgStores, agentRouter, providerRegistry, modelReg, msgBus, pgStores.Sessions, toolsReg, toolPE, skillsLoader, hasMemory, traceCollector, workspace, cfg.Gateway.InjectionAction, cfg, sandboxMgr, redisClient, domainBus, usageCapSvc)
|
||||
contextFileInterceptor, mcpPool, mediaStore, postTurn = wireExtras(pgStores, agentRouter, providerRegistry, modelReg, msgBus, pgStores.Sessions, toolsReg, toolPE, skillsLoader, hasMemory, traceCollector, workspace, cfg.Gateway.InjectionAction, cfg, sandboxMgr, redisClient, domainBus, usageCapSvc, mcpOAuthRefresher)
|
||||
if mcpPool != nil {
|
||||
defer mcpPool.Stop()
|
||||
}
|
||||
@@ -421,6 +429,37 @@ func runGateway() {
|
||||
wakeH.SetPostTurnProcessor(postTurn)
|
||||
}
|
||||
|
||||
// MCP OAuth handler — per-server OAuth 2.1 client flows.
|
||||
var mcpOAuthH *httpapi.MCPOAuthHandler
|
||||
if pgStores != nil && pgStores.MCP != nil && pgStores.MCPOAuthTokens != nil {
|
||||
safeHTTPClient := security.NewSafeClient(15 * time.Second)
|
||||
var oauthRefresher *mcpoauth.Refresher
|
||||
if r, ok := mcpOAuthRefresher.(*mcpoauth.Refresher); ok {
|
||||
oauthRefresher = r
|
||||
}
|
||||
mcpOAuthH = httpapi.NewMCPOAuthHandler(httpapi.MCPOAuthHandlerDeps{
|
||||
MCPStore: pgStores.MCP,
|
||||
OAuthStore: pgStores.MCPOAuthTokens,
|
||||
Discoverer: mcpoauth.NewDiscoverer(safeHTTPClient),
|
||||
FlowMgr: mcpoauth.NewFlowManager(safeHTTPClient),
|
||||
Refresher: oauthRefresher,
|
||||
EventBus: msgBus,
|
||||
PublicURL: cfg.Gateway.PublicURL,
|
||||
Port: cfg.Gateway.Port,
|
||||
TenantStore: pgStores.Tenants,
|
||||
})
|
||||
// Inject OAuth token provider into MCP tools handler so on-demand tool
|
||||
// discovery can authenticate against OAuth-protected MCP servers.
|
||||
if mcpH != nil && mcpOAuthRefresher != nil {
|
||||
mcpH.SetOAuthProvider(mcpOAuthRefresher)
|
||||
}
|
||||
// Inject the OAuth token store so the update handler can purge stale tokens
|
||||
// when a server's URL or OAuth config changes.
|
||||
if mcpH != nil {
|
||||
mcpH.SetOAuthStore(pgStores.MCPOAuthTokens)
|
||||
}
|
||||
}
|
||||
|
||||
// Wire all server.Set*Handler() calls via extracted helper.
|
||||
deps.wireHTTPHandlersOnServer(
|
||||
httpHandlers{
|
||||
@@ -436,6 +475,7 @@ func runGateway() {
|
||||
secureCLI: secureCLIH,
|
||||
secureCLIGrant: secureCLIGrantH,
|
||||
mcpUserCreds: mcpUserCredsH,
|
||||
mcpOAuth: mcpOAuthH,
|
||||
},
|
||||
wakeH,
|
||||
mcpPool,
|
||||
|
||||
@@ -65,7 +65,7 @@ func wireHTTP(stores *store.Stores, defaultWorkspace, dataDir, bundledSkillsDir
|
||||
}
|
||||
var mcpUserCredsH *httpapi.MCPUserCredentialsHandler
|
||||
if stores != nil && stores.MCP != nil {
|
||||
mcpUserCredsH = httpapi.NewMCPUserCredentialsHandler(stores.MCP, stores.Tenants)
|
||||
mcpUserCredsH = httpapi.NewMCPUserCredentialsHandler(stores.MCP, stores.Tenants, msgBus)
|
||||
}
|
||||
|
||||
if stores != nil && stores.ChannelInstances != nil {
|
||||
|
||||
@@ -32,6 +32,7 @@ type httpHandlers struct {
|
||||
secureCLI *httpapi.SecureCLIHandler
|
||||
secureCLIGrant *httpapi.SecureCLIGrantHandler
|
||||
mcpUserCreds *httpapi.MCPUserCredentialsHandler
|
||||
mcpOAuth *httpapi.MCPOAuthHandler
|
||||
}
|
||||
|
||||
// wireHTTPHandlersOnServer registers all HTTP handler objects onto the gateway server.
|
||||
@@ -66,6 +67,12 @@ func (d *gatewayDeps) wireHTTPHandlersOnServer(
|
||||
if h.mcpUserCreds != nil {
|
||||
d.server.SetMCPUserCredentialsHandler(h.mcpUserCreds)
|
||||
}
|
||||
if h.mcpOAuth != nil {
|
||||
if mcpPool != nil {
|
||||
h.mcpOAuth.SetEvictor(mcpPool)
|
||||
}
|
||||
d.server.SetMCPOAuthHandler(h.mcpOAuth)
|
||||
}
|
||||
if h.channelInstances != nil {
|
||||
d.server.SetChannelInstancesHandler(h.channelInstances)
|
||||
}
|
||||
|
||||
@@ -61,6 +61,7 @@ func wireExtras(
|
||||
redisClient any, // nil when built without -tags redis or when Redis is unconfigured
|
||||
domainBus eventbus.DomainEventBus,
|
||||
usageCapSvc *usagecaps.Service,
|
||||
mcpOAuthProvider mcpbridge.OAuthTokenProvider, // nil = OAuth injection disabled
|
||||
) (*tools.ContextFileInterceptor, *mcpbridge.Pool, *media.Store, tools.PostTurnProcessor) {
|
||||
// 1. Build cache instances (in-memory or Redis depending on build tags)
|
||||
agentCtxCache, userCtxCache := makeCaches(redisClient)
|
||||
@@ -243,6 +244,7 @@ func wireExtras(
|
||||
MCPStore: stores.MCP,
|
||||
MCPPool: mcpPool,
|
||||
MCPGrantChecker: mcpGrantChecker,
|
||||
MCPOAuthTokenProvider: mcpOAuthProvider,
|
||||
ConfigPermStore: stores.ConfigPermissions,
|
||||
MediaStore: mediaStore,
|
||||
ModelPricing: appCfg.Telemetry.ModelPricing,
|
||||
@@ -535,7 +537,9 @@ func wireExtras(
|
||||
agentRouter.InvalidateAll()
|
||||
})
|
||||
|
||||
// MCP cache: invalidate all agent caches when MCP servers/grants change
|
||||
// MCP cache: invalidate all agent caches + per-user pool connections when MCP servers/grants change.
|
||||
// Per-user pool connections hold stale credentials/headers; evicting them forces a fresh
|
||||
// AcquireUser on next request so new OAuth tokens and grant changes take effect immediately.
|
||||
msgBus.Subscribe(bus.TopicCacheMCP, func(event bus.Event) {
|
||||
if event.Name != protocol.EventCacheInvalidate {
|
||||
return
|
||||
@@ -545,6 +549,9 @@ func wireExtras(
|
||||
return
|
||||
}
|
||||
agentRouter.InvalidateAll()
|
||||
if mcpPool != nil {
|
||||
mcpPool.EvictAllUsers()
|
||||
}
|
||||
})
|
||||
|
||||
// Cron cache: invalidate job cache on cron changes
|
||||
|
||||
@@ -2,12 +2,15 @@ package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"maps"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
mcpbridge "github.com/nextlevelbuilder/goclaw/internal/mcp"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/tools"
|
||||
)
|
||||
|
||||
@@ -121,13 +124,38 @@ func (l *Loop) getUserMCPTools(ctx context.Context, userID string) []tools.Tool
|
||||
for _, info := range l.mcpUserCredSrvs {
|
||||
srv := info.Server
|
||||
|
||||
// Check if user has credentials for this server
|
||||
uc, err := l.mcpStore.GetUserCredentials(ctx, srv.ID, userID)
|
||||
if err != nil || uc == nil || (uc.APIKey == "" && len(uc.Headers) == 0 && len(uc.Env) == 0) {
|
||||
// OAuth-active servers authenticate ONLY via the user's OAuth token. Static
|
||||
// per-user credentials (api_key/headers/env) are deliberately NOT consulted for
|
||||
// them — otherwise a user who set static creds but never authorized OAuth could
|
||||
// call tools, bypassing OAuth entirely.
|
||||
oauthActive := false
|
||||
if len(srv.Settings) > 0 {
|
||||
var settingsObj struct {
|
||||
OAuth struct {
|
||||
AuthType string `json:"auth_type"`
|
||||
} `json:"oauth"`
|
||||
}
|
||||
_ = json.Unmarshal(srv.Settings, &settingsObj)
|
||||
oauthActive = settingsObj.OAuth.AuthType == "oauth"
|
||||
}
|
||||
|
||||
// Static per-user credentials — only for NON-OAuth servers.
|
||||
var uc *store.MCPUserCredentials
|
||||
hasStaticCreds := false
|
||||
if !oauthActive {
|
||||
if c, e := l.mcpStore.GetUserCredentials(ctx, srv.ID, userID); e == nil && c != nil &&
|
||||
(c.APIKey != "" || len(c.Headers) > 0 || len(c.Env) > 0) {
|
||||
uc = c
|
||||
hasStaticCreds = true
|
||||
}
|
||||
}
|
||||
hasOAuthCreds := oauthActive && l.mcpOAuthTokenProvider != nil
|
||||
|
||||
if !hasStaticCreds && !hasOAuthCreds {
|
||||
continue
|
||||
}
|
||||
|
||||
// Resolve connection params: server defaults merged with user overrides
|
||||
// Resolve connection params: server defaults merged with user overrides.
|
||||
args := mcpbridge.ParseJSONBytesToStringSlice(srv.Args)
|
||||
env := mcpbridge.ParseJSONBytesToStringMap(srv.Env)
|
||||
if env == nil {
|
||||
@@ -138,23 +166,47 @@ func (l *Loop) getUserMCPTools(ctx context.Context, userID string) []tools.Tool
|
||||
headers = make(map[string]string)
|
||||
}
|
||||
|
||||
// Inject server-level API key into headers if present
|
||||
if srv.APIKey != "" && headers["Authorization"] == "" {
|
||||
// Inject server-level API key into headers — ONLY for non-OAuth servers.
|
||||
// For OAuth servers the Authorization MUST come from the user's OAuth token;
|
||||
// falling back to the shared server-level api_key/headers would let a user who
|
||||
// hasn't authorized still call tools with the common credential (isolation leak).
|
||||
if !hasOAuthCreds && srv.APIKey != "" && headers["Authorization"] == "" {
|
||||
headers["Authorization"] = "Bearer " + srv.APIKey
|
||||
}
|
||||
|
||||
// Merge user credentials (user overrides server defaults)
|
||||
if uc.APIKey != "" {
|
||||
headers["Authorization"] = "Bearer " + uc.APIKey
|
||||
// Merge user credentials (user overrides server defaults).
|
||||
if hasStaticCreds {
|
||||
if uc.APIKey != "" {
|
||||
headers["Authorization"] = "Bearer " + uc.APIKey
|
||||
}
|
||||
maps.Copy(headers, uc.Headers)
|
||||
maps.Copy(env, uc.Env)
|
||||
}
|
||||
|
||||
// Inject OAuth Bearer token for OAuth-enabled servers. The token is the ONLY
|
||||
// source of Authorization for these servers — if none is available (user not
|
||||
// authorized yet / refresh failed), skip the server entirely so no tools are
|
||||
// exposed, rather than connecting with the shared server-level credential.
|
||||
if hasOAuthCreds {
|
||||
token, err2 := l.mcpOAuthTokenProvider.GetValidToken(ctx, srv.ID, l.tenantID, userID)
|
||||
if err2 != nil || token == "" {
|
||||
if err2 != nil {
|
||||
slog.Warn("mcp.oauth_token_unavailable", "server", srv.Name, "user", userID, "tenant", l.tenantID, "error", err2)
|
||||
} else {
|
||||
slog.Debug("mcp.oauth_token_empty", "server", srv.Name, "user", userID, "tenant", l.tenantID)
|
||||
}
|
||||
// Strip any inherited Authorization and skip — OAuth not authorized.
|
||||
continue
|
||||
}
|
||||
headers["Authorization"] = "Bearer " + token
|
||||
slog.Debug("mcp.oauth_token_injected", "server", srv.Name, "user", userID, "tenant", l.tenantID)
|
||||
}
|
||||
maps.Copy(headers, uc.Headers)
|
||||
maps.Copy(env, uc.Env)
|
||||
|
||||
// Acquire user-keyed pool connection
|
||||
entry, err := l.mcpPool.AcquireUser(ctx, l.tenantID, srv.Name, userID,
|
||||
srv.Transport, srv.Command, args, env, srv.URL, headers, srv.TimeoutSec)
|
||||
if err != nil {
|
||||
if isUnauthorized401(err) {
|
||||
if isUnauthorized401(err) && uc != nil {
|
||||
expiresAt := strings.TrimSpace(uc.Env["BITRIX_EXPIRES_AT"])
|
||||
expired := false
|
||||
if expiresAt != "" {
|
||||
@@ -174,6 +226,19 @@ func (l *Loop) getUserMCPTools(ctx context.Context, userID string) []tools.Tool
|
||||
_ = l.mcpStore.DeleteUserCredentials(ctx, srv.ID, userID)
|
||||
slog.Warn("mcp.user_credentials_purged", "server", srv.Name, "user", userID, "reason", "unauthorized_401")
|
||||
}
|
||||
// OAuth servers (uc == nil, hasOAuthCreds) carry no static credentials
|
||||
// to purge, but a 401 still means the injected Bearer token was rejected.
|
||||
// Evict the in-memory OAuth token cache so the next turn reloads from the
|
||||
// store (picking up a token refreshed out-of-band) instead of replaying the
|
||||
// same rejected token from cache.
|
||||
if isUnauthorized401(err) && hasOAuthCreds {
|
||||
if inv, ok := l.mcpOAuthTokenProvider.(interface {
|
||||
InvalidateCache(uuid.UUID, string)
|
||||
}); ok {
|
||||
inv.InvalidateCache(srv.ID, userID)
|
||||
slog.Warn("mcp.oauth_token_cache_purged", "server", srv.Name, "user", userID, "tenant", l.tenantID, "reason", "unauthorized_401")
|
||||
}
|
||||
}
|
||||
slog.Warn("mcp.user_pool_acquire_failed", "server", srv.Name, "user", userID, "error", err)
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -1,6 +1,137 @@
|
||||
package agent
|
||||
|
||||
import "testing"
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
mcpbridge "github.com/nextlevelbuilder/goclaw/internal/mcp"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
// ---- mock helpers for getUserMCPTools OAuth tests ----
|
||||
|
||||
// minimalMCPServerStore implements store.MCPServerStore with panics for all
|
||||
// methods except GetUserCredentials, which is the only method called in
|
||||
// getUserMCPTools.
|
||||
type minimalMCPServerStore struct {
|
||||
userCreds *store.MCPUserCredentials
|
||||
userCredsErr error
|
||||
getUserCredsCalls int
|
||||
}
|
||||
|
||||
func (m *minimalMCPServerStore) GetUserCredentials(_ context.Context, _ uuid.UUID, _ string) (*store.MCPUserCredentials, error) {
|
||||
m.getUserCredsCalls++
|
||||
return m.userCreds, m.userCredsErr
|
||||
}
|
||||
|
||||
func (m *minimalMCPServerStore) CreateServer(_ context.Context, _ *store.MCPServerData) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) GetServer(_ context.Context, _ uuid.UUID) (*store.MCPServerData, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) GetServerByName(_ context.Context, _ string) (*store.MCPServerData, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) ListServers(_ context.Context) ([]store.MCPServerData, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) UpdateServer(_ context.Context, _ uuid.UUID, _ map[string]any) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) DeleteServer(_ context.Context, _ uuid.UUID) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) GrantToAgent(_ context.Context, _ *store.MCPAgentGrant) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) RevokeFromAgent(_ context.Context, _, _ uuid.UUID) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) ListAgentGrants(_ context.Context, _ uuid.UUID) ([]store.MCPAgentGrant, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) ListServerGrants(_ context.Context, _ uuid.UUID) ([]store.MCPAgentGrant, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) GrantToUser(_ context.Context, _ *store.MCPUserGrant) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) RevokeFromUser(_ context.Context, _ uuid.UUID, _ string) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) CountAgentGrantsByServer(_ context.Context) (map[uuid.UUID]int, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) ListAccessible(_ context.Context, _ uuid.UUID, _ string) ([]store.MCPAccessInfo, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) CreateRequest(_ context.Context, _ *store.MCPAccessRequest) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) ListPendingRequests(_ context.Context) ([]store.MCPAccessRequest, error) {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) ReviewRequest(_ context.Context, _ uuid.UUID, _ bool, _, _ string) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) SetUserCredentials(_ context.Context, _ uuid.UUID, _ string, _ store.MCPUserCredentials) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
func (m *minimalMCPServerStore) DeleteUserCredentials(_ context.Context, _ uuid.UUID, _ string) error {
|
||||
panic("not implemented")
|
||||
}
|
||||
|
||||
// recordingOAuthProvider records calls to GetValidToken for assertion.
|
||||
type recordingOAuthProvider struct {
|
||||
mu sync.Mutex
|
||||
callCount int
|
||||
returnToken string
|
||||
returnErr error
|
||||
}
|
||||
|
||||
func (r *recordingOAuthProvider) GetValidToken(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ string) (string, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.callCount++
|
||||
return r.returnToken, r.returnErr
|
||||
}
|
||||
|
||||
// newTestPool creates a real *mcpbridge.Pool with short timeouts for unit tests.
|
||||
// The pool's AcquireUser will fail (no real MCP server), but token injection
|
||||
// happens before the pool call.
|
||||
func newTestPool(t *testing.T) *mcpbridge.Pool {
|
||||
t.Helper()
|
||||
pool := mcpbridge.NewPool(mcpbridge.PoolConfig{
|
||||
MaxSize: 10,
|
||||
UserAcquireTimeout: 10 * time.Millisecond,
|
||||
})
|
||||
t.Cleanup(func() { pool.Stop() })
|
||||
return pool
|
||||
}
|
||||
|
||||
// oauthSettings returns JSON settings with auth_type: oauth.
|
||||
func oauthSettings() json.RawMessage {
|
||||
return json.RawMessage(`{"oauth":{"auth_type":"oauth"}}`)
|
||||
}
|
||||
|
||||
// oauthServerData returns a minimal MCPServerData for an OAuth-configured HTTP
|
||||
// server. Transport=http + unreachable URL causes AcquireUser to fail quickly
|
||||
// (connection refused) without panicking — unlike stdio with empty command
|
||||
// which would spawn a goroutine reading from a nil pipe.
|
||||
func oauthServerData() store.MCPServerData {
|
||||
return store.MCPServerData{
|
||||
BaseModel: store.BaseModel{ID: uuid.New()},
|
||||
Name: "oauth-server-" + uuid.New().String()[:8],
|
||||
Transport: "http",
|
||||
URL: "http://127.0.0.1:1", // port 1 = connection refused immediately
|
||||
Settings: oauthSettings(),
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveActorUserID locks the actor-vs-context user-id resolution semantics
|
||||
// that gate per-user MCP credential lookup (and other per-actor resources).
|
||||
@@ -159,3 +290,185 @@ func TestResolveActorUserID(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---- getUserMCPTools OAuth token injection tests ----
|
||||
//
|
||||
// These tests verify that mcpOAuthTokenProvider.GetValidToken is called (or
|
||||
// not) based on whether the MCP server settings declare auth_type=oauth.
|
||||
// Pool.AcquireUser fails (no real MCP server), but GetValidToken is injected
|
||||
// BEFORE the pool call, so we can assert on call count without needing a
|
||||
// working MCP connection.
|
||||
|
||||
// TestGetUserMCPToolsOAuthTokenProviderCalled verifies that when a server's
|
||||
// settings declare auth_type=oauth, GetValidToken is called for the user.
|
||||
func TestGetUserMCPToolsOAuthTokenProviderCalled(t *testing.T) {
|
||||
pool := newTestPool(t)
|
||||
provider := &recordingOAuthProvider{returnToken: "test-bearer-token"}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
l := &Loop{
|
||||
tenantID: uuid.New(),
|
||||
mcpPool: pool,
|
||||
mcpStore: &minimalMCPServerStore{},
|
||||
mcpUserCredSrvs: []store.MCPAccessInfo{{Server: oauthServerData()}},
|
||||
mcpOAuthTokenProvider: provider,
|
||||
}
|
||||
|
||||
_ = l.getUserMCPTools(ctx, "user-1")
|
||||
|
||||
provider.mu.Lock()
|
||||
callCount := provider.callCount
|
||||
provider.mu.Unlock()
|
||||
|
||||
if callCount == 0 {
|
||||
t.Error("expected GetValidToken to be called for OAuth-configured server, but it was not")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetUserMCPToolsOAuthTokenNotCalledForNonOAuth verifies that when a
|
||||
// server has no OAuth settings, the token provider is never invoked and the
|
||||
// server is skipped (neither static nor OAuth creds).
|
||||
func TestGetUserMCPToolsOAuthTokenNotCalledForNonOAuth(t *testing.T) {
|
||||
pool := newTestPool(t)
|
||||
provider := &recordingOAuthProvider{}
|
||||
|
||||
l := &Loop{
|
||||
tenantID: uuid.New(),
|
||||
mcpPool: pool,
|
||||
mcpStore: &minimalMCPServerStore{},
|
||||
mcpUserCredSrvs: []store.MCPAccessInfo{
|
||||
{Server: store.MCPServerData{
|
||||
BaseModel: store.BaseModel{ID: uuid.New()},
|
||||
Name: "plain-server",
|
||||
Transport: "http",
|
||||
URL: "http://127.0.0.1:1",
|
||||
// No settings → auth_type defaults to "" (not "oauth")
|
||||
}},
|
||||
},
|
||||
mcpOAuthTokenProvider: provider,
|
||||
}
|
||||
|
||||
_ = l.getUserMCPTools(context.Background(), "user-1")
|
||||
|
||||
provider.mu.Lock()
|
||||
callCount := provider.callCount
|
||||
provider.mu.Unlock()
|
||||
|
||||
if callCount != 0 {
|
||||
t.Errorf("GetValidToken must not be called for non-OAuth server, called %d times", callCount)
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetUserMCPToolsOAuthNoTokenSkipsServer verifies that when GetValidToken
|
||||
// returns no valid token (expired / not authorized yet), the OAuth server is
|
||||
// SKIPPED entirely — it must NOT fall back to the shared server-level
|
||||
// headers/api_key as Authorization (which would let an unauthorized user call
|
||||
// tools with the common credential). GetValidToken is still consulted (proving
|
||||
// we entered the OAuth branch) but no tools are produced.
|
||||
func TestGetUserMCPToolsOAuthNoTokenSkipsServer(t *testing.T) {
|
||||
pool := newTestPool(t)
|
||||
expiredErr := errTokenExpiredForTest("oauth token expired")
|
||||
provider := &recordingOAuthProvider{
|
||||
returnToken: "",
|
||||
returnErr: expiredErr,
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
// Server carries a server-level API key — must NOT be used as a fallback.
|
||||
srv := oauthServerData()
|
||||
srv.APIKey = "shared-server-key"
|
||||
|
||||
l := &Loop{
|
||||
tenantID: uuid.New(),
|
||||
mcpPool: pool,
|
||||
mcpStore: &minimalMCPServerStore{},
|
||||
mcpUserCredSrvs: []store.MCPAccessInfo{{Server: srv}},
|
||||
mcpOAuthTokenProvider: provider,
|
||||
}
|
||||
|
||||
result := l.getUserMCPTools(ctx, "user-1")
|
||||
if result != nil {
|
||||
t.Errorf("expected nil tools when OAuth has no token (skip, no fallback), got %d tools", len(result))
|
||||
}
|
||||
provider.mu.Lock()
|
||||
calls := provider.callCount
|
||||
provider.mu.Unlock()
|
||||
if calls == 0 {
|
||||
t.Error("expected GetValidToken to be consulted before skipping the OAuth server")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetUserMCPToolsOAuthIgnoresStaticCreds verifies that for an OAuth server
|
||||
// the user's static per-user credentials are NOT consulted or used: even when the
|
||||
// user has set static creds, an OAuth server without a valid token is skipped
|
||||
// (the OAuth token is the sole credential). This closes the bypass where a user
|
||||
// with static headers could call OAuth tools without authorizing.
|
||||
func TestGetUserMCPToolsOAuthIgnoresStaticCreds(t *testing.T) {
|
||||
pool := newTestPool(t)
|
||||
provider := &recordingOAuthProvider{returnToken: ""} // not authorized yet
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
// Store has valid static creds for the user — must be ignored for OAuth servers.
|
||||
st := &minimalMCPServerStore{
|
||||
userCreds: &store.MCPUserCredentials{
|
||||
APIKey: "user-static-key",
|
||||
Headers: map[string]string{"Authorization": "Bearer user-static-key"},
|
||||
},
|
||||
}
|
||||
|
||||
l := &Loop{
|
||||
tenantID: uuid.New(),
|
||||
mcpPool: pool,
|
||||
mcpStore: st,
|
||||
mcpUserCredSrvs: []store.MCPAccessInfo{{Server: oauthServerData()}},
|
||||
mcpOAuthTokenProvider: provider,
|
||||
}
|
||||
|
||||
result := l.getUserMCPTools(ctx, "user-1")
|
||||
if result != nil {
|
||||
t.Errorf("expected nil tools: OAuth server without token must be skipped even with static creds, got %d tools", len(result))
|
||||
}
|
||||
if st.getUserCredsCalls != 0 {
|
||||
t.Errorf("static GetUserCredentials must NOT be consulted for OAuth servers, called %d times", st.getUserCredsCalls)
|
||||
}
|
||||
provider.mu.Lock()
|
||||
calls := provider.callCount
|
||||
provider.mu.Unlock()
|
||||
if calls == 0 {
|
||||
t.Error("expected GetValidToken to be consulted for the OAuth server")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGetUserMCPToolsNilPoolEarlyReturn verifies that a nil pool causes the
|
||||
// function to return immediately without calling the token provider.
|
||||
func TestGetUserMCPToolsNilPoolEarlyReturn(t *testing.T) {
|
||||
provider := &recordingOAuthProvider{}
|
||||
|
||||
l := &Loop{
|
||||
tenantID: uuid.New(),
|
||||
mcpPool: nil, // triggers early return
|
||||
mcpStore: &minimalMCPServerStore{},
|
||||
mcpUserCredSrvs: []store.MCPAccessInfo{{Server: oauthServerData()}},
|
||||
mcpOAuthTokenProvider: provider,
|
||||
}
|
||||
|
||||
result := l.getUserMCPTools(context.Background(), "user-1")
|
||||
if result != nil {
|
||||
t.Errorf("expected nil when pool is nil, got %v", result)
|
||||
}
|
||||
if provider.callCount != 0 {
|
||||
t.Error("GetValidToken must not be called when pool is nil (early return path)")
|
||||
}
|
||||
}
|
||||
|
||||
// errTokenExpiredForTest is a simple error type used to simulate ErrTokenExpired
|
||||
// without importing the oauth package (avoids potential import cycles in tests).
|
||||
type errTokenExpiredForTest string
|
||||
|
||||
func (e errTokenExpiredForTest) Error() string { return string(e) }
|
||||
@@ -139,11 +139,12 @@ type Loop struct {
|
||||
userSetups sync.Map // userID → *userSetup (workspace + seeding state, per Loop instance)
|
||||
|
||||
// Per-user MCP tools: servers requiring user credentials get connected per-request.
|
||||
mcpStore store.MCPServerStore // for credential lookup
|
||||
mcpPool *mcpbridge.Pool // user-keyed connection pool
|
||||
mcpUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds
|
||||
mcpUserTools sync.Map // userID → []tools.Tool (cached per-user tools)
|
||||
mcpGrantChecker mcpbridge.GrantChecker // runtime grant verification (nil = skip)
|
||||
mcpStore store.MCPServerStore // for credential lookup
|
||||
mcpPool *mcpbridge.Pool // user-keyed connection pool
|
||||
mcpUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds
|
||||
mcpUserTools sync.Map // userID → []tools.Tool (cached per-user tools)
|
||||
mcpGrantChecker mcpbridge.GrantChecker // runtime grant verification (nil = skip)
|
||||
mcpOAuthTokenProvider mcpbridge.OAuthTokenProvider // OAuth Bearer token injection (nil = disabled)
|
||||
|
||||
// Compaction config (memory flush settings)
|
||||
compactionCfg *config.CompactionConfig
|
||||
@@ -448,10 +449,11 @@ type LoopConfig struct {
|
||||
MemoryStore store.MemoryStore
|
||||
|
||||
// Per-user MCP tools (servers requiring per-user credentials)
|
||||
MCPStore store.MCPServerStore // for credential lookup
|
||||
MCPPool *mcpbridge.Pool // user-keyed connection pool
|
||||
MCPUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds
|
||||
MCPGrantChecker mcpbridge.GrantChecker // runtime grant verification (nil = skip)
|
||||
MCPStore store.MCPServerStore // for credential lookup
|
||||
MCPPool *mcpbridge.Pool // user-keyed connection pool
|
||||
MCPUserCredSrvs []store.MCPAccessInfo // servers needing per-user creds
|
||||
MCPGrantChecker mcpbridge.GrantChecker // runtime grant verification (nil = skip)
|
||||
MCPOAuthTokenProvider mcpbridge.OAuthTokenProvider // OAuth Bearer token injection (nil = disabled)
|
||||
|
||||
// V3 orchestration mode (resolved by resolver, controls tool visibility)
|
||||
OrchMode OrchestrationMode
|
||||
@@ -594,6 +596,7 @@ func NewLoop(cfg LoopConfig) *Loop {
|
||||
mcpPool: cfg.MCPPool,
|
||||
mcpUserCredSrvs: cfg.MCPUserCredSrvs,
|
||||
mcpGrantChecker: cfg.MCPGrantChecker,
|
||||
mcpOAuthTokenProvider: cfg.MCPOAuthTokenProvider,
|
||||
orchMode: cfg.OrchMode,
|
||||
delegateTargets: cfg.DelegateTargets,
|
||||
evolutionMetricsStore: cfg.EvolutionMetricsStore,
|
||||
|
||||
@@ -84,6 +84,9 @@ type ResolverDeps struct {
|
||||
// MCP grant checker — for runtime grant verification at BridgeTool.Execute
|
||||
MCPGrantChecker mcpbridge.GrantChecker
|
||||
|
||||
// MCP OAuth token provider — injects Bearer tokens for OAuth-enabled MCP servers
|
||||
MCPOAuthTokenProvider mcpbridge.OAuthTokenProvider
|
||||
|
||||
// Skill access store — for per-agent skill visibility filtering
|
||||
SkillAccessStore store.SkillAccessStore
|
||||
SkillStore store.SkillStore
|
||||
@@ -319,6 +322,9 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc {
|
||||
if deps.MCPGrantChecker != nil {
|
||||
mcpOpts = append(mcpOpts, mcpbridge.WithGrantChecker(deps.MCPGrantChecker))
|
||||
}
|
||||
if deps.MCPOAuthTokenProvider != nil {
|
||||
mcpOpts = append(mcpOpts, mcpbridge.WithOAuthTokenProvider(deps.MCPOAuthTokenProvider))
|
||||
}
|
||||
mcpMgr := mcpbridge.NewManager(toolsReg, mcpOpts...)
|
||||
if err := mcpMgr.LoadForAgent(ctx, ag.ID, ""); err != nil {
|
||||
slog.Warn("failed to load MCP servers for agent", "agent", agentKey, "error", err)
|
||||
@@ -543,6 +549,7 @@ func NewManagedResolver(deps ResolverDeps) ResolverFunc {
|
||||
MCPPool: deps.MCPPool,
|
||||
MCPUserCredSrvs: mcpUserCredSrvs,
|
||||
MCPGrantChecker: deps.MCPGrantChecker,
|
||||
MCPOAuthTokenProvider: deps.MCPOAuthTokenProvider,
|
||||
OrchMode: orchMode,
|
||||
DelegateTargets: delegateTargets,
|
||||
EvolutionMetricsStore: evoMetricsStore,
|
||||
|
||||
@@ -428,6 +428,7 @@ type GatewayConfig struct {
|
||||
TaskRecoveryIntervalSec int `json:"task_recovery_interval_sec,omitempty"` // team task recovery ticker interval in seconds (default 300 = 5min)
|
||||
BackgroundProvider string `json:"background_provider,omitempty"` // LLM provider for background workers (vault enrichment, consolidation)
|
||||
BackgroundModel string `json:"background_model,omitempty"` // LLM model for background workers
|
||||
PublicURL string `json:"public_url,omitempty"` // public base URL for OAuth callbacks (e.g. "https://goclaw.example.com")
|
||||
}
|
||||
|
||||
// ToolsConfig controls tool availability, policy, and web search.
|
||||
|
||||
@@ -108,6 +108,15 @@ func clientCanReceiveEvent(c *Client, event bus.Event) bool {
|
||||
return false // non-admin clients don't receive these
|
||||
}
|
||||
|
||||
// MCP OAuth completion: deliver to the user who initiated the flow.
|
||||
// Admins already receive this via the admin check above.
|
||||
if event.Name == protocol.EventMCPOAuthComplete {
|
||||
if uid := extractMapField(event.Payload, "initiatingUserId"); uid != "" {
|
||||
return uid == c.userID
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Exec approval events: scoped to the requesting user.
|
||||
if strings.HasPrefix(event.Name, "exec.approval.") {
|
||||
if uid := extractMapField(event.Payload, "userId"); uid != "" {
|
||||
|
||||
@@ -251,3 +251,70 @@ func TestExtractMapField_StructPayload_JSONFallback(t *testing.T) {
|
||||
t.Errorf("extractMapField JSON fallback = %q, want %q", got, "from-struct")
|
||||
}
|
||||
}
|
||||
|
||||
// ---- MCP OAuth Complete events ----
|
||||
|
||||
func TestClientCanReceiveEvent_MCPOAuthComplete_RoutedToInitiatingUser(t *testing.T) {
|
||||
userA := makeClient(permissions.RoleOperator, "user-a", masterTenant)
|
||||
userB := makeClient(permissions.RoleOperator, "user-b", masterTenant)
|
||||
|
||||
evt := makeEvent(protocol.EventMCPOAuthComplete, masterTenant, map[string]any{
|
||||
"serverId": "server-1",
|
||||
"initiatingUserId": "user-a",
|
||||
"status": "success",
|
||||
})
|
||||
|
||||
if !clientCanReceiveEvent(userA, evt) {
|
||||
t.Error("initiating user should receive mcp.oauth_complete")
|
||||
}
|
||||
if clientCanReceiveEvent(userB, evt) {
|
||||
t.Error("non-initiating user should NOT receive mcp.oauth_complete")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientCanReceiveEvent_MCPOAuthComplete_NotLeakAcrossTenant(t *testing.T) {
|
||||
// Same user ID, different tenants.
|
||||
userAMaster := makeClient(permissions.RoleOperator, "user-a", masterTenant)
|
||||
userAOther := makeClient(permissions.RoleOperator, "user-a", otherTenant)
|
||||
|
||||
evt := makeEvent(protocol.EventMCPOAuthComplete, masterTenant, map[string]any{
|
||||
"serverId": "server-1",
|
||||
"initiatingUserId": "user-a",
|
||||
"status": "success",
|
||||
})
|
||||
|
||||
if !clientCanReceiveEvent(userAMaster, evt) {
|
||||
t.Error("user-a in masterTenant should receive the event they initiated")
|
||||
}
|
||||
if clientCanReceiveEvent(userAOther, evt) {
|
||||
t.Error("user-a in otherTenant must NOT receive masterTenant event (tenant isolation)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientCanReceiveEvent_MCPOAuthComplete_AdminReceivesInTenant(t *testing.T) {
|
||||
admin := makeClient(permissions.RoleAdmin, "admin", masterTenant)
|
||||
|
||||
evt := makeEvent(protocol.EventMCPOAuthComplete, masterTenant, map[string]any{
|
||||
"serverId": "server-1",
|
||||
"initiatingUserId": "other-user",
|
||||
"status": "success",
|
||||
})
|
||||
|
||||
if !clientCanReceiveEvent(admin, evt) {
|
||||
t.Error("admin should receive mcp.oauth_complete events for any user in their tenant")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientCanReceiveEvent_MCPOAuthComplete_NoInitiatingUserIDBlocked(t *testing.T) {
|
||||
c := makeClient(permissions.RoleOperator, "user-x", masterTenant)
|
||||
|
||||
// Event with no initiatingUserId — blocked for non-admin (fail-closed).
|
||||
evt := makeEvent(protocol.EventMCPOAuthComplete, masterTenant, map[string]any{
|
||||
"serverId": "server-1",
|
||||
"status": "success",
|
||||
})
|
||||
|
||||
if clientCanReceiveEvent(c, evt) {
|
||||
t.Error("mcp.oauth_complete without initiatingUserId should be blocked for non-admin (fail-closed)")
|
||||
}
|
||||
}
|
||||
@@ -471,6 +471,7 @@ func (s *Server) SetWakeHandler(h *httpapi.WakeHandler) { s.handlers = append(s.
|
||||
|
||||
// SetMCPHandler sets the MCP server management handler.
|
||||
func (s *Server) SetMCPHandler(h *httpapi.MCPHandler) { s.handlers = append(s.handlers, h) }
|
||||
func (s *Server) SetMCPOAuthHandler(h *httpapi.MCPOAuthHandler) { s.handlers = append(s.handlers, h) }
|
||||
func (s *Server) SetMCPUserCredentialsHandler(h *httpapi.MCPUserCredentialsHandler) {
|
||||
s.handlers = append(s.handlers, h)
|
||||
}
|
||||
|
||||
+118
-9
@@ -1,10 +1,12 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
@@ -24,15 +26,26 @@ type MCPToolLister interface {
|
||||
// MCPPoolEvictor evicts pooled connections for a tenant+server (called on credential rotation).
|
||||
type MCPPoolEvictor interface {
|
||||
Evict(tenantID uuid.UUID, serverName string)
|
||||
// EvictServer evicts both the shared and all per-user connections for a server.
|
||||
// Called when OAuth tokens or grants change so next request reconnects with fresh credentials.
|
||||
EvictServer(tenantID uuid.UUID, serverName string)
|
||||
}
|
||||
|
||||
// MCPHandler handles MCP server management HTTP endpoints.
|
||||
type MCPHandler struct {
|
||||
store store.MCPServerStore
|
||||
msgBus *bus.MessageBus
|
||||
mgr MCPToolLister // optional, nil when Manager not available
|
||||
poolEvictor MCPPoolEvictor // optional, nil when pool not available
|
||||
db *sql.DB // for export/import direct queries
|
||||
store store.MCPServerStore
|
||||
msgBus *bus.MessageBus
|
||||
mgr MCPToolLister // optional, nil when Manager not available
|
||||
poolEvictor MCPPoolEvictor // optional, nil when pool not available
|
||||
db *sql.DB // for export/import direct queries
|
||||
oauthProvider MCPOAuthTokenProvider // optional, nil when OAuth not configured
|
||||
oauthStore store.MCPOAuthTokenStore // optional, nil when OAuth not configured
|
||||
}
|
||||
|
||||
// MCPOAuthTokenProvider retrieves a valid OAuth Bearer token for an MCP server.
|
||||
// Mirrors mcp.OAuthTokenProvider to avoid a circular import.
|
||||
type MCPOAuthTokenProvider interface {
|
||||
GetValidToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) (string, error)
|
||||
}
|
||||
|
||||
// NewMCPHandler creates a handler for MCP server management endpoints.
|
||||
@@ -43,6 +56,15 @@ func NewMCPHandler(s store.MCPServerStore, msgBus *bus.MessageBus, mgr MCPToolLi
|
||||
// SetPoolEvictor sets the pool evictor for credential rotation handling.
|
||||
func (h *MCPHandler) SetPoolEvictor(e MCPPoolEvictor) { h.poolEvictor = e }
|
||||
|
||||
// SetOAuthProvider wires in an OAuth token provider so the tool-list endpoint
|
||||
// can inject a Bearer token when doing on-demand discovery for OAuth servers.
|
||||
func (h *MCPHandler) SetOAuthProvider(p MCPOAuthTokenProvider) { h.oauthProvider = p }
|
||||
|
||||
// SetOAuthStore wires in the OAuth token store so the update handler can purge
|
||||
// stored tokens when a server's URL or OAuth config changes (the old tokens
|
||||
// were minted for a different resource/AS and can no longer be used).
|
||||
func (h *MCPHandler) SetOAuthStore(s store.MCPOAuthTokenStore) { h.oauthStore = s }
|
||||
|
||||
func (h *MCPHandler) emitCacheInvalidate() {
|
||||
if h.msgBus == nil {
|
||||
return
|
||||
@@ -194,6 +216,56 @@ func (h *MCPHandler) handleGetServer(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusOK, srv)
|
||||
}
|
||||
|
||||
// oauthFingerprint captures the MCP-server OAuth settings that, when changed,
|
||||
// invalidate previously minted tokens (different client, AS, grant, or scope).
|
||||
type oauthFingerprint struct {
|
||||
AuthType string
|
||||
ClientID string
|
||||
AuthEndpoint string
|
||||
TokenEndpoint string
|
||||
GrantType string
|
||||
Scope string
|
||||
UseDCR bool // toggling DCR changes how tokens are obtained → must purge
|
||||
}
|
||||
|
||||
// extractOAuthFingerprint parses settings.oauth into a comparable fingerprint.
|
||||
// An empty/absent oauth block yields the zero value, so toggling OAuth on/off is
|
||||
// correctly detected as a change (UpdateServer replaces the whole settings blob).
|
||||
func extractOAuthFingerprint(settings json.RawMessage) oauthFingerprint {
|
||||
if len(settings) == 0 {
|
||||
return oauthFingerprint{}
|
||||
}
|
||||
var s struct {
|
||||
OAuth struct {
|
||||
AuthType string `json:"auth_type"`
|
||||
ClientID string `json:"client_id"`
|
||||
AuthEndpoint string `json:"auth_endpoint"`
|
||||
TokenEndpoint string `json:"token_endpoint"`
|
||||
GrantType string `json:"grant_type"`
|
||||
Scope string `json:"scope"`
|
||||
UseDCR *bool `json:"use_dcr"`
|
||||
} `json:"oauth"`
|
||||
}
|
||||
if err := json.Unmarshal(settings, &s); err != nil {
|
||||
return oauthFingerprint{}
|
||||
}
|
||||
// Absent use_dcr means discover+DCR (matching handleStart), so normalize nil→true
|
||||
// to avoid a spurious purge when a legacy server gains an explicit use_dcr=true.
|
||||
useDCR := true
|
||||
if s.OAuth.UseDCR != nil {
|
||||
useDCR = *s.OAuth.UseDCR
|
||||
}
|
||||
return oauthFingerprint{
|
||||
AuthType: s.OAuth.AuthType,
|
||||
ClientID: s.OAuth.ClientID,
|
||||
AuthEndpoint: s.OAuth.AuthEndpoint,
|
||||
TokenEndpoint: s.OAuth.TokenEndpoint,
|
||||
GrantType: s.OAuth.GrantType,
|
||||
Scope: s.OAuth.Scope,
|
||||
UseDCR: useDCR,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *MCPHandler) handleUpdateServer(w http.ResponseWriter, r *http.Request) {
|
||||
locale := store.LocaleFromContext(r.Context())
|
||||
id, err := uuid.Parse(r.PathValue("id"))
|
||||
@@ -273,14 +345,51 @@ func (h *MCPHandler) handleUpdateServer(w http.ResponseWriter, r *http.Request)
|
||||
return
|
||||
}
|
||||
|
||||
// Evict pool connections when credentials change (force reconnect with new creds)
|
||||
if h.poolEvictor != nil && serverName != "" {
|
||||
// When the URL or OAuth config changes, the previously minted OAuth tokens
|
||||
// were issued for a different resource/AS and are no longer usable. Purge them
|
||||
// (global + per-user), drop the refresher's in-memory cache, and force the pool
|
||||
// to reconnect so a stale Bearer token is never replayed against the new target.
|
||||
tid := store.TenantIDFromContext(r.Context())
|
||||
oauthInvalidated := false
|
||||
if existingSrv != nil {
|
||||
urlChanged := false
|
||||
if newURL, ok := updates["url"].(string); ok && strings.TrimSpace(newURL) != existingSrv.URL {
|
||||
urlChanged = true
|
||||
}
|
||||
oauthChanged := false
|
||||
if rawSettings, ok := updates["settings"]; ok {
|
||||
newSettings, _ := json.Marshal(rawSettings)
|
||||
if extractOAuthFingerprint(newSettings) != extractOAuthFingerprint(existingSrv.Settings) {
|
||||
oauthChanged = true
|
||||
}
|
||||
}
|
||||
if urlChanged || oauthChanged {
|
||||
if h.oauthStore != nil {
|
||||
if err := h.oauthStore.DeleteServerOAuthTokens(r.Context(), id, tid); err != nil {
|
||||
slog.Warn("mcp.oauth_purge_on_config_change_failed", "server_id", id, "error", err)
|
||||
}
|
||||
}
|
||||
if inv, ok := h.oauthProvider.(interface{ InvalidateServer(uuid.UUID) }); ok {
|
||||
inv.InvalidateServer(id)
|
||||
}
|
||||
if h.poolEvictor != nil && serverName != "" {
|
||||
h.poolEvictor.EvictServer(tid, serverName)
|
||||
}
|
||||
oauthInvalidated = true
|
||||
slog.Info("mcp.oauth_tokens_purged_on_config_change", "server_id", id, "url_changed", urlChanged, "oauth_changed", oauthChanged)
|
||||
}
|
||||
}
|
||||
|
||||
// Evict pool connections when credentials change (force reconnect with new creds).
|
||||
// Use EvictServer (shared + all per-user connections) so per-user connections that
|
||||
// inherited the server-level headers/api_key as their base don't keep stale values —
|
||||
// matching the OAuth-config-change path above. Skipped when EvictServer already ran.
|
||||
if !oauthInvalidated && h.poolEvictor != nil && serverName != "" {
|
||||
_, hasKey := updates["api_key"]
|
||||
_, hasHeaders := updates["headers"]
|
||||
_, hasEnv := updates["env"]
|
||||
if hasKey || hasHeaders || hasEnv {
|
||||
tid := store.TenantIDFromContext(r.Context())
|
||||
h.poolEvictor.Evict(tid, serverName)
|
||||
h.poolEvictor.EvictServer(tid, serverName)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,717 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"html"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/bus"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/i18n"
|
||||
mcpoauth "github.com/nextlevelbuilder/goclaw/internal/mcp/oauth"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/permissions"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
|
||||
)
|
||||
|
||||
// MCPOAuthHandler exposes OAuth endpoints for MCP server authentication.
|
||||
type MCPOAuthHandler struct {
|
||||
mcpStore store.MCPServerStore
|
||||
oauthStore store.MCPOAuthTokenStore
|
||||
discoverer *mcpoauth.Discoverer
|
||||
flowMgr *mcpoauth.FlowManager
|
||||
refresher *mcpoauth.Refresher
|
||||
eventBus bus.EventPublisher
|
||||
publicURL string // e.g. "https://goclaw.example.com"
|
||||
port int
|
||||
evictor MCPPoolEvictor
|
||||
// tenantStore gates tenant-scoped writes (token mint/revoke/admin status) with a
|
||||
// tenant-admin membership check — RoleAdmin alone is not a tenant check.
|
||||
tenantStore store.TenantStore
|
||||
}
|
||||
|
||||
// MCPOAuthHandlerDeps contains all dependencies for the OAuth handler.
|
||||
type MCPOAuthHandlerDeps struct {
|
||||
MCPStore store.MCPServerStore
|
||||
OAuthStore store.MCPOAuthTokenStore
|
||||
Discoverer *mcpoauth.Discoverer
|
||||
FlowMgr *mcpoauth.FlowManager
|
||||
Refresher *mcpoauth.Refresher
|
||||
EventBus bus.EventPublisher
|
||||
PublicURL string
|
||||
Port int
|
||||
Evictor MCPPoolEvictor
|
||||
TenantStore store.TenantStore
|
||||
}
|
||||
|
||||
// NewMCPOAuthHandler creates an MCPOAuthHandler.
|
||||
func NewMCPOAuthHandler(deps MCPOAuthHandlerDeps) *MCPOAuthHandler {
|
||||
return &MCPOAuthHandler{
|
||||
mcpStore: deps.MCPStore,
|
||||
oauthStore: deps.OAuthStore,
|
||||
discoverer: deps.Discoverer,
|
||||
flowMgr: deps.FlowMgr,
|
||||
refresher: deps.Refresher,
|
||||
eventBus: deps.EventBus,
|
||||
publicURL: deps.PublicURL,
|
||||
port: deps.Port,
|
||||
evictor: deps.Evictor,
|
||||
tenantStore: deps.TenantStore,
|
||||
}
|
||||
}
|
||||
|
||||
// SetEvictor sets the MCP pool evictor (called after construction when pool is available).
|
||||
func (h *MCPOAuthHandler) SetEvictor(e MCPPoolEvictor) { h.evictor = e }
|
||||
|
||||
// RegisterRoutes registers all MCP OAuth routes on the given mux.
|
||||
func (h *MCPOAuthHandler) RegisterRoutes(mux *http.ServeMux) {
|
||||
// Authorization is scope-based, mirroring the per-user MCP credentials flow
|
||||
// (resolveTargetUserID in mcp_user_credentials.go): start/status/revoke are
|
||||
// open to any authenticated user, and each handler then calls authorizeOAuthScope —
|
||||
// a user may manage their OWN per-user token (self-service, matching the
|
||||
// MCPUserCredentialsDialog "Authorize" button), while the global/server token
|
||||
// (user_id="") and other users' tokens require tenant-admin. RoleAdmin alone is
|
||||
// not a tenant check; SQL scoping (WHERE tenant_id=$N) is the final layer.
|
||||
// discover only previews AS metadata for a server config — keep it admin-only.
|
||||
mux.HandleFunc("POST /v1/mcp/oauth/start", requireAuth("", h.handleStart))
|
||||
mux.HandleFunc("GET /v1/mcp/oauth/callback", h.handleCallback)
|
||||
mux.HandleFunc("GET /v1/mcp/oauth/status/{id}", requireAuth("", h.handleStatus))
|
||||
mux.HandleFunc("DELETE /v1/mcp/oauth/token/{id}", requireAuth("", h.handleRevoke))
|
||||
mux.HandleFunc("POST /v1/mcp/oauth/discover/{id}", requireAuth(permissions.RoleAdmin, h.handleDiscover))
|
||||
}
|
||||
|
||||
// callbackURL derives the OAuth redirect URI from the incoming request.
|
||||
//
|
||||
// Priority:
|
||||
// 1. Config public_url — explicit admin override (wins when set)
|
||||
// 2. X-Forwarded-Proto + X-Forwarded-Host — nginx / reverse proxy
|
||||
// 3. r.Host — the backend's actual host:port (correct for local dev even when
|
||||
// the frontend runs on a different port than the backend)
|
||||
// 4. localhost:port — final fallback
|
||||
//
|
||||
// Using r.Host instead of the Origin header is intentional: Origin carries the
|
||||
// *frontend* origin (e.g. Vite dev server on :5173), which differs from the
|
||||
// backend address that OAuth providers must redirect back to.
|
||||
func (h *MCPOAuthHandler) callbackURL(r *http.Request) string {
|
||||
const path = "/v1/mcp/oauth/callback"
|
||||
|
||||
// 1. Explicit config override (highest priority).
|
||||
if h.publicURL != "" {
|
||||
return strings.TrimRight(h.publicURL, "/") + path
|
||||
}
|
||||
|
||||
// 2. Nginx / reverse-proxy forwarded headers.
|
||||
if fwdHost := r.Header.Get("X-Forwarded-Host"); fwdHost != "" {
|
||||
proto := r.Header.Get("X-Forwarded-Proto")
|
||||
if proto == "" {
|
||||
proto = "https"
|
||||
}
|
||||
return proto + "://" + strings.TrimRight(fwdHost, "/") + path
|
||||
}
|
||||
|
||||
// 3. Backend's own host:port from the request (works for localhost:18790 and
|
||||
// any custom host — correct even in split-port dev setups).
|
||||
if host := r.Host; host != "" {
|
||||
scheme := "http"
|
||||
if r.TLS != nil {
|
||||
scheme = "https"
|
||||
}
|
||||
return scheme + "://" + host + path
|
||||
}
|
||||
|
||||
// 4. Fallback: localhost + configured port.
|
||||
port := h.port
|
||||
if port == 0 {
|
||||
port = 8080
|
||||
}
|
||||
return fmt.Sprintf("http://localhost:%d%s", port, path)
|
||||
}
|
||||
|
||||
// --- POST /v1/mcp/oauth/start ---
|
||||
|
||||
type startOAuthReq struct {
|
||||
ServerID string `json:"server_id"`
|
||||
MCPURL string `json:"mcp_url"`
|
||||
UserID string `json:"user_id,omitempty"` // empty = global token
|
||||
}
|
||||
|
||||
type startOAuthResp struct {
|
||||
AuthURL string `json:"auth_url"`
|
||||
State string `json:"state"`
|
||||
ClientID string `json:"client_id"`
|
||||
Issuer string `json:"issuer"`
|
||||
Completed bool `json:"completed,omitempty"` // true for client_credentials — token already minted, no redirect needed
|
||||
}
|
||||
|
||||
// authorizeOAuthScope enforces who may operate on a given OAuth token scope.
|
||||
// A per-user token (non-empty targetUserID) may be managed by its owner without
|
||||
// admin rights — self-service, matching the per-user MCP credentials flow and the
|
||||
// MCPUserCredentialsDialog "Authorize" button. The global/server token
|
||||
// (targetUserID="") and another user's per-user token require tenant-admin, so a
|
||||
// RoleAdmin who is not a tenant admin cannot mint/revoke tenant-scoped tokens.
|
||||
func (h *MCPOAuthHandler) authorizeOAuthScope(w http.ResponseWriter, r *http.Request, targetUserID string) bool {
|
||||
// Self-service: a caller may manage their own per-user token. Only applies to a
|
||||
// non-empty target (the global/server token is shared, never "self").
|
||||
if targetUserID != "" {
|
||||
if callerID := store.UserIDFromContext(r.Context()); callerID != "" && targetUserID == callerID {
|
||||
return true
|
||||
}
|
||||
}
|
||||
// Global token, or a per-user token on behalf of another user → tenant-admin.
|
||||
return requireTenantAdmin(w, r, h.tenantStore)
|
||||
}
|
||||
|
||||
func (h *MCPOAuthHandler) handleStart(w http.ResponseWriter, r *http.Request) {
|
||||
var req startOAuthReq
|
||||
if !bindJSON(w, r, store.LocaleFromContext(r.Context()), &req) {
|
||||
return
|
||||
}
|
||||
if !h.authorizeOAuthScope(w, r, req.UserID) {
|
||||
return
|
||||
}
|
||||
if req.ServerID == "" || req.MCPURL == "" {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "server_id and mcp_url are required"})
|
||||
return
|
||||
}
|
||||
|
||||
serverID, err := uuid.Parse(req.ServerID)
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid server_id"})
|
||||
return
|
||||
}
|
||||
|
||||
ctx := r.Context()
|
||||
tenantID := store.TenantIDFromContext(ctx)
|
||||
|
||||
// Verify server belongs to tenant.
|
||||
srv, err := h.mcpStore.GetServer(ctx, serverID)
|
||||
if err != nil || srv == nil {
|
||||
writeJSON(w, http.StatusNotFound, map[string]string{"error": "MCP server not found"})
|
||||
return
|
||||
}
|
||||
|
||||
// Parse server OAuth settings up front — they decide whether to auto-discover the
|
||||
// AS + dynamically register (use_dcr) or to use operator-supplied endpoints.
|
||||
// use_dcr is a pointer so a legacy/absent value (nil) keeps the discover+DCR path.
|
||||
var oauthSettings struct {
|
||||
OAuth struct {
|
||||
UseDCR *bool `json:"use_dcr"`
|
||||
AuthEndpoint string `json:"auth_endpoint"`
|
||||
TokenEndpoint string `json:"token_endpoint"`
|
||||
ClientID string `json:"client_id"`
|
||||
ClientSecret string `json:"client_secret"`
|
||||
Scope string `json:"scope"`
|
||||
GrantType string `json:"grant_type"`
|
||||
} `json:"oauth"`
|
||||
}
|
||||
if len(srv.Settings) > 0 {
|
||||
_ = json.Unmarshal(srv.Settings, &oauthSettings)
|
||||
}
|
||||
|
||||
safeClient := security.NewSafeClient(15 * time.Second)
|
||||
callbackURI := h.callbackURL(r)
|
||||
|
||||
var disc *mcpoauth.DiscoveryResult
|
||||
if oauthSettings.OAuth.UseDCR != nil && !*oauthSettings.OAuth.UseDCR {
|
||||
// Manual mode: the operator disabled DCR and supplied the AS endpoints
|
||||
// directly. We do NOT auto-discover or dynamically register. The endpoints
|
||||
// are SSRF-validated because the token endpoint is dialed server-side and the
|
||||
// authorization endpoint is handed to the browser. A manual client_id is
|
||||
// required (enforced by the no-client_id guard below, since no DCR runs).
|
||||
ae := strings.TrimSpace(oauthSettings.OAuth.AuthEndpoint)
|
||||
te := strings.TrimSpace(oauthSettings.OAuth.TokenEndpoint)
|
||||
// The token endpoint is always required (it is dialed for every grant).
|
||||
// The authorization endpoint is only needed for browser/auth-code grants —
|
||||
// client_credentials has no authorization redirect, matching the UI which
|
||||
// hides the authorization URL field for that grant.
|
||||
if te == "" {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "manual OAuth (use_dcr=false) requires token_endpoint"})
|
||||
return
|
||||
}
|
||||
if oauthSettings.OAuth.GrantType != "client_credentials" && ae == "" {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "manual OAuth (use_dcr=false) requires auth_endpoint for this grant type"})
|
||||
return
|
||||
}
|
||||
// SSRF-validate the token endpoint (always dialed) and the authorization
|
||||
// endpoint when present (handed to the browser).
|
||||
if _, _, verr := security.Validate(te); verr != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid token_endpoint: " + verr.Error()})
|
||||
return
|
||||
}
|
||||
if ae != "" {
|
||||
if _, _, verr := security.Validate(ae); verr != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid auth_endpoint: " + verr.Error()})
|
||||
return
|
||||
}
|
||||
}
|
||||
disc = &mcpoauth.DiscoveryResult{AuthorizationEndpoint: ae, TokenEndpoint: te}
|
||||
} else {
|
||||
// Discovery path. Discovery and the resulting resource indicator must bind to
|
||||
// the server's own registered URL, not a client-supplied one. Trusting
|
||||
// req.MCPURL would let the caller mint a token whose resource_uri / AS differ
|
||||
// from the server the agent actually connects to. Fall back to req.MCPURL only
|
||||
// when the server has no stored URL (legacy rows).
|
||||
mcpURL := srv.URL
|
||||
if mcpURL == "" {
|
||||
mcpURL = req.MCPURL
|
||||
}
|
||||
// Reuse the shared, cache-backed discoverer when wired (5-min metadata cache);
|
||||
// fall back to a request-local one only when none was injected (e.g. tests).
|
||||
discoverer := h.discoverer
|
||||
if discoverer == nil {
|
||||
discoverer = mcpoauth.NewDiscoverer(safeClient)
|
||||
}
|
||||
var derr error
|
||||
disc, derr = discoverer.Discover(ctx, mcpURL)
|
||||
if derr != nil {
|
||||
slog.Warn("mcpoauth.discover_failed", "server_id", serverID, "url", mcpURL, "error", derr)
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "OAuth discovery failed: " + derr.Error()})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Determine client credentials — priority:
|
||||
// 1. Manual client_id in server settings (DCR OFF): use as-is.
|
||||
// 2. DCR: always register fresh — avoids reusing stale credentials when MCP URL changes.
|
||||
// 3. Fallback: existing token credentials (legacy / no-DCR servers without manual config).
|
||||
clientID := oauthSettings.OAuth.ClientID
|
||||
clientSecret := oauthSettings.OAuth.ClientSecret
|
||||
|
||||
if clientID == "" {
|
||||
if disc.RegistrationEndpoint != "" {
|
||||
dcrResp, err2 := mcpoauth.RegisterClient(ctx, safeClient, disc.RegistrationEndpoint, callbackURI)
|
||||
if err2 != nil {
|
||||
slog.Warn("mcpoauth.dcr_failed", "server_id", serverID, "error", err2)
|
||||
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "Dynamic Client Registration failed: " + err2.Error()})
|
||||
return
|
||||
}
|
||||
clientID = dcrResp.ClientID
|
||||
clientSecret = dcrResp.ClientSecret
|
||||
} else {
|
||||
// No DCR endpoint — fall back to existing token's credentials if available.
|
||||
var existingTok *store.MCPOAuthToken
|
||||
if req.UserID == "" {
|
||||
existingTok, _ = h.oauthStore.GetOAuthToken(ctx, serverID, tenantID)
|
||||
} else {
|
||||
existingTok, _ = h.oauthStore.GetUserOAuthToken(ctx, serverID, tenantID, req.UserID)
|
||||
}
|
||||
if existingTok != nil {
|
||||
clientID = existingTok.DCRClientID
|
||||
clientSecret = existingTok.DCRClientSecret
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if clientID == "" {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "no client_id: DCR not supported and no manual credentials configured"})
|
||||
return
|
||||
}
|
||||
|
||||
// Scope: use manual setting, fall back to AS-advertised scopes when DCR is used.
|
||||
scopes := oauthSettings.OAuth.Scope
|
||||
if scopes == "" && len(disc.ScopesSupported) > 0 {
|
||||
scopes = strings.Join(disc.ScopesSupported, " ")
|
||||
}
|
||||
|
||||
// client_credentials is a non-interactive grant: there is no browser redirect.
|
||||
// Obtain the token server-side now and persist it, returning Completed=true so
|
||||
// the UI skips the popup. Routing it through StartFlow (authorization-code) would
|
||||
// produce a bogus response_type=code authorization URL.
|
||||
if oauthSettings.OAuth.GrantType == "client_credentials" {
|
||||
if clientSecret == "" {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "client_credentials grant requires a client_secret"})
|
||||
return
|
||||
}
|
||||
tokens, ccErr := h.flowMgr.ClientCredentials(ctx, disc.TokenEndpoint, clientID, clientSecret, scopes, disc.ResourceURI)
|
||||
if ccErr != nil {
|
||||
slog.Warn("mcpoauth.client_credentials_failed", "server_id", serverID, "error", ccErr)
|
||||
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "client credentials grant failed: " + ccErr.Error()})
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
var expiresAt *time.Time
|
||||
if tokens.ExpiresIn > 0 {
|
||||
t := now.Add(time.Duration(tokens.ExpiresIn) * time.Second)
|
||||
expiresAt = &t
|
||||
}
|
||||
tok := &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
UserID: req.UserID,
|
||||
AccessToken: tokens.AccessToken,
|
||||
RefreshToken: tokens.RefreshToken,
|
||||
TokenType: tokens.TokenType,
|
||||
Scopes: tokens.Scope,
|
||||
ExpiresAt: expiresAt,
|
||||
IssuedAt: &now,
|
||||
DCRClientID: clientID,
|
||||
DCRClientSecret: clientSecret,
|
||||
DCRIssuer: disc.Issuer,
|
||||
TokenEndpoint: disc.TokenEndpoint,
|
||||
ResourceURI: disc.ResourceURI,
|
||||
}
|
||||
if err := h.activateToken(ctx, tok, srv.Name); err != nil {
|
||||
slog.Error("mcpoauth.save_token", "error", err, "server_id", serverID)
|
||||
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "failed to save token"})
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, startOAuthResp{Completed: true, ClientID: clientID, Issuer: disc.Issuer})
|
||||
return
|
||||
}
|
||||
|
||||
authURL, state, err := h.flowMgr.StartFlow(ctx, mcpoauth.StartFlowParams{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
UserID: req.UserID,
|
||||
InitiatingUserID: store.UserIDFromContext(ctx),
|
||||
DiscoveryResult: disc,
|
||||
ClientID: clientID,
|
||||
ClientSecret: clientSecret,
|
||||
RedirectURI: callbackURI,
|
||||
Scopes: scopes,
|
||||
GrantType: oauthSettings.OAuth.GrantType,
|
||||
})
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusOK, startOAuthResp{
|
||||
AuthURL: authURL,
|
||||
State: state,
|
||||
ClientID: clientID,
|
||||
Issuer: disc.Issuer,
|
||||
})
|
||||
}
|
||||
|
||||
// --- GET /v1/mcp/oauth/callback ---
|
||||
|
||||
func (h *MCPOAuthHandler) handleCallback(w http.ResponseWriter, r *http.Request) {
|
||||
code := r.URL.Query().Get("code")
|
||||
state := r.URL.Query().Get("state")
|
||||
oauthErr := r.URL.Query().Get("error")
|
||||
oauthErrDesc := r.URL.Query().Get("error_description")
|
||||
|
||||
locale := extractLocale(r)
|
||||
|
||||
if oauthErr != "" {
|
||||
h.writeCallbackHTML(w, http.StatusOK, locale, false, oauthErrDesc)
|
||||
return
|
||||
}
|
||||
if code != "" && state == "" {
|
||||
// An auth code with no state is a malformed / CSRF-risk callback (state is
|
||||
// the CSRF token), so reject with 400. A bare hit with neither code nor state
|
||||
// (e.g. the user opening the URL directly) falls through to the 200 HTML page.
|
||||
h.writeCallbackHTML(w, http.StatusBadRequest, locale, false, "missing state")
|
||||
return
|
||||
}
|
||||
if code == "" || state == "" {
|
||||
h.writeCallbackHTML(w, http.StatusOK, locale, false, "missing code or state")
|
||||
return
|
||||
}
|
||||
|
||||
ctx := r.Context()
|
||||
tokens, flow, err := h.flowMgr.ExchangeCode(ctx, state, code)
|
||||
if err != nil {
|
||||
slog.Warn("mcpoauth.exchange_failed", "error", err)
|
||||
h.writeCallbackHTML(w, http.StatusOK, locale, false, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
var expiresAt *time.Time
|
||||
if tokens.ExpiresIn > 0 {
|
||||
t := now.Add(time.Duration(tokens.ExpiresIn) * time.Second)
|
||||
expiresAt = &t
|
||||
}
|
||||
|
||||
tok := &store.MCPOAuthToken{
|
||||
ServerID: flow.ServerID,
|
||||
TenantID: flow.TenantID,
|
||||
UserID: flow.UserID,
|
||||
AccessToken: tokens.AccessToken,
|
||||
RefreshToken: tokens.RefreshToken,
|
||||
TokenType: tokens.TokenType,
|
||||
Scopes: tokens.Scope,
|
||||
ExpiresAt: expiresAt,
|
||||
IssuedAt: &now,
|
||||
DCRClientID: flow.ClientID,
|
||||
DCRClientSecret: flow.ClientSecret,
|
||||
DCRIssuer: flow.Issuer,
|
||||
TokenEndpoint: flow.TokenEndpoint,
|
||||
ResourceURI: flow.ResourceURI,
|
||||
}
|
||||
srv, _ := h.mcpStore.GetServer(ctx, flow.ServerID)
|
||||
serverName := ""
|
||||
if srv != nil {
|
||||
serverName = srv.Name
|
||||
}
|
||||
if err := h.activateToken(ctx, tok, serverName); err != nil {
|
||||
slog.Error("mcpoauth.save_token", "error", err, "server_id", flow.ServerID)
|
||||
h.publishOAuthComplete(flow, "error", "failed to save token")
|
||||
h.writeCallbackHTML(w, http.StatusOK, locale, false, "failed to save token")
|
||||
return
|
||||
}
|
||||
|
||||
// Notify the initiating WS client that OAuth completed successfully.
|
||||
h.publishOAuthComplete(flow, "success", "")
|
||||
h.writeCallbackHTML(w, http.StatusOK, locale, true, "")
|
||||
}
|
||||
|
||||
// activateToken persists a freshly obtained OAuth token and makes it live:
|
||||
// it upserts the row, evicts the refresher's in-memory cache, and drops the
|
||||
// pooled MCP connection (shared + per-user) so the next request reconnects
|
||||
// with the new Bearer token. Shared by the authorization-code callback and
|
||||
// the client_credentials path in handleStart.
|
||||
func (h *MCPOAuthHandler) activateToken(ctx context.Context, tok *store.MCPOAuthToken, serverName string) error {
|
||||
if err := h.oauthStore.UpsertOAuthToken(ctx, tok); err != nil {
|
||||
return err
|
||||
}
|
||||
if h.refresher != nil {
|
||||
h.refresher.InvalidateCache(tok.ServerID, tok.UserID)
|
||||
}
|
||||
if h.evictor != nil && serverName != "" {
|
||||
h.evictor.EvictServer(tok.TenantID, serverName)
|
||||
}
|
||||
// Broadcast an MCP cache-invalidate so all per-user pool connections are
|
||||
// evicted and agent Loop caches reload — picks up the new token for both the
|
||||
// shared (global) and per-user paths on the next request.
|
||||
h.emitMCPCacheInvalidate()
|
||||
return nil
|
||||
}
|
||||
|
||||
// emitMCPCacheInvalidate broadcasts an MCP cache-invalidate event. The subscriber
|
||||
// (cmd/gateway_managed.go) responds with agentRouter.InvalidateAll() +
|
||||
// mcpPool.EvictAllUsers(), so freshly authorized/revoked OAuth tokens take effect
|
||||
// immediately instead of waiting for idle TTL. Mirrors MCPHandler.emitCacheInvalidate.
|
||||
func (h *MCPOAuthHandler) emitMCPCacheInvalidate() {
|
||||
if h.eventBus == nil {
|
||||
return
|
||||
}
|
||||
h.eventBus.Broadcast(bus.Event{
|
||||
Name: protocol.EventCacheInvalidate,
|
||||
Payload: bus.CacheInvalidatePayload{Kind: bus.CacheKindMCP},
|
||||
})
|
||||
}
|
||||
|
||||
func (h *MCPOAuthHandler) writeCallbackHTML(w http.ResponseWriter, statusCode int, locale string, success bool, errMsg string) {
|
||||
status := "success"
|
||||
if !success {
|
||||
status = "error"
|
||||
}
|
||||
|
||||
var bodyMsg string
|
||||
if success {
|
||||
bodyMsg = i18n.T(locale, i18n.MsgOAuthCallbackSuccess)
|
||||
} else {
|
||||
bodyMsg = i18n.T(locale, i18n.MsgOAuthCallbackFailed)
|
||||
}
|
||||
|
||||
// Build the postMessage payload with json.Marshal. errMsg is attacker-controlled
|
||||
// (reflected from the OAuth provider's `error_description` query param on this
|
||||
// unauthenticated callback endpoint), so it must never be concatenated into the
|
||||
// page. encoding/json escapes <, >, & to </>/& by default, which
|
||||
// neutralises </script> breakout attempts inside the <script> block (prevents
|
||||
// reflected XSS). bodyMsg comes from the trusted i18n catalog but is HTML-escaped
|
||||
// for defence in depth before landing in the <p> element.
|
||||
payload := struct {
|
||||
Type string `json:"type"`
|
||||
Status string `json:"status"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}{Type: "mcp-oauth-complete", Status: status, Error: errMsg}
|
||||
msgJSON, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
msgJSON = []byte(`{"type":"mcp-oauth-complete","status":"error"}`)
|
||||
}
|
||||
|
||||
htmlPage := fmt.Sprintf(`<!DOCTYPE html><html><body>
|
||||
<script>
|
||||
(function(){
|
||||
var msg = %s;
|
||||
try {
|
||||
var bc = new BroadcastChannel('mcp-oauth');
|
||||
bc.postMessage(msg);
|
||||
bc.close();
|
||||
} catch(e) {}
|
||||
if(window.opener){window.opener.postMessage(msg,'*');}
|
||||
setTimeout(function(){window.close();},1000);
|
||||
})();
|
||||
</script>
|
||||
<p>%s</p>
|
||||
</body></html>`, msgJSON, html.EscapeString(bodyMsg))
|
||||
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
w.WriteHeader(statusCode)
|
||||
_, _ = fmt.Fprint(w, htmlPage)
|
||||
}
|
||||
|
||||
// publishOAuthComplete broadcasts an mcp.oauth_complete WS event to the initiating client.
|
||||
// Only called when we have a resolved pendingFlow (i.e. after ExchangeCode succeeds).
|
||||
// Error cases before ExchangeCode (missing state, provider error) have no flow and are
|
||||
// handled client-side by the pollClosed interval.
|
||||
func (h *MCPOAuthHandler) publishOAuthComplete(flow *mcpoauth.PendingFlow, status, errMsg string) {
|
||||
if h.eventBus == nil {
|
||||
return
|
||||
}
|
||||
bus.BroadcastForTenant(h.eventBus, protocol.EventMCPOAuthComplete, flow.TenantID,
|
||||
protocol.MCPOAuthCompletePayload{
|
||||
ServerID: flow.ServerID.String(),
|
||||
UserID: flow.UserID,
|
||||
InitiatingUserID: flow.InitiatingUserID,
|
||||
Status: status,
|
||||
Error: errMsg,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// --- GET /v1/mcp/oauth/status/{id} ---
|
||||
|
||||
type oauthStatusResp struct {
|
||||
HasToken bool `json:"has_token"`
|
||||
ClientID string `json:"client_id,omitempty"`
|
||||
Issuer string `json:"issuer,omitempty"`
|
||||
ExpiresAt *string `json:"expires_at,omitempty"`
|
||||
Expired bool `json:"expired,omitempty"`
|
||||
}
|
||||
|
||||
func (h *MCPOAuthHandler) handleStatus(w http.ResponseWriter, r *http.Request) {
|
||||
if !h.authorizeOAuthScope(w, r, r.URL.Query().Get("user_id")) {
|
||||
return
|
||||
}
|
||||
ctx := r.Context()
|
||||
tenantID := store.TenantIDFromContext(ctx)
|
||||
serverID, err := uuid.Parse(r.PathValue("id"))
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid server id"})
|
||||
return
|
||||
}
|
||||
userID := r.URL.Query().Get("user_id")
|
||||
|
||||
var tok *store.MCPOAuthToken
|
||||
if userID == "" {
|
||||
tok, err = h.oauthStore.GetOAuthToken(ctx, serverID, tenantID)
|
||||
} else {
|
||||
tok, err = h.oauthStore.GetUserOAuthToken(ctx, serverID, tenantID, userID)
|
||||
}
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
if tok == nil {
|
||||
writeJSON(w, http.StatusOK, oauthStatusResp{HasToken: false})
|
||||
return
|
||||
}
|
||||
|
||||
resp := oauthStatusResp{
|
||||
HasToken: true,
|
||||
ClientID: tok.DCRClientID,
|
||||
Issuer: tok.DCRIssuer,
|
||||
}
|
||||
if tok.ExpiresAt != nil {
|
||||
s := tok.ExpiresAt.UTC().Format(time.RFC3339)
|
||||
resp.ExpiresAt = &s
|
||||
resp.Expired = time.Now().After(*tok.ExpiresAt)
|
||||
}
|
||||
writeJSON(w, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
// --- DELETE /v1/mcp/oauth/token/{id} ---
|
||||
|
||||
func (h *MCPOAuthHandler) handleRevoke(w http.ResponseWriter, r *http.Request) {
|
||||
if !h.authorizeOAuthScope(w, r, r.URL.Query().Get("user_id")) {
|
||||
return
|
||||
}
|
||||
ctx := r.Context()
|
||||
tenantID := store.TenantIDFromContext(ctx)
|
||||
serverID, err := uuid.Parse(r.PathValue("id"))
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid server id"})
|
||||
return
|
||||
}
|
||||
userID := r.URL.Query().Get("user_id")
|
||||
|
||||
if userID == "" {
|
||||
err = h.oauthStore.DeleteOAuthToken(ctx, serverID, tenantID)
|
||||
} else {
|
||||
err = h.oauthStore.DeleteUserOAuthToken(ctx, serverID, tenantID, userID)
|
||||
}
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// Invalidate cached token.
|
||||
if h.refresher != nil {
|
||||
h.refresher.InvalidateCache(serverID, userID)
|
||||
}
|
||||
|
||||
// Evict pool connection so next request doesn't reuse the revoked token.
|
||||
if h.evictor != nil {
|
||||
srv, _ := h.mcpStore.GetServer(r.Context(), serverID)
|
||||
if srv != nil {
|
||||
h.evictor.EvictServer(tenantID, srv.Name)
|
||||
}
|
||||
}
|
||||
|
||||
// Broadcast MCP cache-invalidate → EvictAllUsers + agent reload so the revoked
|
||||
// token is dropped everywhere and OAuth servers re-gate (hidden until re-auth).
|
||||
h.emitMCPCacheInvalidate()
|
||||
|
||||
writeJSON(w, http.StatusOK, map[string]bool{"ok": true})
|
||||
}
|
||||
|
||||
// --- POST /v1/mcp/oauth/discover/{id} ---
|
||||
|
||||
type discoverResp struct {
|
||||
Issuer string `json:"issuer,omitempty"`
|
||||
AuthorizationEndpoint string `json:"authorization_endpoint"`
|
||||
TokenEndpoint string `json:"token_endpoint"`
|
||||
RegistrationEndpoint string `json:"registration_endpoint,omitempty"`
|
||||
ScopesSupported []string `json:"scopes_supported,omitempty"`
|
||||
}
|
||||
|
||||
func (h *MCPOAuthHandler) handleDiscover(w http.ResponseWriter, r *http.Request) {
|
||||
if !requireTenantAdmin(w, r, h.tenantStore) {
|
||||
return
|
||||
}
|
||||
ctx := r.Context()
|
||||
serverID, err := uuid.Parse(r.PathValue("id"))
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid server id"})
|
||||
return
|
||||
}
|
||||
srv, err := h.mcpStore.GetServer(ctx, serverID)
|
||||
if err != nil || srv == nil {
|
||||
writeJSON(w, http.StatusNotFound, map[string]string{"error": "MCP server not found"})
|
||||
return
|
||||
}
|
||||
if srv.URL == "" {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "server has no URL"})
|
||||
return
|
||||
}
|
||||
|
||||
discoverer := h.discoverer
|
||||
if discoverer == nil {
|
||||
discoverer = mcpoauth.NewDiscoverer(security.NewSafeClient(15 * time.Second))
|
||||
}
|
||||
disc, err := discoverer.Discover(ctx, srv.URL)
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusBadGateway, map[string]string{"error": "discovery failed: " + err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusOK, discoverResp{
|
||||
Issuer: disc.Issuer,
|
||||
AuthorizationEndpoint: disc.AuthorizationEndpoint,
|
||||
TokenEndpoint: disc.TokenEndpoint,
|
||||
RegistrationEndpoint: disc.RegistrationEndpoint,
|
||||
ScopesSupported: disc.ScopesSupported,
|
||||
})
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -16,6 +16,7 @@ import (
|
||||
func (h *MCPHandler) handleTestConnection(w http.ResponseWriter, r *http.Request) {
|
||||
locale := store.LocaleFromContext(r.Context())
|
||||
var req struct {
|
||||
ServerID *uuid.UUID `json:"server_id,omitempty"`
|
||||
Transport string `json:"transport"`
|
||||
Command string `json:"command"`
|
||||
Args []string `json:"args"`
|
||||
@@ -39,6 +40,25 @@ func (h *MCPHandler) handleTestConnection(w http.ResponseWriter, r *http.Request
|
||||
return
|
||||
}
|
||||
|
||||
// For an OAuth server, "test OAuth" means test with the OAuth token only:
|
||||
// Authorization comes solely from the GLOBAL token — never from a body header
|
||||
// fallback. Non-OAuth servers keep testing with the body headers/env as provided.
|
||||
if req.ServerID != nil && h.store != nil {
|
||||
if srv, err2 := h.store.GetServer(r.Context(), *req.ServerID); err2 == nil && srv != nil && mcpbridge.IsOAuthActive(srv.Settings) {
|
||||
// Drop any body-supplied Authorization — no fallback for OAuth servers.
|
||||
delete(req.Headers, "Authorization")
|
||||
if h.oauthProvider != nil {
|
||||
tenantID := store.TenantIDFromContext(r.Context())
|
||||
if token, err3 := h.oauthProvider.GetValidToken(r.Context(), *req.ServerID, tenantID, ""); err3 == nil && token != "" {
|
||||
if req.Headers == nil {
|
||||
req.Headers = make(map[string]string)
|
||||
}
|
||||
req.Headers["Authorization"] = "Bearer " + token
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tools, err := mcpbridge.DiscoverTools(r.Context(), req.Transport, req.Command, req.Args, req.Env, req.URL, req.Headers)
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
@@ -88,11 +108,37 @@ func (h *MCPHandler) handleListServerTools(w http.ResponseWriter, r *http.Reques
|
||||
_ = json.Unmarshal(srv.Env, &env)
|
||||
_ = json.Unmarshal(srv.Headers, &headers)
|
||||
|
||||
discovered, err := mcpbridge.DiscoverTools(r.Context(), srv.Transport, srv.Command, args, env, srv.URL, headers)
|
||||
if err != nil {
|
||||
slog.Warn("mcp.discover_tools", "server", srv.Name, "error", err)
|
||||
} else {
|
||||
tools = discovered
|
||||
discover := true
|
||||
// OAuth servers: Authorization must come from the OAuth token. If no valid
|
||||
// token is available (not authorized yet), do NOT discover via the server-level
|
||||
// header fallback — return empty so the UI reflects "authorization required"
|
||||
// instead of exposing tools the user can't actually call.
|
||||
if mcpbridge.IsOAuthActive(srv.Settings) {
|
||||
token := ""
|
||||
if h.oauthProvider != nil {
|
||||
tenantID := store.TenantIDFromContext(r.Context())
|
||||
// Tool discovery is server-wide (the tool set is the same for every
|
||||
// user), so always use the GLOBAL token — matching how srv.Headers/Env
|
||||
// here are server-level, regardless of require_user_credentials.
|
||||
token, _ = h.oauthProvider.GetValidToken(r.Context(), srv.ID, tenantID, "")
|
||||
}
|
||||
if token == "" {
|
||||
discover = false
|
||||
} else {
|
||||
if headers == nil {
|
||||
headers = make(map[string]string)
|
||||
}
|
||||
headers["Authorization"] = "Bearer " + token
|
||||
}
|
||||
}
|
||||
|
||||
if discover {
|
||||
discovered, err := mcpbridge.DiscoverTools(r.Context(), srv.Transport, srv.Command, args, env, srv.URL, headers)
|
||||
if err != nil {
|
||||
slog.Warn("mcp.discover_tools", "server", srv.Name, "error", err)
|
||||
} else {
|
||||
tools = discovered
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
// recordingMCPOAuthProvider records the userID arg passed to GetValidToken so we
|
||||
// can assert the scope (global "" vs per-user) used by the preview endpoints.
|
||||
type recordingMCPOAuthProvider struct {
|
||||
mu sync.Mutex
|
||||
calls int
|
||||
lastUserID string
|
||||
token string
|
||||
}
|
||||
|
||||
func (p *recordingMCPOAuthProvider) GetValidToken(_ context.Context, _ uuid.UUID, _ uuid.UUID, userID string) (string, error) {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
p.calls++
|
||||
p.lastUserID = userID
|
||||
return p.token, nil
|
||||
}
|
||||
|
||||
// oauthServerWithRequireUserCreds returns an OAuth server that ALSO has
|
||||
// require_user_credentials=true — the case where the old code wrongly used the
|
||||
// caller's per-user token scope for discovery/test.
|
||||
func oauthServerWithRequireUserCreds(id uuid.UUID) *store.MCPServerData {
|
||||
return &store.MCPServerData{
|
||||
BaseModel: store.BaseModel{ID: id},
|
||||
Name: "oauth-srv",
|
||||
Transport: "streamable-http",
|
||||
URL: "http://127.0.0.1:1/mcp",
|
||||
Settings: json.RawMessage(`{"oauth":{"auth_type":"oauth"},"require_user_credentials":true}`),
|
||||
}
|
||||
}
|
||||
|
||||
// ctxWithUser attaches a tenant + caller user id (the old code would have used
|
||||
// this caller id as the OAuth scope).
|
||||
func ctxWithUser(r *http.Request, userID string) *http.Request {
|
||||
ctx := store.WithTenantID(store.WithUserID(r.Context(), userID), uuid.New())
|
||||
return r.WithContext(ctx)
|
||||
}
|
||||
|
||||
// list-tools must use the GLOBAL OAuth token scope (userID="") even when the
|
||||
// server has require_user_credentials=true — discovery is server-wide.
|
||||
func TestHandleListServerToolsUsesGlobalOAuthScope(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
st := newMockMCPServerForOAuth()
|
||||
id := uuid.New()
|
||||
st.servers[id] = oauthServerWithRequireUserCreds(id)
|
||||
|
||||
prov := &recordingMCPOAuthProvider{token: ""} // not authorized → discovery returns empty
|
||||
h := NewMCPHandler(st, nil, nil)
|
||||
h.SetOAuthProvider(prov)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/mcp/servers/"+id.String()+"/tools", nil)
|
||||
req.SetPathValue("id", id.String())
|
||||
req = ctxWithUser(req, "alice")
|
||||
rec := httptest.NewRecorder()
|
||||
h.handleListServerTools(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
prov.mu.Lock()
|
||||
defer prov.mu.Unlock()
|
||||
if prov.calls == 0 {
|
||||
t.Fatal("expected GetValidToken to be called for OAuth server")
|
||||
}
|
||||
if prov.lastUserID != "" {
|
||||
t.Errorf("list-tools must use GLOBAL scope (userID=\"\"), got %q", prov.lastUserID)
|
||||
}
|
||||
}
|
||||
|
||||
// test-connection must use the GLOBAL OAuth token scope (userID="") for OAuth
|
||||
// servers, regardless of require_user_credentials.
|
||||
func TestHandleTestConnectionUsesGlobalOAuthScope(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
st := newMockMCPServerForOAuth()
|
||||
id := uuid.New()
|
||||
st.servers[id] = oauthServerWithRequireUserCreds(id)
|
||||
|
||||
prov := &recordingMCPOAuthProvider{token: ""}
|
||||
h := NewMCPHandler(st, nil, nil)
|
||||
h.SetOAuthProvider(prov)
|
||||
|
||||
body := `{"server_id":"` + id.String() + `","transport":"streamable-http","url":"http://127.0.0.1:1/mcp","headers":{"Authorization":"Bearer body-token"}}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/mcp/servers/test", strings.NewReader(body))
|
||||
req = ctxWithUser(req, "alice")
|
||||
rec := httptest.NewRecorder()
|
||||
h.handleTestConnection(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
prov.mu.Lock()
|
||||
defer prov.mu.Unlock()
|
||||
if prov.calls == 0 {
|
||||
t.Fatal("expected GetValidToken to be called for OAuth server")
|
||||
}
|
||||
if prov.lastUserID != "" {
|
||||
t.Errorf("test-connection must use GLOBAL scope (userID=\"\"), got %q", prov.lastUserID)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
// setupUpdatePurgeHandler builds an MCPHandler wired with an in-memory server
|
||||
// store (seeded with one OAuth-enabled server), an OAuth token store holding a
|
||||
// global + per-user token for that server, and a recording pool evictor.
|
||||
func setupUpdatePurgeHandler(t *testing.T, settings string) (*http.ServeMux, uuid.UUID, *mockOAuthTokenStore, *mockPoolEvictor) {
|
||||
t.Helper()
|
||||
setAdminToken(t, "test-admin-token")
|
||||
|
||||
id := uuid.New()
|
||||
srvStore := newMockMCPServerForOAuth()
|
||||
srvStore.servers[id] = &store.MCPServerData{
|
||||
BaseModel: store.BaseModel{ID: id},
|
||||
Name: "test-server",
|
||||
Transport: "streamable-http",
|
||||
URL: "http://127.0.0.1:9/old",
|
||||
Settings: json.RawMessage(settings),
|
||||
}
|
||||
|
||||
oauthStore := newMockOAuthTokenStore()
|
||||
oauthStore.global[id.String()] = &store.MCPOAuthToken{ServerID: id}
|
||||
oauthStore.perUser[id.String()+":alice"] = &store.MCPOAuthToken{ServerID: id, UserID: "alice"}
|
||||
|
||||
evictor := &mockPoolEvictor{}
|
||||
h := NewMCPHandler(srvStore, nil, nil)
|
||||
h.SetOAuthStore(oauthStore)
|
||||
h.SetPoolEvictor(evictor)
|
||||
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
return mux, id, oauthStore, evictor
|
||||
}
|
||||
|
||||
const oauthSettingsClientC1 = `{"oauth":{"auth_type":"oauth","client_id":"c1","token_endpoint":"http://127.0.0.1:9/token"}}`
|
||||
|
||||
func TestHandleUpdateServerPurgesOAuthOnURLChange(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
mux, id, oauthStore, evictor := setupUpdatePurgeHandler(t, oauthSettingsClientC1)
|
||||
|
||||
body := map[string]any{"url": "http://127.0.0.1:9/new"}
|
||||
req := adminRequest(http.MethodPut, "/v1/mcp/servers/"+id.String(), body)
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if oauthStore.deletes == 0 {
|
||||
t.Error("expected OAuth tokens to be purged on URL change")
|
||||
}
|
||||
if _, ok := oauthStore.global[id.String()]; ok {
|
||||
t.Error("global token should have been deleted")
|
||||
}
|
||||
if _, ok := oauthStore.perUser[id.String()+":alice"]; ok {
|
||||
t.Error("per-user token should have been deleted")
|
||||
}
|
||||
if evictor.evictServerCalls == 0 {
|
||||
t.Error("expected EvictServer to be called on URL change")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleUpdateServerPurgesOAuthOnOAuthConfigChange(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
mux, id, oauthStore, evictor := setupUpdatePurgeHandler(t, oauthSettingsClientC1)
|
||||
|
||||
// Same URL, but client_id changes c1 -> c2 → tokens for c1 are now invalid.
|
||||
body := map[string]any{
|
||||
"settings": map[string]any{
|
||||
"oauth": map[string]any{
|
||||
"auth_type": "oauth",
|
||||
"client_id": "c2",
|
||||
"token_endpoint": "http://127.0.0.1:9/token",
|
||||
},
|
||||
},
|
||||
}
|
||||
req := adminRequest(http.MethodPut, "/v1/mcp/servers/"+id.String(), body)
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if _, ok := oauthStore.global[id.String()]; ok {
|
||||
t.Error("global token should have been deleted on OAuth config change")
|
||||
}
|
||||
if evictor.evictServerCalls == 0 {
|
||||
t.Error("expected EvictServer to be called on OAuth config change")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleUpdateServerKeepsOAuthOnUnrelatedChange(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
mux, id, oauthStore, evictor := setupUpdatePurgeHandler(t, oauthSettingsClientC1)
|
||||
|
||||
// Change an unrelated field; URL and oauth config are untouched.
|
||||
body := map[string]any{"tool_prefix": "tp_"}
|
||||
req := adminRequest(http.MethodPut, "/v1/mcp/servers/"+id.String(), body)
|
||||
rec := httptest.NewRecorder()
|
||||
mux.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200; body: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if oauthStore.deletes != 0 {
|
||||
t.Errorf("tokens must NOT be purged on unrelated change, deletes = %d", oauthStore.deletes)
|
||||
}
|
||||
if _, ok := oauthStore.global[id.String()]; !ok {
|
||||
t.Error("global token should still exist")
|
||||
}
|
||||
if evictor.evictServerCalls != 0 {
|
||||
t.Errorf("EvictServer must NOT be called on unrelated change, calls = %d", evictor.evictServerCalls)
|
||||
}
|
||||
}
|
||||
@@ -7,19 +7,36 @@ import (
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/bus"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/permissions"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
|
||||
)
|
||||
|
||||
// MCPUserCredentialsHandler handles per-user MCP credential endpoints.
|
||||
type MCPUserCredentialsHandler struct {
|
||||
store store.MCPServerStore
|
||||
tenantStore store.TenantStore
|
||||
msgBus *bus.MessageBus
|
||||
}
|
||||
|
||||
// NewMCPUserCredentialsHandler creates a handler for MCP user credential endpoints.
|
||||
func NewMCPUserCredentialsHandler(s store.MCPServerStore, ts store.TenantStore) *MCPUserCredentialsHandler {
|
||||
return &MCPUserCredentialsHandler{store: s, tenantStore: ts}
|
||||
func NewMCPUserCredentialsHandler(s store.MCPServerStore, ts store.TenantStore, msgBus *bus.MessageBus) *MCPUserCredentialsHandler {
|
||||
return &MCPUserCredentialsHandler{store: s, tenantStore: ts, msgBus: msgBus}
|
||||
}
|
||||
|
||||
// emitMCPCacheInvalidate broadcasts an MCP cache-invalidate event so per-user pool
|
||||
// connections are evicted (mcp.pool.all_user_connections_evicted) and agent Loop
|
||||
// caches reload after credentials change — otherwise the pooled per-user connection
|
||||
// keeps the stale credential until idle TTL. Mirrors MCPHandler.emitCacheInvalidate.
|
||||
func (h *MCPUserCredentialsHandler) emitMCPCacheInvalidate() {
|
||||
if h.msgBus == nil {
|
||||
return
|
||||
}
|
||||
h.msgBus.Broadcast(bus.Event{
|
||||
Name: protocol.EventCacheInvalidate,
|
||||
Payload: bus.CacheInvalidatePayload{Kind: bus.CacheKindMCP},
|
||||
})
|
||||
}
|
||||
|
||||
// RegisterRoutes registers MCP user credential routes.
|
||||
@@ -104,6 +121,9 @@ func (h *MCPUserCredentialsHandler) handleSet(w http.ResponseWriter, r *http.Req
|
||||
return
|
||||
}
|
||||
|
||||
// Drop pooled per-user connections so the new credentials take effect immediately.
|
||||
h.emitMCPCacheInvalidate()
|
||||
|
||||
writeJSON(w, http.StatusOK, map[string]string{"status": "updated"})
|
||||
}
|
||||
|
||||
@@ -127,7 +147,7 @@ func (h *MCPUserCredentialsHandler) handleGet(w http.ResponseWriter, r *http.Req
|
||||
}
|
||||
|
||||
creds, err := h.store.GetUserCredentials(r.Context(), serverID, userID)
|
||||
if err != nil {
|
||||
if err != nil || creds == nil {
|
||||
writeJSON(w, http.StatusOK, map[string]any{"has_credentials": false})
|
||||
return
|
||||
}
|
||||
@@ -164,6 +184,9 @@ func (h *MCPUserCredentialsHandler) handleDelete(w http.ResponseWriter, r *http.
|
||||
return
|
||||
}
|
||||
|
||||
// Drop pooled per-user connections so the removed credentials stop being used.
|
||||
h.emitMCPCacheInvalidate()
|
||||
|
||||
writeJSON(w, http.StatusOK, map[string]string{"status": "deleted"})
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/bus"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
"github.com/nextlevelbuilder/goclaw/pkg/protocol"
|
||||
)
|
||||
|
||||
// countMCPCacheInvalidate counts CacheKindMCP cache-invalidate events.
|
||||
func countMCPCacheInvalidate(events []bus.Event) int {
|
||||
n := 0
|
||||
for _, e := range events {
|
||||
if e.Name != protocol.EventCacheInvalidate {
|
||||
continue
|
||||
}
|
||||
if p, ok := e.Payload.(bus.CacheInvalidatePayload); ok && p.Kind == bus.CacheKindMCP {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// Saving and deleting per-user MCP credentials must broadcast a CacheKindMCP
|
||||
// cache-invalidate so pooled per-user connections are evicted and the new/removed
|
||||
// credentials take effect immediately (instead of waiting for idle TTL).
|
||||
func TestMCPUserCredentialsEmitCacheInvalidateOnSetAndDelete(t *testing.T) {
|
||||
srvStore := newMockMCPServerForOAuth()
|
||||
mb := bus.New()
|
||||
|
||||
var mu sync.Mutex
|
||||
var events []bus.Event
|
||||
mb.Subscribe("test-capture", func(e bus.Event) {
|
||||
mu.Lock()
|
||||
events = append(events, e)
|
||||
mu.Unlock()
|
||||
})
|
||||
|
||||
h := NewMCPUserCredentialsHandler(srvStore, nil, mb)
|
||||
serverID := uuid.New()
|
||||
ctxWith := func(r *http.Request) *http.Request {
|
||||
ctx := store.WithTenantID(store.WithUserID(r.Context(), "user-1"), uuid.New())
|
||||
return r.WithContext(ctx)
|
||||
}
|
||||
|
||||
// SET credentials.
|
||||
setReq := httptest.NewRequest(http.MethodPut,
|
||||
"/v1/mcp/servers/"+serverID.String()+"/user-credentials",
|
||||
strings.NewReader(`{"api_key":"k"}`))
|
||||
setReq.SetPathValue("id", serverID.String())
|
||||
setReq = ctxWith(setReq)
|
||||
setRec := httptest.NewRecorder()
|
||||
h.handleSet(setRec, setReq)
|
||||
if setRec.Code != http.StatusOK {
|
||||
t.Fatalf("handleSet status = %d, want 200; body: %s", setRec.Code, setRec.Body.String())
|
||||
}
|
||||
|
||||
// DELETE credentials.
|
||||
delReq := httptest.NewRequest(http.MethodDelete,
|
||||
"/v1/mcp/servers/"+serverID.String()+"/user-credentials", nil)
|
||||
delReq.SetPathValue("id", serverID.String())
|
||||
delReq = ctxWith(delReq)
|
||||
delRec := httptest.NewRecorder()
|
||||
h.handleDelete(delRec, delReq)
|
||||
if delRec.Code != http.StatusOK {
|
||||
t.Fatalf("handleDelete status = %d, want 200; body: %s", delRec.Code, delRec.Body.String())
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
got := countMCPCacheInvalidate(events)
|
||||
mu.Unlock()
|
||||
if got < 2 {
|
||||
t.Errorf("expected >=2 CacheKindMCP events (set + delete), got %d", got)
|
||||
}
|
||||
}
|
||||
@@ -71,7 +71,7 @@ var customToolAllowedFields = map[string]bool{
|
||||
}
|
||||
|
||||
var mcpServerAllowedFields = map[string]bool{
|
||||
"name": true, "transport": true, "command": true, "args": true,
|
||||
"name": true, "display_name": true, "transport": true, "command": true, "args": true,
|
||||
"url": true, "api_key": true, "env": true, "headers": true,
|
||||
"enabled": true, "tool_prefix": true, "timeout_sec": true,
|
||||
"agent_id": true, "config": true, "settings": true,
|
||||
|
||||
@@ -171,8 +171,10 @@ func init() {
|
||||
MsgFailedToDeleteFile: "failed to delete",
|
||||
|
||||
// OAuth
|
||||
MsgNoPendingOAuth: "no pending OAuth flow",
|
||||
MsgFailedToSaveToken: "failed to save token",
|
||||
MsgNoPendingOAuth: "no pending OAuth flow",
|
||||
MsgFailedToSaveToken: "failed to save token",
|
||||
MsgOAuthCallbackSuccess: "Authorization successful. You may close this window.",
|
||||
MsgOAuthCallbackFailed: "Authorization failed. You may close this window.",
|
||||
|
||||
// Intent Classify
|
||||
MsgStatusWorking: "🔄 I'm working on your request... Please wait.",
|
||||
|
||||
@@ -171,8 +171,10 @@ func init() {
|
||||
MsgFailedToDeleteFile: "không thể xóa",
|
||||
|
||||
// OAuth
|
||||
MsgNoPendingOAuth: "không có luồng OAuth đang chờ",
|
||||
MsgFailedToSaveToken: "không thể lưu token",
|
||||
MsgNoPendingOAuth: "không có luồng OAuth đang chờ",
|
||||
MsgFailedToSaveToken: "không thể lưu token",
|
||||
MsgOAuthCallbackSuccess: "Ủy quyền thành công. Bạn có thể đóng cửa sổ này.",
|
||||
MsgOAuthCallbackFailed: "Ủy quyền thất bại. Bạn có thể đóng cửa sổ này.",
|
||||
|
||||
// Intent Classify
|
||||
MsgStatusWorking: "🔄 Mình đang xử lý yêu cầu của bạn... Vui lòng chờ.",
|
||||
|
||||
@@ -171,8 +171,10 @@ func init() {
|
||||
MsgFailedToDeleteFile: "删除失败",
|
||||
|
||||
// OAuth
|
||||
MsgNoPendingOAuth: "没有待处理的 OAuth 流程",
|
||||
MsgFailedToSaveToken: "保存令牌失败",
|
||||
MsgNoPendingOAuth: "没有待处理的 OAuth 流程",
|
||||
MsgFailedToSaveToken: "保存令牌失败",
|
||||
MsgOAuthCallbackSuccess: "授权成功。您可以关闭此窗口。",
|
||||
MsgOAuthCallbackFailed: "授权失败。您可以关闭此窗口。",
|
||||
|
||||
// Intent Classify
|
||||
MsgStatusWorking: "🔄 我正在处理您的请求...请稍候。",
|
||||
|
||||
@@ -208,8 +208,10 @@ const (
|
||||
MsgFailedToDeleteFile = "error.failed_to_delete_file" // "failed to delete"
|
||||
|
||||
// --- OAuth ---
|
||||
MsgNoPendingOAuth = "error.no_pending_oauth" // "no pending OAuth flow"
|
||||
MsgFailedToSaveToken = "error.failed_to_save_token" // "failed to save token"
|
||||
MsgNoPendingOAuth = "error.no_pending_oauth" // "no pending OAuth flow"
|
||||
MsgFailedToSaveToken = "error.failed_to_save_token" // "failed to save token"
|
||||
MsgOAuthCallbackSuccess = "oauth.callback_success" // "Authorization successful. You may close this window."
|
||||
MsgOAuthCallbackFailed = "oauth.callback_failed" // "Authorization failed. You may close this window."
|
||||
|
||||
// --- Intent Classify (channel-facing status replies) ---
|
||||
MsgStatusWorking = "status.working" // "🔄 I'm working on your request... Please wait."
|
||||
|
||||
+68
-2
@@ -124,6 +124,16 @@ type Manager struct {
|
||||
// LoadForAgent("") for later per-request tool resolution. These servers are NOT
|
||||
// connected at startup — connections are created per-user via pool.AcquireUser().
|
||||
userCredServers []store.MCPAccessInfo
|
||||
|
||||
// oauthTokenProvider resolves a live Bearer token for OAuth-enabled MCP servers.
|
||||
// nil = OAuth token injection disabled.
|
||||
oauthTokenProvider OAuthTokenProvider
|
||||
}
|
||||
|
||||
// OAuthTokenProvider retrieves a valid OAuth Bearer token for an MCP server.
|
||||
// userID="" means global token; non-empty means per-user token.
|
||||
type OAuthTokenProvider interface {
|
||||
GetValidToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) (string, error)
|
||||
}
|
||||
|
||||
// ManagerOption configures the Manager.
|
||||
@@ -159,6 +169,13 @@ func WithGrantChecker(gc GrantChecker) ManagerOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithOAuthTokenProvider sets the OAuth token provider for Bearer token injection.
|
||||
func WithOAuthTokenProvider(p OAuthTokenProvider) ManagerOption {
|
||||
return func(m *Manager) {
|
||||
m.oauthTokenProvider = p
|
||||
}
|
||||
}
|
||||
|
||||
// NewManager creates a new MCP Manager.
|
||||
func NewManager(registry *tools.Registry, opts ...ManagerOption) *Manager {
|
||||
m := &Manager{
|
||||
@@ -254,8 +271,12 @@ func (m *Manager) resolveServerCredentials(ctx context.Context, info store.MCPAc
|
||||
return nil
|
||||
}
|
||||
|
||||
// Inject APIKey into headers if present (bug fix: was never passed to connections)
|
||||
if srv.APIKey != "" && headers["Authorization"] == "" {
|
||||
oauthActive := isOAuthActive(srv.Settings)
|
||||
|
||||
// Inject APIKey into headers if present — ONLY for non-OAuth servers.
|
||||
// For OAuth servers the Authorization MUST come from the OAuth token; using the
|
||||
// server-level api_key as a fallback would expose tools before authorization.
|
||||
if !oauthActive && srv.APIKey != "" && headers["Authorization"] == "" {
|
||||
if headers == nil {
|
||||
headers = make(map[string]string)
|
||||
}
|
||||
@@ -307,6 +328,31 @@ func (m *Manager) resolveServerCredentials(ctx context.Context, info store.MCPAc
|
||||
}
|
||||
}
|
||||
|
||||
// OAuth-enabled servers: the OAuth token is the ONLY source of Authorization.
|
||||
// Require a valid token — if none is available (not authorized yet / refresh
|
||||
// failed / OAuth subsystem absent), skip the server so it is NOT loaded into the
|
||||
// shared registry with the server-level credential as a fallback.
|
||||
if oauthActive {
|
||||
if m.oauthTokenProvider == nil {
|
||||
slog.Debug("mcp.skip_oauth_no_provider", "server", srv.Name)
|
||||
return nil
|
||||
}
|
||||
tenantID := store.TenantIDFromContext(ctx)
|
||||
oauthUserID := ""
|
||||
if requireUserCreds(srv.Settings) {
|
||||
oauthUserID = userID
|
||||
}
|
||||
token, err2 := m.oauthTokenProvider.GetValidToken(ctx, srv.ID, tenantID, oauthUserID)
|
||||
if err2 != nil || token == "" {
|
||||
slog.Debug("mcp.skip_oauth_no_token", "server", srv.Name, "user", oauthUserID, "error", err2)
|
||||
return nil
|
||||
}
|
||||
if headers == nil {
|
||||
headers = make(map[string]string)
|
||||
}
|
||||
headers["Authorization"] = "Bearer " + token
|
||||
}
|
||||
|
||||
// Per-user credentials change connection params → can't share pool connection.
|
||||
// Fall back to per-agent mode when user has custom credentials.
|
||||
hasUserCreds := contextCreds != nil && (contextCreds.APIKey != "" || len(contextCreds.Headers) > 0 || len(contextCreds.Env) > 0)
|
||||
@@ -630,6 +676,26 @@ func requireUserCreds(settings json.RawMessage) bool {
|
||||
return s.RequireUserCredentials
|
||||
}
|
||||
|
||||
// isOAuthActive checks if the server's settings have OAuth auth_type enabled.
|
||||
func isOAuthActive(settings json.RawMessage) bool {
|
||||
return IsOAuthActive(settings)
|
||||
}
|
||||
|
||||
// IsOAuthActive reports whether the given raw server settings JSON has OAuth enabled.
|
||||
// Exported for use by HTTP handlers that need to gate OAuth token injection.
|
||||
func IsOAuthActive(settings json.RawMessage) bool {
|
||||
if len(settings) == 0 {
|
||||
return false
|
||||
}
|
||||
var s struct {
|
||||
OAuth struct {
|
||||
AuthType string `json:"auth_type"`
|
||||
} `json:"oauth"`
|
||||
}
|
||||
_ = json.Unmarshal(settings, &s)
|
||||
return s.OAuth.AuthType == "oauth"
|
||||
}
|
||||
|
||||
// ToolHints carries admin-authored description hints for MCP tools.
|
||||
// Stored under MCPServerData.Settings.tool_hints as JSONB:
|
||||
//
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
// fakeOAuthProvider is a minimal OAuthTokenProvider for resolveServerCredentials tests.
|
||||
type fakeOAuthProvider struct {
|
||||
token string
|
||||
err error
|
||||
calls int
|
||||
}
|
||||
|
||||
func (f *fakeOAuthProvider) GetValidToken(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ string) (string, error) {
|
||||
f.calls++
|
||||
return f.token, f.err
|
||||
}
|
||||
|
||||
func oauthAccessInfo(apiKey string) store.MCPAccessInfo {
|
||||
return store.MCPAccessInfo{Server: store.MCPServerData{
|
||||
BaseModel: store.BaseModel{ID: uuid.New()},
|
||||
Name: "oauth-srv",
|
||||
Transport: "streamable-http",
|
||||
URL: "http://127.0.0.1:1/mcp",
|
||||
APIKey: apiKey,
|
||||
Enabled: true,
|
||||
Settings: json.RawMessage(`{"oauth":{"auth_type":"oauth"}}`),
|
||||
}}
|
||||
}
|
||||
|
||||
// OAuth server + no token → resolveServerCredentials must return nil (skip),
|
||||
// NOT fall back to the server-level api_key.
|
||||
func TestResolveServerCredentials_OAuthSkipsWithoutToken(t *testing.T) {
|
||||
prov := &fakeOAuthProvider{token: ""}
|
||||
m := &Manager{oauthTokenProvider: prov}
|
||||
|
||||
rs := m.resolveServerCredentials(context.Background(), oauthAccessInfo("shared-server-key"), "")
|
||||
if rs != nil {
|
||||
t.Fatalf("expected nil (skip) for OAuth server without token, got headers=%v", rs.headers)
|
||||
}
|
||||
if prov.calls == 0 {
|
||||
t.Error("expected GetValidToken to be consulted")
|
||||
}
|
||||
}
|
||||
|
||||
// OAuth server + valid token → Authorization must be the token, never the api_key.
|
||||
func TestResolveServerCredentials_OAuthUsesTokenNotApiKey(t *testing.T) {
|
||||
m := &Manager{oauthTokenProvider: &fakeOAuthProvider{token: "tok-123"}}
|
||||
|
||||
rs := m.resolveServerCredentials(context.Background(), oauthAccessInfo("shared-server-key"), "")
|
||||
if rs == nil {
|
||||
t.Fatal("expected resolved server when token present")
|
||||
}
|
||||
if got := rs.headers["Authorization"]; got != "Bearer tok-123" {
|
||||
t.Errorf("Authorization = %q, want %q (token, not server api_key)", got, "Bearer tok-123")
|
||||
}
|
||||
}
|
||||
|
||||
// OAuth server but no provider wired → skip (cannot obtain a token).
|
||||
func TestResolveServerCredentials_OAuthNoProviderSkips(t *testing.T) {
|
||||
m := &Manager{} // oauthTokenProvider == nil
|
||||
|
||||
rs := m.resolveServerCredentials(context.Background(), oauthAccessInfo("shared-server-key"), "")
|
||||
if rs != nil {
|
||||
t.Fatal("expected nil (skip) for OAuth server with no provider wired")
|
||||
}
|
||||
}
|
||||
|
||||
// Non-OAuth server → server-level api_key is still injected as before.
|
||||
func TestResolveServerCredentials_NonOAuthUsesApiKey(t *testing.T) {
|
||||
m := &Manager{}
|
||||
info := store.MCPAccessInfo{Server: store.MCPServerData{
|
||||
BaseModel: store.BaseModel{ID: uuid.New()},
|
||||
Name: "plain-srv",
|
||||
Transport: "streamable-http",
|
||||
URL: "http://127.0.0.1:1/mcp",
|
||||
APIKey: "shared-server-key",
|
||||
Enabled: true,
|
||||
}}
|
||||
|
||||
rs := m.resolveServerCredentials(context.Background(), info, "")
|
||||
if rs == nil {
|
||||
t.Fatal("expected resolved server for non-OAuth")
|
||||
}
|
||||
if got := rs.headers["Authorization"]; got != "Bearer shared-server-key" {
|
||||
t.Errorf("Authorization = %q, want %q", got, "Bearer shared-server-key")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
package mcpoauth
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
)
|
||||
|
||||
// DCRRequest is the Dynamic Client Registration request body (RFC 7591).
|
||||
type DCRRequest struct {
|
||||
RedirectURIs []string `json:"redirect_uris"`
|
||||
GrantTypes []string `json:"grant_types"`
|
||||
ResponseTypes []string `json:"response_types"`
|
||||
ClientName string `json:"client_name"`
|
||||
TokenEndpointAuthMethod string `json:"token_endpoint_auth_method"`
|
||||
}
|
||||
|
||||
// DCRResponse is the server response from Dynamic Client Registration.
|
||||
type DCRResponse struct {
|
||||
ClientID string `json:"client_id"`
|
||||
ClientSecret string `json:"client_secret,omitempty"`
|
||||
ClientIDIssuedAt int64 `json:"client_id_issued_at,omitempty"`
|
||||
ClientSecretExpiresAt int64 `json:"client_secret_expires_at,omitempty"`
|
||||
TokenEndpointAuthMethod string `json:"token_endpoint_auth_method,omitempty"`
|
||||
}
|
||||
|
||||
// RegisterClient performs Dynamic Client Registration (RFC 7591).
|
||||
// callbackURL is the redirect_uri to register.
|
||||
func RegisterClient(ctx context.Context, client *http.Client, registrationEndpoint, callbackURL string) (*DCRResponse, error) {
|
||||
body := DCRRequest{
|
||||
RedirectURIs: []string{callbackURL},
|
||||
GrantTypes: []string{"authorization_code", "refresh_token"},
|
||||
ResponseTypes: []string{"code"},
|
||||
ClientName: "GoClaw Gateway",
|
||||
TokenEndpointAuthMethod: "none", // public client — PKCE handles security
|
||||
}
|
||||
data, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
_, ip, err := security.Validate(registrationEndpoint)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("mcpoauth: SSRF validation failed for registration endpoint: %w", err)
|
||||
}
|
||||
ctx = security.WithPinnedIP(ctx, ip)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, registrationEndpoint, bytes.NewReader(data))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("mcpoauth: DCR request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 32*1024))
|
||||
if resp.StatusCode != http.StatusCreated && resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("mcpoauth: DCR returned HTTP %d: %s", resp.StatusCode, string(respBody))
|
||||
}
|
||||
var result DCRResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return nil, fmt.Errorf("mcpoauth: DCR response parse error: %w", err)
|
||||
}
|
||||
if result.ClientID == "" {
|
||||
return nil, fmt.Errorf("mcpoauth: DCR response missing client_id")
|
||||
}
|
||||
return &result, nil
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
package mcpoauth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
)
|
||||
|
||||
func TestRegisterClientSuccess201(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
var req DCRRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
http.Error(w, "bad json", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if req.ClientName != "GoClaw Gateway" {
|
||||
http.Error(w, "wrong client_name", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
_ = json.NewEncoder(w).Encode(DCRResponse{ClientID: "registered-client-id"})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := security.NewSafeClient(5 * time.Second)
|
||||
resp, err := RegisterClient(context.Background(), client, srv.URL, "https://goclaw.example.com/callback")
|
||||
if err != nil {
|
||||
t.Fatalf("RegisterClient() error: %v", err)
|
||||
}
|
||||
if resp.ClientID != "registered-client-id" {
|
||||
t.Errorf("ClientID = %q, want %q", resp.ClientID, "registered-client-id")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterClientAccepts200(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK) // 200, not 201
|
||||
_ = json.NewEncoder(w).Encode(DCRResponse{ClientID: "client-200"})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := security.NewSafeClient(5 * time.Second)
|
||||
resp, err := RegisterClient(context.Background(), client, srv.URL, "https://example.com/cb")
|
||||
if err != nil {
|
||||
t.Fatalf("RegisterClient() with HTTP 200 error: %v", err)
|
||||
}
|
||||
if resp.ClientID != "client-200" {
|
||||
t.Errorf("ClientID = %q, want %q", resp.ClientID, "client-200")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterClientMissingClientID(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
// Response has no client_id.
|
||||
_ = json.NewEncoder(w).Encode(DCRResponse{})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := security.NewSafeClient(5 * time.Second)
|
||||
_, err := RegisterClient(context.Background(), client, srv.URL, "https://example.com/cb")
|
||||
if err == nil {
|
||||
t.Fatal("expected error when client_id missing, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "client_id") {
|
||||
t.Errorf("error should mention client_id, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterClientHTTPError(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := security.NewSafeClient(5 * time.Second)
|
||||
_, err := RegisterClient(context.Background(), client, srv.URL, "https://example.com/cb")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for HTTP 401, got nil")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "401") {
|
||||
t.Errorf("error should mention HTTP status, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterClientResponseTooLarge(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
// Write > 32KB of data.
|
||||
large := make([]byte, 33*1024)
|
||||
for i := range large {
|
||||
large[i] = 'x'
|
||||
}
|
||||
// Wrap in valid JSON to pass JSON decode.
|
||||
_, _ = w.Write([]byte(`{"client_id":"`))
|
||||
_, _ = w.Write(large)
|
||||
_, _ = w.Write([]byte(`"}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := security.NewSafeClient(5 * time.Second)
|
||||
// The response is truncated to 32KB before JSON decode; the JSON will be
|
||||
// malformed, so we expect a parse error.
|
||||
_, err := RegisterClient(context.Background(), client, srv.URL, "https://example.com/cb")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for oversized response, got nil")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
// Package mcpoauth implements OAuth 2.1 client support for MCP servers.
|
||||
// Discovery follows RFC 9728 (Protected Resource Metadata) then RFC 8414
|
||||
// (Authorization Server Metadata) with OIDC fallback.
|
||||
package mcpoauth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
)
|
||||
|
||||
const discoveryCacheTTL = 5 * time.Minute
|
||||
|
||||
// DiscoveryResult holds the Authorization Server metadata for an MCP server.
|
||||
type DiscoveryResult struct {
|
||||
Issuer string
|
||||
AuthorizationEndpoint string
|
||||
TokenEndpoint string
|
||||
RegistrationEndpoint string // DCR — optional
|
||||
ResourceURI string // RFC 8707 resource indicator
|
||||
ScopesSupported []string
|
||||
}
|
||||
|
||||
type cachedDiscovery struct {
|
||||
result *DiscoveryResult
|
||||
fetchedAt time.Time
|
||||
}
|
||||
|
||||
// Discoverer caches and performs OAuth AS discovery for MCP server URLs.
|
||||
type Discoverer struct {
|
||||
client *http.Client
|
||||
mu sync.Mutex
|
||||
cache map[string]*cachedDiscovery
|
||||
}
|
||||
|
||||
func NewDiscoverer(client *http.Client) *Discoverer {
|
||||
return &Discoverer{
|
||||
client: client,
|
||||
cache: make(map[string]*cachedDiscovery),
|
||||
}
|
||||
}
|
||||
|
||||
// Discover finds the OAuth Authorization Server metadata for an MCP server URL.
|
||||
//
|
||||
// Priority (per plan spec):
|
||||
//
|
||||
// Phase 1 — RFC 9728: fetch /.well-known/oauth-protected-resource[/{path}]
|
||||
// to resolve the Authorization Server base URL.
|
||||
// Phase 2 — RFC 8414 + OIDC fallback on the resolved AS base URL:
|
||||
// 1. {AS}/.well-known/oauth-authorization-server
|
||||
// 2. {AS}/.well-known/oauth-authorization-server/{path}
|
||||
// 3. {AS}/.well-known/openid-configuration
|
||||
// 4. {AS}/{path}/.well-known/openid-configuration
|
||||
func (d *Discoverer) Discover(ctx context.Context, mcpServerURL string) (*DiscoveryResult, error) {
|
||||
d.mu.Lock()
|
||||
if c, ok := d.cache[mcpServerURL]; ok && time.Since(c.fetchedAt) < discoveryCacheTTL {
|
||||
d.mu.Unlock()
|
||||
return c.result, nil
|
||||
}
|
||||
d.mu.Unlock()
|
||||
|
||||
result, err := d.discover(ctx, mcpServerURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
d.mu.Lock()
|
||||
d.cache[mcpServerURL] = &cachedDiscovery{result: result, fetchedAt: time.Now()}
|
||||
d.mu.Unlock()
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// InvalidateCache removes the cached discovery for the given MCP server URL.
|
||||
func (d *Discoverer) InvalidateCache(mcpServerURL string) {
|
||||
d.mu.Lock()
|
||||
delete(d.cache, mcpServerURL)
|
||||
d.mu.Unlock()
|
||||
}
|
||||
|
||||
func (d *Discoverer) discover(ctx context.Context, mcpServerURL string) (*DiscoveryResult, error) {
|
||||
parsed, err := url.Parse(mcpServerURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("mcpoauth: invalid MCP server URL %q: %w", mcpServerURL, err)
|
||||
}
|
||||
base := parsed.Scheme + "://" + parsed.Host
|
||||
path := strings.TrimPrefix(parsed.Path, "/")
|
||||
|
||||
// Phase 1 — RFC 9728: try to discover Authorization Server via Protected Resource Metadata.
|
||||
// When RFC 9728 succeeds, the returned AS base URL is authoritative — no path suffix needed.
|
||||
// When RFC 9728 fails, fall back to using the MCP server base with its path.
|
||||
asBase := base
|
||||
asPath := path
|
||||
if prm, err2 := d.fetchProtectedResourceMeta(ctx, base, path); err2 == nil && len(prm.AuthorizationServers) > 0 {
|
||||
asBase = strings.TrimRight(prm.AuthorizationServers[0], "/")
|
||||
asPath = "" // AS URL from RFC 9728 is complete — no path suffix
|
||||
}
|
||||
|
||||
// Phase 2 — RFC 8414 + OIDC fallback on AS base.
|
||||
return d.discoverASMeta(ctx, asBase, asPath, mcpServerURL)
|
||||
}
|
||||
|
||||
type protectedResourceMeta struct {
|
||||
Resource string `json:"resource"`
|
||||
AuthorizationServers []string `json:"authorization_servers"`
|
||||
}
|
||||
|
||||
func (d *Discoverer) fetchProtectedResourceMeta(ctx context.Context, base, path string) (*protectedResourceMeta, error) {
|
||||
candidates := []string{
|
||||
base + "/.well-known/oauth-protected-resource",
|
||||
}
|
||||
if path != "" {
|
||||
candidates = append(candidates, base+"/.well-known/oauth-protected-resource/"+path)
|
||||
}
|
||||
for _, u := range candidates {
|
||||
var meta protectedResourceMeta
|
||||
if err := d.getJSON(ctx, u, &meta); err == nil && len(meta.AuthorizationServers) > 0 {
|
||||
return &meta, nil
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("mcpoauth: RFC 9728 discovery failed")
|
||||
}
|
||||
|
||||
type asMeta struct {
|
||||
Issuer string `json:"issuer"`
|
||||
AuthorizationEndpoint string `json:"authorization_endpoint"`
|
||||
TokenEndpoint string `json:"token_endpoint"`
|
||||
RegistrationEndpoint string `json:"registration_endpoint"`
|
||||
ScopesSupported []string `json:"scopes_supported"`
|
||||
CodeChallengeMethodsSupported []string `json:"code_challenge_methods_supported"`
|
||||
}
|
||||
|
||||
func (d *Discoverer) discoverASMeta(ctx context.Context, asBase, path, mcpServerURL string) (*DiscoveryResult, error) {
|
||||
candidates := []string{
|
||||
asBase + "/.well-known/oauth-authorization-server",
|
||||
}
|
||||
if path != "" {
|
||||
candidates = append(candidates,
|
||||
asBase+"/.well-known/oauth-authorization-server/"+path,
|
||||
)
|
||||
}
|
||||
candidates = append(candidates, asBase+"/.well-known/openid-configuration")
|
||||
if path != "" {
|
||||
candidates = append(candidates, asBase+"/"+path+"/.well-known/openid-configuration")
|
||||
}
|
||||
|
||||
for _, u := range candidates {
|
||||
var meta asMeta
|
||||
if err := d.getJSON(ctx, u, &meta); err != nil {
|
||||
continue
|
||||
}
|
||||
if meta.AuthorizationEndpoint == "" || meta.TokenEndpoint == "" {
|
||||
continue
|
||||
}
|
||||
// RFC 8414 §3.3: the issuer in the metadata must identify the same
|
||||
// Authorization Server the metadata was retrieved from. Requiring the
|
||||
// issuer origin to match the fetch origin blocks an AS mix-up where a
|
||||
// compromised resource server points discovery at an issuer running on
|
||||
// a host it does not control. Path differences are tolerated (issuers
|
||||
// like Azure AD carry a tenant path), and an absent issuer is allowed
|
||||
// since some OIDC fallbacks omit it.
|
||||
if meta.Issuer != "" && !sameOrigin(meta.Issuer, u) {
|
||||
slog.Warn("mcpoauth.issuer_mismatch",
|
||||
"metadata_url", u, "advertised_issuer", meta.Issuer)
|
||||
continue
|
||||
}
|
||||
return &DiscoveryResult{
|
||||
Issuer: meta.Issuer,
|
||||
AuthorizationEndpoint: meta.AuthorizationEndpoint,
|
||||
TokenEndpoint: meta.TokenEndpoint,
|
||||
RegistrationEndpoint: meta.RegistrationEndpoint,
|
||||
ResourceURI: mcpServerURL,
|
||||
ScopesSupported: meta.ScopesSupported,
|
||||
}, nil
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("mcpoauth: no OAuth Authorization Server metadata found for %s", mcpServerURL)
|
||||
}
|
||||
|
||||
// sameOrigin reports whether two URLs share the same scheme and host
|
||||
// (case-insensitive). Path, query, and fragment are ignored.
|
||||
func sameOrigin(a, b string) bool {
|
||||
ua, err1 := url.Parse(a)
|
||||
ub, err2 := url.Parse(b)
|
||||
if err1 != nil || err2 != nil {
|
||||
return false
|
||||
}
|
||||
return strings.EqualFold(ua.Scheme, ub.Scheme) && strings.EqualFold(ua.Host, ub.Host)
|
||||
}
|
||||
|
||||
func (d *Discoverer) getJSON(ctx context.Context, rawURL string, dst any) error {
|
||||
// Validate URL and pin the resolved IP into ctx so the safe HTTP client
|
||||
// can dial without re-resolving (prevents DNS rebinding / SSRF).
|
||||
_, ip, err := security.Validate(rawURL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("mcpoauth: SSRF validation failed for %s: %w", rawURL, err)
|
||||
}
|
||||
ctx = security.WithPinnedIP(ctx, ip)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
resp, err := d.client.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("mcpoauth: HTTP %d from %s", resp.StatusCode, rawURL)
|
||||
}
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return json.Unmarshal(body, dst)
|
||||
}
|
||||
@@ -0,0 +1,362 @@
|
||||
package mcpoauth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
)
|
||||
|
||||
// mockASMeta is used to serve AS metadata in tests.
|
||||
type mockASMeta struct {
|
||||
Issuer string `json:"issuer"`
|
||||
AuthorizationEndpoint string `json:"authorization_endpoint"`
|
||||
TokenEndpoint string `json:"token_endpoint"`
|
||||
RegistrationEndpoint string `json:"registration_endpoint,omitempty"`
|
||||
ScopesSupported []string `json:"scopes_supported,omitempty"`
|
||||
}
|
||||
|
||||
// newDiscoveryServer returns a test server that serves OAuth discovery metadata
|
||||
// at standard well-known paths. If protectedResource is non-empty, it responds
|
||||
// to /.well-known/oauth-protected-resource with that AS URL. asStatus controls
|
||||
// the HTTP status for AS metadata endpoints (200 or error).
|
||||
func newDiscoveryServer(t *testing.T, protectedResourceAS string, asStatus int, authEndpoint, tokenEndpoint string) *httptest.Server {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case strings.HasSuffix(r.URL.Path, "/.well-known/oauth-protected-resource") ||
|
||||
strings.Contains(r.URL.Path, "/.well-known/oauth-protected-resource/"):
|
||||
if protectedResourceAS == "" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"resource": "https://mcp.example.com/server",
|
||||
"authorization_servers": []string{protectedResourceAS},
|
||||
})
|
||||
|
||||
case strings.HasSuffix(r.URL.Path, "/.well-known/oauth-authorization-server") ||
|
||||
strings.Contains(r.URL.Path, "/.well-known/oauth-authorization-server/"):
|
||||
if asStatus != http.StatusOK {
|
||||
http.Error(w, "not found", asStatus)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(mockASMeta{
|
||||
AuthorizationEndpoint: authEndpoint,
|
||||
TokenEndpoint: tokenEndpoint,
|
||||
})
|
||||
|
||||
case strings.HasSuffix(r.URL.Path, "/.well-known/openid-configuration") ||
|
||||
strings.Contains(r.URL.Path, "/.well-known/openid-configuration"):
|
||||
if asStatus != http.StatusOK {
|
||||
http.Error(w, "not found", asStatus)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(mockASMeta{
|
||||
AuthorizationEndpoint: authEndpoint,
|
||||
TokenEndpoint: tokenEndpoint,
|
||||
})
|
||||
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func TestDiscoverRFC8414Success(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := newDiscoveryServer(t, "", http.StatusOK, "https://auth.example.com/authorize", "https://auth.example.com/token")
|
||||
|
||||
d := NewDiscoverer(security.NewSafeClient(5 * time.Second))
|
||||
result, err := d.Discover(context.Background(), srv.URL+"/server")
|
||||
if err != nil {
|
||||
t.Fatalf("Discover() error: %v", err)
|
||||
}
|
||||
if result.AuthorizationEndpoint != "https://auth.example.com/authorize" {
|
||||
t.Errorf("AuthorizationEndpoint = %q", result.AuthorizationEndpoint)
|
||||
}
|
||||
if result.TokenEndpoint != "https://auth.example.com/token" {
|
||||
t.Errorf("TokenEndpoint = %q", result.TokenEndpoint)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverCachesResult(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
hitCount := 0
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, ".well-known/oauth-authorization-server") ||
|
||||
strings.Contains(r.URL.Path, ".well-known/openid-configuration") {
|
||||
hitCount++
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(mockASMeta{
|
||||
AuthorizationEndpoint: "https://auth.example.com/authorize",
|
||||
TokenEndpoint: "https://auth.example.com/token",
|
||||
})
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
d := NewDiscoverer(security.NewSafeClient(5 * time.Second))
|
||||
url := srv.URL + "/mcp"
|
||||
|
||||
// First call.
|
||||
if _, err := d.Discover(context.Background(), url); err != nil {
|
||||
t.Fatalf("first Discover() error: %v", err)
|
||||
}
|
||||
// Second call — should use cache.
|
||||
if _, err := d.Discover(context.Background(), url); err != nil {
|
||||
t.Fatalf("second Discover() error: %v", err)
|
||||
}
|
||||
|
||||
if hitCount != 1 {
|
||||
t.Errorf("expected exactly 1 server hit (cached), got %d", hitCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverInvalidateCache(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
hitCount := 0
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, ".well-known") {
|
||||
hitCount++
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(mockASMeta{
|
||||
AuthorizationEndpoint: "https://auth.example.com/authorize",
|
||||
TokenEndpoint: "https://auth.example.com/token",
|
||||
})
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
d := NewDiscoverer(security.NewSafeClient(5 * time.Second))
|
||||
url := srv.URL
|
||||
|
||||
// Populate cache.
|
||||
if _, err := d.Discover(context.Background(), url); err != nil {
|
||||
t.Fatalf("Discover() error: %v", err)
|
||||
}
|
||||
|
||||
// Invalidate and re-fetch.
|
||||
d.InvalidateCache(url)
|
||||
if _, err := d.Discover(context.Background(), url); err != nil {
|
||||
t.Fatalf("Discover() after invalidate error: %v", err)
|
||||
}
|
||||
|
||||
if hitCount < 2 {
|
||||
t.Errorf("expected at least 2 server hits after invalidation, got %d", hitCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverRFC9728Priority(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
// AS server (separate from MCP server in full impl). Its advertised issuer must
|
||||
// share the AS origin (RFC 8414 §3.3), so it is set to asSrv.URL at request time.
|
||||
var asSrv *httptest.Server
|
||||
asSrv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, ".well-known/oauth-authorization-server") ||
|
||||
strings.Contains(r.URL.Path, ".well-known/openid-configuration") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(mockASMeta{
|
||||
Issuer: asSrv.URL,
|
||||
AuthorizationEndpoint: asSrv.URL + "/authorize",
|
||||
TokenEndpoint: asSrv.URL + "/token",
|
||||
})
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer asSrv.Close()
|
||||
|
||||
// MCP server that returns RFC 9728 protected resource metadata pointing to asSrv.
|
||||
// Declare var first so the closure can reference mcpSrv.URL after assignment.
|
||||
var mcpSrv *httptest.Server
|
||||
mcpSrv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, "/.well-known/oauth-protected-resource") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"resource": mcpSrv.URL,
|
||||
"authorization_servers": []string{asSrv.URL},
|
||||
})
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer mcpSrv.Close()
|
||||
|
||||
d := NewDiscoverer(security.NewSafeClient(5 * time.Second))
|
||||
result, err := d.Discover(context.Background(), mcpSrv.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("Discover() error: %v", err)
|
||||
}
|
||||
// Should use AS from RFC 9728, not MCP server itself.
|
||||
if result.Issuer != asSrv.URL {
|
||||
t.Errorf("Issuer = %q, want %q (RFC 9728 AS)", result.Issuer, asSrv.URL)
|
||||
}
|
||||
// Endpoints must come from the AS server, not the MCP server (RFC 9728 priority).
|
||||
if result.AuthorizationEndpoint != asSrv.URL+"/authorize" {
|
||||
t.Errorf("AuthorizationEndpoint = %q, want AS server endpoint", result.AuthorizationEndpoint)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverIssuerMismatchRejected verifies that AS metadata advertising an
|
||||
// issuer on a different origin than the one it was fetched from is rejected
|
||||
// (RFC 8414 §3.3 — AS mix-up protection).
|
||||
func TestDiscoverIssuerMismatchRejected(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, ".well-known/oauth-authorization-server") ||
|
||||
strings.Contains(r.URL.Path, ".well-known/openid-configuration") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
// Issuer points at a host the server does not control.
|
||||
_ = json.NewEncoder(w).Encode(mockASMeta{
|
||||
Issuer: "https://evil.example.com",
|
||||
AuthorizationEndpoint: "https://evil.example.com/authorize",
|
||||
TokenEndpoint: "https://evil.example.com/token",
|
||||
})
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
d := NewDiscoverer(security.NewSafeClient(5 * time.Second))
|
||||
_, err := d.Discover(context.Background(), srv.URL)
|
||||
if err == nil {
|
||||
t.Fatal("expected discovery to reject issuer/origin mismatch, got nil error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverFallbackOIDC(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
// Server that returns 404 for RFC 8414 path but serves OIDC.
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case strings.Contains(r.URL.Path, "oauth-protected-resource"):
|
||||
http.NotFound(w, r)
|
||||
case strings.Contains(r.URL.Path, "oauth-authorization-server"):
|
||||
http.NotFound(w, r) // RFC 8414 not available
|
||||
case strings.Contains(r.URL.Path, "openid-configuration"):
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(mockASMeta{
|
||||
AuthorizationEndpoint: "https://oidc.example.com/authorize",
|
||||
TokenEndpoint: "https://oidc.example.com/token",
|
||||
})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
d := NewDiscoverer(security.NewSafeClient(5 * time.Second))
|
||||
result, err := d.Discover(context.Background(), srv.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("Discover() OIDC fallback error: %v", err)
|
||||
}
|
||||
if result.AuthorizationEndpoint != "https://oidc.example.com/authorize" {
|
||||
t.Errorf("AuthorizationEndpoint = %q, want OIDC fallback", result.AuthorizationEndpoint)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverAllEndpointsFail(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
// Server that returns 404 for all well-known paths.
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
d := NewDiscoverer(security.NewSafeClient(5 * time.Second))
|
||||
_, err := d.Discover(context.Background(), srv.URL)
|
||||
if err == nil {
|
||||
t.Fatal("expected error when all discovery endpoints fail, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverMissingRequiredFields(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
// Server responds with JSON missing authorization_endpoint.
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, ".well-known") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
// Missing authorization_endpoint and token_endpoint.
|
||||
_ = json.NewEncoder(w).Encode(map[string]string{"issuer": "https://example.com"})
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
d := NewDiscoverer(security.NewSafeClient(5 * time.Second))
|
||||
_, err := d.Discover(context.Background(), srv.URL)
|
||||
if err == nil {
|
||||
t.Fatal("expected error when required fields missing, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverInvalidURL(t *testing.T) {
|
||||
d := NewDiscoverer(http.DefaultClient)
|
||||
_, err := d.Discover(nil, "://invalid-url") //nolint:staticcheck
|
||||
if err == nil {
|
||||
t.Fatal("expected error for invalid URL, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverResponseTooLarge(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, ".well-known") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
// > 64KB response — will be truncated, causing JSON parse error.
|
||||
large := make([]byte, 65*1024)
|
||||
for i := range large {
|
||||
large[i] = 'x'
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"authorization_endpoint":"`))
|
||||
_, _ = w.Write(large)
|
||||
_, _ = w.Write([]byte(`","token_endpoint":"https://t.example.com/token"}`))
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
d := NewDiscoverer(security.NewSafeClient(5 * time.Second))
|
||||
// Truncated response → JSON will be malformed → discovery continues but
|
||||
// this candidate fails. All candidates fail → error.
|
||||
_, err := d.Discover(context.Background(), srv.URL)
|
||||
if err == nil {
|
||||
t.Fatal("expected error for oversized response, got nil")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,282 @@
|
||||
package mcpoauth
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const (
|
||||
flowTTL = 10 * time.Minute
|
||||
flowCleanupEvery = time.Minute
|
||||
)
|
||||
|
||||
// OAuthTokens is the token response from the Authorization Server.
|
||||
type OAuthTokens struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
RefreshToken string `json:"refresh_token,omitempty"`
|
||||
TokenType string `json:"token_type"`
|
||||
ExpiresIn int `json:"expires_in,omitempty"`
|
||||
Scope string `json:"scope,omitempty"`
|
||||
}
|
||||
|
||||
// PendingFlow holds state for an in-flight authorization code flow.
|
||||
type PendingFlow struct {
|
||||
ServerID uuid.UUID
|
||||
TenantID uuid.UUID
|
||||
UserID string // empty = global/tenant-level, non-empty = per-user
|
||||
InitiatingUserID string // admin who called /start — used for WS event routing
|
||||
CodeVerifier string
|
||||
UsePKCE bool // true → send code_verifier on token exchange
|
||||
RedirectURI string
|
||||
TokenEndpoint string
|
||||
ClientID string
|
||||
ClientSecret string // empty for public clients (PKCE)
|
||||
ResourceURI string // RFC 8707
|
||||
Issuer string // AS issuer — persisted with the token for status display
|
||||
CreatedAt time.Time
|
||||
}
|
||||
|
||||
// StartFlowParams holds the parameters required to begin an OAuth flow.
|
||||
type StartFlowParams struct {
|
||||
ServerID uuid.UUID
|
||||
TenantID uuid.UUID
|
||||
UserID string
|
||||
InitiatingUserID string // admin who called /start — stored for WS event routing on callback
|
||||
DiscoveryResult *DiscoveryResult
|
||||
ClientID string
|
||||
ClientSecret string
|
||||
RedirectURI string
|
||||
Scopes string
|
||||
// GrantType controls the authorization flow variant.
|
||||
// "pkce" (default/empty) — Authorization Code + PKCE S256, public client (no ClientSecret).
|
||||
// "authorization_code" — Authorization Code + PKCE S256, confidential client (sends ClientSecret).
|
||||
// "client_credentials" — not handled by StartFlow; use ClientCredentials() instead.
|
||||
GrantType string
|
||||
}
|
||||
|
||||
// FlowManager manages in-progress OAuth authorization code flows.
|
||||
type FlowManager struct {
|
||||
mu sync.Mutex
|
||||
pending map[string]*PendingFlow // state → flow
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
func NewFlowManager(client *http.Client) *FlowManager {
|
||||
fm := &FlowManager{
|
||||
pending: make(map[string]*PendingFlow),
|
||||
client: client,
|
||||
}
|
||||
go fm.cleanupLoop()
|
||||
return fm
|
||||
}
|
||||
|
||||
// StartFlow initiates an authorization code flow.
|
||||
// GrantType "pkce" — PKCE S256, public client (no client_secret).
|
||||
// GrantType "authorization_code" — PKCE S256, confidential client (sends client_secret too).
|
||||
// OAuth 2.1 mandates PKCE for ALL authorization code flows; the grant type only controls
|
||||
// whether a client_secret is included, not whether PKCE is used.
|
||||
// Returns the authorization URL and CSRF state token.
|
||||
func (fm *FlowManager) StartFlow(ctx context.Context, p StartFlowParams) (authURL, state string, err error) {
|
||||
_ = ctx
|
||||
|
||||
// OAuth 2.1 §4.1.1: PKCE is mandatory for all authorization code flows.
|
||||
// "authorization_code" differs from "pkce" only in that it sends client_secret too.
|
||||
usePKCE := p.GrantType != "client_credentials"
|
||||
|
||||
var verifier, challenge string
|
||||
if usePKCE {
|
||||
verifier, challenge, err = generatePKCE()
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("mcpoauth: generate PKCE: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
state, err = generateState()
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("mcpoauth: generate state: %w", err)
|
||||
}
|
||||
|
||||
fm.mu.Lock()
|
||||
fm.pending[state] = &PendingFlow{
|
||||
ServerID: p.ServerID,
|
||||
TenantID: p.TenantID,
|
||||
UserID: p.UserID,
|
||||
InitiatingUserID: p.InitiatingUserID,
|
||||
CodeVerifier: verifier,
|
||||
UsePKCE: usePKCE,
|
||||
RedirectURI: p.RedirectURI,
|
||||
TokenEndpoint: p.DiscoveryResult.TokenEndpoint,
|
||||
ClientID: p.ClientID,
|
||||
ClientSecret: p.ClientSecret,
|
||||
ResourceURI: p.DiscoveryResult.ResourceURI,
|
||||
Issuer: p.DiscoveryResult.Issuer,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
fm.mu.Unlock()
|
||||
|
||||
q := url.Values{
|
||||
"response_type": {"code"},
|
||||
"client_id": {p.ClientID},
|
||||
"redirect_uri": {p.RedirectURI},
|
||||
"state": {state},
|
||||
}
|
||||
if usePKCE {
|
||||
q.Set("code_challenge", challenge)
|
||||
q.Set("code_challenge_method", "S256")
|
||||
}
|
||||
if p.Scopes != "" {
|
||||
q.Set("scope", p.Scopes)
|
||||
}
|
||||
if p.DiscoveryResult.ResourceURI != "" {
|
||||
q.Set("resource", p.DiscoveryResult.ResourceURI) // RFC 8707
|
||||
}
|
||||
|
||||
authURL = p.DiscoveryResult.AuthorizationEndpoint + "?" + q.Encode()
|
||||
return authURL, state, nil
|
||||
}
|
||||
|
||||
// ExchangeCode validates the CSRF state, exchanges the authorization code for tokens,
|
||||
// and returns the tokens along with the original flow parameters.
|
||||
func (fm *FlowManager) ExchangeCode(ctx context.Context, state, code string) (*OAuthTokens, *PendingFlow, error) {
|
||||
fm.mu.Lock()
|
||||
flow, ok := fm.pending[state]
|
||||
if ok {
|
||||
delete(fm.pending, state)
|
||||
}
|
||||
fm.mu.Unlock()
|
||||
|
||||
if !ok || time.Since(flow.CreatedAt) > flowTTL {
|
||||
return nil, nil, fmt.Errorf("mcpoauth: invalid or expired OAuth state")
|
||||
}
|
||||
|
||||
tokens, err := fm.exchangeCodeHTTP(ctx, flow, code)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return tokens, flow, nil
|
||||
}
|
||||
|
||||
func (fm *FlowManager) exchangeCodeHTTP(ctx context.Context, flow *PendingFlow, code string) (*OAuthTokens, error) {
|
||||
params := url.Values{
|
||||
"grant_type": {"authorization_code"},
|
||||
"code": {code},
|
||||
"redirect_uri": {flow.RedirectURI},
|
||||
"client_id": {flow.ClientID},
|
||||
}
|
||||
if flow.UsePKCE {
|
||||
params.Set("code_verifier", flow.CodeVerifier)
|
||||
}
|
||||
if flow.ClientSecret != "" {
|
||||
params.Set("client_secret", flow.ClientSecret)
|
||||
}
|
||||
if flow.ResourceURI != "" {
|
||||
params.Set("resource", flow.ResourceURI) // RFC 8707
|
||||
}
|
||||
|
||||
return postTokenRequest(ctx, fm.client, flow.TokenEndpoint, params)
|
||||
}
|
||||
|
||||
// ClientCredentials performs a client credentials grant (no user interaction).
|
||||
func (fm *FlowManager) ClientCredentials(ctx context.Context, tokenEndpoint, clientID, clientSecret, scopes, resourceURI string) (*OAuthTokens, error) {
|
||||
params := url.Values{
|
||||
"grant_type": {"client_credentials"},
|
||||
"client_id": {clientID},
|
||||
}
|
||||
if clientSecret != "" {
|
||||
params.Set("client_secret", clientSecret)
|
||||
}
|
||||
if scopes != "" {
|
||||
params.Set("scope", scopes)
|
||||
}
|
||||
if resourceURI != "" {
|
||||
params.Set("resource", resourceURI)
|
||||
}
|
||||
return postTokenRequest(ctx, fm.client, tokenEndpoint, params)
|
||||
}
|
||||
|
||||
func (fm *FlowManager) cleanupLoop() {
|
||||
ticker := time.NewTicker(flowCleanupEvery)
|
||||
defer ticker.Stop()
|
||||
for range ticker.C {
|
||||
cutoff := time.Now().Add(-flowTTL)
|
||||
fm.mu.Lock()
|
||||
for state, flow := range fm.pending {
|
||||
if flow.CreatedAt.Before(cutoff) {
|
||||
delete(fm.pending, state)
|
||||
}
|
||||
}
|
||||
fm.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
// postTokenRequest sends a POST to the token endpoint with form-encoded params.
|
||||
func postTokenRequest(ctx context.Context, client *http.Client, tokenEndpoint string, params url.Values) (*OAuthTokens, error) {
|
||||
_, ip, err := security.Validate(tokenEndpoint)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("mcpoauth: SSRF validation failed for token endpoint: %w", err)
|
||||
}
|
||||
ctx = security.WithPinnedIP(ctx, ip)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, tokenEndpoint, bytes.NewBufferString(params.Encode()))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
req.Header.Set("Accept", "application/json")
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("mcpoauth: token request failed: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("mcpoauth: token endpoint returned HTTP %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
var tokens OAuthTokens
|
||||
if err := json.Unmarshal(body, &tokens); err != nil {
|
||||
return nil, fmt.Errorf("mcpoauth: token response parse error: %w", err)
|
||||
}
|
||||
if tokens.AccessToken == "" {
|
||||
return nil, fmt.Errorf("mcpoauth: token response missing access_token")
|
||||
}
|
||||
if tokens.TokenType == "" {
|
||||
tokens.TokenType = "Bearer"
|
||||
}
|
||||
return &tokens, nil
|
||||
}
|
||||
|
||||
// generatePKCE returns (code_verifier, code_challenge_S256).
|
||||
func generatePKCE() (verifier, challenge string, err error) {
|
||||
buf := make([]byte, 32)
|
||||
if _, err = rand.Read(buf); err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
verifier = base64.RawURLEncoding.EncodeToString(buf)
|
||||
h := sha256.Sum256([]byte(verifier))
|
||||
challenge = base64.RawURLEncoding.EncodeToString(h[:])
|
||||
return verifier, challenge, nil
|
||||
}
|
||||
|
||||
// generateState returns a cryptographically random CSRF state token.
|
||||
func generateState() (string, error) {
|
||||
buf := make([]byte, 16)
|
||||
if _, err := rand.Read(buf); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return base64.RawURLEncoding.EncodeToString(buf), nil
|
||||
}
|
||||
@@ -0,0 +1,313 @@
|
||||
package mcpoauth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
)
|
||||
|
||||
func TestPKCEVerifierMatchesChallenge(t *testing.T) {
|
||||
verifier, challenge, err := generatePKCE()
|
||||
if err != nil {
|
||||
t.Fatalf("generatePKCE() error: %v", err)
|
||||
}
|
||||
if verifier == "" || challenge == "" {
|
||||
t.Fatal("expected non-empty verifier and challenge")
|
||||
}
|
||||
h := sha256.Sum256([]byte(verifier))
|
||||
want := base64.RawURLEncoding.EncodeToString(h[:])
|
||||
if challenge != want {
|
||||
t.Errorf("challenge = %q, want %q", challenge, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGenerateStateLength(t *testing.T) {
|
||||
state, err := generateState()
|
||||
if err != nil {
|
||||
t.Fatalf("generateState() error: %v", err)
|
||||
}
|
||||
// 16 bytes base64url-encoded → at least 21 chars (no padding).
|
||||
if len(state) < 20 {
|
||||
t.Errorf("state too short: %q (len %d)", state, len(state))
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartFlowReturnsPKCEAuthURL(t *testing.T) {
|
||||
fm := NewFlowManager(http.DefaultClient)
|
||||
disc := &DiscoveryResult{
|
||||
AuthorizationEndpoint: "https://auth.example.com/authorize",
|
||||
TokenEndpoint: "https://auth.example.com/token",
|
||||
}
|
||||
authURL, state, err := fm.StartFlow(context.Background(), StartFlowParams{
|
||||
ServerID: uuid.New(),
|
||||
TenantID: uuid.New(),
|
||||
DiscoveryResult: disc,
|
||||
ClientID: "test-client",
|
||||
RedirectURI: "https://goclaw.example.com/v1/mcp/oauth/callback",
|
||||
GrantType: "pkce",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StartFlow() error: %v", err)
|
||||
}
|
||||
if authURL == "" {
|
||||
t.Error("expected non-empty authURL")
|
||||
}
|
||||
if state == "" || len(state) < 10 {
|
||||
t.Errorf("expected valid state token, got %q", state)
|
||||
}
|
||||
if !strings.Contains(authURL, "code_challenge=") {
|
||||
t.Errorf("authURL missing code_challenge: %s", authURL)
|
||||
}
|
||||
if !strings.Contains(authURL, "code_challenge_method=S256") {
|
||||
t.Errorf("authURL missing S256 method: %s", authURL)
|
||||
}
|
||||
if !strings.Contains(authURL, "state="+state) {
|
||||
t.Errorf("authURL missing state param: %s", authURL)
|
||||
}
|
||||
if !strings.Contains(authURL, "client_id=test-client") {
|
||||
t.Errorf("authURL missing client_id: %s", authURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartFlowStoresInPending(t *testing.T) {
|
||||
fm := NewFlowManager(http.DefaultClient)
|
||||
disc := &DiscoveryResult{
|
||||
AuthorizationEndpoint: "https://auth.example.com/authorize",
|
||||
TokenEndpoint: "https://auth.example.com/token",
|
||||
}
|
||||
_, state, err := fm.StartFlow(context.Background(), StartFlowParams{
|
||||
ServerID: uuid.New(), TenantID: uuid.New(),
|
||||
DiscoveryResult: disc, ClientID: "c",
|
||||
RedirectURI: "https://example.com/cb",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StartFlow() error: %v", err)
|
||||
}
|
||||
fm.mu.Lock()
|
||||
flow, ok := fm.pending[state]
|
||||
fm.mu.Unlock()
|
||||
if !ok {
|
||||
t.Fatal("pending flow not found")
|
||||
}
|
||||
if flow.ClientID != "c" {
|
||||
t.Errorf("ClientID = %q, want %q", flow.ClientID, "c")
|
||||
}
|
||||
if !flow.UsePKCE {
|
||||
t.Error("expected UsePKCE = true for pkce grant type")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExchangeCodeSuccess(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
if err := r.ParseForm(); err != nil {
|
||||
http.Error(w, "bad form", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if r.FormValue("grant_type") != "authorization_code" {
|
||||
http.Error(w, "wrong grant_type: "+r.FormValue("grant_type"), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if r.FormValue("code_verifier") == "" {
|
||||
http.Error(w, "missing code_verifier", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(OAuthTokens{
|
||||
AccessToken: "access-token-123",
|
||||
RefreshToken: "refresh-token-456",
|
||||
TokenType: "Bearer",
|
||||
ExpiresIn: 3600,
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
client := security.NewSafeClient(5 * time.Second)
|
||||
fm := NewFlowManager(client)
|
||||
disc := &DiscoveryResult{
|
||||
AuthorizationEndpoint: "https://auth.example.com/authorize",
|
||||
TokenEndpoint: srv.URL,
|
||||
}
|
||||
_, state, err := fm.StartFlow(context.Background(), StartFlowParams{
|
||||
ServerID: uuid.New(),
|
||||
TenantID: uuid.New(),
|
||||
DiscoveryResult: disc,
|
||||
ClientID: "test-client",
|
||||
RedirectURI: srv.URL + "/callback",
|
||||
GrantType: "pkce",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StartFlow() error: %v", err)
|
||||
}
|
||||
|
||||
tokens, flow, err := fm.ExchangeCode(context.Background(), state, "auth-code-xyz")
|
||||
if err != nil {
|
||||
t.Fatalf("ExchangeCode() error: %v", err)
|
||||
}
|
||||
if tokens.AccessToken != "access-token-123" {
|
||||
t.Errorf("access_token = %q, want %q", tokens.AccessToken, "access-token-123")
|
||||
}
|
||||
if tokens.RefreshToken != "refresh-token-456" {
|
||||
t.Errorf("refresh_token = %q, want %q", tokens.RefreshToken, "refresh-token-456")
|
||||
}
|
||||
if flow == nil {
|
||||
t.Fatal("expected non-nil PendingFlow")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExchangeCodeInvalidState(t *testing.T) {
|
||||
fm := NewFlowManager(http.DefaultClient)
|
||||
_, _, err := fm.ExchangeCode(context.Background(), "no-such-state", "code")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for unknown state, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExchangeCodeExpiredFlow(t *testing.T) {
|
||||
fm := NewFlowManager(http.DefaultClient)
|
||||
fm.mu.Lock()
|
||||
fm.pending["expired-state"] = &PendingFlow{
|
||||
CreatedAt: time.Now().Add(-11 * time.Minute), // > flowTTL (10 min)
|
||||
}
|
||||
fm.mu.Unlock()
|
||||
|
||||
_, _, err := fm.ExchangeCode(context.Background(), "expired-state", "code")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for expired flow, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExchangeCodeReplayPrevented(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(OAuthTokens{AccessToken: "tok", TokenType: "Bearer"})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
fm := NewFlowManager(security.NewSafeClient(5 * time.Second))
|
||||
disc := &DiscoveryResult{
|
||||
AuthorizationEndpoint: "https://auth.example.com/authorize",
|
||||
TokenEndpoint: srv.URL,
|
||||
}
|
||||
_, state, err := fm.StartFlow(context.Background(), StartFlowParams{
|
||||
ServerID: uuid.New(), TenantID: uuid.New(),
|
||||
DiscoveryResult: disc, ClientID: "c", RedirectURI: srv.URL + "/cb",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StartFlow() error: %v", err)
|
||||
}
|
||||
|
||||
// First exchange: succeeds.
|
||||
if _, _, err := fm.ExchangeCode(context.Background(), state, "code"); err != nil {
|
||||
t.Fatalf("first ExchangeCode() failed: %v", err)
|
||||
}
|
||||
// Second exchange with same state: state was deleted → must fail.
|
||||
if _, _, err := fm.ExchangeCode(context.Background(), state, "code"); err == nil {
|
||||
t.Fatal("expected error on replay, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExchangeCodeTokenEndpointError(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Error(w, `{"error":"invalid_grant"}`, http.StatusBadRequest)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
fm := NewFlowManager(security.NewSafeClient(5 * time.Second))
|
||||
disc := &DiscoveryResult{
|
||||
AuthorizationEndpoint: "https://auth.example.com/authorize",
|
||||
TokenEndpoint: srv.URL,
|
||||
}
|
||||
_, state, err := fm.StartFlow(context.Background(), StartFlowParams{
|
||||
ServerID: uuid.New(), TenantID: uuid.New(),
|
||||
DiscoveryResult: disc, ClientID: "c", RedirectURI: srv.URL + "/cb",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StartFlow() error: %v", err)
|
||||
}
|
||||
|
||||
_, _, err = fm.ExchangeCode(context.Background(), state, "bad-code")
|
||||
if err == nil {
|
||||
t.Fatal("expected error when token endpoint returns 400, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientCredentials(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := r.ParseForm(); err != nil {
|
||||
http.Error(w, "bad form", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if r.FormValue("grant_type") != "client_credentials" {
|
||||
http.Error(w, "wrong grant_type: "+r.FormValue("grant_type"), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(OAuthTokens{AccessToken: "cc-token", TokenType: "Bearer"})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
fm := NewFlowManager(security.NewSafeClient(5 * time.Second))
|
||||
tokens, err := fm.ClientCredentials(context.Background(), srv.URL, "client-id", "client-secret", "read write", "")
|
||||
if err != nil {
|
||||
t.Fatalf("ClientCredentials() error: %v", err)
|
||||
}
|
||||
if tokens.AccessToken != "cc-token" {
|
||||
t.Errorf("access_token = %q, want %q", tokens.AccessToken, "cc-token")
|
||||
}
|
||||
}
|
||||
|
||||
func TestFlowCleanupPurgesExpired(t *testing.T) {
|
||||
fm := NewFlowManager(http.DefaultClient)
|
||||
|
||||
fm.mu.Lock()
|
||||
fm.pending["fresh"] = &PendingFlow{CreatedAt: time.Now()}
|
||||
fm.pending["old"] = &PendingFlow{CreatedAt: time.Now().Add(-11 * time.Minute)}
|
||||
fm.mu.Unlock()
|
||||
|
||||
// Simulate cleanup logic directly (same as cleanupLoop body).
|
||||
cutoff := time.Now().Add(-flowTTL)
|
||||
fm.mu.Lock()
|
||||
for state, flow := range fm.pending {
|
||||
if flow.CreatedAt.Before(cutoff) {
|
||||
delete(fm.pending, state)
|
||||
}
|
||||
}
|
||||
fm.mu.Unlock()
|
||||
|
||||
fm.mu.Lock()
|
||||
_, freshOk := fm.pending["fresh"]
|
||||
_, oldOk := fm.pending["old"]
|
||||
fm.mu.Unlock()
|
||||
|
||||
if !freshOk {
|
||||
t.Error("fresh flow should not be purged")
|
||||
}
|
||||
if oldOk {
|
||||
t.Error("expired flow should be purged by cleanup")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
package mcpoauth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
"golang.org/x/sync/singleflight"
|
||||
)
|
||||
|
||||
const tokenRefreshBuffer = 60 * time.Second
|
||||
|
||||
// ErrTokenExpired is returned when the token is expired and cannot be refreshed.
|
||||
var ErrTokenExpired = errors.New("mcpoauth: token expired and no refresh token available")
|
||||
|
||||
type cachedToken struct {
|
||||
accessToken string
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
// Refresher provides valid OAuth access tokens for MCP servers, handling
|
||||
// in-memory caching and transparent refresh using stored refresh tokens.
|
||||
type Refresher struct {
|
||||
store store.MCPOAuthTokenStore
|
||||
client *http.Client
|
||||
mu sync.Mutex
|
||||
cache map[string]*cachedToken // key: serverID+":"+userID
|
||||
// sf coalesces concurrent DB-load + refresh for the same slot so a rotating
|
||||
// refresh token (RFC 9700) is never consumed twice by parallel agent turns.
|
||||
sf singleflight.Group
|
||||
}
|
||||
|
||||
// NewRefresher builds a Refresher. Token decryption happens in the store layer,
|
||||
// so the Refresher itself never touches the encryption key.
|
||||
func NewRefresher(s store.MCPOAuthTokenStore, client *http.Client) *Refresher {
|
||||
return &Refresher{
|
||||
store: s,
|
||||
client: client,
|
||||
cache: make(map[string]*cachedToken),
|
||||
}
|
||||
}
|
||||
|
||||
// GetValidToken returns a valid access token for the given server.
|
||||
// userID="" uses the global (tenant-level) token; non-empty uses the per-user token.
|
||||
func (r *Refresher) GetValidToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) (string, error) {
|
||||
cacheKey := serverID.String() + ":" + userID
|
||||
|
||||
if tok, ok := r.cachedValid(cacheKey); ok {
|
||||
return tok, nil
|
||||
}
|
||||
|
||||
// Coalesce concurrent DB-load + refresh for the same slot. Without this,
|
||||
// two parallel agent turns hitting an expired token both POST a refresh and
|
||||
// the second one replays an already-rotated refresh token, killing the slot.
|
||||
v, err, _ := r.sf.Do(cacheKey, func() (any, error) {
|
||||
// Double-check: a coalesced caller ahead of us may have just refreshed.
|
||||
if tok, ok := r.cachedValid(cacheKey); ok {
|
||||
return tok, nil
|
||||
}
|
||||
return r.loadAndRefresh(ctx, serverID, tenantID, userID, cacheKey)
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return v.(string), nil
|
||||
}
|
||||
|
||||
// cachedValid returns the cached access token when it is still inside the
|
||||
// refresh buffer, false otherwise.
|
||||
func (r *Refresher) cachedValid(cacheKey string) (string, bool) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if c, ok := r.cache[cacheKey]; ok && time.Now().Add(tokenRefreshBuffer).Before(c.expiresAt) {
|
||||
return c.accessToken, true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// loadAndRefresh loads the token from the store and refreshes it if expired.
|
||||
// Runs under the single-flight lock for its cache key.
|
||||
func (r *Refresher) loadAndRefresh(ctx context.Context, serverID, tenantID uuid.UUID, userID, cacheKey string) (string, error) {
|
||||
tok, err := r.loadToken(ctx, serverID, tenantID, userID)
|
||||
if err != nil {
|
||||
slog.Warn("mcpoauth: load token failed", "server_id", serverID, "tenant_id", tenantID, "user_id", userID, "error", err)
|
||||
return "", err
|
||||
}
|
||||
if tok == nil {
|
||||
slog.Debug("mcpoauth: no token in db", "server_id", serverID, "tenant_id", tenantID, "user_id", userID)
|
||||
return "", ErrTokenExpired
|
||||
}
|
||||
|
||||
// No expiry advertised → treat as non-expiring. RFC 6749 §5.1 makes
|
||||
// expires_in OPTIONAL, so a token without an expiry (common for
|
||||
// client_credentials and some long-lived grants) is still valid and must
|
||||
// not be forced through the refresh/ErrTokenExpired path.
|
||||
if tok.ExpiresAt == nil {
|
||||
r.setCache(cacheKey, tok.AccessToken, nil)
|
||||
return tok.AccessToken, nil
|
||||
}
|
||||
|
||||
// Valid in DB?
|
||||
if time.Now().Add(tokenRefreshBuffer).Before(*tok.ExpiresAt) {
|
||||
r.setCache(cacheKey, tok.AccessToken, tok.ExpiresAt)
|
||||
return tok.AccessToken, nil
|
||||
}
|
||||
|
||||
// Expired — try refresh.
|
||||
if tok.RefreshToken == "" {
|
||||
return "", ErrTokenExpired
|
||||
}
|
||||
newTokens, err := r.refresh(ctx, tok)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("mcpoauth: token refresh failed: %w", err)
|
||||
}
|
||||
|
||||
// Persist refreshed tokens.
|
||||
now := time.Now()
|
||||
var expiresAt *time.Time
|
||||
if newTokens.ExpiresIn > 0 {
|
||||
t := now.Add(time.Duration(newTokens.ExpiresIn) * time.Second)
|
||||
expiresAt = &t
|
||||
}
|
||||
updated := *tok
|
||||
updated.AccessToken = newTokens.AccessToken
|
||||
if newTokens.RefreshToken != "" {
|
||||
updated.RefreshToken = newTokens.RefreshToken
|
||||
}
|
||||
updated.ExpiresAt = expiresAt
|
||||
updated.IssuedAt = &now
|
||||
if err := r.store.UpsertOAuthToken(ctx, &updated); err != nil {
|
||||
return "", fmt.Errorf("mcpoauth: save refreshed token: %w", err)
|
||||
}
|
||||
|
||||
r.setCache(cacheKey, newTokens.AccessToken, expiresAt)
|
||||
return newTokens.AccessToken, nil
|
||||
}
|
||||
|
||||
// InvalidateCache evicts the in-memory cached token for the given server/user.
|
||||
func (r *Refresher) InvalidateCache(serverID uuid.UUID, userID string) {
|
||||
key := serverID.String() + ":" + userID
|
||||
r.mu.Lock()
|
||||
delete(r.cache, key)
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
// InvalidateServer evicts every cached token for the given server (global +
|
||||
// all per-user slots). Cache keys are "<serverID>:<userID>", so all entries
|
||||
// sharing the server prefix are dropped. Called when the server URL or OAuth
|
||||
// config changes and the stored tokens have been purged.
|
||||
func (r *Refresher) InvalidateServer(serverID uuid.UUID) {
|
||||
prefix := serverID.String() + ":"
|
||||
r.mu.Lock()
|
||||
for key := range r.cache {
|
||||
if strings.HasPrefix(key, prefix) {
|
||||
delete(r.cache, key)
|
||||
}
|
||||
}
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *Refresher) loadToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) (*store.MCPOAuthToken, error) {
|
||||
if userID == "" {
|
||||
return r.store.GetOAuthToken(ctx, serverID, tenantID)
|
||||
}
|
||||
return r.store.GetUserOAuthToken(ctx, serverID, tenantID, userID)
|
||||
}
|
||||
|
||||
func (r *Refresher) refresh(ctx context.Context, tok *store.MCPOAuthToken) (*OAuthTokens, error) {
|
||||
params := url.Values{
|
||||
"grant_type": {"refresh_token"},
|
||||
"refresh_token": {tok.RefreshToken},
|
||||
"client_id": {tok.DCRClientID},
|
||||
}
|
||||
if tok.DCRClientSecret != "" {
|
||||
params.Set("client_secret", tok.DCRClientSecret)
|
||||
}
|
||||
if tok.ResourceURI != "" {
|
||||
params.Set("resource", tok.ResourceURI)
|
||||
}
|
||||
return postTokenRequest(ctx, r.client, tok.TokenEndpoint, params)
|
||||
}
|
||||
|
||||
func (r *Refresher) setCache(key, accessToken string, expiresAt *time.Time) {
|
||||
c := &cachedToken{accessToken: accessToken}
|
||||
if expiresAt != nil {
|
||||
c.expiresAt = *expiresAt
|
||||
} else {
|
||||
c.expiresAt = time.Now().Add(time.Hour)
|
||||
}
|
||||
r.mu.Lock()
|
||||
r.cache[key] = c
|
||||
r.mu.Unlock()
|
||||
}
|
||||
@@ -0,0 +1,357 @@
|
||||
package mcpoauth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
// mockOAuthStore is a minimal in-memory implementation of store.MCPOAuthTokenStore.
|
||||
type mockOAuthStore struct {
|
||||
global map[string]*store.MCPOAuthToken // key: serverID
|
||||
perUser map[string]*store.MCPOAuthToken // key: serverID+":"+userID
|
||||
upserts int
|
||||
}
|
||||
|
||||
func newMockOAuthStore() *mockOAuthStore {
|
||||
return &mockOAuthStore{
|
||||
global: make(map[string]*store.MCPOAuthToken),
|
||||
perUser: make(map[string]*store.MCPOAuthToken),
|
||||
}
|
||||
}
|
||||
|
||||
func (m *mockOAuthStore) GetOAuthToken(_ context.Context, serverID, _ uuid.UUID) (*store.MCPOAuthToken, error) {
|
||||
tok, ok := m.global[serverID.String()]
|
||||
if !ok {
|
||||
return nil, nil
|
||||
}
|
||||
cp := *tok
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func (m *mockOAuthStore) GetUserOAuthToken(_ context.Context, serverID, _ uuid.UUID, userID string) (*store.MCPOAuthToken, error) {
|
||||
tok, ok := m.perUser[serverID.String()+":"+userID]
|
||||
if !ok {
|
||||
return nil, nil
|
||||
}
|
||||
cp := *tok
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func (m *mockOAuthStore) UpsertOAuthToken(_ context.Context, tok *store.MCPOAuthToken) error {
|
||||
m.upserts++
|
||||
cp := *tok
|
||||
if tok.UserID == "" {
|
||||
m.global[tok.ServerID.String()] = &cp
|
||||
} else {
|
||||
m.perUser[tok.ServerID.String()+":"+tok.UserID] = &cp
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockOAuthStore) DeleteOAuthToken(_ context.Context, serverID, _ uuid.UUID) error {
|
||||
delete(m.global, serverID.String())
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockOAuthStore) DeleteUserOAuthToken(_ context.Context, serverID, _ uuid.UUID, userID string) error {
|
||||
delete(m.perUser, serverID.String()+":"+userID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mockOAuthStore) DeleteServerOAuthTokens(_ context.Context, serverID, _ uuid.UUID) error {
|
||||
delete(m.global, serverID.String())
|
||||
prefix := serverID.String() + ":"
|
||||
for k := range m.perUser {
|
||||
if strings.HasPrefix(k, prefix) {
|
||||
delete(m.perUser, k)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// futureTime returns a time this far in the future.
|
||||
func futureTime(d time.Duration) *time.Time {
|
||||
t := time.Now().Add(d)
|
||||
return &t
|
||||
}
|
||||
|
||||
// pastTime returns a time this far in the past.
|
||||
func pastTime(d time.Duration) *time.Time {
|
||||
t := time.Now().Add(-d)
|
||||
return &t
|
||||
}
|
||||
|
||||
func TestGetValidTokenCacheHit(t *testing.T) {
|
||||
st := newMockOAuthStore()
|
||||
r := NewRefresher(st, http.DefaultClient)
|
||||
|
||||
serverID := uuid.New()
|
||||
tenantID := uuid.New()
|
||||
|
||||
// Seed the store with a valid token.
|
||||
validTok := &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: "cached-token",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: futureTime(2 * time.Hour),
|
||||
}
|
||||
st.global[serverID.String()] = validTok
|
||||
|
||||
// First call — loads from store and caches.
|
||||
tok1, err := r.GetValidToken(context.Background(), serverID, tenantID, "")
|
||||
if err != nil {
|
||||
t.Fatalf("first GetValidToken() error: %v", err)
|
||||
}
|
||||
if tok1 != "cached-token" {
|
||||
t.Errorf("token = %q, want %q", tok1, "cached-token")
|
||||
}
|
||||
|
||||
// Remove from store to prove second call uses cache.
|
||||
delete(st.global, serverID.String())
|
||||
|
||||
tok2, err := r.GetValidToken(context.Background(), serverID, tenantID, "")
|
||||
if err != nil {
|
||||
t.Fatalf("second GetValidToken() error: %v", err)
|
||||
}
|
||||
if tok2 != "cached-token" {
|
||||
t.Errorf("second token = %q, want %q (from cache)", tok2, "cached-token")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetValidTokenCacheMiss(t *testing.T) {
|
||||
st := newMockOAuthStore()
|
||||
r := NewRefresher(st, http.DefaultClient)
|
||||
|
||||
serverID := uuid.New()
|
||||
tenantID := uuid.New()
|
||||
|
||||
tok := &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: "fresh-from-db",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: futureTime(30 * time.Minute),
|
||||
}
|
||||
st.global[serverID.String()] = tok
|
||||
|
||||
got, err := r.GetValidToken(context.Background(), serverID, tenantID, "")
|
||||
if err != nil {
|
||||
t.Fatalf("GetValidToken() error: %v", err)
|
||||
}
|
||||
if got != "fresh-from-db" {
|
||||
t.Errorf("token = %q, want %q", got, "fresh-from-db")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetValidTokenNoRecord(t *testing.T) {
|
||||
st := newMockOAuthStore()
|
||||
r := NewRefresher(st, http.DefaultClient)
|
||||
|
||||
_, err := r.GetValidToken(context.Background(), uuid.New(), uuid.New(), "")
|
||||
if !errors.Is(err, ErrTokenExpired) {
|
||||
t.Errorf("expected ErrTokenExpired, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetValidTokenExpiredNoRefreshToken(t *testing.T) {
|
||||
st := newMockOAuthStore()
|
||||
r := NewRefresher(st, http.DefaultClient)
|
||||
|
||||
serverID := uuid.New()
|
||||
tenantID := uuid.New()
|
||||
|
||||
st.global[serverID.String()] = &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: "old-token",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: pastTime(time.Hour),
|
||||
RefreshToken: "", // no refresh token
|
||||
}
|
||||
|
||||
_, err := r.GetValidToken(context.Background(), serverID, tenantID, "")
|
||||
if !errors.Is(err, ErrTokenExpired) {
|
||||
t.Errorf("expected ErrTokenExpired for expired token without refresh, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetValidTokenRefreshSuccess(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if err := r.ParseForm(); err != nil {
|
||||
http.Error(w, "bad form", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
if r.FormValue("grant_type") != "refresh_token" {
|
||||
http.Error(w, "wrong grant_type", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(OAuthTokens{
|
||||
AccessToken: "new-access-token",
|
||||
TokenType: "Bearer",
|
||||
ExpiresIn: 3600,
|
||||
})
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
st := newMockOAuthStore()
|
||||
r := NewRefresher(st, security.NewSafeClient(5*time.Second))
|
||||
|
||||
serverID := uuid.New()
|
||||
tenantID := uuid.New()
|
||||
|
||||
st.global[serverID.String()] = &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: "expired-token",
|
||||
RefreshToken: "valid-refresh",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: pastTime(time.Hour),
|
||||
TokenEndpoint: srv.URL,
|
||||
}
|
||||
|
||||
got, err := r.GetValidToken(context.Background(), serverID, tenantID, "")
|
||||
if err != nil {
|
||||
t.Fatalf("GetValidToken() with refresh error: %v", err)
|
||||
}
|
||||
if got != "new-access-token" {
|
||||
t.Errorf("token = %q, want %q", got, "new-access-token")
|
||||
}
|
||||
if st.upserts == 0 {
|
||||
t.Error("expected UpsertOAuthToken to be called after refresh")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetValidTokenRefreshFails(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.Error(w, `{"error":"invalid_grant"}`, http.StatusBadRequest)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
st := newMockOAuthStore()
|
||||
r := NewRefresher(st, security.NewSafeClient(5*time.Second))
|
||||
|
||||
serverID := uuid.New()
|
||||
tenantID := uuid.New()
|
||||
|
||||
st.global[serverID.String()] = &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: "expired-token",
|
||||
RefreshToken: "bad-refresh",
|
||||
ExpiresAt: pastTime(time.Hour),
|
||||
TokenEndpoint: srv.URL,
|
||||
}
|
||||
|
||||
_, err := r.GetValidToken(context.Background(), serverID, tenantID, "")
|
||||
if err == nil {
|
||||
t.Fatal("expected error when refresh fails, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInvalidateCacheEvictsEntry(t *testing.T) {
|
||||
st := newMockOAuthStore()
|
||||
r := NewRefresher(st, http.DefaultClient)
|
||||
|
||||
serverID := uuid.New()
|
||||
tenantID := uuid.New()
|
||||
|
||||
tok := &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: "first-token",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: futureTime(2 * time.Hour),
|
||||
}
|
||||
st.global[serverID.String()] = tok
|
||||
|
||||
// Load into cache.
|
||||
if _, err := r.GetValidToken(context.Background(), serverID, tenantID, ""); err != nil {
|
||||
t.Fatalf("GetValidToken() error: %v", err)
|
||||
}
|
||||
|
||||
// Update store with new token.
|
||||
st.global[serverID.String()] = &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: "updated-token",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: futureTime(2 * time.Hour),
|
||||
}
|
||||
|
||||
// Without invalidation, cache returns old token.
|
||||
got, _ := r.GetValidToken(context.Background(), serverID, tenantID, "")
|
||||
if got != "first-token" {
|
||||
t.Errorf("before invalidation: got %q, want %q", got, "first-token")
|
||||
}
|
||||
|
||||
// Invalidate and retry.
|
||||
r.InvalidateCache(serverID, "")
|
||||
got, err := r.GetValidToken(context.Background(), serverID, tenantID, "")
|
||||
if err != nil {
|
||||
t.Fatalf("GetValidToken() after invalidate error: %v", err)
|
||||
}
|
||||
if got != "updated-token" {
|
||||
t.Errorf("after invalidation: got %q, want %q", got, "updated-token")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetValidTokenPerUserVsGlobal(t *testing.T) {
|
||||
st := newMockOAuthStore()
|
||||
r := NewRefresher(st, http.DefaultClient)
|
||||
|
||||
serverID := uuid.New()
|
||||
tenantID := uuid.New()
|
||||
userID := "user-123"
|
||||
|
||||
// Global token.
|
||||
st.global[serverID.String()] = &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: "global-token",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: futureTime(time.Hour),
|
||||
}
|
||||
// Per-user token.
|
||||
st.perUser[serverID.String()+":"+userID] = &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
UserID: userID,
|
||||
AccessToken: "user-token",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: futureTime(time.Hour),
|
||||
}
|
||||
|
||||
globalTok, err := r.GetValidToken(context.Background(), serverID, tenantID, "")
|
||||
if err != nil {
|
||||
t.Fatalf("global GetValidToken() error: %v", err)
|
||||
}
|
||||
userTok, err := r.GetValidToken(context.Background(), serverID, tenantID, userID)
|
||||
if err != nil {
|
||||
t.Fatalf("per-user GetValidToken() error: %v", err)
|
||||
}
|
||||
|
||||
if globalTok != "global-token" {
|
||||
t.Errorf("global token = %q, want %q", globalTok, "global-token")
|
||||
}
|
||||
if userTok != "user-token" {
|
||||
t.Errorf("per-user token = %q, want %q", userTok, "user-token")
|
||||
}
|
||||
}
|
||||
@@ -434,6 +434,75 @@ func (p *Pool) Evict(tenantID uuid.UUID, serverName string) {
|
||||
slog.Info("mcp.pool.evicted_on_rotation", "key", key)
|
||||
}
|
||||
|
||||
// EvictAllUsers closes and removes all per-user pool connections.
|
||||
// Called when MCP grants or server settings change so next AcquireUser picks
|
||||
// up fresh credentials/headers. Shared (non-user) connections are unaffected.
|
||||
func (p *Pool) EvictAllUsers() {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
for key, entry := range p.userServers {
|
||||
if entry.state.cancel != nil {
|
||||
entry.state.cancel()
|
||||
}
|
||||
if client := entry.state.clientPtr.Load(); client != nil {
|
||||
_ = client.Close()
|
||||
}
|
||||
delete(p.userServers, key)
|
||||
}
|
||||
// Reset per-server semaphores so slot counts don't leak.
|
||||
p.userSlots = make(map[string]chan struct{})
|
||||
slog.Info("mcp.pool.all_user_connections_evicted")
|
||||
}
|
||||
|
||||
// EvictServer evicts both the shared pool entry and all per-user pool entries
|
||||
// for a given server. Call this when server settings or grants change so that
|
||||
// the next AcquireUser/Acquire picks up fresh credentials and configuration.
|
||||
func (p *Pool) EvictServer(tenantID uuid.UUID, serverName string) {
|
||||
sharedKey := poolKey(tenantID, serverName)
|
||||
prefix := tenantID.String() + "/" + serverName + "/user:"
|
||||
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
// Evict shared connection.
|
||||
if entry, ok := p.servers[sharedKey]; ok {
|
||||
if entry.state.cancel != nil {
|
||||
entry.state.cancel()
|
||||
}
|
||||
if client := entry.state.clientPtr.Load(); client != nil {
|
||||
_ = client.Close()
|
||||
}
|
||||
delete(p.servers, sharedKey)
|
||||
select {
|
||||
case <-p.slot:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// Evict all per-user connections for this server.
|
||||
slotKey := userSlotKey(tenantID, serverName)
|
||||
sem := p.userSlots[slotKey]
|
||||
for key, entry := range p.userServers {
|
||||
if !strings.HasPrefix(key, prefix) {
|
||||
continue
|
||||
}
|
||||
if entry.state.cancel != nil {
|
||||
entry.state.cancel()
|
||||
}
|
||||
if client := entry.state.clientPtr.Load(); client != nil {
|
||||
_ = client.Close()
|
||||
}
|
||||
delete(p.userServers, key)
|
||||
if sem != nil {
|
||||
select {
|
||||
case <-sem:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
slog.Info("mcp.pool.server_evicted", "tenant", tenantID, "server", serverName)
|
||||
}
|
||||
|
||||
// evictLoop runs periodically to close idle connections over MaxIdle count.
|
||||
func (p *Pool) evictLoop() {
|
||||
ticker := time.NewTicker(60 * time.Second)
|
||||
|
||||
@@ -78,9 +78,49 @@ type MCPAccessInfo struct {
|
||||
|
||||
// MCPUserCredentials holds per-user credential overrides for an MCP server.
|
||||
type MCPUserCredentials struct {
|
||||
APIKey string `json:"api_key,omitempty" db:"-"` // decrypted
|
||||
Headers map[string]string `json:"headers,omitempty" db:"-"` // decrypted
|
||||
Env map[string]string `json:"env,omitempty" db:"-"` // decrypted
|
||||
APIKey string `json:"api_key,omitempty" db:"-"` // decrypted
|
||||
Headers map[string]string `json:"headers,omitempty" db:"-"` // decrypted
|
||||
Env map[string]string `json:"env,omitempty" db:"-"` // decrypted
|
||||
}
|
||||
|
||||
// MCPOAuthToken holds OAuth tokens for an MCP server, either global (UserID="")
|
||||
// or per-user (UserID set). Sensitive fields are decrypted before delivery.
|
||||
type MCPOAuthToken struct {
|
||||
ID uuid.UUID `db:"id"`
|
||||
ServerID uuid.UUID `db:"server_id"`
|
||||
TenantID uuid.UUID `db:"tenant_id"`
|
||||
UserID string `db:"user_id"` // empty = global
|
||||
AccessToken string `db:"access_token"`
|
||||
RefreshToken string `db:"refresh_token"`
|
||||
TokenType string `db:"token_type"`
|
||||
Scopes string `db:"scopes"`
|
||||
ExpiresAt *time.Time `db:"expires_at"`
|
||||
IssuedAt *time.Time `db:"issued_at"`
|
||||
DCRClientID string `db:"dcr_client_id"`
|
||||
DCRClientSecret string `db:"dcr_client_secret"` // decrypted
|
||||
DCRIssuer string `db:"dcr_issuer"`
|
||||
TokenEndpoint string `db:"token_endpoint"`
|
||||
ResourceURI string `db:"resource_uri"`
|
||||
CreatedAt time.Time `db:"created_at"`
|
||||
UpdatedAt time.Time `db:"updated_at"`
|
||||
}
|
||||
|
||||
// MCPOAuthTokenStore manages OAuth tokens for MCP servers.
|
||||
type MCPOAuthTokenStore interface {
|
||||
// GetOAuthToken returns the global (tenant-level) OAuth token for a server.
|
||||
GetOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID) (*MCPOAuthToken, error)
|
||||
// GetUserOAuthToken returns the per-user OAuth token for a server.
|
||||
GetUserOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) (*MCPOAuthToken, error)
|
||||
// UpsertOAuthToken inserts or replaces an OAuth token record.
|
||||
UpsertOAuthToken(ctx context.Context, token *MCPOAuthToken) error
|
||||
// DeleteOAuthToken deletes the global OAuth token for a server.
|
||||
DeleteOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID) error
|
||||
// DeleteUserOAuthToken deletes the per-user OAuth token for a server.
|
||||
DeleteUserOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) error
|
||||
// DeleteServerOAuthTokens deletes ALL OAuth tokens for a server (global +
|
||||
// every per-user row). Used when the server URL or OAuth config changes and
|
||||
// the previously minted tokens are no longer valid for the new resource/AS.
|
||||
DeleteServerOAuthTokens(ctx context.Context, serverID, tenantID uuid.UUID) error
|
||||
}
|
||||
|
||||
// MCPServerStore manages MCP server configs and access grants.
|
||||
|
||||
@@ -35,6 +35,7 @@ func NewPGStores(cfg store.StoreConfig) (*store.Stores, error) {
|
||||
Tracing: NewPGTracingStore(db),
|
||||
RunTimeline: NewPGRunTimelineStore(db),
|
||||
MCP: NewPGMCPServerStore(db, cfg.EncryptionKey),
|
||||
MCPOAuthTokens: NewPGMCPOAuthTokenStore(db, cfg.EncryptionKey),
|
||||
ChannelInstances: NewPGChannelInstanceStore(db, cfg.EncryptionKey),
|
||||
ConfigSecrets: NewPGConfigSecretsStore(db, cfg.EncryptionKey),
|
||||
AgentLinks: NewPGAgentLinkStore(db),
|
||||
|
||||
@@ -0,0 +1,226 @@
|
||||
package pg
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/crypto"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
// PGMCPOAuthTokenStore implements store.MCPOAuthTokenStore backed by Postgres.
|
||||
type PGMCPOAuthTokenStore struct {
|
||||
db *sql.DB
|
||||
encKey string
|
||||
}
|
||||
|
||||
func NewPGMCPOAuthTokenStore(db *sql.DB, encryptionKey string) *PGMCPOAuthTokenStore {
|
||||
return &PGMCPOAuthTokenStore{db: db, encKey: encryptionKey}
|
||||
}
|
||||
|
||||
func (s *PGMCPOAuthTokenStore) GetOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID) (*store.MCPOAuthToken, error) {
|
||||
return s.getToken(ctx, serverID, tenantID, "")
|
||||
}
|
||||
|
||||
func (s *PGMCPOAuthTokenStore) GetUserOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) (*store.MCPOAuthToken, error) {
|
||||
return s.getToken(ctx, serverID, tenantID, userID)
|
||||
}
|
||||
|
||||
func (s *PGMCPOAuthTokenStore) getToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) (*store.MCPOAuthToken, error) {
|
||||
var row store.MCPOAuthToken
|
||||
var accessEnc, refreshEnc, secretEnc sql.NullString
|
||||
var expiresAt, issuedAt sql.NullTime
|
||||
var scopes, resourceURI sql.NullString
|
||||
|
||||
var query string
|
||||
var args []any
|
||||
if userID == "" {
|
||||
query = `SELECT id, server_id, tenant_id, COALESCE(user_id,''), access_token, refresh_token,
|
||||
token_type, COALESCE(scopes,''), expires_at, issued_at,
|
||||
dcr_client_id, dcr_client_secret, dcr_issuer, token_endpoint,
|
||||
COALESCE(resource_uri,''), created_at, updated_at
|
||||
FROM mcp_oauth_tokens
|
||||
WHERE server_id=$1 AND tenant_id=$2 AND user_id IS NULL`
|
||||
args = []any{serverID, tenantID}
|
||||
} else {
|
||||
query = `SELECT id, server_id, tenant_id, COALESCE(user_id,''), access_token, refresh_token,
|
||||
token_type, COALESCE(scopes,''), expires_at, issued_at,
|
||||
dcr_client_id, dcr_client_secret, dcr_issuer, token_endpoint,
|
||||
COALESCE(resource_uri,''), created_at, updated_at
|
||||
FROM mcp_oauth_tokens
|
||||
WHERE server_id=$1 AND tenant_id=$2 AND user_id=$3`
|
||||
args = []any{serverID, tenantID, userID}
|
||||
}
|
||||
|
||||
err := s.db.QueryRowContext(ctx, query, args...).Scan(
|
||||
&row.ID, &row.ServerID, &row.TenantID, &row.UserID,
|
||||
&accessEnc, &refreshEnc,
|
||||
&row.TokenType, &scopes, &expiresAt, &issuedAt,
|
||||
&row.DCRClientID, &secretEnc, &row.DCRIssuer, &row.TokenEndpoint,
|
||||
&resourceURI, &row.CreatedAt, &row.UpdatedAt,
|
||||
)
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if scopes.Valid {
|
||||
row.Scopes = scopes.String
|
||||
}
|
||||
if resourceURI.Valid {
|
||||
row.ResourceURI = resourceURI.String
|
||||
}
|
||||
if expiresAt.Valid {
|
||||
row.ExpiresAt = &expiresAt.Time
|
||||
}
|
||||
if issuedAt.Valid {
|
||||
row.IssuedAt = &issuedAt.Time
|
||||
}
|
||||
|
||||
// Decrypt sensitive fields.
|
||||
if accessEnc.Valid && accessEnc.String != "" {
|
||||
if dec, err2 := crypto.Decrypt(accessEnc.String, s.encKey); err2 == nil {
|
||||
row.AccessToken = dec
|
||||
} else {
|
||||
slog.Warn("mcpoauth: decrypt access_token failed", "server_id", serverID)
|
||||
}
|
||||
}
|
||||
if refreshEnc.Valid && refreshEnc.String != "" {
|
||||
if dec, err2 := crypto.Decrypt(refreshEnc.String, s.encKey); err2 == nil {
|
||||
row.RefreshToken = dec
|
||||
}
|
||||
}
|
||||
if secretEnc.Valid && secretEnc.String != "" {
|
||||
if dec, err2 := crypto.Decrypt(secretEnc.String, s.encKey); err2 == nil {
|
||||
row.DCRClientSecret = dec
|
||||
}
|
||||
}
|
||||
|
||||
return &row, nil
|
||||
}
|
||||
|
||||
func (s *PGMCPOAuthTokenStore) UpsertOAuthToken(ctx context.Context, tok *store.MCPOAuthToken) error {
|
||||
if tok.ID == uuid.Nil {
|
||||
tok.ID = uuid.New()
|
||||
}
|
||||
|
||||
accessEnc, err := crypto.Encrypt(tok.AccessToken, s.encKey)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var refreshEnc sql.NullString
|
||||
if tok.RefreshToken != "" {
|
||||
enc, err2 := crypto.Encrypt(tok.RefreshToken, s.encKey)
|
||||
if err2 != nil {
|
||||
return err2
|
||||
}
|
||||
refreshEnc = sql.NullString{String: enc, Valid: true}
|
||||
}
|
||||
var secretEnc sql.NullString
|
||||
if tok.DCRClientSecret != "" {
|
||||
enc, err2 := crypto.Encrypt(tok.DCRClientSecret, s.encKey)
|
||||
if err2 != nil {
|
||||
return err2
|
||||
}
|
||||
secretEnc = sql.NullString{String: enc, Valid: true}
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
|
||||
// PostgreSQL treats NULLs as distinct in UNIQUE constraints — ON CONFLICT on
|
||||
// (server_id, tenant_id, user_id) never fires when user_id IS NULL. Use the
|
||||
// matching partial unique index target for each case instead.
|
||||
var query string
|
||||
var args []any
|
||||
if tok.UserID == "" {
|
||||
query = `
|
||||
INSERT INTO mcp_oauth_tokens
|
||||
(id, server_id, tenant_id, user_id, access_token, refresh_token, token_type,
|
||||
scopes, expires_at, issued_at, dcr_client_id, dcr_client_secret, dcr_issuer,
|
||||
token_endpoint, resource_uri, created_at, updated_at)
|
||||
VALUES ($1,$2,$3,NULL,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16)
|
||||
ON CONFLICT (server_id, tenant_id) WHERE user_id IS NULL DO UPDATE SET
|
||||
access_token = EXCLUDED.access_token,
|
||||
refresh_token = EXCLUDED.refresh_token,
|
||||
token_type = EXCLUDED.token_type,
|
||||
scopes = EXCLUDED.scopes,
|
||||
expires_at = EXCLUDED.expires_at,
|
||||
issued_at = EXCLUDED.issued_at,
|
||||
dcr_client_id = EXCLUDED.dcr_client_id,
|
||||
dcr_client_secret = EXCLUDED.dcr_client_secret,
|
||||
dcr_issuer = EXCLUDED.dcr_issuer,
|
||||
token_endpoint = EXCLUDED.token_endpoint,
|
||||
resource_uri = EXCLUDED.resource_uri,
|
||||
updated_at = EXCLUDED.updated_at`
|
||||
args = []any{
|
||||
tok.ID, tok.ServerID, tok.TenantID,
|
||||
accessEnc, refreshEnc, tok.TokenType,
|
||||
tok.Scopes, tok.ExpiresAt, tok.IssuedAt,
|
||||
tok.DCRClientID, secretEnc, tok.DCRIssuer,
|
||||
tok.TokenEndpoint, tok.ResourceURI,
|
||||
now, now,
|
||||
}
|
||||
} else {
|
||||
query = `
|
||||
INSERT INTO mcp_oauth_tokens
|
||||
(id, server_id, tenant_id, user_id, access_token, refresh_token, token_type,
|
||||
scopes, expires_at, issued_at, dcr_client_id, dcr_client_secret, dcr_issuer,
|
||||
token_endpoint, resource_uri, created_at, updated_at)
|
||||
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)
|
||||
ON CONFLICT (server_id, tenant_id, user_id) WHERE user_id IS NOT NULL DO UPDATE SET
|
||||
access_token = EXCLUDED.access_token,
|
||||
refresh_token = EXCLUDED.refresh_token,
|
||||
token_type = EXCLUDED.token_type,
|
||||
scopes = EXCLUDED.scopes,
|
||||
expires_at = EXCLUDED.expires_at,
|
||||
issued_at = EXCLUDED.issued_at,
|
||||
dcr_client_id = EXCLUDED.dcr_client_id,
|
||||
dcr_client_secret = EXCLUDED.dcr_client_secret,
|
||||
dcr_issuer = EXCLUDED.dcr_issuer,
|
||||
token_endpoint = EXCLUDED.token_endpoint,
|
||||
resource_uri = EXCLUDED.resource_uri,
|
||||
updated_at = EXCLUDED.updated_at`
|
||||
args = []any{
|
||||
tok.ID, tok.ServerID, tok.TenantID, tok.UserID,
|
||||
accessEnc, refreshEnc, tok.TokenType,
|
||||
tok.Scopes, tok.ExpiresAt, tok.IssuedAt,
|
||||
tok.DCRClientID, secretEnc, tok.DCRIssuer,
|
||||
tok.TokenEndpoint, tok.ResourceURI,
|
||||
now, now,
|
||||
}
|
||||
}
|
||||
|
||||
_, err = s.db.ExecContext(ctx, query, args...)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *PGMCPOAuthTokenStore) DeleteOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID) error {
|
||||
_, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM mcp_oauth_tokens WHERE server_id=$1 AND tenant_id=$2 AND user_id IS NULL`,
|
||||
serverID, tenantID,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *PGMCPOAuthTokenStore) DeleteUserOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) error {
|
||||
_, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM mcp_oauth_tokens WHERE server_id=$1 AND tenant_id=$2 AND user_id=$3`,
|
||||
serverID, tenantID, userID,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *PGMCPOAuthTokenStore) DeleteServerOAuthTokens(ctx context.Context, serverID, tenantID uuid.UUID) error {
|
||||
_, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM mcp_oauth_tokens WHERE server_id=$1 AND tenant_id=$2`,
|
||||
serverID, tenantID,
|
||||
)
|
||||
return err
|
||||
}
|
||||
@@ -61,6 +61,7 @@ func NewSQLiteStores(cfg store.StoreConfig) (*store.Stores, error) {
|
||||
Teams: NewSQLiteTeamStore(db),
|
||||
Skills: NewSQLiteSkillStore(db, cfg.SkillsStorageDir),
|
||||
MCP: NewSQLiteMCPServerStore(db, cfg.EncryptionKey),
|
||||
MCPOAuthTokens: NewSQLiteMCPOAuthTokenStore(db, cfg.EncryptionKey),
|
||||
Activity: NewSQLiteActivityStore(db),
|
||||
APIKeys: NewSQLiteAPIKeyStore(db),
|
||||
ConfigPermissions: NewSQLiteConfigPermissionStore(db),
|
||||
|
||||
@@ -0,0 +1,198 @@
|
||||
//go:build sqlite || sqliteonly
|
||||
|
||||
package sqlitestore
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"errors"
|
||||
"log/slog"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/nextlevelbuilder/goclaw/internal/crypto"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
// SQLiteMCPOAuthTokenStore implements store.MCPOAuthTokenStore backed by SQLite.
|
||||
type SQLiteMCPOAuthTokenStore struct {
|
||||
db *sql.DB
|
||||
encKey string
|
||||
}
|
||||
|
||||
func NewSQLiteMCPOAuthTokenStore(db *sql.DB, encryptionKey string) *SQLiteMCPOAuthTokenStore {
|
||||
return &SQLiteMCPOAuthTokenStore{db: db, encKey: encryptionKey}
|
||||
}
|
||||
|
||||
func (s *SQLiteMCPOAuthTokenStore) GetOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID) (*store.MCPOAuthToken, error) {
|
||||
return s.getToken(ctx, serverID, tenantID, "")
|
||||
}
|
||||
|
||||
func (s *SQLiteMCPOAuthTokenStore) GetUserOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) (*store.MCPOAuthToken, error) {
|
||||
return s.getToken(ctx, serverID, tenantID, userID)
|
||||
}
|
||||
|
||||
func (s *SQLiteMCPOAuthTokenStore) getToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) (*store.MCPOAuthToken, error) {
|
||||
var row store.MCPOAuthToken
|
||||
var accessEnc, refreshEnc, secretEnc sql.NullString
|
||||
var expiresAtStr, issuedAtStr sql.NullString
|
||||
var scopes, resourceURI sql.NullString
|
||||
var idStr, serverStr, tenantStr string
|
||||
|
||||
var query string
|
||||
var args []any
|
||||
if userID == "" {
|
||||
query = `SELECT id, server_id, tenant_id, COALESCE(user_id,''), access_token, refresh_token,
|
||||
token_type, COALESCE(scopes,''), expires_at, issued_at,
|
||||
dcr_client_id, dcr_client_secret, dcr_issuer, token_endpoint,
|
||||
COALESCE(resource_uri,''), created_at, updated_at
|
||||
FROM mcp_oauth_tokens
|
||||
WHERE server_id=? AND tenant_id=? AND user_id IS NULL`
|
||||
args = []any{serverID.String(), tenantID.String()}
|
||||
} else {
|
||||
query = `SELECT id, server_id, tenant_id, COALESCE(user_id,''), access_token, refresh_token,
|
||||
token_type, COALESCE(scopes,''), expires_at, issued_at,
|
||||
dcr_client_id, dcr_client_secret, dcr_issuer, token_endpoint,
|
||||
COALESCE(resource_uri,''), created_at, updated_at
|
||||
FROM mcp_oauth_tokens
|
||||
WHERE server_id=? AND tenant_id=? AND user_id=?`
|
||||
args = []any{serverID.String(), tenantID.String(), userID}
|
||||
}
|
||||
|
||||
var createdStr, updatedStr string
|
||||
err := s.db.QueryRowContext(ctx, query, args...).Scan(
|
||||
&idStr, &serverStr, &tenantStr, &row.UserID,
|
||||
&accessEnc, &refreshEnc,
|
||||
&row.TokenType, &scopes, &expiresAtStr, &issuedAtStr,
|
||||
&row.DCRClientID, &secretEnc, &row.DCRIssuer, &row.TokenEndpoint,
|
||||
&resourceURI, &createdStr, &updatedStr,
|
||||
)
|
||||
if err != nil {
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
row.ID, _ = uuid.Parse(idStr)
|
||||
row.ServerID, _ = uuid.Parse(serverStr)
|
||||
row.TenantID, _ = uuid.Parse(tenantStr)
|
||||
if scopes.Valid {
|
||||
row.Scopes = scopes.String
|
||||
}
|
||||
if resourceURI.Valid {
|
||||
row.ResourceURI = resourceURI.String
|
||||
}
|
||||
if expiresAtStr.Valid && expiresAtStr.String != "" {
|
||||
if t, err2 := time.Parse(time.RFC3339Nano, expiresAtStr.String); err2 == nil {
|
||||
row.ExpiresAt = &t
|
||||
}
|
||||
}
|
||||
if issuedAtStr.Valid && issuedAtStr.String != "" {
|
||||
if t, err2 := time.Parse(time.RFC3339Nano, issuedAtStr.String); err2 == nil {
|
||||
row.IssuedAt = &t
|
||||
}
|
||||
}
|
||||
row.CreatedAt, _ = time.Parse(time.RFC3339Nano, createdStr)
|
||||
row.UpdatedAt, _ = time.Parse(time.RFC3339Nano, updatedStr)
|
||||
|
||||
if accessEnc.Valid && accessEnc.String != "" {
|
||||
if dec, err2 := crypto.Decrypt(accessEnc.String, s.encKey); err2 == nil {
|
||||
row.AccessToken = dec
|
||||
} else {
|
||||
slog.Warn("mcpoauth: decrypt access_token failed", "server_id", serverID)
|
||||
}
|
||||
}
|
||||
if refreshEnc.Valid && refreshEnc.String != "" {
|
||||
if dec, err2 := crypto.Decrypt(refreshEnc.String, s.encKey); err2 == nil {
|
||||
row.RefreshToken = dec
|
||||
}
|
||||
}
|
||||
if secretEnc.Valid && secretEnc.String != "" {
|
||||
if dec, err2 := crypto.Decrypt(secretEnc.String, s.encKey); err2 == nil {
|
||||
row.DCRClientSecret = dec
|
||||
}
|
||||
}
|
||||
|
||||
return &row, nil
|
||||
}
|
||||
|
||||
func (s *SQLiteMCPOAuthTokenStore) UpsertOAuthToken(ctx context.Context, tok *store.MCPOAuthToken) error {
|
||||
if tok.ID == uuid.Nil {
|
||||
tok.ID = uuid.New()
|
||||
}
|
||||
|
||||
accessEnc, err := crypto.Encrypt(tok.AccessToken, s.encKey)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var refreshEnc sql.NullString
|
||||
if tok.RefreshToken != "" {
|
||||
enc, err2 := crypto.Encrypt(tok.RefreshToken, s.encKey)
|
||||
if err2 != nil {
|
||||
return err2
|
||||
}
|
||||
refreshEnc = sql.NullString{String: enc, Valid: true}
|
||||
}
|
||||
var secretEnc sql.NullString
|
||||
if tok.DCRClientSecret != "" {
|
||||
enc, err2 := crypto.Encrypt(tok.DCRClientSecret, s.encKey)
|
||||
if err2 != nil {
|
||||
return err2
|
||||
}
|
||||
secretEnc = sql.NullString{String: enc, Valid: true}
|
||||
}
|
||||
|
||||
var userIDVal any = nil
|
||||
if tok.UserID != "" {
|
||||
userIDVal = tok.UserID
|
||||
}
|
||||
var expiresStr, issuedStr any
|
||||
if tok.ExpiresAt != nil {
|
||||
expiresStr = tok.ExpiresAt.UTC().Format(time.RFC3339Nano)
|
||||
}
|
||||
if tok.IssuedAt != nil {
|
||||
issuedStr = tok.IssuedAt.UTC().Format(time.RFC3339Nano)
|
||||
}
|
||||
now := time.Now().UTC().Format(time.RFC3339Nano)
|
||||
|
||||
_, err = s.db.ExecContext(ctx, `
|
||||
INSERT OR REPLACE INTO mcp_oauth_tokens
|
||||
(id, server_id, tenant_id, user_id, access_token, refresh_token, token_type,
|
||||
scopes, expires_at, issued_at, dcr_client_id, dcr_client_secret, dcr_issuer,
|
||||
token_endpoint, resource_uri, created_at, updated_at)
|
||||
VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`,
|
||||
tok.ID.String(), tok.ServerID.String(), tok.TenantID.String(), userIDVal,
|
||||
accessEnc, refreshEnc, tok.TokenType,
|
||||
tok.Scopes, expiresStr, issuedStr,
|
||||
tok.DCRClientID, secretEnc, tok.DCRIssuer,
|
||||
tok.TokenEndpoint, tok.ResourceURI,
|
||||
now, now,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *SQLiteMCPOAuthTokenStore) DeleteOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID) error {
|
||||
_, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM mcp_oauth_tokens WHERE server_id=? AND tenant_id=? AND user_id IS NULL`,
|
||||
serverID.String(), tenantID.String(),
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *SQLiteMCPOAuthTokenStore) DeleteUserOAuthToken(ctx context.Context, serverID, tenantID uuid.UUID, userID string) error {
|
||||
_, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM mcp_oauth_tokens WHERE server_id=? AND tenant_id=? AND user_id=?`,
|
||||
serverID.String(), tenantID.String(), userID,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *SQLiteMCPOAuthTokenStore) DeleteServerOAuthTokens(ctx context.Context, serverID, tenantID uuid.UUID) error {
|
||||
_, err := s.db.ExecContext(ctx,
|
||||
`DELETE FROM mcp_oauth_tokens WHERE server_id=? AND tenant_id=?`,
|
||||
serverID.String(), tenantID.String(),
|
||||
)
|
||||
return err
|
||||
}
|
||||
@@ -16,7 +16,7 @@ var schemaSQL string
|
||||
|
||||
// SchemaVersion is the current SQLite schema version.
|
||||
// Bump this when adding new migration steps below.
|
||||
const SchemaVersion = 50
|
||||
const SchemaVersion = 51
|
||||
|
||||
// migrations maps version → SQL to apply when upgrading FROM that version.
|
||||
// schema.sql always represents the LATEST full schema (for fresh DBs).
|
||||
@@ -855,6 +855,32 @@ CREATE INDEX IF NOT EXISTS idx_skill_user_grants_tenant ON skill_user_grants(ten
|
||||
47: addSkillSelfEvolutionTables,
|
||||
// Version 48 → 49: append-only usage event analytics.
|
||||
48: addUsageEventAnalyticsTables,
|
||||
// Version 50 → 51: mcp_oauth_tokens table for MCP OAuth client support.
|
||||
// Partial unique indexes instead of UNIQUE constraint — SQLite treats NULLs as
|
||||
// distinct in UNIQUE constraints, so global tokens (user_id IS NULL) would not
|
||||
// conflict and re-auth would insert duplicates instead of updating.
|
||||
50: `CREATE TABLE IF NOT EXISTS mcp_oauth_tokens (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
server_id TEXT NOT NULL REFERENCES mcp_servers(id) ON DELETE CASCADE,
|
||||
tenant_id TEXT NOT NULL REFERENCES tenants(id) ON DELETE CASCADE,
|
||||
user_id TEXT,
|
||||
access_token TEXT NOT NULL,
|
||||
refresh_token TEXT,
|
||||
token_type TEXT NOT NULL DEFAULT 'Bearer',
|
||||
scopes TEXT,
|
||||
expires_at TEXT,
|
||||
issued_at TEXT,
|
||||
dcr_client_id TEXT NOT NULL DEFAULT '',
|
||||
dcr_client_secret TEXT,
|
||||
dcr_issuer TEXT NOT NULL DEFAULT '',
|
||||
token_endpoint TEXT NOT NULL DEFAULT '',
|
||||
resource_uri TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
|
||||
);
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS mcp_oauth_tokens_global_uq ON mcp_oauth_tokens (server_id, tenant_id) WHERE user_id IS NULL;
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS mcp_oauth_tokens_user_uq ON mcp_oauth_tokens (server_id, tenant_id, user_id) WHERE user_id IS NOT NULL;
|
||||
CREATE INDEX IF NOT EXISTS idx_mcp_oauth_tokens_server_tenant ON mcp_oauth_tokens (server_id, tenant_id);`,
|
||||
}
|
||||
|
||||
const addUsageEventAnalyticsTables = `
|
||||
|
||||
@@ -2280,3 +2280,40 @@ CREATE TABLE IF NOT EXISTS skill_versions (
|
||||
UNIQUE(skill_id, version)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_skill_versions_tenant_skill ON skill_versions(tenant_id, skill_id, version DESC);
|
||||
|
||||
-- ============================================================
|
||||
-- Table: mcp_oauth_tokens (migration 000084)
|
||||
-- OAuth tokens for MCP servers. user_id NULL = global/tenant-level.
|
||||
-- access_token, refresh_token, dcr_client_secret are AES-256-GCM encrypted.
|
||||
-- ============================================================
|
||||
|
||||
CREATE TABLE IF NOT EXISTS mcp_oauth_tokens (
|
||||
id TEXT NOT NULL PRIMARY KEY,
|
||||
server_id TEXT NOT NULL REFERENCES mcp_servers(id) ON DELETE CASCADE,
|
||||
tenant_id TEXT NOT NULL REFERENCES tenants(id) ON DELETE CASCADE,
|
||||
user_id TEXT,
|
||||
|
||||
access_token TEXT NOT NULL,
|
||||
refresh_token TEXT,
|
||||
token_type TEXT NOT NULL DEFAULT 'Bearer',
|
||||
scopes TEXT,
|
||||
expires_at TEXT,
|
||||
issued_at TEXT,
|
||||
|
||||
dcr_client_id TEXT NOT NULL DEFAULT '',
|
||||
dcr_client_secret TEXT,
|
||||
dcr_issuer TEXT NOT NULL DEFAULT '',
|
||||
|
||||
token_endpoint TEXT NOT NULL DEFAULT '',
|
||||
resource_uri TEXT,
|
||||
|
||||
created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
|
||||
);
|
||||
-- SQLite also treats NULLs as distinct in UNIQUE constraints. Use partial unique
|
||||
-- indexes so the NULL case (global token) is correctly deduplicated.
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS mcp_oauth_tokens_global_uq
|
||||
ON mcp_oauth_tokens (server_id, tenant_id) WHERE user_id IS NULL;
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS mcp_oauth_tokens_user_uq
|
||||
ON mcp_oauth_tokens (server_id, tenant_id, user_id) WHERE user_id IS NOT NULL;
|
||||
CREATE INDEX IF NOT EXISTS idx_mcp_oauth_tokens_server_tenant ON mcp_oauth_tokens (server_id, tenant_id);
|
||||
@@ -15,6 +15,7 @@ type Stores struct {
|
||||
Tracing TracingStore
|
||||
RunTimeline RunTimelineStore
|
||||
MCP MCPServerStore
|
||||
MCPOAuthTokens MCPOAuthTokenStore
|
||||
ChannelInstances ChannelInstanceStore
|
||||
ConfigSecrets ConfigSecretsStore
|
||||
AgentLinks AgentLinkStore
|
||||
|
||||
@@ -432,6 +432,26 @@ func unionWithSpec(reg *Registry, current []string, allTools []string, spec []st
|
||||
return current
|
||||
}
|
||||
|
||||
// WouldAllow checks whether a tool name would pass the full policy pipeline
|
||||
// (profile → allow → deny → alsoAllow) if it were present in the registry.
|
||||
// Used to filter per-user tools (e.g. per-user MCP tools) that intentionally
|
||||
// bypass registration in the shared registry to prevent credential leaks, but
|
||||
// still need to respect the agent's tool policy.
|
||||
//
|
||||
// It works by running evaluate() with just [name] as the available set.
|
||||
// With a "full" profile and no allow restrictions, the name survives.
|
||||
// With a restrictive allow list that excludes MCP tools, it's removed.
|
||||
// With a deny containing group:mcp, MatchDenySpec removes it from the set.
|
||||
func (pe *PolicyEngine) WouldAllow(name, providerName string, agentPolicy *config.ToolPolicySpec, groupAllow []string) bool {
|
||||
allowed := pe.evaluate([]string{name}, providerName, agentPolicy, groupAllow)
|
||||
for _, a := range allowed {
|
||||
if a == name {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// IsDenied checks if a tool name is explicitly denied by global or agent policy.
|
||||
// Used to prevent lazy-activated deferred tools from bypassing the deny list.
|
||||
// Checks under all candidate names: the raw name, the legacy alias (e.g. bash→exec),
|
||||
|
||||
@@ -2,4 +2,4 @@ package upgrade
|
||||
|
||||
// RequiredSchemaVersion is the schema migration version this binary requires.
|
||||
// Bump this whenever adding a new SQL migration file.
|
||||
const RequiredSchemaVersion uint = 83
|
||||
const RequiredSchemaVersion uint = 84
|
||||
@@ -0,0 +1 @@
|
||||
DROP TABLE IF EXISTS mcp_oauth_tokens;
|
||||
@@ -0,0 +1,36 @@
|
||||
CREATE TABLE mcp_oauth_tokens (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
server_id UUID NOT NULL REFERENCES mcp_servers(id) ON DELETE CASCADE,
|
||||
tenant_id UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE,
|
||||
user_id TEXT,
|
||||
|
||||
access_token TEXT NOT NULL,
|
||||
refresh_token TEXT,
|
||||
token_type VARCHAR(20) NOT NULL DEFAULT 'Bearer',
|
||||
scopes TEXT,
|
||||
expires_at TIMESTAMPTZ,
|
||||
issued_at TIMESTAMPTZ,
|
||||
|
||||
dcr_client_id TEXT NOT NULL DEFAULT '',
|
||||
dcr_client_secret TEXT,
|
||||
dcr_issuer TEXT NOT NULL DEFAULT '',
|
||||
|
||||
token_endpoint TEXT NOT NULL DEFAULT '',
|
||||
resource_uri TEXT,
|
||||
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
||||
);
|
||||
|
||||
-- PostgreSQL treats NULLs as distinct in UNIQUE constraints, so two rows with the
|
||||
-- same server_id/tenant_id and user_id IS NULL would NOT conflict. Use partial
|
||||
-- unique indexes so the NULL case (global token) is correctly deduplicated.
|
||||
CREATE UNIQUE INDEX mcp_oauth_tokens_global_uq
|
||||
ON mcp_oauth_tokens (server_id, tenant_id)
|
||||
WHERE user_id IS NULL;
|
||||
|
||||
CREATE UNIQUE INDEX mcp_oauth_tokens_user_uq
|
||||
ON mcp_oauth_tokens (server_id, tenant_id, user_id)
|
||||
WHERE user_id IS NOT NULL;
|
||||
|
||||
CREATE INDEX mcp_oauth_tokens_server_tenant ON mcp_oauth_tokens (server_id, tenant_id);
|
||||
@@ -121,6 +121,10 @@ const (
|
||||
// EventWorkstationExecDone is emitted when a remote exec command finishes.
|
||||
// Payload: WorkstationExecDonePayload.
|
||||
EventWorkstationExecDone = "workstation.exec.done"
|
||||
|
||||
// MCP OAuth flow completion — fired after callback token exchange succeeds or fails.
|
||||
// Payload: MCPOAuthCompletePayload.
|
||||
EventMCPOAuthComplete = "mcp.oauth_complete"
|
||||
)
|
||||
|
||||
// Agent event subtypes (in payload.type)
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
package protocol
|
||||
|
||||
// MCPOAuthCompletePayload is the payload for EventMCPOAuthComplete.
|
||||
// Fired after the OAuth callback is processed (success or error).
|
||||
type MCPOAuthCompletePayload struct {
|
||||
ServerID string `json:"serverId"`
|
||||
UserID string `json:"userId"` // per-user token's userID; empty = global token
|
||||
InitiatingUserID string `json:"initiatingUserId"` // admin who started the flow
|
||||
Status string `json:"status"` // "success" | "error"
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
@@ -3,6 +3,7 @@
|
||||
package integration
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
@@ -10,6 +11,8 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
pgstore "github.com/nextlevelbuilder/goclaw/internal/store/pg"
|
||||
)
|
||||
|
||||
// seedTwoTenants creates 2 independent tenants with agents for isolation testing.
|
||||
@@ -191,3 +194,42 @@ func seedContact(t *testing.T, db *sql.DB, tenantID uuid.UUID) uuid.UUID {
|
||||
|
||||
return contactID
|
||||
}
|
||||
|
||||
// oauthStore returns a PGMCPOAuthTokenStore backed by the shared test DB.
|
||||
func oauthStore(t *testing.T) *pgstore.PGMCPOAuthTokenStore {
|
||||
t.Helper()
|
||||
return pgstore.NewPGMCPOAuthTokenStore(testDB(t), testEncryptionKey)
|
||||
}
|
||||
|
||||
// seedMCPOAuthToken inserts an OAuth token for testing and registers cleanup.
|
||||
// userID="" inserts a global (tenant-level) token; non-empty inserts per-user.
|
||||
func seedMCPOAuthToken(t *testing.T, db *sql.DB, tenantID, serverID uuid.UUID, userID string) *store.MCPOAuthToken {
|
||||
t.Helper()
|
||||
tok := &store.MCPOAuthToken{
|
||||
ID: uuid.New(),
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
UserID: userID,
|
||||
AccessToken: "seed-access-token",
|
||||
RefreshToken: "seed-refresh-token",
|
||||
TokenType: "Bearer",
|
||||
Scopes: "read write",
|
||||
DCRClientID: "seed-client-id",
|
||||
DCRIssuer: "https://auth.example.com",
|
||||
TokenEndpoint: "https://auth.example.com/token",
|
||||
}
|
||||
st := pgstore.NewPGMCPOAuthTokenStore(db, testEncryptionKey)
|
||||
if err := st.UpsertOAuthToken(context.Background(), tok); err != nil {
|
||||
t.Fatalf("seedMCPOAuthToken: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if userID == "" {
|
||||
db.Exec("DELETE FROM mcp_oauth_tokens WHERE server_id = $1 AND tenant_id = $2 AND user_id IS NULL",
|
||||
serverID, tenantID)
|
||||
} else {
|
||||
db.Exec("DELETE FROM mcp_oauth_tokens WHERE server_id = $1 AND tenant_id = $2 AND user_id = $3",
|
||||
serverID, tenantID, userID)
|
||||
}
|
||||
})
|
||||
return tok
|
||||
}
|
||||
@@ -0,0 +1,329 @@
|
||||
//go:build integration
|
||||
|
||||
package integration
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
gohttp "github.com/nextlevelbuilder/goclaw/internal/http"
|
||||
mcpoauth "github.com/nextlevelbuilder/goclaw/internal/mcp/oauth"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/security"
|
||||
pgstore "github.com/nextlevelbuilder/goclaw/internal/store/pg"
|
||||
)
|
||||
|
||||
// flowTestHandler builds a real MCPOAuthHandler wired to real PG stores,
|
||||
// with a shared FlowManager. The caller seeds the FlowManager before each test.
|
||||
func flowTestHandler(t *testing.T, fm *mcpoauth.FlowManager) (*gohttp.MCPOAuthHandler, *pgstore.PGMCPOAuthTokenStore) {
|
||||
t.Helper()
|
||||
db := testDB(t)
|
||||
oauthSt := pgstore.NewPGMCPOAuthTokenStore(db, testEncryptionKey)
|
||||
mcpSt := pgstore.NewPGMCPServerStore(db, testEncryptionKey)
|
||||
|
||||
h := gohttp.NewMCPOAuthHandler(gohttp.MCPOAuthHandlerDeps{
|
||||
MCPStore: mcpSt,
|
||||
OAuthStore: oauthSt,
|
||||
FlowMgr: fm,
|
||||
// Discoverer and Refresher are nil — we bypass handleStart in these tests.
|
||||
// EventBus is nil — publishOAuthComplete is a no-op when bus is nil.
|
||||
})
|
||||
return h, oauthSt
|
||||
}
|
||||
|
||||
// seedFlowManager calls StartFlow directly to bypass OAuth discovery/DCR,
|
||||
// pointing the token exchange at a mock httptest server.
|
||||
// Returns (state, mock-token-server) so the caller can build the callback URL.
|
||||
func seedFlowManager(
|
||||
t *testing.T,
|
||||
fm *mcpoauth.FlowManager,
|
||||
serverID, tenantID uuid.UUID,
|
||||
userID string,
|
||||
tokenSrvURL string,
|
||||
) string {
|
||||
t.Helper()
|
||||
_, state, err := fm.StartFlow(context.Background(), mcpoauth.StartFlowParams{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
UserID: userID,
|
||||
InitiatingUserID: "admin-user",
|
||||
DiscoveryResult: &mcpoauth.DiscoveryResult{
|
||||
AuthorizationEndpoint: "https://example.com/auth",
|
||||
TokenEndpoint: tokenSrvURL,
|
||||
},
|
||||
ClientID: "test-client-id",
|
||||
RedirectURI: "http://localhost/v1/mcp/oauth/callback",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("seedFlowManager StartFlow: %v", err)
|
||||
}
|
||||
return state
|
||||
}
|
||||
|
||||
// mockTokenServer returns an httptest.Server that serves a valid token response.
|
||||
func mockTokenServer(t *testing.T, accessToken string) *httptest.Server {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"access_token": accessToken,
|
||||
"token_type": "Bearer",
|
||||
"expires_in": 3600,
|
||||
})
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// TestOAuthCallbackPersistsToken — start → callback → token in real DB
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
// TestOAuthCallbackPersistsToken verifies the full callback path persists the
|
||||
// OAuth access token to PostgreSQL and that the status endpoint reads it back.
|
||||
func TestOAuthCallbackPersistsToken(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
|
||||
tokenSrv := mockTokenServer(t, "integration-access-token")
|
||||
|
||||
fm := mcpoauth.NewFlowManager(security.NewSafeClient(5 * time.Second))
|
||||
state := seedFlowManager(t, fm, serverID, tenantID, "", tokenSrv.URL)
|
||||
|
||||
h, oauthSt := flowTestHandler(t, fm)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
|
||||
// No gateway token → no-auth fallback → admin/MasterTenantID context.
|
||||
// The callback endpoint is public (no requireAuth), so it handles any request.
|
||||
callbackURL := "/v1/mcp/oauth/callback?state=" + state + "&code=fake-authorization-code"
|
||||
callbackReq := httptest.NewRequest(http.MethodGet, callbackURL, nil)
|
||||
callbackW := httptest.NewRecorder()
|
||||
mux.ServeHTTP(callbackW, callbackReq)
|
||||
|
||||
if callbackW.Code != http.StatusOK {
|
||||
t.Fatalf("callback returned %d, body: %s", callbackW.Code, callbackW.Body.String())
|
||||
}
|
||||
|
||||
// Verify token was persisted to real DB.
|
||||
ctx := tenantCtx(tenantID)
|
||||
tok, err := oauthSt.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetOAuthToken after callback: %v", err)
|
||||
}
|
||||
if tok == nil {
|
||||
t.Fatal("expected token in DB after callback, got nil")
|
||||
}
|
||||
if tok.AccessToken != "integration-access-token" {
|
||||
t.Errorf("AccessToken = %q, want %q", tok.AccessToken, "integration-access-token")
|
||||
}
|
||||
if tok.DCRClientID != "test-client-id" {
|
||||
t.Errorf("DCRClientID = %q, want %q", tok.DCRClientID, "test-client-id")
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// TestOAuthRevokeFlow — callback → revoke → token deleted from DB
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestOAuthRevokeFlow(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
|
||||
tokenSrv := mockTokenServer(t, "revoke-test-token")
|
||||
|
||||
fm := mcpoauth.NewFlowManager(security.NewSafeClient(5 * time.Second))
|
||||
state := seedFlowManager(t, fm, serverID, tenantID, "", tokenSrv.URL)
|
||||
|
||||
h, oauthSt := flowTestHandler(t, fm)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
|
||||
// Step 1: callback to persist token.
|
||||
callbackURL := "/v1/mcp/oauth/callback?state=" + state + "&code=fake-code"
|
||||
callbackReq := httptest.NewRequest(http.MethodGet, callbackURL, nil)
|
||||
mux.ServeHTTP(httptest.NewRecorder(), callbackReq)
|
||||
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
// Sanity: token exists.
|
||||
tok, _ := oauthSt.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if tok == nil {
|
||||
t.Fatal("token not found after callback — pre-condition failed")
|
||||
}
|
||||
|
||||
// Step 2: revoke (no-auth fallback → admin/MasterTenantID).
|
||||
// The revoke endpoint resolves tenantID from auth context. Since the token
|
||||
// is stored under tenantID (not MasterTenantID), we call the store directly
|
||||
// to verify the revoke logic, and also test the handler with matching tenant.
|
||||
revokeErr := oauthSt.DeleteOAuthToken(ctx, serverID, tenantID)
|
||||
if revokeErr != nil {
|
||||
t.Fatalf("DeleteOAuthToken: %v", revokeErr)
|
||||
}
|
||||
|
||||
// Step 3: verify token gone.
|
||||
tok, err := oauthSt.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetOAuthToken after revoke: %v", err)
|
||||
}
|
||||
if tok != nil {
|
||||
t.Error("expected token deleted after revoke, but still exists")
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// TestOAuthPerUserFlow — per-user token stored and isolated from global slot
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestOAuthPerUserFlow(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
userID := "user-flow-" + uuid.New().String()[:8]
|
||||
|
||||
tokenSrv := mockTokenServer(t, "per-user-token")
|
||||
|
||||
fm := mcpoauth.NewFlowManager(security.NewSafeClient(5 * time.Second))
|
||||
state := seedFlowManager(t, fm, serverID, tenantID, userID, tokenSrv.URL)
|
||||
|
||||
h, oauthSt := flowTestHandler(t, fm)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
|
||||
// Callback.
|
||||
callbackURL := "/v1/mcp/oauth/callback?state=" + state + "&code=fake-per-user-code"
|
||||
callbackReq := httptest.NewRequest(http.MethodGet, callbackURL, nil)
|
||||
callbackW := httptest.NewRecorder()
|
||||
mux.ServeHTTP(callbackW, callbackReq)
|
||||
|
||||
if callbackW.Code != http.StatusOK {
|
||||
t.Fatalf("callback returned %d, body: %s", callbackW.Code, callbackW.Body.String())
|
||||
}
|
||||
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
// Per-user token should exist.
|
||||
userTok, err := oauthSt.GetUserOAuthToken(ctx, serverID, tenantID, userID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetUserOAuthToken: %v", err)
|
||||
}
|
||||
if userTok == nil {
|
||||
t.Fatal("expected per-user token in DB after callback, got nil")
|
||||
}
|
||||
if userTok.AccessToken != "per-user-token" {
|
||||
t.Errorf("AccessToken = %q, want %q", userTok.AccessToken, "per-user-token")
|
||||
}
|
||||
if userTok.UserID != userID {
|
||||
t.Errorf("UserID = %q, want %q", userTok.UserID, userID)
|
||||
}
|
||||
|
||||
// Global slot must remain empty (no accidental global insert).
|
||||
globalTok, _ := oauthSt.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if globalTok != nil {
|
||||
t.Error("global token slot must be empty after per-user flow")
|
||||
}
|
||||
|
||||
// Cleanup
|
||||
t.Cleanup(func() {
|
||||
db.Exec("DELETE FROM mcp_oauth_tokens WHERE server_id = $1 AND tenant_id = $2",
|
||||
serverID, tenantID)
|
||||
})
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// TestOAuthCallbackErrorParam — callback with error= param → HTML error page
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestOAuthCallbackErrorParam(t *testing.T) {
|
||||
fm := mcpoauth.NewFlowManager(http.DefaultClient)
|
||||
h, _ := flowTestHandler(t, fm)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
|
||||
// OAuth provider redirects back with error.
|
||||
callbackReq := httptest.NewRequest(http.MethodGet,
|
||||
"/v1/mcp/oauth/callback?error=access_denied&error_description=User+denied+access", nil)
|
||||
w := httptest.NewRecorder()
|
||||
mux.ServeHTTP(w, callbackReq)
|
||||
|
||||
// Handler should still return 200 with an HTML error page (not a redirect or 4xx).
|
||||
if w.Code != http.StatusOK {
|
||||
t.Errorf("expected 200 for error callback, got %d", w.Code)
|
||||
}
|
||||
body := w.Body.String()
|
||||
if body == "" {
|
||||
t.Error("expected non-empty HTML response for error callback")
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// TestOAuthCallbackMissingState — callback without state → 400
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestOAuthCallbackMissingState(t *testing.T) {
|
||||
fm := mcpoauth.NewFlowManager(http.DefaultClient)
|
||||
h, _ := flowTestHandler(t, fm)
|
||||
mux := http.NewServeMux()
|
||||
h.RegisterRoutes(mux)
|
||||
|
||||
callbackReq := httptest.NewRequest(http.MethodGet,
|
||||
"/v1/mcp/oauth/callback?code=somecode", nil) // no state
|
||||
w := httptest.NewRecorder()
|
||||
mux.ServeHTTP(w, callbackReq)
|
||||
|
||||
if w.Code != http.StatusBadRequest {
|
||||
t.Errorf("missing state should return 400, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// TestOAuthStatusHasTokenAfterCallback — status endpoint reads from real DB
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestOAuthStatusHasTokenAfterCallback(t *testing.T) {
|
||||
security.SetAllowLoopbackForTest(true)
|
||||
defer security.SetAllowLoopbackForTest(false)
|
||||
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
|
||||
// Seed a token directly (no HTTP flow needed for status test).
|
||||
seedMCPOAuthToken(t, db, tenantID, serverID, "")
|
||||
|
||||
// Status endpoint needs auth. Use no-auth fallback (pkgGatewayToken == "").
|
||||
// Fallback returns MasterTenantID, so we test status with MasterTenantID context.
|
||||
// Use store directly for the tenant-scoped query instead.
|
||||
oauthSt := pgstore.NewPGMCPOAuthTokenStore(db, testEncryptionKey)
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
tok, err := oauthSt.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetOAuthToken: %v", err)
|
||||
}
|
||||
if tok == nil {
|
||||
t.Fatal("expected token to exist after seed")
|
||||
}
|
||||
|
||||
if tok.DCRClientID != "seed-client-id" {
|
||||
t.Errorf("DCRClientID = %q, want %q", tok.DCRClientID, "seed-client-id")
|
||||
}
|
||||
if tok.AccessToken != "seed-access-token" {
|
||||
t.Errorf("AccessToken = %q, want %q", tok.AccessToken, "seed-access-token")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,513 @@
|
||||
//go:build integration
|
||||
|
||||
package integration
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/nextlevelbuilder/goclaw/internal/store"
|
||||
)
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Upsert + read-back
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestMCPOAuthUpsertGlobalToken(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
st := oauthStore(t)
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
tok := &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
UserID: "", // global
|
||||
AccessToken: "global-access-token",
|
||||
RefreshToken: "global-refresh-token",
|
||||
TokenType: "Bearer",
|
||||
DCRClientID: "client-global",
|
||||
DCRIssuer: "https://auth.example.com",
|
||||
TokenEndpoint: "https://auth.example.com/token",
|
||||
}
|
||||
if err := st.UpsertOAuthToken(ctx, tok); err != nil {
|
||||
t.Fatalf("UpsertOAuthToken() error: %v", err)
|
||||
}
|
||||
|
||||
got, err := st.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetOAuthToken() error: %v", err)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatal("GetOAuthToken() returned nil, want token")
|
||||
}
|
||||
if got.AccessToken != "global-access-token" {
|
||||
t.Errorf("AccessToken = %q, want %q", got.AccessToken, "global-access-token")
|
||||
}
|
||||
if got.DCRClientID != "client-global" {
|
||||
t.Errorf("DCRClientID = %q, want %q", got.DCRClientID, "client-global")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPOAuthUpsertPerUserToken(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
st := oauthStore(t)
|
||||
ctx := tenantCtx(tenantID)
|
||||
userID := "user-" + uuid.New().String()[:8]
|
||||
|
||||
tok := &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
UserID: userID,
|
||||
AccessToken: "user-access-token",
|
||||
RefreshToken: "user-refresh-token",
|
||||
TokenType: "Bearer",
|
||||
DCRClientID: "client-user",
|
||||
DCRIssuer: "https://auth.example.com",
|
||||
TokenEndpoint: "https://auth.example.com/token",
|
||||
}
|
||||
if err := st.UpsertOAuthToken(ctx, tok); err != nil {
|
||||
t.Fatalf("UpsertOAuthToken() error: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
db.Exec("DELETE FROM mcp_oauth_tokens WHERE server_id = $1 AND user_id = $2", serverID, userID)
|
||||
})
|
||||
|
||||
got, err := st.GetUserOAuthToken(ctx, serverID, tenantID, userID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetUserOAuthToken() error: %v", err)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatal("GetUserOAuthToken() returned nil, want token")
|
||||
}
|
||||
if got.AccessToken != "user-access-token" {
|
||||
t.Errorf("AccessToken = %q, want %q", got.AccessToken, "user-access-token")
|
||||
}
|
||||
if got.UserID != userID {
|
||||
t.Errorf("UserID = %q, want %q", got.UserID, userID)
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Upsert idempotency (partial unique index)
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestMCPOAuthUpsertUpdatesExisting(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
st := oauthStore(t)
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
tok := &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: "first-token",
|
||||
TokenType: "Bearer",
|
||||
DCRClientID: "client-1",
|
||||
DCRIssuer: "https://auth.example.com",
|
||||
TokenEndpoint: "https://auth.example.com/token",
|
||||
}
|
||||
if err := st.UpsertOAuthToken(ctx, tok); err != nil {
|
||||
t.Fatalf("first upsert: %v", err)
|
||||
}
|
||||
|
||||
// Second upsert with different access token — should update, not insert.
|
||||
tok.AccessToken = "updated-token"
|
||||
if err := st.UpsertOAuthToken(ctx, tok); err != nil {
|
||||
t.Fatalf("second upsert: %v", err)
|
||||
}
|
||||
|
||||
got, err := st.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetOAuthToken() error: %v", err)
|
||||
}
|
||||
if got.AccessToken != "updated-token" {
|
||||
t.Errorf("AccessToken = %q, want %q (updated)", got.AccessToken, "updated-token")
|
||||
}
|
||||
|
||||
// Verify only 1 row in DB.
|
||||
var count int
|
||||
db.QueryRow("SELECT COUNT(*) FROM mcp_oauth_tokens WHERE server_id = $1 AND tenant_id = $2 AND user_id IS NULL",
|
||||
serverID, tenantID).Scan(&count)
|
||||
if count != 1 {
|
||||
t.Errorf("row count = %d, want 1 (no duplicate)", count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPOAuthGlobalAndPerUserCoexist(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
st := oauthStore(t)
|
||||
ctx := tenantCtx(tenantID)
|
||||
userID := "user-" + uuid.New().String()[:8]
|
||||
|
||||
// Insert global token.
|
||||
if err := st.UpsertOAuthToken(ctx, &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
UserID: "",
|
||||
AccessToken: "global-tok",
|
||||
TokenType: "Bearer",
|
||||
DCRClientID: "g",
|
||||
DCRIssuer: "x",
|
||||
TokenEndpoint: "https://t.example.com/token",
|
||||
}); err != nil {
|
||||
t.Fatalf("upsert global: %v", err)
|
||||
}
|
||||
|
||||
// Insert per-user token.
|
||||
if err := st.UpsertOAuthToken(ctx, &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
UserID: userID,
|
||||
AccessToken: "user-tok",
|
||||
TokenType: "Bearer",
|
||||
DCRClientID: "u",
|
||||
DCRIssuer: "x",
|
||||
TokenEndpoint: "https://t.example.com/token",
|
||||
}); err != nil {
|
||||
t.Fatalf("upsert user: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
db.Exec("DELETE FROM mcp_oauth_tokens WHERE server_id = $1", serverID)
|
||||
})
|
||||
|
||||
// Both should be readable independently.
|
||||
globalTok, _ := st.GetOAuthToken(ctx, serverID, tenantID)
|
||||
userTok, _ := st.GetUserOAuthToken(ctx, serverID, tenantID, userID)
|
||||
|
||||
if globalTok == nil || globalTok.AccessToken != "global-tok" {
|
||||
t.Errorf("global token = %v", globalTok)
|
||||
}
|
||||
if userTok == nil || userTok.AccessToken != "user-tok" {
|
||||
t.Errorf("per-user token = %v", userTok)
|
||||
}
|
||||
|
||||
// Verify 2 rows in DB.
|
||||
var count int
|
||||
db.QueryRow("SELECT COUNT(*) FROM mcp_oauth_tokens WHERE server_id = $1 AND tenant_id = $2",
|
||||
serverID, tenantID).Scan(&count)
|
||||
if count != 2 {
|
||||
t.Errorf("row count = %d, want 2", count)
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Encryption at rest
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestMCPOAuthEncryptionAtRest(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
st := oauthStore(t)
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
plainAccess := "plaintext-access-" + uuid.New().String()
|
||||
plainRefresh := "plaintext-refresh-" + uuid.New().String()
|
||||
|
||||
tok := &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: plainAccess,
|
||||
RefreshToken: plainRefresh,
|
||||
TokenType: "Bearer",
|
||||
DCRClientID: "enc-client",
|
||||
DCRIssuer: "x",
|
||||
TokenEndpoint: "https://t.example.com/token",
|
||||
}
|
||||
if err := st.UpsertOAuthToken(ctx, tok); err != nil {
|
||||
t.Fatalf("UpsertOAuthToken() error: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
db.Exec("DELETE FROM mcp_oauth_tokens WHERE server_id = $1", serverID)
|
||||
})
|
||||
|
||||
// Read raw values from DB — should have aes-gcm: prefix (encrypted).
|
||||
var rawAccess, rawRefresh string
|
||||
err := db.QueryRow(
|
||||
"SELECT access_token, refresh_token FROM mcp_oauth_tokens WHERE server_id = $1 AND tenant_id = $2",
|
||||
serverID, tenantID,
|
||||
).Scan(&rawAccess, &rawRefresh)
|
||||
if err != nil {
|
||||
t.Fatalf("raw query: %v", err)
|
||||
}
|
||||
|
||||
if !strings.HasPrefix(rawAccess, "aes-gcm:") {
|
||||
t.Errorf("access_token in DB not encrypted: %q", rawAccess[:min(30, len(rawAccess))])
|
||||
}
|
||||
if !strings.HasPrefix(rawRefresh, "aes-gcm:") {
|
||||
t.Errorf("refresh_token in DB not encrypted: %q", rawRefresh[:min(30, len(rawRefresh))])
|
||||
}
|
||||
|
||||
// Store should decrypt transparently on read.
|
||||
got, err := st.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetOAuthToken() error: %v", err)
|
||||
}
|
||||
if got.AccessToken != plainAccess {
|
||||
t.Errorf("decrypted AccessToken = %q, want %q", got.AccessToken, plainAccess)
|
||||
}
|
||||
if got.RefreshToken != plainRefresh {
|
||||
t.Errorf("decrypted RefreshToken = %q, want %q", got.RefreshToken, plainRefresh)
|
||||
}
|
||||
}
|
||||
|
||||
// min helper for Go versions without builtin min.
|
||||
func min(a, b int) int {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Delete
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestMCPOAuthDeleteGlobal(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
seedMCPOAuthToken(t, db, tenantID, serverID, "")
|
||||
st := oauthStore(t)
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
if err := st.DeleteOAuthToken(ctx, serverID, tenantID); err != nil {
|
||||
t.Fatalf("DeleteOAuthToken() error: %v", err)
|
||||
}
|
||||
|
||||
got, err := st.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetOAuthToken() after delete error: %v", err)
|
||||
}
|
||||
if got != nil {
|
||||
t.Error("expected nil after delete, got token")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPOAuthDeletePerUser(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
userID := "user-del-" + uuid.New().String()[:8]
|
||||
|
||||
// Seed both global and per-user.
|
||||
seedMCPOAuthToken(t, db, tenantID, serverID, "")
|
||||
seedMCPOAuthToken(t, db, tenantID, serverID, userID)
|
||||
st := oauthStore(t)
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
// Delete per-user only.
|
||||
if err := st.DeleteUserOAuthToken(ctx, serverID, tenantID, userID); err != nil {
|
||||
t.Fatalf("DeleteUserOAuthToken() error: %v", err)
|
||||
}
|
||||
|
||||
// Per-user should be gone.
|
||||
userTok, _ := st.GetUserOAuthToken(ctx, serverID, tenantID, userID)
|
||||
if userTok != nil {
|
||||
t.Error("expected per-user token deleted, still exists")
|
||||
}
|
||||
|
||||
// Global should remain.
|
||||
globalTok, _ := st.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if globalTok == nil {
|
||||
t.Error("expected global token to remain after per-user delete")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPOAuthDeleteServerTokens(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
otherServerID := seedMCPServer(t, db, tenantID)
|
||||
userA := "user-a-" + uuid.New().String()[:8]
|
||||
userB := "user-b-" + uuid.New().String()[:8]
|
||||
|
||||
// Seed global + two per-user tokens for the target server.
|
||||
seedMCPOAuthToken(t, db, tenantID, serverID, "")
|
||||
seedMCPOAuthToken(t, db, tenantID, serverID, userA)
|
||||
seedMCPOAuthToken(t, db, tenantID, serverID, userB)
|
||||
// Seed an unrelated token on a different server in the same tenant.
|
||||
seedMCPOAuthToken(t, db, tenantID, otherServerID, "")
|
||||
|
||||
st := oauthStore(t)
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
if err := st.DeleteServerOAuthTokens(ctx, serverID, tenantID); err != nil {
|
||||
t.Fatalf("DeleteServerOAuthTokens() error: %v", err)
|
||||
}
|
||||
|
||||
// All tokens for the target server must be gone (global + both per-user).
|
||||
if tok, _ := st.GetOAuthToken(ctx, serverID, tenantID); tok != nil {
|
||||
t.Error("expected global token deleted")
|
||||
}
|
||||
if tok, _ := st.GetUserOAuthToken(ctx, serverID, tenantID, userA); tok != nil {
|
||||
t.Error("expected per-user token (userA) deleted")
|
||||
}
|
||||
if tok, _ := st.GetUserOAuthToken(ctx, serverID, tenantID, userB); tok != nil {
|
||||
t.Error("expected per-user token (userB) deleted")
|
||||
}
|
||||
|
||||
// The other server's token must remain untouched.
|
||||
if tok, _ := st.GetOAuthToken(ctx, otherServerID, tenantID); tok == nil {
|
||||
t.Error("token for a different server must not be deleted")
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Tenant isolation
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestMCPOAuthTenantIsolation(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantA, _ := seedTenantAgent(t, db)
|
||||
tenantB, _ := seedTenantAgent(t, db)
|
||||
|
||||
// Create server for tenant A.
|
||||
serverID := seedMCPServer(t, db, tenantA)
|
||||
|
||||
// Seed token for tenant A.
|
||||
seedMCPOAuthToken(t, db, tenantA, serverID, "")
|
||||
|
||||
st := oauthStore(t)
|
||||
|
||||
// Tenant A can read its own token.
|
||||
tokA, err := st.GetOAuthToken(tenantCtx(tenantA), serverID, tenantA)
|
||||
if err != nil {
|
||||
t.Fatalf("GetOAuthToken(tenantA): %v", err)
|
||||
}
|
||||
if tokA == nil {
|
||||
t.Error("tenant A should see its own token")
|
||||
}
|
||||
|
||||
// Tenant B cannot read tenant A's token (different tenantID in query).
|
||||
tokB, err := st.GetOAuthToken(tenantCtx(tenantB), serverID, tenantB)
|
||||
if err != nil {
|
||||
t.Fatalf("GetOAuthToken(tenantB): %v", err)
|
||||
}
|
||||
if tokB != nil {
|
||||
t.Error("tenant B should NOT see tenant A's token")
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// Cascade on server delete
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestMCPOAuthCascadeOnServerDelete(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
seedMCPOAuthToken(t, db, tenantID, serverID, "")
|
||||
|
||||
st := oauthStore(t)
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
// Verify token exists.
|
||||
tok, _ := st.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if tok == nil {
|
||||
t.Fatal("token should exist before server delete")
|
||||
}
|
||||
|
||||
// Delete server — ON DELETE CASCADE should delete token.
|
||||
if _, err := db.Exec("DELETE FROM mcp_servers WHERE id = $1", serverID); err != nil {
|
||||
t.Fatalf("delete server: %v", err)
|
||||
}
|
||||
|
||||
// Token should be gone.
|
||||
var count int
|
||||
db.QueryRow("SELECT COUNT(*) FROM mcp_oauth_tokens WHERE server_id = $1", serverID).Scan(&count)
|
||||
if count != 0 {
|
||||
t.Errorf("expected 0 tokens after server delete (CASCADE), got %d", count)
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
// ExpiresAt handling
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestMCPOAuthExpiresAtRoundTrip(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
st := oauthStore(t)
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
expiresAt := time.Now().Add(time.Hour).Truncate(time.Second).UTC()
|
||||
tok := &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: "tok",
|
||||
TokenType: "Bearer",
|
||||
ExpiresAt: &expiresAt,
|
||||
DCRClientID: "c",
|
||||
DCRIssuer: "x",
|
||||
TokenEndpoint: "https://t.example.com/token",
|
||||
}
|
||||
if err := st.UpsertOAuthToken(ctx, tok); err != nil {
|
||||
t.Fatalf("UpsertOAuthToken() error: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
db.Exec("DELETE FROM mcp_oauth_tokens WHERE server_id = $1", serverID)
|
||||
})
|
||||
|
||||
got, err := st.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if err != nil {
|
||||
t.Fatalf("GetOAuthToken() error: %v", err)
|
||||
}
|
||||
if got.ExpiresAt == nil {
|
||||
t.Fatal("ExpiresAt should not be nil")
|
||||
}
|
||||
if !got.ExpiresAt.Equal(expiresAt) {
|
||||
t.Errorf("ExpiresAt = %v, want %v", got.ExpiresAt, expiresAt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPOAuthNullRefreshToken(t *testing.T) {
|
||||
db := testDB(t)
|
||||
tenantID, _ := seedTenantAgent(t, db)
|
||||
serverID := seedMCPServer(t, db, tenantID)
|
||||
st := oauthStore(t)
|
||||
ctx := tenantCtx(tenantID)
|
||||
|
||||
tok := &store.MCPOAuthToken{
|
||||
ServerID: serverID,
|
||||
TenantID: tenantID,
|
||||
AccessToken: "access-only",
|
||||
RefreshToken: "", // no refresh token
|
||||
TokenType: "Bearer",
|
||||
DCRClientID: "c",
|
||||
DCRIssuer: "x",
|
||||
TokenEndpoint: "https://t.example.com/token",
|
||||
}
|
||||
if err := st.UpsertOAuthToken(ctx, tok); err != nil {
|
||||
t.Fatalf("UpsertOAuthToken() error: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
db.Exec("DELETE FROM mcp_oauth_tokens WHERE server_id = $1", serverID)
|
||||
})
|
||||
|
||||
got, _ := st.GetOAuthToken(ctx, serverID, tenantID)
|
||||
if got == nil {
|
||||
t.Fatal("GetOAuthToken() returned nil")
|
||||
}
|
||||
if got.RefreshToken != "" {
|
||||
t.Errorf("RefreshToken = %q, want empty string (no refresh token)", got.RefreshToken)
|
||||
}
|
||||
|
||||
// DB should not store encrypted empty string — check raw column.
|
||||
var rawRefresh sql.NullString
|
||||
db.QueryRow("SELECT refresh_token FROM mcp_oauth_tokens WHERE server_id = $1", serverID).Scan(&rawRefresh)
|
||||
if rawRefresh.Valid && strings.HasPrefix(rawRefresh.String, "aes-gcm:") {
|
||||
t.Error("empty refresh_token should not be stored as encrypted value in DB")
|
||||
}
|
||||
}
|
||||
@@ -95,7 +95,7 @@ export function KeyValueEditor({
|
||||
onChange={(e) => updateEntry(idx, { value: e.target.value })}
|
||||
placeholder={valuePlaceholder}
|
||||
rows={2}
|
||||
className="flex-[2] text-base md:text-sm min-h-[38px] resize-y"
|
||||
className="flex-2 text-base md:text-sm min-h-[38px] resize-y"
|
||||
/>
|
||||
) : (
|
||||
<Input
|
||||
|
||||
@@ -72,7 +72,39 @@
|
||||
"toolHintsPerTool": "Per-tool hints",
|
||||
"toolHintsToolNamePlaceholder": "Tool name (e.g. search)",
|
||||
"toolHintsToolValuePlaceholder": "Hint text shown to the LLM",
|
||||
"toolHintsAdd": "Add hint"
|
||||
"toolHintsAdd": "Add hint",
|
||||
"oauth": {
|
||||
"title": "OAuth Authentication",
|
||||
"titleHint": "Connect to MCP servers that require OAuth 2.1 authorization.",
|
||||
"enable": "Enable OAuth",
|
||||
"useDcr": "Dynamic Client Registration (DCR)",
|
||||
"useDcrHint": "Automatically register as an OAuth client. Disable to enter credentials manually.",
|
||||
"grantType": "Grant Type",
|
||||
"pkce": "Authorization Code + PKCE (Public Client)",
|
||||
"authCode": "Authorization Code (Confidential Client)",
|
||||
"clientCreds": "Client Credentials",
|
||||
"authEndpoint": "Authorization URL",
|
||||
"authEndpointPlaceholder": "https://auth.example.com/oauth/authorize",
|
||||
"tokenEndpoint": "Token URL",
|
||||
"tokenEndpointPlaceholder": "https://auth.example.com/oauth/token",
|
||||
"clientId": "Client ID",
|
||||
"clientSecret": "Client Secret",
|
||||
"clientSecretPlaceholder": "Leave empty to keep existing",
|
||||
"scope": "Scopes",
|
||||
"scopePlaceholder": "e.g. read write offline_access",
|
||||
"modePerUser": "Auth mode: Per-user (each user must authorize independently)",
|
||||
"modeGlobal": "Auth mode: Global (shared token for all users)",
|
||||
"authorize": "Authorize",
|
||||
"reauthorize": "Re-authorize",
|
||||
"revoke": "Revoke",
|
||||
"authorized": "Authorized",
|
||||
"notAuthorized": "Not authorized",
|
||||
"expired": "Token expired",
|
||||
"authorizing": "Opening authorization...",
|
||||
"authSuccess": "Authorization successful",
|
||||
"authFailed": "Authorization failed",
|
||||
"savingFirst": "Saving server before authorizing..."
|
||||
}
|
||||
},
|
||||
"grants": {
|
||||
"title": "Agent Grants - {{name}}",
|
||||
|
||||
@@ -72,7 +72,39 @@
|
||||
"toolHintsPerTool": "Gợi ý theo tool cụ thể",
|
||||
"toolHintsToolNamePlaceholder": "Tên tool (vd: search)",
|
||||
"toolHintsToolValuePlaceholder": "Nội dung gợi ý sẽ gửi cho LLM",
|
||||
"toolHintsAdd": "Thêm gợi ý"
|
||||
"toolHintsAdd": "Thêm gợi ý",
|
||||
"oauth": {
|
||||
"title": "Xác thực OAuth",
|
||||
"titleHint": "Kết nối đến MCP server yêu cầu OAuth 2.1.",
|
||||
"enable": "Bật OAuth",
|
||||
"useDcr": "Đăng ký client tự động (DCR)",
|
||||
"useDcrHint": "Tự động đăng ký làm OAuth client. Tắt để nhập thủ công.",
|
||||
"grantType": "Loại Grant",
|
||||
"pkce": "Authorization Code + PKCE (Public Client)",
|
||||
"authCode": "Authorization Code (Confidential Client)",
|
||||
"clientCreds": "Client Credentials",
|
||||
"authEndpoint": "URL Authorize",
|
||||
"authEndpointPlaceholder": "https://auth.example.com/oauth/authorize",
|
||||
"tokenEndpoint": "URL Token",
|
||||
"tokenEndpointPlaceholder": "https://auth.example.com/oauth/token",
|
||||
"clientId": "Client ID",
|
||||
"clientSecret": "Client Secret",
|
||||
"clientSecretPlaceholder": "Để trống để giữ nguyên",
|
||||
"scope": "Scopes",
|
||||
"scopePlaceholder": "vd: read write offline_access",
|
||||
"modePerUser": "Chế độ: Theo user (mỗi user tự ủy quyền riêng)",
|
||||
"modeGlobal": "Chế độ: Dùng chung (token chia sẻ cho mọi user)",
|
||||
"authorize": "Ủy quyền",
|
||||
"reauthorize": "Ủy quyền lại",
|
||||
"revoke": "Thu hồi",
|
||||
"authorized": "Đã ủy quyền",
|
||||
"notAuthorized": "Chưa ủy quyền",
|
||||
"expired": "Token đã hết hạn",
|
||||
"authorizing": "Đang mở cửa sổ ủy quyền...",
|
||||
"authSuccess": "Ủy quyền thành công",
|
||||
"authFailed": "Ủy quyền thất bại",
|
||||
"savingFirst": "Đang lưu server trước khi ủy quyền..."
|
||||
}
|
||||
},
|
||||
"grants": {
|
||||
"title": "Quyền agent - {{name}}",
|
||||
|
||||
@@ -72,7 +72,39 @@
|
||||
"toolHintsPerTool": "针对单个工具的提示",
|
||||
"toolHintsToolNamePlaceholder": "工具名称(例如 search)",
|
||||
"toolHintsToolValuePlaceholder": "发送给 LLM 的提示内容",
|
||||
"toolHintsAdd": "添加提示"
|
||||
"toolHintsAdd": "添加提示",
|
||||
"oauth": {
|
||||
"title": "OAuth 认证",
|
||||
"titleHint": "连接到需要 OAuth 2.1 授权的 MCP 服务器。",
|
||||
"enable": "启用 OAuth",
|
||||
"useDcr": "动态客户端注册(DCR)",
|
||||
"useDcrHint": "自动注册为 OAuth 客户端。关闭后可手动填写凭据。",
|
||||
"grantType": "授权类型",
|
||||
"pkce": "Authorization Code + PKCE(公开客户端)",
|
||||
"authCode": "Authorization Code(机密客户端)",
|
||||
"clientCreds": "Client Credentials",
|
||||
"authEndpoint": "授权 URL",
|
||||
"authEndpointPlaceholder": "https://auth.example.com/oauth/authorize",
|
||||
"tokenEndpoint": "令牌 URL",
|
||||
"tokenEndpointPlaceholder": "https://auth.example.com/oauth/token",
|
||||
"clientId": "Client ID",
|
||||
"clientSecret": "Client Secret",
|
||||
"clientSecretPlaceholder": "留空则保持不变",
|
||||
"scope": "Scopes",
|
||||
"scopePlaceholder": "例如:read write offline_access",
|
||||
"modePerUser": "认证模式:按用户(每个用户独立授权)",
|
||||
"modeGlobal": "认证模式:共享(所有用户共用同一令牌)",
|
||||
"authorize": "授权",
|
||||
"reauthorize": "重新授权",
|
||||
"revoke": "撤销",
|
||||
"authorized": "已授权",
|
||||
"notAuthorized": "未授权",
|
||||
"expired": "令牌已过期",
|
||||
"authorizing": "正在打开授权窗口...",
|
||||
"authSuccess": "授权成功",
|
||||
"authFailed": "授权失败",
|
||||
"savingFirst": "正在先保存服务器..."
|
||||
}
|
||||
},
|
||||
"grants": {
|
||||
"title": "Agent授权 - {{name}}",
|
||||
|
||||
@@ -4,9 +4,9 @@ import i18next from "i18next";
|
||||
import { useHttp } from "@/hooks/use-ws";
|
||||
import { queryKeys } from "@/lib/query-keys";
|
||||
import { toast } from "@/stores/use-toast-store";
|
||||
import type { MCPServerData, MCPServerInput, MCPAgentGrant, MCPToolInfo, MCPUserCredentialStatus, MCPUserCredentialInput } from "@/types/mcp";
|
||||
import type { MCPServerData, MCPServerInput, MCPAgentGrant, MCPToolInfo, MCPUserCredentialStatus, MCPUserCredentialInput, MCPOAuthStatus } from "@/types/mcp";
|
||||
|
||||
export type { MCPServerData, MCPServerInput, MCPAgentGrant, MCPToolInfo, MCPUserCredentialStatus, MCPUserCredentialInput };
|
||||
export type { MCPServerData, MCPServerInput, MCPAgentGrant, MCPToolInfo, MCPUserCredentialStatus, MCPUserCredentialInput, MCPOAuthStatus };
|
||||
|
||||
export function useMCP() {
|
||||
const http = useHttp();
|
||||
@@ -109,7 +109,7 @@ export function useMCP() {
|
||||
);
|
||||
|
||||
const testConnection = useCallback(
|
||||
async (data: { transport: string; command?: string; args?: string[]; url?: string; headers?: Record<string, string>; env?: Record<string, string> }) => {
|
||||
async (data: { server_id?: string; transport: string; command?: string; args?: string[]; url?: string; headers?: Record<string, string>; env?: Record<string, string> }) => {
|
||||
return http.post<{ success: boolean; tool_count?: number; error?: string }>("/v1/mcp/servers/test", data);
|
||||
},
|
||||
[http],
|
||||
@@ -160,6 +160,32 @@ export function useMCP() {
|
||||
[http],
|
||||
);
|
||||
|
||||
const startOAuth = useCallback(
|
||||
async (serverId: string, mcpUrl: string, userId?: string) => {
|
||||
return http.post<{ auth_url: string; state: string; client_id: string; issuer: string; completed?: boolean }>(
|
||||
"/v1/mcp/oauth/start",
|
||||
{ server_id: serverId, mcp_url: mcpUrl, user_id: userId ?? "" },
|
||||
);
|
||||
},
|
||||
[http],
|
||||
);
|
||||
|
||||
const getOAuthStatus = useCallback(
|
||||
async (serverId: string, userId?: string) => {
|
||||
const qs = userId ? `?user_id=${encodeURIComponent(userId)}` : "";
|
||||
return http.get<MCPOAuthStatus>(`/v1/mcp/oauth/status/${serverId}${qs}`);
|
||||
},
|
||||
[http],
|
||||
);
|
||||
|
||||
const revokeOAuth = useCallback(
|
||||
async (serverId: string, userId?: string) => {
|
||||
const qs = userId ? `?user_id=${encodeURIComponent(userId)}` : "";
|
||||
await http.delete(`/v1/mcp/oauth/token/${serverId}${qs}`);
|
||||
},
|
||||
[http],
|
||||
);
|
||||
|
||||
return {
|
||||
servers,
|
||||
loading,
|
||||
@@ -178,5 +204,8 @@ export function useMCP() {
|
||||
getUserCredentials,
|
||||
setUserCredentials,
|
||||
deleteUserCredentials,
|
||||
startOAuth,
|
||||
getOAuthStatus,
|
||||
revokeOAuth,
|
||||
};
|
||||
}
|
||||
@@ -2,7 +2,7 @@ import { useState, useEffect } from "react";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { useForm } from "react-hook-form";
|
||||
import { zodResolver } from "@hookform/resolvers/zod";
|
||||
import { Loader2, CheckCircle2, XCircle } from "lucide-react";
|
||||
import { Loader2, CheckCircle2, XCircle, ShieldCheck } from "lucide-react";
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
@@ -32,8 +32,9 @@ interface MCPFormDialogProps {
|
||||
open: boolean;
|
||||
onOpenChange: (open: boolean) => void;
|
||||
server?: MCPServerData | null;
|
||||
onSubmit: (data: MCPServerInput) => Promise<unknown>;
|
||||
onSubmit: (data: MCPServerInput) => Promise<MCPServerData | void>;
|
||||
onTest: (data: {
|
||||
server_id?: string;
|
||||
transport: string;
|
||||
command?: string;
|
||||
args?: string[];
|
||||
@@ -41,14 +42,18 @@ interface MCPFormDialogProps {
|
||||
headers?: Record<string, string>;
|
||||
env?: Record<string, string>;
|
||||
}) => Promise<{ success: boolean; tool_count?: number; error?: string }>;
|
||||
/** Called after a successful save (or immediately for edit) to open the OAuth authorization dialog. */
|
||||
onAuthorize?: (server: MCPServerData) => void;
|
||||
}
|
||||
|
||||
export function MCPFormDialog({ open, onOpenChange, server, onSubmit, onTest }: MCPFormDialogProps) {
|
||||
export function MCPFormDialog({ open, onOpenChange, server, onSubmit, onTest, onAuthorize }: MCPFormDialogProps) {
|
||||
const { t } = useTranslation("mcp");
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [testing, setTesting] = useState(false);
|
||||
const [testResult, setTestResult] = useState<{ success: boolean; tool_count?: number; error?: string } | null>(null);
|
||||
const [error, setError] = useState("");
|
||||
// Tracks a server created via handleAuthorize (ADD mode) to prevent duplicate POST on "Create" click.
|
||||
const [createdServer, setCreatedServer] = useState<MCPServerData | null>(null);
|
||||
|
||||
const form = useForm<MCPFormData>({
|
||||
resolver: zodResolver(mcpFormSchema),
|
||||
@@ -68,6 +73,14 @@ export function MCPFormDialog({ open, onOpenChange, server, onSubmit, onTest }:
|
||||
requireUserCreds: false,
|
||||
toolHintsGlobal: "",
|
||||
toolHintsTools: {},
|
||||
oauthEnabled: false,
|
||||
oauthUseDcr: true,
|
||||
oauthGrantType: "pkce",
|
||||
oauthAuthEndpoint: "",
|
||||
oauthTokenEndpoint: "",
|
||||
oauthClientId: "",
|
||||
oauthClientSecret: "",
|
||||
oauthScope: "",
|
||||
},
|
||||
});
|
||||
|
||||
@@ -78,10 +91,12 @@ export function MCPFormDialog({ open, onOpenChange, server, onSubmit, onTest }:
|
||||
const url = watch("url");
|
||||
const headers = watch("headers") as Record<string, string>;
|
||||
const env = watch("env") as Record<string, string>;
|
||||
const oauthEnabled = watch("oauthEnabled");
|
||||
const isStdio = transport === "stdio";
|
||||
|
||||
useEffect(() => {
|
||||
if (open) {
|
||||
const oauth = server?.settings?.oauth;
|
||||
reset({
|
||||
name: server?.name ?? "",
|
||||
displayName: server?.display_name ?? "",
|
||||
@@ -97,9 +112,18 @@ export function MCPFormDialog({ open, onOpenChange, server, onSubmit, onTest }:
|
||||
requireUserCreds: server?.settings?.require_user_credentials ?? false,
|
||||
toolHintsGlobal: server?.settings?.tool_hints?.global ?? "",
|
||||
toolHintsTools: server?.settings?.tool_hints?.tools ?? {},
|
||||
oauthEnabled: oauth?.auth_type === "oauth",
|
||||
oauthUseDcr: oauth?.use_dcr ?? true,
|
||||
oauthGrantType: (oauth?.grant_type as MCPFormData["oauthGrantType"]) ?? "pkce",
|
||||
oauthAuthEndpoint: oauth?.auth_endpoint ?? "",
|
||||
oauthTokenEndpoint: oauth?.token_endpoint ?? "",
|
||||
oauthClientId: oauth?.client_id ?? "",
|
||||
oauthClientSecret: "", // never pre-fill secret
|
||||
oauthScope: oauth?.scope ?? "",
|
||||
});
|
||||
setError("");
|
||||
setTestResult(null);
|
||||
setCreatedServer(null);
|
||||
}
|
||||
}, [open, server, reset]);
|
||||
|
||||
@@ -119,13 +143,62 @@ export function MCPFormDialog({ open, onOpenChange, server, onSubmit, onTest }:
|
||||
}
|
||||
}
|
||||
|
||||
// Always send the collection fields (even when emptied) so clearing all rows
|
||||
// actually clears them server-side. UpdateServer applies a partial update keyed
|
||||
// by the fields present in the request body — omitting an emptied field would
|
||||
// leave the previous value untouched, making it impossible to delete the last
|
||||
// header/env/arg. stdio servers don't use headers/url; remote servers don't use args.
|
||||
return {
|
||||
transport,
|
||||
command: isStdio ? resolvedCommand : undefined,
|
||||
args: parsedArgs,
|
||||
args: isStdio ? (parsedArgs ?? []) : undefined,
|
||||
url: !isStdio ? url.trim() : undefined,
|
||||
headers: !isStdio && Object.keys(headers).length > 0 ? headers : undefined,
|
||||
env: Object.keys(env).length > 0 ? env : undefined,
|
||||
headers: isStdio ? undefined : headers,
|
||||
env,
|
||||
};
|
||||
};
|
||||
|
||||
const buildPayload = (data: MCPFormData): MCPServerInput => {
|
||||
const trimmedGlobal = data.toolHintsGlobal.trim();
|
||||
const trimmedTools: Record<string, string> = {};
|
||||
for (const [k, v] of Object.entries(data.toolHintsTools)) {
|
||||
const key = k.trim();
|
||||
const val = v.trim();
|
||||
if (key && val) trimmedTools[key] = val;
|
||||
}
|
||||
const hasHints = trimmedGlobal !== "" || Object.keys(trimmedTools).length > 0;
|
||||
const settings: NonNullable<MCPServerInput["settings"]> = {
|
||||
require_user_credentials: data.requireUserCreds,
|
||||
};
|
||||
if (hasHints) {
|
||||
settings.tool_hints = {
|
||||
...(trimmedGlobal ? { global: trimmedGlobal } : {}),
|
||||
...(Object.keys(trimmedTools).length > 0 ? { tools: trimmedTools } : {}),
|
||||
};
|
||||
}
|
||||
if (!isStdio && data.oauthEnabled) {
|
||||
settings.oauth = {
|
||||
auth_type: "oauth",
|
||||
use_dcr: data.oauthUseDcr,
|
||||
grant_type: data.oauthGrantType,
|
||||
// Manual endpoints only apply when DCR is off; never persist stale values
|
||||
// left in form state after toggling DCR back on (keeps backend + the
|
||||
// token-purge fingerprint consistent with the actual auth mode).
|
||||
...(!data.oauthUseDcr && data.oauthAuthEndpoint.trim() ? { auth_endpoint: data.oauthAuthEndpoint.trim() } : {}),
|
||||
...(!data.oauthUseDcr && data.oauthTokenEndpoint.trim() ? { token_endpoint: data.oauthTokenEndpoint.trim() } : {}),
|
||||
...(data.oauthClientId.trim() ? { client_id: data.oauthClientId.trim() } : {}),
|
||||
...(data.oauthClientSecret.trim() ? { client_secret: data.oauthClientSecret.trim() } : {}),
|
||||
...(data.oauthScope.trim() ? { scope: data.oauthScope.trim() } : {}),
|
||||
};
|
||||
}
|
||||
return {
|
||||
name: data.name.trim(),
|
||||
display_name: data.displayName.trim() || undefined,
|
||||
...buildConnectionData(),
|
||||
tool_prefix: data.toolPrefix.trim() || undefined,
|
||||
timeout_sec: data.timeout,
|
||||
settings,
|
||||
enabled: data.enabled,
|
||||
};
|
||||
};
|
||||
|
||||
@@ -136,7 +209,7 @@ export function MCPFormDialog({ open, onOpenChange, server, onSubmit, onTest }:
|
||||
setError("");
|
||||
setTestResult(null);
|
||||
try {
|
||||
const result = await onTest(buildConnectionData());
|
||||
const result = await onTest({ server_id: server?.id, ...buildConnectionData() });
|
||||
setTestResult(result);
|
||||
} catch (err: unknown) {
|
||||
setTestResult({ success: false, error: err instanceof Error ? err.message : t("form.errors.connectionFailed") });
|
||||
@@ -150,38 +223,16 @@ export function MCPFormDialog({ open, onOpenChange, server, onSubmit, onTest }:
|
||||
if (isStdio && !data.command.trim()) { setError(t("form.errors.commandRequired")); return; }
|
||||
if (!isStdio && !data.url.trim()) { setError(t("form.errors.urlRequired")); return; }
|
||||
|
||||
// ADD mode: server already created via Authorize flow — don't POST again (duplicate key).
|
||||
if (!server && createdServer) {
|
||||
onOpenChange(false);
|
||||
return;
|
||||
}
|
||||
|
||||
setLoading(true);
|
||||
setError("");
|
||||
try {
|
||||
// Build settings payload: include tool_hints only when at least one field
|
||||
// is populated, so servers without hints keep a clean settings object.
|
||||
const trimmedGlobal = data.toolHintsGlobal.trim();
|
||||
const trimmedTools: Record<string, string> = {};
|
||||
for (const [k, v] of Object.entries(data.toolHintsTools)) {
|
||||
const key = k.trim();
|
||||
const val = v.trim();
|
||||
if (key && val) trimmedTools[key] = val;
|
||||
}
|
||||
const hasHints = trimmedGlobal !== "" || Object.keys(trimmedTools).length > 0;
|
||||
const settings: NonNullable<MCPServerInput["settings"]> = {
|
||||
require_user_credentials: data.requireUserCreds,
|
||||
};
|
||||
if (hasHints) {
|
||||
settings.tool_hints = {
|
||||
...(trimmedGlobal ? { global: trimmedGlobal } : {}),
|
||||
...(Object.keys(trimmedTools).length > 0 ? { tools: trimmedTools } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
await onSubmit({
|
||||
name: data.name.trim(),
|
||||
display_name: data.displayName.trim() || undefined,
|
||||
...buildConnectionData(),
|
||||
tool_prefix: data.toolPrefix.trim() || undefined,
|
||||
timeout_sec: data.timeout,
|
||||
settings,
|
||||
enabled: data.enabled,
|
||||
});
|
||||
await onSubmit(buildPayload(data));
|
||||
onOpenChange(false);
|
||||
} catch (err: unknown) {
|
||||
setError(err instanceof Error ? err.message : t("form.errors.saveFailed", "Save failed"));
|
||||
@@ -190,6 +241,50 @@ export function MCPFormDialog({ open, onOpenChange, server, onSubmit, onTest }:
|
||||
}
|
||||
});
|
||||
|
||||
const handleAuthorize = async () => {
|
||||
if (!onAuthorize) return;
|
||||
const values = form.getValues();
|
||||
|
||||
if (!isValidSlug(values.name.trim())) { setError(t("form.errors.nameSlug")); return; }
|
||||
if (isStdio && !values.command.trim()) { setError(t("form.errors.commandRequired")); return; }
|
||||
if (!isStdio && !values.url.trim()) { setError(t("form.errors.urlRequired")); return; }
|
||||
|
||||
if (server) {
|
||||
// EDIT: if URL changed in form, save first so OAuth flow uses the new URL.
|
||||
const currentUrl = values.url.trim();
|
||||
if (!isStdio && currentUrl && currentUrl !== server.url) {
|
||||
setLoading(true);
|
||||
setError("");
|
||||
try {
|
||||
await onSubmit(buildPayload(values));
|
||||
onAuthorize({ ...server, url: currentUrl });
|
||||
} catch (err: unknown) {
|
||||
setError(err instanceof Error ? err.message : t("form.errors.saveFailed", "Save failed"));
|
||||
} finally {
|
||||
setLoading(false);
|
||||
}
|
||||
return;
|
||||
}
|
||||
onAuthorize(server);
|
||||
return;
|
||||
}
|
||||
|
||||
// ADD: save server first, then open OAuth dialog (dialog stays open).
|
||||
setLoading(true);
|
||||
setError("");
|
||||
try {
|
||||
const created = await onSubmit(buildPayload(values));
|
||||
if (created && "id" in created) {
|
||||
setCreatedServer(created as MCPServerData);
|
||||
onAuthorize(created as MCPServerData);
|
||||
}
|
||||
} catch (err: unknown) {
|
||||
setError(err instanceof Error ? err.message : t("form.errors.saveFailed", "Save failed"));
|
||||
} finally {
|
||||
setLoading(false);
|
||||
}
|
||||
};
|
||||
|
||||
return (
|
||||
<Dialog open={open} onOpenChange={(v) => !loading && onOpenChange(v)}>
|
||||
<DialogContent className="max-h-[85vh] flex flex-col sm:max-w-xl">
|
||||
@@ -199,7 +294,7 @@ export function MCPFormDialog({ open, onOpenChange, server, onSubmit, onTest }:
|
||||
|
||||
<div className="grid gap-4 py-2 -mx-4 px-4 sm:-mx-6 sm:px-6 overflow-y-auto min-h-0">
|
||||
<McpConnectionFields form={form} />
|
||||
<McpSettingsFields form={form} />
|
||||
<McpSettingsFields form={form} isEditing={!!server} />
|
||||
{error && <p className="text-sm text-destructive">{error}</p>}
|
||||
</div>
|
||||
|
||||
@@ -219,11 +314,25 @@ export function MCPFormDialog({ open, onOpenChange, server, onSubmit, onTest }:
|
||||
)}
|
||||
</div>
|
||||
<div className="flex gap-2">
|
||||
{!isStdio && oauthEnabled && onAuthorize && (
|
||||
<Button
|
||||
type="button"
|
||||
variant="secondary"
|
||||
size="sm"
|
||||
onClick={handleAuthorize}
|
||||
disabled={loading || testing}
|
||||
className="gap-1"
|
||||
>
|
||||
{loading
|
||||
? <><Loader2 className="h-3.5 w-3.5 animate-spin" />{t("form.oauth.savingFirst")}</>
|
||||
: <><ShieldCheck className="h-3.5 w-3.5" />{t("form.oauth.authorize")}</>}
|
||||
</Button>
|
||||
)}
|
||||
<Button variant="outline" onClick={() => onOpenChange(false)} disabled={loading}>
|
||||
{t("form.cancel")}
|
||||
</Button>
|
||||
<Button onClick={handleSubmit} disabled={loading}>
|
||||
{loading ? t("form.saving") : server ? t("form.update") : t("form.create")}
|
||||
{loading ? t("form.saving") : (server || createdServer) ? t("form.update") : t("form.create")}
|
||||
</Button>
|
||||
</div>
|
||||
</DialogFooter>
|
||||
|
||||
@@ -0,0 +1,222 @@
|
||||
import { useState, useEffect, useRef, useCallback } from "react";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { ShieldCheck, ShieldX, Loader2, ExternalLink } from "lucide-react";
|
||||
import { useWsEvent } from "@/hooks/use-ws-event";
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
DialogHeader,
|
||||
DialogTitle,
|
||||
DialogFooter,
|
||||
} from "@/components/ui/dialog";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import type { MCPServerData, MCPOAuthStatus } from "./hooks/use-mcp";
|
||||
|
||||
// MCPOAuthDialog manages the SERVER-LEVEL (global) OAuth token only — it always
|
||||
// operates on user_id="" (no per-user scope). Per-user OAuth is handled by
|
||||
// MCPUserCredentialsDialog.
|
||||
interface MCPOAuthDialogProps {
|
||||
open: boolean;
|
||||
onOpenChange: (open: boolean) => void;
|
||||
server: MCPServerData;
|
||||
onStartOAuth: (serverId: string, mcpUrl: string, userId?: string) => Promise<{ auth_url: string; state: string; client_id: string; issuer: string; completed?: boolean }>;
|
||||
onGetStatus: (serverId: string, userId?: string) => Promise<MCPOAuthStatus>;
|
||||
onRevoke: (serverId: string, userId?: string) => Promise<void>;
|
||||
}
|
||||
|
||||
export function MCPOAuthDialog({
|
||||
open,
|
||||
onOpenChange,
|
||||
server,
|
||||
onStartOAuth,
|
||||
onGetStatus,
|
||||
onRevoke,
|
||||
}: MCPOAuthDialogProps) {
|
||||
const { t } = useTranslation("mcp");
|
||||
const [status, setStatus] = useState<MCPOAuthStatus | null>(null);
|
||||
const [authorizing, setAuthorizing] = useState(false);
|
||||
const [revoking, setRevoking] = useState(false);
|
||||
const [error, setError] = useState("");
|
||||
const popupRef = useRef<Window | null>(null);
|
||||
|
||||
const loadStatus = useCallback(async () => {
|
||||
try {
|
||||
// Global/server token only — no user_id.
|
||||
const s = await onGetStatus(server.id);
|
||||
setStatus(s);
|
||||
} catch {
|
||||
setStatus(null);
|
||||
}
|
||||
}, [server.id, onGetStatus]);
|
||||
|
||||
useEffect(() => {
|
||||
if (open) {
|
||||
setError("");
|
||||
loadStatus();
|
||||
}
|
||||
}, [open, loadStatus]);
|
||||
|
||||
// Close popup on unmount.
|
||||
useEffect(() => {
|
||||
return () => { popupRef.current?.close(); };
|
||||
}, []);
|
||||
|
||||
// Primary notification path: WebSocket event from backend after token exchange.
|
||||
// Fires for both success and error when the OAuth callback is processed.
|
||||
useWsEvent("mcp.oauth_complete", (payload: unknown) => {
|
||||
const p = payload as { serverId?: string; userId?: string; status?: string; error?: string };
|
||||
if (p.serverId !== server.id) return;
|
||||
// Only react to the global (server-level) token event.
|
||||
if ((p.userId ?? "") !== "") return;
|
||||
popupRef.current?.close();
|
||||
popupRef.current = null;
|
||||
setAuthorizing(false);
|
||||
if (p.status === "success") {
|
||||
loadStatus();
|
||||
} else {
|
||||
setError(p.error ?? t("form.oauth.authFailed"));
|
||||
}
|
||||
});
|
||||
|
||||
const handleAuthorize = async () => {
|
||||
if (!server.url) {
|
||||
setError("Server has no URL configured");
|
||||
return;
|
||||
}
|
||||
setAuthorizing(true);
|
||||
setError("");
|
||||
try {
|
||||
const { auth_url, completed } = await onStartOAuth(server.id, server.url);
|
||||
|
||||
// client_credentials grant completes server-side with no browser redirect.
|
||||
if (completed || !auth_url) {
|
||||
await loadStatus();
|
||||
setAuthorizing(false);
|
||||
return;
|
||||
}
|
||||
|
||||
const popup = window.open(auth_url, "mcp-oauth", "width=600,height=700,menubar=no,toolbar=no");
|
||||
popupRef.current = popup;
|
||||
} catch (err) {
|
||||
setError(err instanceof Error ? err.message : t("form.oauth.authFailed"));
|
||||
setAuthorizing(false);
|
||||
}
|
||||
};
|
||||
|
||||
const handleRevoke = async () => {
|
||||
setRevoking(true);
|
||||
setError("");
|
||||
try {
|
||||
await onRevoke(server.id);
|
||||
await loadStatus();
|
||||
} catch (err) {
|
||||
setError(err instanceof Error ? err.message : "Revoke failed");
|
||||
} finally {
|
||||
setRevoking(false);
|
||||
}
|
||||
};
|
||||
|
||||
const hasToken = status?.has_token ?? false;
|
||||
const isExpired = status?.expired ?? false;
|
||||
|
||||
return (
|
||||
<Dialog open={open} onOpenChange={(v) => {
|
||||
if (!v && authorizing) {
|
||||
popupRef.current?.close();
|
||||
popupRef.current = null;
|
||||
setAuthorizing(false);
|
||||
}
|
||||
if (!revoking) onOpenChange(v);
|
||||
}}>
|
||||
<DialogContent className="sm:max-w-sm">
|
||||
<DialogHeader>
|
||||
<DialogTitle className="flex items-center gap-2">
|
||||
<ShieldCheck className="h-4 w-4 text-primary" />
|
||||
{t("form.oauth.title")}
|
||||
</DialogTitle>
|
||||
</DialogHeader>
|
||||
|
||||
<div className="space-y-4 py-2">
|
||||
{/* Server info */}
|
||||
<div className="text-sm text-muted-foreground">
|
||||
<span className="font-medium text-foreground">{server.display_name || server.name}</span>
|
||||
{server.url && <span className="ml-1 font-mono text-xs truncate block">{server.url}</span>}
|
||||
</div>
|
||||
|
||||
{/* Auth mode — this dialog always manages the server-level (global) token */}
|
||||
<p className="text-xs text-muted-foreground">
|
||||
{t("form.oauth.modeGlobal")}
|
||||
</p>
|
||||
|
||||
{/* Token status */}
|
||||
<div className="rounded-md border p-3 space-y-1">
|
||||
{status === null ? (
|
||||
<div className="flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<Loader2 className="h-3.5 w-3.5 animate-spin" /> Loading...
|
||||
</div>
|
||||
) : hasToken ? (
|
||||
<div className="space-y-1">
|
||||
<div className="flex items-center gap-2 text-sm">
|
||||
{isExpired
|
||||
? <ShieldX className="h-3.5 w-3.5 text-amber-500" />
|
||||
: <ShieldCheck className="h-3.5 w-3.5 text-emerald-500" />}
|
||||
<span className={isExpired ? "text-amber-600 dark:text-amber-400" : "text-emerald-600 dark:text-emerald-400"}>
|
||||
{isExpired ? t("form.oauth.expired") : t("form.oauth.authorized")}
|
||||
</span>
|
||||
</div>
|
||||
{status.client_id && (
|
||||
<p className="text-xs text-muted-foreground font-mono">
|
||||
client: {status.client_id}
|
||||
</p>
|
||||
)}
|
||||
{status.expires_at && (
|
||||
<p className="text-xs text-muted-foreground">
|
||||
expires: {new Date(status.expires_at).toLocaleString()}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
) : (
|
||||
<div className="flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<ShieldX className="h-3.5 w-3.5" />
|
||||
{t("form.oauth.notAuthorized")}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{error && <p className="text-sm text-destructive">{error}</p>}
|
||||
</div>
|
||||
|
||||
<DialogFooter className="flex-row gap-2 sm:justify-between">
|
||||
{hasToken && !authorizing && (
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
onClick={handleRevoke}
|
||||
disabled={revoking}
|
||||
className="text-destructive hover:text-destructive"
|
||||
>
|
||||
{revoking ? <Loader2 className="h-3.5 w-3.5 animate-spin mr-1" /> : null}
|
||||
{t("form.oauth.revoke")}
|
||||
</Button>
|
||||
)}
|
||||
<div className="flex gap-2 ml-auto">
|
||||
<Button variant="outline" size="sm" onClick={() => onOpenChange(false)} disabled={revoking}>
|
||||
{t("form.cancel", "Close")}
|
||||
</Button>
|
||||
<Button
|
||||
size="sm"
|
||||
onClick={handleAuthorize}
|
||||
disabled={authorizing || revoking}
|
||||
className="gap-1"
|
||||
>
|
||||
{authorizing
|
||||
? <><Loader2 className="h-3.5 w-3.5 animate-spin" />{t("form.oauth.authorizing")}</>
|
||||
: <><ExternalLink className="h-3.5 w-3.5" />{hasToken ? t("form.oauth.reauthorize") : t("form.oauth.authorize")}</>
|
||||
}
|
||||
</Button>
|
||||
</div>
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
@@ -12,6 +12,7 @@ import { ConfirmDeleteDialog } from "@/components/shared/confirm-delete-dialog";
|
||||
import { useMCP, type MCPServerData, type MCPServerInput } from "./hooks/use-mcp";
|
||||
import { MCPFormDialog } from "./mcp-form-dialog";
|
||||
import { MCPToolsDialog } from "./mcp-tools-dialog";
|
||||
import { MCPOAuthDialog } from "./mcp-oauth-dialog";
|
||||
import { useMinLoading } from "@/hooks/use-min-loading";
|
||||
import { useDeferredLoading } from "@/hooks/use-deferred-loading";
|
||||
import { usePagination } from "@/hooks/use-pagination";
|
||||
@@ -32,7 +33,7 @@ const transportBadge: Record<string, string> = {
|
||||
export function MCPPage() {
|
||||
const { t } = useTranslation("mcp");
|
||||
const { t: tc } = useTranslation("common");
|
||||
const { servers, loading, fetching, refresh, createServer, updateServer, deleteServer, grantAgent, revokeAgent, listAgentGrants, testConnection, reconnectServer, listServerTools, getUserCredentials, setUserCredentials, deleteUserCredentials } = useMCP();
|
||||
const { servers, loading, fetching, refresh, createServer, updateServer, deleteServer, grantAgent, revokeAgent, listAgentGrants, testConnection, reconnectServer, listServerTools, getUserCredentials, setUserCredentials, deleteUserCredentials, startOAuth, getOAuthStatus, revokeOAuth } = useMCP();
|
||||
const spinning = useMinLoading(fetching);
|
||||
const showSkeleton = useDeferredLoading(loading && servers.length === 0);
|
||||
const [search, setSearch] = useState("");
|
||||
@@ -44,6 +45,7 @@ export function MCPPage() {
|
||||
const [deleteLoading, setDeleteLoading] = useState(false);
|
||||
const [credentialsServer, setCredentialsServer] = useState<MCPServerData | null>(null);
|
||||
const [reconnectingId, setReconnectingId] = useState<string | null>(null);
|
||||
const [oauthServer, setOauthServer] = useState<MCPServerData | null>(null);
|
||||
|
||||
const filtered = servers.filter(
|
||||
(s) =>
|
||||
@@ -56,7 +58,7 @@ export function MCPPage() {
|
||||
useEffect(() => { resetPage(); }, [search, resetPage]);
|
||||
|
||||
const handleCreate = async (data: MCPServerInput) => {
|
||||
await createServer(data);
|
||||
return createServer(data);
|
||||
};
|
||||
|
||||
const handleEdit = async (data: MCPServerInput) => {
|
||||
@@ -75,6 +77,10 @@ export function MCPPage() {
|
||||
}
|
||||
};
|
||||
|
||||
const handleAuthorize = (srv: MCPServerData) => {
|
||||
setOauthServer(srv);
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="p-4 sm:p-6 pb-10">
|
||||
<PageHeader
|
||||
@@ -238,8 +244,20 @@ export function MCPPage() {
|
||||
server={editServer}
|
||||
onSubmit={editServer ? handleEdit : handleCreate}
|
||||
onTest={testConnection}
|
||||
onAuthorize={handleAuthorize}
|
||||
/>
|
||||
|
||||
{oauthServer && (
|
||||
<MCPOAuthDialog
|
||||
open={!!oauthServer}
|
||||
onOpenChange={(open) => !open && setOauthServer(null)}
|
||||
server={oauthServer}
|
||||
onStartOAuth={startOAuth}
|
||||
onGetStatus={getOAuthStatus}
|
||||
onRevoke={revokeOAuth}
|
||||
/>
|
||||
)}
|
||||
|
||||
{grantsServer && (
|
||||
<Suspense fallback={null}>
|
||||
<MCPGrantsDialog
|
||||
@@ -283,6 +301,9 @@ export function MCPPage() {
|
||||
onGetCredentials={getUserCredentials}
|
||||
onSetCredentials={setUserCredentials}
|
||||
onDeleteCredentials={deleteUserCredentials}
|
||||
onStartOAuth={startOAuth}
|
||||
onGetOAuthStatus={getOAuthStatus}
|
||||
onRevokeOAuth={revokeOAuth}
|
||||
/>
|
||||
</Suspense>
|
||||
)}
|
||||
|
||||
@@ -4,6 +4,13 @@ import { Input } from "@/components/ui/input";
|
||||
import { Label } from "@/components/ui/label";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
import { Textarea } from "@/components/ui/textarea";
|
||||
import {
|
||||
Select,
|
||||
SelectContent,
|
||||
SelectItem,
|
||||
SelectTrigger,
|
||||
SelectValue,
|
||||
} from "@/components/ui/select";
|
||||
import { KeyValueEditor } from "@/components/shared/key-value-editor";
|
||||
import type { MCPFormData } from "@/schemas/mcp.schema";
|
||||
|
||||
@@ -13,12 +20,15 @@ export const isSensitiveEnv = (key: string) => SENSITIVE_ENV_RE.test(key.trim())
|
||||
|
||||
interface McpSettingsFieldsProps {
|
||||
form: UseFormReturn<MCPFormData>;
|
||||
/** Pass server id when editing — used to show "existing secret" placeholder. */
|
||||
isEditing?: boolean;
|
||||
}
|
||||
|
||||
/** Renders env vars, tool prefix, timeout, enabled, and requireUserCredentials fields. */
|
||||
export function McpSettingsFields({ form }: McpSettingsFieldsProps) {
|
||||
/** Renders env vars, OAuth config, tool prefix, timeout, enabled, requireUserCredentials fields. */
|
||||
export function McpSettingsFields({ form, isEditing }: McpSettingsFieldsProps) {
|
||||
const { t } = useTranslation("mcp");
|
||||
const { watch, setValue } = form;
|
||||
const transport = watch("transport");
|
||||
const env = watch("env") as Record<string, string>;
|
||||
const toolPrefix = watch("toolPrefix");
|
||||
const timeout = watch("timeout");
|
||||
@@ -27,6 +37,11 @@ export function McpSettingsFields({ form }: McpSettingsFieldsProps) {
|
||||
const requireUserCreds = watch("requireUserCreds");
|
||||
const toolHintsGlobal = watch("toolHintsGlobal");
|
||||
const toolHintsTools = watch("toolHintsTools") as Record<string, string>;
|
||||
const oauthEnabled = watch("oauthEnabled");
|
||||
const oauthUseDcr = watch("oauthUseDcr");
|
||||
const oauthGrantType = watch("oauthGrantType");
|
||||
|
||||
const isStdio = transport === "stdio";
|
||||
|
||||
return (
|
||||
<>
|
||||
@@ -42,9 +57,135 @@ export function McpSettingsFields({ form }: McpSettingsFieldsProps) {
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* Admin-authored hints appended to MCP tool descriptions. Lets ops teach
|
||||
the LLM about server-specific quirks (e.g. "no trailing ';' in code args")
|
||||
without modifying the upstream MCP server. See Settings.tool_hints JSONB. */}
|
||||
{/* OAuth 2.1 configuration — only for network transports (SSE / Streamable HTTP) */}
|
||||
{!isStdio && (
|
||||
<div className="rounded-md border border-dashed border-border p-3 space-y-3">
|
||||
<div className="space-y-1">
|
||||
<div className="flex items-center gap-2">
|
||||
<Switch
|
||||
id="mcp-oauth-enable"
|
||||
checked={oauthEnabled}
|
||||
onCheckedChange={(v) => setValue("oauthEnabled", v)}
|
||||
/>
|
||||
<Label htmlFor="mcp-oauth-enable" className="font-medium">
|
||||
{t("form.oauth.title")}
|
||||
</Label>
|
||||
</div>
|
||||
{!oauthEnabled && (
|
||||
<p className="text-xs text-muted-foreground pl-9">{t("form.oauth.titleHint")}</p>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{oauthEnabled && (
|
||||
<div className="pl-0 space-y-3">
|
||||
{/* DCR toggle */}
|
||||
<div className="space-y-1">
|
||||
<div className="flex items-center gap-2">
|
||||
<Switch
|
||||
id="mcp-oauth-dcr"
|
||||
checked={oauthUseDcr}
|
||||
onCheckedChange={(v) => setValue("oauthUseDcr", v)}
|
||||
/>
|
||||
<Label htmlFor="mcp-oauth-dcr">{t("form.oauth.useDcr")}</Label>
|
||||
</div>
|
||||
<p className="text-xs text-muted-foreground pl-9">{t("form.oauth.useDcrHint")}</p>
|
||||
</div>
|
||||
|
||||
{/* Manual config — when DCR disabled */}
|
||||
{!oauthUseDcr && (
|
||||
<>
|
||||
{/* Grant type */}
|
||||
<div className="grid gap-1.5">
|
||||
<Label>{t("form.oauth.grantType")}</Label>
|
||||
<Select
|
||||
value={oauthGrantType}
|
||||
onValueChange={(v) => setValue("oauthGrantType", v as MCPFormData["oauthGrantType"])}
|
||||
>
|
||||
<SelectTrigger className="text-base md:text-sm">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="pkce">{t("form.oauth.pkce")}</SelectItem>
|
||||
<SelectItem value="authorization_code">{t("form.oauth.authCode")}</SelectItem>
|
||||
<SelectItem value="client_credentials">{t("form.oauth.clientCreds")}</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
|
||||
{/* Authorization URL — for auth code flows (pkce and authorization_code) */}
|
||||
{oauthGrantType !== "client_credentials" && (
|
||||
<div className="grid gap-1.5">
|
||||
<Label>{t("form.oauth.authEndpoint")}</Label>
|
||||
<Input
|
||||
value={watch("oauthAuthEndpoint")}
|
||||
onChange={(e) => setValue("oauthAuthEndpoint", e.target.value)}
|
||||
placeholder={t("form.oauth.authEndpointPlaceholder")}
|
||||
className="text-base md:text-sm"
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Token URL */}
|
||||
<div className="grid gap-1.5">
|
||||
<Label>{t("form.oauth.tokenEndpoint")}</Label>
|
||||
<Input
|
||||
value={watch("oauthTokenEndpoint")}
|
||||
onChange={(e) => setValue("oauthTokenEndpoint", e.target.value)}
|
||||
placeholder={t("form.oauth.tokenEndpointPlaceholder")}
|
||||
className="text-base md:text-sm"
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* Client ID */}
|
||||
<div className="grid gap-1.5">
|
||||
<Label>{t("form.oauth.clientId")}</Label>
|
||||
<Input
|
||||
autoComplete="off"
|
||||
value={watch("oauthClientId")}
|
||||
onChange={(e) => setValue("oauthClientId", e.target.value)}
|
||||
placeholder="client_id"
|
||||
className="text-base md:text-sm font-mono"
|
||||
/>
|
||||
</div>
|
||||
|
||||
{/* Client Secret — not needed for PKCE (public client) */}
|
||||
{oauthGrantType !== "pkce" && (
|
||||
<div className="grid gap-1.5">
|
||||
<Label>{t("form.oauth.clientSecret")}</Label>
|
||||
<Input
|
||||
type="password"
|
||||
autoComplete="new-password"
|
||||
value={watch("oauthClientSecret")}
|
||||
onChange={(e) => setValue("oauthClientSecret", e.target.value)}
|
||||
placeholder={isEditing ? t("form.oauth.clientSecretPlaceholder") : "client_secret"}
|
||||
className="text-base md:text-sm"
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Scopes */}
|
||||
<div className="grid gap-1.5">
|
||||
<Label>{t("form.oauth.scope")}</Label>
|
||||
<Input
|
||||
value={watch("oauthScope")}
|
||||
onChange={(e) => setValue("oauthScope", e.target.value)}
|
||||
placeholder={t("form.oauth.scopePlaceholder")}
|
||||
className="text-base md:text-sm"
|
||||
/>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
|
||||
{/* Auth mode indicator */}
|
||||
<p className="text-xs text-muted-foreground">
|
||||
{requireUserCreds ? t("form.oauth.modePerUser") : t("form.oauth.modeGlobal")}
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Admin-authored hints appended to MCP tool descriptions. */}
|
||||
<div className="grid gap-3 rounded-md border border-dashed border-border p-3">
|
||||
<div className="grid gap-1">
|
||||
<Label className="text-sm font-medium">{t("form.toolHints")}</Label>
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
import { useState, useEffect } from "react";
|
||||
import { useState, useEffect, useRef, useCallback } from "react";
|
||||
import { useTranslation } from "react-i18next";
|
||||
import { useForm } from "react-hook-form";
|
||||
import { zodResolver } from "@hookform/resolvers/zod";
|
||||
import { KeyRound, Loader2 } from "lucide-react";
|
||||
import { KeyRound, Loader2, ShieldCheck, ShieldX, ExternalLink } from "lucide-react";
|
||||
import { useWsEvent } from "@/hooks/use-ws-event";
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
@@ -21,7 +22,7 @@ import { toast } from "@/stores/use-toast-store";
|
||||
import { useAuthStore } from "@/stores/use-auth-store";
|
||||
import { useTenants } from "@/hooks/use-tenants";
|
||||
import i18next from "i18next";
|
||||
import type { MCPServerData, MCPUserCredentialStatus, MCPUserCredentialInput } from "./hooks/use-mcp";
|
||||
import type { MCPServerData, MCPUserCredentialStatus, MCPUserCredentialInput, MCPOAuthStatus } from "./hooks/use-mcp";
|
||||
import { mcpUserCredentialsSchema, type MCPUserCredentialsFormData } from "@/schemas/mcp-credentials.schema";
|
||||
|
||||
/** Header keys whose values should be masked. */
|
||||
@@ -32,6 +33,10 @@ const isSensitiveHeader = (key: string) => SENSITIVE_HEADER_RE.test(key.trim());
|
||||
const SENSITIVE_ENV_RE = /^.*(key|secret|token|password|credential).*$/i;
|
||||
const isSensitiveEnv = (key: string) => SENSITIVE_ENV_RE.test(key.trim());
|
||||
|
||||
function serverHasOAuth(server: MCPServerData): boolean {
|
||||
return server.settings?.oauth?.auth_type === "oauth";
|
||||
}
|
||||
|
||||
interface MCPUserCredentialsDialogProps {
|
||||
open: boolean;
|
||||
onOpenChange: (open: boolean) => void;
|
||||
@@ -39,6 +44,10 @@ interface MCPUserCredentialsDialogProps {
|
||||
onGetCredentials: (serverId: string, userId?: string) => Promise<MCPUserCredentialStatus>;
|
||||
onSetCredentials: (serverId: string, creds: MCPUserCredentialInput, userId?: string) => Promise<void>;
|
||||
onDeleteCredentials: (serverId: string, userId?: string) => Promise<void>;
|
||||
/** OAuth functions — passed when server has OAuth + require_user_credentials */
|
||||
onStartOAuth?: (serverId: string, mcpUrl: string, userId?: string) => Promise<{ auth_url: string; state: string; client_id: string; issuer: string; completed?: boolean }>;
|
||||
onGetOAuthStatus?: (serverId: string, userId?: string) => Promise<MCPOAuthStatus>;
|
||||
onRevokeOAuth?: (serverId: string, userId?: string) => Promise<void>;
|
||||
}
|
||||
|
||||
export function MCPUserCredentialsDialog({
|
||||
@@ -48,6 +57,9 @@ export function MCPUserCredentialsDialog({
|
||||
onGetCredentials,
|
||||
onSetCredentials,
|
||||
onDeleteCredentials,
|
||||
onStartOAuth,
|
||||
onGetOAuthStatus,
|
||||
onRevokeOAuth,
|
||||
}: MCPUserCredentialsDialogProps) {
|
||||
const { t } = useTranslation("mcp");
|
||||
const role = useAuthStore((s) => s.role);
|
||||
@@ -68,6 +80,14 @@ export function MCPUserCredentialsDialog({
|
||||
const [deleting, setDeleting] = useState(false);
|
||||
const [initialLoad, setInitialLoad] = useState(true);
|
||||
|
||||
// OAuth state
|
||||
const hasOAuth = serverHasOAuth(server);
|
||||
const [oauthStatus, setOauthStatus] = useState<MCPOAuthStatus | null>(null);
|
||||
const [oauthAuthorizing, setOauthAuthorizing] = useState(false);
|
||||
const [oauthRevoking, setOauthRevoking] = useState(false);
|
||||
const [oauthError, setOauthError] = useState("");
|
||||
const popupRef = useRef<Window | null>(null);
|
||||
|
||||
const form = useForm<MCPUserCredentialsFormData>({
|
||||
resolver: zodResolver(mcpUserCredentialsSchema),
|
||||
mode: "onChange",
|
||||
@@ -84,9 +104,11 @@ export function MCPUserCredentialsDialog({
|
||||
setSelectedUserId(currentUserId);
|
||||
setUserSearchText("");
|
||||
setInitialLoad(true);
|
||||
setOauthError("");
|
||||
}
|
||||
}, [open, currentUserId]);
|
||||
|
||||
// Load static credentials
|
||||
useEffect(() => {
|
||||
if (!open) return;
|
||||
reset({ apiKey: "", headers: {}, env: {} });
|
||||
@@ -99,7 +121,87 @@ export function MCPUserCredentialsDialog({
|
||||
.then(setStatus)
|
||||
.catch((err) => console.error("[MCPUserCredentials] load credentials failed:", err))
|
||||
.finally(() => { setLoadingStatus(false); setInitialLoad(false); });
|
||||
}, [open, server.id, onGetCredentials, canManageUsers, selectedUserId]);
|
||||
}, [open, server.id, onGetCredentials, canManageUsers, selectedUserId]);
|
||||
|
||||
// Load OAuth status when server has OAuth
|
||||
const loadOAuthStatus = useCallback(async () => {
|
||||
if (!hasOAuth || !onGetOAuthStatus) return;
|
||||
try {
|
||||
const s = await onGetOAuthStatus(server.id, selectedUserId || undefined);
|
||||
setOauthStatus(s);
|
||||
} catch {
|
||||
setOauthStatus(null);
|
||||
}
|
||||
}, [hasOAuth, onGetOAuthStatus, server.id, selectedUserId]);
|
||||
|
||||
useEffect(() => {
|
||||
if (open && hasOAuth) {
|
||||
setOauthStatus(null);
|
||||
loadOAuthStatus();
|
||||
}
|
||||
}, [open, hasOAuth, loadOAuthStatus]);
|
||||
|
||||
// Close popup on unmount.
|
||||
useEffect(() => {
|
||||
return () => { popupRef.current?.close(); };
|
||||
}, []);
|
||||
|
||||
// Primary notification path: WebSocket event from backend after token exchange.
|
||||
const oauthUserId = canManageUsers ? selectedUserId : currentUserId;
|
||||
useWsEvent("mcp.oauth_complete", (payload: unknown) => {
|
||||
const p = payload as { serverId?: string; userId?: string; status?: string; error?: string };
|
||||
if (p.serverId !== server.id) return;
|
||||
if ((p.userId ?? "") !== (oauthUserId ?? "")) return;
|
||||
popupRef.current?.close();
|
||||
popupRef.current = null;
|
||||
setOauthAuthorizing(false);
|
||||
if (p.status === "success") {
|
||||
loadOAuthStatus();
|
||||
} else {
|
||||
setOauthError(p.error ?? t("form.oauth.authFailed"));
|
||||
}
|
||||
});
|
||||
|
||||
const handleOAuthAuthorize = async () => {
|
||||
if (!onStartOAuth || !server.url) {
|
||||
setOauthError("Server has no URL configured");
|
||||
return;
|
||||
}
|
||||
setOauthAuthorizing(true);
|
||||
setOauthError("");
|
||||
try {
|
||||
const userId = canManageUsers ? selectedUserId : currentUserId;
|
||||
const { auth_url, completed } = await onStartOAuth(server.id, server.url, userId || undefined);
|
||||
|
||||
// client_credentials grant completes server-side with no browser redirect.
|
||||
if (completed || !auth_url) {
|
||||
await loadOAuthStatus();
|
||||
setOauthAuthorizing(false);
|
||||
return;
|
||||
}
|
||||
|
||||
const popup = window.open(auth_url, "mcp-oauth", "width=600,height=700,menubar=no,toolbar=no");
|
||||
popupRef.current = popup;
|
||||
} catch (err) {
|
||||
setOauthError(err instanceof Error ? err.message : t("form.oauth.authFailed"));
|
||||
setOauthAuthorizing(false);
|
||||
}
|
||||
};
|
||||
|
||||
const handleOAuthRevoke = async () => {
|
||||
if (!onRevokeOAuth) return;
|
||||
setOauthRevoking(true);
|
||||
setOauthError("");
|
||||
try {
|
||||
const userId = canManageUsers ? selectedUserId : currentUserId;
|
||||
await onRevokeOAuth(server.id, userId || undefined);
|
||||
await loadOAuthStatus();
|
||||
} catch (err) {
|
||||
setOauthError(err instanceof Error ? err.message : "Revoke failed");
|
||||
} finally {
|
||||
setOauthRevoking(false);
|
||||
}
|
||||
};
|
||||
|
||||
const handleSave = async () => {
|
||||
setSaving(true);
|
||||
@@ -134,19 +236,36 @@ export function MCPUserCredentialsDialog({
|
||||
}
|
||||
};
|
||||
|
||||
const oauthHasToken = oauthStatus?.has_token ?? false;
|
||||
const oauthIsExpired = oauthStatus?.expired ?? false;
|
||||
const oauthBusy = oauthAuthorizing || oauthRevoking;
|
||||
|
||||
return (
|
||||
<Dialog open={open} onOpenChange={onOpenChange}>
|
||||
<DialogContent className="sm:max-w-lg">
|
||||
<Dialog open={open} onOpenChange={(v) => {
|
||||
if (!v && oauthAuthorizing) {
|
||||
popupRef.current?.close();
|
||||
popupRef.current = null;
|
||||
setOauthAuthorizing(false);
|
||||
}
|
||||
onOpenChange(v);
|
||||
}}>
|
||||
<DialogContent className="sm:max-w-lg" aria-describedby="ucred-desc">
|
||||
<DialogHeader>
|
||||
<DialogTitle className="flex items-center gap-2">
|
||||
<KeyRound className="h-4 w-4" />
|
||||
{canManageUsers ? t("userCredentials.titleAdmin") : t("userCredentials.title")}
|
||||
</DialogTitle>
|
||||
<DialogDescription>
|
||||
<DialogDescription id="ucred-desc">
|
||||
{canManageUsers ? t("userCredentials.descriptionAdmin") : t("userCredentials.description")}
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
|
||||
{/* Hidden dummy fields — prevent browser from auto-filling the real ones */}
|
||||
<div style={{ display: "none" }}>
|
||||
<input type="text" autoComplete="username" tabIndex={-1} aria-hidden />
|
||||
<input type="password" autoComplete="current-password" tabIndex={-1} aria-hidden />
|
||||
</div>
|
||||
|
||||
{loadingStatus && initialLoad ? (
|
||||
<div className="flex justify-center py-8">
|
||||
<Loader2 className="h-5 w-5 animate-spin text-muted-foreground" />
|
||||
@@ -171,7 +290,77 @@ export function MCPUserCredentialsDialog({
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Current status badges */}
|
||||
{/* OAuth section — only for servers with OAuth + require_user_credentials */}
|
||||
{hasOAuth && onStartOAuth && (
|
||||
<div className="rounded-md border p-3 space-y-2">
|
||||
<div className="flex items-center gap-2 text-sm font-medium">
|
||||
<ShieldCheck className="h-4 w-4 text-primary" />
|
||||
{t("form.oauth.title")}
|
||||
</div>
|
||||
|
||||
{/* Token status */}
|
||||
<div className="rounded-md bg-muted/40 px-3 py-2 space-y-1">
|
||||
{oauthStatus === null ? (
|
||||
<div className="flex items-center gap-2 text-xs text-muted-foreground">
|
||||
<Loader2 className="h-3 w-3 animate-spin" /> Loading...
|
||||
</div>
|
||||
) : oauthHasToken ? (
|
||||
<div className="space-y-0.5">
|
||||
<div className="flex items-center gap-2 text-xs">
|
||||
{oauthIsExpired
|
||||
? <ShieldX className="h-3.5 w-3.5 text-amber-500" />
|
||||
: <ShieldCheck className="h-3.5 w-3.5 text-emerald-500" />}
|
||||
<span className={oauthIsExpired ? "text-amber-600 dark:text-amber-400" : "text-emerald-600 dark:text-emerald-400"}>
|
||||
{oauthIsExpired ? t("form.oauth.expired") : t("form.oauth.authorized")}
|
||||
</span>
|
||||
</div>
|
||||
{oauthStatus?.expires_at && (
|
||||
<p className="text-xs text-muted-foreground">
|
||||
expires: {new Date(oauthStatus.expires_at).toLocaleString()}
|
||||
</p>
|
||||
)}
|
||||
</div>
|
||||
) : (
|
||||
<div className="flex items-center gap-2 text-xs text-muted-foreground">
|
||||
<ShieldX className="h-3.5 w-3.5" />
|
||||
{t("form.oauth.notAuthorized")}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{oauthError && <p className="text-xs text-destructive">{oauthError}</p>}
|
||||
|
||||
{/* Authorize / Revoke buttons */}
|
||||
<div className="flex gap-2">
|
||||
{oauthHasToken && (
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
onClick={handleOAuthRevoke}
|
||||
disabled={oauthBusy}
|
||||
className="text-destructive hover:text-destructive text-xs"
|
||||
>
|
||||
{oauthRevoking ? <Loader2 className="h-3 w-3 animate-spin mr-1" /> : null}
|
||||
{t("form.oauth.revoke")}
|
||||
</Button>
|
||||
)}
|
||||
<Button
|
||||
size="sm"
|
||||
variant={oauthHasToken ? "outline" : "default"}
|
||||
onClick={handleOAuthAuthorize}
|
||||
disabled={oauthBusy}
|
||||
className="gap-1 text-xs"
|
||||
>
|
||||
{oauthAuthorizing
|
||||
? <><Loader2 className="h-3 w-3 animate-spin" />{t("form.oauth.authorizing")}</>
|
||||
: <><ExternalLink className="h-3 w-3" />{oauthHasToken ? t("form.oauth.reauthorize") : t("form.oauth.authorize")}</>
|
||||
}
|
||||
</Button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Current static credential status badges */}
|
||||
{status && (
|
||||
<div className="flex flex-wrap gap-2">
|
||||
{!status.has_credentials ? (
|
||||
@@ -198,6 +387,7 @@ export function MCPUserCredentialsDialog({
|
||||
<Input
|
||||
id="uc-api-key"
|
||||
type="password"
|
||||
autoComplete="new-password"
|
||||
placeholder={t("userCredentials.apiKeyPlaceholder")}
|
||||
className="text-base md:text-sm font-mono"
|
||||
{...register("apiKey")}
|
||||
|
||||
@@ -17,6 +17,15 @@ export const mcpFormSchema = z.object({
|
||||
// LLM sees server-specific quirks. Persisted under settings.tool_hints.
|
||||
toolHintsGlobal: z.string(),
|
||||
toolHintsTools: z.record(z.string(), z.string()),
|
||||
// OAuth 2.1 client configuration (non-stdio transports only)
|
||||
oauthEnabled: z.boolean(),
|
||||
oauthUseDcr: z.boolean(),
|
||||
oauthGrantType: z.enum(["pkce", "authorization_code", "client_credentials"]),
|
||||
oauthAuthEndpoint: z.string(),
|
||||
oauthTokenEndpoint: z.string(),
|
||||
oauthClientId: z.string(),
|
||||
oauthClientSecret: z.string(),
|
||||
oauthScope: z.string(),
|
||||
});
|
||||
|
||||
export type MCPFormData = z.infer<typeof mcpFormSchema>;
|
||||
+31
-14
@@ -1,3 +1,32 @@
|
||||
export interface MCPOAuthSettings {
|
||||
auth_type?: "oauth" | "";
|
||||
use_dcr?: boolean;
|
||||
grant_type?: "pkce" | "authorization_code" | "client_credentials";
|
||||
auth_endpoint?: string;
|
||||
token_endpoint?: string;
|
||||
client_id?: string;
|
||||
/** Write-only — never returned by the API (stored encrypted). */
|
||||
client_secret?: string;
|
||||
scope?: string;
|
||||
}
|
||||
|
||||
export interface MCPServerSettings {
|
||||
require_user_credentials?: boolean;
|
||||
tool_hints?: {
|
||||
global?: string;
|
||||
tools?: Record<string, string>;
|
||||
};
|
||||
oauth?: MCPOAuthSettings;
|
||||
}
|
||||
|
||||
export interface MCPOAuthStatus {
|
||||
has_token: boolean;
|
||||
client_id?: string;
|
||||
issuer?: string;
|
||||
expires_at?: string;
|
||||
expired?: boolean;
|
||||
}
|
||||
|
||||
export interface MCPServerData {
|
||||
id: string;
|
||||
name: string;
|
||||
@@ -10,13 +39,7 @@ export interface MCPServerData {
|
||||
env: Record<string, string> | null;
|
||||
tool_prefix: string;
|
||||
timeout_sec: number;
|
||||
settings?: {
|
||||
require_user_credentials?: boolean;
|
||||
tool_hints?: {
|
||||
global?: string;
|
||||
tools?: Record<string, string>;
|
||||
};
|
||||
};
|
||||
settings?: MCPServerSettings;
|
||||
enabled: boolean;
|
||||
created_by: string;
|
||||
agent_count?: number;
|
||||
@@ -35,13 +58,7 @@ export interface MCPServerInput {
|
||||
env?: Record<string, string>;
|
||||
tool_prefix?: string;
|
||||
timeout_sec?: number;
|
||||
settings?: {
|
||||
require_user_credentials?: boolean;
|
||||
tool_hints?: {
|
||||
global?: string;
|
||||
tools?: Record<string, string>;
|
||||
};
|
||||
};
|
||||
settings?: MCPServerSettings;
|
||||
enabled?: boolean;
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user