From accadf9e1f50ebaa1f9cbb6e7cb33ac1a6b9b887 Mon Sep 17 00:00:00 2001
From: David Dal Busco <david.dalbusco@outlook.com>
Date: Sun, 19 Nov 2023 15:23:53 +0100
Subject: [PATCH] fix: delete asset stable (#329)

---
 src/satellite/src/db/store.rs                   |  3 ++-
 src/satellite/src/lib.rs                        |  4 ++--
 .../src/storage/certification/impls.rs          |  4 +++-
 src/satellite/src/storage/state.rs              | 17 +++++++++++++++++
 4 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/satellite/src/db/store.rs b/src/satellite/src/db/store.rs
index 1934be25f..8e336c439 100644
--- a/src/satellite/src/db/store.rs
+++ b/src/satellite/src/db/store.rs
@@ -294,7 +294,8 @@ pub fn delete_docs(collection: &CollectionKey) -> Result<(), String> {
 fn delete_docs_impl(collection: &CollectionKey) -> Result<(), String> {
     let rule = get_state_rule(collection)?;
 
-    let keys: Vec<Key> = get_state_docs(collection, &rule)?.iter()
+    let keys: Vec<Key> = get_state_docs(collection, &rule)?
+        .iter()
         .map(|(key, _)| key.clone())
         .collect();
 
diff --git a/src/satellite/src/lib.rs b/src/satellite/src/lib.rs
index ba755fb89..13bef7c4f 100644
--- a/src/satellite/src/lib.rs
+++ b/src/satellite/src/lib.rs
@@ -11,7 +11,7 @@ mod storage;
 mod types;
 
 use crate::controllers::store::get_admin_controllers;
-use crate::db::store::{delete_doc, get_doc as get_doc_store, get_docs, insert_doc, delete_docs};
+use crate::db::store::{delete_doc, delete_docs, get_doc as get_doc_store, get_docs, insert_doc};
 use crate::db::types::interface::{DelDoc, SetDoc};
 use crate::db::types::state::Doc;
 use crate::guards::caller_is_admin_controller;
@@ -31,7 +31,7 @@ use crate::storage::store::{
     commit_batch, create_batch, create_chunk, delete_asset, delete_assets, delete_domain,
     get_config as get_storage_config, get_content_chunks, get_custom_domains, get_public_asset,
     init_certified_assets, list_assets as list_assets_store, set_config as set_storage_config,
-    set_domain
+    set_domain,
 };
 use crate::storage::types::config::StorageConfigRewrites;
 use crate::storage::types::domain::{CustomDomains, DomainName};
diff --git a/src/satellite/src/storage/certification/impls.rs b/src/satellite/src/storage/certification/impls.rs
index 89e0cd69d..ca0e92049 100644
--- a/src/satellite/src/storage/certification/impls.rs
+++ b/src/satellite/src/storage/certification/impls.rs
@@ -15,7 +15,9 @@ use crate::storage::types::config::StorageConfig;
 use crate::storage::types::state::FullPath;
 use crate::storage::types::store::Asset;
 use crate::storage::url::alternative_paths;
-use ic_certification::{fork, fork_hash, labeled, labeled_hash, merge_hash_trees, AsHashTree, Hash, HashTree, pruned};
+use ic_certification::{
+    fork, fork_hash, labeled, labeled_hash, merge_hash_trees, pruned, AsHashTree, Hash, HashTree,
+};
 use sha2::{Digest, Sha256};
 
 impl CertifiedAssetHashes {
diff --git a/src/satellite/src/storage/state.rs b/src/satellite/src/storage/state.rs
index 8ecc459c6..f7d48c1b1 100644
--- a/src/satellite/src/storage/state.rs
+++ b/src/satellite/src/storage/state.rs
@@ -9,6 +9,7 @@ use crate::storage::types::state::{
 };
 use crate::storage::types::store::{Asset, AssetEncoding};
 use crate::types::core::{Blob, CollectionKey};
+use crate::types::state::StableState;
 use shared::serializers::{deserialize_from_bytes, serialize_to_bytes};
 use std::borrow::Cow;
 
@@ -132,6 +133,7 @@ pub fn delete_asset(
             delete_asset_heap(full_path, &mut state.borrow_mut().heap.storage.assets)
         }),
         Memory::Stable => STATE.with(|state| {
+            delete_content_chunks_stable(collection, full_path, &mut state.borrow_mut().stable);
             delete_asset_stable(collection, full_path, &mut state.borrow_mut().stable.assets)
         }),
     }
@@ -183,6 +185,21 @@ fn delete_asset_stable(
     assets.remove(&stable_full_path(collection, full_path))
 }
 
+fn delete_content_chunks_stable(
+    collection: &CollectionKey,
+    full_path: &FullPath,
+    state: &mut StableState,
+) {
+    if let Some(asset) = get_asset_stable(collection, full_path, &state.assets) {
+        for (_, encoding) in asset.encodings.iter() {
+            for chunk in encoding.content_chunks.iter() {
+                let key: StableEncodingChunkKey = deserialize_from_bytes(Cow::Owned(chunk.clone()));
+                state.content_chunks.remove(&key);
+            }
+        }
+    }
+}
+
 fn delete_asset_heap(full_path: &FullPath, assets: &mut AssetsHeap) -> Option<Asset> {
     assets.remove(full_path)
 }