From c51e43f0bb5ab1b0cbab129725267cb646cbcfab Mon Sep 17 00:00:00 2001
From: val06 <valeri.pliskin@datadoghq.com>
Date: Tue, 26 Nov 2024 13:01:21 +0200
Subject: [PATCH] [USM] fix edge-case in protocols testutil (#31430)

---
 .../protocols/testutil/patternscanner.go      | 35 +++++++++++++------
 1 file changed, 25 insertions(+), 10 deletions(-)

diff --git a/pkg/network/protocols/testutil/patternscanner.go b/pkg/network/protocols/testutil/patternscanner.go
index 5b73e706d57c1d..ef95e75b24c1a6 100644
--- a/pkg/network/protocols/testutil/patternscanner.go
+++ b/pkg/network/protocols/testutil/patternscanner.go
@@ -26,6 +26,9 @@ type PatternScanner struct {
 
 	// keep the stdout/err in case of failure
 	buffers []string
+
+	//Buffer for accumulating partial lines
+	lineBuf string
 }
 
 // NewScanner returns a new instance of PatternScanner.
@@ -41,18 +44,30 @@ func NewScanner(pattern *regexp.Regexp, doneChan chan struct{}) *PatternScanner
 // Write implemented io.Writer to be used as a callback for log/string writing.
 // Once we find a match in for the given pattern, we notify the caller.
 func (ps *PatternScanner) Write(p []byte) (n int, err error) {
-	ps.buffers = append(ps.buffers, string(p))
-	n = len(p)
-	err = nil
-
-	if !ps.stopped && ps.pattern.Match(p) {
-		ps.stopOnce.Do(func() {
-			ps.DoneChan <- struct{}{}
-			ps.stopped = true
-		})
+	// Ignore writes after the pattern has been matched.
+	if ps.stopped {
+		return len(p), nil
+	}
+
+	// Append new data to the line buffer.
+	ps.lineBuf += string(p)
+
+	// Split the buffer into lines.
+	lines := strings.Split(ps.lineBuf, "\n")
+	ps.lineBuf = lines[len(lines)-1] // Save the last (possibly incomplete) line.
+
+	// Process all complete lines.
+	for _, line := range lines[:len(lines)-1] {
+		ps.buffers = append(ps.buffers, line) // Save the log line.
+		if !ps.stopped && ps.pattern.MatchString(line) {
+			ps.stopOnce.Do(func() {
+				ps.stopped = true
+				close(ps.DoneChan) // Notify the caller about the match.
+			})
+		}
 	}
 
-	return
+	return len(p), nil
 }
 
 // PrintLogs writes the captured logs into the test logger.