Draft: feat: Allow initial DKG on different subnets

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,10 @@ 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)?;
+            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 +586,16 @@ 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));
+        let args = if let Some(subnet_id) = subnet_id {
+            args.with_subnet_id(subnet_id)
+        } else {
+            args
+        };
+        Encode!(&args).unwrap()
+    }
+
     fn ecdsa_sign_request(key_id: EcdsaKeyId) -> Vec<u8> {
         let args = SignWithECDSAArgs {
             message_hash: [1; 32],
@@ -669,6 +676,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/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/registry/admin/bin/create_subnet.rs

@@ -37,6 +37,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>,
@@ -302,6 +307,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(Into::into),
             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,
@@ -390,6 +396,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,
@@ -505,6 +512,19 @@ 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 cmd = ProposeToCreateSubnetCmd {
+            initial_dkg_subnet_id: Some(initial_dkg_subnet_id),
+            ..empty_propose_to_create_subnet_cmd()
+        };
+        assert_eq!(
+            cmd.new_payload().initial_dkg_subnet_id,
+            Some(initial_dkg_subnet_id.into())
+        );
+    }
+
     #[test]
     #[should_panic(
         expected = "must specify 'pre_signatures_to_create_in_advance' for key ecdsa:Secp256k1:some_key_name"

rs/registry/canister/canister/registry.did

@@ -101,6 +101,7 @@ type CreateSubnetPayload = record {
   replica_version_id : text;
   dkg_interval_length : nat64;
   subnet_id_override : opt principal;
+  initial_dkg_subnet_id : opt principal;
   ssh_backup_access : vec text;
   initial_notary_delay_millis : nat64;
   subnet_type : SubnetType;

rs/registry/canister/canister/registry_test.did

@@ -101,6 +101,7 @@ type CreateSubnetPayload = record {
   replica_version_id : text;
   dkg_interval_length : nat64;
   subnet_id_override : opt principal;
+  initial_dkg_subnet_id : opt principal;
   ssh_backup_access : vec text;
   initial_notary_delay_millis : nat64;
   subnet_type : SubnetType;

rs/registry/canister/src/mutations/do_create_subnet.rs

@@ -61,6 +61,11 @@ impl Registry {
             payload.node_ids.clone(),
             RegistryVersion::new(self.latest_version()),
         );
+        let request = if let Some(initial_dkg_subnet_id) = payload.initial_dkg_subnet_id {
+            request.with_subnet_id(initial_dkg_subnet_id)
+        } else {
+            request
+        };
 
         // 2a. Invoke NI-DKG on ic_00
         let response_bytes = call(
@@ -266,6 +271,9 @@ pub struct CreateSubnetPayload {
     pub node_ids: Vec<NodeId>,
 
     pub subnet_id_override: Option<PrincipalId>,
+    /// Optional subnet that should handle `setup_initial_dkg`.
+    /// If not set, the request is handled by the NNS subnet as before.
+    pub initial_dkg_subnet_id: Option<SubnetId>,
 
     pub max_ingress_bytes_per_message: u64,
     pub max_ingress_bytes_per_block: Option<u64>,

rs/registry/canister/tests/create_subnet.rs

@@ -405,6 +405,7 @@ fn make_create_subnet_payload(node_ids: Vec<NodeId>) -> CreateSubnetPayload {
     CreateSubnetPayload {
         node_ids,
         subnet_id_override: None,
+        initial_dkg_subnet_id: None,
         max_ingress_bytes_per_message: 60 * 1024 * 1024,
         max_ingress_bytes_per_block: None,
         max_ingress_messages_per_block: 1000,

rs/tests/consensus/tecdsa/utils/src/lib.rs

@@ -1079,6 +1079,7 @@ pub async fn create_new_subnet_with_keys(
     let payload = CreateSubnetPayload {
         node_ids,
         subnet_id_override: None,
+        initial_dkg_subnet_id: None,
         max_ingress_bytes_per_message: config.max_ingress_bytes_per_message,
         max_ingress_bytes_per_block: Some(config.max_ingress_bytes_per_block),
         max_ingress_messages_per_block: config.max_ingress_messages_per_block,

rs/tests/driver/src/nns.rs

@@ -721,13 +721,35 @@ pub async fn submit_create_application_subnet_proposal(
     replica_version: ReplicaVersion,
     cost_schedule: Option<CanisterCyclesCostSchedule>,
     max_number_of_canisters: Option<u64>,
+) -> ProposalId {
+    submit_create_application_subnet_proposal_with_initial_dkg_subnet(
+        governance,
+        node_ids,
+        replica_version,
+        cost_schedule,
+        max_number_of_canisters,
+        None,
+    )
+    .await
+}
+
+/// Submits a proposal for creating an application subnet and optionally
+/// selecting the subnet that handles initial DKG.
+pub async fn submit_create_application_subnet_proposal_with_initial_dkg_subnet(
+    governance: &Canister<'_>,
+    node_ids: Vec<NodeId>,
+    replica_version: ReplicaVersion,
+    cost_schedule: Option<CanisterCyclesCostSchedule>,
+    max_number_of_canisters: Option<u64>,
+    initial_dkg_subnet_id: Option<SubnetId>,
 ) -> ProposalId {
     let config =
         subnet_configuration::get_default_config_params(SubnetType::Application, node_ids.len());
     let max_number_of_canisters = max_number_of_canisters.unwrap_or(0);
     let payload = CreateSubnetPayload {
         node_ids,
         subnet_id_override: None,
+        initial_dkg_subnet_id,
         max_ingress_bytes_per_message: config.max_ingress_bytes_per_message,
         max_ingress_bytes_per_block: Some(config.max_ingress_bytes_per_block),
         max_ingress_messages_per_block: config.max_ingress_messages_per_block,

rs/tests/nns/BUILD.bazel

@@ -33,7 +33,7 @@ system_test_nns(
     ],
 )
 
-system_test(
+system_test_nns(
     name = "create_subnet_test",
     tags = [
         "long_test",

rs/tests/nns/create_subnet_test.rs

@@ -30,7 +30,8 @@ use ic_system_test_driver::driver::test_env_api::{
     HasPublicApiUrl, HasTopologySnapshot, IcNodeContainer, NnsInstallationBuilder, SshSession,
 };
 use ic_system_test_driver::nns::{
-    self, get_software_version_from_snapshot, submit_create_application_subnet_proposal,
+    self, get_software_version_from_snapshot,
+    submit_create_application_subnet_proposal_with_initial_dkg_subnet,
 };
 use ic_system_test_driver::nns::{get_subnet_list_from_registry, vote_on_proposal};
 use ic_system_test_driver::systest;
@@ -64,6 +65,10 @@ pub fn setup(env: TestEnv) {
             Subnet::fast(SubnetType::System, NNS_PRE_MASTER)
                 .with_dkg_interval_length(Height::from(DKG_INTERVAL_LENGTH)),
         )
+        .add_subnet(
+            Subnet::fast(SubnetType::Application, APP_PRE_MASTER)
+                .with_dkg_interval_length(Height::from(DKG_INTERVAL_LENGTH)),
+        )
         .with_unassigned_nodes(APP_PRE_MASTER * APP_SUBNETS)
         .setup_and_start(&env)
         .expect("failed to setup IC under test");
@@ -84,6 +89,11 @@ pub fn test(env: TestEnv) {
     install_nns_canisters(&env);
     let topology_snapshot = &env.topology_snapshot();
     let subnet = topology_snapshot.root_subnet();
+    let initial_dkg_subnet_id = topology_snapshot
+        .subnets()
+        .find(|s| s.subnet_id != subnet.subnet_id)
+        .expect("missing second subnet for setup_initial_dkg")
+        .subnet_id;
     let endpoint = subnet.nodes().next().unwrap();
 
     // get IDs of all unassigned nodes
@@ -118,12 +128,13 @@ pub fn test(env: TestEnv) {
                 log,
                 "Submitting proposal to create subnet with nodes: {nodes:?}"
             );
-            let proposal_id = submit_create_application_subnet_proposal(
+            let proposal_id = submit_create_application_subnet_proposal_with_initial_dkg_subnet(
                 &governance,
                 nodes,
                 version.clone(),
                 Some(CanisterCyclesCostSchedule::Normal),
                 Some(0),
+                Some(initial_dkg_subnet_id),
             )
             .await;
             info!(log, "Voting on proposal {proposal_id}");

rs/types/management_canister_types/src/lib.rs

@@ -2608,12 +2608,14 @@ impl<'a> Payload<'a> for CreateCanisterArgs {
 /// record {
 ///   node_ids : vec principal;
 ///   registry_version : nat64;
+///   subnet_id : opt principal;
 /// }
 /// ```
 #[derive(Debug, CandidType, Deserialize)]
 pub struct SetupInitialDKGArgs {
     node_ids: Vec<PrincipalId>,
     registry_version: u64,
+    subnet_id: Option<SubnetId>,
 }
 
 impl Payload<'_> for SetupInitialDKGArgs {}
@@ -2623,6 +2625,14 @@ impl SetupInitialDKGArgs {
         Self {
             node_ids: node_ids.iter().map(|node_id| node_id.get()).collect(),
             registry_version: registry_version.get(),
+            subnet_id: None,
+        }
+    }
+
+    pub fn with_subnet_id(self, subnet_id: SubnetId) -> Self {
+        Self {
+            subnet_id: Some(subnet_id),
+            ..self
         }
     }
 
@@ -2642,6 +2652,10 @@ impl SetupInitialDKGArgs {
     pub fn get_registry_version(&self) -> RegistryVersion {
         RegistryVersion::new(self.registry_version)
     }
+
+    pub fn get_subnet_id(&self) -> Option<SubnetId> {
+        self.subnet_id
+    }
 }
 
 /// Represents the response for a request to setup an initial DKG for a new