wip on oauth2 support over mcp

brendan-kellam · brendan-kellam · commit 871ee1b6dc8f · 2026-03-03T22:04:45.000-08:00
diff --git a/packages/db/prisma/migrations/20260304051539_add_oauth_tables/migration.sql b/packages/db/prisma/migrations/20260304051539_add_oauth_tables/migration.sql
@@ -0,0 +1,47 @@
+-- CreateTable
+CREATE TABLE "OAuthClient" (
+    "id" TEXT NOT NULL,
+    "name" TEXT NOT NULL,
+    "redirectUris" TEXT[],
+    "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+    CONSTRAINT "OAuthClient_pkey" PRIMARY KEY ("id")
+);
+
+-- CreateTable
+CREATE TABLE "OAuthAuthorizationCode" (
+    "codeHash" TEXT NOT NULL,
+    "clientId" TEXT NOT NULL,
+    "userId" TEXT NOT NULL,
+    "redirectUri" TEXT NOT NULL,
+    "codeChallenge" TEXT NOT NULL,
+    "expiresAt" TIMESTAMP(3) NOT NULL,
+    "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+    CONSTRAINT "OAuthAuthorizationCode_pkey" PRIMARY KEY ("codeHash")
+);
+
+-- CreateTable
+CREATE TABLE "OAuthToken" (
+    "hash" TEXT NOT NULL,
+    "clientId" TEXT NOT NULL,
+    "userId" TEXT NOT NULL,
+    "scope" TEXT NOT NULL DEFAULT '',
+    "expiresAt" TIMESTAMP(3) NOT NULL,
+    "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    "lastUsedAt" TIMESTAMP(3),
+
+    CONSTRAINT "OAuthToken_pkey" PRIMARY KEY ("hash")
+);
+
+-- AddForeignKey
+ALTER TABLE "OAuthAuthorizationCode" ADD CONSTRAINT "OAuthAuthorizationCode_clientId_fkey" FOREIGN KEY ("clientId") REFERENCES "OAuthClient"("id") ON DELETE CASCADE ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "OAuthAuthorizationCode" ADD CONSTRAINT "OAuthAuthorizationCode_userId_fkey" FOREIGN KEY ("userId") REFERENCES "User"("id") ON DELETE CASCADE ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "OAuthToken" ADD CONSTRAINT "OAuthToken_clientId_fkey" FOREIGN KEY ("clientId") REFERENCES "OAuthClient"("id") ON DELETE CASCADE ON UPDATE CASCADE;
+
+-- AddForeignKey
+ALTER TABLE "OAuthToken" ADD CONSTRAINT "OAuthToken_userId_fkey" FOREIGN KEY ("userId") REFERENCES "User"("id") ON DELETE CASCADE ON UPDATE CASCADE;
diff --git a/packages/db/prisma/schema.prisma b/packages/db/prisma/schema.prisma
@@ -365,6 +365,9 @@ model User {
   chats Chat[]
   sharedChats ChatAccess[]
 
+  oauthTokens    OAuthToken[]
+  oauthAuthCodes OAuthAuthorizationCode[]
+
   createdAt DateTime @default(now())
   updatedAt DateTime @updatedAt
 
@@ -485,3 +488,45 @@ model ChatAccess {
 
   @@unique([chatId, userId])
 }
+
+// OAuth2 Authorization Server models
+// @see: https://datatracker.ietf.org/doc/html/rfc6749
+
+/// A registered OAuth2 client application (e.g. Claude Desktop, Cursor).
+/// Created via dynamic client registration (RFC 7591) at POST /api/oauth/register.
+model OAuthClient {
+  id           String   @id @default(cuid())
+  name         String
+  redirectUris String[]
+  createdAt    DateTime @default(now())
+
+  authCodes OAuthAuthorizationCode[]
+  tokens    OAuthToken[]
+}
+
+/// A short-lived authorization code issued during the OAuth2 authorization code flow.
+/// Single-use and expires after 10 minutes. Stores the PKCE code challenge.
+model OAuthAuthorizationCode {
+  codeHash      String   @id // hashSecret(rawCode)
+  clientId      String
+  client        OAuthClient @relation(fields: [clientId], references: [id], onDelete: Cascade)
+  userId        String
+  user          User        @relation(fields: [userId], references: [id], onDelete: Cascade)
+  redirectUri   String
+  codeChallenge String // BASE64URL(SHA-256(codeVerifier))
+  expiresAt     DateTime
+  createdAt     DateTime    @default(now())
+}
+
+/// An opaque OAuth2 access token. The raw token is never stored — only its HMAC-SHA256 hash.
+model OAuthToken {
+  hash       String   @id // hashSecret(rawToken secret portion)
+  clientId   String
+  client     OAuthClient @relation(fields: [clientId], references: [id], onDelete: Cascade)
+  userId     String
+  user       User        @relation(fields: [userId], references: [id], onDelete: Cascade)
+  scope      String      @default("")
+  expiresAt  DateTime
+  createdAt  DateTime    @default(now())
+  lastUsedAt DateTime?
+}
diff --git a/packages/shared/src/crypto.ts b/packages/shared/src/crypto.ts
@@ -40,6 +40,16 @@ export function generateApiKey(): { key: string; hash: string } {
     };
 }
 
