Add default account mapping for Wasm contracts.

AstarNetwork · Oct 12, 2023 · 1a6f4e2 · 1a6f4e2
1 parent 3a5d9fd
commit 1a6f4e2
Show file tree

Hide file tree

Showing 7 changed files with 30 additions and 20 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,8 +21,12 @@
 extern crate alloc;
 use alloc::format;
 
-use astar_primitives::xvm::{Context, VmId, XvmCall};
+use astar_primitives::{
+    evm::UnifiedAddressMapper,
+    xvm::{Context, VmId, XvmCall},
+};
 use frame_support::dispatch::Encode;
+use frame_system::RawOrigin;
 use pallet_contracts::chain_extension::{
     ChainExtension, Environment, Ext, InitState, RetVal, ReturnFlags,
 };
@@ -48,18 +52,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
@@ -89,6 +94,20 @@ where
                 // things on behalf of the caller via XVM.
                 let source = env.ext().address();
 
+                // Claim a default account if needed.
+                if UA::to_h160(&source).is_none() {
+                    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,
                     weight_limit,

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")