Merge pull request #235 from utahkanz-ops/feature/133-ws-auth

feat: add Bearer token auth to WebSocket server (#133)

Author: soomtochukwu

agent/__tests__/websocket-auth.test.ts

@@ -0,0 +1,59 @@
+import http from 'http';
+import { NotificationServer } from '../websocket-server.js';
+import WebSocket, { Server as WS } from 'ws';
+
+// Helper to create a mock HTTP server
+function createHttpServer(): http.Server {
+  const server = http.createServer();
+  server.listen(0);
+  return server;
+}
+
+describe('NotificationServer WebSocket authentication', () => {
+  const originalEnv = process.env;
+
+  beforeEach(() => {
+    jest.resetModules();
+    process.env = { ...originalEnv, GEMINI_API_KEY: 'test-secret-key' };
+  });
+
+  afterAll(() => {
+    process.env = originalEnv;
+  });
+
+  test('rejects connection without Authorization header', (done) => {
+    const httpServer = createHttpServer();
+    const ns = new NotificationServer(httpServer, { allowedOrigins: [] });
+    const wsPort = (httpServer.address() as any).port;
+    const ws = new WebSocket(`ws://localhost:${wsPort}?userId=user123`);
+    ws.on('error', (err) => {
+      // Server should close with code 4000 (custom close) -> ws error
+      expect(err).toBeDefined();
+      httpServer.close();
+      done();
+    });
+    ws.on('open', () => {
+      // Should not open
+    });
+  });
+
+  test('accepts connection with correct Bearer token', (done) => {
+    const httpServer = createHttpServer();
+    const ns = new NotificationServer(httpServer, { allowedOrigins: [] });
+    const wsPort = (httpServer.address() as any).port;
+    const ws = new WebSocket(`ws://localhost:${wsPort}?userId=user123`, {
+      headers: { Authorization: 'Bearer test-secret-key' },
+    });
+    ws.on('open', () => {
+      // Connection succeeded
+      ws.close();
+      httpServer.close();
+      done();
+    });
+    ws.on('error', (err) => {
+      // Should not error
+      httpServer.close();
+      done(err);
+    });
+  });
+});

agent/websocket-server.ts

@@ -18,6 +18,7 @@ import {
   DEFAULT_RATE_LIMIT_WINDOW_MS,
 } from "./websocket-limits.js";
 import { projectGoalStatus, type GoalProjection } from "./goal-projection.js";
+import { loadAgentEnv } from "./env.js";
 
 interface ActiveClient {
   ws: WebSocket;
@@ -87,7 +88,11 @@ export class NotificationServer {
   private readonly maxMessagesPerSecond: number;
   private readonly rateLimitWindowMs: number;
 
+  private readonly apiKey: string;
+
   constructor(httpServer: Server, options: NotificationServerOptions = {}) {
+    const env = loadAgentEnv();
+    this.apiKey = env.GEMINI_API_KEY;
     const {
       maxMessageBytes = DEFAULT_MAX_WS_MESSAGE_BYTES,
       allowedOrigins = [],
@@ -107,9 +112,16 @@ export class NotificationServer {
         secure: boolean;
         req: IncomingMessage;
       }) => {
-        // info.origin is empty string '' for non-browser clients in ws@8
         const origin = info.origin === "" ? undefined : info.origin;
-        return isOriginAllowed(origin, this.allowedOrigins);
+        if (!isOriginAllowed(origin, this.allowedOrigins)) {
+          return false;
+        }
+        const authHeader = info.req.headers["authorization"];
+        if (!authHeader || typeof authHeader !== "string" || !authHeader.startsWith("Bearer ")) {
+          return false;
+        }
+        const token = authHeader.slice(7).trim();
+        return token === this.apiKey;
       },
     });
     this.setupConnectionHandlers();