TRA-4331 Implement full data streaming over WebSocket (#819)

* Implement full data streaming over WebSocket

* Fix the linting error

* Make the empty being the criteria

* Use a label to break the nested loop

agent/pkg/api/socket_routes.go

@@ -110,20 +110,31 @@ func websocketHandler(w http.ResponseWriter, r *http.Request, eventHandlers Even
 		logger.Log.Error(err)
 	}
 
+out:
 	for {
-		_, 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)
+		// 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
 			}
+		}
 
-			break
+		enableFullEntries := false
+		if len(params[1]) > 0 {
+			enableFullEntries = true
 		}
 
 		if !isTapper && !isQuerySet {
-			query := string(msg)
+			query := string(params[0])
 			err = basenine.Validate(shared.BasenineHost, shared.BaseninePort, query)
 			if err != nil {
 				toastBytes, _ := models.CreateWebsocketToastMessage(&models.ToastMessage{
@@ -150,10 +161,15 @@ func websocketHandler(w http.ResponseWriter, r *http.Request, eventHandlers Even
 					var entry *tapApi.Entry
 					err = json.Unmarshal(bytes, &entry)
 
-					base := tapApi.Summarize(entry)
+					var message []byte
+					if enableFullEntries {
+						message, _ = models.CreateFullEntryWebSocketMessage(entry)
+					} else {
+						base := tapApi.Summarize(entry)
+						message, _ = models.CreateBaseEntryWebSocketMessage(base)
+					}
 
-					baseEntryBytes, _ := models.CreateBaseEntryWebSocketMessage(base)
-					if err := SendToSocket(socketId, baseEntryBytes); err != nil {
+					if err := SendToSocket(socketId, message); err != nil {
 						logger.Log.Error(err)
 					}
 				}
@@ -185,7 +201,7 @@ func websocketHandler(w http.ResponseWriter, r *http.Request, eventHandlers Even
 
 			connection.Query(query, data, meta)
 		} else {
-			eventHandlers.WebSocketMessage(socketId, msg)
+			eventHandlers.WebSocketMessage(socketId, params[0])
 		}
 	}
 }

agent/pkg/models/models.go

@@ -42,6 +42,11 @@ type WebSocketEntryMessage struct {
 	Data *tapApi.BaseEntry `json:"data,omitempty"`
 }
 
+type WebSocketFullEntryMessage struct {
+	*shared.WebSocketMessageMetadata
+	Data *tapApi.Entry `json:"data,omitempty"`
+}
+
 type WebSocketTappedEntryMessage struct {
 	*shared.WebSocketMessageMetadata
 	Data *tapApi.OutputChannelItem
@@ -88,6 +93,16 @@ func CreateBaseEntryWebSocketMessage(base *tapApi.BaseEntry) ([]byte, error) {
 	return json.Marshal(message)
 }
 
+func CreateFullEntryWebSocketMessage(entry *tapApi.Entry) ([]byte, error) {
+	message := &WebSocketFullEntryMessage{
+		WebSocketMessageMetadata: &shared.WebSocketMessageMetadata{
+			MessageType: shared.WebSocketMessageTypeFullEntry,
+		},
+		Data: entry,
+	}
+	return json.Marshal(message)
+}
+
 func CreateWebsocketTappedEntryMessage(base *tapApi.OutputChannelItem) ([]byte, error) {
 	message := &WebSocketTappedEntryMessage{
 		WebSocketMessageMetadata: &shared.WebSocketMessageMetadata{

shared/models.go

@@ -15,6 +15,7 @@ type WebSocketMessageType string
 
 const (
 	WebSocketMessageTypeEntry         WebSocketMessageType = "entry"
+	WebSocketMessageTypeFullEntry     WebSocketMessageType = "fullEntry"
 	WebSocketMessageTypeTappedEntry   WebSocketMessageType = "tappedEntry"
 	WebSocketMessageTypeUpdateStatus  WebSocketMessageType = "status"
 	WebSocketMessageTypeAnalyzeStatus WebSocketMessageType = "analyzeStatus"

ui/src/components/Pages/TrafficPage/TrafficPage.tsx

@@ -122,6 +122,7 @@ export const TrafficPage: React.FC<TrafficPageProps> = ({setAnalyzeStatus}) => {
         ws.current.onopen = () => {
             setWsConnection(WsConnectionStatus.Connected);
             ws.current.send(query);
+            ws.current.send("");
         }
         ws.current.onclose = () => {
             setWsConnection(WsConnectionStatus.Closed);