Don't charge extra for trie deletions outside of contract execution

core/store/src/trie/mod.rs

@@ -41,6 +41,7 @@ use std::hash::Hash;
 use std::rc::Rc;
 use std::str;
 use std::sync::{Arc, RwLock, RwLockReadGuard};
+pub use trie_recording::RecordingMode;
 
 pub mod accounting_cache;
 mod config;
@@ -696,6 +697,13 @@ impl Trie {
             .unwrap_or_default()
     }
 
+    /// Set recording mode for the current trie recorder.
+    pub fn set_recording_mode(&mut self, recording_mode: RecordingMode) {
+        if let Some(recorder) = self.recorder.as_ref() {
+            recorder.borrow_mut().set_recording_mode(recording_mode);
+        }
+    }
+
     /// Constructs a Trie from the partial storage (i.e. state proof) that
     /// was returned from recorded_storage(). If used to access the same trie
     /// nodes as when the partial storage was generated, this trie will behave

core/store/src/trie/trie_recording.rs

@@ -16,6 +16,16 @@ pub struct TrieRecorder {
     code_len_counter: usize,
     /// Account IDs for which the code should be recorded.
     pub codes_to_record: HashSet<AccountId>,
+    pub recording_mode: RecordingMode,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum RecordingMode {
+    /// For recording internal runtime operations - doesn't charge extra for removals,
+    /// assuming that runtime operations aren't malicious.
+    Runtime,
+    /// Used during contract execution - charges extra for removals, assuming that they are malicious.
+    Contract,
 }
 
 impl TrieRecorder {
@@ -26,6 +36,7 @@ impl TrieRecorder {
             removal_counter: 0,
             code_len_counter: 0,
             codes_to_record: Default::default(),
+            recording_mode: RecordingMode::Runtime,
         }
     }
 
@@ -37,7 +48,9 @@ impl TrieRecorder {
     }
 
     pub fn record_removal(&mut self) {
-        self.removal_counter = self.removal_counter.saturating_add(1)
+        if self.recording_mode == RecordingMode::Contract {
+            self.removal_counter = self.removal_counter.saturating_add(1)
+        }
     }
 
     pub fn record_code_len(&mut self, code_len: usize) {
@@ -66,6 +79,10 @@ impl TrieRecorder {
             .saturating_add(removals_size)
             .saturating_add(self.code_len_counter)
     }
+
+    pub fn set_recording_mode(&mut self, recording_mode: RecordingMode) {
+        self.recording_mode = recording_mode;
+    }
 }
 
 #[cfg(test)]

runtime/runtime/src/lib.rs

@@ -50,6 +50,7 @@ use near_primitives::utils::{
 use near_primitives::version::{ProtocolFeature, ProtocolVersion};
 use near_primitives_core::apply::ApplyChunkReason;
 use near_store::trie::receipts_column_helper::DelayedReceiptQueue;
+use near_store::trie::RecordingMode;
 use near_store::{
     get, get_account, get_postponed_receipt, get_promise_yield_receipt, get_received_data,
     has_received_data, remove_postponed_receipt, remove_promise_yield_receipt, set, set_access_key,
@@ -1510,6 +1511,7 @@ impl Runtime {
             let recorded_storage_size_before = state_update.trie().recorded_storage_size();
             let storage_proof_size_upper_bound_before =
                 state_update.trie().recorded_storage_size_upper_bound();
+            state_update.trie.set_recording_mode(RecordingMode::Contract);
             let result = self.process_receipt(
                 state_update,
                 apply_state,
@@ -1519,6 +1521,7 @@ impl Runtime {
                 &mut stats,
                 epoch_info_provider,
             );
+            state_update.trie.set_recording_mode(RecordingMode::Runtime);
             let node_counter_after = state_update.trie().get_trie_nodes_count();
             tracing::trace!(target: "runtime", ?node_counter_before, ?node_counter_after);