fix(HMS-3369): Implement replacing as beforeSend hook

With implementing the replace functionality as Sentry hook, we can implement it bit more efficiently for only messages that are sent to sentry. It also allows us to use the sentry writer as LevelWriter and leverage WriteLevel function. WriteLevel is bit more streamlined as it allows to skip filtering when not needed.
RHEnVision · Jan 15, 2024 · b448d8b · b448d8b
1 parent 22b063a
commit b448d8b
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 143 deletions.
diff --git a/internal/logging/sentry.go b/internal/logging/sentry.go
@@ -2,28 +2,25 @@ package logging
 
 import (
 	"fmt"
-	"io"
 
 	"github.com/RHEnVision/provisioning-backend/internal/config"
 	sentrywriter "github.com/archdx/zerolog-sentry"
 	"github.com/getsentry/sentry-go"
+	"github.com/rs/zerolog"
 )
 
 // sentryWriter creates a zerolog writer for sentry.
 // Uses github.com/archdx/zerolog-sentry which is very simple wrapper.
-func sentryWriter(dsn string) (io.Writer, func(), error) {
-	sWriter, err := sentrywriter.New(dsn)
+func sentryWriter(dsn string) (zerolog.LevelWriter, func() error, error) {
+	replacer := NewSentryReplacer()
+
+	sWriter, err := sentrywriter.New(dsn, sentrywriter.WithBeforeSend(replacer.Replace))
 	if err != nil {
-		return nil, func() {}, fmt.Errorf("cannot initialize sentry: %w", err)
+		return nil, func() error { return nil }, fmt.Errorf("cannot initialize sentry: %w", err)
 	}
 	sentry.ConfigureScope(func(scope *sentry.Scope) {
 		scope.SetTag("stream", config.BinaryName())
 	})
-	fWriter := NewSentryReplacer(sWriter)
 
-	closeFunc := func() {
-		_ = fWriter.Close()
-		_ = sWriter.Close()
-	}
-	return fWriter, closeFunc, nil
+	return sWriter, sWriter.Close, nil
 }
diff --git a/internal/logging/sentry_replacer.go b/internal/logging/sentry_replacer.go
@@ -1,12 +1,10 @@
 package logging
 
 import (
-	"bytes"
-	"fmt"
-	"io"
 	"regexp"
 	"strings"
-	"sync"
+
+	"github.com/getsentry/sentry-go"
 )
 
 var filters = []string{
@@ -26,70 +24,20 @@ var filters = []string{
 	`\d\d\d\d-\d\d-\d\dT\d\d:\d\d:\d\d.\d+\+\d\d:\d\d`,
 }
 
-var replacement = []byte{'?'}
+var replacement = "?"
 
 type SentryReplacer struct {
-	buf    []byte
-	w      io.Writer
-	re     *regexp.Regexp
-	mu     sync.Mutex
-	closed bool
+	re *regexp.Regexp
 }
 
-func NewSentryReplacer(w io.Writer) *SentryReplacer {
+func NewSentryReplacer() *SentryReplacer {
 	sr := SentryReplacer{
-		w:  w,
 		re: regexp.MustCompile("(" + strings.Join(filters, "|") + ")"),
 	}
 	return &sr
 }
 
-func (sr *SentryReplacer) Write(p []byte) (n int, err error) {
-	sr.mu.Lock()
-	defer sr.mu.Unlock()
-
-	if sr.closed {
-		return 0, io.EOF
-	}
-
-	// Append p to our own buffer, see if there's anything we need to censor.
-	sr.buf = append(sr.buf, p...)
-
-	// If we've appended at least a line, censor it and write it out.
-	for {
-		if len(sr.buf) == 0 {
-			// Buffer flushed out completely.
-			return len(p), nil
-		}
-
-		idx := bytes.IndexRune(sr.buf, '\n')
-		if idx < 0 {
-			// No line yet, just lie to the caller and tell them we wrote p.
-			return len(p), nil
-		}
-
-		var line []byte
-		line, sr.buf = sr.buf[:idx+1], sr.buf[idx+1:]
-		line = sr.re.ReplaceAll(line, replacement)
-
-		_, err := sr.w.Write(line)
-		if err != nil {
-			// This is not strictly the error related to the incoming `p`, but the best we can do.
-			return 0, fmt.Errorf("cannot filter: %w", err)
-		}
-	}
-}
-
-func (sr *SentryReplacer) Close() error {
-	sr.mu.Lock()
-	defer sr.mu.Unlock()
-
-	replaced := sr.re.ReplaceAll(sr.buf, replacement)
-	_, err := sr.w.Write(replaced)
-	if err != nil {
-		return fmt.Errorf("cannot close filter: %w", err)
-	}
-
-	sr.closed = true
-	return nil
+func (sr *SentryReplacer) Replace(event *sentry.Event, hint *sentry.EventHint) *sentry.Event {
+	event.Message = sr.re.ReplaceAllString(event.Message, replacement)
+	return event
 }
diff --git a/internal/logging/sentry_replacer_test.go b/internal/logging/sentry_replacer_test.go
@@ -1,106 +1,71 @@
 package logging
 
 import (
-	"bytes"
 	"testing"
 
+	"github.com/getsentry/sentry-go"
 	"github.com/stretchr/testify/require"
 )
 
 func TestNewline(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("x\nx\n"))
-	require.Equal(t, "x\nx\n", buf.String())
-}
-
-func TestNoNewline(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("x\nx"))
-	require.Equal(t, "x\n", buf.String())
+	evt := sentry.Event{Message: "x\nx\n"}
+	repl := NewSentryReplacer()
+	result := repl.Replace(&evt, nil)
+	require.Equal(t, "x\nx\n", result.Message)
 }
 
 func TestNoNewlineClose(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("x\nx"))
-	repl.Close()
-	require.Equal(t, "x\nx", buf.String())
-}
-
-func TestBufferFlush(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("a\n"))
-	_, _ = repl.Write([]byte("\n"))
-	require.Equal(t, "a\n\n", buf.String())
-	require.Zero(t, len(repl.buf))
+	evt := sentry.Event{Message: "x\nx"}
+	repl := NewSentryReplacer()
+	result := repl.Replace(&evt, nil)
+	require.Equal(t, "x\nx", result.Message)
 }
 
 func TestUUID(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("ca767444-d1f9-11ed-afa1-0242ac120002\n"))
