Disconnect clients when updating RLS rules

Author: mamcx

crates/schema/src/auto_migrate.rs

@@ -1019,13 +1019,26 @@ fn auto_migrate_constraints(plan: &mut AutoMigratePlan, new_tables: &HashSet<&Id
 // Because we can refer to many tables and fields on the row level-security query, we need to remove all of them,
 // then add the new ones, instead of trying to track the graph of dependencies.
 fn auto_migrate_row_level_security(plan: &mut AutoMigratePlan) -> Result<()> {
+    // Track if any RLS rules were changed.
+    let mut old_rls = HashSet::new();
+    let mut new_rls = HashSet::new();
+
     for rls in plan.old.row_level_security() {
+        old_rls.insert(rls.key());
         plan.steps.push(AutoMigrateStep::RemoveRowLevelSecurity(rls.key()));
     }
     for rls in plan.new.row_level_security() {
+        new_rls.insert(rls.key());
         plan.steps.push(AutoMigrateStep::AddRowLevelSecurity(rls.key()));
     }
 
+    // After changing an RLS rule and updating the module, the updated rule does not take effect for clients until they reconnect.
+    // This is because subscriptions are cached and we don't flush the cache when RLS rules are updated.
+    // We can force flush the cache by force disconnecting all clients if an RLS rule has been added, removed, or updated.
+    if old_rls != new_rls {
+        plan.steps.push(AutoMigrateStep::DisconnectAllUsers);
+    }
+
     Ok(())
 }
 
@@ -2245,4 +2258,52 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn change_rls_disconnect_clients() {
+        let old_def = create_module_def(|_builder| {});
+
+        let new_def = create_module_def(|_builder| {});
+
+        let plan = ponder_auto_migrate(&old_def, &new_def).expect("auto migration should succeed");
+        assert!(!plan.disconnects_all_users(), "{plan:#?}");
+
+        let old_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT true;");
+        });
+        let new_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT false;");
+        });
+
+        let plan = ponder_auto_migrate(&old_def, &new_def).expect("auto migration should succeed");
+        assert!(plan.disconnects_all_users(), "{plan:#?}");
+
+        let old_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT true;");
+        });
+
+        let new_def = create_module_def(|_builder| {
+            // Remove RLS
+        });
+        let plan = ponder_auto_migrate(&old_def, &new_def).expect("auto migration should succeed");
+        assert!(plan.disconnects_all_users(), "{plan:#?}");
+
+        let old_def = create_module_def(|_builder| {});
+
+        let new_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT false;");
+        });
+        let plan = ponder_auto_migrate(&old_def, &new_def).expect("auto migration should succeed");
+        assert!(plan.disconnects_all_users(), "{plan:#?}");
+
+        let old_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT true;");
+        });
+
+        let new_def = create_module_def(|builder| {
+            builder.add_row_level_security("SELECT true;");
+        });
+        let plan = ponder_auto_migrate(&old_def, &new_def).expect("auto migration should succeed");
+        assert!(!plan.disconnects_all_users(), "{plan:#?}");
+    }
 }

smoketests/tests/rls.py

@@ -1,3 +1,5 @@
+import logging
+
 from .. import Smoketest, random_string
 
 class Rls(Smoketest):
@@ -78,3 +80,71 @@ def test_publish_fails_for_rls_on_private_table(self):
 
         with self.assertRaises(Exception):
             self.publish_module(name)
+
+class DisconnectRls(Smoketest):
+    AUTOPUBLISH = False
+    
+    MODULE_CODE = """
+use spacetimedb::{Identity, ReducerContext, Table};
+
+#[spacetimedb::table(name = users, public)]
+pub struct Users {
+    name: String,
+    identity: Identity,
+}
+
+#[spacetimedb::reducer]
+pub fn add_user(ctx: &ReducerContext, name: String) {
+    ctx.db.users().insert(Users { name, identity: ctx.sender });
+}
+
+// RLS
+"""
+    
+    ADD_RLS = """ 
+    #[spacetimedb::client_visibility_filter]
+const USER_FILTER: spacetimedb::Filter = spacetimedb::Filter::Sql(
+    "SELECT * FROM users WHERE identity = :sender"
+);
+"""
+    
+    def test_rls_disconnect_if_change(self):
+        """This tests that changing the RLS rules disconnects existing clients"""
+        
+        name = random_string()
+        
+        self.write_module_code(self.MODULE_CODE)
+        
+        self.publish_module(name)
+        logging.info("Initial publish complete")
+        
+        # Now add the RLS rules
+        self.write_module_code(self.MODULE_CODE + self.ADD_RLS)
+        self.publish_module(name, clear=False, break_clients=True)
+        
+        logging.info("Re-publish with RLS complete")
+        
+        logs = self.logs(100)
+        
+        # Validate disconnect + schema migration logs
+        self.assertIn("Disconnecting all users", logs)
+    
+    def test_rls_disconnect_no(self):
+        """This tests that not changing the RLS rules does not disconnect existing clients"""
+        
+        name = random_string()
+        
+        self.write_module_code(self.MODULE_CODE + self.ADD_RLS)
+        
+        self.publish_module(name)
+        logging.info("Initial publish complete")
+        
+        # Now re-publish the same module code
+        self.publish_module(name, clear=False, break_clients=False)
+        
+        logging.info("Re-publish without RLS change complete")
+        
+        logs = self.logs(100)
+        
+        # Validate no disconnect logs
+        self.assertNotIn("Disconnecting all users", logs)