Fix response writer interface (#43)

* Fix response writer interface

go.mod

@@ -4,6 +4,7 @@ go 1.10
 
 require (
 	github.com/agiledragon/gomonkey/v2 v2.3.1
+	github.com/felixge/httpsnoop v1.0.3
 	github.com/prometheus/client_golang v1.0.0
 	github.com/stretchr/testify v1.3.0
 )

go.sum

@@ -8,6 +8,8 @@ github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+Ce
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/felixge/httpsnoop v1.0.3 h1:s/nj+GCswXYzN5v2DpNMuMQYe+0DDwt5WVCU6CWBdXk=
+github.com/felixge/httpsnoop v1.0.3/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
 github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9GBnD5lWE=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=

middleware/interceptor.go

@@ -12,18 +12,15 @@ func Interceptor(options ...InterceptorOption) httpware.Middleware {
 	config := NewInterceptorConfig(options...)
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(writer http.ResponseWriter, req *http.Request) {
-			writerInterceptor, ok := writer.(*ResponseWriterInterceptor)
-			if !ok {
-				writerInterceptor = NewResponseWriterInterceptor(writer)
-			}
+			w := NewResponseWriterInterceptor(writer)
 
 			req.Body = interceptor.NewCopyReadCloser(req.Body)
-			config.CallbackBefore(writerInterceptor, req)
+			config.CallbackBefore(w, req)
 			defer func() {
-				config.CallbackAfter(writerInterceptor, req)
+				config.CallbackAfter(w, req)
 			}()
 
-			next.ServeHTTP(writerInterceptor, req)
+			next.ServeHTTP(w.ResponseWriter, req)
 		})
 	}
 }

middleware/metrics.go

@@ -4,8 +4,8 @@ import (
 	"fmt"
 	"net/http"
 	"strconv"
-	"time"
 
+	"github.com/felixge/httpsnoop"
 	"github.com/gol4ng/httpware/v4"
 	"github.com/gol4ng/httpware/v4/metrics"
 )
@@ -14,31 +14,30 @@ func Metrics(recorder metrics.Recorder, options ... metrics.Option) httpware.Mid
 	config := metrics.NewConfig(recorder, options...)
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(writer http.ResponseWriter, req *http.Request) {
-			writerInterceptor, ok := writer.(*ResponseWriterInterceptor)
-			if !ok {
-				writerInterceptor = NewResponseWriterInterceptor(writer)
-			}
 			handlerName := config.IdentifierProvider(req)
 			if config.MeasureInflightRequests {
 				config.Recorder.AddInflightRequests(req.Context(), handlerName, 1)
 				defer config.Recorder.AddInflightRequests(req.Context(), handlerName, -1)
 			}
 
-			start := time.Now()
+			httpMetrics := httpsnoop.Metrics{}
+
 			defer func() {
-				code := strconv.Itoa(writerInterceptor.StatusCode)
+				code := strconv.Itoa(httpMetrics.Code)
 				if !config.SplitStatus {
-					code = fmt.Sprintf("%dxx", writerInterceptor.StatusCode/100)
+					code = fmt.Sprintf("%dxx", httpMetrics.Code/100)
 				}
 
-				config.Recorder.ObserveHTTPRequestDuration(req.Context(), handlerName, time.Since(start), req.Method, code)
+				config.Recorder.ObserveHTTPRequestDuration(req.Context(), handlerName, httpMetrics.Duration, req.Method, code)
 
 				if config.ObserveResponseSize {
-					config.Recorder.ObserveHTTPResponseSize(req.Context(), handlerName, int64(len(writerInterceptor.Body)), req.Method, code)
+					config.Recorder.ObserveHTTPResponseSize(req.Context(), handlerName, httpMetrics.Written, req.Method, code)
 				}
 			}()
 
-			next.ServeHTTP(writerInterceptor, req)
+			httpMetrics.CaptureMetrics(writer, func(writer http.ResponseWriter) {
+				next.ServeHTTP(writer, req)
+			})
 		})
 	}
 }

middleware/metrics_test.go

@@ -31,7 +31,6 @@ func TestMetrics(t *testing.T) {
 	req = httptest.NewRequest(http.MethodGet, "http://fake-addr", nil)
 	// create handler that set http status to 200 and write some response content
 	handler := func(w http.ResponseWriter, r *http.Request) {
-		assert.IsType(t, middleware.NewResponseWriterInterceptor(nil), w)
 		assert.Equal(t, req, r)
 		w.WriteHeader(http.StatusOK)
 		_, _ = w.Write([]byte(responseBody)) // contentLength=13

middleware/response_writer_interceptor.go

@@ -1,6 +1,7 @@
 package middleware
 
 import (
+	"github.com/felixge/httpsnoop"
 	"net/http"
 )
 
@@ -10,19 +11,29 @@ type ResponseWriterInterceptor struct {
 	Body       []byte
 }
 
-func (w *ResponseWriterInterceptor) WriteHeader(statusCode int) {
-	w.StatusCode = statusCode
-	w.ResponseWriter.WriteHeader(statusCode)
-}
-
-func (w *ResponseWriterInterceptor) Write(p []byte) (int, error) {
-	w.Body = append(w.Body, p...)
-	return w.ResponseWriter.Write(p)
-}
-
 func NewResponseWriterInterceptor(writer http.ResponseWriter) *ResponseWriterInterceptor {
-	return &ResponseWriterInterceptor{
-		StatusCode:     http.StatusOK,
-		ResponseWriter: writer,
+	rw, ok := writer.(*ResponseWriterInterceptor)
+	if ok {
+		return rw
+	}
+	rw = &ResponseWriterInterceptor{
+		StatusCode: http.StatusOK,
 	}
+	wrapper := httpsnoop.Wrap(writer, httpsnoop.Hooks{
+		WriteHeader: func(next httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
+			return func(code int) {
+				next(code)
+				rw.StatusCode = code
+			}
+		},
+		Write: func(next httpsnoop.WriteFunc) httpsnoop.WriteFunc {
+			return func(p []byte) (int, error) {
+				n, err := next(p)
+				rw.Body = append(rw.Body, p...)
+				return n, err
+			}
+		},
+	})
+	rw.ResponseWriter = wrapper
+	return rw
 }