Cache credentials in home dir, not working dir (#1122)
This commit is contained in:
parent
98093e604a
commit
b3e26de862
|
@ -4,12 +4,23 @@
|
||||||
* SPDX-License-Identifier: Apache-2.0
|
* SPDX-License-Identifier: Apache-2.0
|
||||||
*/
|
*/
|
||||||
|
|
||||||
import { describe, it, expect, vi } from 'vitest';
|
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||||
import { webLoginClient } from './oauth2.js';
|
import { getOauthClient } from './oauth2.js';
|
||||||
import { OAuth2Client } from 'google-auth-library';
|
import { OAuth2Client } from 'google-auth-library';
|
||||||
|
import * as fs from 'fs';
|
||||||
|
import * as path from 'path';
|
||||||
import http from 'http';
|
import http from 'http';
|
||||||
import open from 'open';
|
import open from 'open';
|
||||||
import crypto from 'crypto';
|
import crypto from 'crypto';
|
||||||
|
import * as os from 'os';
|
||||||
|
|
||||||
|
vi.mock('os', async (importOriginal) => {
|
||||||
|
const os = await importOriginal<typeof import('os')>();
|
||||||
|
return {
|
||||||
|
...os,
|
||||||
|
homedir: vi.fn(),
|
||||||
|
};
|
||||||
|
});
|
||||||
|
|
||||||
vi.mock('google-auth-library');
|
vi.mock('google-auth-library');
|
||||||
vi.mock('http');
|
vi.mock('http');
|
||||||
|
@ -17,6 +28,18 @@ vi.mock('open');
|
||||||
vi.mock('crypto');
|
vi.mock('crypto');
|
||||||
|
|
||||||
describe('oauth2', () => {
|
describe('oauth2', () => {
|
||||||
|
let tempHomeDir: string;
|
||||||
|
|
||||||
|
beforeEach(() => {
|
||||||
|
tempHomeDir = fs.mkdtempSync(
|
||||||
|
path.join(os.tmpdir(), 'gemini-cli-test-home-'),
|
||||||
|
);
|
||||||
|
vi.mocked(os.homedir).mockReturnValue(tempHomeDir);
|
||||||
|
});
|
||||||
|
afterEach(() => {
|
||||||
|
fs.rmSync(tempHomeDir, { recursive: true, force: true });
|
||||||
|
});
|
||||||
|
|
||||||
it('should perform a web login', async () => {
|
it('should perform a web login', async () => {
|
||||||
const mockAuthUrl = 'https://example.com/auth';
|
const mockAuthUrl = 'https://example.com/auth';
|
||||||
const mockCode = 'test-code';
|
const mockCode = 'test-code';
|
||||||
|
@ -33,16 +56,17 @@ describe('oauth2', () => {
|
||||||
generateAuthUrl: mockGenerateAuthUrl,
|
generateAuthUrl: mockGenerateAuthUrl,
|
||||||
getToken: mockGetToken,
|
getToken: mockGetToken,
|
||||||
setCredentials: mockSetCredentials,
|
setCredentials: mockSetCredentials,
|
||||||
|
credentials: mockTokens,
|
||||||
} as unknown as OAuth2Client;
|
} as unknown as OAuth2Client;
|
||||||
vi.mocked(OAuth2Client).mockImplementation(() => mockOAuth2Client);
|
vi.mocked(OAuth2Client).mockImplementation(() => mockOAuth2Client);
|
||||||
|
|
||||||
vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never);
|
vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never);
|
||||||
vi.mocked(open).mockImplementation(async () => ({}) as never);
|
vi.mocked(open).mockImplementation(async () => ({}) as never);
|
||||||
|
|
||||||
let requestCallback!: (
|
let requestCallback!: http.RequestListener<
|
||||||
req: http.IncomingMessage,
|
typeof http.IncomingMessage,
|
||||||
res: http.ServerResponse,
|
typeof http.ServerResponse
|
||||||
) => void;
|
>;
|
||||||
const mockHttpServer = {
|
const mockHttpServer = {
|
||||||
listen: vi.fn((port: number, callback?: () => void) => {
|
listen: vi.fn((port: number, callback?: () => void) => {
|
||||||
if (callback) {
|
if (callback) {
|
||||||
|
@ -58,14 +82,14 @@ describe('oauth2', () => {
|
||||||
address: () => ({ port: 1234 }),
|
address: () => ({ port: 1234 }),
|
||||||
};
|
};
|
||||||
vi.mocked(http.createServer).mockImplementation((cb) => {
|
vi.mocked(http.createServer).mockImplementation((cb) => {
|
||||||
requestCallback = cb as (
|
requestCallback = cb as http.RequestListener<
|
||||||
req: http.IncomingMessage,
|
typeof http.IncomingMessage,
|
||||||
res: http.ServerResponse,
|
typeof http.ServerResponse
|
||||||
) => void;
|
>;
|
||||||
return mockHttpServer as unknown as http.Server;
|
return mockHttpServer as unknown as http.Server;
|
||||||
});
|
});
|
||||||
|
|
||||||
const clientPromise = webLoginClient();
|
const clientPromise = getOauthClient();
|
||||||
|
|
||||||
// Wait for the server to be created
|
// Wait for the server to be created
|
||||||
await new Promise((resolve) => setTimeout(resolve, 0));
|
await new Promise((resolve) => setTimeout(resolve, 0));
|
||||||
|
@ -78,15 +102,17 @@ describe('oauth2', () => {
|
||||||
end: vi.fn(),
|
end: vi.fn(),
|
||||||
} as unknown as http.ServerResponse;
|
} as unknown as http.ServerResponse;
|
||||||
|
|
||||||
if (requestCallback) {
|
await requestCallback(mockReq, mockRes);
|
||||||
await requestCallback(mockReq, mockRes);
|
|
||||||
}
|
|
||||||
|
|
||||||
const client = await clientPromise;
|
const client = await clientPromise;
|
||||||
|
expect(client).toBe(mockOAuth2Client);
|
||||||
|
|
||||||
expect(open).toHaveBeenCalledWith(mockAuthUrl);
|
expect(open).toHaveBeenCalledWith(mockAuthUrl);
|
||||||
expect(mockGetToken).toHaveBeenCalledWith(mockCode);
|
expect(mockGetToken).toHaveBeenCalledWith(mockCode);
|
||||||
expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens);
|
expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens);
|
||||||
expect(client).toBe(mockOAuth2Client);
|
|
||||||
|
const tokenPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json');
|
||||||
|
const tokenData = JSON.parse(fs.readFileSync(tokenPath, 'utf-8'));
|
||||||
|
expect(tokenData).toEqual(mockTokens);
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
|
@ -12,6 +12,7 @@ import * as net from 'net';
|
||||||
import open from 'open';
|
import open from 'open';
|
||||||
import path from 'node:path';
|
import path from 'node:path';
|
||||||
import { promises as fs } from 'node:fs';
|
import { promises as fs } from 'node:fs';
|
||||||
|
import * as os from 'os';
|
||||||
|
|
||||||
// OAuth Client ID used to initiate OAuth2Client class.
|
// OAuth Client ID used to initiate OAuth2Client class.
|
||||||
const OAUTH_CLIENT_ID =
|
const OAUTH_CLIENT_ID =
|
||||||
|
@ -41,30 +42,8 @@ const SIGN_IN_FAILURE_URL =
|
||||||
const GEMINI_DIR = '.gemini';
|
const GEMINI_DIR = '.gemini';
|
||||||
const CREDENTIAL_FILENAME = 'oauth_creds.json';
|
const CREDENTIAL_FILENAME = 'oauth_creds.json';
|
||||||
|
|
||||||
export async function getCachedCredentialClient(): Promise<OAuth2Client> {
|
|
||||||
try {
|
|
||||||
const creds = await fs.readFile(
|
|
||||||
path.join(process.cwd(), GEMINI_DIR, CREDENTIAL_FILENAME),
|
|
||||||
'utf-8',
|
|
||||||
);
|
|
||||||
|
|
||||||
const oAuth2Client = new OAuth2Client({
|
|
||||||
clientId: OAUTH_CLIENT_ID,
|
|
||||||
clientSecret: OAUTH_CLIENT_SECRET,
|
|
||||||
});
|
|
||||||
oAuth2Client.setCredentials(JSON.parse(creds));
|
|
||||||
// This will either return the existing token or refresh it.
|
|
||||||
await oAuth2Client.getAccessToken();
|
|
||||||
// If we are here, the token is valid.
|
|
||||||
return oAuth2Client;
|
|
||||||
} catch (_) {
|
|
||||||
// Could not load credentials.
|
|
||||||
throw new Error('Could not load credentials');
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
export async function clearCachedCredentials(): Promise<void> {
|
export async function clearCachedCredentials(): Promise<void> {
|
||||||
await fs.rm(path.join(process.cwd(), GEMINI_DIR, CREDENTIAL_FILENAME));
|
await fs.rm(getCachedCredentialPath());
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function getOauthClient(): Promise<OAuth2Client> {
|
export async function getOauthClient(): Promise<OAuth2Client> {
|
||||||
|
@ -72,16 +51,19 @@ export async function getOauthClient(): Promise<OAuth2Client> {
|
||||||
return await getCachedCredentialClient();
|
return await getCachedCredentialClient();
|
||||||
} catch (_) {
|
} catch (_) {
|
||||||
const loggedInClient = await webLoginClient();
|
const loggedInClient = await webLoginClient();
|
||||||
await fs.mkdir(path.join(process.cwd(), GEMINI_DIR), { recursive: true });
|
|
||||||
|
await fs.mkdir(path.dirname(getCachedCredentialPath()), {
|
||||||
|
recursive: true,
|
||||||
|
});
|
||||||
await fs.writeFile(
|
await fs.writeFile(
|
||||||
path.join(process.cwd(), GEMINI_DIR, CREDENTIAL_FILENAME),
|
getCachedCredentialPath(),
|
||||||
JSON.stringify(loggedInClient.credentials, null, 2),
|
JSON.stringify(loggedInClient.credentials, null, 2),
|
||||||
);
|
);
|
||||||
return loggedInClient;
|
return loggedInClient;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function webLoginClient(): Promise<OAuth2Client> {
|
async function webLoginClient(): Promise<OAuth2Client> {
|
||||||
const port = await getAvailablePort();
|
const port = await getAvailablePort();
|
||||||
const oAuth2Client = new OAuth2Client({
|
const oAuth2Client = new OAuth2Client({
|
||||||
clientId: OAUTH_CLIENT_ID,
|
clientId: OAUTH_CLIENT_ID,
|
||||||
|
@ -163,3 +145,26 @@ function getAvailablePort(): Promise<number> {
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async function getCachedCredentialClient(): Promise<OAuth2Client> {
|
||||||
|
try {
|
||||||
|
const creds = await fs.readFile(getCachedCredentialPath(), 'utf-8');
|
||||||
|
|
||||||
|
const oAuth2Client = new OAuth2Client({
|
||||||
|
clientId: OAUTH_CLIENT_ID,
|
||||||
|
clientSecret: OAUTH_CLIENT_SECRET,
|
||||||
|
});
|
||||||
|
oAuth2Client.setCredentials(JSON.parse(creds));
|
||||||
|
// This will either return the existing token or refresh it.
|
||||||
|
await oAuth2Client.getAccessToken();
|
||||||
|
// If we are here, the token is valid.
|
||||||
|
return oAuth2Client;
|
||||||
|
} catch (_) {
|
||||||
|
// Could not load credentials.
|
||||||
|
throw new Error('Could not load credentials');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function getCachedCredentialPath(): string {
|
||||||
|
return path.join(os.homedir(), GEMINI_DIR, CREDENTIAL_FILENAME);
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue