seal_reentrant_count returns contract reentrant count

frame/contracts/fixtures/account_entrance_count_call.wat

@@ -0,0 +1,38 @@
+(module
+	(import "seal0" "seal_input" (func $seal_input (param i32 i32)))
+	(import "__unstable__" "seal_account_entrance_count" (func $seal_account_entrance_count (param i32) (result i32)))
+	(import "env" "memory" (memory 1 1))
+
+	;; [0, 32) buffer where input is copied
+
+	;; [32, 36) size of the input buffer
+	(data (i32.const 32) "\20")
+
+	(func $assert (param i32)
+		(block $ok
+			(br_if $ok
+				(get_local 0)
+			)
+			(unreachable)
+		)
+	)
+	(func (export "call")
+		(local $account_entrance_count i32)
+
+		;; Reading "callee" contract address (which is the address of the caller)
+		(call $seal_input (i32.const 0) (i32.const 32))
+
+		;; assert account_entrance_count == 1
+		(call $assert
+			(i32.eq (call $seal_account_entrance_count (i32.const 0)) (i32.const 1))
+		)
+
+		;; assert account_entrance_count == 0 for another account
+		(call $assert
+			(i32.eq (call $seal_account_entrance_count (i32.const 32)) (i32.const 0))
+		)
+	)
+
+	(func (export "deploy"))
+
+)

frame/contracts/fixtures/reentrant_count_call.wat

@@ -0,0 +1,67 @@
+(module
+	(import "seal0" "seal_input" (func $seal_input (param i32 i32)))
+	(import "seal0" "seal_call" (func $seal_call (param i32 i32 i64 i32 i32 i32 i32 i32 i32) (result i32)))
+	(import "__unstable__" "seal_reentrant_count" (func $seal_reentrant_count (result i32)))
+	(import "env" "memory" (memory 1 1))
+
+	;; [0, 32) buffer where code hash is copied
+
+	;; [32, 36) buffer for the call stack high
+
+	;; [36, 40) size of the input buffer
+	(data (i32.const 36) "\24")
+
+	(func $assert (param i32)
+		(block $ok
+			(br_if $ok
+				(get_local 0)
+			)
+			(unreachable)
+		)
+	)
+	(func (export "call")
+		(local $manual_reentrant_count i32)
+		(local $seal_call_exit_code i32)
+
+		;; Reading input
+		(call $seal_input (i32.const 0) (i32.const 36))
+
+		;; reading manually passed reentrant count
+		(set_local $manual_reentrant_count (i32.load (i32.const 32)))
+
+		;; reentrance count is calculated correctly
+		(call $assert
+			(i32.eq (call $seal_reentrant_count) (get_local $manual_reentrant_count))
+		)
+
+		(i32.eq (call $seal_reentrant_count) (i32.const 5))
+		(if
+			(then) ;; recursion exit case
+			(else
+				;; incrementing manual reentrant count high
+				(i32.store (i32.const 32) (i32.add (i32.load (i32.const 32)) (i32.const 1)))
+
+				;; Call to itself
+				(set_local $seal_call_exit_code
+					(call $seal_call
+						(i32.const 0)	;; Pointer to "callee" address.
+						(i32.const 32)	;; Length of "callee" address.
+						(i64.const 0)	;; How much gas to devote for the execution. 0 = all.
+						(i32.const 0)	;; Pointer to the buffer with value to transfer
+						(i32.const 0)	;; Length of the buffer with value to transfer.
+						(i32.const 0)	;; Pointer to input data buffer address
+						(i32.const 36)	;; Length of input data buffer
+						(i32.const 0xffffffff) ;; u32 max sentinel value: do not copy output
+						(i32.const 0) ;; Ptr to output buffer len
+					)
+				)
+
+				(call $assert
+					(i32.eq (get_local $seal_call_exit_code) (i32.const 0))
+				)
+			)
+		)
+	)
+
+	(func (export "deploy"))
+)

frame/contracts/fixtures/reentrant_count_delegated_call.wat