+export function generateOAuthToken(): { token: string; hash: string } {
+    const secret = crypto.randomBytes(32).toString('hex');
+    const hash = hashSecret(secret);
+
+    return {
+        token: `sourcebot-oauth-${secret}`,
+        hash,
+    };
+}
+
 export function decrypt(iv: string, encryptedText: string): string {
     const encryptionKey = Buffer.from(env.SOURCEBOT_ENCRYPTION_KEY, 'ascii');
 
diff --git a/packages/shared/src/index.server.ts b/packages/shared/src/index.server.ts
@@ -45,6 +45,7 @@ export {
     decrypt,
     hashSecret,
     generateApiKey,
+    generateOAuthToken,
     verifySignature,
     encryptOAuthToken,
     decryptOAuthToken,
diff --git a/packages/web/next.config.mjs b/packages/web/next.config.mjs
@@ -24,6 +24,16 @@ const nextConfig = {
                 source: "/ingest/decide",
                 destination: `https://us.i.posthog.com/decide`,
             },
+            // Expose OAuth discovery documents at canonical RFC paths (without /api prefix)
+            // so MCP clients and OAuth tools can find them via standard discovery.
+            {
+                source: "/.well-known/oauth-authorization-server",
+                destination: "/api/.well-known/oauth-authorization-server",
+            },
+            {
+                source: "/.well-known/oauth-protected-resource",
+                destination: "/api/.well-known/oauth-protected-resource",
+            },
         ];
     },
     // This is required to support PostHog trailing slash API requests
diff --git a/packages/web/src/__mocks__/prisma.ts b/packages/web/src/__mocks__/prisma.ts
@@ -1,5 +1,5 @@
 import { SINGLE_TENANT_ORG_DOMAIN, SINGLE_TENANT_ORG_ID, SINGLE_TENANT_ORG_NAME } from '@/lib/constants';
-import { Account, ApiKey, Org, PrismaClient, User } from '@prisma/client';
+import { Account, ApiKey, OAuthToken, Org, PrismaClient, User } from '@prisma/client';
 import { beforeEach, vi } from 'vitest';
 import { mockDeep, mockReset } from 'vitest-mock-extended';
 
@@ -47,4 +47,15 @@ export const MOCK_USER_WITH_ACCOUNTS: User & { accounts: Account[] } = {
     accounts: [],
 }
 
+export const MOCK_OAUTH_TOKEN: OAuthToken & { user: User & { accounts: Account[] } } = {
+    hash: 'oauthtoken',
+    clientId: 'test-client-id',
+    userId: MOCK_USER_WITH_ACCOUNTS.id,
+    scope: '',
+    expiresAt: new Date(Date.now() + 1000 * 60 * 60 * 24 * 365), // 1 year from now
+    createdAt: new Date(),
+    lastUsedAt: null,
+    user: MOCK_USER_WITH_ACCOUNTS,
+}
+
 export const userScopedPrismaClientExtension = vi.fn();
