From 76921f18a0c593ff05716f6041c22cabd5cbdaef Mon Sep 17 00:00:00 2001
From: Ahmad Kaouk <ahmadkaouk.93@gmail.com>
Date: Thu, 2 Nov 2023 07:54:03 +0100
Subject: [PATCH] add a limit

---
 frame/evm/precompile/clear-storage/src/lib.rs | 11 ++-
 .../evm/precompile/clear-storage/src/tests.rs | 84 ++++++++++++++++++-
 2 files changed, 91 insertions(+), 4 deletions(-)

diff --git a/frame/evm/precompile/clear-storage/src/lib.rs b/frame/evm/precompile/clear-storage/src/lib.rs
index 98a24fdabc..371cf9f487 100644
--- a/frame/evm/precompile/clear-storage/src/lib.rs
+++ b/frame/evm/precompile/clear-storage/src/lib.rs
@@ -43,12 +43,14 @@ impl<Runtime> StorageCleanerPrecompile<Runtime>
 where
 	Runtime: pallet_evm::Config,
 {
-	#[precompile::public("clearSuicidedStorage(address[])")]
+	#[precompile::public("clearSuicidedStorage(address[],uint32)")]
 	fn clear_suicided_storage(
 		handle: &mut impl PrecompileHandle,
 		addresses: BoundedVec<Address, GetArrayLimit>,
+		limit: u32,
 	) -> EvmResult {
 		let addresses: Vec<_> = addresses.into();
+		let mut deleted_entries = 0;
 
 		for address in addresses {
 			// Read Suicided storage item
@@ -58,8 +60,7 @@ where
 				return Err(revert(format!("NotSuicided: {}", address.0)));
 			}
 
-			let mut iter = pallet_evm::Pallet::<Runtime>::iter_account_storages(&address.0).drain();
-
+			let mut iter = pallet_evm::AccountStorages::<Runtime>::drain_prefix(&address.0);
 			loop {
 				handle.record_db_read::<Runtime>(116)?;
 				// Record the gas cost of deleting the storage item
@@ -70,6 +71,10 @@ where
 					Self::clear_suicided_contract(address);
 					break;
 				}
+				deleted_entries += 1;
+				if deleted_entries >= limit {
+					return Ok(());
+				}
 			}
 		}
 
diff --git a/frame/evm/precompile/clear-storage/src/tests.rs b/frame/evm/precompile/clear-storage/src/tests.rs
index f7bfa8b857..9f23aab4cd 100644
--- a/frame/evm/precompile/clear-storage/src/tests.rs
+++ b/frame/evm/precompile/clear-storage/src/tests.rs
@@ -62,13 +62,13 @@ fn test_clear_suicided_contract_succesfull() {
 			let contract_address = mock_contract_with_entries(1, 10);
 			// Add contract to the suicided contracts
 			pallet_evm::Suicided::<Runtime>::insert(contract_address, ());
-
 			precompiles()
 				.prepare_test(
 					Alice,
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: vec![contract_address.into()].into(),
+						limit: u32::MAX,
 					},
 				)
 				.execute_returns(());
@@ -106,6 +106,7 @@ fn test_clear_suicided_contract_failed() {
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: vec![contract_address.into()].into(),
+						limit: u32::MAX,
 					},
 				)
 				.execute_reverts(|output| {
@@ -136,6 +137,7 @@ fn test_clear_suicided_empty_input() {
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: vec![].into(),
+						limit: u32::MAX,
 					},
 				)
 				.execute_returns(());
@@ -178,6 +180,7 @@ fn test_clear_suicided_contract_multiple_addresses() {
 							contract_address3.into(),
 						]
 						.into(),
+						limit: u32::MAX,
 					},
 				)
 				.execute_returns(());
@@ -239,6 +242,7 @@ fn test_clear_suicided_mixed_suicided_and_non_suicided() {
 							contract_address4.into(),
 						]
 						.into(),
+						limit: u32::MAX,
 					},
 				)
 				.execute_reverts(|output| {
@@ -273,6 +277,7 @@ fn test_clear_suicided_no_storage_entries() {
 					Precompile1,
 					PCall::clear_suicided_storage {
 						addresses: vec![contract_address1.into()].into(),
+						limit: u32::MAX,
 					},
 				)
 				.execute_returns(());
@@ -327,3 +332,80 @@ fn test_clear_suicided_no_storage_entries() {
 			);
 		})
 }
+
+// Test that the precompile deletes a maximum of `ENTRY_LIMIT` entries
+#[test]
+fn test_clear_suicided_entry_limit() {
+	ExtBuilder::default()
+		.with_balances(vec![(Alice.into(), 10000000000000000000)])
+		.build()
+		.execute_with(|| {
+			let contract_address1 = mock_contract_with_entries(1, 4);
+			let contract_address2 = mock_contract_with_entries(2, 3);
+			// Add contract to the suicided contracts
+			pallet_evm::Suicided::<Runtime>::insert(contract_address1, ());
+			pallet_evm::Suicided::<Runtime>::insert(contract_address2, ());
+
+			precompiles()
+				.prepare_test(
+					Alice,
+					Precompile1,
+					PCall::clear_suicided_storage {
+						addresses: vec![contract_address1.into(), contract_address2.into()].into(),
+						limit: 5,
+					},
+				)
+				.execute_returns(());
+
+			assert_eq!(
+				pallet_evm::AccountStorages::<Runtime>::iter_prefix(contract_address1).count(),
+				0
+			);
+			assert_eq!(
+				pallet_evm::Suicided::<Runtime>::contains_key(contract_address1),
+				false
+			);
+
+			assert_eq!(
+				pallet_evm::AccountStorages::<Runtime>::iter_prefix(contract_address2).count(),
+				2
+			);
+
+			assert_eq!(
+				pallet_evm::Suicided::<Runtime>::contains_key(contract_address2),
+				true
+			);
+		})
+}
+
+#[test]
+fn test_clear_suicided_entry_limit_1() {
+	ExtBuilder::default()
+		.with_balances(vec![(Alice.into(), 10000000000000000000)])
+		.build()
+		.execute_with(|| {
+			let contract_address1 = mock_contract_with_entries(1, 5);
+			// Add contract to the suicided contracts
+			pallet_evm::Suicided::<Runtime>::insert(contract_address1, ());
+
+			precompiles()
+				.prepare_test(
+					Alice,
+					Precompile1,
+					PCall::clear_suicided_storage {
+						addresses: vec![contract_address1.into()].into(),
+						limit: 4,
+					},
+				)
+				.execute_returns(());
+
+			assert_eq!(
+				pallet_evm::AccountStorages::<Runtime>::iter_prefix(contract_address1).count(),
+				1
+			);
+			assert_eq!(
+				pallet_evm::Suicided::<Runtime>::contains_key(contract_address1),
+				true
+			);
+		})
+}