From ac1492403d4c9dad8c6ad0911ab12a525338a28f Mon Sep 17 00:00:00 2001
From: Leo-Besancon <lb@massa.net>
Date: Fri, 2 Aug 2024 17:27:21 +0200
Subject: [PATCH] Fix wrong LedgerChanges::Delete handling (#4737)

* Fix bad LedgerChanges::Delete handling

* Avoid too much code duplication

* Update ledger_db.rs

* Fix check
---
 massa-ledger-worker/Cargo.toml       |  4 +--
 massa-ledger-worker/src/ledger_db.rs | 41 +++++++++++++++++++---------
 2 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/massa-ledger-worker/Cargo.toml b/massa-ledger-worker/Cargo.toml
index ab37769ceb..c755a97bc6 100644
--- a/massa-ledger-worker/Cargo.toml
+++ b/massa-ledger-worker/Cargo.toml
@@ -5,7 +5,7 @@ authors = ["Massa Labs <info@massa.net>"]
 edition = "2021"
 
 [features]
-test-exports = ["tempfile", "massa_models/test-exports", "massa_ledger_exports/test-exports", "massa_db_worker", "parking_lot"]
+test-exports = ["tempfile", "massa_models/test-exports", "massa_ledger_exports/test-exports", "massa_db_worker"]
 
 [dependencies]
 serde_json = {workspace = true}   # BOM UPGRADE     Revert to "1.0" if problem
@@ -15,7 +15,7 @@ massa_models = {workspace = true}
 massa_serialization = {workspace = true}
 massa_db_exports = {workspace = true}
 massa_db_worker = {workspace = true, "optional" = true}
-parking_lot = {workspace = true, "features" = ["deadlock_detection"], "optional" = true}
+parking_lot = {workspace = true, "features" = ["deadlock_detection"]}
 
 [dev-dependencies]
 massa_signature = {workspace = true}
diff --git a/massa-ledger-worker/src/ledger_db.rs b/massa-ledger-worker/src/ledger_db.rs
index a95a319961..8dca6c1c74 100644
--- a/massa-ledger-worker/src/ledger_db.rs
+++ b/massa-ledger-worker/src/ledger_db.rs
@@ -3,8 +3,8 @@
 //! Module to interact with the disk ledger
 
 use massa_db_exports::{
-    DBBatch, MassaDirection, MassaIteratorMode, ShareableMassaDBController, CRUD_ERROR,
-    KEY_SER_ERROR, LEDGER_PREFIX, STATE_CF,
+    DBBatch, MassaDBController, MassaDirection, MassaIteratorMode, ShareableMassaDBController,
+    CRUD_ERROR, KEY_SER_ERROR, LEDGER_PREFIX, STATE_CF,
 };
 use massa_ledger_exports::*;
 use massa_models::amount::AmountDeserializer;
@@ -16,6 +16,7 @@ use massa_models::{
 use massa_serialization::{
     DeserializeError, Deserializer, Serializer, U64VarIntDeserializer, U64VarIntSerializer,
 };
+use parking_lot::{lock_api::RwLockReadGuard, RawRwLock};
 use std::collections::{BTreeSet, HashMap};
 use std::fmt::Debug;
 
@@ -286,6 +287,9 @@ impl LedgerDB {
     fn put_entry(&self, addr: &Address, ledger_entry: LedgerEntry, batch: &mut DBBatch) {
         let db = self.db.read();
 
+        // Ensures any potential previous entry is fully deleted.
+        delete_datastore_entries(addr, &db, batch);
+
         // Version
         //TODO: Get version number from parameters
         let mut bytes_version = Vec::new();
@@ -456,18 +460,29 @@ impl LedgerDB {
             .expect(KEY_SER_ERROR);
         db.delete_key(batch, serialized_key);
 
-        // datastore
-        let key_prefix = datastore_prefix_from_address(addr, &[]);
+        delete_datastore_entries(addr, &db, batch);
+    }
+}
 
-        for (serialized_key, _) in db
-            .iterator_cf(
-                STATE_CF,
-                MassaIteratorMode::From(&key_prefix, MassaDirection::Forward),
-            )
-            .take_while(|(key, _)| key < &end_prefix(&key_prefix).unwrap())
-        {
-            db.delete_key(batch, serialized_key.to_vec());
-        }
+// Helper function to delete all datastore entries for a given address
+// Needs to be called in put entry and delete entry
+// Note: This function takes a lock on the DB to avoid multiple reads.
+fn delete_datastore_entries(
+    addr: &Address,
+    db: &RwLockReadGuard<RawRwLock, Box<dyn MassaDBController>>,
+    batch: &mut std::collections::BTreeMap<Vec<u8>, Option<Vec<u8>>>,
+) {
+    // datastore
+    let key_prefix = datastore_prefix_from_address(addr, &[]);
+
+    for (serialized_key, _) in db
+        .iterator_cf(
+            STATE_CF,
+            MassaIteratorMode::From(&key_prefix, MassaDirection::Forward),
+        )
+        .take_while(|(key, _)| key < &end_prefix(&key_prefix).unwrap())
+    {
+        db.delete_key(batch, serialized_key.to_vec());
     }
 }