fix: user roles

supertokens · Mar 4, 2024 · d7cbcfa · d7cbcfa
1 parent 3d93ab5
commit d7cbcfa
Show file tree

Hide file tree

Showing 4 changed files with 85 additions and 36 deletions.
diff --git a/src/main/java/io/supertokens/inmemorydb/Start.java b/src/main/java/io/supertokens/inmemorydb/Start.java
@@ -1937,9 +1937,20 @@ public String[] getRolesThatHavePermission(AppIdentifier appIdentifier, String p
     }
 
     @Override
-    public boolean deleteRole(AppIdentifier appIdentifier, String role) throws StorageQueryException {
+    public boolean deleteRole_Transaction(TransactionConnection con, AppIdentifier appIdentifier, String role) throws StorageQueryException {
         try {
-            return UserRolesQueries.deleteRole(this, appIdentifier, role);
+            return UserRolesQueries.deleteRole_Transaction(this, (Connection) con.getConnection(), appIdentifier, role);
+        } catch (SQLException e) {
+            throw new StorageQueryException(e);
+        }
+    }
+
+    @Override
+    public boolean deleteAllUserRoleAssociationsForRole_Transaction(TransactionConnection con,
+                                                                    AppIdentifier appIdentifier, String role)
+            throws StorageQueryException {
+        try {
+            return UserRolesQueries.deleteAllUserRoleAssociationsForRole_Transaction(this, (Connection) con.getConnection(), appIdentifier, role);
         } catch (SQLException e) {
             throw new StorageQueryException(e);
         }

diff --git a/src/main/java/io/supertokens/inmemorydb/queries/UserRolesQueries.java b/src/main/java/io/supertokens/inmemorydb/queries/UserRolesQueries.java
@@ -111,29 +111,15 @@ public static void addPermissionToRoleOrDoNothingIfExists_Transaction(Start star
         });
     }
 
-    public static boolean deleteRole(Start start, AppIdentifier appIdentifier, String role) throws SQLException, StorageQueryException {
-
-        try {
-            return start.startTransaction(con -> {
-                // Row lock must be taken to delete the role, otherwise the table may be locked for delete
-                Connection sqlCon = (Connection) con.getConnection();
-                ((ConnectionWithLocks) sqlCon).lock(appIdentifier.getAppId() + "~" + role + Config.getConfig(start).getRolesTable());
-
-                String QUERY = "DELETE FROM " + getConfig(start).getRolesTable()
-                        + " WHERE app_id = ? AND role = ? ;";
-
-                try {
-                    return update(sqlCon, QUERY, pst -> {
-                        pst.setString(1, appIdentifier.getAppId());
-                        pst.setString(2, role);
-                    }) == 1;
-                } catch (SQLException e) {
-                    throw new StorageTransactionLogicException(e);
-                }
-            });
-        } catch (StorageTransactionLogicException e) {
-            throw new StorageQueryException(e.actualException);
-        }
+    public static boolean deleteRole_Transaction(Start start, Connection sqlCon, AppIdentifier appIdentifier,
+                                                 String role) throws SQLException, StorageQueryException {
+
+        String QUERY = "DELETE FROM " + getConfig(start).getRolesTable()
+                + " WHERE app_id = ? AND role = ? ;";
+        return update(sqlCon, QUERY, pst -> {
+            pst.setString(1, appIdentifier.getAppId());
+            pst.setString(2, role);
+        }) == 1;
     }
 
     public static boolean doesRoleExist(Start start, AppIdentifier appIdentifier, String role)
@@ -339,4 +325,14 @@ public static int deleteAllRolesForUser_Transaction(Connection con, Start start,
             pst.setString(2, userId);
         });
     }
+
+    public static boolean deleteAllUserRoleAssociationsForRole_Transaction(Start start, Connection sqlCon, AppIdentifier appIdentifier, String role)
+            throws SQLException, StorageQueryException {
+        String QUERY = "DELETE FROM " + getConfig(start).getRolesTable()
+                + " WHERE app_id = ? AND role = ? ;";
+        return update(sqlCon, QUERY, pst -> {
+            pst.setString(1, appIdentifier.getAppId());
+            pst.setString(2, role);
+        }) >= 1;
+    }
 }
diff --git a/src/main/java/io/supertokens/userroles/UserRoles.java b/src/main/java/io/supertokens/userroles/UserRoles.java
@@ -292,17 +292,36 @@ public static boolean deleteRole(Main main, AppIdentifier appIdentifier, String
         Storage[] storages = StorageLayer.getStoragesForApp(main, appIdentifier);
         boolean deletedRole = false;
         for (Storage storage : storages) {
-            if (storage.getUserPoolId().equals(appStorage.getUserPoolId())) {
-                continue; // we want to delete this in the end
-            }
             UserRolesSQLStorage userRolesStorage = StorageUtils.getUserRolesStorage(storage);
-            deletedRole = userRolesStorage.deleteRole(appIdentifier, role) || deletedRole;
+
+            try {
+                deletedRole = userRolesStorage.startTransaction(con -> {
+                    return userRolesStorage.deleteAllUserRoleAssociationsForRole_Transaction(con, appIdentifier, role);
+                }) || deletedRole;
+            } catch (StorageTransactionLogicException e) {
+                if (e.actualException instanceof StorageQueryException) {
+                    throw (StorageQueryException) e.actualException;
+                }
+                throw new IllegalStateException(e);
+            }
         }
 
         // Delete the role from the public tenant storage in the end so that the user
         // never sees a role for user that has been deleted while the deletion is in progress
         UserRolesSQLStorage userRolesStorage = StorageUtils.getUserRolesStorage(appStorage);
-        deletedRole = userRolesStorage.deleteRole(appIdentifier, role) || deletedRole;
+        try {
+            deletedRole = userRolesStorage.startTransaction(con -> {
+                boolean deleted = false;
+                deleted = userRolesStorage.deleteAllUserRoleAssociationsForRole_Transaction(con, appIdentifier, role) || deleted;
+                deleted = userRolesStorage.deleteRole_Transaction(con, appIdentifier, role) || deleted;
+                return deleted;
+            }) || deletedRole;
+        } catch (StorageTransactionLogicException e) {
+            if (e.actualException instanceof StorageQueryException) {
+                throw (StorageQueryException) e.actualException;
+            }
+            throw new IllegalStateException(e);
+        }
 
         return deletedRole;
     }

