From 5ea5bbfd4472b0274bb15b04456ed9097efa7630 Mon Sep 17 00:00:00 2001
From: kiyon <kiyonlin@163.com>
Date: Sat, 24 Oct 2020 10:19:40 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=91=B7=20Improve=20csrf=20middleware?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Skip non GET/POST http method
- Delete token if matched
- Use cfg.Expiration instead of cfg.CookieExpires
---
 middleware/csrf/csrf.go | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/middleware/csrf/csrf.go b/middleware/csrf/csrf.go
index 5f976bdf..eab9c034 100644
--- a/middleware/csrf/csrf.go
+++ b/middleware/csrf/csrf.go
@@ -149,7 +149,9 @@ 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) {
+		if (cfg.Next != nil && cfg.Next(c)) ||
+			// Or non GET/POST method
+			(c.Method() != fiber.MethodGet && c.Method() != fiber.MethodPost) {
 			return c.Next()
 		}
 
@@ -162,7 +164,7 @@ func New(config ...Config) fiber.Handler {
 			token = utils.UUID()
 			// Add token with timestamp expiration
 			db.Lock()
-			db.tokens[token] = int64(time.Now().Unix()) + expiration
+			db.tokens[token] = time.Now().Unix() + expiration
 			db.Unlock()
 		} else {
 			// Use the server generated token previously to compare
@@ -187,6 +189,13 @@ func New(config ...Config) fiber.Handler {
 			if !ok || time.Now().Unix() >= t {
 				return fiber.ErrForbidden
 			}
+
+			// Delete token from DB
+			db.Lock()
+			delete(db.tokens, csrf)
+			db.Unlock()
+
+			return c.Next()
 		}
 
 		// Create new cookie to send new CSRF token
@@ -195,7 +204,7 @@ func New(config ...Config) fiber.Handler {
 			Value:    token,
 			Domain:   cfg.Cookie.Domain,
 			Path:     cfg.Cookie.Path,
-			Expires:  time.Now().Add(cfg.CookieExpires),
+			Expires:  time.Now().Add(cfg.Expiration),
 			Secure:   cfg.Cookie.Secure,
 			HTTPOnly: cfg.Cookie.HTTPOnly,
 			SameSite: cfg.Cookie.SameSite,