Merge pull request #314 from rita7lopes/master

Add ping to authenticate

Author: adubovikov

publish/hep.go

@@ -28,6 +28,21 @@ type HEPOutputer struct {
 	msgPing  []byte
 }
 
+func writeAndFlush(client *HEPConn, data []byte, action string) (int, error) {
+	hl, err := client.conn.Write(data)
+	if err != nil {
+		promstats.HepFileFlushesError.Inc()
+		return 0, fmt.Errorf("error writing to socket during %s: %w", action, err)
+	}
+
+	if err := client.writer.Flush(); err != nil {
+		promstats.HepFileFlushesError.Inc()
+		return 0, fmt.Errorf("error flushing writer during %s: %w", action, err)
+	}
+
+	return hl, nil
+}
+
 func NewHEPOutputer(serverAddr string) (*HEPOutputer, error) {
 	a := strings.Split(cutSpace(serverAddr), ",")
 	l := len(a)
@@ -43,6 +58,11 @@ func NewHEPOutputer(serverAddr string) (*HEPOutputer, error) {
 			errCnt++
 		} else {
 			if config.Cfg.HEPBufferEnable {
+				logp.Debug("collector", "send ping packet after disconnect")
+				_, err := writeAndFlush(&h.client[n], h.msgPing, "ping operation establish connection")
+				if err != nil {
+					return nil, err
+				}
 				if _, err := os.Stat(config.Cfg.HEPBufferFile); err == nil {
 					if _, err := h.copyHEPFileOut(n); err != nil {
 						logp.Err("Sending HEP from file error: %v", err)
@@ -161,8 +181,10 @@ func (h *HEPOutputer) Send(msg []byte) {
 			logp.Debug("Connection is not up", fmt.Sprintf("index: %d, Len: %d, once: %v", n, len(h.addr), onceSent))
 			err = fmt.Errorf("connection is broken")
 		} else {
-			h.client[n].writer.Write(msg)
-			err = h.client[n].writer.Flush()
+			_, err := writeAndFlush(&h.client[n], msg, "sending message")
+			if err != nil {
+				logp.Err("Failed to send message: %s", err.Error())
+			}
 		}
 
 		if err != nil {
@@ -183,17 +205,15 @@ func (h *HEPOutputer) Send(msg []byte) {
 				} else {
 					if h.msgPing != nil {
 						logp.Debug("collector", "send ping packet after disconnect")
-						h.client[n].writer.Write(h.msgPing)
-						err = h.client[n].writer.Flush()
+						_, err := writeAndFlush(&h.client[n], h.msgPing, "Error during resend ping packet")
 						if err != nil {
-							logp.Err("Bad during resend ping packet : %v", err)
+							logp.Err("Error sending ping packet: %s", err.Error())
 						}
 					}
 
-					h.client[n].writer.Write(msg)
-					err = h.client[n].writer.Flush()
+					_, err := writeAndFlush(&h.client[n], msg, "Bad resend")
 					if err != nil {
-						logp.Err("Bad resend: %v", err)
+						logp.Err(err.Error())
 						if config.Cfg.HEPBufferEnable && (!onceSent && n == (len(h.addr)-1)) {
 							h.copyHEPbufftoFile(msg)
 						}
@@ -241,11 +261,16 @@ func (h *HEPOutputer) copyHEPFileOut(n int) (int, error) {
 		return 0, fmt.Errorf("Connection is broken")
 	}
 
-	//Send Logged HEP upon reconnect out to backend
-	hl, err := h.client[n].conn.Write(HEPFileData)
+	hl, err := writeAndFlush(&h.client[n], h.msgPing, "ping operation")
+	if err != nil {
+		return 0, err
+	}
+
+	// Send Logged HEP upon reconnect
+	hl, err = writeAndFlush(&h.client[n], HEPFileData, "HEP reconnect")
 	if err != nil {
 		promstats.HepFileFlushesError.Inc()
-		return 0, fmt.Errorf("Bad write to socket")
+		return 0, err
 	}
 
 	err = h.client[n].writer.Flush()

sniffer/sniffer.go

@@ -374,6 +374,12 @@ LOOP:
 					conn, err := sniffer.collectorTCPconn.Accept()
 					if err != nil {
 						logp.Err("Error accepting tcp connection: %s", err.Error())
+						// Close connection if it exists
+						if conn != nil {
+							conn.Close()
+						}
+						// Add delay before next attempt
+						time.Sleep(time.Second)
 						continue
 					}
 
@@ -657,6 +663,9 @@ func (sniffer *SnifferSetup) handleRequestSimple(conn net.Conn) {
 		_, err := conn.Read(message)
 		if err != nil {
 			fmt.Println("Error reading:", err.Error())
+			if conn != nil {
+				conn.Close() // Ensure the connection is closed
+			}
 			break
 		}
 
@@ -674,6 +683,9 @@ func (sniffer *SnifferSetup) handleRequestSimple(conn net.Conn) {
 			_, err := conn.Read(data)
 			if err != nil {
 				fmt.Println("Error reading:", err.Error())
+				if conn != nil {
+					conn.Close() // Ensure the connection is closed
+				}
 				break
 			}
 
@@ -715,6 +727,9 @@ func (sniffer *SnifferSetup) handleRequestExtended(conn net.Conn) {
 		n, err := conn.Read(message)
 		if err != nil {
 			logp.Err("Incoming tcp connection closed during read with error [1]: %s", err.Error())
+			if conn != nil {
+				conn.Close() // Ensure the connection is closed
+			}
 			break
 		}
 
@@ -778,6 +793,9 @@ func (sniffer *SnifferSetup) handleRequestExtended(conn net.Conn) {
 						if err != nil {
 							logp.Err("Incoming tcp connection closed during direct read from buffer with error [2]: %s", err.Error())
 							bufferPool.Reset()
+							if conn != nil {
+								conn.Close() // Ensure the connection is closed
+							}
 							break
 						}