fix: manage node patch serialization for diffs range [1, 256], #6593

Author: federico-stacks

stackslib/src/chainstate/stacks/index/node.rs

@@ -1113,20 +1113,50 @@ impl fmt::Debug for TrieNodePatch {
 }
 
 impl StacksMessageCodec for TrieNodePatch {
+    /// Serializes this [`TrieNodePatch`] to the given writer, with the following format:
+    ///
+    /// 0    1        1+P      2+P              2+P+N
+    /// |----|--------|----------|----------------|
+    ///   id     ptr    diff len     ptr diffs
+    ///   (1)    (P)       (1)          (N)
+    ///
+    /// where:
+    /// - `id` is [`TrieNodeID::Patch`]
+    /// - `ptr` is a compressed [`TriePtr`]
+    /// - `diff len` is the number of diffs, serialized as `len - 1`
+    /// - `ptr diffs` are the patch diffs written in compressed format
+    ///
+    /// # Invariants
+    ///
+    /// The number of diffs must be in the range `1..=256`. A patch is valid only
+    /// if it contains at least one diff (see the factory methods
+    /// [`TrieNodePatch::try_from_nodetype`] and [`TrieNodePatch::try_from_patch`]).
+    ///
+    /// To fit in a `u8`, the diff count is normalized to `len - 1` before
+    /// serialization.
+    ///
+    /// # Errors
+    ///
+    /// Returns `Err(codec_error::SerializeError)` if:
+    /// * Writing to `fd` fails.
+    /// * `ptr` fails to serialize.
+    /// * Any pointer in `ptr diffs` fails to serialize.
+    /// * The diff count is `0` or greater than `256`.
     fn consensus_serialize<W: Write>(&self, fd: &mut W) -> Result<(), codec_error> {
         write_next(fd, &(TrieNodeID::Patch as u8))?;
         self.ptr
             .write_bytes_compressed(fd)
             .map_err(|e| codec_error::SerializeError(format!("Failed to serialize .ptr: {e:?}")))?;
 
         let num_ptrs = self.ptr_diff.len();
-        if num_ptrs >= 256 {
-            return Err(codec_error::SerializeError(
-                "Cannot serialize TrieNodePatch with more than 256 ptrs".to_string(),
-            ));
+        if num_ptrs == 0 || num_ptrs > 256 {
+            return Err(codec_error::SerializeError(format!(
+                "Cannot serialize TrieNodePatch with invalid ptrs len {num_ptrs} (expected 1..=256)"
+            )));
         }
-        // SAFETY: checked that num_ptrs < 256
-        let num_ptrs_u8 = u8::try_from(num_ptrs).expect("infallible");
+        // normalize num_ptrs to range [0, 255] to fit in u8
+        let num_ptrs_norm = num_ptrs.checked_sub(1).expect("infallible");
+        let num_ptrs_u8 = u8::try_from(num_ptrs_norm).expect("infallible");
         write_next(fd, &num_ptrs_u8).map_err(|e| {
             codec_error::SerializeError(format!("Failed to serialize .ptr_diff.len(): {e:?}"))
         })?;
@@ -1139,6 +1169,22 @@ impl StacksMessageCodec for TrieNodePatch {
         Ok(())
     }
 
+    /// Deserializes a [`TrieNodePatch`] from the given reader.
+    ///
+    /// This method expects the byte stream to be in the exact format produced by
+    /// [`TrieNodePatch::consensus_serialize`] (see that method for the detailed
+    /// wire format description)
+    ///
+    /// During deserialization, the stored diff length is de-normalized by
+    /// adding `1`, reversing the `len - 1` normalization applied during
+    /// serialization.
+    ///
+    /// # Errors
+    ///
+    /// Returns `Err(codec_error::DeserializeError)` if:
+    /// * The node identifier does not match [`TrieNodeID::Patch`].
+    /// * Reading from `fd` fails.
+    /// * The pointer or any pointer diff fails to deserialize.
     fn consensus_deserialize<R: Read>(fd: &mut R) -> Result<Self, codec_error> {
         let id: u8 = read_next(fd)?;
         if id != TrieNodeID::Patch as u8 {
@@ -1148,8 +1194,11 @@ impl StacksMessageCodec for TrieNodePatch {
         }
 
         let ptr = TriePtr::read_bytes_compressed(fd)?;
-        let num_ptrs: u8 = read_next(fd)?;
-        let num_ptrs = usize::try_from(num_ptrs).expect("infallible");
+        let num_ptrs_u8: u8 = read_next(fd)?;
+        let num_ptrs_norm = usize::try_from(num_ptrs_u8).expect("infallible");
+        // denormalize num_ptrs to range [1, 256] (reversing the -1 introduced during serialization)
+        let num_ptrs = num_ptrs_norm.checked_add(1).expect("infallible");
+
         let mut ptr_diff: Vec<TriePtr> = Vec::with_capacity(num_ptrs);
         for _ in 0..num_ptrs {
             ptr_diff.push(TriePtr::read_bytes_compressed(fd)?);
@@ -1301,6 +1350,7 @@ impl TrieNodePatch {
         new_node: &TrieNodeType,
     ) -> Option<Self> {
         if clear_ctrl_bits(old_node.id()) != clear_ctrl_bits(new_node.id()) {
+            trace!("Cannot produce TrieNodePatch: old node and new node are not the same type!");
             return None;
         }
 
@@ -1320,9 +1370,11 @@ impl TrieNodePatch {
             (_, _) => None,
         };
         let Some(patch) = patch_opt else {
+            trace!("Cannot produce TrieNodePatch: old node and new node are type leaf!");
             return None;
         };
         if patch.ptr_diff.len() == 0 {
+            trace!("Cannot produce TrieNodePatch: patch has no diffs!");
             return None;
         }
         Some(patch)
@@ -1335,6 +1387,7 @@ impl TrieNodePatch {
         new_node: &TrieNodeType,
     ) -> Option<Self> {
         if clear_ctrl_bits(old_patch.ptr.id) != clear_ctrl_bits(new_node.id()) {
+            trace!("Cannot produce TrieNodePatch: old node and new node are not the same type!");
             return None;
         }
 
@@ -1344,6 +1397,7 @@ impl TrieNodePatch {
             ptr_diff,
         };
         if patch.ptr_diff.len() == 0 {
+            trace!("Cannot produce TrieNodePatch: patch has no diffs!");
             return None;
         }
         return Some(patch);

stackslib/src/chainstate/stacks/index/test/node.rs

@@ -21,6 +21,7 @@ use std::io::Cursor;
 
 use super::*;
 use crate::chainstate::stacks::index::*;
+use crate::codec::{Error as codec_error, StacksMessageCodec};
 
 #[test]
 fn trieptr_to_bytes() {
@@ -5043,3 +5044,122 @@ fn trie_cursor_walk_32() {
 
     dump_trie(&mut trie_io);
 }
+
+#[test]
+fn trie_node_patch_try_from_nodetype_returns_none_when_no_diffs() {
+    let node = TrieNodeType::Node4(TrieNode4::new(&[1]));
+
+    let old_node_ptr = TriePtr::default();
+    let old_node = &node;
+    let new_node = &node;
+    let result = TrieNodePatch::try_from_nodetype(old_node_ptr, old_node, new_node);
+
+    assert!(
+        result.is_none(),
+        "None because the computed patch has no diffs"
+    );
+}
+
+#[test]
+fn trie_node_patch_try_from_patch_returns_none_when_no_diffs() {
+    let old_patch_ptr = TriePtr::new(TrieNodeID::Node4 as u8, 0, 0);
+    let old_patch = TrieNodePatch {
+        ptr: old_patch_ptr.clone(),
+        ptr_diff: vec![],
+    };
+    let new_node = TrieNodeType::Node4(TrieNode4::new(&[1]));
+    let result = TrieNodePatch::try_from_patch(old_patch_ptr, &old_patch, &new_node);
+
+    assert!(
+        result.is_none(),
+        "None because the computed patch has no diffs"
+    );
+}
+
+#[test]
+fn trie_node_patch_serialize_ok() {
+    let patch_node = TrieNodePatch {
+        ptr: TriePtr::new(1, 10, 0),
+        ptr_diff: vec![TriePtr::new(1, 20, 0).clone(); 1],
+    };
+
+    let mut buffer = Cursor::new(Vec::new());
+    patch_node
+        .consensus_serialize(&mut buffer)
+        .expect("serialization should be ok");
+
+    // To fit in 1 byte, diff count is serialized 0-based (where 0 => 1 and 255 => 256)
+    let diff_count = 0u8;
+    assert_eq!(
+        vec![6, 65, 10, 0, 0, 0, 0, diff_count, 65, 20, 0, 0, 0, 0],
+        buffer.into_inner(),
+    );
+}
+
+#[test]
+fn trie_node_patch_serialize_fails_with_ptr_diffs_len_0() {
+    let patch_node = TrieNodePatch {
+        ptr: TriePtr::default(),
+        ptr_diff: vec![],
+    };
+
+    let mut buffer = Cursor::new(Vec::new());
+    let error = patch_node
+        .consensus_serialize(&mut buffer)
+        .expect_err("serialization should fail");
+
+    assert!(
+        matches!(&error, codec_error::SerializeError(msg) if msg.contains("len 0")),
+        "instead got: {error}"
+    );
+}
+
+#[test]
+fn trie_node_patch_serialize_ok_with_ptr_diffs_len_256() {
+    let patch_node = TrieNodePatch {
+        ptr: TriePtr::default(),
+        ptr_diff: vec![TriePtr::default(); 256],
+    };
+
+    let mut buffer = Cursor::new(Vec::new());
+    let result = patch_node.consensus_serialize(&mut buffer);
+    assert!(
+        result.is_ok(),
+        "Got Error: {}",
+        result.unwrap_err().to_string()
+    );
+}
+
+#[test]
+fn trie_node_patch_serialize_fails_with_ptr_diffs_len_257() {
+    let patch_node = TrieNodePatch {
+        ptr: TriePtr::default(),
+        ptr_diff: vec![TriePtr::default(); 257],
+    };
+
+    let mut buffer = Cursor::new(Vec::new());
+    let error = patch_node
+        .consensus_serialize(&mut buffer)
+        .expect_err("serialization should fail");
+
+    assert!(
+        matches!(&error, codec_error::SerializeError(msg) if msg.contains("len 257")),
+        "instead got: {error}"
+    );
+}
+
+#[test]
+fn trie_node_patch_deserialize_ok_with_ptr_diffs_len_1() {
+    // To fit in 1 byte, diff count is serialized 0-based (where 0 => 1 and 255 => 256)
+    let diff_count = 0u8;
+    let mut buffer = Cursor::new(vec![6, 65, 10, 0, 0, 0, 0, diff_count, 65, 20, 0, 0, 0, 0]);
+
+    let patch_node =
+        TrieNodePatch::consensus_deserialize(&mut buffer).expect("deserialization should be ok");
+
+    let expected = TrieNodePatch {
+        ptr: TriePtr::new(1, 10, 0),
+        ptr_diff: vec![TriePtr::new(1, 20, 0); 1],
+    };
+    assert_eq!(expected, patch_node);
+}