feat(sns): Enable upgrading SNS-controlled canisters using chunked WASMs

rs/nervous_system/integration_tests/src/pocket_ic_helpers.rs

@@ -1432,6 +1432,7 @@ pub mod sns {
         use ic_sns_governance::governance::UPGRADE_STEPS_INTERVAL_REFRESH_BACKOFF_SECONDS;
         use ic_sns_governance::pb::v1::get_neuron_response;
         use pocket_ic::ErrorCode;
+        use sns_pb::UpgradeSnsControlledCanister;
 
         pub const EXPECTED_UPGRADE_DURATION_MAX_SECONDS: u64 = 1000;
         pub const EXPECTED_UPGRADE_STEPS_REFRESH_MAX_SECONDS: u64 =
@@ -1726,6 +1727,44 @@ pub mod sns {
             assert!(proposal_data.executed_timestamp_seconds > 0);
         }
 
+        pub async fn propose_to_upgrade_sns_controlled_canister_and_wait(
+            pocket_ic: &PocketIc,
+            sns_governance_canister_id: PrincipalId,
+            upgrade: UpgradeSnsControlledCanister,
+        ) {
+            // Get an ID of an SNS neuron that can submit proposals. We rely on the fact that this
+            // neuron either holds the majority of the voting power or the follow graph is set up
+            // s.t. when this neuron submits a proposal, that proposal gets through without the need
+            // for any voting.
+            let (sns_neuron_id, sns_neuron_principal_id) =
+                find_neuron_with_majority_voting_power(pocket_ic, sns_governance_canister_id)
+                    .await
+                    .expect("cannot find SNS neuron with dissolve delay over 6 months.");
+
+            let proposal_data = propose_and_wait(
+                pocket_ic,
+                sns_governance_canister_id,
+                sns_neuron_principal_id,
+                sns_neuron_id.clone(),
+                sns_pb::Proposal {
+                    title: "Upgrade SNS controlled canister.".to_string(),
+                    summary: "".to_string(),
+                    url: "".to_string(),
+                    action: Some(sns_pb::proposal::Action::UpgradeSnsControlledCanister(
+                        upgrade,
+                    )),
+                },
+            )
+            .await
+            .unwrap();
+
+            // Check 1: The upgrade proposal did not fail.
+            assert_eq!(proposal_data.failure_reason, None);
+
+            // Check 2: The upgrade proposal succeeded.
+            assert!(proposal_data.executed_timestamp_seconds > 0);
+        }
+
         /// Get the neuron with the given ID from the SNS Governance canister.
         #[allow(dead_code)]
         async fn get_neuron(

rs/nervous_system/integration_tests/tests/upgrade_sns_controlled_canister_with_large_wasm.rs

@@ -2,7 +2,6 @@ use std::collections::BTreeSet;
 
 use candid::Principal;
 use canister_test::Wasm;
-use ic_base_types::PrincipalId;
 use ic_management_canister_types::CanisterInstallMode;
 use ic_nervous_system_integration_tests::pocket_ic_helpers::{
     await_with_timeout, install_canister_on_subnet, nns, sns,
@@ -11,10 +10,9 @@ use ic_nervous_system_integration_tests::{
     create_service_nervous_system_builder::CreateServiceNervousSystemBuilder,
     pocket_ic_helpers::{add_wasms_to_sns_wasm, install_nns_canisters},
 };
-use ic_nervous_system_root::change_canister::ChangeCanisterRequest;
-use ic_nervous_system_root::change_canister::ChunkedCanisterWasm;
 use ic_nns_constants::ROOT_CANISTER_ID;
 use ic_nns_test_utils::common::modify_wasm_bytes;
+use ic_sns_governance::pb::v1::{ChunkedCanisterWasm, UpgradeSnsControlledCanister};
 use ic_sns_swap::pb::v1::Lifecycle;
 use pocket_ic::nonblocking::PocketIc;
 use pocket_ic::PocketIcBuilder;
@@ -24,52 +22,21 @@ const MAX_INSTALL_CHUNKED_CODE_TIME_SECONDS: u64 = 5 * 60;
 
 const CHUNK_SIZE: usize = 1024 * 1024; // 1 MiB
 
-#[tokio::test]
-async fn test_store_same_as_target() {
-    let store_same_as_target = true;
-    run_test(store_same_as_target).await;
-}
+// TODO: Figure out how to best support uploading chunks into the target itself, which has
+// SNS Root as the controller but not SNS Governance.
+//
+// #[tokio::test]
+// async fn test_store_same_as_target() {
+//     let store_same_as_target = true;
+//     run_test(store_same_as_target).await;
+// }
 
 #[tokio::test]
 async fn test_store_different_from_target() {
     let store_same_as_target = false;
     run_test(store_same_as_target).await;
 }
 
-mod interim_sns_helpers {
-    use super::*;
-
-    use candid::{Decode, Encode};
-    use pocket_ic::nonblocking::PocketIc;
-    use pocket_ic::WasmResult;
-
-    /// Interim test function for calling Root.change_canister.
-    ///
-    /// This function is not in src/pocket_ic_helpers.rs because it's going to be replaced with
-    /// a proposal with the same effect. It should not be used in any other tests.
-    pub async fn change_canister(
-        pocket_ic: &PocketIc,
-        canister_id: PrincipalId,
-        sender: PrincipalId,
-        request: ChangeCanisterRequest,
-    ) {
-        let result = pocket_ic
-            .update_call(
-                canister_id.into(),
-                sender.into(),
-                "change_canister",
-                Encode!(&request).unwrap(),
-            )
-            .await
-            .unwrap();
-        let result = match result {
-            WasmResult::Reply(result) => result,
-            WasmResult::Reject(s) => panic!("Call to change_canister failed: {:#?}", s),
-        };
-        Decode!(&result, ()).unwrap()
-    }
-}
-
 fn very_large_wasm_bytes() -> Vec<u8> {
     let image_classification_canister_wasm_path =
         std::env::var("IMAGE_CLASSIFICATION_CANISTER_WASM_PATH")
@@ -212,7 +179,7 @@ async fn run_test(store_same_as_target: bool) {
             app_subnet,
             vec![],
             None,
-            vec![sns.root.canister_id],
+            vec![sns.governance.canister_id, sns.root.canister_id],
         )
         .await
     };
@@ -239,25 +206,19 @@ async fn run_test(store_same_as_target: bool) {
     .await;
 
     // 2. Run code under test.
-    interim_sns_helpers::change_canister(
+    sns::governance::propose_to_upgrade_sns_controlled_canister_and_wait(
         &pocket_ic,
-        sns.root.canister_id,
         sns.governance.canister_id,
-        ChangeCanisterRequest {
-            stop_before_installing: true,
-            mode: CanisterInstallMode::Upgrade,
-            canister_id: target_canister_id,
-            // This is the old field being generalized.
-            wasm_module: vec![],
-            // This is the new field we want to test.
+        UpgradeSnsControlledCanister {
+            canister_id: Some(target_canister_id.get()),
+            new_canister_wasm: vec![],
+            canister_upgrade_arg: None,
+            mode: Some(CanisterInstallMode::Upgrade as i32),
             chunked_canister_wasm: Some(ChunkedCanisterWasm {
                 wasm_module_hash: new_wasm_hash.clone().to_vec(),
-                store_canister_id,
+                store_canister_id: Some(store_canister_id.get()),
                 chunk_hashes_list,
             }),
-            arg: vec![],
-            compute_allocation: None,
-            memory_allocation: None,
         },
     )
     .await;

rs/nervous_system/root/src/change_canister.rs

@@ -64,7 +64,7 @@ pub struct ChangeCanisterRequest {
     #[serde(with = "serde_bytes")]
     pub wasm_module: Vec<u8>,
 
-    /// If the entire WASM does not into the 2 MiB ingress limit, then `new_canister_wasm`
+    /// If the entire WASM does not fit into the 2 MiB ingress limit, then `new_canister_wasm`
     /// should be empty, and this field should be set instead.
     pub chunked_canister_wasm: Option<ChunkedCanisterWasm>,
 

rs/sns/governance/api/src/ic_sns_governance.pb.v1.rs

@@ -366,6 +366,29 @@ pub struct Motion {
     #[prost(string, tag = "1")]
     pub motion_text: ::prost::alloc::string::String,
 }
+
+/// Represents a WASM split into smaller chunks, each of which can safely be sent around the ICP.
+#[derive(
+    candid::CandidType,
+    candid::Deserialize,
+    comparable::Comparable,
+    Clone,
+    PartialEq,
+    ::prost::Message,
+)]
+pub struct ChunkedCanisterWasm {
+    /// Obligatory check sum of the overall WASM to be reassembled from chunks.
+    #[prost(bytes = "vec", tag = "1")]
+    pub wasm_module_hash: ::prost::alloc::vec::Vec<u8>,
+    /// Obligatory; indicates which canister stores the WASM chunks.
+    #[prost(message, optional, tag = "2")]
+    pub store_canister_id: ::core::option::Option<::ic_base_types::PrincipalId>,
+    /// Specifies a list of hash values for the chunks that comprise this WASM. Must contain at least
+    /// one chunk.
+    #[prost(bytes = "vec", repeated, tag = "3")]
+    pub chunk_hashes_list: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec<u8>>,
+}
+
 /// A proposal function that upgrades a canister that is controlled by the
 /// SNS governance canister.
 #[derive(
@@ -395,6 +418,10 @@ pub struct UpgradeSnsControlledCanister {
         tag = "4"
     )]
     pub mode: ::core::option::Option<i32>,
+    /// If the entire WASM does not fit into the 2 MiB ingress limit, then `new_canister_wasm` should be
+    /// an empty, and this field should be set instead.
+    #[prost(message, optional, tag = "5")]
+    pub chunked_canister_wasm: ::core::option::Option<ChunkedCanisterWasm>,
 }
 /// A proposal to transfer SNS treasury funds to (optionally a Subaccount of) the
 /// target principal.

rs/sns/governance/canister/governance.did

@@ -711,8 +711,15 @@ type PendingVersion = record {
   target_version : opt Version;
 };
 
+type ChunkedCanisterWasm = record {
+  wasm_module_hash : blob;
+  store_canister_id : opt principal;
+  chunk_hashes_list : vec blob;
+};
+
 type UpgradeSnsControlledCanister = record {
   new_canister_wasm : blob;
+  chunked_canister_wasm : opt ChunkedCanisterWasm;
   mode : opt int32;
   canister_id : opt principal;
   canister_upgrade_arg : opt blob;

rs/sns/governance/canister/governance_test.did

@@ -725,8 +725,15 @@ type PendingVersion = record {
   target_version : opt Version;
 };
 
+type ChunkedCanisterWasm = record {
+  wasm_module_hash : blob;
+  store_canister_id : opt principal;
+  chunk_hashes_list : vec blob;
+};
+
 type UpgradeSnsControlledCanister = record {
   new_canister_wasm : blob;
+  chunked_canister_wasm : opt ChunkedCanisterWasm;
   mode : opt int32;
   canister_id : opt principal;
   canister_upgrade_arg : opt blob;

rs/sns/governance/proto/ic_sns_governance/pb/v1/governance.proto

@@ -312,6 +312,17 @@ message Motion {
   string motion_text = 1;
 }
 
+// Represents a WASM split into smaller chunks, each of which can safely be sent around the ICP.
+message ChunkedCanisterWasm {
+  // Obligatory check sum of the overall WASM to be reassembled from chunks.
+  bytes wasm_module_hash = 1;
+  // Obligatory; indicates which canister stores the WASM chunks.
+  ic_base_types.pb.v1.PrincipalId store_canister_id = 2;
+  // Specifies a list of hash values for the chunks that comprise this WASM. Must contain at least
+  // one chunk.
+  repeated bytes chunk_hashes_list = 3;
+}
+
 // A proposal function that upgrades a canister that is controlled by the
 // SNS governance canister.
 message UpgradeSnsControlledCanister {
@@ -323,6 +334,9 @@ message UpgradeSnsControlledCanister {
   optional bytes canister_upgrade_arg = 3;
   // Canister install_code mode.
   optional types.v1.CanisterInstallMode mode = 4;
+  // If the entire WASM does not fit into the 2 MiB ingress limit, then `new_canister_wasm` should be
+  // an empty, and this field should be set instead.
+  optional ChunkedCanisterWasm chunked_canister_wasm = 5;
 }
 
 // A proposal to transfer SNS treasury funds to (optionally a Subaccount of) the

rs/sns/governance/src/gen/ic_sns_governance.pb.v1.rs

@@ -366,6 +366,27 @@ pub struct Motion {
     #[prost(string, tag = "1")]
     pub motion_text: ::prost::alloc::string::String,
 }
+/// Represents a WASM split into smaller chunks, each of which can safely be sent around the ICP.
+#[derive(
+    candid::CandidType,
+    candid::Deserialize,
+    comparable::Comparable,
+    Clone,
+    PartialEq,
+    ::prost::Message,
+)]
+pub struct ChunkedCanisterWasm {
+    /// Obligatory check sum of the overall WASM to be reassembled from chunks.
+    #[prost(bytes = "vec", tag = "1")]
+    pub wasm_module_hash: ::prost::alloc::vec::Vec<u8>,
+    /// Obligatory; indicates which canister stores the WASM chunks.
+    #[prost(message, optional, tag = "2")]
+    pub store_canister_id: ::core::option::Option<::ic_base_types::PrincipalId>,
+    /// Specifies a list of hash values for the chunks that comprise this WASM. Must contain at least
+    /// one chunk.
+    #[prost(bytes = "vec", repeated, tag = "3")]
+    pub chunk_hashes_list: ::prost::alloc::vec::Vec<::prost::alloc::vec::Vec<u8>>,
+}
 /// A proposal function that upgrades a canister that is controlled by the
 /// SNS governance canister.
 #[derive(
@@ -395,6 +416,10 @@ pub struct UpgradeSnsControlledCanister {
         tag = "4"
     )]
     pub mode: ::core::option::Option<i32>,
+    /// If the entire WASM does not fit into the 2 MiB ingress limit, then `new_canister_wasm` should be
+    /// an empty, and this field should be set instead.
+    #[prost(message, optional, tag = "5")]
+    pub chunked_canister_wasm: ::core::option::Option<ChunkedCanisterWasm>,
 }
 /// A proposal to transfer SNS treasury funds to (optionally a Subaccount of) the
 /// target principal.