From 48f82536120d83bee266ae2ec5d99b2af8888143 Mon Sep 17 00:00:00 2001
From: Rainnny7 <braydonrainnny@gmail.com>
Date: Thu, 19 Sep 2024 23:50:30 -0400
Subject: [PATCH] use backup codes when checking tfa pins

---
 .../cc/pulseapp/api/service/AuthService.java  | 39 +++++++++++++++++--
 .../cc/pulseapp/api/service/UserService.java  |  5 +--
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/src/main/java/cc/pulseapp/api/service/AuthService.java b/src/main/java/cc/pulseapp/api/service/AuthService.java
index e394414..ddb880f 100644
--- a/src/main/java/cc/pulseapp/api/service/AuthService.java
+++ b/src/main/java/cc/pulseapp/api/service/AuthService.java
@@ -128,10 +128,7 @@ public final class AuthService {
             if (pin == null || (pin.isBlank())) { // No TFA pin received
                 throw new BadRequestException(Error.BORDER_CROSSING);
             }
-            // Incorrect TFA pin
-            if (pin.length() != 6 || (!tfaService.getPin(user.getTfa().getSecret()).equals(pin))) {
-                throw new BadRequestException(Error.TFA_PIN_INVALID);
-            }
+            useTfaPin(user, pin); // Attempt to use the pin
         }
         user.setLastLogin(new Date());
         user = userRepository.save(user);
@@ -139,6 +136,39 @@ public final class AuthService {
                 UserDTO.asDTO(user, new Date(snowflakeService.extractCreationTime(user.getSnowflake()))));
     }
 
+    /**
+     * Use a TFA pin for a user.
+     *
+     * @param user the user to use TFA for
+     * @param pin the pin to use
+     * @throws BadRequestException if using TFA fails
+     */
+    public void useTfaPin(@NonNull User user, @NonNull String pin) throws BadRequestException {
+        if (pin.length() != 6) { // Ensure the pin is the correct length
+            throw new BadRequestException(Error.TFA_PIN_INVALID);
+        }
+        if (!user.hasFlag(UserFlag.TFA_ENABLED)) { // Ensure TFA is already on
+            throw new BadRequestException(Error.TFA_NOT_ENABLED);
+        }
+        String encryptedPin = HashUtils.hash(Base64.getDecoder().decode(user.getTfa().getBackupCodesSalt()), pin);
+
+        // Before checking the pin, check the user's backup codes
+        for (String backupCode : user.getTfa().getBackupCodes()) {
+            if (!encryptedPin.equals(backupCode)) {
+                continue;
+            }
+            // The code is a valid backup code, remove it from the user's list
+            user.getTfa().getBackupCodes().remove(backupCode);
+            userRepository.save(user);
+            return;
+        }
+
+        // Check if the TFA pin is valid
+        if (!tfaService.getPin(user.getTfa().getSecret()).equals(pin)) {
+            throw new BadRequestException(Error.TFA_PIN_INVALID);
+        }
+    }
+
     /**
      * Get the authenticated user.
      *
@@ -262,6 +292,7 @@ public final class AuthService {
         USER_NOT_FOUND,
         PASSWORDS_DO_NOT_MATCH,
         BORDER_CROSSING,
+        TFA_NOT_ENABLED,
         TFA_PIN_INVALID,
         EMAIL_ALREADY_USED,
         USER_DISABLED
diff --git a/src/main/java/cc/pulseapp/api/service/UserService.java b/src/main/java/cc/pulseapp/api/service/UserService.java
index 85f853d..dd68d63 100644
--- a/src/main/java/cc/pulseapp/api/service/UserService.java
+++ b/src/main/java/cc/pulseapp/api/service/UserService.java
@@ -238,9 +238,8 @@ public final class UserService {
         if (!user.hasFlag(UserFlag.TFA_ENABLED)) { // Ensure TFA is already on
             throw new BadRequestException(Error.TFA_NOT_ENABLED);
         }
-        if (!tfaService.getPin(user.getTfa().getSecret()).equals(input.getPin())) { // Ensure the pin is valid
-            throw new BadRequestException(Error.TFA_PIN_INVALID);
-        }
+        authService.useTfaPin(user, input.getPin()); // Ensure the pin is valid
+
         // Disable TFA for the user
         user.setTfa(null);
         user.removeFlag(UserFlag.TFA_ENABLED);