Skip to content

Commit

Permalink
fix: totp import api (#915)
Browse files Browse the repository at this point in the history
* fix: totp import api

* fix: refactor
  • Loading branch information
sattvikc committed Jan 29, 2024
1 parent 9ef8d5f commit f884b13
Show file tree
Hide file tree
Showing 3 changed files with 120 additions and 43 deletions.
89 changes: 47 additions & 42 deletions src/main/java/io/supertokens/totp/Totp.java
Original file line number Diff line number Diff line change
Expand Up @@ -81,71 +81,76 @@ public static TOTPDevice registerDevice(Main main, String userId,
}
}

public static TOTPDevice createDevice(Main main, AppIdentifierWithStorage appIdentifierWithStorage, TOTPDevice device)
public static TOTPDevice createDevice(Main main, AppIdentifierWithStorage appIdentifierWithStorage, String userId,
String deviceName, int skew, int period, String secretKey, boolean verified,
long createdAt)
throws DeviceAlreadyExistsException, StorageQueryException, FeatureNotEnabledException,
TenantOrAppNotFoundException {

Mfa.checkForMFAFeature(appIdentifierWithStorage, main);

TOTPSQLStorage totpStorage = appIdentifierWithStorage.getTOTPStorage();
try {
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());
if (deviceName != null) {
TOTPSQLStorage totpStorage = appIdentifierWithStorage.getTOTPStorage();
try {
return totpStorage.startTransaction(con -> {
try {
TOTPDevice existingDevice = totpStorage.getDeviceByName_Transaction(con, appIdentifierWithStorage, userId, deviceName);
if (existingDevice == null) {
return totpStorage.createDevice_Transaction(con, appIdentifierWithStorage, new TOTPDevice(
userId, deviceName, secretKey, period, skew, verified, createdAt
));
} else if (!existingDevice.verified) {
totpStorage.deleteDevice_Transaction(con, appIdentifierWithStorage, userId, deviceName);
return totpStorage.createDevice_Transaction(con, appIdentifierWithStorage, new TOTPDevice(
userId, deviceName, secretKey, period, skew, verified, createdAt
));
} else {
throw new StorageTransactionLogicException(new DeviceAlreadyExistsException());
}
} catch (TenantOrAppNotFoundException | DeviceAlreadyExistsException e) {
throw new StorageTransactionLogicException(e);
}
} catch (TenantOrAppNotFoundException | DeviceAlreadyExistsException e) {
throw new StorageTransactionLogicException(e);
});
} catch (StorageTransactionLogicException e) {
if (e.actualException instanceof DeviceAlreadyExistsException) {
throw (DeviceAlreadyExistsException) e.actualException;
}
});
} catch (StorageTransactionLogicException e) {
if (e.actualException instanceof DeviceAlreadyExistsException) {
throw (DeviceAlreadyExistsException) e.actualException;
throw new StorageQueryException(e.actualException);
}
throw new StorageQueryException(e.actualException);
}
}

