From d1cad8c692a27d3dee948f1474930695b910ca40 Mon Sep 17 00:00:00 2001
From: Julien Tant <785518+JulienTant@users.noreply.github.com>
Date: Fri, 27 Sep 2024 08:35:36 -0700
Subject: [PATCH] add reuse check on update MFA (#28304)

---
 server/platform/shared/mfa/mfa.go      | 62 +++++++++++++++-----------
 server/platform/shared/mfa/mfa_test.go | 41 ++++++++++++++---
 2 files changed, 71 insertions(+), 32 deletions(-)

diff --git a/server/platform/shared/mfa/mfa.go b/server/platform/shared/mfa/mfa.go
index 7f138b47f4..9fae26b4eb 100644
--- a/server/platform/shared/mfa/mfa.go
+++ b/server/platform/shared/mfa/mfa.go
@@ -85,27 +85,25 @@ func (m *MFA) GenerateSecret(siteURL, userEmail, userID string) (string, []byte,
 
 // Activate set the mfa as active and store it with the StoreActive function provided
 func (m *MFA) Activate(userMfaSecret, userID string, token string) error {
-	otpConfig := &dgoogauth.OTPConfig{
-		Secret:      userMfaSecret,
-		WindowSize:  3,
-		HotpCounter: 0,
-	}
-
-	trimmedToken := strings.TrimSpace(token)
-
-	ok, err := otpConfig.Authenticate(trimmedToken)
+	usedTs, err := m.store.GetMfaUsedTimestamps(userID)
 	if err != nil {
-		return errors.Wrap(err, "unable to parse the token")
+		return errors.Wrap(err, "unable to retrieve the DisallowReuse slice")
 	}
 
-	if !ok {
-		return InvalidToken
+	otpConfig, err := m.authenticate(userMfaSecret, usedTs, token)
+	if err != nil {
+		return errors.Wrap(err, "unable to authenticate the token")
 	}
 
-	if err := m.store.UpdateMfaActive(userID, true); err != nil {
+	if err = m.store.UpdateMfaActive(userID, true); err != nil {
 		return errors.Wrap(err, "unable to store mfa active")
 	}
 
+	err = m.store.StoreMfaUsedTimestamps(userID, otpConfig.DisallowReuse)
+	if err != nil {
+		return errors.Wrap(err, "unable to store the DisallowReuse slice")
+	}
+
 	return nil
 }
 
@@ -129,21 +127,14 @@ func (m *MFA) ValidateToken(user *model.User, token string) (bool, error) {
 		return false, errors.Wrap(err, "unable to retrieve the DisallowReuse slice")
 	}
 
-	otpConfig := &dgoogauth.OTPConfig{
-		Secret:        user.MfaSecret,
-		WindowSize:    3,
-		HotpCounter:   0,
-		DisallowReuse: usedTs,
-	}
-
-	trimmedToken := strings.TrimSpace(token)
-	ok, err := otpConfig.Authenticate(trimmedToken)
+	otpConfig, err := m.authenticate(user.MfaSecret, usedTs, token)
 	if err != nil {
+		if err == InvalidToken {
+			return false, nil
+		}
+
 		return false, errors.Wrap(err, "unable to parse the token")
 	}
-	if !ok {
-		return false, nil
-	}
 
 	err = m.store.StoreMfaUsedTimestamps(user.Id, otpConfig.DisallowReuse)
 	if err != nil {
@@ -152,3 +143,24 @@ func (m *MFA) ValidateToken(user *model.User, token string) (bool, error) {
 
 	return true, nil
 }
+
+func (*MFA) authenticate(userMfaSecret string, usedTs []int, token string) (*dgoogauth.OTPConfig, error) {
+	trimmedToken := strings.TrimSpace(token)
+
+	otpConfig := &dgoogauth.OTPConfig{
+		Secret:        userMfaSecret,
+		WindowSize:    3,
+		HotpCounter:   0,
+		DisallowReuse: usedTs,
+	}
+
+	ok, err := otpConfig.Authenticate(trimmedToken)
+	if err != nil {
+		return nil, errors.Wrap(err, "unable to parse the token")
+	}
+	if !ok {
+		return nil, InvalidToken
+	}
+
+	return otpConfig, nil
+}
diff --git a/server/platform/shared/mfa/mfa_test.go b/server/platform/shared/mfa/mfa_test.go
index 49c9386daf..bca03905ee 100644
--- a/server/platform/shared/mfa/mfa_test.go
+++ b/server/platform/shared/mfa/mfa_test.go
@@ -77,19 +77,26 @@ func TestActivate(t *testing.T) {
 	token := dgoogauth.ComputeCode(userMfaSecret, time.Now().UTC().Unix()/30)
 
 	t.Run("fail on wrongly formatted token", func(t *testing.T) {
-		err := New(nil).Activate(userMfaSecret, userID, "invalid-token")
+		storeMock := mocks.UserStore{}
+		storeMock.On("GetMfaUsedTimestamps", userID).Return([]int{}, nil).Once()
+
+		err := New(&storeMock).Activate(userMfaSecret, userID, "invalid-token")
 		require.Error(t, err)
 		require.Contains(t, err.Error(), "unable to parse the token")
 	})
 
 	t.Run("fail on invalid token", func(t *testing.T) {
-		err := New(nil).Activate(userMfaSecret, userID, "000000")
+		storeMock := mocks.UserStore{}
+		storeMock.On("GetMfaUsedTimestamps", userID).Return([]int{}, nil).Once()
+
+		err := New(&storeMock).Activate(userMfaSecret, userID, "000000")
 		require.Error(t, err)
 		require.Contains(t, err.Error(), "invalid mfa token")
 	})
 
 	t.Run("fail on store action fail", func(t *testing.T) {
 		storeMock := mocks.UserStore{}
+		storeMock.On("GetMfaUsedTimestamps", userID).Return([]int{}, nil).Once()
 		storeMock.On("UpdateMfaActive", userID, true).Return(func(userId string, active bool) error {
 			return errors.New("failed to update mfa active")
 		})
@@ -100,14 +107,34 @@ func TestActivate(t *testing.T) {
 	})
 
 	t.Run("Successful activate", func(t *testing.T) {
-		storeMock := mocks.UserStore{}
-		storeMock.On("UpdateMfaActive", userID, true).Return(func(userId string, active bool) error {
-			return nil
-		})
+		userID := model.NewId()
+		secret := newRandomBase32String(mfaSecretSize)
 
-		err := New(&storeMock).Activate(userMfaSecret, userID, fmt.Sprintf("%06d", token))
+		t0 := time.Now().UTC().Unix() / 30
+		code := fmt.Sprintf("%06d", dgoogauth.ComputeCode(secret, t0))
+
+		usMock := mocks.UserStore{}
+		usMock.On("GetMfaUsedTimestamps", userID).Return([]int{}, nil).Once()
+		usMock.On("UpdateMfaActive", userID, true).Return(nil).Once()
+		usMock.On("StoreMfaUsedTimestamps", userID, mock.AnythingOfType("[]int")).Return(nil).Once()
+
+		err := New(&usMock).Activate(secret, userID, code)
 		require.NoError(t, err)
 	})
+
+	t.Run("disallow reuse of totp", func(t *testing.T) {
+		userID := model.NewId()
+		secret := newRandomBase32String(mfaSecretSize)
+
+		t0 := time.Now().UTC().Unix() / 30
+		code := fmt.Sprintf("%06d", dgoogauth.ComputeCode(secret, t0))
+
+		usMock := mocks.UserStore{}
+		usMock.On("GetMfaUsedTimestamps", userID).Return([]int{int(t0)}, nil).Once()
+
+		err := New(&usMock).Activate(secret, userID, code)
+		require.Error(t, err)
+	})
 }
 
 func TestDeactivate(t *testing.T) {