dkforest

commit fc17a5cbfdc0669f7332a8f511cbcf3b5361e298
parent c1cc4c60f3e68fa7c64c05627583f79e6e95802e
Author: n0tr1v <n0tr1v@protonmail.com>
Date:   Wed, 14 Dec 2022 10:39:55 -0500

ensure token is used for the right 2fa step

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

1 file changed, 34 insertions(+), 20 deletions(-)
diff --git a/pkg/web/handlers/handlers.go b/pkg/web/handlers/handlers.go
@@ -199,7 +199,20 @@ func HomeHandler(c echo.Context) error {
 // partialAuthCache keep track of partial auth token -> user id.
 // When a user login and have 2fa enabled, we create a "partial" auth cookie.
 // The token can be used to complete the 2fa authentication.
-var partialAuthCache = cache1.New[database.UserID](10*time.Minute, time.Hour)
+var partialAuthCache = cache1.New[PartialAuthItem](10*time.Minute, time.Hour)
+
+type PartialAuthItem struct {
+	UserID database.UserID
+	Step   PartialAuthStep
+}
+
+type PartialAuthStep string
+
+const (
+	TwoFactorStep PartialAuthStep = "2fa"
+	PgpSignStep   PartialAuthStep = "pgp_sign_2fa"
+	PgpStep       PartialAuthStep = "pgp_2fa"
+)
 
 func LoginHandler(c echo.Context) error {
 
@@ -277,15 +290,16 @@ func loginHandler(c echo.Context) error {
 
 			if user.GpgTwoFactorEnabled {
 				token := utils.GenerateToken32()
-				partialAuthCache.Set(token, user.ID, cache1.DefaultExpiration)
 				if user.GpgTwoFactorMode {
+					partialAuthCache.Set(token, PartialAuthItem{user.ID, PgpSignStep}, cache1.DefaultExpiration)
 					return SessionsGpgSignTwoFactorHandler(c, true, token)
 				}
+				partialAuthCache.Set(token, PartialAuthItem{user.ID, PgpStep}, cache1.DefaultExpiration)
 				return SessionsGpgTwoFactorHandler(c, true, token)
 
 			} else if string(user.TwoFactorSecret) != "" {
 				token := utils.GenerateToken32()
-				partialAuthCache.Set(token, user.ID, cache1.DefaultExpiration)
+				partialAuthCache.Set(token, PartialAuthItem{user.ID, TwoFactorStep}, cache1.DefaultExpiration)
 				return SessionsTwoFactorHandler(c, true, token)
 			}
 
@@ -380,14 +394,14 @@ func LoginCompletedHandler(c echo.Context) error {
 
 // SessionsGpgTwoFactorHandler ...
 func SessionsGpgTwoFactorHandler(c echo.Context, step1 bool, token string) error {
-	userID, found := partialAuthCache.Get(token)
-	if !found {
+	item, found := partialAuthCache.Get(token)
+	if !found || item.Step != PgpStep {
 		return c.Redirect(http.StatusFound, "/")
 	}
 
-	user, err := database.GetUserByID(userID)
+	user, err := database.GetUserByID(item.UserID)
 	if err != nil {
-		logrus.Errorf("failed to get user %d", userID)
+		logrus.Errorf("failed to get user %d", item.UserID)
 		return c.Redirect(http.StatusFound, "/")
 	}
 
@@ -419,7 +433,7 @@ func SessionsGpgTwoFactorHandler(c echo.Context, step1 bool, token string) error
 
 	if string(user.TwoFactorSecret) != "" {
 		token := utils.GenerateToken32()
-		partialAuthCache.Set(token, user.ID, cache1.DefaultExpiration)
+		partialAuthCache.Set(token, PartialAuthItem{user.ID, TwoFactorStep}, cache1.DefaultExpiration)
 		return SessionsTwoFactorHandler(c, true, token)
 	}
 
@@ -428,14 +442,14 @@ func SessionsGpgTwoFactorHandler(c echo.Context, step1 bool, token string) error
 
 // SessionsGpgSignTwoFactorHandler ...
 func SessionsGpgSignTwoFactorHandler(c echo.Context, step1 bool, token string) error {
-	userID, found := partialAuthCache.Get(token)
-	if !found {
+	item, found := partialAuthCache.Get(token)
+	if !found || item.Step != PgpSignStep {
 		return c.Redirect(http.StatusFound, "/")
 	}
 
-	user, err := database.GetUserByID(userID)
+	user, err := database.GetUserByID(item.UserID)
 	if err != nil {
-		logrus.Errorf("failed to get user %d", userID)
+		logrus.Errorf("failed to get user %d", item.UserID)
 		return c.Redirect(http.StatusFound, "/")
 	}
 
@@ -463,7 +477,7 @@ func SessionsGpgSignTwoFactorHandler(c echo.Context, step1 bool, token string) e
 
 	if string(user.TwoFactorSecret) != "" {
 		token := utils.GenerateToken32()
-		partialAuthCache.Set(token, user.ID, cache1.DefaultExpiration)
+		partialAuthCache.Set(token, PartialAuthItem{user.ID, TwoFactorStep}, cache1.DefaultExpiration)
 		return SessionsTwoFactorHandler(c, true, token)
 	}
 
@@ -472,8 +486,8 @@ func SessionsGpgSignTwoFactorHandler(c echo.Context, step1 bool, token string) e
 
 // SessionsTwoFactorHandler ...
 func SessionsTwoFactorHandler(c echo.Context, step1 bool, token string) error {
-	userID, found := partialAuthCache.Get(token)
-	if !found {
+	item, found := partialAuthCache.Get(token)
+	if !found || item.Step != TwoFactorStep {
 		return c.Redirect(http.StatusFound, "/")
 	}
 
@@ -481,9 +495,9 @@ func SessionsTwoFactorHandler(c echo.Context, step1 bool, token string) error {
 	data.Token = token
 	if !step1 {
 		code := c.Request().PostFormValue("code")
-		user, err := database.GetUserByID(userID)
+		user, err := database.GetUserByID(item.UserID)
 		if err != nil {
-			logrus.Errorf("failed to get user %d", userID)
+			logrus.Errorf("failed to get user %d", item.UserID)
 			return c.Redirect(http.StatusFound, "/")
 		}
 		secret := string(user.TwoFactorSecret)
@@ -501,7 +515,7 @@ func SessionsTwoFactorHandler(c echo.Context, step1 bool, token string) error {
 
 // SessionsTwoFactorRecoveryHandler ...
 func SessionsTwoFactorRecoveryHandler(c echo.Context, token string) error {
-	userID, found := partialAuthCache.Get(token)
+	item, found := partialAuthCache.Get(token)
 	if !found {
 		return c.Redirect(http.StatusFound, "/")
 	}
@@ -510,9 +524,9 @@ func SessionsTwoFactorRecoveryHandler(c echo.Context, token string) error {
 	data.Token = token
 	recoveryCode := c.Request().PostFormValue("code")
 	if recoveryCode != "" {
-		user, err := database.GetUserByID(userID)
+		user, err := database.GetUserByID(item.UserID)
 		if err != nil {
-			logrus.Errorf("failed to get user %d", userID)
+			logrus.Errorf("failed to get user %d", item.UserID)
 			return c.Redirect(http.StatusFound, "/")
 		}
 		if err := bcrypt.CompareHashAndPassword([]byte(user.TwoFactorRecovery), []byte(recoveryCode)); err != nil {