From 1bf7ed03488928eb389f9da5551606162f183350 Mon Sep 17 00:00:00 2001
From: Paul Scott <408401+icio@users.noreply.github.com>
Date: Mon, 29 Jul 2024 13:53:01 +0100
Subject: [PATCH] tsweb: add QuietLogging option (#12838)

Allows the use of tsweb.LogHandler exclusively for callbacks describing the
handler HTTP requests.

Fixes #12837

Signed-off-by: Paul Scott <paul@tailscale.com>
---
 tsweb/tsweb.go      |  6 ++++-
 tsweb/tsweb_test.go | 56 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/tsweb/tsweb.go b/tsweb/tsweb.go
index 56680cc92..9ddb3fad5 100644
--- a/tsweb/tsweb.go
+++ b/tsweb/tsweb.go
@@ -276,6 +276,10 @@ type LogOptions struct {
 	// Now is a function giving the current time. Defaults to [time.Now].
 	Now func() time.Time
 
+	// QuietLogging suppresses all logging of handled HTTP requests, even if
+	// there are errors or status codes considered unsuccessful. Use this option
+	// to add your own logging in OnCompletion.
+	QuietLogging bool
 	// QuietLoggingIfSuccessful suppresses logging of handled HTTP requests
 	// where the request's response status code is 200 or 304.
 	QuietLoggingIfSuccessful bool
@@ -569,7 +573,7 @@ func (h logHandler) logRequest(r *http.Request, lw *loggingResponseWriter, msg A
 		}
 	}
 
-	if !h.opts.QuietLoggingIfSuccessful || (msg.Code != http.StatusOK && msg.Code != http.StatusNotModified) {
+	if !h.opts.QuietLogging && !(h.opts.QuietLoggingIfSuccessful && (msg.Code == http.StatusOK || msg.Code == http.StatusNotModified)) {
 		h.opts.Logf("%s", msg)
 	}
 
diff --git a/tsweb/tsweb_test.go b/tsweb/tsweb_test.go
index 7214ff3e6..c348a2598 100644
--- a/tsweb/tsweb_test.go
+++ b/tsweb/tsweb_test.go
@@ -1021,6 +1021,62 @@ func TestStdHandler_OnErrorPanic(t *testing.T) {
 	res.Body.Close()
 }
 
+func TestLogHandler_QuietLogging(t *testing.T) {
+	now := time.Now()
+	var logs []string
+	logf := func(format string, args ...any) {
+		logs = append(logs, fmt.Sprintf(format, args...))
+	}
+
+	var done bool
+	onComp := func(r *http.Request, alr AccessLogRecord) {
+		if done {
+			t.Fatal("expected only one OnCompletion call")
+		}
+		done = true
+
+		want := AccessLogRecord{
+			Time:       now,
+			RemoteAddr: "192.0.2.1:1234",
+			Proto:      "HTTP/1.1",
+			Host:       "example.com",
+			Method:     "GET",
+			RequestURI: "/",
+			Code:       200,
+		}
+		if diff := cmp.Diff(want, alr); diff != "" {
+			t.Fatalf("unexpected OnCompletion AccessLogRecord (-want +got):\n%s", diff)
+		}
+	}
+
+	LogHandler(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(200)
+			w.WriteHeader(201) // loggingResponseWriter will write a warning.
+		}),
+		LogOptions{
+			Logf:         logf,
+			OnCompletion: onComp,
+			QuietLogging: true,
+			Now:          func() time.Time { return now },
+		},
+	).ServeHTTP(
+		httptest.NewRecorder(),
+		httptest.NewRequest("GET", "/", nil),
+	)
+
+	if !done {
+		t.Fatal("OnCompletion call didn't happen")
+	}
+
+	wantLogs := []string{
+		"[unexpected] HTTP handler set statusCode twice (200 and 201)",
+	}
+	if diff := cmp.Diff(wantLogs, logs); diff != "" {
+		t.Fatalf("logs (-want +got):\n%s", diff)
+	}
+}
+
 func TestErrorHandler_Panic(t *testing.T) {
 	// errorHandler should panic when not wrapped in logHandler.
 	defer func() {