lib.rs

// This file is part of Substrate.

// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: Apache-2.0

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// 	http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! # Merkle Mountain Range
//!
//! ## Overview
//!
//! Details on Merkle Mountain Ranges (MMRs) can be found here:
//! <https://github.com/mimblewimble/grin/blob/master/doc/mmr.md>
//!
//! The MMR pallet constructs an MMR from leaves provided by the [`MerkleMountainRangeTree::push`]
//! method. MMR nodes are stored both in:
//! - on-chain storage - hashes only; not full leaf content;
//! - off-chain storage - via Indexing API we push full leaf content (and all internal nodes as
//! well) to the Off-chain DB, so that the data is available for Off-chain workers.
//! Hashing used for MMR is configurable independently from the rest of the runtime (i.e. not using
//! `frame_system::Hashing`) so something compatible with external chains can be used (like
//! Keccak256 for Ethereum compatibility).
//!
//! Depending on the usage context (off-chain vs on-chain) the pallet is able to:
//! - verify MMR leaf proofs (on-chain)
//! - generate leaf proofs (off-chain)
//!
//! See [primitives::Compact] documentation for how you can optimize proof size for leafs that are
//! composed from multiple elements.
//!
//! ## What for?
//!
//! Primary use case for this pallet is to generate MMR root hashes, that can latter on be used by
//! BEEFY protocol (see <https://github.com/paritytech/grandpa-bridge-gadget>).
//! MMR root hashes along with BEEFY will make it possible to build Super Light Clients (SLC) of
//! Substrate-based chains. The SLC will be able to follow finality and can be shown proofs of more
//! details that happened on the source chain.
//! In that case the chain which contains the pallet generates the Root Hashes and Proofs, which
//! are then presented to another chain acting as a light client which can verify them.
//!
//! Secondary use case is to archive historical data, but still be able to retrieve them on-demand
//! if needed. For instance if parent block hashes are stored in the MMR it's possible at any point
//! in time to provide an MMR proof about some past block hash, while this data can be safely pruned
//! from on-chain storage.
//!
//! NOTE This pallet is experimental and not proven to work in production.
#![cfg_attr(not(feature = "std"), no_std)]

use codec::Decode;
use core::marker::PhantomData;
use frame_system::pallet_prelude::{BlockNumberFor, HeaderFor};
use itertools::Itertools;
use log;
use log::trace;
use merkle_mountain_range::{helper::pos_height_in_tree, MMRStore};
use sp_core::H256;
use sp_core::offchain::StorageKind;

use sp_runtime::traits::{self, One};
use sp_std::prelude::*;

use ismp::{
	messaging::{hash_request, Keccak256},
	router::Request,
};
use mmr_primitives::{DataOrHash, FullLeaf, LeafMetadata, MerkleMountainRangeTree};
pub use pallet::*;
pub use sp_mmr_primitives::{
	self as primitives, utils::NodesUtils, Error, LeafDataProvider, LeafIndex, NodeIndex,
};
use sp_mmr_primitives::mmr_lib::leaf_index_to_pos;

pub use mmr::storage::{OffchainStorage, Storage};
use pallet_ismp::NoOpMmrTree;

pub mod mmr;

/// An MMR specific to the pallet.
type ModuleMmr<StorageType, T, I> = mmr::Mmr<StorageType, T, I, LeafOf<T, I>>;

/// Leaf data.
type LeafOf<T, I> = <T as Config<I>>::Leaf;

/// Hashing used for the pallet.
pub(crate) type HashingOf<T, I> = <T as Config<I>>::Hashing;
/// Hash type used for the pallet.
pub(crate) type HashOf<T, I> = <<T as Config<I>>::Hashing as traits::Hash>::Output;

#[frame_support::pallet]
pub mod pallet {
	use super::*;
	use frame_support::pallet_prelude::*;
	use mmr_primitives::ForkIdentifier;

	#[pallet::pallet]
	#[pallet::without_storage_info]
	pub struct Pallet<T, I = ()>(PhantomData<(T, I)>);

	/// This pallet's configuration trait
	#[pallet::config]
	pub trait Config<I: 'static = ()>: frame_system::Config + pallet_ismp::Config {
		/// Prefix for elements stored in the Off-chain DB via Indexing API.
		///
		/// Each node of the MMR is inserted both on-chain and off-chain via Indexing API.
		/// The former does not store full leaf content, just its compact version (hash),
		/// and some of the inner mmr nodes might be pruned from on-chain storage.
		/// The latter will contain all the entries in their full form.
		///
		/// Each node is stored in the Off-chain DB under key derived from the
		/// [`Self::INDEXING_PREFIX`] and its in-tree index (MMR position).
		const INDEXING_PREFIX: &'static [u8];

