diff --git a/modules/evm/src/lib.rs b/modules/evm/src/lib.rs index 4322f34a6b..b951070c6e 100644 --- a/modules/evm/src/lib.rs +++ b/modules/evm/src/lib.rs @@ -1198,21 +1198,22 @@ impl Pallet { /// 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::::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::::remove(&contract_info.code_hash); - *maybe_code_info = None; - } + Accounts::::mutate_exists(&address, |maybe_account| { + if let Some(account) = maybe_account { + if let Some(ContractInfo { code_hash, .. }) = account.contract_info { + CodeInfos::::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::::remove(&code_hash); + account.contract_info = None; + *maybe_code_info = None; + } + } + }); } - }); - } + } + }); if let Some(AccountInfo { contract_info: Some(_), .. @@ -1221,7 +1222,7 @@ impl Pallet { // 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 index cb75279d74..c81a1f6f59 100644 --- 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::::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::::insert(&code_hash, BoundedVec::try_from(code).unwrap()); + CodeInfos::::insert( + &code_hash, + CodeInfo { + code_size: 1, + ref_count: 2, + }, + ); + Accounts::::insert( + &address, + AccountInfo { + nonce: 0, + contract_info: Some(ContractInfo { + code_hash, + maintainer: Default::default(), + deployed: false, + }), + }, + ); + let _ = Pallet::::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::::insert(&code_hash, BoundedVec::try_from(code).unwrap()); + CodeInfos::::insert( + &code_hash, + CodeInfo { + code_size: 1, + ref_count: 1, + }, + ); + Accounts::::insert( + &address, + AccountInfo { + nonce: 0, + contract_info: Some(ContractInfo { + code_hash, + maintainer: Default::default(), + deployed: false, + }), + }, + ); + assert_ok!(Pallet::::remove_account(&address)); + assert_eq!(Accounts::::contains_key(&address), false); + assert_eq!(CodeInfos::::contains_key(&code_hash), false); + assert_eq!(Codes::::contains_key(&code_hash), false); + }); +}