Websocket auth, pass access token in gke configs

src/auth.ts

@@ -6,14 +6,16 @@ export const requireAccessKey = (req, res, next) => {
     const authHeader = req.headers.authorization;
 
     if (authHeader) {
-        const token = authHeader.split(' ')[1];
-
-        if (token !== getAccessKey()) {
+        if (!checkAccessKey(authHeader)) {
             return res.sendStatus(403);
         }
-
         next();
     } else {
         res.sendStatus(401);
     }
 };
+
+export const checkAccessKey = (authHeader) => {
+      const token = authHeader.split(' ')[1];
+      return Boolean(authHeader && (token === getAccessKey()));
+};

src/websocket/connectionManager.ts

@@ -9,6 +9,8 @@ import type { Embeddings } from '@langchain/core/embeddings';
 import type { IncomingMessage } from 'http';
 import logger from '../utils/logger';
 import { ChatOpenAI } from '@langchain/openai';
+import { getAccessKey } from '../config';
+import { checkAccessKey } from '../auth';
 
 export const handleConnection = async (
   ws: WebSocket,
@@ -18,6 +20,20 @@ export const handleConnection = async (
     const searchParams = new URL(request.url, `http://${request.headers.host}`)
       .searchParams;
 
+    if (getAccessKey()) {
+      const securtyProtocolHeader = request.headers['sec-websocket-protocol'];
+      if (!checkAccessKey(securtyProtocolHeader)) {
+        ws.send(
+          JSON.stringify({
+            type: 'error',
+            data: 'Incorrect or missing authentication token.',
+            key: 'FAILED_AUTHORIZATION',
+          }),
+        );
+        ws.close();
+      };
+    }
+
     const [chatModelProviders, embeddingModelProviders] = await Promise.all([
       getAvailableChatModelProviders(),
       getAvailableEmbeddingModelProviders(),