diff --git a/contrib/net/http/make_responsewriter.go b/contrib/net/http/make_responsewriter.go index b9ff8a2e13..13ac9a8a14 100644 --- a/contrib/net/http/make_responsewriter.go +++ b/contrib/net/http/make_responsewriter.go @@ -52,6 +52,10 @@ import "net/http" // // This code is generated because we have to account for all the permutations // of the interfaces. +// +// In case of any new interfaces or methods we didn't consider here, we also +// implement the rwUnwrapper interface, which is used internally by +// the standard library: https://github.com/golang/go/blob/6d89b38ed86e0bfa0ddaba08dc4071e6bb300eea/src/net/http/responsecontroller.go#L42-L44 func wrapResponseWriter(w http.ResponseWriter) (http.ResponseWriter, *responseWriter) { {{- range .Interfaces }} h{{.}}, ok{{.}} := w.(http.{{.}}) @@ -61,6 +65,7 @@ func wrapResponseWriter(w http.ResponseWriter) (http.ResponseWriter, *responseWr type monitoredResponseWriter interface { http.ResponseWriter Status() int + Unwrap() http.ResponseWriter } switch { {{- range .Combinations }} diff --git a/contrib/net/http/trace.go b/contrib/net/http/trace.go index a019f5819e..b04867c931 100644 --- a/contrib/net/http/trace.go +++ b/contrib/net/http/trace.go @@ -111,3 +111,8 @@ func (w *responseWriter) WriteHeader(status int) { w.ResponseWriter.WriteHeader(status) w.status = status } + +// Unwrap returns the underlying wrapped http.ResponseWriter. +func (w *responseWriter) Unwrap() http.ResponseWriter { + return w.ResponseWriter +} diff --git a/contrib/net/http/trace_gen.go b/contrib/net/http/trace_gen.go index 728dac8694..db04144454 100644 --- a/contrib/net/http/trace_gen.go +++ b/contrib/net/http/trace_gen.go @@ -17,6 +17,10 @@ import "net/http" // // This code is generated because we have to account for all the permutations // of the interfaces. +// +// In case of any new interfaces or methods we didn't consider here, we also +// implement the rwUnwrapper interface, which is used internally by +// the standard library: https://github.com/golang/go/blob/6d89b38ed86e0bfa0ddaba08dc4071e6bb300eea/src/net/http/responsecontroller.go#L42-L44 func wrapResponseWriter(w http.ResponseWriter) (http.ResponseWriter, *responseWriter) { hFlusher, okFlusher := w.(http.Flusher) hPusher, okPusher := w.(http.Pusher) @@ -27,6 +31,7 @@ func wrapResponseWriter(w http.ResponseWriter) (http.ResponseWriter, *responseWr type monitoredResponseWriter interface { http.ResponseWriter Status() int + Unwrap() http.ResponseWriter } switch { case okFlusher && okPusher && okCloseNotifier && okHijacker: diff --git a/contrib/net/http/trace_test.go b/contrib/net/http/trace_test.go index f7018976a8..fe8e062973 100644 --- a/contrib/net/http/trace_test.go +++ b/contrib/net/http/trace_test.go @@ -11,6 +11,7 @@ import ( "net/http" "net/http/httptest" "testing" + "time" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext" @@ -18,6 +19,7 @@ import ( "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestTraceAndServe(t *testing.T) { @@ -357,6 +359,48 @@ func TestTraceAndServeHost(t *testing.T) { }) } +// TestUnwrap tests the implementation of the rwUnwrapper interface, which is used internally +// by the standard library: https://github.com/golang/go/blob/6d89b38ed86e0bfa0ddaba08dc4071e6bb300eea/src/net/http/responsecontroller.go#L42-L44 +// See also: https://github.com/DataDog/dd-trace-go/issues/2674 +func TestUnwrap(t *testing.T) { + h := WrapHandler(deadlineHandler, "service-name", "resource-name") + srv := httptest.NewServer(h) + defer srv.Close() + + resp, err := srv.Client().Get(srv.URL + "/") + require.NoError(t, err) + defer resp.Body.Close() + + b, err := io.ReadAll(resp.Body) + require.NoError(t, err) + respText := string(b) + + assert.Equal(t, http.StatusOK, resp.StatusCode) + assert.Equal(t, "OK", respText) +} + +var deadlineHandler = http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + rc := http.NewResponseController(w) + + // Use the SetReadDeadline and SetWriteDeadline methods, which are not explicitly implemented in the trace_gen.go + // generated file. Before the Unwrap change, these methods returned a "feature not supported" error. + + err := rc.SetReadDeadline(time.Now().Add(10 * time.Second)) + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(err.Error())) + return + } + err = rc.SetWriteDeadline(time.Now().Add(10 * time.Second)) + if err != nil { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(err.Error())) + return + } + w.WriteHeader(http.StatusOK) + w.Write([]byte("OK")) +}) + type noopHandler struct{} func (noopHandler) ServeHTTP(_ http.ResponseWriter, _ *http.Request) {}