Handle capture on redefined http.responseWriters

Co-authored-by: Mathieu Lonjaret <mathieu.lonjaret@gmail.com>

pkg/middlewares/accesslog/logger.go

@@ -227,6 +227,15 @@ func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request, next http
 		core[ClientHost] = forwardedFor
 	}
 
+	ctx := req.Context()
+	capt, err := capture.FromContext(ctx)
+	if err != nil {
+		log.FromContext(log.With(ctx, log.Str(log.MiddlewareType, "AccessLogs"))).
+			WithError(err).
+			Errorf("Could not get Capture")
+		return
+	}
+
 	next.ServeHTTP(rw, reqWithDataTable)
 
 	if _, ok := core[ClientUsername]; !ok {
@@ -237,13 +246,6 @@ func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request, next http
 		headers: rw.Header().Clone(),
 	}
 
-	ctx := req.Context()
-	capt, err := capture.FromContext(ctx)
-	if err != nil {
-		log.FromContext(log.With(ctx, log.Str(log.MiddlewareType, "AccessLogs"))).Errorf("Could not get Capture: %v", err)
-		return
-	}
-
 	logDataTable.DownstreamResponse.status = capt.StatusCode()
 	logDataTable.DownstreamResponse.size = capt.ResponseSize()
 	logDataTable.Request.size = capt.RequestSize()

pkg/middlewares/accesslog/logger_test.go

@@ -57,7 +57,7 @@ func TestLogRotation(t *testing.T) {
 	})
 
 	chain := alice.New()
