From 19d6e4c00da0e916c0060153856008bf659def1b Mon Sep 17 00:00:00 2001
From: a5r0n <32464596+a5r0n@users.noreply.github.com>
Date: Fri, 30 Aug 2024 15:37:40 +0300
Subject: [PATCH] feat: set timeout dynamically - Add TimeoutFunc (#337)

closes #313
---
 auth_jwt.go | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/auth_jwt.go b/auth_jwt.go
index 4ca1809..35db57b 100644
--- a/auth_jwt.go
+++ b/auth_jwt.go
@@ -39,6 +39,8 @@ type GinJWTMiddleware struct {
 
 	// Duration that a jwt token is valid. Optional, defaults to one hour.
 	Timeout time.Duration
+	// Callback function that will override the default timeout duration.
+	TimeoutFunc func(data interface{}) time.Duration
 
 	// This field allows clients to refresh their token until MaxRefresh has passed.
 	// Note that clients can refresh their token in the last moment of MaxRefresh.
@@ -313,6 +315,12 @@ func (mw *GinJWTMiddleware) MiddlewareInit() error {
 		mw.Timeout = time.Hour
 	}
 
+	if mw.TimeoutFunc == nil {
+		mw.TimeoutFunc = func(data interface{}) time.Duration {
+			return mw.Timeout
+		}
+	}
+
 	if mw.TimeFunc == nil {
 		mw.TimeFunc = time.Now
 	}
@@ -508,7 +516,7 @@ func (mw *GinJWTMiddleware) LoginHandler(c *gin.Context) {
 		}
 	}
 
-	expire := mw.TimeFunc().Add(mw.Timeout)
+	expire := mw.TimeFunc().Add(mw.TimeoutFunc(claims))
 	claims["exp"] = expire.Unix()
 	claims["orig_iat"] = mw.TimeFunc().Unix()
 	tokenString, err := mw.signedString(token)
@@ -583,7 +591,7 @@ func (mw *GinJWTMiddleware) RefreshToken(c *gin.Context) (string, time.Time, err
 		newClaims[key] = claims[key]
 	}
 
-	expire := mw.TimeFunc().Add(mw.Timeout)
+	expire := mw.TimeFunc().Add(mw.TimeoutFunc(claims))
 	newClaims["exp"] = expire.Unix()
 	newClaims["orig_iat"] = mw.TimeFunc().Unix()
 	tokenString, err := mw.signedString(newToken)
@@ -633,7 +641,7 @@ func (mw *GinJWTMiddleware) TokenGenerator(data interface{}) (string, time.Time,
 		}
 	}
 
-	expire := mw.TimeFunc().Add(mw.Timeout)
+	expire := mw.TimeFunc().Add(mw.TimeoutFunc(claims))
 	claims["exp"] = expire.Unix()
 	claims["orig_iat"] = mw.TimeFunc().Unix()
 	tokenString, err := mw.signedString(token)