From f0c9f1289790d003eb0e2dee18beaa78dd5f7c96 Mon Sep 17 00:00:00 2001
From: Vikrant Gupta <vikrant.thomso@gmail.com>
Date: Tue, 20 Aug 2024 13:17:56 +0530
Subject: [PATCH] fix: do not use relative URLs for ws connections (#5715)

* fix: do not use relative URLs for ws connections

* fix: handle local env better

* chore: update the websocket endpoint

* chore: handle OSS/Docker installations
---
 frontend/src/api/common/getQueryStats.ts      | 25 ++++++++++++++++---
 .../tests/LogsExplorerViews.test.tsx          |  4 ++-
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/frontend/src/api/common/getQueryStats.ts b/frontend/src/api/common/getQueryStats.ts
index 7277a32852..b70c4c722a 100644
--- a/frontend/src/api/common/getQueryStats.ts
+++ b/frontend/src/api/common/getQueryStats.ts
@@ -1,6 +1,7 @@
 import getLocalStorageApi from 'api/browser/localstorage/get';
 import { ENVIRONMENT } from 'constants/env';
 import { LOCALSTORAGE } from 'constants/localStorage';
+import { isEmpty } from 'lodash-es';
 
 export interface WsDataEvent {
 	read_rows: number;
@@ -12,14 +13,30 @@ interface GetQueryStatsProps {
 	setData: React.Dispatch<React.SetStateAction<WsDataEvent | undefined>>;
 }
 
+function getURL(baseURL: string, queryId: string): URL | string {
+	if (baseURL && !isEmpty(baseURL)) {
+		return `${baseURL}/ws/query_progress?q=${queryId}`;
+	}
+	const url = new URL(`/ws/query_progress?q=${queryId}`, window.location.href);
+
+	if (window.location.protocol === 'http:') {
+		url.protocol = 'ws';
+	} else {
+		url.protocol = 'wss';
+	}
+
+	return url;
+}
+
 export function getQueryStats(props: GetQueryStatsProps): void {
 	const { queryId, setData } = props;
 
 	const token = getLocalStorageApi(LOCALSTORAGE.AUTH_TOKEN) || '';
-	const socket = new WebSocket(
-		`${ENVIRONMENT.wsURL}/api/v3/query_progress?q=${queryId}`,
-		token,
-	);
+
+	// https://github.com/whatwg/websockets/issues/20 reason for not using the relative URLs
+	const url = getURL(ENVIRONMENT.wsURL, queryId);
+
+	const socket = new WebSocket(url, token);
 
 	socket.addEventListener('message', (event) => {
 		try {
diff --git a/frontend/src/container/LogsExplorerViews/tests/LogsExplorerViews.test.tsx b/frontend/src/container/LogsExplorerViews/tests/LogsExplorerViews.test.tsx
index 67c3fcfe2c..8262d6f9bc 100644
--- a/frontend/src/container/LogsExplorerViews/tests/LogsExplorerViews.test.tsx
+++ b/frontend/src/container/LogsExplorerViews/tests/LogsExplorerViews.test.tsx
@@ -46,7 +46,9 @@ jest.mock(
 		},
 );
 
-jest.mock('api/common/getQueryStats', () => jest.fn());
+jest.mock('api/common/getQueryStats', () => ({
+	getQueryStats: jest.fn(),
+}));
 
 jest.mock('constants/panelTypes', () => ({
 	AVAILABLE_EXPORT_PANEL_TYPES: ['graph', 'table'],