From 72c84639bea89379090500a58a84f854e33e7d59 Mon Sep 17 00:00:00 2001
From: Fernando Bandeira <fernando.bandeira94@gmail.com>
Date: Sun, 18 Apr 2021 23:08:01 -0300
Subject: [PATCH] :sparkles: Etag middleware support to custom etags

---
 middleware/etag/etag.go      |  4 ++
 middleware/etag/etag_test.go | 77 +++++++++++++++++++++++++++++++++++-
 2 files changed, 79 insertions(+), 2 deletions(-)

diff --git a/middleware/etag/etag.go b/middleware/etag/etag.go
index ef58e45f..61d3f6f5 100644
--- a/middleware/etag/etag.go
+++ b/middleware/etag/etag.go
@@ -39,6 +39,10 @@ func New(config ...Config) fiber.Handler {
 		if len(body) <= 0 {
 			return
 		}
+		// Skip ETag if header is already present
+		if c.Response().Header.PeekBytes(normalizedHeaderETag) != nil {
+			return
+		}
 
 		// Generate ETag for response
 		bb := bytebufferpool.Get()
diff --git a/middleware/etag/etag_test.go b/middleware/etag/etag_test.go
index 2a2c64d8..b0e4bc76 100644
--- a/middleware/etag/etag_test.go
+++ b/middleware/etag/etag_test.go
@@ -1,14 +1,14 @@
 package etag
 
 import (
+	"bytes"
 	"io/ioutil"
 	"net/http/httptest"
 	"testing"
 
-	"github.com/valyala/fasthttp"
-
 	"github.com/gofiber/fiber/v2"
 	"github.com/gofiber/fiber/v2/utils"
+	"github.com/valyala/fasthttp"
 )
 
 // go test -run Test_ETag_Next
@@ -166,6 +166,79 @@ func testETagWeakEtag(t *testing.T, headerIfNoneMatch, matched bool) {
 	}
 }
 
+// go test -run Test_ETag_CustomEtag
+func Test_ETag_CustomEtag(t *testing.T) {
+	t.Run("without HeaderIfNoneMatch", func(t *testing.T) {
+		testETagCustomEtag(t, false, false)
+	})
+	t.Run("with HeaderIfNoneMatch and not matched", func(t *testing.T) {
+		testETagCustomEtag(t, true, false)
+	})
+	t.Run("with HeaderIfNoneMatch and matched", func(t *testing.T) {
+		testETagCustomEtag(t, true, true)
+	})
+}
+
+func testETagCustomEtag(t *testing.T, headerIfNoneMatch, matched bool) {
+	app := fiber.New()
+
+	app.Use(New())
+
+	app.Get("/", func(c *fiber.Ctx) error {
+		c.Set(fiber.HeaderETag, `"custom"`)
+		if bytes.Equal(c.Request().Header.Peek(fiber.HeaderIfNoneMatch), []byte(`"custom"`)) {
+			return c.SendStatus(fiber.StatusNotModified)
+		}
+		return c.SendString("Hello, World!")
+	})
+
+	req := httptest.NewRequest("GET", "/", nil)
+	if headerIfNoneMatch {
+		etag := `"non-match"`
+		if matched {
+			etag = `"custom"`
+		}
+		req.Header.Set(fiber.HeaderIfNoneMatch, etag)
+	}
+
+	resp, err := app.Test(req)
+	utils.AssertEqual(t, nil, err)
+
+	if !headerIfNoneMatch || !matched {
+		utils.AssertEqual(t, fiber.StatusOK, resp.StatusCode)
+		utils.AssertEqual(t, `"custom"`, resp.Header.Get(fiber.HeaderETag))
+		return
+	}
+
+	if matched {
+		utils.AssertEqual(t, fiber.StatusNotModified, resp.StatusCode)
+		b, err := ioutil.ReadAll(resp.Body)
+		utils.AssertEqual(t, nil, err)
+		utils.AssertEqual(t, 0, len(b))
+	}
+}
+
+// go test -run Test_ETag_CustomEtagPut
+func Test_ETag_CustomEtagPut(t *testing.T) {
+	app := fiber.New()
+
+	app.Use(New())
+
+	app.Put("/", func(c *fiber.Ctx) error {
+		c.Set(fiber.HeaderETag, `"custom"`)
+		if !bytes.Equal(c.Request().Header.Peek(fiber.HeaderIfMatch), []byte(`"custom"`)) {
+			return c.SendStatus(fiber.StatusPreconditionFailed)
+		}
+		return c.SendString("Hello, World!")
+	})
+
+	req := httptest.NewRequest("PUT", "/", nil)
+	req.Header.Set(fiber.HeaderIfMatch, `"non-match"`)
+	resp, err := app.Test(req)
+	utils.AssertEqual(t, nil, err)
+	utils.AssertEqual(t, fiber.StatusPreconditionFailed, resp.StatusCode)
+}
+
 // go test -v -run=^$ -bench=Benchmark_Etag -benchmem -count=4
 func Benchmark_Etag(b *testing.B) {
 	app := fiber.New()