		/// A hasher type for MMR.
		///
		/// To construct trie nodes that result in merging (bagging) two peaks, depending on the
		/// node kind we take either:
		/// - The node (hash) itself if it's an inner node.
		/// - The hash of SCALE-encoding of the leaf data if it's a leaf node.
		///
		/// Then we create a tuple of these two hashes, SCALE-encode it (concatenate) and
		/// hash, to obtain a new MMR inner node - the new peak.
		type Hashing: traits::Hash;

		/// Generic leaf type to be inserted into the MMR.
		type Leaf: mmr_primitives::FullLeaf + scale_info::TypeInfo;

		/// A type that returns a hash unique to every block as a fork identifer for offchain keys
		type ForkIdentifierProvider: ForkIdentifier<Self>;

		/// Leaves count to prune
		const LEAF_COUNT_THRESHOLD: u64;
	}

	/// Latest MMR Root hash.
	#[pallet::storage]
	#[pallet::getter(fn mmr_root_hash)]
	pub type RootHash<T: Config<I>, I: 'static = ()> = StorageValue<_, HashOf<T, I>, ValueQuery>;

	/// Current size of the MMR (number of leaves).
	#[pallet::storage]
	#[pallet::getter(fn leaf_count)]
	pub type NumberOfLeaves<T: Config<I>, I: 'static = ()> = StorageValue<_, LeafIndex, ValueQuery>;

	/// Height at which the pallet started inserting leaves into offchain storage.
	#[pallet::storage]
	#[pallet::getter(fn initial_height)]
	pub type InitialHeight<T: Config<I>, I: 'static = ()> =
		StorageValue<_, BlockNumberFor<T>, OptionQuery>;

	/// Temporary leaf storage for while the block is still executing.
	#[pallet::storage]
	#[pallet::getter(fn intermediate_leaves)]
	pub type IntermediateLeaves<T: Config<I>, I: 'static = ()> =
		CountedStorageMap<_, Identity, NodeIndex, T::Leaf, OptionQuery>;

	/// Hashes of the nodes in the MMR.
	///
	/// Note this collection only contains MMR peaks, the inner nodes (and leaves)
	/// are pruned and only stored in the Offchain DB.
	#[pallet::storage]
	#[pallet::getter(fn mmr_peak)]
	pub type Nodes<T: Config<I>, I: 'static = ()> =
		CountedStorageMap<_, Identity, NodeIndex, HashOf<T, I>, OptionQuery>;

	// Set the initial height at which leaves were pushed to the offchain db for the offchain
	// mmr gadget. Since this is in on_initialize, then the leaves were set in a previous block.
	#[pallet::hooks]
	impl<T: Config<I>, I: 'static> Hooks<BlockNumberFor<T>> for Pallet<T, I>
	where
		HashOf<T, I>: Into<H256>,
	{
		fn on_initialize(_n: BlockNumberFor<T>) -> Weight {
			if NumberOfLeaves::<T, I>::get() > 0 && InitialHeight::<T, I>::get().is_none() {
				InitialHeight::<T, I>::put(frame_system::Pallet::<T>::block_number() - One::one())
			}

			Default::default()
		}
		fn on_idle(_n: BlockNumberFor<T>, _remaining_weight: Weight) -> Weight {
			Self::prune_mmr_leaves().unwrap(); // It should not panic
			Default::default()
		}
	}
}

