Merge branch 'abk/run-738-clear-chunk-store' into 'master'

RUN-738: Implement clear_chunk_store

Add the canister manager method to clear the Wasm chunk store for a
specific canister. The only conditions on this method are that it has to
be called by a controller of the canister. 

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

rs/execution_environment/src/canister_manager.rs

@@ -343,10 +343,11 @@ impl CanisterManager {
     ) -> Result<(), UserError> {
         let method_name = ingress.method_name();
         let sender = ingress.sender();
+        let method = Ic00Method::from_str(ingress.method_name());
         // The message is targeted towards the management canister. The
         // actual type of the method will determine if the message should be
         // accepted or not.
-        match Ic00Method::from_str(ingress.method_name()) {
+        match method {
             // The method is either invalid or it is of a type that users
             // are not allowed to send.
             Err(_)
@@ -383,7 +384,27 @@ impl CanisterManager {
             | Ok(Ic00Method::DeleteCanister)
             | Ok(Ic00Method::UpdateSettings)
             | Ok(Ic00Method::InstallCode)
-            | Ok(Ic00Method::InstallChunkedCode) => {
+            | Ok(Ic00Method::InstallChunkedCode)
+            | Ok(Ic00Method::UploadChunk)
+            | Ok(Ic00Method::StoredChunks)
+            | Ok(Ic00Method::DeleteChunks)
+            | Ok(Ic00Method::ClearChunkStore) => {
+                // Reject large install methods if the flag is not enabled, or
+                // they are not implemented.
+                match method {
+                    Ok(Ic00Method::UploadChunk)
+                    | Ok(Ic00Method::ClearChunkStore)
+                    | Ok(Ic00Method::InstallChunkedCode) if self.config.wasm_chunk_store == FlagStatus::Enabled => {}
+                    Ok(Ic00Method::UploadChunk)
+                    | Ok(Ic00Method::StoredChunks)
+                    | Ok(Ic00Method::DeleteChunks)
+                    | Ok(Ic00Method::ClearChunkStore)
+                    | Ok(Ic00Method::InstallChunkedCode) => return Err(UserError::new(
+                        ErrorCode::CanisterRejectedMessage,
+                        "Chunked upload API is not yet implemented"
+                    )),
+                    _ => {}
+                };
                 match effective_canister_id {
                     Some(canister_id) => {
                         let canister = state.canister_state(&canister_id).ok_or_else(|| UserError::new(
@@ -420,16 +441,6 @@ impl CanisterManager {
                     ))
                 }
             },
-            Ok(Ic00Method::UploadChunk)  if self.config.wasm_chunk_store == FlagStatus::Enabled => {
-                Ok(())
-            }
-            Ok(Ic00Method::UploadChunk) |
-            Ok(Ic00Method::StoredChunks) |
-            Ok(Ic00Method::DeleteChunks) |
-            Ok(Ic00Method::ClearChunkStore) => Err(UserError::new(
-                ErrorCode::CanisterRejectedMessage,
-                "Chunked upload API is not yet implemented"
-            )),
         }
     }
 
@@ -1514,6 +1525,22 @@ impl CanisterManager {
             hash: hash.to_vec(),
         })
     }
+
+    pub(crate) fn clear_chunk_store(
+        &self,
+        sender: PrincipalId,
+        canister: &mut CanisterState,
+    ) -> Result<(), 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 = WasmChunkStore::new(Arc::clone(&self.fd_factory));
+        Ok(())
+    }
 }
 
 #[derive(Debug, PartialEq, Eq)]

rs/execution_environment/src/canister_manager/tests.rs

@@ -23,8 +23,9 @@ use ic_error_types::{ErrorCode, UserError};
 use ic_ic00_types::{
     CanisterChange, CanisterChangeDetails, CanisterChangeOrigin, CanisterIdRecord,
     CanisterInstallMode, CanisterInstallModeV2, CanisterSettingsArgsBuilder,
-    CanisterStatusResultV2, CanisterStatusType, CreateCanisterArgs, EmptyBlob, InstallCodeArgsV2,
-    Method, Payload, SkipPreUpgrade, UpdateSettingsArgs, UploadChunkArgs, UploadChunkReply,
+    CanisterStatusResultV2, CanisterStatusType, ClearChunkStoreArgs, CreateCanisterArgs, EmptyBlob,
+    InstallCodeArgsV2, Method, Payload, SkipPreUpgrade, UpdateSettingsArgs, UploadChunkArgs,
+    UploadChunkReply,
 };
 use ic_interfaces::execution_environment::{
     ExecutionComplexity, ExecutionMode, HypervisorError, SubnetAvailableMemory,
@@ -7170,3 +7171,92 @@ fn upload_chunk_fails_when_it_exceeds_chunk_size() {
         initial_subnet_available_memory
     );
 }
+
+#[test]
+fn clear_chunk_store_works() {
+    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 chunk = vec![1, 2, 3, 4, 5];
+    let hash = ic_crypto_sha2::Sha256::hash(&chunk);
+    let initial_memory_usage = test.canister_state(canister_id).memory_usage();
+
+    // After uploading the chunk, it is present in the store and the memory
+    // usage is positive.
+    let upload_args = UploadChunkArgs {
+        canister_id: canister_id.into(),
+        chunk,
+    };
+    test.subnet_message("upload_chunk", upload_args.encode())
+        .unwrap();
+    assert!(test.canister_state(canister_id).memory_usage() > initial_memory_usage);
+    assert!(test
+        .canister_state(canister_id)
+        .system_state
+        .wasm_chunk_store
+        .get_chunk_data(&hash)
+        .is_some());
+
+    // After clearing, the chunk should be absent and memory usage should be
+    // zero.
+    let clear_args = ClearChunkStoreArgs {
+        canister_id: canister_id.into(),
+    };
+    test.subnet_message("clear_chunk_store", clear_args.encode())
+        .unwrap();
+    assert_eq!(
+        test.canister_state(canister_id).memory_usage(),
+        initial_memory_usage
+    );
+    assert!(test
+        .canister_state(canister_id)
+        .system_state
+        .wasm_chunk_store
+        .get_chunk_data(&hash)
+        .is_none());
+}
+
+#[test]
+fn clear_chunk_store_fails_from_non_controller() {
+    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 uc = test
+        .canister_from_cycles_and_binary(CYCLES, UNIVERSAL_CANISTER_WASM.into())
+        .unwrap();
+
+    let args = ClearChunkStoreArgs {
+        canister_id: canister_id.into(),
+    };
+
+    let clear_store = wasm()
+        .call_with_cycles(
+            CanisterId::ic_00(),
+            Method::ClearChunkStore,
+            call_args()
+                .other_side(args.encode())
+                .on_reject(wasm().reject_message().reject()),
+            Cycles::new(CYCLES.get() / 2),
+        )
+        .build();
+
+    let result = test.ingress(uc, "update", clear_store);
+    let expected_err = format!(
+        "Only the controllers of the canister {} can control it.",
+        canister_id
+    );
+    match result {
+        Ok(WasmResult::Reject(reject)) => {
+            assert!(
+                reject.contains(&expected_err),
+                "Reject \"{reject}\" does not contain expected error \"{expected_err}\""
+            );
+        }
+        other => panic!("Expected reject, but got {:?}", other),
+    }
+}

rs/execution_environment/src/execution_environment.rs

@@ -31,7 +31,7 @@ use ic_cycles_account_manager::{
 use ic_error_types::{ErrorCode, RejectCode, UserError};
 use ic_ic00_types::{
     CanisterChangeOrigin, CanisterHttpRequestArgs, CanisterIdRecord, CanisterInfoRequest,
-    CanisterInfoResponse, CanisterSettingsArgs, CanisterStatusType,
+    CanisterInfoResponse, CanisterSettingsArgs, CanisterStatusType, ClearChunkStoreArgs,
     ComputeInitialEcdsaDealingsArgs, CreateCanisterArgs, ECDSAPublicKeyArgs,
     ECDSAPublicKeyResponse, EcdsaKeyId, EmptyBlob, InstallChunkedCodeArgs, InstallCodeArgsV2,
     Method as Ic00Method, Payload as Ic00Payload, ProvisionalCreateCanisterWithCyclesArgs,
@@ -1044,9 +1044,16 @@ impl ExecutionEnvironment {
                 Some((res, msg.take_cycles()))
             }
 
+            Ok(Ic00Method::ClearChunkStore) => {
+                let res = match ClearChunkStoreArgs::decode(payload) {
+                    Err(err) => Err(err),
+                    Ok(request) => self.clear_chunk_store(*msg.sender(), &mut state, request),
+                };
+                Some((res, msg.take_cycles()))
+            }
+
             Ok(Ic00Method::StoredChunks)
             | Ok(Ic00Method::DeleteChunks)
-            | Ok(Ic00Method::ClearChunkStore)
             | Ok(Ic00Method::InstallChunkedCode) => Some((
                 Err(UserError::new(
                     ErrorCode::CanisterRejectedMessage,
@@ -1547,6 +1554,19 @@ impl ExecutionEnvironment {
             .map_err(|err| err.into())
     }
 
+    fn clear_chunk_store(
+        &self,
+        sender: PrincipalId,
+        state: &mut ReplicatedState,
+        args: ClearChunkStoreArgs,
+    ) -> Result<Vec<u8>, UserError> {
+        let canister = get_canister_mut(args.get_canister_id(), state)?;
+        self.canister_manager
+            .clear_chunk_store(sender, canister)
+            .map(|()| EmptyBlob.encode())
+            .map_err(|err| err.into())
+    }
+
     // Executes an inter-canister response.
     //
     // Returns a tuple with the result, along with a flag indicating whether or

rs/system_api/src/routing.rs

@@ -6,7 +6,7 @@ use ic_btc_interface::NetworkInRequest as BitcoinNetwork;
 use ic_error_types::UserError;
 use ic_ic00_types::{
     BitcoinGetBalanceArgs, BitcoinGetCurrentFeePercentilesArgs, BitcoinGetUtxosArgs,
-    BitcoinSendTransactionArgs, CanisterIdRecord, CanisterInfoRequest,
+    BitcoinSendTransactionArgs, CanisterIdRecord, CanisterInfoRequest, ClearChunkStoreArgs,
     ComputeInitialEcdsaDealingsArgs, ECDSAPublicKeyArgs, EcdsaKeyId, InstallChunkedCodeArgs,
     InstallCodeArgsV2, Method as Ic00Method, Payload, ProvisionalTopUpCanisterArgs,
     SignWithECDSAArgs, UninstallCodeArgs, UpdateSettingsArgs, UploadChunkArgs,
@@ -211,9 +211,21 @@ pub(super) fn resolve_destination(
                     ResolveDestinationError::SubnetNotFound(canister_id, Ic00Method::UploadChunk)
                 })
         }
-        Ok(Ic00Method::StoredChunks)
-        | Ok(Ic00Method::DeleteChunks)
-        | Ok(Ic00Method::ClearChunkStore) => {
+        Ok(Ic00Method::ClearChunkStore) => {
+            let args = ClearChunkStoreArgs::decode(payload)?;
+            let canister_id = args.get_canister_id();
+            network_topology
+                .routing_table
+                .route(canister_id.get())
+                .map(|subnet_id| subnet_id.get())
+                .ok_or({
+                    ResolveDestinationError::SubnetNotFound(
+                        canister_id,
+                        Ic00Method::ClearChunkStore,
+                    )
+                })
+        }
+        Ok(Ic00Method::StoredChunks) | Ok(Ic00Method::DeleteChunks) => {
             Err(ResolveDestinationError::UserError(UserError::new(
                 ic_error_types::ErrorCode::CanisterRejectedMessage,
                 "Chunked upload API is not yet implemented",

rs/types/ic00_types/src/lib.rs

@@ -2276,3 +2276,20 @@ impl InstallChunkedCodeArgs {
             .map(|p| CanisterId::unchecked_from_principal(p))
     }
 }
+
+/// Struct used for encoding/decoding
+/// `(record {
+///     canister_id: principal;
+/// })`
+#[derive(Default, Clone, CandidType, Deserialize, Debug)]
+pub struct ClearChunkStoreArgs {
+    pub canister_id: PrincipalId,
+}
+
+impl Payload<'_> for ClearChunkStoreArgs {}
+
+impl ClearChunkStoreArgs {
+    pub fn get_canister_id(&self) -> CanisterId {
+        CanisterId::unchecked_from_principal(self.canister_id)
+    }
+}

rs/types/types/src/messages/ingress_messages.rs

@@ -11,8 +11,8 @@ use crate::{
 };
 use ic_error_types::{ErrorCode, UserError};
 use ic_ic00_types::{
-    CanisterIdRecord, CanisterInfoRequest, InstallChunkedCodeArgs, InstallCodeArgsV2, Method,
-    Payload, UpdateSettingsArgs, UploadChunkArgs, IC_00,
+    CanisterIdRecord, CanisterInfoRequest, ClearChunkStoreArgs, InstallChunkedCodeArgs,
+    InstallCodeArgsV2, Method, Payload, UpdateSettingsArgs, UploadChunkArgs, IC_00,
 };
 use ic_protobuf::{
     log::ingress_message_log_entry::v1::IngressMessageLogEntry,
@@ -486,7 +486,11 @@ pub fn extract_effective_canister_id(
             Ok(record) => Ok(Some(record.get_canister_id())),
             Err(err) => Err(ParseIngressError::InvalidSubnetPayload(err.to_string())),
         },
-        Ok(Method::StoredChunks) | Ok(Method::DeleteChunks) | Ok(Method::ClearChunkStore) => {
+        Ok(Method::ClearChunkStore) => match ClearChunkStoreArgs::decode(ingress.arg()) {
+            Ok(record) => Ok(Some(record.get_canister_id())),
+            Err(err) => Err(ParseIngressError::InvalidSubnetPayload(err.to_string())),
+        },
+        Ok(Method::StoredChunks) | Ok(Method::DeleteChunks) => {
             Err(ParseIngressError::UnknownSubnetMethod)
         }
         Ok(Method::CreateCanister)

rs/types/types/src/messages/inter_canister.rs

@@ -3,8 +3,9 @@ use ic_error_types::{RejectCode, TryFromError, UserError};
 #[cfg(test)]
 use ic_exhaustive_derive::ExhaustiveSet;
 use ic_ic00_types::{
-    CanisterIdRecord, CanisterInfoRequest, InstallChunkedCodeArgs, InstallCodeArgsV2, Method,
-    Payload as _, ProvisionalTopUpCanisterArgs, UpdateSettingsArgs, UploadChunkArgs,
+    CanisterIdRecord, CanisterInfoRequest, ClearChunkStoreArgs, InstallChunkedCodeArgs,
+    InstallCodeArgsV2, Method, Payload as _, ProvisionalTopUpCanisterArgs, UpdateSettingsArgs,
+    UploadChunkArgs,
 };
 use ic_protobuf::{
     proxy::{try_from_option_field, ProxyDecodeError},
@@ -132,9 +133,13 @@ impl Request {
                 Ok(record) => Some(record.get_canister_id()),
                 Err(_) => None,
             },
-            Ok(Method::StoredChunks) | Ok(Method::DeleteChunks) | Ok(Method::ClearChunkStore) => {
-                None
+            Ok(Method::ClearChunkStore) => {
+                match ClearChunkStoreArgs::decode(&self.method_payload) {
+                    Ok(record) => Some(record.get_canister_id()),
+                    Err(_) => None,
+                }
             }
+            Ok(Method::StoredChunks) | Ok(Method::DeleteChunks) => None,
             Ok(Method::CreateCanister)
             | Ok(Method::SetupInitialDKG)
             | Ok(Method::HttpRequest)