From 2da7c0f0ed245a26e5f284828dd1b840329bb927 Mon Sep 17 00:00:00 2001
From: RamiBerm <54766858+RamiBerm@users.noreply.github.com>
Date: Mon, 27 Dec 2021 15:11:36 +0200
Subject: [PATCH] TRA-4065 fix mizu tapping (#561)

* Update main.go and passive_tapper.go

* Update main.go
---
 agent/main.go         | 18 ++++++++++++++----
 tap/passive_tapper.go |  3 ++-
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/agent/main.go b/agent/main.go
index 6f1df5b3c..06812d4e0 100644
--- a/agent/main.go
+++ b/agent/main.go
@@ -441,11 +441,21 @@ func handleIncomingMessageAsTapper(socketConnection *websocket.Conn) {
 				return
 			}
 		} else {
-			var tapConfigMessage *shared.WebSocketTapConfigMessage
-			if err := json.Unmarshal(message, &tapConfigMessage); err != nil {
-				logger.Log.Errorf("received unknown message from socket connection: %s, err: %s, (%v,%+v)", string(message), err, err, err)
+			var socketMessageBase shared.WebSocketMessageMetadata
+			if err := json.Unmarshal(message, &socketMessageBase); err != nil {
+				logger.Log.Errorf("Could not unmarshal websocket message %v", err)
 			} else {
-				tap.UpdateTapTargets(tapConfigMessage.TapTargets)
+				switch socketMessageBase.MessageType {
+				case shared.WebSocketMessageTypeTapConfig:
+					var tapConfigMessage *shared.WebSocketTapConfigMessage
+					if err := json.Unmarshal(message, &tapConfigMessage); err != nil {
+						logger.Log.Errorf("received unknown message from socket connection: %s, err: %s, (%v,%+v)", string(message), err, err, err)
+					} else {
+						tap.UpdateTapTargets(tapConfigMessage.TapTargets)
+					}
+				default:
+					logger.Log.Warningf("Received socket message of type %s for which no handlers are defined", socketMessageBase.MessageType)
+				}
 			}
 		}
 	}
diff --git a/tap/passive_tapper.go b/tap/passive_tapper.go
index 95f4e1c4d..fcf0a9eda 100644
--- a/tap/passive_tapper.go
+++ b/tap/passive_tapper.go
@@ -194,11 +194,12 @@ func startPassiveTapper(opts *TapOpts, outputItems chan *api.OutputChannelItem)
 	diagnose.InitializeErrorsMap(*debug, *verbose, *quiet)
 	diagnose.InitializeTapperInternalStats()
 
+	mainPacketInputChan = make(chan source.TcpPacketInfo)
+
 	if err := initializePacketSources(); err != nil {
 		logger.Log.Fatal(err)
 	}
 
-	mainPacketInputChan = make(chan source.TcpPacketInfo)
 	assembler := NewTcpAssembler(outputItems, streamsMap, opts)
 
 	diagnose.AppStats.SetStartTime(time.Now())