fix: totp APIs and flaky tests fix (#879)

* fix: totp APIs and flaky tests fix

* fix: PR comment

src/main/java/io/supertokens/totp/Totp.java

@@ -187,120 +187,118 @@ private static void checkAndStoreCode(TenantIdentifierWithStorage tenantIdentifi
 
         TOTPSQLStorage totpSQLStorage = tenantIdentifierWithStorage.getTOTPStorage();
 
-        while (true) {
-            try {
-                totpSQLStorage.startTransaction(con -> {
-                    try {
-                        TOTPUsedCode[] usedCodes = totpSQLStorage.getAllUsedCodesDescOrder_Transaction(con,
-                                tenantIdentifierWithStorage,
-                                userId);
-
-                        // N represents # of invalid attempts that will trigger rate limiting:
-                        int N = Config.getConfig(tenantIdentifierWithStorage, main).getTotpMaxAttempts(); // (Default 5)
-                        // Count # of contiguous invalids in latest N attempts (stop at first valid):
-                        long invalidOutOfN = Arrays.stream(usedCodes).limit(N).takeWhile(usedCode -> !usedCode.isValid)
-                                .count();
-                        int rateLimitResetTimeInMs = Config.getConfig(tenantIdentifierWithStorage, main)
-                                .getTotpRateLimitCooldownTimeSec() *
-                                1000; // (Default 15 mins)
-
-                        // Check if the user has been rate limited:
-                        if (invalidOutOfN == N) {
-                            // All of the latest N attempts were invalid:
-                            long latestInvalidCodeCreatedTime = usedCodes[0].createdTime;
-                            long now = System.currentTimeMillis();
-
-                            if (now - latestInvalidCodeCreatedTime < rateLimitResetTimeInMs) {
-                                // Less than rateLimitResetTimeInMs (default = 15 mins) time has elasped since
-                                // the last invalid code:
-                                long timeLeftMs = (rateLimitResetTimeInMs - (now - latestInvalidCodeCreatedTime));
-                                throw new StorageTransactionLogicException(new LimitReachedException(timeLeftMs));
-
-                                // If we insert the used code here, then it will further delay the user from
-                                // being able to login. So not inserting it here.
-                            }
-                        }
+        try {
+            totpSQLStorage.startTransaction(con -> {
+                try {
+                    TOTPUsedCode[] usedCodes = totpSQLStorage.getAllUsedCodesDescOrder_Transaction(con,
+                            tenantIdentifierWithStorage,
+                            userId);
+
+                    // N represents # of invalid attempts that will trigger rate limiting:
+                    int N = Config.getConfig(tenantIdentifierWithStorage, main).getTotpMaxAttempts(); // (Default 5)
+                    // Count # of contiguous invalids in latest N attempts (stop at first valid):
+                    long invalidOutOfN = Arrays.stream(usedCodes).limit(N).takeWhile(usedCode -> !usedCode.isValid)
+                            .count();
+                    int rateLimitResetTimeInMs = Config.getConfig(tenantIdentifierWithStorage, main)
+                            .getTotpRateLimitCooldownTimeSec() *
+                            1000; // (Default 15 mins)
+
+                    // Check if the user has been rate limited:
+                    if (invalidOutOfN == N) {
+                        // All of the latest N attempts were invalid:
+                        long latestInvalidCodeCreatedTime = usedCodes[0].createdTime;
+                        long now = System.currentTimeMillis();
 
-                        // Check if the code is valid for any device:
-                        boolean isValid = false;
-                        TOTPDevice matchingDevice = null;
-                        for (TOTPDevice device : devices) {
-                            // Check if the code is valid for this device:
-                            if (checkCode(device, code)) {
-                                isValid = true;
-                                matchingDevice = device;
-                                break;
-                            }
-                        }
+                        if (now - latestInvalidCodeCreatedTime < rateLimitResetTimeInMs) {
+                            // Less than rateLimitResetTimeInMs (default = 15 mins) time has elasped since
+                            // the last invalid code:
+                            long timeLeftMs = (rateLimitResetTimeInMs - (now - latestInvalidCodeCreatedTime));
+                            throw new StorageTransactionLogicException(new LimitReachedException(timeLeftMs, (int)invalidOutOfN, N));
 
-                        // Check if the code has been previously used by the user and it was valid (and
-                        // is still valid). If so, this could be a replay attack. So reject it.
-                        if (isValid) {
-                            for (TOTPUsedCode usedCode : usedCodes) {
-                                // One edge case is that if the user has 2 devices, and they are used back to
-                                // back (within 90 seconds) such that the code of the first device was
-                                // regenerated by the second device, then it won't allow the second device's
-                                // code to be used until it is expired.
-                                // But this would be rare so we can ignore it for now.
-                                if (usedCode.code.equals(code) && usedCode.isValid
-                                        && usedCode.expiryTime > System.currentTimeMillis()) {
-                                    isValid = false;
-                                    // We found a matching device but the code
-                                    // will be considered invalid here.
-                                }
-                            }
+                            // If we insert the used code here, then it will further delay the user from
+                            // being able to login. So not inserting it here.
                         }
+                    }
 
-                        // Insert the code into the list of used codes:
-
-                        // If device is found, calculate used code expiry time for that device (based on
-                        // its period and skew). Otherwise, use the max used code expiry time of all the
-                        // devices.
-                        int maxUsedCodeExpiry = Arrays.stream(devices)
-                                .mapToInt(device -> device.period * (2 * device.skew + 1))
-                                .max()
-                                .orElse(0);
-                        int expireInSec = (matchingDevice != null)
-                                ? matchingDevice.period * (2 * matchingDevice.skew + 1)
-                                : maxUsedCodeExpiry;
-
-                        long now = System.currentTimeMillis();
-                        TOTPUsedCode newCode = new TOTPUsedCode(userId,
-                                code,
-                                isValid, now + 1000 * expireInSec, now);
-                        try {
-                            totpSQLStorage.insertUsedCode_Transaction(con, tenantIdentifierWithStorage, newCode);
-                            totpSQLStorage.commitTransaction(con);
-                        } catch (UsedCodeAlreadyExistsException | UnknownTotpUserIdException e) {
-                            throw new StorageTransactionLogicException(e);
+                    // Check if the code is valid for any device:
+                    boolean isValid = false;
+                    TOTPDevice matchingDevice = null;
+                    for (TOTPDevice device : devices) {
+                        // Check if the code is valid for this device:
+                        if (checkCode(device, code)) {
+                            isValid = true;
+                            matchingDevice = device;
+                            break;
                         }
+                    }
 
-                        if (!isValid) {
-                            // transaction has been committed, so we can directly throw the exception:
-                            throw new StorageTransactionLogicException(new InvalidTotpException());
+                    // Check if the code has been previously used by the user and it was valid (and
+                    // is still valid). If so, this could be a replay attack. So reject it.
+                    if (isValid) {
+                        for (TOTPUsedCode usedCode : usedCodes) {
+                            // One edge case is that if the user has 2 devices, and they are used back to
+                            // back (within 90 seconds) such that the code of the first device was
+                            // regenerated by the second device, then it won't allow the second device's
+                            // code to be used until it is expired.
+                            // But this would be rare so we can ignore it for now.
+                            if (usedCode.code.equals(code) && usedCode.isValid
+                                    && usedCode.expiryTime > System.currentTimeMillis()) {
+                                isValid = false;
+                                // We found a matching device but the code
+                                // will be considered invalid here.
+                            }
                         }
+                    }
 
-                        return null;
-                    } catch (TenantOrAppNotFoundException e) {
+                    // Insert the code into the list of used codes:
+
+                    // If device is found, calculate used code expiry time for that device (based on
+                    // its period and skew). Otherwise, use the max used code expiry time of all the
+                    // devices.
+                    int maxUsedCodeExpiry = Arrays.stream(devices)
+                            .mapToInt(device -> device.period * (2 * device.skew + 1))
+                            .max()
+                            .orElse(0);
+                    int expireInSec = (matchingDevice != null)
+                            ? matchingDevice.period * (2 * matchingDevice.skew + 1)
+                            : maxUsedCodeExpiry;
+
+                    long now = System.currentTimeMillis();
+                    TOTPUsedCode newCode = new TOTPUsedCode(userId,
+                            code,
+                            isValid, now + 1000L * expireInSec, now);
+                    try {
+                        totpSQLStorage.insertUsedCode_Transaction(con, tenantIdentifierWithStorage, newCode);
+                        totpSQLStorage.commitTransaction(con);
+                    } catch (UnknownTotpUserIdException e) {
                         throw new StorageTransactionLogicException(e);
+                    } catch (UsedCodeAlreadyExistsException e) {
+                        throw new StorageTransactionLogicException(new InvalidTotpException((int) invalidOutOfN, N));
                     }
-                });
-                return; // exit the while loop
-            } catch (StorageTransactionLogicException e) {
-                // throwing errors will also help exit the while loop:
-                if (e.actualException instanceof TenantOrAppNotFoundException) {
-                    throw (TenantOrAppNotFoundException) e.actualException;
-                } else if (e.actualException instanceof LimitReachedException) {
-                    throw (LimitReachedException) e.actualException;
-                } else if (e.actualException instanceof InvalidTotpException) {
-                    throw (InvalidTotpException) e.actualException;
-                } else if (e.actualException instanceof UnknownTotpUserIdException) {
-                    throw (UnknownTotpUserIdException) e.actualException;
-                } else if (e.actualException instanceof UsedCodeAlreadyExistsException) {
-                    throw new InvalidTotpException();
-                } else {
-                    throw e;
+
+                    if (!isValid) {
+                        // transaction has been committed, so we can directly throw the exception:
+                        throw new StorageTransactionLogicException(new InvalidTotpException((int)invalidOutOfN+1, N));
+                    }
+
+                    return null;
+                } catch (TenantOrAppNotFoundException e) {
+                    throw new StorageTransactionLogicException(e);
                 }
+            });
+            return; // exit the while loop
+        } catch (StorageTransactionLogicException e) {
+            // throwing errors will also help exit the while loop:
+            if (e.actualException instanceof TenantOrAppNotFoundException) {
+                throw (TenantOrAppNotFoundException) e.actualException;
+            } else if (e.actualException instanceof LimitReachedException) {
+                throw (LimitReachedException) e.actualException;
+            } else if (e.actualException instanceof InvalidTotpException) {
+                throw (InvalidTotpException) e.actualException;
+            } else if (e.actualException instanceof UnknownTotpUserIdException) {
+                throw (UnknownTotpUserIdException) e.actualException;
+            } else {
+                throw e;
             }
         }
     }

src/main/java/io/supertokens/totp/exceptions/InvalidTotpException.java

@@ -1,5 +1,12 @@
 package io.supertokens.totp.exceptions;
 
 public class InvalidTotpException extends Exception {
+    public int currentAttempts;
+    public int maxAttempts;
 
+    public InvalidTotpException(int currentAttempts, int maxAttempts) {
+        super("Invalid totp");
+        this.currentAttempts = currentAttempts;
+        this.maxAttempts = maxAttempts;
+    }
 }

src/main/java/io/supertokens/totp/exceptions/LimitReachedException.java

@@ -3,9 +3,13 @@
 public class LimitReachedException extends Exception {
 
     public long retryAfterMs;
+    public int currentAttempts;
+    public int maxAttempts;
 
-    public LimitReachedException(long retryAfterMs) {
+    public LimitReachedException(long retryAfterMs, int currentAttempts, int maxAttempts) {
         super("Retry in " + retryAfterMs + " ms");
         this.retryAfterMs = retryAfterMs;
+        this.currentAttempts = currentAttempts;
+        this.maxAttempts = maxAttempts;
     }
 }

src/main/java/io/supertokens/webserver/api/totp/VerifyTotpAPI.java

@@ -18,6 +18,7 @@
 import io.supertokens.totp.exceptions.InvalidTotpException;
 import io.supertokens.totp.exceptions.LimitReachedException;
 import io.supertokens.useridmapping.UserIdType;
+import io.supertokens.utils.SemVer;
 import io.supertokens.webserver.InputParser;
 import io.supertokens.webserver.WebserverAPI;
 import jakarta.servlet.ServletException;
@@ -78,12 +79,20 @@ protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws I
         } catch (InvalidTotpException e) {
             result.addProperty("status", "INVALID_TOTP_ERROR");
             super.sendJsonResponse(200, result, resp);
+            if (getVersionFromRequest(req).greaterThanOrEqualTo(SemVer.v4_1)) {
+                result.addProperty("currentNumberOfFailedAttempts", e.currentAttempts);
+                result.addProperty("maxNumberOfFailedAttempts", e.maxAttempts);
+            }
         } catch (UnknownTotpUserIdException e) {
             result.addProperty("status", "UNKNOWN_USER_ID_ERROR");
             super.sendJsonResponse(200, result, resp);
         } catch (LimitReachedException e) {
             result.addProperty("status", "LIMIT_REACHED_ERROR");
             result.addProperty("retryAfterMs", e.retryAfterMs);
+            if (getVersionFromRequest(req).greaterThanOrEqualTo(SemVer.v4_1)) {
+                result.addProperty("currentNumberOfFailedAttempts", e.currentAttempts);
+                result.addProperty("maxNumberOfFailedAttempts", e.maxAttempts);
+            }
             super.sendJsonResponse(200, result, resp);
         } catch (StorageQueryException | StorageTransactionLogicException | FeatureNotEnabledException |
                  TenantOrAppNotFoundException e) {

src/main/java/io/supertokens/webserver/api/totp/VerifyTotpDeviceAPI.java

@@ -17,6 +17,7 @@
 import io.supertokens.totp.exceptions.InvalidTotpException;
 import io.supertokens.totp.exceptions.LimitReachedException;
 import io.supertokens.useridmapping.UserIdType;
+import io.supertokens.utils.SemVer;
 import io.supertokens.webserver.InputParser;
 import io.supertokens.webserver.WebserverAPI;
 import jakarta.servlet.ServletException;
@@ -83,10 +84,19 @@ protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws I
             super.sendJsonResponse(200, result, resp);
         } catch (InvalidTotpException e) {
             result.addProperty("status", "INVALID_TOTP_ERROR");
+
+            if (getVersionFromRequest(req).greaterThanOrEqualTo(SemVer.v4_1)) {
+                result.addProperty("currentNumberOfFailedAttempts", e.currentAttempts);
+                result.addProperty("maxNumberOfFailedAttempts", e.maxAttempts);
+            }
             super.sendJsonResponse(200, result, resp);
         } catch (LimitReachedException e) {
             result.addProperty("status", "LIMIT_REACHED_ERROR");
             result.addProperty("retryAfterMs", e.retryAfterMs);
+            if (getVersionFromRequest(req).greaterThanOrEqualTo(SemVer.v4_1)) {
+                result.addProperty("currentNumberOfFailedAttempts", e.currentAttempts);
+                result.addProperty("maxNumberOfFailedAttempts", e.maxAttempts);
+            }
             super.sendJsonResponse(200, result, resp);
         } catch (StorageQueryException | StorageTransactionLogicException | TenantOrAppNotFoundException e) {
             throw new ServletException(e);