diff --git a/packages/pocket-ic/CHANGELOG.md b/packages/pocket-ic/CHANGELOG.md index 61ec1048076..ec269d0d139 100644 --- a/packages/pocket-ic/CHANGELOG.md +++ b/packages/pocket-ic/CHANGELOG.md @@ -15,6 +15,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - The function `get_default_effective_canister_id` to retrieve a default effective canister id for canister creation on a PocketIC instance. - The function `PocketIc::get_controllers` to get the controllers of a canister. - Functions `PocketIc::take_canister_snapshot`, `PocketIc::load_canister_snapshot`, `PocketIc::list_canister_snapshots`, and `PocketIc::delete_canister_snapshot` to manage canister snapshots. +- Functions `PocketIc::upload_chunk`, `PocketIc::stored_chunks`, and `PocketIc::clear_chunk_store` to manage the WASM chunk store of a canister. +- The function `PocketIc::install_chunked_canister` to install a canister from WASM chunks in the WASM chunk store of a canister. ### Removed - Functions `PocketIc::from_config`, `PocketIc::from_config_and_max_request_time`, and `PocketIc::from_config_and_server_url`. diff --git a/packages/pocket-ic/src/lib.rs b/packages/pocket-ic/src/lib.rs index a5f4975216f..1e2c224ffbd 100644 --- a/packages/pocket-ic/src/lib.rs +++ b/packages/pocket-ic/src/lib.rs @@ -41,7 +41,7 @@ use crate::{ InstanceId, MockCanisterHttpResponse, RawEffectivePrincipal, RawMessageId, SubnetId, SubnetKind, SubnetSpec, Topology, }, - management_canister::{CanisterId, CanisterStatusResult, Snapshot}, + management_canister::{CanisterId, CanisterInstallMode, CanisterStatusResult, Snapshot}, nonblocking::PocketIc as PocketIcAsync, }; use candid::{ @@ -739,6 +739,73 @@ impl PocketIc { }) } + /// Upload a WASM chunk to the WASM chunk store of a canister. + /// Returns the WASM chunk hash. + #[instrument(skip(self), fields(instance_id=self.pocket_ic.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))] + pub fn upload_chunk( + &self, + canister_id: CanisterId, + sender: Option, + chunk: Vec, + ) -> Result, CallError> { + let runtime = self.runtime.clone(); + runtime.block_on(async { + self.pocket_ic + .upload_chunk(canister_id, sender, chunk) + .await + }) + } + + /// List WASM chunk hashes in the WASM chunk store of a canister. + #[instrument(skip(self), fields(instance_id=self.pocket_ic.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))] + pub fn stored_chunks( + &self, + canister_id: CanisterId, + sender: Option, + ) -> Result>, CallError> { + let runtime = self.runtime.clone(); + runtime.block_on(async { self.pocket_ic.stored_chunks(canister_id, sender).await }) + } + + /// Clear the WASM chunk store of a canister. + #[instrument(skip(self), fields(instance_id=self.pocket_ic.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))] + pub fn clear_chunk_store( + &self, + canister_id: CanisterId, + sender: Option, + ) -> Result<(), CallError> { + let runtime = self.runtime.clone(); + runtime.block_on(async { self.pocket_ic.clear_chunk_store(canister_id, sender).await }) + } + + /// Install a WASM module assembled from chunks on an existing canister. + #[instrument(skip(self, mode, chunk_hashes_list, wasm_module_hash, arg), fields(instance_id=self.pocket_ic.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string(), store_canister_id = %store_canister_id.to_string(), arg_len = %arg.len()))] + pub fn install_chunked_canister( + &self, + canister_id: CanisterId, + sender: Option, + mode: CanisterInstallMode, + store_canister_id: CanisterId, + chunk_hashes_list: Vec>, + wasm_module_hash: Vec, + arg: Vec, + ) -> Result<(), CallError> { + let runtime = self.runtime.clone(); + runtime.block_on(async { + self.pocket_ic + .install_chunked_canister( + canister_id, + sender, + mode, + store_canister_id, + chunk_hashes_list, + wasm_module_hash, + arg, + ) + .await + }) + } + /// Install a WASM module on an existing canister. #[instrument(skip(self, wasm_module, arg), fields(instance_id=self.pocket_ic.instance_id, canister_id = %canister_id.to_string(), wasm_module_len = %wasm_module.len(), arg_len = %arg.len(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))] pub fn install_canister( diff --git a/packages/pocket-ic/src/nonblocking.rs b/packages/pocket-ic/src/nonblocking.rs index f5b2e54e2a6..1aa336bcab3 100644 --- a/packages/pocket-ic/src/nonblocking.rs +++ b/packages/pocket-ic/src/nonblocking.rs @@ -12,7 +12,8 @@ use crate::management_canister::{ CanisterInstallModeUpgradeInnerWasmMemoryPersistenceInner, CanisterSettings, CanisterStatusResult, ChunkHash, DeleteCanisterSnapshotArgs, InstallChunkedCodeArgs, InstallCodeArgs, LoadCanisterSnapshotArgs, ProvisionalCreateCanisterWithCyclesArgs, Snapshot, - TakeCanisterSnapshotArgs, UpdateSettingsArgs, UploadChunkArgs, + StoredChunksResult, TakeCanisterSnapshotArgs, UpdateSettingsArgs, UploadChunkArgs, + UploadChunkResult, }; pub use crate::DefaultEffectiveCanisterIdError; use crate::{CallError, PocketIcBuilder, UserError, WasmResult}; @@ -730,6 +731,98 @@ impl PocketIc { canister_id } + /// Upload a WASM chunk to the WASM chunk store of a canister. + /// Returns the WASM chunk hash. + #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))] + pub async fn upload_chunk( + &self, + canister_id: CanisterId, + sender: Option, + chunk: Vec, + ) -> Result, CallError> { + call_candid_as::<_, (UploadChunkResult,)>( + self, + Principal::management_canister(), + RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()), + sender.unwrap_or(Principal::anonymous()), + "upload_chunk", + (UploadChunkArgs { canister_id, chunk },), + ) + .await + .map(|responses| responses.0.hash) + } + + /// List WASM chunk hashes in the WASM chunk store of a canister. + #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))] + pub async fn stored_chunks( + &self, + canister_id: CanisterId, + sender: Option, + ) -> Result>, CallError> { + call_candid_as::<_, (StoredChunksResult,)>( + self, + Principal::management_canister(), + RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()), + sender.unwrap_or(Principal::anonymous()), + "stored_chunks", + (CanisterIdRecord { canister_id },), + ) + .await + .map(|responses| responses.0.into_iter().map(|chunk| chunk.hash).collect()) + } + + /// Clear the WASM chunk store of a canister. + #[instrument(skip(self), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string()))] + pub async fn clear_chunk_store( + &self, + canister_id: CanisterId, + sender: Option, + ) -> Result<(), CallError> { + call_candid_as( + self, + Principal::management_canister(), + RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()), + sender.unwrap_or(Principal::anonymous()), + "clear_chunk_store", + (CanisterIdRecord { canister_id },), + ) + .await + } + + /// Install a WASM module assembled from chunks on an existing canister. + #[instrument(skip(self, mode, chunk_hashes_list, wasm_module_hash, arg), fields(instance_id=self.instance_id, canister_id = %canister_id.to_string(), sender = %sender.unwrap_or(Principal::anonymous()).to_string(), store_canister_id = %store_canister_id.to_string(), arg_len = %arg.len()))] + pub async fn install_chunked_canister( + &self, + canister_id: CanisterId, + sender: Option, + mode: CanisterInstallMode, + store_canister_id: CanisterId, + chunk_hashes_list: Vec>, + wasm_module_hash: Vec, + arg: Vec, + ) -> Result<(), CallError> { + call_candid_as( + self, + Principal::management_canister(), + RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()), + sender.unwrap_or(Principal::anonymous()), + "install_chunked_code", + (InstallChunkedCodeArgs { + mode, + target_canister: canister_id, + store_canister: Some(store_canister_id), + chunk_hashes_list: chunk_hashes_list + .into_iter() + .map(|hash| ChunkHash { hash }) + .collect(), + wasm_module_hash, + arg, + sender_canister_version: None, + },), + ) + .await + } + async fn install_canister_helper( &self, mode: CanisterInstallMode, @@ -755,60 +848,27 @@ impl PocketIc { ) .await } else { + self.clear_chunk_store(canister_id, sender).await.unwrap(); let chunks: Vec<_> = wasm_module.chunks(INSTALL_CODE_CHUNK_SIZE).collect(); - let hashes: Vec<_> = chunks - .iter() - .map(|c| { - let mut hasher = Sha256::new(); - hasher.update(c); - ChunkHash { - hash: hasher.finalize().to_vec(), - } - }) - .collect(); - call_candid_as::<_, ()>( - self, - Principal::management_canister(), - RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()), - sender.unwrap_or(Principal::anonymous()), - "clear_chunk_store", - (CanisterIdRecord { canister_id },), - ) - .await - .unwrap(); + let mut hashes = vec![]; for chunk in chunks { - call_candid_as::<_, ()>( - self, - Principal::management_canister(), - RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()), - sender.unwrap_or(Principal::anonymous()), - "upload_chunk", - (UploadChunkArgs { - canister_id, - chunk: chunk.to_vec(), - },), - ) - .await - .unwrap(); + let hash = self + .upload_chunk(canister_id, sender, chunk.to_vec()) + .await + .unwrap(); + hashes.push(hash); } let mut hasher = Sha256::new(); hasher.update(wasm_module); let wasm_module_hash = hasher.finalize().to_vec(); - call_candid_as::<_, ()>( - self, - Principal::management_canister(), - RawEffectivePrincipal::CanisterId(canister_id.as_slice().to_vec()), - sender.unwrap_or(Principal::anonymous()), - "install_chunked_code", - (InstallChunkedCodeArgs { - mode, - target_canister: canister_id, - store_canister: None, - chunk_hashes_list: hashes, - wasm_module_hash, - arg, - sender_canister_version: None, - },), + self.install_chunked_canister( + canister_id, + sender, + mode, + canister_id, + hashes, + wasm_module_hash, + arg, ) .await } diff --git a/packages/pocket-ic/tests/tests.rs b/packages/pocket-ic/tests/tests.rs index bc8c612e3bc..671e5931cc1 100644 --- a/packages/pocket-ic/tests/tests.rs +++ b/packages/pocket-ic/tests/tests.rs @@ -1,8 +1,8 @@ use candid::{decode_one, encode_one, CandidType, Decode, Deserialize, Encode, Principal}; use pocket_ic::management_canister::{ - CanisterId, CanisterIdRecord, CanisterSettings, EcdsaPublicKeyResult, HttpRequestResult, - ProvisionalCreateCanisterWithCyclesArgs, SchnorrAlgorithm, SchnorrPublicKeyArgsKeyId, - SchnorrPublicKeyResult, + CanisterId, CanisterIdRecord, CanisterInstallMode, CanisterSettings, EcdsaPublicKeyResult, + HttpRequestResult, ProvisionalCreateCanisterWithCyclesArgs, SchnorrAlgorithm, + SchnorrPublicKeyArgsKeyId, SchnorrPublicKeyResult, }; use pocket_ic::{ common::rest::{ @@ -1699,3 +1699,70 @@ fn test_canister_snapshots() { .unwrap(); pic.take_canister_snapshot(canister_id, None, None).unwrap(); } + +#[test] +fn test_wasm_chunk_store() { + let pic = PocketIc::new(); + + // We create an empty canister and top it up with cycles (WASM chunk store operations cost cycles). + let canister_id = pic.create_canister(); + pic.add_cycles(canister_id, INIT_CYCLES); + + // There should be no chunks in the WASM chunk store yet. + let stored_chunks = pic.stored_chunks(canister_id, None).unwrap(); + assert!(stored_chunks.is_empty()); + + // Chunk the test canister into two chunks. + let mut first_chunk = test_canister_wasm(); + let second_chunk = first_chunk.split_off(first_chunk.len() / 2); + assert!(!first_chunk.is_empty()); + assert!(!second_chunk.is_empty()); + + // We upload a bogus chunk to the WASM chunk store and confirm that the returned hash + // matches the actual hash of the chunk. + let first_chunk_hash = pic + .upload_chunk(canister_id, None, first_chunk.clone()) + .unwrap(); + let mut hasher = Sha256::new(); + hasher.update(first_chunk.clone()); + assert_eq!(first_chunk_hash, hasher.finalize().to_vec()); + + // We upload the same chunk once more and get the same hash back. + let same_chunk_hash = pic + .upload_chunk(canister_id, None, first_chunk.clone()) + .unwrap(); + assert_eq!(first_chunk_hash, same_chunk_hash); + + // We upload a different chunk. + let second_chunk_hash = pic.upload_chunk(canister_id, None, second_chunk).unwrap(); + + // Now the two chunks should be stored in the WASM chunk store. + let stored_chunks = pic.stored_chunks(canister_id, None).unwrap(); + assert_eq!(stored_chunks.len(), 2); + assert!(stored_chunks.contains(&first_chunk_hash)); + assert!(stored_chunks.contains(&second_chunk_hash)); + + // We create a new canister and install it from chunks. + let test_canister = pic.create_canister(); + pic.add_cycles(test_canister, INIT_CYCLES); + let mut hasher = Sha256::new(); + hasher.update(test_canister_wasm()); + let test_canister_wasm_hash = hasher.finalize().to_vec(); + pic.install_chunked_canister( + test_canister, + None, + CanisterInstallMode::Install, + canister_id, + vec![first_chunk_hash, second_chunk_hash], + test_canister_wasm_hash, + Encode!(&()).unwrap(), + ) + .unwrap(); + + // We clear the WASM chunk store. + pic.clear_chunk_store(canister_id, None).unwrap(); + + // There should be no more chunks in the WASM chunk store. + let stored_chunks = pic.stored_chunks(canister_id, None).unwrap(); + assert!(stored_chunks.is_empty()); +}