fix: totp import api

supertokens · Jan 29, 2024 · 77cca5e · 77cca5e
1 parent 9ef8d5f
commit 77cca5e
Show file tree

Hide file tree

Showing 3 changed files with 107 additions and 35 deletions.
diff --git a/src/main/java/io/supertokens/totp/Totp.java b/src/main/java/io/supertokens/totp/Totp.java
@@ -87,46 +87,34 @@ public static TOTPDevice createDevice(Main main, AppIdentifierWithStorage appIde
 
         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 (device.deviceName != null) {
+            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());
+                        }
+                    } 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);
+        TOTPDevice[] devices = totpStorage.getDevices(appIdentifierWithStorage, device.userId);
         int verifiedDevicesCount = Arrays.stream(devices).filter(d -> d.verified).toArray().length;
 
         while (true) {
@@ -146,6 +134,17 @@ public static TOTPDevice registerDevice(AppIdentifierWithStorage appIdentifierWi
         }
     }
 
+    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());
+
+        return createDevice(main, appIdentifierWithStorage, device);
+    }
+
     private static void checkAndStoreCode(TenantIdentifierWithStorage tenantIdentifierWithStorage, Main main,
             String userId, TOTPDevice[] devices,
             String code)

diff --git a/src/main/java/io/supertokens/webserver/api/totp/ImportTotpDeviceAPI.java b/src/main/java/io/supertokens/webserver/api/totp/ImportTotpDeviceAPI.java
@@ -89,9 +89,12 @@ 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,
+                    new TOTPDevice(userId, deviceName,
+                            secretKey, period, skew, 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");

diff --git a/src/test/java/io/supertokens/test/totp/api/ImportTotpDeviceAPITest.java b/src/test/java/io/supertokens/test/totp/api/ImportTotpDeviceAPITest.java
@@ -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(
@@ -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");
+        }
+    }
 }