otelhttp: Remove custom wrapper after handling request

Signed-off-by: Janusz Marcinkiewicz <[email protected]>

Author: VirrageS

CHANGELOG.md

@@ -11,6 +11,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 ### Changed
 
 - Jaeger remote sampler's probabilistic strategy now uses the same sampling algorithm as `trace.TraceIDRatioBased` in `go.opentelemetry.io/contrib/samplers/jaegerremote`. (#6892)
+- Remove the custom body wrapper from the request's body after the request is processed to allow body type comparisons with the original type in `go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp` and `go.opentelemetry.io/contrib/instrumentation/github.com/gorilla/mux/otelmux`. (#6914)
 
 ### Removed
 

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

@@ -138,7 +138,12 @@ func (tw traceware) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	// ReadCloser fulfills a certain interface and it is indeed nil or NoBody.
 	bw := request.NewBodyWrapper(r.Body, readRecordFunc)
 	if r.Body != nil && r.Body != http.NoBody {
+		prevBody := r.Body
 		r.Body = bw
+
+		// Restore the original body after the request is processed to avoid issues
+		// with extra wrapper since `http/server.go` later checks type of `r.Body`.
+		defer func() { r.Body = prevBody }()
 	}
 
 	writeRecordFunc := func(int64) {}

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

@@ -16,6 +16,7 @@ import (
 	"github.com/gorilla/mux"
 	"github.com/stretchr/testify/assert"
 
+	"go.opentelemetry.io/contrib/instrumentation/github.com/gorilla/mux/otelmux/internal/request"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
@@ -28,6 +29,23 @@ var sc = trace.NewSpanContext(trace.SpanContextConfig{
 	TraceFlags: trace.FlagsSampled,
 })
 
+func TestRequestBodyWrapper(t *testing.T) {
+	router := mux.NewRouter()
+	router.Use(Middleware("foobar"))
+	router.HandleFunc("/user/{id}", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}))
+
+	r := httptest.NewRequest("POST", "/user/123", strings.NewReader(`{"name":"John Doe","age":30}`))
+	r = r.WithContext(trace.ContextWithRemoteSpanContext(context.Background(), sc))
+	w := httptest.NewRecorder()
+
+	router.ServeHTTP(w, r)
+
+	_, ok := r.Body.(*request.BodyWrapper)
+	assert.Falsef(t, ok, "body should not be wrapped after request is processed")
+}
+
 func TestPassthroughSpanFromGlobalTracer(t *testing.T) {
 	var called bool
 	router := mux.NewRouter()

instrumentation/net/http/otelhttp/handler.go

@@ -140,7 +140,12 @@ func (h *middleware) serveHTTP(w http.ResponseWriter, r *http.Request, next http
 	// ReadCloser fulfills a certain interface and it is indeed nil or NoBody.
 	bw := request.NewBodyWrapper(r.Body, readRecordFunc)
 	if r.Body != nil && r.Body != http.NoBody {
+		prevBody := r.Body
 		r.Body = bw
+
+		// Restore the original body after the request is processed to avoid issues
+		// with extra wrapper since `http/server.go` later checks type of `r.Body`.
+		defer func() { r.Body = prevBody }()
 	}
 
 	writeRecordFunc := func(int64) {}

instrumentation/net/http/otelhttp/handler_test.go

@@ -17,6 +17,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp/internal/request"
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/propagation"
@@ -99,6 +100,8 @@ func TestHandler(t *testing.T) {
 			rr := httptest.NewRecorder()
 			tc.handler(t).ServeHTTP(rr, r)
 			assert.Equal(t, tc.expectedStatusCode, rr.Result().StatusCode) //nolint:bodyclose // False positive for httptest.ResponseRecorder: https://github.com/timakin/bodyclose/issues/59.
+			_, ok := r.Body.(*request.BodyWrapper)
+			assert.Falsef(t, ok, "body should not be wrapped after request is processed")
 		})
 	}
 }