Skip to content

Commit

Permalink
[USM] fix edge-case in protocols testutil (DataDog#31430)
Browse files Browse the repository at this point in the history
  • Loading branch information
val06 authored Nov 26, 2024
1 parent 00429a1 commit c51e43f
Showing 1 changed file with 25 additions and 10 deletions.
35 changes: 25 additions & 10 deletions pkg/network/protocols/testutil/patternscanner.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand Down

0 comments on commit c51e43f

Please sign in to comment.