From 15bc38dd976e1dd5a52dfb0e1af86e01e0fd6906 Mon Sep 17 00:00:00 2001
From: imstar15 <imstar15@qq.com>
Date: Thu, 9 Nov 2023 17:00:28 +0800
Subject: [PATCH] Make sure that the execution fee is the token of this chain
 or the token of the recipient chain

---
 pallets/automation-time/src/lib.rs | 58 +++++++++++++++++++++++-------
 1 file changed, 46 insertions(+), 12 deletions(-)

diff --git a/pallets/automation-time/src/lib.rs b/pallets/automation-time/src/lib.rs
index 6937feb6..00c3e1c1 100644
--- a/pallets/automation-time/src/lib.rs
+++ b/pallets/automation-time/src/lib.rs
@@ -59,7 +59,10 @@ use frame_support::{
 	weights::constants::WEIGHT_REF_TIME_PER_SECOND,
 };
 use frame_system::pallet_prelude::*;
-use orml_traits::{FixedConversionRateProvider, MultiCurrency};
+use orml_traits::{
+	location::{Parse, Reserve},
+	FixedConversionRateProvider, MultiCurrency,
+};
 use pallet_parachain_staking::DelegatorActions;
 use pallet_timestamp::{self as timestamp};
 pub use pallet_xcmp_handler::InstructionSequence;
@@ -197,6 +200,14 @@ pub mod pallet {
 			BalanceOf<Self>,
 			<Self as frame_system::Config>::RuntimeCall,
 		>;
+
+		/// The way to retreave the reserve of a MultiAsset. This can be
+		/// configured to accept absolute or relative paths for self tokens
+		type ReserveProvider: Reserve;
+
+		/// Self chain location.
+		#[pallet::constant]
+		type SelfLocation: Get<MultiLocation>;
 	}
 
 	const STORAGE_VERSION: StorageVersion = StorageVersion::new(2);
@@ -267,7 +278,10 @@ pub mod pallet {
 		/// The version of the `VersionedMultiLocation` value used is not able
 		/// to be interpreted.
 		BadVersion,
+		UnsupportedFeePayment,
 		CannotReanchor,
+		/// Invalid asset location.
+		InvalidAssetLocation,
 	}
 
 	#[pallet::event]
@@ -374,6 +388,8 @@ pub mod pallet {
 		/// * `DuplicateTask`: There can be no duplicate tasks.
 		/// * `TimeTooFarOut`: Execution time or frequency are past the max time horizon.
 		/// * `TimeSlotFull`: Time slot is full. No more tasks can be scheduled for this time.
+		/// * `UnsupportedFeePayment`: Unsupported fee payment.
+		/// * `InvalidAssetLocation` Invalid asset location.
 		#[pallet::call_index(2)]
 		#[pallet::weight(<T as Config>::WeightInfo::schedule_xcmp_task_full(schedule.number_of_executions()))]
 		pub fn schedule_xcmp_task(
@@ -391,10 +407,18 @@ pub mod pallet {
 				MultiLocation::try_from(*destination).map_err(|()| Error::<T>::BadVersion)?;
 			let schedule_fee =
 				MultiLocation::try_from(*schedule_fee).map_err(|()| Error::<T>::BadVersion)?;
+
+			let execution_fee: AssetPayment = *execution_fee;
+			let execution_fee_location =
+				MultiLocation::try_from(execution_fee.clone().asset_location)
+					.map_err(|()| Error::<T>::BadVersion)?;
+
+			Self::ensure_supported_execution_fee_location(&execution_fee_location, &destination)?;
+
 			let action = Action::XCMP {
 				destination,
 				schedule_fee,
-				execution_fee: *execution_fee,
+				execution_fee,
 				encoded_call,
 				encoded_call_weight,
 				overall_weight,
@@ -1354,14 +1378,11 @@ pub mod pallet {
 			abort_errors: Vec<Vec<u8>>,
 		) -> DispatchResult {
 			match action.clone() {
-				Action::XCMP { execution_fee, instruction_sequence, .. } => {
+				Action::XCMP { execution_fee, .. } => {
 					let asset_location = MultiLocation::try_from(execution_fee.asset_location)
 						.map_err(|()| Error::<T>::BadVersion)?;
-					let asset_location = asset_location
-						.reanchored(
-							&MultiLocation::new(1, X1(Parachain(T::SelfParaId::get().into()))),
-							T::UniversalLocation::get(),
-						)
+					let _asset_location = asset_location
+						.reanchored(&T::SelfLocation::get().into(), T::UniversalLocation::get())
 						.map_err(|_| Error::<T>::CannotReanchor)?;
 				},
 				_ => (),
@@ -1521,10 +1542,7 @@ pub mod pallet {
 
 			let schedule_fee_location = action.schedule_fee_location::<T>();
 			let schedule_fee_location = schedule_fee_location
-				.reanchored(
-					&MultiLocation::new(1, X1(Parachain(T::SelfParaId::get().into()))),
-					T::UniversalLocation::get(),
-				)
+				.reanchored(&T::SelfLocation::get().into(), T::UniversalLocation::get())
 				.map_err(|_| Error::<T>::CannotReanchor)?;
 
 			let fee = if schedule_fee_location == MultiLocation::default() {
@@ -1542,6 +1560,22 @@ pub mod pallet {
 
 			Ok(fee)
 		}
+
+		pub fn ensure_supported_execution_fee_location(
+			exeuction_fee_location: &MultiLocation,
+			destination: &MultiLocation,
+		) -> Result<(), DispatchError> {
+			if exeuction_fee_location.chain_part().is_none() {
+				return Err(Error::<T>::InvalidAssetLocation.into())
+			}
+
+			let self_location = T::SelfLocation::get();
+			if exeuction_fee_location != &self_location && exeuction_fee_location != destination {
+				return Err(Error::<T>::UnsupportedFeePayment.into())
+			}
+
+			Ok(())
+		}
 	}
 
 	impl<T: Config> pallet_valve::Shutdown for Pallet<T> {