feat: calculate aal without transaction (#1437)

## What kind of change does this PR introduce?

First of a few refactoring PRs:
- Change `NewSession` to take in a `UserID` and `FactorID` 
- Load identities outside of AAL and AMR calculation so
`CalculateAALAndAMR` doesn't touch the database.

This will help with ensuring that the transaction doesn't run for too
long and aid with the Hooks Implementation.

---------

Co-authored-by: joel <joel@joels-MacBook-Pro.local>

internal/api/auth_test.go

@@ -75,9 +75,8 @@ func (ts *AuthTestSuite) TestMaybeLoadUserOrSession() {
 	u, err := models.FindUserByEmailAndAudience(ts.API.db, "test@example.com", ts.Config.JWT.Aud)
 	require.NoError(ts.T(), err)
 
-	s, err := models.NewSession()
+	s, err := models.NewSession(u.ID, nil)
 	require.NoError(ts.T(), err)
-	s.UserID = u.ID
 	require.NoError(ts.T(), ts.API.db.Create(s))
 
 	require.NoError(ts.T(), ts.API.db.Load(s))

internal/api/mfa_test.go

@@ -60,10 +60,8 @@ func (ts *MFATestSuite) SetupTest() {
 	f := models.NewFactor(u, "test_factor", models.TOTP, models.FactorStateUnverified, "secretkey")
 	require.NoError(ts.T(), ts.API.db.Create(f), "Error saving new test factor")
 	// Create corresponding session
-	s, err := models.NewSession()
+	s, err := models.NewSession(u.ID, &f.ID)
 	require.NoError(ts.T(), err, "Error creating test session")
-	s.UserID = u.ID
-	s.FactorID = &f.ID
 	require.NoError(ts.T(), ts.API.db.Create(s), "Error saving test session")
 
 	u, err = models.FindUserByEmailAndAudience(ts.API.db, ts.TestEmail, ts.Config.JWT.Aud)
@@ -223,10 +221,8 @@ func (ts *MFATestSuite) TestMFAVerifyFactor() {
 			require.NoError(ts.T(), ts.API.db.Update(f), "Error updating new test factor")
 
 			// Create session to be invalidated
-			secondarySession, err := models.NewSession()
+			secondarySession, err := models.NewSession(ts.TestUser.ID, &f.ID)
 			require.NoError(ts.T(), err, "Error creating test session")
-			secondarySession.UserID = ts.TestUser.ID
-			secondarySession.FactorID = &f.ID
 			require.NoError(ts.T(), ts.API.db.Create(secondarySession), "Error saving test session")
 
 			token := ts.generateToken(ts.TestUser, r.SessionId)
@@ -304,10 +300,8 @@ func (ts *MFATestSuite) TestUnenrollVerifiedFactor() {
 			factors, err := models.FindFactorsByUser(ts.API.db, ts.TestUser)
 			require.NoError(ts.T(), err, "error finding factors")
 			f := factors[0]
-			secondarySession, err = models.NewSession()
+			secondarySession, err = models.NewSession(ts.TestUser.ID, &f.ID)
 			require.NoError(ts.T(), err, "Error creating test session")
-			secondarySession.UserID = ts.TestUser.ID
-			secondarySession.FactorID = &f.ID
 			require.NoError(ts.T(), ts.API.db.Create(secondarySession), "Error saving test session")
 
 			sharedSecret := ts.TestOTPKey.Secret()
@@ -342,10 +336,8 @@ func (ts *MFATestSuite) TestUnenrollVerifiedFactor() {
 func (ts *MFATestSuite) TestUnenrollUnverifiedFactor() {
 	var secondarySession *models.Session
 	f := ts.TestUser.Factors[0]
-	secondarySession, err := models.NewSession()
+	secondarySession, err := models.NewSession(ts.TestUser.ID, &f.ID)
 	require.NoError(ts.T(), err, "Error creating test session")
-	secondarySession.UserID = ts.TestUser.ID
-	secondarySession.FactorID = &f.ID
 	require.NoError(ts.T(), ts.API.db.Create(secondarySession), "Error saving test session")
 
 	sharedSecret := ts.TestOTPKey.Secret()

internal/api/token.go

@@ -311,7 +311,7 @@ func (a *API) generateAccessToken(ctx context.Context, tx *storage.Connection, u
 		if terr != nil {
 			return "", 0, terr
 		}
-		aal, amr, terr = session.CalculateAALAndAMR(tx)
+		aal, amr, terr = session.CalculateAALAndAMR(user)
 		if terr != nil {
 			return "", 0, terr
 		}
@@ -447,12 +447,15 @@ func (a *API) updateMFASessionAndClaims(r *http.Request, tx *storage.Connection,
 		if terr != nil {
 			return terr
 		}
+		if err := tx.Load(user, "Identities"); err != nil {
+			return err
+		}
 		// Swap to ensure current token is the latest one
 		refreshToken, terr = models.GrantRefreshTokenSwap(r, tx, user, currentToken)
 		if terr != nil {
 			return terr
 		}
-		aal, _, terr := session.CalculateAALAndAMR(tx)
+		aal, _, terr := session.CalculateAALAndAMR(user)
 		if terr != nil {
 			return terr
 		}

internal/models/refresh_token.go

@@ -127,17 +127,11 @@ func createRefreshToken(tx *storage.Connection, user *User, oldToken *RefreshTok
 	}
 
 	if token.SessionId == nil {
-		session, err := NewSession()
+		session, err := NewSession(user.ID, params.FactorID)
 		if err != nil {
 			return nil, errors.Wrap(err, "error instantiating new session object")
 		}
 
-		session.UserID = user.ID
-
-		if params.FactorID != nil {
-			session.FactorID = params.FactorID
-		}
-
 		if params.SessionNotAfter != nil {
 			session.NotAfter = params.SessionNotAfter
 		}

internal/models/sessions.go

@@ -154,14 +154,16 @@ func (s *Session) DetermineTag(tags []string) string {
 	return tags[0]
 }
 
-func NewSession() (*Session, error) {
+func NewSession(userID uuid.UUID, factorID *uuid.UUID) (*Session, error) {
 	id := uuid.Must(uuid.NewV4())
 
 	defaultAAL := AAL1.String()
 
 	session := &Session{
-		ID:  id,
-		AAL: &defaultAAL,
+		ID:       id,
+		AAL:      &defaultAAL,
+		UserID:   userID,
+		FactorID: factorID,
 	}
 
 	return session, nil
@@ -280,7 +282,7 @@ func (s *Session) UpdateAssociatedAAL(tx *storage.Connection, aal string) error
 	return tx.Update(s)
 }
 
-func (s *Session) CalculateAALAndAMR(tx *storage.Connection) (aal string, amr []AMREntry, err error) {
+func (s *Session) CalculateAALAndAMR(user *User) (aal string, amr []AMREntry, err error) {
 	amr, aal = []AMREntry{}, AAL1.String()
 	for _, claim := range s.AMRClaims {
 		if *claim.AuthenticationMethod == TOTPSignIn.String() {
@@ -290,40 +292,22 @@ func (s *Session) CalculateAALAndAMR(tx *storage.Connection) (aal string, amr []
 	}
 
 	// makes sure that the AMR claims are always ordered most-recent first
-
-	// sort in ascending order
-	sort.Sort(sortAMREntries{
-		Array: amr,
-	})
-
-	// now reverse for descending order
-	_ = sort.Reverse(sortAMREntries{
+	sort.Sort(sort.Reverse(sortAMREntries{
 		Array: amr,
-	})
-
-	lastIndex := len(amr) - 1
-
-	if lastIndex > -1 && amr[lastIndex].Method == SSOSAML.String() {
-		// initial AMR claim is from sso/saml, we need to add information
-		// about the provider that was used for the authentication
-		identities, err := FindIdentitiesByUserID(tx, s.UserID)
-		if err != nil {
-			return aal, amr, err
-		}
+	}))
 
-		if len(identities) == 1 {
-			identity := identities[0]
-
-			if strings.HasPrefix(identity.Provider, "sso:") {
-				amr[lastIndex].Provider = strings.TrimPrefix(identity.Provider, "sso:")
-			}
-		}
-
-		// otherwise we can't identify that this user account has only
-		// one SSO identity, so we are not encoding the provider at
-		// this time
+	if len(amr) > 0 && amr[len(amr)-1].Method == SSOSAML.String() {
+		return aal, amr, nil
 	}
-
+	// initial AMR claim is from sso/saml, we need to add information
+	// about the provider that was used for the authentication
+	identities := user.Identities
+	if len(identities) == 1 && identities[0].IsForSSOProvider() {
+		amr[len(amr)-1].Provider = strings.TrimPrefix(identities[0].Provider, "sso:")
+	}
+	// otherwise we can't identify that this user account has only
+	// one SSO identity, so we are not encoding the provider at
+	// this time
 	return aal, amr, nil
 }
 

internal/models/sessions_test.go

@@ -43,9 +43,8 @@ func TestSession(t *testing.T) {
 func (ts *SessionsTestSuite) TestFindBySessionIDWithForUpdate() {
 	u, err := FindUserByEmailAndAudience(ts.db, "test@example.com", ts.Config.JWT.Aud)
 	require.NoError(ts.T(), err)
-	session, err := NewSession()
+	session, err := NewSession(u.ID, nil)
 	require.NoError(ts.T(), err)
-	session.UserID = u.ID
 	require.NoError(ts.T(), ts.db.Create(session))
 
 	found, err := FindSessionByID(ts.db, session.ID, true)
@@ -58,9 +57,8 @@ func (ts *SessionsTestSuite) TestCalculateAALAndAMR() {
 	totalDistinctClaims := 2
 	u, err := FindUserByEmailAndAudience(ts.db, "test@example.com", ts.Config.JWT.Aud)
 	require.NoError(ts.T(), err)
-	session, err := NewSession()
+	session, err := NewSession(u.ID, nil)
 	require.NoError(ts.T(), err)
-	session.UserID = u.ID
 	require.NoError(ts.T(), ts.db.Create(session))
 
 	err = AddClaimToSession(ts.db, session.ID, PasswordGrant)
@@ -72,7 +70,7 @@ func (ts *SessionsTestSuite) TestCalculateAALAndAMR() {
 	session, err = FindSessionByID(ts.db, session.ID, false)
 	require.NoError(ts.T(), err)
 
-	aal, amr, err := session.CalculateAALAndAMR(ts.db)
+	aal, amr, err := session.CalculateAALAndAMR(u)
 	require.NoError(ts.T(), err)
 	require.Equal(ts.T(), AAL2.String(), aal)
 	require.Equal(ts.T(), totalDistinctClaims, len(amr))
@@ -83,7 +81,7 @@ func (ts *SessionsTestSuite) TestCalculateAALAndAMR() {
 	session, err = FindSessionByID(ts.db, session.ID, false)
 	require.NoError(ts.T(), err)
 
-	aal, amr, err = session.CalculateAALAndAMR(ts.db)
+	aal, amr, err = session.CalculateAALAndAMR(u)
 	require.NoError(ts.T(), err)
 
 	require.Equal(ts.T(), AAL2.String(), aal)