Websocket auth, pass access token in gke configs

deploy/gcp/Makefile

@@ -15,8 +15,8 @@ APP_IMAGE_TAG=$(GCP_REPO)/$(GCP_PROJECT_ID)/$(PREFIX)-app:latest
 CLUSTER_NAME=$(PREFIX)-cluster
 
 
-.PHONY: build-deplpy
-build-deplpy: docker-build-all deploy
+.PHONY: build-deploy
+build-deploy: docker-build-all deploy
 
 
 .PHONY: docker-build-all
@@ -34,7 +34,8 @@ show_config:
 	&& echo $(APP_IMAGE_TAG) \
 	&& echo $(SEARCH_PORT) \
 	&& echo $(BACKEND_PORT) \
-	&& echo $(OPENAI)
+	&& echo $(OPENAI) \
+	&& echo $(SUPER_SECRET_KEY)
 
 .PHONY: docker-build-push-searxng
 docker-build-push-searxng:
@@ -72,6 +73,7 @@ deploy:
 	&& export TF_VAR_search_port=$(SEARCH_PORT) \
 	&& export TF_VAR_backend_port=$(BACKEND_PORT) \
 	&& export TF_VAR_open_ai=$(OPENAI) \
+	&& export TF_VAR_secret_key=$(SUPER_SECRET_KEY) \
 	&& terraform apply
 
 

deploy/gcp/main.tf

@@ -136,6 +136,11 @@ resource "kubernetes_deployment" "backend" {
           name = "PORT"
           value = var.backend_port 
         }
+        env {
+          # Access key for backend
+          name = "SUPER_SECRET_KEY"
+          value = var.secret_key
+        }
         }
       }
     }
@@ -205,6 +210,11 @@ variable "open_ai" {
   type        = string
 }
 
+variable "secret_key" {
+  description = "Access key to secure backend endpoints"
+  type        = string
+}
+
 variable "search_port" {
   description = "Port for searxng service"
   type        = number

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(),

ui/components/ChatWindow.tsx

@@ -7,6 +7,7 @@ import Chat from './Chat';
 import EmptyChat from './EmptyChat';
 import { toast } from 'sonner';
 import { clientFetch } from '@/lib/utils';
+import { getAccessKey } from '@/lib/config';
 
 export type Message = {
   id: string;
@@ -98,7 +99,14 @@ const useSocket = (url: string) => {
 
         wsURL.search = searchParams.toString();
 
-        const ws = new WebSocket(wsURL.toString());
+        let protocols: any[] = [];
+        const secretToken = getAccessKey();
+
+        if (secretToken) {
+          protocols = ["Authorization", `${secretToken}`];
+        };
+
+        const ws = new WebSocket(wsURL.toString(), protocols);
 
         ws.onopen = () => {
           console.log('[DEBUG] open');

ui/lib/utils.ts

@@ -24,14 +24,14 @@ export const formatTimeDifference = (date1: Date, date2: Date): string => {
 export const clientFetch = async (path: string, payload: any): Promise<any> => {
   let headers = payload.headers;
   const url = `${getBackendURL()}${path}`;
-  const secret_token = getAccessKey();
+  const secretToken = getAccessKey();
 
-  if (secret_token) {
+  if (secretToken) {
     if (headers == null) {
       headers = {};
     };
 
-    headers['Authorization'] = `Bearer ${secret_token}`;
+    headers['Authorization'] = `Bearer ${secretToken}`;
     payload.headers = headers;
   };