From 18e9dbdfeb1342f42845c059a05ec0514692369e Mon Sep 17 00:00:00 2001
From: joel <joel@joels-MacBook-Pro.local>
Date: Fri, 9 Feb 2024 11:27:26 +0800
Subject: [PATCH] fix: update error codes

---
 internal/api/token.go | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/internal/api/token.go b/internal/api/token.go
index 18233b047..0b3453c2e 100644
--- a/internal/api/token.go
+++ b/internal/api/token.go
@@ -433,9 +433,17 @@ func (a *API) updateMFASessionAndClaims(r *http.Request, tx *storage.Connection,
 	var tokenString string
 	var expiresAt int64
 	var refreshToken *models.RefreshToken
-	session := getSession(ctx)
-	err := tx.Transaction(func(tx *storage.Connection) error {
-		if terr := models.AddClaimToSession(tx, session.ID, authenticationMethod); terr != nil {
+	currentClaims := getClaims(ctx)
+	sessionId, err := uuid.FromString(currentClaims.SessionId)
+	if err != nil {
+		return nil, internalServerError("Cannot read SessionId claim as UUID").WithInternalError(err)
+	}
+	err = tx.Transaction(func(tx *storage.Connection) error {
+		if terr := models.AddClaimToSession(tx, sessionId, authenticationMethod); terr != nil {
+			return terr
+		}
+		session, terr := models.FindSessionByID(tx, sessionId, false)
+		if terr != nil {
 			return terr
 		}
 		currentToken, terr := models.FindTokenBySessionID(tx, &session.ID)
@@ -458,8 +466,8 @@ func (a *API) updateMFASessionAndClaims(r *http.Request, tx *storage.Connection,
 		if err := session.UpdateAssociatedAAL(tx, aal); err != nil {
 			return err
 		}
+		tokenString, expiresAt, terr = a.generateAccessToken(ctx, tx, user, &sessionId, models.TOTPSignIn)
 
-		tokenString, expiresAt, terr = a.generateAccessToken(ctx, tx, user, &session.ID, models.TOTPSignIn)
 		if terr != nil {
 			httpErr, ok := terr.(*HTTPError)
 			if ok {