feat(station,upgrader): support large station and upgrader wasm

apps/wallet/src/components/system-upgrade/SystemUpgradeActionBtn.vue

@@ -160,6 +160,8 @@ const submitUpgrade = async (model: SystemUpgradeFormProps['modelValue']): Promi
           ? [new Uint8Array(hexStringToArrayBuffer(model.wasmInitArg))]
           : [],
       module: new Uint8Array(fileBuffer),
+      // TODO: Add support for extra chunks once the control-panel supports it
+      module_extra_chunks: [],
       target: assertAndReturn(model.target, 'model.target is required'),
     },
     {

apps/wallet/src/generated/station/station.did

@@ -555,11 +555,22 @@ type SystemUpgradeTarget = variant {
   UpgradeUpgrader;
 };
 
+type WasmModuleExtraChunks = record {
+  // The asset canister from which the chunks are to be retrieved.
+  store_canister : principal;
+  // The list of chunk hashes in the order they should be appended to the wasm module.
+  chunk_hashes_list : vec blob;
+  // The hash of the assembled wasm module.
+  wasm_module_hash : blob;
+};
+
 type SystemUpgradeOperationInput = record {
   // The target to change.
   target : SystemUpgradeTarget;
   // The wasm module to install.
   module : blob;
+  // Additional wasm module chunks to append to the wasm module.
+  module_extra_chunks : opt WasmModuleExtraChunks;
   // The initial argument passed to the new wasm module.
   arg : opt blob;
 };

apps/wallet/src/generated/station/station.did.d.ts

@@ -1063,6 +1063,7 @@ export interface SystemUpgradeOperation {
 }
 export interface SystemUpgradeOperationInput {
   'arg' : [] | [Uint8Array | number[]],
+  'module_extra_chunks' : [] | [WasmModuleExtraChunks],
   'target' : SystemUpgradeTarget,
   'module' : Uint8Array | number[],
 }
@@ -1165,6 +1166,11 @@ export type UserStatus = { 'Inactive' : null } |
   { 'Active' : null };
 export type ValidationMethodResourceTarget = { 'No' : null } |
   { 'ValidationMethod' : CanisterMethod };
+export interface WasmModuleExtraChunks {
+  'wasm_module_hash' : Uint8Array | number[],
+  'chunk_hashes_list' : Array<Uint8Array | number[]>,
+  'store_canister' : Principal,
+}
 export interface _SERVICE {
   'canister_status' : ActorMethod<[CanisterStatusInput], CanisterStatusResult>,
   'capabilities' : ActorMethod<[], CapabilitiesResult>,

apps/wallet/src/generated/station/station.did.js

@@ -336,12 +336,18 @@ export const idlFactory = ({ IDL }) => {
     'policy_id' : UUID,
   });
   const RemoveRequestPolicyOperationInput = IDL.Record({ 'policy_id' : UUID });
+  const WasmModuleExtraChunks = IDL.Record({
+    'wasm_module_hash' : IDL.Vec(IDL.Nat8),
+    'chunk_hashes_list' : IDL.Vec(IDL.Vec(IDL.Nat8)),
+    'store_canister' : IDL.Principal,
+  });
   const SystemUpgradeTarget = IDL.Variant({
     'UpgradeUpgrader' : IDL.Null,
     'UpgradeStation' : IDL.Null,
   });
   const SystemUpgradeOperationInput = IDL.Record({
     'arg' : IDL.Opt(IDL.Vec(IDL.Nat8)),
+    'module_extra_chunks' : IDL.Opt(WasmModuleExtraChunks),
     'target' : SystemUpgradeTarget,
     'module' : IDL.Vec(IDL.Nat8),
   });

core/station/api/spec.did

@@ -555,11 +555,22 @@ type SystemUpgradeTarget = variant {
   UpgradeUpgrader;
 };
 
+type WasmModuleExtraChunks = record {
+  // The asset canister from which the chunks are to be retrieved.
+  store_canister : principal;
+  // The list of chunk hashes in the order they should be appended to the wasm module.
+  chunk_hashes_list : vec blob;
+  // The hash of the assembled wasm module.
+  wasm_module_hash : blob;
+};
+
 type SystemUpgradeOperationInput = record {
   // The target to change.
   target : SystemUpgradeTarget;
   // The wasm module to install.
   module : blob;
+  // Additional wasm module chunks to append to the wasm module.
+  module_extra_chunks : opt WasmModuleExtraChunks;
   // The initial argument passed to the new wasm module.
   arg : opt blob;
 };

core/station/api/src/system.rs

@@ -1,6 +1,7 @@
 use super::TimestampRfc3339;
 use crate::{DisasterRecoveryCommitteeDTO, MetadataDTO, Sha256HashDTO, UuidDTO};
 use candid::{CandidType, Deserialize, Principal};
+use orbit_essentials::types::WasmModuleExtraChunks;
 
 #[derive(CandidType, serde::Serialize, Deserialize, Clone, Debug)]
 pub struct SystemInfoDTO {
@@ -116,6 +117,7 @@ pub struct SystemUpgradeOperationInput {
     pub target: SystemUpgradeTargetDTO,
     #[serde(with = "serde_bytes")]
     pub module: Vec<u8>,
+    pub module_extra_chunks: Option<WasmModuleExtraChunks>,
     #[serde(deserialize_with = "orbit_essentials::deserialize::deserialize_option_blob")]
     pub arg: Option<Vec<u8>>,
 }

core/station/impl/src/factories/requests/change_external_canister.rs

@@ -80,6 +80,7 @@ impl Execute for ChangeExternalCanisterRequestExecute<'_, '_> {
                 self.operation.input.canister_id,
                 self.operation.input.mode.clone(),
                 &self.operation.input.module,
+                &None,
                 self.operation.input.arg.clone(),
             )
             .await

core/station/impl/src/factories/requests/system_upgrade.rs

@@ -36,9 +36,13 @@ impl Create<SystemUpgradeOperationInput> for SystemUpgradeRequestCreate {
                     hasher.finalize().to_vec()
                 }),
                 module_checksum: {
-                    let mut hasher = Sha256::new();
-                    hasher.update(&operation_input.module);
-                    hasher.finalize().to_vec()
+                    if let Some(ref module_extra_chunks) = operation_input.module_extra_chunks {
+                        module_extra_chunks.wasm_module_hash.clone()
+                    } else {
+                        let mut hasher = Sha256::new();
+                        hasher.update(&operation_input.module);
+                        hasher.finalize().to_vec()
+                    }
                 },
                 input: operation_input.into(),
             }),
