From bdf4c6723fba3de8a4e01ee621c7e2b70d2b42f5 Mon Sep 17 00:00:00 2001
From: mpl <mathieu.lonjaret@gmail.com>
Date: Tue, 10 Dec 2019 18:18:04 +0100
Subject: [PATCH] detect CloseNotify capability in accesslog and metrics

---
 .../accesslog/capture_response_writer.go      | 24 ++++++++++++++
 pkg/middlewares/accesslog/field_middleware.go |  2 +-
 pkg/middlewares/accesslog/logger.go           |  2 +-
 pkg/middlewares/metrics/metrics.go            |  4 +--
 pkg/middlewares/metrics/recorder.go           | 31 +++++++++++++++++++
 5 files changed, 59 insertions(+), 4 deletions(-)

diff --git a/pkg/middlewares/accesslog/capture_response_writer.go b/pkg/middlewares/accesslog/capture_response_writer.go
index 094b5c922..da4991a72 100644
--- a/pkg/middlewares/accesslog/capture_response_writer.go
+++ b/pkg/middlewares/accesslog/capture_response_writer.go
@@ -13,6 +13,20 @@ var (
 	_ middlewares.Stateful = &captureResponseWriter{}
 )
 
+type capturer interface {
+	http.ResponseWriter
+	Size() int64
+	Status() int
+}
+
+func newCaptureResponseWriter(rw http.ResponseWriter) capturer {
+	capt := &captureResponseWriter{rw: rw}
+	if _, ok := rw.(http.CloseNotifier); !ok {
+		return capt
+	}
+	return captureResponseWriterWithCloseNotify{capt}
+}
+
 // captureResponseWriter is a wrapper of type http.ResponseWriter
 // that tracks request status and size
 type captureResponseWriter struct {
@@ -21,6 +35,16 @@ type captureResponseWriter struct {
 	size   int64
 }
 
+type captureResponseWriterWithCloseNotify struct {
+	*captureResponseWriter
+}
+
+// CloseNotify returns a channel that receives at most a
+// single value (true) when the client connection has gone away.
+func (r *captureResponseWriterWithCloseNotify) CloseNotify() <-chan bool {
+	return r.rw.(http.CloseNotifier).CloseNotify()
+}
+
 func (crw *captureResponseWriter) Header() http.Header {
 	return crw.rw.Header()
 }
diff --git a/pkg/middlewares/accesslog/field_middleware.go b/pkg/middlewares/accesslog/field_middleware.go
index 19b455939..f201ba53a 100644
--- a/pkg/middlewares/accesslog/field_middleware.go
+++ b/pkg/middlewares/accesslog/field_middleware.go
@@ -49,7 +49,7 @@ func AddServiceFields(rw http.ResponseWriter, req *http.Request, next http.Handl
 
 // AddOriginFields add origin fields
 func AddOriginFields(rw http.ResponseWriter, req *http.Request, next http.Handler, data *LogData) {
-	crw := &captureResponseWriter{rw: rw}
+	crw := newCaptureResponseWriter(rw)
 	start := time.Now().UTC()
 
 	next.ServeHTTP(crw, req)
diff --git a/pkg/middlewares/accesslog/logger.go b/pkg/middlewares/accesslog/logger.go
index 5f9a3c9f9..8114d9319 100644
--- a/pkg/middlewares/accesslog/logger.go
+++ b/pkg/middlewares/accesslog/logger.go
@@ -200,7 +200,7 @@ func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request, next http
 		core[ClientHost] = forwardedFor
 	}
 
-	crw := &captureResponseWriter{rw: rw}
+	crw := newCaptureResponseWriter(rw)
 
 	next.ServeHTTP(crw, reqWithDataTable)
 
diff --git a/pkg/middlewares/metrics/metrics.go b/pkg/middlewares/metrics/metrics.go
index b5f36640a..34d9dccc8 100644
--- a/pkg/middlewares/metrics/metrics.go
+++ b/pkg/middlewares/metrics/metrics.go
@@ -89,10 +89,10 @@ func (m *metricsMiddleware) ServeHTTP(rw http.ResponseWriter, req *http.Request)
 	}(labels)
 
 	start := time.Now()
-	recorder := &responseRecorder{rw, http.StatusOK}
+	recorder := newResponseRecorder(rw)
 	m.next.ServeHTTP(recorder, req)
 
-	labels = append(labels, "code", strconv.Itoa(recorder.statusCode))
+	labels = append(labels, "code", strconv.Itoa(recorder.getCode()))
 	m.reqsCounter.With(labels...).Add(1)
 	m.reqDurationHistogram.With(labels...).Observe(time.Since(start).Seconds())
 }
diff --git a/pkg/middlewares/metrics/recorder.go b/pkg/middlewares/metrics/recorder.go
index 7ada2d988..4206558c7 100644
--- a/pkg/middlewares/metrics/recorder.go
+++ b/pkg/middlewares/metrics/recorder.go
@@ -6,6 +6,23 @@ import (
 	"net/http"
 )
 
+type recorder interface {
+	http.ResponseWriter
+	http.Flusher
+	getCode() int
+}
+
+func newResponseRecorder(rw http.ResponseWriter) recorder {
+	rec := &responseRecorder{
+		ResponseWriter: rw,
+		statusCode:     http.StatusOK,
+	}
+	if _, ok := rw.(http.CloseNotifier); !ok {
+		return rec
+	}
+	return responseRecorderWithCloseNotify{rec}
+}
+
 // responseRecorder captures information from the response and preserves it for
 // later analysis.
 type responseRecorder struct {
@@ -13,6 +30,20 @@ type responseRecorder struct {
 	statusCode int
 }
 
+type responseRecorderWithCloseNotify struct {
+	*responseRecorder
+}
+
+// CloseNotify returns a channel that receives at most a
+// single value (true) when the client connection has gone away.
+func (r *responseRecorderWithCloseNotify) CloseNotify() <-chan bool {
+	return r.ResponseWriter.(http.CloseNotifier).CloseNotify()
+}
+
+func (r *responseRecorder) getCode() int {
+	return r.statusCode
+}
+
 // WriteHeader captures the status code for later retrieval.
 func (r *responseRecorder) WriteHeader(status int) {
 	r.ResponseWriter.WriteHeader(status)