From e45db1ecc113b9698d4089a831b57210b85e670e Mon Sep 17 00:00:00 2001 From: Emeric Chevalier Date: Tue, 25 Oct 2022 18:23:31 +0200 Subject: [PATCH 1/2] limit is not compatible with hash_db prefix usage --- .../reference-trie/src/substrate_like.rs | 23 +++++++++++++++---- trie-db/test/src/triedbmut.rs | 19 +++++++++++++++ 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/test-support/reference-trie/src/substrate_like.rs b/test-support/reference-trie/src/substrate_like.rs index e9b67489..59fa2487 100644 --- a/test-support/reference-trie/src/substrate_like.rs +++ b/test-support/reference-trie/src/substrate_like.rs @@ -16,6 +16,7 @@ use super::{CodecError as Error, NodeCodec as NodeCodecT, *}; use trie_db::node::Value; +use trie_db::nibble_ops; /// No extension trie with no hashed value. pub struct HashedValueNoExt; @@ -45,7 +46,8 @@ impl TrieLayout for HashedValueNoExtThreshold { /// Constants specific to encoding with external value node support. pub mod trie_constants { const FIRST_PREFIX: u8 = 0b_00 << 6; - pub const NIBBLE_SIZE_BOUND: usize = u16::max_value() as usize; + //pub const NIBBLE_SIZE_BOUND: usize = u16::max_value() as usize - 1; + pub const NIBBLE_SIZE_BOUND: usize = 4; pub const LEAF_PREFIX_MASK: u8 = 0b_01 << 6; pub const BRANCH_WITHOUT_MASK: u8 = 0b_10 << 6; pub const BRANCH_WITH_MASK: u8 = 0b_11 << 6; @@ -259,11 +261,11 @@ where /// Encode and allocate node type header (type and size), and partial value. /// It uses an iterator over encoded partial bytes as input. fn partial_from_iterator_encode>( - partial: I, - nibble_count: usize, + mut partial: I, + total_nibble_count: usize, node_kind: NodeKind, ) -> Vec { - let nibble_count = std::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, nibble_count); + let nibble_count = std::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, total_nibble_count); let mut output = Vec::with_capacity(4 + (nibble_count / nibble_ops::NIBBLE_PER_BYTE)); match node_kind { @@ -275,7 +277,18 @@ fn partial_from_iterator_encode>( NodeKind::HashedValueBranch => NodeHeader::HashedValueBranch(nibble_count).encode_to(&mut output), }; - output.extend(partial); + if total_nibble_count != nibble_count { + // truncate key content. + let was_aligned = total_nibble_count % nibble_ops::NIBBLE_PER_BYTE == 0; + debug_assert!(nibble_count % nibble_ops::NIBBLE_PER_BYTE == 0); + let mut nb_bytes = trie_constants::NIBBLE_SIZE_BOUND / nibble_ops::NIBBLE_PER_BYTE; + if trie_constants::NIBBLE_SIZE_BOUND % nibble_ops::NIBBLE_PER_BYTE > 0 { + nb_bytes += 1; + } + output.extend(partial.take(nb_bytes)); + } else { + output.extend(partial); + } output } diff --git a/trie-db/test/src/triedbmut.rs b/trie-db/test/src/triedbmut.rs index cb7c0d28..0b2d2c52 100644 --- a/trie-db/test/src/triedbmut.rs +++ b/trie-db/test/src/triedbmut.rs @@ -226,6 +226,25 @@ fn insert_replace_root_internal() { ); } +#[test] +fn oversized_key() { + use reference_trie::HashedValueNoExtThreshold as Layout; + let check = |keysize: usize, oversized: Option| { + let mut memdb = PrefixedMemoryDB::::default(); + let mut root = Default::default(); + let mut t = TrieDBMutBuilder::::new(&mut memdb, &mut root).build(); + t.insert(&vec![0x01u8; keysize][..], &[0x01u8, 0x23]).unwrap(); + std::mem::drop(t); + let t = TrieDBBuilder::::new(&memdb, &root).build(); + let mut truncated_key = vec![0x01u8; keysize - oversized.unwrap_or(0)]; + assert_eq!(t.get(&truncated_key[..]).unwrap(), Some(vec![0x01u8, 0x23])); + }; + check(4 as usize / 2, None); // ok + check(4 as usize / 2 + 1, Some(1)); // Oversized, rem last two nibble + //check(u16::MAX as usize / 2, None); // ok + //check(u16::MAX as usize / 2 + 1, Some(1)); // Oversized, rem last two nibble +} + test_layouts!(insert_make_branch_root, insert_make_branch_root_internal); fn insert_make_branch_root_internal() { let mut memdb = PrefixedMemoryDB::::default(); From 04c84002a5d5ab36f9d1559ec0f5998ec72827c7 Mon Sep 17 00:00:00 2001 From: Emeric Chevalier Date: Wed, 26 Oct 2022 12:26:58 +0200 Subject: [PATCH 2/2] Remove key size limit from codec. --- .../reference-trie/src/substrate_like.rs | 31 ++++--------------- trie-db/test/src/triedbmut.rs | 19 ------------ 2 files changed, 6 insertions(+), 44 deletions(-) diff --git a/test-support/reference-trie/src/substrate_like.rs b/test-support/reference-trie/src/substrate_like.rs index 59fa2487..53f78a5c 100644 --- a/test-support/reference-trie/src/substrate_like.rs +++ b/test-support/reference-trie/src/substrate_like.rs @@ -15,8 +15,7 @@ //! Codec and layout configuration similar to upstream default substrate one. use super::{CodecError as Error, NodeCodec as NodeCodecT, *}; -use trie_db::node::Value; -use trie_db::nibble_ops; +use trie_db::{nibble_ops, node::Value}; /// No extension trie with no hashed value. pub struct HashedValueNoExt; @@ -46,8 +45,6 @@ impl TrieLayout for HashedValueNoExtThreshold { /// Constants specific to encoding with external value node support. pub mod trie_constants { const FIRST_PREFIX: u8 = 0b_00 << 6; - //pub const NIBBLE_SIZE_BOUND: usize = u16::max_value() as usize - 1; - pub const NIBBLE_SIZE_BOUND: usize = 4; pub const LEAF_PREFIX_MASK: u8 = 0b_01 << 6; pub const BRANCH_WITHOUT_MASK: u8 = 0b_10 << 6; pub const BRANCH_WITH_MASK: u8 = 0b_11 << 6; @@ -261,12 +258,10 @@ where /// Encode and allocate node type header (type and size), and partial value. /// It uses an iterator over encoded partial bytes as input. fn partial_from_iterator_encode>( - mut partial: I, - total_nibble_count: usize, + partial: I, + nibble_count: usize, node_kind: NodeKind, ) -> Vec { - let nibble_count = std::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, total_nibble_count); - let mut output = Vec::with_capacity(4 + (nibble_count / nibble_ops::NIBBLE_PER_BYTE)); match node_kind { NodeKind::Leaf => NodeHeader::Leaf(nibble_count).encode_to(&mut output), @@ -277,18 +272,7 @@ fn partial_from_iterator_encode>( NodeKind::HashedValueBranch => NodeHeader::HashedValueBranch(nibble_count).encode_to(&mut output), }; - if total_nibble_count != nibble_count { - // truncate key content. - let was_aligned = total_nibble_count % nibble_ops::NIBBLE_PER_BYTE == 0; - debug_assert!(nibble_count % nibble_ops::NIBBLE_PER_BYTE == 0); - let mut nb_bytes = trie_constants::NIBBLE_SIZE_BOUND / nibble_ops::NIBBLE_PER_BYTE; - if trie_constants::NIBBLE_SIZE_BOUND % nibble_ops::NIBBLE_PER_BYTE > 0 { - nb_bytes += 1; - } - output.extend(partial.take(nb_bytes)); - } else { - output.extend(partial); - } + output.extend(partial); output } @@ -391,8 +375,6 @@ pub(crate) fn size_and_prefix_iterator( prefix: u8, prefix_mask: usize, ) -> impl Iterator { - let size = std::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, size); - let max_value = 255u8 >> prefix_mask; let l1 = std::cmp::min(max_value as usize - 1, size); let (first_byte, mut rem) = if size == l1 { @@ -435,14 +417,13 @@ fn decode_size(first: u8, input: &mut impl Input, prefix_mask: usize) -> Result< return Ok(result) } result -= 1; - while result <= trie_constants::NIBBLE_SIZE_BOUND { + loop { let n = input.read_byte()? as usize; if n < 255 { return Ok(result + n + 1) } result += 255; } - Ok(trie_constants::NIBBLE_SIZE_BOUND) } /// Reference implementation of a `TrieStream` without extension. @@ -454,7 +435,7 @@ pub struct ReferenceTrieStreamNoExt { /// Create a leaf/branch node, encoding a number of nibbles. fn fuse_nibbles_node<'a>(nibbles: &'a [u8], kind: NodeKind) -> impl Iterator + 'a { - let size = std::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, nibbles.len()); + let size = nibbles.len(); let iter_start = match kind { NodeKind::Leaf => size_and_prefix_iterator(size, trie_constants::LEAF_PREFIX_MASK, 2), diff --git a/trie-db/test/src/triedbmut.rs b/trie-db/test/src/triedbmut.rs index 0b2d2c52..cb7c0d28 100644 --- a/trie-db/test/src/triedbmut.rs +++ b/trie-db/test/src/triedbmut.rs @@ -226,25 +226,6 @@ fn insert_replace_root_internal() { ); } -#[test] -fn oversized_key() { - use reference_trie::HashedValueNoExtThreshold as Layout; - let check = |keysize: usize, oversized: Option| { - let mut memdb = PrefixedMemoryDB::::default(); - let mut root = Default::default(); - let mut t = TrieDBMutBuilder::::new(&mut memdb, &mut root).build(); - t.insert(&vec![0x01u8; keysize][..], &[0x01u8, 0x23]).unwrap(); - std::mem::drop(t); - let t = TrieDBBuilder::::new(&memdb, &root).build(); - let mut truncated_key = vec![0x01u8; keysize - oversized.unwrap_or(0)]; - assert_eq!(t.get(&truncated_key[..]).unwrap(), Some(vec![0x01u8, 0x23])); - }; - check(4 as usize / 2, None); // ok - check(4 as usize / 2 + 1, Some(1)); // Oversized, rem last two nibble - //check(u16::MAX as usize / 2, None); // ok - //check(u16::MAX as usize / 2 + 1, Some(1)); // Oversized, rem last two nibble -} - test_layouts!(insert_make_branch_root, insert_make_branch_root_internal); fn insert_make_branch_root_internal() { let mut memdb = PrefixedMemoryDB::::default();