fix remove contract (#1739)

* fix remove contract * avoid clone * add assert_debug msg
AcalaNetwork · Jan 17, 2022 · 2a49dbc · 2a49dbc
1 parent 96dd094
commit 2a49dbc
Show file tree

Hide file tree

Showing 2 changed files with 86 additions and 15 deletions.
diff --git a/modules/evm/src/lib.rs b/modules/evm/src/lib.rs
@@ -1198,21 +1198,22 @@ impl<T: Config> Pallet<T> {
 	/// Removes an account from Accounts and AccountStorages.
 	pub fn remove_account(address: &EvmAddress) -> DispatchResult {
 		// Deref code, and remove it if ref count is zero.
-		if let Some(AccountInfo {
-			contract_info: Some(contract_info),
-			..
-		}) = Self::accounts(address)
-		{
-			CodeInfos::<T>::mutate_exists(&contract_info.code_hash, |maybe_code_info| {
-				if let Some(code_info) = maybe_code_info.as_mut() {
-					code_info.ref_count = code_info.ref_count.saturating_sub(1);
-					if code_info.ref_count == 0 {
-						Codes::<T>::remove(&contract_info.code_hash);
-						*maybe_code_info = None;
-					}
+		Accounts::<T>::mutate_exists(&address, |maybe_account| {
+			if let Some(account) = maybe_account {
+				if let Some(ContractInfo { code_hash, .. }) = account.contract_info {
+					CodeInfos::<T>::mutate_exists(&code_hash, |maybe_code_info| {
+						if let Some(code_info) = maybe_code_info {
+							code_info.ref_count = code_info.ref_count.saturating_sub(1);
+							if code_info.ref_count == 0 {
+								Codes::<T>::remove(&code_hash);
+								account.contract_info = None;
+								*maybe_code_info = None;
+							}
+						}
+					});
 				}
-			});
-		}
+			}
+		});
 
 		if let Some(AccountInfo {
 			contract_info: Some(_), ..
@@ -1221,7 +1222,7 @@ impl<T: Config> Pallet<T> {
 			// remove_account can only be called when account is killed. i.e. providers == 0
 			// but contract_info should maintain a provider
 			// so this should never happen
-			debug_assert!(false);
+			debug_assert!(false, "removed account while is still linked to contract info");
 		}
 
 		Ok(())

diff --git a/modules/evm/src/tests.rs b/modules/evm/src/tests.rs
@@ -1885,3 +1885,73 @@ fn convert_decimals_should_not_work() {
 		}));
 	});
 }
+
+#[test]
+fn remove_empty_account_works() {
+	new_test_ext().execute_with(|| {
+		let address = H160::from([1; 20]);
+		assert_ok!(Pallet::<Runtime>::remove_account(&address));
+	});
+}
+
+#[test]
+#[should_panic(expected = "removed account while is still linked to contract info")]
+fn remove_account_with_provides_should_panic() {
+	new_test_ext().execute_with(|| {
+		let address = H160::from([1; 20]);
+		let code = vec![0x00];
+		let code_hash = code_hash(&code);
+		Codes::<Runtime>::insert(&code_hash, BoundedVec::try_from(code).unwrap());
+		CodeInfos::<Runtime>::insert(
+			&code_hash,
+			CodeInfo {
+				code_size: 1,
+				ref_count: 2,
+			},
+		);
+		Accounts::<Runtime>::insert(
+			&address,
+			AccountInfo {
+				nonce: 0,
+				contract_info: Some(ContractInfo {
+					code_hash,
+					maintainer: Default::default(),
+					deployed: false,
+				}),
+			},
+		);
+		let _ = Pallet::<Runtime>::remove_account(&address);
+	});
+}
+
+#[test]
+fn remove_account_works() {
+	new_test_ext().execute_with(|| {
+		let address = H160::from([1; 20]);
+		let code = vec![0x00];
+		let code_hash = code_hash(&code);
+		Codes::<Runtime>::insert(&code_hash, BoundedVec::try_from(code).unwrap());
+		CodeInfos::<Runtime>::insert(
+			&code_hash,
+			CodeInfo {
+				code_size: 1,
+				ref_count: 1,
+			},
+		);
+		Accounts::<Runtime>::insert(
+			&address,
+			AccountInfo {
+				nonce: 0,
+				contract_info: Some(ContractInfo {
+					code_hash,
+					maintainer: Default::default(),
+					deployed: false,
+				}),
+			},
+		);
+		assert_ok!(Pallet::<Runtime>::remove_account(&address));
+		assert_eq!(Accounts::<Runtime>::contains_key(&address), false);
+		assert_eq!(CodeInfos::<Runtime>::contains_key(&code_hash), false);
+		assert_eq!(Codes::<Runtime>::contains_key(&code_hash), false);
+	});
+}