Repository: stellar/rs-soroban-env
Pull request #986 (2023-08-08 20:01:41)

Make `has` checks to properly populate the storage map in recording mode
### What

Make `has` checks to properly populate the storage map in recording mode

### Why

The enforcing mode invariant is that every footprint entry is present in the map (with `None` for non-existent entries). Recording mode should have the same invariant - every entry in the map should be a footprint entry, but `has` check bypassed it.

I've refactored the code to use unified logic. This also will help us to reduce the metering inconsistencies between recording and enforcing modes in the future.

### Known limitations

N/A



Patch:
diff --git a/soroban-env-host/src/host.rs b/soroban-env-host/src/host.rs
index cf86f60c..edbc6c17 100644
--- a/soroban-env-host/src/host.rs
+++ b/soroban-env-host/src/host.rs
@@ -62,6 +62,8 @@ use self::{
     metered_xdr::metered_write_xdr,
 };
 use crate::impl_bignum_host_fns;
+#[cfg(any(test, feature = "testutils"))]
+use crate::storage::{AccessType, Footprint};
 use crate::Compare;
 #[cfg(any(test, feature = "testutils"))]
 use crate::TryIntoVal;
@@ -653,14 +655,48 @@ impl Host {
     }
 
     // Writes an arbitrary ledger entry to storage.
-    // "testutils" is not covered by budget metering.
+    // "testutils" are not covered by budget metering.
     #[cfg(any(test, feature = "testutils"))]
     pub fn add_ledger_entry(
         &self,
         key: &Rc<LedgerKey>,
         val: &Rc<soroban_env_common::xdr::LedgerEntry>,
     ) -> Result<(), HostError> {
-        self.with_mut_storage(|storage| storage.put(key, val, self.as_budget()))
+        self.as_budget().with_free_budget(|| {
+            self.with_mut_storage(|storage| storage.put(key, val, self.as_budget()))
+        })
+    }
+
+    // Performs the necessary setup to access the provided ledger key/entry in
+    // enforcing storage mode.
+    // "testutils" are not covered by budget metering.
+    #[cfg(any(test, feature = "testutils"))]
+    pub fn setup_storage_entry(
+        &self,
+        key: Rc<LedgerKey>,
+        val: Option<Rc<soroban_env_common::xdr::LedgerEntry>>,
+        access_type: AccessType,
+    ) -> Result<(), HostError> {
+        self.as_budget().with_free_budget(|| {
+            self.with_mut_storage(|storage| {
+                storage
+                    .footprint
+                    .record_access(&key, access_type, self.as_budget())?;
+                storage.map = storage.map.insert(key, val, self.as_budget())?;
+                Ok(())
+            })
+        })
+    }
+
+    // Performs the necessary setup to access all the entries in provided
+    // footprint in enforcing mode.
+    // "testutils" are not covered by budget metering.
+    #[cfg(any(test, feature = "testutils"))]
+    pub fn setup_storage_footprint(&self, footprint: Footprint) -> Result<(), HostError> {
+        for (key, access_type) in footprint.0.map {
+            self.setup_storage_entry(key, None, access_type)?;
+        }
+        Ok(())
     }
 
     // Returns the authorizations that have been authenticated for the last
@@ -2091,14 +2127,14 @@ impl VmCallerEnv for Host {
                         ContractDataEntryBody::DataEntry(data) => Ok(self.to_host_val(&data.val)?),
                         _ => Err(self.err(
                             ScErrorType::Storage,
-                            ScErrorCode::UnexpectedType,
+                            ScErrorCode::InternalError,
                             "expected DataEntry",
                             &[],
                         )),
                     },
                     _ => Err(self.err(
                         ScErrorType::Storage,
-                        ScErrorCode::UnexpectedType,
+                        ScErrorCode::InternalError,
                         "expected contract data ledger entry",
                         &[],
                     )),
