From bd875800011f42736c1851f9d9a28617b164c20b Mon Sep 17 00:00:00 2001
From: MrMelon54 <github@mrmelon54.com>
Date: Sun, 30 Mar 2025 00:07:13 +0000
Subject: [PATCH] Fix using no-cache query parameter with url paths

---
 serve/serve.go      | 10 ++++++++--
 serve/serve_test.go | 14 ++++++++++++--
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/serve/serve.go b/serve/serve.go
index 6d00559..5a3d5ea 100644
--- a/serve/serve.go
+++ b/serve/serve.go
@@ -43,7 +43,7 @@ const (
 	BetaSwitchResetQuery = "__bluebell-reset-beta"
 	BetaExpiry           = 24 * time.Hour
 
-	NoCacheQuery = "/?__bluebell-no-cache="
+	NoCacheQuery = "__bluebell-no-cache"
 )
 
 type serveQueries interface {
@@ -60,11 +60,17 @@ type Handler struct {
 }
 
 func cacheBuster(rw http.ResponseWriter, req *http.Request) {
+	q := req.URL.Query()
+	q.Del(BetaSwitchQuery)
+	q.Del(BetaSwitchResetQuery)
+	q.Set(NoCacheQuery, strconv.FormatInt(time.Now().Unix(), 16))
+	req.URL.RawQuery = q.Encode()
+
 	header := rw.Header()
 	header.Set("Cache-Control", "no-cache, no-store, must-revalidate")
 	header.Set("Pragma", "no-cache")
 	header.Set("Expires", "0")
-	http.Redirect(rw, req, NoCacheQuery+strconv.FormatInt(time.Now().Unix(), 16), http.StatusFound)
+	http.Redirect(rw, req, req.URL.RequestURI(), http.StatusFound)
 }
 
 func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
diff --git a/serve/serve_test.go b/serve/serve_test.go
index 6756c8d..f9bf527 100644
--- a/serve/serve_test.go
+++ b/serve/serve_test.go
@@ -46,7 +46,7 @@ func TestHandler_ServeHTTP(t *testing.T) {
 		})
 
 		for _, testPath := range []string{"/", "/posts/test", "/this-is-definitely-a-page-that-I-really-need-to-view-right-now"} {
-			t.Run("switch to "+branch+" branch", func(t *testing.T) {
+			t.Run("switch to "+branch+" branch - "+testPath, func(t *testing.T) {
 				h := New(afero.NewMemMapFs(), &fakeServeDB{})
 
 				req := httptest.NewRequest(http.MethodGet, httpPrefix+"example.com"+testPath+"?__bluebell-switch-beta="+branch, nil)
@@ -58,9 +58,19 @@ func TestHandler_ServeHTTP(t *testing.T) {
 				assert.NotNil(t, res.Body)
 				all, err := io.ReadAll(res.Body)
 				assert.NoError(t, err)
-				assert.Contains(t, string(all), "<a href=\"/?__bluebell-no-cache=")
+				assert.Contains(t, string(all), "<a href=\""+testPath+"?__bluebell-no-cache=")
 				assert.Contains(t, string(all), "\">Found</a>.\n\n")
 
+				location, err := res.Location()
+				if err != nil {
+					return
+				}
+				assert.Equal(t, testPath, location.Path)
+				q := location.Query()
+				assert.True(t, q.Has("__bluebell-no-cache"))
+				assert.False(t, q.Has("__bluebell-switch-beta"))
+				assert.False(t, q.Has("__bluebell-reset-beta"))
+
 				cookies := res.Cookies()
 				assert.Len(t, cookies, 1)
 				assert.Equal(t, cookies[0].Name, "__bluebell-site-beta")