From f3c70a45ebc61c48e125e73b286724fa248f9792 Mon Sep 17 00:00:00 2001
From: QxBytes <39818795+QxBytes@users.noreply.github.com>
Date: Wed, 1 May 2024 13:41:52 -0700
Subject: [PATCH] backport: fix: close connection on any return from
 azure-vnet-telemetry (#2711) (#2719)

fix: close connection on any return from azure-vnet-telemetry (#2711)

* move error handling code to defer

* invert conditional

* address linter issues

* add close connection on telemetry unmarshal error unit test

* Revert "add close connection on telemetry unmarshal error unit test"

This reverts commit e8447f9d0d2370b5c742697ffa25d1532eedbac2.
---
 telemetry/telemetrybuffer.go | 75 ++++++++++++++++++++----------------
 1 file changed, 41 insertions(+), 34 deletions(-)

diff --git a/telemetry/telemetrybuffer.go b/telemetry/telemetrybuffer.go
index 83db69d5f0..0ea13f4e0a 100644
--- a/telemetry/telemetrybuffer.go
+++ b/telemetry/telemetrybuffer.go
@@ -103,47 +103,54 @@ func (tb *TelemetryBuffer) StartServer() error {
 				tb.connections = append(tb.connections, conn)
 				tb.mutex.Unlock()
 				go func() {
+					defer func() {
+						var index int
+						var value net.Conn
+						var found bool
+
+						tb.mutex.Lock()
+						defer tb.mutex.Unlock()
+
+						for index, value = range tb.connections {
+							if value == conn {
+								conn.Close()
+								found = true
+								break
+							}
+						}
+
+						if found {
+							tb.connections = remove(tb.connections, index)
+						}
+					}()
+
 					for {
 						reportStr, err := read(conn)
-						if err == nil {
-							var tmp map[string]interface{}
-							err = json.Unmarshal(reportStr, &tmp)
+						if err != nil {
+							return
+						}
+						var tmp map[string]interface{}
+						err = json.Unmarshal(reportStr, &tmp)
+						if err != nil {
+							log.Logf("StartServer: unmarshal error:%v", err)
+							return
+						}
+						if _, ok := tmp["CniSucceeded"]; ok {
+							var cniReport CNIReport
+							err = json.Unmarshal([]byte(reportStr), &cniReport)
 							if err != nil {
-								log.Logf("StartServer: unmarshal error:%v", err)
 								return
 							}
-							if _, ok := tmp["CniSucceeded"]; ok {
-								var cniReport CNIReport
-								json.Unmarshal([]byte(reportStr), &cniReport)
-								tb.data <- cniReport
-							} else if _, ok := tmp["Metric"]; ok {
-								var aiMetric AIMetric
-								json.Unmarshal([]byte(reportStr), &aiMetric)
-								tb.data <- aiMetric
-							} else {
-								log.Logf("StartServer: default case:%+v...", tmp)
+							tb.data <- cniReport
+						} else if _, ok := tmp["Metric"]; ok {
+							var aiMetric AIMetric
+							err = json.Unmarshal([]byte(reportStr), &aiMetric)
+							if err != nil {
+								return
 							}
+							tb.data <- aiMetric
 						} else {
-							var index int
-							var value net.Conn
-							var found bool
-
-							tb.mutex.Lock()
-							defer tb.mutex.Unlock()
-
-							for index, value = range tb.connections {
-								if value == conn {
-									conn.Close()
-									found = true
-									break
-								}
-							}
-
-							if found {
-								tb.connections = remove(tb.connections, index)
-							}
-
-							return
+							log.Logf("StartServer: default case:%+v...", tmp)
 						}
 					}
 				}()