impl<T, I> MerkleMountainRangeTree for Pallet<T, I>
where
	I: 'static,
	T: Config<I>,
	HashOf<T, I>: Into<H256>,
{
	type Leaf = T::Leaf;

	fn leaf_count() -> LeafIndex {
		NumberOfLeaves::<T, I>::get()
	}

	fn generate_proof(
		indices: Vec<LeafIndex>,
	) -> Result<(Vec<Self::Leaf>, primitives::Proof<H256>), Error> {
		let (leaves, proof) = Pallet::<T, I>::generate_proof(indices)?;
		let proof_nodes = proof.items.into_iter().map(Into::into).collect();
		let new_proof = primitives::Proof {
			leaf_indices: proof.leaf_indices,
			leaf_count: proof.leaf_count,
			items: proof_nodes,
		};

		Ok((leaves, new_proof))
	}

	fn push(leaf: T::Leaf) -> LeafMetadata {
		let temp_count = IntermediateLeaves::<T, I>::count() as u64;
		let index = NumberOfLeaves::<T, I>::get() + temp_count;
		IntermediateLeaves::<T, I>::insert(temp_count, leaf);
		let position = leaf_index_to_pos(index);
		LeafMetadata { position, index }
	}

	fn finalize() -> Result<H256, Error> {
		let buffer_len = IntermediateLeaves::<T, I>::count() as u64;
		// no new leaves? early return
		if buffer_len == 0 {
			return Ok(RootHash::<T, I>::get().into());
		}

		let leaves = NumberOfLeaves::<T, I>::get();
		let mut mmr: ModuleMmr<mmr::storage::RuntimeStorage, T, I> = mmr::Mmr::new(leaves);

		// append new leaves to MMR
		let range = 0u64..buffer_len;
		for index in range {
			let leaf = IntermediateLeaves::<T, I>::get(index)
				.expect("Infallible: Leaf was inserted in this block");
			// Mmr push should never fail
			match mmr.push(leaf) {
				None => {
					log::error!(target: "pallet-mmr", "MMR push failed ");
					// MMR push never fails, but better safe than sorry.
					Err(Error::Push)?
				},
				Some(position) => {
					log::trace!(target: "pallet-mmr", "MMR push {position}");
				},
			}
		}

		// Update the size, `mmr.finalize()` should also never fail.
		let (leaves, root) = match mmr.finalize() {
			Ok((leaves, root)) => (leaves, root),
			Err(e) => {
				log::error!(target: "pallet-mmr", "MMR finalize failed: {:?}", e);
				Err(Error::Commit)?
			},
		};

		let _ = IntermediateLeaves::<T, I>::clear(buffer_len as u32, None);
		NumberOfLeaves::<T, I>::put(leaves);
		RootHash::<T, I>::put(root);

		Ok(root.into())
	}

	fn get_leaf(pos: NodeIndex) -> Result<Option<Self::Leaf>, Error> {
		let store = Storage::<OffchainStorage, T, _, Self::Leaf>::default();
		store
			.get_elem(pos)
			.map(|val| {
				val.and_then(|inner| match inner {
					DataOrHash::Data(leaf) => Some(leaf),
					_ => None,
				})
			})
			.map_err(|_| Error::LeafNotFound)
	}

	// fetch the peaks and under each peak see if latest leaves (i.e right most leaf (leafIndex) is
	// still valid, meaning has not timedout or processed. if its invalid then prune all leaves and
	// inner under the peak. for now only prune the earliest peak
	fn prune_mmr_leaves() -> Result<(), Error> {
		if Self::leaf_count() < T::LEAF_COUNT_THRESHOLD {
			Ok(())?
		}

		let peaks_indexs = Nodes::<T, I>::iter()
			.sorted_by_key(|(k, _)| *k)
			.map(|(k, _v)| k)
			.collect::<Vec<NodeIndex>>();
		// if there is only 1 peak meaning the tree has no of leaves 2^n we are pruning inner nodes
		// i.e the left most inner node on the (h-1) layer, where h = height of the tree
		// if there is more than 1 peak, then prune all the nodes on the first peak
		if peaks_indexs.len() == 1 {
			todo!()
		} else {
			if let Some(peak_index) = peaks_indexs.iter().next() {
				// get the last leaf under the peak and check
				let last_leaf_index = *peak_index - pos_height_in_tree(*peak_index) as u64;
				if let Some(leaf_type) = Self::get_leaf(last_leaf_index)? {
					let last_leaf_to_delete = {
						let encoded_inner_leaf = leaf_type.preimage();
						let leaf_request: ismp::router::Request =
							Decode::decode(&mut &encoded_inner_leaf[..])
								.map_err(|_| Error::LeafNotFound)?;

						let claimed =
							pallet_ismp::child_trie::RequestCommitments::<T>::get(hash_request::<
								pallet_ismp::Pallet<T>,
							>(&leaf_request))
							.ok_or(Error::LeafNotFound)?
							.claimed;

						match leaf_request {
							Request::Post(post) => {
								// check if it has timedout and if fees has been claimed by this
								// request
								let _timeout = post.timeout_timestamp;
								claimed
							},
							Request::Get(ref get) => {
								// check if it has timedout and if fees has been claimed by this
								// request
								let _timeout = get.timeout_timestamp;
								claimed
							},
						}
					};

					// if we can delete the last leaf we can delete all nodes under the peak
					if last_leaf_to_delete {
						for node_index in 0..*peak_index {
							let leaf = Self::get_leaf(node_index)?.ok_or(Error::LeafNotFound)?;
							let commitment = pallet_ismp::Pallet::<T>::keccak256(&leaf.preimage()[..]);
							// delete the node
							let offchain_key = NoOpMmrTree::<T>::offchain_key(commitment);
							sp_io::offchain::local_storage_clear(StorageKind::PERSISTENT, &offchain_key)
						}
					}else{
						trace!(target: "mmr:pruning","No nodes to prune")
					}
				}
			}
		}

		Ok(())
	}
}
/// Stateless MMR proof verification for batch of leaves.
///
/// This function can be used to verify received MMR [primitives::Proof] (`proof`)
/// for given leaves set (`leaves`) against a known MMR root hash (`root`).
/// Note, the leaves should be sorted such that corresponding leaves and leaf indices have the
/// same position in both the `leaves` vector and the `leaf_indices` vector contained in the
/// [primitives::Proof].
pub fn verify_leaves_proof<H, L>(
	root: H::Output,
	leaves: Vec<mmr::Node<H, L>>,
	proof: primitives::Proof<H::Output>,
) -> Result<(), primitives::Error>
where
	H: traits::Hash,
	L: mmr_primitives::FullLeaf,
{
	let is_valid = mmr::verify_leaves_proof::<H, L>(root, leaves, proof)?;
	if is_valid {
		Ok(())
	} else {
		Err(primitives::Error::Verify.log_debug(("The proof is incorrect.", root)))
	}
}