-	require.Equal(t, "?\n", buf.String())
-}
-
-func TestUUIDSplit(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("ca767444-d1f9-"))
-	_, _ = repl.Write([]byte("11ed-afa1-"))
-	_, _ = repl.Write([]byte(""))
-	_, _ = repl.Write([]byte("0242ac120002\n"))
-	require.Equal(t, "?\n", buf.String())
+	evt := sentry.Event{Message: "ca767444-d1f9-11ed-afa1-0242ac120002\n"}
+	repl := NewSentryReplacer()
+	result := repl.Replace(&evt, nil)
+	require.Equal(t, "?\n", result.Message)
 }
 
 func TestARN(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("arn:aws:iam::4328974392798432:role/my-role-123\n"))
-	require.Equal(t, "?\n", buf.String())
-}
-
-func TestARNSplit(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("arn:aws:iam::"))
-	_, _ = repl.Write([]byte("4328974392798432:role/my-role-123\n"))
-	require.Equal(t, "?\n", buf.String())
+	evt := sentry.Event{Message: "arn:aws:iam::4328974392798432:role/my-role-123\n"}
+	repl := NewSentryReplacer()
+	result := repl.Replace(&evt, nil)
+	require.Equal(t, "?\n", result.Message)
 }
 
 func TestIPv4(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("read tcp 10.128.24.14:42094->10.0.217.126:6379: i/o timeout\n"))
-	require.Equal(t, "read tcp ?->?: i/o timeout\n", buf.String())
+	evt := sentry.Event{Message: "read tcp 10.128.24.14:42094->10.0.217.126:6379: i/o timeout\n"}
+	repl := NewSentryReplacer()
+	result := repl.Replace(&evt, nil)
+	require.Equal(t, "read tcp ?->?: i/o timeout\n", result.Message)
 }
 
 func TestFingerprint(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("pubkey with fingerprint (57:d4:13:ff:c0:74:51:50:41:ec:e1:cd:f1:88:b0:61)\n"))
-	require.Equal(t, "pubkey with fingerprint (?)\n", buf.String())
+	evt := sentry.Event{Message: "pubkey with fingerprint (57:d4:13:ff:c0:74:51:50:41:ec:e1:cd:f1:88:b0:61)\n"}
+	repl := NewSentryReplacer()
+	result := repl.Replace(&evt, nil)
+	require.Equal(t, "pubkey with fingerprint (?)\n", result.Message)
 }
 
 func TestAWSResourceID(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("instance ID 'i-0fe8a8adc1403f5b1' does not exist\n\n\n"))
-	require.Equal(t, "instance ID '?' does not exist\n\n\n", buf.String())
+	evt := sentry.Event{Message: "instance ID 'i-0fe8a8adc1403f5b1' does not exist\n\n\n"}
+	repl := NewSentryReplacer()
+	result := repl.Replace(&evt, nil)
+	require.Equal(t, "instance ID '?' does not exist\n\n\n", result.Message)
 }
 
 func TestGoogleProject(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("The resource 'projects/xxx' was not found\n"))
-	require.Equal(t, "The resource ? was not found\n", buf.String())
+	evt := sentry.Event{Message: "The resource 'projects/xxx' was not found\n"}
+	repl := NewSentryReplacer()
+	result := repl.Replace(&evt, nil)
+	require.Equal(t, "The resource ? was not found\n", result.Message)
 }
 
 func TestAzureTime(t *testing.T) {
-	buf := bytes.NewBufferString("")
-	repl := NewSentryReplacer(buf)
-	_, _ = repl.Write([]byte("'start time': '2023-06-24T19:34:34.2581206+00:00'\n"))
-	require.Equal(t, "'start time': '?'\n", buf.String())
+	evt := sentry.Event{Message: "'start time': '2023-06-24T19:34:34.2581206+00:00'\n"}
+	repl := NewSentryReplacer()
+	result := repl.Replace(&evt, nil)
+	require.Equal(t, "'start time': '?'\n", result.Message)
 }
diff --git a/internal/logging/zerolog.go b/internal/logging/zerolog.go
@@ -129,7 +129,7 @@ func InitializeLogger() (zerolog.Logger, func()) {
 			fn()
 		}
 	}
-	logger := decorate(zerolog.New(io.MultiWriter(writers...)))
+	logger := decorate(zerolog.New(zerolog.MultiLevelWriter(writers...)))
 	log.Logger = logger
 	zerolog.DefaultContextLogger = &logger
 	return logger, closeFn