use httpsnoop to ensure http.ResponseWriter additional interfaces are preserved (#388)

* add a test to make sure the recordingResponseWriter preserves interfaces implemented by the wrapped writer

* use httpsnoop to wrap the response writer so we don't lose interfaces

* don't bother checking for deprecated interfaces

* run precommit

Author: toshok

instrumentation/github.com/gorilla/mux/otelmux/example/go.sum

@@ -13,6 +13,8 @@ github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs
 github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98=
 github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
+github.com/felixge/httpsnoop v1.0.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8SPQ=
+github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
 github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=

instrumentation/github.com/gorilla/mux/otelmux/go.mod

@@ -8,6 +8,7 @@ replace (
 )
 
 require (
+	github.com/felixge/httpsnoop v1.0.1
 	github.com/gorilla/mux v1.8.0
 	github.com/stretchr/testify v1.6.1
 	go.opentelemetry.io/contrib v0.12.0

instrumentation/github.com/gorilla/mux/otelmux/go.sum

@@ -1,6 +1,8 @@
 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.1 h1:lvB5Jl89CsZtGIWuTcDM1E/vkVs49/Ml7JJe07l8SPQ=
+github.com/felixge/httpsnoop v1.0.1/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U=
 github.com/google/go-cmp v0.5.2 h1:X2ev0eStA3AbceY54o37/0PQ/UWqKEiiO2dKL5OPaFM=
 github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI=

instrumentation/github.com/gorilla/mux/otelmux/mux.go

@@ -19,6 +19,7 @@ import (
 	"net/http"
 	"sync"
 
+	"github.com/felixge/httpsnoop"
 	"github.com/gorilla/mux"
 
 	otelcontrib "go.opentelemetry.io/contrib"
@@ -74,27 +75,6 @@ type recordingResponseWriter struct {
 	status  int
 }
 
-func (w *recordingResponseWriter) Header() http.Header {
-	return w.writer.Header()
-}
-
-func (w *recordingResponseWriter) Write(slice []byte) (int, error) {
-	w.writeHeader(http.StatusOK)
-	return w.writer.Write(slice)
-}
-
-func (w *recordingResponseWriter) WriteHeader(statusCode int) {
-	w.writeHeader(statusCode)
-	w.writer.WriteHeader(statusCode)
-}
-
-func (w *recordingResponseWriter) writeHeader(statusCode int) {
-	if !w.written {
-		w.written = true
-		w.status = statusCode
-	}
-}
-
 var rrwPool = &sync.Pool{
 	New: func() interface{} {
 		return &recordingResponseWriter{}
@@ -105,7 +85,26 @@ func getRRW(writer http.ResponseWriter) *recordingResponseWriter {
 	rrw := rrwPool.Get().(*recordingResponseWriter)
 	rrw.written = false
 	rrw.status = 0
-	rrw.writer = writer
+	rrw.writer = httpsnoop.Wrap(writer, httpsnoop.Hooks{
+		Write: func(next httpsnoop.WriteFunc) httpsnoop.WriteFunc {
+			return func(b []byte) (int, error) {
+				if !rrw.written {
+					rrw.written = true
+					rrw.status = http.StatusOK
+				}
+				return next(b)
+			}
+		},
+		WriteHeader: func(next httpsnoop.WriteHeaderFunc) httpsnoop.WriteHeaderFunc {
+			return func(statusCode int) {
+				if !rrw.written {
+					rrw.written = true
+					rrw.status = statusCode
+				}
+				next(statusCode)
+			}
+		},
+	})
 	return rrw
 }
 
@@ -145,7 +144,7 @@ func (tw traceware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	r2 := r.WithContext(ctx)
 	rrw := getRRW(w)
 	defer putRRW(rrw)
-	tw.handler.ServeHTTP(rrw, r2)
+	tw.handler.ServeHTTP(rrw.writer, r2)
 	attrs := semconv.HTTPAttributesFromHTTPStatusCode(rrw.status)
 	spanStatus, spanMessage := semconv.SpanStatusFromHTTPStatusCode(rrw.status)
 	span.SetAttributes(attrs...)

instrumentation/github.com/gorilla/mux/otelmux/mux_test.go

@@ -15,7 +15,10 @@
 package otelmux
 
 import (
+	"bufio"
 	"context"
+	"io"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -183,3 +186,58 @@ func TestPropagationWithCustomPropagators(t *testing.T) {
 
 	router.ServeHTTP(w, r)
 }
+
+type testResponseWriter struct {
+	writer http.ResponseWriter
+}
+
+func (rw *testResponseWriter) Header() http.Header {
+	return rw.writer.Header()
+}
+func (rw *testResponseWriter) Write(b []byte) (int, error) {
+	return rw.writer.Write(b)
+}
+func (rw *testResponseWriter) WriteHeader(statusCode int) {
+	rw.writer.WriteHeader(statusCode)
+}
+
+// implement Hijacker
+func (rw *testResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	return nil, nil, nil
+}
+
+// implement Pusher
+func (rw *testResponseWriter) Push(target string, opts *http.PushOptions) error {
+	return nil
+}
+
+// implement Flusher
+func (rw *testResponseWriter) Flush() {
+}
+
+// implement io.ReaderFrom
+func (rw *testResponseWriter) ReadFrom(r io.Reader) (n int64, err error) {
+	return 0, nil
+}
+
+func TestResponseWriterInterfaces(t *testing.T) {
+	// make sure the recordingResponseWriter preserves interfaces implemented by the wrapped writer
+	provider, _ := mocktrace.NewTracerProviderAndTracer(tracerName)
+
+	router := mux.NewRouter()
+	router.Use(Middleware("foobar", WithTracerProvider(provider)))
+	router.HandleFunc("/user/{id}", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Implements(t, (*http.Hijacker)(nil), w)
+		assert.Implements(t, (*http.Pusher)(nil), w)
+		assert.Implements(t, (*http.Flusher)(nil), w)
+		assert.Implements(t, (*io.ReaderFrom)(nil), w)
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	r := httptest.NewRequest("GET", "/user/123", nil)
+	w := &testResponseWriter{
+		writer: httptest.NewRecorder(),
+	}
+
+	router.ServeHTTP(w, r)
+}