From dd4dd698786aafa0e40d0757607263f9cc766299 Mon Sep 17 00:00:00 2001
From: Tomasz Elendt <tomasz.elendt@gmail.com>
Date: Thu, 4 Jun 2015 23:24:54 +0200
Subject: [PATCH] Solve "The Proxy Problem" of InstrumentHandler
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

InstrumentHandler provides proxy object over given http.ResponseWriter
- responseWriterDelegator that only implements the bare minimum
required of an http.ResponseWriter and doesn’t implement any interface
upgrades (http.Flusher, http.Hijacker, etc.). This commit fixes it by
providing fancyResponseWriterDelegator with all the fancy bells and
whistles if standard library http.ResponseWriter is detected.

Heavily inspired by Goji's middleware.

https://avtok.com/2014/11/05/interface-upgrades.html#the-proxy-problem
---
 prometheus/http.go | 41 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 40 insertions(+), 1 deletion(-)

diff --git a/prometheus/http.go b/prometheus/http.go
index e391af4..eabe602 100644
--- a/prometheus/http.go
+++ b/prometheus/http.go
@@ -14,6 +14,9 @@
 package prometheus
 
 import (
+	"bufio"
+	"io"
+	"net"
 	"net/http"
 	"strconv"
 	"strings"
@@ -141,7 +144,18 @@ func InstrumentHandlerFuncWithOpts(opts SummaryOpts, handlerFunc func(http.Respo
 			urlLen = len(r.URL.String())
 		}
 		go computeApproximateRequestSize(r, out, urlLen)
-		handlerFunc(delegate, r)
+
+		_, cn := w.(http.CloseNotifier)
+		_, fl := w.(http.Flusher)
+		_, hj := w.(http.Hijacker)
+		_, rf := w.(io.ReaderFrom)
+		var rw http.ResponseWriter
+		if cn && fl && hj && rf {
+			rw = &fancyResponseWriterDelegator{delegate}
+		} else {
+			rw = delegate
+		}
+		handlerFunc(rw, r)
 
 		elapsed := float64(time.Since(now)) / float64(time.Microsecond)
 
@@ -197,6 +211,31 @@ func (r *responseWriterDelegator) Write(b []byte) (int, error) {
 	return n, err
 }
 
+type fancyResponseWriterDelegator struct {
+	*responseWriterDelegator
+}
+
+func (f *fancyResponseWriterDelegator) CloseNotify() <-chan bool {
+	return f.ResponseWriter.(http.CloseNotifier).CloseNotify()
+}
+
+func (f *fancyResponseWriterDelegator) Flush() {
+	f.ResponseWriter.(http.Flusher).Flush()
+}
+
+func (f *fancyResponseWriterDelegator) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return f.ResponseWriter.(http.Hijacker).Hijack()
+}
+
+func (f *fancyResponseWriterDelegator) ReadFrom(r io.Reader) (int64, error) {
+	if !f.wroteHeader {
+		f.WriteHeader(http.StatusOK)
+	}
+	n, err := f.ResponseWriter.(io.ReaderFrom).ReadFrom(r)
+	f.written += n
+	return n, err
+}
+
 func sanitizeMethod(m string) string {
 	switch m {
 	case "GET", "get":