diff --git a/packages/core/src/mcp/token-storage/hybrid-token-storage.test.ts b/packages/core/src/mcp/token-storage/hybrid-token-storage.test.ts new file mode 100644 index 0000000000..5303d84770 --- /dev/null +++ b/packages/core/src/mcp/token-storage/hybrid-token-storage.test.ts @@ -0,0 +1,274 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import { HybridTokenStorage } from './hybrid-token-storage.js'; +import { KeychainTokenStorage } from './keychain-token-storage.js'; +import { FileTokenStorage } from './file-token-storage.js'; +import { type OAuthCredentials, TokenStorageType } from './types.js'; + +vi.mock('./keychain-token-storage.js', () => ({ + KeychainTokenStorage: vi.fn().mockImplementation(() => ({ + isAvailable: vi.fn(), + getCredentials: vi.fn(), + setCredentials: vi.fn(), + deleteCredentials: vi.fn(), + listServers: vi.fn(), + getAllCredentials: vi.fn(), + clearAll: vi.fn(), + })), +})); + +vi.mock('./file-token-storage.js', () => ({ + FileTokenStorage: vi.fn().mockImplementation(() => ({ + getCredentials: vi.fn(), + setCredentials: vi.fn(), + deleteCredentials: vi.fn(), + listServers: vi.fn(), + getAllCredentials: vi.fn(), + clearAll: vi.fn(), + })), +})); + +interface MockStorage { + isAvailable?: ReturnType; + getCredentials: ReturnType; + setCredentials: ReturnType; + deleteCredentials: ReturnType; + listServers: ReturnType; + getAllCredentials: ReturnType; + clearAll: ReturnType; +} + +describe('HybridTokenStorage', () => { + let storage: HybridTokenStorage; + let mockKeychainStorage: MockStorage; + let mockFileStorage: MockStorage; + const originalEnv = process.env; + + beforeEach(() => { + vi.clearAllMocks(); + process.env = { ...originalEnv }; + + // Create mock instances before creating HybridTokenStorage + mockKeychainStorage = { + isAvailable: vi.fn(), + getCredentials: vi.fn(), + setCredentials: vi.fn(), + deleteCredentials: vi.fn(), + listServers: vi.fn(), + getAllCredentials: vi.fn(), + clearAll: vi.fn(), + }; + + mockFileStorage = { + getCredentials: vi.fn(), + setCredentials: vi.fn(), + deleteCredentials: vi.fn(), + listServers: vi.fn(), + getAllCredentials: vi.fn(), + clearAll: vi.fn(), + }; + + ( + KeychainTokenStorage as unknown as ReturnType + ).mockImplementation(() => mockKeychainStorage); + ( + FileTokenStorage as unknown as ReturnType + ).mockImplementation(() => mockFileStorage); + + storage = new HybridTokenStorage('test-service'); + }); + + afterEach(() => { + process.env = originalEnv; + }); + + describe('storage selection', () => { + it('should use keychain when available', async () => { + mockKeychainStorage.isAvailable!.mockResolvedValue(true); + mockKeychainStorage.getCredentials.mockResolvedValue(null); + + await storage.getCredentials('test-server'); + + expect(mockKeychainStorage.isAvailable).toHaveBeenCalled(); + expect(mockKeychainStorage.getCredentials).toHaveBeenCalledWith( + 'test-server', + ); + expect(await storage.getStorageType()).toBe(TokenStorageType.KEYCHAIN); + }); + + it('should use file storage when GEMINI_FORCE_FILE_STORAGE is set', async () => { + process.env['GEMINI_FORCE_FILE_STORAGE'] = 'true'; + mockFileStorage.getCredentials.mockResolvedValue(null); + + await storage.getCredentials('test-server'); + + expect(mockKeychainStorage.isAvailable).not.toHaveBeenCalled(); + expect(mockFileStorage.getCredentials).toHaveBeenCalledWith( + 'test-server', + ); + expect(await storage.getStorageType()).toBe( + TokenStorageType.ENCRYPTED_FILE, + ); + }); + + it('should fall back to file storage when keychain is unavailable', async () => { + mockKeychainStorage.isAvailable!.mockResolvedValue(false); + mockFileStorage.getCredentials.mockResolvedValue(null); + + await storage.getCredentials('test-server'); + + expect(mockKeychainStorage.isAvailable).toHaveBeenCalled(); + expect(mockFileStorage.getCredentials).toHaveBeenCalledWith( + 'test-server', + ); + expect(await storage.getStorageType()).toBe( + TokenStorageType.ENCRYPTED_FILE, + ); + }); + + it('should fall back to file storage when keychain throws error', async () => { + mockKeychainStorage.isAvailable!.mockRejectedValue( + new Error('Keychain error'), + ); + mockFileStorage.getCredentials.mockResolvedValue(null); + + await storage.getCredentials('test-server'); + + expect(mockKeychainStorage.isAvailable).toHaveBeenCalled(); + expect(mockFileStorage.getCredentials).toHaveBeenCalledWith( + 'test-server', + ); + expect(await storage.getStorageType()).toBe( + TokenStorageType.ENCRYPTED_FILE, + ); + }); + + it('should cache storage selection', async () => { + mockKeychainStorage.isAvailable!.mockResolvedValue(true); + mockKeychainStorage.getCredentials.mockResolvedValue(null); + + await storage.getCredentials('test-server'); + await storage.getCredentials('another-server'); + + expect(mockKeychainStorage.isAvailable).toHaveBeenCalledTimes(1); + }); + }); + + describe('getCredentials', () => { + it('should delegate to selected storage', async () => { + const credentials: OAuthCredentials = { + serverName: 'test-server', + token: { + accessToken: 'access-token', + tokenType: 'Bearer', + }, + updatedAt: Date.now(), + }; + + mockKeychainStorage.isAvailable!.mockResolvedValue(true); + mockKeychainStorage.getCredentials.mockResolvedValue(credentials); + + const result = await storage.getCredentials('test-server'); + + expect(result).toEqual(credentials); + expect(mockKeychainStorage.getCredentials).toHaveBeenCalledWith( + 'test-server', + ); + }); + }); + + describe('setCredentials', () => { + it('should delegate to selected storage', async () => { + const credentials: OAuthCredentials = { + serverName: 'test-server', + token: { + accessToken: 'access-token', + tokenType: 'Bearer', + }, + updatedAt: Date.now(), + }; + + mockKeychainStorage.isAvailable!.mockResolvedValue(true); + mockKeychainStorage.setCredentials.mockResolvedValue(undefined); + + await storage.setCredentials(credentials); + + expect(mockKeychainStorage.setCredentials).toHaveBeenCalledWith( + credentials, + ); + }); + }); + + describe('deleteCredentials', () => { + it('should delegate to selected storage', async () => { + mockKeychainStorage.isAvailable!.mockResolvedValue(true); + mockKeychainStorage.deleteCredentials.mockResolvedValue(undefined); + + await storage.deleteCredentials('test-server'); + + expect(mockKeychainStorage.deleteCredentials).toHaveBeenCalledWith( + 'test-server', + ); + }); + }); + + describe('listServers', () => { + it('should delegate to selected storage', async () => { + const servers = ['server1', 'server2']; + mockKeychainStorage.isAvailable!.mockResolvedValue(true); + mockKeychainStorage.listServers.mockResolvedValue(servers); + + const result = await storage.listServers(); + + expect(result).toEqual(servers); + expect(mockKeychainStorage.listServers).toHaveBeenCalled(); + }); + }); + + describe('getAllCredentials', () => { + it('should delegate to selected storage', async () => { + const credentialsMap = new Map([ + [ + 'server1', + { + serverName: 'server1', + token: { accessToken: 'token1', tokenType: 'Bearer' }, + updatedAt: Date.now(), + }, + ], + [ + 'server2', + { + serverName: 'server2', + token: { accessToken: 'token2', tokenType: 'Bearer' }, + updatedAt: Date.now(), + }, + ], + ]); + + mockKeychainStorage.isAvailable!.mockResolvedValue(true); + mockKeychainStorage.getAllCredentials.mockResolvedValue(credentialsMap); + + const result = await storage.getAllCredentials(); + + expect(result).toEqual(credentialsMap); + expect(mockKeychainStorage.getAllCredentials).toHaveBeenCalled(); + }); + }); + + describe('clearAll', () => { + it('should delegate to selected storage', async () => { + mockKeychainStorage.isAvailable!.mockResolvedValue(true); + mockKeychainStorage.clearAll.mockResolvedValue(undefined); + + await storage.clearAll(); + + expect(mockKeychainStorage.clearAll).toHaveBeenCalled(); + }); + }); +}); diff --git a/packages/core/src/mcp/token-storage/hybrid-token-storage.ts b/packages/core/src/mcp/token-storage/hybrid-token-storage.ts new file mode 100644 index 0000000000..70edac07cd --- /dev/null +++ b/packages/core/src/mcp/token-storage/hybrid-token-storage.ts @@ -0,0 +1,97 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { BaseTokenStorage } from './base-token-storage.js'; +import { FileTokenStorage } from './file-token-storage.js'; +import type { TokenStorage, OAuthCredentials } from './types.js'; +import { TokenStorageType } from './types.js'; + +const FORCE_FILE_STORAGE_ENV_VAR = 'GEMINI_FORCE_FILE_STORAGE'; + +export class HybridTokenStorage extends BaseTokenStorage { + private storage: TokenStorage | null = null; + private storageType: TokenStorageType | null = null; + private storageInitPromise: Promise | null = null; + + constructor(serviceName: string) { + super(serviceName); + } + + private async initializeStorage(): Promise { + const forceFileStorage = process.env[FORCE_FILE_STORAGE_ENV_VAR] === 'true'; + + if (!forceFileStorage) { + try { + const { KeychainTokenStorage } = await import( + './keychain-token-storage.js' + ); + const keychainStorage = new KeychainTokenStorage(this.serviceName); + + const isAvailable = await keychainStorage.isAvailable(); + if (isAvailable) { + this.storage = keychainStorage; + this.storageType = TokenStorageType.KEYCHAIN; + return this.storage; + } + } catch (_e) { + // Fallback to file storage if keychain fails to initialize + } + } + + this.storage = new FileTokenStorage(this.serviceName); + this.storageType = TokenStorageType.ENCRYPTED_FILE; + return this.storage; + } + + private async getStorage(): Promise { + if (this.storage !== null) { + return this.storage; + } + + // Use a single initialization promise to avoid race conditions + if (!this.storageInitPromise) { + this.storageInitPromise = this.initializeStorage(); + } + + // Wait for initialization to complete + return await this.storageInitPromise; + } + + async getCredentials(serverName: string): Promise { + const storage = await this.getStorage(); + return storage.getCredentials(serverName); + } + + async setCredentials(credentials: OAuthCredentials): Promise { + const storage = await this.getStorage(); + await storage.setCredentials(credentials); + } + + async deleteCredentials(serverName: string): Promise { + const storage = await this.getStorage(); + await storage.deleteCredentials(serverName); + } + + async listServers(): Promise { + const storage = await this.getStorage(); + return storage.listServers(); + } + + async getAllCredentials(): Promise> { + const storage = await this.getStorage(); + return storage.getAllCredentials(); + } + + async clearAll(): Promise { + const storage = await this.getStorage(); + await storage.clearAll(); + } + + async getStorageType(): Promise { + await this.getStorage(); + return this.storageType!; + } +} diff --git a/packages/core/src/mcp/token-storage/keychain-token-storage.test.ts b/packages/core/src/mcp/token-storage/keychain-token-storage.test.ts new file mode 100644 index 0000000000..5b34ed01b5 --- /dev/null +++ b/packages/core/src/mcp/token-storage/keychain-token-storage.test.ts @@ -0,0 +1,352 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import type { KeychainTokenStorage } from './keychain-token-storage.js'; +import type { OAuthCredentials } from './types.js'; + +// Hoist the mock to be available in the vi.mock factory +const mockKeytar = vi.hoisted(() => ({ + getPassword: vi.fn(), + setPassword: vi.fn(), + deletePassword: vi.fn(), + findCredentials: vi.fn(), +})); + +const mockServiceName = 'service-name'; +const mockCryptoRandomBytesString = 'random-string'; + +// Mock the dynamic import of 'keytar' +vi.mock('keytar', () => ({ + default: mockKeytar, +})); + +vi.mock('node:crypto', () => ({ + randomBytes: vi.fn(() => ({ + toString: vi.fn(() => mockCryptoRandomBytesString), + })), +})); + +describe('KeychainTokenStorage', () => { + let storage: KeychainTokenStorage; + + beforeEach(async () => { + vi.resetAllMocks(); + // Reset the internal state of the keychain-token-storage module + vi.resetModules(); + const { KeychainTokenStorage } = await import( + './keychain-token-storage.js' + ); + storage = new KeychainTokenStorage(mockServiceName); + }); + + afterEach(() => { + vi.restoreAllMocks(); + vi.useRealTimers(); + }); + + const validCredentials = { + serverName: 'test-server', + token: { + accessToken: 'access-token', + tokenType: 'Bearer', + expiresAt: Date.now() + 3600000, + }, + updatedAt: Date.now(), + } as OAuthCredentials; + + describe('checkKeychainAvailability', () => { + it('should return true if keytar is available and functional', async () => { + mockKeytar.setPassword.mockResolvedValue(undefined); + mockKeytar.getPassword.mockResolvedValue('test'); + mockKeytar.deletePassword.mockResolvedValue(true); + + const isAvailable = await storage.checkKeychainAvailability(); + expect(isAvailable).toBe(true); + expect(mockKeytar.setPassword).toHaveBeenCalledWith( + mockServiceName, + `__keychain_test__${mockCryptoRandomBytesString}`, + 'test', + ); + expect(mockKeytar.getPassword).toHaveBeenCalledWith( + mockServiceName, + `__keychain_test__${mockCryptoRandomBytesString}`, + ); + expect(mockKeytar.deletePassword).toHaveBeenCalledWith( + mockServiceName, + `__keychain_test__${mockCryptoRandomBytesString}`, + ); + }); + + it('should return false if keytar fails to set password', async () => { + mockKeytar.setPassword.mockRejectedValue(new Error('write error')); + const isAvailable = await storage.checkKeychainAvailability(); + expect(isAvailable).toBe(false); + }); + + it('should return false if retrieved password does not match', async () => { + mockKeytar.setPassword.mockResolvedValue(undefined); + mockKeytar.getPassword.mockResolvedValue('wrong-password'); + mockKeytar.deletePassword.mockResolvedValue(true); + const isAvailable = await storage.checkKeychainAvailability(); + expect(isAvailable).toBe(false); + }); + + it('should cache the availability result', async () => { + mockKeytar.setPassword.mockResolvedValue(undefined); + mockKeytar.getPassword.mockResolvedValue('test'); + mockKeytar.deletePassword.mockResolvedValue(true); + + await storage.checkKeychainAvailability(); + await storage.checkKeychainAvailability(); + + expect(mockKeytar.setPassword).toHaveBeenCalledTimes(1); + }); + }); + + describe('with keychain unavailable', () => { + beforeEach(async () => { + // Force keychain to be unavailable + mockKeytar.setPassword.mockRejectedValue(new Error('keychain error')); + await storage.checkKeychainAvailability(); + }); + + it('getCredentials should throw', async () => { + await expect(storage.getCredentials('server')).rejects.toThrow( + 'Keychain is not available', + ); + }); + + it('setCredentials should throw', async () => { + await expect(storage.setCredentials(validCredentials)).rejects.toThrow( + 'Keychain is not available', + ); + }); + + it('deleteCredentials should throw', async () => { + await expect(storage.deleteCredentials('server')).rejects.toThrow( + 'Keychain is not available', + ); + }); + + it('listServers should throw', async () => { + await expect(storage.listServers()).rejects.toThrow( + 'Keychain is not available', + ); + }); + + it('getAllCredentials should throw', async () => { + await expect(storage.getAllCredentials()).rejects.toThrow( + 'Keychain is not available', + ); + }); + }); + + describe('with keychain available', () => { + beforeEach(async () => { + mockKeytar.setPassword.mockResolvedValue(undefined); + mockKeytar.getPassword.mockResolvedValue('test'); + mockKeytar.deletePassword.mockResolvedValue(true); + await storage.checkKeychainAvailability(); + // Reset mocks after availability check + vi.resetAllMocks(); + }); + + describe('getCredentials', () => { + it('should return null if no credentials are found', async () => { + mockKeytar.getPassword.mockResolvedValue(null); + const result = await storage.getCredentials('test-server'); + expect(result).toBeNull(); + expect(mockKeytar.getPassword).toHaveBeenCalledWith( + mockServiceName, + 'test-server', + ); + }); + + it('should return credentials if found and not expired', async () => { + mockKeytar.getPassword.mockResolvedValue( + JSON.stringify(validCredentials), + ); + const result = await storage.getCredentials('test-server'); + expect(result).toEqual(validCredentials); + }); + + it('should return null if credentials have expired', async () => { + const expiredCreds = { + ...validCredentials, + token: { ...validCredentials.token, expiresAt: Date.now() - 1000 }, + }; + mockKeytar.getPassword.mockResolvedValue(JSON.stringify(expiredCreds)); + const result = await storage.getCredentials('test-server'); + expect(result).toBeNull(); + }); + + it('should throw if stored data is corrupted JSON', async () => { + mockKeytar.getPassword.mockResolvedValue('not-json'); + await expect(storage.getCredentials('test-server')).rejects.toThrow( + 'Failed to parse stored credentials for test-server', + ); + }); + }); + + describe('setCredentials', () => { + it('should save credentials to keychain', async () => { + vi.useFakeTimers(); + mockKeytar.setPassword.mockResolvedValue(undefined); + await storage.setCredentials(validCredentials); + expect(mockKeytar.setPassword).toHaveBeenCalledWith( + mockServiceName, + 'test-server', + JSON.stringify({ ...validCredentials, updatedAt: Date.now() }), + ); + }); + + it('should throw if saving to keychain fails', async () => { + mockKeytar.setPassword.mockRejectedValue( + new Error('keychain write error'), + ); + await expect(storage.setCredentials(validCredentials)).rejects.toThrow( + 'keychain write error', + ); + }); + }); + + describe('deleteCredentials', () => { + it('should delete credentials from keychain', async () => { + mockKeytar.deletePassword.mockResolvedValue(true); + await storage.deleteCredentials('test-server'); + expect(mockKeytar.deletePassword).toHaveBeenCalledWith( + mockServiceName, + 'test-server', + ); + }); + + it('should throw if no credentials were found to delete', async () => { + mockKeytar.deletePassword.mockResolvedValue(false); + await expect(storage.deleteCredentials('test-server')).rejects.toThrow( + 'No credentials found for test-server', + ); + }); + + it('should throw if deleting from keychain fails', async () => { + mockKeytar.deletePassword.mockRejectedValue( + new Error('keychain delete error'), + ); + await expect(storage.deleteCredentials('test-server')).rejects.toThrow( + 'keychain delete error', + ); + }); + }); + + describe('listServers', () => { + it('should return a list of server names', async () => { + mockKeytar.findCredentials.mockResolvedValue([ + { account: 'server1', password: '' }, + { account: 'server2', password: '' }, + ]); + const result = await storage.listServers(); + expect(result).toEqual(['server1', 'server2']); + }); + + it('should not include internal test keys in the server list', async () => { + mockKeytar.findCredentials.mockResolvedValue([ + { account: 'server1', password: '' }, + { + account: `__keychain_test__${mockCryptoRandomBytesString}`, + password: '', + }, + { account: 'server2', password: '' }, + ]); + const result = await storage.listServers(); + expect(result).toEqual(['server1', 'server2']); + }); + + it('should return an empty array on error', async () => { + mockKeytar.findCredentials.mockRejectedValue(new Error('find error')); + const result = await storage.listServers(); + expect(result).toEqual([]); + }); + }); + + describe('getAllCredentials', () => { + it('should return a map of all valid credentials', async () => { + const creds2 = { + ...validCredentials, + serverName: 'server2', + }; + const expiredCreds = { + ...validCredentials, + serverName: 'expired-server', + token: { ...validCredentials.token, expiresAt: Date.now() - 1000 }, + }; + const structurallyInvalidCreds = { + serverName: 'invalid-server', + }; + + mockKeytar.findCredentials.mockResolvedValue([ + { + account: 'test-server', + password: JSON.stringify(validCredentials), + }, + { account: 'server2', password: JSON.stringify(creds2) }, + { + account: 'expired-server', + password: JSON.stringify(expiredCreds), + }, + { account: 'bad-server', password: 'not-json' }, + { + account: 'invalid-server', + password: JSON.stringify(structurallyInvalidCreds), + }, + ]); + + const result = await storage.getAllCredentials(); + expect(result.size).toBe(2); + expect(result.get('test-server')).toEqual(validCredentials); + expect(result.get('server2')).toEqual(creds2); + expect(result.has('expired-server')).toBe(false); + expect(result.has('bad-server')).toBe(false); + expect(result.has('invalid-server')).toBe(false); + }); + }); + + describe('clearAll', () => { + it('should delete all credentials for the service', async () => { + mockKeytar.findCredentials.mockResolvedValue([ + { account: 'server1', password: '' }, + { account: 'server2', password: '' }, + ]); + mockKeytar.deletePassword.mockResolvedValue(true); + + await storage.clearAll(); + + expect(mockKeytar.deletePassword).toHaveBeenCalledTimes(2); + expect(mockKeytar.deletePassword).toHaveBeenCalledWith( + mockServiceName, + 'server1', + ); + expect(mockKeytar.deletePassword).toHaveBeenCalledWith( + mockServiceName, + 'server2', + ); + }); + + it('should throw an aggregated error if deletions fail', async () => { + mockKeytar.findCredentials.mockResolvedValue([ + { account: 'server1', password: '' }, + { account: 'server2', password: '' }, + ]); + mockKeytar.deletePassword + .mockResolvedValueOnce(true) + .mockRejectedValueOnce(new Error('delete failed')); + + await expect(storage.clearAll()).rejects.toThrow( + 'Failed to clear some credentials: delete failed', + ); + }); + }); + }); +}); diff --git a/packages/core/src/mcp/token-storage/keychain-token-storage.ts b/packages/core/src/mcp/token-storage/keychain-token-storage.ts new file mode 100644 index 0000000000..70eccbadf5 --- /dev/null +++ b/packages/core/src/mcp/token-storage/keychain-token-storage.ts @@ -0,0 +1,251 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as crypto from 'node:crypto'; +import { BaseTokenStorage } from './base-token-storage.js'; +import type { OAuthCredentials } from './types.js'; + +interface Keytar { + getPassword(service: string, account: string): Promise; + setPassword( + service: string, + account: string, + password: string, + ): Promise; + deletePassword(service: string, account: string): Promise; + findCredentials( + service: string, + ): Promise>; +} + +const KEYCHAIN_TEST_PREFIX = '__keychain_test__'; + +export class KeychainTokenStorage extends BaseTokenStorage { + private keychainAvailable: boolean | null = null; + private keytarModule: Keytar | null = null; + private keytarLoadAttempted = false; + + async getKeytar(): Promise { + // If we've already tried loading (successfully or not), return the result + if (this.keytarLoadAttempted) { + return this.keytarModule; + } + + this.keytarLoadAttempted = true; + + try { + // Try to import keytar without any timeout - let the OS handle it + const moduleName = 'keytar'; + const module = await import(moduleName); + this.keytarModule = module.default || module; + } catch (error) { + console.error(error); + } + return this.keytarModule; + } + + async getCredentials(serverName: string): Promise { + if (!(await this.checkKeychainAvailability())) { + throw new Error('Keychain is not available'); + } + + const keytar = await this.getKeytar(); + if (!keytar) { + throw new Error('Keytar module not available'); + } + + try { + const sanitizedName = this.sanitizeServerName(serverName); + const data = await keytar.getPassword(this.serviceName, sanitizedName); + + if (!data) { + return null; + } + + const credentials = JSON.parse(data) as OAuthCredentials; + + if (this.isTokenExpired(credentials)) { + return null; + } + + return credentials; + } catch (error) { + if (error instanceof SyntaxError) { + throw new Error(`Failed to parse stored credentials for ${serverName}`); + } + throw error; + } + } + + async setCredentials(credentials: OAuthCredentials): Promise { + if (!(await this.checkKeychainAvailability())) { + throw new Error('Keychain is not available'); + } + + const keytar = await this.getKeytar(); + if (!keytar) { + throw new Error('Keytar module not available'); + } + + this.validateCredentials(credentials); + + const sanitizedName = this.sanitizeServerName(credentials.serverName); + const updatedCredentials: OAuthCredentials = { + ...credentials, + updatedAt: Date.now(), + }; + + const data = JSON.stringify(updatedCredentials); + await keytar.setPassword(this.serviceName, sanitizedName, data); + } + + async deleteCredentials(serverName: string): Promise { + if (!(await this.checkKeychainAvailability())) { + throw new Error('Keychain is not available'); + } + + const keytar = await this.getKeytar(); + if (!keytar) { + throw new Error('Keytar module not available'); + } + + const sanitizedName = this.sanitizeServerName(serverName); + const deleted = await keytar.deletePassword( + this.serviceName, + sanitizedName, + ); + + if (!deleted) { + throw new Error(`No credentials found for ${serverName}`); + } + } + + async listServers(): Promise { + if (!(await this.checkKeychainAvailability())) { + throw new Error('Keychain is not available'); + } + + const keytar = await this.getKeytar(); + if (!keytar) { + throw new Error('Keytar module not available'); + } + + try { + const credentials = await keytar.findCredentials(this.serviceName); + return credentials + .filter((cred) => !cred.account.startsWith(KEYCHAIN_TEST_PREFIX)) + .map((cred: { account: string }) => cred.account); + } catch (error) { + console.error('Failed to list servers from keychain:', error); + return []; + } + } + + async getAllCredentials(): Promise> { + if (!(await this.checkKeychainAvailability())) { + throw new Error('Keychain is not available'); + } + + const keytar = await this.getKeytar(); + if (!keytar) { + throw new Error('Keytar module not available'); + } + + const result = new Map(); + try { + const credentials = ( + await keytar.findCredentials(this.serviceName) + ).filter((c) => !c.account.startsWith(KEYCHAIN_TEST_PREFIX)); + + for (const cred of credentials) { + try { + const data = JSON.parse(cred.password) as OAuthCredentials; + if (!this.isTokenExpired(data)) { + result.set(cred.account, data); + } + } catch (error) { + console.error( + `Failed to parse credentials for ${cred.account}:`, + error, + ); + } + } + } catch (error) { + console.error('Failed to get all credentials from keychain:', error); + } + + return result; + } + + async clearAll(): Promise { + if (!(await this.checkKeychainAvailability())) { + throw new Error('Keychain is not available'); + } + + const servers = this.keytarModule + ? await this.keytarModule + .findCredentials(this.serviceName) + .then((creds) => creds.map((c) => c.account)) + .catch((error: Error) => { + throw new Error( + `Failed to list servers for clearing: ${error.message}`, + ); + }) + : []; + const errors: Error[] = []; + + for (const server of servers) { + try { + await this.deleteCredentials(server); + } catch (error) { + errors.push(error as Error); + } + } + + if (errors.length > 0) { + throw new Error( + `Failed to clear some credentials: ${errors.map((e) => e.message).join(', ')}`, + ); + } + } + + // Checks whether or not a set-get-delete cycle with the keychain works. + // Returns false if any operation fails. + async checkKeychainAvailability(): Promise { + if (this.keychainAvailable !== null) { + return this.keychainAvailable; + } + + try { + const keytar = await this.getKeytar(); + if (!keytar) { + this.keychainAvailable = false; + return false; + } + + const testAccount = `${KEYCHAIN_TEST_PREFIX}${crypto.randomBytes(8).toString('hex')}`; + const testPassword = 'test'; + + await keytar.setPassword(this.serviceName, testAccount, testPassword); + const retrieved = await keytar.getPassword(this.serviceName, testAccount); + const deleted = await keytar.deletePassword( + this.serviceName, + testAccount, + ); + + const success = deleted && retrieved === testPassword; + this.keychainAvailable = success; + return success; + } catch (_error) { + this.keychainAvailable = false; + return false; + } + } + + async isAvailable(): Promise { + return this.checkKeychainAvailability(); + } +} diff --git a/packages/core/src/mcp/token-storage/types.ts b/packages/core/src/mcp/token-storage/types.ts index 4b7bca66bf..1e95a975e0 100644 --- a/packages/core/src/mcp/token-storage/types.ts +++ b/packages/core/src/mcp/token-storage/types.ts @@ -35,3 +35,8 @@ export interface TokenStorage { getAllCredentials(): Promise>; clearAll(): Promise; } + +export enum TokenStorageType { + KEYCHAIN = 'keychain', + ENCRYPTED_FILE = 'encrypted_file', +}