diff --git a/Cargo.lock b/Cargo.lock index 6e7db9de3187..ed21812b174f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -21395,6 +21395,7 @@ dependencies = [ "rstest", "serde", "serde_json", + "strum 0.26.3", "tempfile", "tokio", "url", diff --git a/rs/orchestrator/src/firewall.rs b/rs/orchestrator/src/firewall.rs index 60ef00ff1491..644c69baaa4e 100644 --- a/rs/orchestrator/src/firewall.rs +++ b/rs/orchestrator/src/firewall.rs @@ -841,7 +841,7 @@ mod tests { SubnetRecordBuilder, add_single_subnet_record, add_subnet_list_record, }; use ic_test_utilities_types::ids::{SUBNET_1, node_test_id, subnet_test_id}; - use rstest::rstest; + use strum::IntoEnumIterator; use super::*; @@ -983,28 +983,22 @@ mod tests { ); } - #[rstest] - fn nftables_golden_assigned_replica_test( + #[test] + fn nftables_golden_assigned_replica_test() { // For assigned replicas, only Type4 (cloud engine) nodes have a different firewall - #[values( - None, - Some(NodeRewardType::Unspecified), - Some(NodeRewardType::Type0), - Some(NodeRewardType::Type1), - Some(NodeRewardType::Type2), - Some(NodeRewardType::Type3), - Some(NodeRewardType::Type3dot1), - Some(NodeRewardType::Type1dot1) - )] - reward_type: Option, - ) { - golden_test( - Role::AssignedReplica(SUBNET_ID), - node_test_id(0), - reward_type, - NFTABLES_ASSIGNED_REPLICA_GOLDEN_BYTES, - "assigned_replica", - ); + for reward_type in NodeRewardType::iter() + .filter(|reward_type| *reward_type != NodeRewardType::Type4) + .map(Some) + .chain(std::iter::once(None)) + { + golden_test( + Role::AssignedReplica(SUBNET_ID), + node_test_id(0), + reward_type, + NFTABLES_ASSIGNED_REPLICA_GOLDEN_BYTES, + "assigned_replica", + ); + } } #[test] @@ -1018,28 +1012,22 @@ mod tests { ); } - #[rstest] - fn nftables_unassigned_replica_golden_test( + #[test] + fn nftables_unassigned_replica_golden_test() { // For unassigned replicas, only Type4 (cloud engine) nodes have a different firewall - #[values( - None, - Some(NodeRewardType::Unspecified), - Some(NodeRewardType::Type0), - Some(NodeRewardType::Type1), - Some(NodeRewardType::Type2), - Some(NodeRewardType::Type3), - Some(NodeRewardType::Type3dot1), - Some(NodeRewardType::Type1dot1) - )] - reward_type: Option, - ) { - golden_test( - Role::UnassignedReplica, - node_test_id(0), - reward_type, - NFTABLES_UNASSIGNED_REPLICA_GOLDEN_BYTES, - "unassigned_replica", - ); + for reward_type in NodeRewardType::iter() + .filter(|reward_type| *reward_type != NodeRewardType::Type4) + .map(Some) + .chain(std::iter::once(None)) + { + golden_test( + Role::UnassignedReplica, + node_test_id(0), + reward_type, + NFTABLES_UNASSIGNED_REPLICA_GOLDEN_BYTES, + "unassigned_replica", + ); + } } #[test] @@ -1053,64 +1041,48 @@ mod tests { ); } - #[rstest] - fn nftables_golden_boundary_node_system_subnet_test( - // For boundary nodes, the node reward type has no effect on the firewall - #[values( - None, - Some(NodeRewardType::Unspecified), - Some(NodeRewardType::Type0), - Some(NodeRewardType::Type1), - Some(NodeRewardType::Type2), - Some(NodeRewardType::Type3), - Some(NodeRewardType::Type3dot1), - Some(NodeRewardType::Type1dot1), - Some(NodeRewardType::Type4) - )] - reward_type: Option, - ) { + #[test] + fn nftables_golden_boundary_node_system_subnet_test() { // pick the node id such that the API BN's SOCKS proxy serves system subnet nodes // the assert checks that let api_bn_id_for_system_subnet = node_test_id(0); assert!(api_bn_id_for_system_subnet < node_test_id(API_BOUNDARY_NODE_ID)); - golden_test( - Role::BoundaryNode, - api_bn_id_for_system_subnet, - reward_type, - NFTABLES_BOUNDARY_NODE_SYSTEM_SUBNET_GOLDEN_BYTES, - "boundary_node_system_subnet", - ); + // For boundary nodes, the node reward type has no effect on the firewall + for reward_type in NodeRewardType::iter() + .map(Some) + .chain(std::iter::once(None)) + { + golden_test( + Role::BoundaryNode, + api_bn_id_for_system_subnet, + reward_type, + NFTABLES_BOUNDARY_NODE_SYSTEM_SUBNET_GOLDEN_BYTES, + "boundary_node_system_subnet", + ); + } } - #[rstest] - fn nftables_golden_boundary_node_app_subnet_test( - // For boundary nodes, the node reward type has no effect on the firewall - #[values( - None, - Some(NodeRewardType::Unspecified), - Some(NodeRewardType::Type0), - Some(NodeRewardType::Type1), - Some(NodeRewardType::Type2), - Some(NodeRewardType::Type3), - Some(NodeRewardType::Type3dot1), - Some(NodeRewardType::Type1dot1), - Some(NodeRewardType::Type4) - )] - reward_type: Option, - ) { + #[test] + fn nftables_golden_boundary_node_app_subnet_test() { // pick the node id such that the API BN's SOCKS proxy serves app subnet nodes // the assert checks that let api_bn_id_for_app_subnet = node_test_id(1234); assert!(api_bn_id_for_app_subnet > node_test_id(API_BOUNDARY_NODE_ID)); - golden_test( - Role::BoundaryNode, - api_bn_id_for_app_subnet, - reward_type, - NFTABLES_BOUNDARY_NODE_APP_SUBNET_GOLDEN_BYTES, - "boundary_node_app_subnet", - ); + // For boundary nodes, the node reward type has no effect on the firewall + for reward_type in NodeRewardType::iter() + .map(Some) + .chain(std::iter::once(None)) + { + golden_test( + Role::BoundaryNode, + api_bn_id_for_app_subnet, + reward_type, + NFTABLES_BOUNDARY_NODE_APP_SUBNET_GOLDEN_BYTES, + "boundary_node_app_subnet", + ); + } } /// Runs [`Firewall::check_for_firewall_config`] and compares the output against the specified diff --git a/rs/protobuf/generator/src/lib.rs b/rs/protobuf/generator/src/lib.rs index 2f8f9b1cd434..94511a6bedbb 100644 --- a/rs/protobuf/generator/src/lib.rs +++ b/rs/protobuf/generator/src/lib.rs @@ -197,6 +197,10 @@ fn build_registry_proto(def: &Path, out: &Path) { ".registry.node", "#[derive(serde::Serialize, serde::Deserialize)]", ); + config.type_attribute( + ".registry.node.v1.NodeRewardType", + "#[derive(strum::EnumIter)]", + ); config.type_attribute( ".registry.firewall", "#[derive(candid::CandidType, serde::Serialize, serde::Deserialize)]", diff --git a/rs/protobuf/src/gen/registry/registry.node.v1.rs b/rs/protobuf/src/gen/registry/registry.node.v1.rs index 00aef3c5e072..d49ff69e247e 100644 --- a/rs/protobuf/src/gen/registry/registry.node.v1.rs +++ b/rs/protobuf/src/gen/registry/registry.node.v1.rs @@ -72,6 +72,7 @@ pub struct NodeRecord { #[derive( serde::Serialize, serde::Deserialize, + strum::EnumIter, Clone, Copy, Debug, diff --git a/rs/registry/canister/BUILD.bazel b/rs/registry/canister/BUILD.bazel index cf6220622ede..bfb32e17312c 100644 --- a/rs/registry/canister/BUILD.bazel +++ b/rs/registry/canister/BUILD.bazel @@ -61,6 +61,7 @@ DEPENDENCIES = [ "@crate_index//:maplit", "@crate_index//:prost", "@crate_index//:serde", + "@crate_index//:strum", "@crate_index//:url", ] diff --git a/rs/registry/canister/Cargo.toml b/rs/registry/canister/Cargo.toml index f5896d35aa61..4bdbb480ca1b 100644 --- a/rs/registry/canister/Cargo.toml +++ b/rs/registry/canister/Cargo.toml @@ -59,6 +59,7 @@ maplit = "1.0.2" on_wire = { path = "../../rust_canisters/on_wire" } prost = { workspace = true } serde = { workspace = true } +strum = { workspace = true } url = { workspace = true } canbench-rs = { workspace = true, optional = true } rand = { workspace = true } diff --git a/rs/registry/canister/src/invariants/subnet/tests.rs b/rs/registry/canister/src/invariants/subnet/tests.rs index 3b1fe19e3344..93fd0bc1ece7 100644 --- a/rs/registry/canister/src/invariants/subnet/tests.rs +++ b/rs/registry/canister/src/invariants/subnet/tests.rs @@ -13,6 +13,7 @@ use ic_protobuf::{ }; use ic_registry_keys::{make_node_record_key, make_subnet_list_record_key, make_subnet_record_key}; use ic_test_utilities_types::ids::{node_test_id, subnet_test_id, user_test_id}; +use strum::IntoEnumIterator; #[test] fn only_application_subnets_and_engines_can_be_free_cycles_cost_schedule() { @@ -290,16 +291,11 @@ fn cloud_engine_subnets_must_have_type4_nodes() { check_subnet_invariants(&snapshot).unwrap(); // Sad case: CloudEngine subnet with nodes that have a non-Type4 reward type. - for reward_type in [ - None, - Some(NodeRewardType::Unspecified), - Some(NodeRewardType::Type0), - Some(NodeRewardType::Type1), - Some(NodeRewardType::Type2), - Some(NodeRewardType::Type3), - Some(NodeRewardType::Type3dot1), - Some(NodeRewardType::Type1dot1), - ] { + for reward_type in NodeRewardType::iter() + .filter(|reward_type| *reward_type != NodeRewardType::Type4) + .map(Some) + .chain(std::iter::once(None)) + { println!( "Ensuring that a node with reward type {reward_type:?} cannot be part of a CloudEngine subnet" );