@@ -89,7 +93,11 @@ impl Execute for SystemUpgradeRequestExecute<'_, '_> {
                 let arg = self.operation.input.arg.as_ref().unwrap_or(&default_arg);
                 let out = self
                     .system_service
-                    .upgrade_station(&self.operation.input.module, arg)
+                    .upgrade_station(
+                        &self.operation.input.module,
+                        &self.operation.input.module_extra_chunks,
+                        arg,
+                    )
                     .await
                     .map_err(|err| RequestExecuteError::Failed {
                         reason: format!("failed to upgrade station: {}", err),
@@ -110,11 +118,12 @@ impl Execute for SystemUpgradeRequestExecute<'_, '_> {
                 self.system_service
                     .upgrade_upgrader(
                         &self.operation.input.module,
+                        &self.operation.input.module_extra_chunks,
                         self.operation.input.arg.clone(),
                     )
                     .await
                     .map_err(|err| RequestExecuteError::Failed {
-                        reason: format!("failed to upgrade upgrader: {}", err),
+                        reason: format!("failed to upgrade upgrader: {} ({:?})", err, err.details),
                     })?;
 
                 // The upgrader might have just gained the ability to perform disaster recovery, so sync it now.

core/station/impl/src/mappers/request_operation.rs

@@ -29,7 +29,7 @@ use crate::{
         RemoveRequestPolicyOperation, RemoveRequestPolicyOperationInput, RemoveUserGroupOperation,
         RequestOperation, SetDisasterRecoveryOperation, SetDisasterRecoveryOperationInput,
         SystemUpgradeOperation, SystemUpgradeOperationInput, SystemUpgradeTarget,
-        TransferOperation, User,
+        TransferOperation, User, WasmModuleExtraChunks,
     },
     repositories::{
         AccountRepository, AddressBookRepository, UserRepository, ACCOUNT_REPOSITORY,
@@ -345,11 +345,32 @@ impl From<station_api::SystemUpgradeTargetDTO> for SystemUpgradeTarget {
     }
 }
 
+impl From<orbit_essentials::types::WasmModuleExtraChunks> for WasmModuleExtraChunks {
+    fn from(input: orbit_essentials::types::WasmModuleExtraChunks) -> WasmModuleExtraChunks {
+        WasmModuleExtraChunks {
+            store_canister: input.store_canister,
+            chunk_hashes_list: input.chunk_hashes_list,
+            wasm_module_hash: input.wasm_module_hash,
+        }
+    }
+}
+
+impl From<WasmModuleExtraChunks> for orbit_essentials::types::WasmModuleExtraChunks {
+    fn from(input: WasmModuleExtraChunks) -> orbit_essentials::types::WasmModuleExtraChunks {
+        orbit_essentials::types::WasmModuleExtraChunks {
+            store_canister: input.store_canister,
+            chunk_hashes_list: input.chunk_hashes_list,
+            wasm_module_hash: input.wasm_module_hash,
+        }
+    }
+}
+
 impl From<SystemUpgradeOperationInput> for station_api::SystemUpgradeOperationInput {
     fn from(input: SystemUpgradeOperationInput) -> station_api::SystemUpgradeOperationInput {
         station_api::SystemUpgradeOperationInput {
             target: input.target.into(),
             module: input.module,
+            module_extra_chunks: input.module_extra_chunks.map(|c| c.into()),
             arg: input.arg,
         }
     }
@@ -360,6 +381,7 @@ impl From<station_api::SystemUpgradeOperationInput> for SystemUpgradeOperationIn
         SystemUpgradeOperationInput {
             target: input.target.into(),
             module: input.module,
+            module_extra_chunks: input.module_extra_chunks.map(|c| c.into()),
             arg: input.arg,
         }
     }

core/station/impl/src/models/request_operation.rs

@@ -261,12 +261,21 @@ pub enum SystemUpgradeTarget {
     UpgradeUpgrader,
 }
 
+#[storable]
+#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
+pub struct WasmModuleExtraChunks {
+    pub store_canister: Principal,
+    pub chunk_hashes_list: Vec<Vec<u8>>,
+    pub wasm_module_hash: Vec<u8>,
+}
+
 #[storable]
 #[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct SystemUpgradeOperationInput {
     pub target: SystemUpgradeTarget,
     /// The module is only available while the operation is not finalized.
     pub module: Vec<u8>,
+    pub module_extra_chunks: Option<WasmModuleExtraChunks>,
     pub arg: Option<Vec<u8>>,
 }
 

core/station/impl/src/services/change_canister.rs

@@ -1,11 +1,12 @@
-use crate::{errors::ChangeCanisterError, models::CanisterInstallMode};
-use candid::Principal;
-use ic_cdk::api::management_canister::{
-    main::{self as mgmt, InstallCodeArgument},
-    provisional::CanisterIdRecord,
+use crate::{
+    errors::ChangeCanisterError,
+    models::{CanisterInstallMode, WasmModuleExtraChunks},
 };
+use candid::Principal;
+use ic_cdk::api::management_canister::{main as mgmt, provisional::CanisterIdRecord};
 use lazy_static::lazy_static;
 use orbit_essentials::api::ServiceResult;
+use orbit_essentials::install_chunked_code::install_chunked_code;
 use std::sync::Arc;
 
 lazy_static! {
@@ -27,6 +28,7 @@ impl ChangeCanisterService {
         canister_id: Principal,
         mode: CanisterInstallMode,
         module: &[u8],
+        module_extra_chunks: &Option<WasmModuleExtraChunks>,
         arg: Option<Vec<u8>>,
     ) -> ServiceResult<(), ChangeCanisterError> {
         use candid::Encode;
@@ -55,16 +57,15 @@ impl ChangeCanisterService {
 
         // Install or upgrade canister
         let default_bytes = Encode!(&()).unwrap();
-        let install_code_result = mgmt::install_code(InstallCodeArgument {
-            mode: mode.into(),
-            canister_id: canister_id.to_owned(),
-            wasm_module: module.to_owned(),
-            arg: arg.unwrap_or(default_bytes),
-        })
+        let install_code_result = install_chunked_code(
+            canister_id,
+            mode.into(),
+            module.to_owned(),
+            module_extra_chunks.as_ref().map(|c| c.clone().into()),
+            arg.unwrap_or(default_bytes),
+        )
         .await
-        .map_err(|(_, err)| ChangeCanisterError::Failed {
-            reason: err.to_string(),
-        });
+        .map_err(|err| ChangeCanisterError::Failed { reason: err });
 
         // Restart canister (regardless of whether the upgrade succeeded or not)
         mgmt::start_canister(CanisterIdRecord {

core/station/impl/src/services/system.rs

@@ -13,7 +13,7 @@ use crate::{
         system::{DisasterRecoveryCommittee, SystemInfo, SystemState},
         CanisterInstallMode, CanisterUpgradeModeArgs, CycleObtainStrategy,
         ManageSystemInfoOperationInput, RequestId, RequestKey, RequestOperation, RequestStatus,
-        SystemUpgradeTarget,
+        SystemUpgradeTarget, WasmModuleExtraChunks,
     },
     repositories::{
         permission::PERMISSION_REPOSITORY, RequestRepository, REQUEST_REPOSITORY,
@@ -137,14 +137,20 @@ impl SystemService {
     }
 
     /// Execute an upgrade of the station by requesting the upgrader to perform it on our behalf.
-    pub async fn upgrade_station(&self, module: &[u8], arg: &[u8]) -> ServiceResult<()> {
+    pub async fn upgrade_station(
+        &self,
+        module: &[u8],
+        module_extra_chunks: &Option<WasmModuleExtraChunks>,
+        arg: &[u8],
+    ) -> ServiceResult<()> {
         let upgrader_canister_id = self.get_upgrader_canister_id();
 
         ic_cdk::call(
             upgrader_canister_id,
             "trigger_upgrade",
             (UpgradeParams {
                 module: module.to_owned(),
+                module_extra_chunks: module_extra_chunks.clone().map(|c| c.into()),
                 arg: arg.to_owned(),
             },),
         )
@@ -157,13 +163,19 @@ impl SystemService {
     }
 
     /// Execute an upgrade of the upgrader canister.
-    pub async fn upgrade_upgrader(&self, module: &[u8], arg: Option<Vec<u8>>) -> ServiceResult<()> {
+    pub async fn upgrade_upgrader(
+        &self,
+        module: &[u8],
+        module_extra_chunks: &Option<WasmModuleExtraChunks>,
+        arg: Option<Vec<u8>>,
+    ) -> ServiceResult<()> {
         let upgrader_canister_id = self.get_upgrader_canister_id();
         self.change_canister_service
             .install_canister(
                 upgrader_canister_id,
                 CanisterInstallMode::Upgrade(CanisterUpgradeModeArgs {}),
                 module,
+                module_extra_chunks,
                 arg,
             )
             .await

core/upgrader/api/spec.did

@@ -12,9 +12,20 @@ type InitArg = record {
   target_canister : principal;
 };
 
+type WasmModuleExtraChunks = record {
+  // The asset canister from which the chunks are to be retrieved.
+  store_canister : principal;
+  // The list of chunk hashes in the order they should be appended to the wasm module.
+  chunk_hashes_list : vec blob;
+  // The hash of the assembled wasm module.
+  wasm_module_hash : blob;
+};
+
 type UpgradeParams = record {
   // The wasm module to upgrade to.
   module : blob;
+  // Additional wasm module chunks to append to the wasm module.
+  module_extra_chunks : opt WasmModuleExtraChunks;
   // The argument to be passed to upgrade.
   arg : blob;
 };

core/upgrader/api/src/lib.rs

@@ -1,11 +1,13 @@
 use candid::{CandidType, Deserialize, Principal};
+use orbit_essentials::types::WasmModuleExtraChunks;
 use station_api::TimestampRfc3339;
 pub use station_api::{MetadataDTO, UuidDTO};
 
-#[derive(Clone, Debug, CandidType, serde::Serialize, Deserialize, PartialEq)]
+#[derive(Clone, Debug, CandidType, serde::Serialize, Deserialize)]
 pub struct UpgradeParams {
     #[serde(with = "serde_bytes")]
     pub module: Vec<u8>,
+    pub module_extra_chunks: Option<WasmModuleExtraChunks>,
     #[serde(with = "serde_bytes")]
     pub arg: Vec<u8>,
 }

core/upgrader/impl/Cargo.toml

@@ -24,6 +24,7 @@ ic-stable-structures = { workspace = true }
 lazy_static = { workspace = true }
 mockall = { workspace = true }
 serde = { workspace = true, features = ['derive'] }
+serde_bytes = { workspace = true }
 serde_cbor = { workspace = true }
 serde_json = { workspace = true }
 station-api = { path = '../../station/api', version = '0.0.2-alpha.4' }

core/upgrader/impl/src/lib.rs

@@ -78,6 +78,7 @@ lazy_static! {
 async fn trigger_upgrade(params: upgrader_api::UpgradeParams) -> Result<(), TriggerUpgradeError> {
     let input: UpgradeParams = UpgradeParams {
         module: params.module,
+        module_extra_chunks: params.module_extra_chunks,
         arg: params.arg,
         install_mode: CanisterInstallMode::Upgrade(None),
     };