From f05ac635d277ecd9fcdae63f27bf9786bee38d51 Mon Sep 17 00:00:00 2001 From: Rehan Pasha Date: Tue, 10 Sep 2024 10:13:11 -0400 Subject: [PATCH] Squashed commit of the following: commit 9fa866688216dd5f08a22c5f3c01da3337fc400f Merge: 8b0f94e6 8a49875a Author: Pasha, Rehan Date: Sat Sep 7 21:58:48 2024 +0530 Merge branch 'main' into fix-5402 Signed-off-by: Pasha, Rehan commit 8b0f94e6b500c5e2d65d1c65bbcbb63087afe7e7 Author: Pasha, Rehan Date: Sat Sep 7 21:57:41 2024 +0530 Update CHANGELOG.md Signed-off-by: Pasha, Rehan commit f76eb1d70480828fb0626b3285a46a82cd6c13a7 Author: Pasha, Rehan Date: Sat Sep 7 20:35:39 2024 +0530 Delete instrgen/driver/driver Signed-off-by: Pasha, Rehan commit 158e36cae0eff704203349929c919c0d8f4e669d Author: Rehan Pasha Date: Sat Sep 7 11:03:24 2024 -0400 Updating test case and lint Signed-off-by: Rehan Pasha commit 256b7de4f36c5ccaf11c7ceb8117aa8c7cb6529e Author: Rehan Pasha Date: Tue Aug 27 06:28:33 2024 -0400 Adding test case for otelmux and fixing lint Signed-off-by: Rehan Pasha Signed-off-by: Rehan Pasha --- .../github.com/gorilla/mux/otelmux/mux.go | 9 +++ .../gorilla/mux/otelmux/test/mux_test.go | 57 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/instrumentation/github.com/gorilla/mux/otelmux/mux.go b/instrumentation/github.com/gorilla/mux/otelmux/mux.go index c7b2355eca8..0e1168ee20e 100644 --- a/instrumentation/github.com/gorilla/mux/otelmux/mux.go +++ b/instrumentation/github.com/gorilla/mux/otelmux/mux.go @@ -4,7 +4,9 @@ package otelmux // import "go.opentelemetry.io/contrib/instrumentation/github.com/gorilla/mux/otelmux" import ( + "bufio" "fmt" + "net" "net/http" "sync" @@ -76,6 +78,13 @@ type recordingResponseWriter struct { status int } +func (h *recordingResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + if hijacker, ok := h.writer.(http.Hijacker); ok { + return hijacker.Hijack() + } + return nil, nil, fmt.Errorf("underlying ResponseWriter does not support hijacking") +} + var rrwPool = &sync.Pool{ New: func() interface{} { return &recordingResponseWriter{} diff --git a/instrumentation/github.com/gorilla/mux/otelmux/test/mux_test.go b/instrumentation/github.com/gorilla/mux/otelmux/test/mux_test.go index d5e489b2eb6..859b61fcff5 100644 --- a/instrumentation/github.com/gorilla/mux/otelmux/test/mux_test.go +++ b/instrumentation/github.com/gorilla/mux/otelmux/test/mux_test.go @@ -23,6 +23,63 @@ import ( "go.opentelemetry.io/otel/trace" ) +func TestRecordingResponseWriterHijackWithMiddleware(t *testing.T) { + // Create a mock HTTP handler + mockHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hj, ok := w.(http.Hijacker) + require.True(t, ok, "ResponseWriter does not implement http.Hijacker") + + conn, rw, err := hj.Hijack() + require.NoError(t, err) + assert.NotNil(t, conn) + assert.NotNil(t, rw) + + err = conn.Close() + require.NoError(t, err) + }) + + // Wrap the handler with otelmux.Middleware + router := mux.NewRouter() + router.Use(otelmux.Middleware("test-service")) + router.Handle("/hijack", mockHandler) + + // Create a mock HTTP request and response writer + req := httptest.NewRequest("GET", "http://example.com/hijack", nil) + rr := httptest.NewRecorder() + + // Serve the HTTP request using the wrapped handler + router.ServeHTTP(rr, req) + + // Verify the response status + assert.Equal(t, http.StatusOK, rr.Code) +} + +func TestRecordingResponseWriterHijackNonHijackerWithMiddleware(t *testing.T) { + // Create a mock HTTP handler + mockHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + hj, ok := w.(http.Hijacker) + require.False(t, ok, "ResponseWriter should not implement http.Hijacker") + + _, _, err := hj.Hijack() + assert.Error(t, err) + }) + + // Wrap the handler with otelmux.Middleware + router := mux.NewRouter() + router.Use(otelmux.Middleware("test-service")) + router.Handle("/non-hijack", mockHandler) + + // Create a mock HTTP request and response writer + req := httptest.NewRequest("GET", "http://example.com/non-hijack", nil) + rr := httptest.NewRecorder() + + // Serve the HTTP request using the wrapped handler + router.ServeHTTP(rr, req) + + // Verify the response status + assert.Equal(t, http.StatusOK, rr.Code) +} + func TestCustomSpanNameFormatter(t *testing.T) { exporter := tracetest.NewInMemoryExporter()