From 8ae82573abbcc6d591d91fdb0ff5d57b2104f55b Mon Sep 17 00:00:00 2001 From: wizardchen Date: Thu, 23 Jul 2026 15:33:20 +0800 Subject: [PATCH] fix(mcp): coordinate OAuth token refresh and authorization polling Prevent concurrent refresh-token rotation across instances, distinguish expired-but-refreshable tokens from active authorization, and bind OAuth popup polling to a per-attempt ID so stale tokens cannot satisfy new consent. --- docs/docs.go | 43 ++- docs/swagger.json | 43 ++- docs/swagger.yaml | 32 +- frontend/src/api/embed/index.ts | 14 +- frontend/src/api/mcp-service.ts | 48 ++- frontend/src/i18n/locales/en-US.ts | 1 + frontend/src/i18n/locales/zh-CN.ts | 1 + .../components/McpOAuthCard.oauth.test.mjs | 39 +++ .../views/chat/components/McpOAuthCard.vue | 9 +- .../settings/components/McpServiceDialog.vue | 27 +- internal/agent/tools/mcp_oauth.go | 5 + internal/application/repository/mcp_oauth.go | 48 +++ .../application/repository/mcp_oauth_test.go | 43 +++ internal/handler/mcp_oauth.go | 48 ++- internal/mcp/client.go | 61 +++- internal/mcp/oauth_lifecycle.go | 280 ++++++++++++++++ internal/mcp/oauth_lifecycle_test.go | 316 ++++++++++++++++++ internal/mcp/oauth_manager.go | 98 ++++-- internal/mcp/oauth_principal_test.go | 48 +++ internal/mcp/oauth_state.go | 116 ++++++- internal/mcp/oauth_state_test.go | 71 ++++ internal/mcp/oauth_tokenstore.go | 37 +- internal/middleware/logger.go | 32 +- internal/middleware/logger_test.go | 13 + internal/types/interfaces/mcp_oauth.go | 20 ++ internal/types/mcp_oauth.go | 9 +- migrations/sqlite/000000_init.up.sql | 6 + .../sqlite/000001_kb_activity_scope.down.sql | 4 - .../sqlite/000001_kb_activity_scope.up.sql | 5 - .../000074_mcp_oauth_refresh_lease.down.sql | 3 + .../000074_mcp_oauth_refresh_lease.up.sql | 7 + 31 files changed, 1418 insertions(+), 109 deletions(-) create mode 100644 frontend/src/views/chat/components/McpOAuthCard.oauth.test.mjs create mode 100644 internal/mcp/oauth_lifecycle.go create mode 100644 internal/mcp/oauth_lifecycle_test.go create mode 100644 internal/mcp/oauth_state_test.go delete mode 100644 migrations/sqlite/000001_kb_activity_scope.down.sql delete mode 100644 migrations/sqlite/000001_kb_activity_scope.up.sql create mode 100644 migrations/versioned/000074_mcp_oauth_refresh_lease.down.sql create mode 100644 migrations/versioned/000074_mcp_oauth_refresh_lease.up.sql diff --git a/docs/docs.go b/docs/docs.go index 3ae135a6..ca5d87bf 100644 --- a/docs/docs.go +++ b/docs/docs.go @@ -4119,10 +4119,16 @@ const docTemplate = `{ }, { "type": "integer", - "description": "标签ID筛选(seq_id)", + "description": "标签ID筛选(seq_id)，兼容旧版单标签", "name": "tag_id", "in": "query" }, + { + "type": "string", + "description": "标签UUID筛选，逗号分隔（OR语义）", + "name": "tag_ids", + "in": "query" + }, { "type": "string", "description": "关键词搜索", @@ -4285,12 +4291,13 @@ const docTemplate = `{ "ApiKeyAuth": [] } ], - "description": "将所有FAQ条目导出为CSV文件", + "description": "将所有FAQ条目导出为 CSV（默认）或 JSON。?format=json 返回与 FAQEntryPayload 结构兼容的数组。", "consumes": [ "application/json" ], "produces": [ - "text/csv" + "text/csv", + "application/json" ], "tags": [ "FAQ管理" @@ -4303,11 +4310,17 @@ const docTemplate = `{ "name": "id", "in": "path", "required": true + }, + { + "type": "string", + "description": "导出格式：csv（默认）或 json", + "name": "format", + "in": "query" } ], "responses": { "200": { - "description": "CSV文件", + "description": "导出文件", "schema": { "type": "file" } @@ -8147,7 +8160,7 @@ const docTemplate = `{ ], "responses": { "200": { - "description": "{authorization_url: string}", + "description": "{authorization_url: string, authorization_attempt: string}", "schema": { "type": "object", "additionalProperties": true @@ -8169,7 +8182,7 @@ const docTemplate = `{ "Bearer": [] } ], - "description": "返回当前用户对指定 MCP 服务是否已完成 OAuth 授权", + "description": "返回当前用户的 OAuth Token 生命周期状态；传 authorization_attempt 时只检查本次授权流程", "produces": [ "application/json" ], @@ -8184,11 +8197,17 @@ const docTemplate = `{ "name": "id", "in": "path", "required": true + }, + { + "type": "string", + "description": "本次授权尝试 ID；传入后不会接受历史 Token", + "name": "authorization_attempt", + "in": "query" } ], "responses": { "200": { - "description": "{authorized: bool}", + "description": "{authorized: bool, state: string, refresh_available: bool, expires_at?: string}", "schema": { "type": "object", "additionalProperties": true @@ -14950,7 +14969,10 @@ const docTemplate = `{ "kb.share_added", "kb.share_permission_changed", "kb.share_removed", - "wiki.content_changed" + "wiki.content_changed", + "faq.import_started", + "faq.import_completed", + "faq.import_failed" ], "x-enum-varnames": [ "AuditActionMemberAdded", @@ -15009,7 +15031,10 @@ const docTemplate = `{ "AuditActionKBShareAdded", "AuditActionKBSharePermissionChanged", "AuditActionKBShareRemoved", - "AuditActionWikiContentChanged" + "AuditActionWikiContentChanged", + "AuditActionFAQImportStarted", + "AuditActionFAQImportCompleted", + "AuditActionFAQImportFailed" ] }, "github_com_Tencent_WeKnora_internal_types.AuditLog": { diff --git a/docs/swagger.json b/docs/swagger.json index bb2457d8..099e1239 100644 --- a/docs/swagger.json +++ b/docs/swagger.json @@ -4112,10 +4112,16 @@ }, { "type": "integer", - "description": "标签ID筛选(seq_id)", + "description": "标签ID筛选(seq_id)，兼容旧版单标签", "name": "tag_id", "in": "query" }, + { + "type": "string", + "description": "标签UUID筛选，逗号分隔（OR语义）", + "name": "tag_ids", + "in": "query" + }, { "type": "string", "description": "关键词搜索", @@ -4278,12 +4284,13 @@ "ApiKeyAuth": [] } ], - "description": "将所有FAQ条目导出为CSV文件", + "description": "将所有FAQ条目导出为 CSV（默认）或 JSON。?format=json 返回与 FAQEntryPayload 结构兼容的数组。", "consumes": [ "application/json" ], "produces": [ - "text/csv" + "text/csv", + "application/json" ], "tags": [ "FAQ管理" @@ -4296,11 +4303,17 @@ "name": "id", "in": "path", "required": true + }, + { + "type": "string", + "description": "导出格式：csv（默认）或 json", + "name": "format", + "in": "query" } ], "responses": { "200": { - "description": "CSV文件", + "description": "导出文件", "schema": { "type": "file" } @@ -8140,7 +8153,7 @@ ], "responses": { "200": { - "description": "{authorization_url: string}", + "description": "{authorization_url: string, authorization_attempt: string}", "schema": { "type": "object", "additionalProperties": true @@ -8162,7 +8175,7 @@ "Bearer": [] } ], - "description": "返回当前用户对指定 MCP 服务是否已完成 OAuth 授权", + "description": "返回当前用户的 OAuth Token 生命周期状态；传 authorization_attempt 时只检查本次授权流程", "produces": [ "application/json" ], @@ -8177,11 +8190,17 @@ "name": "id", "in": "path", "required": true + }, + { + "type": "string", + "description": "本次授权尝试 ID；传入后不会接受历史 Token", + "name": "authorization_attempt", + "in": "query" } ], "responses": { "200": { - "description": "{authorized: bool}", + "description": "{authorized: bool, state: string, refresh_available: bool, expires_at?: string}", "schema": { "type": "object", "additionalProperties": true @@ -14943,7 +14962,10 @@ "kb.share_added", "kb.share_permission_changed", "kb.share_removed", - "wiki.content_changed" + "wiki.content_changed", + "faq.import_started", + "faq.import_completed", + "faq.import_failed" ], "x-enum-varnames": [ "AuditActionMemberAdded", @@ -15002,7 +15024,10 @@ "AuditActionKBShareAdded", "AuditActionKBSharePermissionChanged", "AuditActionKBShareRemoved", - "AuditActionWikiContentChanged" + "AuditActionWikiContentChanged", + "AuditActionFAQImportStarted", + "AuditActionFAQImportCompleted", + "AuditActionFAQImportFailed" ] }, "github_com_Tencent_WeKnora_internal_types.AuditLog": { diff --git a/docs/swagger.yaml b/docs/swagger.yaml index 76cbb8a1..072b233e 100644 --- a/docs/swagger.yaml +++ b/docs/swagger.yaml @@ -331,6 +331,9 @@ definitions: - kb.share_permission_changed - kb.share_removed - wiki.content_changed + - faq.import_started + - faq.import_completed + - faq.import_failed type: string x-enum-varnames: - AuditActionMemberAdded @@ -390,6 +393,9 @@ definitions: - AuditActionKBSharePermissionChanged - AuditActionKBShareRemoved - AuditActionWikiContentChanged + - AuditActionFAQImportStarted + - AuditActionFAQImportCompleted + - AuditActionFAQImportFailed github_com_Tencent_WeKnora_internal_types.AuditLog: properties: action: @@ -8328,10 +8334,14 @@ paths: in: query name: page_size type: integer - - description: 标签ID筛选(seq_id) + - description: 标签ID筛选(seq_id)，兼容旧版单标签 in: query name: tag_id type: integer + - description: 标签UUID筛选，逗号分隔（OR语义） + in: query + name: tag_ids + type: string - description: 关键词搜索 in: query name: keyword @@ -8525,18 +8535,23 @@ paths: get: consumes: - application/json - description: 将所有FAQ条目导出为CSV文件 + description: 将所有FAQ条目导出为 CSV（默认）或 JSON。?format=json 返回与 FAQEntryPayload 结构兼容的数组。 parameters: - description: 知识库ID in: path name: id required: true type: string + - description: 导出格式：csv（默认）或 json + in: query + name: format + type: string produces: - text/csv + - application/json responses: "200": - description: CSV文件 + description: 导出文件 schema: type: file "400": @@ -10900,7 +10915,7 @@ paths: - application/json responses: "200": - description: '{authorization_url: string}' + description: '{authorization_url: string, authorization_attempt: string}' schema: additionalProperties: true type: object @@ -10915,18 +10930,23 @@ paths: - MCP服务 /mcp-services/{id}/oauth/status: get: - description: 返回当前用户对指定 MCP 服务是否已完成 OAuth 授权 + description: 返回当前用户的 OAuth Token 生命周期状态；传 authorization_attempt 时只检查本次授权流程 parameters: - description: MCP 服务 ID in: path name: id required: true type: string + - description: 本次授权尝试 ID；传入后不会接受历史 Token + in: query + name: authorization_attempt + type: string produces: - application/json responses: "200": - description: '{authorized: bool}' + description: '{authorized: bool, state: string, refresh_available: bool, + expires_at?: string}' schema: additionalProperties: true type: object diff --git a/frontend/src/api/embed/index.ts b/frontend/src/api/embed/index.ts index 2c6f72e6..95658995 100644 --- a/frontend/src/api/embed/index.ts +++ b/frontend/src/api/embed/index.ts @@ -365,13 +365,17 @@ export async function getEmbedMCPOAuthAuthorizeURL( visitorId: string, serviceId: string, body: { redirect_uri: string; frontend_redirect?: string }, -): Promise { +): Promise<{ authorizationUrl: string; authorizationAttempt: string }> { const response: any = await post( `/api/v1/embed/${channelId}/sessions/${encodeURIComponent(sessionId)}/mcp-services/${encodeURIComponent(serviceId)}/oauth/authorize-url`, body, { headers: embedSessionHeaders(token, sessionSig, visitorId) }, ) - return (response.data ?? response)?.authorization_url ?? '' + const data = response.data ?? response + return { + authorizationUrl: data?.authorization_url ?? '', + authorizationAttempt: data?.authorization_attempt ?? '', + } } export async function getEmbedMCPOAuthStatus( @@ -381,9 +385,13 @@ export async function getEmbedMCPOAuthStatus( sessionSig: string, visitorId: string, serviceId: string, + authorizationAttempt?: string, ): Promise { + const query = authorizationAttempt + ? `?authorization_attempt=${encodeURIComponent(authorizationAttempt)}` + : '' const response: any = await get( - `/api/v1/embed/${channelId}/sessions/${encodeURIComponent(sessionId)}/mcp-services/${encodeURIComponent(serviceId)}/oauth/status`, + `/api/v1/embed/${channelId}/sessions/${encodeURIComponent(sessionId)}/mcp-services/${encodeURIComponent(serviceId)}/oauth/status${query}`, { headers: embedSessionHeaders(token, sessionSig, visitorId) }, ) return Boolean((response.data ?? response)?.authorized) diff --git a/frontend/src/api/mcp-service.ts b/frontend/src/api/mcp-service.ts index 448de1ca..e630cc39 100644 --- a/frontend/src/api/mcp-service.ts +++ b/frontend/src/api/mcp-service.ts @@ -194,21 +194,59 @@ export async function deleteMCPCredentialField( // to avoid a route conflict, and allow-listed for no-auth in the backend). export const MCP_OAUTH_CALLBACK_PATH = '/api/v1/mcp-oauth/callback' -// Begin authorization for the current user. Returns the URL to open in a popup. +export interface MCPOAuthAuthorization { + authorizationUrl: string + authorizationAttempt: string +} + +export type MCPOAuthTokenState = 'authorized' | 'refreshable' | 'reauth_required' | 'pending' + +export interface MCPOAuthStatus { + authorized: boolean + state: MCPOAuthTokenState + refresh_available: boolean + expires_at?: string +} + +// Begin authorization for the current user. The attempt id binds polling to +// this popup, so an older stored token cannot be mistaken for fresh consent. export async function getMCPOAuthAuthorizeURL( serviceId: string, body: { redirect_uri: string; frontend_redirect?: string } -): Promise { +): Promise { const response: any = await post(`/api/v1/mcp-services/${serviceId}/oauth/authorize-url`, body) - return (response.data ?? response)?.authorization_url ?? '' + const data = response.data ?? response + return { + authorizationUrl: data?.authorization_url ?? '', + authorizationAttempt: data?.authorization_attempt ?? '', + } } // Whether the current user has authorized this service. -export async function getMCPOAuthStatus(serviceId: string): Promise { - const response: any = await get(`/api/v1/mcp-services/${serviceId}/oauth/status`) +export async function getMCPOAuthStatus( + serviceId: string, + authorizationAttempt?: string, +): Promise { + const query = authorizationAttempt + ? `?authorization_attempt=${encodeURIComponent(authorizationAttempt)}` + : '' + const response: any = await get(`/api/v1/mcp-services/${serviceId}/oauth/status${query}`) return Boolean((response.data ?? response)?.authorized) } +// Full lifecycle status for management surfaces. Expired access tokens with a +// refresh token are "refreshable", not falsely presented as already usable. +export async function getMCPOAuthAuthorizationStatus(serviceId: string): Promise { + const response: any = await get(`/api/v1/mcp-services/${serviceId}/oauth/status`) + const data = response.data ?? response + return { + authorized: Boolean(data?.authorized), + state: data?.state ?? 'reauth_required', + refresh_available: Boolean(data?.refresh_available), + expires_at: data?.expires_at, + } +} + // Revoke the current user's token (forces re-authorization). export async function revokeMCPOAuthToken(serviceId: string): Promise { await del(`/api/v1/mcp-services/${serviceId}/oauth/token`) diff --git a/frontend/src/i18n/locales/en-US.ts b/frontend/src/i18n/locales/en-US.ts index 49b02646..0e96b2c0 100755 --- a/frontend/src/i18n/locales/en-US.ts +++ b/frontend/src/i18n/locales/en-US.ts @@ -4997,6 +4997,7 @@ export default { authTypeNone: 'None / Custom Header', authTypeApiKey: 'API Key / Token', authTypeOAuth: 'OAuth 2.0 (authorize on first connect)', + oauthRefreshable: 'Token expired; it will refresh automatically on next use', oauthAuthorizeHint: 'Clicking "Authorize" saves the current config first, then starts authorization (each user authorizes individually).', apiKeyHeader: 'Header Name', apiKeyHeaderDesc: 'Defaults to X-API-Key. For Bearer, set Authorization and put "Bearer " in the value below; for raw-token services use Authorization with the raw token.', diff --git a/frontend/src/i18n/locales/zh-CN.ts b/frontend/src/i18n/locales/zh-CN.ts index 146cdc2e..33492abc 100755 --- a/frontend/src/i18n/locales/zh-CN.ts +++ b/frontend/src/i18n/locales/zh-CN.ts @@ -5045,6 +5045,7 @@ export default { oauthScopes: "Scopes（可选，空格分隔）", oauthAuthorization: "授权状态", oauthAuthorized: "已授权", + oauthRefreshable: "Token 已过期，将在下次使用时自动刷新", oauthUnauthorized: "未授权", oauthAuthorize: "去授权", oauthReauthorize: "重新授权", diff --git a/frontend/src/views/chat/components/McpOAuthCard.oauth.test.mjs b/frontend/src/views/chat/components/McpOAuthCard.oauth.test.mjs new file mode 100644 index 00000000..23b9cc3b --- /dev/null +++ b/frontend/src/views/chat/components/McpOAuthCard.oauth.test.mjs @@ -0,0 +1,39 @@ +import assert from 'node:assert/strict' +import { readFileSync } from 'node:fs' +import { dirname, join } from 'node:path' +import { fileURLToPath } from 'node:url' +import test from 'node:test' + +const here = dirname(fileURLToPath(import.meta.url)) +const cardSource = readFileSync(join(here, 'McpOAuthCard.vue'), 'utf8') +const settingsSource = readFileSync(join(here, '../../settings/components/McpServiceDialog.vue'), 'utf8') +const apiSource = readFileSync(join(here, '../../../api/mcp-service.ts'), 'utf8') + +test('in-chat OAuth polling is bound to the newly opened authorization attempt', () => { + assert.match(cardSource, /authorization\.authorizationAttempt/) + assert.match( + cardSource, + /getMCPOAuthStatus\(props\.serviceId, authorization\.authorizationAttempt\)/, + ) + assert.match( + cardSource, + /props\.serviceId,\s*authorization\.authorizationAttempt,\s*\)/, + ) +}) + +test('settings OAuth polling cannot accept a pre-existing token as fresh authorization', () => { + assert.match( + settingsSource, + /getMCPOAuthStatus\(serviceId, authorization\.authorizationAttempt\)/, + ) +}) + +test('OAuth status API sends the attempt id to the backend', () => { + assert.match(apiSource, /authorization_attempt=\$\{encodeURIComponent\(authorizationAttempt\)\}/) +}) + +test('settings distinguishes refreshable tokens from usable authorization', () => { + assert.match(settingsSource, /getMCPOAuthAuthorizationStatus\(props\.service\.id\)/) + assert.match(settingsSource, /oauthTokenState === 'refreshable'/) + assert.match(apiSource, /state: data\?\.state \?\? 'reauth_required'/) +}) diff --git a/frontend/src/views/chat/components/McpOAuthCard.vue b/frontend/src/views/chat/components/McpOAuthCard.vue index 0efb990d..62bfa110 100644 --- a/frontend/src/views/chat/components/McpOAuthCard.vue +++ b/frontend/src/views/chat/components/McpOAuthCard.vue @@ -175,7 +175,7 @@ const authorize = async () => { const frontendRedirect = useEmbedOAuth() ? window.location.origin + window.location.pathname + window.location.search : window.location.origin + '/' - const authUrl = useEmbedOAuth() + const authorization = useEmbedOAuth() ? await getEmbedMCPOAuthAuthorizeURL( props.embedChannelId!, props.embedToken!, @@ -189,12 +189,12 @@ const authorize = async () => { redirect_uri: redirectUri, frontend_redirect: frontendRedirect, }) - if (!authUrl) { + if (!authorization.authorizationUrl || !authorization.authorizationAttempt) { MessagePlugin.error(t('agentStream.mcpOAuth.startFailed')) authorizing.value = false return } - const popup = window.open(authUrl, 'mcp_oauth', 'width=600,height=720') + const popup = window.open(authorization.authorizationUrl, 'mcp_oauth', 'width=600,height=720') poll = window.setInterval(async () => { const closed = !popup || popup.closed let ok = false @@ -207,8 +207,9 @@ const authorize = async () => { props.embedSessionSig!, props.embedVisitorId!, props.serviceId, + authorization.authorizationAttempt, ) - : await getMCPOAuthStatus(props.serviceId) + : await getMCPOAuthStatus(props.serviceId, authorization.authorizationAttempt) } catch { /* transient; keep polling */ } diff --git a/frontend/src/views/settings/components/McpServiceDialog.vue b/frontend/src/views/settings/components/McpServiceDialog.vue index 546961d4..08bc6cc6 100644 --- a/frontend/src/views/settings/components/McpServiceDialog.vue +++ b/frontend/src/views/settings/components/McpServiceDialog.vue @@ -220,9 +220,12 @@
- + {{ t('mcpServiceDialog.oauthAuthorized', '已授权') }} + + {{ t('mcpServiceDialog.oauthRefreshable', 'Token 已过期，将自动刷新') }} + {{ t('mcpServiceDialog.oauthUnauthorized', '未授权') }} @@ -232,10 +235,10 @@ :loading="oauthAuthorizing || oauthChecking || submitting" @click="handleAuthorize" > - {{ oauthAuthorized ? t('mcpServiceDialog.oauthReauthorize', '重新授权') : t('mcpServiceDialog.oauthAuthorize', '去授权') }} + {{ oauthTokenState === 'reauth_required' ? t('mcpServiceDialog.oauthAuthorize', '去授权') : t('mcpServiceDialog.oauthReauthorize', '重新授权') }} [ // ---- OAuth authorization state (edit mode only) ---- const oauthAuthorized = ref(false) +const oauthTokenState = ref('reauth_required') const oauthChecking = ref(false) const oauthAuthorizing = ref(false) @@ -614,7 +620,9 @@ async function refreshOAuthStatus() { if (props.mode !== 'edit' || !props.service?.id || !isOAuth.value) return oauthChecking.value = true try { - oauthAuthorized.value = await getMCPOAuthStatus(props.service.id) + const status = await getMCPOAuthAuthorizationStatus(props.service.id) + oauthAuthorized.value = status.authorized + oauthTokenState.value = status.state } catch (e) { console.error('Failed to query MCP OAuth status:', e) } finally { @@ -634,21 +642,22 @@ async function startAuthorize(serviceId: string) { // app root is harmless; the popup is closed by the opener below once the // authorization status flips, so this page is only shown briefly. const frontendRedirect = window.location.origin + '/' - const authUrl = await getMCPOAuthAuthorizeURL(serviceId, { + const authorization = await getMCPOAuthAuthorizeURL(serviceId, { redirect_uri: redirectUri, frontend_redirect: frontendRedirect, }) - if (!authUrl) { + if (!authorization.authorizationUrl || !authorization.authorizationAttempt) { MessagePlugin.error(t('mcpServiceDialog.toasts.authorizeFailed', '发起授权失败') as string) oauthAuthorizing.value = false return } - const popup = window.open(authUrl, 'mcp_oauth', 'width=600,height=720') + const popup = window.open(authorization.authorizationUrl, 'mcp_oauth', 'width=600,height=720') // Poll for completion: either the popup closes or the status flips. const timer = window.setInterval(async () => { const closed = !popup || popup.closed try { - oauthAuthorized.value = await getMCPOAuthStatus(serviceId) + oauthAuthorized.value = await getMCPOAuthStatus(serviceId, authorization.authorizationAttempt) + if (oauthAuthorized.value) oauthTokenState.value = 'authorized' } catch { /* keep polling */ } if (oauthAuthorized.value || closed) { window.clearInterval(timer) @@ -704,6 +713,7 @@ async function handleRevokeOAuth() { try { await revokeMCPOAuthToken(props.service.id) oauthAuthorized.value = false + oauthTokenState.value = 'reauth_required' MessagePlugin.success(t('mcpServiceDialog.toasts.revoked', '已撤销授权') as string) } catch (e) { console.error('Failed to revoke MCP OAuth token:', e) @@ -961,6 +971,7 @@ watch( }, } oauthAuthorized.value = false + oauthTokenState.value = 'reauth_required' refreshOAuthStatus() } else { resetForm() diff --git a/internal/agent/tools/mcp_oauth.go b/internal/agent/tools/mcp_oauth.go index df8b70b1..eb59d531 100644 --- a/internal/agent/tools/mcp_oauth.go +++ b/internal/agent/tools/mcp_oauth.go @@ -2,6 +2,7 @@ package tools import ( "context" + "errors" "fmt" "strings" "time" @@ -271,6 +272,10 @@ func isAuthorizationRequired(err error) bool { if mcpclient.IsOAuthAuthorizationRequiredError(err) || mcpclient.IsAuthorizationRequiredError(err) { return true } + var reauth *mcp.OAuthReauthorizationRequiredError + if errors.As(err, &reauth) { + return true + } msg := err.Error() return strings.Contains(msg, "authorization required") || strings.Contains(msg, "no valid token") || diff --git a/internal/application/repository/mcp_oauth.go b/internal/application/repository/mcp_oauth.go index be0a2349..8a68257f 100644 --- a/internal/application/repository/mcp_oauth.go +++ b/internal/application/repository/mcp_oauth.go @@ -140,3 +140,51 @@ func (r *mcpOAuthRepository) DeleteTokenForPrincipal( ). Delete(&types.MCPOAuthToken{}).Error } + +func (r *mcpOAuthRepository) TryAcquireTokenRefreshLease( + ctx context.Context, + tenantID uint64, + principal types.Principal, + serviceID, leaseID string, + leaseUntil time.Time, +) (bool, error) { + principal = principal.Normalize() + if !principal.Valid() || leaseID == "" { + return false, nil + } + now := time.Now() + result := r.db.WithContext(ctx). + Model(&types.MCPOAuthToken{}). + Where( + "tenant_id = ? AND principal_type = ? AND principal_id = ? AND service_id = ?", + tenantID, principal.Type, principal.ID, serviceID, + ). + Where("refresh_lease_until IS NULL OR refresh_lease_until < ?", now). + Updates(map[string]interface{}{ + "refresh_lease_id": leaseID, + "refresh_lease_until": leaseUntil, + }) + return result.RowsAffected == 1, result.Error +} + +func (r *mcpOAuthRepository) ReleaseTokenRefreshLease( + ctx context.Context, + tenantID uint64, + principal types.Principal, + serviceID, leaseID string, +) error { + principal = principal.Normalize() + if !principal.Valid() || leaseID == "" { + return nil + } + return r.db.WithContext(ctx). + Model(&types.MCPOAuthToken{}). + Where( + "tenant_id = ? AND principal_type = ? AND principal_id = ? AND service_id = ? AND refresh_lease_id = ?", + tenantID, principal.Type, principal.ID, serviceID, leaseID, + ). + Updates(map[string]interface{}{ + "refresh_lease_id": "", + "refresh_lease_until": nil, + }).Error +} diff --git a/internal/application/repository/mcp_oauth_test.go b/internal/application/repository/mcp_oauth_test.go index 592af7c3..cbada316 100644 --- a/internal/application/repository/mcp_oauth_test.go +++ b/internal/application/repository/mcp_oauth_test.go @@ -79,3 +79,46 @@ func TestMCPOAuthRepositoryLegacyUserTokenUsesWebPrincipal(t *testing.T) { require.NoError(t, err) require.Equal(t, "legacy-token", token.AccessToken) } + +func TestMCPOAuthRepositoryRefreshLeaseHasSingleOwner(t *testing.T) { + repo := newMCPOAuthTestRepo(t) + ctx := context.Background() + principal := types.Principal{Type: types.PrincipalWebUser, ID: "u1"} + require.NoError(t, repo.SaveTokenForPrincipal(ctx, &types.MCPOAuthToken{ + TenantID: 7, + UserID: principal.StorageID(), + PrincipalType: principal.Type, + PrincipalID: principal.ID, + ServiceID: "svc1", + AccessToken: "access", + RefreshToken: "refresh", + ExpiresAt: time.Now().Add(-time.Minute), + })) + + first, err := repo.TryAcquireTokenRefreshLease( + ctx, 7, principal, "svc1", "lease-1", time.Now().Add(time.Minute), + ) + require.NoError(t, err) + require.True(t, first) + + second, err := repo.TryAcquireTokenRefreshLease( + ctx, 7, principal, "svc1", "lease-2", time.Now().Add(time.Minute), + ) + require.NoError(t, err) + require.False(t, second) + + // A non-owner cannot release the current owner's lease. + require.NoError(t, repo.ReleaseTokenRefreshLease(ctx, 7, principal, "svc1", "lease-2")) + second, err = repo.TryAcquireTokenRefreshLease( + ctx, 7, principal, "svc1", "lease-2", time.Now().Add(time.Minute), + ) + require.NoError(t, err) + require.False(t, second) + + require.NoError(t, repo.ReleaseTokenRefreshLease(ctx, 7, principal, "svc1", "lease-1")) + second, err = repo.TryAcquireTokenRefreshLease( + ctx, 7, principal, "svc1", "lease-2", time.Now().Add(time.Minute), + ) + require.NoError(t, err) + require.True(t, second) +} diff --git a/internal/handler/mcp_oauth.go b/internal/handler/mcp_oauth.go index 19520783..fe92e928 100644 --- a/internal/handler/mcp_oauth.go +++ b/internal/handler/mcp_oauth.go @@ -65,7 +65,7 @@ type mcpOAuthAuthorizeRequest struct { // @Produce json // @Param id path string true "MCP 服务 ID" // @Param request body map[string]interface{} true "{redirect_uri: string, frontend_redirect?: string}" -// @Success 200 {object} map[string]interface{} "{authorization_url: string}" +// @Success 200 {object} map[string]interface{} "{authorization_url: string, authorization_attempt: string}" // @Failure 400 {object} errors.AppError // @Security Bearer // @Router /mcp-services/{id}/oauth/authorize-url [post] @@ -103,7 +103,9 @@ func (h *MCPOAuthHandler) AuthorizeURL(c *gin.Context) { return } - authURL, err := h.oauth.StartAuthorization(ctx, service, tenantID, principal, req.RedirectURI, req.FrontendRedirect) + authURL, attemptID, err := h.oauth.StartAuthorization( + ctx, service, tenantID, principal, req.RedirectURI, req.FrontendRedirect, + ) if err != nil { logger.ErrorWithFields(ctx, err, map[string]interface{}{ "service_id": secutils.SanitizeForLog(serviceID), @@ -112,7 +114,10 @@ func (h *MCPOAuthHandler) AuthorizeURL(c *gin.Context) { return } - c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{"authorization_url": authURL}}) + c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{ + "authorization_url": authURL, + "authorization_attempt": attemptID, + }}) } // Callback receives the authorization-server redirect. It is registered as a @@ -146,7 +151,7 @@ func (h *MCPOAuthHandler) Callback(c *gin.Context) { return } - frontendRedirect, err := h.oauth.CompleteAuthorization(ctx, state, code) + frontendRedirect, serviceID, err := h.oauth.CompleteAuthorization(ctx, state, code) if frontendRedirect == "" { frontendRedirect = fallbackRedirect } @@ -155,6 +160,12 @@ func (h *MCPOAuthHandler) Callback(c *gin.Context) { c.Redirect(http.StatusFound, frontendRedirect+"#mcp_oauth_error="+urlQueryEscape("authorization_failed")) return } + // The old transport may have been created with an OAuth client registration + // that was invalidated together with the refresh token. Recreate it against + // the freshly persisted token/client on next use. + if h.mcpManager != nil && serviceID != "" { + _ = h.mcpManager.CloseClient(serviceID) + } c.Redirect(http.StatusFound, frontendRedirect+"#mcp_oauth_result=success") } @@ -162,11 +173,12 @@ func (h *MCPOAuthHandler) Callback(c *gin.Context) { // // Status godoc // @Summary 查询 MCP OAuth 授权状态 -// @Description 返回当前用户对指定 MCP 服务是否已完成 OAuth 授权 +// @Description 返回当前用户的 OAuth Token 生命周期状态；传 authorization_attempt 时只检查本次授权流程 // @Tags MCP服务 // @Produce json // @Param id path string true "MCP 服务 ID" -// @Success 200 {object} map[string]interface{} "{authorized: bool}" +// @Param authorization_attempt query string false "本次授权尝试 ID；传入后不会接受历史 Token" +// @Success 200 {object} map[string]interface{} "{authorized: bool, state: string, refresh_available: bool, expires_at?: string}" // @Security Bearer // @Router /mcp-services/{id}/oauth/status [get] func (h *MCPOAuthHandler) Status(c *gin.Context) { @@ -179,12 +191,32 @@ func (h *MCPOAuthHandler) Status(c *gin.Context) { return } - authorized, err := h.oauth.IsAuthorized(ctx, tenantID, principal, serviceID) + attemptID := strings.TrimSpace(c.Query("authorization_attempt")) + if attemptID != "" { + authorized, err := h.oauth.IsAuthorizationAttemptComplete( + ctx, tenantID, principal, serviceID, attemptID, + ) + if err != nil { + c.Error(errors.NewInternalServerError("failed to query authorization status: " + err.Error())) + return + } + state := "pending" + if authorized { + state = "authorized" + } + c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{ + "authorized": authorized, + "state": state, + }}) + return + } + + status, err := h.oauth.AuthorizationStatus(ctx, tenantID, principal, serviceID) if err != nil { c.Error(errors.NewInternalServerError("failed to query authorization status: " + err.Error())) return } - c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{"authorized": authorized}}) + c.JSON(http.StatusOK, gin.H{"success": true, "data": status}) } // Revoke removes the current user's stored token and recycles connections. diff --git a/internal/mcp/client.go b/internal/mcp/client.go index 61935882..63b2bb94 100644 --- a/internal/mcp/client.go +++ b/internal/mcp/client.go @@ -66,6 +66,7 @@ type ClientConfig struct { type mcpGoClient struct { service *types.MCPService client *client.Client + oauth *oauthRuntime connected bool initialized bool } @@ -223,6 +224,16 @@ func NewMCPClient(config *ClientConfig) (MCPClient, error) { service: config.Service, client: mcpClient, } + if useOAuth { + instance.oauth = newOAuthRuntime( + config.OAuthRepo, + config.TenantID, + config.Principal, + config.Service.ID, + *config.Service.URL, + oauthConfig, + ) + } mcpClient.OnConnectionLost(instance.onConnectionLost) return instance, nil } @@ -246,10 +257,11 @@ func buildOAuthConfig(config *ClientConfig, httpClient *http.Client) (transport. if !principal.Valid() { return transport.OAuthConfig{}, false, fmt.Errorf("principal context is required to connect to an OAuth MCP service") } + config.Principal = principal oauthCfg := transport.OAuthConfig{ Scopes: svc.AuthConfig.Scopes, - TokenStore: newDBTokenStore(config.OAuthRepo, config.TenantID, principal, svc.ID), + TokenStore: newManagedTokenStore(config.OAuthRepo, config.TenantID, principal, svc.ID), PKCEEnabled: true, AuthServerMetadataURL: svc.AuthConfig.AuthServerMetadataURL, HTTPClient: httpClient, @@ -288,14 +300,37 @@ func (c *mcpGoClient) checkErrorAndDisconnectIfNeeded(err error) { } } +// oauthCall runs one MCP operation with WeKnora-owned token lifecycle checks. +// A resource-server 401 forces exactly one refresh and one retry. Other errors +// are never retried, which avoids duplicating tool side effects after ambiguous +// network failures. +func oauthCall[T any](ctx context.Context, c *mcpGoClient, operation func() (T, error)) (T, error) { + var zero T + if c.oauth != nil { + if err := c.oauth.ensureFresh(ctx, false, nil); err != nil { + return zero, err + } + } + result, err := operation() + if err == nil || c.oauth == nil || !isOAuthAuthorizationFailure(err) { + return result, err + } + if refreshErr := c.oauth.ensureFresh(ctx, true, client.GetOAuthHandler(err)); refreshErr != nil { + return zero, refreshErr + } + return operation() +} + // Connect establishes connection to the MCP service func (c *mcpGoClient) Connect(ctx context.Context) error { if c.connected { return ErrAlreadyConnected } - // Start the client - if err := c.client.Start(ctx); err != nil { + _, err := oauthCall(ctx, c, func() (struct{}, error) { + return struct{}{}, c.client.Start(ctx) + }) + if err != nil { if oerr := asOAuthRequired(err); oerr != nil { return oerr } @@ -344,7 +379,9 @@ func (c *mcpGoClient) Initialize(ctx context.Context) (*InitializeResult, error) }, } - result, err := c.client.Initialize(ctx, req) + result, err := oauthCall(ctx, c, func() (*mcp.InitializeResult, error) { + return c.client.Initialize(ctx, req) + }) if err != nil { c.checkErrorAndDisconnectIfNeeded(err) if oerr := asOAuthRequired(err); oerr != nil { @@ -373,7 +410,9 @@ func (c *mcpGoClient) ListTools(ctx context.Context) ([]*types.MCPTool, error) { } req := mcp.ListToolsRequest{} - result, err := c.client.ListTools(ctx, req) + result, err := oauthCall(ctx, c, func() (*mcp.ListToolsResult, error) { + return c.client.ListTools(ctx, req) + }) if err != nil { c.checkErrorAndDisconnectIfNeeded(err) return nil, fmt.Errorf("failed to list tools: %w", err) @@ -400,7 +439,9 @@ func (c *mcpGoClient) ListResources(ctx context.Context) ([]*types.MCPResource, } req := mcp.ListResourcesRequest{} - result, err := c.client.ListResources(ctx, req) + result, err := oauthCall(ctx, c, func() (*mcp.ListResourcesResult, error) { + return c.client.ListResources(ctx, req) + }) if err != nil { c.checkErrorAndDisconnectIfNeeded(err) return nil, fmt.Errorf("failed to list resources: %w", err) @@ -433,7 +474,9 @@ func (c *mcpGoClient) CallTool(ctx context.Context, name string, args map[string }, } - result, err := c.client.CallTool(ctx, req) + result, err := oauthCall(ctx, c, func() (*mcp.CallToolResult, error) { + return c.client.CallTool(ctx, req) + }) if err != nil { c.checkErrorAndDisconnectIfNeeded(err) return nil, fmt.Errorf("failed to call tool: %w", err) @@ -474,7 +517,9 @@ func (c *mcpGoClient) ReadResource(ctx context.Context, uri string) (*ReadResour }, } - result, err := c.client.ReadResource(ctx, req) + result, err := oauthCall(ctx, c, func() (*mcp.ReadResourceResult, error) { + return c.client.ReadResource(ctx, req) + }) if err != nil { c.checkErrorAndDisconnectIfNeeded(err) return nil, fmt.Errorf("failed to read resource: %w", err) diff --git a/internal/mcp/oauth_lifecycle.go b/internal/mcp/oauth_lifecycle.go new file mode 100644 index 00000000..ec7ed05f --- /dev/null +++ b/internal/mcp/oauth_lifecycle.go @@ -0,0 +1,280 @@ +package mcp + +import ( + "context" + "errors" + "fmt" + "strings" + "time" + + "github.com/Tencent/WeKnora/internal/logger" + "github.com/Tencent/WeKnora/internal/types" + "github.com/Tencent/WeKnora/internal/types/interfaces" + "github.com/google/uuid" + "github.com/mark3labs/mcp-go/client" + "github.com/mark3labs/mcp-go/client/transport" +) + +const ( + oauthRefreshSkew = 30 * time.Second + oauthRefreshLease = 45 * time.Second + oauthRefreshPoll = 100 * time.Millisecond + oauthStateAuthorized = "authorized" + oauthStateRefreshable = "refreshable" + oauthStateReauthNeeded = "reauth_required" +) + +// OAuthAuthorizationStatus distinguishes a currently usable access token from +// an expired token that can still be refreshed. This prevents a stale database +// row from being presented as an already successful authorization. +type OAuthAuthorizationStatus struct { + Authorized bool `json:"authorized"` + State string `json:"state"` + RefreshAvailable bool `json:"refresh_available"` + ExpiresAt *time.Time `json:"expires_at,omitempty"` +} + +// OAuthReauthorizationRequiredError means no usable access token can be +// recovered without interactive user consent. +type OAuthReauthorizationRequiredError struct { + Reason string +} + +func (e *OAuthReauthorizationRequiredError) Error() string { + if e.Reason == "" { + return "MCP OAuth authorization required" + } + return "MCP OAuth authorization required: " + e.Reason +} + +// OAuthRefreshTemporaryError preserves a token when refresh failed for a +// transient reason. Callers must surface/retry this as an operational failure, +// not open a new consent popup. +type OAuthRefreshTemporaryError struct { + Err error +} + +func (e *OAuthRefreshTemporaryError) Error() string { + return fmt.Sprintf("MCP OAuth token refresh temporarily failed: %v", e.Err) +} + +func (e *OAuthRefreshTemporaryError) Unwrap() error { return e.Err } + +type oauthRuntime struct { + repo interfaces.MCPOAuthRepository + tenantID uint64 + principal types.Principal + serviceID string + handler *transport.OAuthHandler + leaseDuration time.Duration +} + +func newOAuthRuntime( + repo interfaces.MCPOAuthRepository, + tenantID uint64, + principal types.Principal, + serviceID, baseURL string, + cfg transport.OAuthConfig, +) *oauthRuntime { + h := transport.NewOAuthHandler(cfg) + h.SetBaseURL(baseURL) + leaseDuration := oauthRefreshLease + if cfg.HTTPClient != nil && cfg.HTTPClient.Timeout > 0 && cfg.HTTPClient.Timeout+15*time.Second > leaseDuration { + leaseDuration = cfg.HTTPClient.Timeout + 15*time.Second + } + return &oauthRuntime{ + repo: repo, + tenantID: tenantID, + principal: principal.Normalize(), + serviceID: serviceID, + handler: h, + leaseDuration: leaseDuration, + } +} + +func tokenStatus(token *types.MCPOAuthToken, now time.Time) OAuthAuthorizationStatus { + status := OAuthAuthorizationStatus{State: oauthStateReauthNeeded} + if token == nil || token.AccessToken == "" { + return status + } + status.RefreshAvailable = token.RefreshToken != "" + if !token.ExpiresAt.IsZero() { + expiresAt := token.ExpiresAt + status.ExpiresAt = &expiresAt + } + if token.ExpiresAt.IsZero() || token.ExpiresAt.After(now) { + status.Authorized = true + status.State = oauthStateAuthorized + return status + } + if status.RefreshAvailable { + status.State = oauthStateRefreshable + } + return status +} + +func (r *oauthRuntime) ensureFresh(ctx context.Context, force bool, override *transport.OAuthHandler) error { + row, err := r.repo.GetTokenForPrincipal(ctx, r.tenantID, r.principal, r.serviceID) + if err != nil { + return fmt.Errorf("load MCP OAuth token: %w", err) + } + if row == nil || row.AccessToken == "" { + return &OAuthReauthorizationRequiredError{Reason: "no token is stored"} + } + now := time.Now() + if !force { + if row.ExpiresAt.IsZero() || row.ExpiresAt.After(now.Add(oauthRefreshSkew)) { + return nil + } + // Tokens issued without refresh_token remain usable through their actual + // expiry; the refresh skew must not shorten their lifetime. + if row.RefreshToken == "" && row.ExpiresAt.After(now) { + return nil + } + } + if row.RefreshToken == "" { + _ = r.repo.DeleteTokenForPrincipal(ctx, r.tenantID, r.principal, r.serviceID) + return &OAuthReauthorizationRequiredError{Reason: "the access token expired and no refresh token is available"} + } + return r.refreshWithLease(ctx, row, override) +} + +func (r *oauthRuntime) refreshWithLease( + ctx context.Context, observed *types.MCPOAuthToken, override *transport.OAuthHandler, +) error { + for { + leaseID := uuid.NewString() + leaseDuration := r.leaseDuration + if leaseDuration <= 0 { + leaseDuration = oauthRefreshLease + } + leaseUntil := time.Now().Add(leaseDuration) + acquired, err := r.repo.TryAcquireTokenRefreshLease( + ctx, r.tenantID, r.principal, r.serviceID, leaseID, leaseUntil, + ) + if err != nil { + return fmt.Errorf("claim MCP OAuth token refresh: %w", err) + } + if acquired { + return r.refreshAsLeaseOwner(ctx, observed, leaseID, override) + } + + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(oauthRefreshPoll): + } + current, err := r.repo.GetTokenForPrincipal(ctx, r.tenantID, r.principal, r.serviceID) + if err != nil { + return fmt.Errorf("reload MCP OAuth token after concurrent refresh: %w", err) + } + if current == nil || current.AccessToken == "" { + return &OAuthReauthorizationRequiredError{Reason: "the refresh token is no longer valid"} + } + if oauthTokenMaterialChanged(current, observed) { + if current.ExpiresAt.IsZero() || current.ExpiresAt.After(time.Now().Add(oauthRefreshSkew)) { + return nil + } + observed = current + } + } +} + +func (r *oauthRuntime) refreshAsLeaseOwner( + ctx context.Context, observed *types.MCPOAuthToken, leaseID string, override *transport.OAuthHandler, +) error { + defer func() { + releaseCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + if err := r.repo.ReleaseTokenRefreshLease( + releaseCtx, r.tenantID, r.principal, r.serviceID, leaseID, + ); err != nil { + logger.GetLogger(releaseCtx).Warnf("failed to release MCP OAuth refresh lease: %v", err) + } + }() + + current, err := r.repo.GetTokenForPrincipal(ctx, r.tenantID, r.principal, r.serviceID) + if err != nil { + return fmt.Errorf("reload MCP OAuth token before refresh: %w", err) + } + if current == nil || current.AccessToken == "" { + return &OAuthReauthorizationRequiredError{Reason: "no token is stored"} + } + // Another owner may have completed a refresh immediately before this lease + // was acquired. Never consume its newly rotated refresh token unnecessarily. + if oauthTokenMaterialChanged(current, observed) && + (current.ExpiresAt.IsZero() || current.ExpiresAt.After(time.Now().Add(oauthRefreshSkew))) { + return nil + } + if current.RefreshToken == "" { + return r.invalidateToken(ctx, false, "no refresh token is available") + } + + handler := override + if handler == nil { + handler = r.handler + } + refreshed, refreshErr := handler.RefreshToken(ctx, current.RefreshToken) + if refreshErr == nil && refreshed != nil && refreshed.AccessToken != "" { + logger.GetLogger(ctx).Infof("MCP OAuth token refreshed: service=%s principal=%s", r.serviceID, r.principal.StorageID()) + return nil + } + if refreshErr == nil { + refreshErr = errors.New("authorization server returned an empty access token") + } + permanent, resetClient := permanentRefreshFailure(refreshErr) + if permanent { + return r.invalidateToken(ctx, resetClient, "the refresh token or OAuth client is no longer valid") + } + return &OAuthRefreshTemporaryError{Err: refreshErr} +} + +func oauthTokenMaterialChanged(current, observed *types.MCPOAuthToken) bool { + if current == nil || observed == nil { + return current != observed + } + return current.AccessToken != observed.AccessToken || + current.RefreshToken != observed.RefreshToken || + !current.ExpiresAt.Equal(observed.ExpiresAt) +} + +func (r *oauthRuntime) invalidateToken(ctx context.Context, resetClient bool, reason string) error { + if err := r.repo.DeleteTokenForPrincipal(ctx, r.tenantID, r.principal, r.serviceID); err != nil { + return fmt.Errorf("delete invalid MCP OAuth token: %w", err) + } + if resetClient { + if err := r.repo.DeleteClient(ctx, r.tenantID, r.serviceID); err != nil { + return fmt.Errorf("delete invalid MCP OAuth client registration: %w", err) + } + } + return &OAuthReauthorizationRequiredError{Reason: reason} +} + +func permanentRefreshFailure(err error) (permanent bool, resetClient bool) { + var oauthErr transport.OAuthError + if errors.As(err, &oauthErr) { + switch strings.ToLower(oauthErr.ErrorCode) { + case "invalid_grant", "invalid_token", "bad_refresh_token", "expired_token": + return true, false + case "invalid_client", "unauthorized_client": + return true, true + } + } + lower := strings.ToLower(err.Error()) + if strings.Contains(lower, "status 400") { + return true, false + } + if strings.Contains(lower, "status 401") { + return true, true + } + return false, false +} + +func isOAuthAuthorizationFailure(err error) bool { + if err == nil { + return false + } + var reauth *OAuthReauthorizationRequiredError + return errors.As(err, &reauth) || client.IsOAuthAuthorizationRequiredError(err) || + client.IsAuthorizationRequiredError(err) +} diff --git a/internal/mcp/oauth_lifecycle_test.go b/internal/mcp/oauth_lifecycle_test.go new file mode 100644 index 00000000..f44d8d9f --- /dev/null +++ b/internal/mcp/oauth_lifecycle_test.go @@ -0,0 +1,316 @@ +package mcp + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/mark3labs/mcp-go/client/transport" + "github.com/stretchr/testify/require" +) + +type lockedOAuthRepo struct { + *fakeOAuthRepo + mu sync.Mutex +} + +func newLockedOAuthRepo() *lockedOAuthRepo { + return &lockedOAuthRepo{fakeOAuthRepo: newFakeOAuthRepo()} +} + +func cloneOAuthToken(token *types.MCPOAuthToken) *types.MCPOAuthToken { + if token == nil { + return nil + } + clone := *token + if token.RefreshLeaseUntil != nil { + leaseUntil := *token.RefreshLeaseUntil + clone.RefreshLeaseUntil = &leaseUntil + } + return &clone +} + +func (r *lockedOAuthRepo) GetTokenForPrincipal( + _ context.Context, tenantID uint64, principal types.Principal, serviceID string, +) (*types.MCPOAuthToken, error) { + r.mu.Lock() + defer r.mu.Unlock() + return cloneOAuthToken(r.tokens[fakeOAuthKey(tenantID, principal, serviceID)]), nil +} + +func (r *lockedOAuthRepo) SaveTokenForPrincipal(_ context.Context, token *types.MCPOAuthToken) error { + r.mu.Lock() + defer r.mu.Unlock() + principal := types.Principal{Type: token.PrincipalType, ID: token.PrincipalID}.Normalize() + token = cloneOAuthToken(token) + token.UpdatedAt = time.Now() + r.tokens[fakeOAuthKey(token.TenantID, principal, token.ServiceID)] = token + return nil +} + +func (r *lockedOAuthRepo) DeleteTokenForPrincipal( + _ context.Context, tenantID uint64, principal types.Principal, serviceID string, +) error { + r.mu.Lock() + defer r.mu.Unlock() + delete(r.tokens, fakeOAuthKey(tenantID, principal, serviceID)) + return nil +} + +func (r *lockedOAuthRepo) TryAcquireTokenRefreshLease( + _ context.Context, + tenantID uint64, + principal types.Principal, + serviceID, leaseID string, + leaseUntil time.Time, +) (bool, error) { + r.mu.Lock() + defer r.mu.Unlock() + row := r.tokens[fakeOAuthKey(tenantID, principal, serviceID)] + if row == nil || (row.RefreshLeaseUntil != nil && row.RefreshLeaseUntil.After(time.Now())) { + return false, nil + } + row.RefreshLeaseID = leaseID + row.RefreshLeaseUntil = &leaseUntil + row.UpdatedAt = time.Now() + return true, nil +} + +func (r *lockedOAuthRepo) ReleaseTokenRefreshLease( + _ context.Context, + tenantID uint64, + principal types.Principal, + serviceID, leaseID string, +) error { + r.mu.Lock() + defer r.mu.Unlock() + row := r.tokens[fakeOAuthKey(tenantID, principal, serviceID)] + if row != nil && row.RefreshLeaseID == leaseID { + row.RefreshLeaseID = "" + row.RefreshLeaseUntil = nil + row.UpdatedAt = time.Now() + } + return nil +} + +func newOAuthLifecycleFixture( + t *testing.T, tokenStatus int, tokenBody map[string]any, +) (*oauthRuntime, *lockedOAuthRepo, *atomic.Int32, func()) { + t.Helper() + var requests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + switch req.URL.Path { + case "/metadata": + _ = json.NewEncoder(w).Encode(map[string]any{ + "issuer": "http://" + req.Host, + "authorization_endpoint": "http://" + req.Host + "/authorize", + "token_endpoint": "http://" + req.Host + "/token", + "response_types_supported": []string{"code"}, + "token_endpoint_auth_methods_supported": []string{"none"}, + }) + case "/token": + requests.Add(1) + require.NoError(t, req.ParseForm()) + require.Equal(t, "refresh_token", req.Form.Get("grant_type")) + require.Equal(t, "old-refresh", req.Form.Get("refresh_token")) + w.WriteHeader(tokenStatus) + _ = json.NewEncoder(w).Encode(tokenBody) + default: + http.NotFound(w, req) + } + })) + + repo := newLockedOAuthRepo() + principal := types.Principal{Type: types.PrincipalWebUser, ID: "user-1"} + row := &types.MCPOAuthToken{ + TenantID: 7, + PrincipalType: principal.Type, + PrincipalID: principal.ID, + UserID: principal.StorageID(), + ServiceID: "svc-1", + AccessToken: "old-access", + RefreshToken: "old-refresh", + TokenType: "Bearer", + ExpiresAt: time.Now().Add(-time.Minute), + UpdatedAt: time.Now().Add(-time.Hour), + } + repo.tokens[fakeOAuthKey(7, principal, "svc-1")] = row + store := newDBTokenStore(repo, 7, principal, "svc-1") + runtime := newOAuthRuntime(repo, 7, principal, "svc-1", server.URL, transport.OAuthConfig{ + ClientID: "client-1", + AuthServerMetadataURL: server.URL + "/metadata", + TokenStore: store, + HTTPClient: server.Client(), + }) + return runtime, repo, &requests, server.Close +} + +func TestOAuthRuntimeRefreshesExpiredToken(t *testing.T) { + runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusOK, map[string]any{ + "access_token": "new-access", + "refresh_token": "rotated-refresh", + "token_type": "Bearer", + "expires_in": 3600, + }) + defer closeServer() + + require.NoError(t, runtime.ensureFresh(context.Background(), false, nil)) + require.EqualValues(t, 1, requests.Load()) + row, err := repo.GetTokenForPrincipal(context.Background(), 7, runtime.principal, "svc-1") + require.NoError(t, err) + require.Equal(t, "new-access", row.AccessToken) + require.Equal(t, "rotated-refresh", row.RefreshToken) + require.True(t, row.ExpiresAt.After(time.Now())) +} + +func TestOAuthRuntimeDeletesPermanentlyInvalidRefreshToken(t *testing.T) { + runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusBadRequest, map[string]any{ + "error": "invalid_grant", + "error_description": "refresh token expired", + }) + defer closeServer() + + err := runtime.ensureFresh(context.Background(), false, nil) + var reauth *OAuthReauthorizationRequiredError + require.ErrorAs(t, err, &reauth) + require.EqualValues(t, 1, requests.Load()) + row, getErr := repo.GetTokenForPrincipal(context.Background(), 7, runtime.principal, "svc-1") + require.NoError(t, getErr) + require.Nil(t, row) +} + +func TestOAuthRuntimePreservesTokenOnTemporaryRefreshFailure(t *testing.T) { + runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusServiceUnavailable, map[string]any{ + "error": "temporarily_unavailable", + }) + defer closeServer() + + err := runtime.ensureFresh(context.Background(), false, nil) + var temporary *OAuthRefreshTemporaryError + require.ErrorAs(t, err, &temporary) + require.EqualValues(t, 1, requests.Load()) + row, getErr := repo.GetTokenForPrincipal(context.Background(), 7, runtime.principal, "svc-1") + require.NoError(t, getErr) + require.NotNil(t, row) + require.Equal(t, "old-refresh", row.RefreshToken) +} + +func TestOAuthRuntimeSerializesRotatingRefreshToken(t *testing.T) { + runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusOK, map[string]any{ + "access_token": "new-access", + "refresh_token": "rotated-refresh", + "token_type": "Bearer", + "expires_in": 3600, + }) + defer closeServer() + + const callers = 12 + start := make(chan struct{}) + errs := make(chan error, callers) + var wg sync.WaitGroup + for range callers { + wg.Add(1) + go func() { + defer wg.Done() + <-start + errs <- runtime.ensureFresh(context.Background(), false, nil) + }() + } + close(start) + wg.Wait() + close(errs) + for err := range errs { + require.NoError(t, err) + } + require.EqualValues(t, 1, requests.Load(), "a rotating refresh token must be consumed once") + row, err := repo.GetTokenForPrincipal(context.Background(), 7, runtime.principal, "svc-1") + require.NoError(t, err) + require.Equal(t, "rotated-refresh", row.RefreshToken) +} + +func TestOAuthCallRefreshesAndRetriesResource401Once(t *testing.T) { + runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusOK, map[string]any{ + "access_token": "new-access", + "refresh_token": "rotated-refresh", + "token_type": "Bearer", + "expires_in": 3600, + }) + defer closeServer() + row := repo.tokens[fakeOAuthKey(7, runtime.principal, "svc-1")] + row.ExpiresAt = time.Now().Add(time.Hour) + + calls := 0 + result, err := oauthCall(context.Background(), &mcpGoClient{oauth: runtime}, func() (string, error) { + calls++ + if calls == 1 { + return "", &transport.OAuthAuthorizationRequiredError{Handler: runtime.handler} + } + return "ok", nil + }) + require.NoError(t, err) + require.Equal(t, "ok", result) + require.Equal(t, 2, calls) + require.EqualValues(t, 1, requests.Load()) +} + +func TestOAuthCallDoesNotRetryMoreThanOnce(t *testing.T) { + runtime, repo, requests, closeServer := newOAuthLifecycleFixture(t, http.StatusOK, map[string]any{ + "access_token": "new-access", + "refresh_token": "rotated-refresh", + "token_type": "Bearer", + "expires_in": 3600, + }) + defer closeServer() + row := repo.tokens[fakeOAuthKey(7, runtime.principal, "svc-1")] + row.ExpiresAt = time.Now().Add(time.Hour) + + calls := 0 + _, err := oauthCall(context.Background(), &mcpGoClient{oauth: runtime}, func() (string, error) { + calls++ + return "", &transport.OAuthAuthorizationRequiredError{Handler: runtime.handler} + }) + require.Error(t, err) + require.Equal(t, 2, calls) + require.EqualValues(t, 1, requests.Load()) +} + +func TestTokenStatusDoesNotTreatExpiredRowAsAuthorized(t *testing.T) { + expired := &types.MCPOAuthToken{ + AccessToken: "stale-access", + RefreshToken: "refresh", + ExpiresAt: time.Now().Add(-time.Minute), + } + status := tokenStatus(expired, time.Now()) + require.False(t, status.Authorized) + require.Equal(t, oauthStateRefreshable, status.State) + require.True(t, status.RefreshAvailable) + + expired.RefreshToken = "" + status = tokenStatus(expired, time.Now()) + require.False(t, status.Authorized) + require.Equal(t, oauthStateReauthNeeded, status.State) +} + +func TestOAuthRuntimeDoesNotExpireNonRefreshableTokenEarly(t *testing.T) { + repo := newLockedOAuthRepo() + principal := types.Principal{Type: types.PrincipalWebUser, ID: "user-1"} + repo.tokens[fakeOAuthKey(7, principal, "svc-1")] = &types.MCPOAuthToken{ + TenantID: 7, + PrincipalType: principal.Type, + PrincipalID: principal.ID, + ServiceID: "svc-1", + AccessToken: "access", + ExpiresAt: time.Now().Add(10 * time.Second), + } + runtime := &oauthRuntime{repo: repo, tenantID: 7, principal: principal, serviceID: "svc-1"} + require.NoError(t, runtime.ensureFresh(context.Background(), false, nil)) + row, err := repo.GetTokenForPrincipal(context.Background(), 7, principal, "svc-1") + require.NoError(t, err) + require.NotNil(t, row) +} diff --git a/internal/mcp/oauth_manager.go b/internal/mcp/oauth_manager.go index df24f067..761e5fd1 100644 --- a/internal/mcp/oauth_manager.go +++ b/internal/mcp/oauth_manager.go @@ -69,7 +69,7 @@ func (m *OAuthManager) newHandler( } // StartAuthorization performs discovery + (one-time) dynamic client -// registration, then returns the authorization URL the browser should visit. +// registration, then returns the authorization URL and an opaque attempt ID. // redirectURI is the backend callback URL registered with the auth server; // frontendRedirect is where the callback bounces the browser when finished. func (m *OAuthManager) StartAuthorization( @@ -78,29 +78,29 @@ func (m *OAuthManager) StartAuthorization( tenantID uint64, principal types.Principal, redirectURI, frontendRedirect string, -) (string, error) { +) (authorizationURL, attemptID string, err error) { if !service.AuthConfig.IsOAuth() { - return "", fmt.Errorf("MCP service %s does not use OAuth", service.ID) + return "", "", fmt.Errorf("MCP service %s does not use OAuth", service.ID) } principal = principal.Normalize() if !principal.Valid() { - return "", fmt.Errorf("principal context is required to authorize OAuth MCP service %s", service.ID) + return "", "", fmt.Errorf("principal context is required to authorize OAuth MCP service %s", service.ID) } h, err := m.newHandler(ctx, service, tenantID, principal, redirectURI) if err != nil { - return "", err + return "", "", err } // Register a client dynamically if we don't have one yet for this service. existing, _ := m.repo.GetClient(ctx, tenantID, service.ID) if existing == nil { if err := h.RegisterClient(ctx, clientRegistrationName); err != nil { - return "", fmt.Errorf("dynamic client registration failed: %w", err) + return "", "", fmt.Errorf("dynamic client registration failed: %w", err) } clientID := h.GetClientID() if clientID == "" { - return "", fmt.Errorf("dynamic client registration returned an empty client_id") + return "", "", fmt.Errorf("dynamic client registration returned an empty client_id") } if err := m.repo.SaveClient(ctx, &types.MCPOAuthClient{ TenantID: tenantID, @@ -114,17 +114,17 @@ func (m *OAuthManager) StartAuthorization( verifier, err := transport.GenerateCodeVerifier() if err != nil { - return "", fmt.Errorf("failed to generate PKCE verifier: %w", err) + return "", "", fmt.Errorf("failed to generate PKCE verifier: %w", err) } challenge := transport.GenerateCodeChallenge(verifier) state, err := transport.GenerateState() if err != nil { - return "", fmt.Errorf("failed to generate state: %w", err) + return "", "", fmt.Errorf("failed to generate state: %w", err) } authURL, err := h.GetAuthorizationURL(ctx, state, challenge) if err != nil { - return "", fmt.Errorf("failed to build authorization URL: %w", err) + return "", "", fmt.Errorf("failed to build authorization URL: %w", err) } if err := m.states.Put(ctx, state, OAuthState{ @@ -137,10 +137,10 @@ func (m *OAuthManager) StartAuthorization( RedirectURI: redirectURI, FrontendRedirect: frontendRedirect, }); err != nil { - return "", fmt.Errorf("failed to persist authorization state: %w", err) + return "", "", fmt.Errorf("failed to persist authorization state: %w", err) } - return authURL, nil + return authURL, state, nil } // StartAuthorizationForService loads the MCP service by ID and starts the @@ -160,67 +160,109 @@ func (m *OAuthManager) StartAuthorizationForService( if service == nil { return "", fmt.Errorf("MCP service not found") } - return m.StartAuthorization(ctx, service, tenantID, principal, redirectURI, frontendRedirect) + authURL, _, err := m.StartAuthorization(ctx, service, tenantID, principal, redirectURI, frontendRedirect) + return authURL, err } // CompleteAuthorization handles the provider callback: it validates state, // exchanges the code for tokens (PKCE), and persists the per-user token. -// Returns the frontend redirect URL recorded at StartAuthorization time. +// Returns the frontend redirect URL and service ID recorded at +// StartAuthorization time so the caller can recycle any cached transport that +// still carries the previous OAuth client registration. func (m *OAuthManager) CompleteAuthorization( ctx context.Context, state, code string, -) (frontendRedirect string, err error) { +) (frontendRedirect, serviceID string, err error) { ctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), oauthCallbackTimeout) defer cancel() st, err := m.states.Take(ctx, state) if err != nil { - return "", err + return "", "", err } frontendRedirect = st.FrontendRedirect + serviceID = st.ServiceID principal := st.Principal.Normalize() if !principal.Valid() && st.UserID != "" { principal = types.Principal{Type: types.PrincipalWebUser, ID: st.UserID}.Normalize() } if !principal.Valid() { - return frontendRedirect, fmt.Errorf("principal context is missing from OAuth state") + return frontendRedirect, serviceID, fmt.Errorf("principal context is missing from OAuth state") } service, err := m.serviceRepo.GetByID(ctx, st.TenantID, st.ServiceID) if err != nil { - return frontendRedirect, fmt.Errorf("failed to load MCP service: %w", err) + return frontendRedirect, serviceID, fmt.Errorf("failed to load MCP service: %w", err) } if service == nil { - return frontendRedirect, fmt.Errorf("MCP service not found") + return frontendRedirect, serviceID, fmt.Errorf("MCP service not found") } h, err := m.newHandler(ctx, service, st.TenantID, principal, st.RedirectURI) if err != nil { - return frontendRedirect, err + return frontendRedirect, serviceID, err } // Re-prime the expected state so the library's CSRF check passes after // reconstructing the handler in this separate request. h.SetExpectedState(state) if err := h.ProcessAuthorizationResponse(ctx, code, state, st.CodeVerifier); err != nil { - return frontendRedirect, fmt.Errorf("token exchange failed: %w", err) + return frontendRedirect, serviceID, fmt.Errorf("token exchange failed: %w", err) + } + if err := m.states.CompleteAttempt(ctx, state); err != nil { + return frontendRedirect, serviceID, fmt.Errorf("failed to record authorization completion: %w", err) } // ProcessAuthorizationResponse persists the token via the TokenStore. logger.GetLogger(ctx).Infof( "MCP OAuth authorized: service=%s principal=%s", st.ServiceID, principal.StorageID(), ) - return frontendRedirect, nil + return frontendRedirect, serviceID, nil } -// IsAuthorized reports whether the given principal has a stored (non-empty) token -// for the service. -func (m *OAuthManager) IsAuthorized( - ctx context.Context, tenantID uint64, principal types.Principal, serviceID string, +// IsAuthorizationAttemptComplete reports whether this exact authorization +// attempt completed for the requested principal and service. A pre-existing +// token must never satisfy a newly opened OAuth popup. +func (m *OAuthManager) IsAuthorizationAttemptComplete( + ctx context.Context, + tenantID uint64, + principal types.Principal, + serviceID, attemptID string, ) (bool, error) { - tok, err := m.repo.GetTokenForPrincipal(ctx, tenantID, principal, serviceID) + attempt, err := m.states.Attempt(ctx, attemptID) if err != nil { return false, err } - return tok != nil && tok.AccessToken != "", nil + principal = principal.Normalize() + attemptPrincipal := attempt.Principal.Normalize() + if attempt.TenantID != tenantID || attempt.ServiceID != serviceID || + attemptPrincipal.Type != principal.Type || attemptPrincipal.ID != principal.ID { + return false, fmt.Errorf("oauth authorization attempt does not match the current principal or service") + } + return attempt.Completed, nil +} + +// AuthorizationStatus reports whether the stored access token is usable now, +// or is expired but still has a refresh token that runtime use can rotate. +func (m *OAuthManager) AuthorizationStatus( + ctx context.Context, tenantID uint64, principal types.Principal, serviceID string, +) (OAuthAuthorizationStatus, error) { + tok, err := m.repo.GetTokenForPrincipal(ctx, tenantID, principal, serviceID) + if err != nil { + return OAuthAuthorizationStatus{}, err + } + return tokenStatus(tok, time.Now()), nil +} + +// IsAuthorized reports whether the given principal has an access token that is +// usable now. An expired row is not authorization success merely because its +// encrypted token columns remain non-empty. +func (m *OAuthManager) IsAuthorized( + ctx context.Context, tenantID uint64, principal types.Principal, serviceID string, +) (bool, error) { + status, err := m.AuthorizationStatus(ctx, tenantID, principal, serviceID) + if err != nil { + return false, err + } + return status.Authorized, nil } // Revoke removes the principal's stored token for the service. diff --git a/internal/mcp/oauth_principal_test.go b/internal/mcp/oauth_principal_test.go index c819dcd7..79dd74c3 100644 --- a/internal/mcp/oauth_principal_test.go +++ b/internal/mcp/oauth_principal_test.go @@ -81,6 +81,36 @@ func (r *fakeOAuthRepo) DeleteTokenForPrincipal( return nil } +func (r *fakeOAuthRepo) TryAcquireTokenRefreshLease( + _ context.Context, + tenantID uint64, + principal types.Principal, + serviceID, leaseID string, + leaseUntil time.Time, +) (bool, error) { + row := r.tokens[fakeOAuthKey(tenantID, principal, serviceID)] + if row == nil || (row.RefreshLeaseUntil != nil && row.RefreshLeaseUntil.After(time.Now())) { + return false, nil + } + row.RefreshLeaseID = leaseID + row.RefreshLeaseUntil = &leaseUntil + return true, nil +} + +func (r *fakeOAuthRepo) ReleaseTokenRefreshLease( + _ context.Context, + tenantID uint64, + principal types.Principal, + serviceID, leaseID string, +) error { + row := r.tokens[fakeOAuthKey(tenantID, principal, serviceID)] + if row != nil && row.RefreshLeaseID == leaseID { + row.RefreshLeaseID = "" + row.RefreshLeaseUntil = nil + } + return nil +} + func TestDBTokenStoreUsesPrincipal(t *testing.T) { repo := newFakeOAuthRepo() principal := types.Principal{Type: types.PrincipalAPIExternalUser, ID: "7:external-42"} @@ -108,6 +138,24 @@ func TestDBTokenStoreUsesPrincipal(t *testing.T) { require.Equal(t, expiresAt, token.ExpiresAt) } +func TestManagedTokenStoreLeavesRefreshToOAuthRuntime(t *testing.T) { + repo := newFakeOAuthRepo() + principal := types.Principal{Type: types.PrincipalWebUser, ID: "user-1"} + store := newManagedTokenStore(repo, 7, principal, "svc-1") + require.NoError(t, store.SaveToken(context.Background(), &transport.Token{ + AccessToken: "access", + RefreshToken: "refresh", + ExpiresAt: time.Now().Add(-time.Minute), + })) + + token, err := store.GetToken(context.Background()) + require.NoError(t, err) + require.True(t, token.ExpiresAt.IsZero(), "mcp-go must not race WeKnora's coordinated refresh") + row, err := repo.GetTokenForPrincipal(context.Background(), 7, principal, "svc-1") + require.NoError(t, err) + require.False(t, row.ExpiresAt.IsZero(), "the database must retain the real expiry for preflight checks") +} + func TestOAuthCacheKeyUsesPrincipalForOAuthServices(t *testing.T) { service := &types.MCPService{ ID: "svc-1", diff --git a/internal/mcp/oauth_state.go b/internal/mcp/oauth_state.go index 76ea7369..03ab6b02 100644 --- a/internal/mcp/oauth_state.go +++ b/internal/mcp/oauth_state.go @@ -41,8 +41,9 @@ type OAuthState struct { type oauthStateStore struct { rdb *redis.Client - mu sync.Mutex - mem map[string]memStateEntry + mu sync.Mutex + mem map[string]memStateEntry + attempts map[string]memAttemptEntry } type memStateEntry struct { @@ -50,8 +51,28 @@ type memStateEntry struct { expiresAt time.Time } +// OAuthAttempt is the non-secret, authenticated status of one authorization +// flow. It is kept separately from OAuthState because Take consumes the PKCE +// state before the token exchange finishes, while the opener still needs to +// distinguish this callback from an older token stored for the same service. +type OAuthAttempt struct { + TenantID uint64 `json:"tenant_id"` + Principal types.Principal `json:"principal"` + ServiceID string `json:"service_id"` + Completed bool `json:"completed"` +} + +type memAttemptEntry struct { + value OAuthAttempt + expiresAt time.Time +} + func newOAuthStateStore(rdb *redis.Client) *oauthStateStore { - s := &oauthStateStore{rdb: rdb, mem: make(map[string]memStateEntry)} + s := &oauthStateStore{ + rdb: rdb, + mem: make(map[string]memStateEntry), + attempts: make(map[string]memAttemptEntry), + } if rdb == nil { go s.gcLoop() } @@ -66,21 +87,101 @@ func (s *oauthStateStore) key(state string) string { return "weknora:mcp_oauth_state:" + state } +func (s *oauthStateStore) attemptKey(state string) string { + return s.key(state) + ":attempt" +} + // Put stores a state with a fixed TTL. func (s *oauthStateStore) Put(ctx context.Context, state string, value OAuthState) error { + attempt := OAuthAttempt{ + TenantID: value.TenantID, + Principal: value.Principal.Normalize(), + ServiceID: value.ServiceID, + } if s.rdb != nil { data, err := json.Marshal(value) if err != nil { return err } - return s.rdb.Set(ctx, s.key(state), data, oauthStateTTL).Err() + attemptData, err := json.Marshal(attempt) + if err != nil { + return err + } + pipe := s.rdb.TxPipeline() + pipe.Set(ctx, s.key(state), data, oauthStateTTL) + pipe.Set(ctx, s.attemptKey(state), attemptData, oauthStateTTL) + _, err = pipe.Exec(ctx) + return err } s.mu.Lock() defer s.mu.Unlock() - s.mem[state] = memStateEntry{value: value, expiresAt: time.Now().Add(oauthStateTTL)} + expiresAt := time.Now().Add(oauthStateTTL) + s.mem[state] = memStateEntry{value: value, expiresAt: expiresAt} + s.attempts[state] = memAttemptEntry{value: attempt, expiresAt: expiresAt} return nil } +// CompleteAttempt marks an authorization attempt complete only after the code +// exchange has successfully persisted a token. +func (s *oauthStateStore) CompleteAttempt(ctx context.Context, state string) error { + if s.rdb != nil { + data, err := s.rdb.Get(ctx, s.attemptKey(state)).Bytes() + if err != nil { + if err == redis.Nil { + return fmt.Errorf("oauth attempt not found or expired") + } + return err + } + var attempt OAuthAttempt + if err := json.Unmarshal(data, &attempt); err != nil { + return err + } + attempt.Completed = true + data, err = json.Marshal(attempt) + if err != nil { + return err + } + return s.rdb.Set(ctx, s.attemptKey(state), data, oauthStateTTL).Err() + } + s.mu.Lock() + defer s.mu.Unlock() + entry, ok := s.attempts[state] + if !ok || time.Now().After(entry.expiresAt) { + delete(s.attempts, state) + return fmt.Errorf("oauth attempt not found or expired") + } + entry.value.Completed = true + entry.expiresAt = time.Now().Add(oauthStateTTL) + s.attempts[state] = entry + return nil +} + +// Attempt returns the status record for one authorization flow. +func (s *oauthStateStore) Attempt(ctx context.Context, state string) (OAuthAttempt, error) { + if s.rdb != nil { + data, err := s.rdb.Get(ctx, s.attemptKey(state)).Bytes() + if err != nil { + if err == redis.Nil { + return OAuthAttempt{}, fmt.Errorf("oauth attempt not found or expired") + } + return OAuthAttempt{}, err + } + var attempt OAuthAttempt + if err := json.Unmarshal(data, &attempt); err != nil { + return OAuthAttempt{}, err + } + return attempt, nil + } + s.mu.Lock() + defer s.mu.Unlock() + entry, ok := s.attempts[state] + if !ok || time.Now().After(entry.expiresAt) { + delete(s.attempts, state) + return OAuthAttempt{}, fmt.Errorf("oauth attempt not found or expired") + } + return entry.value, nil +} + // Take retrieves and deletes a state (single-use). Returns an error if the // state is unknown or expired. func (s *oauthStateStore) Take(ctx context.Context, state string) (OAuthState, error) { @@ -122,6 +223,11 @@ func (s *oauthStateStore) gcLoop() { delete(s.mem, k) } } + for k, v := range s.attempts { + if now.After(v.expiresAt) { + delete(s.attempts, k) + } + } s.mu.Unlock() } } diff --git a/internal/mcp/oauth_state_test.go b/internal/mcp/oauth_state_test.go new file mode 100644 index 00000000..b7648b2c --- /dev/null +++ b/internal/mcp/oauth_state_test.go @@ -0,0 +1,71 @@ +package mcp + +import ( + "context" + "testing" + + "github.com/Tencent/WeKnora/internal/types" + "github.com/stretchr/testify/require" +) + +func TestOAuthAttemptCompletesOnlyAfterCallbackExchange(t *testing.T) { + ctx := context.Background() + store := newOAuthStateStore(nil) + principal := types.Principal{Type: types.PrincipalWebUser, ID: "user-1"} + state := OAuthState{ + TenantID: 7, + Principal: principal, + ServiceID: "service-1", + } + + require.NoError(t, store.Put(ctx, "attempt-1", state)) + + attempt, err := store.Attempt(ctx, "attempt-1") + require.NoError(t, err) + require.False(t, attempt.Completed) + + // Consuming the OAuth state means the provider callback started. It must + // not look successful until the code exchange and token persistence finish. + _, err = store.Take(ctx, "attempt-1") + require.NoError(t, err) + attempt, err = store.Attempt(ctx, "attempt-1") + require.NoError(t, err) + require.False(t, attempt.Completed) + + require.NoError(t, store.CompleteAttempt(ctx, "attempt-1")) + attempt, err = store.Attempt(ctx, "attempt-1") + require.NoError(t, err) + require.True(t, attempt.Completed) + require.Equal(t, uint64(7), attempt.TenantID) + require.Equal(t, principal.Normalize(), attempt.Principal) + require.Equal(t, "service-1", attempt.ServiceID) +} + +func TestAuthorizationAttemptStatusIsScopedToPrincipalAndService(t *testing.T) { + ctx := context.Background() + manager := &OAuthManager{states: newOAuthStateStore(nil)} + principal := types.Principal{Type: types.PrincipalWebUser, ID: "user-1"} + + require.NoError(t, manager.states.Put(ctx, "attempt-1", OAuthState{ + TenantID: 7, + Principal: principal, + ServiceID: "service-1", + })) + require.NoError(t, manager.states.CompleteAttempt(ctx, "attempt-1")) + + completed, err := manager.IsAuthorizationAttemptComplete( + ctx, 7, principal, "service-1", "attempt-1", + ) + require.NoError(t, err) + require.True(t, completed) + + _, err = manager.IsAuthorizationAttemptComplete( + ctx, 7, types.Principal{Type: types.PrincipalWebUser, ID: "user-2"}, "service-1", "attempt-1", + ) + require.ErrorContains(t, err, "does not match") + + _, err = manager.IsAuthorizationAttemptComplete( + ctx, 7, principal, "service-2", "attempt-1", + ) + require.ErrorContains(t, err, "does not match") +} diff --git a/internal/mcp/oauth_tokenstore.go b/internal/mcp/oauth_tokenstore.go index 5fbece72..f03df0f7 100644 --- a/internal/mcp/oauth_tokenstore.go +++ b/internal/mcp/oauth_tokenstore.go @@ -2,6 +2,7 @@ package mcp import ( "context" + "fmt" "time" "github.com/Tencent/WeKnora/internal/types" @@ -11,9 +12,9 @@ import ( // dbTokenStore is a transport.TokenStore backed by the MCPOAuthRepository, // scoped to a single (tenant, principal, service) tuple. The mcp-go OAuth handler -// calls GetToken before each request (refreshing via refresh_token when -// expired) and SaveToken after a successful authorization or refresh, so this -// store transparently persists refreshed tokens back to the database. +// calls SaveToken after a successful authorization or refresh. Runtime MCP +// transports receive the managedTokenStore wrapper below so refresh decisions +// stay in WeKnora's coordinated lifecycle instead of the dependency. type dbTokenStore struct { repo interfaces.MCPOAuthRepository tenantID uint64 @@ -21,6 +22,30 @@ type dbTokenStore struct { serviceID string } +// managedTokenStore hides local expiry from mcp-go transports. WeKnora checks +// the persisted ExpiresAt before every operation and performs the coordinated +// refresh itself; allowing the dependency to also auto-refresh would bypass +// the cross-instance lease and collapse refresh failures into a generic +// authorization-required error. +type managedTokenStore struct { + *dbTokenStore +} + +func newManagedTokenStore( + repo interfaces.MCPOAuthRepository, tenantID uint64, principal types.Principal, serviceID string, +) *managedTokenStore { + return &managedTokenStore{dbTokenStore: newDBTokenStore(repo, tenantID, principal, serviceID)} +} + +func (s *managedTokenStore) GetToken(ctx context.Context) (*transport.Token, error) { + token, err := s.dbTokenStore.GetToken(ctx) + if err != nil { + return nil, err + } + token.ExpiresAt = time.Time{} + return token, nil +} + // newDBTokenStore creates a per-principal, per-service token store. func newDBTokenStore( repo interfaces.MCPOAuthRepository, tenantID uint64, principal types.Principal, serviceID string, @@ -59,6 +84,12 @@ func (s *dbTokenStore) SaveToken(ctx context.Context, token *transport.Token) er if err := ctx.Err(); err != nil { return err } + if token == nil || token.AccessToken == "" { + return fmt.Errorf("OAuth token response did not contain an access_token") + } + if token.TokenType == "" { + token.TokenType = "Bearer" + } expiresAt := token.ExpiresAt if expiresAt.IsZero() && token.ExpiresIn > 0 { expiresAt = time.Now().Add(time.Duration(token.ExpiresIn) * time.Second) diff --git a/internal/middleware/logger.go b/internal/middleware/logger.go index 22eb3fe1..ad7c5b92 100644 --- a/internal/middleware/logger.go +++ b/internal/middleware/logger.go @@ -4,6 +4,7 @@ import ( "bytes" "context" "io" + "net/url" "regexp" "strings" "time" @@ -44,7 +45,8 @@ func (r loggerResponseBodyWriter) Write(b []byte) (int, error) { var sensitiveFieldRegex = regexp.MustCompile( `(?i)("(?:new[_-]?password|old[_-]?password|password|passwd|token|access[_-]?token|` + `refresh[_-]?token|id[_-]?token|authorization|auth[_-]?token|api[_-]?key|` + - `api[_-]?secret|secret[_-]?key|client[_-]?secret|private[_-]?key|secret)")\s*:\s*"[^"]*"`, + `api[_-]?secret|secret[_-]?key|client[_-]?secret|private[_-]?key|secret|` + + `authorization[_-]?url|authorization[_-]?attempt)")\s*:\s*"[^"]*"`, ) // sanitizeBody 清理敏感信息 @@ -52,6 +54,32 @@ func sanitizeBody(body string) string { return sensitiveFieldRegex.ReplaceAllString(body, `$1:"***"`) } +var sensitiveQueryFields = map[string]struct{}{ + "access_token": {}, + "authorization_attempt": {}, + "code": {}, + "id_token": {}, + "refresh_token": {}, + "state": {}, + "token": {}, +} + +// sanitizeQuery prevents OAuth authorization codes and CSRF/attempt state from +// being copied into access logs. Parsing the query also covers repeated and +// percent-encoded parameters without relying on fragile string replacement. +func sanitizeQuery(raw string) string { + values, err := url.ParseQuery(raw) + if err != nil { + return "[invalid query omitted]" + } + for key := range values { + if _, sensitive := sensitiveQueryFields[strings.ToLower(key)]; sensitive { + values[key] = []string{"***"} + } + } + return values.Encode() +} + // readRequestBody 读取请求体（限制大小用于日志，但完整读取用于重置） func readRequestBody(c *gin.Context) string { if c.Request.Body == nil { @@ -172,7 +200,7 @@ func Logger() gin.HandlerFunc { method := c.Request.Method if raw != "" { - path = path + "?" + raw + path = path + "?" + sanitizeQuery(raw) } // 读取响应体 diff --git a/internal/middleware/logger_test.go b/internal/middleware/logger_test.go index caa1069a..e27ba295 100644 --- a/internal/middleware/logger_test.go +++ b/internal/middleware/logger_test.go @@ -53,6 +53,11 @@ func TestSanitizeBody(t *testing.T) { in: `{"baseUrl":"https://example.com","modelName":"gpt"}`, want: `{"baseUrl":"https://example.com","modelName":"gpt"}`, }, + { + name: "OAuth authorization response fields", + in: `{"authorization_url":"https://idp.example/authorize?state=secret","authorization_attempt":"secret-state"}`, + want: `{"authorization_url":"***","authorization_attempt":"***"}`, + }, } for _, tc := range cases { @@ -64,3 +69,11 @@ func TestSanitizeBody(t *testing.T) { }) } } + +func TestSanitizeQuery(t *testing.T) { + got := sanitizeQuery("code=secret-code&state=secret-state&next=%2Fsettings&state=second") + want := "code=%2A%2A%2A&next=%2Fsettings&state=%2A%2A%2A" + if got != want { + t.Fatalf("sanitizeQuery() = %q, want %q", got, want) + } +} diff --git a/internal/types/interfaces/mcp_oauth.go b/internal/types/interfaces/mcp_oauth.go index ad91bff1..8160390e 100644 --- a/internal/types/interfaces/mcp_oauth.go +++ b/internal/types/interfaces/mcp_oauth.go @@ -2,6 +2,7 @@ package interfaces import ( "context" + "time" "github.com/Tencent/WeKnora/internal/types" ) @@ -41,4 +42,23 @@ type MCPOAuthRepository interface { // DeleteTokenForPrincipal removes the per-principal token for a service. DeleteTokenForPrincipal(ctx context.Context, tenantID uint64, principal types.Principal, serviceID string) error + + // TryAcquireTokenRefreshLease atomically claims the right to use a refresh + // token. This prevents concurrent instances from presenting the same + // rotating refresh token to the authorization server. + TryAcquireTokenRefreshLease( + ctx context.Context, + tenantID uint64, + principal types.Principal, + serviceID, leaseID string, + leaseUntil time.Time, + ) (bool, error) + + // ReleaseTokenRefreshLease releases a lease only when leaseID still owns it. + ReleaseTokenRefreshLease( + ctx context.Context, + tenantID uint64, + principal types.Principal, + serviceID, leaseID string, + ) error } diff --git a/internal/types/mcp_oauth.go b/internal/types/mcp_oauth.go index c49bc983..696b9134 100644 --- a/internal/types/mcp_oauth.go +++ b/internal/types/mcp_oauth.go @@ -85,8 +85,13 @@ type MCPOAuthToken struct { RefreshToken string `json:"-" gorm:"type:text"` TokenType string `json:"token_type" gorm:"type:varchar(32)"` ExpiresAt time.Time `json:"expires_at"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` + // RefreshLeaseID / RefreshLeaseUntil coordinate refresh-token rotation + // across application instances. They are operational fields only and are + // never exposed through the API. + RefreshLeaseID string `json:"-" gorm:"type:varchar(36)"` + RefreshLeaseUntil *time.Time `json:"-"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` } // TableName pins the table name (see MCPOAuthClient.TableName); the default diff --git a/migrations/sqlite/000000_init.up.sql b/migrations/sqlite/000000_init.up.sql index 24f4b687..80c82d0d 100644 --- a/migrations/sqlite/000000_init.up.sql +++ b/migrations/sqlite/000000_init.up.sql @@ -351,6 +351,8 @@ CREATE TABLE IF NOT EXISTS audit_logs ( request_method VARCHAR(16) NOT NULL DEFAULT '', outcome VARCHAR(16) NOT NULL DEFAULT 'success', details TEXT NOT NULL DEFAULT '{}', + scope_type VARCHAR(32) NOT NULL DEFAULT '', + scope_id VARCHAR(64) NOT NULL DEFAULT '', created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP ); CREATE INDEX IF NOT EXISTS idx_audit_logs_tenant_id_desc @@ -361,6 +363,8 @@ CREATE INDEX IF NOT EXISTS idx_audit_logs_tenant_action ON audit_logs(tenant_id, action); CREATE INDEX IF NOT EXISTS idx_audit_logs_created_at ON audit_logs(created_at); +CREATE INDEX IF NOT EXISTS idx_audit_logs_tenant_scope_desc + ON audit_logs(tenant_id, scope_type, scope_id, id DESC); -- user_resource_favorites — sqlite mirror of migration 000047. Same -- composite PK (user_id, tenant_id, resource_type, resource_id) so the @@ -497,6 +501,8 @@ CREATE TABLE IF NOT EXISTS mcp_oauth_tokens ( refresh_token TEXT, token_type VARCHAR(32), expires_at DATETIME, + refresh_lease_id VARCHAR(36), + refresh_lease_until DATETIME, created_at DATETIME DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (service_id) REFERENCES mcp_services(id) ON DELETE CASCADE diff --git a/migrations/sqlite/000001_kb_activity_scope.down.sql b/migrations/sqlite/000001_kb_activity_scope.down.sql deleted file mode 100644 index 545d925c..00000000 --- a/migrations/sqlite/000001_kb_activity_scope.down.sql +++ /dev/null @@ -1,4 +0,0 @@ -DROP INDEX IF EXISTS idx_audit_logs_tenant_scope_desc; - -ALTER TABLE audit_logs DROP COLUMN scope_id; -ALTER TABLE audit_logs DROP COLUMN scope_type; diff --git a/migrations/sqlite/000001_kb_activity_scope.up.sql b/migrations/sqlite/000001_kb_activity_scope.up.sql deleted file mode 100644 index 7b2ab303..00000000 --- a/migrations/sqlite/000001_kb_activity_scope.up.sql +++ /dev/null @@ -1,5 +0,0 @@ -ALTER TABLE audit_logs ADD COLUMN scope_type VARCHAR(32) NOT NULL DEFAULT ''; -ALTER TABLE audit_logs ADD COLUMN scope_id VARCHAR(64) NOT NULL DEFAULT ''; - -CREATE INDEX IF NOT EXISTS idx_audit_logs_tenant_scope_desc - ON audit_logs(tenant_id, scope_type, scope_id, id DESC); diff --git a/migrations/versioned/000074_mcp_oauth_refresh_lease.down.sql b/migrations/versioned/000074_mcp_oauth_refresh_lease.down.sql new file mode 100644 index 00000000..8eb0443b --- /dev/null +++ b/migrations/versioned/000074_mcp_oauth_refresh_lease.down.sql @@ -0,0 +1,3 @@ +ALTER TABLE mcp_oauth_tokens + DROP COLUMN IF EXISTS refresh_lease_until, + DROP COLUMN IF EXISTS refresh_lease_id; diff --git a/migrations/versioned/000074_mcp_oauth_refresh_lease.up.sql b/migrations/versioned/000074_mcp_oauth_refresh_lease.up.sql new file mode 100644 index 00000000..bbc23c05 --- /dev/null +++ b/migrations/versioned/000074_mcp_oauth_refresh_lease.up.sql @@ -0,0 +1,7 @@ +DO $$ BEGIN RAISE NOTICE '[Migration 000074] Adding MCP OAuth refresh lease...'; END $$; + +ALTER TABLE mcp_oauth_tokens + ADD COLUMN IF NOT EXISTS refresh_lease_id VARCHAR(36), + ADD COLUMN IF NOT EXISTS refresh_lease_until TIMESTAMP WITH TIME ZONE; + +DO $$ BEGIN RAISE NOTICE '[Migration 000074] MCP OAuth refresh lease ready'; END $$;