1extern crate alloc;
19
20use alloc::{vec, vec::Vec};
21
22use crate::Config;
23use codec::{Decode, DecodeWithMemTracking, Encode};
24use frame_support::{
25 dispatch::DispatchInfo, pallet_prelude::TransactionSource, RuntimeDebugNoBound,
26};
27use scale_info::TypeInfo;
28use sp_runtime::{
29 traits::{
30 AsSystemOriginSigner, CheckedAdd, DispatchInfoOf, Dispatchable, One, PostDispatchInfoOf,
31 TransactionExtension, ValidateResult, Zero,
32 },
33 transaction_validity::{
34 InvalidTransaction, TransactionLongevity, TransactionValidityError, ValidTransaction,
35 },
36 DispatchResult, Saturating,
37};
38use sp_weights::Weight;
39
40#[derive(Encode, Decode, DecodeWithMemTracking, Clone, Eq, PartialEq, TypeInfo)]
53#[scale_info(skip_type_params(T))]
54pub struct CheckNonce<T: Config>(#[codec(compact)] pub T::Nonce);
55
56pub struct ValidNonceInfo {
58 pub provides: Vec<Vec<u8>>,
60 pub requires: Vec<Vec<u8>>,
62}
63
64impl<T: Config> CheckNonce<T> {
65 pub fn from(nonce: T::Nonce) -> Self {
67 Self(nonce)
68 }
69
70 pub fn validate_nonce_for_account(
73 who: &T::AccountId,
74 nonce: T::Nonce,
75 ) -> Result<ValidNonceInfo, TransactionValidityError> {
76 let account = crate::Account::<T>::get(who);
77 if account.providers.is_zero() && account.sufficients.is_zero() {
78 return Err(InvalidTransaction::Payment.into())
80 }
81 if nonce < account.nonce {
82 return Err(InvalidTransaction::Stale.into())
83 }
84
85 let provides = vec![Encode::encode(&(who.clone(), nonce))];
86 let requires = if account.nonce < nonce {
87 vec![Encode::encode(&(who.clone(), nonce.saturating_sub(One::one())))]
88 } else {
89 vec![]
90 };
91
92 Ok(ValidNonceInfo { provides, requires })
93 }
94
95 pub fn prepare_nonce_for_account(
97 who: &T::AccountId,
98 mut nonce: T::Nonce,
99 ) -> Result<(), TransactionValidityError> {
100 let account = crate::Account::<T>::get(who);
101 if nonce > account.nonce {
102 return Err(InvalidTransaction::Future.into())
103 }
104 nonce = nonce.checked_add(&T::Nonce::one()).unwrap_or(T::Nonce::zero());
105 crate::Account::<T>::mutate(who, |account| account.nonce = nonce);
106 Ok(())
107 }
108}
109
110impl<T: Config> core::fmt::Debug for CheckNonce<T> {
111 #[cfg(feature = "std")]
112 fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
113 write!(f, "CheckNonce({})", self.0)
114 }
115
116 #[cfg(not(feature = "std"))]
117 fn fmt(&self, _: &mut core::fmt::Formatter) -> core::fmt::Result {
118 Ok(())
119 }
120}
121
122#[derive(RuntimeDebugNoBound)]
124pub enum Val<T: Config> {
125 CheckNonce(T::AccountId),
127 Refund(Weight),
129}
130
131#[derive(RuntimeDebugNoBound)]
134pub enum Pre {
135 NonceChecked,
137 Refund(Weight),
139}
140
141impl<T: Config> TransactionExtension<T::RuntimeCall> for CheckNonce<T>
142where
143 T::RuntimeCall: Dispatchable<Info = DispatchInfo>,
144 <T::RuntimeCall as Dispatchable>::RuntimeOrigin: AsSystemOriginSigner<T::AccountId> + Clone,
145{
146 const IDENTIFIER: &'static str = "CheckNonce";
147 type Implicit = ();
148 type Val = Val<T>;
149 type Pre = Pre;
150
151 fn weight(&self, _: &T::RuntimeCall) -> sp_weights::Weight {
152 <T::ExtensionsWeightInfo as super::WeightInfo>::check_nonce()
153 }
154
155 fn validate(
156 &self,
157 origin: <T as Config>::RuntimeOrigin,
158 call: &T::RuntimeCall,
159 _info: &DispatchInfoOf<T::RuntimeCall>,
160 _len: usize,
161 _self_implicit: Self::Implicit,
162 _inherited_implication: &impl Encode,
163 _source: TransactionSource,
164 ) -> ValidateResult<Self::Val, T::RuntimeCall> {
165 let Some(who) = origin.as_system_origin_signer() else {
166 return Ok((Default::default(), Val::Refund(self.weight(call)), origin))
167 };
168 let ValidNonceInfo { provides, requires } = Self::validate_nonce_for_account(who, self.0)?;
169
170 let validity = ValidTransaction {
171 priority: 0,
172 requires,
173 provides,
174 longevity: TransactionLongevity::max_value(),
175 propagate: true,
176 };
177
178 Ok((validity, Val::CheckNonce(who.clone()), origin))
179 }
180
181 fn prepare(
182 self,
183 val: Self::Val,
184 _origin: &T::RuntimeOrigin,
185 _call: &T::RuntimeCall,
186 _info: &DispatchInfoOf<T::RuntimeCall>,
187 _len: usize,
188 ) -> Result<Self::Pre, TransactionValidityError> {
189 let (who, nonce) = match val {
190 Val::CheckNonce(who) => (who, self.0),
191 Val::Refund(weight) => return Ok(Pre::Refund(weight)),
192 };
193 Self::prepare_nonce_for_account(&who, nonce).map(|_| Pre::NonceChecked)
194 }
195
196 fn post_dispatch_details(
197 pre: Self::Pre,
198 _info: &DispatchInfo,
199 _post_info: &PostDispatchInfoOf<T::RuntimeCall>,
200 _len: usize,
201 _result: &DispatchResult,
202 ) -> Result<Weight, TransactionValidityError> {
203 match pre {
204 Pre::NonceChecked => Ok(Weight::zero()),
205 Pre::Refund(weight) => Ok(weight),
206 }
207 }
208}
209
210#[cfg(test)]
211mod tests {
212 use super::*;
213 use crate::mock::{new_test_ext, RuntimeCall, Test, CALL};
214 use frame_support::{
215 assert_ok, assert_storage_noop, dispatch::GetDispatchInfo, traits::OriginTrait,
216 };
217 use sp_runtime::{
218 traits::{AsTransactionAuthorizedOrigin, DispatchTransaction, TxBaseImplication},
219 transaction_validity::TransactionSource::External,
220 };
221
222 #[test]
223 fn signed_ext_check_nonce_works() {
224 new_test_ext().execute_with(|| {
225 crate::Account::<Test>::insert(
226 1,
227 crate::AccountInfo {
228 nonce: 1u64.into(),
229 consumers: 0,
230 providers: 1,
231 sufficients: 0,
232 data: 0,
233 },
234 );
235 let info = DispatchInfo::default();
236 let len = 0_usize;
237 assert_storage_noop!({
239 assert_eq!(
240 CheckNonce::<Test>(0u64.into())
241 .validate_only(Some(1).into(), CALL, &info, len, External, 0)
242 .unwrap_err(),
243 TransactionValidityError::Invalid(InvalidTransaction::Stale)
244 );
245 assert_eq!(
246 CheckNonce::<Test>(0u64.into())
247 .validate_and_prepare(Some(1).into(), CALL, &info, len, 0)
248 .unwrap_err(),
249 TransactionValidityError::Invalid(InvalidTransaction::Stale)
250 );
251 });
252 assert_ok!(CheckNonce::<Test>(1u64.into()).validate_only(
254 Some(1).into(),
255 CALL,
256 &info,
257 len,
258 External,
259 0,
260 ));
261 assert_ok!(CheckNonce::<Test>(1u64.into()).validate_and_prepare(
262 Some(1).into(),
263 CALL,
264 &info,
265 len,
266 0,
267 ));
268 assert_ok!(CheckNonce::<Test>(5u64.into()).validate_only(
270 Some(1).into(),
271 CALL,
272 &info,
273 len,
274 External,
275 0,
276 ));
277 assert_eq!(
278 CheckNonce::<Test>(5u64.into())
279 .validate_and_prepare(Some(1).into(), CALL, &info, len, 0)
280 .unwrap_err(),
281 TransactionValidityError::Invalid(InvalidTransaction::Future)
282 );
283 })
284 }
285
286 #[test]
287 fn signed_ext_check_nonce_requires_provider() {
288 new_test_ext().execute_with(|| {
289 crate::Account::<Test>::insert(
290 2,
291 crate::AccountInfo {
292 nonce: 1u64.into(),
293 consumers: 0,
294 providers: 1,
295 sufficients: 0,
296 data: 0,
297 },
298 );
299 crate::Account::<Test>::insert(
300 3,
301 crate::AccountInfo {
302 nonce: 1u64.into(),
303 consumers: 0,
304 providers: 0,
305 sufficients: 1,
306 data: 0,
307 },
308 );
309 let info = DispatchInfo::default();
310 let len = 0_usize;
311 assert_storage_noop!({
313 assert_eq!(
314 CheckNonce::<Test>(1u64.into())
315 .validate_only(Some(1).into(), CALL, &info, len, External, 0)
316 .unwrap_err(),
317 TransactionValidityError::Invalid(InvalidTransaction::Payment)
318 );
319 assert_eq!(
320 CheckNonce::<Test>(1u64.into())
321 .validate_and_prepare(Some(1).into(), CALL, &info, len, 0)
322 .unwrap_err(),
323 TransactionValidityError::Invalid(InvalidTransaction::Payment)
324 );
325 });
326 assert_ok!(CheckNonce::<Test>(1u64.into()).validate_only(
328 Some(2).into(),
329 CALL,
330 &info,
331 len,
332 External,
333 0,
334 ));
335 assert_ok!(CheckNonce::<Test>(1u64.into()).validate_and_prepare(
336 Some(2).into(),
337 CALL,
338 &info,
339 len,
340 0,
341 ));
342 assert_ok!(CheckNonce::<Test>(1u64.into()).validate_only(
344 Some(3).into(),
345 CALL,
346 &info,
347 len,
348 External,
349 0,
350 ));
351 assert_ok!(CheckNonce::<Test>(1u64.into()).validate_and_prepare(
352 Some(3).into(),
353 CALL,
354 &info,
355 len,
356 0,
357 ));
358 })
359 }
360
361 #[test]
362 fn unsigned_check_nonce_works() {
363 new_test_ext().execute_with(|| {
364 let info = DispatchInfo::default();
365 let len = 0_usize;
366 let (_, val, origin) = CheckNonce::<Test>(1u64.into())
367 .validate(None.into(), CALL, &info, len, (), &TxBaseImplication(CALL), External)
368 .unwrap();
369 assert!(!origin.is_transaction_authorized());
370 assert_ok!(CheckNonce::<Test>(1u64.into()).prepare(val, &origin, CALL, &info, len));
371 })
372 }
373
374 #[test]
375 fn check_nonce_preserves_account_data() {
376 new_test_ext().execute_with(|| {
377 crate::Account::<Test>::insert(
378 1,
379 crate::AccountInfo {
380 nonce: 1u64.into(),
381 consumers: 0,
382 providers: 1,
383 sufficients: 0,
384 data: 0,
385 },
386 );
387 let info = DispatchInfo::default();
388 let len = 0_usize;
389 let (_, val, origin) = CheckNonce::<Test>(1u64.into())
391 .validate(Some(1).into(), CALL, &info, len, (), &TxBaseImplication(CALL), External)
392 .unwrap();
393 crate::Account::<Test>::mutate(1, |info| {
395 info.data = 42;
396 });
397 assert_ok!(CheckNonce::<Test>(1u64.into()).prepare(val, &origin, CALL, &info, len));
399 let expected_info = crate::AccountInfo {
401 nonce: 2u64.into(),
402 consumers: 0,
403 providers: 1,
404 sufficients: 0,
405 data: 42,
406 };
407 assert_eq!(crate::Account::<Test>::get(1), expected_info);
408 })
409 }
410
411 #[test]
412 fn check_nonce_skipped_and_refund_for_other_origins() {
413 new_test_ext().execute_with(|| {
414 let ext = CheckNonce::<Test>(1u64.into());
415
416 let mut info = CALL.get_dispatch_info();
417 info.extension_weight = ext.weight(CALL);
418
419 assert!(info.extension_weight != Weight::zero());
421
422 let len = CALL.encoded_size();
423
424 let origin = crate::RawOrigin::Root.into();
425 let (pre, origin) = ext.validate_and_prepare(origin, CALL, &info, len, 0).unwrap();
426
427 assert!(origin.as_system_ref().unwrap().is_root());
428
429 let pd_res = Ok(());
430 let mut post_info = frame_support::dispatch::PostDispatchInfo {
431 actual_weight: Some(info.total_weight()),
432 pays_fee: Default::default(),
433 };
434
435 <CheckNonce<Test> as TransactionExtension<RuntimeCall>>::post_dispatch(
436 pre,
437 &info,
438 &mut post_info,
439 len,
440 &pd_res,
441 )
442 .unwrap();
443
444 assert_eq!(post_info.actual_weight, Some(info.call_weight));
445 })
446 }
447}