From 73f3e448cf90faee0be4117ae6772524de70c690 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=2E=20Mert=20Y=C4=B1ld=C4=B1ran?= <mehmet@up9.com>
Date: Tue, 22 Feb 2022 15:27:50 +0300
Subject: [PATCH] Fix the issue in the WebSocket that's introduced by #819
 (#839)

---
 agent/pkg/api/socket_routes.go                | 42 +++++++++----------
 .../Pages/TrafficPage/TrafficPage.tsx         |  3 +-
 2 files changed, 22 insertions(+), 23 deletions(-)

diff --git a/agent/pkg/api/socket_routes.go b/agent/pkg/api/socket_routes.go
index ad6983bcd..f3d810395 100644
--- a/agent/pkg/api/socket_routes.go
+++ b/agent/pkg/api/socket_routes.go
@@ -30,6 +30,11 @@ type SocketConnection struct {
 	isTapper      bool
 }
 
+type WebSocketParams struct {
+	Query             string `json:"query"`
+	EnableFullEntries bool   `json:"enableFullEntries"`
+}
+
 var (
 	websocketUpgrader = websocket.Upgrader{
 		ReadBufferSize:  1024,
@@ -110,31 +115,26 @@ func websocketHandler(w http.ResponseWriter, r *http.Request, eventHandlers Even
 		logger.Log.Error(err)
 	}
 
-out:
+	var params WebSocketParams
+
 	for {
-		// params[0]: query
-		// params[1]: enableFullEntries (empty: disable, non-empty: enable)
-		params := make([][]byte, 2)
-		for i := range params {
-			_, params[i], err = ws.ReadMessage()
-			if err != nil {
-				if _, ok := err.(*websocket.CloseError); ok {
-					logger.Log.Debugf("Received websocket close message, socket id: %d", socketId)
-				} else {
-					logger.Log.Errorf("Error reading message, socket id: %d, error: %v", socketId, err)
-				}
-
-				break out
+		_, msg, err := ws.ReadMessage()
+		if err != nil {
+			if _, ok := err.(*websocket.CloseError); ok {
+				logger.Log.Debugf("Received websocket close message, socket id: %d", socketId)
+			} else {
+				logger.Log.Errorf("Error reading message, socket id: %d, error: %v", socketId, err)
 			}
-		}
 
-		enableFullEntries := false
-		if len(params[1]) > 0 {
-			enableFullEntries = true
+			break
 		}
 
 		if !isTapper && !isQuerySet {
-			query := string(params[0])
+			if err := json.Unmarshal(msg, &params); err != nil {
+				logger.Log.Errorf("Error: %v", socketId, err)
+			}
+
+			query := params.Query
 			err = basenine.Validate(shared.BasenineHost, shared.BaseninePort, query)
 			if err != nil {
 				toastBytes, _ := models.CreateWebsocketToastMessage(&models.ToastMessage{
@@ -162,7 +162,7 @@ out:
 					err = json.Unmarshal(bytes, &entry)
 
 					var message []byte
-					if enableFullEntries {
+					if params.EnableFullEntries {
 						message, _ = models.CreateFullEntryWebSocketMessage(entry)
 					} else {
 						base := tapApi.Summarize(entry)
@@ -201,7 +201,7 @@ out:
 
 			connection.Query(query, data, meta)
 		} else {
-			eventHandlers.WebSocketMessage(socketId, params[0])
+			eventHandlers.WebSocketMessage(socketId, msg)
 		}
 	}
 }
diff --git a/ui/src/components/Pages/TrafficPage/TrafficPage.tsx b/ui/src/components/Pages/TrafficPage/TrafficPage.tsx
index c8732378f..2b71bed5f 100644
--- a/ui/src/components/Pages/TrafficPage/TrafficPage.tsx
+++ b/ui/src/components/Pages/TrafficPage/TrafficPage.tsx
@@ -121,8 +121,7 @@ export const TrafficPage: React.FC<TrafficPageProps> = ({setAnalyzeStatus}) => {
         ws.current = new WebSocket(MizuWebsocketURL);
         ws.current.onopen = () => {
             setWsConnection(WsConnectionStatus.Connected);
-            ws.current.send(query);
-            ws.current.send("");
+            ws.current.send(JSON.stringify({"query": query, "enableFullEntries": false}));
         }
         ws.current.onclose = () => {
             setWsConnection(WsConnectionStatus.Closed);