Replace read_compactsize and write_compactsize with CompactSizeMessage (#3014)

* Replace read_compactsize and write_compactsize with CompactSizeMessage

* Add tests for CompactSize64

* Add compact size range and conversion tests

zebra-chain/src/block/header.rs

@@ -4,7 +4,7 @@ use chrono::{DateTime, Duration, Utc};
 use thiserror::Error;
 
 use crate::{
-    serialization::{TrustedPreallocate, MAX_PROTOCOL_MESSAGE_LEN},
+    serialization::{CompactSizeMessage, TrustedPreallocate, MAX_PROTOCOL_MESSAGE_LEN},
     work::{difficulty::CompactDifficulty, equihash::Solution},
 };
 
@@ -129,10 +129,11 @@ impl Header {
 pub struct CountedHeader {
     /// The header for a block
     pub header: Header,
+
     /// The number of transactions that come after the header
     ///
     /// TODO: should this always be zero? (#1924)
-    pub transaction_count: usize,
+    pub transaction_count: CompactSizeMessage,
 }
 
 /// The serialized size of a Zcash block header.

zebra-chain/src/block/serialize.rs

@@ -5,8 +5,7 @@ use chrono::{TimeZone, Utc};
 
 use crate::{
     serialization::{
-        ReadZcashExt, SerializationError, WriteZcashExt, ZcashDeserialize, ZcashDeserializeInto,
+        ReadZcashExt, SerializationError, ZcashDeserialize, ZcashDeserializeInto, ZcashSerialize,
-        ZcashSerialize,
     },
     work::{difficulty::CompactDifficulty, equihash},
 };
@@ -89,7 +88,7 @@ impl ZcashDeserialize for Header {
 impl ZcashSerialize for CountedHeader {
     fn zcash_serialize<W: io::Write>(&self, mut writer: W) -> Result<(), io::Error> {
         self.header.zcash_serialize(&mut writer)?;
-        writer.write_compactsize(self.transaction_count as u64)?;
+        self.transaction_count.zcash_serialize(&mut writer)?;
         Ok(())
     }
 }
@@ -98,7 +97,7 @@ impl ZcashDeserialize for CountedHeader {
     fn zcash_deserialize<R: io::Read>(mut reader: R) -> Result<Self, SerializationError> {
         Ok(CountedHeader {
             header: (&mut reader).zcash_deserialize_into()?,
-            transaction_count: reader.read_compactsize()?.try_into().unwrap(),
+            transaction_count: (&mut reader).zcash_deserialize_into()?,
         })
     }
 }

zebra-chain/src/block/tests/preallocate.rs

@@ -1,16 +1,17 @@
 //! Tests for trusted preallocation during deserialization.
 
+use std::convert::TryInto;
+
+use proptest::prelude::*;
+
 use crate::{
     block::{
         header::MIN_COUNTED_HEADER_LEN, CountedHeader, Hash, Header, BLOCK_HASH_SIZE,
-        MAX_BLOCK_BYTES, MAX_PROTOCOL_MESSAGE_LEN,
+        MAX_PROTOCOL_MESSAGE_LEN,
     },
-    serialization::{TrustedPreallocate, ZcashSerialize},
+    serialization::{CompactSizeMessage, TrustedPreallocate, ZcashSerialize},
 };
 
-use proptest::prelude::*;
-use std::convert::TryInto;
-
 proptest! {
     /// Verify that the serialized size of a block hash used to calculate the allocation limit is correct
     #[test]
@@ -51,10 +52,10 @@ proptest! {
     /// Confirm that each counted header takes at least COUNTED_HEADER_LEN bytes when serialized.
     /// This verifies that our calculated [`TrustedPreallocate::max_allocation`] is indeed an upper bound.
     #[test]
-    fn counted_header_min_length(header in any::<Header>(), transaction_count in (0..MAX_BLOCK_BYTES)) {
+    fn counted_header_min_length(header in any::<Header>(), transaction_count in any::<CompactSizeMessage>()) {
         let header = CountedHeader {
             header,
-            transaction_count: transaction_count.try_into().expect("Must run test on platform with at least 32 bit address space"),
+            transaction_count,
         };
         let serialized_header = header.zcash_serialize_to_vec().expect("Serialization to vec must succeed");
         prop_assert!(serialized_header.len() >= MIN_COUNTED_HEADER_LEN)
@@ -71,7 +72,7 @@ proptest! {
     fn counted_header_max_allocation(header in any::<Header>()) {
         let header = CountedHeader {
             header,
-            transaction_count: 0,
+            transaction_count: 0.try_into().expect("zero fits in MAX_PROTOCOL_MESSAGE_LEN"),
         };
         let max_allocation: usize = CountedHeader::max_allocation().try_into().unwrap();
         let mut smallest_disallowed_vec = Vec::with_capacity(max_allocation + 1);

zebra-chain/src/serialization/read_zcash.rs

@@ -1,97 +1,14 @@
 use std::{
-    convert::TryInto,
     io,
     net::{IpAddr, Ipv6Addr, SocketAddr},
 };
 
-use byteorder::{BigEndian, LittleEndian, ReadBytesExt};
+use byteorder::{BigEndian, ReadBytesExt};
-
-use super::{SerializationError, MAX_PROTOCOL_MESSAGE_LEN};
 
 /// Extends [`Read`] with methods for writing Zcash/Bitcoin types.
 ///
 /// [`Read`]: https://doc.rust-lang.org/std/io/trait.Read.html
 pub trait ReadZcashExt: io::Read {
-    /// Reads a `u64` using the Bitcoin `CompactSize` encoding.
-    ///
-    /// # Security
-    ///
-    /// Deserialized sizes must be validated before being used.
-    ///
-    /// Preallocating vectors using untrusted `CompactSize`s allows memory
-    /// denial of service attacks. Valid sizes must be less than
-    /// `MAX_PROTOCOL_MESSAGE_LEN / min_serialized_item_bytes` (or a lower limit
-    /// specified by the Zcash consensus rules or Bitcoin network protocol).
-    ///
-    /// As a defence-in-depth for memory preallocation attacks,
-    /// Zebra rejects sizes greater than the protocol message length limit.
-    /// (These sizes should be impossible, because each array items takes at
-    /// least one byte.)
-    ///
-    /// # Examples
-    ///
-    /// ```
-    /// use zebra_chain::serialization::ReadZcashExt;
-    ///
-    /// use std::io::Cursor;
-    /// assert_eq!(
-    ///     0x12,
-    ///     Cursor::new(b"\x12")
-    ///         .read_compactsize().unwrap()
-    /// );
-    /// assert_eq!(
-    ///     0xfd,
-    ///     Cursor::new(b"\xfd\xfd\x00")
-    ///         .read_compactsize().unwrap()
-    /// );
-    /// assert_eq!(
-    ///     0xaafd,
-    ///     Cursor::new(b"\xfd\xfd\xaa")
-    ///         .read_compactsize().unwrap()
-    /// );
-    /// ```
-    ///
-    /// Sizes greater than the maximum network message length are invalid,
-    /// they return a `Parse` error:
-    /// ```
-    /// # use zebra_chain::serialization::ReadZcashExt;
-    /// # use std::io::Cursor;
-    /// Cursor::new(b"\xfe\xfd\xaa\xbb\x00").read_compactsize().unwrap_err();
-    /// Cursor::new(b"\xff\xfd\xaa\xbb\xcc\x22\x00\x00\x00").read_compactsize().unwrap_err();
-    /// ```
-    #[inline]
-    fn read_compactsize(&mut self) -> Result<u64, SerializationError> {
-        use SerializationError::Parse;
-        let flag_byte = self.read_u8()?;
-        let size = match flag_byte {
-            n @ 0x00..=0xfc => Ok(n as u64),
-            0xfd => match self.read_u16::<LittleEndian>()? {
-                n @ 0x0000_00fd..=0x0000_ffff => Ok(n as u64),
-                _ => Err(Parse("non-canonical CompactSize")),
-            },
-            0xfe => match self.read_u32::<LittleEndian>()? {
-                n @ 0x0001_0000..=0xffff_ffff => Ok(n as u64),
-                _ => Err(Parse("non-canonical CompactSize")),
-            },
-            0xff => match self.read_u64::<LittleEndian>()? {
-                n @ 0x1_0000_0000..=0xffff_ffff_ffff_ffff => Ok(n),
-                _ => Err(Parse("non-canonical CompactSize")),
-            },
-        }?;
-
-        // # Security
-        // Defence-in-depth for memory DoS via preallocation.
-        if size
-            > MAX_PROTOCOL_MESSAGE_LEN
-                .try_into()
-                .expect("usize fits in u64")
-        {
-            Err(Parse("CompactSize larger than protocol message limit"))?;
-        }
-
-        Ok(size)
-    }
-
     /// Read an IP address in Bitcoin format.
     #[inline]
     fn read_ip_addr(&mut self) -> io::Result<IpAddr> {

zebra-chain/src/serialization/tests/prop.rs

@@ -1,32 +1,40 @@
 //! Property-based tests for basic serialization primitives.
 
+use std::{convert::TryFrom, io::Cursor};
+
 use proptest::prelude::*;
 
-use std::io::Cursor;
-
 use crate::{
-    serialization::{ReadZcashExt, WriteZcashExt, ZcashSerialize},
+    serialization::{
+        CompactSize64, CompactSizeMessage, ZcashDeserialize, ZcashDeserializeInto, ZcashSerialize,
+        MAX_PROTOCOL_MESSAGE_LEN,
+    },
     transaction::UnminedTx,
 };
 
 proptest! {
     #[test]
-    fn compactsize_write_then_read_round_trip(s in 0u64..0x2_0000u64) {
+    fn compact_size_message_write_then_read_round_trip(size in any::<CompactSizeMessage>()) {
         zebra_test::init();
 
+        let buf = size.zcash_serialize_to_vec().unwrap();
         // Maximum encoding size of a CompactSize is 9 bytes.
-        let mut buf = [0u8; 8+1];
+        prop_assert!(buf.len() <= 9);
-        Cursor::new(&mut buf[..]).write_compactsize(s).unwrap();
+
-        let expect_s = Cursor::new(&buf[..]).read_compactsize().unwrap();
+        let expect_size: CompactSizeMessage = buf.zcash_deserialize_into().unwrap();
-        prop_assert_eq!(s, expect_s);
+        prop_assert_eq!(size, expect_size);
+
+        // Also check the range is correct
+        let size: usize = size.into();
+        prop_assert!(size <= MAX_PROTOCOL_MESSAGE_LEN);
     }
 
     #[test]
-    fn compactsize_read_then_write_round_trip(bytes in prop::array::uniform9(0u8..)) {
+    fn compact_size_message_read_then_write_round_trip(bytes in prop::array::uniform9(0u8..)) {
         zebra_test::init();
 
         // Only do the test if the bytes were valid.
-        if let Ok(s) = Cursor::new(&bytes[..]).read_compactsize() {
+        if let Ok(size) = CompactSizeMessage::zcash_deserialize(Cursor::new(&bytes[..])) {
             // The CompactSize encoding is variable-length, so we may not even
             // read all of the input bytes, and therefore we can't expect that
             // the encoding will reproduce bytes that were never read. Instead,
@@ -34,11 +42,70 @@ proptest! {
             // so that if the encoding is different, we'll catch it on the part
             // that's written.
             let mut expect_bytes = bytes;
-            Cursor::new(&mut expect_bytes[..]).write_compactsize(s).unwrap();
+            size.zcash_serialize(Cursor::new(&mut expect_bytes[..])).unwrap();
+
+            prop_assert_eq!(bytes, expect_bytes);
+
+            // Also check the range is correct
+            let size: usize = size.into();
+            prop_assert!(size <= MAX_PROTOCOL_MESSAGE_LEN);
+        }
+    }
+
+    #[test]
+    fn compact_size_message_range(size in any::<usize>()) {
+        zebra_test::init();
+
+        if let Ok(compact_size) = CompactSizeMessage::try_from(size) {
+            prop_assert!(size <= MAX_PROTOCOL_MESSAGE_LEN);
+
+            let compact_size: usize = compact_size.into();
+            prop_assert_eq!(size, compact_size);
+        } else {
+            prop_assert!(size > MAX_PROTOCOL_MESSAGE_LEN);
+        }
+    }
+
+    #[test]
+    fn compact_size_64_write_then_read_round_trip(size in any::<CompactSize64>()) {
+        zebra_test::init();
+
+        let buf = size.zcash_serialize_to_vec().unwrap();
+        // Maximum encoding size of a CompactSize is 9 bytes.
+        prop_assert!(buf.len() <= 9);
+
+        let expect_size: CompactSize64 = buf.zcash_deserialize_into().unwrap();
+        prop_assert_eq!(size, expect_size);
+    }
+
+    #[test]
+    fn compact_size_64_read_then_write_round_trip(bytes in prop::array::uniform9(0u8..)) {
+        zebra_test::init();
+
+        // Only do the test if the bytes were valid.
+        if let Ok(s) = CompactSize64::zcash_deserialize(Cursor::new(&bytes[..])) {
+            // The CompactSize encoding is variable-length, so we may not even
+            // read all of the input bytes, and therefore we can't expect that
+            // the encoding will reproduce bytes that were never read. Instead,
+            // copy the input bytes, and overwrite them with the encoding of s,
+            // so that if the encoding is different, we'll catch it on the part
+            // that's written.
+            let mut expect_bytes = bytes;
+            s.zcash_serialize(Cursor::new(&mut expect_bytes[..])).unwrap();
+
             prop_assert_eq!(bytes, expect_bytes);
         }
     }
 
+    #[test]
+    fn compact_size_64_conversion(size in any::<u64>()) {
+        zebra_test::init();
+
+        let compact_size = CompactSize64::from(size);
+        let compact_size: u64 = compact_size.into();
+        prop_assert_eq!(size, compact_size);
+    }
+
     #[test]
     fn transaction_serialized_size(transaction in any::<UnminedTx>()) {
         zebra_test::init();

zebra-chain/src/serialization/write_zcash.rs

@@ -1,81 +1,14 @@
 use std::{
-    convert::TryInto,
     io,
     net::{IpAddr, SocketAddr},
 };
 
-use byteorder::{BigEndian, LittleEndian, WriteBytesExt};
+use byteorder::{BigEndian, WriteBytesExt};
-
-use super::MAX_PROTOCOL_MESSAGE_LEN;
 
 /// Extends [`Write`] with methods for writing Zcash/Bitcoin types.
 ///
 /// [`Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
 pub trait WriteZcashExt: io::Write {
-    /// Writes a `u64` using the Bitcoin `CompactSize` encoding.
-    ///
-    /// # Panics
-    ///
-    /// Zebra panics on sizes greater than the protocol message length limit.
-    /// This is a defence-in-depth for memory preallocation attacks.
-    /// (These sizes should be impossible, because each array item takes at
-    /// least one byte.)
-    ///
-    /// # Examples
-    ///
-    /// ```
-    /// use zebra_chain::serialization::WriteZcashExt;
-    ///
-    /// let mut buf = Vec::new();
-    /// buf.write_compactsize(0x12).unwrap();
-    /// assert_eq!(buf, b"\x12");
-    ///
-    /// let mut buf = Vec::new();
-    /// buf.write_compactsize(0xfd).unwrap();
-    /// assert_eq!(buf, b"\xfd\xfd\x00");
-    ///
-    /// let mut buf = Vec::new();
-    /// buf.write_compactsize(0xaafd).unwrap();
-    /// assert_eq!(buf, b"\xfd\xfd\xaa");
-    /// ```
-    ///
-    /// Sizes greater than the maximum network message length are invalid
-    /// and cause a panic:
-    /// ```should_panic
-    /// # use zebra_chain::serialization::WriteZcashExt;
-    /// let mut buf = Vec::new();
-    /// buf.write_compactsize(0xbbaafd).unwrap_err();
-    /// buf.write_compactsize(0x22ccbbaafd).unwrap_err();
-    /// ```
-    #[inline]
-    fn write_compactsize(&mut self, n: u64) -> io::Result<()> {
-        // # Security
-        // Defence-in-depth for memory DoS via preallocation.
-        //
-        assert!(
-            n <= MAX_PROTOCOL_MESSAGE_LEN
-                .try_into()
-                .expect("usize fits in u64"),
-            "CompactSize larger than protocol message limit"
-        );
-
-        match n {
-            0x0000_0000..=0x0000_00fc => self.write_u8(n as u8),
-            0x0000_00fd..=0x0000_ffff => {
-                self.write_u8(0xfd)?;
-                self.write_u16::<LittleEndian>(n as u16)
-            }
-            0x0001_0000..=0xffff_ffff => {
-                self.write_u8(0xfe)?;
-                self.write_u32::<LittleEndian>(n as u32)
-            }
-            _ => {
-                self.write_u8(0xff)?;
-                self.write_u64::<LittleEndian>(n)
-            }
-        }
-    }
-
     /// Write an `IpAddr` in Bitcoin format.
     #[inline]
     fn write_ip_addr(&mut self, addr: IpAddr) -> io::Result<()> {

zebra-chain/src/transaction/serialize.rs

@@ -15,9 +15,9 @@ use crate::{
         Groth16Proof, Halo2Proof, ZkSnarkProof,
     },
     serialization::{
-        zcash_deserialize_external_count, zcash_serialize_external_count, AtLeastOne, ReadZcashExt,
+        zcash_deserialize_external_count, zcash_serialize_empty_list,
-        SerializationError, TrustedPreallocate, WriteZcashExt, ZcashDeserialize,
+        zcash_serialize_external_count, AtLeastOne, ReadZcashExt, SerializationError,
-        ZcashDeserializeInto, ZcashSerialize,
+        TrustedPreallocate, ZcashDeserialize, ZcashDeserializeInto, ZcashSerialize,
     },
     sprout,
 };
@@ -69,32 +69,30 @@ impl ZcashDeserialize for pallas::Base {
 
 impl<P: ZkSnarkProof> ZcashSerialize for JoinSplitData<P> {
     fn zcash_serialize<W: io::Write>(&self, mut writer: W) -> Result<(), io::Error> {
-        writer.write_compactsize(self.joinsplits().count() as u64)?;
+        let joinsplits: Vec<_> = self.joinsplits().cloned().collect();
-        for joinsplit in self.joinsplits() {
+        joinsplits.zcash_serialize(&mut writer)?;
-            joinsplit.zcash_serialize(&mut writer)?;
+
-        }
         writer.write_all(&<[u8; 32]>::from(self.pub_key)[..])?;
         writer.write_all(&<[u8; 64]>::from(self.sig)[..])?;
         Ok(())
     }
 }
 
-impl<P: ZkSnarkProof> ZcashDeserialize for Option<JoinSplitData<P>> {
+impl<P> ZcashDeserialize for Option<JoinSplitData<P>>
+where
+    P: ZkSnarkProof,
+    sprout::JoinSplit<P>: TrustedPreallocate,
+{
     fn zcash_deserialize<R: io::Read>(mut reader: R) -> Result<Self, SerializationError> {
-        let num_joinsplits = reader.read_compactsize()?;
+        let joinsplits: Vec<sprout::JoinSplit<P>> = (&mut reader).zcash_deserialize_into()?;
-        match num_joinsplits {
+        match joinsplits.split_first() {
-            0 => Ok(None),
+            None => Ok(None),
-            n => {
+            Some((first, rest)) => {
-                let first = sprout::JoinSplit::zcash_deserialize(&mut reader)?;
-                let mut rest = Vec::with_capacity((n - 1) as usize);
-                for _ in 0..(n - 1) {
-                    rest.push(sprout::JoinSplit::zcash_deserialize(&mut reader)?);
-                }
                 let pub_key = reader.read_32_bytes()?.into();
                 let sig = reader.read_64_bytes()?.into();
                 Ok(Some(JoinSplitData {
-                    first,
+                    first: first.clone(),
-                    rest,
+                    rest: rest.to_vec(),
                     pub_key,
                     sig,
                 }))
@@ -112,9 +110,9 @@ impl ZcashSerialize for Option<sapling::ShieldedData<SharedAnchor>> {
         match self {
             None => {
                 // nSpendsSapling
-                writer.write_compactsize(0)?;
+                zcash_serialize_empty_list(&mut writer)?;
                 // nOutputsSapling
-                writer.write_compactsize(0)?;
+                zcash_serialize_empty_list(&mut writer)?;
             }
             Some(sapling_shielded_data) => {
                 sapling_shielded_data.zcash_serialize(&mut writer)?;
@@ -254,7 +252,8 @@ impl ZcashSerialize for Option<orchard::ShieldedData> {
         match self {
             None => {
                 // nActionsOrchard
-                writer.write_compactsize(0)?;
+                zcash_serialize_empty_list(writer)?;
+
                 // We don't need to write anything else here.
                 // "The fields flagsOrchard, valueBalanceOrchard, anchorOrchard, sizeProofsOrchard,
                 // proofsOrchard , and bindingSigOrchard are present if and only if nActionsOrchard > 0."
@@ -399,7 +398,7 @@ impl ZcashSerialize for Transaction {
                 lock_time.zcash_serialize(&mut writer)?;
                 match joinsplit_data {
                     // Write 0 for nJoinSplits to signal no JoinSplitData.
-                    None => writer.write_compactsize(0)?,
+                    None => zcash_serialize_empty_list(writer)?,
                     Some(jsd) => jsd.zcash_serialize(&mut writer)?,
                 }
             }
@@ -417,7 +416,7 @@ impl ZcashSerialize for Transaction {
                 writer.write_u32::<LittleEndian>(expiry_height.0)?;
                 match joinsplit_data {
                     // Write 0 for nJoinSplits to signal no JoinSplitData.
-                    None => writer.write_compactsize(0)?,
+                    None => zcash_serialize_empty_list(writer)?,
                     Some(jsd) => jsd.zcash_serialize(&mut writer)?,
                 }
             }
@@ -448,30 +447,28 @@ impl ZcashSerialize for Transaction {
                         // Signal no value balance.
                         writer.write_i64::<LittleEndian>(0)?;
                         // Signal no shielded spends and no shielded outputs.
-                        writer.write_compactsize(0)?;
+                        zcash_serialize_empty_list(&mut writer)?;
-                        writer.write_compactsize(0)?;
+                        zcash_serialize_empty_list(&mut writer)?;
                     }
                     Some(sapling_shielded_data) => {
                         sapling_shielded_data
                             .value_balance
                             .zcash_serialize(&mut writer)?;
-                        writer.write_compactsize(sapling_shielded_data.spends().count() as u64)?;
+
-                        for spend in sapling_shielded_data.spends() {
+                        let spends: Vec<_> = sapling_shielded_data.spends().cloned().collect();
-                            spend.zcash_serialize(&mut writer)?;
+                        spends.zcash_serialize(&mut writer)?;
-                        }
+
-                        writer.write_compactsize(sapling_shielded_data.outputs().count() as u64)?;
+                        let outputs: Vec<_> = sapling_shielded_data
-                        for output in sapling_shielded_data
                             .outputs()
                             .cloned()
                             .map(sapling::OutputInTransactionV4)
-                        {
+                            .collect();
-                            output.zcash_serialize(&mut writer)?;
+                        outputs.zcash_serialize(&mut writer)?;
-                        }
                     }
                 }
 
                 match joinsplit_data {
-                    None => writer.write_compactsize(0)?,
+                    None => zcash_serialize_empty_list(&mut writer)?,
                     Some(jsd) => jsd.zcash_serialize(&mut writer)?,
                 }
 

zebra-chain/src/transparent/script.rs

@@ -1,11 +1,11 @@
 //! Bitcoin script for Zebra
 
-#![allow(clippy::unit_arg)]
-
-use crate::serialization::{SerializationError, WriteZcashExt, ZcashDeserialize, ZcashSerialize};
-
 use std::{fmt, io};
 
+use crate::serialization::{
+    zcash_serialize_bytes, SerializationError, ZcashDeserialize, ZcashSerialize,
+};
+
 /// An encoding of a Bitcoin script.
 #[derive(Clone, Eq, PartialEq, Serialize, Deserialize, Hash)]
 #[cfg_attr(
@@ -42,16 +42,14 @@ impl fmt::Debug for Script {
 }
 
 impl ZcashSerialize for Script {
-    fn zcash_serialize<W: io::Write>(&self, mut writer: W) -> Result<(), io::Error> {
+    fn zcash_serialize<W: io::Write>(&self, writer: W) -> Result<(), io::Error> {
-        writer.write_compactsize(self.0.len() as u64)?;
+        zcash_serialize_bytes(&self.0, writer)
-        writer.write_all(&self.0[..])?;
-        Ok(())
     }
 }
 
 impl ZcashDeserialize for Script {
     fn zcash_deserialize<R: io::Read>(reader: R) -> Result<Self, SerializationError> {
-        Ok(Script(Vec::zcash_deserialize(reader)?))
+        Vec::zcash_deserialize(reader).map(Script)
     }
 }
 

zebra-chain/src/transparent/serialize.rs

@@ -5,14 +5,22 @@ use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 use crate::{
     block,
     serialization::{
-        ReadZcashExt, SerializationError, WriteZcashExt, ZcashDeserialize, ZcashDeserializeInto,
+        zcash_serialize_bytes, ReadZcashExt, SerializationError, ZcashDeserialize,
-        ZcashSerialize,
+        ZcashDeserializeInto, ZcashSerialize,
     },
     transaction,
 };
 
 use super::{CoinbaseData, Input, OutPoint, Output, Script};
 
+/// The maximum length of the coinbase data.
+/// Includes the encoded coinbase height, if any.
+///
+/// > The number of bytes in the coinbase script, up to a maximum of 100 bytes.
+///
+/// https://developer.bitcoin.org/reference/transactions.html#coinbase-input-the-input-of-the-first-transaction-in-a-block
+pub const MAX_COINBASE_DATA_LEN: usize = 100;
+
 /// The coinbase data for a genesis block.
 ///
 /// Zcash uses the same coinbase data for the Mainnet, Testnet, and Regtest
@@ -116,27 +124,6 @@ fn parse_coinbase_height(
     }
 }
 
-/// Return the encoded length of `height`, as a prefix to the coinbase data.
-fn coinbase_height_len(height: block::Height) -> usize {
-    // We can't write this as a match statement on stable until exclusive range
-    // guards are stabilized.
-    if let 0 = height.0 {
-        0
-    } else if let _h @ 1..=16 = height.0 {
-        1
-    } else if let _h @ 17..=127 = height.0 {
-        2
-    } else if let _h @ 128..=32767 = height.0 {
-        3
-    } else if let _h @ 32768..=8_388_607 = height.0 {
-        4
-    } else if let _h @ 8_388_608..=block::Height::MAX_AS_U32 = height.0 {
-        5
-    } else {
-        panic!("Invalid coinbase height");
-    }
-}
-
 /// Encode `height` into a block height, as a prefix of the coinbase data.
 /// Does not write `coinbase_data`.
 ///
@@ -223,11 +210,12 @@ impl ZcashSerialize for Input {
             } => {
                 writer.write_all(&[0; 32][..])?;
                 writer.write_u32::<LittleEndian>(0xffff_ffff)?;
-                let height_len = coinbase_height_len(*height);
+
-                let total_len = height_len + data.as_ref().len();
+                let mut height_and_data = Vec::new();
-                writer.write_compactsize(total_len as u64)?;
+                write_coinbase_height(*height, data, &mut height_and_data)?;
-                write_coinbase_height(*height, data, &mut writer)?;
+                height_and_data.extend(&data.0);
-                writer.write_all(data.as_ref())?;
+                zcash_serialize_bytes(&height_and_data, &mut writer)?;
+
                 writer.write_u32::<LittleEndian>(*sequence)?;
             }
         }
@@ -244,15 +232,15 @@ impl ZcashDeserialize for Input {
             if reader.read_u32::<LittleEndian>()? != 0xffff_ffff {
                 return Err(SerializationError::Parse("wrong index in coinbase"));
             }
-            let len = reader.read_compactsize()?;
+
-            if len > 100 {
+            let data: Vec<u8> = (&mut reader).zcash_deserialize_into()?;
+            if data.len() > MAX_COINBASE_DATA_LEN {
                 return Err(SerializationError::Parse("coinbase has too much data"));
             }
-            // Memory Denial of Service: this length has just been checked
-            let mut data = vec![0; len as usize];
-            reader.read_exact(&mut data[..])?;
             let (height, data) = parse_coinbase_height(data)?;
+
             let sequence = reader.read_u32::<LittleEndian>()?;
+
             Ok(Input::Coinbase {
                 height,
                 data,

zebra-chain/src/work/equihash.rs

@@ -1,12 +1,15 @@
 //! Equihash Solution and related items.
 
-use crate::block::Header;
-use crate::serialization::{
-    serde_helpers, ReadZcashExt, SerializationError, WriteZcashExt, ZcashDeserialize,
-    ZcashSerialize,
-};
 use std::{fmt, io};
 
+use crate::{
+    block::Header,
+    serialization::{
+        serde_helpers, zcash_serialize_bytes, SerializationError, ZcashDeserialize,
+        ZcashDeserializeInto, ZcashSerialize,
+    },
+};
+
 /// The error type for Equihash
 #[non_exhaustive]
 #[derive(Debug, thiserror::Error)]
@@ -84,58 +87,25 @@ impl Clone for Solution {
 impl Eq for Solution {}
 
 impl ZcashSerialize for Solution {
-    fn zcash_serialize<W: io::Write>(&self, mut writer: W) -> Result<(), io::Error> {
+    fn zcash_serialize<W: io::Write>(&self, writer: W) -> Result<(), io::Error> {
-        writer.write_compactsize(SOLUTION_SIZE as u64)?;
+        zcash_serialize_bytes(&self.0.to_vec(), writer)
-        writer.write_all(&self.0[..])?;
-        Ok(())
     }
 }
 
 impl ZcashDeserialize for Solution {
     fn zcash_deserialize<R: io::Read>(mut reader: R) -> Result<Self, SerializationError> {
-        let solution_size = reader.read_compactsize()?;
+        let solution: Vec<u8> = (&mut reader).zcash_deserialize_into()?;
-        if solution_size != (SOLUTION_SIZE as u64) {
+
+        if solution.len() != SOLUTION_SIZE {
             return Err(SerializationError::Parse(
                 "incorrect equihash solution size",
             ));
         }
+
         let mut bytes = [0; SOLUTION_SIZE];
-        reader.read_exact(&mut bytes[..])?;
+        // Won't panic, because we just checked the length.
+        bytes.copy_from_slice(&solution);
+
         Ok(Self(bytes))
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::block::MAX_BLOCK_BYTES;
-
-    static EQUIHASH_SIZE_TESTS: &[u64] = &[
-        0,
-        1,
-        (SOLUTION_SIZE - 1) as u64,
-        SOLUTION_SIZE as u64,
-        (SOLUTION_SIZE + 1) as u64,
-        MAX_BLOCK_BYTES - 1,
-        MAX_BLOCK_BYTES,
-    ];
-
-    #[test]
-    fn equihash_solution_size_field() {
-        zebra_test::init();
-
-        for size in EQUIHASH_SIZE_TESTS {
-            let mut data = Vec::new();
-            data.write_compactsize(*size as u64)
-                .expect("Compact size should serialize");
-            data.resize(data.len() + SOLUTION_SIZE, 0);
-            let result = Solution::zcash_deserialize(data.as_slice());
-            if *size == (SOLUTION_SIZE as u64) {
-                result.expect("Correct size field in EquihashSolution should deserialize");
-            } else {
-                result
-                    .expect_err("Wrong size field in EquihashSolution should fail on deserialize");
-            }
-        }
-    }
-}

zebra-chain/src/work/tests/vectors.rs

@@ -1,10 +1,13 @@
-use super::super::*;
+use std::convert::TryInto;
 
 use crate::{
-    block::Block,
+    block::{Block, MAX_BLOCK_BYTES},
-    serialization::{ZcashDeserialize, ZcashDeserializeInto, ZcashSerialize},
+    serialization::{CompactSizeMessage, ZcashDeserialize, ZcashDeserializeInto, ZcashSerialize},
+    work::equihash::{Solution, SOLUTION_SIZE},
 };
 
+use super::super::*;
+
 /// Includes the 32-byte nonce.
 const EQUIHASH_SOLUTION_BLOCK_OFFSET: usize = equihash::Solution::INPUT_LENGTH + 32;
 
@@ -45,3 +48,36 @@ fn equihash_solution_test_vectors_are_valid() -> color_eyre::eyre::Result<()> {
 
     Ok(())
 }
+
+static EQUIHASH_SIZE_TESTS: &[usize] = &[
+    0,
+    1,
+    SOLUTION_SIZE - 1,
+    SOLUTION_SIZE,
+    SOLUTION_SIZE + 1,
+    (MAX_BLOCK_BYTES - 1) as usize,
+    MAX_BLOCK_BYTES as usize,
+];
+
+#[test]
+fn equihash_solution_size_field() {
+    zebra_test::init();
+
+    for size in EQUIHASH_SIZE_TESTS.iter().copied() {
+        let mut data = Vec::new();
+
+        let size: CompactSizeMessage = size
+            .try_into()
+            .expect("test size fits in MAX_PROTOCOL_MESSAGE_LEN");
+        size.zcash_serialize(&mut data)
+            .expect("CompactSize should serialize");
+        data.resize(data.len() + SOLUTION_SIZE, 0);
+
+        let result = Solution::zcash_deserialize(data.as_slice());
+        if size == SOLUTION_SIZE.try_into().unwrap() {
+            result.expect("Correct size field in EquihashSolution should deserialize");
+        } else {
+            result.expect_err("Wrong size field in EquihashSolution should fail on deserialize");
+        }
+    }
+}

zebra-state/src/service.rs

@@ -790,7 +790,11 @@ impl Service<Request> for StateService {
                             .best_block(hash.into())
                             .expect("block for found hash is in the best chain");
                         block::CountedHeader {
-                            transaction_count: block.transactions.len(),
+                            transaction_count: block
+                                .transactions
+                                .len()
+                                .try_into()
+                                .expect("transaction count has already been validated"),
                             header: block.header,
                         }
                     })

zebra-state/src/service/tests.rs

@@ -62,7 +62,11 @@ async fn test_populated_state_responds_correctly(
         .iter()
         .map(|block| CountedHeader {
             header: block.header,
-            transaction_count: block.transactions.len(),
+            transaction_count: block
+                .transactions
+                .len()
+                .try_into()
+                .expect("test block transaction counts are valid"),
         })
         .collect();