From 14bc26e065dc572b839c6cbd44b1a6ed26d80613 Mon Sep 17 00:00:00 2001
From: Mauricio Siu <47042324+Siumauricio@users.noreply.github.com>
Date: Sun, 6 Apr 2025 00:07:41 -0600
Subject: [PATCH] feat(websocket): enhance WebSocket server with request
 validation and client instantiation

- Added request validation to ensure user authentication before establishing WebSocket connections.
- Refactored WebSocket client instantiation to simplify connection management.
---
 apps/dokploy/server/wss/drawer-logs.ts | 14 +++++++++----
 apps/dokploy/utils/api.ts              | 27 +++++++-------------------
 2 files changed, 17 insertions(+), 24 deletions(-)

diff --git a/apps/dokploy/server/wss/drawer-logs.ts b/apps/dokploy/server/wss/drawer-logs.ts
index dcdeaad7d..0202ae521 100644
--- a/apps/dokploy/server/wss/drawer-logs.ts
+++ b/apps/dokploy/server/wss/drawer-logs.ts
@@ -3,6 +3,7 @@ import { applyWSSHandler } from "@trpc/server/adapters/ws";
 import { WebSocketServer } from "ws";
 import { appRouter } from "../api/root";
 import { createTRPCContext } from "../api/trpc";
+import { validateRequest } from "@dokploy/server/lib/auth";
 
 export const setupDrawerLogsWebSocketServer = (
 	server: http.Server<typeof http.IncomingMessage, typeof http.ServerResponse>,
@@ -32,8 +33,13 @@ export const setupDrawerLogsWebSocketServer = (
 		}
 	});
 
-	// Return cleanup function
-	return () => {
-		wssTerm.close();
-	};
+	wssTerm.on("connection", async (ws, req) => {
+		const _url = new URL(req.url || "", `http://${req.headers.host}`);
+		const { user, session } = await validateRequest(req);
+
+		if (!user || !session) {
+			ws.close();
+			return;
+		}
+	});
 };
diff --git a/apps/dokploy/utils/api.ts b/apps/dokploy/utils/api.ts
index 7c003f481..56197528b 100644
--- a/apps/dokploy/utils/api.ts
+++ b/apps/dokploy/utils/api.ts
@@ -27,28 +27,15 @@ const getWsUrl = () => {
 	const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
 	const host = window.location.host;
 
-	// Use the base URL for all tRPC WebSocket connections
 	return `${protocol}${host}/drawer-logs`;
 };
 
-// Singleton WebSocket client instance
-let wsClientInstance: ReturnType<typeof createWSClient> | null = null;
-
-const getWsClient = () => {
-	if (typeof window === "undefined") return null;
-
-	if (!wsClientInstance) {
-		wsClientInstance = createWSClient({
-			url: getWsUrl() || "",
-			onClose: () => {
-				// Reset the instance when connection closes so it can be recreated
-				wsClientInstance = null;
-			},
-		});
-	}
-
-	return wsClientInstance;
-};
+const wsClient =
+	typeof window !== "undefined"
+		? createWSClient({
+				url: getWsUrl() || "",
+			})
+		: null;
 
 /** A set of type-safe react-query hooks for your tRPC API. */
 export const api = createTRPCNext<AppRouter>({
@@ -70,7 +57,7 @@ export const api = createTRPCNext<AppRouter>({
 				splitLink({
 					condition: (op) => op.type === "subscription",
 					true: wsLink({
-						client: getWsClient()!,
+						client: wsClient!,
 					}),
 					false: splitLink({
 						condition: (op) => op.input instanceof FormData,