Merge branch 'abk/run-801-check-freezing-threshold-in-upload-chunk' i…

…nto 'master'

RUN-801: Check freezing threshold on upload_chunk

When uploading to the chunk store, we should check that a canister
doesn't go over the freezing threshold. Otherwise the canister could be
uninstalled with much less warning than expected. 

See merge request dfinity-lab/public/ic!15512

rs/execution_environment/src/canister_manager.rs

@@ -1437,16 +1437,27 @@ impl CanisterManager {
         canister: &mut CanisterState,
         chunk: &[u8],
         subnet_available_memory: &mut SubnetAvailableMemory,
+        subnet_size: usize,
     ) -> Result<UploadChunkReply, CanisterManagerError> {
         if self.config.wasm_chunk_store == FlagStatus::Disabled {
             return Err(CanisterManagerError::WasmChunkStoreError {
                 message: "Wasm chunk store not enabled".to_string(),
             });
         }
+
         validate_controller(canister, &sender)?;
+
+        canister
+            .system_state
+            .wasm_chunk_store
+            .can_insert_chunk(chunk)
+            .map_err(|err| CanisterManagerError::WasmChunkStoreError { message: err })?;
+
+        let new_memory_usage = canister.memory_usage() + wasm_chunk_store::chunk_size();
+
         match canister.memory_allocation() {
             MemoryAllocation::Reserved(bytes) => {
-                if bytes < canister.memory_usage() + wasm_chunk_store::chunk_size() {
+                if bytes < new_memory_usage {
                     return Err(CanisterManagerError::NotEnoughMemoryAllocationGiven {
                         memory_allocation_given: canister.memory_allocation(),
                         memory_usage_needed: canister.memory_usage()
@@ -1455,6 +1466,26 @@ impl CanisterManager {
                 }
             }
             MemoryAllocation::BestEffort => {
+                // Memory usage will increase by the chunk size, so we need to
+                // check that it doesn't bump us over the freezing threshold.
+                let threshold = self.cycles_account_manager.freeze_threshold_cycles(
+                    canister.system_state.freeze_threshold,
+                    canister.memory_allocation(),
+                    new_memory_usage,
+                    canister.message_memory_usage(),
+                    canister.compute_allocation(),
+                    subnet_size,
+                    canister.system_state.reserved_balance(),
+                );
+                if threshold > canister.system_state.balance() {
+                    return Err(CanisterManagerError::WasmChunkStoreError {
+                        message: format!(
+                            "Cannot upload chunk. At least {} additional cycles are required.",
+                            threshold - canister.system_state.balance()
+                        ),
+                    });
+                }
+
                 subnet_available_memory
                     .try_decrement(
                         wasm_chunk_store::chunk_size(),
@@ -1471,11 +1502,14 @@ impl CanisterManager {
                     )?;
             }
         }
+
+        // We initially checked that this chunk can be inserted, so the unwarp
+        // here is guaranteed to succeed.
         let hash = canister
             .system_state
             .wasm_chunk_store
             .insert_chunk(chunk)
-            .map_err(|err| CanisterManagerError::WasmChunkStoreError { message: err })?;
+            .expect("Error: Insert chunk cannot fail after checking `can_insert_chunk`");
         Ok(UploadChunkReply {
             hash: hash.to_vec(),
         })

rs/execution_environment/src/canister_manager/tests.rs

@@ -6936,6 +6936,7 @@ fn upload_chunk_fails_when_allocation_exceeded() {
     };
     let result = test.subnet_message("upload_chunk", upload_args.encode());
     assert!(result.is_ok());
+    let initial_subnet_available_memory = test.subnet_available_memory();
 
     // Second chunk upload fails
     let chunk = vec![4, 4];
@@ -6946,6 +6947,11 @@ fn upload_chunk_fails_when_allocation_exceeded() {
     let result = test.subnet_message("upload_chunk", upload_args.encode());
     let error_code = result.unwrap_err().code();
     assert_eq!(error_code, ErrorCode::InsufficientMemoryAllocation);
+
+    assert_eq!(
+        test.subnet_available_memory(),
+        initial_subnet_available_memory
+    );
 }
 
 #[test]
@@ -6972,6 +6978,7 @@ fn upload_chunk_fails_when_subnet_memory_exceeded() {
     };
     let result = test.subnet_message("upload_chunk", upload_args.encode());
     assert!(result.is_ok());
+    let initial_subnet_available_memory = test.subnet_available_memory();
 
     // Second chunk upload fails
     let chunk = vec![4, 4];
@@ -6982,6 +6989,11 @@ fn upload_chunk_fails_when_subnet_memory_exceeded() {
     let result = test.subnet_message("upload_chunk", upload_args.encode());
     let error_code = result.unwrap_err().code();
     assert_eq!(error_code, ErrorCode::SubnetOversubscribed);
+
+    assert_eq!(
+        test.subnet_available_memory(),
+        initial_subnet_available_memory
+    );
 }
 
 #[test]
@@ -7040,8 +7052,8 @@ fn upload_chunk_fails_with_feature_disabled() {
     const CYCLES: Cycles = Cycles::new(1_000_000_000_000_000);
 
     let mut test = ExecutionTestBuilder::new().build();
-
     let canister_id = test.create_canister(CYCLES);
+    let initial_subnet_available_memory = test.subnet_available_memory();
 
     let chunk = vec![1, 2, 3, 4, 5];
 
@@ -7053,14 +7065,18 @@ fn upload_chunk_fails_with_feature_disabled() {
     let result = test.subnet_message("upload_chunk", upload_args.encode());
     let error_code = result.unwrap_err().code();
     assert_eq!(error_code, ErrorCode::CanisterContractViolation);
+
+    assert_eq!(
+        test.subnet_available_memory(),
+        initial_subnet_available_memory
+    );
 }
 
 #[test]
 fn uninstall_clears_wasm_chunk_store() {
     const CYCLES: Cycles = Cycles::new(1_000_000_000_000_000);
 
     let mut test = ExecutionTestBuilder::new().with_wasm_chunk_store().build();
-
     let canister_id = test.create_canister(CYCLES);
 
     // Upload a chunk
@@ -7086,3 +7102,71 @@ fn uninstall_clears_wasm_chunk_store() {
         NumBytes::from(0)
     );
 }
+
+#[test]
+fn upload_chunk_fails_when_freeze_threshold_triggered() {
+    const CYCLES: Cycles = Cycles::new(1_000_000_000_000_000);
+
+    let mut test = ExecutionTestBuilder::new().with_wasm_chunk_store().build();
+    let canister_id = test.create_canister(CYCLES);
+    let initial_subnet_available_memory = test.subnet_available_memory();
+
+    // Make balance just a bit higher than freezing threshold so upload_chunk
+    // fails.
+    let threshold = test.freezing_threshold(canister_id);
+    let new_balance = threshold + Cycles::from(1_000_u128);
+    let to_remove = test.canister_state(canister_id).system_state.balance() - new_balance;
+    test.canister_state_mut(canister_id)
+        .system_state
+        .remove_cycles(to_remove, CyclesUseCase::BurnedCycles);
+
+    // Upload a chunk
+    let upload_args = UploadChunkArgs {
+        canister_id: canister_id.into(),
+        chunk: vec![42; 10],
+    };
+    let error = test
+        .subnet_message("upload_chunk", upload_args.encode())
+        .unwrap_err();
+
+    assert_eq!(error.code(), ErrorCode::CanisterContractViolation);
+    assert!(error
+        .description()
+        .contains("additional cycles are required"));
+
+    assert_eq!(
+        test.subnet_available_memory(),
+        initial_subnet_available_memory
+    );
+}
+
+#[test]
+fn upload_chunk_fails_when_it_exceeds_chunk_size() {
+    const CYCLES: Cycles = Cycles::new(1_000_000_000_000_000);
+
+    let mut test = ExecutionTestBuilder::new().with_wasm_chunk_store().build();
+    let canister_id = test.create_canister(CYCLES);
+    let initial_subnet_available_memory = test.subnet_available_memory();
+
+    let max_chunk_size = wasm_chunk_store::chunk_size().get() as usize;
+
+    // Upload a chunk that is too large
+    let upload_args = UploadChunkArgs {
+        canister_id: canister_id.into(),
+        chunk: vec![42; max_chunk_size + 1],
+    };
+    let error = test
+        .subnet_message("upload_chunk", upload_args.encode())
+        .unwrap_err();
+
+    assert_eq!(error.code(), ErrorCode::CanisterContractViolation);
+    assert_eq!(
+        error .description(),
+        "Error from Wasm chunk store: Wasm chunk size 1048577 exceeds the maximum chunk size of 1048576"
+    );
+
+    assert_eq!(
+        test.subnet_available_memory(),
+        initial_subnet_available_memory
+    );
+}

rs/execution_environment/src/execution_environment.rs

@@ -1038,6 +1038,7 @@ impl ExecutionEnvironment {
                         &mut state,
                         request,
                         &mut round_limits.subnet_available_memory,
+                        registry_settings.subnet_size,
                     ),
                 };
                 Some((res, msg.take_cycles()))
@@ -1531,10 +1532,17 @@ impl ExecutionEnvironment {
         state: &mut ReplicatedState,
         args: UploadChunkArgs,
         subnet_available_memory: &mut SubnetAvailableMemory,
+        subnet_size: usize,
     ) -> Result<Vec<u8>, UserError> {
         let canister = get_canister_mut(args.get_canister_id(), state)?;
         self.canister_manager
-            .upload_chunk(sender, canister, &args.chunk, subnet_available_memory)
+            .upload_chunk(
+                sender,
+                canister,
+                &args.chunk,
+                subnet_available_memory,
+                subnet_size,
+            )
             .map(|reply| reply.encode())
             .map_err(|err| err.into())
     }

rs/interfaces/src/execution_environment.rs

@@ -159,7 +159,7 @@ impl fmt::Display for ExecutionComplexity {
 /// Note that there are situations where execution available memory is smaller than
 /// the wasm custom sections memory, i.e. when the memory is consumed by something
 /// other than wasm custom sections.
-#[derive(Serialize, Deserialize, Clone, Copy, Debug, Default)]
+#[derive(Serialize, Deserialize, Clone, Copy, Debug, Default, PartialEq)]
 pub struct SubnetAvailableMemory {
     /// The execution memory available on the subnet, i.e. the canister memory
     /// (Wasm binary, Wasm memory, stable memory) without message memory.

rs/replicated_state/src/canister_state/system_state/wasm_chunk_store.rs

@@ -89,12 +89,20 @@ impl WasmChunkStore {
             })
     }
 
-    pub fn insert_chunk(&mut self, chunk: &[u8]) -> Result<[u8; 32], String> {
+    /// Check all conditions for inserting this chunk are satisfied.  Invariant:
+    /// If this returns [`Ok`], then [`Self::insert_chunk`] is guaranteed to
+    /// succeed.
+    pub fn can_insert_chunk(&self, chunk: &[u8]) -> Result<(), String> {
         if chunk.len() > CHUNK_SIZE as usize {
             return Err(
                 format! {"Wasm chunk size {} exceeds the maximum chunk size of {}", chunk.len(), CHUNK_SIZE},
             );
         }
+        Ok(())
+    }
+
+    pub fn insert_chunk(&mut self, chunk: &[u8]) -> Result<[u8; 32], String> {
+        self.can_insert_chunk(chunk)?;
 
         let hash = ic_crypto_sha2::Sha256::hash(chunk);
 
@@ -201,6 +209,7 @@ impl TryFrom<pb::WasmChunkStoreMetadata> for WasmChunkStoreMetadata {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use assert_matches::assert_matches;
 
     fn get_chunk_as_vec(store: &WasmChunkStore, hash: WasmChunkHash) -> Vec<u8> {
         store
@@ -233,14 +242,22 @@ mod tests {
     #[test]
     fn error_when_chunk_exceeds_size_limit() {
         let mut store = WasmChunkStore::new_for_testing();
-        let contents = vec![0xab; 1024 * 1024 + 1];
+        let contents = vec![0xab; chunk_size().get() as usize + 1];
         let result = store.insert_chunk(&contents);
         assert_eq!(
             result,
             Err("Wasm chunk size 1048577 exceeds the maximum chunk size of 1048576".to_string())
         );
     }
 
+    #[test]
+    fn can_insert_chunk_up_to_max_size() {
+        let mut store = WasmChunkStore::new_for_testing();
+        let contents = vec![0xab; chunk_size().get() as usize];
+        let result = store.insert_chunk(&contents);
+        assert_matches!(result, Ok(_));
+    }
+
     #[test]
     fn can_insert_and_retrieve_multiple_chunks() {
         let mut store = WasmChunkStore::new_for_testing();
@@ -255,4 +272,29 @@ mod tests {
         let round_trip_contents2 = get_chunk_as_vec(&store, hash2);
         assert_eq!(contents2, round_trip_contents2);
     }
+
+    mod proptest_tests {
+        use super::*;
+        use proptest::collection::vec as prop_vec;
+        use proptest::prelude::*;
+
+        const MB: usize = 1024 * 1024;
+
+        proptest! {
+            #[test]
+            fn insert_result_matches_can_insert(vecs in prop_vec((any::<u8>(), 0..2 * MB), 100)) {
+                let mut store = WasmChunkStore::new_for_testing();
+                for (byte, length) in vecs {
+                    let chunk = vec![byte; length];
+                    let check = store.can_insert_chunk(&chunk);
+                    let hash = store.insert_chunk(&chunk);
+                    if hash.is_ok() {
+                        assert_eq!(check, Ok(()));
+                    } else {
+                        assert_eq!(check.unwrap_err(), hash.unwrap_err());
+                    }
+                }
+            }
+        }
+    }
 }