From c999b7e35408e3c363057534163ab9753a97f9e8 Mon Sep 17 00:00:00 2001 From: shishu314 Date: Tue, 16 Sep 2025 10:05:29 -0400 Subject: [PATCH] feat(security) - Encrypted oauth flag (#8101) Co-authored-by: Shi Shu --- .../oauth-credential-storage.test.ts | 94 +- .../code_assist/oauth-credential-storage.ts | 16 +- packages/core/src/code_assist/oauth2.test.ts | 1801 +++++++++-------- packages/core/src/code_assist/oauth2.ts | 30 +- 4 files changed, 1100 insertions(+), 841 deletions(-) diff --git a/packages/core/src/code_assist/oauth-credential-storage.test.ts b/packages/core/src/code_assist/oauth-credential-storage.test.ts index 2927d31e75..c555b923e6 100644 --- a/packages/core/src/code_assist/oauth-credential-storage.test.ts +++ b/packages/core/src/code_assist/oauth-credential-storage.test.ts @@ -7,7 +7,6 @@ import { type Credentials } from 'google-auth-library'; import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { OAuthCredentialStorage } from './oauth-credential-storage.js'; -import { HybridTokenStorage } from '../mcp/token-storage/hybrid-token-storage.js'; import type { OAuthCredentials } from '../mcp/token-storage/types.js'; import * as path from 'node:path'; @@ -15,7 +14,14 @@ import * as os from 'node:os'; import { promises as fs } from 'node:fs'; // Mock external dependencies -vi.mock('../mcp/token-storage/hybrid-token-storage.js'); +const mockHybridTokenStorage = vi.hoisted(() => ({ + getCredentials: vi.fn(), + setCredentials: vi.fn(), + deleteCredentials: vi.fn(), +})); +vi.mock('../mcp/token-storage/hybrid-token-storage.js', () => ({ + HybridTokenStorage: vi.fn(() => mockHybridTokenStorage), +})); vi.mock('node:fs', () => ({ promises: { readFile: vi.fn(), @@ -26,9 +32,6 @@ vi.mock('node:os'); vi.mock('node:path'); describe('OAuthCredentialStorage', () => { - let storage: HybridTokenStorage; - let oauthStorage: OAuthCredentialStorage; - const mockCredentials: Credentials = { access_token: 'mock_access_token', refresh_token: 'mock_refresh_token', @@ -52,12 +55,13 @@ describe('OAuthCredentialStorage', () => { const oldFilePath = '/mock/home/.gemini/oauth.json'; beforeEach(() => { - storage = new HybridTokenStorage(''); - oauthStorage = new OAuthCredentialStorage(storage); - - vi.spyOn(storage, 'getCredentials').mockResolvedValue(null); - vi.spyOn(storage, 'setCredentials').mockResolvedValue(undefined); - vi.spyOn(storage, 'deleteCredentials').mockResolvedValue(undefined); + vi.spyOn(mockHybridTokenStorage, 'getCredentials').mockResolvedValue(null); + vi.spyOn(mockHybridTokenStorage, 'setCredentials').mockResolvedValue( + undefined, + ); + vi.spyOn(mockHybridTokenStorage, 'deleteCredentials').mockResolvedValue( + undefined, + ); vi.spyOn(fs, 'readFile').mockRejectedValue(new Error('File not found')); vi.spyOn(fs, 'rm').mockResolvedValue(undefined); @@ -72,25 +76,33 @@ describe('OAuthCredentialStorage', () => { describe('loadCredentials', () => { it('should load credentials from HybridTokenStorage if available', async () => { - vi.spyOn(storage, 'getCredentials').mockResolvedValue(mockMcpCredentials); + vi.spyOn(mockHybridTokenStorage, 'getCredentials').mockResolvedValue( + mockMcpCredentials, + ); - const result = await oauthStorage.loadCredentials(); + const result = await OAuthCredentialStorage.loadCredentials(); - expect(storage.getCredentials).toHaveBeenCalledWith('main-account'); + expect(mockHybridTokenStorage.getCredentials).toHaveBeenCalledWith( + 'main-account', + ); expect(result).toEqual(mockCredentials); }); it('should fallback to migrateFromFileStorage if no credentials in HybridTokenStorage', async () => { - vi.spyOn(storage, 'getCredentials').mockResolvedValue(null); + vi.spyOn(mockHybridTokenStorage, 'getCredentials').mockResolvedValue( + null, + ); vi.spyOn(fs, 'readFile').mockResolvedValue( JSON.stringify(mockCredentials), ); - const result = await oauthStorage.loadCredentials(); + const result = await OAuthCredentialStorage.loadCredentials(); - expect(storage.getCredentials).toHaveBeenCalledWith('main-account'); + expect(mockHybridTokenStorage.getCredentials).toHaveBeenCalledWith( + 'main-account', + ); expect(fs.readFile).toHaveBeenCalledWith(oldFilePath, 'utf-8'); - expect(storage.setCredentials).toHaveBeenCalled(); // Verify credentials were saved + expect(mockHybridTokenStorage.setCredentials).toHaveBeenCalled(); // Verify credentials were saved expect(fs.rm).toHaveBeenCalledWith(oldFilePath, { force: true }); // Verify old file was removed expect(result).toEqual(mockCredentials); }); @@ -101,41 +113,47 @@ describe('OAuthCredentialStorage', () => { code: 'ENOENT', }); - const result = await oauthStorage.loadCredentials(); + const result = await OAuthCredentialStorage.loadCredentials(); expect(result).toBeNull(); }); it('should throw an error if loading fails', async () => { - vi.spyOn(storage, 'getCredentials').mockRejectedValue( + vi.spyOn(mockHybridTokenStorage, 'getCredentials').mockRejectedValue( new Error('Loading error'), ); - await expect(oauthStorage.loadCredentials()).rejects.toThrow( + await expect(OAuthCredentialStorage.loadCredentials()).rejects.toThrow( 'Failed to load OAuth credentials', ); }); it('should throw an error if read file fails', async () => { - vi.spyOn(storage, 'getCredentials').mockResolvedValue(null); + vi.spyOn(mockHybridTokenStorage, 'getCredentials').mockResolvedValue( + null, + ); vi.spyOn(fs, 'readFile').mockRejectedValue( new Error('Permission denied'), ); - await expect(oauthStorage.loadCredentials()).rejects.toThrow( + await expect(OAuthCredentialStorage.loadCredentials()).rejects.toThrow( 'Failed to load OAuth credentials', ); }); it('should not throw error if migration file removal failed', async () => { - vi.spyOn(storage, 'getCredentials').mockResolvedValue(null); + vi.spyOn(mockHybridTokenStorage, 'getCredentials').mockResolvedValue( + null, + ); vi.spyOn(fs, 'readFile').mockResolvedValue( JSON.stringify(mockCredentials), ); - vi.spyOn(oauthStorage, 'saveCredentials').mockResolvedValue(undefined); + vi.spyOn(OAuthCredentialStorage, 'saveCredentials').mockResolvedValue( + undefined, + ); vi.spyOn(fs, 'rm').mockRejectedValue(new Error('Deletion failed')); - const result = await oauthStorage.loadCredentials(); + const result = await OAuthCredentialStorage.loadCredentials(); expect(result).toEqual(mockCredentials); }); @@ -143,9 +161,11 @@ describe('OAuthCredentialStorage', () => { describe('saveCredentials', () => { it('should save credentials to HybridTokenStorage', async () => { - await oauthStorage.saveCredentials(mockCredentials); + await OAuthCredentialStorage.saveCredentials(mockCredentials); - expect(storage.setCredentials).toHaveBeenCalledWith(mockMcpCredentials); + expect(mockHybridTokenStorage.setCredentials).toHaveBeenCalledWith( + mockMcpCredentials, + ); }); it('should throw an error if access_token is missing', async () => { @@ -154,7 +174,7 @@ describe('OAuthCredentialStorage', () => { access_token: undefined, }; await expect( - oauthStorage.saveCredentials(invalidCredentials), + OAuthCredentialStorage.saveCredentials(invalidCredentials), ).rejects.toThrow( 'Attempted to save credentials without an access token.', ); @@ -163,13 +183,15 @@ describe('OAuthCredentialStorage', () => { describe('clearCredentials', () => { it('should delete credentials from HybridTokenStorage', async () => { - await oauthStorage.clearCredentials(); + await OAuthCredentialStorage.clearCredentials(); - expect(storage.deleteCredentials).toHaveBeenCalledWith('main-account'); + expect(mockHybridTokenStorage.deleteCredentials).toHaveBeenCalledWith( + 'main-account', + ); }); it('should attempt to remove the old file-based storage', async () => { - await oauthStorage.clearCredentials(); + await OAuthCredentialStorage.clearCredentials(); expect(fs.rm).toHaveBeenCalledWith(oldFilePath, { force: true }); }); @@ -177,15 +199,17 @@ describe('OAuthCredentialStorage', () => { it('should not throw an error if deleting old file fails', async () => { vi.spyOn(fs, 'rm').mockRejectedValue(new Error('File deletion failed')); - await expect(oauthStorage.clearCredentials()).resolves.toBeUndefined(); + await expect( + OAuthCredentialStorage.clearCredentials(), + ).resolves.toBeUndefined(); }); it('should throw an error if clearing from HybridTokenStorage fails', async () => { - vi.spyOn(storage, 'deleteCredentials').mockRejectedValue( + vi.spyOn(mockHybridTokenStorage, 'deleteCredentials').mockRejectedValue( new Error('Deletion error'), ); - await expect(oauthStorage.clearCredentials()).rejects.toThrow( + await expect(OAuthCredentialStorage.clearCredentials()).rejects.toThrow( 'Failed to clear OAuth credentials', ); }); diff --git a/packages/core/src/code_assist/oauth-credential-storage.ts b/packages/core/src/code_assist/oauth-credential-storage.ts index 9c6f085f3f..30c940628b 100644 --- a/packages/core/src/code_assist/oauth-credential-storage.ts +++ b/packages/core/src/code_assist/oauth-credential-storage.ts @@ -17,16 +17,14 @@ const KEYCHAIN_SERVICE_NAME = 'gemini-cli-oauth'; const MAIN_ACCOUNT_KEY = 'main-account'; export class OAuthCredentialStorage { - constructor( - private readonly storage: HybridTokenStorage = new HybridTokenStorage( - KEYCHAIN_SERVICE_NAME, - ), - ) {} + private static storage: HybridTokenStorage = new HybridTokenStorage( + KEYCHAIN_SERVICE_NAME, + ); /** * Load cached OAuth credentials */ - async loadCredentials(): Promise { + static async loadCredentials(): Promise { try { const credentials = await this.storage.getCredentials(MAIN_ACCOUNT_KEY); @@ -59,7 +57,7 @@ export class OAuthCredentialStorage { /** * Save OAuth credentials */ - async saveCredentials(credentials: Credentials): Promise { + static async saveCredentials(credentials: Credentials): Promise { if (!credentials.access_token) { throw new Error('Attempted to save credentials without an access token.'); } @@ -83,7 +81,7 @@ export class OAuthCredentialStorage { /** * Clear cached OAuth credentials */ - async clearCredentials(): Promise { + static async clearCredentials(): Promise { try { await this.storage.deleteCredentials(MAIN_ACCOUNT_KEY); @@ -99,7 +97,7 @@ export class OAuthCredentialStorage { /** * Migrate credentials from old file-based storage to keychain */ - private async migrateFromFileStorage(): Promise { + private static async migrateFromFileStorage(): Promise { const oldFilePath = path.join(os.homedir(), GEMINI_DIR, OAUTH_FILE); let credsJson: string; diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts index 01ec0fe6b9..50c0f5e351 100644 --- a/packages/core/src/code_assist/oauth2.test.ts +++ b/packages/core/src/code_assist/oauth2.test.ts @@ -4,6 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +import type { Credentials } from 'google-auth-library'; import type { Mock } from 'vitest'; import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { @@ -23,6 +24,7 @@ import * as os from 'node:os'; import { AuthType } from '../core/contentGenerator.js'; import type { Config } from '../config/config.js'; import readline from 'node:readline'; +import { FORCE_ENCRYPTED_FILE_ENV_VAR } from '../mcp/token-storage/index.js'; vi.mock('os', async (importOriginal) => { const os = await importOriginal(); @@ -41,6 +43,14 @@ vi.mock('../utils/browser.js', () => ({ shouldAttemptBrowserLaunch: () => true, })); +vi.mock('./oauth-credential-storage.js', () => ({ + OAuthCredentialStorage: { + saveCredentials: vi.fn(), + loadCredentials: vi.fn(), + clearCredentials: vi.fn(), + }, +})); + const mockConfig = { getNoBrowser: () => false, getProxy: () => 'http://test.proxy.com:8080', @@ -51,700 +61,24 @@ const mockConfig = { global.fetch = vi.fn(); describe('oauth2', () => { - let tempHomeDir: string; - - beforeEach(() => { - tempHomeDir = fs.mkdtempSync( - path.join(os.tmpdir(), 'gemini-cli-test-home-'), - ); - (os.homedir as Mock).mockReturnValue(tempHomeDir); - }); - afterEach(() => { - fs.rmSync(tempHomeDir, { recursive: true, force: true }); - vi.clearAllMocks(); - resetOauthClientForTesting(); - vi.unstubAllEnvs(); - }); - - it('should perform a web login', async () => { - const mockAuthUrl = 'https://example.com/auth'; - const mockCode = 'test-code'; - const mockState = 'test-state'; - const mockTokens = { - access_token: 'test-access-token', - refresh_token: 'test-refresh-token', - }; - - const mockGenerateAuthUrl = vi.fn().mockReturnValue(mockAuthUrl); - const mockGetToken = vi.fn().mockResolvedValue({ tokens: mockTokens }); - const mockSetCredentials = vi.fn(); - const mockGetAccessToken = vi - .fn() - .mockResolvedValue({ token: 'mock-access-token' }); - const mockOAuth2Client = { - generateAuthUrl: mockGenerateAuthUrl, - getToken: mockGetToken, - setCredentials: mockSetCredentials, - getAccessToken: mockGetAccessToken, - credentials: mockTokens, - on: vi.fn(), - } as unknown as OAuth2Client; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockOAuth2Client, - ); - - vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never); - (open as Mock).mockImplementation(async () => ({ on: vi.fn() }) as never); - - // Mock the UserInfo API response - (global.fetch as Mock).mockResolvedValue({ - ok: true, - json: vi - .fn() - .mockResolvedValue({ email: 'test-google-account@gmail.com' }), - } as unknown as Response); - - let requestCallback!: http.RequestListener< - typeof http.IncomingMessage, - typeof http.ServerResponse - >; - - let serverListeningCallback: (value: unknown) => void; - const serverListeningPromise = new Promise( - (resolve) => (serverListeningCallback = resolve), - ); - - let capturedPort = 0; - const mockHttpServer = { - listen: vi.fn((port: number, _host: string, callback?: () => void) => { - capturedPort = port; - if (callback) { - callback(); - } - serverListeningCallback(undefined); - }), - close: vi.fn((callback?: () => void) => { - if (callback) { - callback(); - } - }), - on: vi.fn(), - address: () => ({ port: capturedPort }), - }; - (http.createServer as Mock).mockImplementation((cb) => { - requestCallback = cb as http.RequestListener< - typeof http.IncomingMessage, - typeof http.ServerResponse - >; - return mockHttpServer as unknown as http.Server; - }); - - const clientPromise = getOauthClient( - AuthType.LOGIN_WITH_GOOGLE, - mockConfig, - ); - - // wait for server to start listening. - await serverListeningPromise; - - const mockReq = { - url: `/oauth2callback?code=${mockCode}&state=${mockState}`, - } as http.IncomingMessage; - const mockRes = { - writeHead: vi.fn(), - end: vi.fn(), - } as unknown as http.ServerResponse; - - await requestCallback(mockReq, mockRes); - - const client = await clientPromise; - expect(client).toBe(mockOAuth2Client); - - expect(open).toHaveBeenCalledWith(mockAuthUrl); - expect(mockGetToken).toHaveBeenCalledWith({ - code: mockCode, - redirect_uri: `http://localhost:${capturedPort}/oauth2callback`, - }); - expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens); - - // Verify Google Account was cached - const googleAccountPath = path.join( - tempHomeDir, - '.gemini', - 'google_accounts.json', - ); - expect(fs.existsSync(googleAccountPath)).toBe(true); - const cachedGoogleAccount = fs.readFileSync(googleAccountPath, 'utf-8'); - expect(JSON.parse(cachedGoogleAccount)).toEqual({ - active: 'test-google-account@gmail.com', - old: [], - }); - - // Verify the getCachedGoogleAccount function works - const userAccountManager = new UserAccountManager(); - expect(userAccountManager.getCachedGoogleAccount()).toBe( - 'test-google-account@gmail.com', - ); - }); - - it('should perform login with user code', async () => { - const mockConfigWithNoBrowser = { - getNoBrowser: () => true, - getProxy: () => 'http://test.proxy.com:8080', - isBrowserLaunchSuppressed: () => true, - } as unknown as Config; - - const mockCodeVerifier = { - codeChallenge: 'test-challenge', - codeVerifier: 'test-verifier', - }; - const mockAuthUrl = 'https://example.com/auth-user-code'; - const mockCode = 'test-user-code'; - const mockTokens = { - access_token: 'test-access-token-user-code', - refresh_token: 'test-refresh-token-user-code', - }; - - const mockGenerateAuthUrl = vi.fn().mockReturnValue(mockAuthUrl); - const mockGetToken = vi.fn().mockResolvedValue({ tokens: mockTokens }); - const mockSetCredentials = vi.fn(); - const mockGenerateCodeVerifierAsync = vi - .fn() - .mockResolvedValue(mockCodeVerifier); - - const mockOAuth2Client = { - generateAuthUrl: mockGenerateAuthUrl, - getToken: mockGetToken, - setCredentials: mockSetCredentials, - generateCodeVerifierAsync: mockGenerateCodeVerifierAsync, - on: vi.fn(), - } as unknown as OAuth2Client; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockOAuth2Client, - ); - - const mockReadline = { - question: vi.fn((_query, callback) => callback(mockCode)), - close: vi.fn(), - }; - (readline.createInterface as Mock).mockReturnValue(mockReadline); - - const consoleLogSpy = vi.spyOn(console, 'log').mockImplementation(() => {}); - - const client = await getOauthClient( - AuthType.LOGIN_WITH_GOOGLE, - mockConfigWithNoBrowser, - ); - - expect(client).toBe(mockOAuth2Client); - - // Verify the auth flow - expect(mockGenerateCodeVerifierAsync).toHaveBeenCalled(); - expect(mockGenerateAuthUrl).toHaveBeenCalled(); - expect(consoleLogSpy).toHaveBeenCalledWith( - expect.stringContaining(mockAuthUrl), - ); - expect(mockReadline.question).toHaveBeenCalledWith( - 'Enter the authorization code: ', - expect.any(Function), - ); - expect(mockGetToken).toHaveBeenCalledWith({ - code: mockCode, - codeVerifier: mockCodeVerifier.codeVerifier, - redirect_uri: 'https://codeassist.google.com/authcode', - }); - expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens); - - consoleLogSpy.mockRestore(); - }); - - describe('in Cloud Shell', () => { - const mockGetAccessToken = vi.fn(); - let mockComputeClient: Compute; + describe('with encrypted flag false', () => { + let tempHomeDir: string; beforeEach(() => { - mockGetAccessToken.mockResolvedValue({ token: 'test-access-token' }); - mockComputeClient = { - credentials: { refresh_token: 'test-refresh-token' }, - getAccessToken: mockGetAccessToken, - } as unknown as Compute; - - (Compute as unknown as Mock).mockImplementation(() => mockComputeClient); + process.env[FORCE_ENCRYPTED_FILE_ENV_VAR] = 'false'; + tempHomeDir = fs.mkdtempSync( + path.join(os.tmpdir(), 'gemini-cli-test-home-'), + ); + (os.homedir as Mock).mockReturnValue(tempHomeDir); + }); + afterEach(() => { + fs.rmSync(tempHomeDir, { recursive: true, force: true }); + vi.clearAllMocks(); + resetOauthClientForTesting(); + vi.unstubAllEnvs(); }); - it('should attempt to load cached credentials first', async () => { - const cachedCreds = { refresh_token: 'cached-token' }; - const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); - await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); - await fs.promises.writeFile(credsPath, JSON.stringify(cachedCreds)); - - const mockClient = { - setCredentials: vi.fn(), - getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }), - getTokenInfo: vi.fn().mockResolvedValue({}), - on: vi.fn(), - }; - - // To mock the new OAuth2Client() inside the function - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockClient as unknown as OAuth2Client, - ); - - await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); - - expect(mockClient.setCredentials).toHaveBeenCalledWith(cachedCreds); - expect(mockClient.getAccessToken).toHaveBeenCalled(); - expect(mockClient.getTokenInfo).toHaveBeenCalled(); - expect(Compute).not.toHaveBeenCalled(); // Should not fetch new client if cache is valid - }); - - it('should use Compute to get a client if no cached credentials exist', async () => { - await getOauthClient(AuthType.CLOUD_SHELL, mockConfig); - - expect(Compute).toHaveBeenCalledWith({}); - expect(mockGetAccessToken).toHaveBeenCalled(); - }); - - it('should not cache the credentials after fetching them via ADC', async () => { - const newCredentials = { refresh_token: 'new-adc-token' }; - mockComputeClient.credentials = newCredentials; - mockGetAccessToken.mockResolvedValue({ token: 'new-adc-token' }); - - await getOauthClient(AuthType.CLOUD_SHELL, mockConfig); - - const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); - expect(fs.existsSync(credsPath)).toBe(false); - }); - - it('should return the Compute client on successful ADC authentication', async () => { - const client = await getOauthClient(AuthType.CLOUD_SHELL, mockConfig); - expect(client).toBe(mockComputeClient); - }); - - it('should throw an error if ADC fails', async () => { - const testError = new Error('ADC Failed'); - mockGetAccessToken.mockRejectedValue(testError); - - await expect( - getOauthClient(AuthType.CLOUD_SHELL, mockConfig), - ).rejects.toThrow( - 'Could not authenticate using Cloud Shell credentials. Please select a different authentication method or ensure you are in a properly configured environment. Error: ADC Failed', - ); - }); - }); - - describe('credential loading order', () => { - it('should prioritize default cached credentials over GOOGLE_APPLICATION_CREDENTIALS', async () => { - // Setup default cached credentials - const defaultCreds = { refresh_token: 'default-cached-token' }; - const defaultCredsPath = path.join( - tempHomeDir, - '.gemini', - 'oauth_creds.json', - ); - await fs.promises.mkdir(path.dirname(defaultCredsPath), { - recursive: true, - }); - await fs.promises.writeFile( - defaultCredsPath, - JSON.stringify(defaultCreds), - ); - - // Setup credentials via environment variable - const envCreds = { refresh_token: 'env-var-token' }; - const envCredsPath = path.join(tempHomeDir, 'env_creds.json'); - await fs.promises.writeFile(envCredsPath, JSON.stringify(envCreds)); - vi.stubEnv('GOOGLE_APPLICATION_CREDENTIALS', envCredsPath); - - const mockClient = { - setCredentials: vi.fn(), - getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }), - getTokenInfo: vi.fn().mockResolvedValue({}), - on: vi.fn(), - }; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockClient as unknown as OAuth2Client, - ); - - await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); - - // Assert the correct credentials were used - expect(mockClient.setCredentials).toHaveBeenCalledWith(defaultCreds); - expect(mockClient.setCredentials).not.toHaveBeenCalledWith(envCreds); - }); - - it('should fall back to GOOGLE_APPLICATION_CREDENTIALS if default cache is missing', async () => { - // Setup credentials via environment variable - const envCreds = { refresh_token: 'env-var-token' }; - const envCredsPath = path.join(tempHomeDir, 'env_creds.json'); - await fs.promises.writeFile(envCredsPath, JSON.stringify(envCreds)); - vi.stubEnv('GOOGLE_APPLICATION_CREDENTIALS', envCredsPath); - - const mockClient = { - setCredentials: vi.fn(), - getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }), - getTokenInfo: vi.fn().mockResolvedValue({}), - on: vi.fn(), - }; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockClient as unknown as OAuth2Client, - ); - - await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); - - // Assert the correct credentials were used - expect(mockClient.setCredentials).toHaveBeenCalledWith(envCreds); - }); - }); - - describe('with GCP environment variables', () => { - it('should use GOOGLE_CLOUD_ACCESS_TOKEN when GOOGLE_GENAI_USE_GCA is true', async () => { - vi.stubEnv('GOOGLE_GENAI_USE_GCA', 'true'); - vi.stubEnv('GOOGLE_CLOUD_ACCESS_TOKEN', 'gcp-access-token'); - - const mockSetCredentials = vi.fn(); - const mockGetAccessToken = vi - .fn() - .mockResolvedValue({ token: 'gcp-access-token' }); - const mockOAuth2Client = { - setCredentials: mockSetCredentials, - getAccessToken: mockGetAccessToken, - on: vi.fn(), - } as unknown as OAuth2Client; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockOAuth2Client, - ); - - // Mock the UserInfo API response for fetchAndCacheUserInfo - (global.fetch as Mock).mockResolvedValue({ - ok: true, - json: vi - .fn() - .mockResolvedValue({ email: 'test-gcp-account@gmail.com' }), - } as unknown as Response); - - const client = await getOauthClient( - AuthType.LOGIN_WITH_GOOGLE, - mockConfig, - ); - - expect(client).toBe(mockOAuth2Client); - expect(mockSetCredentials).toHaveBeenCalledWith({ - access_token: 'gcp-access-token', - }); - - // Verify fetchAndCacheUserInfo was effectively called - expect(mockGetAccessToken).toHaveBeenCalled(); - expect(global.fetch).toHaveBeenCalledWith( - 'https://www.googleapis.com/oauth2/v2/userinfo', - { - headers: { - Authorization: 'Bearer gcp-access-token', - }, - }, - ); - - // Verify Google Account was cached - const googleAccountPath = path.join( - tempHomeDir, - '.gemini', - 'google_accounts.json', - ); - const cachedContent = fs.readFileSync(googleAccountPath, 'utf-8'); - expect(JSON.parse(cachedContent)).toEqual({ - active: 'test-gcp-account@gmail.com', - old: [], - }); - }); - - it('should not use GCP token if GOOGLE_CLOUD_ACCESS_TOKEN is not set', async () => { - vi.stubEnv('GOOGLE_GENAI_USE_GCA', 'true'); - - const mockSetCredentials = vi.fn(); - const mockGetAccessToken = vi - .fn() - .mockResolvedValue({ token: 'cached-access-token' }); - const mockGetTokenInfo = vi.fn().mockResolvedValue({}); - const mockOAuth2Client = { - setCredentials: mockSetCredentials, - getAccessToken: mockGetAccessToken, - getTokenInfo: mockGetTokenInfo, - on: vi.fn(), - } as unknown as OAuth2Client; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockOAuth2Client, - ); - - // Make it fall through to cached credentials path - const cachedCreds = { refresh_token: 'cached-token' }; - const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); - await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); - await fs.promises.writeFile(credsPath, JSON.stringify(cachedCreds)); - - await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); - - // It should be called with the cached credentials, not the GCP access token. - expect(mockSetCredentials).toHaveBeenCalledTimes(1); - expect(mockSetCredentials).toHaveBeenCalledWith(cachedCreds); - }); - - it('should not use GCP token if GOOGLE_GENAI_USE_GCA is not set', async () => { - vi.stubEnv('GOOGLE_CLOUD_ACCESS_TOKEN', 'gcp-access-token'); - - const mockSetCredentials = vi.fn(); - const mockGetAccessToken = vi - .fn() - .mockResolvedValue({ token: 'cached-access-token' }); - const mockGetTokenInfo = vi.fn().mockResolvedValue({}); - const mockOAuth2Client = { - setCredentials: mockSetCredentials, - getAccessToken: mockGetAccessToken, - getTokenInfo: mockGetTokenInfo, - on: vi.fn(), - } as unknown as OAuth2Client; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockOAuth2Client, - ); - - // Make it fall through to cached credentials path - const cachedCreds = { refresh_token: 'cached-token' }; - const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); - await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); - await fs.promises.writeFile(credsPath, JSON.stringify(cachedCreds)); - - await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); - - // It should be called with the cached credentials, not the GCP access token. - expect(mockSetCredentials).toHaveBeenCalledTimes(1); - expect(mockSetCredentials).toHaveBeenCalledWith(cachedCreds); - }); - }); - - describe('error handling', () => { - it('should handle browser launch failure with FatalAuthenticationError', async () => { - const mockError = new Error('Browser launch failed'); - (open as Mock).mockRejectedValue(mockError); - - const mockOAuth2Client = { - generateAuthUrl: vi.fn().mockReturnValue('https://example.com/auth'), - on: vi.fn(), - } as unknown as OAuth2Client; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockOAuth2Client, - ); - - await expect( - getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig), - ).rejects.toThrow('Failed to open browser: Browser launch failed'); - }); - - it('should handle authentication timeout with proper error message', async () => { - const mockAuthUrl = 'https://example.com/auth'; - const mockOAuth2Client = { - generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), - on: vi.fn(), - } as unknown as OAuth2Client; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockOAuth2Client, - ); - - (open as Mock).mockImplementation(async () => ({ on: vi.fn() }) as never); - - const mockHttpServer = { - listen: vi.fn(), - close: vi.fn(), - on: vi.fn(), - address: () => ({ port: 3000 }), - }; - (http.createServer as Mock).mockImplementation( - () => mockHttpServer as unknown as http.Server, - ); - - // Mock setTimeout to trigger timeout immediately - const originalSetTimeout = global.setTimeout; - global.setTimeout = vi.fn( - (callback) => (callback(), {} as unknown as NodeJS.Timeout), - ) as unknown as typeof setTimeout; - - await expect( - getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig), - ).rejects.toThrow( - 'Authentication timed out after 5 minutes. The browser tab may have gotten stuck in a loading state. Please try again or use NO_BROWSER=true for manual authentication.', - ); - - global.setTimeout = originalSetTimeout; - }); - - it('should handle OAuth callback errors with descriptive messages', async () => { - const mockAuthUrl = 'https://example.com/auth'; - const mockOAuth2Client = { - generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), - on: vi.fn(), - } as unknown as OAuth2Client; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockOAuth2Client, - ); - - (open as Mock).mockImplementation(async () => ({ on: vi.fn() }) as never); - - let requestCallback!: http.RequestListener; - let serverListeningCallback: (value: unknown) => void; - const serverListeningPromise = new Promise( - (resolve) => (serverListeningCallback = resolve), - ); - - const mockHttpServer = { - listen: vi.fn((_port: number, _host: string, callback?: () => void) => { - if (callback) callback(); - serverListeningCallback(undefined); - }), - close: vi.fn(), - on: vi.fn(), - address: () => ({ port: 3000 }), - }; - (http.createServer as Mock).mockImplementation((cb) => { - requestCallback = cb; - return mockHttpServer as unknown as http.Server; - }); - - const clientPromise = getOauthClient( - AuthType.LOGIN_WITH_GOOGLE, - mockConfig, - ); - await serverListeningPromise; - - // Test OAuth error with description - const mockReq = { - url: '/oauth2callback?error=access_denied&error_description=User+denied+access', - } as http.IncomingMessage; - const mockRes = { - writeHead: vi.fn(), - end: vi.fn(), - } as unknown as http.ServerResponse; - - await expect(async () => { - await requestCallback(mockReq, mockRes); - await clientPromise; - }).rejects.toThrow( - 'Google OAuth error: access_denied. User denied access', - ); - }); - - it('should handle OAuth error without description', async () => { - const mockAuthUrl = 'https://example.com/auth'; - const mockOAuth2Client = { - generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), - on: vi.fn(), - } as unknown as OAuth2Client; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockOAuth2Client, - ); - - (open as Mock).mockImplementation(async () => ({ on: vi.fn() }) as never); - - let requestCallback!: http.RequestListener; - let serverListeningCallback: (value: unknown) => void; - const serverListeningPromise = new Promise( - (resolve) => (serverListeningCallback = resolve), - ); - - const mockHttpServer = { - listen: vi.fn((_port: number, _host: string, callback?: () => void) => { - if (callback) callback(); - serverListeningCallback(undefined); - }), - close: vi.fn(), - on: vi.fn(), - address: () => ({ port: 3000 }), - }; - (http.createServer as Mock).mockImplementation((cb) => { - requestCallback = cb; - return mockHttpServer as unknown as http.Server; - }); - - const clientPromise = getOauthClient( - AuthType.LOGIN_WITH_GOOGLE, - mockConfig, - ); - await serverListeningPromise; - - // Test OAuth error without description - const mockReq = { - url: '/oauth2callback?error=server_error', - } as http.IncomingMessage; - const mockRes = { - writeHead: vi.fn(), - end: vi.fn(), - } as unknown as http.ServerResponse; - - await expect(async () => { - await requestCallback(mockReq, mockRes); - await clientPromise; - }).rejects.toThrow( - 'Google OAuth error: server_error. No additional details provided', - ); - }); - - it('should handle token exchange failure with descriptive error', async () => { - const mockAuthUrl = 'https://example.com/auth'; - const mockCode = 'test-code'; - const mockState = 'test-state'; - - const mockOAuth2Client = { - generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), - getToken: vi.fn().mockRejectedValue(new Error('Token exchange failed')), - on: vi.fn(), - } as unknown as OAuth2Client; - (OAuth2Client as unknown as Mock).mockImplementation( - () => mockOAuth2Client, - ); - - vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never); - (open as Mock).mockImplementation(async () => ({ on: vi.fn() }) as never); - - let requestCallback!: http.RequestListener; - let serverListeningCallback: (value: unknown) => void; - const serverListeningPromise = new Promise( - (resolve) => (serverListeningCallback = resolve), - ); - - const mockHttpServer = { - listen: vi.fn((_port: number, _host: string, callback?: () => void) => { - if (callback) callback(); - serverListeningCallback(undefined); - }), - close: vi.fn(), - on: vi.fn(), - address: () => ({ port: 3000 }), - }; - (http.createServer as Mock).mockImplementation((cb) => { - requestCallback = cb; - return mockHttpServer as unknown as http.Server; - }); - - const clientPromise = getOauthClient( - AuthType.LOGIN_WITH_GOOGLE, - mockConfig, - ); - await serverListeningPromise; - - const mockReq = { - url: `/oauth2callback?code=${mockCode}&state=${mockState}`, - } as http.IncomingMessage; - const mockRes = { - writeHead: vi.fn(), - end: vi.fn(), - } as unknown as http.ServerResponse; - - await expect(async () => { - await requestCallback(mockReq, mockRes); - await clientPromise; - }).rejects.toThrow( - 'Failed to exchange authorization code for tokens: Token exchange failed', - ); - }); - - it('should handle fetchAndCacheUserInfo failure gracefully', async () => { + it('should perform a web login', async () => { const mockAuthUrl = 'https://example.com/auth'; const mockCode = 'test-code'; const mockState = 'test-state'; @@ -753,13 +87,18 @@ describe('oauth2', () => { refresh_token: 'test-refresh-token', }; + const mockGenerateAuthUrl = vi.fn().mockReturnValue(mockAuthUrl); + const mockGetToken = vi.fn().mockResolvedValue({ tokens: mockTokens }); + const mockSetCredentials = vi.fn(); + const mockGetAccessToken = vi + .fn() + .mockResolvedValue({ token: 'mock-access-token' }); const mockOAuth2Client = { - generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), - getToken: vi.fn().mockResolvedValue({ tokens: mockTokens }), - setCredentials: vi.fn(), - getAccessToken: vi - .fn() - .mockResolvedValue({ token: 'test-access-token' }), + generateAuthUrl: mockGenerateAuthUrl, + getToken: mockGetToken, + setCredentials: mockSetCredentials, + getAccessToken: mockGetAccessToken, + credentials: mockTokens, on: vi.fn(), } as unknown as OAuth2Client; (OAuth2Client as unknown as Mock).mockImplementation( @@ -769,34 +108,46 @@ describe('oauth2', () => { vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never); (open as Mock).mockImplementation(async () => ({ on: vi.fn() }) as never); - // Mock fetch to fail + // Mock the UserInfo API response (global.fetch as Mock).mockResolvedValue({ - ok: false, - status: 500, - statusText: 'Internal Server Error', + ok: true, + json: vi + .fn() + .mockResolvedValue({ email: 'test-google-account@gmail.com' }), } as unknown as Response); - const consoleErrorSpy = vi - .spyOn(console, 'error') - .mockImplementation(() => {}); + let requestCallback!: http.RequestListener< + typeof http.IncomingMessage, + typeof http.ServerResponse + >; - let requestCallback!: http.RequestListener; let serverListeningCallback: (value: unknown) => void; const serverListeningPromise = new Promise( (resolve) => (serverListeningCallback = resolve), ); + let capturedPort = 0; const mockHttpServer = { - listen: vi.fn((_port: number, _host: string, callback?: () => void) => { - if (callback) callback(); + listen: vi.fn((port: number, _host: string, callback?: () => void) => { + capturedPort = port; + if (callback) { + callback(); + } serverListeningCallback(undefined); }), - close: vi.fn(), + close: vi.fn((callback?: () => void) => { + if (callback) { + callback(); + } + }), on: vi.fn(), - address: () => ({ port: 3000 }), + address: () => ({ port: capturedPort }), }; (http.createServer as Mock).mockImplementation((cb) => { - requestCallback = cb; + requestCallback = cb as http.RequestListener< + typeof http.IncomingMessage, + typeof http.ServerResponse + >; return mockHttpServer as unknown as http.Server; }); @@ -804,6 +155,8 @@ describe('oauth2', () => { AuthType.LOGIN_WITH_GOOGLE, mockConfig, ); + + // wait for server to start listening. await serverListeningPromise; const mockReq = { @@ -815,35 +168,67 @@ describe('oauth2', () => { } as unknown as http.ServerResponse; await requestCallback(mockReq, mockRes); + const client = await clientPromise; - - // Authentication should succeed even if fetchAndCacheUserInfo fails expect(client).toBe(mockOAuth2Client); - expect(consoleErrorSpy).toHaveBeenCalledWith( - 'Failed to fetch user info:', - 500, - 'Internal Server Error', - ); - consoleErrorSpy.mockRestore(); + expect(open).toHaveBeenCalledWith(mockAuthUrl); + expect(mockGetToken).toHaveBeenCalledWith({ + code: mockCode, + redirect_uri: `http://localhost:${capturedPort}/oauth2callback`, + }); + expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens); + + // Verify Google Account was cached + const googleAccountPath = path.join( + tempHomeDir, + '.gemini', + 'google_accounts.json', + ); + expect(fs.existsSync(googleAccountPath)).toBe(true); + const cachedGoogleAccount = fs.readFileSync(googleAccountPath, 'utf-8'); + expect(JSON.parse(cachedGoogleAccount)).toEqual({ + active: 'test-google-account@gmail.com', + old: [], + }); + + // Verify the getCachedGoogleAccount function works + const userAccountManager = new UserAccountManager(); + expect(userAccountManager.getCachedGoogleAccount()).toBe( + 'test-google-account@gmail.com', + ); }); - it('should handle user code authentication failure with descriptive error', async () => { + it('should perform login with user code', async () => { const mockConfigWithNoBrowser = { getNoBrowser: () => true, getProxy: () => 'http://test.proxy.com:8080', isBrowserLaunchSuppressed: () => true, } as unknown as Config; + const mockCodeVerifier = { + codeChallenge: 'test-challenge', + codeVerifier: 'test-verifier', + }; + const mockAuthUrl = 'https://example.com/auth-user-code'; + const mockCode = 'test-user-code'; + const mockTokens = { + access_token: 'test-access-token-user-code', + refresh_token: 'test-refresh-token-user-code', + }; + + const mockGenerateAuthUrl = vi.fn().mockReturnValue(mockAuthUrl); + const mockGetToken = vi.fn().mockResolvedValue({ tokens: mockTokens }); + const mockSetCredentials = vi.fn(); + const mockGenerateCodeVerifierAsync = vi + .fn() + .mockResolvedValue(mockCodeVerifier); + const mockOAuth2Client = { - generateCodeVerifierAsync: vi.fn().mockResolvedValue({ - codeChallenge: 'test-challenge', - codeVerifier: 'test-verifier', - }), - generateAuthUrl: vi.fn().mockReturnValue('https://example.com/auth'), - getToken: vi - .fn() - .mockRejectedValue(new Error('Invalid authorization code')), + generateAuthUrl: mockGenerateAuthUrl, + getToken: mockGetToken, + setCredentials: mockSetCredentials, + generateCodeVerifierAsync: mockGenerateCodeVerifierAsync, on: vi.fn(), } as unknown as OAuth2Client; (OAuth2Client as unknown as Mock).mockImplementation( @@ -851,7 +236,7 @@ describe('oauth2', () => { ); const mockReadline = { - question: vi.fn((_query, callback) => callback('invalid-code')), + question: vi.fn((_query, callback) => callback(mockCode)), close: vi.fn(), }; (readline.createInterface as Mock).mockReturnValue(mockReadline); @@ -859,96 +244,922 @@ describe('oauth2', () => { const consoleLogSpy = vi .spyOn(console, 'log') .mockImplementation(() => {}); - const consoleErrorSpy = vi - .spyOn(console, 'error') - .mockImplementation(() => {}); - await expect( - getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfigWithNoBrowser), - ).rejects.toThrow('Failed to authenticate with user code.'); - - expect(consoleErrorSpy).toHaveBeenCalledWith( - 'Failed to authenticate with authorization code:', - 'Invalid authorization code', + const client = await getOauthClient( + AuthType.LOGIN_WITH_GOOGLE, + mockConfigWithNoBrowser, ); + expect(client).toBe(mockOAuth2Client); + + // Verify the auth flow + expect(mockGenerateCodeVerifierAsync).toHaveBeenCalled(); + expect(mockGenerateAuthUrl).toHaveBeenCalled(); + expect(consoleLogSpy).toHaveBeenCalledWith( + expect.stringContaining(mockAuthUrl), + ); + expect(mockReadline.question).toHaveBeenCalledWith( + 'Enter the authorization code: ', + expect.any(Function), + ); + expect(mockGetToken).toHaveBeenCalledWith({ + code: mockCode, + codeVerifier: mockCodeVerifier.codeVerifier, + redirect_uri: 'https://codeassist.google.com/authcode', + }); + expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens); + consoleLogSpy.mockRestore(); - consoleErrorSpy.mockRestore(); + }); + + describe('in Cloud Shell', () => { + const mockGetAccessToken = vi.fn(); + let mockComputeClient: Compute; + + beforeEach(() => { + mockGetAccessToken.mockResolvedValue({ token: 'test-access-token' }); + mockComputeClient = { + credentials: { refresh_token: 'test-refresh-token' }, + getAccessToken: mockGetAccessToken, + } as unknown as Compute; + + (Compute as unknown as Mock).mockImplementation( + () => mockComputeClient, + ); + }); + + it('should attempt to load cached credentials first', async () => { + const cachedCreds = { refresh_token: 'cached-token' }; + const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); + await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); + await fs.promises.writeFile(credsPath, JSON.stringify(cachedCreds)); + + const mockClient = { + setCredentials: vi.fn(), + getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }), + getTokenInfo: vi.fn().mockResolvedValue({}), + on: vi.fn(), + }; + + // To mock the new OAuth2Client() inside the function + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockClient as unknown as OAuth2Client, + ); + + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + + expect(mockClient.setCredentials).toHaveBeenCalledWith(cachedCreds); + expect(mockClient.getAccessToken).toHaveBeenCalled(); + expect(mockClient.getTokenInfo).toHaveBeenCalled(); + expect(Compute).not.toHaveBeenCalled(); // Should not fetch new client if cache is valid + }); + + it('should use Compute to get a client if no cached credentials exist', async () => { + await getOauthClient(AuthType.CLOUD_SHELL, mockConfig); + + expect(Compute).toHaveBeenCalledWith({}); + expect(mockGetAccessToken).toHaveBeenCalled(); + }); + + it('should not cache the credentials after fetching them via ADC', async () => { + const newCredentials = { refresh_token: 'new-adc-token' }; + mockComputeClient.credentials = newCredentials; + mockGetAccessToken.mockResolvedValue({ token: 'new-adc-token' }); + + await getOauthClient(AuthType.CLOUD_SHELL, mockConfig); + + const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); + expect(fs.existsSync(credsPath)).toBe(false); + }); + + it('should return the Compute client on successful ADC authentication', async () => { + const client = await getOauthClient(AuthType.CLOUD_SHELL, mockConfig); + expect(client).toBe(mockComputeClient); + }); + + it('should throw an error if ADC fails', async () => { + const testError = new Error('ADC Failed'); + mockGetAccessToken.mockRejectedValue(testError); + + await expect( + getOauthClient(AuthType.CLOUD_SHELL, mockConfig), + ).rejects.toThrow( + 'Could not authenticate using Cloud Shell credentials. Please select a different authentication method or ensure you are in a properly configured environment. Error: ADC Failed', + ); + }); + }); + + describe('credential loading order', () => { + it('should prioritize default cached credentials over GOOGLE_APPLICATION_CREDENTIALS', async () => { + // Setup default cached credentials + const defaultCreds = { refresh_token: 'default-cached-token' }; + const defaultCredsPath = path.join( + tempHomeDir, + '.gemini', + 'oauth_creds.json', + ); + await fs.promises.mkdir(path.dirname(defaultCredsPath), { + recursive: true, + }); + await fs.promises.writeFile( + defaultCredsPath, + JSON.stringify(defaultCreds), + ); + + // Setup credentials via environment variable + const envCreds = { refresh_token: 'env-var-token' }; + const envCredsPath = path.join(tempHomeDir, 'env_creds.json'); + await fs.promises.writeFile(envCredsPath, JSON.stringify(envCreds)); + vi.stubEnv('GOOGLE_APPLICATION_CREDENTIALS', envCredsPath); + + const mockClient = { + setCredentials: vi.fn(), + getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }), + getTokenInfo: vi.fn().mockResolvedValue({}), + on: vi.fn(), + }; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockClient as unknown as OAuth2Client, + ); + + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + + // Assert the correct credentials were used + expect(mockClient.setCredentials).toHaveBeenCalledWith(defaultCreds); + expect(mockClient.setCredentials).not.toHaveBeenCalledWith(envCreds); + }); + + it('should fall back to GOOGLE_APPLICATION_CREDENTIALS if default cache is missing', async () => { + // Setup credentials via environment variable + const envCreds = { refresh_token: 'env-var-token' }; + const envCredsPath = path.join(tempHomeDir, 'env_creds.json'); + await fs.promises.writeFile(envCredsPath, JSON.stringify(envCreds)); + vi.stubEnv('GOOGLE_APPLICATION_CREDENTIALS', envCredsPath); + + const mockClient = { + setCredentials: vi.fn(), + getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }), + getTokenInfo: vi.fn().mockResolvedValue({}), + on: vi.fn(), + }; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockClient as unknown as OAuth2Client, + ); + + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + + // Assert the correct credentials were used + expect(mockClient.setCredentials).toHaveBeenCalledWith(envCreds); + }); + }); + + describe('with GCP environment variables', () => { + it('should use GOOGLE_CLOUD_ACCESS_TOKEN when GOOGLE_GENAI_USE_GCA is true', async () => { + vi.stubEnv('GOOGLE_GENAI_USE_GCA', 'true'); + vi.stubEnv('GOOGLE_CLOUD_ACCESS_TOKEN', 'gcp-access-token'); + + const mockSetCredentials = vi.fn(); + const mockGetAccessToken = vi + .fn() + .mockResolvedValue({ token: 'gcp-access-token' }); + const mockOAuth2Client = { + setCredentials: mockSetCredentials, + getAccessToken: mockGetAccessToken, + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + // Mock the UserInfo API response for fetchAndCacheUserInfo + (global.fetch as Mock).mockResolvedValue({ + ok: true, + json: vi + .fn() + .mockResolvedValue({ email: 'test-gcp-account@gmail.com' }), + } as unknown as Response); + + const client = await getOauthClient( + AuthType.LOGIN_WITH_GOOGLE, + mockConfig, + ); + + expect(client).toBe(mockOAuth2Client); + expect(mockSetCredentials).toHaveBeenCalledWith({ + access_token: 'gcp-access-token', + }); + + // Verify fetchAndCacheUserInfo was effectively called + expect(mockGetAccessToken).toHaveBeenCalled(); + expect(global.fetch).toHaveBeenCalledWith( + 'https://www.googleapis.com/oauth2/v2/userinfo', + { + headers: { + Authorization: 'Bearer gcp-access-token', + }, + }, + ); + + // Verify Google Account was cached + const googleAccountPath = path.join( + tempHomeDir, + '.gemini', + 'google_accounts.json', + ); + const cachedContent = fs.readFileSync(googleAccountPath, 'utf-8'); + expect(JSON.parse(cachedContent)).toEqual({ + active: 'test-gcp-account@gmail.com', + old: [], + }); + }); + + it('should not use GCP token if GOOGLE_CLOUD_ACCESS_TOKEN is not set', async () => { + vi.stubEnv('GOOGLE_GENAI_USE_GCA', 'true'); + + const mockSetCredentials = vi.fn(); + const mockGetAccessToken = vi + .fn() + .mockResolvedValue({ token: 'cached-access-token' }); + const mockGetTokenInfo = vi.fn().mockResolvedValue({}); + const mockOAuth2Client = { + setCredentials: mockSetCredentials, + getAccessToken: mockGetAccessToken, + getTokenInfo: mockGetTokenInfo, + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + // Make it fall through to cached credentials path + const cachedCreds = { refresh_token: 'cached-token' }; + const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); + await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); + await fs.promises.writeFile(credsPath, JSON.stringify(cachedCreds)); + + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + + // It should be called with the cached credentials, not the GCP access token. + expect(mockSetCredentials).toHaveBeenCalledTimes(1); + expect(mockSetCredentials).toHaveBeenCalledWith(cachedCreds); + }); + + it('should not use GCP token if GOOGLE_GENAI_USE_GCA is not set', async () => { + vi.stubEnv('GOOGLE_CLOUD_ACCESS_TOKEN', 'gcp-access-token'); + + const mockSetCredentials = vi.fn(); + const mockGetAccessToken = vi + .fn() + .mockResolvedValue({ token: 'cached-access-token' }); + const mockGetTokenInfo = vi.fn().mockResolvedValue({}); + const mockOAuth2Client = { + setCredentials: mockSetCredentials, + getAccessToken: mockGetAccessToken, + getTokenInfo: mockGetTokenInfo, + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + // Make it fall through to cached credentials path + const cachedCreds = { refresh_token: 'cached-token' }; + const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); + await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); + await fs.promises.writeFile(credsPath, JSON.stringify(cachedCreds)); + + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + + // It should be called with the cached credentials, not the GCP access token. + expect(mockSetCredentials).toHaveBeenCalledTimes(1); + expect(mockSetCredentials).toHaveBeenCalledWith(cachedCreds); + }); + }); + + describe('error handling', () => { + it('should handle browser launch failure with FatalAuthenticationError', async () => { + const mockError = new Error('Browser launch failed'); + (open as Mock).mockRejectedValue(mockError); + + const mockOAuth2Client = { + generateAuthUrl: vi.fn().mockReturnValue('https://example.com/auth'), + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + await expect( + getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig), + ).rejects.toThrow('Failed to open browser: Browser launch failed'); + }); + + it('should handle authentication timeout with proper error message', async () => { + const mockAuthUrl = 'https://example.com/auth'; + const mockOAuth2Client = { + generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + (open as Mock).mockImplementation( + async () => ({ on: vi.fn() }) as never, + ); + + const mockHttpServer = { + listen: vi.fn(), + close: vi.fn(), + on: vi.fn(), + address: () => ({ port: 3000 }), + }; + (http.createServer as Mock).mockImplementation( + () => mockHttpServer as unknown as http.Server, + ); + + // Mock setTimeout to trigger timeout immediately + const originalSetTimeout = global.setTimeout; + global.setTimeout = vi.fn( + (callback) => (callback(), {} as unknown as NodeJS.Timeout), + ) as unknown as typeof setTimeout; + + await expect( + getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig), + ).rejects.toThrow( + 'Authentication timed out after 5 minutes. The browser tab may have gotten stuck in a loading state. Please try again or use NO_BROWSER=true for manual authentication.', + ); + + global.setTimeout = originalSetTimeout; + }); + + it('should handle OAuth callback errors with descriptive messages', async () => { + const mockAuthUrl = 'https://example.com/auth'; + const mockOAuth2Client = { + generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + (open as Mock).mockImplementation( + async () => ({ on: vi.fn() }) as never, + ); + + let requestCallback!: http.RequestListener; + let serverListeningCallback: (value: unknown) => void; + const serverListeningPromise = new Promise( + (resolve) => (serverListeningCallback = resolve), + ); + + const mockHttpServer = { + listen: vi.fn( + (_port: number, _host: string, callback?: () => void) => { + if (callback) callback(); + serverListeningCallback(undefined); + }, + ), + close: vi.fn(), + on: vi.fn(), + address: () => ({ port: 3000 }), + }; + (http.createServer as Mock).mockImplementation((cb) => { + requestCallback = cb; + return mockHttpServer as unknown as http.Server; + }); + + const clientPromise = getOauthClient( + AuthType.LOGIN_WITH_GOOGLE, + mockConfig, + ); + await serverListeningPromise; + + // Test OAuth error with description + const mockReq = { + url: '/oauth2callback?error=access_denied&error_description=User+denied+access', + } as http.IncomingMessage; + const mockRes = { + writeHead: vi.fn(), + end: vi.fn(), + } as unknown as http.ServerResponse; + + await expect(async () => { + await requestCallback(mockReq, mockRes); + await clientPromise; + }).rejects.toThrow( + 'Google OAuth error: access_denied. User denied access', + ); + }); + + it('should handle OAuth error without description', async () => { + const mockAuthUrl = 'https://example.com/auth'; + const mockOAuth2Client = { + generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + (open as Mock).mockImplementation( + async () => ({ on: vi.fn() }) as never, + ); + + let requestCallback!: http.RequestListener; + let serverListeningCallback: (value: unknown) => void; + const serverListeningPromise = new Promise( + (resolve) => (serverListeningCallback = resolve), + ); + + const mockHttpServer = { + listen: vi.fn( + (_port: number, _host: string, callback?: () => void) => { + if (callback) callback(); + serverListeningCallback(undefined); + }, + ), + close: vi.fn(), + on: vi.fn(), + address: () => ({ port: 3000 }), + }; + (http.createServer as Mock).mockImplementation((cb) => { + requestCallback = cb; + return mockHttpServer as unknown as http.Server; + }); + + const clientPromise = getOauthClient( + AuthType.LOGIN_WITH_GOOGLE, + mockConfig, + ); + await serverListeningPromise; + + // Test OAuth error without description + const mockReq = { + url: '/oauth2callback?error=server_error', + } as http.IncomingMessage; + const mockRes = { + writeHead: vi.fn(), + end: vi.fn(), + } as unknown as http.ServerResponse; + + await expect(async () => { + await requestCallback(mockReq, mockRes); + await clientPromise; + }).rejects.toThrow( + 'Google OAuth error: server_error. No additional details provided', + ); + }); + + it('should handle token exchange failure with descriptive error', async () => { + const mockAuthUrl = 'https://example.com/auth'; + const mockCode = 'test-code'; + const mockState = 'test-state'; + + const mockOAuth2Client = { + generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), + getToken: vi + .fn() + .mockRejectedValue(new Error('Token exchange failed')), + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never); + (open as Mock).mockImplementation( + async () => ({ on: vi.fn() }) as never, + ); + + let requestCallback!: http.RequestListener; + let serverListeningCallback: (value: unknown) => void; + const serverListeningPromise = new Promise( + (resolve) => (serverListeningCallback = resolve), + ); + + const mockHttpServer = { + listen: vi.fn( + (_port: number, _host: string, callback?: () => void) => { + if (callback) callback(); + serverListeningCallback(undefined); + }, + ), + close: vi.fn(), + on: vi.fn(), + address: () => ({ port: 3000 }), + }; + (http.createServer as Mock).mockImplementation((cb) => { + requestCallback = cb; + return mockHttpServer as unknown as http.Server; + }); + + const clientPromise = getOauthClient( + AuthType.LOGIN_WITH_GOOGLE, + mockConfig, + ); + await serverListeningPromise; + + const mockReq = { + url: `/oauth2callback?code=${mockCode}&state=${mockState}`, + } as http.IncomingMessage; + const mockRes = { + writeHead: vi.fn(), + end: vi.fn(), + } as unknown as http.ServerResponse; + + await expect(async () => { + await requestCallback(mockReq, mockRes); + await clientPromise; + }).rejects.toThrow( + 'Failed to exchange authorization code for tokens: Token exchange failed', + ); + }); + + it('should handle fetchAndCacheUserInfo failure gracefully', async () => { + const mockAuthUrl = 'https://example.com/auth'; + const mockCode = 'test-code'; + const mockState = 'test-state'; + const mockTokens = { + access_token: 'test-access-token', + refresh_token: 'test-refresh-token', + }; + + const mockOAuth2Client = { + generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), + getToken: vi.fn().mockResolvedValue({ tokens: mockTokens }), + setCredentials: vi.fn(), + getAccessToken: vi + .fn() + .mockResolvedValue({ token: 'test-access-token' }), + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never); + (open as Mock).mockImplementation( + async () => ({ on: vi.fn() }) as never, + ); + + // Mock fetch to fail + (global.fetch as Mock).mockResolvedValue({ + ok: false, + status: 500, + statusText: 'Internal Server Error', + } as unknown as Response); + + const consoleErrorSpy = vi + .spyOn(console, 'error') + .mockImplementation(() => {}); + + let requestCallback!: http.RequestListener; + let serverListeningCallback: (value: unknown) => void; + const serverListeningPromise = new Promise( + (resolve) => (serverListeningCallback = resolve), + ); + + const mockHttpServer = { + listen: vi.fn( + (_port: number, _host: string, callback?: () => void) => { + if (callback) callback(); + serverListeningCallback(undefined); + }, + ), + close: vi.fn(), + on: vi.fn(), + address: () => ({ port: 3000 }), + }; + (http.createServer as Mock).mockImplementation((cb) => { + requestCallback = cb; + return mockHttpServer as unknown as http.Server; + }); + + const clientPromise = getOauthClient( + AuthType.LOGIN_WITH_GOOGLE, + mockConfig, + ); + await serverListeningPromise; + + const mockReq = { + url: `/oauth2callback?code=${mockCode}&state=${mockState}`, + } as http.IncomingMessage; + const mockRes = { + writeHead: vi.fn(), + end: vi.fn(), + } as unknown as http.ServerResponse; + + await requestCallback(mockReq, mockRes); + const client = await clientPromise; + + // Authentication should succeed even if fetchAndCacheUserInfo fails + expect(client).toBe(mockOAuth2Client); + expect(consoleErrorSpy).toHaveBeenCalledWith( + 'Failed to fetch user info:', + 500, + 'Internal Server Error', + ); + + consoleErrorSpy.mockRestore(); + }); + + it('should handle user code authentication failure with descriptive error', async () => { + const mockConfigWithNoBrowser = { + getNoBrowser: () => true, + getProxy: () => 'http://test.proxy.com:8080', + isBrowserLaunchSuppressed: () => true, + } as unknown as Config; + + const mockOAuth2Client = { + generateCodeVerifierAsync: vi.fn().mockResolvedValue({ + codeChallenge: 'test-challenge', + codeVerifier: 'test-verifier', + }), + generateAuthUrl: vi.fn().mockReturnValue('https://example.com/auth'), + getToken: vi + .fn() + .mockRejectedValue(new Error('Invalid authorization code')), + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + const mockReadline = { + question: vi.fn((_query, callback) => callback('invalid-code')), + close: vi.fn(), + }; + (readline.createInterface as Mock).mockReturnValue(mockReadline); + + const consoleLogSpy = vi + .spyOn(console, 'log') + .mockImplementation(() => {}); + const consoleErrorSpy = vi + .spyOn(console, 'error') + .mockImplementation(() => {}); + + await expect( + getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfigWithNoBrowser), + ).rejects.toThrow('Failed to authenticate with user code.'); + + expect(consoleErrorSpy).toHaveBeenCalledWith( + 'Failed to authenticate with authorization code:', + 'Invalid authorization code', + ); + + consoleLogSpy.mockRestore(); + consoleErrorSpy.mockRestore(); + }); + }); + + describe('clearCachedCredentialFile', () => { + it('should clear cached credentials and Google account', async () => { + const cachedCreds = { refresh_token: 'test-token' }; + const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); + await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); + await fs.promises.writeFile(credsPath, JSON.stringify(cachedCreds)); + + const googleAccountPath = path.join( + tempHomeDir, + '.gemini', + 'google_accounts.json', + ); + const accountData = { active: 'test@example.com', old: [] }; + await fs.promises.writeFile( + googleAccountPath, + JSON.stringify(accountData), + ); + const userAccountManager = new UserAccountManager(); + + expect(fs.existsSync(credsPath)).toBe(true); + expect(fs.existsSync(googleAccountPath)).toBe(true); + expect(userAccountManager.getCachedGoogleAccount()).toBe( + 'test@example.com', + ); + + await clearCachedCredentialFile(); + expect(fs.existsSync(credsPath)).toBe(false); + expect(userAccountManager.getCachedGoogleAccount()).toBeNull(); + const updatedAccountData = JSON.parse( + fs.readFileSync(googleAccountPath, 'utf-8'), + ); + expect(updatedAccountData.active).toBeNull(); + expect(updatedAccountData.old).toContain('test@example.com'); + }); + + it('should clear the in-memory OAuth client cache', async () => { + const mockSetCredentials = vi.fn(); + const mockGetAccessToken = vi + .fn() + .mockResolvedValue({ token: 'test-token' }); + const mockGetTokenInfo = vi.fn().mockResolvedValue({}); + const mockOAuth2Client = { + setCredentials: mockSetCredentials, + getAccessToken: mockGetAccessToken, + getTokenInfo: mockGetTokenInfo, + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + // Pre-populate credentials to make getOauthClient resolve quickly + const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); + await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); + await fs.promises.writeFile( + credsPath, + JSON.stringify({ refresh_token: 'token' }), + ); + + // First call, should create a client + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + expect(OAuth2Client).toHaveBeenCalledTimes(1); + + // Second call, should use cached client + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + expect(OAuth2Client).toHaveBeenCalledTimes(1); + + clearOauthClientCache(); + + // Third call, after clearing cache, should create a new client + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + expect(OAuth2Client).toHaveBeenCalledTimes(2); + }); }); }); - describe('clearCachedCredentialFile', () => { - it('should clear cached credentials and Google account', async () => { - const cachedCreds = { refresh_token: 'test-token' }; - const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); - await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); - await fs.promises.writeFile(credsPath, JSON.stringify(cachedCreds)); - - const googleAccountPath = path.join( - tempHomeDir, - '.gemini', - 'google_accounts.json', + describe('with encrypted flag true', () => { + let tempHomeDir: string; + beforeEach(() => { + process.env[FORCE_ENCRYPTED_FILE_ENV_VAR] = 'true'; + tempHomeDir = fs.mkdtempSync( + path.join(os.tmpdir(), 'gemini-cli-test-home-'), ); - const accountData = { active: 'test@example.com', old: [] }; - await fs.promises.writeFile( - googleAccountPath, - JSON.stringify(accountData), - ); - const userAccountManager = new UserAccountManager(); - - expect(fs.existsSync(credsPath)).toBe(true); - expect(fs.existsSync(googleAccountPath)).toBe(true); - expect(userAccountManager.getCachedGoogleAccount()).toBe( - 'test@example.com', - ); - - await clearCachedCredentialFile(); - expect(fs.existsSync(credsPath)).toBe(false); - expect(userAccountManager.getCachedGoogleAccount()).toBeNull(); - const updatedAccountData = JSON.parse( - fs.readFileSync(googleAccountPath, 'utf-8'), - ); - expect(updatedAccountData.active).toBeNull(); - expect(updatedAccountData.old).toContain('test@example.com'); + (os.homedir as Mock).mockReturnValue(tempHomeDir); }); - it('should clear the in-memory OAuth client cache', async () => { - const mockSetCredentials = vi.fn(); - const mockGetAccessToken = vi - .fn() - .mockResolvedValue({ token: 'test-token' }); - const mockGetTokenInfo = vi.fn().mockResolvedValue({}); + afterEach(() => { + fs.rmSync(tempHomeDir, { recursive: true, force: true }); + vi.clearAllMocks(); + resetOauthClientForTesting(); + vi.unstubAllEnvs(); + }); + + it('should save credentials using OAuthCredentialStorage during web login', async () => { + const { OAuthCredentialStorage } = await import( + './oauth-credential-storage.js' + ); + const mockAuthUrl = 'https://example.com/auth'; + const mockCode = 'test-code'; + const mockState = 'test-state'; + const mockTokens = { + access_token: 'test-access-token', + refresh_token: 'test-refresh-token', + }; + + let onTokensCallback: (tokens: Credentials) => void = () => {}; + const mockOn = vi.fn((event, callback) => { + if (event === 'tokens') { + onTokensCallback = callback; + } + }); + + const mockGetToken = vi.fn().mockImplementation(async () => { + onTokensCallback(mockTokens); + return { tokens: mockTokens }; + }); + const mockOAuth2Client = { - setCredentials: mockSetCredentials, - getAccessToken: mockGetAccessToken, - getTokenInfo: mockGetTokenInfo, - on: vi.fn(), + generateAuthUrl: vi.fn().mockReturnValue(mockAuthUrl), + getToken: mockGetToken, + setCredentials: vi.fn(), + getAccessToken: vi + .fn() + .mockResolvedValue({ token: 'mock-access-token' }), + on: mockOn, + credentials: mockTokens, } as unknown as OAuth2Client; (OAuth2Client as unknown as Mock).mockImplementation( () => mockOAuth2Client, ); - // Pre-populate credentials to make getOauthClient resolve quickly - const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); - await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); - await fs.promises.writeFile( - credsPath, - JSON.stringify({ refresh_token: 'token' }), + vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never); + (open as Mock).mockImplementation(async () => ({ on: vi.fn() }) as never); + + (global.fetch as Mock).mockResolvedValue({ + ok: true, + json: vi + .fn() + .mockResolvedValue({ email: 'test-google-account@gmail.com' }), + } as unknown as Response); + + let requestCallback!: http.RequestListener; + let serverListeningCallback: (value: unknown) => void; + const serverListeningPromise = new Promise( + (resolve) => (serverListeningCallback = resolve), ); - // First call, should create a client - await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); - expect(OAuth2Client).toHaveBeenCalledTimes(1); + let capturedPort = 0; + const mockHttpServer = { + listen: vi.fn((port: number, _host: string, callback?: () => void) => { + capturedPort = port; + if (callback) { + callback(); + } + serverListeningCallback(undefined); + }), + close: vi.fn((callback?: () => void) => { + if (callback) { + callback(); + } + }), + on: vi.fn(), + address: () => ({ port: capturedPort }), + }; + (http.createServer as Mock).mockImplementation((cb) => { + requestCallback = cb as http.RequestListener; + return mockHttpServer as unknown as http.Server; + }); - // Second call, should use cached client - await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); - expect(OAuth2Client).toHaveBeenCalledTimes(1); + const clientPromise = getOauthClient( + AuthType.LOGIN_WITH_GOOGLE, + mockConfig, + ); - clearOauthClientCache(); + await serverListeningPromise; + + const mockReq = { + url: `/oauth2callback?code=${mockCode}&state=${mockState}`, + } as http.IncomingMessage; + const mockRes = { + writeHead: vi.fn(), + end: vi.fn(), + } as unknown as http.ServerResponse; + + requestCallback(mockReq, mockRes); + + await clientPromise; + + expect( + OAuthCredentialStorage.saveCredentials as Mock, + ).toHaveBeenCalledWith(mockTokens); + const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); + expect(fs.existsSync(credsPath)).toBe(false); + }); + + it('should load credentials using OAuthCredentialStorage and not from file', async () => { + const { OAuthCredentialStorage } = await import( + './oauth-credential-storage.js' + ); + const cachedCreds = { refresh_token: 'cached-encrypted-token' }; + (OAuthCredentialStorage.loadCredentials as Mock).mockResolvedValue( + cachedCreds, + ); + + // Create a dummy unencrypted credential file. + // If the logic is correct, this file should be ignored. + const unencryptedCreds = { refresh_token: 'unencrypted-token' }; + const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); + await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); + await fs.promises.writeFile(credsPath, JSON.stringify(unencryptedCreds)); + + const mockClient = { + setCredentials: vi.fn(), + getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }), + getTokenInfo: vi.fn().mockResolvedValue({}), + on: vi.fn(), + }; + + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockClient as unknown as OAuth2Client, + ); - // Third call, after clearing cache, should create a new client await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); - expect(OAuth2Client).toHaveBeenCalledTimes(2); + + expect(OAuthCredentialStorage.loadCredentials as Mock).toHaveBeenCalled(); + expect(mockClient.setCredentials).toHaveBeenCalledWith(cachedCreds); + expect(mockClient.setCredentials).not.toHaveBeenCalledWith( + unencryptedCreds, + ); + }); + + it('should clear credentials using OAuthCredentialStorage', async () => { + const { OAuthCredentialStorage } = await import( + './oauth-credential-storage.js' + ); + + // Create a dummy unencrypted credential file. It should not be deleted. + const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); + await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); + await fs.promises.writeFile(credsPath, '{}'); + + await clearCachedCredentialFile(); + + expect( + OAuthCredentialStorage.clearCredentials as Mock, + ).toHaveBeenCalled(); + expect(fs.existsSync(credsPath)).toBe(true); // The unencrypted file should remain }); }); }); diff --git a/packages/core/src/code_assist/oauth2.ts b/packages/core/src/code_assist/oauth2.ts index 142a3791f0..b88877058b 100644 --- a/packages/core/src/code_assist/oauth2.ts +++ b/packages/core/src/code_assist/oauth2.ts @@ -23,6 +23,8 @@ import { UserAccountManager } from '../utils/userAccountManager.js'; import { AuthType } from '../core/contentGenerator.js'; import readline from 'node:readline'; import { Storage } from '../config/storage.js'; +import { OAuthCredentialStorage } from './oauth-credential-storage.js'; +import { FORCE_ENCRYPTED_FILE_ENV_VAR } from '../mcp/token-storage/index.js'; const userAccountManager = new UserAccountManager(); @@ -63,6 +65,10 @@ export interface OauthWebLogin { const oauthClientPromises = new Map>(); +function getUseEncryptedStorageFlag() { + return process.env[FORCE_ENCRYPTED_FILE_ENV_VAR] === 'true'; +} + async function initOauthClient( authType: AuthType, config: Config, @@ -74,6 +80,7 @@ async function initOauthClient( proxy: config.getProxy(), }, }); + const useEncryptedStorage = getUseEncryptedStorageFlag(); if ( process.env['GOOGLE_GENAI_USE_GCA'] && @@ -87,7 +94,11 @@ async function initOauthClient( } client.on('tokens', async (tokens: Credentials) => { - await cacheCredentials(tokens); + if (useEncryptedStorage) { + await OAuthCredentialStorage.saveCredentials(tokens); + } else { + await cacheCredentials(tokens); + } }); // If there are cached creds on disk, they always take precedence @@ -419,6 +430,16 @@ export function getAvailablePort(): Promise { } async function loadCachedCredentials(client: OAuth2Client): Promise { + const useEncryptedStorage = getUseEncryptedStorageFlag(); + if (useEncryptedStorage) { + const credentials = await OAuthCredentialStorage.loadCredentials(); + if (credentials) { + client.setCredentials(credentials); + return true; + } + return false; + } + const pathsToTry = [ Storage.getOAuthCredsPath(), process.env['GOOGLE_APPLICATION_CREDENTIALS'], @@ -470,7 +491,12 @@ export function clearOauthClientCache() { export async function clearCachedCredentialFile() { try { - await fs.rm(Storage.getOAuthCredsPath(), { force: true }); + const useEncryptedStorage = getUseEncryptedStorageFlag(); + if (useEncryptedStorage) { + await OAuthCredentialStorage.clearCredentials(); + } else { + await fs.rm(Storage.getOAuthCredsPath(), { force: true }); + } // Clear the Google Account ID cache when credentials are cleared await userAccountManager.clearCachedGoogleAccount(); // Clear the in-memory OAuth client cache to force re-authentication