Fix/percistent mcp connection (#392)
* feat(oauth): add support for refresh tokens in OAuth flow and update database schema * feat(prompts): add MCP prompts and corresponding functionality for prompt retrieval * feat(auth): enhance error handling for refresh token operations and validation
This commit is contained in:
@@ -0,0 +1,270 @@
|
||||
import { describe, it, expect, vi, beforeEach } from 'vitest'
|
||||
import { createQueuedMockSupabase } from '@/tests/helpers'
|
||||
|
||||
const mocks = vi.hoisted(() => ({
|
||||
supabaseFactory: vi.fn(),
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/auth/api-keys', async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import('@/lib/auth/api-keys')>()
|
||||
return {
|
||||
...actual,
|
||||
createServiceClientNoCookies: () => mocks.supabaseFactory(),
|
||||
}
|
||||
})
|
||||
|
||||
vi.mock('@/lib/auth/oauth-codes', () => ({
|
||||
decryptAuthCode: vi.fn(),
|
||||
verifyPkce: vi.fn(),
|
||||
hashAuthCode: vi.fn(() => 'auth-code-hash'),
|
||||
}))
|
||||
|
||||
vi.mock('@/lib/company/context', () => ({
|
||||
requireCompanyId: vi.fn().mockResolvedValue('company-1'),
|
||||
}))
|
||||
|
||||
import { POST } from '../route'
|
||||
import { decryptAuthCode, verifyPkce } from '@/lib/auth/oauth-codes'
|
||||
import { generateRefreshToken } from '@/lib/auth/api-keys'
|
||||
|
||||
function formRequest(body: Record<string, string>) {
|
||||
return new Request('http://localhost/api/mcp-oauth/token', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/x-www-form-urlencoded' },
|
||||
body: new URLSearchParams(body).toString(),
|
||||
})
|
||||
}
|
||||
|
||||
describe('POST /api/mcp-oauth/token', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
describe('grant_type validation', () => {
|
||||
it('rejects unknown grant types', async () => {
|
||||
const res = await POST(formRequest({ grant_type: 'password' }))
|
||||
expect(res.status).toBe(400)
|
||||
const body = await res.json()
|
||||
expect(body.error).toBe('unsupported_grant_type')
|
||||
})
|
||||
|
||||
it('rejects unsupported content type', async () => {
|
||||
const req = new Request('http://localhost/api/mcp-oauth/token', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'text/plain' },
|
||||
body: 'grant_type=authorization_code',
|
||||
})
|
||||
const res = await POST(req)
|
||||
expect(res.status).toBe(400)
|
||||
})
|
||||
})
|
||||
|
||||
describe('authorization_code grant', () => {
|
||||
it('returns access_token, refresh_token, and expires_in on success', async () => {
|
||||
vi.mocked(decryptAuthCode).mockReturnValue({
|
||||
userId: 'user-1',
|
||||
codeChallenge: 'challenge',
|
||||
redirectUri: 'https://claude.ai/api/cb',
|
||||
exp: Date.now() + 60_000,
|
||||
})
|
||||
vi.mocked(verifyPkce).mockReturnValue(true)
|
||||
|
||||
const { supabase, enqueueMany } = createQueuedMockSupabase()
|
||||
mocks.supabaseFactory.mockReturnValue(supabase)
|
||||
enqueueMany([
|
||||
{ data: null, error: null }, // insert into oauth_used_codes
|
||||
{ data: null, error: null }, // delete expired codes (best-effort)
|
||||
{ data: null, error: null }, // insert into api_keys
|
||||
])
|
||||
|
||||
const res = await POST(
|
||||
formRequest({
|
||||
grant_type: 'authorization_code',
|
||||
code: 'ciphertext',
|
||||
code_verifier: 'verifier',
|
||||
redirect_uri: 'https://claude.ai/api/cb',
|
||||
})
|
||||
)
|
||||
expect(res.status).toBe(200)
|
||||
const body = await res.json()
|
||||
expect(body.access_token).toMatch(/^gnubok_sk_/)
|
||||
expect(body.refresh_token).toMatch(/^gnubok_rt_/)
|
||||
expect(body.token_type).toBe('Bearer')
|
||||
expect(body.expires_in).toBe(3600)
|
||||
})
|
||||
|
||||
it('rejects an already-used auth code (replay)', async () => {
|
||||
vi.mocked(decryptAuthCode).mockReturnValue({
|
||||
userId: 'user-1',
|
||||
codeChallenge: 'challenge',
|
||||
redirectUri: 'https://claude.ai/api/cb',
|
||||
exp: Date.now() + 60_000,
|
||||
})
|
||||
vi.mocked(verifyPkce).mockReturnValue(true)
|
||||
|
||||
const { supabase, enqueue } = createQueuedMockSupabase()
|
||||
mocks.supabaseFactory.mockReturnValue(supabase)
|
||||
enqueue({ data: null, error: { message: 'unique violation' } })
|
||||
|
||||
const res = await POST(
|
||||
formRequest({
|
||||
grant_type: 'authorization_code',
|
||||
code: 'ciphertext',
|
||||
code_verifier: 'verifier',
|
||||
redirect_uri: 'https://claude.ai/api/cb',
|
||||
})
|
||||
)
|
||||
expect(res.status).toBe(400)
|
||||
const body = await res.json()
|
||||
expect(body.error).toBe('invalid_grant')
|
||||
})
|
||||
|
||||
it('rejects when PKCE verification fails', async () => {
|
||||
vi.mocked(decryptAuthCode).mockReturnValue({
|
||||
userId: 'user-1',
|
||||
codeChallenge: 'challenge',
|
||||
redirectUri: 'https://claude.ai/api/cb',
|
||||
exp: Date.now() + 60_000,
|
||||
})
|
||||
vi.mocked(verifyPkce).mockReturnValue(false)
|
||||
|
||||
const res = await POST(
|
||||
formRequest({
|
||||
grant_type: 'authorization_code',
|
||||
code: 'ciphertext',
|
||||
code_verifier: 'wrong',
|
||||
redirect_uri: 'https://claude.ai/api/cb',
|
||||
})
|
||||
)
|
||||
expect(res.status).toBe(400)
|
||||
const body = await res.json()
|
||||
expect(body.error).toBe('invalid_grant')
|
||||
expect(body.error_description).toContain('PKCE')
|
||||
})
|
||||
})
|
||||
|
||||
describe('refresh_token grant', () => {
|
||||
it('rotates both tokens and returns a fresh access_token', async () => {
|
||||
const { token: refreshToken } = generateRefreshToken()
|
||||
|
||||
const { supabase, enqueueMany } = createQueuedMockSupabase()
|
||||
mocks.supabaseFactory.mockReturnValue(supabase)
|
||||
enqueueMany([
|
||||
{ data: { id: 'key-1', revoked_at: null }, error: null }, // SELECT
|
||||
{ data: [{ id: 'key-1' }], error: null }, // UPDATE ... RETURNING
|
||||
])
|
||||
|
||||
const res = await POST(
|
||||
formRequest({
|
||||
grant_type: 'refresh_token',
|
||||
refresh_token: refreshToken,
|
||||
})
|
||||
)
|
||||
|
||||
expect(res.status).toBe(200)
|
||||
const body = await res.json()
|
||||
expect(body.access_token).toMatch(/^gnubok_sk_/)
|
||||
expect(body.refresh_token).toMatch(/^gnubok_rt_/)
|
||||
expect(body.refresh_token).not.toBe(refreshToken) // rotated
|
||||
expect(body.expires_in).toBe(3600)
|
||||
})
|
||||
|
||||
it('returns 400 when refresh_token is missing', async () => {
|
||||
const res = await POST(formRequest({ grant_type: 'refresh_token' }))
|
||||
expect(res.status).toBe(400)
|
||||
const body = await res.json()
|
||||
expect(body.error).toBe('invalid_request')
|
||||
})
|
||||
|
||||
it('returns 400 when refresh_token is unknown', async () => {
|
||||
const { supabase, enqueue } = createQueuedMockSupabase()
|
||||
mocks.supabaseFactory.mockReturnValue(supabase)
|
||||
enqueue({ data: null, error: null }) // SELECT — no row
|
||||
|
||||
const res = await POST(
|
||||
formRequest({
|
||||
grant_type: 'refresh_token',
|
||||
refresh_token: 'gnubok_rt_unknown',
|
||||
})
|
||||
)
|
||||
expect(res.status).toBe(400)
|
||||
const body = await res.json()
|
||||
expect(body.error).toBe('invalid_grant')
|
||||
})
|
||||
|
||||
it('returns 400 when the api_key is revoked', async () => {
|
||||
const { supabase, enqueue } = createQueuedMockSupabase()
|
||||
mocks.supabaseFactory.mockReturnValue(supabase)
|
||||
enqueue({
|
||||
data: { id: 'key-1', revoked_at: '2026-05-01T00:00:00Z' },
|
||||
error: null,
|
||||
})
|
||||
|
||||
const res = await POST(
|
||||
formRequest({
|
||||
grant_type: 'refresh_token',
|
||||
refresh_token: 'gnubok_rt_anything',
|
||||
})
|
||||
)
|
||||
expect(res.status).toBe(400)
|
||||
const body = await res.json()
|
||||
expect(body.error).toBe('invalid_grant')
|
||||
expect(body.error_description).toContain('revoked')
|
||||
})
|
||||
|
||||
it('returns 500 when the lookup fails with a DB error', async () => {
|
||||
const { supabase, enqueue } = createQueuedMockSupabase()
|
||||
mocks.supabaseFactory.mockReturnValue(supabase)
|
||||
enqueue({ data: null, error: { message: 'connection reset' } })
|
||||
|
||||
const res = await POST(
|
||||
formRequest({
|
||||
grant_type: 'refresh_token',
|
||||
refresh_token: 'gnubok_rt_anything',
|
||||
})
|
||||
)
|
||||
expect(res.status).toBe(500)
|
||||
const body = await res.json()
|
||||
expect(body.error).toBe('server_error')
|
||||
})
|
||||
|
||||
it('returns 500 when the rotation update fails with a DB error', async () => {
|
||||
const { supabase, enqueueMany } = createQueuedMockSupabase()
|
||||
mocks.supabaseFactory.mockReturnValue(supabase)
|
||||
enqueueMany([
|
||||
{ data: { id: 'key-1', revoked_at: null }, error: null }, // SELECT
|
||||
{ data: null, error: { message: 'deadlock detected' } }, // UPDATE — DB error
|
||||
])
|
||||
|
||||
const res = await POST(
|
||||
formRequest({
|
||||
grant_type: 'refresh_token',
|
||||
refresh_token: 'gnubok_rt_anything',
|
||||
})
|
||||
)
|
||||
expect(res.status).toBe(500)
|
||||
const body = await res.json()
|
||||
expect(body.error).toBe('server_error')
|
||||
})
|
||||
|
||||
it('returns 400 when the CAS update affects 0 rows (concurrent reuse)', async () => {
|
||||
const { supabase, enqueueMany } = createQueuedMockSupabase()
|
||||
mocks.supabaseFactory.mockReturnValue(supabase)
|
||||
enqueueMany([
|
||||
{ data: { id: 'key-1', revoked_at: null }, error: null }, // SELECT
|
||||
{ data: [], error: null }, // UPDATE — 0 rows (lost the CAS race)
|
||||
])
|
||||
|
||||
const res = await POST(
|
||||
formRequest({
|
||||
grant_type: 'refresh_token',
|
||||
refresh_token: 'gnubok_rt_anything',
|
||||
})
|
||||
)
|
||||
expect(res.status).toBe(400)
|
||||
const body = await res.json()
|
||||
expect(body.error).toBe('invalid_grant')
|
||||
expect(body.error_description).toContain('already used')
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -1,17 +1,25 @@
|
||||
import { NextResponse } from 'next/server'
|
||||
import { decryptAuthCode, verifyPkce, hashAuthCode } from '@/lib/auth/oauth-codes'
|
||||
import { generateApiKey, createServiceClientNoCookies, ALL_SCOPES } from '@/lib/auth/api-keys'
|
||||
import {
|
||||
generateApiKey,
|
||||
generateRefreshToken,
|
||||
hashRefreshToken,
|
||||
createServiceClientNoCookies,
|
||||
ALL_SCOPES,
|
||||
} from '@/lib/auth/api-keys'
|
||||
import { requireCompanyId } from '@/lib/company/context'
|
||||
|
||||
const ACCESS_TOKEN_TTL_SECONDS = 3600
|
||||
|
||||
/**
|
||||
* OAuth 2.0 Token Endpoint.
|
||||
*
|
||||
* Exchanges an authorization code for an API key (access token).
|
||||
* 1. Decrypts the stateless auth code
|
||||
* 2. Checks for replay (single-use enforcement per OAuth 2.1 §4.1.2)
|
||||
* 3. Verifies PKCE (S256 only)
|
||||
* 4. Creates the API key (deferred from /authorize to prevent orphaned keys)
|
||||
* 5. Returns the key as a bearer token
|
||||
* Supports two grant types:
|
||||
* - authorization_code: exchange a PKCE-protected auth code for a fresh
|
||||
* api_key (access_token) plus a refresh_token.
|
||||
* - refresh_token: rotate the refresh_token and return the same api_key
|
||||
* with a fresh expires_in. The api_key itself does not expire
|
||||
* server-side; expires_in is a hint so clients refresh on a cadence.
|
||||
*/
|
||||
export async function POST(request: Request) {
|
||||
let params: URLSearchParams
|
||||
@@ -28,17 +36,29 @@ export async function POST(request: Request) {
|
||||
}
|
||||
|
||||
const grantType = params.get('grant_type')
|
||||
|
||||
if (grantType === 'authorization_code') {
|
||||
return handleAuthorizationCodeGrant(params)
|
||||
}
|
||||
|
||||
if (grantType === 'refresh_token') {
|
||||
return handleRefreshTokenGrant(params)
|
||||
}
|
||||
|
||||
return NextResponse.json(
|
||||
{
|
||||
error: 'unsupported_grant_type',
|
||||
error_description: 'Only authorization_code and refresh_token are supported',
|
||||
},
|
||||
{ status: 400 }
|
||||
)
|
||||
}
|
||||
|
||||
async function handleAuthorizationCodeGrant(params: URLSearchParams) {
|
||||
const code = params.get('code')
|
||||
const codeVerifier = params.get('code_verifier')
|
||||
const redirectUri = params.get('redirect_uri')
|
||||
|
||||
if (grantType !== 'authorization_code') {
|
||||
return NextResponse.json(
|
||||
{ error: 'unsupported_grant_type', error_description: 'Only authorization_code is supported' },
|
||||
{ status: 400 }
|
||||
)
|
||||
}
|
||||
|
||||
if (!code) {
|
||||
return NextResponse.json(
|
||||
{ error: 'invalid_request', error_description: 'Missing code parameter' },
|
||||
@@ -46,7 +66,6 @@ export async function POST(request: Request) {
|
||||
)
|
||||
}
|
||||
|
||||
// Decrypt the auth code
|
||||
const payload = decryptAuthCode(code)
|
||||
if (!payload) {
|
||||
return NextResponse.json(
|
||||
@@ -55,7 +74,6 @@ export async function POST(request: Request) {
|
||||
)
|
||||
}
|
||||
|
||||
// Verify redirect_uri matches
|
||||
if (redirectUri && redirectUri !== payload.redirectUri) {
|
||||
return NextResponse.json(
|
||||
{ error: 'invalid_grant', error_description: 'redirect_uri mismatch' },
|
||||
@@ -63,7 +81,6 @@ export async function POST(request: Request) {
|
||||
)
|
||||
}
|
||||
|
||||
// Verify PKCE (S256 only)
|
||||
if (!codeVerifier) {
|
||||
return NextResponse.json(
|
||||
{ error: 'invalid_request', error_description: 'code_verifier is required' },
|
||||
@@ -78,7 +95,6 @@ export async function POST(request: Request) {
|
||||
)
|
||||
}
|
||||
|
||||
// Single-use enforcement: check and mark code as used (atomically via unique constraint)
|
||||
const codeHash = hashAuthCode(code)
|
||||
const supabase = createServiceClientNoCookies()
|
||||
|
||||
@@ -87,7 +103,6 @@ export async function POST(request: Request) {
|
||||
.insert({ code_hash: codeHash })
|
||||
|
||||
if (replayError) {
|
||||
// Unique constraint violation = code already used
|
||||
return NextResponse.json(
|
||||
{ error: 'invalid_grant', error_description: 'Authorization code already used' },
|
||||
{ status: 400 }
|
||||
@@ -101,11 +116,10 @@ export async function POST(request: Request) {
|
||||
.lt('created_at', new Date(Date.now() - 10 * 60 * 1000).toISOString())
|
||||
.then(() => {})
|
||||
|
||||
// Resolve company context for the user
|
||||
const companyId = await requireCompanyId(supabase, payload.userId)
|
||||
|
||||
// Create the API key now (after PKCE verification — prevents orphaned keys)
|
||||
const { key, hash, prefix } = generateApiKey()
|
||||
const refresh = generateRefreshToken()
|
||||
|
||||
const { error: insertError } = await supabase
|
||||
.from('api_keys')
|
||||
@@ -116,6 +130,7 @@ export async function POST(request: Request) {
|
||||
key_prefix: prefix,
|
||||
name: 'MCP-klient (OAuth)',
|
||||
scopes: ALL_SCOPES,
|
||||
refresh_token_hash: refresh.hash,
|
||||
})
|
||||
|
||||
if (insertError) {
|
||||
@@ -128,6 +143,91 @@ export async function POST(request: Request) {
|
||||
return NextResponse.json({
|
||||
access_token: key,
|
||||
token_type: 'Bearer',
|
||||
expires_in: ACCESS_TOKEN_TTL_SECONDS,
|
||||
refresh_token: refresh.token,
|
||||
scope: 'mcp',
|
||||
})
|
||||
}
|
||||
|
||||
async function handleRefreshTokenGrant(params: URLSearchParams) {
|
||||
const refreshToken = params.get('refresh_token')
|
||||
if (!refreshToken) {
|
||||
return NextResponse.json(
|
||||
{ error: 'invalid_request', error_description: 'refresh_token is required' },
|
||||
{ status: 400 }
|
||||
)
|
||||
}
|
||||
|
||||
const supabase = createServiceClientNoCookies()
|
||||
const presentedHash = hashRefreshToken(refreshToken)
|
||||
|
||||
// Look up the api_key row by refresh_token_hash. The hash is unique among
|
||||
// non-null values, so there's at most one match.
|
||||
const { data: row, error: lookupError } = await supabase
|
||||
.from('api_keys')
|
||||
.select('id, revoked_at')
|
||||
.eq('refresh_token_hash', presentedHash)
|
||||
.maybeSingle()
|
||||
|
||||
if (lookupError) {
|
||||
return NextResponse.json(
|
||||
{ error: 'server_error', error_description: 'Failed to look up refresh token' },
|
||||
{ status: 500 }
|
||||
)
|
||||
}
|
||||
|
||||
if (!row) {
|
||||
return NextResponse.json(
|
||||
{ error: 'invalid_grant', error_description: 'Invalid refresh token' },
|
||||
{ status: 400 }
|
||||
)
|
||||
}
|
||||
|
||||
if (row.revoked_at) {
|
||||
return NextResponse.json(
|
||||
{ error: 'invalid_grant', error_description: 'Refresh token revoked' },
|
||||
{ status: 400 }
|
||||
)
|
||||
}
|
||||
|
||||
// Rotate both tokens atomically. OAuth 2.1 §6.1 recommends rotating the
|
||||
// refresh token; we also rotate the api_key because key_hash is one-way
|
||||
// and we cannot recover the original plaintext to return to the client.
|
||||
// The .eq('refresh_token_hash', presentedHash) guard makes this a CAS:
|
||||
// a concurrent refresh with the same token will affect 0 rows.
|
||||
const rotated = generateRefreshToken()
|
||||
const { key: newKey, hash: newKeyHash, prefix: newKeyPrefix } = generateApiKey()
|
||||
|
||||
const { data: updated, error: updateError } = await supabase
|
||||
.from('api_keys')
|
||||
.update({
|
||||
refresh_token_hash: rotated.hash,
|
||||
key_hash: newKeyHash,
|
||||
key_prefix: newKeyPrefix,
|
||||
})
|
||||
.eq('id', row.id)
|
||||
.eq('refresh_token_hash', presentedHash)
|
||||
.select('id')
|
||||
|
||||
if (updateError) {
|
||||
return NextResponse.json(
|
||||
{ error: 'server_error', error_description: 'Failed to rotate refresh token' },
|
||||
{ status: 500 }
|
||||
)
|
||||
}
|
||||
|
||||
if (!updated || updated.length === 0) {
|
||||
return NextResponse.json(
|
||||
{ error: 'invalid_grant', error_description: 'Refresh token already used' },
|
||||
{ status: 400 }
|
||||
)
|
||||
}
|
||||
|
||||
return NextResponse.json({
|
||||
access_token: newKey,
|
||||
token_type: 'Bearer',
|
||||
expires_in: ACCESS_TOKEN_TTL_SECONDS,
|
||||
refresh_token: rotated.token,
|
||||
scope: 'mcp',
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user