From 6d3bcfedecc865fbb17fe123bfdbb27a5b229e87 Mon Sep 17 00:00:00 2001
From: Sergey Pepyakin <s.pepyakin@gmail.com>
Date: Mon, 6 Aug 2018 18:02:32 +0300
Subject: [PATCH] If contract reaches max depth, return Err

---
 substrate/runtime/contract/src/exec.rs  | 16 ++++++++----
 substrate/runtime/contract/src/tests.rs | 34 ++++++++++++++++++++++++-
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/substrate/runtime/contract/src/exec.rs b/substrate/runtime/contract/src/exec.rs
index 403fe050e6e13..7755582245cce 100644
--- a/substrate/runtime/contract/src/exec.rs
+++ b/substrate/runtime/contract/src/exec.rs
@@ -14,14 +14,14 @@
 // You should have received a copy of the GNU General Public License
 // along with Substrate. If not, see <http://www.gnu.org/licenses/>.
 
-use super::{CodeOf, ContractAddressFor, Module, Trait};
+use super::{CodeOf, MaxDepth, ContractAddressFor, Module, Trait};
 use account_db::{AccountDb, OverlayAccountDb};
 use gas::GasMeter;
 use vm;
 
 use rstd::prelude::*;
 use runtime_primitives::traits::{Zero, CheckedAdd, CheckedSub};
-use runtime_support::StorageMap;
+use runtime_support::{StorageMap, StorageValue};
 use staking;
 use system;
 
@@ -50,15 +50,17 @@ impl<'a, T: Trait> ExecutionContext<'a, T> {
 		gas_meter: &mut GasMeter<T>,
 		_data: &[u8],
 	) -> Result<CallReceipt, &'static str> {
-		let dest_code = <CodeOf<T>>::get(&dest);
-
-		// TODO: check the new depth
+		if self.depth == <MaxDepth<T>>::get() as usize {
+			return Err("reached maximum depth, cannot make a call");
+		}
 
 		let call_base_fee = <Module<T>>::call_base_fee();
 		if gas_meter.charge(call_base_fee).is_out_of_gas() {
 			return Err("not enough gas to pay base call fee");
 		}
 
+		let dest_code = <CodeOf<T>>::get(&dest);
+
 		let (exec_result, change_set) = {
 			let mut overlay = OverlayAccountDb::new(&self.overlay);
 
@@ -112,6 +114,10 @@ impl<'a, T: Trait> ExecutionContext<'a, T> {
 		ctor: &[u8],
 		_data: &[u8],
 	) -> Result<CreateReceipt<T>, &'static str> {
+		if self.depth == <MaxDepth<T>>::get() as usize {
+			return Err("reached maximum depth, cannot create");
+		}
+
 		let create_base_fee = <Module<T>>::create_base_fee();
 		if gas_meter.charge(create_base_fee).is_out_of_gas() {
 			return Err("not enough gas to pay base create fee");
diff --git a/substrate/runtime/contract/src/tests.rs b/substrate/runtime/contract/src/tests.rs
index 503fb780c3010..fa7832df53e41 100644
--- a/substrate/runtime/contract/src/tests.rs
+++ b/substrate/runtime/contract/src/tests.rs
@@ -121,7 +121,7 @@ fn new_test_ext(existential_deposit: u64, gas_price: u64) -> runtime_io::TestExt
 			call_base_fee: 135,
 			create_base_fee: 175,
 			gas_price,
-			max_depth: 1024,
+			max_depth: 100,
 		}.build_storage()
 			.unwrap(),
 	);
@@ -221,6 +221,38 @@ fn contract_transfer_oog() {
 
 }
 
+#[test]
+fn contract_transfer_max_depth() {
+	const CONTRACT_SHOULD_TRANSFER_TO: u64 = 9;
+
+	let code_transfer = wabt::wat2wasm(CODE_TRANSFER).unwrap();
+
+	with_externalities(&mut new_test_ext(0, 2), || {
+		<CodeOf<Test>>::insert(CONTRACT_SHOULD_TRANSFER_TO, code_transfer.to_vec());
+
+		Staking::set_free_balance(&0, 100_000_000);
+		Staking::set_free_balance(&CONTRACT_SHOULD_TRANSFER_TO, 11);
+
+		assert_err!(
+			Contract::call(&0, CONTRACT_SHOULD_TRANSFER_TO, 3, 100_000, Vec::new()),
+			"vm execute returned error while call"
+		);
+
+		assert_eq!(
+			Staking::free_balance(&0),
+			// 3 - value sent with the transaction
+			// 2 * 6 * 100 - gas used by the contract (6) multiplied by gas price (2)
+			//               multiplied by max depth (100).
+			// 2 * 135 * 100 - base gas fee for call (by transaction) multiplied by max depth (100).
+			100_000_000 - (2 * 135 * 100) - (2 * 6 * 100),
+		);
+		assert_eq!(
+			Staking::free_balance(&CONTRACT_SHOULD_TRANSFER_TO),
+			11,
+		);
+	});
+}
+
 /// Convert a byte slice to a string with hex values.
 ///
 /// Each value is preceeded with a `\` character.