[pallet-transaction-storage] Improved check_proof check + tests + docs

prdoc/pr_10153.prdoc

@@ -0,0 +1,16 @@
+title: '[pallet-transaction-storage] Improved `check_proof` check + tests + docs'
+doc:
+- audience: Runtime Dev
+  description: |-
+    **This PR:**
+
+    * Fixes `check_proof` and its `binary_search_by_key` chunk
+    * Adds the `ensure_chunk_proof_works` test, which covers all possible chunk index build/verify proof roundtrips (to catch all corner cases)
+    * Improves docs around `pallet-transaction-storage`
+crates:
+- name: pallet-transaction-storage
+  bump: patch
+- name: sp-transaction-storage-proof
+  bump: patch
+- name: sp-io
+  bump: patch

substrate/frame/transaction-storage/src/lib.rs

@@ -43,7 +43,7 @@ use frame_support::{
 };
 use sp_runtime::traits::{BlakeTwo256, Dispatchable, Hash, One, Saturating, Zero};
 use sp_transaction_storage_proof::{
-	encode_index, random_chunk, InherentError, TransactionStorageProof, CHUNK_SIZE,
+	encode_index, num_chunks, random_chunk, InherentError, TransactionStorageProof, CHUNK_SIZE,
 	INHERENT_IDENTIFIER,
 };
 
@@ -80,14 +80,13 @@ pub struct TransactionInfo {
 	/// Size of indexed data in bytes.
 	size: u32,
 	/// Total number of chunks added in the block with this transaction. This
-	/// is used find transaction info by block chunk index using binary search.
+	/// is used to find transaction info by block chunk index using binary search.
+	///
+	/// Cumulative value of all previous transactions in the block; the last transaction holds the
+	/// total chunk value.
 	block_chunks: u32,
 }
 
