From 2863e49db951f1dd89a86889466ad140027ff75f Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Sat, 4 Apr 2020 08:55:10 -0700
Subject: [PATCH] tsweb: don't flush, treat no-op Handler as 200, like Go

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 tsweb/tsweb.go      | 28 ++++++++++------------------
 tsweb/tsweb_test.go | 17 ++++++-----------
 2 files changed, 16 insertions(+), 29 deletions(-)

diff --git a/tsweb/tsweb.go b/tsweb/tsweb.go
index 00b666342..f77438053 100644
--- a/tsweb/tsweb.go
+++ b/tsweb/tsweb.go
@@ -198,10 +198,16 @@ func (h handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		Referer:    r.Referer(),
 	}
 
-	lw := loggingResponseWriter{ResponseWriter: w, logf: h.logf}
-	err := h.h.ServeHTTPErr(&lw, r)
+	lw := &loggingResponseWriter{ResponseWriter: w, logf: h.logf}
+	err := h.h.ServeHTTPErr(lw, r)
 	hErr, hErrOK := err.(HTTPError)
 
+	if lw.code == 0 && err == nil && !lw.hijacked {
+		// If the handler didn't write and didn't send a header, that still means 200.
+		// (See https://play.golang.org/p/4P7nx_Tap7p)
+		lw.code = 200
+	}
+
 	msg.Seconds = h.timeNow().Sub(msg.When).Seconds()
 	msg.Code = lw.code
 	msg.Bytes = lw.bytes
@@ -231,31 +237,17 @@ func (h handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			h.logf("[unexpected] HTTPError %v did not contain an HTTP status code, sending internal server error", hErr)
 			msg.Code = http.StatusInternalServerError
 		}
-		http.Error(&lw, hErr.Msg, msg.Code)
+		http.Error(lw, hErr.Msg, msg.Code)
 	case err != nil:
 		// Handler returned a generic error. Serve an internal server
 		// error, if necessary.
 		msg.Err = err.Error()
 		if lw.code == 0 {
 			msg.Code = http.StatusInternalServerError
-			http.Error(&lw, "internal server error", msg.Code)
+			http.Error(lw, "internal server error", msg.Code)
 		}
-	case lw.code == 0:
-		// Handler exited successfully, but didn't generate a
-		// response. Synthesize an internal server error.
-		msg.Code = http.StatusInternalServerError
-		msg.Err = "[unexpected] handler did not respond to the client"
-		http.Error(&lw, "internal server error", msg.Code)
 	}
 
-	// Cleanup below is common to all success and error paths. msg has
-	// been populated with relevant information either way.
-
-	// TODO(danderson): needed? Copied from existing code, but
-	// doesn't HTTPServer do this by itself?
-	if f, _ := w.(http.Flusher); !lw.hijacked && f != nil {
-		f.Flush()
-	}
 	h.logf("%s", msg)
 }
 
diff --git a/tsweb/tsweb_test.go b/tsweb/tsweb_test.go
index 6fa3ac2bf..be49c9694 100644
--- a/tsweb/tsweb_test.go
+++ b/tsweb/tsweb_test.go
@@ -192,7 +192,7 @@ func TestStdHandler(t *testing.T) {
 			name:     "handler does nothing",
 			h:        HandlerFunc(func(http.ResponseWriter, *http.Request) error { return nil }),
 			r:        req(bgCtx, "http://example.com/foo"),
-			wantCode: 500,
+			wantCode: 200,
 			wantLog: AccessLogRecord{
 				When:       clock.Start,
 				Seconds:    1.0,
@@ -200,8 +200,7 @@ func TestStdHandler(t *testing.T) {
 				Host:       "example.com",
 				Method:     "GET",
 				RequestURI: "/foo",
-				Code:       500,
-				Err:        "[unexpected] handler did not respond to the client",
+				Code:       200,
 			},
 		},
 
@@ -215,7 +214,7 @@ func TestStdHandler(t *testing.T) {
 				return err
 			}),
 			r:        req(bgCtx, "http://example.com/foo"),
-			wantCode: 0,
+			wantCode: 200,
 			wantLog: AccessLogRecord{
 				When:    clock.Start,
 				Seconds: 1.0,
@@ -242,15 +241,11 @@ func TestStdHandler(t *testing.T) {
 			clock.Reset()
 
 			rec := noopHijacker{httptest.NewRecorder(), false}
-			// ResponseRecorder defaults Code to 200, grump.
-			rec.Code = 0
 			h := stdHandler(test.h, logf, clock.Now)
 			h.ServeHTTP(&rec, test.r)
-			if rec.Code != test.wantCode {
-				t.Errorf("HTTP code = %v, want %v", rec.Code, test.wantCode)
-			}
-			if !rec.hijacked && !rec.Flushed {
-				t.Errorf("handler didn't flush")
+			res := rec.Result()
+			if res.StatusCode != test.wantCode {
+				t.Errorf("HTTP code = %v, want %v", res.StatusCode, test.wantCode)
 			}
 			if len(logs) != 1 {
 				t.Errorf("handler didn't write a request log")