Add host function that enforces transaction nesting depth

primitives/externalities/src/lib.rs

@@ -188,6 +188,11 @@ pub trait Externalities: ExtensionStore {
  /// Changes made without any open transaction are committed immediately.
  fn storage_start_transaction(&mut self);
 
+ /// Same as `storage_start_transaction` but with enforcing a maximum nesting depth.
+ ///
+ /// Will return an `Err` when it would otherwise violate the supplied `max_depth`.
+ fn storage_start_transaction_with_limit(&mut self, max_depth: u32) -> Result<(), ()>;
+
  /// Rollback the last transaction started by `storage_start_transaction`.
  ///
  /// Any changes made during that storage transaction are discarded. Returns an error when

primitives/io/src/lib.rs

@@ -233,6 +233,15 @@ pub trait Storage {
  self.storage_start_transaction();
  }
 
+ /// Start a new nested transaction while enforcing a maximum nesting depth.
+ ///
+ /// This function is equal to `start_transaction` with the exception that it
+ /// will return an error then it would otherwise violate the supplied `max_depth`. The
+ /// parameter describes the maximum allowed nesting depth.
+ fn start_transaction_with_limit(&mut self, max_depth: u32) -> Result<(), ()> {
+ self.storage_start_transaction_with_limit(max_depth)
+ }
+
  /// Rollback the last transaction started by `start_transaction`.
  ///
  /// Any changes made during that transaction are discarded.

primitives/state-machine/src/basic.rs

@@ -323,6 +323,10 @@ impl Externalities for BasicExternalities {
  unimplemented!("Transactions are not supported by BasicExternalities");
  }
 
+ fn storage_start_transaction_with_limit(&mut self, _: u32) -> Result<(), ()> {
+ unimplemented!("Transactions are not supported by BasicExternalities");
+ }
+
  fn storage_rollback_transaction(&mut self) -> Result<(), ()> {
  unimplemented!("Transactions are not supported by BasicExternalities");
  }

primitives/state-machine/src/ext.rs

@@ -644,6 +644,10 @@ where
  self.overlay.start_transaction()
  }
 
+ fn storage_start_transaction_with_limit(&mut self, max_depth: u32) -> Result<(), ()> {
+ self.overlay.start_transaction_with_limit(max_depth).map_err(|_| ())
+ }
+
  fn storage_rollback_transaction(&mut self) -> Result<(), ()> {
  self.mark_dirty();
  self.overlay.rollback_transaction().map_err(|_| ())
@@ -1121,4 +1125,49 @@ mod tests {
 
  assert_eq!(Vec::<u32>::decode(&mut &data[..]).unwrap(), vec![1, 2]);
  }
+
+ #[test]
+ fn transaction_limit_works() {
+ let mut cache = StorageTransactionCache::default();
+ let mut overlay = OverlayedChanges::default();
+ let backend = InMemoryBackend::default();
+ let mut ext = TestExt::new(&mut overlay, &mut cache, &backend, None);
+
+ // Zero will always error out
+ assert!(ext.storage_start_transaction_with_limit(0).is_err());
+
+ // No transaction created, yet
+ assert!(ext.storage_start_transaction_with_limit(1).is_ok());
+ assert!(ext.storage_start_transaction_with_limit(1).is_err());
+ assert!(ext.storage_start_transaction_with_limit(0).is_err());
+
+ // We only spawned one transaction until here
+ assert!(ext.storage_start_transaction_with_limit(2).is_ok());
+ assert!(ext.storage_start_transaction_with_limit(2).is_err());
+
+ // The unchecked version can always spawn more transactions
+ ext.storage_start_transaction();
+ assert!(ext.storage_start_transaction_with_limit(3).is_err());
+ assert!(ext.storage_start_transaction_with_limit(4).is_ok());
+ assert!(ext.storage_start_transaction_with_limit(4).is_err());
+
+ // Smaller numbers won't work either
+ assert!(ext.storage_start_transaction_with_limit(3).is_err());
+ assert!(ext.storage_start_transaction_with_limit(2).is_err());
+ assert!(ext.storage_start_transaction_with_limit(1).is_err());
+ assert!(ext.storage_start_transaction_with_limit(0).is_err());
+
+ // We are at 4 open transactions
+ assert!(ext.storage_start_transaction_with_limit(4).is_err());
+ assert!(ext.storage_commit_transaction().is_ok());
+ assert!(ext.storage_start_transaction_with_limit(4).is_ok());
+ assert!(ext.storage_start_transaction_with_limit(4).is_err());
+ assert!(ext.storage_rollback_transaction().is_ok());
+ assert!(ext.storage_start_transaction_with_limit(4).is_ok());
+ assert!(ext.storage_rollback_transaction().is_ok());
+ assert!(ext.storage_commit_transaction().is_ok());
+ assert!(ext.storage_rollback_transaction().is_ok());
+ assert!(ext.storage_commit_transaction().is_ok());
+ assert!(ext.storage_start_transaction_with_limit(1).is_ok());
+ }
 }