-fn num_chunks(bytes: u32) -> u32 {
-	(bytes as u64).div_ceil(CHUNK_SIZE as u64) as u32
-}
-
 #[frame_support::pallet]
 pub mod pallet {
 	use super::*;
@@ -133,7 +132,7 @@ pub mod pallet {
 		NotConfigured,
 		/// Renewed extrinsic is not found.
 		RenewedNotFound,
-		/// Attempting to store empty transaction
+		/// Attempting to store an empty transaction
 		EmptyTransaction,
 		/// Proof was not expected in this block.
 		UnexpectedProof,
@@ -225,7 +224,7 @@ pub mod pallet {
 			let mut index = 0;
 			BlockTransactions::<T>::mutate(|transactions| {
 				if transactions.len() + 1 > T::MaxBlockTransactions::get() as usize {
-					return Err(Error::<T>::TooManyTransactions)
+					return Err(Error::<T>::TooManyTransactions);
 				}
 				let total_chunks = transactions.last().map_or(0, |t| t.block_chunks) + chunk_count;
 				index = transactions.len() as u32;
@@ -269,7 +268,7 @@ pub mod pallet {
 			let mut index = 0;
 			BlockTransactions::<T>::mutate(|transactions| {
 				if transactions.len() + 1 > T::MaxBlockTransactions::get() as usize {
-					return Err(Error::<T>::TooManyTransactions)
+					return Err(Error::<T>::TooManyTransactions);
 				}
 				let chunks = num_chunks(info.size);
 				let total_chunks = transactions.last().map_or(0, |t| t.block_chunks) + chunks;
@@ -301,39 +300,24 @@ pub mod pallet {
 		) -> DispatchResultWithPostInfo {
 			ensure_none(origin)?;
 			ensure!(!ProofChecked::<T>::get(), Error::<T>::DoubleCheck);
+
+			// Get the target block metadata.
 			let number = frame_system::Pallet::<T>::block_number();
 			let period = StoragePeriod::<T>::get();
 			let target_number = number.saturating_sub(period);
 			ensure!(!target_number.is_zero(), Error::<T>::UnexpectedProof);
 			let total_chunks = ChunkCount::<T>::get(target_number);
-			ensure!(total_chunks != 0, Error::<T>::UnexpectedProof);
+			let transactions =
+				Transactions::<T>::get(target_number).ok_or(Error::<T>::MissingStateData)?;
+
+			// Verify the proof with a "random" chunk (randomness is based on the parent hash).
 			let parent_hash = frame_system::Pallet::<T>::parent_hash();
-			let selected_chunk_index = random_chunk(parent_hash.as_ref(), total_chunks);
-			let (info, chunk_index) = match Transactions::<T>::get(target_number) {
-				Some(infos) => {
-					let index = match infos
-						.binary_search_by_key(&selected_chunk_index, |info| info.block_chunks)
-					{
-						Ok(index) => index,
-						Err(index) => index,
-					};
-					let info = infos.get(index).ok_or(Error::<T>::MissingStateData)?.clone();
-					let chunks = num_chunks(info.size);
-					let prev_chunks = info.block_chunks - chunks;
-					(info, selected_chunk_index - prev_chunks)
-				},
-				None => return Err(Error::<T>::MissingStateData.into()),
-			};
-			ensure!(
-				sp_io::trie::blake2_256_verify_proof(
-					info.chunk_root,
-					&proof.proof,
-					&encode_index(chunk_index),
-					&proof.chunk,
-					sp_runtime::StateVersion::V1,
-				),
-				Error::<T>::InvalidProof
-			);
+			Self::ensure_chunk_proof(
+				proof,
+				parent_hash.as_ref(),
+				transactions.to_vec(),
+				total_chunks,
+			)?;
 			ProofChecked::<T>::put(true);
 			Self::deposit_event(Event::ProofChecked);
 			Ok(().into())
@@ -466,5 +450,60 @@ pub mod pallet {
 			T::FeeDestination::on_unbalanced(credit);
 			Ok(())
 		}
+
+		pub(crate) fn ensure_chunk_proof(
+			proof: TransactionStorageProof,
+			random_hash: &[u8],
+			infos: Vec<TransactionInfo>,
+			total_chunks: u32,
+		) -> Result<(), Error<T>> {
+			ensure!(total_chunks != 0, Error::<T>::UnexpectedProof);
+			ensure!(!infos.is_empty(), Error::<T>::UnexpectedProof);
+
+			// Get the random chunk index (from all transactions in the block = 0..total_chunks).
+			let selected_block_chunk_index = random_chunk(random_hash, total_chunks as _);
+
+			// Let's find the corresponding transaction and its "local" chunk index for "global"
+			// `selected_block_chunk_index`.
+			let (tx_info, tx_chunk_index) = {
+				// Binary search for the transaction that owns this `selected_block_chunk_index`
+				// chunk.
+				let tx_index = infos
+					.binary_search_by_key(&selected_block_chunk_index, |info| {
+						// Each `info.block_chunks` is cumulative count,
+						// so last chunk index = count - 1.
+						info.block_chunks.saturating_sub(1)
+					})
+					.unwrap_or_else(|tx_index| tx_index);
+
+				// Get the transaction and its local chunk index.
+				let tx_info = infos.get(tx_index).ok_or(Error::<T>::MissingStateData)?;
+				// We shouldn't reach this point; we rely on the fact that `fn store` does not allow
+				// empty transactions. Without this check, it would fail anyway below with
+				// `InvalidProof`.
+				ensure!(!tx_info.block_chunks.is_zero(), Error::<T>::EmptyTransaction);
+
+				// Convert a global chunk index into a transaction-local one.
+				let tx_chunks = num_chunks(tx_info.size);
+				let prev_chunks = tx_info.block_chunks - tx_chunks;
+				let tx_chunk_index = selected_block_chunk_index - prev_chunks;
+
+				(tx_info, tx_chunk_index)
+			};
+
+			// Verify the tx chunk proof.
+			ensure!(
+				sp_io::trie::blake2_256_verify_proof(
+					tx_info.chunk_root,
+					&proof.proof,
+					&encode_index(tx_chunk_index),
+					&proof.chunk,
+					sp_runtime::StateVersion::V1,
+				),
+				Error::<T>::InvalidProof
+			);
+
+			Ok(())
+		}
 	}
 }

substrate/frame/transaction-storage/src/tests.rs

@@ -22,7 +22,7 @@ use crate::mock::*;
 use frame_support::{assert_noop, assert_ok};
 use frame_system::RawOrigin;
 use sp_runtime::{DispatchError, TokenError::FundsUnavailable};
-use sp_transaction_storage_proof::registration::build_proof;
+use sp_transaction_storage_proof::{registration::build_proof, CHUNK_SIZE};
 
 const MAX_DATA_SIZE: u32 = DEFAULT_MAX_TRANSACTION_SIZE;
 
@@ -114,6 +114,62 @@ fn checks_proof() {
 	});
 }
 
+#[test]
+fn ensure_chunk_proof_works() {
+	new_test_ext().execute_with(|| {
+		// Prepare a bunch of transactions with variable chunk sizes.
+		let transactions = vec![
+			vec![0u8; CHUNK_SIZE - 1],
+			vec![1u8; CHUNK_SIZE],
+			vec![2u8; CHUNK_SIZE + 1],
+			vec![3u8; 2 * CHUNK_SIZE - 1],
+			vec![3u8; 2 * CHUNK_SIZE],
+			vec![3u8; 2 * CHUNK_SIZE + 1],
+			vec![4u8; 7 * CHUNK_SIZE - 1],
+			vec![4u8; 7 * CHUNK_SIZE],
+			vec![4u8; 7 * CHUNK_SIZE + 1],
+		];
+		let expected_total_chunks =
+			transactions.iter().map(|t| t.len().div_ceil(CHUNK_SIZE) as u32).sum::<u32>();
+
+		// Store a couple of transactions in one block.
+		run_to_block(1, || None);
+		let caller = 1;
+		for transaction in transactions.clone() {
+			assert_ok!(TransactionStorage::<Test>::store(
+				RawOrigin::Signed(caller).into(),
+				transaction
+			));
+		}
+		run_to_block(2, || None);
+
+		// Read all the block transactions metadata.
+		let total_chunks = ChunkCount::<Test>::get(1);
+		let tx_infos = Transactions::<Test>::get(1).unwrap();
+		assert_eq!(expected_total_chunks, total_chunks);
+		assert_eq!(9, tx_infos.len());
+
+		// Verify proofs for all possible chunk indexes.
+		for chunk_index in 0..total_chunks {
+			// chunk index randomness
+			let mut random_hash = [0u8; 32];
+			random_hash[..8].copy_from_slice(&(chunk_index as u64).to_be_bytes());
+			let selected_chunk_index = random_chunk(random_hash.as_ref(), total_chunks);
+			assert_eq!(selected_chunk_index, chunk_index);
+
+			// build/check chunk proof roundtrip
+			let proof =
+				build_proof(random_hash.as_ref(), transactions.clone()).expect("valid proof");
+			assert_ok!(TransactionStorage::<Test>::ensure_chunk_proof(
+				proof,
+				random_hash.as_ref(),
+				tx_infos.to_vec(),
+				total_chunks
+			));
+		}
+	});
+}
+
 #[test]
 fn renews_data() {
 	new_test_ext().execute_with(|| {

substrate/primitives/io/src/lib.rs

@@ -1486,7 +1486,7 @@ pub trait Hashing {
 /// Interface that provides transaction indexing API.
 #[runtime_interface]
 pub trait TransactionIndex {
-	/// Add transaction index. Returns indexed content hash.
+	/// Indexes the specified transaction for the given `extrinsic` and `context_hash`.
 	fn index(
 		&mut self,
 		extrinsic: u32,
@@ -1496,7 +1496,8 @@ pub trait TransactionIndex {
 		self.storage_index_transaction(extrinsic, &context_hash, size);
 	}
 
-	/// Conduct a 512-bit Keccak hash.
+	/// Renews the transaction index entry for the given `extrinsic` using the provided
+	/// `context_hash`.
 	fn renew(&mut self, extrinsic: u32, context_hash: PassPointerAndReadCopy<[u8; 32], 32>) {
 		self.storage_renew_transaction_index(extrinsic, &context_hash);
 	}

substrate/primitives/transaction-storage-proof/src/lib.rs

@@ -17,6 +17,8 @@
 
 //! Storage proof primitives. Contains types and basic code to extract storage
 //! proofs for indexed transactions.
+//!
+//! Note: We use `u32` for both `total_chunks` and the chunk index.
 
 #![cfg_attr(not(feature = "std"), no_std)]
 
@@ -104,7 +106,7 @@ impl sp_inherents::InherentDataProvider for InherentDataProvider {
 		mut error: &[u8],
 	) -> Option<Result<(), Error>> {
 		if *identifier != INHERENT_IDENTIFIER {
-			return None
+			return None;
 		}
 
 		let error = InherentError::decode(&mut error).ok()?;
@@ -114,14 +116,25 @@ impl sp_inherents::InherentDataProvider for InherentDataProvider {
 }
 
 /// A utility function to extract a chunk index from the source of randomness.
+///
+/// Note: The caller must ensure that `total_chunks` is not 0.
 pub fn random_chunk(random_hash: &[u8], total_chunks: u32) -> u32 {
 	let mut buf = [0u8; 8];
 	buf.copy_from_slice(&random_hash[0..8]);
 	let random_u64 = u64::from_be_bytes(buf);
 	(random_u64 % total_chunks as u64) as u32
 }
 
-/// A utility function to encode transaction index as trie key.
+/// A utility function to calculate the number of chunks.
+///
+/// * `bytes` - number of bytes
+pub fn num_chunks(bytes: u32) -> u32 {
+	(bytes as u64).div_ceil(CHUNK_SIZE as u64) as u32
+}
+
+/// A utility function to encode the transaction index as a trie key.
+///
+/// * `input` - chunk index.
 pub fn encode_index(input: u32) -> Vec<u8> {
 	codec::Encode::encode(&codec::Compact(input))
 }
@@ -163,13 +176,12 @@ pub mod registration {
 			.saturating_sub(DEFAULT_STORAGE_PERIOD.into());
 		if number.is_zero() {
 			// Too early to collect proofs.
-			return Ok(InherentDataProvider::new(None))
+			return Ok(InherentDataProvider::new(None));
 		}
 
 		let proof = match client.block_indexed_body(number)? {
-			Some(transactions) if !transactions.is_empty() =>
-				Some(build_proof(parent.as_ref(), transactions)?),
-			Some(_) | None => {
+			Some(transactions) => Some(build_proof(parent.as_ref(), transactions)?),
+			None => {
 				// Nothing was indexed in that block.
 				None
 			},
@@ -182,19 +194,20 @@ pub mod registration {
 		random_hash: &[u8],
 		transactions: Vec<Vec<u8>>,
 	) -> Result<TransactionStorageProof, Error> {
-		let mut db = sp_trie::MemoryDB::<Hasher>::default();
+		// Get total chunks, we will need it to generate a random chunk index.
+		let total_chunks: u32 = transactions.iter().map(|t| num_chunks(t.len() as u32)).sum();
+		if total_chunks.is_zero() {
+			const MSG: &str = "No chunks to build proof for.";
+			return Err(Error::Application(Box::from(MSG)));
+		}
+		let target_chunk_index = random_chunk(random_hash.as_ref(), total_chunks);
 
+		let mut db = sp_trie::MemoryDB::<Hasher>::default();
 		let mut target_chunk = None;
 		let mut target_root = Default::default();
 		let mut target_chunk_key = Default::default();
 		let mut chunk_proof = Default::default();
 
-		let total_chunks: u64 =
-			transactions.iter().map(|t| t.len().div_ceil(CHUNK_SIZE) as u64).sum();
-		let mut buf = [0u8; 8];
-		buf.copy_from_slice(&random_hash[0..8]);
-		let random_u64 = u64::from_be_bytes(buf);
-		let target_chunk_index = random_u64 % total_chunks;
 		// Generate tries for each transaction.
 		let mut chunk_index = 0;
 		for transaction in transactions {
@@ -244,5 +257,9 @@ pub mod registration {
 			&[(encode_index(0), Some(proof.chunk))],
 		)
 		.unwrap();
+
+		// Fail for empty transactions/chunks.
+		assert!(build_proof(&random, vec![]).is_err());
+		assert!(build_proof(&random, vec![vec![]]).is_err());
 	}
 }