Refactor p2p reader

core/Cargo.toml

@@ -27,6 +27,7 @@ siphasher = "0.3"
 log = "0.4"
 chrono = { version = "0.4.11", features = ["serde"] }
 zeroize = { version = "1.1", features =["zeroize_derive"] }
+bytes = "0.5"
 
 keychain = { package = "grin_keychain", path = "../keychain", version = "4.1.0-alpha.1" }
 util = { package = "grin_util", path = "../util", version = "4.1.0-alpha.1" }

core/src/global.rs

@@ -26,7 +26,7 @@ use crate::core::block::HeaderVersion;
 use crate::{
 	pow::{
 		self, new_cuckaroo_ctx, new_cuckarood_ctx, new_cuckaroom_ctx, new_cuckarooz_ctx,
-		new_cuckatoo_ctx, PoWContext,
+		new_cuckatoo_ctx, BitVec, PoWContext,
 	},
 	ser::ProtocolVersion,
 };
@@ -411,3 +411,51 @@ where
 	last_n.reverse();
 	last_n
 }
+
+/// Calculates the size of a header (in bytes) given a number of edge bits in the PoW
+#[inline]
+pub fn header_size_bytes(edge_bits: u8) -> usize {
+	let size = 2 + 2 * 8 + 5 * 32 + 32 + 2 * 8;
+	let proof_size = 8 + 4 + 8 + 1 + BitVec::bytes_len(edge_bits as usize * proofsize());
+	size + proof_size
+}
+
+#[cfg(test)]
+mod test {
+	use super::*;
+	use crate::core::Block;
+	use crate::genesis::*;
+	use crate::pow::mine_genesis_block;
+	use crate::ser::{BinWriter, Writeable};
+
+	fn test_header_len(genesis: Block) {
+		let mut raw = Vec::<u8>::with_capacity(1_024);
+		let mut writer = BinWriter::new(&mut raw, ProtocolVersion::local());
+		genesis.header.write(&mut writer).unwrap();
+		assert_eq!(raw.len(), header_size_bytes(genesis.header.pow.edge_bits()));
+	}
+
+	#[test]
+	fn automated_testing_header_len() {
+		set_local_chain_type(ChainTypes::AutomatedTesting);
+		test_header_len(mine_genesis_block().unwrap());
+	}
+
+	#[test]
+	fn user_testing_header_len() {
+		set_local_chain_type(ChainTypes::UserTesting);
+		test_header_len(mine_genesis_block().unwrap());
+	}
+
+	#[test]
+	fn floonet_header_len() {
+		set_local_chain_type(ChainTypes::Floonet);
+		test_header_len(genesis_floo());
+	}
+
+	#[test]
+	fn mainnet_header_len() {
+		set_local_chain_type(ChainTypes::Mainnet);
+		test_header_len(genesis_main());
+	}
+}

core/src/pow/types.rs

@@ -471,15 +471,17 @@ impl Writeable for Proof {
 	}
 }
 
+/// A bit vector
 // TODO this could likely be optimized by writing whole bytes (or even words)
 // in the `BitVec` at once, dealing with the truncation, instead of bits by bits