diff --git a/packages/web/src/app/api/(server)/.well-known/oauth-authorization-server/route.ts b/packages/web/src/app/api/(server)/.well-known/oauth-authorization-server/route.ts
@@ -0,0 +1,20 @@
+import { apiHandler } from '@/lib/apiHandler';
+import { env } from '@sourcebot/shared';
+
+// RFC 8414: OAuth 2.0 Authorization Server Metadata
+// @see: https://datatracker.ietf.org/doc/html/rfc8414
+export const GET = apiHandler(async () => {
+    const issuer = env.AUTH_URL.replace(/\/$/, '');
+
+    return Response.json({
+        issuer,
+        authorization_endpoint: `${issuer}/oauth/authorize`,
+        token_endpoint: `${issuer}/api/oauth/token`,
+        registration_endpoint: `${issuer}/api/oauth/register`,
+        revocation_endpoint: `${issuer}/api/oauth/revoke`,
+        response_types_supported: ['code'],
+        grant_types_supported: ['authorization_code'],
+        code_challenge_methods_supported: ['S256'],
+        token_endpoint_auth_methods_supported: ['none'],
+    });
+}, { track: false });
diff --git a/packages/web/src/app/api/(server)/.well-known/oauth-protected-resource/route.ts b/packages/web/src/app/api/(server)/.well-known/oauth-protected-resource/route.ts
@@ -0,0 +1,14 @@
+import { apiHandler } from '@/lib/apiHandler';
+import { env } from '@sourcebot/shared';
+
+// RFC 9728: OAuth 2.0 Protected Resource Metadata
+// Tells OAuth clients which authorization server protects this resource.
+// @see: https://datatracker.ietf.org/doc/html/rfc9728
+export const GET = apiHandler(async () => {
+    const issuer = env.AUTH_URL.replace(/\/$/, '');
+
+    return Response.json({
+        resource: `${issuer}/api/mcp`,
+        authorization_servers: [issuer],
+    });
+}, { track: false });
diff --git a/packages/web/src/app/api/(server)/mcp/route.ts b/packages/web/src/app/api/(server)/mcp/route.ts
@@ -9,6 +9,23 @@ import { StatusCodes } from 'http-status-codes';
 import { NextRequest } from 'next/server';
 import { sew } from '@/actions';
 import { apiHandler } from '@/lib/apiHandler';
