Repository: stellar/rs-soroban-env
Pull request #1381 (2024-04-03 23:28:49)

Charge for BTreeSet
### What

Charge for BTreeSet creation in the `ParsedModule`/`ModueCache`

### Why

The BTreeSet was the right structure to use, but its cost was not charged for by the budget. 
Here we account for the cost of its creation (memory allocation, copy).

### Known limitations

Using `MeteredMap` would be more correct from metering perspective, which accounts for all look-ups and insertions. However, it is too inconvenient to use (every insertion creates a new map), and the extra lookup/insertion costs should be limited to worth the churn. 


Patch:
diff --git a/soroban-env-host/src/host/declared_size.rs b/soroban-env-host/src/host/declared_size.rs
index d6ab4a43..32a9981a 100644
--- a/soroban-env-host/src/host/declared_size.rs
+++ b/soroban-env-host/src/host/declared_size.rs
@@ -62,6 +62,7 @@ impl_declared_size_type!(i64, 8);
 impl_declared_size_type!(u128, 16);
 impl_declared_size_type!(i128, 16);
 impl_declared_size_type!(usize, 8);
+impl_declared_size_type!(&str, 16);
 
 // Val-wrapping types
 impl_declared_size_type!(Val, 8);
@@ -343,6 +344,7 @@ mod test {
         expect!["16"].assert_eq(size_of::<u128>().to_string().as_str());
         expect!["16"].assert_eq(size_of::<i128>().to_string().as_str());
         expect!["8"].assert_eq(size_of::<usize>().to_string().as_str());
+        expect!["16"].assert_eq(size_of::<&str>().to_string().as_str());
 
         // Val-wrapping types
         expect!["8"].assert_eq(size_of::<Val>().to_string().as_str());
@@ -574,6 +576,7 @@ mod test {
         assert_mem_size_le_declared_size!(u128);
         assert_mem_size_le_declared_size!(i128);
         assert_mem_size_le_declared_size!(usize);
+        assert_mem_size_le_declared_size!(&str);
 
         // Val-wrapping types
         assert_mem_size_le_declared_size!(Val);
diff --git a/soroban-env-host/src/vm.rs b/soroban-env-host/src/vm.rs
index edf1759e..adc75aed 100644
--- a/soroban-env-host/src/vm.rs
+++ b/soroban-env-host/src/vm.rs
@@ -160,7 +160,7 @@ impl Vm {
             //    to return the same error.
             let _span0 = tracy_span!("define host functions");
             let ledger_proto = host.with_ledger_info(|li| Ok(li.protocol_version))?;
-            parsed_module.with_import_symbols(|module_symbols| {
+            parsed_module.with_import_symbols(host, |module_symbols| {
                 for hf in HOST_FUNCTIONS {
                     if !module_symbols.contains(&(hf.mod_str, hf.fn_str)) {
                         continue;
@@ -229,7 +229,7 @@ impl Vm {
         if let Some(linker) = &*host.try_borrow_linker()? {
             Self::instantiate(host, contract_id, parsed_module, linker)
         } else {
-            let linker = parsed_module.make_linker()?;
+            let linker = parsed_module.make_linker(host)?;
             Self::instantiate(host, contract_id, parsed_module, &linker)
         }
     }
@@ -271,7 +271,7 @@ impl Vm {
         let config = get_wasmi_config(host.as_budget())?;
         let engine = wasmi::Engine::new(&config);
         let parsed_module = Self::parse_module(host, &engine, wasm, cost_inputs)?;
-        let linker = parsed_module.make_linker()?;
+        let linker = parsed_module.make_linker(host)?;
         Self::instantiate(host, contract_id, parsed_module, &linker)
     }
 
diff --git a/soroban-env-host/src/vm/module_cache.rs b/soroban-env-host/src/vm/module_cache.rs
index a0215441..1ed3e2db 100644
--- a/soroban-env-host/src/vm/module_cache.rs
+++ b/soroban-env-host/src/vm/module_cache.rs
@@ -4,7 +4,7 @@ use super::{
 };
 use crate::{
     budget::{get_wasmi_config, AsBudget},
-    host::metered_clone::MeteredClone,
+    host::metered_clone::{MeteredClone, MeteredContainer},
     xdr::{Hash, ScErrorCode, ScErrorType},
     Host, HostError, MeteredOrdMap,
 };
@@ -101,7 +101,7 @@ impl ModuleCache {
     ) -> Result<T, HostError> {
         let mut import_symbols = BTreeSet::new();
         for module in self.modules.values(host)? {
-            module.with_import_symbols(|module_symbols| {
+            module.with_import_symbols(host, |module_symbols| {
                 for hf in HOST_FUNCTIONS {
                     let sym = (hf.mod_str, hf.fn_str);
                     if module_symbols.contains(&sym) {
@@ -111,6 +111,14 @@ impl ModuleCache {
                 Ok(())
             })?;
         }
+        // We approximate the cost of `BTreeSet` with the cost of initializng a
+        // `Vec` with the same elements, and we are doing it after the set has
+        // been created. The element count has been limited/charged during the
+        // parsing phase, so there is no DOS factor. We don't charge for
+        // insertion/lookups, since they should be cheap and number of
+        // operations on the set is limited (only used during `Linker`
+        // creation).
+        Vec::<(&str, &str)>::charge_bulk_init_cpy(import_symbols.len() as u64, host)?;
         callback(&import_symbols)
     }
 
diff --git a/soroban-env-host/src/vm/parsed_module.rs b/soroban-env-host/src/vm/parsed_module.rs
index 30782100..0b9228eb 100644
--- a/soroban-env-host/src/vm/parsed_module.rs
+++ b/soroban-env-host/src/vm/parsed_module.rs
@@ -1,5 +1,6 @@
 use crate::{
     err,
+    host::metered_clone::MeteredContainer,
     meta::{self, get_ledger_protocol_version},
     xdr::{ContractCostType, Limited, ReadXdr, ScEnvMetaEntry, ScErrorCode, ScErrorType},
     Host, HostError, DEFAULT_XDR_RW_LIMITS,
@@ -165,6 +166,7 @@ impl ParsedModule {
 
     pub fn with_import_symbols<T>(
         &self,
+        host: &Host,
         callback: impl FnOnce(&BTreeSet<(&str, &str)>) -> Result<T, HostError>,
     ) -> Result<T, HostError> {
         // Cap symbols we're willing to import at 10 characters for each of
@@ -186,11 +188,24 @@ impl ParsedModule {
                 None
             })
             .collect();
+        // We approximate the cost of `BTreeSet` with the cost of initializng a
+        // `Vec` with the same elements, and we are doing it after the set has
+        // been created. The element count has been limited/charged during the
+        // parsing phase, so there is no DOS factor. We don't charge for
+        // insertion/lookups, since they should be cheap and number of
+        // operations on the set is limited.
+        if cfg!(feature = "next")
+            && host.get_ledger_protocol_version()? >= ModuleCache::MIN_LEDGER_VERSION
+        {
+            Vec::<(&str, &str)>::charge_bulk_init_cpy(symbols.len() as u64, host)?;
+        }
         callback(&symbols)
     }
 
-    pub fn make_linker(&self) -> Result<wasmi::Linker<Host>, HostError> {
-        self.with_import_symbols(|symbols| Host::make_linker(self.module.engine(), symbols))
+    pub fn make_linker(&self, host: &Host) -> Result<wasmi::Linker<Host>, HostError> {
+        self.with_import_symbols(host, |symbols| {
+            Host::make_linker(self.module.engine(), symbols)
+        })
     }
 
     #[cfg(feature = "bench")]


Test patch:
diff --git a/soroban-env-host/src/test/hostile.rs b/soroban-env-host/src/test/hostile.rs
index 09ab71e1..b6c80223 100644
--- a/soroban-env-host/src/test/hostile.rs
+++ b/soroban-env-host/src/test/hostile.rs
@@ -533,13 +533,13 @@ fn excessive_logging() -> Result<(), HostError> {
     #[cfg(feature = "next")]
     let expected_budget = expect![[r#"
         =================================================================
-        Cpu limit: 2000000; used: 212908
-        Mem limit: 500000; used: 166460
+        Cpu limit: 2000000; used: 214501
+        Mem limit: 500000; used: 166628
         =================================================================
         CostType                           cpu_insns      mem_bytes      
         WasmInsnExec                       300            0              
-        MemAlloc                           15292          67040          
-        MemCpy                             2331           0              
+        MemAlloc                           16609          67208          
+        MemCpy                             2607           0              
         MemCmp                             416            0              
         DispatchHostFunction               310            0              
         VisitObject                        244            0