dkforest

commit 647b0c28f9a1ab2b33cac31133e4d89197df2392
parent c56a8bf9efa353ba0bb73ad4a1cd8244ee6f33bd
Author: n0tr1v <n0tr1v@protonmail.com>
Date:   Wed, 14 Dec 2022 21:17:28 -0500

cache SetD helper

Diffstat:
M
pkg/cache/cache.go
 | 
7
+++++++
M
pkg/web/handlers/handlers.go
 | 
20
++++++++++----------

2 files changed, 17 insertions(+), 10 deletions(-)
diff --git a/pkg/cache/cache.go b/pkg/cache/cache.go
@@ -85,6 +85,13 @@ func (c *Cache[K, V]) Set(k K, v V, d time.Duration) {
 	c.mtx.Unlock()
 }
 
+// SetD same as Set, but use the DefaultExpiration automatically
+func (c *Cache[K, V]) SetD(k K, v V) {
+	c.mtx.Lock()
+	c.set(k, v, c.defaultExpiration)
+	c.mtx.Unlock()
+}
+
 // Add an item to the cache only if an item doesn't already exist for the given
 // key, or if the existing item has expired. Returns an error otherwise.
 func (c *Cache[K, V]) Add(k K, v V, d time.Duration) error {
diff --git a/pkg/web/handlers/handlers.go b/pkg/web/handlers/handlers.go
@@ -326,15 +326,15 @@ func loginHandler(c echo.Context) error {
 			if user.GpgTwoFactorEnabled {
 				token := utils.GenerateToken32()
 				if user.GpgTwoFactorMode {
-					partialAuthCache.Set(token, NewPartialAuthItem(user.ID, PgpSignStep), cache.DefaultExpiration)
+					partialAuthCache.SetD(token, NewPartialAuthItem(user.ID, PgpSignStep))
 					return SessionsGpgSignTwoFactorHandler(c, true, token)
 				}
-				partialAuthCache.Set(token, NewPartialAuthItem(user.ID, PgpStep), cache.DefaultExpiration)
+				partialAuthCache.SetD(token, NewPartialAuthItem(user.ID, PgpStep))
 				return SessionsGpgTwoFactorHandler(c, true, token)
 
 			} else if string(user.TwoFactorSecret) != "" {
 				token := utils.GenerateToken32()
-				partialAuthCache.Set(token, NewPartialAuthItem(user.ID, TwoFactorStep), cache.DefaultExpiration)
+				partialAuthCache.SetD(token, NewPartialAuthItem(user.ID, TwoFactorStep))
 				return SessionsTwoFactorHandler(c, true, token)
 			}
 
@@ -468,7 +468,7 @@ func SessionsGpgTwoFactorHandler(c echo.Context, step1 bool, token string) error
 
 	if string(user.TwoFactorSecret) != "" {
 		token := utils.GenerateToken32()
-		partialAuthCache.Set(token, NewPartialAuthItem(user.ID, TwoFactorStep), cache.DefaultExpiration)
+		partialAuthCache.SetD(token, NewPartialAuthItem(user.ID, TwoFactorStep))
 		return SessionsTwoFactorHandler(c, true, token)
 	}
 
@@ -512,7 +512,7 @@ func SessionsGpgSignTwoFactorHandler(c echo.Context, step1 bool, token string) e
 
 	if string(user.TwoFactorSecret) != "" {
 		token := utils.GenerateToken32()
-		partialAuthCache.Set(token, NewPartialAuthItem(user.ID, TwoFactorStep), cache.DefaultExpiration)
+		partialAuthCache.SetD(token, NewPartialAuthItem(user.ID, TwoFactorStep))
 		return SessionsTwoFactorHandler(c, true, token)
 	}
 
@@ -746,7 +746,7 @@ func SignalCss(c echo.Context) error {
 		info.HelvaticaLoaded = true
 	}
 	info.UpdatedAt = time.Now().Format(time.RFC3339)
-	signupCache.Set(token, info, cache.DefaultExpiration)
+	signupCache.SetD(token, info)
 	return c.NoContent(http.StatusOK)
 }
 
@@ -882,7 +882,7 @@ func signupHandler(c echo.Context) error {
 
 		signupInfo.hasSolvedCaptcha = true
 		data.HasSolvedCaptcha = signupInfo.hasSolvedCaptcha
-		signupCache.Set(signupToken, signupInfo, cache.DefaultExpiration)
+		signupCache.SetD(signupToken, signupInfo)
 
 		config.SignupSucceed.Inc()
 
@@ -1050,7 +1050,7 @@ func forgotPasswordHandler(c echo.Context) error {
 		}
 
 		token := utils.GenerateToken32()
-		partialRecoveryCache.Set(token, PartialRecoveryItem{user.ID, RecoveryCaptchaCompleted}, cache.DefaultExpiration)
+		partialRecoveryCache.SetD(token, PartialRecoveryItem{user.ID, RecoveryCaptchaCompleted})
 
 		data.Token = token
 		data.Step = 2
@@ -1092,7 +1092,7 @@ func forgotPasswordHandler(c echo.Context) error {
 		}
 
 		pgpTokenCache.Delete(userID)
-		partialRecoveryCache.Set(token, PartialRecoveryItem{userID, RecoveryGpgValidated}, cache.DefaultExpiration)
+		partialRecoveryCache.SetD(token, PartialRecoveryItem{userID, RecoveryGpgValidated})
 
 		data.Token = token
 		data.Step = 3
@@ -3632,7 +3632,7 @@ func TwoFactorAuthenticationVerifyHandler(c echo.Context) error {
 	data.QRCode = getImgStr(img)
 	data.Secret = key.Secret()
 	data.RecoveryCode = recovery
-	twoFactorCache.Set(authUser.ID, twoFactorObj{key, recovery}, cache.DefaultExpiration)
+	twoFactorCache.SetD(authUser.ID, twoFactorObj{key, recovery})
 	return c.Render(http.StatusOK, "two-factor-authentication-verify", data)
 }