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.
This commit is contained in:
Mauricio Siu
2025-04-06 00:07:41 -06:00
parent 6c8eb3b711
commit 14bc26e065
2 changed files with 17 additions and 24 deletions

View File

@@ -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;
}
});
};

View File

@@ -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,