mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-04 14:57:42 +00:00
fix: Close MCP OAuth route tenant gaps
This commit is contained in:
Danny Avila
parent
9b9c2c8646
commit
483b9b05d6
5 changed files with 235 additions and 28 deletions
|
|
@ -1,5 +1,5 @@
|
|||
const mongoose = require('mongoose');
|
||||
const { logger, webSearchKeys } = require('@librechat/data-schemas');
|
||||
const { logger, getTenantId, webSearchKeys } = require('@librechat/data-schemas');
|
||||
const {
|
||||
getNewS3URL,
|
||||
needsRefresh,
|
||||
|
|
@ -407,11 +407,24 @@ const clearStoredMCPOAuthState = async (userId, serverName) => {
|
|||
try {
|
||||
const flowsCache = getLogStores(CacheKeys.FLOWS);
|
||||
const flowManager = getFlowStateManager(flowsCache);
|
||||
const flowId = MCPOAuthHandler.generateFlowId(userId, serverName);
|
||||
const results = await Promise.allSettled([
|
||||
flowManager.deleteFlow(flowId, 'mcp_get_tokens'),
|
||||
flowManager.deleteFlow(flowId, 'mcp_oauth'),
|
||||
]);
|
||||
const baseFlowId = MCPOAuthHandler.generateFlowId(userId, serverName);
|
||||
const tenantId = getTenantId();
|
||||
const tokenFlowId = MCPOAuthHandler.generateTokenFlowId(userId, serverName, tenantId);
|
||||
const oauthFlowId = MCPOAuthHandler.generateFlowId(userId, serverName, tenantId);
|
||||
const flowDeletes = [
|
||||
[tokenFlowId, 'mcp_get_tokens'],
|
||||
[oauthFlowId, 'mcp_oauth'],
|
||||
[baseFlowId, 'mcp_get_tokens'],
|
||||
[baseFlowId, 'mcp_oauth'],
|
||||
].filter(
|
||||
([flowId, type], index, deletes) =>
|
||||
deletes.findIndex(([candidateId, candidateType]) => {
|
||||
return candidateId === flowId && candidateType === type;
|
||||
}) === index,
|
||||
);
|
||||
const results = await Promise.allSettled(
|
||||
flowDeletes.map(([flowId, type]) => flowManager.deleteFlow(flowId, type)),
|
||||
);
|
||||
for (const result of results) {
|
||||
if (result.status === 'rejected') {
|
||||
logger.warn(
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ const mockGetMCPServersRegistry = jest.fn();
|
|||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: { error: jest.fn(), info: jest.fn(), warn: jest.fn() },
|
||||
getTenantId: jest.fn(),
|
||||
webSearchKeys: [],
|
||||
}));
|
||||
|
||||
|
|
@ -22,7 +23,14 @@ jest.mock('librechat-data-provider', () => ({
|
|||
|
||||
jest.mock('@librechat/api', () => ({
|
||||
MCPOAuthHandler: {
|
||||
generateFlowId: jest.fn(() => 'user-1:test-server'),
|
||||
generateFlowId: jest.fn((userId, serverName, tenantId) => {
|
||||
const flowId = `${userId}:${serverName}`;
|
||||
return tenantId ? `tenant:${encodeURIComponent(tenantId)}:${flowId}` : flowId;
|
||||
}),
|
||||
generateTokenFlowId: jest.fn((userId, serverName, tenantId) => {
|
||||
const flowId = `${userId}:${serverName}`;
|
||||
return tenantId ? `tenant:${encodeURIComponent(tenantId)}:${flowId}` : flowId;
|
||||
}),
|
||||
revokeOAuthToken: jest.fn(),
|
||||
},
|
||||
MCPTokenStorage: {
|
||||
|
|
@ -78,7 +86,7 @@ jest.mock('~/cache', () => ({
|
|||
getLogStores: (...args) => mockGetLogStores(...args),
|
||||
}));
|
||||
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const { logger, getTenantId } = require('@librechat/data-schemas');
|
||||
const { MCPTokenStorage, MCPOAuthHandler } = require('@librechat/api');
|
||||
const { updateUserPluginsController } = require('~/server/controllers/UserController');
|
||||
|
||||
|
|
@ -138,6 +146,7 @@ function setupMCPMocks() {
|
|||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
getTenantId.mockReturnValue(undefined);
|
||||
});
|
||||
|
||||
describe('updateUserPluginsController MCP OAuth cleanup', () => {
|
||||
|
|
@ -231,6 +240,27 @@ describe('updateUserPluginsController MCP OAuth cleanup', () => {
|
|||
expect(MCPOAuthHandler.revokeOAuthToken).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('clears tenant-scoped and legacy OAuth flow state when tenant context exists', async () => {
|
||||
const { flowManager } = setupMCPMocks();
|
||||
getTenantId.mockReturnValue('tenant-a');
|
||||
MCPTokenStorage.getClientInfoAndMetadata.mockResolvedValue(null);
|
||||
|
||||
const res = createResponse();
|
||||
await updateUserPluginsController(createRequest(), res);
|
||||
|
||||
expect(res.status).toHaveBeenCalledWith(200);
|
||||
expect(flowManager.deleteFlow).toHaveBeenCalledWith(
|
||||
'tenant:tenant-a:user-1:test-server',
|
||||
'mcp_get_tokens',
|
||||
);
|
||||
expect(flowManager.deleteFlow).toHaveBeenCalledWith(
|
||||
'tenant:tenant-a:user-1:test-server',
|
||||
'mcp_oauth',
|
||||
);
|
||||
expect(flowManager.deleteFlow).toHaveBeenCalledWith('user-1:test-server', 'mcp_get_tokens');
|
||||
expect(flowManager.deleteFlow).toHaveBeenCalledWith('user-1:test-server', 'mcp_oauth');
|
||||
});
|
||||
|
||||
it('clears stored OAuth token state when server config is missing', async () => {
|
||||
const { flowManager, registry } = setupMCPMocks();
|
||||
registry.getServerConfig.mockResolvedValue(undefined);
|
||||
|
|
|
|||
|
|
@ -13,9 +13,11 @@ const mockDeleteTokens = jest.fn();
|
|||
const mockLoggerInfo = jest.fn();
|
||||
const mockLoggerWarn = jest.fn();
|
||||
const mockLoggerError = jest.fn();
|
||||
const mockGetTenantId = jest.fn();
|
||||
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: { info: mockLoggerInfo, warn: mockLoggerWarn, error: mockLoggerError },
|
||||
getTenantId: (...args) => mockGetTenantId(...args),
|
||||
webSearchKeys: [],
|
||||
}));
|
||||
|
||||
|
|
@ -23,7 +25,14 @@ jest.mock('@librechat/api', () => {
|
|||
return {
|
||||
MCPOAuthHandler: {
|
||||
revokeOAuthToken: (...args) => mockRevokeOAuthToken(...args),
|
||||
generateFlowId: (userId, serverName) => `${userId}:${serverName}`,
|
||||
generateFlowId: (userId, serverName, tenantId) => {
|
||||
const flowId = `${userId}:${serverName}`;
|
||||
return tenantId ? `tenant:${encodeURIComponent(tenantId)}:${flowId}` : flowId;
|
||||
},
|
||||
generateTokenFlowId: (userId, serverName, tenantId) => {
|
||||
const flowId = `${userId}:${serverName}`;
|
||||
return tenantId ? `tenant:${encodeURIComponent(tenantId)}:${flowId}` : flowId;
|
||||
},
|
||||
},
|
||||
MCPTokenStorage: {
|
||||
getTokens: (...args) => mockGetTokens(...args),
|
||||
|
|
@ -151,6 +160,7 @@ function setupOAuthServerFound() {
|
|||
describe('maybeUninstallOAuthMCP', () => {
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockGetTenantId.mockReturnValue(undefined);
|
||||
});
|
||||
|
||||
test('is a no-op when pluginKey is not an MCP key', async () => {
|
||||
|
|
@ -205,6 +215,20 @@ describe('maybeUninstallOAuthMCP', () => {
|
|||
);
|
||||
});
|
||||
|
||||
test('clears tenant-scoped and legacy flow state when tenant context exists', async () => {
|
||||
setupOAuthServerFound();
|
||||
mockGetTenantId.mockReturnValue('tenant-a');
|
||||
mockGetClientInfoAndMetadata.mockResolvedValue(null);
|
||||
|
||||
await maybeUninstallOAuthMCP(userId, pluginKey, appConfig);
|
||||
|
||||
expect(mockDeleteFlow).toHaveBeenCalledTimes(4);
|
||||
expect(mockDeleteFlow).toHaveBeenCalledWith('tenant:tenant-a:user-123:acme', 'mcp_get_tokens');
|
||||
expect(mockDeleteFlow).toHaveBeenCalledWith('tenant:tenant-a:user-123:acme', 'mcp_oauth');
|
||||
expect(mockDeleteFlow).toHaveBeenCalledWith('user-123:acme', 'mcp_get_tokens');
|
||||
expect(mockDeleteFlow).toHaveBeenCalledWith('user-123:acme', 'mcp_oauth');
|
||||
});
|
||||
|
||||
test('revokes both tokens and runs cleanup on happy path', async () => {
|
||||
setupOAuthServerFound();
|
||||
mockGetTokens.mockResolvedValue({
|
||||
|
|
|
|||
|
|
@ -458,6 +458,26 @@ describe('MCP Routes', () => {
|
|||
expect(response.headers.location).toBe(`${basePath}/oauth/error?error=invalid_client`);
|
||||
expect(mockFlowManager.failFlow).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should redirect instead of hanging when OAuth error flow ID is malformed', async () => {
|
||||
const mockFlowManager = {
|
||||
failFlow: jest.fn(),
|
||||
};
|
||||
|
||||
getLogStores.mockReturnValueOnce({});
|
||||
require('~/config').getFlowStateManager.mockReturnValueOnce(mockFlowManager);
|
||||
MCPOAuthHandler.resolveStateToFlowId.mockResolvedValueOnce('malformed-flow-id');
|
||||
|
||||
const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({
|
||||
error: 'invalid_client',
|
||||
state: 'opaque-state',
|
||||
});
|
||||
const basePath = getBasePath();
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
expect(response.headers.location).toBe(`${basePath}/oauth/error?error=invalid_client`);
|
||||
expect(mockFlowManager.failFlow).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
it('should redirect to error page when code is missing', async () => {
|
||||
|
|
@ -800,6 +820,74 @@ describe('MCP Routes', () => {
|
|||
expect(mockFlowManager.deleteFlow).toHaveBeenCalledWith(flowId, 'mcp_get_tokens');
|
||||
});
|
||||
|
||||
it('should complete pending token flow waiters after storing callback tokens', async () => {
|
||||
const mockFlowManager = {
|
||||
getFlowState: jest.fn().mockImplementation((id, type) => {
|
||||
if (type === 'mcp_get_tokens' && id === 'tenant:tenant-a:test-user-id:test-server') {
|
||||
return Promise.resolve({
|
||||
type: 'mcp_get_tokens',
|
||||
status: 'PENDING',
|
||||
});
|
||||
}
|
||||
return Promise.resolve({ status: 'PENDING' });
|
||||
}),
|
||||
completeFlow: jest.fn().mockResolvedValue(true),
|
||||
deleteFlow: jest.fn().mockResolvedValue(true),
|
||||
};
|
||||
const mockFlowState = {
|
||||
serverName: 'test-server',
|
||||
userId: 'test-user-id',
|
||||
metadata: {},
|
||||
clientInfo: {},
|
||||
codeVerifier: 'test-verifier',
|
||||
tenantId: 'tenant-a',
|
||||
};
|
||||
const mockTokens = {
|
||||
access_token: 'fresh-access-token',
|
||||
refresh_token: 'fresh-refresh-token',
|
||||
};
|
||||
|
||||
MCPOAuthHandler.getFlowState.mockResolvedValue(mockFlowState);
|
||||
MCPOAuthHandler.completeOAuthFlow.mockResolvedValue(mockTokens);
|
||||
MCPTokenStorage.storeTokens.mockResolvedValue();
|
||||
getLogStores.mockReturnValue({});
|
||||
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||
require('~/config').getOAuthReconnectionManager.mockReturnValue({
|
||||
clearReconnection: jest.fn(),
|
||||
});
|
||||
require('~/config').getMCPManager.mockReturnValue({
|
||||
getUserConnection: jest.fn().mockResolvedValue({
|
||||
fetchTools: jest.fn().mockResolvedValue([]),
|
||||
}),
|
||||
});
|
||||
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
|
||||
getCachedTools.mockResolvedValue({});
|
||||
setCachedTools.mockResolvedValue();
|
||||
|
||||
const flowId = 'test-user-id:test-server';
|
||||
const csrfToken = generateTestCsrfToken(flowId);
|
||||
|
||||
const response = await request(app)
|
||||
.get('/api/mcp/test-server/oauth/callback')
|
||||
.set('Cookie', [`oauth_csrf=${csrfToken}`])
|
||||
.query({
|
||||
code: 'test-auth-code',
|
||||
state: flowId,
|
||||
});
|
||||
|
||||
expect(response.status).toBe(302);
|
||||
expect(mockFlowManager.completeFlow).toHaveBeenCalledWith(
|
||||
'tenant:tenant-a:test-user-id:test-server',
|
||||
'mcp_get_tokens',
|
||||
mockTokens,
|
||||
);
|
||||
expect(mockFlowManager.deleteFlow).not.toHaveBeenCalledWith(
|
||||
'tenant:tenant-a:test-user-id:test-server',
|
||||
'mcp_get_tokens',
|
||||
);
|
||||
expect(mockFlowManager.deleteFlow).toHaveBeenCalledWith(flowId, 'mcp_get_tokens');
|
||||
});
|
||||
|
||||
it('should use oauthHeaders from flow state when present', async () => {
|
||||
const mockFlowManager = {
|
||||
getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING' }),
|
||||
|
|
@ -1211,6 +1299,7 @@ describe('MCP Routes', () => {
|
|||
});
|
||||
|
||||
it('should return tokens for a tenant-prefixed flow owned by the user', async () => {
|
||||
const { getTenantId } = require('@librechat/data-schemas');
|
||||
const mockFlowManager = {
|
||||
getFlowState: jest.fn().mockResolvedValue({
|
||||
status: 'COMPLETED',
|
||||
|
|
@ -1220,6 +1309,7 @@ describe('MCP Routes', () => {
|
|||
}),
|
||||
};
|
||||
|
||||
getTenantId.mockReturnValue('tenant-a');
|
||||
getLogStores.mockReturnValue({});
|
||||
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||
|
||||
|
|
@ -1239,6 +1329,18 @@ describe('MCP Routes', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('should reject tenant-prefixed token flow access from another tenant', async () => {
|
||||
const { getTenantId } = require('@librechat/data-schemas');
|
||||
getTenantId.mockReturnValue('tenant-b');
|
||||
|
||||
const response = await request(app).get(
|
||||
'/api/mcp/oauth/tokens/tenant:tenant-a:test-user-id:test-server',
|
||||
);
|
||||
|
||||
expect(response.status).toBe(403);
|
||||
expect(response.body).toEqual({ error: 'Access denied' });
|
||||
});
|
||||
|
||||
it('should return 401 when user is not authenticated', async () => {
|
||||
const unauthApp = express();
|
||||
unauthApp.use(express.json());
|
||||
|
|
@ -1332,6 +1434,7 @@ describe('MCP Routes', () => {
|
|||
});
|
||||
|
||||
it('should return flow status for a tenant-prefixed flow owned by the user', async () => {
|
||||
const { getTenantId } = require('@librechat/data-schemas');
|
||||
const mockFlowManager = {
|
||||
getFlowState: jest.fn().mockResolvedValue({
|
||||
status: 'PENDING',
|
||||
|
|
@ -1339,6 +1442,7 @@ describe('MCP Routes', () => {
|
|||
}),
|
||||
};
|
||||
|
||||
getTenantId.mockReturnValue('tenant-a');
|
||||
getLogStores.mockReturnValue({});
|
||||
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
|
||||
|
||||
|
|
@ -1359,6 +1463,18 @@ describe('MCP Routes', () => {
|
|||
);
|
||||
});
|
||||
|
||||
it('should reject tenant-prefixed status access from another tenant', async () => {
|
||||
const { getTenantId } = require('@librechat/data-schemas');
|
||||
getTenantId.mockReturnValue('tenant-b');
|
||||
|
||||
const response = await request(app).get(
|
||||
'/api/mcp/oauth/status/tenant:tenant-a:test-user-id:test-server',
|
||||
);
|
||||
|
||||
expect(response.status).toBe(403);
|
||||
expect(response.body).toEqual({ error: 'Access denied' });
|
||||
});
|
||||
|
||||
it('should return 403 when flowId does not match authenticated user', async () => {
|
||||
const response = await request(app).get('/api/mcp/oauth/status/other-user-id:test-server');
|
||||
|
||||
|
|
|
|||
|
|
@ -55,11 +55,24 @@ const OAUTH_CSRF_COOKIE_PATH = '/api/mcp';
|
|||
const getOAuthFlowId = (userId, serverName) =>
|
||||
MCPOAuthHandler.generateFlowId(userId, serverName, getTenantId());
|
||||
|
||||
const getFlowUserId = (flowId) => MCPOAuthHandler.parseFlowId(flowId)?.userId;
|
||||
|
||||
const canAccessOAuthFlow = (flowId, userId) => {
|
||||
const flowUserId = getFlowUserId(flowId);
|
||||
return flowUserId === userId || flowUserId === 'system';
|
||||
const parsed = MCPOAuthHandler.parseFlowId(flowId);
|
||||
if (!parsed) {
|
||||
return false;
|
||||
}
|
||||
if (parsed.tenantId && parsed.tenantId !== getTenantId()) {
|
||||
return false;
|
||||
}
|
||||
return parsed.userId === userId || parsed.userId === 'system';
|
||||
};
|
||||
|
||||
const clearGetTokensFlow = async ({ flowManager, flowId, tokens }) => {
|
||||
const state = await flowManager.getFlowState(flowId, 'mcp_get_tokens');
|
||||
if (state?.type === 'mcp_get_tokens' && state.status === 'PENDING') {
|
||||
await flowManager.completeFlow(flowId, 'mcp_get_tokens', tokens);
|
||||
return;
|
||||
}
|
||||
await flowManager.deleteFlow(flowId, 'mcp_get_tokens');
|
||||
};
|
||||
|
||||
const checkMCPUsePermissions = generateCheckAccess({
|
||||
|
|
@ -173,18 +186,21 @@ router.get('/:serverName/oauth/callback', async (req, res) => {
|
|||
const flowManager = getFlowStateManager(flowsCache);
|
||||
const flowId = await MCPOAuthHandler.resolveStateToFlowId(state, flowManager);
|
||||
if (flowId) {
|
||||
const flowUserId = getFlowUserId(flowId);
|
||||
if (!flowUserId) {
|
||||
return;
|
||||
}
|
||||
const hasCsrf = validateOAuthCsrf(req, res, flowId, OAUTH_CSRF_COOKIE_PATH);
|
||||
const hasSession = !hasCsrf && validateOAuthSession(req, flowUserId);
|
||||
if (hasCsrf || hasSession) {
|
||||
await flowManager.failFlow(flowId, 'mcp_oauth', String(oauthError));
|
||||
logger.debug('[MCP OAuth] Marked flow as FAILED with OAuth error', {
|
||||
const parsed = MCPOAuthHandler.parseFlowId(flowId);
|
||||
if (!parsed) {
|
||||
logger.warn('[MCP OAuth] Invalid flow ID format for OAuth error callback', {
|
||||
flowId,
|
||||
error: oauthError,
|
||||
});
|
||||
} else {
|
||||
const hasCsrf = validateOAuthCsrf(req, res, flowId, OAUTH_CSRF_COOKIE_PATH);
|
||||
const hasSession = !hasCsrf && validateOAuthSession(req, parsed.userId);
|
||||
if (hasCsrf || hasSession) {
|
||||
await flowManager.failFlow(flowId, 'mcp_oauth', String(oauthError));
|
||||
logger.debug('[MCP OAuth] Marked flow as FAILED with OAuth error', {
|
||||
flowId,
|
||||
error: oauthError,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (err) {
|
||||
|
|
@ -216,14 +232,14 @@ router.get('/:serverName/oauth/callback', async (req, res) => {
|
|||
}
|
||||
logger.debug('[MCP OAuth] Resolved flow ID from state', { flowId });
|
||||
|
||||
const flowUserId = getFlowUserId(flowId);
|
||||
if (!flowUserId) {
|
||||
const parsedFlowId = MCPOAuthHandler.parseFlowId(flowId);
|
||||
if (!parsedFlowId) {
|
||||
logger.error('[MCP OAuth] Invalid flow ID format', { flowId });
|
||||
return res.redirect(`${basePath}/oauth/error?error=invalid_state`);
|
||||
}
|
||||
|
||||
const hasCsrf = validateOAuthCsrf(req, res, flowId, OAUTH_CSRF_COOKIE_PATH);
|
||||
const hasSession = !hasCsrf && validateOAuthSession(req, flowUserId);
|
||||
const hasSession = !hasCsrf && validateOAuthSession(req, parsedFlowId.userId);
|
||||
let hasActiveFlow = false;
|
||||
if (!hasCsrf && !hasSession) {
|
||||
const pendingFlow = await flowManager.getFlowState(flowId, 'mcp_oauth');
|
||||
|
|
@ -345,9 +361,17 @@ router.get('/:serverName/oauth/callback', async (req, res) => {
|
|||
serverName,
|
||||
flowState.tenantId,
|
||||
);
|
||||
await flowManager.deleteFlow(tokenFlowId, 'mcp_get_tokens');
|
||||
await clearGetTokensFlow({
|
||||
flowManager,
|
||||
flowId: tokenFlowId,
|
||||
tokens,
|
||||
});
|
||||
if (tokenFlowId !== flowId) {
|
||||
await flowManager.deleteFlow(flowId, 'mcp_get_tokens');
|
||||
await clearGetTokensFlow({
|
||||
flowManager,
|
||||
flowId,
|
||||
tokens,
|
||||
});
|
||||
}
|
||||
} catch (error) {
|
||||
logger.warn('[MCP OAuth] Failed to clear cached token flow state', error);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue