Add read_list to ReadZcashExt

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<T>` with a leading count in a safer manner.
+    #[inline]
+    fn read_list<T: ZcashDeserialize>(
+        &mut self,
+        max_count: usize,
+    ) -> Result<Vec<T>, 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.

zebra-network/src/protocol/codec.rs

@@ -383,26 +383,11 @@ impl Codec {
     fn read_addr<R: Read>(&self, mut reader: R) -> Result<Message, Error> {
         use crate::meta_addr::MetaAddr;
 
-        // XXX we may want to factor this logic out into
-        // fn read_vec<R: Read, T: ZcashDeserialize>(reader: R) -> Result<Vec<T>, 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<MetaAddr> = reader.read_list(max_count)?;
 
         Ok(Message::Addr(addrs))
     }
@@ -483,43 +468,31 @@ impl Codec {
         })
     }
 
-    fn _read_generic_inventory_hash_vector<R: Read>(
-        &self,
-        mut reader: R,
-    ) -> Result<Vec<InventoryHash>, 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<R: Read>(&self, mut reader: R) -> Result<Message, Error> {
         // 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<R: Read>(&self, reader: R) -> Result<Message, Error> {
-        let hashes = self._read_generic_inventory_hash_vector(reader)?;
+        let hashes: Vec<InventoryHash> = reader.read_list(max_count)?;
         Ok(Message::Inv(hashes))
     }
 
-    fn read_getdata<R: Read>(&self, reader: R) -> Result<Message, Error> {
-        let hashes = self._read_generic_inventory_hash_vector(reader)?;
+    fn read_getdata<R: Read>(&self, mut reader: R) -> Result<Message, Error> {
+        // 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<InventoryHash> = reader.read_list(max_count)?;
         Ok(Message::GetData(hashes))
     }
 
-    fn read_notfound<R: Read>(&self, reader: R) -> Result<Message, Error> {
-        let hashes = self._read_generic_inventory_hash_vector(reader)?;
+    fn read_notfound<R: Read>(&self, mut reader: R) -> Result<Message, Error> {
+        // 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<InventoryHash> = reader.read_list(max_count)?;
+
         Ok(Message::GetData(hashes))
     }