fix: use httpsnoop to wrap response writer

uptrace · Apr 14, 2022 · 3d86932 · 3d86932
1 parent d9be2c4
commit 3d86932
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 38 deletions.
diff --git a/example/basic/go.mod b/example/basic/go.mod
@@ -13,6 +13,7 @@ require (
 
 require (
 	github.com/fatih/color v1.13.0 // indirect
+	github.com/felixge/httpsnoop v1.0.2 // indirect
 	github.com/mattn/go-colorable v0.1.12 // indirect
 	github.com/mattn/go-isatty v0.0.14 // indirect
 	golang.org/x/sys v0.0.0-20220319134239-a9b59b0215f8 // indirect

diff --git a/example/basic/go.sum b/example/basic/go.sum
@@ -2,6 +2,8 @@ github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
 github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
+github.com/felixge/httpsnoop v1.0.2 h1:+nS9g82KMXccJ/wp0zyRW9ZBHFETmMGtkk+2CTTrW4o=
+github.com/felixge/httpsnoop v1.0.2/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
 github.com/mattn/go-colorable v0.1.12 h1:jF+Du6AlPIjs2BiUiQlKOX0rt3SujHxPnksPKZbaA40=
 github.com/mattn/go-colorable v0.1.12/go.mod h1:u5H1YNBxpqRaxsYJYSkiCWKzEfiAb1Gb520KVy5xxl4=

diff --git a/extra/reqlog/middleware.go b/extra/reqlog/middleware.go
@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"github.com/fatih/color"
+	"github.com/felixge/httpsnoop"
 
 	"github.com/uptrace/bunrouter"
 )
@@ -67,10 +68,10 @@ func (m *middleware) Middleware(next bunrouter.HandlerFunc) bunrouter.HandlerFun
 	}
 
 	return func(w http.ResponseWriter, req bunrouter.Request) error {
-		rec := NewLoggingResponseWriter(w)
+		rec := NewResponseWriter(w)
 
 		now := time.Now()
-		err := next(rec, req)
+		err := next(rec.Wrapped, req)
 		dur := time.Since(now)
 		statusCode := rec.StatusCode()
 
@@ -103,45 +104,31 @@ func (m *middleware) Middleware(next bunrouter.HandlerFunc) bunrouter.HandlerFun
 
 //------------------------------------------------------------------------------
 
-type LoggingResponseWriter interface {
-	http.ResponseWriter
-	StatusCode() int
-}
-
-func NewLoggingResponseWriter(w http.ResponseWriter) LoggingResponseWriter {
-	rec := &statusCodeRecorder{
-		ResponseWriter: w,
-		statusCode:     http.StatusOK,
-	}
-
-	if _, ok := w.(http.Flusher); ok {
-		return flusher{rec}
-	}
-	return rec
-}
-
-type statusCodeRecorder struct {
-	http.ResponseWriter
+type ResponseWriter struct {
+	Wrapped    http.ResponseWriter
 	statusCode int
 }
 
-func (rec *statusCodeRecorder) StatusCode() int {
-	return rec.statusCode
-}
-
-func (rec *statusCodeRecorder) WriteHeader(statusCode int) {
-	rec.statusCode = statusCode
-	rec.ResponseWriter.WriteHeader(statusCode)
-}
-
-type flusher struct {
-	*statusCodeRecorder
-}
-
-var _ http.Flusher = (*flusher)(nil)
-
-func (f flusher) Flush() {
-	f.ResponseWriter.(http.Flusher).Flush()
+func NewResponseWriter(w http.ResponseWriter) *ResponseWriter {
+	var rw ResponseWriter
+	rw.Wrapped = httpsnoop.Wrap(w, httpsnoop.Hooks{
+		WriteHeader: func(next httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
+			return func(statusCode int) {
+				if rw.statusCode == 0 {
+					rw.statusCode = statusCode
+				}
+				next(statusCode)
+			}
+		},
+	})
+	return &rw
+}
+
+func (w *ResponseWriter) StatusCode() int {
+	if w.statusCode != 0 {
+		return w.statusCode
+	}
+	return http.StatusOK
 }
 
 //------------------------------------------------------------------------------