Fix add size_limit to the proof_size of weight info

darwinia-network · Nov 2, 2023 · d65d147 · d65d147
1 parent 9e9c31b
commit d65d147
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 12 deletions.
diff --git 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
@@ -108,13 +108,15 @@ 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 {
 			return Err(ExitError::OutOfGas);
 		}
 		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)]