AgentServer/apps/server/test/mcp-oauth.test.ts

// MCP 支持完善 · 阶段 D：OAuth 授权（按服务共享，docs/design/16-mcp-support.md §5）
import crypto from 'node:crypto'
import fs from 'node:fs'
import http from 'node:http'
import type { AddressInfo } from 'node:net'
import path from 'node:path'
import pino from 'pino'
import { afterAll, beforeAll, describe, expect, it } from 'vitest'
import { Client } from '@modelcontextprotocol/sdk/client/index.js'
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'
import { createApp, type App } from '../src/app'
import { config } from '../src/config'
import { buildServer } from '../src/http/server'
import type { GatewayRunContext } from '../src/mcp-gateway/tokens'
import { resolveAccess } from '../src/services/access-policy'

let app: App
let server: Awaited<ReturnType<typeof buildServer>>
let base: string
let upstream: http.Server
let origin: string

/** 测试用的授权服务器 + 受保护的 MCP 服务（同一个端口） */
const valid = new Set<string>()
const codes = new Map<string, { challenge: string; redirectUri: string }>()
const refreshTokens = new Set<string>()
const grants: string[] = []
function startUpstream(): Promise<void> {
  upstream = http.createServer(async (req, res) => {
    const url = new URL(req.url!, origin)
    let body = ''
    req.on('data', c => (body += c))
    await new Promise(r => req.on('end', r))
    const json = (status: number, v: unknown) => res.writeHead(status, { 'content-type': 'application/json' }).end(JSON.stringify(v))
    if (url.pathname.startsWith('/.well-known/oauth-protected-resource')) return json(200, { resource: `${origin}/mcp`, authorization_servers: [origin] })
    if (url.pathname === '/.well-known/oauth-authorization-server') {
      return json(200, {
        issuer: origin, authorization_endpoint: `${origin}/authorize`, token_endpoint: `${origin}/token`, registration_endpoint: `${origin}/register`,
        response_types_supported: ['code'], grant_types_supported: ['authorization_code', 'refresh_token'], code_challenge_methods_supported: ['S256'],
        token_endpoint_auth_methods_supported: ['none', 'client_secret_post'],
      })
    }
    if (url.pathname === '/register') return json(201, { ...JSON.parse(body), client_id: 'dyn-client', client_id_issued_at: Math.floor(Date.now() / 1000) })
    if (url.pathname === '/authorize') {
      // 模拟用户登录并同意：直接带授权码跳回
      const code = crypto.randomBytes(8).toString('hex')
      codes.set(code, { challenge: url.searchParams.get('code_challenge')!, redirectUri: url.searchParams.get('redirect_uri')! })
      const back = new URL(url.searchParams.get('redirect_uri')!)
      back.searchParams.set('code', code)
      back.searchParams.set('state', url.searchParams.get('state')!)
      return res.writeHead(302, { location: back.toString() }).end()
    }
    if (url.pathname === '/token') {
      const p = new URLSearchParams(body)
      grants.push(p.get('grant_type')!)
      if (p.get('grant_type') === 'authorization_code') {
        const c = codes.get(p.get('code')!)
        const challenge = crypto.createHash('sha256').update(p.get('code_verifier') ?? '').digest('base64url')
        if (!c || c.challenge !== challenge || c.redirectUri !== p.get('redirect_uri')) return json(400, { error: 'invalid_grant' })
        codes.delete(p.get('code')!)
      } else if (p.get('grant_type') === 'refresh_token') {
        if (!refreshTokens.has(p.get('refresh_token')!)) return json(400, { error: 'invalid_grant' })
      } else return json(400, { error: 'unsupported_grant_type' })
      const access = crypto.randomBytes(8).toString('hex')
      const refresh = crypto.randomBytes(8).toString('hex')
      valid.add(access)
      refreshTokens.add(refresh)
      return json(200, { access_token: access, token_type: 'Bearer', expires_in: 3600, refresh_token: refresh })
    }
    if (url.pathname === '/mcp') {
      const token = req.headers.authorization?.replace(/^Bearer /, '')
      if (!token || !valid.has(token)) {
        return res.writeHead(401, { 'www-authenticate': `Bearer resource_metadata="${origin}/.well-known/oauth-protected-resource/mcp"` }).end()
      }
      const mcp = new McpServer({ name: 'secure', version: '1.0.0' })
      mcp.registerTool('whoami', { description: '当前账号', inputSchema: {}, annotations: { readOnlyHint: true } }, async () => ({ content: [{ type: 'text', text: 'account:alice' }] }))
      const t = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined })
      res.on('close', () => { void t.close(); void mcp.close() })
      await mcp.connect(t)
      return t.handleRequest(req, res, body ? JSON.parse(body) : undefined)
    }
    res.writeHead(404).end()
  })
  return new Promise(r => upstream.listen(0, '127.0.0.1', () => {
    origin = `http://127.0.0.1:${(upstream.address() as AddressInfo).port}`
    r()
  }))
}