diff --git a/soroban-env-host/src/storage.rs b/soroban-env-host/src/storage.rs
index 68073ccf..5f7e5f68 100644
--- a/soroban-env-host/src/storage.rs
+++ b/soroban-env-host/src/storage.rs
@@ -205,22 +205,7 @@ impl Storage {
         budget: &Budget,
     ) -> Result<Rc<LedgerEntry>, HostError> {
         let _span = tracy_span!("storage get");
-        let ty = AccessType::ReadOnly;
-        match self.mode {
-            FootprintMode::Recording(ref src) => {
-                self.footprint.record_access(key, ty, budget)?;
-                // In recording mode we treat the map as a cache
-                // that misses read-through to the underlying src.
-                if !self.map.contains_key::<Rc<LedgerKey>>(key, budget)? {
-                    self.map = self
-                        .map
-                        .insert(Rc::clone(key), Some(src.get(key)?), budget)?;
-                }
-            }
-            FootprintMode::Enforcing => {
-                self.footprint.enforce_access(key, ty, budget)?;
-            }
-        };
+        self.prepare_read_only_access(key, budget)?;
         match self.map.get::<Rc<LedgerKey>>(key, budget)? {
             None | Some(None) => Err((ScErrorType::Storage, ScErrorCode::MissingValue).into()),
             Some(Some(val)) => Ok(Rc::clone(val)),
@@ -292,27 +277,14 @@ impl Storage {
     /// declared in the [Footprint].
     pub fn has(&mut self, key: &Rc<LedgerKey>, budget: &Budget) -> Result<bool, HostError> {
         let _span = tracy_span!("storage has");
-        let ty = AccessType::ReadOnly;
-        match self.mode {
-            FootprintMode::Recording(ref src) => {
-                self.footprint.record_access(key, ty, budget)?;
-                // We don't cache has() calls but we do
-                // consult the cache before answering them.
-                match self.map.get::<Rc<LedgerKey>>(key, budget)? {
-                    Some(None) => Ok(false),
-                    Some(Some(_)) => Ok(true),
-                    None => src.has(key),
-                }
-            }
-            FootprintMode::Enforcing => {
-                self.footprint.enforce_access(key, ty, budget)?;
-                match self.map.get::<Rc<LedgerKey>>(key, budget)? {
-                    Some(None) => Ok(false),
-                    Some(Some(_)) => Ok(true),
-                    None => Ok(false),
-                }
-            }
-        }
+        self.prepare_read_only_access(key, budget)?;
+        Ok(self
+            .map
+            .get::<Rc<LedgerKey>>(key, budget)?
+            // Key has to be present in storage at this point, so not having it
+            // would be an internal error.
+            .ok_or_else(|| HostError::from((ScErrorType::Storage, ScErrorCode::InternalError)))?
+            .is_some())
     }
 
     /// Bumps `key` to live for at least `bump_by_ledgers` from now (not
@@ -415,6 +387,33 @@ impl Storage {
         }
         Ok(())
     }
+
+    fn prepare_read_only_access(
+        &mut self,
+        key: &Rc<LedgerKey>,
+        budget: &Budget,
+    ) -> Result<(), HostError> {
+        let ty = AccessType::ReadOnly;
+        match self.mode {
+            FootprintMode::Recording(ref src) => {
+                self.footprint.record_access(key, ty, budget)?;
+                // In recording mode we treat the map as a cache
+                // that misses read-through to the underlying src.
+                if !self.map.contains_key::<Rc<LedgerKey>>(key, budget)? {
+                    let value = if src.has(&key)? {
+                        Some(src.get(key)?)
+                    } else {
+                        None
+                    };
+                    self.map = self.map.insert(Rc::clone(key), value, budget)?;
+                }
+            }
+            FootprintMode::Enforcing => {
+                self.footprint.enforce_access(key, ty, budget)?;
+            }
+        };
+        Ok(())
+    }
 }
 
 #[cfg(test)]


Test patch:
diff --git a/soroban-env-host/src/test/complex.rs b/soroban-env-host/src/test/complex.rs
index 951ee983..3025e6fa 100644
--- a/soroban-env-host/src/test/complex.rs
+++ b/soroban-env-host/src/test/complex.rs
@@ -1,5 +1,8 @@
 use crate::{
-    budget::Budget, host_object::HostVec, storage::Storage, Host, HostError, MeteredOrdMap,
+    budget::Budget,
+    host_object::HostVec,
+    storage::{Footprint, Storage},
+    Host, HostError, MeteredOrdMap,
 };
 use soroban_env_common::{Env, Symbol};
 use soroban_test_wasms::COMPLEX;
@@ -39,9 +42,14 @@ fn run_complex() -> Result<(), HostError> {
 
     // Run 2: enforce preflight footprint, with empty map -- contract should only write.
     {
-        let store = Storage::with_enforcing_footprint_and_map(foot, MeteredOrdMap::default());
+        let store = Storage::with_enforcing_footprint_and_map(
+            Footprint::default(),
+            MeteredOrdMap::default(),
+        );
         let host = Host::with_storage_and_budget(store, Budget::default());
         host.set_ledger_info(info)?;
+        host.setup_storage_footprint(foot)?;
+
         let contract_id_obj =
             host.register_test_contract_wasm_from_source_account(COMPLEX, account_id, salt);
         host.call(
diff --git a/soroban-env-host/src/test/lifecycle.rs b/soroban-env-host/src/test/lifecycle.rs
index 417dcbfe..2b28a0e5 100644
--- a/soroban-env-host/src/test/lifecycle.rs
+++ b/soroban-env-host/src/test/lifecycle.rs
@@ -103,23 +103,17 @@ fn test_create_contract_from_source_account(host: &Host, wasm: &[u8]) -> Hash {
 
     let wasm_hash = Hash(Sha256::digest(wasm).try_into().unwrap());
 
-    host.with_mut_storage(|s: &mut Storage| {
-        s.footprint
-            .record_access(
-                &host.contract_instance_ledger_key(&contract_id).unwrap(),
-                AccessType::ReadWrite,
-                host.as_budget(),
-            )
-            .unwrap();
-        s.footprint
-            .record_access(
-                &host.contract_code_ledger_key(&wasm_hash).unwrap(),
-                AccessType::ReadWrite,
-                host.as_budget(),
-            )
-            .unwrap();
-        Ok(())
-    })
+    host.setup_storage_entry(
+        host.contract_instance_ledger_key(&contract_id).unwrap(),
+        None,
+        AccessType::ReadWrite,
+    )
+    .unwrap();
+    host.setup_storage_entry(
+        host.contract_code_ledger_key(&wasm_hash).unwrap(),
+        None,
+        AccessType::ReadWrite,
+    )
     .unwrap();
 
     // Create contract
@@ -185,23 +179,17 @@ fn create_contract_using_parent_id_test() {
     // Install the code for the child contract.
     let wasm_hash = xdr::Hash(Sha256::digest(&child_wasm).try_into().unwrap());
     // Add the contract code and code reference access to the footprint.
-    host.with_mut_storage(|s: &mut Storage| {
-        s.footprint
-            .record_access(
-                &host.contract_instance_ledger_key(&child_id).unwrap(),
-                AccessType::ReadWrite,
-                host.as_budget(),
-            )
-            .unwrap();
-        s.footprint
-            .record_access(
-                &host.contract_code_ledger_key(&wasm_hash).unwrap(),
-                AccessType::ReadWrite,
-                host.as_budget(),
-            )
-            .unwrap();
-        Ok(())
-    })
+    host.setup_storage_entry(
+        host.contract_instance_ledger_key(&child_id).unwrap(),
+        None,
+        AccessType::ReadWrite,
+    )
+    .unwrap();
+    host.setup_storage_entry(
+        host.contract_code_ledger_key(&wasm_hash).unwrap(),
+        None,
+        AccessType::ReadWrite,
+    )
     .unwrap();
 
     // Prepare arguments for the factory contract call.
diff --git a/soroban-env-host/src/test/metering_benchmark.rs b/soroban-env-host/src/test/metering_benchmark.rs
index b064c17f..5eb30488 100644
--- a/soroban-env-host/src/test/metering_benchmark.rs
+++ b/soroban-env-host/src/test/metering_benchmark.rs
@@ -1,6 +1,8 @@
 use crate::{
-    budget::Budget, host_object::HostVec, storage::Storage, Host, HostError, LedgerInfo,
-    MeteredOrdMap,
+    budget::Budget,
+    host_object::HostVec,
+    storage::{Footprint, Storage},
+    Host, HostError, LedgerInfo, MeteredOrdMap,
 };
 use soroban_env_common::{Env, Symbol};
 use soroban_test_wasms::{ADD_I32, COMPLEX};
@@ -54,9 +56,13 @@ fn run_add_i32() -> Result<(), HostError> {
     };
     // Run 2: enforce preflight footprint
     {
-        let store = Storage::with_enforcing_footprint_and_map(foot, MeteredOrdMap::default());
+        let store = Storage::with_enforcing_footprint_and_map(
+            Footprint::default(),
+            MeteredOrdMap::default(),
+        );
         let host = Host::with_storage_and_budget(store, Budget::default());
         host.set_ledger_info(LEDGER_INFO)?;
+        host.setup_storage_footprint(foot)?;
         let contract_id_obj =
             host.register_test_contract_wasm_from_source_account(ADD_I32, account_id, salt);
         host.measured_call(
@@ -91,9 +97,13 @@ fn run_complex() -> Result<(), HostError> {
 
     // Run 2: enforce preflight footprint, with empty map -- contract should only write.
     {
-        let store = Storage::with_enforcing_footprint_and_map(foot, MeteredOrdMap::default());
+        let store = Storage::with_enforcing_footprint_and_map(
+            Footprint::default(),
+            MeteredOrdMap::default(),
+        );
         let host = Host::with_storage_and_budget(store, Budget::default());
         host.set_ledger_info(LEDGER_INFO)?;
+        host.setup_storage_footprint(foot)?;
         let contract_id_obj =
             host.register_test_contract_wasm_from_source_account(COMPLEX, account_id, salt);
         host.measured_call(