dkforest

commit d899773a5c062b1a0a132e482aebd7c05127d6f6
parent 511d59e61be09fc23cf73a3a7c88bad0bed2099c
Author: n0tr1v <n0tr1v@protonmail.com>
Date:   Tue,  6 Jun 2023 21:29:32 -0700

cleanup

Diffstat:
M
pkg/web/handlers/handlers.go
 | 
33
+++++++++++++++++++--------------

1 file changed, 19 insertions(+), 14 deletions(-)
diff --git a/pkg/web/handlers/handlers.go b/pkg/web/handlers/handlers.go
@@ -455,6 +455,11 @@ func SessionsGpgTwoFactorHandler(c echo.Context, step1 bool, token string) error
 		return c.Redirect(http.StatusFound, "/")
 	}
 
+	cleanup := func() {
+		pgpTokenCache.Delete(user.ID)
+		partialAuthCache.Delete(token)
+	}
+
 	var data sessionsGpgTwoFactorData
 	data.Token = token
 
@@ -477,15 +482,13 @@ func SessionsGpgTwoFactorHandler(c echo.Context, step1 bool, token string) error
 	if data.Code != pgpToken.Value {
 		item.Attempt++
 		if item.Attempt >= 4 {
-			pgpTokenCache.Delete(user.ID)
-			partialAuthCache.Delete(token)
+			cleanup()
 			return c.Redirect(http.StatusFound, "/")
 		}
 		data.ErrorCode = "invalid code"
 		return c.Render(http.StatusOK, "sessions-gpg-two-factor", data)
 	}
-	pgpTokenCache.Delete(user.ID)
-	partialAuthCache.Delete(token)
+	cleanup()
 
 	if user.HasTotpEnabled() {
 		token := utils.GenerateToken32()
@@ -510,6 +513,11 @@ func SessionsGpgSignTwoFactorHandler(c echo.Context, step1 bool, token string) e
 		return c.Redirect(http.StatusFound, "/")
 	}
 
+	cleanup := func() {
+		pgpTokenCache.Delete(user.ID)
+		partialAuthCache.Delete(token)
+	}
+
 	var data sessionsGpgSignTwoFactorData
 	data.Token = token
 
@@ -528,15 +536,13 @@ func SessionsGpgSignTwoFactorHandler(c echo.Context, step1 bool, token string) e
 	if !utils.PgpCheckSignMessage(pgpToken.PKey, pgpToken.Value, data.SignedMessage) {
 		item.Attempt++
 		if item.Attempt >= 4 {
-			pgpTokenCache.Delete(user.ID)
-			partialAuthCache.Delete(token)
+			cleanup()
 			return c.Redirect(http.StatusFound, "/")
 		}
 		data.ErrorSignedMessage = "invalid signature"
 		return c.Render(http.StatusOK, "sessions-gpg-sign-two-factor", data)
 	}
-	pgpTokenCache.Delete(user.ID)
-	partialAuthCache.Delete(token)
+	cleanup()
 
 	if user.HasTotpEnabled() {
 		token := utils.GenerateToken32()
@@ -554,6 +560,7 @@ func SessionsTwoFactorHandler(c echo.Context, step1 bool, token string) error {
 	if !found || item.Step != TwoFactorStep {
 		return c.Redirect(http.StatusFound, "/")
 	}
+	cleanup := func() { partialAuthCache.Delete(token) }
 
 	var data sessionsTwoFactorData
 	data.Token = token
@@ -568,15 +575,14 @@ func SessionsTwoFactorHandler(c echo.Context, step1 bool, token string) error {
 		if !totp.Validate(code, secret) {
 			item.Attempt++
 			if item.Attempt >= 4 {
-				partialAuthCache.Delete(token)
+				cleanup()
 				return c.Redirect(http.StatusFound, "/")
 			}
 			data.Error = "Two-factor authentication failed."
 			return c.Render(http.StatusOK, "sessions-two-factor", data)
 		}
 
-		partialAuthCache.Delete(token)
-
+		cleanup()
 		return completeLogin(c, user, item.SessionDuration)
 	}
 	return c.Render(http.StatusOK, "sessions-two-factor", data)
@@ -589,6 +595,7 @@ func SessionsTwoFactorRecoveryHandler(c echo.Context, token string) error {
 	if !found {
 		return c.Redirect(http.StatusFound, "/")
 	}
+	cleanup := func() { partialAuthCache.Delete(token) }
 
 	var data sessionsTwoFactorRecoveryData
 	data.Token = token
@@ -603,9 +610,7 @@ func SessionsTwoFactorRecoveryHandler(c echo.Context, token string) error {
 			data.Error = "Recovery code authentication failed"
 			return c.Render(http.StatusOK, "sessions-two-factor-recovery", data)
 		}
-
-		partialAuthCache.Delete(token)
-
+		cleanup()
 		return completeLogin(c, user, item.SessionDuration)
 	}
 	return c.Render(http.StatusOK, "sessions-two-factor-recovery", data)