Replace read_compactsize and write_compactsize with CompactSizeMessage

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,
-        ZcashSerialize,
+        ReadZcashExt, SerializationError, ZcashDeserialize, ZcashDeserializeInto, 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 super::{SerializationError, MAX_PROTOCOL_MESSAGE_LEN};
+use byteorder::{BigEndian, ReadBytesExt};
 
 /// 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,44 +1,111 @@
 //! Property-based tests for basic serialization primitives.
 
-use proptest::prelude::*;
+use std::{convert::TryFrom, io::Cursor};
 
-use std::io::Cursor;
+use proptest::prelude::*;
 
 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];
-        Cursor::new(&mut buf[..]).write_compactsize(s).unwrap();
-        let expect_s = Cursor::new(&buf[..]).read_compactsize().unwrap();
-        prop_assert_eq!(s, expect_s);
+        prop_assert!(buf.len() <= 9);
+
+        let expect_size: CompactSizeMessage = buf.zcash_deserialize_into().unwrap();
+        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,
             // 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;
-            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 super::MAX_PROTOCOL_MESSAGE_LEN;
+use byteorder::{BigEndian, WriteBytesExt};
 
 /// 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<()> {