Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 62 additions & 9 deletions apps/sim/app/api/mcp/oauth/callback/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,46 @@ const logger = createLogger('McpOauthCallbackAPI')

export const dynamic = 'force-dynamic'

class OauthCallbackStepTimeout extends Error {
constructor(step: string, ms: number) {
super(`MCP OAuth callback step "${step}" did not settle within ${ms}ms`)
this.name = 'OauthCallbackStepTimeout'
}
}

/**
* Times and bounds one awaited step of the callback so a stalled operation
* surfaces as a labeled, logged error instead of hanging the request forever.
* The losing promise is not cancelled (a wedged DB/socket op can't be), so it
* settles in the background with its rejection swallowed; the point is that the
* request stops waiting on it and the logs name the exact step that stalled.
*/
async function timedStep<T>(step: string, ms: number, fn: () => Promise<T>): Promise<T> {
const start = Date.now()
logger.info(`OAuth callback step start: ${step}`)
const work = Promise.resolve(fn())
work.catch(() => {})
Comment thread
waleedlatif1 marked this conversation as resolved.
let timer: ReturnType<typeof setTimeout> | undefined
try {
const value = await Promise.race([
work,
new Promise<never>((_, reject) => {
timer = setTimeout(() => reject(new OauthCallbackStepTimeout(step, ms)), ms)
timer.unref?.()
}),
])
logger.info(`OAuth callback step done: ${step} (${Date.now() - start}ms)`)
return value
} catch (error) {
logger.error(`OAuth callback step failed: ${step} (${Date.now() - start}ms)`, {
error: toError(error).message,
})
throw error
} finally {
clearTimeout(timer)
}
}