@@ -0,0 +1,69 @@
+(module
+	(import "seal0" "seal_input" (func $seal_input (param i32 i32)))
+	(import "seal0" "seal_set_storage" (func $seal_set_storage (param i32 i32 i32)))
+	(import "seal0" "seal_delegate_call" (func $seal_delegate_call (param i32 i32 i32 i32 i32 i32) (result i32)))
+	(import "__unstable__" "seal_reentrant_count" (func $seal_reentrant_count (result i32)))
+	(import "env" "memory" (memory 1 1))
+
+	;; [0, 32) buffer where code hash is copied
+
+	;; [32, 36) buffer for the call stack high
+
+	;; [36, 40) size of the input buffer
+	(data (i32.const 36) "\24")
+
+	(func $assert (param i32)
+		(block $ok
+			(br_if $ok
+				(get_local 0)
+			)
+			(unreachable)
+		)
+	)
+	(func (export "call")
+		(local $callstack_high i32)
+		(local $delegate_call_exit_code i32)
+
+		;; Reading input
+		(call $seal_input (i32.const 0) (i32.const 36))
+
+		;; reading passed callstack high
+		(set_local $callstack_high (i32.load (i32.const 32)))
+
+		;; incrementing callstack high
+		(i32.store (i32.const 32) (i32.add (i32.load (i32.const 32)) (i32.const 1)))
+
+		;; reentrance count stays 0
+		(call $assert
+			(i32.eq (call $seal_reentrant_count) (i32.const 0))
+		)
+
+		(i32.eq (get_local $callstack_high) (i32.const 5))
+		(if
+			(then) ;; exit recursion case
+			(else
+				;; Call to itself
+				(set_local $delegate_call_exit_code
+					(call $seal_delegate_call
+						(i32.const 0)	;; Set no call flags
+						(i32.const 0)	;; Pointer to "callee" code_hash.
+						(i32.const 0)	;; Input is ingored
+						(i32.const 36)	;; Length of the input
+						(i32.const 4294967295)	;; u32 max sentinel value: do not copy output
+						(i32.const 0)	;; Length is ignored in this case
+					)
+				)
+
+				(call $assert
+					(i32.eq (get_local $delegate_call_exit_code) (i32.const 0))
+				)
+			)
+		)
+
+		(call $assert
+			(i32.le_s (get_local $callstack_high) (i32.const 5))
+		)
+	)
+
+	(func (export "deploy"))
+)

frame/contracts/src/benchmarking/mod.rs

@@ -2086,6 +2086,59 @@ benchmarks! {
 		let origin = RawOrigin::Signed(instance.caller.clone());
 	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
 
+	seal_reentrant_count {
+		let r in 0 .. API_BENCHMARK_BATCHES;
+		let code = WasmModule::<T>::from(ModuleDefinition {
+			memory: Some(ImportedMemory::max::<T>()),
+			imported_functions: vec![ImportedFunction {
+				module: "__unstable__",
+				name: "seal_reentrant_count",
+				params: vec![],
+				return_type: Some(ValueType::I32),
+			}],
+			call_body: Some(body::repeated(r * API_BENCHMARK_BATCH_SIZE, &[
+				Instruction::Call(0),
+				Instruction::Drop,
+			])),
+			.. Default::default()
+		});
+		let instance = Contract::<T>::new(code, vec![])?;
+		let origin = RawOrigin::Signed(instance.caller.clone());
+	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
+
+	seal_account_entrance_count {
+		let r in 0 .. API_BENCHMARK_BATCHES;
+		let dummy_code = WasmModule::<T>::dummy_with_bytes(0);
+		let accounts = (0..r * API_BENCHMARK_BATCH_SIZE)
+			.map(|i| Contract::with_index(i + 1, dummy_code.clone(), vec![]))
+			.collect::<Result<Vec<_>, _>>()?;
+		let account_id_len = accounts.get(0).map(|i| i.account_id.encode().len()).unwrap_or(0);
+		let account_id_bytes = accounts.iter().flat_map(|x| x.account_id.encode()).collect();
+		let code = WasmModule::<T>::from(ModuleDefinition {
+			memory: Some(ImportedMemory::max::<T>()),
+			imported_functions: vec![ImportedFunction {
+				module: "__unstable__",
+				name: "seal_account_entrance_count",
+				params: vec![ValueType::I32],
+				return_type: Some(ValueType::I32),
+			}],
+			data_segments: vec![
+				DataSegment {
+					offset: 0,
+					value: account_id_bytes,
+				},
+			],
+			call_body: Some(body::repeated_dyn(r * API_BENCHMARK_BATCH_SIZE, vec![
+				Counter(0, account_id_len as u32), // account_ptr
+				Regular(Instruction::Call(0)),
+				Regular(Instruction::Drop),
+			])),
+			.. Default::default()
+		});
+		let instance = Contract::<T>::new(code, vec![])?;
+		let origin = RawOrigin::Signed(instance.caller.clone());
+	}: call(origin, instance.addr, 0u32.into(), Weight::MAX, None, vec![])
+
 	// We make the assumption that pushing a constant and dropping a value takes roughly
 	// the same amount of time. We follow that `t.load` and `drop` both have the weight
 	// of this benchmark / 2. We need to make this assumption because there is no way

frame/contracts/src/exec.rs

@@ -296,6 +296,15 @@ pub trait Ext: sealing::Sealed {
 
 	/// Sets new code hash for existing contract.
 	fn set_code_hash(&mut self, hash: CodeHash<Self::T>) -> Result<(), DispatchError>;
+
+	/// Returns the number of times the currently executing contract exists on the call stack in
+	/// addition to the calling instance. A value of 0 means no reentrancy.
+	fn reentrant_count(&self) -> u32;
+
+	/// Returns the number of times the specified contract exists on the call stack. Delegated calls
+	/// are not calculated as separate entrance.
+	/// A value of 0 means it does not exist on the call stack.
+	fn account_entrance_count(&self, account_id: &AccountIdOf<Self::T>) -> u32;
 }
 
 /// Describes the different functions that can be exported by an [`Executable`].
@@ -1305,6 +1314,17 @@ where
 		});
 		Ok(())
 	}
