SMT storage key hashing (FuelLabs#1207)

This PR uses an upcoming release of the `fuel-vm 0.34`.

It brings a new [important
feature](FuelLabs/fuel-vm#485) - hashing the
leaf key for the SMT to prevent tree structure manipulation.

As an example: Hashed storage key decreases the number of nodes in the
SMT from 1.3M to 70K for 30K leaves in the `run_contract_large_state`
e2e test. It improves the test time from 200 seconds to 13 seconds in
Debug mode and to 1 second in the release(instead of 20 seconds).

So it improves security and performance.

The change is breaking because it affects the `state_root` field ->
generated `ContractId`.

Fixes FuelLabs#1143

---------

Co-authored-by: Brandon Vrooman <brandon.vrooman@fuel.sh>

Cargo.toml

@@ -75,7 +75,7 @@ fuel-core-tests = { version = "0.0.0", path = "./tests" }
 fuel-core-xtask = { version = "0.0.0", path = "./xtask" }
 
 # Fuel dependencies
-fuel-vm-private = { version = "0.33", package = "fuel-vm" }
+fuel-vm-private = { version = "0.34", package = "fuel-vm" }
 
 # Common dependencies
 anyhow = "1.0"

bin/e2e-test-client/src/tests/script.rs

@@ -69,17 +69,18 @@ pub async fn dry_run(ctx: &TestContext) -> Result<(), Failed> {
 // Maybe deploy a contract with large state and execute the script
 pub async fn run_contract_large_state(ctx: &TestContext) -> Result<(), Failed> {
     let contract_config = include_bytes!("test_data/large_state/contract.json");
-    let contract_config: ContractConfig =
+    let mut contract_config: ContractConfig =
         serde_json::from_slice(contract_config.as_ref())
             .expect("Should be able do decode the ContractConfig");
     let dry_run = include_bytes!("test_data/large_state/tx.json");
     let dry_run: Transaction = serde_json::from_slice(dry_run.as_ref())
         .expect("Should be able do decode the Transaction");
 
-    // Optimization to run test fastly. If the contract changed, you need to update the
-    // `1bfd51cb31b8d0bc7d93d38f97ab771267d8786ab87073e0c2b8f9ddc44b274e` in the
-    // `test_data/large_state/contract.json`.
-    // contract_config.calculate_contract_id();
+    // If the contract changed, you need to update the
+    // `f4292fe50d21668e140636ab69c7d4b3d069f66eb9ef3da4b0a324409cc36b8c` in the
+    // `test_data/large_state/contract.json` together with:
+    // 244, 41, 47, 229, 13, 33, 102, 142, 20, 6, 54, 171, 105, 199, 212, 179, 208, 105, 246, 110, 185, 239, 61, 164, 176, 163, 36, 64, 156, 195, 107, 140,
+    contract_config.calculate_contract_id();
     let contract_id = contract_config.contract_id;
     println!("\nThe `contract_id` of the contract with large state: {contract_id}");
 

bin/e2e-test-client/src/tests/test_data/large_state/tx.json

@@ -3386,7 +3386,7 @@
       0,
       0,
       1,
-      27, 253, 81, 203, 49, 184, 208, 188, 125, 147, 211, 143, 151, 171, 119, 18, 103, 216, 120, 106, 184, 112, 115, 224, 194, 184, 249, 221, 196, 75, 39, 78,
+      244, 41, 47, 229, 13, 33, 102, 142, 20, 6, 54, 171, 105, 199, 212, 179, 208, 105, 246, 110, 185, 239, 61, 164, 176, 163, 36, 64, 156, 195, 107, 140,
       0,
       0,
       0,
@@ -4113,7 +4113,7 @@
             "block_height": 0,
             "tx_index": 0
           },
-          "contract_id": "1bfd51cb31b8d0bc7d93d38f97ab771267d8786ab87073e0c2b8f9ddc44b274e"
+          "contract_id": "f4292fe50d21668e140636ab69c7d4b3d069f66eb9ef3da4b0a324409cc36b8c"
         }
       },
       {

crates/fuel-core/src/database.rs

@@ -3,6 +3,7 @@ use crate::{
     state::{
         in_memory::memory_store::MemoryStore,
         DataSource,
+        WriteOperation,
     },
 };
 use fuel_core_chain_config::{
@@ -20,6 +21,7 @@ use fuel_core_storage::{
     Result as StorageResult,
 };
 use fuel_core_types::fuel_types::BlockHeight;
+use itertools::Itertools;
 use serde::{
     de::DeserializeOwned,
     Serialize,
@@ -235,6 +237,32 @@ impl Database {
         }
     }
 
+    fn batch_insert<K: AsRef<[u8]>, V: Serialize, S>(
+        &self,
+        column: Column,
+        set: S,
+    ) -> DatabaseResult<()>
+    where
+        S: Iterator<Item = (K, V)>,
+    {
+        let set: Vec<_> = set
+            .map(|(key, value)| {
+                let value =
+                    postcard::to_stdvec(&value).map_err(|_| DatabaseError::Codec)?;
+
+                let tuple = (
+                    key.as_ref().to_vec(),
+                    column,
+                    WriteOperation::Insert(Arc::new(value)),
+                );
+
+                Ok::<_, DatabaseError>(tuple)
+            })
+            .try_collect()?;
+
+        self.data.batch_write(&mut set.into_iter())
+    }
+
     fn remove<V: DeserializeOwned>(
         &self,
         key: &[u8],

crates/fuel-core/src/database/balances.rs

@@ -2,13 +2,15 @@ use crate::database::{
     storage::{
         ContractsAssetsMerkleData,
         ContractsAssetsMerkleMetadata,
+        DatabaseColumn,
         SparseMerkleMetadata,
     },
     Column,
     Database,
 };
 use fuel_core_storage::{
     tables::ContractsAssets,
+    ContractsAssetKey,
     Error as StorageError,
     Mappable,
     MerkleRoot,
@@ -19,15 +21,21 @@ use fuel_core_storage::{
     StorageMutate,
 };
 use fuel_core_types::{
+    fuel_asm::Word,
     fuel_merkle::{
         sparse,
         sparse::{
             in_memory,
             MerkleTree,
+            MerkleTreeKey,
         },
     },
-    fuel_types::ContractId,
+    fuel_types::{
+        AssetId,
+        ContractId,
+    },
 };
+use itertools::Itertools;
 use std::{
     borrow::{
         BorrowMut,
@@ -77,9 +85,10 @@ impl StorageMutate<ContractsAssets> for Database {
             MerkleTree::load(storage, &root)
                 .map_err(|err| StorageError::Other(err.into()))?;
 
+        let asset_id = *key.asset_id().deref();
         // Update the contact's key-value dataset. The key is the asset id and the
         // value the Word
-        tree.update(key.asset_id().deref(), value.to_be_bytes().as_slice())
+        tree.update(MerkleTreeKey::new(asset_id), value.to_be_bytes().as_slice())
             .map_err(|err| StorageError::Other(err.into()))?;
 
         // Generate new metadata for the updated tree
@@ -112,9 +121,10 @@ impl StorageMutate<ContractsAssets> for Database {
                 MerkleTree::load(storage, &root)
                     .map_err(|err| StorageError::Other(err.into()))?;
 
+            let asset_id = *key.asset_id().deref();
             // Update the contract's key-value dataset. The key is the asset id and
             // the value is the Word
-            tree.delete(key.asset_id().deref())
+            tree.delete(MerkleTreeKey::new(asset_id))
                 .map_err(|err| StorageError::Other(err.into()))?;
 
             let root = tree.root();
@@ -146,6 +156,53 @@ impl MerkleRootStorage<ContractId, ContractsAssets> for Database {
     }
 }
 
+impl Database {
+    /// Initialize the balances of the contract from the all leafs.
+    /// This method is more performant than inserting balances one by one.
+    pub fn init_contract_balances<S>(
+        &mut self,
+        contract_id: &ContractId,
+        balances: S,
+    ) -> Result<(), StorageError>
+    where
+        S: Iterator<Item = (AssetId, Word)>,
+    {
+        if self
+            .storage::<ContractsAssetsMerkleMetadata>()
+            .contains_key(contract_id)?
+        {
+            return Err(
+                anyhow::anyhow!("The contract balances is already initialized").into(),
+            )
+        }
+
+        let balances = balances.collect_vec();
+
+        // Keys and values should be original without any modifications.
+        // Key is `ContractId` ++ `AssetId`
+        self.batch_insert(
+            Column::ContractsAssets,
+            balances.clone().into_iter().map(|(asset, value)| {
+                (ContractsAssetKey::new(contract_id, &asset), value)
+            }),
+        )?;
+
+        // Merkle data:
+        // - Asset key should be converted into `MerkleTreeKey` by `new` function that hashes them.
+        // - The balance value are original.
+        let balances = balances
+            .into_iter()
+            .map(|(asset, value)| (MerkleTreeKey::new(asset), value.to_be_bytes()));
+        let (root, nodes) = in_memory::MerkleTree::nodes_from_set(balances);
+        self.batch_insert(ContractsAssetsMerkleData::column(), nodes.into_iter())?;
+        let metadata = SparseMerkleMetadata { root };
+        self.storage::<ContractsAssetsMerkleMetadata>()
+            .insert(contract_id, &metadata)?;
+
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -157,6 +214,16 @@ mod tests {
         AssetId,
         Word,
     };
+    use rand::Rng;
+
+    fn random_asset_id<R>(rng: &mut R) -> AssetId
+    where
+        R: Rng + ?Sized,
+    {
+        let mut bytes = [0u8; 32];
+        rng.fill(bytes.as_mut());
+        bytes.into()
+    }
 
     #[test]
     fn get() {
@@ -371,6 +438,61 @@ mod tests {
         assert_eq!(root_0, root_2);
     }
 
+    #[test]
+    fn init_contract_balances_works() {
+        use rand::{
+            rngs::StdRng,
+            RngCore,
+            SeedableRng,
+        };
+
+        let rng = &mut StdRng::seed_from_u64(1234);
+        let gen = || Some((random_asset_id(rng), rng.next_u64()));
+        let data = core::iter::from_fn(gen).take(5_000).collect::<Vec<_>>();
+
+        let contract_id = ContractId::from([1u8; 32]);
+        let init_database = &mut Database::default();
+
+        init_database
+            .init_contract_balances(&contract_id, data.clone().into_iter())
+            .expect("Should init contract");
+        let init_root = init_database
+            .storage::<ContractsAssets>()
+            .root(&contract_id)
+            .expect("Should get root");
+
+        let seq_database = &mut Database::default();
+        for (asset, value) in data.iter() {
+            seq_database
+                .storage::<ContractsAssets>()
+                .insert(&ContractsAssetKey::new(&contract_id, asset), value)
+                .expect("Should insert a state");
+        }
+        let seq_root = seq_database
+            .storage::<ContractsAssets>()
+            .root(&contract_id)
+            .expect("Should get root");
+
+        assert_eq!(init_root, seq_root);
+
+        for (asset, value) in data.into_iter() {
+            let init_value = init_database
+                .storage::<ContractsAssets>()
+                .get(&ContractsAssetKey::new(&contract_id, &asset))
+                .expect("Should get a state from init database")
+                .unwrap()
+                .into_owned();
+            let seq_value = seq_database
+                .storage::<ContractsAssets>()
+                .get(&ContractsAssetKey::new(&contract_id, &asset))
+                .expect("Should get a state from seq database")
+                .unwrap()
+                .into_owned();
+            assert_eq!(init_value, value);
+            assert_eq!(seq_value, value);
+        }
+    }
+
     #[test]
     fn remove_deletes_merkle_metadata_when_empty() {
         let contract_id = ContractId::from([1u8; 32]);