modification as per feedback

rudderlabs · Jun 8, 2023 · dc1e219 · dc1e219
1 parent b152bde
commit dc1e219
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 12 deletions.
diff --git a/chiware/stats.go b/chiware/stats.go
@@ -12,11 +12,28 @@ import (
 	"github.com/rudderlabs/rudder-go-kit/stats"
 )
 
-const (
-	UnknownPath = "/unknown"
-)
+type config struct {
+	redactUnknownPaths bool
+}
+
+type Option func(*config)
+
+// RedactUnknownPaths sets the redactUnknownPaths flag.
+// If set to true, the path will be redacted if the route is not found.
+// If set to false, the path will be used as is.
+func RedactUnknownPaths(redactUnknownPaths bool) Option {
+	return func(c *config) {
+		c.redactUnknownPaths = redactUnknownPaths
+	}
+}
 
-func StatMiddleware(ctx context.Context, router chi.Router, s stats.Stats, component string) func(http.Handler) http.Handler {
+func StatMiddleware(ctx context.Context, router chi.Router, s stats.Stats, component string, options ...Option) func(http.Handler) http.Handler {
+	conf := config{
+		redactUnknownPaths: true,
+	}
+	for _, option := range options {
+		option(&conf)
+	}
 	var concurrentRequests int32
 	activeClientCount := s.NewStat(fmt.Sprintf("%s.concurrent_requests_count", component), stats.GaugeType)
 	go func() {
@@ -37,6 +54,9 @@ func StatMiddleware(ctx context.Context, router chi.Router, s stats.Stats, compo
 		if path := chi.RouteContext(r.Context()).RoutePattern(); path != "" {
 			return path
 		}
+		if conf.redactUnknownPaths {
+			return "/redacted"
+		}
 		return r.URL.Path
 	}
 	return func(next http.Handler) http.Handler {
@@ -47,15 +67,11 @@ func StatMiddleware(ctx context.Context, router chi.Router, s stats.Stats, compo
 			defer atomic.AddInt32(&concurrentRequests, -1)
 
 			next.ServeHTTP(sw, r)
-			reqType := getPath(r)
-			if !chi.Routes(router).Match(chi.RouteContext(r.Context()), r.Method, reqType) {
-				reqType = UnknownPath
-			}
 			s.NewSampledTaggedStat(
 				fmt.Sprintf("%s.response_time", component),
 				stats.TimerType,
 				map[string]string{
-					"reqType": reqType,
+					"reqType": getPath(r),
 					"method":  r.Method,
 					"code":    strconv.Itoa(sw.status),
 				}).Since(start)

diff --git a/chiware/stats_test.go b/chiware/stats_test.go
@@ -18,7 +18,7 @@ import (
 
 func TestStatsMiddleware(t *testing.T) {
 	component := "test"
-	testCase := func(expectedStatusCode int, pathTemplate, requestPath, expectedMethod string) func(t *testing.T) {
+	testCase := func(expectedStatusCode int, pathTemplate, requestPath, expectedMethod string, options ...chiware.Option) func(t *testing.T) {
 		return func(t *testing.T) {
 			ctrl := gomock.NewController(t)
 			mockStats := mock_stats.NewMockStats(ctrl)
@@ -40,7 +40,7 @@ func TestStatsMiddleware(t *testing.T) {
 			defer cancel()
 			router := chi.NewRouter()
 			router.Use(
-				chiware.StatMiddleware(ctx, router, mockStats, component),
+				chiware.StatMiddleware(ctx, router, mockStats, component, options...),
 			)
 			router.MethodFunc(expectedMethod, pathTemplate, handler)
 
@@ -53,6 +53,8 @@ func TestStatsMiddleware(t *testing.T) {
 
 	t.Run("template with param in path", testCase(http.StatusNotFound, "/v1/{param}", "/v1/abc", "GET"))
 	t.Run("template without param in path", testCase(http.StatusNotFound, "/v1/some-other/key", "/v1/some-other/key", "GET"))
-	t.Run("template with unknown path ", testCase(http.StatusNotFound, chiware.UnknownPath, "/a/b/c", "GET"))
+	t.Run("template with unknown path ", testCase(http.StatusNotFound, "/a/b/c", "/a/b/c", "GET", chiware.RedactUnknownPaths(false)))
+	t.Run("template with unknown path ", testCase(http.StatusNotFound, "/redacted", "/a/b/c", "GET", chiware.RedactUnknownPaths(true)))
+	t.Run("template with unknown path ", testCase(http.StatusNotFound, "/redacted", "/a/b/c", "GET"))
 
 }