Skip to content

Commit

Permalink
fix: core fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
@sattvikc
sattvikc committed Sep 28, 2023
1 parent 7ef46d1 commit 29f9643
Show file tree
Hide file tree
Showing 5 changed files with 85 additions and 71 deletions.
51 changes: 31 additions & 20 deletions src/main/java/io/supertokens/inmemorydb/Start.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand All @@ -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(),
Expand Down
90 changes: 46 additions & 44 deletions src/main/java/io/supertokens/totp/Totp.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
}

Expand All @@ -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
Expand Down Expand Up @@ -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);
}

Expand All @@ -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;
}
Expand Down Expand Up @@ -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();
}
Expand All @@ -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,
Expand All @@ -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 {

Expand All @@ -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:
Expand All @@ -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;
Expand Down
4 changes: 2 additions & 2 deletions src/main/java/io/supertokens/webserver/api/totp/VerifyTotpAPI.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down
7 changes: 4 additions & 3 deletions src/test/java/io/supertokens/test/StorageLayerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
4 changes: 2 additions & 2 deletions src/test/java/io/supertokens/test/totp/TOTPStorageTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down

0 comments on commit 29f9643

Please sign in to comment.