From 92ddf0542f3a0f6d1cb571a841b55503f67e4f0d Mon Sep 17 00:00:00 2001 From: Henry de Valence Date: Fri, 20 Dec 2019 16:43:19 -0800 Subject: [PATCH] Provide impl Zcash[De]Serialize for Vec. This replaces the read_list function and makes the code significantly cleaner. The only downside is that it loses exact preallocation, but this is probably not a big deal. --- zebra-chain/src/serialization.rs | 56 ++++----- zebra-network/src/protocol/external/codec.rs | 125 ++++--------------- 2 files changed, 48 insertions(+), 133 deletions(-) diff --git a/zebra-chain/src/serialization.rs b/zebra-chain/src/serialization.rs index 1e169b8f..80e96fd7 100644 --- a/zebra-chain/src/serialization.rs +++ b/zebra-chain/src/serialization.rs @@ -55,6 +55,31 @@ pub trait ZcashDeserialize: Sized { fn zcash_deserialize(reader: R) -> Result; } +impl ZcashSerialize for Vec { + fn zcash_serialize(&self, mut writer: W) -> Result<(), SerializationError> { + writer.write_compactsize(self.len() as u64)?; + for x in self { + x.zcash_serialize(&mut writer)?; + } + Ok(()) + } +} + +impl ZcashDeserialize for Vec { + fn zcash_deserialize(mut reader: R) -> Result { + let len = reader.read_compactsize()?; + // We're given len, so we could preallocate. But blindly preallocating + // without a size bound can allow DOS attacks, and there's no way to + // pass a size bound in a ZcashDeserialize impl, so instead we allocate + // as we read from the reader. + let mut vec = Vec::new(); + for _ in 0..len { + vec.push(T::zcash_deserialize(&mut reader)?); + } + Ok(vec) + } +} + /// Extends [`Write`] with methods for writing Zcash/Bitcoin types. /// /// [`Write`]: https://doc.rust-lang.org/std/io/trait.Write.html @@ -248,37 +273,6 @@ pub trait ReadZcashExt: io::Read { self.read_exact(&mut bytes)?; Ok(bytes) } - - /// Convenience method to read a `Vec` with a leading count in - /// a safer manner. - /// - /// This method preallocates a buffer, performing a single - /// allocation in the honest case. It's possible for someone to - /// send a short message with a large count field, so if we - /// naively trust the count field we could be tricked into - /// preallocating a large buffer. Instead, we rely on the passed - /// maximum count for a valid message and select the min of the - /// two values. - #[inline] - fn read_list( - &mut self, - max_count: usize, - ) -> Result, SerializationError> { - // This prevents the inferred type for zcash_deserialize from - // taking ownership of &mut self. This wouldn't really be an - // issue if the target impl's `Copy`, but we need to own it. - let mut self2 = self; - - let count = self2.read_compactsize()? as usize; - - let mut items = Vec::with_capacity(std::cmp::min(count, max_count)); - - for _ in 0..count { - items.push(T::zcash_deserialize(&mut self2)?); - } - - return Ok(items); - } } /// Mark all types implementing `Read` as implementing the extension. diff --git a/zebra-network/src/protocol/external/codec.rs b/zebra-network/src/protocol/external/codec.rs index 1e73a791..82e74f4f 100644 --- a/zebra-network/src/protocol/external/codec.rs +++ b/zebra-network/src/protocol/external/codec.rs @@ -9,7 +9,7 @@ use chrono::{TimeZone, Utc}; use tokio_util::codec::{Decoder, Encoder}; use zebra_chain::{ - block::{Block, BlockHeader, BlockHeaderHash}, + block::{Block, BlockHeaderHash}, serialization::{ ReadZcashExt, SerializationError as Error, WriteZcashExt, ZcashDeserialize, ZcashSerialize, }, @@ -20,7 +20,6 @@ use zebra_chain::{ use crate::{constants, types::Network}; use super::{ - inv::InventoryHash, message::{Message, RejectReason}, types::*, }; @@ -202,12 +201,7 @@ impl Codec { writer.write_string(&reason)?; writer.write_all(&data.unwrap())?; } - Addr(ref addrs) => { - writer.write_compactsize(addrs.len() as u64)?; - for addr in addrs { - addr.zcash_serialize(&mut writer)?; - } - } + Addr(ref addrs) => addrs.zcash_serialize(&mut writer)?, GetAddr => { /* Empty payload -- no-op */ } Block { ref version, @@ -222,10 +216,7 @@ impl Codec { ref hash_stop, } => { writer.write_u32::(version.0)?; - writer.write_compactsize(block_locator_hashes.len() as u64)?; - for hash in block_locator_hashes { - hash.zcash_serialize(&mut writer)?; - } + block_locator_hashes.zcash_serialize(&mut writer)?; hash_stop.zcash_serialize(&mut writer)?; } GetHeaders { @@ -234,36 +225,13 @@ impl Codec { ref hash_stop, } => { writer.write_u32::(version.0)?; - writer.write_compactsize(block_locator_hashes.len() as u64)?; - for hash in block_locator_hashes { - hash.zcash_serialize(&mut writer)?; - } + block_locator_hashes.zcash_serialize(&mut writer)?; hash_stop.zcash_serialize(&mut writer)?; } - Headers(ref headers) => { - writer.write_compactsize(headers.len() as u64)?; - for header in headers { - header.zcash_serialize(&mut writer)?; - } - } - Inv(ref hashes) => { - writer.write_compactsize(hashes.len() as u64)?; - for hash in hashes { - hash.zcash_serialize(&mut writer)?; - } - } - GetData(ref hashes) => { - writer.write_compactsize(hashes.len() as u64)?; - for hash in hashes { - hash.zcash_serialize(&mut writer)?; - } - } - NotFound(ref hashes) => { - writer.write_compactsize(hashes.len() as u64)?; - for hash in hashes { - hash.zcash_serialize(&mut writer)?; - } - } + Headers(ref headers) => headers.zcash_serialize(&mut writer)?, + Inv(ref hashes) => hashes.zcash_serialize(&mut writer)?, + GetData(ref hashes) => hashes.zcash_serialize(&mut writer)?, + NotFound(ref hashes) => hashes.zcash_serialize(&mut writer)?, Tx { ref version, ref transaction, @@ -488,16 +456,8 @@ impl Codec { }) } - fn read_addr(&self, mut reader: R) -> Result { - use crate::meta_addr::MetaAddr; - - // addrs are encoded as: timestamp + services + ipv6 + port - const ENCODED_ADDR_SIZE: usize = 4 + 8 + 16 + 2; - let max_count = self.builder.max_len / ENCODED_ADDR_SIZE; - - let addrs: Vec = reader.read_list(max_count)?; - - Ok(Message::Addr(addrs)) + fn read_addr(&self, reader: R) -> Result { + Ok(Message::Addr(Vec::zcash_deserialize(reader)?)) } fn read_getaddr(&self, mut _reader: R) -> Result { @@ -507,24 +467,15 @@ impl Codec { fn read_block(&self, mut reader: R) -> Result { Ok(Message::Block { version: Version(reader.read_u32::()?), - block: Block::zcash_deserialize(&mut reader)?, }) } fn read_getblocks(&self, mut reader: R) -> Result { - let version = Version(reader.read_u32::()?); - - let max_count = self.builder.max_len / 32; - - let block_locator_hashes: Vec = reader.read_list(max_count)?; - - let hash_stop = BlockHeaderHash(reader.read_32_bytes()?); - Ok(Message::GetBlocks { - version, - block_locator_hashes, - hash_stop, + version: Version(reader.read_u32::()?), + block_locator_hashes: Vec::zcash_deserialize(&mut reader)?, + hash_stop: BlockHeaderHash(reader.read_32_bytes()?), }) } @@ -534,62 +485,32 @@ impl Codec { /// /// [Zcash block header](https://zips.z.cash/protocol/protocol.pdf#page=84) fn read_headers(&self, mut reader: R) -> Result { - const ENCODED_HEADER_SIZE: usize = 4 + 32 + 32 + 32 + 4 + 4 + 32 + 3 + 1344; - let max_count = self.builder.max_len / ENCODED_HEADER_SIZE; - - let headers: Vec = reader.read_list(max_count)?; - - Ok(Message::Headers(headers)) + Ok(Message::Headers(Vec::zcash_deserialize(&mut reader)?)) } fn read_getheaders(&self, mut reader: R) -> Result { - let version = Version(reader.read_u32::()?); - - let max_count = self.builder.max_len / 32; - - let block_locator_hashes: Vec = reader.read_list(max_count)?; - - let hash_stop = BlockHeaderHash(reader.read_32_bytes()?); - Ok(Message::GetHeaders { - version, - block_locator_hashes, - hash_stop, + version: Version(reader.read_u32::()?), + block_locator_hashes: Vec::zcash_deserialize(&mut reader)?, + hash_stop: BlockHeaderHash(reader.read_32_bytes()?), }) } - fn read_inv(&self, mut reader: R) -> Result { - // encoding: 4 byte type tag + 32 byte hash - const ENCODED_INVHASH_SIZE: usize = 4 + 32; - let max_count = self.builder.max_len / ENCODED_INVHASH_SIZE; - - let hashes: Vec = reader.read_list(max_count)?; - Ok(Message::Inv(hashes)) + fn read_inv(&self, reader: R) -> Result { + Ok(Message::Inv(Vec::zcash_deserialize(reader)?)) } - fn read_getdata(&self, mut reader: R) -> Result { - // encoding: 4 byte type tag + 32 byte hash - const ENCODED_INVHASH_SIZE: usize = 4 + 32; - let max_count = self.builder.max_len / ENCODED_INVHASH_SIZE; - - let hashes: Vec = reader.read_list(max_count)?; - Ok(Message::GetData(hashes)) + fn read_getdata(&self, reader: R) -> Result { + Ok(Message::GetData(Vec::zcash_deserialize(reader)?)) } - fn read_notfound(&self, mut reader: R) -> Result { - // encoding: 4 byte type tag + 32 byte hash - const ENCODED_INVHASH_SIZE: usize = 4 + 32; - let max_count = self.builder.max_len / ENCODED_INVHASH_SIZE; - - let hashes: Vec = reader.read_list(max_count)?; - - Ok(Message::GetData(hashes)) + fn read_notfound(&self, reader: R) -> Result { + Ok(Message::GetData(Vec::zcash_deserialize(reader)?)) } fn read_tx(&self, mut reader: R) -> Result { Ok(Message::Tx { version: Version(reader.read_u32::()?), - transaction: Transaction::zcash_deserialize(&mut reader)?, }) }