1#![cfg_attr(not(feature = "std"), no_std)]
37
38use codec::{Decode, DecodeWithMemTracking, Encode};
39use frame_support::{
40	dispatch::{DispatchInfo, DispatchResult, PostDispatchInfo},
41	pallet_prelude::{TransactionSource, Weight},
42	traits::{
43		tokens::{
44			fungibles::{Balanced, Credit, Inspect},
45			WithdrawConsequence,
46		},
47		IsType,
48	},
49	DefaultNoBound,
50};
51use pallet_transaction_payment::OnChargeTransaction;
52use scale_info::TypeInfo;
53use sp_runtime::{
54	traits::{
55		AsSystemOriginSigner, DispatchInfoOf, Dispatchable, PostDispatchInfoOf, RefundWeight,
56		TransactionExtension, Zero,
57	},
58	transaction_validity::{InvalidTransaction, TransactionValidityError, ValidTransaction},
59};
60
61#[cfg(test)]
62mod mock;
63#[cfg(test)]
64mod tests;
65
66#[cfg(feature = "runtime-benchmarks")]
67mod benchmarking;
68
69mod payment;
70pub mod weights;
71
72pub use payment::*;
73pub use weights::WeightInfo;
74
75pub(crate) type OnChargeTransactionOf<T> =
77	<T as pallet_transaction_payment::Config>::OnChargeTransaction;
78pub(crate) type BalanceOf<T> = <OnChargeTransactionOf<T> as OnChargeTransaction<T>>::Balance;
80pub(crate) type LiquidityInfoOf<T> =
82	<OnChargeTransactionOf<T> as OnChargeTransaction<T>>::LiquidityInfo;
83
84pub(crate) type AssetBalanceOf<T> =
87	<<T as Config>::Fungibles as Inspect<<T as frame_system::Config>::AccountId>>::Balance;
88pub(crate) type AssetIdOf<T> =
90	<<T as Config>::Fungibles as Inspect<<T as frame_system::Config>::AccountId>>::AssetId;
91
92pub(crate) type ChargeAssetBalanceOf<T> =
95	<<T as Config>::OnChargeAssetTransaction as OnChargeAssetTransaction<T>>::Balance;
96pub(crate) type ChargeAssetIdOf<T> =
98	<<T as Config>::OnChargeAssetTransaction as OnChargeAssetTransaction<T>>::AssetId;
99pub(crate) type ChargeAssetLiquidityOf<T> =
101	<<T as Config>::OnChargeAssetTransaction as OnChargeAssetTransaction<T>>::LiquidityInfo;
102
103#[derive(Encode, Decode, DefaultNoBound, TypeInfo)]
105pub enum InitialPayment<T: Config> {
106	#[default]
108	Nothing,
109	Native(LiquidityInfoOf<T>),
111	Asset(Credit<T::AccountId, T::Fungibles>),
113}
114
115pub use pallet::*;
116
117#[frame_support::pallet]
118pub mod pallet {
119	use super::*;
120
121	#[pallet::config]
122	pub trait Config: frame_system::Config + pallet_transaction_payment::Config {
123		#[allow(deprecated)]
125		type RuntimeEvent: From<Event<Self>> + IsType<<Self as frame_system::Config>::RuntimeEvent>;
126		type Fungibles: Balanced<Self::AccountId>;
128		type OnChargeAssetTransaction: OnChargeAssetTransaction<Self>;
130		type WeightInfo: WeightInfo;
132		#[cfg(feature = "runtime-benchmarks")]
134		type BenchmarkHelper: BenchmarkHelperTrait<
135			Self::AccountId,
136			<<Self as Config>::Fungibles as Inspect<Self::AccountId>>::AssetId,
137			<<Self as Config>::OnChargeAssetTransaction as OnChargeAssetTransaction<Self>>::AssetId,
138		>;
139	}
140
141	#[pallet::pallet]
142	pub struct Pallet<T>(_);
143
144	#[cfg(feature = "runtime-benchmarks")]
145	pub trait BenchmarkHelperTrait<AccountId, FunAssetIdParameter, AssetIdParameter> {
147		fn create_asset_id_parameter(id: u32) -> (FunAssetIdParameter, AssetIdParameter);
149		fn setup_balances_and_pool(asset_id: FunAssetIdParameter, account: AccountId);
152	}
153
154	#[pallet::event]
155	#[pallet::generate_deposit(pub(super) fn deposit_event)]
156	pub enum Event<T: Config> {
157		AssetTxFeePaid {
160			who: T::AccountId,
161			actual_fee: AssetBalanceOf<T>,
162			tip: AssetBalanceOf<T>,
163			asset_id: Option<ChargeAssetIdOf<T>>,
164		},
165	}
166}
167
168#[derive(Encode, Decode, DecodeWithMemTracking, Clone, Eq, PartialEq, TypeInfo)]
174#[scale_info(skip_type_params(T))]
175pub struct ChargeAssetTxPayment<T: Config> {
176	#[codec(compact)]
177	tip: BalanceOf<T>,
178	asset_id: Option<ChargeAssetIdOf<T>>,
179}
180
181impl<T: Config> ChargeAssetTxPayment<T>
182where
183	T::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
184	AssetBalanceOf<T>: Send + Sync,
185	BalanceOf<T>: Send + Sync + IsType<ChargeAssetBalanceOf<T>>,
186	ChargeAssetIdOf<T>: Send + Sync,
187	Credit<T::AccountId, T::Fungibles>: IsType<ChargeAssetLiquidityOf<T>>,
188{
189	pub fn from(tip: BalanceOf<T>, asset_id: Option<ChargeAssetIdOf<T>>) -> Self {
191		Self { tip, asset_id }
192	}
193
194	fn withdraw_fee(
197		&self,
198		who: &T::AccountId,
199		call: &T::RuntimeCall,
200		info: &DispatchInfoOf<T::RuntimeCall>,
201		fee: BalanceOf<T>,
202	) -> Result<(BalanceOf<T>, InitialPayment<T>), TransactionValidityError> {
203		debug_assert!(self.tip <= fee, "tip should be included in the computed fee");
204		if fee.is_zero() {
205			Ok((fee, InitialPayment::Nothing))
206		} else if let Some(asset_id) = self.asset_id.clone() {
207			T::OnChargeAssetTransaction::withdraw_fee(
208				who,
209				call,
210				info,
211				asset_id,
212				fee.into(),
213				self.tip.into(),
214			)
215			.map(|i| (fee, InitialPayment::Asset(i.into())))
216		} else {
217			<OnChargeTransactionOf<T> as OnChargeTransaction<T>>::withdraw_fee(
218				who, call, info, fee, self.tip,
219			)
220			.map(|i| (fee, InitialPayment::Native(i)))
221			.map_err(|_| -> TransactionValidityError { InvalidTransaction::Payment.into() })
222		}
223	}
224
225	fn can_withdraw_fee(
228		&self,
229		who: &T::AccountId,
230		call: &T::RuntimeCall,
231		info: &DispatchInfoOf<T::RuntimeCall>,
232		fee: BalanceOf<T>,
233	) -> Result<(), TransactionValidityError> {
234		debug_assert!(self.tip <= fee, "tip should be included in the computed fee");
235		if fee.is_zero() {
236			Ok(())
237		} else if let Some(asset_id) = self.asset_id.clone() {
238			T::OnChargeAssetTransaction::can_withdraw_fee(
239				who,
240				call,
241				info,
242				asset_id,
243				fee.into(),
244				self.tip.into(),
245			)
246		} else {
247			<OnChargeTransactionOf<T> as OnChargeTransaction<T>>::can_withdraw_fee(
248				who, call, info, fee, self.tip,
249			)
250			.map_err(|_| -> TransactionValidityError { InvalidTransaction::Payment.into() })
251		}
252	}
253}
254
255impl<T: Config> core::fmt::Debug for ChargeAssetTxPayment<T> {
256	#[cfg(feature = "std")]
257	fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
258		write!(f, "ChargeAssetTxPayment<{:?}, {:?}>", self.tip, self.asset_id.encode())
259	}
260	#[cfg(not(feature = "std"))]
261	fn fmt(&self, _: &mut core::fmt::Formatter) -> core::fmt::Result {
262		Ok(())
263	}
264}
265
266pub enum Val<T: Config> {
268	Charge {
269		tip: BalanceOf<T>,
270		who: T::AccountId,
272		fee: BalanceOf<T>,
274	},
275	NoCharge,
276}
277
278pub enum Pre<T: Config> {
281	Charge {
282		tip: BalanceOf<T>,
283		who: T::AccountId,
285		initial_payment: InitialPayment<T>,
287		asset_id: Option<ChargeAssetIdOf<T>>,
289		weight: Weight,
291	},
292	NoCharge {
293		refund: Weight,
295	},
296}
297
298impl<T: Config> TransactionExtension<T::RuntimeCall> for ChargeAssetTxPayment<T>
299where
300	T::RuntimeCall: Dispatchable<Info = DispatchInfo, PostInfo = PostDispatchInfo>,
301	AssetBalanceOf<T>: Send + Sync,
302	BalanceOf<T>: Send + Sync + From<u64> + IsType<ChargeAssetBalanceOf<T>>,
303	ChargeAssetIdOf<T>: Send + Sync,
304	Credit<T::AccountId, T::Fungibles>: IsType<ChargeAssetLiquidityOf<T>>,
305	<T::RuntimeCall as Dispatchable>::RuntimeOrigin: AsSystemOriginSigner<T::AccountId> + Clone,
306{
307	const IDENTIFIER: &'static str = "ChargeAssetTxPayment";
308	type Implicit = ();
309	type Val = Val<T>;
310	type Pre = Pre<T>;
311
312	fn weight(&self, _: &T::RuntimeCall) -> Weight {
313		if self.asset_id.is_some() {
314			<T as Config>::WeightInfo::charge_asset_tx_payment_asset()
315		} else {
316			<T as Config>::WeightInfo::charge_asset_tx_payment_native()
317		}
318	}
319
320	fn validate(
321		&self,
322		origin: <T::RuntimeCall as Dispatchable>::RuntimeOrigin,
323		call: &T::RuntimeCall,
324		info: &DispatchInfoOf<T::RuntimeCall>,
325		len: usize,
326		_self_implicit: Self::Implicit,
327		_inherited_implication: &impl Encode,
328		_source: TransactionSource,
329	) -> Result<
330		(ValidTransaction, Self::Val, <T::RuntimeCall as Dispatchable>::RuntimeOrigin),
331		TransactionValidityError,
332	> {
333		use pallet_transaction_payment::ChargeTransactionPayment;
334		let Some(who) = origin.as_system_origin_signer() else {
335			return Ok((ValidTransaction::default(), Val::NoCharge, origin))
336		};
337		let fee = pallet_transaction_payment::Pallet::<T>::compute_fee(len as u32, info, self.tip);
339		self.can_withdraw_fee(&who, call, info, fee)?;
340		let priority = ChargeTransactionPayment::<T>::get_priority(info, len, self.tip, fee);
341		let val = Val::Charge { tip: self.tip, who: who.clone(), fee };
342		let validity = ValidTransaction { priority, ..Default::default() };
343		Ok((validity, val, origin))
344	}
345
346	fn prepare(
347		self,
348		val: Self::Val,
349		_origin: &<T::RuntimeCall as Dispatchable>::RuntimeOrigin,
350		call: &T::RuntimeCall,
351		info: &DispatchInfoOf<T::RuntimeCall>,
352		_len: usize,
353	) -> Result<Self::Pre, TransactionValidityError> {
354		match val {
355			Val::Charge { tip, who, fee } => {
356				let (_fee, initial_payment) = self.withdraw_fee(&who, call, info, fee)?;
358				Ok(Pre::Charge {
359					tip,
360					who,
361					initial_payment,
362					asset_id: self.asset_id.clone(),
363					weight: self.weight(call),
364				})
365			},
366			Val::NoCharge => Ok(Pre::NoCharge { refund: self.weight(call) }),
367		}
368	}
369
370	fn post_dispatch_details(
371		pre: Self::Pre,
372		info: &DispatchInfoOf<T::RuntimeCall>,
373		post_info: &PostDispatchInfoOf<T::RuntimeCall>,
374		len: usize,
375		result: &DispatchResult,
376	) -> Result<Weight, TransactionValidityError> {
377		let (tip, who, initial_payment, asset_id, extension_weight) = match pre {
378			Pre::Charge { tip, who, initial_payment, asset_id, weight } =>
379				(tip, who, initial_payment, asset_id, weight),
380			Pre::NoCharge { refund } => {
381				return Ok(refund)
383			},
384		};
385
386		match initial_payment {
387			InitialPayment::Native(already_withdrawn) => {
388				let actual_ext_weight = <T as Config>::WeightInfo::charge_asset_tx_payment_native();
391				let unspent_weight = extension_weight.saturating_sub(actual_ext_weight);
392				let mut actual_post_info = *post_info;
393				actual_post_info.refund(unspent_weight);
394				pallet_transaction_payment::ChargeTransactionPayment::<T>::post_dispatch_details(
395					pallet_transaction_payment::Pre::Charge {
396						tip,
397						who,
398						imbalance: already_withdrawn,
399					},
400					info,
401					&actual_post_info,
402					len,
403					result,
404				)?;
405				Ok(unspent_weight)
406			},
407			InitialPayment::Asset(already_withdrawn) => {
408				let actual_ext_weight = <T as Config>::WeightInfo::charge_asset_tx_payment_asset();
409				let unspent_weight = extension_weight.saturating_sub(actual_ext_weight);
410				let mut actual_post_info = *post_info;
411				actual_post_info.refund(unspent_weight);
412				let actual_fee = pallet_transaction_payment::Pallet::<T>::compute_actual_fee(
413					len as u32,
414					info,
415					&actual_post_info,
416					tip,
417				);
418
419				let (converted_fee, converted_tip) =
420					T::OnChargeAssetTransaction::correct_and_deposit_fee(
421						&who,
422						info,
423						&actual_post_info,
424						actual_fee.into(),
425						tip.into(),
426						already_withdrawn.into(),
427					)?;
428				Pallet::<T>::deposit_event(Event::<T>::AssetTxFeePaid {
429					who,
430					actual_fee: converted_fee,
431					tip: converted_tip,
432					asset_id,
433				});
434				Ok(unspent_weight)
435			},
436			InitialPayment::Nothing => {
437				debug_assert!(tip.is_zero(), "tip should be zero if initial fee was zero.");
442				Ok(extension_weight
443					.saturating_sub(<T as Config>::WeightInfo::charge_asset_tx_payment_zero()))
444			},
445		}
446	}
447}