From 12f584fff8b197006b9c10f2491144512317a8ac Mon Sep 17 00:00:00 2001 From: Shreya Keshive Date: Mon, 15 Sep 2025 18:49:15 -0400 Subject: [PATCH] feat(vscode-ide-companion): add auth token validation to IDE server (#8491) --- docs/ide-companion-spec.md | 12 +- .../src/ide-server.test.ts | 109 ++++++++++++++++++ .../vscode-ide-companion/src/ide-server.ts | 93 +++++++++++---- 3 files changed, 188 insertions(+), 26 deletions(-) diff --git a/docs/ide-companion-spec.md b/docs/ide-companion-spec.md index 39f7355ce1..b13622af52 100644 --- a/docs/ide-companion-spec.md +++ b/docs/ide-companion-spec.md @@ -37,7 +37,17 @@ For Gemini CLI to connect, it needs to discover which IDE instance it's running - `workspacePath` (string): A list of all open workspace root paths, delimited by the OS-specific path separator (`:` for Linux/macOS, `;` for Windows). The CLI uses this path to ensure it's running in the same project folder that's open in the IDE. If the CLI's current working directory is not a sub-directory of `workspacePath`, the connection will be rejected. Your extension **MUST** provide the correct, absolute path(s) to the root of the open workspace(s). - **Tie-Breaking with Environment Variables (Recommended):** For the most reliable experience, your extension **SHOULD** both create the discovery file and set the `GEMINI_CLI_IDE_SERVER_PORT` and `GEMINI_CLI_IDE_WORKSPACE_PATH` environment variables in the integrated terminal. The file serves as the primary discovery mechanism, but the environment variables are crucial for tie-breaking. If a user has multiple IDE windows open for the same workspace, the CLI uses the `GEMINI_CLI_IDE_SERVER_PORT` variable to identify and connect to the correct window's server. - For prototyping, you may opt to _only_ set the environment variables. However, this is not a robust solution for a production extension, as environment variables may not be reliably set in all terminal sessions (e.g., restored terminals), which can lead to connection failures. -- **Authentication:** (TBD) +- **Authentication:** To secure the connection, the extension **SHOULD** generate a unique, secret token and include it in the discovery file. The CLI will then include this token in all requests to the MCP server. + - **Token Generation:** The extension should generate a random string to be used as a bearer token. + - **Discovery File Content:** The `authToken` field must be added to the JSON object in the discovery file: + ```json + { + "port": 12345, + "workspacePath": "/path/to/project", + "authToken": "a-very-secret-token" + } + ``` + - **Request Authorization:** The CLI will read the `authToken` from the file and include it in the `Authorization` header for all HTTP requests to the MCP server (e.g., `Authorization: Bearer a-very-secret-token`). Your server **MUST** validate this token on every request and reject any that are unauthorized. ## II. The Context Interface diff --git a/packages/vscode-ide-companion/src/ide-server.test.ts b/packages/vscode-ide-companion/src/ide-server.test.ts index 5dbbd3e936..2367ac1551 100644 --- a/packages/vscode-ide-companion/src/ide-server.test.ts +++ b/packages/vscode-ide-companion/src/ide-server.test.ts @@ -12,6 +12,10 @@ import * as path from 'node:path'; import { IDEServer } from './ide-server.js'; import type { DiffManager } from './diff-manager.js'; +vi.mock('node:crypto', () => ({ + randomUUID: vi.fn(() => 'test-auth-token'), +})); + const mocks = vi.hoisted(() => ({ diffManager: { onDidChange: vi.fn(() => ({ dispose: vi.fn() })), @@ -21,6 +25,7 @@ const mocks = vi.hoisted(() => ({ vi.mock('node:fs/promises', () => ({ writeFile: vi.fn(() => Promise.resolve(undefined)), unlink: vi.fn(() => Promise.resolve(undefined)), + chmod: vi.fn(() => Promise.resolve(undefined)), })); vi.mock('node:os', async (importOriginal) => { @@ -134,6 +139,7 @@ describe('IDEServer', () => { port: parseInt(port, 10), workspacePath: expectedWorkspacePaths, ppid: process.ppid, + authToken: 'test-auth-token', }); expect(fs.writeFile).toHaveBeenCalledWith( expectedPortFile, @@ -143,6 +149,8 @@ describe('IDEServer', () => { expectedPpidPortFile, expectedContent, ); + expect(fs.chmod).toHaveBeenCalledWith(expectedPortFile, 0o600); + expect(fs.chmod).toHaveBeenCalledWith(expectedPpidPortFile, 0o600); }); it('should set a single folder path', async () => { @@ -169,6 +177,7 @@ describe('IDEServer', () => { port: parseInt(port, 10), workspacePath: '/foo/bar', ppid: process.ppid, + authToken: 'test-auth-token', }); expect(fs.writeFile).toHaveBeenCalledWith( expectedPortFile, @@ -178,6 +187,8 @@ describe('IDEServer', () => { expectedPpidPortFile, expectedContent, ); + expect(fs.chmod).toHaveBeenCalledWith(expectedPortFile, 0o600); + expect(fs.chmod).toHaveBeenCalledWith(expectedPpidPortFile, 0o600); }); it('should set an empty string if no folders are open', async () => { @@ -204,6 +215,7 @@ describe('IDEServer', () => { port: parseInt(port, 10), workspacePath: '', ppid: process.ppid, + authToken: 'test-auth-token', }); expect(fs.writeFile).toHaveBeenCalledWith( expectedPortFile, @@ -213,6 +225,8 @@ describe('IDEServer', () => { expectedPpidPortFile, expectedContent, ); + expect(fs.chmod).toHaveBeenCalledWith(expectedPortFile, 0o600); + expect(fs.chmod).toHaveBeenCalledWith(expectedPpidPortFile, 0o600); }); it('should update the path when workspace folders change', async () => { @@ -253,6 +267,7 @@ describe('IDEServer', () => { port: parseInt(port, 10), workspacePath: expectedWorkspacePaths, ppid: process.ppid, + authToken: 'test-auth-token', }); expect(fs.writeFile).toHaveBeenCalledWith( expectedPortFile, @@ -262,6 +277,8 @@ describe('IDEServer', () => { expectedPpidPortFile, expectedContent, ); + expect(fs.chmod).toHaveBeenCalledWith(expectedPortFile, 0o600); + expect(fs.chmod).toHaveBeenCalledWith(expectedPpidPortFile, 0o600); // Simulate removing a folder vscodeMock.workspace.workspaceFolders = [{ uri: { fsPath: '/baz/qux' } }]; @@ -275,6 +292,7 @@ describe('IDEServer', () => { port: parseInt(port, 10), workspacePath: '/baz/qux', ppid: process.ppid, + authToken: 'test-auth-token', }); expect(fs.writeFile).toHaveBeenCalledWith( expectedPortFile, @@ -284,6 +302,8 @@ describe('IDEServer', () => { expectedPpidPortFile, expectedContent2, ); + expect(fs.chmod).toHaveBeenCalledWith(expectedPortFile, 0o600); + expect(fs.chmod).toHaveBeenCalledWith(expectedPpidPortFile, 0o600); }); it('should clear env vars and delete port file on stop', async () => { @@ -335,6 +355,7 @@ describe('IDEServer', () => { port: parseInt(port, 10), workspacePath: expectedWorkspacePaths, ppid: process.ppid, + authToken: 'test-auth-token', }); expect(fs.writeFile).toHaveBeenCalledWith( expectedPortFile, @@ -344,6 +365,94 @@ describe('IDEServer', () => { expectedPpidPortFile, expectedContent, ); + expect(fs.chmod).toHaveBeenCalledWith(expectedPortFile, 0o600); + expect(fs.chmod).toHaveBeenCalledWith(expectedPpidPortFile, 0o600); }, ); + + describe('auth token', () => { + let port: number; + + beforeEach(async () => { + await ideServer.start(mockContext); + port = (ideServer as unknown as { port: number }).port; + }); + + it('should allow request without auth token for backwards compatibility', async () => { + const response = await fetch(`http://localhost:${port}/mcp`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + jsonrpc: '2.0', + method: 'initialize', + params: {}, + id: 1, + }), + }); + expect(response.status).not.toBe(401); + }); + + it('should allow request with valid auth token', async () => { + const response = await fetch(`http://localhost:${port}/mcp`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer test-auth-token`, + }, + body: JSON.stringify({ + jsonrpc: '2.0', + method: 'initialize', + params: {}, + id: 1, + }), + }); + expect(response.status).not.toBe(401); + }); + + it('should reject request with invalid auth token', async () => { + const response = await fetch(`http://localhost:${port}/mcp`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: 'Bearer invalid-token', + }, + body: JSON.stringify({ + jsonrpc: '2.0', + method: 'initialize', + params: {}, + id: 1, + }), + }); + expect(response.status).toBe(401); + const body = await response.text(); + expect(body).toBe('Unauthorized'); + }); + + it('should reject request with malformed auth token', async () => { + const malformedHeaders = [ + 'Bearer', + 'invalid-token', + 'Bearer token extra', + ]; + + for (const header of malformedHeaders) { + const response = await fetch(`http://localhost:${port}/mcp`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: header, + }, + body: JSON.stringify({ + jsonrpc: '2.0', + method: 'initialize', + params: {}, + id: 1, + }), + }); + expect(response.status, `Failed for header: ${header}`).toBe(401); + const body = await response.text(); + expect(body, `Failed for header: ${header}`).toBe('Unauthorized'); + } + }); + }); }); diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index fe86b20f5d..883012d1d3 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -27,13 +27,23 @@ const MCP_SESSION_ID_HEADER = 'mcp-session-id'; const IDE_SERVER_PORT_ENV_VAR = 'GEMINI_CLI_IDE_SERVER_PORT'; const IDE_WORKSPACE_PATH_ENV_VAR = 'GEMINI_CLI_IDE_WORKSPACE_PATH'; -async function writePortAndWorkspace( - context: vscode.ExtensionContext, - port: number, - portFile: string, - ppidPortFile: string, - log: (message: string) => void, -): Promise { +interface WritePortAndWorkspaceArgs { + context: vscode.ExtensionContext; + port: number; + portFile: string; + ppidPortFile: string; + authToken: string; + log: (message: string) => void; +} + +async function writePortAndWorkspace({ + context, + port, + portFile, + ppidPortFile, + authToken, + log, +}: WritePortAndWorkspaceArgs): Promise { const workspaceFolders = vscode.workspace.workspaceFolders; const workspacePath = workspaceFolders && workspaceFolders.length > 0 @@ -49,15 +59,22 @@ async function writePortAndWorkspace( workspacePath, ); - const content = JSON.stringify({ port, workspacePath, ppid: process.ppid }); + const content = JSON.stringify({ + port, + workspacePath, + ppid: process.ppid, + authToken, + }); log(`Writing port file to: ${portFile}`); log(`Writing ppid port file to: ${ppidPortFile}`); try { await Promise.all([ - fs.writeFile(portFile, content), - fs.writeFile(ppidPortFile, content), + fs.writeFile(portFile, content).then(() => fs.chmod(portFile, 0o600)), + fs + .writeFile(ppidPortFile, content) + .then(() => fs.chmod(ppidPortFile, 0o600)), ]); } catch (err) { const message = err instanceof Error ? err.message : String(err); @@ -95,6 +112,7 @@ export class IDEServer { private portFile: string | undefined; private ppidPortFile: string | undefined; private port: number | undefined; + private authToken: string | undefined; private transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; private openFilesManager: OpenFilesManager | undefined; @@ -108,10 +126,30 @@ export class IDEServer { start(context: vscode.ExtensionContext): Promise { return new Promise((resolve) => { this.context = context; + this.authToken = randomUUID(); const sessionsWithInitialNotification = new Set(); const app = express(); app.use(express.json({ limit: '10mb' })); + app.use((req, res, next) => { + const authHeader = req.headers.authorization; + if (authHeader) { + const parts = authHeader.split(' '); + if (parts.length !== 2 || parts[0] !== 'Bearer') { + this.log('Malformed Authorization header. Rejecting request.'); + res.status(401).send('Unauthorized'); + return; + } + const token = parts[1]; + if (token !== this.authToken) { + this.log('Invalid auth token provided. Rejecting request.'); + res.status(401).send('Unauthorized'); + return; + } + } + next(); + }); + const mcpServer = createMcpServer(this.diffManager); this.openFilesManager = new OpenFilesManager(context); @@ -250,13 +288,16 @@ export class IDEServer { ); this.log(`IDE server listening on port ${this.port}`); - await writePortAndWorkspace( - context, - this.port, - this.portFile, - this.ppidPortFile, - this.log, - ); + if (this.authToken) { + await writePortAndWorkspace({ + context, + port: this.port, + portFile: this.portFile, + ppidPortFile: this.ppidPortFile, + authToken: this.authToken, + log: this.log, + }); + } } resolve(); }); @@ -282,15 +323,17 @@ export class IDEServer { this.server && this.port && this.portFile && - this.ppidPortFile + this.ppidPortFile && + this.authToken ) { - await writePortAndWorkspace( - this.context, - this.port, - this.portFile, - this.ppidPortFile, - this.log, - ); + await writePortAndWorkspace({ + context: this.context, + port: this.port, + portFile: this.portFile, + ppidPortFile: this.ppidPortFile, + authToken: this.authToken, + log: this.log, + }); this.broadcastIdeContextUpdate(); } }