public static TOTPDevice registerDevice(AppIdentifierWithStorage appIdentifierWithStorage, Main main, String userId,
String deviceName, int skew, int period)
throws StorageQueryException, DeviceAlreadyExistsException, NoSuchAlgorithmException,
FeatureNotEnabledException, TenantOrAppNotFoundException, StorageTransactionLogicException {

String secret = generateSecret();
TOTPDevice device = new TOTPDevice(userId, deviceName, secret, period, skew, false, System.currentTimeMillis());
TOTPSQLStorage totpStorage = appIdentifierWithStorage.getTOTPStorage();

if (deviceName != null) {
return createDevice(main, 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;

while (true) {
try {
return createDevice(main, appIdentifierWithStorage, new TOTPDevice(
device.userId,
return createDevice(main, appIdentifierWithStorage,
userId,
"TOTP Device " + verifiedDevicesCount,
device.secretKey,
device.period,
device.skew,
device.verified,
device.createdAt
));
skew,
period,
secretKey,
verified,
createdAt
);
} catch (DeviceAlreadyExistsException e){
}
verifiedDevicesCount++;
}
}

public static TOTPDevice registerDevice(AppIdentifierWithStorage appIdentifierWithStorage, Main main, String userId,
String deviceName, int skew, int period)
throws StorageQueryException, DeviceAlreadyExistsException, NoSuchAlgorithmException,
FeatureNotEnabledException, TenantOrAppNotFoundException, StorageTransactionLogicException {

String secretKey = generateSecret();

return createDevice(main, appIdentifierWithStorage, userId, deviceName, skew, period, secretKey, false,
System.currentTimeMillis());
}

private static void checkAndStoreCode(TenantIdentifierWithStorage tenantIdentifierWithStorage, Main main,
String userId, TOTPDevice[] devices,
String code)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -89,9 +89,11 @@ protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws I
appIdentifierWithStorage = getAppIdentifierWithStorage(req);
}

Totp.createDevice(super.main, appIdentifierWithStorage, new TOTPDevice(userId, deviceName, secretKey, period, skew, true, System.currentTimeMillis()));
TOTPDevice createdDevice = Totp.createDevice(super.main, appIdentifierWithStorage,
userId, deviceName, skew, period, secretKey, true, System.currentTimeMillis());

result.addProperty("status", "OK");
result.addProperty("deviceName", createdDevice.deviceName);
super.sendJsonResponse(200, result, resp);
} catch (DeviceAlreadyExistsException e) {
result.addProperty("status", "DEVICE_ALREADY_EXISTS_ERROR");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,7 @@ public void testApi() throws Exception {
Utils.getCdiVersionStringLatestForTests(),
"totp");
assert res.get("status").getAsString().equals("OK");
assertEquals("d1", res.get("deviceName").getAsString());

// try again with same device name:
JsonObject res2 = HttpRequestForTesting.sendJsonPOSTRequest(
Expand Down Expand Up @@ -188,4 +189,73 @@ public void testApi() throws Exception {
process.kill();
assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STOPPED));
}

@Test
public void testApiWithoutDeviceName() throws Exception {
String[] args = {"../"};

TestingProcessManager.TestingProcess process = TestingProcessManager.start(args);
assertNotNull(process.checkOrWaitForEvent(ProcessState.PROCESS_STATE.STARTED));

if (StorageLayer.getStorage(process.getProcess()).getType() != STORAGE_TYPE.SQL) {
return;
}

FeatureFlag.getInstance(process.main)
.setLicenseKeyAndSyncFeatures(TotpLicenseTest.OPAQUE_KEY_WITH_MFA_FEATURE);
FeatureFlagTestContent.getInstance(process.main)
.setKeyValue(FeatureFlagTestContent.ENABLED_FEATURES, new EE_FEATURES[]{EE_FEATURES.MFA});

if (StorageLayer.getStorage(process.getProcess()).getType() != STORAGE_TYPE.SQL) {
return;
}

{
String secret = "ZNPARPDTO6BFVSOFM3BPJGORPYTNTDSF";

JsonObject body = new JsonObject();
body.addProperty("secretKey", "");
body.addProperty("userId", "user-id");
body.addProperty("secretKey", secret);
body.addProperty("skew", 0);
body.addProperty("period", 30);

JsonObject res = HttpRequestForTesting.sendJsonPOSTRequest(
process.getProcess(),
"",
"http://localhost:3567/recipe/totp/device/import",
body,
1000,
1000,
null,
Utils.getCdiVersionStringLatestForTests(),
"totp");
assert res.get("status").getAsString().equals("OK");
assertEquals("TOTP Device 0", res.get("deviceName").getAsString());
}

{ // Check for device already exists
String secret = "ZNPARPDTO6BFVSOFM3BPJGORPYTNTDSF";

JsonObject body = new JsonObject();
body.addProperty("secretKey", "");
body.addProperty("userId", "user-id");
body.addProperty("secretKey", secret);
body.addProperty("skew", 0);
body.addProperty("period", 30);
body.addProperty("deviceName", "TOTP Device 0");

JsonObject res = HttpRequestForTesting.sendJsonPOSTRequest(
process.getProcess(),
"",
"http://localhost:3567/recipe/totp/device/import",
body,
1000,
1000,
null,
Utils.getCdiVersionStringLatestForTests(),
"totp");
assert res.get("status").getAsString().equals("DEVICE_ALREADY_EXISTS_ERROR");
}
}
}

0 comments on commit f884b13

Please sign in to comment.