From c56a058a740cbb2a2a276e2b4a4399b4878dcad8 Mon Sep 17 00:00:00 2001
From: Hristo <53634432+izo0x90@users.noreply.github.com>
Date: Fri, 10 May 2024 19:32:35 -0400
Subject: [PATCH] Websocket auth, pass access token in gke configs

---
 deploy/gcp/Makefile                |  8 +++++---
 deploy/gcp/main.tf                 | 10 ++++++++++
 src/auth.ts                        | 10 ++++++----
 src/websocket/connectionManager.ts | 16 ++++++++++++++++
 ui/components/ChatWindow.tsx       | 10 +++++++++-
 ui/lib/utils.ts                    |  6 +++---
 6 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/deploy/gcp/Makefile b/deploy/gcp/Makefile
index 9d959f2..35c9384 100644
--- a/deploy/gcp/Makefile
+++ b/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
 
 
diff --git a/deploy/gcp/main.tf b/deploy/gcp/main.tf
index 9d53397..ace5d90 100644
--- a/deploy/gcp/main.tf
+++ b/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
diff --git a/src/auth.ts b/src/auth.ts
index ecb88d4..c6f3cef 100644
--- a/src/auth.ts
+++ b/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()));
+};
diff --git a/src/websocket/connectionManager.ts b/src/websocket/connectionManager.ts
index 5cb075b..f504c10 100644
--- a/src/websocket/connectionManager.ts
+++ b/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(),
diff --git a/ui/components/ChatWindow.tsx b/ui/components/ChatWindow.tsx
index 87a8ad3..3a58b9f 100644
--- a/ui/components/ChatWindow.tsx
+++ b/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');
diff --git a/ui/lib/utils.ts b/ui/lib/utils.ts
index 3b25633..6d59e79 100644
--- a/ui/lib/utils.ts
+++ b/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;
   };