From 29f96430e67359ad67b89a1fe7c23b69b9ab5290 Mon Sep 17 00:00:00 2001 From: Sattvik Chakravarthy Date: Thu, 28 Sep 2023 15:29:52 +0530 Subject: [PATCH] fix: core fixes --- .../java/io/supertokens/inmemorydb/Start.java | 51 ++++++----- src/main/java/io/supertokens/totp/Totp.java | 90 ++++++++++--------- .../webserver/api/totp/VerifyTotpAPI.java | 4 +- .../io/supertokens/test/StorageLayerTest.java | 7 +- .../test/totp/TOTPStorageTest.java | 4 +- 5 files changed, 85 insertions(+), 71 deletions(-) diff --git a/src/main/java/io/supertokens/inmemorydb/Start.java b/src/main/java/io/supertokens/inmemorydb/Start.java index f2c8bfada..6becde092 100644 --- a/src/main/java/io/supertokens/inmemorydb/Start.java +++ b/src/main/java/io/supertokens/inmemorydb/Start.java @@ -71,7 +71,7 @@ import io.supertokens.pluginInterface.totp.TOTPUsedCode; import io.supertokens.pluginInterface.totp.exception.DeviceAlreadyExistsException; import io.supertokens.pluginInterface.totp.exception.UnknownDeviceException; -import io.supertokens.pluginInterface.totp.exception.UnknownUserIdTotpException; +import io.supertokens.pluginInterface.totp.exception.UnknownTotpUserIdException; import io.supertokens.pluginInterface.totp.exception.UsedCodeAlreadyExistsException; import io.supertokens.pluginInterface.totp.sqlStorage.TOTPSQLStorage; import io.supertokens.pluginInterface.useridmapping.UserIdMapping; @@ -2601,29 +2601,40 @@ public void revokeExpiredSessions() throws StorageQueryException { // TOTP recipe: @Override public void createDevice(AppIdentifier appIdentifier, TOTPDevice device) - throws StorageQueryException, DeviceAlreadyExistsException, TenantOrAppNotFoundException { + throws DeviceAlreadyExistsException, TenantOrAppNotFoundException, StorageQueryException { try { - TOTPQueries.createDevice(this, appIdentifier, device); - } catch (StorageTransactionLogicException e) { - if (e.actualException instanceof SQLiteException) { - String errMsg = e.actualException.getMessage(); - - if (isPrimaryKeyError(errMsg, Config.getConfig(this).getTotpUserDevicesTable(), - new String[]{"app_id", "user_id", "device_name"})) { - throw new DeviceAlreadyExistsException(); - } else if (isForeignKeyConstraintError( - errMsg, - Config.getConfig(this).getAppsTable(), - new String[]{"app_id"}, - new Object[]{appIdentifier.getAppId()})) { - throw new TenantOrAppNotFoundException(appIdentifier); + startTransaction(con -> { + try { + createDevice_Transaction(con, new AppIdentifier(null, null), device); + } catch (DeviceAlreadyExistsException | TenantOrAppNotFoundException e) { + throw new StorageTransactionLogicException(e); } + return null; + }); + } catch (StorageTransactionLogicException e) { + if (e.actualException instanceof DeviceAlreadyExistsException) { + throw (DeviceAlreadyExistsException) e.actualException; + } else if (e.actualException instanceof TenantOrAppNotFoundException) { + throw (TenantOrAppNotFoundException) e.actualException; + } else if (e.actualException instanceof StorageQueryException) { + throw (StorageQueryException) e.actualException; } - - throw new StorageQueryException(e.actualException); } } + @Override + public TOTPDevice createDevice_Transaction(TransactionConnection con, AppIdentifier appIdentifier, TOTPDevice device) + throws DeviceAlreadyExistsException, TenantOrAppNotFoundException, StorageQueryException { + // TODO + return null; + } + + @Override + public TOTPDevice getDeviceByName_Transaction(TransactionConnection con, AppIdentifier appIdentifier, String userId, String deviceName) throws StorageQueryException { + // TODO + return null; + } + @Override public void markDeviceAsVerified(AppIdentifier appIdentifier, String userId, String deviceName) throws StorageQueryException, UnknownDeviceException { @@ -2717,7 +2728,7 @@ public TOTPDevice[] getDevices_Transaction(TransactionConnection con, AppIdentif @Override public void insertUsedCode_Transaction(TransactionConnection con, TenantIdentifier tenantIdentifier, TOTPUsedCode usedCodeObj) - throws StorageQueryException, UnknownUserIdTotpException, UsedCodeAlreadyExistsException, + throws StorageQueryException, UnknownTotpUserIdException, UsedCodeAlreadyExistsException, TenantOrAppNotFoundException { Connection sqlCon = (Connection) con.getConnection(); try { @@ -2732,7 +2743,7 @@ public void insertUsedCode_Transaction(TransactionConnection con, TenantIdentifi Config.getConfig(this).getTotpUsersTable(), new String[]{"app_id", "user_id"}, new Object[]{tenantIdentifier.getAppId(), usedCodeObj.userId})) { - throw new UnknownUserIdTotpException(); + throw new UnknownTotpUserIdException(); } else if (isForeignKeyConstraintError( e.getMessage(), diff --git a/src/main/java/io/supertokens/totp/Totp.java b/src/main/java/io/supertokens/totp/Totp.java index bff58e656..63e7b08f0 100644 --- a/src/main/java/io/supertokens/totp/Totp.java +++ b/src/main/java/io/supertokens/totp/Totp.java @@ -16,7 +16,7 @@ import io.supertokens.pluginInterface.totp.TOTPUsedCode; import io.supertokens.pluginInterface.totp.exception.DeviceAlreadyExistsException; import io.supertokens.pluginInterface.totp.exception.UnknownDeviceException; -import io.supertokens.pluginInterface.totp.exception.UnknownUserIdTotpException; +import io.supertokens.pluginInterface.totp.exception.UnknownTotpUserIdException; import io.supertokens.pluginInterface.totp.exception.UsedCodeAlreadyExistsException; import io.supertokens.pluginInterface.totp.sqlStorage.TOTPSQLStorage; import io.supertokens.storageLayer.StorageLayer; @@ -94,34 +94,30 @@ public static TOTPDevice registerDevice(Main main, String userId, } } - private static TOTPDevice registerDeviceRecursive(AppIdentifierWithStorage appIdentifierWithStorage, TOTPDevice device, int deviceNameCounter) throws StorageQueryException, DeviceAlreadyExistsException, TenantOrAppNotFoundException, StorageTransactionLogicException { + private static TOTPDevice createDevice(AppIdentifierWithStorage appIdentifierWithStorage, TOTPDevice device) + throws DeviceAlreadyExistsException, StorageQueryException { TOTPSQLStorage totpStorage = appIdentifierWithStorage.getTOTPStorage(); - TOTPDevice newDevice = new TOTPDevice(device.userId, "TOTP Device " + deviceNameCounter, device.secretKey, device.period, device.skew, false); try { - totpStorage.createDevice(appIdentifierWithStorage, newDevice); - return newDevice; - } catch (DeviceAlreadyExistsException e) { - TOTPDevice[] devices = totpStorage.getDevices(appIdentifierWithStorage, device.userId); - // iterate through all devices to find device with same name - TOTPDevice existingDevice = Arrays.stream(devices).filter(d -> d.deviceName.equals(newDevice.deviceName)) - .findFirst().orElse(null); - - if (existingDevice != null) { - if (existingDevice.verified) { - // device with same name exists and is verified - return registerDeviceRecursive(appIdentifierWithStorage, device, deviceNameCounter + 1); - } else { - // device with same name exists but is not verified - // delete the device and retry - totpStorage.startTransaction(con -> { - totpStorage.deleteDevice_Transaction(con, appIdentifierWithStorage, newDevice.userId, newDevice.deviceName); - totpStorage.commitTransaction(con); - return null; - }); - return registerDeviceRecursive(appIdentifierWithStorage, device, deviceNameCounter); + return totpStorage.startTransaction(con -> { + try { + TOTPDevice existingDevice = totpStorage.getDeviceByName_Transaction(con, appIdentifierWithStorage, device.userId, device.deviceName); + if (existingDevice == null) { + return totpStorage.createDevice_Transaction(con, appIdentifierWithStorage, device); + } else if (!existingDevice.verified) { + totpStorage.deleteDevice_Transaction(con, appIdentifierWithStorage, device.userId, device.deviceName); + return totpStorage.createDevice_Transaction(con, appIdentifierWithStorage, device); + } else { + throw new StorageTransactionLogicException(new DeviceAlreadyExistsException()); + } + } catch (TenantOrAppNotFoundException | DeviceAlreadyExistsException e) { + throw new StorageTransactionLogicException(e); } + }); + } catch (StorageTransactionLogicException e) { + if (e.actualException instanceof DeviceAlreadyExistsException) { + throw (DeviceAlreadyExistsException) e.actualException; } - throw e; + throw new StorageQueryException(e.actualException); } } @@ -141,21 +137,33 @@ public static TOTPDevice registerDevice(AppIdentifierWithStorage appIdentifierWi TOTPSQLStorage totpStorage = appIdentifierWithStorage.getTOTPStorage(); if (deviceName != null) { - totpStorage.createDevice(appIdentifierWithStorage, device); - return device; + return createDevice(appIdentifierWithStorage, device); } // Find number of existing devices to set device name TOTPDevice[] devices = totpStorage.getDevices(appIdentifierWithStorage, userId); int verifiedDevicesCount = Arrays.stream(devices).filter(d -> d.verified).toArray().length; - return registerDeviceRecursive(appIdentifierWithStorage, device, verifiedDevicesCount + 1); + while (true) { + try { + return createDevice(appIdentifierWithStorage, new TOTPDevice( + device.userId, + "TOTP Device " + verifiedDevicesCount, + device.secretKey, + device.period, + device.skew, + device.verified + )); + } catch (DeviceAlreadyExistsException e){ + } + verifiedDevicesCount++; + } } private static void checkAndStoreCode(TenantIdentifierWithStorage tenantIdentifierWithStorage, Main main, String userId, TOTPDevice[] devices, String code) - throws InvalidTotpException, UnknownUserIdTotpException, + throws InvalidTotpException, UnknownTotpUserIdException, LimitReachedException, StorageQueryException, StorageTransactionLogicException, TenantOrAppNotFoundException { // Note that the TOTP cron runs every 1 hour, so all the expired tokens can stay @@ -277,7 +285,7 @@ private static void checkAndStoreCode(TenantIdentifierWithStorage tenantIdentifi try { totpSQLStorage.insertUsedCode_Transaction(con, tenantIdentifierWithStorage, newCode); totpSQLStorage.commitTransaction(con); - } catch (UsedCodeAlreadyExistsException | UnknownUserIdTotpException e) { + } catch (UsedCodeAlreadyExistsException | UnknownTotpUserIdException e) { throw new StorageTransactionLogicException(e); } @@ -300,16 +308,10 @@ private static void checkAndStoreCode(TenantIdentifierWithStorage tenantIdentifi throw (LimitReachedException) e.actualException; } else if (e.actualException instanceof InvalidTotpException) { throw (InvalidTotpException) e.actualException; - } else if (e.actualException instanceof UnknownUserIdTotpException) { - throw (UnknownUserIdTotpException) e.actualException; + } else if (e.actualException instanceof UnknownTotpUserIdException) { + throw (UnknownTotpUserIdException) e.actualException; } else if (e.actualException instanceof UsedCodeAlreadyExistsException) { - // retry the transaction after a small delay: - int delayInMs = (int) (Math.random() * 10 + 1); - try { - Thread.sleep(delayInMs); - } catch (InterruptedException ignored) { - // ignore the error and retry - } + throw new InvalidTotpException(); } else { throw e; } @@ -375,7 +377,7 @@ public static boolean verifyDevice(TenantIdentifierWithStorage tenantIdentifierW // This behaviour is okay so we can ignore it. try { checkAndStoreCode(tenantIdentifierWithStorage, main, userId, new TOTPDevice[] { matchingDevice }, code); - } catch (UnknownUserIdTotpException e) { + } catch (UnknownTotpUserIdException e) { // User must have deleted the device in parallel. throw new UnknownDeviceException(); } @@ -386,7 +388,7 @@ public static boolean verifyDevice(TenantIdentifierWithStorage tenantIdentifierW @TestOnly public static void verifyCode(Main main, String userId, String code) - throws InvalidTotpException, UnknownUserIdTotpException, LimitReachedException, + throws InvalidTotpException, UnknownTotpUserIdException, LimitReachedException, StorageQueryException, StorageTransactionLogicException, FeatureNotEnabledException { try { verifyCode(new TenantIdentifierWithStorage(null, null, null, StorageLayer.getStorage(main)), main, @@ -397,7 +399,7 @@ public static void verifyCode(Main main, String userId, String code) } public static void verifyCode(TenantIdentifierWithStorage tenantIdentifierWithStorage, Main main, String userId, String code) - throws InvalidTotpException, UnknownUserIdTotpException, LimitReachedException, + throws InvalidTotpException, UnknownTotpUserIdException, LimitReachedException, StorageQueryException, StorageTransactionLogicException, FeatureNotEnabledException, TenantOrAppNotFoundException { @@ -413,7 +415,7 @@ public static void verifyCode(TenantIdentifierWithStorage tenantIdentifierWithSt TOTPDevice[] devices = totpStorage.getDevices(tenantIdentifierWithStorage.toAppIdentifier(), userId); if (devices.length == 0) { // No devices found. So we can't verify the code anyway. - throw new UnknownUserIdTotpException(); + throw new UnknownTotpUserIdException(); } // Filter out unverified devices: @@ -425,7 +427,7 @@ public static void verifyCode(TenantIdentifierWithStorage tenantIdentifierWithSt // can ignore it. try { checkAndStoreCode(tenantIdentifierWithStorage, main, userId, devices, code); - } catch (UnknownUserIdTotpException e) { + } catch (UnknownTotpUserIdException e) { // User must have deleted the device in parallel // since they cannot un-verify a device (no API exists) throw e; diff --git a/src/main/java/io/supertokens/webserver/api/totp/VerifyTotpAPI.java b/src/main/java/io/supertokens/webserver/api/totp/VerifyTotpAPI.java index be575ab3a..b413332df 100644 --- a/src/main/java/io/supertokens/webserver/api/totp/VerifyTotpAPI.java +++ b/src/main/java/io/supertokens/webserver/api/totp/VerifyTotpAPI.java @@ -13,7 +13,7 @@ import io.supertokens.pluginInterface.exceptions.StorageTransactionLogicException; import io.supertokens.pluginInterface.multitenancy.TenantIdentifierWithStorage; import io.supertokens.pluginInterface.multitenancy.exceptions.TenantOrAppNotFoundException; -import io.supertokens.pluginInterface.totp.exception.UnknownUserIdTotpException; +import io.supertokens.pluginInterface.totp.exception.UnknownTotpUserIdException; import io.supertokens.totp.Totp; import io.supertokens.totp.exceptions.InvalidTotpException; import io.supertokens.totp.exceptions.LimitReachedException; @@ -78,7 +78,7 @@ protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws I } catch (InvalidTotpException e) { result.addProperty("status", "INVALID_TOTP_ERROR"); super.sendJsonResponse(200, result, resp); - } catch (UnknownUserIdTotpException e) { + } catch (UnknownTotpUserIdException e) { result.addProperty("status", "UNKNOWN_USER_ID_ERROR"); super.sendJsonResponse(200, result, resp); } catch (LimitReachedException e) { diff --git a/src/test/java/io/supertokens/test/StorageLayerTest.java b/src/test/java/io/supertokens/test/StorageLayerTest.java index 9ff015a38..df64313bd 100644 --- a/src/test/java/io/supertokens/test/StorageLayerTest.java +++ b/src/test/java/io/supertokens/test/StorageLayerTest.java @@ -11,8 +11,9 @@ import io.supertokens.pluginInterface.multitenancy.exceptions.TenantOrAppNotFoundException; import io.supertokens.pluginInterface.totp.TOTPDevice; import io.supertokens.pluginInterface.totp.TOTPUsedCode; +import io.supertokens.pluginInterface.totp.exception.DeviceAlreadyExistsException; import io.supertokens.pluginInterface.totp.exception.UnknownDeviceException; -import io.supertokens.pluginInterface.totp.exception.UnknownUserIdTotpException; +import io.supertokens.pluginInterface.totp.exception.UnknownTotpUserIdException; import io.supertokens.pluginInterface.totp.exception.UsedCodeAlreadyExistsException; import io.supertokens.pluginInterface.totp.sqlStorage.TOTPSQLStorage; import io.supertokens.storageLayer.StorageLayer; @@ -48,7 +49,7 @@ public static void insertUsedCodeUtil(TOTPSQLStorage storage, TOTPUsedCode usedC storage.insertUsedCode_Transaction(con, new TenantIdentifier(null, null, null), usedCode); storage.commitTransaction(con); return null; - } catch (UnknownUserIdTotpException | UsedCodeAlreadyExistsException e) { + } catch (UnknownTotpUserIdException | UsedCodeAlreadyExistsException e) { throw new StorageTransactionLogicException(e); } catch (TenantOrAppNotFoundException e) { throw new IllegalStateException(e); @@ -85,7 +86,7 @@ public void totpCodeLengthTest() throws Exception { TOTPDevice d1 = new TOTPDevice("user", "d1", "secret", 30, 1, false); storage.createDevice(new AppIdentifier(null, null), d1); - + // Try code with length > 8 try { TOTPUsedCode code = new TOTPUsedCode("user", "123456789", true, nextDay, now); diff --git a/src/test/java/io/supertokens/test/totp/TOTPStorageTest.java b/src/test/java/io/supertokens/test/totp/TOTPStorageTest.java index 5ef2a3c42..e03169d2a 100644 --- a/src/test/java/io/supertokens/test/totp/TOTPStorageTest.java +++ b/src/test/java/io/supertokens/test/totp/TOTPStorageTest.java @@ -15,7 +15,7 @@ import io.supertokens.pluginInterface.totp.TOTPUsedCode; import io.supertokens.pluginInterface.totp.exception.DeviceAlreadyExistsException; import io.supertokens.pluginInterface.totp.exception.UnknownDeviceException; -import io.supertokens.pluginInterface.totp.exception.UnknownUserIdTotpException; +import io.supertokens.pluginInterface.totp.exception.UnknownTotpUserIdException; import io.supertokens.pluginInterface.totp.exception.UsedCodeAlreadyExistsException; import io.supertokens.pluginInterface.totp.sqlStorage.TOTPSQLStorage; import io.supertokens.storageLayer.StorageLayer; @@ -95,7 +95,7 @@ public static void insertUsedCodesUtil(TOTPSQLStorage storage, TOTPUsedCode[] us for (TOTPUsedCode usedCode : usedCodes) { storage.insertUsedCode_Transaction(con, new TenantIdentifier(null, null, null), usedCode); } - } catch (UnknownUserIdTotpException | UsedCodeAlreadyExistsException e) { + } catch (UnknownTotpUserIdException | UsedCodeAlreadyExistsException e) { throw new StorageTransactionLogicException(e); } catch (TenantOrAppNotFoundException e) { throw new IllegalStateException(e);