From d65d14760c66a9fb94be6c7bd123c1b04e1632f0 Mon Sep 17 00:00:00 2001
From: bear <boundless.forest@outlook.com>
Date: Thu, 2 Nov 2023 14:58:02 +0800
Subject: [PATCH] Fix add `size_limit` to the proof_size of weight info

---
 frame/evm/src/runner/stack.rs | 34 ++++++++++++++++++++++------------
 primitives/evm/src/lib.rs     | 19 +++++++++++++++++++
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/frame/evm/src/runner/stack.rs b/frame/evm/src/runner/stack.rs
index 3d20134c2d..eccb11422d 100644
--- a/frame/evm/src/runner/stack.rs
+++ b/frame/evm/src/runner/stack.rs
@@ -986,12 +986,17 @@ where
 						if let Some(meta) = <AccountCodesMetadata<T>>::get(address) {
 							weight_info.try_record_proof_size_or_fail(meta.size)?;
 						} else {
-							// If it does not exist, try to record `create_contract_limit` first.
-							weight_info.try_record_proof_size_or_fail(size_limit)?;
-							let meta = Pallet::<T>::account_code_metadata(address);
-							let actual_size = meta.size;
-							// Refund if applies
-							weight_info.refund_proof_size(size_limit.saturating_sub(actual_size));
+							if let Some(remaining_proof_size) = weight_info.remaining_proof_size() {
+								let pre_size = remaining_proof_size.min(size_limit);
+								weight_info.try_record_proof_size_or_fail(pre_size)?;
+
+								let actual_size = Pallet::<T>::account_code_metadata(address).size;
+								if actual_size > pre_size {
+									return Err(ExitError::OutOfGas);
+								}
+								// Refund if applies
+								weight_info.refund_proof_size(pre_size.saturating_sub(actual_size));
+							}
 						}
 						recorded.account_codes.push(address);
 					}
@@ -1081,12 +1086,17 @@ where
 			if let Some(meta) = <AccountCodesMetadata<T>>::get(address) {
 				weight_info.try_record_proof_size_or_fail(meta.size)?;
 			} else {
-				// If it does not exist, try to record `create_contract_limit` first.
-				weight_info.try_record_proof_size_or_fail(size_limit)?;
-				let meta = Pallet::<T>::account_code_metadata(address);
-				let actual_size = meta.size;
-				// Refund if applies
-				weight_info.refund_proof_size(size_limit.saturating_sub(actual_size));
+				if let Some(remaining_proof_size) = weight_info.remaining_proof_size() {
+					let pre_size = remaining_proof_size.min(size_limit);
+					weight_info.try_record_proof_size_or_fail(pre_size)?;
+
+					let actual_size = Pallet::<T>::account_code_metadata(address).size;
+					if actual_size > pre_size {
+						return Err(ExitError::OutOfGas);
+					}
+					// Refund if applies
+					weight_info.refund_proof_size(pre_size.saturating_sub(actual_size));
+				}
 			}
 			recorded.account_codes.push(address);
 			// Already recorded, return
diff --git a/primitives/evm/src/lib.rs b/primitives/evm/src/lib.rs
index bdbb2792f0..6cff8ed901 100644
--- a/primitives/evm/src/lib.rs
+++ b/primitives/evm/src/lib.rs
@@ -108,6 +108,7 @@ impl WeightInfo {
 			_ => return Err("must provide Some valid weight limit or None"),
 		})
 	}
+	
 	fn try_consume(&self, cost: u64, limit: u64, usage: u64) -> Result<u64, ExitError> {
 		let usage = usage.checked_add(cost).ok_or(ExitError::OutOfGas)?;
 		if usage > limit {
@@ -115,6 +116,7 @@ impl WeightInfo {
 		}
 		Ok(usage)
 	}
+	
 	pub fn try_record_ref_time_or_fail(&mut self, cost: u64) -> Result<(), ExitError> {
 		if let (Some(ref_time_usage), Some(ref_time_limit)) =
 			(self.ref_time_usage, self.ref_time_limit)
@@ -127,6 +129,7 @@ impl WeightInfo {
 		}
 		Ok(())
 	}
+
 	pub fn try_record_proof_size_or_fail(&mut self, cost: u64) -> Result<(), ExitError> {
 		if let (Some(proof_size_usage), Some(proof_size_limit)) =
 			(self.proof_size_usage, self.proof_size_limit)
@@ -139,18 +142,34 @@ impl WeightInfo {
 		}
 		Ok(())
 	}
+
 	pub fn refund_proof_size(&mut self, amount: u64) {
 		if let Some(proof_size_usage) = self.proof_size_usage {
 			let proof_size_usage = proof_size_usage.saturating_sub(amount);
 			self.proof_size_usage = Some(proof_size_usage);
 		}
 	}
+
 	pub fn refund_ref_time(&mut self, amount: u64) {
 		if let Some(ref_time_usage) = self.ref_time_usage {
 			let ref_time_usage = ref_time_usage.saturating_sub(amount);
 			self.ref_time_usage = Some(ref_time_usage);
 		}
 	}
+	pub fn remaining_proof_size(&self) -> Option<u64> {
+		if let (Some(proof_size_usage), Some(proof_size_limit)) = (self.proof_size_usage, self.proof_size_limit) {
+			return Some(proof_size_limit.saturating_sub(proof_size_usage));
+		}
+		None
+	}
+
+	pub fn remaining_ref_time(&self) -> Option<u64> {
+		if let (Some(ref_time_usage), Some(ref_time_limit)) = (self.ref_time_usage, self.ref_time_limit) {
+			return Some(ref_time_limit.saturating_sub(ref_time_usage));
+		}
+		None
+	}
+
 }
 
 #[derive(Clone, Eq, PartialEq, Debug, Encode, Decode, TypeInfo)]