diff --git a/zebra-chain/src/serialization.rs b/zebra-chain/src/serialization.rs index 971b7d1d7..08bec886e 100644 --- a/zebra-chain/src/serialization.rs +++ b/zebra-chain/src/serialization.rs @@ -244,6 +244,36 @@ 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. + #[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; + + // Preallocate a buffer, performing a single allocation in the + // honest case. Although the size of the received data buffer + // is bounded by the codec's max_len field, it's still + // 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, calculate the maximum count for a valid + // message from the codec's max_len using encoded_type_size. + 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/codec.rs b/zebra-network/src/protocol/codec.rs index 3e5455da9..a8f29f03e 100644 --- a/zebra-network/src/protocol/codec.rs +++ b/zebra-network/src/protocol/codec.rs @@ -383,26 +383,11 @@ impl Codec { fn read_addr(&self, mut reader: R) -> Result { use crate::meta_addr::MetaAddr; - // XXX we may want to factor this logic out into - // fn read_vec(reader: R) -> Result, Error> - // on ReadZcashExt (and similarly for WriteZcashExt) - let count = reader.read_compactsize()? as usize; - // Preallocate a buffer, performing a single allocation in the honest - // case. Although the size of the recieved data buffer is bounded by the - // codec's max_len field, it's still possible for someone to send a - // short addr message with a large count field, so if we naively trust - // the count field we could be tricked into preallocating a large - // buffer. Instead, calculate the maximum count for a valid message from - // the codec's max_len using ENCODED_ADDR_SIZE. - // // 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 mut addrs = Vec::with_capacity(std::cmp::min(count, max_count)); - for _ in 0..count { - addrs.push(MetaAddr::zcash_deserialize(&mut reader)?); - } + let addrs: Vec = reader.read_list(max_count)?; Ok(Message::Addr(addrs)) } @@ -483,43 +468,31 @@ impl Codec { }) } - fn _read_generic_inventory_hash_vector( - &self, - mut reader: R, - ) -> Result, Error> { - let count = reader.read_compactsize()? as usize; - // Preallocate a buffer, performing a single allocation in the honest - // case. Although the size of the received data buffer is bounded by the - // codec's max_len field, it's still 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, calculate the maximum count for a valid message from - // the codec's max_len using ENCODED_INVHASH_SIZE. - // + 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 mut hashes = Vec::with_capacity(std::cmp::min(count, max_count)); - for _ in 0..count { - hashes.push(InventoryHash::zcash_deserialize(&mut reader)?); - } - - return Ok(hashes); - } - - fn read_inv(&self, reader: R) -> Result { - let hashes = self._read_generic_inventory_hash_vector(reader)?; + let hashes: Vec = reader.read_list(max_count)?; Ok(Message::Inv(hashes)) } - fn read_getdata(&self, reader: R) -> Result { - let hashes = self._read_generic_inventory_hash_vector(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_notfound(&self, reader: R) -> Result { - let hashes = self._read_generic_inventory_hash_vector(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)) }