Add Limit to Tranasctional Layers

frame/contracts/src/exec.rs

@@ -766,14 +766,27 @@ where
 
  // All changes performed by the contract are executed under a storage transaction.
  // This allows for roll back on error. Changes to the cached contract_info are
- // comitted or rolled back when popping the frame.
- let (success, output) = with_transaction(|| {
- let output = do_transaction();
- match &output {
- Ok(result) if !result.did_revert() => TransactionOutcome::Commit((true, output)),
- _ => TransactionOutcome::Rollback((false, output)),
- }
- });
+ // committed or rolled back when popping the frame.
+ //
+ // `with_transactional` may return an error caused by a limit in the
+ // transactional storage depth.
+ let transaction_outcome =
+ with_transaction(|| -> TransactionOutcome<Result<_, DispatchError>> {
+ let output = do_transaction();
+ match &output {
+ Ok(result) if !result.did_revert() =>
+ TransactionOutcome::Commit(Ok((true, output))),
+ _ => TransactionOutcome::Rollback(Ok((false, output))),
+ }
+ });
+
+ let (success, output) = match transaction_outcome {
+ // `with_transactional` executed successfully, and we have the expected output.
+ Ok((success, output)) => (success, output),
+ // `with_transactional` returned an error, and we propagate that error and note no state
+ // has changed.
+ Err(error) => (false, Err(error.into())),
+ };
  self.pop_frame(success);
  output
  }

frame/support/procedural/src/transactional.rs

@@ -49,7 +49,9 @@ pub fn require_transactional(_attr: TokenStream, input: TokenStream) -> Result<T
  let output = quote! {
  #(#attrs)*
  #vis #sig {
- #crate_::storage::require_transaction();
+ if !#crate_::storage::is_transactional() {
+ return Err(#crate_::sp_runtime::TransactionalError::NoLayer.into());
+ }
  #block
  }
  };

frame/support/src/storage/mod.rs

@@ -27,8 +27,11 @@ use crate::{
 };
 use codec::{Decode, Encode, EncodeLike, FullCodec, FullEncode};
 use sp_core::storage::ChildInfo;
-use sp_runtime::generic::{Digest, DigestItem};
 pub use sp_runtime::TransactionOutcome;
+use sp_runtime::{
+ generic::{Digest, DigestItem},
+ DispatchError, TransactionalError,
+};
 use sp_std::prelude::*;
 pub use types::Key;
 
@@ -44,71 +47,121 @@ pub mod types;
 pub mod unhashed;
 pub mod weak_bounded_vec;
 
