Add api to migrate scope and usage keys

Author: ales-otf

pallets/rate-limiting/src/lib.rs

@@ -370,6 +370,76 @@ pub mod pallet {
             true
         }
 
+        /// Migrates a stored rate limit configuration from one scope to another.
+        ///
+        /// Returns `true` when an entry was moved. Passing identical `from`/`to` scopes simply
+        /// checks that a configuration exists.
+        pub fn migrate_limit_scope(
+            identifier: &TransactionIdentifier,
+            from: Option<<T as Config<I>>::LimitScope>,
+            to: Option<<T as Config<I>>::LimitScope>,
+        ) -> bool {
+            if from == to {
+                return Limits::<T, I>::contains_key(identifier);
+            }
+
+            let mut migrated = false;
+            Limits::<T, I>::mutate(identifier, |maybe_config| {
+                if let Some(config) = maybe_config {
+                    match (from.as_ref(), to.as_ref()) {
+                        (None, Some(target)) => {
+                            if let RateLimit::Global(kind) = config {
+                                *config = RateLimit::scoped_single(target.clone(), *kind);
+                                migrated = true;
+                            }
+                        }
+                        (Some(source), Some(target)) => {
+                            if let RateLimit::Scoped(map) = config {
+                                if let Some(kind) = map.remove(source) {
+                                    map.insert(target.clone(), kind);
+                                    migrated = true;
+                                }
+                            }
+                        }
+                        (Some(source), None) => {
+                            if let RateLimit::Scoped(map) = config {
+                                if map.len() == 1 && map.contains_key(source) {
+                                    if let Some(kind) = map.remove(source) {
+                                        *config = RateLimit::global(kind);
+                                        migrated = true;
+                                    }
+                                }
+                            }
+                        }
+                        _ => {}
+                    }
+                }
+            });
+
+            migrated
+        }
+
+        /// Migrates the cached usage information for a rate-limited call to a new key.
+        ///
+        /// Returns `true` when an entry was moved. Passing identical keys simply checks that an
+        /// entry exists.
+        pub fn migrate_usage_key(
+            identifier: &TransactionIdentifier,
+            from: Option<<T as Config<I>>::UsageKey>,
+            to: Option<<T as Config<I>>::UsageKey>,
+        ) -> bool {
+            if from == to {
+                return LastSeen::<T, I>::contains_key(identifier, &to);
+            }
+
+            let Some(block) = LastSeen::<T, I>::take(identifier, from) else {
+                return false;
+            };
+
+            LastSeen::<T, I>::insert(identifier, to, block);
+            true
+        }
+
         /// Returns the configured limit for the specified pallet/extrinsic names, if any.
         pub fn limit_for_call_names(
             pallet_name: &str,

pallets/rate-limiting/src/tests.rs

@@ -176,6 +176,110 @@ fn is_within_limit_true_after_required_span() {
     });
 }
 
+#[test]
+fn migrate_limit_scope_global_to_scoped() {
+    new_test_ext().execute_with(|| {
+        let target_call =
+            RuntimeCall::System(frame_system::Call::<Test>::remark { remark: Vec::new() });
+        let identifier = identifier_for(&target_call);
+
+        Limits::<Test, ()>::insert(identifier, RateLimit::global(RateLimitKind::Exact(3)));
+
+        assert!(RateLimiting::migrate_limit_scope(
+            &identifier,
+            None,
+            Some(9)
+        ));
+
+        match RateLimiting::limits(identifier).expect("config") {
+            RateLimit::Scoped(map) => {
+                assert_eq!(map.len(), 1);
+                assert_eq!(map.get(&9), Some(&RateLimitKind::Exact(3)));
+            }
+            other => panic!("unexpected config: {:?}", other),
+        }
+    });
+}
+
+#[test]
+fn migrate_limit_scope_scoped_to_scoped() {
+    new_test_ext().execute_with(|| {
+        let target_call =
+            RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 });
+        let identifier = identifier_for(&target_call);
+
+        let mut map = sp_std::collections::btree_map::BTreeMap::new();
+        map.insert(1u16, RateLimitKind::Exact(4));
+        map.insert(2u16, RateLimitKind::Exact(6));
+        Limits::<Test, ()>::insert(identifier, RateLimit::Scoped(map));
+
+        assert!(RateLimiting::migrate_limit_scope(
+            &identifier,
+            Some(1),
+            Some(3)
+        ));
+
+        match RateLimiting::limits(identifier).expect("config") {
+            RateLimit::Scoped(map) => {
+                assert!(map.get(&1).is_none());
+                assert_eq!(map.get(&3), Some(&RateLimitKind::Exact(4)));
+                assert_eq!(map.get(&2), Some(&RateLimitKind::Exact(6)));
+            }
+            other => panic!("unexpected config: {:?}", other),
+        }
+    });
+}
+
+#[test]
+fn migrate_limit_scope_scoped_to_global() {
+    new_test_ext().execute_with(|| {
+        let target_call =
+            RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 });
+        let identifier = identifier_for(&target_call);
+
+        let mut map = sp_std::collections::btree_map::BTreeMap::new();
+        map.insert(7u16, RateLimitKind::Exact(8));
+        Limits::<Test, ()>::insert(identifier, RateLimit::Scoped(map));
+
+        assert!(RateLimiting::migrate_limit_scope(
+            &identifier,
+            Some(7),
+            None
+        ));
+
+        match RateLimiting::limits(identifier).expect("config") {
+            RateLimit::Global(kind) => assert_eq!(kind, RateLimitKind::Exact(8)),
+            other => panic!("unexpected config: {:?}", other),
+        }
+    });
+}
+
+#[test]
+fn migrate_usage_key_moves_entry() {
+    new_test_ext().execute_with(|| {
+        let target_call =
+            RuntimeCall::RateLimiting(RateLimitingCall::set_default_rate_limit { block_span: 0 });
+        let identifier = identifier_for(&target_call);
+
+        LastSeen::<Test, ()>::insert(identifier, Some(5u16), 11);
+
+        assert!(RateLimiting::migrate_usage_key(
+            &identifier,
+            Some(5),
+            Some(6)
+        ));
+        assert!(LastSeen::<Test, ()>::get(identifier, Some(5u16)).is_none());
+        assert_eq!(LastSeen::<Test, ()>::get(identifier, Some(6u16)), Some(11));
+
+        assert!(RateLimiting::migrate_usage_key(&identifier, Some(6), None));
+        assert!(LastSeen::<Test, ()>::get(identifier, Some(6u16)).is_none());
+        assert_eq!(
+            LastSeen::<Test, ()>::get(identifier, None::<UsageKey>),
+            Some(11)
+        );
+    });
+}
+
 #[test]
 fn set_rate_limit_updates_storage_and_emits_event() {
     new_test_ext().execute_with(|| {