+import { env } from '@sourcebot/shared';
+
+// On 401, tell MCP clients where to find the OAuth protected resource metadata (RFC 9728)
+// so they can discover the authorization server and initiate the authorization code flow.
+// @see: https://modelcontextprotocol.io/specification/2025-03-26/basic/authentication
+// @see: https://datatracker.ietf.org/doc/html/rfc9728
+function mcpErrorResponse(error: ServiceError): Response {
+    const response = serviceErrorResponse(error);
+    if (error.statusCode === StatusCodes.UNAUTHORIZED) {
+        const issuer = env.AUTH_URL.replace(/\/$/, '');
+        response.headers.set(
+            'WWW-Authenticate',
+            `Bearer realm="Sourcebot", resource_metadata_uri="${issuer}/.well-known/oauth-protected-resource"`
+        );
+    }
+    return response;
+}
 
 // @see: https://modelcontextprotocol.io/specification/2025-11-25/basic/transports#session-management
 interface McpSession {
@@ -66,7 +83,7 @@ export const POST = apiHandler(async (request: NextRequest) => {
     );
 
     if (isServiceError(response)) {
-        return serviceErrorResponse(response);
+        return mcpErrorResponse(response);
     }
 
     return response;
@@ -99,7 +116,7 @@ export const DELETE = apiHandler(async (request: NextRequest) => {
     );
 
     if (isServiceError(result)) {
-        return serviceErrorResponse(result);
+        return mcpErrorResponse(result);
     }
 
     return result;
@@ -132,7 +149,7 @@ export const GET = apiHandler(async (request: NextRequest) => {
     );
 
     if (isServiceError(result)) {
-        return serviceErrorResponse(result);
+        return mcpErrorResponse(result);
     }
 
     return result;
diff --git a/packages/web/src/app/api/(server)/oauth/register/route.ts b/packages/web/src/app/api/(server)/oauth/register/route.ts
@@ -0,0 +1,47 @@
+import { apiHandler } from '@/lib/apiHandler';
+import { requestBodySchemaValidationError, serviceErrorResponse } from '@/lib/serviceError';
+import { prisma } from '@/prisma';
+import { NextRequest } from 'next/server';
+import { z } from 'zod';
+
+// RFC 7591: OAuth 2.0 Dynamic Client Registration
+// @see: https://datatracker.ietf.org/doc/html/rfc7591
+const registerRequestSchema = z.object({
+    client_name: z.string().min(1),
+    redirect_uris: z.array(z.string().url()).min(1),
+});
+
+export const POST = apiHandler(async (request: NextRequest) => {
+    const body = await request.json();
+    const parsed = registerRequestSchema.safeParse(body);
+
+    if (!parsed.success) {
+        return serviceErrorResponse(requestBodySchemaValidationError(parsed.error));
+    }
+
+    const { client_name, redirect_uris } = parsed.data;
+
+    // Reject wildcard redirect URIs per security best practices
+    if (redirect_uris.some((uri) => uri.includes('*'))) {
+        return Response.json(
+            { error: 'invalid_redirect_uri', error_description: 'Wildcard redirect URIs are not allowed.' },
+            { status: 400 }
+        );
+    }
+
+    const client = await prisma.oAuthClient.create({
+        data: {
+            name: client_name,
+            redirectUris: redirect_uris,
+        },
+    });
+
+    return Response.json(
+        {
+            client_id: client.id,
+            client_name: client.name,
+            redirect_uris: client.redirectUris,
+        },
+        { status: 201 }
+    );
+});
diff --git a/packages/web/src/app/api/(server)/oauth/revoke/route.ts b/packages/web/src/app/api/(server)/oauth/revoke/route.ts
@@ -0,0 +1,17 @@
+import { revokeToken } from '@/features/oauth/server';
+import { apiHandler } from '@/lib/apiHandler';
+import { NextRequest } from 'next/server';
+
+// RFC 7009: OAuth 2.0 Token Revocation
+// Always returns 200 regardless of whether the token existed.
+// @see: https://datatracker.ietf.org/doc/html/rfc7009
+export const POST = apiHandler(async (request: NextRequest) => {
+    const formData = await request.formData();
+    const token = formData.get('token');
+
+    if (token) {
+        await revokeToken(token.toString());
+    }
+
+    return new Response(null, { status: 200 });
+});
diff --git a/packages/web/src/app/api/(server)/oauth/token/route.ts b/packages/web/src/app/api/(server)/oauth/token/route.ts
@@ -0,0 +1,51 @@
+import { verifyAndExchangeCode, ACCESS_TOKEN_TTL_SECONDS } from '@/features/oauth/server';
+import { apiHandler } from '@/lib/apiHandler';
+import { NextRequest } from 'next/server';
+
+// OAuth 2.0 Token Endpoint
+// Supports grant_type=authorization_code with PKCE (RFC 7636).
+// @see: https://datatracker.ietf.org/doc/html/rfc6749#section-3.2
+export const POST = apiHandler(async (request: NextRequest) => {
+    const formData = await request.formData();
+
+    const grantType = formData.get('grant_type');
+    if (grantType !== 'authorization_code') {
+        return Response.json(
+            { error: 'unsupported_grant_type', error_description: 'Only authorization_code is supported.' },
+            { status: 400 }
+        );
+    }
+
+    const code = formData.get('code');
+    const clientId = formData.get('client_id');
+    const redirectUri = formData.get('redirect_uri');
+    const codeVerifier = formData.get('code_verifier');
+
+    if (!code || !clientId || !redirectUri || !codeVerifier) {
+        return Response.json(
+            { error: 'invalid_request', error_description: 'Missing required parameters: code, client_id, redirect_uri, code_verifier.' },
+            { status: 400 }
+        );
+    }
+
+    const result = await verifyAndExchangeCode({
+        rawCode: code.toString(),
+        clientId: clientId.toString(),
+        redirectUri: redirectUri.toString(),
+        codeVerifier: codeVerifier.toString(),
+    });
+
+    if ('error' in result) {
+        return Response.json(
+            { error: result.error, error_description: result.errorDescription },
+            { status: 400 }
+        );
+    }
+
+    return Response.json({
+        access_token: result.token,
+        token_type: 'Bearer',
+        expires_in: ACCESS_TOKEN_TTL_SECONDS,
+        scope: '',
+    });
+});
diff --git a/packages/web/src/app/oauth/authorize/page.tsx b/packages/web/src/app/oauth/authorize/page.tsx
diff --git a/packages/web/src/features/oauth/server.ts b/packages/web/src/features/oauth/server.ts
diff --git a/packages/web/src/proxy.ts b/packages/web/src/proxy.ts
diff --git a/packages/web/src/withAuthV2.ts b/packages/web/src/withAuthV2.ts