-struct BitVec {
+pub struct BitVec {
 	bits: Vec<u8>,
 }
 
 impl BitVec {
 	/// Number of bytes required to store the provided number of bits
-	fn bytes_len(bits_len: usize) -> usize {
+	#[inline]
+	pub fn bytes_len(bits_len: usize) -> usize {
 		(bits_len + 7) / 8
 	}
 

core/src/ser.rs

@@ -22,6 +22,7 @@
 use crate::core::hash::{DefaultHashable, Hash, Hashed};
 use crate::global::PROTOCOL_VERSION;
 use byteorder::{BigEndian, ByteOrder, ReadBytesExt};
+use bytes::Buf;
 use keychain::{BlindingFactor, Identifier, IDENTIFIER_SIZE};
 use std::convert::TryInto;
 use std::fmt::{self, Debug};
@@ -79,6 +80,12 @@ impl From<io::Error> for Error {
 	}
 }
 
+impl From<io::ErrorKind> for Error {
+	fn from(e: io::ErrorKind) -> Error {
+		Error::IOErr(format!("{}", io::Error::from(e)), e)
+	}
+}
+
 impl fmt::Display for Error {
 	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 		match *self {
@@ -575,6 +582,109 @@ impl<'a> Reader for StreamingReader<'a> {
 	}
 }
 
+/// Protocol version-aware wrapper around a `Buf` impl
+pub struct BufReader<'a, B: Buf> {
+	inner: &'a mut B,
+	version: ProtocolVersion,
+	bytes_read: usize,
+}
+
+impl<'a, B: Buf> BufReader<'a, B> {
+	/// Construct a new BufReader
+	pub fn new(buf: &'a mut B, version: ProtocolVersion) -> Self {
+		Self {
+			inner: buf,
+			version,
+			bytes_read: 0,
+		}
+	}
+
+	/// Check whether the buffer has enough bytes remaining to perform a read
+	fn has_remaining(&mut self, len: usize) -> Result<(), Error> {
+		if self.inner.remaining() >= len {
+			self.bytes_read += len;
+			Ok(())
+		} else {
+			Err(io::ErrorKind::UnexpectedEof.into())
+		}
+	}
+
+	/// The total bytes read
+	pub fn bytes_read(&self) -> u64 {
+		self.bytes_read as u64
+	}
+
+	/// Convenience function to read from the buffer and deserialize
+	pub fn body<T: Readable>(&mut self) -> Result<T, Error> {
+		T::read(self)
+	}
+}
+
+impl<'a, B: Buf> Reader for BufReader<'a, B> {
+	fn read_u8(&mut self) -> Result<u8, Error> {
+		self.has_remaining(1)?;
+		Ok(self.inner.get_u8())
+	}
+
+	fn read_u16(&mut self) -> Result<u16, Error> {
+		self.has_remaining(2)?;
+		Ok(self.inner.get_u16())
+	}
+
+	fn read_u32(&mut self) -> Result<u32, Error> {
+		self.has_remaining(4)?;
+		Ok(self.inner.get_u32())
+	}
+
+	fn read_u64(&mut self) -> Result<u64, Error> {
+		self.has_remaining(8)?;
+		Ok(self.inner.get_u64())
+	}
+
+	fn read_i32(&mut self) -> Result<i32, Error> {
+		self.has_remaining(4)?;
+		Ok(self.inner.get_i32())
+	}
+
+	fn read_i64(&mut self) -> Result<i64, Error> {
+		self.has_remaining(8)?;
+		Ok(self.inner.get_i64())
+	}
+
+	fn read_bytes_len_prefix(&mut self) -> Result<Vec<u8>, Error> {
+		let len = self.read_u64()?;
+		self.read_fixed_bytes(len as usize)
+	}
+
+	fn read_fixed_bytes(&mut self, len: usize) -> Result<Vec<u8>, Error> {
+		// not reading more than 100k bytes in a single read
+		if len > 100_000 {
+			return Err(Error::TooLargeReadErr);
+		}
+		self.has_remaining(len)?;
+
+		let mut buf = vec![0; len];
+		self.inner.copy_to_slice(&mut buf[..]);
+		Ok(buf)
+	}
+
+	fn expect_u8(&mut self, val: u8) -> Result<u8, Error> {
+		let b = self.read_u8()?;
+		if b == val {
+			Ok(b)
+		} else {
+			Err(Error::UnexpectedData {
+				expected: vec![val],
+				received: vec![b],
+			})
+		}
+	}
+
+	fn protocol_version(&self) -> ProtocolVersion {
+		self.version
+	}
+}
+
 impl Readable for Commitment {
 	fn read<R: Reader>(reader: &mut R) -> Result<Commitment, Error> {
 		let a = reader.read_fixed_bytes(PEDERSEN_COMMITMENT_SIZE)?;

p2p/Cargo.toml

@@ -20,6 +20,7 @@ serde_derive = "1"
 tempfile = "3.1"
 log = "0.4"
 chrono = { version = "0.4.11", features = ["serde"] }
+bytes = "0.5"
 
 grin_core = { path = "../core", version = "4.1.0-alpha.1" }
 grin_store = { path = "../store", version = "4.1.0-alpha.1" }