XVM: default account mapping for Wasm contracts (#1056)

* Add default account mapping for Wasm contracts. * Only map accounts for Wasm if payable call. * Charge weight for claim default address call. * Fix & docs update. * To H160 DB read.
AstarNetwork · Oct 16, 2023 · 25eec60 · 25eec60
1 parent 58780ad
commit 25eec60
Show file tree

Hide file tree

Showing 9 changed files with 48 additions and 28 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/chain-extensions/xvm/Cargo.toml b/chain-extensions/xvm/Cargo.toml
@@ -15,6 +15,7 @@ log = { workspace = true }
 num-traits = { workspace = true }
 pallet-contracts = { workspace = true }
 pallet-contracts-primitives = { workspace = true }
+pallet-unified-accounts = { workspace = true }
 parity-scale-codec = { workspace = true }
 scale-info = { workspace = true }
 sp-core = { workspace = true }
@@ -34,6 +35,7 @@ std = [
 	"num-traits/std",
 	"pallet-contracts/std",
 	"pallet-contracts-primitives/std",
+	"pallet-unified-accounts/std",
 	"scale-info/std",
 	"sp-std/std",
 	"sp-core/std",

diff --git a/chain-extensions/xvm/src/lib.rs b/chain-extensions/xvm/src/lib.rs
@@ -21,11 +21,16 @@
 extern crate alloc;
 use alloc::format;
 
-use astar_primitives::xvm::{Context, VmId, XvmCall};
-use frame_support::dispatch::Encode;
+use astar_primitives::{
+    evm::UnifiedAddressMapper,
+    xvm::{Context, VmId, XvmCall},
+};
+use frame_support::{dispatch::Encode, traits::Get, weights::Weight};
+use frame_system::RawOrigin;
 use pallet_contracts::chain_extension::{
     ChainExtension, Environment, Ext, InitState, RetVal, ReturnFlags,
 };
+use pallet_unified_accounts::WeightInfo;
 use sp_runtime::DispatchError;
 use sp_std::marker::PhantomData;
 use xvm_chain_extension_types::{XvmCallArgs, XvmExecutionResult};
@@ -48,18 +53,19 @@ impl TryFrom<u16> for XvmFuncId {
 }
 
 /// XVM chain extension.
-pub struct XvmExtension<T, XC>(PhantomData<(T, XC)>);
+pub struct XvmExtension<T, XC, UA>(PhantomData<(T, XC, UA)>);
 
-impl<T, XC> Default for XvmExtension<T, XC> {
+impl<T, XC, UA> Default for XvmExtension<T, XC, UA> {
     fn default() -> Self {
         XvmExtension(PhantomData)
     }
 }
 
-impl<T, XC> ChainExtension<T> for XvmExtension<T, XC>
+impl<T, XC, UA> ChainExtension<T> for XvmExtension<T, XC, UA>
 where
-    T: pallet_contracts::Config,
+    T: pallet_contracts::Config + pallet_unified_accounts::Config,
     XC: XvmCall<T::AccountId>,
+    UA: UnifiedAddressMapper<T::AccountId>,
 {
     fn call<E: Ext>(&mut self, env: Environment<E, InitState>) -> Result<RetVal, DispatchError>
     where
@@ -87,7 +93,30 @@ where
                 // Similar to EVM behavior, the `source` should be (limited to) the
                 // contract address. Otherwise contracts would be able to do arbitrary
                 // things on behalf of the caller via XVM.
-                let source = env.ext().address();
+                let source = env.ext().address().clone();
+
+                // Claim the default evm address if needed.
+                let mut actual_weight = Weight::zero();
+                if value > 0 {
+                    // `UA::to_h160` 1 DB read.
+                    actual_weight.saturating_accrue(T::DbWeight::get().reads(1));
+
+                    if UA::to_h160(&source).is_none() {
+                        let weight_of_claim = <T as pallet_unified_accounts::Config>::WeightInfo::claim_default_evm_address();
+                        actual_weight.saturating_accrue(weight_of_claim);
+
+                        let claim_result =
+                            pallet_unified_accounts::Pallet::<T>::claim_default_evm_address(
+                                RawOrigin::Signed(source.clone()).into(),
+                            );
+                        if claim_result.is_err() {
+                            return Ok(RetVal::Diverging {
+                                flags: ReturnFlags::REVERT,
+                                data: format!("{:?}", claim_result.err()).into(),
+                            });
+                        }
+                    }
+                }
 
                 let xvm_context = Context {
                     source_vm_id: VmId::Wasm,
@@ -104,13 +133,13 @@ where
                         }
                     }
                 };
-                let call_result =
-                    XC::call(xvm_context, vm_id, source.clone(), to, input, value, None);
+                let call_result = XC::call(xvm_context, vm_id, source, to, input, value, None);
 
-                let actual_weight = match call_result {
+                let used_weight = match call_result {
                     Ok(ref info) => info.used_weight,
                     Err(ref err) => err.used_weight,
                 };
+                actual_weight.saturating_accrue(used_weight);
                 env.adjust_weight(charged_weight, actual_weight);
 
                 match call_result {

diff --git a/runtime/local/src/chain_extensions.rs b/runtime/local/src/chain_extensions.rs
@@ -16,7 +16,7 @@
 // You should have received a copy of the GNU General Public License
 // along with Astar. If not, see <http://www.gnu.org/licenses/>.
 
-use super::{Runtime, Xvm};
+use super::{Runtime, UnifiedAccounts, Xvm};
 
 /// Registered WASM contracts chain extensions.
 pub use pallet_chain_extension_assets::AssetsExtension;
@@ -31,7 +31,7 @@ impl RegisteredChainExtension<Runtime> for DappsStakingExtension<Runtime> {
     const ID: u16 = 00;
 }
 
-impl RegisteredChainExtension<Runtime> for XvmExtension<Runtime, Xvm> {
+impl RegisteredChainExtension<Runtime> for XvmExtension<Runtime, Xvm, UnifiedAccounts> {
     const ID: u16 = 01;
 }
 

diff --git a/runtime/local/src/lib.rs b/runtime/local/src/lib.rs
@@ -844,7 +844,7 @@ impl pallet_contracts::Config for Runtime {
     type WeightInfo = pallet_contracts::weights::SubstrateWeight<Self>;
     type ChainExtension = (
         DappsStakingExtension<Self>,
-        XvmExtension<Self, Xvm>,
+        XvmExtension<Self, Xvm, UnifiedAccounts>,
         AssetsExtension<Self, pallet_chain_extension_assets::weights::SubstrateWeight<Self>>,
     );
     type Schedule = Schedule;

diff --git a/runtime/shibuya/src/chain_extensions.rs b/runtime/shibuya/src/chain_extensions.rs
@@ -16,7 +16,7 @@
 // You should have received a copy of the GNU General Public License
 // along with Astar. If not, see <http://www.gnu.org/licenses/>.
 
-use super::{Runtime, Xvm};
+use super::{Runtime, UnifiedAccounts, Xvm};
 
 /// Registered WASM contracts chain extensions.
 pub use pallet_chain_extension_assets::AssetsExtension;
@@ -31,7 +31,7 @@ impl RegisteredChainExtension<Runtime> for DappsStakingExtension<Runtime> {
     const ID: u16 = 00;
 }
 
-impl RegisteredChainExtension<Runtime> for XvmExtension<Runtime, Xvm> {
+impl RegisteredChainExtension<Runtime> for XvmExtension<Runtime, Xvm, UnifiedAccounts> {
     const ID: u16 = 01;
 }
 

diff --git a/runtime/shibuya/src/lib.rs b/runtime/shibuya/src/lib.rs
@@ -666,7 +666,7 @@ impl pallet_contracts::Config for Runtime {
     type WeightInfo = pallet_contracts::weights::SubstrateWeight<Self>;
     type ChainExtension = (
         DappsStakingExtension<Self>,
-        XvmExtension<Self, Xvm>,
+        XvmExtension<Self, Xvm, UnifiedAccounts>,
         AssetsExtension<Self, pallet_chain_extension_assets::weights::SubstrateWeight<Self>>,
     );
     type Schedule = Schedule;

diff --git a/tests/integration/src/setup.rs b/tests/integration/src/setup.rs
@@ -119,14 +119,6 @@ mod shibuya {
             get_evm_signature(who, secret)
         ));
     }
-
-    pub fn claim_default_accounts(account: AccountId) {
-        let default_h160 = UnifiedAccounts::to_default_h160(&account);
-        assert_ok!(UnifiedAccounts::claim_default_evm_address(
-            RuntimeOrigin::signed(account.clone())
-        ));
-        assert_eq!(UnifiedAccounts::to_h160(&account).unwrap(), default_h160);
-    }
 }
 
 #[cfg(feature = "shiden")]

diff --git a/tests/integration/src/xvm.rs b/tests/integration/src/xvm.rs
@@ -330,10 +330,6 @@ fn calling_evm_payable_from_wasm_works() {
         let wasm_caller_addr = deploy_wasm_contract(CALL_EVM_PAYBLE_NAME);
 
         let value = UNIT;
-
-        // claim the default mappings for wasm contract
-        claim_default_accounts(wasm_caller_addr.clone());
-
         let evm_payable = evm_payable_callee_addr.as_ref().to_vec();
         let deposit_func = hex::decode("d0e30db0").expect("invalid deposit function hex");
         let input = hex::decode("0000002a")