From 05182fc7d58922b793d1290e9faf1fa0150feed3 Mon Sep 17 00:00:00 2001
From: Kris <krousey@google.com>
Date: Tue, 3 Nov 2015 14:06:51 -0800
Subject: [PATCH] Properly defaulting the status in log code

---
 pkg/httplog/log.go      | 37 ++++++++++++++++++++++++++-----------
 pkg/httplog/log_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 11 deletions(-)

diff --git a/pkg/httplog/log.go b/pkg/httplog/log.go
index 77b2dbe8e77..440c9e7b04e 100644
--- a/pkg/httplog/log.go
+++ b/pkg/httplog/log.go
@@ -53,10 +53,12 @@ type logger interface {
 // the http.ResponseWriter. We can recover panics from go-restful, and
 // the logging value is questionable.
 type respLogger struct {
-	status      int
-	statusStack string
-	addedInfo   string
-	startTime   time.Time
+	hijacked       bool
+	statusRecorded bool
+	status         int
+	statusStack    string
+	addedInfo      string
+	startTime      time.Time
 
 	req *http.Request
 	w   http.ResponseWriter
@@ -155,7 +157,11 @@ func (rl *respLogger) Addf(format string, data ...interface{}) {
 func (rl *respLogger) Log() {
 	latency := time.Since(rl.startTime)
 	if glog.V(2) {
-		glog.InfoDepth(1, fmt.Sprintf("%s %s: (%v) %v%v%v [%s %s]", rl.req.Method, rl.req.RequestURI, latency, rl.status, rl.statusStack, rl.addedInfo, rl.req.Header["User-Agent"], rl.req.RemoteAddr))
+		if !rl.hijacked {
+			glog.InfoDepth(1, fmt.Sprintf("%s %s: (%v) %v%v%v [%s %s]", rl.req.Method, rl.req.RequestURI, latency, rl.status, rl.statusStack, rl.addedInfo, rl.req.Header["User-Agent"], rl.req.RemoteAddr))
+		} else {
+			glog.InfoDepth(1, fmt.Sprintf("%s %s: (%v) hijacked [%s %s]", rl.req.Method, rl.req.RequestURI, latency, rl.req.Header["User-Agent"], rl.req.RemoteAddr))
+		}
 	}
 }
 
@@ -166,6 +172,9 @@ func (rl *respLogger) Header() http.Header {
 
 // Write implements http.ResponseWriter.
 func (rl *respLogger) Write(b []byte) (int, error) {
+	if !rl.statusRecorded {
+		rl.recordStatus(http.StatusOK) // Default if WriteHeader hasn't been called
+	}
 	return rl.w.Write(b)
 }
 
@@ -181,7 +190,19 @@ func (rl *respLogger) Flush() {
 
 // WriteHeader implements http.ResponseWriter.
 func (rl *respLogger) WriteHeader(status int) {
+	rl.recordStatus(status)
+	rl.w.WriteHeader(status)
+}
+
+// Hijack implements http.Hijacker.
+func (rl *respLogger) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	rl.hijacked = true
+	return rl.w.(http.Hijacker).Hijack()
+}
+
+func (rl *respLogger) recordStatus(status int) {
 	rl.status = status
+	rl.statusRecorded = true
 	if rl.logStacktracePred(status) {
 		// Only log stacks for errors
 		stack := make([]byte, 2048)
@@ -190,10 +211,4 @@ func (rl *respLogger) WriteHeader(status int) {
 	} else {
 		rl.statusStack = ""
 	}
-	rl.w.WriteHeader(status)
-}
-
-// Hijack implements http.Hijacker.
-func (rl *respLogger) Hijack() (net.Conn, *bufio.ReadWriter, error) {
-	return rl.w.(http.Hijacker).Hijack()
 }
diff --git a/pkg/httplog/log_test.go b/pkg/httplog/log_test.go
index a24a194188a..a5d72b1ccbc 100644
--- a/pkg/httplog/log_test.go
+++ b/pkg/httplog/log_test.go
@@ -129,3 +129,33 @@ func TestUnlogged(t *testing.T) {
 		handler(w, req)
 	}
 }
+
+type testResponseWriter struct{}
+
+func (*testResponseWriter) Header() http.Header       { return nil }
+func (*testResponseWriter) Write([]byte) (int, error) { return 0, nil }
+func (*testResponseWriter) WriteHeader(int)           {}
+
+func TestLoggedStatus(t *testing.T) {
+	req, err := http.NewRequest("GET", "http://example.com", nil)
+	if err != nil {
+		t.Errorf("unexpected error: %v", err)
+	}
+
+	var tw http.ResponseWriter = new(testResponseWriter)
+	logger := NewLogged(req, &tw)
+	logger.Write(nil)
+
+	if logger.status != http.StatusOK {
+		t.Errorf("expected status after write to be %v, got %v", http.StatusOK, logger.status)
+	}
+
+	tw = new(testResponseWriter)
+	logger = NewLogged(req, &tw)
+	logger.WriteHeader(http.StatusForbidden)
+	logger.Write(nil)
+
+	if logger.status != http.StatusForbidden {
+		t.Errorf("expected status after write to remain %v, got %v", http.StatusForbidden, logger.status)
+	}
+}