function runCtx(serverId: string): GatewayRunContext {
  const ws = path.join(config.dataDir, 'workspaces', `ses_oauth_${Math.random().toString(36).slice(2)}`)
  fs.mkdirSync(ws, { recursive: true })
  const settings = app.settings.get()
  return {
    sessionId: `ses_${Math.random().toString(36).slice(2)}`, runId: 'r', provider: 'claude', workspaceDir: ws, permissionMode: 'auto', settings,
    access: resolveAccess({ policy: { mode: 'restricted', read: [], write: [], deny: [] }, knowledgeRoots: [], selectedKnowledgeRootIds: [], workspaceDir: ws, settings }),
    knowledgeRoots: [], mcpServerIds: [serverId], audit: () => {}, interact: async () => ({ resolution: 'allowed', scope: 'once' }),
  }
}
async function callWhoami(serverId: string) {
  const client = new Client({ name: 'test', version: '1' })
  await client.connect(new StreamableHTTPClientTransport(new URL(`${base}/mcp`), { requestInit: { headers: { Authorization: `Bearer ${app.gatewayTokens.issue(runCtx(serverId))}` } } }))
  try {
    return await client.callTool({ name: 'secure__whoami', arguments: {} }) as { content: { text: string }[]; isError?: boolean }
  } finally {
    await client.close()
  }
}

beforeAll(async () => {
  const log = pino({ level: 'silent' })
  app = createApp(log, { providers: {} })
  server = await buildServer(app, { logger: log })
  base = await server.listen({ port: 0, host: '127.0.0.1' })
  await startUpstream()
})
afterAll(async () => {
  await app.mcpUpstream.closeAll()
  await server.close()
  upstream.closeAllConnections()
  upstream.close()
})

describe('OAuth 授权', () => {
  it('未授权时连接给出明确提示；发起授权（动态注册 + PKCE）→ 回调换取令牌 → 工具可用；令牌失效后自动续期；可以清除', async () => {
    const s = app.mcpServers.create({
      name: 'secure', title: '受保护服务', description: '', transport: 'http', enabled: true, url: `${origin}/mcp`, headers: [],
      oauth: { enabled: true, clientId: '', clientSecret: null, scope: 'read' }, command: '', args: [], env: [], allowNetwork: false, readPaths: [], writePaths: [],
    })
    expect(s.oauthStatus).toMatchObject({ authorized: false, registered: false })
    const refreshed = await app.mcpServers.refresh(s.id)
    expect(refreshed.lastError).toMatch(/需要 OAuth 授权/)

    const { authorizationUrl } = await app.mcpOAuth.start(s.id, base)
    const auth = new URL(authorizationUrl!)
    expect(auth.origin).toBe(origin)
    expect(auth.searchParams.get('client_id')).toBe('dyn-client')
    expect(auth.searchParams.get('redirect_uri')).toBe(`${base}/admin/v1/mcp/oauth/callback`)
    expect(auth.searchParams.get('code_challenge_method')).toBe('S256')
    expect(auth.searchParams.get('scope')).toBe('read')

    // 浏览器打开授权网址，授权服务器跳回回调地址（不带后台登录）
    const redirect = await fetch(authorizationUrl!, { redirect: 'manual' })
    const back = new URL(redirect.headers.get('location')!)
    const cb = await server.inject({ method: 'GET', url: `${back.pathname}${back.search}` })
    expect(cb.statusCode).toBe(200)
    expect(cb.body).toContain('授权成功')
    // state 只能用一次
    expect((await server.inject({ method: 'GET', url: `${back.pathname}${back.search}` })).body).toContain('授权未完成')

    const dto = app.mcpServers.get(s.id)
    expect(dto.oauthStatus).toMatchObject({ authorized: true, registered: true, lastError: null })
    expect(JSON.stringify(dto)).not.toMatch(/access_token|refresh_token|dyn-client/)
    await expect.poll(() => app.mcpServers.get(s.id).tools.map(t => t.name)).toContain('whoami')
    expect((await callWhoami(s.id)).content[0]!.text).toBe('account:alice')

    // 访问令牌失效：SDK 用刷新令牌续期后重试
    valid.clear()
    await app.mcpUpstream.close(s.id)
    expect((await callWhoami(s.id)).content[0]!.text).toBe('account:alice')
    expect(grants).toEqual(['authorization_code', 'refresh_token'])

    // 刷新令牌也失效：提示需要重新授权
    valid.clear()
    refreshTokens.clear()
    await app.mcpUpstream.close(s.id)
    const failed = await callWhoami(s.id)
    expect(failed.isError).toBe(true)
    expect(failed.content[0]!.text).toMatch(/需要 OAuth 授权/)
    expect(app.mcpServers.get(s.id).oauthStatus!.lastError).toMatch(/重新授权/)

    await app.mcpOAuth.revoke(s.id)
    expect(app.mcpServers.get(s.id).oauthStatus).toMatchObject({ authorized: false })
  })

  it('回调参数无效或授权被拒绝时给出说明', async () => {
    expect((await server.inject({ method: 'GET', url: '/admin/v1/mcp/oauth/callback?state=nope&code=x' })).body).toContain('授权请求无效或已过期')
    const id = app.mcpServers.list().find(x => x.name === 'secure')!.id
    const { authorizationUrl } = await app.mcpOAuth.start(id, base)
    const state = new URL(authorizationUrl!).searchParams.get('state')!
    const r = await server.inject({ method: 'GET', url: `/admin/v1/mcp/oauth/callback?state=${state}&error=access_denied&error_description=${encodeURIComponent('用户拒绝')}` })
    expect(r.body).toContain('授权未完成：用户拒绝')
  })
})