cfx_storage/impls/merkle_patricia_trie/
compressed_path.rs

1// Copyright 2019 Conflux Foundation. All rights reserved.
2// Conflux is free software and distributed under GNU General Public License.
3// See http://www.gnu.org/licenses/
4
5pub trait CompressedPathTrait: Debug {
6    fn path_slice(&self) -> &[u8];
7    fn path_mask(&self) -> u8;
8
9    fn path_size(&self) -> u16 { self.path_slice().len() as u16 }
10
11    fn path_steps(&self) -> u16 {
12        CompressedPathRaw::calculate_path_steps(
13            self.path_size(),
14            self.path_mask(),
15        )
16    }
17
18    fn as_ref(&self) -> CompressedPathRef<'_> {
19        CompressedPathRef {
20            path_slice: self.path_slice(),
21            path_mask: self.path_mask(),
22        }
23    }
24
25    fn rlp_append(&self, s: &mut RlpStream) {
26        s.begin_list(2);
27        s.append(&self.path_mask()).append(&self.path_slice());
28    }
29}
30
31impl CompressedPathTrait for [u8] {
32    fn path_slice(&self) -> &[u8] { self }
33
34    fn path_mask(&self) -> u8 { CompressedPathRaw::NO_MISSING_NIBBLE }
35}
36
37impl<'a> CompressedPathTrait for &'a [u8] {
38    fn path_slice(&self) -> &[u8] { self }
39
40    fn path_mask(&self) -> u8 { CompressedPathRaw::NO_MISSING_NIBBLE }
41}
42
43#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
44pub struct CompressedPathRef<'a> {
45    pub path_slice: &'a [u8],
46    path_mask: u8,
47}
48
49impl<'a> CompressedPathRef<'a> {
50    pub fn new(path_slice: &'a [u8], path_mask: u8) -> Self {
51        Self {
52            path_slice,
53            path_mask,
54        }
55    }
56}
57
58#[derive(Default)]
59pub struct CompressedPathRaw {
60    path_size: u16,
61    path: MaybeInPlaceByteArray,
62    path_mask: u8,
63    pub byte_array_memory_manager:
64        FieldsOffsetMaybeInPlaceByteArrayMemoryManager<
65            u16,
66            TrivialSizeFieldConverterU16,
67            CompressedPathRawByteArrayMemoryManager,
68            CompressedPathRawByteArrayMemoryManager,
69        >,
70}
71
72/// CompressedPathRaw is Send + Sync.
73unsafe impl Send for CompressedPathRaw {}
74unsafe impl Sync for CompressedPathRaw {}
75
76#[cfg(test)]
77mod tests {
78    use super::{super::maybe_in_place_byte_array::*, *};
79    use rlp::{Decodable, Rlp, RlpStream};
80
81    #[test]
82    fn test_compressed_path_raw_memory_manager_size() {
83        assert_eq!(
84            std::mem::size_of::<
85                FieldsOffsetMaybeInPlaceByteArrayMemoryManager<
86                    u16,
87                    TrivialSizeFieldConverterU16,
88                    CompressedPathRawByteArrayMemoryManager,
89                    CompressedPathRawByteArrayMemoryManager,
90                >,
91            >(),
92            0
93        );
94    }
95
96    /// RLP-encode a raw `(mask, slice)` pair, bypassing the constructors so
97    /// tests can craft over-long / malformed inputs they can't produce.
98    fn encode_path(path_mask: u8, path_slice: &[u8]) -> Vec<u8> {
99        let mut s = RlpStream::new_list(2);
100        s.append(&path_mask);
101        s.append(&path_slice);
102        s.out().to_vec()
103    }
104
105    fn decode_path(bytes: &[u8]) -> Result<CompressedPathRaw, DecoderError> {
106        CompressedPathRaw::decode(&Rlp::new(bytes))
107    }
108
109    #[test]
110    fn test_decode_normal_path_roundtrips() {
111        let slice = vec![0xabu8; 32];
112        let bytes = encode_path(CompressedPathRaw::NO_MISSING_NIBBLE, &slice);
113        let decoded = decode_path(&bytes).expect("valid path decodes");
114        assert_eq!(decoded.path_slice(), slice.as_slice());
115    }
116
117    #[test]
118    fn test_decode_path_at_cap_is_accepted_without_truncation() {
119        let slice = vec![0u8; CompressedPathRaw::MAX_PATH_BYTES];
120        let bytes = encode_path(CompressedPathRaw::NO_MISSING_NIBBLE, &slice);
121        let decoded = decode_path(&bytes).expect("path at cap decodes");
122        assert_eq!(
123            decoded.path_slice().len(),
124            CompressedPathRaw::MAX_PATH_BYTES
125        );
126    }
127
128    #[test]
129    fn test_decode_over_cap_path_is_rejected() {
130        let slice = vec![0u8; CompressedPathRaw::MAX_PATH_BYTES + 1];
131        let bytes = encode_path(CompressedPathRaw::NO_MISSING_NIBBLE, &slice);
132        assert!(decode_path(&bytes).is_err());
133    }
134
135    #[test]
136    fn test_decode_mask_handling() {
137        // Empty slice + non-zero mask underflows path_steps()/walk(); reject.
138        let bytes = encode_path(CompressedPathRaw::second_nibble_mask(), &[]);
139        assert!(decode_path(&bytes).is_err());
140
141        let bytes = encode_path(CompressedPathRaw::NO_MISSING_NIBBLE, &[]);
142        assert!(decode_path(&bytes).is_ok());
143
144        // The empty-only restriction must not reject a non-empty masked path.
145        let bytes = encode_path(0xff, &[1u8]);
146        assert!(decode_path(&bytes).is_ok());
147    }
148
149    #[test]
150    fn test_path_steps_at_cap_does_not_overflow() {
151        let path = CompressedPathRaw::new_zeroed(
152            CompressedPathRaw::MAX_PATH_BYTES as u16,
153            CompressedPathRaw::NO_MISSING_NIBBLE,
154        );
155        assert_eq!(
156            path.path_steps(),
157            (CompressedPathRaw::MAX_PATH_BYTES as u16) * 2
158        );
159    }
160}
161
162make_parallel_field_maybe_in_place_byte_array_memory_manager!(
163    CompressedPathRawByteArrayMemoryManager,
164    CompressedPathRaw,
165    byte_array_memory_manager,
166    path,
167    path_size: u16,
168    TrivialSizeFieldConverterU16,
169);
170
171impl CompressedPathRaw {
172    const BITS_0_3_MASK: u8 = 0x0f;
173    const BITS_4_7_MASK: u8 = 0xf0;
174    /// Maximum bytes in a compressed path / MPT key.
175    ///
176    /// A byte is two nibbles and `path_steps()` (the nibble count) is a `u16`,
177    /// so a path longer than `u16::MAX / 2` overflows `path_size * 2` in
178    /// [`Self::calculate_path_steps`]; one over 65535 bytes also truncates the
179    /// `u16` `path_size`, making `Drop` free a `Layout` it never allocated
180    /// (UB). Such a path can't exist in a validly-built snapshot anyway — its
181    /// merkle would overflow `path_steps` on the honest builder — so capping
182    /// the receive path rejects only malformed input, not the tens-of-bytes
183    /// keys real state uses.
184    pub const MAX_PATH_BYTES: usize = (u16::MAX / 2) as usize;
185    pub const NO_MISSING_NIBBLE: u8 = 0;
186
187    /// Validate an untrusted `(path_size, path_mask)` before it reaches the
188    /// path machinery. Used by both [`CompressedPathRaw`] decoders (RLP,
189    /// serde).
190    fn check_path_encoding(
191        path_size: usize, path_mask: u8,
192    ) -> Result<(), &'static str> {
193        if path_size > Self::MAX_PATH_BYTES {
194            return Err("CompressedPathRaw path too long");
195        }
196        // An empty path has no nibble, so a non-zero begin/end mask makes
197        // `calculate_path_steps`/`walk` compute `0 - 1` and panic.
198        if path_size == 0 && path_mask != Self::NO_MISSING_NIBBLE {
199            return Err("CompressedPathRaw empty path with non-zero mask");
200        }
201        Ok(())
202    }
203}
204
205impl<'a> CompressedPathTrait for CompressedPathRef<'a> {
206    fn path_slice(&self) -> &[u8] { self.path_slice }
207
208    fn path_mask(&self) -> u8 { self.path_mask }
209
210    fn path_size(&self) -> u16 { self.path_slice.len() as u16 }
211}
212
213impl CompressedPathTrait for CompressedPathRaw {
214    fn path_slice(&self) -> &[u8] {
215        self.path.get_slice(self.path_size as usize)
216    }
217
218    fn path_mask(&self) -> u8 { self.path_mask }
219}
220
221impl<'a> From<&'a [u8]> for CompressedPathRaw {
222    fn from(x: &'a [u8]) -> Self {
223        CompressedPathRaw::new(x, Self::NO_MISSING_NIBBLE)
224    }
225}
226
227impl<'a> From<CompressedPathRef<'a>> for CompressedPathRaw {
228    fn from(x: CompressedPathRef<'a>) -> Self {
229        CompressedPathRaw::new(x.path_slice, x.path_mask)
230    }
231}
232
233impl CompressedPathRaw {
234    /// Create a new CompressedPathRaw from valid (path_slice, path_mask)
235    /// combination.
236    pub fn new(path_slice: &[u8], path_mask: u8) -> Self {
237        let path_size = path_slice.len();
238
239        Self {
240            path_size: path_size as u16,
241            path: MaybeInPlaceByteArray::copy_from(path_slice, path_size),
242            path_mask,
243            byte_array_memory_manager: Default::default(),
244        }
245    }
246
247    #[inline]
248    fn last_byte_mut(&mut self) -> &mut u8 {
249        // Safe, because the index is valid.
250        unsafe {
251            self.path
252                .get_slice_mut(self.path_size as usize)
253                .get_unchecked_mut(self.path_size as usize - 1)
254        }
255    }
256
257    pub fn new_and_apply_mask(path_slice: &[u8], path_mask: u8) -> Self {
258        let path_size = path_slice.len();
259        let mut ret = Self {
260            path_size: path_size as u16,
261            path: MaybeInPlaceByteArray::copy_from(path_slice, path_size),
262            path_mask,
263            byte_array_memory_manager: Default::default(),
264        };
265        if path_size > 0 {
266            // 0xf* -> no second nibble
267            // 0x0* -> has second nibble
268            // 0xf0 -> 0x0f -> &= 0xf0
269            // 0x00 -> 0xff -> &= 0xff
270            *ret.last_byte_mut() &= !Self::first_nibble(path_mask);
271        }
272
273        ret
274    }
275
276    pub fn new_zeroed(path_size: u16, path_mask: u8) -> Self {
277        Self {
278            path_size,
279            path: MaybeInPlaceByteArray::new_zeroed(path_size as usize),
280            path_mask,
281            byte_array_memory_manager: Default::default(),
282        }
283    }
284
285    #[inline]
286    pub const fn first_nibble_mask() -> u8 { Self::BITS_0_3_MASK }
287
288    #[inline]
289    pub const fn second_nibble_mask() -> u8 { Self::BITS_4_7_MASK }
290
291    #[inline]
292    fn calculate_path_steps(path_size: u16, path_mask: u8) -> u16 {
293        path_size * 2
294            - (Self::clear_second_nibble(path_mask) != 0) as u16
295            - (Self::second_nibble(path_mask) != 0) as u16
296    }
297
298    #[inline]
299    pub fn from_first_nibble(x: u8) -> u8 { x << 4 }
300
301    #[inline]
302    pub fn first_nibble(x: u8) -> u8 { x >> 4 }
303
304    #[inline]
305    pub fn clear_second_nibble(x: u8) -> u8 { x & Self::BITS_4_7_MASK }
306
307    #[inline]
308    pub fn second_nibble(x: u8) -> u8 { x & Self::BITS_0_3_MASK }
309
310    #[inline]
311    pub fn set_second_nibble(x: u8, second_nibble: u8) -> u8 {
312        Self::clear_second_nibble(x) | second_nibble
313    }
314
315    #[inline]
316    pub fn has_second_nibble(path_mask: u8) -> bool {
317        Self::clear_second_nibble(path_mask)
318            == CompressedPathRaw::NO_MISSING_NIBBLE
319    }
320
321    #[inline]
322    pub fn no_second_nibble(path_mask: u8) -> bool {
323        Self::clear_second_nibble(path_mask)
324            != CompressedPathRaw::NO_MISSING_NIBBLE
325    }
326
327    pub fn extend_path<X: CompressedPathTrait>(x: &X, child_index: u8) -> Self {
328        let new_size;
329        let path_mask;
330        // Need to extend the length.
331        let x_path_mask = x.path_mask();
332        if Self::has_second_nibble(x_path_mask) {
333            new_size = x.path_size() + 1;
334            path_mask = x_path_mask | Self::second_nibble_mask();
335        } else {
336            new_size = x.path_size();
337            path_mask = Self::second_nibble(x_path_mask);
338        }
339        let mut ret = Self::new_zeroed(new_size, path_mask);
340        ret.path.get_slice_mut(new_size as usize)[0..x.path_size() as usize]
341            .copy_from_slice(x.path_slice());
342        // The last byte will be a half-byte.
343        if Self::has_second_nibble(x_path_mask) {
344            *ret.last_byte_mut() = Self::from_first_nibble(child_index);
345        } else {
346            let last_byte = *ret.last_byte_mut();
347            *ret.last_byte_mut() =
348                Self::set_second_nibble(last_byte, child_index);
349        }
350
351        ret
352    }
353
354    /// y must be a valid path following x. i.e. when x ends with a full byte, y
355    /// must be non-empty and start with nibble child_index.
356    pub fn join_connected_paths<
357        X: CompressedPathTrait,
358        Y: CompressedPathTrait,
359    >(
360        x: &X, child_index: u8, y: &Y,
361    ) -> Self {
362        let x_slice = x.path_slice();
363        let x_slice_len = x_slice.len();
364        let x_path_mask = x.path_mask();
365        let y_slice = y.path_slice();
366
367        // TODO(yz): it happens to be the same no matter what end_mask of x is,
368        // because u8 = 2 nibbles. When we switch to u32 as path unit
369        // the concated size may vary.
370        let size = x_slice_len + y_slice.len();
371
372        let mut path;
373        {
374            let slice;
375            // TODO: resolve warnings in unsafe code.
376            #[allow(clippy::uninit_vec)]
377            unsafe {
378                if size > MaybeInPlaceByteArray::MAX_INPLACE_SIZE {
379                    // Create uninitialized vector.
380                    let mut value = Vec::with_capacity(size);
381                    value.set_len(size);
382                    let mut value_box = value.into_boxed_slice();
383
384                    let ptr = value_box.as_mut_ptr();
385                    // Don't free the buffer since it's stored in the return
386                    // value.
387                    let _ = Box::into_raw(value_box);
388                    path = MaybeInPlaceByteArray { ptr };
389                    slice = std::slice::from_raw_parts_mut(ptr, size);
390                } else {
391                    let in_place: [u8;
392                        MaybeInPlaceByteArray::MAX_INPLACE_SIZE] =
393                        Default::default();
394                    path = MaybeInPlaceByteArray { in_place };
395                    slice = &mut path.in_place[0..size];
396                }
397            }
398
399            if Self::has_second_nibble(x_path_mask) {
400                slice[0..x_slice_len].copy_from_slice(x_slice);
401            } else {
402                slice[0..x_slice_len - 1]
403                    .copy_from_slice(&x_slice[0..x_slice_len - 1]);
404                slice[x_slice_len - 1] = CompressedPathRaw::set_second_nibble(
405                    x_slice[x_slice_len - 1],
406                    child_index,
407                );
408            }
409            slice[x_slice_len..].copy_from_slice(y_slice);
410        }
411
412        Self {
413            path_size: size as u16,
414            path,
415            path_mask: Self::set_second_nibble(
416                y.path_mask(),
417                CompressedPathRaw::second_nibble(x_path_mask),
418            ),
419            byte_array_memory_manager: Default::default(),
420        }
421    }
422}
423
424impl Clone for CompressedPathRaw {
425    fn clone(&self) -> Self {
426        Self {
427            path_mask: self.path_mask,
428            path_size: self.path_size,
429            path: MaybeInPlaceByteArray::clone(
430                &self.path,
431                self.path_size as usize,
432            ),
433            byte_array_memory_manager: Default::default(),
434        }
435    }
436}
437
438impl<'a> Encodable for CompressedPathRef<'a> {
439    fn rlp_append(&self, s: &mut RlpStream) {
440        CompressedPathTrait::rlp_append(self, s);
441    }
442}
443
444impl Encodable for CompressedPathRaw {
445    fn rlp_append(&self, s: &mut RlpStream) {
446        CompressedPathTrait::rlp_append(self, s);
447    }
448}
449
450impl Decodable for CompressedPathRaw {
451    fn decode(rlp: &Rlp) -> Result<Self, DecoderError> {
452        let path_mask = rlp.val_at::<u8>(0)?;
453        let path_slice = rlp.val_at::<Vec<u8>>(1)?;
454        Self::check_path_encoding(path_slice.len(), path_mask)
455            .map_err(DecoderError::Custom)?;
456        Ok(CompressedPathRaw::new(path_slice.as_slice(), path_mask))
457    }
458}
459
460impl Serialize for CompressedPathRaw {
461    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
462    where S: Serializer {
463        let path_mask = self.path_mask();
464        let path_mask = format!("0x{:x}", path_mask);
465        let path_slice = self.path_slice();
466        let mut struc = serializer.serialize_struct("CompressedPathRaw", 2)?;
467        struc.serialize_field("pathMask", &path_mask)?;
468        struc.serialize_field(
469            "pathSlice",
470            &("0x".to_owned() + path_slice.to_hex::<String>().as_ref()),
471        )?;
472
473        struc.end()
474    }
475}
476
477impl<'a> Deserialize<'a> for CompressedPathRaw {
478    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
479    where D: Deserializer<'a> {
480        let (path_mask, path_slice) = deserializer.deserialize_struct(
481            "CompressedPathRaw",
482            FIELDS,
483            CompressedPathRawVisitor,
484        )?;
485
486        Self::check_path_encoding(path_slice.len(), path_mask)
487            .map_err(de::Error::custom)?;
488        Ok(CompressedPathRaw::new(&path_slice[..], path_mask))
489    }
490}
491
492const FIELDS: &'static [&'static str] = &["pathMask", "pathSlice"];
493
494enum Field {
495    Mask,
496    Slice,
497}
498
499impl<'de> Deserialize<'de> for Field {
500    fn deserialize<D>(deserializer: D) -> Result<Field, D::Error>
501    where D: Deserializer<'de> {
502        struct FieldVisitor;
503
504        impl<'de> Visitor<'de> for FieldVisitor {
505            type Value = Field;
506
507            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
508                formatter.write_str("`pathMask` or `pathSlice`")
509            }
510
511            fn visit_str<E>(self, value: &str) -> Result<Field, E>
512            where E: de::Error {
513                match value {
514                    "pathMask" => Ok(Field::Mask),
515                    "pathSlice" => Ok(Field::Slice),
516                    _ => Err(de::Error::unknown_field(value, FIELDS)),
517                }
518            }
519        }
520
521        deserializer.deserialize_identifier(FieldVisitor)
522    }
523}
524
525struct CompressedPathRawVisitor;
526impl<'a> Visitor<'a> for CompressedPathRawVisitor {
527    type Value = (u8, Vec<u8>);
528
529    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
530        write!(formatter, "Struct CompressedPath")
531    }
532
533    fn visit_map<V>(self, mut visitor: V) -> Result<Self::Value, V::Error>
534    where V: MapAccess<'a> {
535        let mut path_mask = None;
536        let mut path_slice = None;
537
538        while let Some(key) = visitor.next_key()? {
539            match key {
540                Field::Mask => {
541                    if path_mask.is_some() {
542                        return Err(de::Error::duplicate_field("pathMask"));
543                    }
544
545                    path_mask = Some(visitor.next_value()?);
546                }
547                Field::Slice => {
548                    if path_slice.is_some() {
549                        return Err(de::Error::duplicate_field("pathSlice"));
550                    }
551
552                    path_slice = Some(visitor.next_value()?);
553                }
554            }
555        }
556
557        let path_mask: String =
558            path_mask.ok_or_else(|| de::Error::missing_field("pathMask"))?;
559
560        let path_mask = if let Some(s) = path_mask.strip_prefix("0x") {
561            u8::from_str_radix(&s, 16).map_err(|e| {
562                de::Error::custom(format!("pathMask: invalid hex: {}", e))
563            })?
564        } else {
565            return Err(de::Error::custom(
566                "pathMask: invalid format. Expected a 0x-prefixed hex string",
567            ));
568        };
569
570        let path_slice: String =
571            path_slice.ok_or_else(|| de::Error::missing_field("pathSlice"))?;
572
573        let path_slice: Vec<u8> = if let (Some(s), true) =
574            (path_slice.strip_prefix("0x"), path_slice.len() & 1 == 0)
575        {
576            FromHex::from_hex(s).map_err(|e| {
577                de::Error::custom(format!("pathSlice: invalid hex: {}", e))
578            })?
579        } else {
580            return Err(de::Error::custom("pathSlice: invalid format. Expected a 0x-prefixed hex string with even length"));
581        };
582
583        Ok((path_mask, path_slice))
584    }
585}
586
587impl PartialEq<Self> for CompressedPathRaw {
588    fn eq(&self, other: &Self) -> bool { self.as_ref().eq(&other.as_ref()) }
589}
590
591impl Eq for CompressedPathRaw {}
592
593impl Hash for CompressedPathRaw {
594    fn hash<H: Hasher>(&self, state: &mut H) { self.as_ref().hash(state) }
595}
596
597impl Debug for CompressedPathRaw {
598    fn fmt(&self, f: &mut Formatter) -> Result<(), Error> {
599        self.as_ref().fmt(f)
600    }
601}
602
603impl CompressedPathRaw {
604    pub fn path_slice_mut(&mut self) -> &mut [u8] {
605        self.path.get_slice_mut(self.path_size as usize)
606    }
607}
608
609impl<'a> PartialEq<Self> for dyn CompressedPathTrait + 'a {
610    fn eq(&self, other: &(dyn CompressedPathTrait + 'a)) -> bool {
611        self.as_ref().eq(&other.as_ref())
612    }
613}
614
615impl<'a> Eq for dyn CompressedPathTrait + 'a {}
616
617impl<'a> Hash for dyn CompressedPathTrait + 'a {
618    fn hash<H: Hasher>(&self, state: &mut H) { self.as_ref().hash(state) }
619}
620
621impl<'a> Borrow<dyn CompressedPathTrait + 'a> for CompressedPathRaw {
622    fn borrow(&self) -> &(dyn CompressedPathTrait + 'a) { self }
623}
624
625use super::maybe_in_place_byte_array::*;
626use rlp::*;
627use rustc_hex::{FromHex, ToHex};
628use serde::{
629    de::{self, MapAccess, Visitor},
630    ser::SerializeStruct,
631    Deserialize, Deserializer, Serialize, Serializer,
632};
633
634use std::{
635    borrow::Borrow,
636    fmt::{self, Debug, Error, Formatter},
637    hash::{Hash, Hasher},
638    result::Result,
639};