feat: CON-1716 Allow SetupInitialDkg requests to be answered by subnets other than NNS

rs/embedders/src/wasmtime_embedder/system_api/routing.rs

@@ -13,10 +13,10 @@ use ic_management_canister_types_private::{
     LoadCanisterSnapshotArgs, MasterPublicKeyId, Method as Ic00Method, NodeMetricsHistoryArgs,
     Payload, ProvisionalTopUpCanisterArgs, ReadCanisterSnapshotDataArgs,
     ReadCanisterSnapshotMetadataArgs, RenameCanisterArgs, ReshareChainKeyArgs,
-    SchnorrPublicKeyArgs, SignWithECDSAArgs, SignWithSchnorrArgs, StoredChunksArgs, SubnetInfoArgs,
-    TakeCanisterSnapshotArgs, UninstallCodeArgs, UpdateSettingsArgs,
-    UploadCanisterSnapshotDataArgs, UploadCanisterSnapshotMetadataArgs, UploadChunkArgs,
-    VetKdDeriveKeyArgs, VetKdPublicKeyArgs,
+    SchnorrPublicKeyArgs, SetupInitialDKGArgs, SignWithECDSAArgs, SignWithSchnorrArgs,
+    StoredChunksArgs, SubnetInfoArgs, TakeCanisterSnapshotArgs, UninstallCodeArgs,
+    UpdateSettingsArgs, UploadCanisterSnapshotDataArgs, UploadCanisterSnapshotMetadataArgs,
+    UploadChunkArgs, VetKdDeriveKeyArgs, VetKdPublicKeyArgs,
 };
 use ic_replicated_state::NetworkTopology;
 use itertools::Itertools;
