Merge pull request #553 from Chia-Network/back-reference-length

back-reference path length

Author: arvidn

fuzz/Cargo.toml

@@ -98,3 +98,9 @@ name = "incremental-serializer"
 path = "fuzz_targets/incremental_serializer.rs"
 test = false
 doc = false
+
+[[bin]]
+name = "serializer-cmp"
+path = "fuzz_targets/serializer_cmp.rs"
+test = false
+doc = false

fuzz/fuzz_targets/serializer_cmp.rs

@@ -0,0 +1,103 @@
+#![no_main]
+
+mod make_tree;
+mod node_eq;
+
+use clvmr::allocator::{Allocator, NodePtr, SExp};
+use clvmr::serde::node_from_bytes_backrefs;
+use clvmr::serde::write_atom::write_atom;
+use clvmr::serde::ReadCacheLookup;
+use clvmr::serde::{serialized_length, treehash, ObjectCache};
+use std::io;
+use std::io::Cursor;
+use std::io::Write;
+
+use node_eq::node_eq;
+
+use libfuzzer_sys::fuzz_target;
+
+const BACK_REFERENCE: u8 = 0xfe;
+const CONS_BOX_MARKER: u8 = 0xff;
+
+#[derive(PartialEq, Eq)]
+enum ReadOp {
+    Parse,
+    Cons,
+}
+
+// make sure back-references returned by ReadCacheLookup are smaller than the
+// node they reference
+pub fn compare_back_references(allocator: &Allocator, node: NodePtr) -> io::Result<Vec<u8>> {
+    let mut f = Cursor::new(Vec::new());
+
+    let mut read_op_stack: Vec<ReadOp> = vec![ReadOp::Parse];
+    let mut write_stack: Vec<NodePtr> = vec![node];
+
+    let mut read_cache_lookup = ReadCacheLookup::new();
+
+    let mut thc = ObjectCache::new(treehash);
+    let mut slc = ObjectCache::new(serialized_length);
+
+    while let Some(node_to_write) = write_stack.pop() {
+        let op = read_op_stack.pop();
+        assert!(op == Some(ReadOp::Parse));
+
+        let node_serialized_length = *slc
+            .get_or_calculate(allocator, &node_to_write, None)
+            .expect("couldn't calculate serialized length");
+        let node_tree_hash = thc
+            .get_or_calculate(allocator, &node_to_write, None)
+            .expect("can't get treehash");
+
+        let result1 = read_cache_lookup.find_path(node_tree_hash, node_serialized_length);
+        match result1 {
+            Some(path) => {
+                f.write_all(&[BACK_REFERENCE])?;
+                write_atom(&mut f, &path)?;
+                read_cache_lookup.push(*node_tree_hash);
+                {
+                    // make sure the path is never encoded as more bytes than
+                    // the node we're referencing
+                    use std::io::Write;
+                    let mut temp = Cursor::new(Vec::<u8>::new());
+                    temp.write_all(&[BACK_REFERENCE])?;
+                    write_atom(&mut temp, &path)?;
+                    let temp = temp.into_inner();
+                    assert!(temp.len() < node_serialized_length as usize);
+                }
+            }
+            None => match allocator.sexp(node_to_write) {
+                SExp::Pair(left, right) => {
+                    f.write_all(&[CONS_BOX_MARKER])?;
+                    write_stack.push(right);
+                    write_stack.push(left);
+                    read_op_stack.push(ReadOp::Cons);
+                    read_op_stack.push(ReadOp::Parse);
+                    read_op_stack.push(ReadOp::Parse);
+                }
+                SExp::Atom => {
+                    let atom = allocator.atom(node_to_write);
+                    write_atom(&mut f, atom.as_ref())?;
+                    read_cache_lookup.push(*node_tree_hash);
+                }
+            },
+        }
+        while let Some(ReadOp::Cons) = read_op_stack.last() {
+            read_op_stack.pop();
+            read_cache_lookup.pop2_and_cons();
+        }
+    }
+    Ok(f.into_inner())
+}
+
+// serializing with the regular compressed serializer should yield the same
+// result as using the incremental one (as long as it's in a single add() call).
+fuzz_target!(|data: &[u8]| {
+    let mut unstructured = arbitrary::Unstructured::new(data);
+    let mut allocator = Allocator::new();
+    let (program, _) = make_tree::make_tree(&mut allocator, &mut unstructured);
+
+    let b1 = compare_back_references(&allocator, program).unwrap();
+    let b2 = node_from_bytes_backrefs(&mut allocator, &b1).unwrap();
+    assert!(node_eq(&allocator, b2, program));
+});

src/serde/mod.rs

@@ -26,6 +26,7 @@ pub use de_tree::{parse_triples, ParsedTriple};
 pub use identity_hash::RandomState;
 pub use incremental::{Serializer, UndoState};
 pub use object_cache::{serialized_length, treehash, ObjectCache};
+pub use read_cache_lookup::ReadCacheLookup;
 pub use ser::{node_to_bytes, node_to_bytes_limit};
 pub use ser_br::{node_to_bytes_backrefs, node_to_bytes_backrefs_limit};
 pub use serialized_length::{serialized_length_atom, serialized_length_small_number};

src/serde/read_cache_lookup.rs

@@ -21,6 +21,7 @@ use bitvec::vec::BitVec;
 use std::collections::{HashMap, HashSet};
 
 use super::bytes32::{hash_blob, hash_blobs, Bytes32};
+use super::serialized_length::atom_length_bits;
 
 #[derive(Debug, Clone)]
 pub struct ReadCacheLookup {
@@ -153,7 +154,14 @@ impl ReadCacheLookup {
             let mut new_partial_paths = vec![];
             for (node, path) in partial_paths.iter_mut() {
                 if *node == self.root_hash {
-                    possible_responses.push(reversed_path_to_vec_u8(path));
+                    // make sure we never return a path that needs more (or the
+                    // same) bytes to serialize than the node we're referencing.
+                    if let Some(path_len) = atom_length_bits(path.len() as u64) {
+                        if path_len < max_bytes_for_path_encoding {
+                            let p = reversed_path_to_vec_u8(path);
+                            possible_responses.push(p);
+                        }
+                    }
                     continue;
                 }
 
@@ -162,12 +170,14 @@ impl ReadCacheLookup {
                     for (parent, direction) in items.iter() {
                         if *(self.count.get(parent).unwrap_or(&0)) > 0 && !seen_ids.contains(parent)
                         {
-                            if path.len() + 1 > max_path_length {
+                            if path.len() > max_path_length {
                                 return possible_responses;
                             }
-                            let mut new_path = path.clone();
-                            new_path.push(*direction);
-                            new_partial_paths.push((*parent, new_path));
+                            if path.len() < max_path_length {
+                                let mut new_path = path.clone();
+                                new_path.push(*direction);
+                                new_partial_paths.push((*parent, new_path));
+                            }
                         }
                         seen_ids.insert(parent);
                     }

src/serde/serialized_length.rs

@@ -21,6 +21,26 @@ pub fn serialized_length_small_number(val: u32) -> u32 {
     len_for_value(val) as u32 + 1
 }
 
+// given an atom with num_bits (counting from the most significant set bit)
+// return the number of bytes we need to serialized this atom
+pub fn atom_length_bits(num_bits: u64) -> Option<u64> {
+    if num_bits < 8 {
+        return Some(1);
+    }
+    let num_bytes = (num_bits + 7) / 8;
+    match num_bytes {
+        1..0x40 => Some(1 + num_bytes),
+        0x40..0x2000 => Some(2 + num_bytes),
+        0x2000..0x10_0000 => Some(3 + num_bytes),
+        0x10_0000..0x800_0000 => Some(4 + num_bytes),
+        0x800_0000..0x4_0000_0000 => Some(5 + num_bytes),
+        _ => {
+            assert!(num_bits >= 0x4_0000_0000 * 8 - 7);
+            None
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -53,4 +73,22 @@ mod tests {
     fn test_serialized_length_small_number(#[case] value: u32, #[case] expect: u32) {
         assert_eq!(serialized_length_small_number(value), expect);
     }
+
+    #[rstest]
+    #[case(0, Some(1))]
+    #[case(1, Some(1))]
+    #[case(7, Some(1))]
+    #[case(8, Some(2))]
+    #[case(9, Some(3))]
+    #[case(504, Some(1+63))]
+    #[case(505, Some(2+64))]
+    #[case(0xfff8, Some(2+0x1fff))]
+    #[case(0xfff9, Some(3+0x2000))]
+    #[case(0x3ffffff8, Some(4 + (0x3ffffff8 + 7) / 8))]
+    #[case(0x3ffffff9, Some(5 + (0x3ffffff9 + 7) / 8))]
+    #[case(0x1ffffffff8, Some(5 + (0x1ffffffff8 + 7) / 8))]
+    #[case(0x1ffffffff9, None)]
+    fn test_atom_length_bits(#[case] num_bits: u64, #[case] expect: Option<u64>) {
+        assert_eq!(atom_length_bits(num_bits), expect);
+    }
 }