From 0e98e3481aa465efca4eedc38cb6f0db7f3a8c98 Mon Sep 17 00:00:00 2001
From: Pijus Kamandulis <pkpjuklas@gmail.com>
Date: Mon, 28 Oct 2024 20:20:52 +0200
Subject: [PATCH] Strip trailing slash using middleware

---
 .../middleware/strip_trailing_slashes.go      | 18 +++++++++++++++
 .../middleware/trailing_slash_stripper.go     | 17 --------------
 api/router.go                                 | 23 ++-----------------
 api/tests/documents_trailingslash_test.go     |  9 ++++++--
 4 files changed, 27 insertions(+), 40 deletions(-)
 create mode 100644 api/handlers/middleware/strip_trailing_slashes.go
 delete mode 100644 api/handlers/middleware/trailing_slash_stripper.go

diff --git a/api/handlers/middleware/strip_trailing_slashes.go b/api/handlers/middleware/strip_trailing_slashes.go
new file mode 100644
index 0000000..ff26ed0
--- /dev/null
+++ b/api/handlers/middleware/strip_trailing_slashes.go
@@ -0,0 +1,18 @@
+package middleware
+
+import (
+	"github.com/gin-gonic/gin"
+)
+
+func StripTrailingSlashes(r *gin.Engine) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		path := c.Request.URL.Path
+		if len(path) > 1 && path[len(path)-1] == '/' {
+			c.Request.URL.Path = path[:len(path)-1]
+			r.HandleContext(c)
+			c.Abort()
+			return
+		}
+		c.Next()
+	}
+}
diff --git a/api/handlers/middleware/trailing_slash_stripper.go b/api/handlers/middleware/trailing_slash_stripper.go
deleted file mode 100644
index 8877f32..0000000
--- a/api/handlers/middleware/trailing_slash_stripper.go
+++ /dev/null
@@ -1,17 +0,0 @@
-package middleware
-
-import (
-	"strings"
-
-	"github.com/gin-gonic/gin"
-)
-
-func TrailingSlashStripper() gin.HandlerFunc {
-	return func(c *gin.Context) {
-		if (len(c.Request.URL.Path)) > 1 { //dont strip root dir slash, path="/"
-			var stripped_path = strings.TrimSuffix(c.Request.URL.Path, "/")
-			c.Request.URL.Path = stripped_path
-		}
-		c.Next()
-	}
-}
diff --git a/api/router.go b/api/router.go
index ed84aab..5408f17 100644
--- a/api/router.go
+++ b/api/router.go
@@ -13,15 +13,15 @@ import (
 )
 
 func CreateRouter() *gin.Engine {
-
 	router := gin.Default(func(e *gin.Engine) {
-		e.RemoveExtraSlash = true
+		e.RedirectTrailingSlash = false
 	})
 
 	if config.Config.Debug {
 		router.Use(middleware.RequestLogger())
 	}
 
+	router.Use(middleware.StripTrailingSlashes(router))
 	router.Use(middleware.Authentication())
 
 	router.GET("/dbs/:databaseId/colls/:collId/pkranges", handlers.GetPartitionKeyRanges)
@@ -52,30 +52,11 @@ func CreateRouter() *gin.Engine {
 
 	router.GET("/cosmium/export", handlers.CosmiumExport)
 
-	addRoutesForTrailingSlashes(router)
-
 	handlers.RegisterExplorerHandlers(router)
 
 	return router
 }
 
-func addRoutesForTrailingSlashes(router *gin.Engine) {
-	trailingSlashGroup := router.Group("/")
-	//prepend, so slash is stripped before authentication middleware reads path
-	trailingSlashGroup.Handlers = prepend(trailingSlashGroup.Handlers, middleware.TrailingSlashStripper())
-
-	for _, route := range router.Routes() {
-		if route.Path != "/" { //don't append slash to root path, already handled by RemoveExtraSlash
-			trailingSlashGroup.Handle(route.Method, route.Path+"/", route.HandlerFunc)
-		}
-	}
-}
-
-func prepend[T any](a []T, e T) []T {
-	a = append([]T{e}, a...)
-	return a
-}
-
 func StartAPI() {
 	if !config.Config.Debug {
 		gin.SetMode(gin.ReleaseMode)
diff --git a/api/tests/documents_trailingslash_test.go b/api/tests/documents_trailingslash_test.go
index 0557266..04d461f 100644
--- a/api/tests/documents_trailingslash_test.go
+++ b/api/tests/documents_trailingslash_test.go
@@ -27,10 +27,15 @@ func Test_Documents_Read_Trailing_Slash(t *testing.T) {
 		req, _ := http.NewRequest("GET", testUrl, nil)
 		req.Header.Add("x-ms-date", date)
 		req.Header.Add("authorization", "sig="+url.QueryEscape(signature))
-		_, err := httpClient.Do(req)
+		res, err := httpClient.Do(req)
 
 		assert.Nil(t, err)
 
+		if res != nil {
+			defer res.Body.Close()
+			assert.Equal(t, http.StatusOK, res.StatusCode, "Expected HTTP status 200 OK")
+		} else {
+			t.FailNow()
+		}
 	})
-
 }