@@ -72,13 +72,16 @@ pub(super) fn resolve_destination(
         | Ok(Ic00Method::FlexibleHttpRequest)
         | Ok(Ic00Method::BitcoinSendTransactionInternal)
         | Ok(Ic00Method::BitcoinGetSuccessors) => Ok(own_subnet.get()),
-        // This message needs to be routed to the NNS subnet.  We assume that
-        // this message can only be sent by canisters on the NNS subnet hence
-        // returning `own_subnet` here is fine.
-        //
-        // It might be cleaner to pipe in the actual NNS subnet id to this
-        // function and return that instead.
-        Ok(Ic00Method::SetupInitialDKG) => Ok(own_subnet.get()),
+        Ok(Ic00Method::SetupInitialDKG) => {
+            let args = SetupInitialDKGArgs::decode(payload)?;
+            // This message should be routed to the NNS subnet by default. We assume that
+            // this message can only be sent by canisters on the NNS subnet hence
+            // defaulting to `own_subnet` here is fine.
+            //
+            // It might be cleaner to pipe in the actual NNS subnet id to this
+            // function and return that instead.
+            Ok(args.get_subnet_id().unwrap_or(own_subnet).get())
+        }
         Ok(Ic00Method::UpdateSettings) => {
             // Find the destination canister from the payload.
             let args = UpdateSettingsArgs::decode(payload)?;
@@ -589,6 +592,12 @@ mod tests {
         Encode!(&args).unwrap()
     }
 
+    fn setup_initial_dkg_request(subnet_id: Option<SubnetId>) -> Vec<u8> {
+        let args =
+            SetupInitialDKGArgs::new(vec![node_test_id(0)], RegistryVersion::from(100), subnet_id);
+        Encode!(&args).unwrap()
+    }
+
     fn ecdsa_sign_request(key_id: EcdsaKeyId) -> Vec<u8> {
         let args = SignWithECDSAArgs {
             message_hash: [1; 32],
@@ -669,6 +678,45 @@ mod tests {
         }
     }
 
+    #[test]
+    fn resolve_setup_initial_dkg_defaults_to_own_subnet() {
+        let logger = no_op_logger();
+        let own_subnet = subnet_test_id(2);
+        assert_eq!(
+            resolve_destination(
+                &network_with_ecdsa_subnets(),
+                &Ic00Method::SetupInitialDKG.to_string(),
+                &setup_initial_dkg_request(None),
+                own_subnet,
+                canister_test_id(1),
+                false,
+                &logger,
+            )
+            .unwrap(),
+            own_subnet.get()
+        );
+    }
+
+    #[test]
+    fn resolve_setup_initial_dkg_routes_to_requested_subnet() {
+        let logger = no_op_logger();
+        let own_subnet = subnet_test_id(2);
+        let requested_subnet = subnet_test_id(1);
+        assert_eq!(
+            resolve_destination(
+                &network_with_ecdsa_subnets(),
+                &Ic00Method::SetupInitialDKG.to_string(),
+                &setup_initial_dkg_request(Some(requested_subnet)),
+                own_subnet,
+                canister_test_id(1),
+                false,
+                &logger,
+            )
+            .unwrap(),
+            requested_subnet.get()
+        );
+    }
+
     #[test]
     fn resolve_reshare_chain_key_key_not_held_error() {
         let logger = no_op_logger();

rs/execution_environment/src/execution_environment/tests.rs

@@ -2273,6 +2273,7 @@ fn create_canister_xnet_called_from_nns() {
 fn setup_initial_dkg_sender_on_nns() {
     let own_subnet = subnet_test_id(1);
     let nns_subnet = subnet_test_id(2);
+    let other_subnet = subnet_test_id(3);
     let nns_canister = canister_test_id(1);
     let mut test = ExecutionTestBuilder::new()
         .with_subnet_type(SubnetType::System)
@@ -2281,12 +2282,15 @@ fn setup_initial_dkg_sender_on_nns() {
         .with_caller(nns_subnet, nns_canister)
         .build();
     let nodes = vec![node_test_id(1)];
-    let args = ic00::SetupInitialDKGArgs::new(nodes, RegistryVersion::new(1));
-    test.inject_call_to_ic00(
-        Method::SetupInitialDKG,
-        args.encode(),
-        test.canister_creation_fee().real(),
-    );
+    for subnet_id in [None, Some(own_subnet), Some(other_subnet), Some(nns_subnet)] {
+        let args =
+            ic00::SetupInitialDKGArgs::new(nodes.clone(), RegistryVersion::new(1), subnet_id);
+        test.inject_call_to_ic00(
+            Method::SetupInitialDKG,
+            args.encode(),
+            test.canister_creation_fee().real(),
+        );
+    }
     test.execute_all();
     assert_eq!(0, test.xnet_messages().len());
 }
@@ -2303,33 +2307,38 @@ fn setup_initial_dkg_sender_not_on_nns() {
         .with_caller(other_subnet, other_canister)
         .build();
     let nodes = vec![node_test_id(1)];
-    let args = ic00::SetupInitialDKGArgs::new(nodes, RegistryVersion::new(1));
-    test.inject_call_to_ic00(
-        Method::SetupInitialDKG,
-        args.encode(),
-        test.canister_creation_fee().real(),
-    );
+    for subnet_id in [None, Some(own_subnet), Some(other_subnet), Some(nns_subnet)] {
+        let args =
+            ic00::SetupInitialDKGArgs::new(nodes.clone(), RegistryVersion::new(1), subnet_id);
+        test.inject_call_to_ic00(
+            Method::SetupInitialDKG,
+            args.encode(),
+            test.canister_creation_fee().real(),
+        );
+    }
     test.execute_all();
-    let response = test.xnet_messages()[0].clone();
-    assert_eq!(
-        response,
-        Response {
-            originator: other_canister,
-            respondent: CanisterId::from(own_subnet),
-            originator_reply_callback: CallbackId::new(0),
-            refund: test.canister_creation_fee().real(),
-            response_payload: Payload::Reject(RejectContext::new(
-                RejectCode::CanisterError,
-                format!(
-                    "{} is called by {}. It can only be called by NNS.",
-                    ic00::Method::SetupInitialDKG,
-                    other_canister,
-                )
-            )),
-            deadline: NO_DEADLINE,
-        }
-        .into()
-    );
+    assert_eq!(test.xnet_messages().len(), 4);
+    for response in test.xnet_messages().iter() {
+        assert_eq!(
+            response.clone(),
+            Response {
+                originator: other_canister,
+                respondent: CanisterId::from(own_subnet),
+                originator_reply_callback: CallbackId::new(0),
+                refund: test.canister_creation_fee().real(),
+                response_payload: Payload::Reject(RejectContext::new(
+                    RejectCode::CanisterError,
+                    format!(
+                        "{} is called by {}. It can only be called by NNS.",
+                        ic00::Method::SetupInitialDKG,
+                        other_canister,
+                    )
+                )),
+                deadline: NO_DEADLINE,
+            }
+            .into()
+        );
+    }
 }
 
 #[test]

rs/nns/governance/src/proposals/fulfill_subnet_rental_request.rs

@@ -239,6 +239,7 @@ impl ValidFulfillSubnetRentalRequest {
 
             subnet_type: SubnetType::Application,
             subnet_id_override: None,
+            initial_dkg_subnet_id: None,
             start_as_nns: false,
             is_halted: false,
             chain_key_config: None,

rs/recovery/src/admin_helper.rs

@@ -195,6 +195,7 @@ impl AdminHelper {
         subnet_id: SubnetId,
         checkpoint_height: Height,
         state_hash: String,
+        initial_dkg_subnet_id: Option<SubnetId>,
         chain_key_config: Option<(ChainKeyConfig, SubnetId)>,
         replacement_nodes: &[NodeId],
         registry_params: Option<RegistryParams>,
@@ -208,6 +209,10 @@ impl AdminHelper {
             .add_argument("height", checkpoint_height)
             .add_argument("state-hash", state_hash);
 
+        if let Some(initial_dkg_subnet_id) = initial_dkg_subnet_id {
+            ic_admin.add_argument("initial-dkg-subnet-id", initial_dkg_subnet_id);
+        }
+
         if let Some((config, subnet_id)) = chain_key_config {
             let key_requests = config
                 .key_configs
@@ -538,6 +543,7 @@ mod tests {
                 Height::from(666),
                 "fake_state_hash".to_string(),
                 None,
+                None,
                 &[],
                 None,
                 UNIX_EPOCH + Duration::from_nanos(123456),
@@ -596,6 +602,7 @@ mod tests {
                 subnet_id_from_str(FAKE_SUBNET_ID_1),
                 Height::from(666),
                 "fake_state_hash".to_string(),
+                Some(subnet_id_from_str(FAKE_SUBNET_ID_2)),
                 Some((chain_key_config, subnet_id_from_str(FAKE_SUBNET_ID_2))),
                 &[node_id_from_str(FAKE_NODE_ID)],
                 Some(RegistryParams {
@@ -620,6 +627,7 @@ mod tests {
             --subnet-index gpvux-2ejnk-3hgmh-cegwf-iekfc-b7rzs-hrvep-5euo2-3ywz3-k3hcb-cqe \
             --height 666 \
             --state-hash fake_state_hash \
+            --initial-dkg-subnet-id mklno-zzmhy-zutel-oujwg-dzcli-h6nfy-2serg-gnwru-vuwck-hcxit-wqe \
             --initial-chain-key-configs-to-request '[\
                 {\"subnet_id\":\"mklno-zzmhy-zutel-oujwg-dzcli-h6nfy-2serg-gnwru-vuwck-hcxit-wqe\",\"key_id\":\"ecdsa:Secp256k1:test_key_1\",\"pre_signatures_to_create_in_advance\":\"77\",\"max_queue_size\":\"30\"},\
                 {\"subnet_id\":\"mklno-zzmhy-zutel-oujwg-dzcli-h6nfy-2serg-gnwru-vuwck-hcxit-wqe\",\"key_id\":\"schnorr:Bip340Secp256k1:test_key_2\",\"pre_signatures_to_create_in_advance\":\"12\",\"max_queue_size\":\"32\"},\

rs/recovery/src/app_subnet_recovery.rs

@@ -185,6 +185,11 @@ pub struct AppSubnetRecoveryArgs {
     #[clap(long, value_parser=crate::util::subnet_id_from_str)]
     pub chain_key_subnet_id: Option<SubnetId>,
 
+    /// Optional subnet used to run `setup_initial_dkg` during recovery CUP proposal.
+    /// If not set, the request is handled by the NNS subnet.
+    #[clap(long, value_parser=crate::util::subnet_id_from_str)]
+    pub initial_dkg_subnet_id: Option<SubnetId>,
+
     /// If present the tool will start execution for the provided step, skipping the initial ones
     #[clap(long = "resume")]
     pub next_step: Option<StepType>,
@@ -364,6 +369,12 @@ impl RecoveryIterator<StepType, StepTypeIter> for AppSubnetRecovery {
                         "Enter ID of subnet to reshare Chain keys from: ",
                     );
                 }
+                if self.params.initial_dkg_subnet_id.is_none() {
+                    self.params.initial_dkg_subnet_id = read_optional_subnet_id(
+                        &self.logger,
+                        "Enter ID of subnet to setup initial DKG on (default: NNS): ",
+                    );
+                }
             }
 
             StepType::UploadState => {
@@ -566,6 +577,7 @@ impl RecoveryIterator<StepType, StepTypeIter> for AppSubnetRecovery {
                     state_params.hash,
                     self.params.replacement_nodes.as_ref().unwrap_or(&default),
                     None,
+                    self.params.initial_dkg_subnet_id,
                     self.params.chain_key_subnet_id,
                 )?))
             }

rs/recovery/src/lib.rs

@@ -784,6 +784,7 @@ impl Recovery {
         state_hash: String,
         replacement_nodes: &[NodeId],
         registry_params: Option<RegistryParams>,
+        initial_dkg_subnet_id: Option<SubnetId>,
         chain_key_subnet_id: Option<SubnetId>,
     ) -> RecoveryResult<impl Step + use<>> {
         let chain_key_config = chain_key_subnet_id
@@ -811,6 +812,7 @@ impl Recovery {
                     subnet_id,
                     checkpoint_height,
                     state_hash,
+                    initial_dkg_subnet_id,
                     chain_key_config,
                     replacement_nodes,
                     registry_params,

rs/recovery/src/nns_recovery_failover_nodes.rs

@@ -412,6 +412,7 @@ impl RecoveryIterator<StepType, StepTypeIter> for NNSRecoveryFailoverNodes {
                         &[],
                         Some(registry_params),
                         None,
+                        None,
                     )?))
                 } else {
                     Err(RecoveryError::StepSkipped)

rs/recovery/src/recovery_state.rs

@@ -197,6 +197,7 @@ mod tests {
                 upload_method: None,
                 wait_for_cup_node: None,
                 chain_key_subnet_id: Some(fake_subnet_id()),
+                initial_dkg_subnet_id: None,
                 next_step: None,
                 upgrade_image_url: None,
                 upgrade_image_hash: None,

rs/recovery/subnet_splitting/src/subnet_splitting.rs

@@ -326,6 +326,7 @@ impl SubnetSplitting {
             state_hash,
             /*replacement_nodes=*/ &[],
             /*registry_params=*/ None,
+            /*initial_dkg_subnet_id=*/ None,
             /*chain_key_subnet_id=*/ None,
         )
     }

rs/registry/admin/bin/create_subnet.rs

@@ -14,7 +14,7 @@ use ic_protobuf::registry::subnet::v1::SubnetFeatures as SubnetFeaturesPb;
 use ic_registry_resource_limits::ResourceLimits;
 use ic_registry_subnet_features::SubnetFeatures;
 use ic_registry_subnet_type::SubnetType;
-use ic_types::{NodeId, PrincipalId, ReplicaVersion};
+use ic_types::{NodeId, PrincipalId, ReplicaVersion, SubnetId};
 use registry_canister::mutations::do_create_subnet;
 use registry_canister::mutations::do_create_subnet::CanisterCyclesCostSchedule;
 use std::collections::BTreeMap;
@@ -38,6 +38,11 @@ pub(crate) struct ProposeToCreateSubnetCmd {
     // Assigns this subnet ID to the newly created subnet
     pub subnet_id_override: Option<PrincipalId>,
 
+    #[clap(long)]
+    /// Optional subnet that should handle `setup_initial_dkg` for subnet creation.
+    /// If not set, handling defaults to the NNS subnet.
+    pub initial_dkg_subnet_id: Option<PrincipalId>,
+
     #[clap(long)]
     /// Maximum amount of bytes per message. This is a hard cap.
     pub max_ingress_bytes_per_message: Option<u64>,
@@ -314,6 +319,7 @@ impl ProposeToCreateSubnetCmd {
         do_create_subnet::CreateSubnetPayload {
             node_ids,
             subnet_id_override: self.subnet_id_override,
+            initial_dkg_subnet_id: self.initial_dkg_subnet_id.map(SubnetId::from),
             max_ingress_bytes_per_message: self.max_ingress_bytes_per_message.unwrap_or_default(),
             max_ingress_messages_per_block: self.max_ingress_messages_per_block.unwrap_or_default(),
             max_ingress_bytes_per_block: self.max_ingress_bytes_per_block,
@@ -401,6 +407,7 @@ mod tests {
             summary_file: None,
             subnet_handler_id: None,
             subnet_id_override: None,
+            initial_dkg_subnet_id: None,
             max_ingress_bytes_per_message: None,
             max_ingress_messages_per_block: None,
             max_ingress_bytes_per_block: None,
@@ -519,6 +526,20 @@ mod tests {
         );
     }
 
+    #[test]
+    fn cli_to_payload_conversion_includes_initial_dkg_subnet_id() {
+        let initial_dkg_subnet_id = PrincipalId::from_str("gxevo-lhkam-aaaaa-aaaap-yai").unwrap();
+        let mut cmd = ProposeToCreateSubnetCmd {
+            initial_dkg_subnet_id: Some(initial_dkg_subnet_id),
+            ..empty_propose_to_create_subnet_cmd()
+        };
+        cmd.apply_defaults_for_unset_fields();
+        assert_eq!(
+            cmd.new_payload().initial_dkg_subnet_id,
+            Some(SubnetId::from(initial_dkg_subnet_id))
+        );
+    }
+
     #[test]
     #[should_panic(
         expected = "must specify 'pre_signatures_to_create_in_advance' for key ecdsa:Secp256k1:some_key_name"