feat(nervous-system): Enable Root to upgrade canisters using chunked …

…Wasms (#3300)

This PR enables Root to upgrade canisters using chunked Wasms. This is
relevant to both the NNS and the SNSs.

---------

Co-authored-by: max-dfinity <[email protected]>

rs/nervous_system/common/test_utils/BUILD.bazel

@@ -8,6 +8,7 @@ package_group(
     name = "ic_nervous_system_common_test_utils_visibility",
     packages = [
         "//rs/nervous_system/common/...",
+        "//rs/nervous_system/integration_tests/...",
         "//rs/nns/...",
         "//rs/sns/...",
     ],

rs/nervous_system/common/test_utils/src/wasm_helpers.rs

@@ -1,5 +1,6 @@
 use ic_wasm;
 use libflate::gzip;
+use std::io::Read;
 
 /// A small, valid WASM suitable for tests.
 pub const SMALLEST_VALID_WASM_BYTES: &[u8; 8] = &[0, 0x61, 0x73, 0x6D, 1, 0, 0, 0];
@@ -25,7 +26,7 @@ pub fn annotate_wasm_with_metadata(
     wasm_module.emit_wasm()
 }
 
-// Gzips a wasm, returning the hash of its compressed representation.
+/// Gzips a wasm, returning the hash of its compressed representation.
 pub fn gzip_wasm(wasm: &[u8]) -> Vec<u8> {
     let mut encoder = gzip::Encoder::new(Vec::new()).expect("Failed to create gzip encoder.");
     std::io::copy(&mut &wasm[..], &mut encoder).expect("Failed to copy WASM bytes.");
@@ -34,3 +35,13 @@ pub fn gzip_wasm(wasm: &[u8]) -> Vec<u8> {
         .into_result()
         .expect("Failed to finish gzip encoding.")
 }
+
+/// Decompresses a previously gzipped wasm.
+pub fn ungzip_wasm(gzipped_bytes: &[u8]) -> Vec<u8> {
+    let mut decoder = gzip::Decoder::new(gzipped_bytes).expect("Failed to create gzip decoder.");
+    let mut wasm_buf = Vec::new();
+    decoder
+        .read_to_end(&mut wasm_buf)
+        .expect("Failed decoding Wasm.");
+    wasm_buf
+}

rs/nervous_system/integration_tests/BUILD.bazel

@@ -41,6 +41,7 @@ BASE_DEPENDENCIES = [
         "//packages/pocket-ic",
         "//rs/crypto/sha2",
         "//rs/nervous_system/common/test_keys",
+        "//rs/nervous_system/common/test_utils",
         "//rs/nns/constants",
         "//rs/protobuf",
         "//rs/registry/canister",
@@ -100,6 +101,7 @@ DEV_DATA = [
     "//rs/sns/root:sns-root-canister",
     "//rs/sns/swap:sns-swap-canister",
     "//rs/universal_canister/impl:universal_canister.wasm.gz",
+    "//testnet/prebuilt-canisters:image-classification",
     "@cycles-ledger.wasm.gz//file",
     "@mainnet_ic-icrc1-archive//file",
     "@mainnet_ic-icrc1-index-ng//file",
@@ -124,6 +126,7 @@ DEV_ENV = {
     "IC_ICRC1_ARCHIVE_WASM_PATH": "$(rootpath //rs/ledger_suite/icrc1/archive:archive_canister)",
     "IC_ICRC1_INDEX_NG_WASM_PATH": "$(rootpath //rs/ledger_suite/icrc1/index-ng:index_ng_canister)",
     "IC_ICRC1_LEDGER_WASM_PATH": "$(rootpath //rs/ledger_suite/icrc1/ledger:ledger_canister)",
+    "IMAGE_CLASSIFICATION_CANISTER_WASM_PATH": "$(rootpath //testnet/prebuilt-canisters:image-classification)",
     "LEDGER_CANISTER_WASM_PATH": "$(rootpath //rs/ledger_suite/icp/ledger:ledger-canister-wasm)",
     "LEDGER_CANISTER_NOTIFY_METHOD_WASM_PATH": "$(rootpath //rs/ledger_suite/icp/ledger:ledger-canister-wasm-notify-method)",
     "LEDGER_ARCHIVE_NODE_CANISTER_WASM_PATH": "$(rootpath //rs/ledger_suite/icp/archive:ledger-archive-node-canister-wasm)",
@@ -178,6 +181,7 @@ rust_test_suite_with_extra_srcs(
             "tests/deploy_fresh_sns_test.rs",
             "tests/sns_release_qualification_legacy.rs",
             "tests/sns_upgrade_test_utils_legacy.rs",
+            "tests/upgrade_sns_controlled_canister_with_large_wasm.rs",
         ],
     ),
     aliases = ALIASES,
@@ -279,3 +283,20 @@ rust_test(
     ],
     deps = [":nervous_system_integration_tests"] + DEPENDENCIES_WITH_TEST_FEATURES + DEV_DEPENDENCIES,
 )
+
+rust_test(
+    name = "upgrade_sns_controlled_canister_with_large_wasm",
+    timeout = "long",
+    srcs = [
+        "tests/upgrade_sns_controlled_canister_with_large_wasm.rs",
+    ],
+    aliases = ALIASES,
+    data = DEV_DATA,
+    env = DEV_ENV | {"RUST_TEST_NOCAPTURE": "1"},
+    flaky = True,
+    proc_macro_deps = MACRO_DEPENDENCIES + MACRO_DEV_DEPENDENCIES,
+    tags = [
+        "cpu:4",
+    ],
+    deps = [":nervous_system_integration_tests"] + DEPENDENCIES_WITH_TEST_FEATURES + DEV_DEPENDENCIES,
+)

rs/nervous_system/integration_tests/Cargo.toml

@@ -49,6 +49,7 @@ ic-icrc1-index-ng = { path = "../../ledger_suite/icrc1/index-ng" }
 ic-icrc1-tokens-u64 = { path = "../../ledger_suite/icrc1/tokens_u64" }
 ic-management-canister-types = { path = "../../types/management_canister_types" }
 ic-nervous-system-common-test-keys = { path = "../common/test_keys" }
+ic-nervous-system-common-test-utils = { path = "../../nervous_system/common/test_utils" }
 ic-nervous-system-root = { path = "../root" }
 ic-nns-constants = { path = "../../nns/constants" }
 ic-nns-gtc = { path = "../../nns/gtc" }

rs/nervous_system/integration_tests/src/pocket_ic_helpers.rs

@@ -185,6 +185,33 @@ pub async fn install_canister_with_controllers(
     );
 }
 
+pub async fn install_canister_on_subnet(
+    pocket_ic: &PocketIc,
+    subnet_id: Principal,
+    arg: Vec<u8>,
+    wasm: Option<Wasm>,
+    controllers: Vec<PrincipalId>,
+) -> CanisterId {
+    let controllers = controllers.into_iter().map(|c| c.0).collect::<Vec<_>>();
+    let controller_principal = controllers.first().cloned();
+    let settings = Some(CanisterSettings {
+        controllers: Some(controllers),
+        ..Default::default()
+    });
+    let canister_id = pocket_ic
+        .create_canister_on_subnet(None, settings, subnet_id)
+        .await;
+    pocket_ic
+        .add_cycles(canister_id, STARTING_CYCLES_PER_CANISTER)
+        .await;
+    if let Some(wasm) = wasm {
+        pocket_ic
+            .install_canister(canister_id, wasm.bytes(), arg, controller_principal)
+            .await;
+    }
+    CanisterId::unchecked_from_principal(canister_id.into())
+}
+
 // TODO migrate this to nns::governance
 pub async fn add_wasm_via_nns_proposal(
     pocket_ic: &PocketIc,