feat(security) - Use hybrid token storage when flag is enabled (#8010)

Co-authored-by: Shi Shu <shii@google.com>

packages/core/src/mcp/oauth-token-storage.test.ts

@@ -4,13 +4,15 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import { promises as fs } from 'node:fs';
 import * as path from 'node:path';
 import { MCPOAuthTokenStorage } from './oauth-token-storage.js';
-import type { OAuthToken, OAuthCredentials } from './token-storage/types.js';
+import { FORCE_ENCRYPTED_FILE_ENV_VAR } from './token-storage/index.js';
+import type { OAuthCredentials, OAuthToken } from './token-storage/types.js';
+import { GEMINI_DIR } from '../utils/paths.js';
 
-// Mock file system operations
+// Mock dependencies
 vi.mock('node:fs', () => ({
   promises: {
     readFile: vi.fn(),
@@ -18,13 +20,33 @@ vi.mock('node:fs', () => ({
     mkdir: vi.fn(),
     unlink: vi.fn(),
   },
-  mkdirSync: vi.fn(),
 }));
 
-vi.mock('node:os', () => ({
+vi.mock('node:path', () => ({
-  homedir: vi.fn(() => '/mock/home'),
+  dirname: vi.fn(),
+  join: vi.fn(),
 }));
 
+vi.mock('../config/storage.js', () => ({
+  Storage: {
+    getMcpOAuthTokensPath: vi.fn(),
+  },
+}));
+
+const mockHybridTokenStorage = {
+  listServers: vi.fn(),
+  setCredentials: vi.fn(),
+  getCredentials: vi.fn(),
+  deleteCredentials: vi.fn(),
+  clearAll: vi.fn(),
+  getAllCredentials: vi.fn(),
+};
+vi.mock('./token-storage/hybrid-token-storage.js', () => ({
+  HybridTokenStorage: vi.fn(() => mockHybridTokenStorage),
+}));
+
+const ONE_HR_MS = 3600000;
+
 describe('MCPOAuthTokenStorage', () => {
   let tokenStorage: MCPOAuthTokenStorage;
 
@@ -33,7 +55,7 @@ describe('MCPOAuthTokenStorage', () => {
     refreshToken: 'refresh_token_456',
     tokenType: 'Bearer',
     scope: 'read write',
-    expiresAt: Date.now() + 3600000, // 1 hour from now
+    expiresAt: Date.now() + ONE_HR_MS,
   };
 
   const mockCredentials: OAuthCredentials = {
@@ -44,13 +66,17 @@ describe('MCPOAuthTokenStorage', () => {
     updatedAt: Date.now(),
   };
 
+  describe('with encrypted flag false', () => {
     beforeEach(() => {
+      vi.stubEnv(FORCE_ENCRYPTED_FILE_ENV_VAR, 'false');
       tokenStorage = new MCPOAuthTokenStorage();
+
       vi.clearAllMocks();
-    vi.spyOn(console, 'error').mockImplementation(() => {});
+      vi.spyOn(console, 'error');
     });
 
     afterEach(() => {
+      vi.unstubAllEnvs();
       vi.restoreAllMocks();
     });
 
@@ -73,7 +99,7 @@ describe('MCPOAuthTokenStorage', () => {
         expect(tokens.size).toBe(1);
         expect(tokens.get('test-server')).toEqual(mockCredentials);
         expect(fs.readFile).toHaveBeenCalledWith(
-        path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'),
+          path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'),
           'utf-8',
         );
       });
@@ -116,18 +142,18 @@ describe('MCPOAuthTokenStorage', () => {
         );
 
         expect(fs.mkdir).toHaveBeenCalledWith(
-        path.join('/mock/home', '.gemini'),
+          path.join('/mock/home', GEMINI_DIR),
           { recursive: true },
         );
         expect(fs.writeFile).toHaveBeenCalledWith(
-        path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'),
+          path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'),
           expect.stringContaining('test-server'),
           { mode: 0o600 },
         );
       });
 
       it('should update existing token for same server', async () => {
-      const existingCredentials = {
+        const existingCredentials: OAuthCredentials = {
           ...mockCredentials,
           serverName: 'existing-server',
         };
@@ -136,11 +162,16 @@ describe('MCPOAuthTokenStorage', () => {
         );
         vi.mocked(fs.writeFile).mockResolvedValue(undefined);
 
-      const newToken = { ...mockToken, accessToken: 'new_access_token' };
+        const newToken: OAuthToken = {
+          ...mockToken,
+          accessToken: 'new_access_token',
+        };
         await tokenStorage.saveToken('existing-server', newToken);
 
         const writeCall = vi.mocked(fs.writeFile).mock.calls[0];
-      const savedData = JSON.parse(writeCall[1] as string);
+        const savedData = JSON.parse(
+          writeCall[1] as string,
+        ) as OAuthCredentials[];
 
         expect(savedData).toHaveLength(1);
         expect(savedData[0].token.accessToken).toBe('new_access_token');
@@ -195,8 +226,14 @@ describe('MCPOAuthTokenStorage', () => {
 
     describe('deleteCredentials', () => {
       it('should remove token for specific server', async () => {
-      const credentials1 = { ...mockCredentials, serverName: 'server1' };
+        const credentials1: OAuthCredentials = {
-      const credentials2 = { ...mockCredentials, serverName: 'server2' };
+          ...mockCredentials,
+          serverName: 'server1',
+        };
+        const credentials2: OAuthCredentials = {
+          ...mockCredentials,
+          serverName: 'server2',
+        };
         vi.mocked(fs.readFile).mockResolvedValue(
           JSON.stringify([credentials1, credentials2]),
         );
@@ -220,7 +257,7 @@ describe('MCPOAuthTokenStorage', () => {
         await tokenStorage.deleteCredentials('test-server');
 
         expect(fs.unlink).toHaveBeenCalledWith(
-        path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'),
+          path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'),
         );
         expect(fs.writeFile).not.toHaveBeenCalled();
       });
@@ -252,7 +289,7 @@ describe('MCPOAuthTokenStorage', () => {
 
     describe('isTokenExpired', () => {
       it('should return false for token without expiry', () => {
-      const tokenWithoutExpiry = { ...mockToken };
+        const tokenWithoutExpiry: OAuthToken = { ...mockToken };
         delete tokenWithoutExpiry.expiresAt;
 
         const result = tokenStorage.isTokenExpired(tokenWithoutExpiry);
@@ -261,9 +298,9 @@ describe('MCPOAuthTokenStorage', () => {
       });
 
       it('should return false for valid token', () => {
-      const futureToken = {
+        const futureToken: OAuthToken = {
           ...mockToken,
-        expiresAt: Date.now() + 3600000, // 1 hour from now
+          expiresAt: Date.now() + ONE_HR_MS,
         };
 
         const result = tokenStorage.isTokenExpired(futureToken);
@@ -272,9 +309,9 @@ describe('MCPOAuthTokenStorage', () => {
       });
 
       it('should return true for expired token', () => {
-      const expiredToken = {
+        const expiredToken: OAuthToken = {
           ...mockToken,
-        expiresAt: Date.now() - 3600000, // 1 hour ago
+          expiresAt: Date.now() - ONE_HR_MS,
         };
 
         const result = tokenStorage.isTokenExpired(expiredToken);
@@ -283,7 +320,7 @@ describe('MCPOAuthTokenStorage', () => {
       });
 
       it('should return true for token expiring within buffer time', () => {
-      const soonToExpireToken = {
+        const soonToExpireToken: OAuthToken = {
           ...mockToken,
           expiresAt: Date.now() + 60000, // 1 minute from now (within 5-minute buffer)
         };
@@ -301,7 +338,7 @@ describe('MCPOAuthTokenStorage', () => {
         await tokenStorage.clearAll();
 
         expect(fs.unlink).toHaveBeenCalledWith(
-        path.join('/mock/home', '.gemini', 'mcp-oauth-tokens.json'),
+          path.join('/mock/home', GEMINI_DIR, 'mcp-oauth-tokens.json'),
         );
       });
 
@@ -323,4 +360,91 @@ describe('MCPOAuthTokenStorage', () => {
         );
       });
     });
+  });
+
+  describe('with encrypted flag true', () => {
+    beforeEach(() => {
+      vi.stubEnv(FORCE_ENCRYPTED_FILE_ENV_VAR, 'true');
+      tokenStorage = new MCPOAuthTokenStorage();
+
+      vi.clearAllMocks();
+      vi.spyOn(console, 'error');
+    });
+
+    afterEach(() => {
+      vi.unstubAllEnvs();
+      vi.restoreAllMocks();
+    });
+
+    it('should use HybridTokenStorage to list all credentials', async () => {
+      mockHybridTokenStorage.getAllCredentials.mockResolvedValue(new Map());
+      const servers = await tokenStorage.getAllCredentials();
+      expect(mockHybridTokenStorage.getAllCredentials).toHaveBeenCalled();
+      expect(servers).toEqual(new Map());
+    });
+
+    it('should use HybridTokenStorage to list servers', async () => {
+      mockHybridTokenStorage.listServers.mockResolvedValue(['server1']);
+      const servers = await tokenStorage.listServers();
+      expect(mockHybridTokenStorage.listServers).toHaveBeenCalled();
+      expect(servers).toEqual(['server1']);
+    });
+
+    it('should use HybridTokenStorage to set credentials', async () => {
+      await tokenStorage.setCredentials(mockCredentials);
+      expect(mockHybridTokenStorage.setCredentials).toHaveBeenCalledWith(
+        mockCredentials,
+      );
+    });
+
+    it('should use HybridTokenStorage to save a token', async () => {
+      const serverName = 'server1';
+      const now = Date.now();
+      vi.spyOn(Date, 'now').mockReturnValue(now);
+
+      await tokenStorage.saveToken(
+        serverName,
+        mockToken,
+        'clientId',
+        'tokenUrl',
+        'mcpUrl',
+      );
+
+      const expectedCredential: OAuthCredentials = {
+        serverName,
+        token: mockToken,
+        clientId: 'clientId',
+        tokenUrl: 'tokenUrl',
+        mcpServerUrl: 'mcpUrl',
+        updatedAt: now,
+      };
+
+      expect(mockHybridTokenStorage.setCredentials).toHaveBeenCalledWith(
+        expectedCredential,
+      );
+      expect(path.dirname).toHaveBeenCalled();
+      expect(fs.mkdir).toHaveBeenCalled();
+    });
+
+    it('should use HybridTokenStorage to get credentials', async () => {
+      mockHybridTokenStorage.getCredentials.mockResolvedValue(mockCredentials);
+      const result = await tokenStorage.getCredentials('server1');
+      expect(mockHybridTokenStorage.getCredentials).toHaveBeenCalledWith(
+        'server1',
+      );
+      expect(result).toBe(mockCredentials);
+    });
+
+    it('should use HybridTokenStorage to delete credentials', async () => {
+      await tokenStorage.deleteCredentials('server1');
+      expect(mockHybridTokenStorage.deleteCredentials).toHaveBeenCalledWith(
+        'server1',
+      );
+    });
+
+    it('should use HybridTokenStorage to clear all tokens', async () => {
+      await tokenStorage.clearAll();
+      expect(mockHybridTokenStorage.clearAll).toHaveBeenCalled();
+    });
+  });
 });

packages/core/src/mcp/oauth-token-storage.ts

@@ -13,11 +13,22 @@ import type {
   OAuthCredentials,
   TokenStorage,
 } from './token-storage/types.js';
+import { HybridTokenStorage } from './token-storage/hybrid-token-storage.js';
+import {
+  DEFAULT_SERVICE_NAME,
+  FORCE_ENCRYPTED_FILE_ENV_VAR,
+} from './token-storage/index.js';
 
 /**
  * Class for managing MCP OAuth token storage and retrieval.
  */
 export class MCPOAuthTokenStorage implements TokenStorage {
+  private readonly hybridTokenStorage = new HybridTokenStorage(
+    DEFAULT_SERVICE_NAME,
+  );
+  private readonly useEncryptedFile =
+    process.env[FORCE_ENCRYPTED_FILE_ENV_VAR] === 'true';
+
   /**
    * Get the path to the token storage file.
    *
@@ -41,6 +52,9 @@ export class MCPOAuthTokenStorage implements TokenStorage {
    * @returns A map of server names to credentials
    */
   async getAllCredentials(): Promise<Map<string, OAuthCredentials>> {
+    if (this.useEncryptedFile) {
+      return this.hybridTokenStorage.getAllCredentials();
+    }
     const tokenMap = new Map<string, OAuthCredentials>();
 
     try {
@@ -64,11 +78,17 @@ export class MCPOAuthTokenStorage implements TokenStorage {
   }
 
   async listServers(): Promise<string[]> {
+    if (this.useEncryptedFile) {
+      return this.hybridTokenStorage.listServers();
+    }
     const tokens = await this.getAllCredentials();
     return Array.from(tokens.keys());
   }
 
   async setCredentials(credentials: OAuthCredentials): Promise<void> {
+    if (this.useEncryptedFile) {
+      return this.hybridTokenStorage.setCredentials(credentials);
+    }
     const tokens = await this.getAllCredentials();
     tokens.set(credentials.serverName, credentials);
 
@@ -116,6 +136,9 @@ export class MCPOAuthTokenStorage implements TokenStorage {
       updatedAt: Date.now(),
     };
 
+    if (this.useEncryptedFile) {
+      return this.hybridTokenStorage.setCredentials(credential);
+    }
     await this.setCredentials(credential);
   }
 
@@ -126,6 +149,9 @@ export class MCPOAuthTokenStorage implements TokenStorage {
    * @returns The stored credentials or null if not found
    */
   async getCredentials(serverName: string): Promise<OAuthCredentials | null> {
+    if (this.useEncryptedFile) {
+      return this.hybridTokenStorage.getCredentials(serverName);
+    }
     const tokens = await this.getAllCredentials();
     return tokens.get(serverName) || null;
   }
@@ -136,6 +162,9 @@ export class MCPOAuthTokenStorage implements TokenStorage {
    * @param serverName The name of the MCP server
    */
   async deleteCredentials(serverName: string): Promise<void> {
+    if (this.useEncryptedFile) {
+      return this.hybridTokenStorage.deleteCredentials(serverName);
+    }
     const tokens = await this.getAllCredentials();
 
     if (tokens.delete(serverName)) {
@@ -179,6 +208,9 @@ export class MCPOAuthTokenStorage implements TokenStorage {
    * Clear all stored MCP OAuth tokens.
    */
   async clearAll(): Promise<void> {
+    if (this.useEncryptedFile) {
+      return this.hybridTokenStorage.clearAll();
+    }
     try {
       const tokenFile = this.getTokenFilePath();
       await fs.unlink(tokenFile);

packages/core/src/mcp/token-storage/index.ts

@@ -0,0 +1,14 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export * from './types.js';
+export * from './base-token-storage.js';
+export * from './file-token-storage.js';
+export * from './hybrid-token-storage.js';
+
+export const DEFAULT_SERVICE_NAME = 'gemini-cli-oauth';
+export const FORCE_ENCRYPTED_FILE_ENV_VAR =
+  'GEMINI_FORCE_ENCRYPTED_FILE_STORAGE';