ajust cache middleware config.Next position

gofiber · Jan 26, 2021 · 08a2e7a · 08a2e7a
1 parent ba7fc03
commit 08a2e7a
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 5 deletions.
diff --git a/middleware/cache/cache.go b/middleware/cache/cache.go
@@ -43,11 +43,6 @@ func New(config ...Config) fiber.Handler {
 
 	// Return new handler
 	return func(c *fiber.Ctx) error {
-		// Don't execute middleware if Next returns true
-		if cfg.Next != nil && cfg.Next(c) {
-			return c.Next()
-		}
-
 		// Only cache GET methods
 		if c.Method() != fiber.MethodGet {
 			return c.Next()
@@ -105,6 +100,11 @@ func New(config ...Config) fiber.Handler {
 			return err
 		}
 
+		// Don't cache response if Next returns true
+		if cfg.Next != nil && cfg.Next(c) {
+			return nil
+		}
+
 		// Cache response
 		e.body = utils.SafeBytes(c.Response().Body())
 		e.status = c.Response().StatusCode()

diff --git a/middleware/cache/cache_test.go b/middleware/cache/cache_test.go
@@ -231,6 +231,44 @@ func Test_Cache_NothingToCache(t *testing.T) {
 	}
 }
 
+func Test_Cache_CustomNext(t *testing.T) {
+	app := fiber.New()
+
+	app.Use(New(Config{
+		Next: func(c *fiber.Ctx) bool {
+			return !(c.Response().StatusCode() == fiber.StatusOK)
+		},
+		CacheControl: true,
+	}))
+
+	app.Get("/", func(c *fiber.Ctx) error {
+		return c.SendString(time.Now().String())
+	})
+
+	app.Get("/error", func(c *fiber.Ctx) error {
+		return c.Status(fiber.StatusInternalServerError).SendString(time.Now().String())
+	})
+
+	resp, err := app.Test(httptest.NewRequest("GET", "/", nil))
+	utils.AssertEqual(t, nil, err)
+	body, err := ioutil.ReadAll(resp.Body)
+	utils.AssertEqual(t, nil, err)
+
+	respCached, err := app.Test(httptest.NewRequest("GET", "/", nil))
+	utils.AssertEqual(t, nil, err)
+	bodyCached, err := ioutil.ReadAll(respCached.Body)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, true, bytes.Equal(body, bodyCached))
+	utils.AssertEqual(t, true, respCached.Header.Get(fiber.HeaderCacheControl) != "")
+
+	_, err = app.Test(httptest.NewRequest("GET", "/error", nil))
+	utils.AssertEqual(t, nil, err)
+
+	errRespCached, err := app.Test(httptest.NewRequest("GET", "/error", nil))
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, true, errRespCached.Header.Get(fiber.HeaderCacheControl) == "")
+}
+
 func Test_CustomKey(t *testing.T) {
 	app := fiber.New()
 	var called bool