From b3e094bc4002b6674b2ee8480c9cf05abbac9d8c Mon Sep 17 00:00:00 2001
From: Henry de Valence <hdevalence@hdevalence.ca>
Date: Thu, 19 Sep 2019 06:30:20 -0700
Subject: [PATCH] Clean parsing via ReadZcashExt read-array helpers.

This adds convenience methods to `ReadZcashExt` that read 4 and 12 byte
fixed size arrays from the `Reader`, making the actual parsing code more
legible.

Closes #10.
---
 zebra-chain/src/serialization.rs | 16 ++++++++++++++++
 zebra-network/src/message.rs     | 27 ++++++++-------------------
 2 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/zebra-chain/src/serialization.rs b/zebra-chain/src/serialization.rs
index bdae4d68..aa08866c 100644
--- a/zebra-chain/src/serialization.rs
+++ b/zebra-chain/src/serialization.rs
@@ -220,6 +220,22 @@ pub trait ReadZcashExt: io::Read {
         self.read_exact(&mut buf)?;
         String::from_utf8(buf).map_err(|_| io::ErrorKind::InvalidData.into())
     }
+
+    /// Convenience method to read a `[u8; 4]`.
+    #[inline]
+    fn read_4_bytes(&mut self) -> io::Result<[u8; 4]> {
+        let mut bytes = [0; 4];
+        self.read_exact(&mut bytes)?;
+        Ok(bytes)
+    }
+
+    /// Convenience method to read a `[u8; 12]`.
+    #[inline]
+    fn read_12_bytes(&mut self) -> io::Result<[u8; 12]> {
+        let mut bytes = [0; 12];
+        self.read_exact(&mut bytes)?;
+        Ok(bytes)
+    }
 }
 
 /// Mark all types implementing `Read` as implementing the extension.
diff --git a/zebra-network/src/message.rs b/zebra-network/src/message.rs
index d0e39e22..0ce4b09d 100644
--- a/zebra-network/src/message.rs
+++ b/zebra-network/src/message.rs
@@ -392,29 +392,18 @@ impl Message {
         version: Version,
     ) -> Result<Self, SerializationError> {
         use SerializationError::ParseError;
-        let message_magic = {
-            let mut bytes = [0u8; 4];
-            reader.read_exact(&mut bytes)?;
-            Magic(bytes)
-        };
+
+        // Read header data
+        let message_magic = Magic(reader.read_4_bytes()?);
+        let command = reader.read_12_bytes()?;
+        let body_len = reader.read_u32::<LittleEndian>()? as usize;
+        let checksum = Sha256dChecksum(reader.read_4_bytes()?);
+
         if magic != message_magic {
             return Err(ParseError("Message has incorrect magic value"));
         }
 
-        let command = {
-            let mut bytes = [0u8; 12];
-            reader.read_exact(&mut bytes)?;
-            bytes
-        };
-
-        let body_len = reader.read_u32::<LittleEndian>()? as usize;
-        // XXX ugly
-        let checksum = {
-            let mut bytes = [0u8; 4];
-            reader.read_exact(&mut bytes)?;
-            Sha256dChecksum(bytes)
-        };
-
+        // XXX bound the body_len value to avoid large attacker-controlled allocs
         // XXX add a ChecksumReader<R: Read>(R) wrapper and avoid this
         let body = {
             let mut bytes = vec![0; body_len];