From d852c616c6047f711197f376b7f3b3927ddf68d6 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Wed, 8 Nov 2023 11:48:36 -0800
Subject: [PATCH] logtail: fix Logger.Write return result

io.Writer says you need to write completely on err=nil. (the result
int should be the same as the input buffer length)

We weren't doing that. We used to, but at some point the verbose
filtering was modifying buf before the final return of len(buf).

We've been getting lucky probably, that callers haven't looked at our
results and turned us into a short write error.

Updates #cleanup
Updates tailscale/corp#15664

Change-Id: I01e765ba35b86b759819e38e0072eceb9d10d75c
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 logtail/logtail.go      |  4 +++-
 logtail/logtail_test.go | 27 +++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/logtail/logtail.go b/logtail/logtail.go
index 4544af9d7..f17fd11d5 100644
--- a/logtail/logtail.go
+++ b/logtail/logtail.go
@@ -735,6 +735,8 @@ func (l *Logger) Write(buf []byte) (int, error) {
 	if len(buf) == 0 {
 		return 0, nil
 	}
+	inLen := len(buf) // length as provided to us, before modifications to downstream writers
+
 	level, buf := parseAndRemoveLogLevel(buf)
 	if l.stderr != nil && l.stderr != io.Discard && int64(level) <= atomic.LoadInt64(&l.stderrLevel) {
 		if buf[len(buf)-1] == '\n' {
@@ -752,7 +754,7 @@ func (l *Logger) Write(buf []byte) (int, error) {
 
 	b := l.encodeLocked(buf, level)
 	_, err := l.sendLocked(b)
-	return len(buf), err
+	return inLen, err
 }
 
 var (
diff --git a/logtail/logtail_test.go b/logtail/logtail_test.go
index f145d7a35..f0d3f36b2 100644
--- a/logtail/logtail_test.go
+++ b/logtail/logtail_test.go
@@ -379,3 +379,30 @@ func TestEncode(t *testing.T) {
 		}
 	}
 }
+
+// Test that even if Logger.Write modifies the input buffer, we still return the
+// length of the input buffer, not what we shrank it down to. Otherwise the
+// caller will think we did a short write, violating the io.Writer contract.
+func TestLoggerWriteResult(t *testing.T) {
+	buf := NewMemoryBuffer(100)
+	lg := &Logger{
+		clock:  tstest.NewClock(tstest.ClockOpts{Start: time.Unix(123, 0)}),
+		buffer: buf,
+	}
+
+	const in = "[v1] foo"
+	n, err := lg.Write([]byte(in))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got, want := n, len(in); got != want {
+		t.Errorf("Write = %v; want %v", got, want)
+	}
+	back, err := buf.TryReadLine()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got, want := string(back), `{"logtail": {"client_time": "1970-01-01T00:02:03Z"}, "v":1,"text": "foo"}`+"\n"; got != want {
+		t.Errorf("mismatch.\n got: %#q\nwant: %#q", back, want)
+	}
+}