From d411ec10d6d7d4646288d914d28e4f9ba871c85f Mon Sep 17 00:00:00 2001
From: Jesse Quinn <me@jessequinn.info>
Date: Mon, 24 Jan 2022 12:36:46 -0300
Subject: [PATCH] added HEAD method to caching (#1730)

* added HEAD method to caching

* changed key due to head and get sharing same key

* Update cache.go

- add a improvement task for later

* Update cache.go

correct comment

Co-authored-by: Jesse Quinn <jesse.quinn@zpesystems.com>
Co-authored-by: RW <rene@gofiber.io>
---
 middleware/cache/cache.go      |  7 ++--
 middleware/cache/cache_test.go | 61 ++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/middleware/cache/cache.go b/middleware/cache/cache.go
index 84652b28c0..f27b7638ab 100644
--- a/middleware/cache/cache.go
+++ b/middleware/cache/cache.go
@@ -57,14 +57,15 @@ func New(config ...Config) fiber.Handler {
 
 	// Return new handler
 	return func(c *fiber.Ctx) error {
-		// Only cache GET methods
-		if c.Method() != fiber.MethodGet {
+		// Only cache GET and HEAD methods
+		if c.Method() != fiber.MethodGet && c.Method() != fiber.MethodHead {
 			c.Set(cfg.CacheHeader, cacheUnreachable)
 			return c.Next()
 		}
 
 		// Get key from request
-		key := cfg.KeyGenerator(c)
+		// TODO(allocation optimization): try to minimize the allocation from 2 to 1
+		key := cfg.KeyGenerator(c) + "_" + c.Method()
 
 		// Get entry from pool
 		e := manager.get(key)
diff --git a/middleware/cache/cache_test.go b/middleware/cache/cache_test.go
index e8c272b990..9b301ad1bc 100644
--- a/middleware/cache/cache_test.go
+++ b/middleware/cache/cache_test.go
@@ -341,6 +341,67 @@ func Test_CacheHeader(t *testing.T) {
 	utils.AssertEqual(t, cacheUnreachable, errRespCached.Header.Get("X-Cache"))
 }
 
+func Test_Cache_WithHead(t *testing.T) {
+	app := fiber.New()
+	app.Use(New())
+
+	app.Get("/", func(c *fiber.Ctx) error {
+		now := fmt.Sprintf("%d", time.Now().UnixNano())
+		return c.SendString(now)
+	})
+
+	req := httptest.NewRequest("HEAD", "/", nil)
+	resp, err := app.Test(req)
+	utils.AssertEqual(t, cacheMiss, resp.Header.Get("X-Cache"))
+
+	cachedReq := httptest.NewRequest("HEAD", "/", nil)
+	cachedResp, err := app.Test(cachedReq)
+	utils.AssertEqual(t, cacheHit, cachedResp.Header.Get("X-Cache"))
+
+	body, err := ioutil.ReadAll(resp.Body)
+	utils.AssertEqual(t, nil, err)
+	cachedBody, err := ioutil.ReadAll(cachedResp.Body)
+	utils.AssertEqual(t, nil, err)
+
+	utils.AssertEqual(t, cachedBody, body)
+}
+
+func Test_Cache_WithHeadThenGet(t *testing.T) {
+	app := fiber.New()
+	app.Use(New())
+	app.Get("/get", func(c *fiber.Ctx) error {
+		return c.SendString(c.Query("cache"))
+	})
+
+	headResp, err := app.Test(httptest.NewRequest("HEAD", "/head?cache=123", nil))
+	utils.AssertEqual(t, nil, err)
+	headBody, err := ioutil.ReadAll(headResp.Body)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, "", string(headBody))
+	utils.AssertEqual(t, cacheMiss, headResp.Header.Get("X-Cache"))
+
+	headResp, err = app.Test(httptest.NewRequest("HEAD", "/head?cache=123", nil))
+	utils.AssertEqual(t, nil, err)
+	headBody, err = ioutil.ReadAll(headResp.Body)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, "", string(headBody))
+	utils.AssertEqual(t, cacheHit, headResp.Header.Get("X-Cache"))
+
+	getResp, err := app.Test(httptest.NewRequest("GET", "/get?cache=123", nil))
+	utils.AssertEqual(t, nil, err)
+	getBody, err := ioutil.ReadAll(getResp.Body)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, "123", string(getBody))
+	utils.AssertEqual(t, cacheMiss, getResp.Header.Get("X-Cache"))
+
+	getResp, err = app.Test(httptest.NewRequest("GET", "/get?cache=123", nil))
+	utils.AssertEqual(t, nil, err)
+	getBody, err = ioutil.ReadAll(getResp.Body)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, "123", string(getBody))
+	utils.AssertEqual(t, cacheHit, getResp.Header.Get("X-Cache"))
+}
+
 func Test_CustomCacheHeader(t *testing.T) {
 	app := fiber.New()