-#[cfg(all(feature = "std", any(test, debug_assertions)))]
-mod debug_helper {
- use std::cell::RefCell;
+mod transaction_level_tracker {
+ use core::sync::atomic::{AtomicU32, Ordering};
+
+ type Layer = u32;
+ static NUM_LEVELS: AtomicU32 = AtomicU32::new(0);
+ const TRANSACTIONAL_LIMIT: Layer = 255;
 
- thread_local! {
- static TRANSACTION_LEVEL: RefCell<u32> = RefCell::new(0);
+ pub fn get_transaction_level() -> Layer {
+ NUM_LEVELS.load(Ordering::SeqCst)
  }
 
- pub fn require_transaction() {
- let level = TRANSACTION_LEVEL.with(|v| *v.borrow());
- if level == 0 {
- panic!("Require transaction not called within with_transaction");
- }
+ /// Increments the transaction level. Returns an error if levels go past the limit.
+ ///
+ /// Returns a guard that when dropped decrements the transaction level automatically.
+ pub fn inc_transaction_level() -> Result<StorageLayerGuard, ()> {
+ NUM_LEVELS
+ .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |existing_levels| {
+ if existing_levels >= TRANSACTIONAL_LIMIT {
+ return None
+ }
+ // Cannot overflow because of check above.
+ Some(existing_levels + 1)
+ })
+ .map_err(|_| ())?;
+ Ok(StorageLayerGuard)
  }
 
- pub struct TransactionLevelGuard;
+ fn dec_transaction_level() {
+ NUM_LEVELS
+ .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |existing_levels| {
+ if existing_levels == 0 {
+ log::warn!(
+ "We are underflowing with calculating transactional levels. Not great, but let's not panic...",
+ );
+ None
+ } else {
+ // Cannot underflow because of checks above.
+ Some(existing_levels - 1)
+ }
+ })
+ .ok();
+ }
 
- impl Drop for TransactionLevelGuard {
- fn drop(&mut self) {
- TRANSACTION_LEVEL.with(|v| *v.borrow_mut() -= 1);
- }
+ pub fn is_transactional() -> bool {
+ get_transaction_level() > 0
  }
 
- /// Increments the transaction level.
- ///
- /// Returns a guard that when dropped decrements the transaction level automatically.
- pub fn inc_transaction_level() -> TransactionLevelGuard {
- TRANSACTION_LEVEL.with(|v| {
- let mut val = v.borrow_mut();
- *val += 1;
- if *val > 10 {
- log::warn!(
- "Detected with_transaction with nest level {}. Nested usage of with_transaction is not recommended.",
- *val
- );
- }
- });
+ pub struct StorageLayerGuard;
 
- TransactionLevelGuard
+ impl Drop for StorageLayerGuard {
+ fn drop(&mut self) {
+ dec_transaction_level()
+ }
  }
 }
 
-/// Assert this method is called within a storage transaction.
-/// This will **panic** if is not called within a storage transaction.
-///
-/// This assertion is enabled for native execution and when `debug_assertions` are enabled.
-pub fn require_transaction() {
- #[cfg(all(feature = "std", any(test, debug_assertions)))]
- debug_helper::require_transaction();
+/// Check if the current call is within a transactional layer.
+pub fn is_transactional() -> bool {
+ transaction_level_tracker::is_transactional()
 }
 
 /// Execute the supplied function in a new storage transaction.
 ///
 /// All changes to storage performed by the supplied function are discarded if the returned
 /// outcome is `TransactionOutcome::Rollback`.
 ///
-/// Transactions can be nested to any depth. Commits happen to the parent transaction.
-pub fn with_transaction<R>(f: impl FnOnce() -> TransactionOutcome<R>) -> R {
+/// Transactions can be nested up to `TRANSACTIONAL_LIMIT` times; more than that will result in an
+/// error.
+///
+/// Commits happen to the parent transaction.
+pub fn with_transaction<T, E>(f: impl FnOnce() -> TransactionOutcome<Result<T, E>>) -> Result<T, E>
+where
+ E: From<DispatchError>,
+{
  use sp_io::storage::{commit_transaction, rollback_transaction, start_transaction};
  use TransactionOutcome::*;
 
+ let _guard = transaction_level_tracker::inc_transaction_level()
+ .map_err(|()| TransactionalError::LimitReached.into())?;
+
  start_transaction();
 
- #[cfg(all(feature = "std", any(test, debug_assertions)))]
- let _guard = debug_helper::inc_transaction_level();
+ match f() {
+ Commit(res) => {
+ commit_transaction();
+ res
+ },
+ Rollback(res) => {
+ rollback_transaction();
+ res
+ },
+ }
+}
+
+/// Same as [`with_transaction`] but without a limit check on nested transactional layers.
+///
+/// This is mostly for backwards compatibility before there was a transactional layer limit.
+/// It is recommended to only use [`with_transaction`] to avoid users from generating too many
+/// transactional layers.
+pub fn with_transaction_unchecked<R>(f: impl FnOnce() -> TransactionOutcome<R>) -> R {
+ use sp_io::storage::{commit_transaction, rollback_transaction, start_transaction};
+ use TransactionOutcome::*;
+
+ let maybe_guard = transaction_level_tracker::inc_transaction_level();
+
+ if maybe_guard.is_err() {
+ log::warn!(
+ "The transactional layer limit has been reached, and new transactional layers are being
+ spawned with `with_transaction_unchecked`. This could be caused by someone trying to
+ attack your chain, and you should investigate usage of `with_transaction_unchecked` and
+ potentially migrate to `with_transaction`, which enforces a transactional limit.",
+ );
+ }
+
+ start_transaction();
 
  match f() {
  Commit(res) => {
@@ -1421,12 +1474,13 @@ pub fn storage_prefix(pallet_name: &[u8], storage_name: &[u8]) -> [u8; 32] {
 #[cfg(test)]
 mod test {
  use super::*;
- use crate::{assert_ok, hash::Identity, Twox128};
+ use crate::{assert_noop, assert_ok, hash::Identity, Twox128};
  use bounded_vec::BoundedVec;
  use frame_support::traits::ConstU32;
  use generator::StorageValue as _;
  use sp_core::hashing::twox_128;
  use sp_io::TestExternalities;
+ use sp_runtime::DispatchResult;
  use weak_bounded_vec::WeakBoundedVec;
 
  #[test]
@@ -1538,25 +1592,67 @@ mod test {
  }
 
  #[test]
- #[should_panic(expected = "Require transaction not called within with_transaction")]
- fn require_transaction_should_panic() {
+ fn is_transactional_should_return_false() {
  TestExternalities::default().execute_with(|| {
- require_transaction();
+ assert!(!is_transactional());
  });
  }
 
  #[test]
- fn require_transaction_should_not_panic_in_with_transaction() {
+ fn is_transactional_should_not_error_in_with_transaction() {
  TestExternalities::default().execute_with(|| {
- with_transaction(|| {
- require_transaction();
- TransactionOutcome::Commit(())
- });
-
- with_transaction(|| {
- require_transaction();
- TransactionOutcome::Rollback(())
- });
+ assert_ok!(with_transaction(|| -> TransactionOutcome<DispatchResult> {
+ assert!(is_transactional());
+ TransactionOutcome::Commit(Ok(()))
+ }));
+
+ assert_noop!(
+ with_transaction(|| -> TransactionOutcome<DispatchResult> {
+ assert!(is_transactional());
+ TransactionOutcome::Rollback(Err("revert".into()))
+ }),
+ "revert"
+ );
+ });
+ }
+
+ fn recursive_transactional(num: u32) -> DispatchResult {
+ if num == 0 {
+ return Ok(())
+ }
+
+ with_transaction(|| -> TransactionOutcome<DispatchResult> {
+ let res = recursive_transactional(num - 1);
+ TransactionOutcome::Commit(res)
+ })
+ }
+
+ #[test]
+ fn transaction_limit_should_work() {
+ TestExternalities::default().execute_with(|| {
+ assert_eq!(transaction_level_tracker::get_transaction_level(), 0);
+
+ assert_ok!(with_transaction(|| -> TransactionOutcome<DispatchResult> {
+ assert_eq!(transaction_level_tracker::get_transaction_level(), 1);
+ TransactionOutcome::Commit(Ok(()))
+ }));
+
+ assert_ok!(with_transaction(|| -> TransactionOutcome<DispatchResult> {
+ assert_eq!(transaction_level_tracker::get_transaction_level(), 1);
+ let res = with_transaction(|| -> TransactionOutcome<DispatchResult> {
+ assert_eq!(transaction_level_tracker::get_transaction_level(), 2);
+ TransactionOutcome::Commit(Ok(()))
+ });
+ TransactionOutcome::Commit(res)
+ }));
+
+ assert_ok!(recursive_transactional(255));
+ assert_noop!(
+ recursive_transactional(256),
+ sp_runtime::TransactionalError::LimitReached
+ );
+
+ assert_eq!(transaction_level_tracker::get_transaction_level(), 0);
  });
  }