Embed the Hijacker interface in intercepting writer

go.mod

@@ -17,6 +17,7 @@ require (
 	github.com/davecgh/go-spew v1.1.1
 	github.com/dgrijalva/jwt-go v3.2.0+incompatible
 	github.com/edsrzf/mmap-go v1.0.0 // indirect
+	github.com/felixge/httpsnoop v1.0.2
 	github.com/franela/goblin v0.0.0-20200105215937-c9ffbefa60db // indirect
 	github.com/franela/goreq v0.0.0-20171204163338-bcd34c9993f8 // indirect
 	github.com/go-logfmt/logfmt v0.5.0

go.sum

@@ -72,6 +72,8 @@ github.com/envoyproxy/go-control-plane v0.6.9/go.mod h1:SBwIajubJHhxtWwsL9s8ss4s
 github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
 github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4=
+github.com/felixge/httpsnoop v1.0.2 h1:+nS9g82KMXccJ/wp0zyRW9ZBHFETmMGtkk+2CTTrW4o=
+github.com/felixge/httpsnoop v1.0.2/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/franela/goblin v0.0.0-20200105215937-c9ffbefa60db h1:gb2Z18BhTPJPpLQWj4T+rfKHYCHxRHCtRxhKKjRidVw=
 github.com/franela/goblin v0.0.0-20200105215937-c9ffbefa60db/go.mod h1:7dvUGVsVBjqR7JHJk0brhHOZYGmfBYOrK0ZhYMEtBr4=
 github.com/franela/goreq v0.0.0-20171204163338-bcd34c9993f8 h1:a9ENSRDFBUPkJ5lCgVZh26+ZbGyoVJG7yb5SSzF5H54=

transport/http/server.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"net/http"
 
+	"github.com/felixge/httpsnoop"
 	"github.com/go-kit/kit/endpoint"
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/transport"
@@ -96,15 +97,31 @@ func (s Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 
 	if len(s.finalizer) > 0 {
-		iw := &interceptingWriter{w, http.StatusOK, 0}
+		metrics := httpsnoop.Metrics{Code: http.StatusOK}
+
+		w = httpsnoop.Wrap(w, httpsnoop.Hooks{
+			WriteHeader: func(whf httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
+				return func(code int) {
+					metrics.Code = code
+					whf(code)
+				}
+			},
+			Write: func(wf httpsnoop.WriteFunc) httpsnoop.WriteFunc {
+				return func(b []byte) (int, error) {
+					n, err := wf(b)
+					metrics.Written += int64(n)
+					return n, err
+				}
+			},
+		})
+
 		defer func() {
-			ctx = context.WithValue(ctx, ContextKeyResponseHeaders, iw.Header())
-			ctx = context.WithValue(ctx, ContextKeyResponseSize, iw.written)
+			ctx = context.WithValue(ctx, ContextKeyResponseHeaders, w.Header())
+			ctx = context.WithValue(ctx, ContextKeyResponseSize, metrics.Written)
 			for _, f := range s.finalizer {
-				f(ctx, iw.code, r)
+				f(ctx, metrics.Code, r)
 			}
 		}()
-		w = iw
 	}
 
 	for _, f := range s.before {
@@ -223,22 +240,3 @@ type StatusCoder interface {
 type Headerer interface {
 	Headers() http.Header
 }
-
-type interceptingWriter struct {
-	http.ResponseWriter
-	code    int
-	written int64
-}
-
-// WriteHeader may not be explicitly called, so care must be taken to
-// initialize w.code to its default value of http.StatusOK.
-func (w *interceptingWriter) WriteHeader(code int) {
-	w.code = code
-	w.ResponseWriter.WriteHeader(code)
-}
-
-func (w *interceptingWriter) Write(p []byte) (int, error) {
-	n, err := w.ResponseWriter.Write(p)
-	w.written += int64(n)
-	return n, err
-}