package usecase_test import ( "context" "crypto/rand" "testing" "time" "github.com/alicebob/miniredis/v2" "github.com/stretchr/testify/require" "github.com/zeromicro/go-zero/core/stores/redis" libcrypto "gateway/internal/library/crypto" redislib "gateway/internal/library/redis" memberconfig "gateway/internal/model/member/config" member "gateway/internal/model/member/domain" "gateway/internal/model/member/repository" libtotp "gateway/internal/model/member/totp" "gateway/internal/model/member/usecase" ) func newTOTPFixture(t *testing.T) *usecase.TOTPUseCaseParam { t.Helper() mr := miniredis.RunT(t) rds, err := redislib.NewClient(redis.RedisConf{Host: mr.Addr(), Type: testRedisTypeNode}) require.NoError(t, err) key := make([]byte, 32) _, err = rand.Read(key) require.NoError(t, err) cipher, err := libcrypto.NewAESGCM(key) require.NoError(t, err) now := time.Unix(1_716_000_000, 0) return &usecase.TOTPUseCaseParam{ Profile: repository.NewMemoryTOTPProfileRepository(), Enroll: repository.NewRedisTOTPEnrollStore(rds), Replay: repository.NewRedisTOTPReplayStore(rds), Cipher: cipher, Config: memberconfig.Config{}.Defaults(), Now: func() time.Time { return now }, } } // currentCode reads the staged secret from the enroll store, decrypts it, // and returns a freshly generated TOTP code for the configured `now`. func currentCode(t *testing.T, param *usecase.TOTPUseCaseParam, tenantID, uid string) string { t.Helper() cipherBlob, err := param.Enroll.Get(context.Background(), tenantID, uid) require.NoError(t, err) secret, err := param.Cipher.Decrypt(cipherBlob) require.NoError(t, err) code, err := libtotp.Generate(secret, param.Now(), time.Duration(param.Config.TOTP.PeriodSeconds)*time.Second, param.Config.TOTP.Digits) require.NoError(t, err) return code } func TestTOTPUseCase_EnrollConfirmAndStatus(t *testing.T) { param := newTOTPFixture(t) uc := usecase.MustTOTPUseCase(*param) dto, err := uc.StartEnroll(context.Background(), "t1", "u1", "u1@example.com") require.NoError(t, err) require.NotEmpty(t, dto.OtpauthURL) require.Equal(t, 6, dto.Digits) require.Equal(t, 30, dto.PeriodSec) code := currentCode(t, param, "t1", "u1") backup, err := uc.ConfirmEnroll(context.Background(), "t1", "u1", code) require.NoError(t, err) require.Len(t, backup, param.Config.TOTP.BackupCodeCount) status, err := uc.Status(context.Background(), "t1", "u1") require.NoError(t, err) require.True(t, status.Enrolled) require.Equal(t, len(backup), status.BackupCodesRemaining) } func TestTOTPUseCase_StartEnroll_AlreadyEnrolled(t *testing.T) { param := newTOTPFixture(t) uc := usecase.MustTOTPUseCase(*param) _, err := uc.StartEnroll(context.Background(), "t1", "u1", "") require.NoError(t, err) code := currentCode(t, param, "t1", "u1") _, err = uc.ConfirmEnroll(context.Background(), "t1", "u1", code) require.NoError(t, err) _, err = uc.StartEnroll(context.Background(), "t1", "u1", "") require.Error(t, err) } func TestTOTPUseCase_ConfirmEnroll_MissingStash(t *testing.T) { param := newTOTPFixture(t) uc := usecase.MustTOTPUseCase(*param) _, err := uc.ConfirmEnroll(context.Background(), "t1", "u1", "123456") require.Error(t, err) } func TestTOTPUseCase_ConfirmEnroll_BadCode(t *testing.T) { param := newTOTPFixture(t) uc := usecase.MustTOTPUseCase(*param) _, err := uc.StartEnroll(context.Background(), "t1", "u1", "") require.NoError(t, err) _, err = uc.ConfirmEnroll(context.Background(), "t1", "u1", "000000") require.Error(t, err) } func TestTOTPUseCase_VerifyCode_SuccessAndReplay(t *testing.T) { param := newTOTPFixture(t) uc := usecase.MustTOTPUseCase(*param) _, err := uc.StartEnroll(context.Background(), "t1", "u1", "") require.NoError(t, err) code := currentCode(t, param, "t1", "u1") _, err = uc.ConfirmEnroll(context.Background(), "t1", "u1", code) require.NoError(t, err) // Use a fresh TOTP code for verification; with `now` static the same code // should validate then be rejected on replay. verifyCode, err := libtotp.Generate(decryptStored(t, param, "t1", "u1"), param.Now(), 30*time.Second, 6) require.NoError(t, err) require.NoError(t, uc.VerifyCode(context.Background(), "t1", "u1", verifyCode)) err = uc.VerifyCode(context.Background(), "t1", "u1", verifyCode) require.Error(t, err) } func TestTOTPUseCase_VerifyCode_BackupCode(t *testing.T) { param := newTOTPFixture(t) uc := usecase.MustTOTPUseCase(*param) _, err := uc.StartEnroll(context.Background(), "t1", "u1", "") require.NoError(t, err) code := currentCode(t, param, "t1", "u1") backup, err := uc.ConfirmEnroll(context.Background(), "t1", "u1", code) require.NoError(t, err) require.NotEmpty(t, backup) require.NoError(t, uc.VerifyCode(context.Background(), "t1", "u1", backup[0])) // Same backup code cannot be reused. require.Error(t, uc.VerifyCode(context.Background(), "t1", "u1", backup[0])) status, err := uc.Status(context.Background(), "t1", "u1") require.NoError(t, err) require.Equal(t, len(backup)-1, status.BackupCodesRemaining) } func TestTOTPUseCase_VerifyCode_NotEnrolled(t *testing.T) { param := newTOTPFixture(t) uc := usecase.MustTOTPUseCase(*param) err := uc.VerifyCode(context.Background(), "t1", "u1", "123456") require.Error(t, err) } func TestTOTPUseCase_Disable(t *testing.T) { param := newTOTPFixture(t) uc := usecase.MustTOTPUseCase(*param) _, err := uc.StartEnroll(context.Background(), "t1", "u1", "") require.NoError(t, err) code := currentCode(t, param, "t1", "u1") _, err = uc.ConfirmEnroll(context.Background(), "t1", "u1", code) require.NoError(t, err) require.NoError(t, uc.Disable(context.Background(), "t1", "u1")) status, err := uc.Status(context.Background(), "t1", "u1") require.NoError(t, err) require.False(t, status.Enrolled) } func TestTOTPUseCase_RegenerateBackupCodes(t *testing.T) { param := newTOTPFixture(t) uc := usecase.MustTOTPUseCase(*param) _, err := uc.StartEnroll(context.Background(), "t1", "u1", "") require.NoError(t, err) code := currentCode(t, param, "t1", "u1") first, err := uc.ConfirmEnroll(context.Background(), "t1", "u1", code) require.NoError(t, err) second, err := uc.RegenerateBackupCodes(context.Background(), "t1", "u1") require.NoError(t, err) require.Len(t, second, len(first)) require.NotEqual(t, first, second) require.Error(t, uc.VerifyCode(context.Background(), "t1", "u1", first[0])) require.NoError(t, uc.VerifyCode(context.Background(), "t1", "u1", second[0])) } func TestTOTPUseCase_RegenerateBackupCodes_NotEnrolled(t *testing.T) { param := newTOTPFixture(t) uc := usecase.MustTOTPUseCase(*param) _, err := uc.RegenerateBackupCodes(context.Background(), "t1", "u1") require.Error(t, err) } // decryptStored is a helper for tests to read the encrypted secret directly // from the profile after enrollment to compute fresh verification codes. func decryptStored(t *testing.T, param *usecase.TOTPUseCaseParam, tenantID, uid string) []byte { t.Helper() rec, err := param.Profile.Get(context.Background(), tenantID, uid) require.NoError(t, err) secret, err := param.Cipher.Decrypt(rec.SecretCipher) require.NoError(t, err) return secret } // silence the unused `member` import; intentional retained for symmetry with // otp_usecase_test.go which references the package for sentinels. var _ = member.ErrTOTPNotEnrolled