From bcb814a7a3161584a6d73d189d8d9f1eb179c6b2 Mon Sep 17 00:00:00 2001
From: Keith <keith@opentensor.dev>
Date: Mon, 28 Oct 2024 12:45:30 +0900
Subject: [PATCH] Check whether caller is in whitelist before allowing creation

---
 frame/evm/src/lib.rs          |  6 ++++++
 frame/evm/src/runner/mod.rs   |  2 ++
 frame/evm/src/runner/stack.rs | 16 ++++++++++++++++
 3 files changed, 24 insertions(+)

diff --git a/frame/evm/src/lib.rs b/frame/evm/src/lib.rs
index 4821c8e737..49abf403cb 100644
--- a/frame/evm/src/lib.rs
+++ b/frame/evm/src/lib.rs
@@ -307,6 +307,7 @@ pub mod pallet {
 		) -> DispatchResultWithPostInfo {
 			T::CallOrigin::ensure_address_origin(&source, origin)?;
 
+			let whitelist = <WhitelistedCreators<T>>::get();
 			let is_transactional = true;
 			let validate = true;
 			let info = match T::Runner::create(
@@ -318,6 +319,7 @@ pub mod pallet {
 				max_priority_fee_per_gas,
 				nonce,
 				access_list,
+				whitelist,
 				is_transactional,
 				validate,
 				None,
@@ -394,6 +396,7 @@ pub mod pallet {
 		) -> DispatchResultWithPostInfo {
 			T::CallOrigin::ensure_address_origin(&source, origin)?;
 
+			let whitelist = <WhitelistedCreators<T>>::get();
 			let is_transactional = true;
 			let validate = true;
 			let info = match T::Runner::create2(
@@ -406,6 +409,7 @@ pub mod pallet {
 				max_priority_fee_per_gas,
 				nonce,
 				access_list,
+				whitelist,
 				is_transactional,
 				validate,
 				None,
@@ -516,6 +520,8 @@ pub mod pallet {
 		TransactionMustComeFromEOA,
 		/// Undefined error.
 		Undefined,
+		/// Origin is not allowed to perform the operation.
+		NotAllowed,
 	}
 
 	impl<T> From<TransactionValidationError> for Error<T> {
diff --git a/frame/evm/src/runner/mod.rs b/frame/evm/src/runner/mod.rs
index 45ed14299e..0ade55d868 100644
--- a/frame/evm/src/runner/mod.rs
+++ b/frame/evm/src/runner/mod.rs
@@ -73,6 +73,7 @@ pub trait Runner<T: Config> {
 		max_priority_fee_per_gas: Option<U256>,
 		nonce: Option<U256>,
 		access_list: Vec<(H160, Vec<H256>)>,
+		whitelist: Vec<H160>,
 		is_transactional: bool,
 		validate: bool,
 		weight_limit: Option<Weight>,
@@ -90,6 +91,7 @@ pub trait Runner<T: Config> {
 		max_priority_fee_per_gas: Option<U256>,
 		nonce: Option<U256>,
 		access_list: Vec<(H160, Vec<H256>)>,
+		whitelist: Vec<H160>,
 		is_transactional: bool,
 		validate: bool,
 		weight_limit: Option<Weight>,
diff --git a/frame/evm/src/runner/stack.rs b/frame/evm/src/runner/stack.rs
index 0707159cd2..975350fa98 100644
--- a/frame/evm/src/runner/stack.rs
+++ b/frame/evm/src/runner/stack.rs
@@ -462,6 +462,7 @@ where
 		max_priority_fee_per_gas: Option<U256>,
 		nonce: Option<U256>,
 		access_list: Vec<(H160, Vec<H256>)>,
+		whitelist: Vec<H160>,
 		is_transactional: bool,
 		validate: bool,
 		weight_limit: Option<Weight>,
@@ -469,6 +470,13 @@ where
 		config: &evm::Config,
 	) -> Result<CreateInfo, RunnerError<Self::Error>> {
 		if validate {
+			if !whitelist.contains(&source) {
+				return Err(RunnerError {
+					error: Error::<T>::NotAllowed,
+					weight: Weight::zero(),
+				});
+			}
+
 			Self::validate(
 				source,
 				None,
@@ -517,6 +525,7 @@ where
 		max_priority_fee_per_gas: Option<U256>,
 		nonce: Option<U256>,
 		access_list: Vec<(H160, Vec<H256>)>,
+		whitelist: Vec<H160>,
 		is_transactional: bool,
 		validate: bool,
 		weight_limit: Option<Weight>,
@@ -524,6 +533,13 @@ where
 		config: &evm::Config,
 	) -> Result<CreateInfo, RunnerError<Self::Error>> {
 		if validate {
+			if !whitelist.contains(&source) {
+				return Err(RunnerError {
+					error: Error::<T>::NotAllowed,
+					weight: Weight::zero(),
+				});
+			}
+
 			Self::validate(
 				source,
 				None,