From f34221715868e67112c0f07284819e24d0a42236 Mon Sep 17 00:00:00 2001
From: Todd Treece <360020+toddtreece@users.noreply.github.com>
Date: Wed, 8 May 2024 11:08:43 -0400
Subject: [PATCH] K8s: Add path rewriter filter (#87456)

---
 .../endpoints/filters/path_rewriter.go        | 31 ++++++++++++
 .../endpoints/filters/path_rewriter_test.go   | 47 +++++++++++++++++++
 2 files changed, 78 insertions(+)
 create mode 100644 pkg/apiserver/endpoints/filters/path_rewriter.go
 create mode 100644 pkg/apiserver/endpoints/filters/path_rewriter_test.go

diff --git a/pkg/apiserver/endpoints/filters/path_rewriter.go b/pkg/apiserver/endpoints/filters/path_rewriter.go
new file mode 100644
index 00000000000..f1de341c9c9
--- /dev/null
+++ b/pkg/apiserver/endpoints/filters/path_rewriter.go
@@ -0,0 +1,31 @@
+package filters
+
+import (
+	"net/http"
+	"regexp"
+)
+
+type PathRewriter struct {
+	Pattern     *regexp.Regexp
+	ReplaceFunc func([]string) string
+}
+
+func (r *PathRewriter) Rewrite(path string) (string, bool) {
+	matches := r.Pattern.FindStringSubmatch(path)
+	if matches == nil {
+		return path, false
+	}
+	return r.ReplaceFunc(r.Pattern.FindStringSubmatch(path)), true
+}
+
+func WithPathRewriters(handler http.Handler, rewriters []PathRewriter) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		for _, rewriter := range rewriters {
+			if newPath, ok := rewriter.Rewrite(req.URL.Path); ok {
+				req.URL.Path = newPath
+				break
+			}
+		}
+		handler.ServeHTTP(w, req)
+	})
+}
diff --git a/pkg/apiserver/endpoints/filters/path_rewriter_test.go b/pkg/apiserver/endpoints/filters/path_rewriter_test.go
new file mode 100644
index 00000000000..2340c3ff1f8
--- /dev/null
+++ b/pkg/apiserver/endpoints/filters/path_rewriter_test.go
@@ -0,0 +1,47 @@
+package filters
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"regexp"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_WithPathRewriters(t *testing.T) {
+	mockHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, err := w.Write([]byte(r.URL.Path))
+		require.NoError(t, err)
+	})
+
+	rewriters := []PathRewriter{
+		{
+			Pattern: regexp.MustCompile(`(/apis/scope.grafana.app/.*/query)(/.*)`),
+			ReplaceFunc: func(matches []string) string {
+				return matches[1]
+			},
+		},
+	}
+	handler := WithPathRewriters(mockHandler, rewriters)
+
+	t.Run("should rewrite path", func(t *testing.T) {
+		req, err := http.NewRequest("GET", "/apis/scope.grafana.app/namespaces/stack-1234/query/blah", nil)
+		assert.NoError(t, err)
+		rr := httptest.NewRecorder()
+		handler.ServeHTTP(rr, req)
+		assert.Equal(t, http.StatusOK, rr.Code)
+		assert.Equal(t, "/apis/scope.grafana.app/namespaces/stack-1234/query", rr.Body.String())
+	})
+
+	t.Run("should ignore requests that don't match", func(t *testing.T) {
+		req, err := http.NewRequest("GET", "/apis/scope.grafana.app/namespaces/stack-1234/scopes/1", nil)
+		assert.NoError(t, err)
+		rr := httptest.NewRecorder()
+		handler.ServeHTTP(rr, req)
+		assert.Equal(t, http.StatusOK, rr.Code)
+		assert.Equal(t, "/apis/scope.grafana.app/namespaces/stack-1234/scopes/1", rr.Body.String())
+	})
+}