diff --git a/src/test/java/io/supertokens/test/userRoles/UserRolesStorageTest.java b/src/test/java/io/supertokens/test/userRoles/UserRolesStorageTest.java
@@ -18,6 +18,7 @@
 
 import io.supertokens.ProcessState;
 import io.supertokens.pluginInterface.STORAGE_TYPE;
+import io.supertokens.pluginInterface.StorageUtils;
 import io.supertokens.pluginInterface.exceptions.StorageQueryException;
 import io.supertokens.pluginInterface.exceptions.StorageTransactionLogicException;
 import io.supertokens.pluginInterface.multitenancy.AppIdentifier;
@@ -118,9 +119,15 @@ public void testDeletingARoleWhileItIsBeingRemovedFromAUser() throws Exception {
             }
             // delete the role
             try {
-                boolean wasRoleDeleted = storage.deleteRole(new AppIdentifier(null, null), role);
+                UserRolesSQLStorage userRolesStorage = StorageUtils.getUserRolesStorage(storage);
+                boolean wasRoleDeleted = userRolesStorage.startTransaction(con -> {
+                    boolean wasDeleted = storage.deleteAllUserRoleAssociationsForRole_Transaction(con, new AppIdentifier(null, null), role);
+                    wasDeleted = storage.deleteRole_Transaction(con, new AppIdentifier(null, null), role) || wasDeleted;
+                    return wasDeleted;
+                });
                 r2_success.set(wasRoleDeleted);
-            } catch (StorageQueryException e) {
+            } catch (StorageQueryException | StorageTransactionLogicException e) {
+                throw new RuntimeException(e);
                 // should not come here
             }
         };
@@ -225,10 +232,16 @@ public void testCreatingAndAddingAPermissionToARoleInTransactionAndDeletingRole(
             }
             // delete the newly created role
             try {
-                storage.deleteRole(new AppIdentifier(null, null), role);
+                UserRolesSQLStorage userRolesStorage = StorageUtils.getUserRolesStorage(storage);
+                boolean wasRoleDeleted = userRolesStorage.startTransaction(con -> {
+                    boolean wasDeleted = storage.deleteAllUserRoleAssociationsForRole_Transaction(con, new AppIdentifier(null, null), role);
+                    wasDeleted = storage.deleteRole_Transaction(con, new AppIdentifier(null, null), role) || wasDeleted;
+                    return wasDeleted;
+                });
                 r2_success.set(true);
-            } catch (StorageQueryException e) {
+            } catch (StorageQueryException | StorageTransactionLogicException e) {
                 // should not come here
+                throw new RuntimeException(e);
             }
         };
 
@@ -826,7 +839,12 @@ public void testDeletingRoleResponses() throws Exception {
             UserRoles.createNewRoleOrModifyItsPermissions(process.main, role, null);
 
             // delete role
-            boolean didRoleExist = storage.deleteRole(new AppIdentifier(null, null), role);
+            UserRolesSQLStorage userRolesStorage = StorageUtils.getUserRolesStorage(storage);
+            boolean didRoleExist = userRolesStorage.startTransaction(con -> {
+                boolean wasDeleted = storage.deleteAllUserRoleAssociationsForRole_Transaction(con, new AppIdentifier(null, null), role);
+                wasDeleted = storage.deleteRole_Transaction(con, new AppIdentifier(null, null), role) || wasDeleted;
+                return wasDeleted;
+            });
             assertTrue(didRoleExist);
 
             // check that role doesnt exist
@@ -835,7 +853,12 @@ public void testDeletingRoleResponses() throws Exception {
         {
             // delete a role which doesnt exist
 
-            boolean didRoleExist = storage.deleteRole(new AppIdentifier(null, null), role);
+            UserRolesSQLStorage userRolesStorage = StorageUtils.getUserRolesStorage(storage);
+            boolean didRoleExist = userRolesStorage.startTransaction(con -> {
+                boolean wasDeleted = storage.deleteAllUserRoleAssociationsForRole_Transaction(con, new AppIdentifier(null, null), role);
+                wasDeleted = storage.deleteRole_Transaction(con, new AppIdentifier(null, null), role) || wasDeleted;
+                return wasDeleted;
+            });
             assertFalse(didRoleExist);
         }