feat(vscode-ide-companion): add auth token validation to IDE server (#8491)

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
 

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');
+      }
+    });
+  });
 });

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(
+interface WritePortAndWorkspaceArgs {
-  context: vscode.ExtensionContext,
+  context: vscode.ExtensionContext;
-  port: number,
+  port: number;
-  portFile: string,
+  portFile: string;
-  ppidPortFile: string,
+  ppidPortFile: string;
-  log: (message: string) => void,
+  authToken: string;
-): Promise<void> {
+  log: (message: string) => void;
+}
+
+async function writePortAndWorkspace({
+  context,
+  port,
+  portFile,
+  ppidPortFile,
+  authToken,
+  log,
+}: WritePortAndWorkspaceArgs): Promise<void> {
   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(portFile, content).then(() => fs.chmod(portFile, 0o600)),
-      fs.writeFile(ppidPortFile, content),
+      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<void> {
     return new Promise((resolve) => {
       this.context = context;
+      this.authToken = randomUUID();
       const sessionsWithInitialNotification = new Set<string>();
 
       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(
+          if (this.authToken) {
-            context,
+            await writePortAndWorkspace({
-            this.port,
+              context,
-            this.portFile,
+              port: this.port,
-            this.ppidPortFile,
+              portFile: this.portFile,
-            this.log,
+              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(
+      await writePortAndWorkspace({
-        this.context,
+        context: this.context,
-        this.port,
+        port: this.port,
-        this.portFile,
+        portFile: this.portFile,
-        this.ppidPortFile,
+        ppidPortFile: this.ppidPortFile,
-        this.log,
+        authToken: this.authToken,
-      );
+        log: this.log,
+      });
       this.broadcastIdeContextUpdate();
     }
   }