-	chain = chain.Append(capture.WrapHandler(&capture.Handler{}))
+	chain = chain.Append(capture.Wrap)
 	chain = chain.Append(WrapHandler(logHandler))
 	handler, err := chain.Then(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		rw.WriteHeader(http.StatusOK)
@@ -210,7 +210,7 @@ func TestLoggerHeaderFields(t *testing.T) {
 			}
 
 			chain := alice.New()
-			chain = chain.Append(capture.WrapHandler(&capture.Handler{}))
+			chain = chain.Append(capture.Wrap)
 			chain = chain.Append(WrapHandler(logger))
 			handler, err := chain.Then(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 				rw.WriteHeader(http.StatusOK)
@@ -784,7 +784,7 @@ func doLoggingTLSOpt(t *testing.T, config *types.AccessLog, enableTLS bool) {
 	}
 
 	chain := alice.New()
-	chain = chain.Append(capture.WrapHandler(&capture.Handler{}))
+	chain = chain.Append(capture.Wrap)
 	chain = chain.Append(WrapHandler(logger))
 	handler, err := chain.Then(http.HandlerFunc(logWriterTestHandlerFunc))
 	require.NoError(t, err)

pkg/middlewares/capture/capture.go

@@ -3,9 +3,8 @@
 // For another middleware to get those attributes of a request/response, this middleware
 // should be added before in the middleware chain.
 //
-//	handler, _ := NewHandler()
 //	chain := alice.New().
-//	     Append(WrapHandler(handler)).
+//	     Append(capture.Wrap).
 //	     Append(myOtherMiddleware).
 //	     then(...)
 //
@@ -33,7 +32,6 @@ import (
 	"net"
 	"net/http"
 
-	"github.com/containous/alice"
 	"github.com/traefik/traefik/v2/pkg/middlewares"
 )
 
@@ -41,62 +39,67 @@ type key string
 
 const capturedData key = "capturedData"
 
-// Handler will store each request data to its context.
-type Handler struct{}
-
-// WrapHandler wraps capture handler into an Alice Constructor.
-func WrapHandler(handler *Handler) alice.Constructor {
-	return func(next http.Handler) (http.Handler, error) {
-		return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
-			handler.ServeHTTP(rw, req, next)
-		}), nil
-	}
+// Wrap returns a new handler that inserts a Capture into the given handler.
+// It satisfies the alice.Constructor type.
+func Wrap(handler http.Handler) (http.Handler, error) {
+	c := Capture{}
+	return c.Reset(handler), nil
 }
 
-func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request, next http.Handler) {
-	c := Capture{}
-	if req.Body != nil {
-		readCounter := &readCounter{source: req.Body}
-		c.rr = readCounter
-		req.Body = readCounter
+// FromContext returns the Capture value found in ctx, or an empty Capture otherwise.
+func FromContext(ctx context.Context) (Capture, error) {
+	c := ctx.Value(capturedData)
+	if c == nil {
+		return Capture{}, errors.New("value not found in context")
 	}
-	responseWriter := newResponseWriter(rw)
-	c.rw = responseWriter
-	ctx := context.WithValue(req.Context(), capturedData, &c)
-	next.ServeHTTP(responseWriter, req.WithContext(ctx))
+	capt, ok := c.(*Capture)
+	if !ok {
+		return Capture{}, errors.New("value stored in context is not a *Capture")
+	}
+	return *capt, nil
 }
 
 // Capture is the object populated by the capture middleware,
-// allowing to gather information about the request and response.
+// holding probes that allow to gather information about the request and response.
 type Capture struct {
 	rr *readCounter
 	rw responseWriter
 }
 
-// FromContext returns the Capture value found in ctx, or an empty Capture otherwise.
-func FromContext(ctx context.Context) (*Capture, error) {
-	c := ctx.Value(capturedData)
-	if c == nil {
-		return nil, errors.New("value not found")
-	}
-	capt, ok := c.(*Capture)
-	if !ok {
-		return nil, errors.New("value stored in Context is not a *Capture")
-	}
-	return capt, nil
+// NeedsReset returns whether the given http.ResponseWriter is the capture's probe.
+func (c *Capture) NeedsReset(rw http.ResponseWriter) bool {
+	return c.rw != rw
 }
 
-func (c Capture) ResponseSize() int64 {
+// Reset returns a new handler that renews the Capture's probes, and inserts
+// them when deferring to next.
+func (c *Capture) Reset(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		ctx := context.WithValue(req.Context(), capturedData, c)
+		newReq := req.WithContext(ctx)
+
+		if newReq.Body != nil {
+			readCounter := &readCounter{source: newReq.Body}
+			c.rr = readCounter
+			newReq.Body = readCounter
+		}
+		c.rw = newResponseWriter(rw)
+
+		next.ServeHTTP(c.rw, newReq)
+	})
+}
+
+func (c *Capture) ResponseSize() int64 {
 	return c.rw.Size()
 }
 
-func (c Capture) StatusCode() int {
+func (c *Capture) StatusCode() int {
 	return c.rw.Status()
 }
 
 // RequestSize returns the size of the request's body if it applies,
 // zero otherwise.
-func (c Capture) RequestSize() int64 {
+func (c *Capture) RequestSize() int64 {
 	if c.rr == nil {
 		return 0
 	}

pkg/middlewares/capture/capture_test.go

@@ -38,9 +38,8 @@ func TestCapture(t *testing.T) {
 		assert.Equal(t, "bar", string(all))
 	})
 
-	wrapped := WrapHandler(&Handler{})
 	chain := alice.New()
-	chain = chain.Append(wrapped)
+	chain = chain.Append(Wrap)
 	chain = chain.Append(wrapMiddleware)
 	handlers, err := chain.Then(handler)
 	require.NoError(t, err)
@@ -142,8 +141,7 @@ func BenchmarkCapture(b *testing.B) {
 
 			chain := alice.New()
 			if test.capture || test.body {
-				captureWrapped := WrapHandler(&Handler{})
-				chain = chain.Append(captureWrapped)
+				chain = chain.Append(Wrap)
 			}
 			handlers, err := chain.Then(next)
 			require.NoError(b, err)

pkg/middlewares/metrics/metrics.go

@@ -24,6 +24,7 @@ const (
 	protoWebsocket = "websocket"
 	typeName       = "Metrics"
 	nameEntrypoint = "metrics-entrypoint"
+	nameRouter     = "metrics-router"
 	nameService    = "metrics-service"
 )
 
@@ -56,7 +57,7 @@ func NewEntryPointMiddleware(ctx context.Context, next http.Handler, registry me
 
 // NewRouterMiddleware creates a new metrics middleware for a Router.
 func NewRouterMiddleware(ctx context.Context, next http.Handler, registry metrics.Registry, routerName string, serviceName string) http.Handler {
-	log.FromContext(middlewares.GetLoggerCtx(ctx, nameEntrypoint, typeName)).Debug("Creating middleware")
+	log.FromContext(middlewares.GetLoggerCtx(ctx, nameRouter, typeName)).Debug("Creating middleware")
 
 	return &metricsMiddleware{
 		next:                 next,
@@ -125,17 +126,25 @@ func (m *metricsMiddleware) ServeHTTP(rw http.ResponseWriter, req *http.Request)
 		m.reqsTLSCounter.With(tlsLabels...).Add(1)
 	}
 
-	start := time.Now()
-
-	m.next.ServeHTTP(rw, req)
-
 	ctx := req.Context()
+
 	capt, err := capture.FromContext(ctx)
 	if err != nil {
-		log.FromContext(middlewares.GetLoggerCtx(ctx, nameEntrypoint, typeName)).Errorf("Could not get Capture: %w", err)
+		for i := 0; i < len(m.baseLabels); i += 2 {
+			ctx = log.With(ctx, log.Str(m.baseLabels[i], m.baseLabels[i+1]))
+		}
+		log.FromContext(ctx).WithError(err).Errorf("Could not get Capture")
 		return
 	}
 
+	next := m.next
+	if capt.NeedsReset(rw) {
+		next = capt.Reset(m.next)
+	}
+
+	start := time.Now()
+	next.ServeHTTP(rw, req)
+
 	labels = append(labels, "code", strconv.Itoa(capt.StatusCode()))
 	m.reqDurationHistogram.With(labels...).ObserveFromStart(start)
 	m.reqsCounter.With(labels...).Add(1)