function escapeHtml(value: string): string {
return value
.replace(/&/g, '&amp;')
Expand Down Expand Up @@ -145,8 +185,9 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
serverId
)
}
const serverUrl = server.url
try {
assertSafeOauthServerUrl(server.url)
assertSafeOauthServerUrl(serverUrl)
} catch {
return respond(
'MCP OAuth requires https (or http://localhost for development).',
Expand All @@ -157,16 +198,22 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
}

// Burn state before token exchange so a replayed callback cannot reuse it.
await clearState(row.id, 'callback:burn-before-exchange')
await timedStep('clearState(burn)', 10_000, () =>
clearState(row.id, 'callback:burn-before-exchange')
)

const preregistered = await loadPreregisteredClient(server.id)
const preregistered = await timedStep('loadPreregisteredClient', 15_000, () =>
loadPreregisteredClient(server.id)
)
const provider = new SimMcpOauthProvider({ row, preregistered })
let result: Awaited<ReturnType<typeof mcpAuthGuarded>>
try {
result = await mcpAuthGuarded(provider, {
serverUrl: server.url,
authorizationCode: code,
})
result = await timedStep('mcpAuthGuarded', 120_000, () =>
mcpAuthGuarded(provider, {
serverUrl,
authorizationCode: code,
})
)
Comment thread
waleedlatif1 marked this conversation as resolved.
} catch (e) {
logger.error('Token exchange failed during MCP OAuth callback', e)
return respond(
Expand All @@ -176,7 +223,11 @@ export const GET = withRouteHandler(async (request: NextRequest) => {
server.id
)
} finally {
await clearVerifier(row.id)
await timedStep('clearVerifier', 10_000, () => clearVerifier(row.id)).catch((e) =>
logger.error('Failed to clear PKCE verifier after MCP OAuth callback', {
error: toError(e).message,
})
)
}

if (result !== 'AUTHORIZED') {
Expand All @@ -185,7 +236,9 @@ export const GET = withRouteHandler(async (request: NextRequest) => {

try {
// forceRefresh: skip any stale cache from before re-auth.
await mcpService.discoverServerTools(session.user.id, server.id, server.workspaceId, true)
await timedStep('discoverServerTools', 60_000, () =>
mcpService.discoverServerTools(session.user.id, server.id, server.workspaceId, true)
)
} catch (e) {
logger.warn('Post-auth tools refresh failed', toError(e).message)
}
Expand Down
4 changes: 4 additions & 0 deletions apps/sim/lib/mcp/oauth/storage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -170,18 +170,22 @@ export async function saveClientInformation(
info: OAuthClientInformationMixed
): Promise<void> {
const encrypted = await encryptClientInformation(info)
logger.info('Persisting MCP OAuth client information', { rowId })
await db
.update(mcpServerOauth)
.set({ clientInformation: encrypted, updatedAt: new Date() })
.where(eq(mcpServerOauth.id, rowId))
logger.info('Persisted MCP OAuth client information', { rowId })
}

export async function saveTokens(rowId: string, tokens: OAuthTokens): Promise<void> {
const encrypted = await encryptTokens(tokens)
logger.info('Persisting MCP OAuth tokens', { rowId })
await db
.update(mcpServerOauth)
.set({ tokens: encrypted, lastRefreshedAt: new Date(), updatedAt: new Date() })
.where(eq(mcpServerOauth.id, rowId))
logger.info('Persisted MCP OAuth tokens', { rowId })
}

export async function saveCodeVerifier(rowId: string, verifier: string): Promise<void> {
Expand Down
23 changes: 21 additions & 2 deletions apps/sim/lib/mcp/pinned-fetch.ts
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
import type { FetchLike } from '@modelcontextprotocol/sdk/shared/transport.js'
import { createLogger } from '@sim/logger'
import type { Agent } from 'undici'
import { createPinnedFetchWithDispatcher } from '@/lib/core/security/input-validation.server'
import { validateMcpServerSsrf } from '@/lib/mcp/domain-check'
import { McpError } from '@/lib/mcp/types'

const logger = createLogger('McpOauthFetch')

/** Pinned fetch for the live MCP transport, plus a handle to release its sockets. */
export interface PinnedMcpFetch {
fetch: typeof fetch
Expand Down Expand Up @@ -151,13 +154,17 @@ function releaseStreamOnSettle(
export function createSsrfGuardedMcpFetch(timeoutMs: number = OAUTH_FETCH_TIMEOUT_MS): FetchLike {
return (async (url, init) => {
const target = typeof url === 'string' ? url : url.href
const host = URL.canParse(target) ? new URL(target).host : target
const startedAt = Date.now()
const timeoutSignal = AbortSignal.timeout(timeoutMs)
// Bound every phase — validation, request, body read — by the deadline + caller signal.
const signal = init?.signal ? AbortSignal.any([init.signal, timeoutSignal]) : timeoutSignal
// Per-request Agent must be torn down (finally): a one-shot leg never reuses its socket.
let dispatcher: Agent | undefined
try {
logger.info('OAuth guarded fetch: validating', { host })
const resolvedIP = await withDeadline(validateMcpServerSsrf(target), signal)
logger.info('OAuth guarded fetch: requesting', { host, pinned: Boolean(resolvedIP) })
let response: Response
if (resolvedIP) {
const pinned = createPinnedFetchWithDispatcher(resolvedIP, {
Expand All @@ -175,11 +182,23 @@ export function createSsrfGuardedMcpFetch(timeoutMs: number = OAUTH_FETCH_TIMEOU
if (contentType.includes('text/event-stream')) {
const streamed = releaseStreamOnSettle(response, dispatcher, signal)
dispatcher = undefined // teardown ownership moved to releaseStreamOnSettle
logger.info('OAuth guarded fetch: streaming response', {
host,
status: response.status,
ms: Date.now() - startedAt,
})
return streamed
}
return await bufferUnderDeadline(response, signal)
logger.info('OAuth guarded fetch: reading body', { host, status: response.status })
const buffered = await bufferUnderDeadline(response, signal)
logger.info('OAuth guarded fetch: done', {
host,
status: response.status,
ms: Date.now() - startedAt,
})
return buffered
} catch (error) {
const host = URL.canParse(target) ? new URL(target).host : target
logger.warn('OAuth guarded fetch: failed', { host, ms: Date.now() - startedAt })
// Relabel only our own deadline — by reason identity, not signal state (which may
// abort independently just after the deadline).
if (timeoutSignal.aborted && error === timeoutSignal.reason) {
Expand Down
Loading