refactor: simplify MFA checks

internal/api/mfa.go

@@ -4,7 +4,6 @@ import (
 	"bytes"
 	"encoding/json"
 	"net/http"
-	"time"
 
 	"net/url"
 
@@ -78,8 +77,7 @@ func (a *API) EnrollFactor(w http.ResponseWriter, r *http.Request) error {
 		return unprocessableEntityError("MFA enrollment only supported for non-SSO users at this time")
 	}
 
-	factorType := params.FactorType
-	if factorType != models.TOTP {
+	if params.FactorType != models.TOTP {
 		return badRequestError("factor_type needs to be totp")
 	}
 
@@ -105,7 +103,7 @@ func (a *API) EnrollFactor(w http.ResponseWriter, r *http.Request) error {
 
 	numVerifiedFactors := 0
 	for _, factor := range factors {
-		if factor.Status == models.FactorStateVerified.String() {
+		if factor.IsVerified() {
 			numVerifiedFactors += 1
 		}
 	}
@@ -190,11 +188,9 @@ func (a *API) ChallengeFactor(w http.ResponseWriter, r *http.Request) error {
 		return err
 	}
 
-	creationTime := challenge.CreatedAt
-	expiryTime := creationTime.Add(time.Second * time.Duration(config.MFA.ChallengeExpiryDuration))
 	return sendJSON(w, http.StatusOK, &ChallengeFactorResponse{
 		ID:        challenge.ID,
-		ExpiresAt: expiryTime.Unix(),
+		ExpiresAt: challenge.GetExpiryTime(config.MFA.ChallengeExpiryDuration).Unix(),
 	})
 }
 
@@ -217,7 +213,7 @@ func (a *API) VerifyFactor(w http.ResponseWriter, r *http.Request) error {
 		return badRequestError("invalid body: unable to parse JSON").WithInternalError(err)
 	}
 
-	if factor.UserID != user.ID {
+	if !factor.IsOwnedBy(user) {
 		return internalServerError(InvalidFactorOwnerErrorMessage)
 	}
 
@@ -233,8 +229,7 @@ func (a *API) VerifyFactor(w http.ResponseWriter, r *http.Request) error {
 		return badRequestError("Challenge and verify IP addresses mismatch")
 	}
 
-	hasExpired := time.Now().After(challenge.CreatedAt.Add(time.Second * time.Duration(config.MFA.ChallengeExpiryDuration)))
-	if hasExpired {
+	if challenge.HasExpired(config.MFA.ChallengeExpiryDuration) {
 		err := a.db.Transaction(func(tx *storage.Connection) error {
 			if terr := tx.Destroy(challenge); terr != nil {
 				return internalServerError("Database error deleting challenge").WithInternalError(terr)
@@ -264,7 +259,7 @@ func (a *API) VerifyFactor(w http.ResponseWriter, r *http.Request) error {
 		if terr = challenge.Verify(tx); terr != nil {
 			return terr
 		}
-		if factor.Status != models.FactorStateVerified.String() {
+		if !factor.IsVerified() {
 			if terr = factor.UpdateStatus(tx, models.FactorStateVerified); terr != nil {
 				return terr
 			}
@@ -306,10 +301,10 @@ func (a *API) UnenrollFactor(w http.ResponseWriter, r *http.Request) error {
 	factor := getFactor(ctx)
 	session := getSession(ctx)
 
-	if factor.Status == models.FactorStateVerified.String() && session.GetAAL() != models.AAL2.String() {
+	if factor.IsVerified() && !session.IsAAL2() {
 		return badRequestError("AAL2 required to unenroll verified factor")
 	}
-	if factor.UserID != user.ID {
+	if !factor.IsOwnedBy(user) {
 		return internalServerError(InvalidFactorOwnerErrorMessage)
 	}
 

internal/api/mfa_test.go

@@ -137,7 +137,7 @@ func (ts *MFATestSuite) TestEnrollFactor() {
 			factors, err := models.FindFactorsByUser(ts.API.db, user)
 			ts.Require().NoError(err)
 			latestFactor := factors[len(factors)-1]
-			require.Equal(ts.T(), models.FactorStateUnverified.String(), latestFactor.Status)
+			require.False(ts.T(), latestFactor.IsVerified())
 			if c.friendlyName != "" && c.expectedCode == http.StatusOK {
 				require.Equal(ts.T(), c.friendlyName, latestFactor.FriendlyName)
 			}
@@ -402,7 +402,7 @@ func (ts *MFATestSuite) TestSessionsMaintainAALOnRefresh() {
 	require.NoError(ts.T(), err)
 	ctx, err = ts.API.maybeLoadUserOrSession(ctx)
 	require.NoError(ts.T(), err)
-	require.Equal(ts.T(), models.AAL2.String(), getSession(ctx).GetAAL())
+	require.True(ts.T(), getSession(ctx).IsAAL2())
 }
 
 // Performing MFA Verification followed by a sign in should return an AAL1 session and an AAL2 session
@@ -431,7 +431,7 @@ func (ts *MFATestSuite) TestMFAFollowedByPasswordSignIn() {
 	require.Equal(ts.T(), models.AAL1.String(), getSession(ctx).GetAAL())
 	session, err := models.FindSessionByUserID(ts.API.db, token.User.ID)
 	require.NoError(ts.T(), err)
-	require.Equal(ts.T(), models.AAL2.String(), session.GetAAL())
+	require.True(ts.T(), session.IsAAL2())
 }
 
 func signUp(ts *MFATestSuite, email, password string) (signUpResp AccessTokenResponse) {

internal/models/challenge.go

@@ -22,8 +22,6 @@ func (Challenge) TableName() string {
 	return tableName
 }
 
-const ChallengePrefix = "challenge"
-
 func NewChallenge(factor *Factor, ipAddress string) (*Challenge, error) {
 	id := uuid.Must(uuid.NewV4())
 
@@ -44,10 +42,18 @@ func FindChallengeByChallengeID(tx *storage.Connection, challengeID uuid.UUID) (
 }
 
 // Update the verification timestamp
-func (f *Challenge) Verify(tx *storage.Connection) error {
+func (c *Challenge) Verify(tx *storage.Connection) error {
 	now := time.Now()
-	f.VerifiedAt = &now
-	return tx.UpdateOnly(f, "verified_at")
+	c.VerifiedAt = &now
+	return tx.UpdateOnly(c, "verified_at")
+}
+
+func (c *Challenge) HasExpired(expiryDuration float64) bool {
+	return time.Now().After(c.GetExpiryTime(expiryDuration))
+}
+
+func (c *Challenge) GetExpiryTime(expiryDuration float64) time.Time {
+	return c.CreatedAt.Add(time.Second * time.Duration(expiryDuration))
 }
 
 func findChallenge(tx *storage.Connection, query string, args ...interface{}) (*Challenge, error) {

internal/models/factor.go

@@ -199,6 +199,14 @@ func (f *Factor) DowngradeSessionsToAAL1(tx *storage.Connection) error {
 	return updateFactorAssociatedSessions(tx, f.UserID, f.ID, AAL1.String())
 }
 
+func (f *Factor) IsOwnedBy(user *User) bool {
+	return f.UserID == user.ID
+}
+
+func (f *Factor) IsVerified() bool {
+	return f.Status == FactorStateVerified.String()
+}
+
 func DeleteFactorsByUserId(tx *storage.Connection, userId uuid.UUID) error {
 	if err := tx.RawQuery("DELETE FROM "+(&pop.Model{Value: Factor{}}).TableName()+" WHERE user_id = ?", userId).Exec(); err != nil {
 		return err

internal/models/sessions.go

@@ -201,3 +201,7 @@ func (s *Session) GetAAL() string {
 	}
 	return *(s.AAL)
 }
+
+func (s *Session) IsAAL2() bool {
+	return s.GetAAL() == AAL2.String()
+}