RUN-875: Charge for chunked install on hash mismatch

rs/execution_environment/src/canister_manager.rs

@@ -49,7 +49,7 @@ use ic_types::{
     InvalidMemoryAllocationError, MemoryAllocation, NumBytes, NumInstructions, PrincipalId,
     SubnetId, Time,
 };
-use ic_wasm_types::CanisterModule;
+use ic_wasm_types::{CanisterModule, WasmHash};
 use num_traits::cast::ToPrimitive;
 use prometheus::IntCounter;
 use serde::{Deserialize, Serialize};
@@ -153,12 +153,109 @@ impl CanisterMgrConfig {
     }
 }
 
+#[derive(Clone, Debug)]
+pub enum WasmSource {
+    CanisterModule(CanisterModule),
+    ChunkStore {
+        wasm_chunk_store: WasmChunkStore,
+        chunk_hashes_list: Vec<Vec<u8>>,
+        wasm_module_hash: WasmHash,
+    },
+}
+
+impl From<&WasmSource> for WasmHash {
+    fn from(item: &WasmSource) -> Self {
+        match item {
+            WasmSource::CanisterModule(canister_module) => {
+                Self::from(canister_module.module_hash())
+            }
+            WasmSource::ChunkStore {
+                wasm_module_hash, ..
+            } => wasm_module_hash.clone(),
+        }
+    }
+}
+
+impl WasmSource {
+    pub fn module_hash(&self) -> [u8; 32] {
+        WasmHash::from(self).to_slice()
+    }
+
+    /// The number of instructions to be charged each time we try to convert to
+    /// a canister module.
+    pub fn instructions_to_assemble(&self) -> NumInstructions {
+        match self {
+            Self::CanisterModule(_module) => NumInstructions::from(0),
+            // Charge one instruction per byte, assuming each chunk is the
+            // maximum size.
+            Self::ChunkStore {
+                chunk_hashes_list, ..
+            } => NumInstructions::from(
+                (wasm_chunk_store::chunk_size() * chunk_hashes_list.len() as u64).get(),
+            ),
+        }
+    }
+
+    /// Convert the source to a canister module (assembling chunks if required).
+    pub(crate) fn into_canister_module(self) -> Result<CanisterModule, CanisterManagerError> {
+        match self {
+            Self::CanisterModule(module) => Ok(module),
+            Self::ChunkStore {
+                wasm_chunk_store,
+                chunk_hashes_list,
+                wasm_module_hash,
+            } => {
+                // Assume each chunk uses the full chunk size even though the actual
+                // size might be smaller.
+                let mut wasm_module = Vec::with_capacity(
+                    chunk_hashes_list.len() * wasm_chunk_store::chunk_size().get() as usize,
+                );
+                for hash in chunk_hashes_list {
+                    let hash = hash.as_slice().try_into().map_err(|_| {
+                        CanisterManagerError::WasmChunkStoreError {
+                            message: "Chunk hash is invalid. The length is not 32".to_string(),
+                        }
+                    })?;
+                    for page in wasm_chunk_store.get_chunk_data(&hash).ok_or_else(|| {
+                        CanisterManagerError::WasmChunkStoreError {
+                            message: format!("Chunk hash {:?} was not found", &hash[..32]),
+                        }
+                    })? {
+                        wasm_module.extend_from_slice(page)
+                    }
+                }
+                let canister_module = CanisterModule::new(wasm_module);
+
+                if canister_module.module_hash()[..] != wasm_module_hash.to_slice() {
+                    return Err(CanisterManagerError::WasmChunkStoreError {
+                        message: format!(
+                            "Wasm module hash {:?} does not match given hash {:?}",
+                            canister_module.module_hash(),
+                            wasm_module_hash
+                        ),
+                    });
+                }
+                Ok(canister_module)
+            }
+        }
+    }
+
+    #[allow(dead_code)]
+    /// Only used for tests.
+    fn unwrap_as_slice_for_testing(&self) -> &[u8] {
+        match self {
+            Self::CanisterModule(module) => module.as_slice(),
+            Self::ChunkStore { .. } => panic!("Can't convert WasmSource::ChunkStore to slice"),
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct InstallCodeContext {
     pub origin: CanisterChangeOrigin,
     pub mode: CanisterInstallModeV2,
     pub canister_id: CanisterId,
-    pub wasm_module: CanisterModule,
+    pub wasm_source: WasmSource,
     pub arg: Vec<u8>,
     pub compute_allocation: Option<ComputeAllocation>,
     pub memory_allocation: Option<MemoryAllocation>,
@@ -232,38 +329,22 @@ impl InstallCodeContext {
         store: &WasmChunkStore,
     ) -> Result<Self, InstallCodeContextError> {
         let canister_id = args.target_canister_id();
-        // Assume each chunk uses the full chunk size even though the actual
-        // size might be smaller.
-        let mut wasm_module = Vec::with_capacity(
-            args.chunk_hashes_list.len() * wasm_chunk_store::chunk_size().get() as usize,
-        );
-        for hash in args.chunk_hashes_list {
-            let hash = hash.as_slice().try_into().map_err(|_| {
-                InstallCodeContextError::InvalidHash(
-                    "Chunk hash is invalid. The length is not 32".to_string(),
-                )
-            })?;
-            for page in store.get_chunk_data(&hash).ok_or_else(|| {
-                InstallCodeContextError::InvalidHash(format!(
-                    "Chunk hash {:?} was not found",
-                    &hash[..32]
-                ))
-            })? {
-                wasm_module.extend_from_slice(page)
-            }
-        }
-        let hash = ic_crypto_sha2::Sha256::hash(&wasm_module);
-        if hash[..] != args.wasm_module_hash {
-            return Err(InstallCodeContextError::InvalidHash(format!(
-                "Wasm module hash {:?} does not match given hash {:?}",
-                hash, args.wasm_module_hash
-            )));
-        }
+        let wasm_module_hash = args.wasm_module_hash.try_into().map_err(|hash| {
+            InstallCodeContextError::InvalidHash(format!("Invalid wasm hash {:?}", hash))
+        })?;
         Ok(InstallCodeContext {
             origin,
             mode: args.mode,
             canister_id,
-            wasm_module: CanisterModule::new(wasm_module),
+            wasm_source: WasmSource::ChunkStore {
+                wasm_chunk_store: store.clone(),
+                chunk_hashes_list: args
+                    .chunk_hashes_list
+                    .into_iter()
+                    .map(|h| h.to_vec())
+                    .collect(),
+                wasm_module_hash,
+            },
             arg: args.arg,
             compute_allocation: None,
             memory_allocation: None,
@@ -300,7 +381,7 @@ impl TryFrom<(CanisterChangeOrigin, InstallCodeArgsV2)> for InstallCodeContext {
             origin,
             mode: args.mode,
             canister_id,
-            wasm_module: CanisterModule::new(args.wasm_module),
+            wasm_source: WasmSource::CanisterModule(CanisterModule::new(args.wasm_module)),
             arg: args.arg,
             compute_allocation,
             memory_allocation,