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

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

@@ -111,11 +111,21 @@ pub enum SystemUpgradeTargetDTO {
     UpgradeUpgrader,
 }
 
+#[derive(CandidType, serde::Serialize, Deserialize, Debug, Clone)]
+pub struct WasmModuleExtraChunks {
+    pub store_canister: Principal,
+    // TODO: deserialize_with
+    pub chunk_hashes_list: Vec<Vec<u8>>,
+    #[serde(with = "serde_bytes")]
+    pub wasm_module_hash: Vec<u8>,
+}
+
 #[derive(CandidType, serde::Serialize, Deserialize, Debug, Clone)]
 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/system_upgrade.rs

@@ -89,7 +89,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),
@@ -107,6 +111,12 @@ impl Execute for SystemUpgradeRequestExecute<'_, '_> {
             }
 
             SystemUpgradeTarget::UpgradeUpgrader => {
+                if self.operation.input.module_extra_chunks.is_some() {
+                    return Err(RequestExecuteError::Failed {
+                        reason: "Installing upgrader from chunks is not supported.".to_string(),
+                    });
+                }
+
                 self.system_service
                     .upgrade_upgrader(
                         &self.operation.input.module,

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,42 @@ impl From<station_api::SystemUpgradeTargetDTO> for SystemUpgradeTarget {
     }
 }
 
+impl From<WasmModuleExtraChunks> for station_api::WasmModuleExtraChunks {
+    fn from(input: WasmModuleExtraChunks) -> station_api::WasmModuleExtraChunks {
+        station_api::WasmModuleExtraChunks {
+            store_canister: input.store_canister,
+            chunk_hashes_list: input.chunk_hashes_list,
+            wasm_module_hash: input.wasm_module_hash,
+        }
+    }
+}
+
+impl From<WasmModuleExtraChunks> for upgrader_api::WasmModuleExtraChunks {
+    fn from(input: WasmModuleExtraChunks) -> upgrader_api::WasmModuleExtraChunks {
+        upgrader_api::WasmModuleExtraChunks {
+            store_canister: input.store_canister,
+            chunk_hashes_list: input.chunk_hashes_list,
+            wasm_module_hash: input.wasm_module_hash,
+        }
+    }
+}
+
+impl From<station_api::WasmModuleExtraChunks> for WasmModuleExtraChunks {
+    fn from(input: station_api::WasmModuleExtraChunks) -> WasmModuleExtraChunks {
+        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 +391,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/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(),
             },),
         )

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

@@ -2,10 +2,20 @@ use candid::{CandidType, Deserialize, Principal};
 use station_api::TimestampRfc3339;
 pub use station_api::{MetadataDTO, UuidDTO};
 
+#[derive(Clone, Debug, CandidType, serde::Serialize, Deserialize, PartialEq)]
+pub struct WasmModuleExtraChunks {
+    pub store_canister: Principal,
+    // TODO: deserialize_with
+    pub chunk_hashes_list: Vec<Vec<u8>>,
+    #[serde(with = "serde_bytes")]
+    pub wasm_module_hash: Vec<u8>,
+}
+
 #[derive(Clone, Debug, CandidType, serde::Serialize, Deserialize, PartialEq)]
 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),
     };