primitives/state-machine/src/overlayed_changes/changeset.rs

@@ -214,6 +214,11 @@ impl<K: Ord + Hash + Clone, V> OverlayedMap<K, V> {
  self.changes.is_empty()
  }
 
+ /// Returns how many storage transactions are currently open.
+ pub fn depth(&self) -> usize {
+ self.dirty_keys.len()
+ }
+
  /// Get an optional reference to the value stored for the specified key.
  pub fn get<Q>(&self, key: &Q) -> Option<&OverlayedEntry<V>>
  where

primitives/state-machine/src/overlayed_changes/mod.rs

@@ -44,6 +44,12 @@ use std::{
 
 pub use self::changeset::{AlreadyInRuntime, NoOpenTransaction, NotInRuntime, OverlayedValue};
 
+/// Error returned when spawning a new transaction would vialoate the amount of
+/// transactions that are allowed to exit.
+#[derive(Debug)]
+#[cfg_attr(test, derive(PartialEq))]
+pub struct TransactionLimitReached;
+
 /// Changes that are made outside of extrinsics are marked with this index;
 pub const NO_EXTRINSIC_INDEX: u32 = 0xffffffff;
 
@@ -359,6 +365,21 @@ impl OverlayedChanges {
  self.offchain.overlay_mut().start_transaction();
  }
 
+ /// Same as `start_transaction` but with enforcing a maximum nesting depth.
+ ///
+ /// Will return an `Err` when it would otherwise violate the supplied `max_depth`.
+ pub fn start_transaction_with_limit(
+ &mut self,
+ max_depth: u32,
+ ) -> Result<(), TransactionLimitReached> {
+ let depth: u32 = self.top.depth().try_into().unwrap_or(u32::MAX);
+ if depth >= max_depth {
+ return Err(TransactionLimitReached)
+ }
+ self.start_transaction();
+ Ok(())
+ }
+
  /// Rollback the last transaction started by `start_transaction`.
  ///
  /// Any changes made during that transaction are discarded. Returns an error if

primitives/state-machine/src/read_only.rs

@@ -161,6 +161,10 @@ impl<'a, H: Hasher, B: 'a + Backend<H>> Externalities for ReadOnlyExternalities<
  unimplemented!("Transactions are not supported by ReadOnlyExternalities");
  }
 
+ fn storage_start_transaction_with_limit(&mut self, _: u32) -> Result<(), ()> {
+ unimplemented!("Transactions are not supported by ReadOnlyExternalities");
+ }
+
  fn storage_rollback_transaction(&mut self) -> Result<(), ()> {
  unimplemented!("Transactions are not supported by ReadOnlyExternalities");
  }

primitives/tasks/src/async_externalities.rs

@@ -142,6 +142,10 @@ impl Externalities for AsyncExternalities {
  unimplemented!("Transactions are not supported by AsyncExternalities");
  }
 
+ fn storage_start_transaction_with_limit(&mut self, _: u32) -> Result<(), ()> {
+ unimplemented!("Transactions are not supported by AsyncExternalities");
+ }
+
  fn storage_rollback_transaction(&mut self) -> Result<(), ()> {
  unimplemented!("Transactions are not supported by AsyncExternalities");
  }