From 4c27262402d0c5db2f1edbe9b274c3b902b55026 Mon Sep 17 00:00:00 2001
From: winebarrel <sugawara@winebarrel.jp>
Date: Sat, 8 Feb 2025 16:40:01 +0900
Subject: [PATCH] contrib/gorilla/mux: Add WithStatusCheck()

Add WithStatusCheck() to contrib/gorilla/mux so that the function can be
passed to httptrace.ServeConfig.IsStatusError.

Fixes #2390
---
 contrib/gorilla/mux/mux.go      | 15 ++++++------
 contrib/gorilla/mux/mux_test.go | 42 +++++++++++++++++++++++++++++++++
 contrib/gorilla/mux/option.go   |  9 +++++++
 3 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/contrib/gorilla/mux/mux.go b/contrib/gorilla/mux/mux.go
index 22affde4c9..50008a7494 100644
--- a/contrib/gorilla/mux/mux.go
+++ b/contrib/gorilla/mux/mux.go
@@ -110,13 +110,14 @@ func (r *Router) ServeHTTP(w http.ResponseWriter, req *http.Request) {
 	spanopts = append(spanopts, httptraceinternal.HeaderTagsFromRequest(req, r.config.headerTags))
 	resource := r.config.resourceNamer(r, req)
 	httptrace.TraceAndServe(r.Router, w, req, &httptrace.ServeConfig{
-		Service:     r.config.serviceName,
-		Resource:    resource,
-		FinishOpts:  r.config.finishOpts,
-		SpanOpts:    spanopts,
-		QueryParams: r.config.queryParams,
-		RouteParams: match.Vars,
-		Route:       route,
+		Service:       r.config.serviceName,
+		Resource:      resource,
+		FinishOpts:    r.config.finishOpts,
+		SpanOpts:      spanopts,
+		QueryParams:   r.config.queryParams,
+		RouteParams:   match.Vars,
+		Route:         route,
+		IsStatusError: r.config.isStatusError,
 	})
 }
 
diff --git a/contrib/gorilla/mux/mux_test.go b/contrib/gorilla/mux/mux_test.go
index 8d08812805..e5aceb5637 100644
--- a/contrib/gorilla/mux/mux_test.go
+++ b/contrib/gorilla/mux/mux_test.go
@@ -225,6 +225,48 @@ func TestWithQueryParams(t *testing.T) {
 	assert.Equal("http://localhost/200?<redacted>&id=3&name=5", mt.FinishedSpans()[0].Tags()[ext.HTTPURL])
 }
 
+func TestWithStatusCheck(t *testing.T) {
+	for _, ht := range []struct {
+		name          string
+		code          int
+		hasErr        bool
+		isStatusError func(statusCode int) bool
+	}{
+		{
+			name:          "without-statuscheck",
+			code:          http.StatusInternalServerError,
+			hasErr:        true,
+			isStatusError: nil,
+		},
+		{
+			name:          "with-statuscheck",
+			code:          http.StatusInternalServerError,
+			hasErr:        false,
+			isStatusError: func(statusCode int) bool { return false },
+		},
+	} {
+		t.Run(ht.name, func(t *testing.T) {
+			assert := assert.New(t)
+			mt := mocktracer.Start()
+			defer mt.Stop()
+
+			r := httptest.NewRequest("GET", "/500", nil)
+			w := httptest.NewRecorder()
+			mux := NewRouter(WithStatusCheck(ht.isStatusError))
+			mux.Handle("/500", errorHandler(http.StatusInternalServerError))
+			mux.ServeHTTP(w, r)
+			assert.Equal(ht.code, http.StatusInternalServerError)
+
+			spans := mt.FinishedSpans()
+			assert.Equal(1, len(spans))
+
+			s := spans[0]
+			_, ok := s.Tag(ext.Error).(error)
+			assert.Equal(ht.hasErr, ok)
+		})
+	}
+}
+
 func TestSpanOptions(t *testing.T) {
 	assert := assert.New(t)
 	mt := mocktracer.Start()
diff --git a/contrib/gorilla/mux/option.go b/contrib/gorilla/mux/option.go
index 1346f263f9..65b5ee3624 100644
--- a/contrib/gorilla/mux/option.go
+++ b/contrib/gorilla/mux/option.go
@@ -29,6 +29,7 @@ type routerConfig struct {
 	ignoreRequest func(*http.Request) bool
 	queryParams   bool
 	headerTags    *internal.LockMap
+	isStatusError func(statusCode int) bool
 }
 
 // RouterOption represents an option that can be passed to NewRouter.
@@ -140,3 +141,11 @@ func WithQueryParams() RouterOption {
 		cfg.queryParams = true
 	}
 }
+
+// WithStatusCheck specifies a function fn which reports whether the passed
+// statusCode should be considered an error.
+func WithStatusCheck(fn func(statusCode int) bool) RouterOption {
+	return func(cfg *routerConfig) {
+		cfg.isStatusError = fn
+	}
+}