diff --git a/substrate/frame/assets/src/benchmarking.rs b/substrate/frame/assets/src/benchmarking.rs index 242041bab349..ffd7359d662e 100644 --- a/substrate/frame/assets/src/benchmarking.rs +++ b/substrate/frame/assets/src/benchmarking.rs @@ -26,8 +26,7 @@ use frame_benchmarking::v1::{ }; use frame_support::traits::{EnsureOrigin, Get, UnfilteredDispatchable}; use frame_system::RawOrigin as SystemOrigin; -use sp_runtime::traits::Bounded; -use sp_runtime::traits::Hash; +use sp_runtime::traits::{Bounded, Hash}; use crate::Pallet as Assets; diff --git a/substrate/frame/assets/src/functions.rs b/substrate/frame/assets/src/functions.rs index cf36d55d5fbd..026dfd283e4e 100644 --- a/substrate/frame/assets/src/functions.rs +++ b/substrate/frame/assets/src/functions.rs @@ -475,10 +475,16 @@ impl, I: 'static> Pallet { pub(super) fn do_claim_distribution( distribution_id: DistributionCounter, merkle_proof: Vec, + hashes: u32, ) -> DispatchResult { let proof = codec::Decode::decode(&mut &merkle_proof[..]).map_err(|_| Error::::BadProof)?; + let expected_hashes = T::VerifyExistenceProof::proof_to_hashes(&proof) + .map_err(|_| Error::::BadProof)?; + + ensure!(hashes >= expected_hashes, Error::::TooManyHashes); + let DistributionInfo { asset_id, merkle_root, active } = MerklizedDistribution::::get(distribution_id).ok_or(Error::::Unknown)?; diff --git a/substrate/frame/assets/src/lib.rs b/substrate/frame/assets/src/lib.rs index 009aa9dd7bbc..b5cff7254c79 100644 --- a/substrate/frame/assets/src/lib.rs +++ b/substrate/frame/assets/src/lib.rs @@ -188,8 +188,8 @@ use frame_support::{ Preservation::{Expendable, Preserve}, WithdrawConsequence, }, - BalanceStatus::Reserved, ProofToHashes, - Currency, EnsureOriginWithArg, Incrementable, ReservableCurrency, StoredMap, + BalanceStatus::Reserved, + Currency, EnsureOriginWithArg, Incrementable, ProofToHashes, ReservableCurrency, StoredMap, VerifyExistenceProof, }, }; @@ -299,7 +299,6 @@ pub mod pallet { type Extra = (); type CallbackHandle = (); type WeightInfo = (); - type VerifyExistenceProof = NoTrie; #[cfg(feature = "runtime-benchmarks")] type BenchmarkHelper = (); } @@ -408,7 +407,9 @@ pub mod pallet { type CallbackHandle: AssetsCallback; /// A type used to verify merkle proofs used for distributions. - type VerifyExistenceProof: VerifyExistenceProof::Hash> + ProofToHashes>; + #[pallet::no_default] + type VerifyExistenceProof: VerifyExistenceProof + + ProofToHashes>; /// Weight information for extrinsics in this pallet. type WeightInfo: WeightInfo; @@ -736,8 +737,10 @@ pub mod pallet { DistributionActive, /// The proof provided could not be verified. BadProof, - /// The a leaf node was extracted from the proof, but it did not match the expected format. + /// The leaf node was extracted from the proof, but it did not match the expected format. CannotDecodeLeaf, + /// The proof will require more hashes than expected. + TooManyHashes, } #[pallet::call(weight(>::WeightInfo))] @@ -1884,17 +1887,15 @@ pub mod pallet { /// /// Weight: `O(P)` where `P` is the size of the merkle proof. #[pallet::call_index(34)] - #[pallet::weight({ - let hashes = T::VerifyExistenceProof::proof_to_hashes(merkle_proof); - T::WeightInfo::trie_hash(hashes) - })] + #[pallet::weight(T::WeightInfo::trie_hash(*hashes))] pub fn claim_distribution( origin: OriginFor, distribution_id: DistributionCounter, merkle_proof: Vec, + hashes: u32, ) -> DispatchResult { ensure_signed(origin)?; - Self::do_claim_distribution(distribution_id, merkle_proof)?; + Self::do_claim_distribution(distribution_id, merkle_proof, hashes)?; Ok(()) } diff --git a/substrate/frame/assets/src/tests.rs b/substrate/frame/assets/src/tests.rs index 21a98a958910..e54d95d196f6 100644 --- a/substrate/frame/assets/src/tests.rs +++ b/substrate/frame/assets/src/tests.rs @@ -1963,24 +1963,42 @@ fn merklized_distribution_works() { _, >(flat_distribution, 6); + let hashes = + ::VerifyExistenceProof::proof_to_hashes(&proof_for_1).unwrap(); + // Use this trie root for the distribution assert_ok!(Assets::mint_distribution(RuntimeOrigin::signed(1), 0, root)); // Now users claim their distributions permissionlessly with a proof. - assert_ok!(Assets::claim_distribution(RuntimeOrigin::signed(1), 0, proof_for_1.encode())); + assert_ok!(Assets::claim_distribution( + RuntimeOrigin::signed(1), + 0, + proof_for_1.encode(), + hashes + )); assert_eq!(Assets::balance(0, 1), 1337); // Other users can claim their tokens. - assert_ok!(Assets::claim_distribution(RuntimeOrigin::signed(55), 0, proof_for_69.encode())); + assert_ok!(Assets::claim_distribution( + RuntimeOrigin::signed(55), + 0, + proof_for_69.encode(), + hashes + )); assert_eq!(Assets::balance(0, 69), 69); // Owner (or anyone) can also distribute on behalf of the other users. - assert_ok!(Assets::claim_distribution(RuntimeOrigin::signed(1), 0, proof_for_6.encode())); + assert_ok!(Assets::claim_distribution( + RuntimeOrigin::signed(1), + 0, + proof_for_6.encode(), + hashes + )); assert_eq!(Assets::balance(0, 6), 6); // You cannot double claim. assert_noop!( - Assets::claim_distribution(RuntimeOrigin::signed(6), 0, proof_for_6.encode()), + Assets::claim_distribution(RuntimeOrigin::signed(6), 0, proof_for_6.encode(), hashes), Error::::AlreadyClaimed ); }); diff --git a/substrate/frame/assets/src/types.rs b/substrate/frame/assets/src/types.rs index 5ed12cb4f083..3c3bd59e7be8 100644 --- a/substrate/frame/assets/src/types.rs +++ b/substrate/frame/assets/src/types.rs @@ -324,8 +324,6 @@ pub type DistributionProofOf = pub type DistributionHashOf = <>::VerifyExistenceProof as VerifyExistenceProof>::Hash; -pub type HashOf = ::Hash; - #[derive(Eq, PartialEq, Copy, Clone, RuntimeDebug, Encode, Decode, TypeInfo, MaxEncodedLen)] pub struct DistributionInfo { // The asset id we are distributing.