impl<T: Config<I>, I: 'static> Pallet<T, I> {
	/// Build offchain key from a combination of a fork resistant hash, position and indexing prefix
	///
	/// This combination makes the offchain (key,value) entry resilient to chain forks.
	fn node_temp_offchain_key(
		pos: NodeIndex,
		fork_identifier: <T as frame_system::Config>::Hash,
	) -> sp_std::prelude::Vec<u8> {
		NodesUtils::node_temp_offchain_key::<HeaderFor<T>>(
			&T::INDEXING_PREFIX,
			pos,
			fork_identifier,
		)
	}

	/// Build canonical offchain key for node `pos` in MMR.
	///
	/// Used for nodes added by now finalized blocks.
	/// Never read keys using `node_canon_offchain_key` unless you sure that
	/// there's no `node_offchain_key` key in the storage.
	fn node_canon_offchain_key(pos: NodeIndex) -> sp_std::prelude::Vec<u8> {
		NodesUtils::node_canon_offchain_key(&T::INDEXING_PREFIX, pos)
	}

	/// Return the on-chain MMR root hash.
	pub fn mmr_root() -> HashOf<T, I> {
		RootHash::<T, I>::get()
	}

	/// Generate an MMR proof for the given `leaf_indices`.
	/// Generates a proof for the MMR at the current block height.
	///
	/// Note this method can only be used from an off-chain context
	/// (Offchain Worker or Runtime API call), since it requires
	/// all the leaves to be present.
	/// It may return an error or panic if used incorrectly.
	pub fn generate_proof(
		indices: Vec<LeafIndex>,
	) -> Result<(Vec<LeafOf<T, I>>, primitives::Proof<HashOf<T, I>>), primitives::Error> {
		let leaves_count = NumberOfLeaves::<T, I>::get();
		let mmr: ModuleMmr<mmr::storage::OffchainStorage, T, I> = mmr::Mmr::new(leaves_count);
		mmr.generate_proof(indices)
	}

	/// Verify MMR proof for given `leaves`.
	///
	/// This method is safe to use within the runtime code.
	/// It will return `Ok(())` if the proof is valid
	/// and an `Err(..)` if MMR is inconsistent (some leaves are missing)
	/// or the proof is invalid.
	pub fn verify_leaves(
		leaves: Vec<LeafOf<T, I>>,
		proof: primitives::Proof<HashOf<T, I>>,
	) -> Result<(), primitives::Error> {
		if proof.leaf_count > NumberOfLeaves::<T, I>::get() ||
			proof.leaf_count == 0 ||
			(proof.items.len().saturating_add(leaves.len())) as u64 > proof.leaf_count
		{
			return Err(primitives::Error::Verify
				.log_debug("The proof has incorrect number of leaves or proof items."));
		}

		let mmr: ModuleMmr<mmr::storage::OffchainStorage, T, I> = mmr::Mmr::new(proof.leaf_count);
		let is_valid = mmr.verify_leaves_proof(leaves, proof)?;
		if is_valid {
			Ok(())
		} else {
			Err(primitives::Error::Verify.log_debug("The proof is incorrect."))
		}
	}
}