Merge branch 'main' into ochafik/add-elicitation-tool
This commit is contained in:
Olivier Chafik
commit
b6e39748aa
22 changed files with 3933 additions and 17905 deletions
87
src/modules/auth/auth/auth-core.ts
Normal file
87
src/modules/auth/auth/auth-core.ts
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
import crypto from "crypto";
|
||||
import { OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
|
||||
|
||||
/**
|
||||
* Generates a PKCE code challenge from a verifier string.
|
||||
* Uses S256 method as specified in RFC 7636.
|
||||
* @param verifier The code verifier string
|
||||
* @returns Base64url-encoded SHA256 hash of the verifier
|
||||
*/
|
||||
export function generatePKCEChallenge(verifier: string): string {
|
||||
const buffer = Buffer.from(verifier);
|
||||
const hash = crypto.createHash("sha256").update(buffer);
|
||||
return hash.digest("base64url");
|
||||
}
|
||||
|
||||
/**
|
||||
* Generates a cryptographically secure random token.
|
||||
* @returns 64-character hexadecimal string
|
||||
*/
|
||||
export function generateToken(): string {
|
||||
return crypto.randomBytes(32).toString("hex");
|
||||
}
|
||||
|
||||
/**
|
||||
* Computes SHA256 hash of input data.
|
||||
* @param data The string to hash
|
||||
* @returns Hexadecimal representation of the hash
|
||||
*/
|
||||
export function sha256(data: string): string {
|
||||
return crypto.createHash("sha256").update(data).digest("hex");
|
||||
}
|
||||
|
||||
/**
|
||||
* Encrypts a string using AES-256-CBC encryption.
|
||||
* @param text The plaintext to encrypt
|
||||
* @param key The encryption key (64 hex characters)
|
||||
* @returns Encrypted string in format "iv:ciphertext"
|
||||
*/
|
||||
export function encryptString({ text, key }: { text: string; key: string }): string {
|
||||
const iv = crypto.randomBytes(16);
|
||||
const cipher = crypto.createCipheriv("aes-256-cbc", Buffer.from(key, "hex"), iv);
|
||||
let encrypted = cipher.update(text, "utf-8", "hex");
|
||||
encrypted += cipher.final("hex");
|
||||
return `${iv.toString("hex")}:${encrypted}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Decrypts a string encrypted with encryptString.
|
||||
* @param encryptedText The encrypted string in format "iv:ciphertext"
|
||||
* @param key The encryption key (64 hex characters)
|
||||
* @returns Decrypted plaintext
|
||||
*/
|
||||
export function decryptString({
|
||||
encryptedText,
|
||||
key,
|
||||
}: {
|
||||
encryptedText: string;
|
||||
key: string;
|
||||
}): string {
|
||||
const [ivHex, encrypted] = encryptedText.split(":");
|
||||
const iv = Buffer.from(ivHex, "hex");
|
||||
const decipher = crypto.createDecipheriv("aes-256-cbc", Buffer.from(key, "hex"), iv);
|
||||
let decrypted = decipher.update(encrypted, "hex", "utf-8");
|
||||
decrypted += decipher.final("utf-8");
|
||||
return decrypted;
|
||||
}
|
||||
|
||||
/**
|
||||
* Access token expiry time in seconds (1 hour)
|
||||
*/
|
||||
export const ACCESS_TOKEN_EXPIRY_SEC = 60 * 60;
|
||||
|
||||
/**
|
||||
* Generates a complete set of MCP OAuth tokens.
|
||||
* @returns OAuth tokens with access token, refresh token, and expiry
|
||||
*/
|
||||
export function generateMcpTokens(): OAuthTokens {
|
||||
const mcpAccessToken = generateToken();
|
||||
const mcpRefreshToken = generateToken();
|
||||
|
||||
return {
|
||||
access_token: mcpAccessToken,
|
||||
refresh_token: mcpRefreshToken,
|
||||
expires_in: ACCESS_TOKEN_EXPIRY_SEC,
|
||||
token_type: "Bearer",
|
||||
};
|
||||
}
|
||||
412
src/modules/auth/auth/provider.test.ts
Normal file
412
src/modules/auth/auth/provider.test.ts
Normal file
|
|
@ -0,0 +1,412 @@
|
|||
import { jest, describe, beforeEach, it, expect } from '@jest/globals';
|
||||
import { Response } from "express";
|
||||
import { AuthorizationParams } from "@modelcontextprotocol/sdk/server/auth/provider.js";
|
||||
import { OAuthClientInformationFull, OAuthTokenRevocationRequest, OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
|
||||
import { InvalidTokenError } from "@modelcontextprotocol/sdk/server/auth/errors.js";
|
||||
import { MockRedisClient, setRedisClient } from "../../shared/redis.js";
|
||||
import { McpInstallation, PendingAuthorization, TokenExchange } from "../types.js";
|
||||
import { FeatureReferenceAuthProvider, FeatureReferenceOAuthClientsStore } from "./provider.js";
|
||||
import * as authService from "../services/auth.js";
|
||||
|
||||
// Helper function to create sample client
|
||||
function createTestClient(): OAuthClientInformationFull {
|
||||
return {
|
||||
client_id: "test-client-id",
|
||||
client_name: "Test Client",
|
||||
client_uri: "https://example.com",
|
||||
redirect_uris: ["https://example.com/callback"]
|
||||
};
|
||||
}
|
||||
|
||||
// Helper function to create a mock Response object
|
||||
function createMockResponse() {
|
||||
const res = {
|
||||
redirect: jest.fn().mockReturnThis(),
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn().mockReturnThis(),
|
||||
send: jest.fn().mockReturnThis(),
|
||||
setHeader: jest.fn().mockReturnThis()
|
||||
};
|
||||
return res as unknown as jest.Mocked<Response>;
|
||||
}
|
||||
|
||||
|
||||
function getMockAuthValues() {
|
||||
const client = createTestClient();
|
||||
// Use properly generated tokens for encryption
|
||||
const accessToken = authService.generateToken();
|
||||
const newTokens: OAuthTokens = {
|
||||
access_token: authService.generateToken(),
|
||||
refresh_token: authService.generateToken(),
|
||||
token_type: "bearer",
|
||||
expires_in: 3600,
|
||||
};
|
||||
const mcpInstallation: McpInstallation = {
|
||||
mockUpstreamInstallation: {
|
||||
mockUpstreamAccessToken: "fake-upstream-access-token",
|
||||
mockUpstreamRefreshToken: "fake-upstream-refresh-token",
|
||||
},
|
||||
mcpTokens: {
|
||||
access_token: accessToken,
|
||||
token_type: "Bearer",
|
||||
refresh_token: authService.generateToken(),
|
||||
expires_in: 3600,
|
||||
},
|
||||
clientId: client.client_id,
|
||||
issuedAt: Date.now() / 1000,
|
||||
userId: "test-user-id",
|
||||
};
|
||||
|
||||
return {
|
||||
client,
|
||||
accessToken,
|
||||
newTokens,
|
||||
mcpInstallation,
|
||||
}
|
||||
}
|
||||
|
||||
describe("FeatureReferenceOAuthClientsStore", () => {
|
||||
let clientsStore: FeatureReferenceOAuthClientsStore;
|
||||
|
||||
beforeEach(() => {
|
||||
const mockRedis = new MockRedisClient();
|
||||
setRedisClient(mockRedis);
|
||||
jest.resetAllMocks();
|
||||
clientsStore = new FeatureReferenceOAuthClientsStore();
|
||||
});
|
||||
|
||||
describe("getClient", () => {
|
||||
it("returns undefined for non-existent client", async () => {
|
||||
const result = await clientsStore.getClient("non-existent");
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it("returns client information for existing client", async () => {
|
||||
const client = createTestClient();
|
||||
// First save the client
|
||||
await clientsStore.registerClient(client);
|
||||
|
||||
// Then retrieve it
|
||||
const result = await clientsStore.getClient(client.client_id);
|
||||
|
||||
expect(result).toEqual(client);
|
||||
});
|
||||
});
|
||||
|
||||
describe("registerClient", () => {
|
||||
it("saves and returns client information", async () => {
|
||||
const client = createTestClient();
|
||||
|
||||
const result = await clientsStore.registerClient(client);
|
||||
|
||||
expect(result).toEqual(client);
|
||||
|
||||
// Verify it was saved
|
||||
const retrieved = await clientsStore.getClient(client.client_id);
|
||||
expect(retrieved).toEqual(client);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("FeatureReferenceAuthProvider", () => {
|
||||
let provider: FeatureReferenceAuthProvider;
|
||||
let mockRedis: MockRedisClient;
|
||||
|
||||
beforeEach(() => {
|
||||
jest.resetAllMocks();
|
||||
|
||||
mockRedis = new MockRedisClient();
|
||||
setRedisClient(mockRedis);
|
||||
|
||||
provider = new FeatureReferenceAuthProvider();
|
||||
});
|
||||
|
||||
describe("authorize", () => {
|
||||
it("saves pending authorization and sends HTML response", async () => {
|
||||
const client = createTestClient();
|
||||
// Use a type assertion to make TypeScript ignore the mismatch
|
||||
const params = {
|
||||
redirectUri: "https://example.com/callback",
|
||||
codeChallenge: "test-challenge",
|
||||
codeChallengeMethod: "S256",
|
||||
} as unknown as AuthorizationParams;
|
||||
const res = createMockResponse();
|
||||
|
||||
await provider.authorize(client, params, res);
|
||||
|
||||
// Verify HTML sent with redirect
|
||||
expect(res.send).toHaveBeenCalled();
|
||||
const sentHtml = (res.send as jest.Mock).mock.calls[0][0];
|
||||
expect(sentHtml).toContain('MCP Server Authorization');
|
||||
expect(sentHtml).toContain('Authorization Required');
|
||||
expect(sentHtml).toContain('mock-upstream-idp/authorize?redirect_uri=/mock-upstream-idp/callback&state=');
|
||||
});
|
||||
});
|
||||
|
||||
describe("challengeForAuthorizationCode", () => {
|
||||
it("returns code challenge for valid authorization code", async () => {
|
||||
const client = createTestClient();
|
||||
const pendingAuth: PendingAuthorization = {
|
||||
redirectUri: "https://example.com/callback",
|
||||
codeChallenge: "test-challenge",
|
||||
codeChallengeMethod: "S256",
|
||||
clientId: client.client_id,
|
||||
};
|
||||
|
||||
// First save the pending authorization
|
||||
const authCode = authService.generateToken();
|
||||
await authService.savePendingAuthorization(authCode, pendingAuth);
|
||||
|
||||
const result = await provider.challengeForAuthorizationCode(client, authCode);
|
||||
|
||||
expect(result).toBe("test-challenge");
|
||||
});
|
||||
|
||||
it("throws error for non-existent authorization code", async () => {
|
||||
const client = createTestClient();
|
||||
|
||||
await expect(provider.challengeForAuthorizationCode(client, "test-code"))
|
||||
.rejects
|
||||
.toThrow("Authorization code not found");
|
||||
});
|
||||
|
||||
it("throws error when client ID doesn't match", async () => {
|
||||
const client = createTestClient();
|
||||
const pendingAuth: PendingAuthorization = {
|
||||
redirectUri: "https://example.com/callback",
|
||||
codeChallenge: "test-challenge",
|
||||
codeChallengeMethod: "S256",
|
||||
clientId: "different-client-id",
|
||||
};
|
||||
|
||||
// Save pending auth with different client ID
|
||||
const authCode = authService.generateToken();
|
||||
await authService.savePendingAuthorization(authCode, pendingAuth);
|
||||
|
||||
await expect(provider.challengeForAuthorizationCode(client, authCode))
|
||||
.rejects
|
||||
.toThrow("Authorization code does not match client");
|
||||
});
|
||||
});
|
||||
|
||||
describe("exchangeAuthorizationCode", () => {
|
||||
it("returns tokens when exchange is successful", async () => {
|
||||
const client = createTestClient();
|
||||
const { mcpInstallation } = getMockAuthValues();
|
||||
|
||||
// Setup: save token exchange and installation
|
||||
const authCode = authService.generateToken();
|
||||
const tokenExchange: TokenExchange = {
|
||||
mcpAccessToken: mcpInstallation.mcpTokens.access_token,
|
||||
alreadyUsed: false,
|
||||
};
|
||||
await authService.saveTokenExchange(authCode, tokenExchange);
|
||||
await authService.saveMcpInstallation(mcpInstallation.mcpTokens.access_token, mcpInstallation);
|
||||
|
||||
const result = await provider.exchangeAuthorizationCode(client, authCode);
|
||||
|
||||
expect(result).toEqual({
|
||||
access_token: mcpInstallation.mcpTokens.access_token,
|
||||
expires_in: mcpInstallation.mcpTokens.expires_in,
|
||||
refresh_token: mcpInstallation.mcpTokens.refresh_token,
|
||||
token_type: "Bearer",
|
||||
});
|
||||
});
|
||||
|
||||
it("throws error for invalid authorization code", async () => {
|
||||
const client = createTestClient();
|
||||
|
||||
await expect(provider.exchangeAuthorizationCode(client, "test-code"))
|
||||
.rejects
|
||||
.toThrow("Invalid authorization code");
|
||||
});
|
||||
});
|
||||
|
||||
describe("exchangeRefreshToken", () => {
|
||||
it("returns new tokens when refresh token is valid", async () => {
|
||||
const {
|
||||
client,
|
||||
accessToken,
|
||||
mcpInstallation,
|
||||
} = getMockAuthValues();
|
||||
|
||||
// Setup: save refresh token and installation
|
||||
const refreshToken = authService.generateToken();
|
||||
await authService.saveRefreshToken(refreshToken, accessToken);
|
||||
await authService.saveMcpInstallation(accessToken, mcpInstallation);
|
||||
|
||||
const result = await provider.exchangeRefreshToken(client, refreshToken);
|
||||
|
||||
// Should return new tokens
|
||||
expect(result).toHaveProperty('access_token');
|
||||
expect(result).toHaveProperty('refresh_token');
|
||||
expect(result).toHaveProperty('expires_in', 3600);
|
||||
expect(result).toHaveProperty('token_type', 'Bearer');
|
||||
});
|
||||
|
||||
it("throws error for invalid refresh token", async () => {
|
||||
const client = createTestClient();
|
||||
|
||||
await expect(provider.exchangeRefreshToken(client, "test-refresh-token"))
|
||||
.rejects
|
||||
.toThrow("Invalid refresh token");
|
||||
});
|
||||
|
||||
it("throws error for refresh token with no installation", async () => {
|
||||
const client = createTestClient();
|
||||
const refreshToken = authService.generateToken();
|
||||
const accessToken = authService.generateToken();
|
||||
|
||||
// Only save refresh token, not the installation
|
||||
await authService.saveRefreshToken(refreshToken, accessToken);
|
||||
|
||||
await expect(provider.exchangeRefreshToken(client, refreshToken))
|
||||
.rejects
|
||||
.toThrow("Invalid refresh token");
|
||||
});
|
||||
|
||||
it("throws error when client ID doesn't match", async () => {
|
||||
const client = createTestClient();
|
||||
const accessToken = authService.generateToken();
|
||||
const refreshToken = authService.generateToken();
|
||||
|
||||
const mcpInstallation: McpInstallation = {
|
||||
mockUpstreamInstallation: {
|
||||
mockUpstreamAccessToken: "fake-upstream-access-token",
|
||||
mockUpstreamRefreshToken: "fake-upstream-refresh-token",
|
||||
},
|
||||
mcpTokens: {
|
||||
access_token: accessToken,
|
||||
token_type: "Bearer",
|
||||
expires_in: 3600,
|
||||
},
|
||||
clientId: "different-client-id",
|
||||
issuedAt: Date.now() / 1000,
|
||||
userId: "test-user-id",
|
||||
};
|
||||
|
||||
await authService.saveRefreshToken(refreshToken, accessToken);
|
||||
await authService.saveMcpInstallation(accessToken, mcpInstallation);
|
||||
|
||||
await expect(provider.exchangeRefreshToken(client, refreshToken))
|
||||
.rejects
|
||||
.toThrow("Invalid client");
|
||||
});
|
||||
|
||||
it("works correctly even when the refresh token has expired from Redis", async () => {
|
||||
const client = createTestClient();
|
||||
|
||||
// Simulate the refresh token not being found in Redis (expired)
|
||||
await expect(provider.exchangeRefreshToken(client, "expired-refresh-token"))
|
||||
.rejects
|
||||
.toThrow("Invalid refresh token");
|
||||
});
|
||||
});
|
||||
|
||||
describe("verifyAccessToken", () => {
|
||||
it("returns auth info for valid token", async () => {
|
||||
const accessToken = authService.generateToken();
|
||||
const mcpInstallation: McpInstallation = {
|
||||
mockUpstreamInstallation: {
|
||||
mockUpstreamAccessToken: "fake-upstream-access-token",
|
||||
mockUpstreamRefreshToken: "fake-upstream-refresh-token",
|
||||
},
|
||||
mcpTokens: {
|
||||
access_token: accessToken,
|
||||
token_type: "Bearer",
|
||||
expires_in: 3600,
|
||||
},
|
||||
clientId: "client-id",
|
||||
issuedAt: Date.now() / 1000,
|
||||
userId: "test-user-id",
|
||||
};
|
||||
|
||||
await authService.saveMcpInstallation(accessToken, mcpInstallation);
|
||||
|
||||
const result = await provider.verifyAccessToken(accessToken);
|
||||
|
||||
expect(result).toEqual({
|
||||
token: accessToken,
|
||||
clientId: mcpInstallation.clientId,
|
||||
scopes: ['mcp'],
|
||||
expiresAt: mcpInstallation.mcpTokens.expires_in! + mcpInstallation.issuedAt,
|
||||
extra: {
|
||||
userId: "test-user-id"
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
it("throws error for invalid token", async () => {
|
||||
await expect(provider.verifyAccessToken("invalid-token"))
|
||||
.rejects
|
||||
.toThrow("Invalid access token");
|
||||
});
|
||||
|
||||
it("throws InvalidTokenError for expired token", async () => {
|
||||
const accessToken = authService.generateToken();
|
||||
const oneDayInSeconds = 24 * 60 * 60;
|
||||
const twoDaysAgoInSeconds = Math.floor(Date.now() / 1000) - (2 * oneDayInSeconds);
|
||||
|
||||
const mcpInstallation: McpInstallation = {
|
||||
mockUpstreamInstallation: {
|
||||
mockUpstreamAccessToken: "fake-upstream-access-token",
|
||||
mockUpstreamRefreshToken: "fake-upstream-refresh-token",
|
||||
},
|
||||
mcpTokens: {
|
||||
access_token: accessToken,
|
||||
token_type: "Bearer",
|
||||
expires_in: oneDayInSeconds,
|
||||
},
|
||||
clientId: "client-id",
|
||||
issuedAt: twoDaysAgoInSeconds, // 2 days ago, with 1-day expiry
|
||||
userId: "test-user-id",
|
||||
};
|
||||
|
||||
await authService.saveMcpInstallation(accessToken, mcpInstallation);
|
||||
|
||||
await expect(provider.verifyAccessToken(accessToken))
|
||||
.rejects
|
||||
.toThrow(InvalidTokenError);
|
||||
});
|
||||
});
|
||||
|
||||
describe("revokeToken", () => {
|
||||
it("revokes the installation when given an access token", async () => {
|
||||
const client = createTestClient();
|
||||
const accessToken = authService.generateToken();
|
||||
const mcpInstallation: McpInstallation = {
|
||||
mockUpstreamInstallation: {
|
||||
mockUpstreamAccessToken: "fake-upstream-access-token",
|
||||
mockUpstreamRefreshToken: "fake-upstream-refresh-token",
|
||||
},
|
||||
mcpTokens: {
|
||||
access_token: accessToken,
|
||||
token_type: "Bearer",
|
||||
expires_in: 3600,
|
||||
},
|
||||
clientId: client.client_id,
|
||||
issuedAt: Date.now() / 1000,
|
||||
userId: "test-user-id",
|
||||
};
|
||||
|
||||
// Save the installation
|
||||
await authService.saveMcpInstallation(accessToken, mcpInstallation);
|
||||
|
||||
// Verify it exists
|
||||
const saved = await authService.readMcpInstallation(accessToken);
|
||||
expect(saved).toBeTruthy();
|
||||
|
||||
// Revoke it
|
||||
const request: OAuthTokenRevocationRequest = {
|
||||
token: accessToken,
|
||||
token_type_hint: "access_token"
|
||||
};
|
||||
|
||||
await provider.revokeToken(client, request);
|
||||
|
||||
// Verify it was revoked
|
||||
const revoked = await authService.readMcpInstallation(accessToken);
|
||||
expect(revoked).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
378
src/modules/auth/auth/provider.ts
Normal file
378
src/modules/auth/auth/provider.ts
Normal file
|
|
@ -0,0 +1,378 @@
|
|||
import { Response } from 'express';
|
||||
import { OAuthServerProvider, AuthorizationParams } from '@modelcontextprotocol/sdk/server/auth/provider.js';
|
||||
import { OAuthRegisteredClientsStore } from '@modelcontextprotocol/sdk/server/auth/clients.js';
|
||||
import { OAuthClientInformationFull, OAuthTokenRevocationRequest, OAuthTokens } from '@modelcontextprotocol/sdk/shared/auth.js';
|
||||
import { AuthInfo } from '@modelcontextprotocol/sdk/server/auth/types.js';
|
||||
import {
|
||||
exchangeToken,
|
||||
generateToken,
|
||||
getClientRegistration,
|
||||
readPendingAuthorization,
|
||||
readMcpInstallation,
|
||||
revokeMcpInstallation,
|
||||
saveClientRegistration,
|
||||
savePendingAuthorization,
|
||||
readRefreshToken,
|
||||
generateMcpTokens,
|
||||
saveMcpInstallation,
|
||||
saveRefreshToken,
|
||||
} from '../services/auth.js';
|
||||
import { InvalidTokenError } from '@modelcontextprotocol/sdk/server/auth/errors.js';
|
||||
import { logger } from '../../shared/logger.js';
|
||||
|
||||
/**
|
||||
* Implementation of the OAuthRegisteredClientsStore interface using the existing client registration system
|
||||
*/
|
||||
export class FeatureReferenceOAuthClientsStore implements OAuthRegisteredClientsStore {
|
||||
async getClient(clientId: string): Promise<OAuthClientInformationFull | undefined> {
|
||||
const registration = await getClientRegistration(clientId);
|
||||
if (!registration) {
|
||||
return undefined;
|
||||
}
|
||||
return registration;
|
||||
}
|
||||
|
||||
async registerClient(client: OAuthClientInformationFull): Promise<OAuthClientInformationFull> {
|
||||
await saveClientRegistration(client.client_id, client);
|
||||
return client;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Implementation of the OAuthServerProvider interface for upstream authentication
|
||||
*/
|
||||
export class FeatureReferenceAuthProvider implements OAuthServerProvider {
|
||||
private _clientsStore: FeatureReferenceOAuthClientsStore;
|
||||
|
||||
constructor() {
|
||||
this._clientsStore = new FeatureReferenceOAuthClientsStore();
|
||||
}
|
||||
|
||||
get clientsStore(): OAuthRegisteredClientsStore {
|
||||
return this._clientsStore;
|
||||
}
|
||||
|
||||
async authorize(client: OAuthClientInformationFull, params: AuthorizationParams, res: Response): Promise<void> {
|
||||
|
||||
// Client is validated by the MCP sdk.
|
||||
|
||||
// Generate authorization code
|
||||
const authorizationCode = generateToken();
|
||||
|
||||
// Save the pending authorization with code challenge and state
|
||||
await savePendingAuthorization(authorizationCode, {
|
||||
redirectUri: params.redirectUri,
|
||||
codeChallenge: params.codeChallenge,
|
||||
codeChallengeMethod: 'S256', // Currently only support S256
|
||||
clientId: client.client_id,
|
||||
state: params.state,
|
||||
});
|
||||
|
||||
logger.debug('Saved pending authorization', {
|
||||
authorizationCode: authorizationCode.substring(0, 8) + '...',
|
||||
clientId: client.client_id,
|
||||
state: params.state?.substring(0, 8) + '...'
|
||||
});
|
||||
|
||||
// TODO: should we use a different key, other than the authorization code, to store the pending authorization?
|
||||
|
||||
// You can redirect to another page, or you can send an html response directly
|
||||
// res.redirect(new URL(`mock-upstream-idp/authorize?metadata=${authorizationCode}`, BASE_URI).href);
|
||||
|
||||
// Set permissive CSP for styling
|
||||
res.setHeader('Content-Security-Policy', [
|
||||
"default-src 'self'",
|
||||
"style-src 'self' 'unsafe-inline'",
|
||||
"script-src 'self' 'unsafe-inline'",
|
||||
"img-src 'self' data:",
|
||||
"object-src 'none'",
|
||||
"frame-ancestors 'none'",
|
||||
"form-action 'self'",
|
||||
"base-uri 'self'"
|
||||
].join('; '));
|
||||
|
||||
res.send(`
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>MCP Server Authorization</title>
|
||||
<style>
|
||||
* {
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
box-sizing: border-box;
|
||||
}
|
||||
|
||||
body {
|
||||
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
|
||||
background: #000000;
|
||||
color: #ffffff;
|
||||
min-height: 100vh;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
padding: 20px;
|
||||
}
|
||||
|
||||
.auth-container {
|
||||
background: #ffffff;
|
||||
color: #000000;
|
||||
border-radius: 16px;
|
||||
box-shadow: 0 20px 40px rgba(255, 255, 255, 0.1);
|
||||
padding: 40px;
|
||||
max-width: 500px;
|
||||
width: 100%;
|
||||
text-align: center;
|
||||
border: 1px solid #e2e8f0;
|
||||
}
|
||||
|
||||
.logo-container {
|
||||
margin-bottom: 32px;
|
||||
}
|
||||
|
||||
.logo {
|
||||
width: 80px;
|
||||
height: 80px;
|
||||
margin: 0 auto 16px;
|
||||
filter: invert(1);
|
||||
}
|
||||
|
||||
.mcp-title {
|
||||
font-size: 24px;
|
||||
font-weight: 700;
|
||||
color: #000000;
|
||||
margin-bottom: 8px;
|
||||
letter-spacing: 2px;
|
||||
}
|
||||
|
||||
h1 {
|
||||
color: #000000;
|
||||
font-size: 32px;
|
||||
font-weight: 800;
|
||||
margin-bottom: 12px;
|
||||
line-height: 1.2;
|
||||
}
|
||||
|
||||
.subtitle {
|
||||
color: #4a5568;
|
||||
font-size: 18px;
|
||||
margin-bottom: 32px;
|
||||
line-height: 1.5;
|
||||
}
|
||||
|
||||
.client-info {
|
||||
background: #f8f9fa;
|
||||
border-radius: 12px;
|
||||
padding: 24px;
|
||||
margin-bottom: 32px;
|
||||
border: 2px solid #e2e8f0;
|
||||
}
|
||||
|
||||
.client-info h3 {
|
||||
color: #2d3748;
|
||||
font-size: 18px;
|
||||
font-weight: 600;
|
||||
margin-bottom: 16px;
|
||||
}
|
||||
|
||||
.client-id {
|
||||
background: white;
|
||||
border: 2px solid #e2e8f0;
|
||||
border-radius: 8px;
|
||||
padding: 12px;
|
||||
font-family: 'Courier New', monospace;
|
||||
font-size: 14px;
|
||||
color: #4a5568;
|
||||
word-break: break-all;
|
||||
}
|
||||
|
||||
.auth-flow-info {
|
||||
background: #f8f9fa;
|
||||
border-radius: 12px;
|
||||
padding: 24px;
|
||||
margin-bottom: 32px;
|
||||
border-left: 4px solid #000000;
|
||||
}
|
||||
|
||||
.auth-flow-info h3 {
|
||||
color: #2d3748;
|
||||
font-size: 18px;
|
||||
font-weight: 600;
|
||||
margin-bottom: 12px;
|
||||
}
|
||||
|
||||
.auth-flow-info p {
|
||||
color: #4a5568;
|
||||
font-size: 14px;
|
||||
line-height: 1.5;
|
||||
}
|
||||
|
||||
.btn-primary {
|
||||
background: #000000;
|
||||
color: #ffffff;
|
||||
font-size: 18px;
|
||||
font-weight: 700;
|
||||
padding: 18px 36px;
|
||||
border: none;
|
||||
border-radius: 8px;
|
||||
cursor: pointer;
|
||||
transition: all 0.2s ease;
|
||||
width: 100%;
|
||||
text-decoration: none;
|
||||
display: inline-block;
|
||||
text-align: center;
|
||||
letter-spacing: 1px;
|
||||
}
|
||||
|
||||
.btn-primary:hover {
|
||||
background: #333333;
|
||||
transform: translateY(-2px);
|
||||
box-shadow: 0 8px 25px rgba(0, 0, 0, 0.3);
|
||||
}
|
||||
|
||||
.branding {
|
||||
margin-top: 24px;
|
||||
padding-top: 24px;
|
||||
border-top: 1px solid #e2e8f0;
|
||||
color: #718096;
|
||||
font-size: 12px;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="auth-container">
|
||||
<div class="logo-container">
|
||||
<img src="/mcp-logo.png" alt="MCP Logo" class="logo">
|
||||
<div class="mcp-title">MCP</div>
|
||||
</div>
|
||||
|
||||
<h1>Authorization Required</h1>
|
||||
<p class="subtitle">This client wants to connect to your MCP server</p>
|
||||
|
||||
<div class="client-info">
|
||||
<h3>Client Application</h3>
|
||||
<div class="client-id">${client.client_id}</div>
|
||||
</div>
|
||||
|
||||
<div class="auth-flow-info">
|
||||
<h3>What happens next?</h3>
|
||||
<p>You'll be redirected to authenticate with the upstream provider. Once verified, you'll be granted access to this MCP server's resources.</p>
|
||||
</div>
|
||||
|
||||
<a href="/mock-upstream-idp/authorize?redirect_uri=/mock-upstream-idp/callback&state=${authorizationCode}" class="btn-primary">
|
||||
Continue to Authentication
|
||||
</a>
|
||||
|
||||
<div class="branding">
|
||||
Model Context Protocol (MCP) Server
|
||||
</div>
|
||||
</div>
|
||||
</body>
|
||||
</html>
|
||||
`);
|
||||
}
|
||||
|
||||
async challengeForAuthorizationCode(client: OAuthClientInformationFull, authorizationCode: string): Promise<string> {
|
||||
const pendingAuth = await readPendingAuthorization(authorizationCode);
|
||||
if (!pendingAuth) {
|
||||
throw new Error('Authorization code not found');
|
||||
}
|
||||
|
||||
if (pendingAuth.clientId !== client.client_id) {
|
||||
throw new Error('Authorization code does not match client');
|
||||
}
|
||||
|
||||
return pendingAuth.codeChallenge;
|
||||
}
|
||||
|
||||
async exchangeAuthorizationCode(client: OAuthClientInformationFull, authorizationCode: string): Promise<OAuthTokens> {
|
||||
const tokenData = await exchangeToken(authorizationCode);
|
||||
if (!tokenData) {
|
||||
throw new Error('Invalid authorization code');
|
||||
}
|
||||
|
||||
// Get the MCP installation to retrieve the full token data including refresh token
|
||||
const mcpInstallation = await readMcpInstallation(tokenData.mcpAccessToken);
|
||||
if (!mcpInstallation) {
|
||||
throw new Error('Failed to retrieve MCP installation');
|
||||
}
|
||||
|
||||
// Return the full token data including refresh token
|
||||
return {
|
||||
access_token: mcpInstallation.mcpTokens.access_token,
|
||||
refresh_token: mcpInstallation.mcpTokens.refresh_token,
|
||||
expires_in: mcpInstallation.mcpTokens.expires_in,
|
||||
token_type: 'Bearer',
|
||||
};
|
||||
}
|
||||
|
||||
async exchangeRefreshToken(client: OAuthClientInformationFull, refreshToken: string, _scopes?: string[]): Promise<OAuthTokens> {
|
||||
const accessToken = await readRefreshToken(refreshToken);
|
||||
|
||||
if (!accessToken) {
|
||||
throw new Error('Invalid refresh token');
|
||||
}
|
||||
|
||||
const mcpInstallation = await readMcpInstallation(accessToken);
|
||||
|
||||
if (!mcpInstallation) {
|
||||
throw new Error('Invalid refresh token');
|
||||
}
|
||||
|
||||
// Check the client_id
|
||||
if (mcpInstallation.clientId !== client.client_id) {
|
||||
throw new Error('Invalid client');
|
||||
}
|
||||
|
||||
const newTokens = generateMcpTokens();
|
||||
|
||||
if (newTokens.refresh_token) {
|
||||
await saveRefreshToken(newTokens.refresh_token, newTokens.access_token);
|
||||
}
|
||||
|
||||
// Update the installation with the new tokens
|
||||
await saveMcpInstallation(newTokens.access_token, {
|
||||
...mcpInstallation,
|
||||
mcpTokens: newTokens,
|
||||
issuedAt: Date.now() / 1000,
|
||||
userId: mcpInstallation.userId, // Preserve the user ID
|
||||
});
|
||||
|
||||
return newTokens;
|
||||
}
|
||||
|
||||
async verifyAccessToken(token: string): Promise<AuthInfo> {
|
||||
const installation = await readMcpInstallation(token);
|
||||
if (!installation) {
|
||||
throw new InvalidTokenError("Invalid access token");
|
||||
}
|
||||
|
||||
const expiresAt = (
|
||||
installation.mcpTokens.expires_in
|
||||
? installation.mcpTokens.expires_in + installation.issuedAt
|
||||
: undefined
|
||||
);
|
||||
|
||||
// This can be removed once in the SDK
|
||||
// Check if the token is expired
|
||||
if (!!expiresAt && expiresAt < Date.now() / 1000) {
|
||||
throw new InvalidTokenError("Token has expired");
|
||||
}
|
||||
|
||||
return {
|
||||
token,
|
||||
clientId: installation.clientId,
|
||||
scopes: ['mcp'],
|
||||
expiresAt,
|
||||
extra: {
|
||||
userId: installation.userId
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
async revokeToken(client: OAuthClientInformationFull, request: OAuthTokenRevocationRequest): Promise<void> {
|
||||
await revokeMcpInstallation(request.token);
|
||||
}
|
||||
}
|
||||
382
src/modules/auth/handlers/mock-upstream-idp.ts
Normal file
382
src/modules/auth/handlers/mock-upstream-idp.ts
Normal file
|
|
@ -0,0 +1,382 @@
|
|||
import { Request, Response } from "express";
|
||||
import { generateMcpTokens, readPendingAuthorization, saveMcpInstallation, saveRefreshToken, saveTokenExchange } from "../services/auth.js";
|
||||
import { McpInstallation } from "../types.js";
|
||||
import { logger } from "../../shared/logger.js";
|
||||
|
||||
/**
|
||||
* ============================================================================
|
||||
* MOCK UPSTREAM IDENTITY PROVIDER - FOR DEMONSTRATION ONLY
|
||||
* ============================================================================
|
||||
*
|
||||
* This file simulates what happens when an OAuth server delegates user
|
||||
* authentication to an external identity provider. In production, this would be:
|
||||
*
|
||||
* - Google OAuth (accounts.google.com)
|
||||
* - GitHub OAuth (github.com/login)
|
||||
* - Corporate SSO (SAML, LDAP, Active Directory)
|
||||
* - Auth0/Okta user database
|
||||
*
|
||||
* The mock implementation:
|
||||
* - Shows a user selection UI
|
||||
* - Generates random user IDs for testing
|
||||
* - Simulates the redirect flow back to the OAuth server
|
||||
*
|
||||
* In production, users would see their actual identity provider's login page
|
||||
* (Google's login, GitHub's login, corporate SSO portal, etc.)
|
||||
*
|
||||
* ============================================================================
|
||||
*/
|
||||
|
||||
/**
|
||||
* Mock authorization endpoint - simulates external IDP login page
|
||||
* In production, this would be replaced by redirecting to:
|
||||
* - https://accounts.google.com/oauth/authorize (Google)
|
||||
* - https://github.com/login/oauth/authorize (GitHub)
|
||||
* - https://login.microsoftonline.com (Azure AD)
|
||||
* - Your corporate SSO login page
|
||||
*/
|
||||
export async function handleMockUpstreamAuthorize(req: Request, res: Response) {
|
||||
// get the redirect_uri and state from the query params
|
||||
const { redirect_uri, state } = req.query;
|
||||
|
||||
// Set a more permissive CSP for auth pages to allow inline styles and scripts
|
||||
res.setHeader('Content-Security-Policy', [
|
||||
"default-src 'self'",
|
||||
"style-src 'self' 'unsafe-inline'", // Allow inline styles for auth page styling
|
||||
"script-src 'self' 'unsafe-inline'", // Allow inline scripts for auth page functionality
|
||||
"object-src 'none'",
|
||||
"frame-ancestors 'none'",
|
||||
"form-action 'self'",
|
||||
"base-uri 'self'"
|
||||
].join('; '));
|
||||
|
||||
res.send(`
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Upstream Provider Authentication</title>
|
||||
<style>
|
||||
* {
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
box-sizing: border-box;
|
||||
}
|
||||
|
||||
body {
|
||||
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
|
||||
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
|
||||
min-height: 100vh;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
padding: 20px;
|
||||
}
|
||||
|
||||
.auth-container {
|
||||
background: white;
|
||||
border-radius: 16px;
|
||||
box-shadow: 0 20px 40px rgba(0, 0, 0, 0.1);
|
||||
padding: 40px;
|
||||
max-width: 480px;
|
||||
width: 100%;
|
||||
text-align: center;
|
||||
}
|
||||
|
||||
.logo {
|
||||
width: 64px;
|
||||
height: 64px;
|
||||
background: linear-gradient(135deg, #667eea, #764ba2);
|
||||
border-radius: 16px;
|
||||
margin: 0 auto 24px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
font-size: 24px;
|
||||
font-weight: bold;
|
||||
color: white;
|
||||
}
|
||||
|
||||
h1 {
|
||||
color: #1a202c;
|
||||
font-size: 28px;
|
||||
font-weight: 700;
|
||||
margin-bottom: 8px;
|
||||
}
|
||||
|
||||
.subtitle {
|
||||
color: #718096;
|
||||
font-size: 16px;
|
||||
margin-bottom: 32px;
|
||||
}
|
||||
|
||||
.user-section {
|
||||
background: #f7fafc;
|
||||
border-radius: 12px;
|
||||
padding: 24px;
|
||||
margin-bottom: 32px;
|
||||
border: 2px solid #e2e8f0;
|
||||
}
|
||||
|
||||
.user-section h3 {
|
||||
color: #2d3748;
|
||||
font-size: 18px;
|
||||
font-weight: 600;
|
||||
margin-bottom: 16px;
|
||||
}
|
||||
|
||||
.user-id-display {
|
||||
background: white;
|
||||
border: 2px solid #e2e8f0;
|
||||
border-radius: 8px;
|
||||
padding: 12px;
|
||||
font-family: 'Courier New', monospace;
|
||||
font-size: 14px;
|
||||
color: #4a5568;
|
||||
margin-bottom: 16px;
|
||||
word-break: break-all;
|
||||
}
|
||||
|
||||
.user-actions {
|
||||
display: flex;
|
||||
gap: 12px;
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
|
||||
.btn {
|
||||
flex: 1;
|
||||
min-width: 120px;
|
||||
padding: 12px 24px;
|
||||
border: none;
|
||||
border-radius: 8px;
|
||||
font-size: 14px;
|
||||
font-weight: 600;
|
||||
cursor: pointer;
|
||||
transition: all 0.2s ease;
|
||||
text-decoration: none;
|
||||
display: inline-flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
}
|
||||
|
||||
.btn-secondary {
|
||||
background: #e2e8f0;
|
||||
color: #4a5568;
|
||||
}
|
||||
|
||||
.btn-secondary:hover {
|
||||
background: #cbd5e0;
|
||||
}
|
||||
|
||||
.btn-primary {
|
||||
background: linear-gradient(135deg, #4299e1, #3182ce);
|
||||
color: white;
|
||||
font-size: 16px;
|
||||
padding: 16px 32px;
|
||||
margin-top: 16px;
|
||||
width: 100%;
|
||||
}
|
||||
|
||||
.btn-primary:hover {
|
||||
background: linear-gradient(135deg, #3182ce, #2c5282);
|
||||
transform: translateY(-1px);
|
||||
box-shadow: 0 4px 12px rgba(66, 153, 225, 0.3);
|
||||
}
|
||||
|
||||
.help-text {
|
||||
color: #718096;
|
||||
font-size: 14px;
|
||||
margin-top: 24px;
|
||||
line-height: 1.5;
|
||||
}
|
||||
|
||||
.help-text strong {
|
||||
color: #4a5568;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="auth-container">
|
||||
<div class="logo">🔒</div>
|
||||
<h1>Upstream Authentication</h1>
|
||||
<p class="subtitle">Please verify your identity with the upstream provider</p>
|
||||
|
||||
<div class="user-section">
|
||||
<h3>Your User Identity</h3>
|
||||
<div class="user-id-display" id="userIdDisplay">Loading...</div>
|
||||
<div class="user-actions">
|
||||
<button class="btn btn-secondary" onclick="generateNewUserId()">Generate New ID</button>
|
||||
<button class="btn btn-secondary" onclick="editUserId()">Edit ID</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<button class="btn btn-primary" onclick="authorize()">
|
||||
Complete Authentication
|
||||
</button>
|
||||
|
||||
<div class="help-text">
|
||||
<strong>Testing Multiple Users:</strong> Open this page in different browser windows or incognito tabs to simulate different users. Each will have their own unique User ID and separate MCP sessions.
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
// Generate UUID v4
|
||||
function generateUUID() {
|
||||
return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, function(c) {
|
||||
var r = Math.random() * 16 | 0, v = c == 'x' ? r : (r & 0x3 | 0x8);
|
||||
return v.toString(16);
|
||||
});
|
||||
}
|
||||
|
||||
// Get or create user ID
|
||||
function getUserId() {
|
||||
let userId = localStorage.getItem('mcpUserId');
|
||||
if (!userId) {
|
||||
userId = generateUUID();
|
||||
localStorage.setItem('mcpUserId', userId);
|
||||
}
|
||||
return userId;
|
||||
}
|
||||
|
||||
// Update the display
|
||||
function updateDisplay() {
|
||||
const userId = getUserId();
|
||||
document.getElementById('userIdDisplay').textContent = userId;
|
||||
}
|
||||
|
||||
// Generate new user ID
|
||||
function generateNewUserId() {
|
||||
const newId = generateUUID();
|
||||
localStorage.setItem('mcpUserId', newId);
|
||||
updateDisplay();
|
||||
}
|
||||
|
||||
// Edit user ID
|
||||
function editUserId() {
|
||||
const currentId = getUserId();
|
||||
const newId = prompt('Enter new User ID:', currentId);
|
||||
if (newId && newId.trim()) {
|
||||
localStorage.setItem('mcpUserId', newId.trim());
|
||||
updateDisplay();
|
||||
}
|
||||
}
|
||||
|
||||
// Authorize with current user ID
|
||||
function authorize() {
|
||||
const userId = getUserId();
|
||||
// Handle relative URLs by making them absolute
|
||||
const redirectUri = '${redirect_uri}';
|
||||
const baseUrl = redirectUri.startsWith('http') ? redirectUri : window.location.origin + redirectUri;
|
||||
const url = new URL(baseUrl);
|
||||
url.searchParams.set('state', '${state}');
|
||||
url.searchParams.set('code', 'mock-auth-code');
|
||||
url.searchParams.set('userId', userId);
|
||||
window.location.href = url.toString();
|
||||
}
|
||||
|
||||
// Initialize on page load
|
||||
updateDisplay();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
`);
|
||||
}
|
||||
|
||||
|
||||
// This is the callback URL that the upstream auth server will redirect to after authorization
|
||||
export async function handleMockUpstreamCallback(req: Request, res: Response) {
|
||||
const {
|
||||
// The state returned from the upstream auth server is actually the authorization code
|
||||
state: mcpAuthorizationCode,
|
||||
code: upstreamAuthorizationCode,
|
||||
userId, // User ID from the authorization flow
|
||||
} = req.query;
|
||||
|
||||
logger.debug('Mock upstream IDP callback received', {
|
||||
mcpAuthorizationCode: typeof mcpAuthorizationCode === 'string' ? mcpAuthorizationCode.substring(0, 8) + '...' : mcpAuthorizationCode,
|
||||
upstreamAuthorizationCode: typeof upstreamAuthorizationCode === 'string' ? upstreamAuthorizationCode.substring(0, 8) + '...' : upstreamAuthorizationCode,
|
||||
userId
|
||||
});
|
||||
|
||||
// This is where you'd exchange the upstreamAuthorizationCode for access/refresh tokens
|
||||
// In this case, we're just going to fake it
|
||||
const upstreamTokens = await mockUpstreamTokenExchange(upstreamAuthorizationCode as string);
|
||||
|
||||
// Validate that it's a string
|
||||
if (typeof mcpAuthorizationCode !== "string") {
|
||||
throw new Error("Invalid authorization code");
|
||||
}
|
||||
|
||||
const pendingAuth = await readPendingAuthorization(mcpAuthorizationCode);
|
||||
logger.debug('Reading pending authorization', {
|
||||
mcpAuthorizationCode: mcpAuthorizationCode.substring(0, 8) + '...',
|
||||
found: !!pendingAuth
|
||||
});
|
||||
|
||||
if (!pendingAuth) {
|
||||
throw new Error("No matching authorization found");
|
||||
}
|
||||
|
||||
logger.debug('Generating MCP tokens');
|
||||
const mcpTokens = generateMcpTokens();
|
||||
logger.debug('MCP tokens generated', {
|
||||
hasAccessToken: !!mcpTokens.access_token,
|
||||
hasRefreshToken: !!mcpTokens.refresh_token
|
||||
});
|
||||
|
||||
const mcpInstallation: McpInstallation = {
|
||||
mockUpstreamInstallation: {
|
||||
mockUpstreamAccessToken: upstreamTokens.access_token,
|
||||
mockUpstreamRefreshToken: upstreamTokens.refresh_token,
|
||||
},
|
||||
mcpTokens,
|
||||
clientId: pendingAuth.clientId,
|
||||
issuedAt: Date.now() / 1000,
|
||||
userId: (userId as string) || 'anonymous-user', // Include user ID from auth flow
|
||||
}
|
||||
|
||||
logger.debug('Saving MCP installation');
|
||||
// Store the upstream authorization data
|
||||
await saveMcpInstallation(mcpTokens.access_token, mcpInstallation);
|
||||
logger.debug('MCP installation saved');
|
||||
|
||||
// Store the refresh token -> access token mapping
|
||||
if (mcpTokens.refresh_token) {
|
||||
logger.debug('Saving refresh token mapping');
|
||||
await saveRefreshToken(mcpTokens.refresh_token, mcpTokens.access_token);
|
||||
logger.debug('Refresh token mapping saved');
|
||||
}
|
||||
|
||||
logger.debug('Saving token exchange data');
|
||||
// Store the token exchange data
|
||||
await saveTokenExchange(mcpAuthorizationCode, {
|
||||
mcpAccessToken: mcpTokens.access_token,
|
||||
alreadyUsed: false,
|
||||
});
|
||||
logger.debug('Token exchange data saved');
|
||||
|
||||
// Redirect back to the original application with the authorization code and state
|
||||
const redirectUrl = pendingAuth.state ?
|
||||
`${pendingAuth.redirectUri}?code=${mcpAuthorizationCode}&state=${pendingAuth.state}` :
|
||||
`${pendingAuth.redirectUri}?code=${mcpAuthorizationCode}`;
|
||||
|
||||
logger.debug('Redirecting to callback', {
|
||||
redirectUrl,
|
||||
hasState: !!pendingAuth.state
|
||||
});
|
||||
res.redirect(redirectUrl);
|
||||
logger.debug('Redirect completed');
|
||||
};
|
||||
|
||||
function mockUpstreamTokenExchange(
|
||||
authorizationCode: string,
|
||||
): Promise<{ access_token: string; refresh_token: string }> {
|
||||
// just return the authorization code with a suffix
|
||||
return new Promise((resolve) => {
|
||||
resolve({
|
||||
access_token: `${authorizationCode}-exchanged-for-access-token`,
|
||||
refresh_token: `${authorizationCode}-exchanged-for-refresh-token`,
|
||||
});
|
||||
});
|
||||
}
|
||||
140
src/modules/auth/index.ts
Normal file
140
src/modules/auth/index.ts
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
/**
|
||||
* Auth Module - Self-contained authentication module
|
||||
*
|
||||
* This module encapsulates all OAuth/authentication functionality.
|
||||
* In internal mode, it runs in-process but maintains architectural separation.
|
||||
* It acts as a stand-in for an external OAuth server (Auth0, Okta, etc).
|
||||
*
|
||||
* IMPORTANT: This is NOT using the deprecated MCP SDK integrated auth pattern.
|
||||
* Even in internal mode, the auth module is architecturally separate from MCP.
|
||||
*/
|
||||
|
||||
import { Router } from 'express';
|
||||
import express from 'express';
|
||||
import path from 'path';
|
||||
import { fileURLToPath } from 'url';
|
||||
import rateLimit from 'express-rate-limit';
|
||||
import { mcpAuthRouter } from '@modelcontextprotocol/sdk/server/auth/router.js';
|
||||
import { FeatureReferenceAuthProvider } from './auth/provider.js';
|
||||
import { handleMockUpstreamAuthorize, handleMockUpstreamCallback } from './handlers/mock-upstream-idp.js';
|
||||
import { TokenIntrospectionResponse } from '../../interfaces/auth-validator.js';
|
||||
import { logger } from '../shared/logger.js';
|
||||
|
||||
const __filename = fileURLToPath(import.meta.url);
|
||||
const __dirname = path.dirname(__filename);
|
||||
|
||||
export interface AuthConfig {
|
||||
baseUri: string;
|
||||
authServerUrl?: string; // For metadata purposes
|
||||
redisUrl?: string;
|
||||
}
|
||||
|
||||
export class AuthModule {
|
||||
private provider: FeatureReferenceAuthProvider;
|
||||
private router: Router;
|
||||
|
||||
constructor(private config: AuthConfig) {
|
||||
this.provider = new FeatureReferenceAuthProvider();
|
||||
this.router = this.setupRouter();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get Express router with all auth endpoints
|
||||
*/
|
||||
getRouter(): Router {
|
||||
return this.router;
|
||||
}
|
||||
|
||||
/**
|
||||
* Direct token introspection for internal mode
|
||||
* Returns the same format as the /introspect endpoint would
|
||||
*/
|
||||
async introspectToken(token: string): Promise<TokenIntrospectionResponse> {
|
||||
try {
|
||||
const authInfo = await this.provider.verifyAccessToken(token);
|
||||
|
||||
// Return RFC 7662 compliant introspection response
|
||||
return {
|
||||
active: true,
|
||||
client_id: authInfo.clientId,
|
||||
scope: authInfo.scopes.join(' '),
|
||||
exp: authInfo.expiresAt,
|
||||
sub: String(authInfo.extra?.userId || 'unknown'),
|
||||
username: authInfo.extra?.username as string | undefined,
|
||||
aud: this.config.baseUri,
|
||||
iss: this.config.authServerUrl || this.config.baseUri,
|
||||
token_type: 'Bearer'
|
||||
};
|
||||
} catch (error) {
|
||||
logger.debug('Token introspection failed', { error: (error as Error).message });
|
||||
return { active: false };
|
||||
}
|
||||
}
|
||||
|
||||
private setupRouter(): Router {
|
||||
const router = Router();
|
||||
|
||||
// Rate limiters for different route types
|
||||
const authLimiter = rateLimit({
|
||||
windowMs: 60 * 1000, // 1 minute
|
||||
max: 20, // 20 requests per minute for auth endpoints
|
||||
message: 'Too many authentication attempts',
|
||||
standardHeaders: true,
|
||||
legacyHeaders: false,
|
||||
});
|
||||
|
||||
const staticAssetLimiter = rateLimit({
|
||||
windowMs: 60 * 1000, // 1 minute
|
||||
max: 100, // 100 requests per minute for static assets
|
||||
message: 'Too many requests for static assets',
|
||||
standardHeaders: true,
|
||||
legacyHeaders: false,
|
||||
});
|
||||
|
||||
// OAuth endpoints via SDK's mcpAuthRouter
|
||||
router.use(mcpAuthRouter({
|
||||
provider: this.provider,
|
||||
issuerUrl: new URL(this.config.authServerUrl || this.config.baseUri),
|
||||
tokenOptions: {
|
||||
rateLimit: { windowMs: 5000, limit: 100 }
|
||||
},
|
||||
clientRegistrationOptions: {
|
||||
rateLimit: { windowMs: 60000, limit: 10 }
|
||||
}
|
||||
}));
|
||||
|
||||
// Token introspection endpoint (RFC 7662)
|
||||
// This endpoint exists for external mode compatibility
|
||||
router.post('/introspect', express.urlencoded({ extended: false }), async (req, res) => {
|
||||
try {
|
||||
const { token } = req.body;
|
||||
|
||||
if (!token) {
|
||||
return res.status(400).json({
|
||||
error: 'invalid_request',
|
||||
error_description: 'Missing token parameter'
|
||||
});
|
||||
}
|
||||
|
||||
const result = await this.introspectToken(token);
|
||||
res.json(result);
|
||||
|
||||
} catch (error) {
|
||||
logger.error('Introspection endpoint error', error as Error);
|
||||
res.json({ active: false });
|
||||
}
|
||||
});
|
||||
|
||||
// Mock upstream IDP endpoints (for demo purposes)
|
||||
router.get('/mock-upstream-idp/authorize', authLimiter, handleMockUpstreamAuthorize);
|
||||
router.get('/mock-upstream-idp/callback', authLimiter, handleMockUpstreamCallback);
|
||||
|
||||
// Static assets for auth pages
|
||||
router.get('/mcp-logo.png', staticAssetLimiter, (req, res) => {
|
||||
const logoPath = path.join(__dirname, 'static', 'mcp.png');
|
||||
res.sendFile(logoPath);
|
||||
});
|
||||
|
||||
return router;
|
||||
}
|
||||
}
|
||||
292
src/modules/auth/services/auth.test.ts
Normal file
292
src/modules/auth/services/auth.test.ts
Normal file
|
|
@ -0,0 +1,292 @@
|
|||
import { jest } from '@jest/globals';
|
||||
import crypto from "crypto";
|
||||
import {
|
||||
decryptString,
|
||||
exchangeToken,
|
||||
generateMcpTokens,
|
||||
generatePKCEChallenge,
|
||||
generateToken,
|
||||
getClientRegistration,
|
||||
readMcpInstallation,
|
||||
readPendingAuthorization,
|
||||
readRefreshToken,
|
||||
revokeMcpInstallation,
|
||||
saveClientRegistration,
|
||||
saveMcpInstallation,
|
||||
savePendingAuthorization,
|
||||
saveRefreshToken,
|
||||
saveTokenExchange,
|
||||
} from "./auth.js";
|
||||
import { MockRedisClient, setRedisClient } from "../../shared/redis.js";
|
||||
import { McpInstallation, PendingAuthorization, TokenExchange } from "../types.js";
|
||||
import { OAuthClientInformationFull } from "@modelcontextprotocol/sdk/shared/auth.js";
|
||||
|
||||
|
||||
describe("auth utils", () => {
|
||||
let mockRedis: MockRedisClient;
|
||||
|
||||
beforeEach(() => {
|
||||
mockRedis = new MockRedisClient();
|
||||
setRedisClient(mockRedis);
|
||||
|
||||
jest.resetAllMocks();
|
||||
});
|
||||
|
||||
describe("generateToken", () => {
|
||||
it("generates a 64-character hex string", () => {
|
||||
const token = generateToken();
|
||||
expect(token).toMatch(/^[0-9a-f]{64}$/);
|
||||
});
|
||||
|
||||
it("generates unique tokens", () => {
|
||||
const token1 = generateToken();
|
||||
const token2 = generateToken();
|
||||
expect(token1).not.toBe(token2);
|
||||
});
|
||||
});
|
||||
|
||||
describe("generateMcpTokens", () => {
|
||||
it("generates valid token structure", () => {
|
||||
const tokens = generateMcpTokens();
|
||||
|
||||
expect(tokens.access_token).toBeDefined();
|
||||
expect(tokens.access_token).toMatch(/^[0-9a-f]{64}$/);
|
||||
|
||||
expect(tokens.refresh_token).toBeDefined();
|
||||
expect(tokens.refresh_token).toMatch(/^[0-9a-f]{64}$/);
|
||||
|
||||
expect(tokens.token_type).toBe("Bearer");
|
||||
expect(tokens.expires_in).toBeDefined();
|
||||
});
|
||||
|
||||
it("generates unique tokens", () => {
|
||||
const tokens1 = generateMcpTokens();
|
||||
const tokens2 = generateMcpTokens();
|
||||
|
||||
expect(tokens1.access_token).not.toBe(tokens2.access_token);
|
||||
expect(tokens1.refresh_token).not.toBe(tokens2.refresh_token);
|
||||
});
|
||||
});
|
||||
|
||||
describe("client registration", () => {
|
||||
it("saves and retrieves client registration", async () => {
|
||||
const clientId = "test-client-id";
|
||||
const registration: OAuthClientInformationFull = {
|
||||
client_id: clientId,
|
||||
client_name: "Test Client",
|
||||
client_uri: "https://example.com",
|
||||
redirect_uris: ["https://example.com/callback"]
|
||||
};
|
||||
|
||||
await saveClientRegistration(clientId, registration);
|
||||
const retrieved = await getClientRegistration(clientId);
|
||||
|
||||
expect(retrieved).toEqual(registration);
|
||||
});
|
||||
|
||||
it("returns undefined for non-existent client", async () => {
|
||||
const result = await getClientRegistration("non-existent");
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("pending authorization", () => {
|
||||
it("saves and retrieves pending authorization", async () => {
|
||||
const authCode = generateToken();
|
||||
const pendingAuth: PendingAuthorization = {
|
||||
redirectUri: "https://example.com/callback",
|
||||
codeChallenge: "test-challenge",
|
||||
codeChallengeMethod: "S256",
|
||||
clientId: "test-client-id"
|
||||
};
|
||||
|
||||
await savePendingAuthorization(authCode, pendingAuth);
|
||||
const retrieved = await readPendingAuthorization(authCode);
|
||||
|
||||
expect(retrieved).toEqual(pendingAuth);
|
||||
});
|
||||
|
||||
it("returns undefined for non-existent code", async () => {
|
||||
const result = await readPendingAuthorization("non-existent");
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("token exchange", () => {
|
||||
it("saves and retrieves token exchange data", async () => {
|
||||
const authCode = generateToken();
|
||||
const tokenExchange: TokenExchange = {
|
||||
mcpAccessToken: generateToken(),
|
||||
alreadyUsed: false
|
||||
};
|
||||
|
||||
// For this test, we'll directly manipulate Redis to check the token
|
||||
// instead of using exchangeToken which changes the value
|
||||
await saveTokenExchange(authCode, tokenExchange);
|
||||
|
||||
// Get the key used by saveTokenExchange (now with auth: prefix)
|
||||
const key = "auth:exch:" + crypto.createHash("sha256").update(authCode).digest("hex");
|
||||
|
||||
// Get the encrypted data
|
||||
const encryptedData = await mockRedis.get(key);
|
||||
expect(encryptedData).not.toBeNull();
|
||||
|
||||
// Decrypt it manually to verify
|
||||
const decoded = JSON.parse(decryptString({
|
||||
encryptedText: encryptedData!,
|
||||
key: authCode
|
||||
}));
|
||||
|
||||
expect(decoded).toEqual(tokenExchange);
|
||||
});
|
||||
|
||||
it("prevents duplicate token exchange", async () => {
|
||||
const authCode = generateToken();
|
||||
const tokenExchange: TokenExchange = {
|
||||
mcpAccessToken: generateToken(),
|
||||
alreadyUsed: false
|
||||
};
|
||||
|
||||
await saveTokenExchange(authCode, tokenExchange);
|
||||
|
||||
// First exchange succeeds
|
||||
const first = await exchangeToken(authCode);
|
||||
expect(first).toBeDefined();
|
||||
|
||||
// Mock console.error to suppress expected error message
|
||||
const consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation(() => {});
|
||||
|
||||
// Second exchange throws
|
||||
await expect(exchangeToken(authCode)).rejects.toThrow(
|
||||
"Duplicate use of authorization code detected"
|
||||
);
|
||||
|
||||
// Restore console.error
|
||||
consoleErrorSpy.mockRestore();
|
||||
});
|
||||
|
||||
it("returns undefined for non-existent code", async () => {
|
||||
const result = await exchangeToken("non-existent");
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("MCP installation", () => {
|
||||
it("returns undefined for missing tokens", async () => {
|
||||
const accessToken = generateToken();
|
||||
const result = await readMcpInstallation(accessToken);
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
|
||||
it("saves and retrieves installation data", async () => {
|
||||
const accessToken = generateToken();
|
||||
|
||||
const mcpInstallation: McpInstallation = {
|
||||
mockUpstreamInstallation: {
|
||||
mockUpstreamAccessToken: "fake-upstream-access-token",
|
||||
mockUpstreamRefreshToken: "fake-upstream-refresh-token",
|
||||
},
|
||||
mcpTokens: {
|
||||
access_token: accessToken,
|
||||
token_type: "Bearer",
|
||||
expires_in: 3600,
|
||||
},
|
||||
clientId: "client-id",
|
||||
issuedAt: Date.now() / 1000,
|
||||
userId: "test-user-id",
|
||||
}
|
||||
|
||||
await saveMcpInstallation(accessToken, mcpInstallation);
|
||||
|
||||
const result = await readMcpInstallation(accessToken);
|
||||
expect(result).toEqual(mcpInstallation);
|
||||
});
|
||||
});
|
||||
|
||||
describe("refresh token", () => {
|
||||
it("saves and retrieves refresh token mapping", async () => {
|
||||
const refreshToken = generateToken();
|
||||
const accessToken = generateToken();
|
||||
|
||||
await saveRefreshToken(refreshToken, accessToken);
|
||||
const retrieved = await readRefreshToken(refreshToken);
|
||||
|
||||
expect(retrieved).toBe(accessToken);
|
||||
});
|
||||
|
||||
it("returns undefined for non-existent refresh token", async () => {
|
||||
const result = await readRefreshToken("non-existent");
|
||||
expect(result).toBeUndefined();
|
||||
});
|
||||
});
|
||||
|
||||
describe("revokeMcpInstallation", () => {
|
||||
it("revokes token for valid installation", async () => {
|
||||
// For this test, we'll simply test if the WebClient is called with our mocked token
|
||||
|
||||
// Create mock installation
|
||||
const accessToken = generateToken();
|
||||
|
||||
// Save it to Redis with actual function
|
||||
await saveMcpInstallation(accessToken, {
|
||||
mockUpstreamInstallation: {
|
||||
mockUpstreamAccessToken: "fake-upstream-access-token",
|
||||
mockUpstreamRefreshToken: "fake-upstream-refresh-token",
|
||||
},
|
||||
mcpTokens: {
|
||||
access_token: accessToken,
|
||||
token_type: "Bearer",
|
||||
expires_in: 3600,
|
||||
},
|
||||
clientId: "client-id",
|
||||
issuedAt: Date.now() / 1000,
|
||||
userId: "test-user-id",
|
||||
});
|
||||
|
||||
const getDel = jest.spyOn(mockRedis, 'getDel').mockImplementationOnce(() => {
|
||||
// Need to return encrypted data for successful decryption in the revoke function
|
||||
// Create encrypted data using our access token
|
||||
const mcpInstallation = {
|
||||
mcpTokens: {
|
||||
access_token: accessToken,
|
||||
token_type: "Bearer",
|
||||
expires_in: 3600,
|
||||
},
|
||||
clientId: "client-id",
|
||||
issuedAt: Date.now() / 1000,
|
||||
userId: "test-user-id",
|
||||
};
|
||||
const value = JSON.stringify(mcpInstallation);
|
||||
const iv = crypto.randomBytes(16);
|
||||
const cipher = crypto.createCipheriv("aes-256-cbc", Buffer.from(accessToken, "hex"), iv);
|
||||
let encrypted = cipher.update(value, "utf-8", "hex");
|
||||
encrypted += cipher.final("hex");
|
||||
return Promise.resolve(`${iv.toString("hex")}:${encrypted}`);
|
||||
});
|
||||
|
||||
await revokeMcpInstallation(accessToken);
|
||||
|
||||
// Should have called getDel with the correct key (now auth:installation:)
|
||||
expect(getDel).toHaveBeenCalledWith(expect.stringContaining("auth:installation:"));
|
||||
|
||||
});
|
||||
|
||||
it("handles non-existent installation without error", async () => {
|
||||
await expect(revokeMcpInstallation("non-existent")).resolves.not.toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe("generatePKCEChallenge", () => {
|
||||
it("generates base64url-encoded SHA256 hash", () => {
|
||||
const verifier = "test_verifier";
|
||||
const challenge = generatePKCEChallenge(verifier);
|
||||
expect(challenge).toBe("0Ku4rR8EgR1w3HyHLBCxVLtPsAAks5HOlpmTEt0XhVA");
|
||||
});
|
||||
|
||||
it("generates different challenges for different verifiers", () => {
|
||||
const challenge1 = generatePKCEChallenge("verifier1");
|
||||
const challenge2 = generatePKCEChallenge("verifier2");
|
||||
expect(challenge1).not.toBe(challenge2);
|
||||
});
|
||||
});
|
||||
});
|
||||
86
src/modules/auth/services/auth.ts
Normal file
86
src/modules/auth/services/auth.ts
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
import { redisClient } from "../../shared/redis.js";
|
||||
import { McpInstallation, PendingAuthorization, TokenExchange } from "../types.js";
|
||||
import { OAuthClientInformationFull } from "@modelcontextprotocol/sdk/shared/auth.js";
|
||||
|
||||
// Re-export from auth-core module
|
||||
export {
|
||||
generatePKCEChallenge,
|
||||
generateToken,
|
||||
decryptString,
|
||||
generateMcpTokens
|
||||
} from "../auth/auth-core.js";
|
||||
|
||||
import * as sharedRedisAuth from "./redis-auth.js";
|
||||
|
||||
// Wrapper functions that pass redisClient to shared module functions
|
||||
|
||||
export async function saveClientRegistration(
|
||||
clientId: string,
|
||||
registration: OAuthClientInformationFull,
|
||||
) {
|
||||
return sharedRedisAuth.saveClientRegistration(redisClient, clientId, registration);
|
||||
}
|
||||
|
||||
export async function getClientRegistration(
|
||||
clientId: string,
|
||||
): Promise<OAuthClientInformationFull | undefined> {
|
||||
return sharedRedisAuth.getClientRegistration(redisClient, clientId);
|
||||
}
|
||||
|
||||
export async function savePendingAuthorization(
|
||||
authorizationCode: string,
|
||||
pendingAuthorization: PendingAuthorization,
|
||||
) {
|
||||
return sharedRedisAuth.savePendingAuthorization(redisClient, authorizationCode, pendingAuthorization);
|
||||
}
|
||||
|
||||
export async function readPendingAuthorization(
|
||||
authorizationCode: string,
|
||||
): Promise<PendingAuthorization | undefined> {
|
||||
return sharedRedisAuth.readPendingAuthorization(redisClient, authorizationCode);
|
||||
}
|
||||
|
||||
export async function saveMcpInstallation(
|
||||
mcpAccessToken: string,
|
||||
installation: McpInstallation,
|
||||
) {
|
||||
return sharedRedisAuth.saveMcpInstallation(redisClient, mcpAccessToken, installation);
|
||||
}
|
||||
|
||||
export async function readMcpInstallation(
|
||||
mcpAccessToken: string,
|
||||
): Promise<McpInstallation | undefined> {
|
||||
return sharedRedisAuth.readMcpInstallation(redisClient, mcpAccessToken);
|
||||
}
|
||||
|
||||
export async function saveRefreshToken(
|
||||
refreshToken: string,
|
||||
mcpAccessToken: string,
|
||||
) {
|
||||
return sharedRedisAuth.saveRefreshToken(redisClient, refreshToken, mcpAccessToken);
|
||||
}
|
||||
|
||||
export async function readRefreshToken(
|
||||
refreshToken: string,
|
||||
): Promise<string | undefined> {
|
||||
return sharedRedisAuth.readRefreshToken(redisClient, refreshToken);
|
||||
}
|
||||
|
||||
export async function revokeMcpInstallation(
|
||||
mcpAccessToken: string,
|
||||
): Promise<void> {
|
||||
return sharedRedisAuth.revokeMcpInstallation(redisClient, mcpAccessToken);
|
||||
}
|
||||
|
||||
export async function saveTokenExchange(
|
||||
authorizationCode: string,
|
||||
tokenExchange: TokenExchange,
|
||||
) {
|
||||
return sharedRedisAuth.saveTokenExchange(redisClient, authorizationCode, tokenExchange);
|
||||
}
|
||||
|
||||
export async function exchangeToken(
|
||||
authorizationCode: string,
|
||||
): Promise<TokenExchange | undefined> {
|
||||
return sharedRedisAuth.exchangeToken(redisClient, authorizationCode);
|
||||
}
|
||||
285
src/modules/auth/services/redis-auth.ts
Normal file
285
src/modules/auth/services/redis-auth.ts
Normal file
|
|
@ -0,0 +1,285 @@
|
|||
import { SetOptions } from "@redis/client";
|
||||
import { OAuthClientInformationFull } from "@modelcontextprotocol/sdk/shared/auth.js";
|
||||
import { RedisClient } from "../../shared/redis.js";
|
||||
import { McpInstallation, PendingAuthorization, TokenExchange } from "../types.js";
|
||||
import { sha256, encryptString, decryptString } from "../auth/auth-core.js";
|
||||
import { logger } from "../../shared/logger.js";
|
||||
|
||||
/**
|
||||
* Redis key prefixes for different data types
|
||||
* All auth-related keys use "auth:" prefix to avoid collision with MCP session keys
|
||||
*/
|
||||
export const REDIS_KEY_PREFIXES = {
|
||||
CLIENT_REGISTRATION: "auth:client:",
|
||||
PENDING_AUTHORIZATION: "auth:pending:",
|
||||
MCP_AUTHORIZATION: "auth:installation:", // Changed from "mcp:" to avoid collision
|
||||
TOKEN_EXCHANGE: "auth:exch:",
|
||||
REFRESH_TOKEN: "auth:refresh:",
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Redis key expiry times in seconds
|
||||
*/
|
||||
export const REDIS_EXPIRY_TIMES = {
|
||||
CLIENT_REGISTRATION: 30 * 24 * 60 * 60, // 30 days - client app credentials
|
||||
PENDING_AUTHORIZATION: 10 * 60, // 10 minutes - authorization code -> PendingAuthorization
|
||||
TOKEN_EXCHANGE: 10 * 60, // 10 minutes - authorization code -> MCP access token
|
||||
UPSTREAM_INSTALLATION: 7 * 24 * 60 * 60, // 7 days - MCP access token -> UpstreamInstallation
|
||||
REFRESH_TOKEN: 7 * 24 * 60 * 60, // 7 days - MCP refresh token -> access token
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Saves encrypted data to Redis with optional expiry.
|
||||
*/
|
||||
async function saveEncrypted<T>(
|
||||
redisClient: RedisClient,
|
||||
{
|
||||
prefix,
|
||||
key,
|
||||
data,
|
||||
options,
|
||||
}: {
|
||||
prefix: string;
|
||||
key: string;
|
||||
data: T;
|
||||
options?: SetOptions;
|
||||
}
|
||||
): Promise<string | null> {
|
||||
const value = encryptString({
|
||||
text: JSON.stringify(data),
|
||||
key: key,
|
||||
});
|
||||
|
||||
return await redisClient.set(prefix + sha256(key), value, options);
|
||||
}
|
||||
|
||||
/**
|
||||
* Reads and decrypts data from Redis.
|
||||
*/
|
||||
async function readEncrypted<T>(
|
||||
redisClient: RedisClient,
|
||||
{
|
||||
prefix,
|
||||
key,
|
||||
del = false,
|
||||
}: {
|
||||
prefix: string;
|
||||
key: string;
|
||||
del?: boolean;
|
||||
}
|
||||
): Promise<T | undefined> {
|
||||
const data = del
|
||||
? await redisClient.getDel(prefix + sha256(key))
|
||||
: await redisClient.get(prefix + sha256(key));
|
||||
|
||||
if (!data) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const decoded = decryptString({
|
||||
encryptedText: data,
|
||||
key: key,
|
||||
});
|
||||
|
||||
return JSON.parse(decoded);
|
||||
}
|
||||
|
||||
/**
|
||||
* Saves a client registration to Redis.
|
||||
*/
|
||||
export async function saveClientRegistration(
|
||||
redisClient: RedisClient,
|
||||
clientId: string,
|
||||
registration: OAuthClientInformationFull
|
||||
): Promise<void> {
|
||||
await redisClient.set(
|
||||
REDIS_KEY_PREFIXES.CLIENT_REGISTRATION + clientId,
|
||||
JSON.stringify(registration),
|
||||
{ EX: REDIS_EXPIRY_TIMES.CLIENT_REGISTRATION }
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves a client registration from Redis.
|
||||
*/
|
||||
export async function getClientRegistration(
|
||||
redisClient: RedisClient,
|
||||
clientId: string
|
||||
): Promise<OAuthClientInformationFull | undefined> {
|
||||
const data = await redisClient.get(REDIS_KEY_PREFIXES.CLIENT_REGISTRATION + clientId);
|
||||
if (!data) {
|
||||
return undefined;
|
||||
}
|
||||
return JSON.parse(data);
|
||||
}
|
||||
|
||||
/**
|
||||
* Saves a pending authorization to Redis.
|
||||
*/
|
||||
export async function savePendingAuthorization(
|
||||
redisClient: RedisClient,
|
||||
authorizationCode: string,
|
||||
pendingAuthorization: PendingAuthorization
|
||||
): Promise<void> {
|
||||
await saveEncrypted(redisClient, {
|
||||
prefix: REDIS_KEY_PREFIXES.PENDING_AUTHORIZATION,
|
||||
key: authorizationCode,
|
||||
data: pendingAuthorization,
|
||||
options: { EX: REDIS_EXPIRY_TIMES.PENDING_AUTHORIZATION },
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Reads a pending authorization from Redis.
|
||||
*/
|
||||
export async function readPendingAuthorization(
|
||||
redisClient: RedisClient,
|
||||
authorizationCode: string
|
||||
): Promise<PendingAuthorization | undefined> {
|
||||
return readEncrypted<PendingAuthorization>(redisClient, {
|
||||
prefix: REDIS_KEY_PREFIXES.PENDING_AUTHORIZATION,
|
||||
key: authorizationCode,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Saves an MCP installation to Redis.
|
||||
*/
|
||||
export async function saveMcpInstallation(
|
||||
redisClient: RedisClient,
|
||||
mcpAccessToken: string,
|
||||
installation: McpInstallation
|
||||
): Promise<void> {
|
||||
await saveEncrypted(redisClient, {
|
||||
prefix: REDIS_KEY_PREFIXES.MCP_AUTHORIZATION,
|
||||
key: mcpAccessToken,
|
||||
data: installation,
|
||||
options: { EX: REDIS_EXPIRY_TIMES.UPSTREAM_INSTALLATION },
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Reads an MCP installation from Redis.
|
||||
*/
|
||||
export async function readMcpInstallation(
|
||||
redisClient: RedisClient,
|
||||
mcpAccessToken: string
|
||||
): Promise<McpInstallation | undefined> {
|
||||
return readEncrypted<McpInstallation>(redisClient, {
|
||||
prefix: REDIS_KEY_PREFIXES.MCP_AUTHORIZATION,
|
||||
key: mcpAccessToken,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Links a refresh token to an MCP access token.
|
||||
*/
|
||||
export async function saveRefreshToken(
|
||||
redisClient: RedisClient,
|
||||
refreshToken: string,
|
||||
mcpAccessToken: string
|
||||
): Promise<void> {
|
||||
await saveEncrypted(redisClient, {
|
||||
prefix: REDIS_KEY_PREFIXES.REFRESH_TOKEN,
|
||||
key: refreshToken,
|
||||
data: mcpAccessToken,
|
||||
options: { EX: REDIS_EXPIRY_TIMES.REFRESH_TOKEN },
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Reads the access token associated with a refresh token.
|
||||
*/
|
||||
export async function readRefreshToken(
|
||||
redisClient: RedisClient,
|
||||
refreshToken: string
|
||||
): Promise<string | undefined> {
|
||||
return readEncrypted<string>(redisClient, {
|
||||
prefix: REDIS_KEY_PREFIXES.REFRESH_TOKEN,
|
||||
key: refreshToken,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Revokes an MCP installation.
|
||||
*/
|
||||
export async function revokeMcpInstallation(
|
||||
redisClient: RedisClient,
|
||||
mcpAccessToken: string
|
||||
): Promise<void> {
|
||||
const installation = await readEncrypted<McpInstallation>(redisClient, {
|
||||
prefix: REDIS_KEY_PREFIXES.MCP_AUTHORIZATION,
|
||||
key: mcpAccessToken,
|
||||
del: true,
|
||||
});
|
||||
|
||||
if (!installation) {
|
||||
return;
|
||||
}
|
||||
// In production, would revoke upstream tokens here
|
||||
}
|
||||
|
||||
/**
|
||||
* Saves a token exchange record.
|
||||
*/
|
||||
export async function saveTokenExchange(
|
||||
redisClient: RedisClient,
|
||||
authorizationCode: string,
|
||||
tokenExchange: TokenExchange
|
||||
): Promise<void> {
|
||||
await saveEncrypted(redisClient, {
|
||||
prefix: REDIS_KEY_PREFIXES.TOKEN_EXCHANGE,
|
||||
key: authorizationCode,
|
||||
data: tokenExchange,
|
||||
options: { EX: REDIS_EXPIRY_TIMES.TOKEN_EXCHANGE },
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Exchanges a temporary authorization code for an MCP access token.
|
||||
* Will only succeed the first time to prevent replay attacks.
|
||||
*/
|
||||
export async function exchangeToken(
|
||||
redisClient: RedisClient,
|
||||
authorizationCode: string
|
||||
): Promise<TokenExchange | undefined> {
|
||||
const data = await redisClient.get(
|
||||
REDIS_KEY_PREFIXES.TOKEN_EXCHANGE + sha256(authorizationCode)
|
||||
);
|
||||
|
||||
if (!data) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const decoded = decryptString({
|
||||
encryptedText: data,
|
||||
key: authorizationCode,
|
||||
});
|
||||
|
||||
const tokenExchange: TokenExchange = JSON.parse(decoded);
|
||||
if (tokenExchange.alreadyUsed) {
|
||||
logger.error('Duplicate use of authorization code detected; revoking tokens', undefined, {
|
||||
authorizationCode: authorizationCode.substring(0, 8) + '...'
|
||||
});
|
||||
await revokeMcpInstallation(redisClient, tokenExchange.mcpAccessToken);
|
||||
throw new Error("Duplicate use of authorization code detected; tokens revoked");
|
||||
}
|
||||
|
||||
const rereadData = await saveEncrypted(redisClient, {
|
||||
prefix: REDIS_KEY_PREFIXES.TOKEN_EXCHANGE,
|
||||
key: authorizationCode,
|
||||
data: { ...tokenExchange, alreadyUsed: true },
|
||||
options: { KEEPTTL: true, GET: true },
|
||||
});
|
||||
|
||||
if (rereadData !== data) {
|
||||
// Data concurrently changed while we were updating it. This necessarily means a duplicate use.
|
||||
logger.error('Duplicate use of authorization code detected (concurrent update); revoking tokens', undefined, {
|
||||
authorizationCode: authorizationCode.substring(0, 8) + '...'
|
||||
});
|
||||
await revokeMcpInstallation(redisClient, tokenExchange.mcpAccessToken);
|
||||
throw new Error("Duplicate use of authorization code detected; tokens revoked");
|
||||
}
|
||||
|
||||
return tokenExchange;
|
||||
}
|
||||
BIN
src/modules/auth/static/mcp.png
Normal file
BIN
src/modules/auth/static/mcp.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 40 KiB |
92
src/modules/auth/types.ts
Normal file
92
src/modules/auth/types.ts
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
import { OAuthTokens } from "@modelcontextprotocol/sdk/shared/auth.js";
|
||||
|
||||
/**
|
||||
* Represents a pending OAuth authorization that hasn't been exchanged for tokens yet.
|
||||
* Stored in Redis with the authorization code as the key.
|
||||
*/
|
||||
export interface PendingAuthorization {
|
||||
/** The redirect URI where the client expects to receive the authorization code */
|
||||
redirectUri: string;
|
||||
/** PKCE code challenge - a derived value from the code verifier */
|
||||
codeChallenge: string;
|
||||
/** Method used to derive the code challenge (currently only S256 supported) */
|
||||
codeChallengeMethod: string;
|
||||
/** The OAuth client ID that initiated the authorization */
|
||||
clientId: string;
|
||||
/** Optional state parameter for CSRF protection */
|
||||
state?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents the exchange of an authorization code for an MCP access token.
|
||||
* Used to prevent replay attacks by tracking if a code has been used.
|
||||
*/
|
||||
export interface TokenExchange {
|
||||
/** The MCP access token that was issued for this authorization code */
|
||||
mcpAccessToken: string;
|
||||
/** Whether this authorization code has already been exchanged for tokens */
|
||||
alreadyUsed: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents mock upstream identity provider tokens for demonstration purposes.
|
||||
* In production, this would contain real upstream provider tokens.
|
||||
*/
|
||||
export interface MockUpstreamInstallation {
|
||||
/** Access token from the mock upstream identity provider */
|
||||
mockUpstreamAccessToken: string;
|
||||
/** Refresh token from the mock upstream identity provider */
|
||||
mockUpstreamRefreshToken: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* The complete installation object stored in Redis, containing both
|
||||
* upstream provider information and MCP-specific tokens.
|
||||
* This object is encrypted using the MCP access token as the key.
|
||||
*/
|
||||
export interface McpInstallation {
|
||||
/** Information from the upstream authentication provider */
|
||||
mockUpstreamInstallation: MockUpstreamInstallation;
|
||||
/** MCP OAuth tokens issued to the client */
|
||||
mcpTokens: OAuthTokens;
|
||||
/** The OAuth client ID associated with this installation */
|
||||
clientId: string;
|
||||
/** Unix timestamp (seconds) when the tokens were issued */
|
||||
issuedAt: number;
|
||||
/** Unique identifier for the user (not the OAuth client) */
|
||||
userId: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* OAuth 2.0 Token Introspection Response
|
||||
* Based on RFC 7662: https://tools.ietf.org/html/rfc7662
|
||||
* Used when validating tokens with an external authorization server.
|
||||
*/
|
||||
export interface TokenIntrospectionResponse {
|
||||
/** Whether the token is currently active */
|
||||
active: boolean;
|
||||
/** Space-separated list of scopes associated with the token */
|
||||
scope?: string;
|
||||
/** Client identifier for the OAuth client that requested the token */
|
||||
client_id?: string;
|
||||
/** Human-readable identifier for the resource owner */
|
||||
username?: string;
|
||||
/** Type of the token (e.g., "Bearer") */
|
||||
token_type?: string;
|
||||
/** Expiration time as seconds since Unix epoch */
|
||||
exp?: number;
|
||||
/** Time at which the token was issued as seconds since Unix epoch */
|
||||
iat?: number;
|
||||
/** Time before which the token is not valid as seconds since Unix epoch */
|
||||
nbf?: number;
|
||||
/** Subject identifier for the resource owner */
|
||||
sub?: string;
|
||||
/** Intended audience for the token */
|
||||
aud?: string | string[];
|
||||
/** Issuer of the token */
|
||||
iss?: string;
|
||||
/** Unique identifier for the token */
|
||||
jti?: string;
|
||||
/** Custom field for our implementation to store user ID */
|
||||
userId?: string;
|
||||
}
|
||||
694
src/modules/mcp/handlers/shttp.integration.test.ts
Normal file
694
src/modules/mcp/handlers/shttp.integration.test.ts
Normal file
|
|
@ -0,0 +1,694 @@
|
|||
import { jest } from '@jest/globals';
|
||||
import { Request, Response } from 'express';
|
||||
import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { MockRedisClient, setRedisClient } from '../../shared/redis.js';
|
||||
import { handleStreamableHTTP } from './shttp.js';
|
||||
import { AuthInfo } from '@modelcontextprotocol/sdk/server/auth/types.js';
|
||||
// import { randomUUID } from 'crypto'; // Currently unused but may be needed for future tests
|
||||
import { shutdownSession } from '../services/redisTransport.js';
|
||||
|
||||
// Type for MCP initialization response
|
||||
interface MCPInitResponse {
|
||||
jsonrpc: string;
|
||||
id: string | number;
|
||||
result?: {
|
||||
_meta?: {
|
||||
sessionId?: string;
|
||||
};
|
||||
[key: string]: unknown;
|
||||
};
|
||||
}
|
||||
|
||||
describe('Streamable HTTP Handler Integration Tests', () => {
|
||||
let mockRedis: MockRedisClient;
|
||||
let mockReq: Partial<Request>;
|
||||
let mockRes: Partial<Response>;
|
||||
|
||||
beforeEach(() => {
|
||||
mockRedis = new MockRedisClient();
|
||||
setRedisClient(mockRedis);
|
||||
jest.resetAllMocks();
|
||||
|
||||
// Create mock response with chainable methods
|
||||
mockRes = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn().mockReturnThis(),
|
||||
on: jest.fn().mockReturnThis(),
|
||||
once: jest.fn().mockReturnThis(),
|
||||
emit: jest.fn().mockReturnThis(),
|
||||
headersSent: false,
|
||||
setHeader: jest.fn().mockReturnThis(),
|
||||
writeHead: jest.fn().mockReturnThis(),
|
||||
write: jest.fn().mockReturnThis(),
|
||||
end: jest.fn().mockReturnThis(),
|
||||
getHeader: jest.fn(),
|
||||
removeHeader: jest.fn().mockReturnThis(),
|
||||
socket: {
|
||||
setTimeout: jest.fn(),
|
||||
},
|
||||
} as unknown as Partial<Response>;
|
||||
|
||||
// Create mock request
|
||||
mockReq = {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'content-type': 'application/json',
|
||||
'accept': 'application/json, text/event-stream',
|
||||
'mcp-protocol-version': '2024-11-05',
|
||||
},
|
||||
body: {},
|
||||
};
|
||||
});
|
||||
|
||||
// Helper function to trigger cleanup after handleStreamableHTTP calls
|
||||
const triggerResponseCleanup = async () => {
|
||||
// Find all finish handlers registered during the test
|
||||
const finishHandlers = (mockRes.on as jest.Mock).mock.calls
|
||||
.filter(([event]) => event === 'finish')
|
||||
.map(([, handler]) => handler);
|
||||
|
||||
// Trigger all finish handlers
|
||||
for (const handler of finishHandlers) {
|
||||
if (typeof handler === 'function') {
|
||||
await handler();
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Helper to extract session ID from test context
|
||||
const getSessionIdFromTest = (): string | undefined => {
|
||||
// Try to get from response headers first
|
||||
const setHeaderCalls = (mockRes.setHeader as jest.Mock).mock.calls;
|
||||
const sessionIdHeader = setHeaderCalls.find(([name]) => name === 'mcp-session-id');
|
||||
if (sessionIdHeader?.[1]) {
|
||||
return sessionIdHeader[1] as string;
|
||||
}
|
||||
|
||||
// Fall back to extracting from Redis channels
|
||||
const allChannels = Array.from(mockRedis.subscribers.keys());
|
||||
const serverChannel = allChannels.find(channel => channel.includes('mcp:shttp:toserver:'));
|
||||
return serverChannel?.split(':')[3];
|
||||
};
|
||||
|
||||
afterEach(async () => {
|
||||
// Always trigger cleanup for any MCP servers created during tests
|
||||
await triggerResponseCleanup();
|
||||
mockRedis.clear();
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('Redis Subscription Cleanup', () => {
|
||||
it('should clean up Redis subscriptions after shttp response completes', async () => {
|
||||
// Set up initialization request (no session ID for new initialization)
|
||||
const initRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'test-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
|
||||
mockReq.body = initRequest;
|
||||
mockReq.auth = {
|
||||
clientId: 'test-client-123',
|
||||
token: 'test-token',
|
||||
scopes: ['mcp'],
|
||||
extra: { userId: 'test-user-123' }
|
||||
} as AuthInfo;
|
||||
|
||||
// Call the handler
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Wait longer for async initialization to complete
|
||||
await new Promise(resolve => setTimeout(resolve, 200));
|
||||
|
||||
// get the sessionId from the response
|
||||
const sessionId = getSessionIdFromTest();
|
||||
expect(sessionId).toBeDefined();
|
||||
|
||||
// Check if any subscriptions were created on any channels
|
||||
// Since we don't know the exact sessionId generated, check all channels
|
||||
const allChannels = Array.from(mockRedis.subscribers.keys());
|
||||
const totalSubscriptions = allChannels.reduce((sum, channel) => sum + (mockRedis.subscribers.get(channel)?.length || 0), 0);
|
||||
|
||||
// Should have created at least one subscription (server channel)
|
||||
expect(totalSubscriptions).toBeGreaterThan(0);
|
||||
expect(allChannels.some(channel => channel.includes('mcp:shttp:toserver:'))).toBe(true);
|
||||
|
||||
// Find the finish handler that was registered
|
||||
const finishHandler = (mockRes.on as jest.Mock).mock.calls.find(
|
||||
([event]) => event === 'finish'
|
||||
)?.[1] as (() => Promise<void>) | undefined;
|
||||
|
||||
expect(finishHandler).toBeDefined();
|
||||
|
||||
// Simulate response completion to trigger cleanup
|
||||
if (finishHandler) {
|
||||
await finishHandler();
|
||||
}
|
||||
|
||||
// Verify cleanup handler was registered
|
||||
expect(mockRes.on).toHaveBeenCalledWith('finish', expect.any(Function));
|
||||
|
||||
if (sessionId) {
|
||||
await shutdownSession(sessionId)
|
||||
}
|
||||
});
|
||||
|
||||
it('should handle cleanup errors gracefully', async () => {
|
||||
const initRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'test-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
|
||||
mockReq.body = initRequest;
|
||||
mockReq.auth = {
|
||||
clientId: 'test-client-123',
|
||||
token: 'test-token',
|
||||
scopes: ['mcp'],
|
||||
extra: { userId: 'test-user-123' }
|
||||
} as AuthInfo;
|
||||
|
||||
// Call the handler
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Use Redis error simulation to test error handling
|
||||
const finishHandler = (mockRes.on as jest.Mock).mock.calls.find(
|
||||
([event]) => event === 'finish'
|
||||
)?.[1] as (() => Promise<void>) | undefined;
|
||||
|
||||
// Simulate error during cleanup
|
||||
|
||||
// Cleanup should not throw error even if Redis operations fail
|
||||
if (finishHandler) {
|
||||
await expect(finishHandler()).resolves.not.toThrow();
|
||||
}
|
||||
|
||||
// Clean up the MCP server by sending DELETE request
|
||||
const cleanupSessionId = getSessionIdFromTest();
|
||||
|
||||
if (cleanupSessionId) {
|
||||
// Send DELETE request to clean up MCP server
|
||||
jest.clearAllMocks();
|
||||
mockReq.method = 'DELETE';
|
||||
if (mockReq.headers) {
|
||||
mockReq.headers['mcp-session-id'] = cleanupSessionId;
|
||||
}
|
||||
mockReq.body = {};
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Wait a bit for cleanup to complete
|
||||
await new Promise(resolve => setTimeout(resolve, 50));
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('DELETE Request Session Cleanup', () => {
|
||||
it('should trigger onsessionclosed callback which sends shutdown control message', async () => {
|
||||
// First, create a session with an initialization request
|
||||
const initRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'test-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
|
||||
mockReq.body = initRequest;
|
||||
mockReq.auth = {
|
||||
clientId: 'test-client-123',
|
||||
token: 'test-token',
|
||||
scopes: ['mcp'],
|
||||
extra: { userId: 'test-user-123' }
|
||||
} as AuthInfo;
|
||||
|
||||
// Initialize session
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Wait for async initialization
|
||||
await new Promise(resolve => setTimeout(resolve, 100));
|
||||
|
||||
// For initialization requests with StreamableHTTPServerTransport,
|
||||
// the handler might not immediately return a response if using SSE mode
|
||||
// Let's check different possible locations for the session ID
|
||||
|
||||
// Check JSON responses
|
||||
const jsonCalls = (mockRes.json as jest.Mock).mock.calls;
|
||||
let sessionId: string | undefined;
|
||||
|
||||
if (jsonCalls.length > 0) {
|
||||
const response = jsonCalls[0][0] as MCPInitResponse;
|
||||
if (response?.result?._meta?.sessionId) {
|
||||
sessionId = response.result._meta.sessionId;
|
||||
}
|
||||
}
|
||||
|
||||
// Check write calls (for SSE responses)
|
||||
if (!sessionId) {
|
||||
const writeCalls = (mockRes.write as jest.Mock).mock.calls;
|
||||
for (const [data] of writeCalls) {
|
||||
if (typeof data === 'string' && data.includes('sessionId')) {
|
||||
try {
|
||||
// SSE data format: "data: {...}\n\n"
|
||||
const jsonStr = data.replace(/^data: /, '').trim();
|
||||
const parsed = JSON.parse(jsonStr) as MCPInitResponse;
|
||||
if (parsed?.result?._meta?.sessionId) {
|
||||
sessionId = parsed.result._meta.sessionId;
|
||||
}
|
||||
} catch {
|
||||
// Not valid JSON, continue
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to getting from Redis channels
|
||||
if (!sessionId) {
|
||||
sessionId = getSessionIdFromTest();
|
||||
}
|
||||
|
||||
expect(sessionId).toBeDefined();
|
||||
|
||||
// Reset mocks but keep the session
|
||||
jest.clearAllMocks();
|
||||
|
||||
// Now test DELETE request
|
||||
mockReq.method = 'DELETE';
|
||||
mockReq.headers = {
|
||||
...mockReq.headers,
|
||||
'mcp-session-id': sessionId
|
||||
};
|
||||
mockReq.body = {};
|
||||
|
||||
// Track control messages sent to Redis
|
||||
const publishSpy = jest.spyOn(mockRedis, 'publish');
|
||||
|
||||
// Call DELETE handler - StreamableHTTPServerTransport should handle it
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Wait for async processing and onsessionclosed callback
|
||||
await new Promise(resolve => setTimeout(resolve, 100));
|
||||
|
||||
// The StreamableHTTPServerTransport should handle the DELETE and trigger onsessionclosed
|
||||
// which calls shutdownSession, sending the control message
|
||||
const controlCalls = publishSpy.mock.calls.filter(call =>
|
||||
call[0] === `mcp:control:${sessionId}`
|
||||
);
|
||||
|
||||
expect(controlCalls.length).toBeGreaterThan(0);
|
||||
|
||||
// Verify the control message content
|
||||
const controlCall = publishSpy.mock.calls.find(call =>
|
||||
call[0] === `mcp:control:${sessionId}`
|
||||
);
|
||||
if (controlCall) {
|
||||
const message = JSON.parse(controlCall[1]);
|
||||
expect(message.type).toBe('control');
|
||||
expect(message.action).toBe('SHUTDOWN');
|
||||
}
|
||||
});
|
||||
|
||||
it('should return 401 for DELETE request with wrong user', async () => {
|
||||
// First, create a session as user1
|
||||
const initRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'test-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
|
||||
mockReq.body = initRequest;
|
||||
mockReq.auth = {
|
||||
clientId: 'test-client-123',
|
||||
token: 'test-token',
|
||||
scopes: ['mcp'],
|
||||
extra: { userId: 'user1' }
|
||||
} as AuthInfo;
|
||||
|
||||
// Initialize session as user1
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Wait for async initialization
|
||||
await new Promise(resolve => setTimeout(resolve, 100));
|
||||
|
||||
// Get the session ID from response
|
||||
let sessionId: string | undefined;
|
||||
|
||||
// Check JSON responses
|
||||
const jsonCalls = (mockRes.json as jest.Mock).mock.calls;
|
||||
if (jsonCalls.length > 0) {
|
||||
const response = jsonCalls[0][0] as MCPInitResponse;
|
||||
if (response?.result?._meta?.sessionId) {
|
||||
sessionId = response.result._meta.sessionId;
|
||||
}
|
||||
}
|
||||
|
||||
// Check write calls (for SSE responses)
|
||||
if (!sessionId) {
|
||||
const writeCalls = (mockRes.write as jest.Mock).mock.calls;
|
||||
for (const [data] of writeCalls) {
|
||||
if (typeof data === 'string' && data.includes('sessionId')) {
|
||||
try {
|
||||
const jsonStr = data.replace(/^data: /, '').trim();
|
||||
const parsed = JSON.parse(jsonStr) as MCPInitResponse;
|
||||
if (parsed?.result?._meta?.sessionId) {
|
||||
sessionId = parsed.result._meta.sessionId;
|
||||
}
|
||||
} catch {
|
||||
// Ignore JSON parse errors
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!sessionId) {
|
||||
sessionId = getSessionIdFromTest();
|
||||
}
|
||||
|
||||
// Reset mocks
|
||||
jest.clearAllMocks();
|
||||
|
||||
// Now test DELETE request as user2
|
||||
mockReq.method = 'DELETE';
|
||||
mockReq.headers = {
|
||||
...mockReq.headers,
|
||||
'mcp-session-id': sessionId
|
||||
};
|
||||
mockReq.body = {};
|
||||
mockReq.auth = {
|
||||
clientId: 'test-client-456',
|
||||
token: 'test-token-2',
|
||||
scopes: ['mcp'],
|
||||
extra: { userId: 'user2' }
|
||||
} as AuthInfo;
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Should return 401 for unauthorized access to another user's session
|
||||
expect(mockRes.status).toHaveBeenCalledWith(401);
|
||||
|
||||
// shutdown the session
|
||||
if (sessionId) {
|
||||
await shutdownSession(sessionId)
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('User Session Isolation', () => {
|
||||
it('should prevent users from accessing sessions created by other users', async () => {
|
||||
// Create session for user 1
|
||||
const user1Auth: AuthInfo = {
|
||||
clientId: 'user1-client',
|
||||
token: 'user1-token',
|
||||
scopes: ['mcp'],
|
||||
extra: { userId: 'user1' }
|
||||
};
|
||||
|
||||
const user2Auth: AuthInfo = {
|
||||
clientId: 'user2-client',
|
||||
token: 'user2-token',
|
||||
scopes: ['mcp'],
|
||||
extra: { userId: 'user2' }
|
||||
};
|
||||
|
||||
const initRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'user1-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
|
||||
// User 1 creates session
|
||||
mockReq.body = initRequest;
|
||||
mockReq.auth = user1Auth;
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Wait for async initialization to complete
|
||||
await new Promise(resolve => setTimeout(resolve, 100));
|
||||
|
||||
// Get the actual session ID from response
|
||||
let actualSessionId: string | undefined;
|
||||
|
||||
// Check JSON responses
|
||||
const jsonCalls = (mockRes.json as jest.Mock).mock.calls;
|
||||
if (jsonCalls.length > 0) {
|
||||
const response = jsonCalls[0][0] as MCPInitResponse;
|
||||
if (response?.result?._meta?.sessionId) {
|
||||
actualSessionId = response.result._meta.sessionId;
|
||||
}
|
||||
}
|
||||
|
||||
// Check write calls (for SSE responses)
|
||||
if (!actualSessionId) {
|
||||
const writeCalls = (mockRes.write as jest.Mock).mock.calls;
|
||||
for (const [data] of writeCalls) {
|
||||
if (typeof data === 'string' && data.includes('sessionId')) {
|
||||
try {
|
||||
const jsonStr = data.replace(/^data: /, '').trim();
|
||||
const parsed = JSON.parse(jsonStr) as MCPInitResponse;
|
||||
if (parsed?.result?._meta?.sessionId) {
|
||||
actualSessionId = parsed.result._meta.sessionId;
|
||||
}
|
||||
} catch {
|
||||
// Ignore JSON parse errors
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!actualSessionId) {
|
||||
actualSessionId = getSessionIdFromTest();
|
||||
}
|
||||
|
||||
expect(actualSessionId).toBeDefined();
|
||||
|
||||
// Store finish handler before clearing mocks
|
||||
const finishHandler1 = (mockRes.on as jest.Mock).mock.calls.find(
|
||||
([event]) => event === 'finish'
|
||||
)?.[1] as (() => Promise<void>) | undefined;
|
||||
|
||||
// Reset mocks
|
||||
jest.clearAllMocks();
|
||||
|
||||
// Trigger cleanup for the MCP server created in this step
|
||||
if (finishHandler1) {
|
||||
await finishHandler1();
|
||||
}
|
||||
|
||||
// User 2 tries to access user 1's session
|
||||
mockReq.headers = {
|
||||
...mockReq.headers,
|
||||
'mcp-session-id': actualSessionId
|
||||
};
|
||||
mockReq.body = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'user2-request',
|
||||
method: 'tools/list',
|
||||
params: {}
|
||||
};
|
||||
mockReq.auth = user2Auth;
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Should return 401 for unauthorized access to another user's session
|
||||
expect(mockRes.status).toHaveBeenCalledWith(401);
|
||||
|
||||
// Clean up the MCP server by sending DELETE request
|
||||
if (actualSessionId) {
|
||||
jest.clearAllMocks();
|
||||
mockReq.method = 'DELETE';
|
||||
mockReq.headers['mcp-session-id'] = actualSessionId;
|
||||
mockReq.body = {};
|
||||
mockReq.auth = user1Auth; // Use user1's auth to delete their session
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Wait a bit for cleanup to complete
|
||||
await new Promise(resolve => setTimeout(resolve, 50));
|
||||
}
|
||||
});
|
||||
|
||||
it('should allow users to create separate sessions with same session ID pattern', async () => {
|
||||
// This test shows that different users should be able to use sessions
|
||||
// without interfering with each other, even if session IDs might collide
|
||||
|
||||
const user1Auth: AuthInfo = {
|
||||
clientId: 'user1-client',
|
||||
token: 'user1-token',
|
||||
scopes: ['mcp']
|
||||
};
|
||||
|
||||
const user2Auth: AuthInfo = {
|
||||
clientId: 'user2-client',
|
||||
token: 'user2-token',
|
||||
scopes: ['mcp']
|
||||
};
|
||||
|
||||
const initRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'test-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
|
||||
// User 1 creates session
|
||||
mockReq.body = initRequest;
|
||||
mockReq.auth = user1Auth;
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Store finish handler before clearing mocks
|
||||
const finishHandler1 = (mockRes.on as jest.Mock).mock.calls.find(
|
||||
([event]) => event === 'finish'
|
||||
)?.[1] as (() => Promise<void>) | undefined;
|
||||
|
||||
// Reset for user 2
|
||||
jest.clearAllMocks();
|
||||
|
||||
// Trigger cleanup for User 1's MCP server
|
||||
if (finishHandler1) {
|
||||
await finishHandler1();
|
||||
}
|
||||
|
||||
// User 2 creates their own session
|
||||
mockReq.body = {
|
||||
...initRequest,
|
||||
id: 'init-2',
|
||||
params: {
|
||||
...initRequest.params,
|
||||
clientInfo: { name: 'user2-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
mockReq.auth = user2Auth;
|
||||
delete mockReq.headers!['mcp-session-id']; // New initialization
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Trigger cleanup for User 2's MCP server
|
||||
const finishHandler2 = (mockRes.on as jest.Mock).mock.calls.find(
|
||||
([event]) => event === 'finish'
|
||||
)?.[1] as (() => Promise<void>) | undefined;
|
||||
|
||||
if (finishHandler2) {
|
||||
await finishHandler2();
|
||||
}
|
||||
|
||||
// Both users should be able to create sessions successfully
|
||||
// Sessions should be isolated in Redis using user-scoped keys
|
||||
expect(mockRes.status).not.toHaveBeenCalledWith(400);
|
||||
expect(mockRes.status).not.toHaveBeenCalledWith(403);
|
||||
});
|
||||
|
||||
it('should clean up only the requesting user\'s session on DELETE', async () => {
|
||||
// Create sessions for both users
|
||||
const user1Auth: AuthInfo = {
|
||||
clientId: 'user1-client',
|
||||
token: 'user1-token',
|
||||
scopes: ['mcp']
|
||||
};
|
||||
|
||||
const user2Auth: AuthInfo = {
|
||||
clientId: 'user2-client',
|
||||
token: 'user2-token',
|
||||
scopes: ['mcp']
|
||||
};
|
||||
|
||||
// Create session for user 1
|
||||
const initRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'user1-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
|
||||
mockReq.body = initRequest;
|
||||
mockReq.auth = user1Auth;
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Trigger cleanup for User 1's MCP server
|
||||
const finishHandler1 = (mockRes.on as jest.Mock).mock.calls.find(
|
||||
([event]) => event === 'finish'
|
||||
)?.[1] as (() => Promise<void>) | undefined;
|
||||
|
||||
if (finishHandler1) {
|
||||
await finishHandler1();
|
||||
}
|
||||
|
||||
// Track session 1 ID (would be returned in response headers)
|
||||
const session1Id = 'user1-session-id'; // In real implementation, extract from response
|
||||
|
||||
// Create session for user 2
|
||||
jest.clearAllMocks();
|
||||
mockReq.body = {
|
||||
...initRequest,
|
||||
id: 'init-2',
|
||||
params: {
|
||||
...initRequest.params,
|
||||
clientInfo: { name: 'user2-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
mockReq.auth = user2Auth;
|
||||
delete mockReq.headers!['mcp-session-id'];
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Trigger cleanup for User 2's MCP server
|
||||
const finishHandler2 = (mockRes.on as jest.Mock).mock.calls.find(
|
||||
([event]) => event === 'finish'
|
||||
)?.[1] as (() => Promise<void>) | undefined;
|
||||
|
||||
if (finishHandler2) {
|
||||
await finishHandler2();
|
||||
}
|
||||
|
||||
// Track session 2 ID (placeholder for actual implementation)
|
||||
|
||||
// User 1 deletes their session
|
||||
jest.clearAllMocks();
|
||||
mockReq.method = 'DELETE';
|
||||
mockReq.headers = {
|
||||
...mockReq.headers,
|
||||
'mcp-session-id': session1Id
|
||||
};
|
||||
mockReq.body = {};
|
||||
mockReq.auth = user1Auth;
|
||||
|
||||
await handleStreamableHTTP(mockReq as Request, mockRes as Response);
|
||||
|
||||
// Only user 1's session should be cleaned up
|
||||
// User 2's session should remain active
|
||||
// This test documents expected behavior for proper user isolation
|
||||
});
|
||||
});
|
||||
});
|
||||
229
src/modules/mcp/handlers/shttp.test.ts
Normal file
229
src/modules/mcp/handlers/shttp.test.ts
Normal file
|
|
@ -0,0 +1,229 @@
|
|||
import { jest } from '@jest/globals';
|
||||
import { Request, Response } from 'express';
|
||||
import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { MockRedisClient, setRedisClient } from '../../shared/redis.js';
|
||||
|
||||
describe('Streamable HTTP Handler', () => {
|
||||
let mockRedis: MockRedisClient;
|
||||
|
||||
beforeEach(() => {
|
||||
mockRedis = new MockRedisClient();
|
||||
setRedisClient(mockRedis);
|
||||
jest.resetAllMocks();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
mockRedis.clear();
|
||||
});
|
||||
|
||||
describe('Helper function tests', () => {
|
||||
it('should verify Redis mock is working', async () => {
|
||||
await mockRedis.set('test-key', 'test-value');
|
||||
const value = await mockRedis.get('test-key');
|
||||
expect(value).toBe('test-value');
|
||||
});
|
||||
|
||||
it('should handle Redis pub/sub', async () => {
|
||||
const messageHandler = jest.fn();
|
||||
const cleanup = await mockRedis.createSubscription(
|
||||
'test-channel',
|
||||
messageHandler,
|
||||
jest.fn()
|
||||
);
|
||||
|
||||
await mockRedis.publish('test-channel', 'test-message');
|
||||
|
||||
expect(messageHandler).toHaveBeenCalledWith('test-message');
|
||||
|
||||
await cleanup();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Request validation', () => {
|
||||
it('should identify initialize requests correctly', async () => {
|
||||
const { isInitializeRequest } = await import('@modelcontextprotocol/sdk/types.js');
|
||||
|
||||
const initRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 1,
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'test', version: '1.0' }
|
||||
}
|
||||
};
|
||||
|
||||
const nonInitRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 2,
|
||||
method: 'tools/list',
|
||||
params: {}
|
||||
};
|
||||
|
||||
expect(isInitializeRequest(initRequest)).toBe(true);
|
||||
expect(isInitializeRequest(nonInitRequest)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('HTTP response mock behavior', () => {
|
||||
it('should create proper response mock with chainable methods', () => {
|
||||
const mockRes = {
|
||||
status: jest.fn().mockReturnThis(),
|
||||
json: jest.fn().mockReturnThis(),
|
||||
on: jest.fn().mockReturnThis(),
|
||||
headersSent: false,
|
||||
} as Partial<Response>;
|
||||
|
||||
// Test chaining
|
||||
const result = mockRes.status!(400).json!({
|
||||
jsonrpc: '2.0',
|
||||
error: { code: -32000, message: 'Bad Request' },
|
||||
id: null
|
||||
});
|
||||
|
||||
expect(mockRes.status).toHaveBeenCalledWith(400);
|
||||
expect(mockRes.json).toHaveBeenCalledWith({
|
||||
jsonrpc: '2.0',
|
||||
error: { code: -32000, message: 'Bad Request' },
|
||||
id: null
|
||||
});
|
||||
expect(result).toBe(mockRes);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Session ID generation', () => {
|
||||
it('should generate valid UUIDs', async () => {
|
||||
const { randomUUID } = await import('crypto');
|
||||
|
||||
const sessionId = randomUUID();
|
||||
|
||||
// UUID v4 format: xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx
|
||||
expect(sessionId).toMatch(/^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Redis channel naming', () => {
|
||||
it('should create correct channel names for server communication', () => {
|
||||
const sessionId = 'test-session-123';
|
||||
const requestId = 'req-456';
|
||||
|
||||
const toServerChannel = `mcp:shttp:toserver:${sessionId}`;
|
||||
const toClientChannel = `mcp:shttp:toclient:${sessionId}:${requestId}`;
|
||||
const notificationChannel = `mcp:shttp:toclient:${sessionId}:__GET_stream`;
|
||||
|
||||
expect(toServerChannel).toBe('mcp:shttp:toserver:test-session-123');
|
||||
expect(toClientChannel).toBe('mcp:shttp:toclient:test-session-123:req-456');
|
||||
expect(notificationChannel).toBe('mcp:shttp:toclient:test-session-123:__GET_stream');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Error response formatting', () => {
|
||||
it('should format JSON-RPC error responses correctly', () => {
|
||||
const errorResponse = {
|
||||
jsonrpc: '2.0',
|
||||
error: {
|
||||
code: -32000,
|
||||
message: 'Bad Request: No valid session ID provided',
|
||||
},
|
||||
id: null,
|
||||
};
|
||||
|
||||
expect(errorResponse.jsonrpc).toBe('2.0');
|
||||
expect(errorResponse.error.code).toBe(-32000);
|
||||
expect(errorResponse.error.message).toContain('Bad Request');
|
||||
expect(errorResponse.id).toBe(null);
|
||||
});
|
||||
|
||||
it('should format internal error responses correctly', () => {
|
||||
const internalErrorResponse = {
|
||||
jsonrpc: '2.0',
|
||||
error: {
|
||||
code: -32603,
|
||||
message: 'Internal error',
|
||||
},
|
||||
id: null,
|
||||
};
|
||||
|
||||
expect(internalErrorResponse.jsonrpc).toBe('2.0');
|
||||
expect(internalErrorResponse.error.code).toBe(-32603);
|
||||
expect(internalErrorResponse.error.message).toBe('Internal error');
|
||||
expect(internalErrorResponse.id).toBe(null);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Request/Response patterns', () => {
|
||||
it('should handle typical MCP message structures', () => {
|
||||
const initializeRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 1,
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: { name: 'test-client', version: '1.0.0' }
|
||||
}
|
||||
};
|
||||
|
||||
const toolsListRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 2,
|
||||
method: 'tools/list',
|
||||
params: {}
|
||||
};
|
||||
|
||||
const toolsListResponse: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 2,
|
||||
result: {
|
||||
tools: [
|
||||
{
|
||||
name: 'echo',
|
||||
description: 'Echo the input',
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
text: { type: 'string' }
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
};
|
||||
|
||||
expect(initializeRequest.method).toBe('initialize');
|
||||
expect(toolsListRequest.method).toBe('tools/list');
|
||||
expect(toolsListResponse.result).toBeDefined();
|
||||
expect(Array.isArray(toolsListResponse.result?.tools)).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('HTTP header handling', () => {
|
||||
it('should extract session ID from headers', () => {
|
||||
const mockReq = {
|
||||
headers: {
|
||||
'mcp-session-id': 'test-session-123',
|
||||
'content-type': 'application/json'
|
||||
},
|
||||
body: {}
|
||||
} as Partial<Request>;
|
||||
|
||||
const sessionId = mockReq.headers!['mcp-session-id'] as string;
|
||||
|
||||
expect(sessionId).toBe('test-session-123');
|
||||
});
|
||||
|
||||
it('should handle missing session ID in headers', () => {
|
||||
const mockReq = {
|
||||
headers: {
|
||||
'content-type': 'application/json'
|
||||
},
|
||||
body: {}
|
||||
} as Partial<Request>;
|
||||
|
||||
const sessionId = mockReq.headers!['mcp-session-id'] as string | undefined;
|
||||
|
||||
expect(sessionId).toBeUndefined();
|
||||
});
|
||||
});
|
||||
});
|
||||
177
src/modules/mcp/handlers/shttp.ts
Normal file
177
src/modules/mcp/handlers/shttp.ts
Normal file
|
|
@ -0,0 +1,177 @@
|
|||
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
|
||||
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
|
||||
import { Request, Response } from "express";
|
||||
import { getShttpTransport, isSessionOwnedBy, redisRelayToMcpServer, ServerRedisTransport, setSessionOwner, shutdownSession } from "../services/redisTransport.js";
|
||||
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
|
||||
import { randomUUID } from "crypto";
|
||||
import { createMcpServer } from "../services/mcp.js";
|
||||
import { logger } from "../../shared/logger.js";
|
||||
|
||||
|
||||
declare module "express-serve-static-core" {
|
||||
interface Request {
|
||||
/**
|
||||
* Information about the validated access token, if the `requireBearerAuth` middleware was used.
|
||||
*/
|
||||
auth?: AuthInfo;
|
||||
}
|
||||
}
|
||||
|
||||
function getUserIdFromAuth(auth?: AuthInfo): string | null {
|
||||
return auth?.extra?.userId as string || null;
|
||||
}
|
||||
|
||||
// TODO: Document Streamable HTTP implementation choices:
|
||||
// 1. STATEFUL: Requires clients to initialize sessions and track session IDs
|
||||
// - First request must be 'initialize' without Mcp-Session-Id header
|
||||
// - Server returns session ID, client must include it in subsequent requests
|
||||
// - Alternative: Could implement STATELESS mode (each request independent)
|
||||
// 2. SSE RESPONSES: Returns results via Server-Sent Events stream, not JSON responses
|
||||
// - Requires Accept: application/json, text/event-stream header
|
||||
// - Responses formatted as: event: message\ndata: {...}
|
||||
// - Alternative: Could use JSON response mode (check StreamableHTTPServerTransport options)
|
||||
|
||||
export async function handleStreamableHTTP(req: Request, res: Response) {
|
||||
let shttpTransport: StreamableHTTPServerTransport | undefined = undefined;
|
||||
|
||||
res.on('finish', async () => {
|
||||
await shttpTransport?.close();
|
||||
});
|
||||
|
||||
const onsessionclosed = async (sessionId: string) => {
|
||||
logger.info('Session closed callback triggered', {
|
||||
sessionId,
|
||||
userId: getUserIdFromAuth(req.auth)
|
||||
});
|
||||
await shutdownSession(sessionId);
|
||||
}
|
||||
|
||||
try {
|
||||
// Check for existing session ID
|
||||
const sessionId = req.headers['mcp-session-id'] as string | undefined;
|
||||
const userId = getUserIdFromAuth(req.auth);
|
||||
|
||||
logger.debug('SHTTP request received', {
|
||||
method: req.method,
|
||||
sessionId,
|
||||
userId,
|
||||
hasAuth: !!req.auth,
|
||||
authExtra: req.auth?.extra
|
||||
});
|
||||
|
||||
// if no userid, return 401, we shouldn't get here ideally
|
||||
if (!userId) {
|
||||
logger.warning('Request without user ID', {
|
||||
sessionId,
|
||||
hasAuth: !!req.auth
|
||||
});
|
||||
res.status(401).json({
|
||||
"jsonrpc": "2.0",
|
||||
"error": {
|
||||
"code": -32002,
|
||||
"message": "User ID required"
|
||||
}
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
const isGetRequest = req.method === 'GET';
|
||||
|
||||
// incorrect session for the authed user, return 401
|
||||
if (sessionId) {
|
||||
if (!(await isSessionOwnedBy(sessionId, userId))) {
|
||||
logger.warning('Session ownership mismatch', {
|
||||
sessionId,
|
||||
userId,
|
||||
requestMethod: req.method
|
||||
});
|
||||
res.status(401).json({
|
||||
"jsonrpc": "2.0",
|
||||
"error": {
|
||||
"code": -32001,
|
||||
"message": "Session not found or access denied"
|
||||
}
|
||||
});
|
||||
return;
|
||||
}
|
||||
// Reuse existing transport for owned session
|
||||
logger.info('Reusing existing session', {
|
||||
sessionId,
|
||||
userId,
|
||||
isGetRequest
|
||||
});
|
||||
shttpTransport = await getShttpTransport(sessionId, onsessionclosed, isGetRequest);
|
||||
} else if (isInitializeRequest(req.body)) {
|
||||
// New initialization request - use JSON response mode
|
||||
logger.debug('Processing initialize request', {
|
||||
body: req.body,
|
||||
userId,
|
||||
headerSessionId: sessionId, // This is the sessionId from header (should be undefined for init)
|
||||
isInitializeRequest: true
|
||||
});
|
||||
|
||||
const onsessioninitialized = async (sessionId: string) => {
|
||||
logger.info('Initializing new session', {
|
||||
sessionId,
|
||||
userId
|
||||
});
|
||||
|
||||
const { server, cleanup: mcpCleanup } = createMcpServer();
|
||||
|
||||
const serverRedisTransport = new ServerRedisTransport(sessionId);
|
||||
serverRedisTransport.onclose = mcpCleanup;
|
||||
await server.connect(serverRedisTransport)
|
||||
|
||||
// Set session ownership
|
||||
await setSessionOwner(sessionId, userId);
|
||||
|
||||
logger.info('Session initialized successfully', {
|
||||
sessionId,
|
||||
userId
|
||||
});
|
||||
}
|
||||
|
||||
const newSessionId = randomUUID();
|
||||
shttpTransport = new StreamableHTTPServerTransport({
|
||||
sessionIdGenerator: () => newSessionId,
|
||||
onsessionclosed,
|
||||
onsessioninitialized,
|
||||
});
|
||||
shttpTransport.onclose = await redisRelayToMcpServer(newSessionId, shttpTransport);
|
||||
} else {
|
||||
// Invalid request - no session ID and not initialization request
|
||||
logger.warning('Invalid request: no session ID and not initialization', {
|
||||
hasSessionId: !!sessionId,
|
||||
isInitRequest: false,
|
||||
userId,
|
||||
method: req.method
|
||||
});
|
||||
res.status(400).json({
|
||||
"jsonrpc": "2.0",
|
||||
"error": {
|
||||
"code": -32600,
|
||||
"message": "Invalid request method for existing session"
|
||||
}
|
||||
});
|
||||
return;
|
||||
}
|
||||
// Handle the request with existing transport - no need to reconnect
|
||||
await shttpTransport.handleRequest(req, res, req.body);
|
||||
} catch (error) {
|
||||
logger.error('Error handling MCP request', error as Error, {
|
||||
sessionId: req.headers['mcp-session-id'] as string | undefined,
|
||||
method: req.method,
|
||||
userId: getUserIdFromAuth(req.auth)
|
||||
});
|
||||
|
||||
if (!res.headersSent) {
|
||||
res.status(500).json({
|
||||
"jsonrpc": "2.0",
|
||||
"error": {
|
||||
"code": -32603,
|
||||
"message": "Internal error during request processing"
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
100
src/modules/mcp/handlers/sse.ts
Normal file
100
src/modules/mcp/handlers/sse.ts
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
|
||||
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
|
||||
import contentType from "content-type";
|
||||
import { Request, Response } from "express";
|
||||
import { redisClient } from "../../shared/redis.js";
|
||||
import { createMcpServer } from "../services/mcp.js";
|
||||
import { logger, logMcpMessage } from "../../shared/logger.js";
|
||||
|
||||
declare module "express-serve-static-core" {
|
||||
interface Request {
|
||||
/**
|
||||
* Information about the validated access token, if the `requireBearerAuth` middleware was used.
|
||||
*/
|
||||
auth?: AuthInfo;
|
||||
}
|
||||
}
|
||||
|
||||
function redisChannelForSession(sessionId: string): string {
|
||||
return `mcp:${sessionId}`;
|
||||
}
|
||||
|
||||
export async function handleSSEConnection(req: Request, res: Response) {
|
||||
const { server: mcpServer, cleanup: mcpCleanup } = createMcpServer();
|
||||
const transport = new SSEServerTransport("/message", res);
|
||||
logger.info('Received MCP SSE connection', {
|
||||
sessionId: transport.sessionId
|
||||
});
|
||||
|
||||
const redisCleanup = await redisClient.createSubscription(
|
||||
redisChannelForSession(transport.sessionId),
|
||||
(json) => {
|
||||
// TODO handle DELETE messages
|
||||
// TODO set timeout to kill the session
|
||||
|
||||
const message = JSON.parse(json);
|
||||
logMcpMessage(message, transport.sessionId);
|
||||
transport.handleMessage(message).catch((error) => {
|
||||
logger.error('Error handling message', error as Error, {
|
||||
sessionId: transport.sessionId
|
||||
});
|
||||
});
|
||||
},
|
||||
(error) => {
|
||||
logger.error('Disconnecting due to error in Redis subscriber', error as Error, {
|
||||
sessionId: transport.sessionId
|
||||
});
|
||||
transport
|
||||
.close()
|
||||
.catch((error) =>
|
||||
logger.error('Error closing transport', error as Error, {
|
||||
sessionId: transport.sessionId
|
||||
}),
|
||||
);
|
||||
},
|
||||
);
|
||||
|
||||
const cleanup = () => {
|
||||
void mcpCleanup();
|
||||
redisCleanup().catch((error) =>
|
||||
logger.error('Error disconnecting Redis subscriber', error as Error, {
|
||||
sessionId: transport.sessionId
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// Clean up Redis subscription when the connection closes
|
||||
mcpServer.onclose = cleanup
|
||||
|
||||
logger.info('Listening on Redis channel', {
|
||||
sessionId: transport.sessionId,
|
||||
channel: redisChannelForSession(transport.sessionId)
|
||||
});
|
||||
await mcpServer.connect(transport);
|
||||
}
|
||||
|
||||
export async function handleMessage(req: Request, res: Response) {
|
||||
const sessionId = req.query.sessionId;
|
||||
let body: string;
|
||||
try {
|
||||
if (typeof sessionId !== "string") {
|
||||
throw new Error("Only one sessionId allowed");
|
||||
}
|
||||
|
||||
const ct = contentType.parse(req.headers["content-type"] ?? "");
|
||||
if (ct.type !== "application/json") {
|
||||
throw new Error(`Unsupported content-type: ${ct}`);
|
||||
}
|
||||
|
||||
body = JSON.stringify(req.body);
|
||||
} catch (error) {
|
||||
res.status(400).json({ error: "Bad request" });
|
||||
logger.error('Bad POST request', error as Error, {
|
||||
sessionId,
|
||||
contentType: req.headers['content-type']
|
||||
});
|
||||
return;
|
||||
}
|
||||
await redisClient.publish(redisChannelForSession(sessionId), body);
|
||||
res.status(202).end();
|
||||
}
|
||||
112
src/modules/mcp/index.ts
Normal file
112
src/modules/mcp/index.ts
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
/**
|
||||
* MCP Module - Self-contained MCP server module
|
||||
*
|
||||
* This module provides all MCP protocol functionality (tools, resources, prompts).
|
||||
* It depends only on the ITokenValidator interface for authentication, not on
|
||||
* any specific auth implementation.
|
||||
*
|
||||
* This clean separation means the MCP module works identically whether auth
|
||||
* is internal (in-process) or external (HTTP).
|
||||
*/
|
||||
|
||||
import { Router, Request, Response, NextFunction } from 'express';
|
||||
import cors from 'cors';
|
||||
import path from 'path';
|
||||
import { fileURLToPath } from 'url';
|
||||
import rateLimit from 'express-rate-limit';
|
||||
import { BearerAuthMiddlewareOptions, requireBearerAuth } from '@modelcontextprotocol/sdk/server/auth/middleware/bearerAuth.js';
|
||||
import { getOAuthProtectedResourceMetadataUrl } from '@modelcontextprotocol/sdk/server/auth/router.js';
|
||||
import { ITokenValidator } from '../../interfaces/auth-validator.js';
|
||||
import { handleStreamableHTTP } from './handlers/shttp.js';
|
||||
import { handleMessage, handleSSEConnection } from './handlers/sse.js';
|
||||
|
||||
const __filename = fileURLToPath(import.meta.url);
|
||||
const __dirname = path.dirname(__filename);
|
||||
|
||||
export interface MCPConfig {
|
||||
baseUri: string;
|
||||
redisUrl?: string;
|
||||
}
|
||||
|
||||
export class MCPModule {
|
||||
private router: Router;
|
||||
|
||||
constructor(
|
||||
private config: MCPConfig,
|
||||
private tokenValidator: ITokenValidator
|
||||
) {
|
||||
this.router = this.setupRouter();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get Express router with all MCP endpoints
|
||||
*/
|
||||
getRouter(): Router {
|
||||
return this.router;
|
||||
}
|
||||
|
||||
private setupRouter(): Router {
|
||||
const router = Router();
|
||||
|
||||
// Rate limiter for static assets
|
||||
const staticAssetLimiter = rateLimit({
|
||||
windowMs: 60 * 1000, // 1 minute
|
||||
max: 100, // 100 requests per minute for static assets
|
||||
message: 'Too many requests for static assets',
|
||||
standardHeaders: true,
|
||||
legacyHeaders: false,
|
||||
});
|
||||
|
||||
// CORS configuration for MCP endpoints
|
||||
// Intentionally permissive CORS for public MCP reference server
|
||||
// This allows any MCP client to test against this reference implementation
|
||||
const corsOptions = {
|
||||
origin: true, // Allow any origin
|
||||
methods: ['GET', 'POST', 'DELETE'],
|
||||
allowedHeaders: ['Content-Type', 'Authorization', 'Mcp-Protocol-Version', 'Mcp-Protocol-Id'],
|
||||
exposedHeaders: ['Mcp-Protocol-Version', 'Mcp-Protocol-Id'],
|
||||
credentials: true
|
||||
};
|
||||
|
||||
// Security headers for MCP endpoints
|
||||
const securityHeaders = (req: Request, res: Response, next: NextFunction) => {
|
||||
res.setHeader('X-Content-Type-Options', 'nosniff');
|
||||
res.setHeader('X-Frame-Options', 'SAMEORIGIN');
|
||||
next();
|
||||
};
|
||||
|
||||
// SSE-specific headers
|
||||
const sseHeaders = (req: Request, res: Response, next: NextFunction) => {
|
||||
res.setHeader('Content-Type', 'text/event-stream');
|
||||
res.setHeader('Cache-Control', 'no-store, max-age=0');
|
||||
res.setHeader('Connection', 'keep-alive');
|
||||
next();
|
||||
};
|
||||
|
||||
// Bearer auth middleware using our token validator
|
||||
// This works the same whether the validator is internal or external
|
||||
const bearerAuthOptions: BearerAuthMiddlewareOptions = {
|
||||
verifier: this.tokenValidator,
|
||||
resourceMetadataUrl: getOAuthProtectedResourceMetadataUrl(new URL(this.config.baseUri))
|
||||
};
|
||||
const bearerAuth = requireBearerAuth(bearerAuthOptions);
|
||||
|
||||
// MCP endpoints - Streamable HTTP transport (recommended)
|
||||
router.get('/mcp', cors(corsOptions), bearerAuth, securityHeaders, handleStreamableHTTP);
|
||||
router.post('/mcp', cors(corsOptions), bearerAuth, securityHeaders, handleStreamableHTTP);
|
||||
router.delete('/mcp', cors(corsOptions), bearerAuth, securityHeaders, handleStreamableHTTP);
|
||||
|
||||
// MCP endpoints - SSE transport (legacy)
|
||||
router.get('/sse', cors(corsOptions), bearerAuth, sseHeaders, handleSSEConnection);
|
||||
router.post('/message', cors(corsOptions), bearerAuth, securityHeaders, handleMessage);
|
||||
|
||||
// Static files for MCP
|
||||
router.get('/styles.css', staticAssetLimiter, (req, res) => {
|
||||
const cssPath = path.join(__dirname, '../../static', 'styles.css');
|
||||
res.setHeader('Content-Type', 'text/css');
|
||||
res.sendFile(cssPath);
|
||||
});
|
||||
|
||||
return router;
|
||||
}
|
||||
}
|
||||
786
src/modules/mcp/services/mcp.ts
Normal file
786
src/modules/mcp/services/mcp.ts
Normal file
File diff suppressed because one or more lines are too long
291
src/modules/mcp/services/redisTransport.integration.test.ts
Normal file
291
src/modules/mcp/services/redisTransport.integration.test.ts
Normal file
|
|
@ -0,0 +1,291 @@
|
|||
import { jest } from '@jest/globals';
|
||||
import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { MockRedisClient, setRedisClient } from '../../shared/redis.js';
|
||||
import {
|
||||
ServerRedisTransport,
|
||||
redisRelayToMcpServer,
|
||||
shutdownSession
|
||||
} from './redisTransport.js';
|
||||
import { createMcpServer } from './mcp.js';
|
||||
import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
|
||||
|
||||
describe('Redis Transport Integration', () => {
|
||||
let mockRedis: MockRedisClient;
|
||||
|
||||
beforeEach(() => {
|
||||
mockRedis = new MockRedisClient();
|
||||
setRedisClient(mockRedis);
|
||||
jest.resetAllMocks();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
mockRedis.clear();
|
||||
});
|
||||
|
||||
describe('MCP Initialization Flow', () => {
|
||||
const sessionId = 'test-init-session';
|
||||
|
||||
it('should relay initialization request from client to server through Redis', async () => {
|
||||
// 1. Start the server listening to Redis
|
||||
const { server, cleanup: serverCleanup } = createMcpServer();
|
||||
const serverTransport = new ServerRedisTransport(sessionId);
|
||||
serverTransport.onclose = serverCleanup;
|
||||
await server.connect(serverTransport);
|
||||
|
||||
// 2. Create a mock client transport (simulating the streamable HTTP client side)
|
||||
const mockClientTransport: Transport = {
|
||||
onmessage: undefined,
|
||||
onclose: undefined,
|
||||
onerror: undefined,
|
||||
send: jest.fn(() => Promise.resolve()),
|
||||
close: jest.fn(() => Promise.resolve()),
|
||||
start: jest.fn(() => Promise.resolve())
|
||||
};
|
||||
|
||||
// 3. Set up the Redis relay (this is what happens in the HTTP handler)
|
||||
const cleanup = await redisRelayToMcpServer(sessionId, mockClientTransport);
|
||||
|
||||
// Track messages received by server
|
||||
const serverReceivedMessages: JSONRPCMessage[] = [];
|
||||
const originalServerOnMessage = serverTransport.onmessage;
|
||||
serverTransport.onmessage = (message, extra) => {
|
||||
serverReceivedMessages.push(message);
|
||||
originalServerOnMessage?.(message, extra);
|
||||
};
|
||||
|
||||
// 4. Simulate client sending initialization request
|
||||
const initMessage: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
method: 'initialize',
|
||||
params: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {},
|
||||
clientInfo: {
|
||||
name: 'test-client',
|
||||
version: '1.0.0'
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Trigger the client transport onmessage (simulates HTTP request)
|
||||
mockClientTransport.onmessage?.(initMessage);
|
||||
|
||||
// Wait for message to be relayed through Redis
|
||||
await new Promise(resolve => setTimeout(resolve, 50));
|
||||
|
||||
// 5. Verify server received the init message
|
||||
expect(serverReceivedMessages).toHaveLength(1);
|
||||
expect(serverReceivedMessages[0]).toMatchObject({
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
method: 'initialize'
|
||||
});
|
||||
|
||||
// 6. Simulate server responding (this should get relayed back to client)
|
||||
const initResponse: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
result: {
|
||||
protocolVersion: '2024-11-05',
|
||||
capabilities: {
|
||||
tools: {},
|
||||
prompts: {},
|
||||
resources: {}
|
||||
},
|
||||
serverInfo: {
|
||||
name: 'example-server',
|
||||
version: '1.0.0'
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
await serverTransport.send(initResponse, { relatedRequestId: 'init-1' });
|
||||
|
||||
// Wait for response to be relayed back
|
||||
await new Promise(resolve => setTimeout(resolve, 50));
|
||||
|
||||
// 7. Verify client transport received the response
|
||||
expect(mockClientTransport.send).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
jsonrpc: '2.0',
|
||||
id: 'init-1',
|
||||
result: expect.objectContaining({
|
||||
protocolVersion: '2024-11-05',
|
||||
serverInfo: expect.objectContaining({
|
||||
name: 'example-server'
|
||||
})
|
||||
})
|
||||
}),
|
||||
{ relatedRequestId: 'init-1' }
|
||||
);
|
||||
|
||||
// Cleanup
|
||||
await cleanup();
|
||||
await shutdownSession(sessionId);
|
||||
serverCleanup(); // Clean up MCP server intervals
|
||||
|
||||
// Ensure server transport is closed
|
||||
await serverTransport.close();
|
||||
|
||||
await new Promise(resolve => setTimeout(resolve, 10));
|
||||
});
|
||||
|
||||
it('should handle tools/list request through Redis relay', async () => {
|
||||
// Set up server and mock client
|
||||
const { server, cleanup: serverCleanup } = createMcpServer();
|
||||
const serverTransport = new ServerRedisTransport(sessionId);
|
||||
serverTransport.onclose = serverCleanup;
|
||||
await server.connect(serverTransport);
|
||||
|
||||
const mockClientTransport: Transport = {
|
||||
onmessage: undefined,
|
||||
onclose: undefined,
|
||||
onerror: undefined,
|
||||
send: jest.fn(() => Promise.resolve()),
|
||||
close: jest.fn(() => Promise.resolve()),
|
||||
start: jest.fn(() => Promise.resolve())
|
||||
};
|
||||
|
||||
const cleanup = await redisRelayToMcpServer(sessionId, mockClientTransport);
|
||||
|
||||
// Send tools/list request
|
||||
const toolsListMessage: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'tools-1',
|
||||
method: 'tools/list',
|
||||
params: {}
|
||||
};
|
||||
|
||||
mockClientTransport.onmessage?.(toolsListMessage);
|
||||
|
||||
// Wait for processing and response
|
||||
await new Promise(resolve => setTimeout(resolve, 100));
|
||||
|
||||
// Verify client received a response with tools
|
||||
expect(mockClientTransport.send).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
jsonrpc: '2.0',
|
||||
id: 'tools-1',
|
||||
result: expect.objectContaining({
|
||||
tools: expect.any(Array)
|
||||
})
|
||||
}),
|
||||
undefined
|
||||
);
|
||||
|
||||
// Cleanup
|
||||
await cleanup();
|
||||
await shutdownSession(sessionId);
|
||||
serverCleanup(); // Clean up MCP server intervals
|
||||
|
||||
// Ensure server transport is closed
|
||||
await serverTransport.close();
|
||||
|
||||
await new Promise(resolve => setTimeout(resolve, 10));
|
||||
});
|
||||
|
||||
it('should handle notifications through Redis relay', async () => {
|
||||
// Set up server and mock client
|
||||
const { server, cleanup: serverCleanup } = createMcpServer();
|
||||
const serverTransport = new ServerRedisTransport(sessionId);
|
||||
serverTransport.onclose = serverCleanup;
|
||||
await server.connect(serverTransport);
|
||||
|
||||
const mockClientTransport: Transport = {
|
||||
onmessage: undefined,
|
||||
onclose: undefined,
|
||||
onerror: undefined,
|
||||
send: jest.fn(() => Promise.resolve()),
|
||||
close: jest.fn(() => Promise.resolve()),
|
||||
start: jest.fn(() => Promise.resolve())
|
||||
};
|
||||
|
||||
const cleanup = await redisRelayToMcpServer(sessionId, mockClientTransport);
|
||||
|
||||
// Set up notification subscription manually since notifications don't have an id
|
||||
const notificationChannel = `mcp:shttp:toclient:${sessionId}:__GET_stream`;
|
||||
const notificationCleanup = await mockRedis.createSubscription(notificationChannel, async (redisMessageJson) => {
|
||||
const redisMessage = JSON.parse(redisMessageJson);
|
||||
if (redisMessage.type === 'mcp') {
|
||||
await mockClientTransport.send(redisMessage.message, redisMessage.options);
|
||||
}
|
||||
}, (error) => {
|
||||
mockClientTransport.onerror?.(error);
|
||||
});
|
||||
|
||||
// Send a notification from server (notifications don't have an id)
|
||||
const notification: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'notifications/message',
|
||||
params: {
|
||||
level: 'info',
|
||||
logger: 'test',
|
||||
data: 'Test notification'
|
||||
}
|
||||
};
|
||||
|
||||
await serverTransport.send(notification);
|
||||
|
||||
// Wait for notification to be delivered
|
||||
await new Promise(resolve => setTimeout(resolve, 50));
|
||||
|
||||
// Verify client received the notification
|
||||
expect(mockClientTransport.send).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
jsonrpc: '2.0',
|
||||
method: 'notifications/message',
|
||||
params: expect.objectContaining({
|
||||
level: 'info',
|
||||
data: 'Test notification'
|
||||
})
|
||||
}),
|
||||
undefined
|
||||
);
|
||||
|
||||
// Cleanup notification subscription
|
||||
await notificationCleanup();
|
||||
|
||||
// Cleanup
|
||||
await cleanup();
|
||||
await shutdownSession(sessionId);
|
||||
serverCleanup(); // Clean up MCP server intervals
|
||||
|
||||
// Ensure server transport is closed
|
||||
await serverTransport.close();
|
||||
|
||||
await new Promise(resolve => setTimeout(resolve, 10));
|
||||
});
|
||||
|
||||
it('should not create response subscriptions for notifications', async () => {
|
||||
const mockClientTransport: Transport = {
|
||||
onmessage: undefined,
|
||||
onclose: undefined,
|
||||
onerror: undefined,
|
||||
send: jest.fn(() => Promise.resolve()),
|
||||
close: jest.fn(() => Promise.resolve()),
|
||||
start: jest.fn(() => Promise.resolve())
|
||||
};
|
||||
|
||||
const cleanup = await redisRelayToMcpServer(sessionId, mockClientTransport);
|
||||
|
||||
// Send a notification (no id field)
|
||||
const notification: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'notifications/initialized',
|
||||
params: {}
|
||||
};
|
||||
|
||||
mockClientTransport.onmessage?.(notification);
|
||||
|
||||
// Wait for processing
|
||||
await new Promise(resolve => setTimeout(resolve, 50));
|
||||
|
||||
// Should not create any response channel subscriptions for notifications
|
||||
// (we can't easily test this directly, but we can ensure no errors occur)
|
||||
expect(mockClientTransport.send).not.toHaveBeenCalled();
|
||||
|
||||
await cleanup();
|
||||
});
|
||||
});
|
||||
});
|
||||
536
src/modules/mcp/services/redisTransport.test.ts
Normal file
536
src/modules/mcp/services/redisTransport.test.ts
Normal file
|
|
@ -0,0 +1,536 @@
|
|||
import { jest } from '@jest/globals';
|
||||
import { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { MockRedisClient, setRedisClient } from '../../shared/redis.js';
|
||||
import {
|
||||
ServerRedisTransport,
|
||||
redisRelayToMcpServer,
|
||||
isLive,
|
||||
shutdownSession,
|
||||
setSessionOwner,
|
||||
getSessionOwner,
|
||||
validateSessionOwnership,
|
||||
isSessionOwnedBy
|
||||
} from './redisTransport.js';
|
||||
import { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
|
||||
|
||||
describe('Redis Transport', () => {
|
||||
let mockRedis: MockRedisClient;
|
||||
|
||||
beforeEach(() => {
|
||||
mockRedis = new MockRedisClient();
|
||||
setRedisClient(mockRedis);
|
||||
jest.resetAllMocks();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
// Clear all Redis data and subscriptions
|
||||
mockRedis.clear();
|
||||
});
|
||||
|
||||
describe('ServerRedisTransport', () => {
|
||||
let transport: ServerRedisTransport;
|
||||
const sessionId = 'test-session-123';
|
||||
|
||||
beforeEach(() => {
|
||||
transport = new ServerRedisTransport(sessionId);
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
if (transport) {
|
||||
await transport.close();
|
||||
}
|
||||
});
|
||||
|
||||
it('should create transport with session ID', () => {
|
||||
expect(transport).toBeInstanceOf(ServerRedisTransport);
|
||||
});
|
||||
|
||||
it('should send response messages to request-specific channels', async () => {
|
||||
const responseMessage: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 123,
|
||||
result: { data: 'test response' }
|
||||
};
|
||||
|
||||
const mockSubscriber = jest.fn();
|
||||
await mockRedis.createSubscription(
|
||||
`mcp:shttp:toclient:${sessionId}:123`,
|
||||
mockSubscriber,
|
||||
jest.fn()
|
||||
);
|
||||
|
||||
await transport.send(responseMessage, { relatedRequestId: 123 });
|
||||
|
||||
expect(mockSubscriber).toHaveBeenCalledWith(
|
||||
JSON.stringify({
|
||||
type: 'mcp',
|
||||
message: responseMessage,
|
||||
options: { relatedRequestId: 123 }
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it('should send notification messages to notification channel', async () => {
|
||||
const notificationMessage: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'notifications/message',
|
||||
params: { message: 'test notification' }
|
||||
};
|
||||
|
||||
const mockSubscriber = jest.fn();
|
||||
await mockRedis.createSubscription(
|
||||
`mcp:shttp:toclient:${sessionId}:__GET_stream`,
|
||||
mockSubscriber,
|
||||
jest.fn()
|
||||
);
|
||||
|
||||
await transport.send(notificationMessage);
|
||||
|
||||
expect(mockSubscriber).toHaveBeenCalledWith(
|
||||
JSON.stringify({
|
||||
type: 'mcp',
|
||||
message: notificationMessage,
|
||||
options: undefined
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle close gracefully', async () => {
|
||||
const onCloseMock = jest.fn();
|
||||
transport.onclose = onCloseMock;
|
||||
|
||||
await transport.close();
|
||||
|
||||
expect(onCloseMock).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should respond to shutdown control messages', async () => {
|
||||
await transport.start();
|
||||
|
||||
const onCloseMock = jest.fn();
|
||||
transport.onclose = onCloseMock;
|
||||
|
||||
// Send a shutdown control message
|
||||
await shutdownSession(sessionId);
|
||||
|
||||
// Wait for async processing
|
||||
await new Promise(resolve => setTimeout(resolve, 10));
|
||||
|
||||
expect(onCloseMock).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should receive MCP messages from clients and call onmessage', async () => {
|
||||
const onMessageMock = jest.fn();
|
||||
transport.onmessage = onMessageMock;
|
||||
|
||||
await transport.start();
|
||||
|
||||
// Simulate client sending a message to server
|
||||
const clientMessage: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'test-req',
|
||||
method: 'tools/list',
|
||||
params: {}
|
||||
};
|
||||
|
||||
await mockRedis.publish(
|
||||
`mcp:shttp:toserver:${sessionId}`,
|
||||
JSON.stringify({
|
||||
type: 'mcp',
|
||||
message: clientMessage,
|
||||
extra: { authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } }
|
||||
})
|
||||
);
|
||||
|
||||
// Wait for async processing
|
||||
await new Promise(resolve => setTimeout(resolve, 10));
|
||||
|
||||
expect(onMessageMock).toHaveBeenCalledWith(
|
||||
clientMessage,
|
||||
{ authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } }
|
||||
);
|
||||
|
||||
await transport.close();
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
describe('redisRelayToMcpServer', () => {
|
||||
let mockTransport: Transport;
|
||||
const sessionId = 'test-session-456';
|
||||
|
||||
beforeEach(() => {
|
||||
mockTransport = {
|
||||
onmessage: undefined,
|
||||
onclose: undefined,
|
||||
onerror: undefined,
|
||||
send: jest.fn(() => Promise.resolve()),
|
||||
close: jest.fn(() => Promise.resolve()),
|
||||
start: jest.fn(() => Promise.resolve())
|
||||
};
|
||||
});
|
||||
|
||||
it('should set up message relay from transport to server', async () => {
|
||||
const cleanup = await redisRelayToMcpServer(sessionId, mockTransport);
|
||||
|
||||
// Simulate a message from the transport
|
||||
const requestMessage: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'req-123',
|
||||
method: 'tools/list',
|
||||
params: {}
|
||||
};
|
||||
|
||||
// Trigger the onmessage handler
|
||||
mockTransport.onmessage?.(requestMessage, { authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } });
|
||||
|
||||
// Wait a bit for async processing
|
||||
await new Promise(resolve => setTimeout(resolve, 10));
|
||||
|
||||
// Check that message was published to server channel
|
||||
const serverSubscriber = jest.fn();
|
||||
await mockRedis.createSubscription(
|
||||
`mcp:shttp:toserver:${sessionId}`,
|
||||
serverSubscriber,
|
||||
jest.fn()
|
||||
);
|
||||
|
||||
// The message should have been published
|
||||
expect(mockRedis.numsub(`mcp:shttp:toserver:${sessionId}`)).resolves.toBe(1);
|
||||
|
||||
await cleanup();
|
||||
});
|
||||
|
||||
it('should subscribe to response channel for request messages', async () => {
|
||||
const cleanup = await redisRelayToMcpServer(sessionId, mockTransport);
|
||||
|
||||
const requestMessage: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'req-456',
|
||||
method: 'tools/call',
|
||||
params: { name: 'echo', arguments: { text: 'hello' } }
|
||||
};
|
||||
|
||||
// Trigger the onmessage handler
|
||||
mockTransport.onmessage?.(requestMessage, { authInfo: { token: 'test-token', clientId: 'test-client', scopes: [] } });
|
||||
|
||||
// Wait for subscription setup
|
||||
await new Promise(resolve => setTimeout(resolve, 10));
|
||||
|
||||
// Now simulate a response from the server
|
||||
const responseMessage: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'req-456',
|
||||
result: { content: [{ type: 'text', text: 'hello' }] }
|
||||
};
|
||||
|
||||
await mockRedis.publish(
|
||||
`mcp:shttp:toclient:${sessionId}:req-456`,
|
||||
JSON.stringify({
|
||||
type: 'mcp',
|
||||
message: responseMessage,
|
||||
options: undefined
|
||||
})
|
||||
);
|
||||
|
||||
// Check that the response was sent back to the transport
|
||||
expect(mockTransport.send).toHaveBeenCalledWith(responseMessage, undefined);
|
||||
|
||||
await cleanup();
|
||||
});
|
||||
|
||||
it('should not subscribe for notification messages (no id)', async () => {
|
||||
const cleanup = await redisRelayToMcpServer(sessionId, mockTransport);
|
||||
|
||||
const notificationMessage: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
method: 'notifications/message',
|
||||
params: { message: 'test' }
|
||||
};
|
||||
|
||||
// Trigger the onmessage handler
|
||||
mockTransport.onmessage?.(notificationMessage);
|
||||
|
||||
// Wait a bit
|
||||
await new Promise(resolve => setTimeout(resolve, 10));
|
||||
|
||||
// Should not create any response channel subscriptions
|
||||
expect(await mockRedis.numsub(`mcp:shttp:toclient:${sessionId}:undefined`)).toBe(0);
|
||||
|
||||
await cleanup();
|
||||
});
|
||||
});
|
||||
|
||||
describe('isLive', () => {
|
||||
const sessionId = 'test-session-789';
|
||||
|
||||
it('should return true when session has active subscribers', async () => {
|
||||
// Create a subscription to the server channel
|
||||
await mockRedis.createSubscription(
|
||||
`mcp:shttp:toserver:${sessionId}`,
|
||||
jest.fn(),
|
||||
jest.fn()
|
||||
);
|
||||
|
||||
expect(await isLive(sessionId)).toBe(true);
|
||||
});
|
||||
|
||||
it('should return false when session has no subscribers', async () => {
|
||||
expect(await isLive(sessionId)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Session Ownership', () => {
|
||||
const sessionId = 'test-session-ownership';
|
||||
const userId = 'test-user-123';
|
||||
|
||||
it('should set and get session owner', async () => {
|
||||
await setSessionOwner(sessionId, userId);
|
||||
const owner = await getSessionOwner(sessionId);
|
||||
expect(owner).toBe(userId);
|
||||
});
|
||||
|
||||
it('should validate session ownership correctly', async () => {
|
||||
await setSessionOwner(sessionId, userId);
|
||||
|
||||
expect(await validateSessionOwnership(sessionId, userId)).toBe(true);
|
||||
expect(await validateSessionOwnership(sessionId, 'different-user')).toBe(false);
|
||||
});
|
||||
|
||||
it('should check if session is owned by user including liveness', async () => {
|
||||
// Session not live yet
|
||||
expect(await isSessionOwnedBy(sessionId, userId)).toBe(false);
|
||||
|
||||
// Make session live
|
||||
await mockRedis.createSubscription(
|
||||
`mcp:shttp:toserver:${sessionId}`,
|
||||
jest.fn(),
|
||||
jest.fn()
|
||||
);
|
||||
|
||||
// Still false because no owner set
|
||||
expect(await isSessionOwnedBy(sessionId, userId)).toBe(false);
|
||||
|
||||
// Set owner
|
||||
await setSessionOwner(sessionId, userId);
|
||||
|
||||
// Now should be true
|
||||
expect(await isSessionOwnedBy(sessionId, userId)).toBe(true);
|
||||
|
||||
// False for different user
|
||||
expect(await isSessionOwnedBy(sessionId, 'different-user')).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Integration: Redis message flow', () => {
|
||||
const sessionId = 'integration-test-session';
|
||||
|
||||
it('should relay messages between client and server through Redis', async () => {
|
||||
// Set up client-side transport simulation
|
||||
const clientTransport: Transport = {
|
||||
onmessage: undefined,
|
||||
onclose: undefined,
|
||||
onerror: undefined,
|
||||
send: jest.fn(() => Promise.resolve()),
|
||||
close: jest.fn(() => Promise.resolve()),
|
||||
start: jest.fn(() => Promise.resolve())
|
||||
};
|
||||
|
||||
const cleanup = await redisRelayToMcpServer(sessionId, clientTransport);
|
||||
|
||||
// Client sends a request
|
||||
const listToolsRequest: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'integration-req-1',
|
||||
method: 'tools/list',
|
||||
params: {}
|
||||
};
|
||||
|
||||
// Set up subscription to simulate server receiving the message
|
||||
const serverSubscriber = jest.fn();
|
||||
await mockRedis.createSubscription(
|
||||
`mcp:shttp:toserver:${sessionId}`,
|
||||
serverSubscriber,
|
||||
jest.fn()
|
||||
);
|
||||
|
||||
// Simulate client sending request
|
||||
clientTransport.onmessage?.(listToolsRequest);
|
||||
|
||||
// Wait for async processing
|
||||
await new Promise(resolve => setTimeout(resolve, 10));
|
||||
|
||||
// Verify the message was published to server channel
|
||||
expect(serverSubscriber).toHaveBeenCalledWith(
|
||||
JSON.stringify({
|
||||
type: 'mcp',
|
||||
message: listToolsRequest,
|
||||
extra: undefined,
|
||||
options: undefined
|
||||
})
|
||||
);
|
||||
|
||||
// Simulate server sending response
|
||||
const serverResponse: JSONRPCMessage = {
|
||||
jsonrpc: '2.0',
|
||||
id: 'integration-req-1',
|
||||
result: { tools: [{ name: 'echo', description: 'Echo tool' }] }
|
||||
};
|
||||
|
||||
await mockRedis.publish(
|
||||
`mcp:shttp:toclient:${sessionId}:integration-req-1`,
|
||||
JSON.stringify({
|
||||
type: 'mcp',
|
||||
message: serverResponse,
|
||||
options: undefined
|
||||
})
|
||||
);
|
||||
|
||||
// Wait for response processing
|
||||
await new Promise(resolve => setTimeout(resolve, 10));
|
||||
|
||||
// Verify the response was sent back to client
|
||||
expect(clientTransport.send).toHaveBeenCalledWith(serverResponse, undefined);
|
||||
|
||||
await cleanup();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Control Messages', () => {
|
||||
const sessionId = 'test-control-session';
|
||||
|
||||
it('should send shutdown control messages', async () => {
|
||||
const controlSubscriber = jest.fn();
|
||||
await mockRedis.createSubscription(
|
||||
`mcp:control:${sessionId}`,
|
||||
controlSubscriber,
|
||||
jest.fn()
|
||||
);
|
||||
|
||||
await shutdownSession(sessionId);
|
||||
|
||||
const callArgs = controlSubscriber.mock.calls[0][0] as string;
|
||||
const message = JSON.parse(callArgs);
|
||||
|
||||
expect(message.type).toBe('control');
|
||||
expect(message.action).toBe('SHUTDOWN');
|
||||
expect(typeof message.timestamp).toBe('number');
|
||||
});
|
||||
|
||||
it('should properly shutdown server transport via control message', async () => {
|
||||
const transport = new ServerRedisTransport(sessionId);
|
||||
const onCloseMock = jest.fn();
|
||||
transport.onclose = onCloseMock;
|
||||
|
||||
await transport.start();
|
||||
|
||||
// Send shutdown signal
|
||||
await shutdownSession(sessionId);
|
||||
|
||||
// Wait for async processing
|
||||
await new Promise(resolve => setTimeout(resolve, 10));
|
||||
|
||||
expect(onCloseMock).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('Inactivity Timeout', () => {
|
||||
const sessionId = 'test-inactivity-session';
|
||||
|
||||
beforeEach(() => {
|
||||
jest.useFakeTimers({ doNotFake: ['setImmediate', 'nextTick'] });
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
jest.useRealTimers();
|
||||
});
|
||||
|
||||
it('should shutdown session after 5 minutes of inactivity', async () => {
|
||||
const transport = new ServerRedisTransport(sessionId);
|
||||
const shutdownSpy = jest.spyOn(mockRedis, 'publish');
|
||||
|
||||
await transport.start();
|
||||
|
||||
// Fast-forward time by 5 minutes
|
||||
jest.advanceTimersByTime(5 * 60 * 1000);
|
||||
|
||||
// Should have published shutdown control message
|
||||
expect(shutdownSpy).toHaveBeenCalledWith(
|
||||
`mcp:control:${sessionId}`,
|
||||
expect.stringContaining('"action":"SHUTDOWN"')
|
||||
);
|
||||
|
||||
await transport.close();
|
||||
});
|
||||
|
||||
it('should reset timeout when message is received', async () => {
|
||||
const transport = new ServerRedisTransport(sessionId);
|
||||
const onMessageMock = jest.fn();
|
||||
transport.onmessage = onMessageMock;
|
||||
|
||||
await transport.start();
|
||||
|
||||
// Fast-forward 4 minutes
|
||||
jest.advanceTimersByTime(4 * 60 * 1000);
|
||||
|
||||
// Manually publish a message to trigger the subscription handler
|
||||
const testMessage = { jsonrpc: '2.0', method: 'ping' };
|
||||
await mockRedis.publish(
|
||||
`mcp:shttp:toserver:${sessionId}`,
|
||||
JSON.stringify({
|
||||
type: 'mcp',
|
||||
message: testMessage
|
||||
})
|
||||
);
|
||||
|
||||
// Wait for message to be processed
|
||||
await new Promise(resolve => setImmediate(resolve));
|
||||
|
||||
// Verify message was received
|
||||
expect(onMessageMock).toHaveBeenCalledWith(testMessage, undefined);
|
||||
|
||||
// Clear the publish spy to check only future calls
|
||||
const shutdownSpy = jest.spyOn(mockRedis, 'publish');
|
||||
shutdownSpy.mockClear();
|
||||
|
||||
// Fast-forward 4 more minutes (total 8, but only 4 since last message)
|
||||
jest.advanceTimersByTime(4 * 60 * 1000);
|
||||
|
||||
// Should not have shutdown yet
|
||||
expect(shutdownSpy).not.toHaveBeenCalledWith(
|
||||
`mcp:control:${sessionId}`,
|
||||
expect.stringContaining('"action":"SHUTDOWN"')
|
||||
);
|
||||
|
||||
// Fast-forward 2 more minutes to exceed timeout
|
||||
jest.advanceTimersByTime(2 * 60 * 1000);
|
||||
|
||||
// Now should have shutdown
|
||||
expect(shutdownSpy).toHaveBeenCalledWith(
|
||||
`mcp:control:${sessionId}`,
|
||||
expect.stringContaining('"action":"SHUTDOWN"')
|
||||
);
|
||||
|
||||
await transport.close();
|
||||
}, 10000);
|
||||
|
||||
it('should clear timeout on close', async () => {
|
||||
const transport = new ServerRedisTransport(sessionId);
|
||||
const shutdownSpy = jest.spyOn(mockRedis, 'publish');
|
||||
|
||||
await transport.start();
|
||||
|
||||
// Close transport before timeout
|
||||
await transport.close();
|
||||
|
||||
// Fast-forward past timeout
|
||||
jest.advanceTimersByTime(10 * 60 * 1000);
|
||||
|
||||
// Should not have triggered shutdown
|
||||
expect(shutdownSpy).not.toHaveBeenCalledWith(
|
||||
`mcp:control:${sessionId}`,
|
||||
expect.stringContaining('"action":"SHUTDOWN"')
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
351
src/modules/mcp/services/redisTransport.ts
Normal file
351
src/modules/mcp/services/redisTransport.ts
Normal file
|
|
@ -0,0 +1,351 @@
|
|||
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
|
||||
import { redisClient } from "../../shared/redis.js";
|
||||
import { Transport, TransportSendOptions } from "@modelcontextprotocol/sdk/shared/transport.js";
|
||||
import { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
|
||||
import { JSONRPCMessage, MessageExtraInfo } from "@modelcontextprotocol/sdk/types.js";
|
||||
import { logger } from "../../shared/logger.js";
|
||||
|
||||
let redisTransportCounter = 0;
|
||||
const notificationStreamId = "__GET_stream";
|
||||
|
||||
// Message types for Redis transport
|
||||
type RedisMessage =
|
||||
| {
|
||||
type: 'mcp';
|
||||
message: JSONRPCMessage;
|
||||
extra?: MessageExtraInfo;
|
||||
options?: TransportSendOptions;
|
||||
}
|
||||
| {
|
||||
type: 'control';
|
||||
action: 'SHUTDOWN' | 'PING' | 'STATUS';
|
||||
timestamp?: number;
|
||||
};
|
||||
|
||||
function sendToMcpServer(sessionId: string, message: JSONRPCMessage, extra?: { authInfo?: AuthInfo; }, options?: TransportSendOptions): Promise<void> {
|
||||
const toServerChannel = getToServerChannel(sessionId);
|
||||
|
||||
logger.debug('Sending message to MCP server via Redis', {
|
||||
sessionId,
|
||||
channel: toServerChannel,
|
||||
method: ('method' in message ? message.method : undefined),
|
||||
id: ('id' in message ? message.id : undefined)
|
||||
});
|
||||
|
||||
const redisMessage: RedisMessage = { type: 'mcp', message, extra, options };
|
||||
return redisClient.publish(toServerChannel, JSON.stringify(redisMessage));
|
||||
}
|
||||
|
||||
function getToServerChannel(sessionId: string): string {
|
||||
return `mcp:shttp:toserver:${sessionId}`;
|
||||
}
|
||||
|
||||
function getToClientChannel(sessionId: string, relatedRequestId: string): string {
|
||||
return `mcp:shttp:toclient:${sessionId}:${relatedRequestId}`;
|
||||
}
|
||||
|
||||
function getControlChannel(sessionId: string): string {
|
||||
return `mcp:control:${sessionId}`;
|
||||
}
|
||||
|
||||
function sendControlMessage(sessionId: string, action: 'SHUTDOWN' | 'PING' | 'STATUS'): Promise<void> {
|
||||
const controlChannel = getControlChannel(sessionId);
|
||||
const redisMessage: RedisMessage = {
|
||||
type: 'control',
|
||||
action,
|
||||
timestamp: Date.now()
|
||||
};
|
||||
return redisClient.publish(controlChannel, JSON.stringify(redisMessage));
|
||||
}
|
||||
|
||||
export async function shutdownSession(sessionId: string): Promise<void> {
|
||||
logger.info('Sending shutdown control message', { sessionId });
|
||||
return sendControlMessage(sessionId, 'SHUTDOWN');
|
||||
}
|
||||
|
||||
export async function isLive(sessionId: string): Promise<boolean> {
|
||||
// Check if the session is live by checking if the key exists in Redis
|
||||
const numSubs = await redisClient.numsub(getToServerChannel(sessionId));
|
||||
return numSubs > 0;
|
||||
}
|
||||
|
||||
export async function setSessionOwner(sessionId: string, userId: string): Promise<void> {
|
||||
logger.debug('Setting session owner', { sessionId, userId });
|
||||
await redisClient.set(`session:${sessionId}:owner`, userId);
|
||||
}
|
||||
|
||||
export async function getSessionOwner(sessionId: string): Promise<string | null> {
|
||||
return await redisClient.get(`session:${sessionId}:owner`);
|
||||
}
|
||||
|
||||
export async function validateSessionOwnership(sessionId: string, userId: string): Promise<boolean> {
|
||||
const owner = await getSessionOwner(sessionId);
|
||||
return owner === userId;
|
||||
}
|
||||
|
||||
export async function isSessionOwnedBy(sessionId: string, userId: string): Promise<boolean> {
|
||||
const isLiveSession = await isLive(sessionId);
|
||||
if (!isLiveSession) {
|
||||
logger.debug('Session not live', { sessionId });
|
||||
return false;
|
||||
}
|
||||
const isOwned = await validateSessionOwnership(sessionId, userId);
|
||||
logger.debug('Session ownership check', { sessionId, userId, isOwned });
|
||||
return isOwned;
|
||||
}
|
||||
|
||||
|
||||
export async function redisRelayToMcpServer(sessionId: string, transport: Transport, isGetRequest: boolean = false): Promise<() => Promise<void>> {
|
||||
logger.debug('Setting up Redis relay to MCP server', {
|
||||
sessionId,
|
||||
isGetRequest
|
||||
});
|
||||
|
||||
let redisCleanup: (() => Promise<void>) | undefined = undefined;
|
||||
const cleanup = async () => {
|
||||
// TODO: solve race conditions where we call cleanup while the subscription is being created / before it is created
|
||||
if (redisCleanup) {
|
||||
logger.debug('Cleaning up Redis relay', { sessionId });
|
||||
await redisCleanup();
|
||||
}
|
||||
}
|
||||
|
||||
const subscribe = async (requestId: string) => {
|
||||
const toClientChannel = getToClientChannel(sessionId, requestId);
|
||||
|
||||
logger.debug('Subscribing to client channel', {
|
||||
sessionId,
|
||||
requestId,
|
||||
channel: toClientChannel
|
||||
});
|
||||
|
||||
redisCleanup = await redisClient.createSubscription(toClientChannel, async (redisMessageJson) => {
|
||||
const redisMessage = JSON.parse(redisMessageJson) as RedisMessage;
|
||||
if (redisMessage.type === 'mcp') {
|
||||
logger.debug('Relaying message from Redis to client', {
|
||||
sessionId,
|
||||
requestId,
|
||||
method: ('method' in redisMessage.message ? redisMessage.message.method : undefined)
|
||||
});
|
||||
await transport.send(redisMessage.message, redisMessage.options);
|
||||
}
|
||||
}, (error) => {
|
||||
logger.error('Error in Redis relay subscription', error, {
|
||||
sessionId,
|
||||
channel: toClientChannel
|
||||
});
|
||||
transport.onerror?.(error);
|
||||
});
|
||||
}
|
||||
|
||||
if (isGetRequest) {
|
||||
await subscribe(notificationStreamId);
|
||||
} else {
|
||||
const messagePromise = new Promise<JSONRPCMessage>((resolve) => {
|
||||
transport.onmessage = async (message, extra) => {
|
||||
// First, set up response subscription if needed
|
||||
if ("id" in message) {
|
||||
logger.debug('Setting up response subscription', {
|
||||
sessionId,
|
||||
messageId: message.id,
|
||||
method: ('method' in message ? message.method : undefined)
|
||||
});
|
||||
await subscribe(message.id.toString());
|
||||
}
|
||||
// Now send the message to the MCP server
|
||||
await sendToMcpServer(sessionId, message, extra);
|
||||
resolve(message);
|
||||
}
|
||||
});
|
||||
|
||||
messagePromise.catch((error) => {
|
||||
transport.onerror?.(error);
|
||||
cleanup();
|
||||
});
|
||||
}
|
||||
return cleanup;
|
||||
}
|
||||
|
||||
|
||||
// New Redis transport for server->client messages using request-id based channels
|
||||
export class ServerRedisTransport implements Transport {
|
||||
private counter: number;
|
||||
private _sessionId: string;
|
||||
private controlCleanup?: (() => Promise<void>);
|
||||
private serverCleanup?: (() => Promise<void>);
|
||||
private shouldShutdown = false;
|
||||
private inactivityTimeout?: NodeJS.Timeout;
|
||||
private readonly INACTIVITY_TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes
|
||||
|
||||
onclose?: (() => void) | undefined;
|
||||
onerror?: ((error: Error) => void) | undefined;
|
||||
onmessage?: ((message: JSONRPCMessage, extra?: { authInfo?: AuthInfo; }) => void) | undefined;
|
||||
|
||||
constructor(sessionId: string) {
|
||||
this.counter = redisTransportCounter++;
|
||||
this._sessionId = sessionId;
|
||||
}
|
||||
|
||||
private resetInactivityTimer(): void {
|
||||
// Clear existing timeout if any
|
||||
if (this.inactivityTimeout) {
|
||||
clearTimeout(this.inactivityTimeout);
|
||||
}
|
||||
|
||||
// Set new timeout
|
||||
this.inactivityTimeout = setTimeout(() => {
|
||||
logger.info('Session timed out due to inactivity', {
|
||||
sessionId: this._sessionId,
|
||||
timeoutMs: this.INACTIVITY_TIMEOUT_MS
|
||||
});
|
||||
void shutdownSession(this._sessionId);
|
||||
}, this.INACTIVITY_TIMEOUT_MS);
|
||||
}
|
||||
|
||||
private clearInactivityTimer(): void {
|
||||
if (this.inactivityTimeout) {
|
||||
clearTimeout(this.inactivityTimeout);
|
||||
this.inactivityTimeout = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
async start(): Promise<void> {
|
||||
logger.info('Starting ServerRedisTransport', {
|
||||
sessionId: this._sessionId,
|
||||
inactivityTimeoutMs: this.INACTIVITY_TIMEOUT_MS
|
||||
});
|
||||
|
||||
// Start inactivity timer
|
||||
this.resetInactivityTimer();
|
||||
|
||||
// Subscribe to MCP messages from clients
|
||||
const serverChannel = getToServerChannel(this._sessionId);
|
||||
logger.debug('Subscribing to server channel', {
|
||||
sessionId: this._sessionId,
|
||||
channel: serverChannel
|
||||
});
|
||||
|
||||
this.serverCleanup = await redisClient.createSubscription(
|
||||
serverChannel,
|
||||
(messageJson) => {
|
||||
const redisMessage = JSON.parse(messageJson) as RedisMessage;
|
||||
if (redisMessage.type === 'mcp') {
|
||||
// Reset inactivity timer on each message from client
|
||||
this.resetInactivityTimer();
|
||||
|
||||
logger.debug('Received MCP message from client', {
|
||||
sessionId: this._sessionId,
|
||||
method: ('method' in redisMessage.message ? redisMessage.message.method : undefined),
|
||||
id: ('id' in redisMessage.message ? redisMessage.message.id : undefined)
|
||||
});
|
||||
|
||||
this.onmessage?.(redisMessage.message, redisMessage.extra);
|
||||
}
|
||||
},
|
||||
(error) => {
|
||||
logger.error('Error in server channel subscription', error, {
|
||||
sessionId: this._sessionId,
|
||||
channel: serverChannel
|
||||
});
|
||||
this.onerror?.(error);
|
||||
}
|
||||
);
|
||||
|
||||
// Subscribe to control messages for shutdown
|
||||
const controlChannel = getControlChannel(this._sessionId);
|
||||
logger.debug('Subscribing to control channel', {
|
||||
sessionId: this._sessionId,
|
||||
channel: controlChannel
|
||||
});
|
||||
|
||||
this.controlCleanup = await redisClient.createSubscription(
|
||||
controlChannel,
|
||||
(messageJson) => {
|
||||
const redisMessage = JSON.parse(messageJson) as RedisMessage;
|
||||
if (redisMessage.type === 'control') {
|
||||
logger.info('Received control message', {
|
||||
sessionId: this._sessionId,
|
||||
action: redisMessage.action
|
||||
});
|
||||
|
||||
if (redisMessage.action === 'SHUTDOWN') {
|
||||
logger.info('Shutting down transport due to control message', {
|
||||
sessionId: this._sessionId
|
||||
});
|
||||
this.shouldShutdown = true;
|
||||
this.close();
|
||||
}
|
||||
}
|
||||
},
|
||||
(error) => {
|
||||
logger.error('Error in control channel subscription', error, {
|
||||
sessionId: this._sessionId,
|
||||
channel: controlChannel
|
||||
});
|
||||
this.onerror?.(error);
|
||||
}
|
||||
);
|
||||
|
||||
}
|
||||
|
||||
async send(message: JSONRPCMessage, options?: TransportSendOptions): Promise<void> {
|
||||
const relatedRequestId = options?.relatedRequestId?.toString() ?? ("id" in message ? message.id?.toString() : notificationStreamId);
|
||||
const channel = getToClientChannel(this._sessionId, relatedRequestId)
|
||||
|
||||
logger.debug('Sending message to client', {
|
||||
sessionId: this._sessionId,
|
||||
channel,
|
||||
method: ('method' in message ? message.method : undefined),
|
||||
id: ('id' in message ? message.id : undefined),
|
||||
relatedRequestId
|
||||
});
|
||||
|
||||
const redisMessage: RedisMessage = { type: 'mcp', message, options };
|
||||
const messageStr = JSON.stringify(redisMessage);
|
||||
await redisClient.publish(channel, messageStr);
|
||||
}
|
||||
|
||||
async close(): Promise<void> {
|
||||
logger.info('Closing ServerRedisTransport', {
|
||||
sessionId: this._sessionId,
|
||||
wasShutdown: this.shouldShutdown
|
||||
});
|
||||
|
||||
// Clear inactivity timer
|
||||
this.clearInactivityTimer();
|
||||
|
||||
// Clean up server message subscription
|
||||
if (this.serverCleanup) {
|
||||
await this.serverCleanup();
|
||||
this.serverCleanup = undefined;
|
||||
}
|
||||
|
||||
// Clean up control message subscription
|
||||
if (this.controlCleanup) {
|
||||
await this.controlCleanup();
|
||||
this.controlCleanup = undefined;
|
||||
}
|
||||
|
||||
this.onclose?.();
|
||||
}
|
||||
}
|
||||
|
||||
export async function getShttpTransport(sessionId: string, onsessionclosed: (sessionId: string) => void | Promise<void>, isGetRequest: boolean = false): Promise<StreamableHTTPServerTransport> {
|
||||
logger.debug('Getting StreamableHTTPServerTransport for existing session', {
|
||||
sessionId,
|
||||
isGetRequest
|
||||
});
|
||||
|
||||
// Giving undefined here and setting the sessionId means the
|
||||
// transport wont try to create a new session.
|
||||
const shttpTransport = new StreamableHTTPServerTransport({
|
||||
sessionIdGenerator: undefined,
|
||||
onsessionclosed,
|
||||
})
|
||||
shttpTransport.sessionId = sessionId;
|
||||
|
||||
// Use the new request-id based relay approach
|
||||
const cleanup = await redisRelayToMcpServer(sessionId, shttpTransport, isGetRequest);
|
||||
shttpTransport.onclose = cleanup;
|
||||
return shttpTransport;
|
||||
}
|
||||
33
src/modules/mcp/types.ts
Normal file
33
src/modules/mcp/types.ts
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
/**
|
||||
* OAuth 2.0 Token Introspection Response
|
||||
* Based on RFC 7662: https://tools.ietf.org/html/rfc7662
|
||||
* Used when validating tokens with the external authorization server.
|
||||
*/
|
||||
export interface TokenIntrospectionResponse {
|
||||
/** Whether the token is currently active */
|
||||
active: boolean;
|
||||
/** Space-separated list of scopes associated with the token */
|
||||
scope?: string;
|
||||
/** Client identifier for the OAuth client that requested the token */
|
||||
client_id?: string;
|
||||
/** Human-readable identifier for the resource owner */
|
||||
username?: string;
|
||||
/** Type of the token (e.g., "Bearer") */
|
||||
token_type?: string;
|
||||
/** Expiration time as seconds since Unix epoch */
|
||||
exp?: number;
|
||||
/** Time at which the token was issued as seconds since Unix epoch */
|
||||
iat?: number;
|
||||
/** Time before which the token is not valid as seconds since Unix epoch */
|
||||
nbf?: number;
|
||||
/** Subject identifier for the resource owner */
|
||||
sub?: string;
|
||||
/** Intended audience for the token */
|
||||
aud?: string | string[];
|
||||
/** Issuer of the token */
|
||||
iss?: string;
|
||||
/** Unique identifier for the token */
|
||||
jti?: string;
|
||||
/** Custom field for our implementation to store user ID */
|
||||
userId?: string;
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue