fix: write only defined bytes into concurrent Merkle tree changelog

program-libs/concurrent-merkle-tree/src/changelog.rs

@@ -4,6 +4,8 @@ use light_bounded_vec::BoundedVec;
 
 use crate::errors::ConcurrentMerkleTreeError;
 
+const _: () = assert!(std::mem::size_of::<Option<[u8; 32]>>() == 33);
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 #[repr(transparent)]
 pub struct ChangelogPath<const HEIGHT: usize>(pub [Option<[u8; 32]>; HEIGHT]);

program-libs/concurrent-merkle-tree/src/lib.rs

@@ -17,7 +17,7 @@ use std::{
     alloc::{self, handle_alloc_error, Layout},
     iter::Skip,
     marker::PhantomData,
-    mem,
+    mem, ptr,
 };
 
 use changelog::ChangelogPath;
@@ -228,7 +228,7 @@ where
         // Initialize changelog.
         let path = ChangelogPath::from_fn(|i| Some(H::zero_bytes()[i]));
         let changelog_entry = ChangelogEntry { path, index: 0 };
-        self.changelog.push(changelog_entry);
+        self.push_changelog_entry(changelog_entry);
 
         // Initialize filled subtrees.
         for i in 0..self.height {
@@ -252,6 +252,27 @@ where
         self.changelog.last_index()
     }
 
+    /// Pushes `entry` so that every byte of its slot is defined.
+    ///
+    /// `CyclicBoundedVec::push` copies the struct with `ptr::write`, which
+    /// also copies the undefined value bytes of `None` nodes and the struct
+    /// padding between `path` and `index` from the stack into the account.
+    /// Instead, the slot is zeroed and only defined bytes are written.
+    fn push_changelog_entry(&mut self, entry: ChangelogEntry<HEIGHT>) {
+        self.changelog.push(ChangelogEntry::default_with_index(0));
+        if let Some(slot) = self.changelog.last_mut() {
+            // SAFETY: All-zero bytes are a valid `ChangelogEntry` (all `None`
+            // nodes, index 0). This also zeroes the padding before `index`.
+            unsafe { ptr::write_bytes(slot as *mut ChangelogEntry<HEIGHT>, 0, 1) };
+            slot.index = entry.index;
+            for (dst, src) in slot.path.iter_mut().zip(entry.path.iter()) {
+                if src.is_some() {
+                    *dst = *src;
+                }
+            }
+        }
+    }
+
     /// Returns the index of the current root in the tree's root buffer.
     pub fn root_index(&self) -> usize {
         self.roots.last_index()
@@ -448,7 +469,7 @@ where
                 self.set_rightmost_leaf(new_leaf);
             }
         }
-        self.changelog.push(changelog_entry);
+        self.push_changelog_entry(changelog_entry);
 
         if self.canopy_depth > 0 {
             self.update_canopy(self.changelog.last_index(), 1);
@@ -569,8 +590,7 @@ where
         for (leaf_i, leaf) in leaves.iter().enumerate() {
             let mut current_index = self.next_index();
 
-            self.changelog
-                .push(ChangelogEntry::<HEIGHT>::default_with_index(current_index));
+            self.push_changelog_entry(ChangelogEntry::<HEIGHT>::default_with_index(current_index));
             let changelog_index = self.changelog_index();
 
             let mut current_node = **leaf;

program-libs/concurrent-merkle-tree/tests/tests.rs

@@ -1,4 +1,4 @@
-use std::cmp;
+use std::{cmp, mem::size_of};
 
 use ark_bn254::Fr;
 use ark_ff::{BigInteger, PrimeField, UniformRand};
@@ -11,6 +11,7 @@ use light_concurrent_merkle_tree::{
 };
 use light_hash_set::HashSet;
 use light_hasher::{Hasher, Keccak, Poseidon, Sha256};
+use memoffset::offset_of;
 use num_bigint::BigUint;
 use num_traits::FromBytes;
 use rand::{
@@ -3546,3 +3547,98 @@ fn test_update_with_canopy_poseidon() {
 fn test_update_with_canopy_sha256() {
     update_with_canopy::<Sha256>()
 }
+
+/// The on-chain changelog layout the tree relies on: a node is a tag byte
+/// (0 = `None`, 1 = `Some`) followed by the 32 value bytes, and `index` is
+/// the last `u64` of a `repr(C)` entry.
+#[test]
+fn test_changelog_layout() {
+    assert_eq!(size_of::<Option<[u8; 32]>>(), 33);
+    let value = [0xAB; 32];
+    // SAFETY: Same size; `Some` has every byte initialized.
+    let some_bytes: [u8; 33] = unsafe { std::mem::transmute(Some(value)) };
+    let (tag, some_value) = some_bytes.split_first().unwrap();
+    assert_eq!(*tag, 1);
+    assert_eq!(some_value, value);
+    // Only the tag byte of `None` is defined, so read just that byte.
+    let none = None::<[u8; 32]>;
+    // SAFETY: The tag byte is always initialized and lies within `none`.
+    let none_tag = unsafe { *(&none as *const Option<[u8; 32]> as *const u8) };
+    assert_eq!(none_tag, 0);
+
+    assert_eq!(size_of::<ChangelogEntry<22>>(), 736);
+    assert_eq!(size_of::<ChangelogEntry<26>>(), 872);
+    assert_eq!(size_of::<ChangelogEntry<32>>(), 1064);
+    assert_eq!(size_of::<ChangelogEntry<40>>(), 1328);
+    assert_eq!(offset_of!(ChangelogEntry<22>, index), 736 - 8);
+    assert_eq!(offset_of!(ChangelogEntry<26>, index), 872 - 8);
+    assert_eq!(offset_of!(ChangelogEntry<32>, index), 1064 - 8);
+    assert_eq!(offset_of!(ChangelogEntry<40>, index), 1328 - 8);
+}
+
+/// Every byte of every changelog entry written into the account buffer must
+/// be defined: `None` nodes must be all-zero and the struct padding between
+/// `path` and `index` must be zero. The buffer is pre-filled with a marker so
+/// any byte that is merely left untouched (instead of written) is detected.
+#[test]
+fn test_changelog_bytes_are_defined() {
+    // 33 * 10 = 330 bytes of path, leaving 6 bytes of padding before `index`.
+    const HEIGHT: usize = 10;
+    const CHANGELOG: usize = 8;
+    const ROOTS: usize = 8;
+    const CANOPY: usize = 0;
+    let path_size = size_of::<ChangelogPath<HEIGHT>>();
+    let index_offset = offset_of!(ChangelogEntry<HEIGHT>, index);
+    assert_eq!(index_offset - path_size, 6);
+
+    let mut bytes = vec![
+        0xFFu8;
+        ConcurrentMerkleTree::<Sha256, HEIGHT>::size_in_account(
+            HEIGHT, CHANGELOG, ROOTS, CANOPY
+        )
+    ];
+    let mut merkle_tree =
+        ConcurrentMerkleTreeZeroCopyMut::<Sha256, HEIGHT>::from_bytes_zero_copy_init(
+            bytes.as_mut_slice(),
+            HEIGHT,
+            CANOPY,
+            CHANGELOG,
+            ROOTS,
+        )
+        .unwrap();
+    // `init` writes a full path, `append_batch` writes partial paths with
+    // `None` nodes.
+    merkle_tree.init().unwrap();
+    merkle_tree
+        .append_batch(&[&[1; 32], &[2; 32], &[3; 32]])
+        .unwrap();
+
+    for changelog_index in 0..merkle_tree.changelog.len() {
+        let entry = merkle_tree.changelog.get(changelog_index).unwrap();
+        // SAFETY: The entry lives in `bytes`, which was fully initialized
+        // with the marker before the tree was created.
+        let entry_bytes = unsafe {
+            std::slice::from_raw_parts(
+                entry as *const ChangelogEntry<HEIGHT> as *const u8,
+                size_of::<ChangelogEntry<HEIGHT>>(),
+            )
+        };
+        let (path_bytes, rest) = entry_bytes.split_at(path_size);
+        let (padding, _index) = rest.split_at(index_offset - path_size);
+        for (level, node) in path_bytes.chunks_exact(33).enumerate() {
+            let (tag, value) = node.split_first().unwrap();
+            match *tag {
+                1 => {}
+                0 => assert!(
+                    value.iter().all(|b| *b == 0),
+                    "entry {changelog_index} level {level}: None node has non-zero value bytes"
+                ),
+                tag => panic!("entry {changelog_index} level {level}: invalid tag {tag}"),
+            }
+        }
+        assert!(
+            padding.iter().all(|b| *b == 0),
+            "entry {changelog_index}: padding bytes are not zero"
+        );
+    }
+}

program-libs/hash-set/src/lib.rs

@@ -16,7 +16,7 @@
     cmp::Ordering,
     marker::Send,
     mem,
-    ptr::NonNull,
+    ptr::{self, NonNull},
 };
 
 use light_hasher::{bigint::bigint_to_be_bytes_array, HasherError};
@@ -73,6 +73,84 @@
     pub sequence_number: Option<usize>,
 }
 
+// `Option<HashSetCell>` uses the unused discriminants of the nested
+// `Option<usize>` for its outer `None`. These are the deployed v1 account
+// bytes, so fail compilation if the Rust layout changes.
+const _: () = {
+    assert!(mem::size_of::<usize>() == 8);
+    assert!(mem::size_of::<Option<usize>>() == 16);
+    assert!(mem::size_of::<HashSetCell>() == 48);
+    assert!(mem::align_of::<HashSetCell>() == 8);
+    assert!(mem::offset_of!(HashSetCell, sequence_number) == 0);
+    assert!(mem::offset_of!(HashSetCell, value) == 16);
+    assert!(mem::size_of::<Option<HashSetCell>>() == 48);
+    assert!(mem::align_of::<Option<HashSetCell>>() == 8);
+};
+
+const UNMARKED_BUCKET_TAG: usize = 0;
+const MARKED_BUCKET_TAG: usize = 1;
+const EMPTY_BUCKET_TAG: usize = 2;
+
+const fn bucket_tag(bucket: &Option<HashSetCell>) -> usize {
+    // SAFETY: The layout assertions above pin the enum tag to the first
+    // `usize`, which rustc always initializes.
+    unsafe { *(bucket as *const Option<HashSetCell> as *const usize) }
+}
+
+const _: () = {
+    let cell = HashSetCell {
+        value: [0; 32],
+        sequence_number: None,
+    };
+    assert!(bucket_tag(&Some(cell)) == UNMARKED_BUCKET_TAG);
+    let cell = HashSetCell {
+        value: [0; 32],
+        sequence_number: Some(0),
+    };
+    assert!(bucket_tag(&Some(cell)) == MARKED_BUCKET_TAG);
+    assert!(bucket_tag(&None) == EMPTY_BUCKET_TAG);
+};
+
+#[repr(C)]
+struct RawHashSetCell {
+    tag: usize,
+    sequence_number: usize,
+    value: [u8; 32],
+}
+
+const _: () = {
+    assert!(mem::size_of::<RawHashSetCell>() == mem::size_of::<Option<HashSetCell>>());
+    assert!(mem::align_of::<RawHashSetCell>() == mem::align_of::<Option<HashSetCell>>());
+    assert!(mem::offset_of!(RawHashSetCell, tag) == 0);
+    assert!(mem::offset_of!(RawHashSetCell, sequence_number) == 8);
+    assert!(mem::offset_of!(RawHashSetCell, value) == 16);
+};
+
+/// Writes every byte of a bucket without copying undefined enum payload bytes.
+unsafe fn write_bucket(
+    bucket: *mut Option<HashSetCell>,
+    tag: usize,
+    sequence_number: usize,
+    value: [u8; 32],
+) {
+    ptr::write(
+        bucket.cast::<RawHashSetCell>(),
+        RawHashSetCell {
+            tag,
+            sequence_number,
+            value,
+        },
+    );
+}
+
+unsafe fn write_empty_bucket(bucket: *mut Option<HashSetCell>) {
+    write_bucket(bucket, EMPTY_BUCKET_TAG, 0, [0; 32]);
+}
+
+unsafe fn write_unmarked_bucket(bucket: *mut Option<HashSetCell>, value: [u8; 32]) {
+    write_bucket(bucket, UNMARKED_BUCKET_TAG, 0, value);
+}
+
 unsafe impl Send for HashSet {}
 
 impl HashSetCell {
@@ -145,14 +223,11 @@
 
     /// Size which needs to be allocated on Solana account to fit the hash set.
     pub fn size_in_account(capacity_values: usize) -> usize {
-        let dyn_fields_size = Self::non_dyn_fields_size();
-
-        let buckets_size_unaligned = mem::size_of::<Option<HashSetCell>>() * capacity_values;
-        // Make sure that alignment of `values` matches the alignment of `usize`.
-        let buckets_size = buckets_size_unaligned + mem::align_of::<usize>()
-            - (buckets_size_unaligned % mem::align_of::<usize>());
+        Self::buckets_offset() + mem::size_of::<Option<HashSetCell>>() * capacity_values
+    }
 
-        dyn_fields_size + buckets_size
+    pub(crate) fn buckets_offset() -> usize {
+        Self::non_dyn_fields_size() + mem::size_of::<usize>()
     }
 
     // Create a new hash set with the given capacity
@@ -166,7 +241,7 @@
         let values = NonNull::new(values_ptr).unwrap();
         for i in 0..capacity_values {
             unsafe {
-                std::ptr::write(values_ptr.add(i), None);
+                write_empty_bucket(values_ptr.add(i));
             }
         }
 
@@ -213,11 +288,7 @@
             handle_alloc_error(buckets_layout);
         }
         let buckets = NonNull::new(buckets_dst_ptr).unwrap();
-        for i in 0..capacity {
-            std::ptr::write(buckets_dst_ptr.add(i), None);
-        }
-
-        let offset = Self::non_dyn_fields_size() + mem::size_of::<usize>();
+        let offset = Self::buckets_offset();
         let buckets_src_ptr = bytes.as_ptr().add(offset) as *const Option<HashSetCell>;
         std::ptr::copy(buckets_src_ptr, buckets_dst_ptr, capacity);
 
@@ -286,25 +357,23 @@
         // PANICS: We trust the bounds of `value_index` here.
         let bucket = self.get_bucket_mut(value_index).unwrap();
 
-        match bucket {
+        match *bucket {
             // The cell in the value array is already taken.
-            Some(bucket) => {
+            Some(cell) => {
                 // We can overwrite that cell only if the element
                 // is expired - when the difference between its
                 // sequence number and provided sequence number is
                 // greater than the threshold.
-                if let Some(element_sequence_number) = bucket.sequence_number {
+                if let Some(element_sequence_number) = cell.sequence_number {
                     if current_sequence_number >= element_sequence_number {
-                        *bucket = HashSetCell {
-                            value: bigint_to_be_bytes_array(value)?,
-                            sequence_number: None,
-                        };
+                        let value = bigint_to_be_bytes_array(value)?;
+                        unsafe { write_unmarked_bucket(bucket, value) };
                         return Ok(true);
                     }
                 }
                 // Otherwise, we need to prevent having multiple valid
                 // elements with the same value.
-                if &BigUint::from_be_bytes(bucket.value.as_slice()) == value {
+                if &BigUint::from_be_bytes(cell.value.as_slice()) == value {
                     return Err(HashSetError::ElementAlreadyExists);
                 }
             }
@@ -347,10 +416,8 @@
                 // PANICS: We trust the bounds of `index`.
                 let bucket = self.get_bucket_mut(index).unwrap();
 
-                *bucket = Some(HashSetCell {
-                    value: bigint_to_be_bytes_array(value)?,
-                    sequence_number: None,
-                });
+                let value = bigint_to_be_bytes_array(value)?;
+                unsafe { write_unmarked_bucket(bucket, value) };
                 return Ok(index);
             }
         }