From e0419c2e58bfa6767a4dc7e2386c93e1f194fc21 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20Pa=C3=9F?=
 <22845248+mpass99@users.noreply.github.com>
Date: Sat, 11 Mar 2023 13:19:55 +0000
Subject: [PATCH] Fix Sentry Debug Regex

that was ignoring composed messages including a newline.
Also, add regression test.
---
 internal/nomad/sentry_debug_writer.go      | 25 ++++++++++++++--------
 internal/nomad/sentry_debug_writer_test.go | 12 +++++++++++
 tests/e2e/websocket_test.go                |  4 ++--
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/internal/nomad/sentry_debug_writer.go b/internal/nomad/sentry_debug_writer.go
index 497880e..cc262b6 100644
--- a/internal/nomad/sentry_debug_writer.go
+++ b/internal/nomad/sentry_debug_writer.go
@@ -19,7 +19,7 @@ var (
 	timeDebugMessageFormatEnd = `%s; ec=$?; ` + timeDebugMessageFormat + ` && exit $ec`
 
 	timeDebugMessagePattern = regexp.MustCompile(
-		`(?P<before>.*)\x1EPoseidon (?P<text>.+) (?P<time>\d{13})\x1E(?P<after>.*)`)
+		`(?P<before>[\S\s]*?)\x1EPoseidon (?P<text>[^\x1E]+) (?P<time>\d{13})\x1E(?P<after>.*)`)
 	timeDebugMessagePatternStart = regexp.MustCompile(`\x1EPoseidon`)
 )
 
@@ -57,15 +57,22 @@ func (s *SentryDebugWriter) Write(p []byte) (n int, err error) {
 		return 0, nil
 	}
 
-	go s.handleTimeDebugMessage(match)
-
-	lengthRemainingData := len(match["before"]) + len(match["after"])
-	if lengthRemainingData > 0 {
-		count, err := s.Write(append(match["before"], match["after"]...))
-		return len(p) - (lengthRemainingData - count), err
-	} else {
-		return len(p), nil
+	if len(match["before"]) > 0 {
+		var count int
+		count, err = s.Write(match["before"])
+		n += count
 	}
+
+	go s.handleTimeDebugMessage(match)
+	n += len(p) - len(match["before"]) - len(match["after"])
+
+	if len(match["after"]) > 0 {
+		var count int
+		count, err = s.Write(match["after"])
+		n += count
+	}
+
+	return n, err
 }
 
 func (s *SentryDebugWriter) Close(exitCode int) {
diff --git a/internal/nomad/sentry_debug_writer_test.go b/internal/nomad/sentry_debug_writer_test.go
index 45f762c..2d34b9f 100644
--- a/internal/nomad/sentry_debug_writer_test.go
+++ b/internal/nomad/sentry_debug_writer_test.go
@@ -21,6 +21,18 @@ func TestSentryDebugWriter_Write(t *testing.T) {
 	assert.NotContains(t, buf.String(), description)
 }
 
+func TestSentryDebugWriter_WriteComposed(t *testing.T) {
+	buf := &bytes.Buffer{}
+	w := SentryDebugWriter{Target: buf, Ctx: context.Background()}
+
+	data := "Hello World!\r\n\x1EPoseidon unset 1678540012404\x1E\x1EPoseidon /sbin/setuser user 1678540012408\x1E"
+	count, err := w.Write([]byte(data))
+
+	require.NoError(t, err)
+	assert.Equal(t, len(data), count)
+	assert.Contains(t, buf.String(), "Hello World!")
+}
+
 func TestSentryDebugWriter_Close(t *testing.T) {
 	buf := &bytes.Buffer{}
 	s := NewSentryDebugWriter(buf, context.Background())
diff --git a/tests/e2e/websocket_test.go b/tests/e2e/websocket_test.go
index c6fb551..1ccb6b1 100644
--- a/tests/e2e/websocket_test.go
+++ b/tests/e2e/websocket_test.go
@@ -89,7 +89,7 @@ func (s *E2ETestSuite) TestUserNomad() {
 	s.Run("privileged", func() {
 		stdout, _, _ := ExecuteNonInteractive(&s.Suite, tests.DefaultEnvironmentIDAsInteger,
 			&dto.ExecutionRequest{Command: "id --name --user", PrivilegedExecution: true}, nil)
-		s.Require().Equal("root\r\n", stdout)
+		s.Contains(stdout, "root")
 	})
 }
 
@@ -111,7 +111,7 @@ func (s *E2ETestSuite) TestCommandHead() {
 	s.Require().Error(err)
 	s.Equal(err, &websocket.CloseError{Code: websocket.CloseNormalClosure})
 	stdout, _, _ := helpers.WebSocketOutputMessages(messages)
-	s.Contains(stdout, fmt.Sprintf("%s\r\n%s\r\n", hello, hello))
+	s.Regexp(fmt.Sprintf(`(%s\r\n?){2}`, hello), stdout)
 }
 
 func (s *E2ETestSuite) TestCommandMake() {