+
+	fn reentrant_count(&self) -> u32 {
+		let id: &AccountIdOf<Self::T> = &self.top_frame().account_id;
+		self.account_entrance_count(id).saturating_sub(1)
+	}
+
+	fn account_entrance_count(&self, account_id: &AccountIdOf<Self::T>) -> u32 {
+		self.frames()
+			.filter(|f| f.delegate_caller.is_none() && &f.account_id == account_id)
+			.count() as u32
+	}
 }
 
 fn deposit_event<T: Config>(topics: Vec<T::Hash>, event: Event<T>) {

frame/contracts/src/schedule.rs

@@ -416,6 +416,12 @@ pub struct HostFnWeights<T: Config> {
 	/// Weight of calling `seal_ecdsa_to_eth_address`.
 	pub ecdsa_to_eth_address: Weight,
 
+	/// Weight of calling `seal_reentrant_count`.
+	pub reentrant_count: Weight,
+
+	/// Weight of calling `seal_account_entrance_count`.
+	pub account_entrance_count: Weight,
+
 	/// The type parameter is used in the default implementation.
 	#[codec(skip)]
 	pub _phantom: PhantomData<T>,
@@ -651,6 +657,8 @@ impl<T: Config> Default for HostFnWeights<T> {
 			hash_blake2_128_per_byte: cost_byte_batched!(seal_hash_blake2_128_per_kb),
 			ecdsa_recover: cost_batched!(seal_ecdsa_recover),
 			ecdsa_to_eth_address: cost_batched!(seal_ecdsa_to_eth_address),
+			reentrant_count: cost_batched!(seal_reentrant_count),
+			account_entrance_count: cost_batched!(seal_account_entrance_count),
 			_phantom: PhantomData,
 		}
 	}

frame/contracts/src/tests.rs

@@ -3272,3 +3272,92 @@ fn set_code_hash() {
 		);
 	});
 }
+
+#[test]
+#[cfg(feature = "unstable-interface")]
+fn reentrant_count_works_with_call() {
+	let (wasm, code_hash) = compile_module::<Test>("reentrant_count_call").unwrap();
+	let contract_addr = Contracts::contract_address(&ALICE, &code_hash, &[]);
+
+	ExtBuilder::default().existential_deposit(100).build().execute_with(|| {
+		let _ = Balances::deposit_creating(&ALICE, 1_000_000);
+
+		assert_ok!(Contracts::instantiate_with_code(
+			Origin::signed(ALICE),
+			300_000,
+			GAS_LIMIT,
+			None,
+			wasm,
+			vec![],
+			vec![],
+		));
+
+		// adding reentrant count to the input
+		let input = (contract_addr.clone(), 0).encode();
+
+		Contracts::bare_call(ALICE, contract_addr, 0, GAS_LIMIT, None, input, true)
+			.result
+			.unwrap();
+	});
+}
+
+#[test]
+#[cfg(feature = "unstable-interface")]
+fn reentrant_count_works_with_delegated_call() {
+	let (wasm, code_hash) = compile_module::<Test>("reentrant_count_delegated_call").unwrap();
+	let contract_addr = Contracts::contract_address(&ALICE, &code_hash, &[]);
+
+	ExtBuilder::default().existential_deposit(100).build().execute_with(|| {
+		let _ = Balances::deposit_creating(&ALICE, 1_000_000);
+
+		assert_ok!(Contracts::instantiate_with_code(
+			Origin::signed(ALICE),
+			300_000,
+			GAS_LIMIT,
+			None,
+			wasm,
+			vec![],
+			vec![],
+		));
+
+		// adding a callstack height to the input
+		let input = (code_hash, 1).encode();
+
+		Contracts::bare_call(ALICE, contract_addr.clone(), 0, GAS_LIMIT, None, input, true)
+			.result
+			.unwrap();
+	});
+}
+
+#[test]
+#[cfg(feature = "unstable-interface")]
+fn account_entrance_count_works() {
+	let (wasm, code_hash) = compile_module::<Test>("account_entrance_count_call").unwrap();
+	let contract_addr = Contracts::contract_address(&ALICE, &code_hash, &[]);
+
+	ExtBuilder::default().existential_deposit(100).build().execute_with(|| {
+		let _ = Balances::deposit_creating(&ALICE, 1_000_000);
+
+		assert_ok!(Contracts::instantiate_with_code(
+			Origin::signed(ALICE),
+			300_000,
+			GAS_LIMIT,
+			None,
+			wasm,
+			vec![],
+			vec![],
+		));
+
+		Contracts::bare_call(
+			ALICE,
+			contract_addr.clone(),
+			0,
+			GAS_LIMIT,
+			None,
+			contract_addr.encode(),
+			true,
+		)
+		.result
+		.unwrap();
+	});
+}