refactor: adding options to allow all without specifying all of the members in whitelisting for node swapping feature flags

rs/registry/canister/src/flags.rs

@@ -1,5 +1,7 @@
 use std::{
     cell::{Cell, RefCell},
+    collections::HashSet,
+    hash::Hash,
     str::FromStr,
 };
 
@@ -21,7 +23,7 @@ thread_local! {
     // These are needed for the phased rollout approach in order
     // allow granular rolling out of the feature to specific subnets
     // to specific subset of callers.
-    static NODE_SWAPPING_WHITELISTED_CALLERS: RefCell<Vec<PrincipalId>> = RefCell::new(
+    static NODE_SWAPPING_CALLERS_POLICY: RefCell<WhitelistFeatureAccessPolicy<PrincipalId>> = RefCell::new(WhitelistFeatureAccessPolicy::Only(
         [
             "xph6u-z3z2t-s7hh7-gtlxh-bbgbx-aatlm-eab4o-bsank-nqruh-3ub4q-sae",
             "lgp6d-brhlv-35izu-khc6p-rfszo-zdwng-xbtkh-xyvjg-y3due-7ha7t-uae",
@@ -44,9 +46,9 @@ thread_local! {
             }
         })
         .collect(),
-    );
+    ));
 
-    static NODE_SWAPPING_ENABLED_SUBNETS: RefCell<Vec<SubnetId>> = RefCell::new(
+    static NODE_SWAPPING_SUBNETS_POLICY: RefCell<WhitelistFeatureAccessPolicy<SubnetId>> = RefCell::new(WhitelistFeatureAccessPolicy::Only(
         [
             "2fq7c-slacv-26cgz-vzbx2-2jrcs-5edph-i5s2j-tck77-c3rlz-iobzx-mqe",
             "2zs4v-uoqha-xsuun-lveyr-i4ktc-5y3ju-aysud-niobd-gxnqa-ctqem-hae",
@@ -102,7 +104,7 @@ thread_local! {
             }
         })
         .collect(),
-    );
+    ));
 }
 
 pub(crate) fn is_chunkifying_large_values_enabled() -> bool {
@@ -142,20 +144,77 @@ pub mod temporary_overrides {
     }
 
     pub fn test_set_swapping_whitelisted_callers(override_callers: Vec<PrincipalId>) {
-        NODE_SWAPPING_WHITELISTED_CALLERS.replace(override_callers.into_iter().collect());
+        let policy = WhitelistFeatureAccessPolicy::Only(override_callers.into_iter().collect());
+        NODE_SWAPPING_CALLERS_POLICY.replace(policy);
     }
 
     pub fn test_set_swapping_enabled_subnets(override_subnets: Vec<SubnetId>) {
-        NODE_SWAPPING_ENABLED_SUBNETS.replace(override_subnets.into_iter().collect());
+        let policy = WhitelistFeatureAccessPolicy::Only(override_subnets.into_iter().collect());
+        NODE_SWAPPING_SUBNETS_POLICY.replace(policy);
     }
 }
 
 pub(crate) fn is_node_swapping_enabled_on_subnet(subnet_id: SubnetId) -> bool {
-    NODE_SWAPPING_ENABLED_SUBNETS
-        .with_borrow(|enabled_subnets| enabled_subnets.contains(&subnet_id))
+    NODE_SWAPPING_SUBNETS_POLICY.with_borrow(|subnet_policy| subnet_policy.is_allowed(&subnet_id))
 }
 
 pub(crate) fn is_node_swapping_enabled_for_caller(caller: PrincipalId) -> bool {
-    NODE_SWAPPING_WHITELISTED_CALLERS
-        .with_borrow(|enabled_callers| enabled_callers.contains(&caller))
+    NODE_SWAPPING_CALLERS_POLICY.with_borrow(|caller_policy| caller_policy.is_allowed(&caller))
+}
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+#[allow(dead_code)]
+enum WhitelistFeatureAccessPolicy<T>
+where
+    T: Clone + Eq + Hash,
+{
+    Only(HashSet<T>),
+    All,
+}
+
+impl<T> WhitelistFeatureAccessPolicy<T>
+where
+    T: Clone + Eq + Hash,
+{
+    fn is_allowed(&self, item: &T) -> bool {
+        match &self {
+            WhitelistFeatureAccessPolicy::Only(items) => items.contains(item),
+            WhitelistFeatureAccessPolicy::All => true,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_whitelist_allow_all() {
+        let policy = WhitelistFeatureAccessPolicy::All;
+
+        assert!(policy.is_allowed(&PrincipalId::new_user_test_id(1)));
+        assert!(policy.is_allowed(&PrincipalId::new_anonymous()));
+    }
+
+    #[test]
+    fn test_whitelist_deny_all() {
+        let policy = WhitelistFeatureAccessPolicy::Only(HashSet::new());
+
+        assert!(!policy.is_allowed(&PrincipalId::new_user_test_id(1)));
+        assert!(!policy.is_allowed(&PrincipalId::new_anonymous()));
+    }
+
+    #[test]
+    fn test_whitelist_allow_some() {
+        let user_1 = PrincipalId::new_user_test_id(1);
+        let policy = WhitelistFeatureAccessPolicy::Only(
+            [user_1, PrincipalId::new_user_test_id(2)]
+                .into_iter()
+                .collect(),
+        );
+
+        assert!(policy.is_allowed(&user_1));
+        assert!(!policy.is_allowed(&PrincipalId::new_user_test_id(999)));
+        assert!(!policy.is_allowed(&PrincipalId::new_anonymous()));
+    }
 }

rs/registry/canister/src/mutations/do_swap_node_in_subnet_directly.rs

@@ -391,6 +391,7 @@ mod tests {
     use ic_types::{NodeId, PrincipalId, SubnetId};
     use itertools::Itertools;
 
+    use crate::flags::{is_node_swapping_enabled_for_caller, is_node_swapping_enabled_on_subnet};
     use crate::{
         common::test_helpers::{
             get_invariant_compliant_subnet_record, invariant_compliant_registry,
@@ -1075,4 +1076,37 @@ mod tests {
         assert!(members.contains(&new_node_id));
         assert!(!members.contains(&old_node_id));
     }
+
+    #[test]
+    fn feature_flag_subnet_test() {
+        let subnet_1 = SubnetId::new(PrincipalId::new_subnet_test_id(1));
+        let subnet_2 = SubnetId::new(PrincipalId::new_subnet_test_id(2));
+
+        // Ensure that no subnets are whitelisted
+        test_set_swapping_enabled_subnets(vec![]);
+
+        assert!(!is_node_swapping_enabled_on_subnet(subnet_1));
+
+        // Now add the subnet to the whitelisted ones
+        test_set_swapping_enabled_subnets(vec![subnet_1]);
+        assert!(is_node_swapping_enabled_on_subnet(subnet_1));
+        assert!(!is_node_swapping_enabled_on_subnet(subnet_2));
+    }
+
+    #[test]
+    fn feature_flag_node_operator_test() {
+        let node_operator_1 = PrincipalId::new_user_test_id(1);
+        let node_operator_2 = PrincipalId::new_user_test_id(2);
+
+        // Ensure that no callers are whitelisted
+        test_set_swapping_whitelisted_callers(vec![]);
+
+        assert!(!is_node_swapping_enabled_for_caller(node_operator_1));
+        assert!(!is_node_swapping_enabled_for_caller(node_operator_2));
+
+        // Now add the node operator to the whitelisted ones
+        test_set_swapping_whitelisted_callers(vec![node_operator_1]);
+        assert!(is_node_swapping_enabled_for_caller(node_operator_1));
+        assert!(!is_node_swapping_enabled_for_caller(node_operator_2));
+    }
 }

rs/registry/canister/tests/swap_node_in_subnet_directly.rs

@@ -166,6 +166,9 @@ async fn caller_not_whitelisted() {
     });
     builder.enable_swapping_feature_globally();
     builder.enable_swapping_feature_for_subnet(subnet_id);
+    // In order to avoid the feature being enabled for all node
+    // operators there needs to be some other caller whitelisted.
+    builder.whitelist_swapping_feature_caller(PrincipalId::new_user_test_id(999));
 
     install_registry_canister_with_payload_builder(&pocket_ic, builder.build(), true).await;
 
@@ -221,6 +224,10 @@ async fn subnet_not_whitelisted() {
 
     builder.enable_swapping_feature_globally();
     builder.whitelist_swapping_feature_caller(node_operator_id);
+    // There needs to be at least one subnet specified as whitelisted
+    // to actually check the subnet whitelisting, otherwise they will
+    // all be deemed as whitelisted.
+    builder.enable_swapping_feature_for_subnet(SubnetId::new(PrincipalId::new_subnet_test_id(999)));
 
     install_registry_canister_with_payload_builder(&pocket_ic, builder.build(), true).await;