feat: allowing engine controller canister to perform directly membership changes (#10625)

In order to allow people to replace nodes in their engines we need to
allow the engine controller to proxy the calls on behalf of engine
owners.

---------

Co-authored-by: r-birkner <103420898+r-birkner@users.noreply.github.com>

Author: NikolaMilosa

rs/engine_controller/canister/canister.rs

@@ -7,8 +7,8 @@ use candid::Principal;
 use ic_base_types::{NodeId, PrincipalId, SubnetId};
 use ic_cdk::{api::msg_caller, call::Call, init, post_upgrade, println, update};
 use ic_engine_controller::{
-    CreateEngineArgs, DeleteEngineArgs, DeployGuestosToAllSubnetNodesPayload,
-    EngineControllerInitArgs, NewSubnet, UpdateSubnetPayload,
+    ChangeSubnetMembershipPayload, CreateEngineArgs, DeleteEngineArgs,
+    DeployGuestosToAllSubnetNodesPayload, EngineControllerInitArgs, NewSubnet, UpdateSubnetPayload,
 };
 use ic_nns_constants::REGISTRY_CANISTER_ID;
 use ic_protobuf::registry::subnet::v1::SubnetFeatures;
@@ -361,6 +361,23 @@ async fn deploy_guestos_to_all_subnet_nodes(
     Ok(())
 }
 
+/// Proxies to the registry's `change_subnet_membership` endpoint. The
+/// registry enforces that, when invoked by the engine controller, the target
+/// subnet must be of type `CloudEngine`.
+#[update]
+async fn change_subnet_membership(payload: ChangeSubnetMembershipPayload) -> Result<(), String> {
+    ensure_authorized()?;
+
+    Call::unbounded_wait(REGISTRY_CANISTER_ID.into(), "change_subnet_membership")
+        .with_arg(payload)
+        .await
+        .map_err(|e| format!("registry.change_subnet_membership call failed: {e:?}"))?
+        .candid::<()>()
+        .map_err(|e| format!("Failed to decode registry response: {e}"))?;
+
+    Ok(())
+}
+
 fn main() {
     // This block is intentionally left blank.
 }

rs/engine_controller/engine_controller.did

@@ -107,9 +107,20 @@ type DeployGuestosToAllSubnetNodesPayload = record {
   replica_version_id : text;
 };
 
+// Mirror of `registry_canister::mutations::do_change_subnet_membership::
+// ChangeSubnetMembershipPayload`. Forwarded unchanged to the registry; the
+// registry enforces that the target subnet is of type `CloudEngine` when the
+// caller is the engine controller canister.
+type ChangeSubnetMembershipPayload = record {
+  subnet_id : principal;
+  node_ids_add : vec principal;
+  node_ids_remove : vec principal;
+};
+
 service : (opt EngineControllerInitArgs) -> {
   create_engine : (CreateEngineArgs) -> (CreateEngineResult);
   delete_engine : (DeleteEngineArgs) -> (Result);
   update_subnet : (UpdateSubnetPayload) -> (Result);
+  change_subnet_membership : (ChangeSubnetMembershipPayload) -> (Result);
   deploy_guestos_to_all_subnet_nodes : (DeployGuestosToAllSubnetNodesPayload) -> (Result);
 }

rs/engine_controller/src/lib.rs

@@ -13,6 +13,7 @@ pub use registry_canister::mutations::do_create_subnet::NewSubnet;
 // Re-export the payload types accepted by the proxy endpoints
 // (`update_subnet` and `deploy_guestos_to_all_subnet_nodes`) so clients
 // don't have to depend on `registry-canister` directly to construct them.
+pub use registry_canister::mutations::do_change_subnet_membership::ChangeSubnetMembershipPayload;
 pub use registry_canister::mutations::do_deploy_guestos_to_all_subnet_nodes::DeployGuestosToAllSubnetNodesPayload;
 pub use registry_canister::mutations::do_update_subnet::UpdateSubnetPayload;
 

rs/engine_controller/unreleased_changelog.md

@@ -0,0 +1,24 @@
+# How This File Is Used
+
+In general, upcoming/unreleased behavior changes are described here. For details
+on the process that this file is part of, see
+`rs/nervous_system/changelog_process.md`.
+
+
+# Next Upgrade Proposal
+
+## Added
+
+* New `change_subnet_membership` update method that proxies to the registry's
+  `change_subnet_membership` endpoint. The registry restricts this to
+  `CloudEngine` subnets when the caller is the engine controller canister.
+
+## Changed
+
+## Deprecated
+
+## Removed
+
+## Fixed
+
+## Security

rs/registry/canister/canister/canister.rs

@@ -712,7 +712,7 @@ fn remove_nodes_from_subnet_(payload: RemoveNodesFromSubnetPayload) {
 
 #[unsafe(export_name = "canister_update change_subnet_membership")]
 fn change_subnet_membership() {
-    check_caller_is_governance_and_log("change_subnet_membership");
+    check_caller_is_governance_or_engine_controller_and_log("change_subnet_membership");
     over(candid_one, |payload: ChangeSubnetMembershipPayload| {
         change_subnet_membership_(payload)
     });

rs/registry/canister/src/mutations/do_change_subnet_membership.rs

@@ -6,23 +6,38 @@ use candid::{CandidType, Deserialize};
 #[cfg(target_arch = "wasm32")]
 use dfn_core::println;
 use ic_base_types::{NodeId, PrincipalId, SubnetId};
+use ic_nns_constants::ENGINE_CONTROLLER_CANISTER_ID;
+use ic_protobuf::registry::subnet::v1::SubnetType;
 use ic_registry_keys::make_subnet_record_key;
 use ic_registry_transport::upsert;
 use prost::Message;
 use serde::Serialize;
 
 impl Registry {
     /// Changes membership of nodes in a subnet record in the registry.
-    ///
-    /// This method is called by the governance canister, after a proposal
-    /// for modifying a subnet by changing the membership (adding/removing) has been accepted.
     pub fn do_change_subnet_membership(&mut self, payload: ChangeSubnetMembershipPayload) {
-        println!("{LOG_PREFIX}do_change_subnet_membership started: {payload:?}");
+        let caller = dfn_core::api::caller();
+        println!("{LOG_PREFIX}do_change_subnet_membership started (caller: {caller}): {payload:?}");
 
-        let nodes_to_add = payload.node_ids_add.clone();
         let subnet_id = SubnetId::from(payload.subnet_id);
         let mut subnet_record = self.get_subnet_or_panic(subnet_id);
 
+        // When invoked by the engine controller canister, restrict the operation to
+        // `CloudEngine` subnets. Governance retains full ability to change
+        // membership of any subnet.
+        if caller == ENGINE_CONTROLLER_CANISTER_ID.get() {
+            assert_eq!(
+                subnet_record.subnet_type,
+                i32::from(SubnetType::CloudEngine),
+                "{LOG_PREFIX}do_change_subnet_membership: the engine controller may \
+                 only change membership of CloudEngine subnets; subnet {subnet_id} \
+                 has subnet_type {}",
+                subnet_record.subnet_type
+            );
+        }
+
+        let nodes_to_add = payload.node_ids_add.clone();
+
         let current_subnet_nodes: Vec<NodeId> = subnet_record
             .membership
             .iter()
@@ -64,7 +79,9 @@ impl Registry {
         // Check the invariants and apply the mutations if invariants are satisfied
         self.maybe_apply_mutation_internal(mutations);
 
-        println!("{LOG_PREFIX}do_change_subnet_membership finished: {payload:?}");
+        println!(
+            "{LOG_PREFIX}do_change_subnet_membership finished (caller: {caller}): {payload:?}"
+        );
     }
 }
 

rs/registry/canister/tests/change_subnet_membership.rs

@@ -3,24 +3,44 @@ use std::{collections::BTreeMap, convert::TryFrom};
 
 use assert_matches::assert_matches;
 
-use candid::Encode;
+use candid::{Decode, Encode, Principal};
 use canister_test::PrincipalId;
 use dfn_candid::candid;
 use ic_base_types::NodeId;
+use ic_nervous_system_integration_tests::pocket_ic_helpers::nns::registry::decode_registry_value;
+use ic_nns_constants::{
+    ENGINE_CONTROLLER_CANISTER_ID, GOVERNANCE_CANISTER_ID, REGISTRY_CANISTER_ID,
+};
 use ic_nns_test_utils::{
     itest_helpers::{
         forward_call_via_universal_canister, set_up_registry_canister, set_up_universal_canister,
         state_machine_test_on_nns_subnet,
     },
-    registry::{get_value_or_panic, prepare_registry, prepare_registry_with_two_node_sets},
+    registry::{
+        INITIAL_MUTATION_ID, get_value_or_panic, invariant_compliant_mutation_as_atomic_req,
+        prepare_registry, prepare_registry_with_two_node_sets,
+    },
+};
+use ic_protobuf::registry::{
+    node::v1::{NodeRecord, NodeRewardType},
+    subnet::v1::{SubnetListRecord, SubnetRecord},
 };
-use ic_protobuf::registry::subnet::v1::{SubnetListRecord, SubnetRecord};
 use ic_registry_keys::{make_subnet_list_record_key, make_subnet_record_key};
+use pocket_ic::PocketIcBuilder;
+use pocket_ic::RejectResponse;
+use pocket_ic::nonblocking::PocketIc;
 use registry_canister::{
     init::RegistryCanisterInitPayloadBuilder,
     mutations::do_change_subnet_membership::ChangeSubnetMembershipPayload,
 };
 
+mod common;
+
+use common::test_helpers::{
+    install_registry_canister_with_payload_builder, prepare_registry_with_cloud_engine_subnet,
+    prepare_registry_with_nodes_from_template,
+};
+
 #[test]
 fn test_the_anonymous_user_cannot_change_subnet_membership() {
     state_machine_test_on_nns_subnet(|runtime| {
@@ -410,3 +430,160 @@ fn test_change_subnet_membership_duplicate_nodes() {
         }
     });
 }
+
+/// Sets up a `PocketIc` with the registry canister installed, containing a
+/// CloudEngine subnet plus one extra unassigned type-4 node that can be
+/// swapped into the subnet.
+async fn setup_cloud_engine_registry() -> (PocketIc, Principal, NodeId, NodeId) {
+    let pocket_ic = PocketIcBuilder::new().with_nns_subnet().build_async().await;
+
+    let (cloud_engine_mutate, cloud_engine_subnet_id) =
+        prepare_registry_with_cloud_engine_subnet(4, INITIAL_MUTATION_ID);
+
+    // Prepare an extra unassigned type-4 node to swap into the subnet.
+    // Type-4 is required because the CloudEngine subnet invariant enforces
+    // that all members are type-4 nodes.
+    let extra_node_template = NodeRecord {
+        node_operator_id: PrincipalId::new_user_test_id(999).into_vec(),
+        node_reward_type: Some(NodeRewardType::Type4 as i32),
+        ..Default::default()
+    };
+    // Use a starting mutation id well outside the range consumed by the
+    // CloudEngine setup so IPs / test ids do not collide.
+    let (extra_node_mutate, extra_node_ids_and_pks) = prepare_registry_with_nodes_from_template(
+        1,
+        INITIAL_MUTATION_ID + 100,
+        extra_node_template,
+    );
+    let extra_node_id = *extra_node_ids_and_pks
+        .keys()
+        .next()
+        .expect("expected one extra unassigned node");
+
+    // Any current member of the CloudEngine subnet can be picked as the one
+    // to remove; take the first one from the subnet record.
+    let mut builder = RegistryCanisterInitPayloadBuilder::new();
+    builder.push_init_mutate_request(invariant_compliant_mutation_as_atomic_req(0));
+    builder.push_init_mutate_request(cloud_engine_mutate);
+    builder.push_init_mutate_request(extra_node_mutate);
+    install_registry_canister_with_payload_builder(&pocket_ic, builder.build(), true).await;
+
+    // Read a current member of the CloudEngine subnet to use as the removal
+    // target.
+    let subnet_record = decode_registry_value::<SubnetRecord>(
+        &pocket_ic,
+        make_subnet_record_key(cloud_engine_subnet_id),
+    )
+    .await;
+    let member_to_remove = NodeId::from(
+        PrincipalId::try_from(subnet_record.membership.first().unwrap().as_slice()).unwrap(),
+    );
+
+    (
+        pocket_ic,
+        cloud_engine_subnet_id.get().0,
+        extra_node_id,
+        member_to_remove,
+    )
+}
+
+#[tokio::test]
+async fn test_engine_controller_can_change_membership_of_cloud_engine_subnet() {
+    let (pocket_ic, cloud_engine_subnet_id, node_to_add, node_to_remove) =
+        setup_cloud_engine_registry().await;
+
+    let payload = ChangeSubnetMembershipPayload {
+        subnet_id: PrincipalId::from(cloud_engine_subnet_id),
+        node_ids_add: vec![node_to_add],
+        node_ids_remove: vec![node_to_remove],
+    };
+
+    // The engine controller performs a 1-for-1 swap on the CloudEngine subnet.
+    // This is expected to succeed.
+    let response: Vec<u8> = pocket_ic
+        .update_call(
+            REGISTRY_CANISTER_ID.get().0,
+            ENGINE_CONTROLLER_CANISTER_ID.get().0,
+            "change_subnet_membership",
+            Encode!(&payload).unwrap(),
+        )
+        .await
+        .unwrap();
+    Decode!(&response, ()).unwrap();
+
+    // Verify the membership actually changed.
+    let subnet_record = decode_registry_value::<SubnetRecord>(
+        &pocket_ic,
+        make_subnet_record_key(ic_base_types::SubnetId::from(PrincipalId::from(
+            cloud_engine_subnet_id,
+        ))),
+    )
+    .await;
+    let members = subnet_record
+        .membership
+        .iter()
+        .map(|bytes| NodeId::from(PrincipalId::try_from(bytes.as_slice()).unwrap()))
+        .collect::<Vec<NodeId>>();
+    assert!(
+        members.contains(&node_to_add),
+        "swapped-in node {node_to_add} should now be a subnet member, got {members:?}"
+    );
+    assert!(
+        !members.contains(&node_to_remove),
+        "swapped-out node {node_to_remove} should no longer be a subnet member, got {members:?}"
+    );
+}
+
+#[tokio::test]
+async fn test_engine_controller_cannot_change_membership_of_non_cloud_engine_subnet() {
+    // Set up a registry that contains only the invariant-compliant system
+    // subnet (which is not a CloudEngine).
+    let pocket_ic = PocketIcBuilder::new().with_nns_subnet().build_async().await;
+    let mut builder = RegistryCanisterInitPayloadBuilder::new();
+    builder.push_init_mutate_request(invariant_compliant_mutation_as_atomic_req(0));
+    install_registry_canister_with_payload_builder(&pocket_ic, builder.build(), true).await;
+
+    // Identify the (only) non-CloudEngine subnet.
+    let subnet_list =
+        decode_registry_value::<SubnetListRecord>(&pocket_ic, make_subnet_list_record_key()).await;
+    let non_cloud_engine_subnet =
+        Principal::try_from(subnet_list.subnets.first().unwrap().as_slice()).unwrap();
+
+    let payload = ChangeSubnetMembershipPayload {
+        subnet_id: PrincipalId::from(non_cloud_engine_subnet),
+        node_ids_add: vec![],
+        node_ids_remove: vec![],
+    };
+
+    // First, establish that this payload is itself valid by having governance
+    // (which is exempt from the CloudEngine restriction) successfully apply
+    // it as a no-op change against the non-CloudEngine subnet.
+    let response: Vec<u8> = pocket_ic
+        .update_call(
+            REGISTRY_CANISTER_ID.get().0,
+            GOVERNANCE_CANISTER_ID.get().0,
+            "change_subnet_membership",
+            Encode!(&payload).unwrap(),
+        )
+        .await
+        .unwrap();
+    Decode!(&response, ()).unwrap();
+
+    // Now the engine controller submits the exact same payload; since the
+    // target subnet is not a CloudEngine, the call must be rejected.
+    let response: Result<Vec<u8>, RejectResponse> = pocket_ic
+        .update_call(
+            REGISTRY_CANISTER_ID.get().0,
+            ENGINE_CONTROLLER_CANISTER_ID.get().0,
+            "change_subnet_membership",
+            Encode!(&payload).unwrap(),
+        )
+        .await;
+
+    let err = response.unwrap_err();
+    assert!(
+        err.reject_message
+            .contains("may only change membership of CloudEngine subnets"),
+        "expected a CloudEngine-specific rejection message, got {err:?}"
+    );
+}

rs/registry/canister/unreleased_changelog.md

@@ -19,6 +19,10 @@ on the process that this file is part of, see
   (e.g. on-demand cloud provisioning). All other node providers remain subject
   to the standard limits, and the per-IP `add_node` rate limit continues to
   apply to everyone.
+* `change_subnet_membership` may now be called by the engine controller canister
+  in addition to the governance canister. When invoked by the engine controller,
+  the target subnet must be of type `CloudEngine`; governance retains
+  unrestricted access to any subnet.
 
 ## Deprecated