import {Response} from 'express';
import {randomUUID} from 'crypto';
import type {
OAuthServerProvider,
AuthorizationParams,
} from '@modelcontextprotocol/sdk/server/auth/provider.js';
import type {OAuthRegisteredClientsStore} from '@modelcontextprotocol/sdk/server/auth/clients.js';
import type {
OAuthClientInformationFull,
OAuthTokens,
OAuthTokenRevocationRequest,
} from '@modelcontextprotocol/sdk/shared/auth.js';
import type {AuthInfo} from '@modelcontextprotocol/sdk/server/auth/types.js';
import {InvalidTokenError} from '@modelcontextprotocol/sdk/server/auth/errors.js';
import {OAuthHandler} from './oauth-handler.js';
import type {OAuthCredentials, OAuthTokens as GoogleOAuthTokens} from '../oauth.js';
import {
auditOauthStarted,
auditOauthSuccess,
auditTokenExchangeSuccess,
auditTokenExchangeFailure,
auditTokenRevoked,
auditAccessDenied,
maskEmail,
maskToken,
} from '../../audit.js';
import {encryptToken, decryptToken, TokenPayload} from '../token-encryption.js';
import {logger} from '../../logger.js';
interface PendingAuthRequest {
clientId: string;
redirectUri: string;
codeChallenge: string;
state?: string;
createdAt: Date;
}
interface AuthCodeData {
clientId: string;
codeChallenge: string;
googleTokens: GoogleOAuthTokens;
email: string;
expiresAt: Date;
}
const AUTH_CODE_EXPIRATION_MS = 10 * 60 * 1000;
const PENDING_REQUEST_EXPIRATION_MS = 5 * 60 * 1000;
const TOKEN_REFRESH_THRESHOLD_MS = 5 * 60 * 1000;
const JWE_TOKEN_LIFETIME_SECONDS = 86400 * 365;
export class GmailOAuthServerProvider implements OAuthServerProvider {
private readonly clients: Map<string, OAuthClientInformationFull> = new Map();
private readonly pendingRequests: Map<string, PendingAuthRequest> = new Map();
private readonly authCodes: Map<string, AuthCodeData> = new Map();
private readonly credentials: OAuthCredentials;
private readonly oauthHandler: OAuthHandler;
private readonly encryptionSecret: Uint8Array;
constructor(
credentials: OAuthCredentials,
oauthHandler: OAuthHandler,
encryptionSecret: Uint8Array
) {
this.credentials = credentials;
this.oauthHandler = oauthHandler;
this.encryptionSecret = encryptionSecret;
}
get clientsStore(): OAuthRegisteredClientsStore {
return {
getClient: (clientId: string) => this.clients.get(clientId),
registerClient: (
clientMetadata: Omit<
OAuthClientInformationFull,
'client_id' | 'client_id_issued_at'
>
) => {
const clientId = randomUUID();
const client: OAuthClientInformationFull = {
...clientMetadata,
client_id: clientId,
client_id_issued_at: Math.floor(Date.now() / 1000),
};
this.clients.set(clientId, client);
return client;
},
};
}
async authorize(
client: OAuthClientInformationFull,
params: AuthorizationParams,
res: Response
): Promise<void> {
const registeredUris = client.redirect_uris?.map(u => u.toString()) ?? [];
if (!registeredUris.includes(params.redirectUri)) {
throw new Error('Invalid redirect URI');
}
const requestId = randomUUID();
this.pendingRequests.set(requestId, {
clientId: client.client_id,
redirectUri: params.redirectUri,
codeChallenge: params.codeChallenge,
state: params.state,
createdAt: new Date(),
});
const authUrl = this.oauthHandler.generateAuthUrl(requestId);
auditOauthStarted({
clientId: maskToken(client.client_id),
redirectUri: params.redirectUri,
});
res.redirect(authUrl);
}
async challengeForAuthorizationCode(
_client: OAuthClientInformationFull,
authorizationCode: string
): Promise<string> {
const codeData = this.authCodes.get(authorizationCode);
if (!codeData) {
throw new InvalidTokenError('Invalid authorization code');
}
if (codeData.expiresAt < new Date()) {
this.authCodes.delete(authorizationCode);
throw new InvalidTokenError('Authorization code expired');
}
return codeData.codeChallenge;
}
async exchangeAuthorizationCode(
client: OAuthClientInformationFull,
authorizationCode: string,
_codeVerifier?: string,
_redirectUri?: string,
_resource?: URL
): Promise<OAuthTokens> {
const codeData = this.authCodes.get(authorizationCode);
if (!codeData) {
auditTokenExchangeFailure({reason: 'invalid_code'});
throw new InvalidTokenError('Invalid authorization code');
}
if (codeData.expiresAt < new Date()) {
this.authCodes.delete(authorizationCode);
auditTokenExchangeFailure({reason: 'code_expired'});
throw new InvalidTokenError('Authorization code expired');
}
if (codeData.clientId !== client.client_id) {
auditTokenExchangeFailure({reason: 'client_mismatch', clientId: maskToken(client.client_id)});
throw new InvalidTokenError('Client ID mismatch');
}
this.authCodes.delete(authorizationCode);
const payload: TokenPayload = {
access_token: codeData.googleTokens.access_token,
refresh_token: codeData.googleTokens.refresh_token,
expiry_date: codeData.googleTokens.expiry_date,
email: codeData.email,
};
const jweToken = await encryptToken(payload, this.encryptionSecret);
auditTokenExchangeSuccess({
clientId: maskToken(client.client_id),
});
return {
access_token: jweToken,
token_type: 'Bearer',
expires_in: JWE_TOKEN_LIFETIME_SECONDS,
};
}
async exchangeRefreshToken(
_client: OAuthClientInformationFull,
_refreshToken: string,
_scopes?: string[],
_resource?: URL
): Promise<OAuthTokens> {
throw new InvalidTokenError(
'Refresh tokens not supported - re-authenticate'
);
}
async verifyAccessToken(token: string): Promise<AuthInfo> {
let payload: TokenPayload;
try {
payload = await decryptToken(token, this.encryptionSecret);
} catch {
auditAccessDenied({reason: 'invalid_jwe_token'});
throw new InvalidTokenError('Invalid access token');
}
if (payload.expiry_date < Date.now() + TOKEN_REFRESH_THRESHOLD_MS) {
const oauth2Client = this.oauthHandler.createOAuth2Client({
access_token: payload.access_token,
refresh_token: payload.refresh_token,
expiry_date: payload.expiry_date,
scope: 'https://www.googleapis.com/auth/gmail.readonly',
token_type: 'Bearer',
});
logger.debug({email: maskEmail(payload.email)}, 'Starting token refresh');
try {
const {credentials} = await oauth2Client.refreshAccessToken();
payload.access_token = credentials.access_token!;
payload.expiry_date = credentials.expiry_date!;
logger.debug({email: maskEmail(payload.email)}, 'Token refresh completed');
} catch (error) {
auditAccessDenied({
reason: 'token_refresh_failed',
userId: maskEmail(payload.email),
error: error instanceof Error ? error.message : 'Unknown error',
});
throw new InvalidTokenError('Failed to refresh token - re-authenticate');
}
}
return {
token,
clientId: payload.email,
scopes: ['gmail.readonly'],
expiresAt: Math.floor(Date.now() / 1000) + JWE_TOKEN_LIFETIME_SECONDS,
extra: {
googleAccessToken: payload.access_token,
email: payload.email,
},
};
}
async revokeToken(
_client: OAuthClientInformationFull,
request: OAuthTokenRevocationRequest
): Promise<void> {
try {
const payload = await decryptToken(request.token, this.encryptionSecret);
auditTokenRevoked({userId: maskEmail(payload.email)});
} catch {
auditTokenRevoked({userId: 'unknown'});
}
}
async handleGoogleCallback(
requestId: string,
googleCode: string
): Promise<{ redirectUrl: string }> {
const pendingRequest = this.pendingRequests.get(requestId);
if (!pendingRequest) {
throw new Error('Invalid or expired authorization request');
}
if (
new Date().getTime() - pendingRequest.createdAt.getTime() >
PENDING_REQUEST_EXPIRATION_MS
) {
this.pendingRequests.delete(requestId);
throw new Error('Authorization request expired');
}
const tokens = await this.oauthHandler.exchangeCodeForTokens(googleCode);
const userInfo = await this.oauthHandler.getUserInfo(tokens);
auditOauthSuccess({userId: maskEmail(userInfo.email)});
const authCode = randomUUID();
this.authCodes.set(authCode, {
clientId: pendingRequest.clientId,
codeChallenge: pendingRequest.codeChallenge,
googleTokens: tokens,
email: userInfo.email,
expiresAt: new Date(Date.now() + AUTH_CODE_EXPIRATION_MS),
});
this.pendingRequests.delete(requestId);
const redirectUrl = new URL(pendingRequest.redirectUri);
redirectUrl.searchParams.set('code', authCode);
if (pendingRequest.state) {
redirectUrl.searchParams.set('state', pendingRequest.state);
}
return {redirectUrl: redirectUrl.toString()};
}
cleanup(): void {
const now = new Date();
for (const [requestId, request] of this.pendingRequests.entries()) {
if (
now.getTime() - request.createdAt.getTime() >
PENDING_REQUEST_EXPIRATION_MS
) {
this.pendingRequests.delete(requestId);
}
}
for (const [code, data] of this.authCodes.entries()) {
if (data.expiresAt < now) {
this.authCodes.delete(code);
}
}
}
}
