From 770653366c85b6514ebe672c89bc2e76725804c2 Mon Sep 17 00:00:00 2001
From: Maximilian Pass <22845248+mpass99@users.noreply.github.com>
Date: Thu, 12 Sep 2024 15:01:05 +0200
Subject: [PATCH] Fix Sentry Debug Message for split debug messages

by buffering a message when the message does not contain both start and end of a debug message.
---
 internal/nomad/sentry_debug_writer.go      | 35 +++++++++++++++----
 internal/nomad/sentry_debug_writer_test.go | 39 ++++++++++++++++++++++
 2 files changed, 67 insertions(+), 7 deletions(-)

diff --git a/internal/nomad/sentry_debug_writer.go b/internal/nomad/sentry_debug_writer.go
index cda283d0..2fedc281 100644
--- a/internal/nomad/sentry_debug_writer.go
+++ b/internal/nomad/sentry_debug_writer.go
@@ -1,6 +1,7 @@
 package nomad
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"io"
@@ -22,8 +23,9 @@ var (
 
 	timeDebugMessagePattern = regexp.MustCompile(
 		`(?P<before>[\S\s]*?)\x1EPoseidon (?P<text>[^\x1E]+?) (?P<time>\d{13})\x1E(?P<after>[\S\s]*)`)
-	timeDebugMessagePatternStart = regexp.MustCompile(`\x1EPoseidon`)
-	timeDebugFallbackDescription = "<empty>"
+	timeDebugMessagePatternStart    = regexp.MustCompile(`\x1EPoseidon`)
+	timeDebugMessagePatternStartEnd = regexp.MustCompile(`\x1EPoseidon.*\x1E`)
+	timeDebugFallbackDescription    = "<empty>"
 )
 
 // SentryDebugWriter is scanning the input for the debug message pattern.
@@ -32,8 +34,9 @@ var (
 type SentryDebugWriter struct {
 	Target io.Writer
 	//nolint:containedctx // See #630.
-	Ctx      context.Context
-	lastSpan *sentry.Span
+	Ctx       context.Context
+	lastSpan  *sentry.Span
+	remaining *bytes.Buffer
 }
 
 func NewSentryDebugWriter(ctx context.Context, target io.Writer) *SentryDebugWriter {
@@ -41,9 +44,10 @@ func NewSentryDebugWriter(ctx context.Context, target io.Writer) *SentryDebugWri
 	span := sentry.StartSpan(ctx, "nomad.execute.connect")
 	span.Description = "/bin/bash -c"
 	return &SentryDebugWriter{
-		Target:   target,
-		Ctx:      ctx,
-		lastSpan: span,
+		Target:    target,
+		Ctx:       ctx,
+		lastSpan:  span,
+		remaining: &bytes.Buffer{},
 	}
 }
 
@@ -59,6 +63,15 @@ func (s *SentryDebugWriter) Write(debugData []byte) (n int, err error) {
 	}
 	log.WithContext(s.lastSpan.Context()).WithField("data", fmt.Sprintf("%q", debugData)).Trace("Received data from Nomad container")
 
+	if s.remaining.Len() > 0 {
+		n -= s.remaining.Len()
+		joinedDebugData := make([]byte, 0, s.remaining.Len()+len(debugData))
+		joinedDebugData = append(joinedDebugData, s.remaining.Bytes()...)
+		joinedDebugData = append(joinedDebugData, debugData...)
+		s.remaining.Reset()
+		debugData = joinedDebugData
+	}
+
 	if !timeDebugMessagePatternStart.Match(debugData) {
 		count, err := s.Target.Write(debugData)
 		if err != nil {
@@ -67,6 +80,14 @@ func (s *SentryDebugWriter) Write(debugData []byte) (n int, err error) {
 		return count, err
 	}
 
+	if !timeDebugMessagePatternStartEnd.Match(debugData) {
+		count, err := s.remaining.Write(debugData)
+		if err != nil {
+			err = fmt.Errorf("SentryDebugWriter failed to buffer data: %w", err)
+		}
+		return count, err
+	}
+
 	match := matchAndMapTimeDebugMessage(debugData)
 	if match == nil {
 		log.WithContext(s.lastSpan.Context()).WithField("data", fmt.Sprintf("%q", debugData)).Warn("Exec debug message could not be read completely")
diff --git a/internal/nomad/sentry_debug_writer_test.go b/internal/nomad/sentry_debug_writer_test.go
index 52c44fc5..8532c7c5 100644
--- a/internal/nomad/sentry_debug_writer_test.go
+++ b/internal/nomad/sentry_debug_writer_test.go
@@ -47,6 +47,45 @@ func (s *MainTestSuite) TestSentryDebugWriter_regression_593_empty_command() {
 	s.Equal(timeDebugFallbackDescription, debugWriter.lastSpan.Description)
 }
 
+func (s *MainTestSuite) TestSentryDebugWriter_regression_593_split_debug_message() {
+	buf := &bytes.Buffer{}
+	debugWriter := NewSentryDebugWriter(s.TestCtx, buf)
+
+	s.Run("single debug message", func() {
+		output := "Output"
+		payload := output + "\x1ePoseidon exit 0 1"
+		payload2 := "725462286087\x1E"
+
+		count, err := debugWriter.Write([]byte(payload))
+		s.Require().NoError(err)
+		s.Equal(len(payload), count)
+		s.NotEqual("exit 0", debugWriter.lastSpan.Description)
+
+		count, err = debugWriter.Write([]byte(payload2))
+		s.Require().NoError(err)
+		s.Equal(len(payload2), count)
+		s.Equal("exit 0", debugWriter.lastSpan.Description)
+		s.Equal(output, buf.String())
+	})
+	buf.Reset()
+	s.Run("multiple debug messages", func() {
+		output := "Output"
+		payload := "\x1ePoseidon env 1725462286085\x1E" + output + "\x1ePoseidon exit 0 1"
+		payload2 := "725462286087\x1E"
+
+		count, err := debugWriter.Write([]byte(payload))
+		s.Require().NoError(err)
+		s.Equal(len(payload), count)
+		s.NotEqual("exit 0", debugWriter.lastSpan.Description)
+
+		count, err = debugWriter.Write([]byte(payload2))
+		s.Require().NoError(err)
+		s.Equal(len(payload2), count)
+		s.Equal("exit 0", debugWriter.lastSpan.Description)
+		s.Equal(output, buf.String())
+	})
+}
+
 func (s *MainTestSuite) TestSentryDebugWriter_regression_issue_678() {
 	buf := &bytes.Buffer{}
 	debugWriter := NewSentryDebugWriter(s.TestCtx, buf)