diff --git a/src/blockdata/block.rs b/src/blockdata/block.rs index 09b544b2..4fdd2306 100644 --- a/src/blockdata/block.rs +++ b/src/blockdata/block.rs @@ -26,7 +26,7 @@ use util; use util::Error::{SpvBadTarget, SpvBadProofOfWork}; use util::hash::{BitcoinHash, MerkleRoot, bitcoin_merkle_root}; use util::uint::Uint256; -use consensus::encode::{VarInt, Encodable}; +use consensus::encode::Encodable; use network::constants::Network; use blockdata::transaction::Transaction; use blockdata::constants::max_target; @@ -116,17 +116,6 @@ impl MerkleRoot for Block { } } -/// A block header with txcount attached, which is given in the `headers` -/// network message. -#[derive(PartialEq, Eq, Clone, Debug)] -pub struct LoneBlockHeader { - /// The actual block header - pub header: BlockHeader, - /// The number of transactions in the block. This will always be zero - /// when the LoneBlockHeader is returned as part of a `headers` message. - pub tx_count: VarInt -} - impl BlockHeader { /// Computes the target [0, T] that a blockhash must land in to be valid pub fn target(&self) -> Uint256 { @@ -218,7 +207,6 @@ impl BitcoinHash for Block { impl_consensus_encoding!(BlockHeader, version, prev_blockhash, merkle_root, time, bits, nonce); impl_consensus_encoding!(Block, header, txdata); -impl_consensus_encoding!(LoneBlockHeader, header, tx_count); #[cfg(test)] mod tests { diff --git a/src/network/message.rs b/src/network/message.rs index e5969215..58d792b3 100644 --- a/src/network/message.rs +++ b/src/network/message.rs @@ -19,7 +19,7 @@ //! also defines (de)serialization routines for many primitives. //! -use std::iter; +use std::{iter, mem}; use std::io::Cursor; use blockdata::block; @@ -29,8 +29,9 @@ use network::message_network; use network::message_blockdata; use network::message_filter; use consensus::encode::{Decodable, Encodable}; -use consensus::encode::CheckedData; +use consensus::encode::{CheckedData, VarInt}; use consensus::encode::{self, serialize, Encoder, Decoder}; +use consensus::encode::MAX_VEC_SIZE; /// Serializer for command string #[derive(PartialEq, Eq, Clone, Debug)] @@ -97,7 +98,7 @@ pub enum NetworkMessage { /// `block` Block(block::Block), /// `headers` - Headers(Vec), + Headers(Vec), /// `getaddr` GetAddr, // TODO: checkorder, @@ -155,6 +156,19 @@ impl RawNetworkMessage { } } +struct HeaderSerializationWrapper<'a>(&'a Vec); +impl <'a, S: Encoder> Encodable for HeaderSerializationWrapper<'a> { + #[inline] + fn consensus_encode(&self, s: &mut S) -> Result<(), encode::Error> { + VarInt(self.0.len() as u64).consensus_encode(s)?; + for header in self.0.iter() { + header.consensus_encode(s)?; + 0u8.consensus_encode(s)?; + } + Ok(()) + } +} + impl Encodable for RawNetworkMessage { fn consensus_encode(&self, s: &mut S) -> Result<(), encode::Error> { self.magic.consensus_encode(s)?; @@ -169,7 +183,7 @@ impl Encodable for RawNetworkMessage { NetworkMessage::GetHeaders(ref dat) => serialize(dat), NetworkMessage::Tx(ref dat) => serialize(dat), NetworkMessage::Block(ref dat) => serialize(dat), - NetworkMessage::Headers(ref dat) => serialize(dat), + NetworkMessage::Headers(ref dat) => serialize(&HeaderSerializationWrapper(dat)), NetworkMessage::Ping(ref dat) => serialize(dat), NetworkMessage::Pong(ref dat) => serialize(dat), NetworkMessage::GetCFilters(ref dat) => serialize(dat), @@ -186,6 +200,28 @@ impl Encodable for RawNetworkMessage { } } +struct HeaderDeserializationWrapper(Vec); +impl Decodable for HeaderDeserializationWrapper { + #[inline] + fn consensus_decode(d: &mut D) -> Result { + let len = VarInt::consensus_decode(d)?.0; + let byte_size = (len as usize) + .checked_mul(mem::size_of::()) + .ok_or(encode::Error::ParseFailed("Invalid length"))?; + if byte_size > MAX_VEC_SIZE { + return Err(encode::Error::OversizedVectorAllocation { requested: byte_size, max: MAX_VEC_SIZE }) + } + let mut ret = Vec::with_capacity(len as usize); + for _ in 0..len { + ret.push(Decodable::consensus_decode(d)?); + if >::consensus_decode(d)? != 0u8 { + return Err(encode::Error::ParseFailed("Headers message should not contain transactions")); + } + } + Ok(HeaderDeserializationWrapper(ret)) + } +} + impl Decodable for RawNetworkMessage { fn consensus_decode(d: &mut D) -> Result { let magic = Decodable::consensus_decode(d)?; @@ -204,7 +240,9 @@ impl Decodable for RawNetworkMessage { "getheaders" => NetworkMessage::GetHeaders(Decodable::consensus_decode(&mut mem_d)?), "mempool" => NetworkMessage::MemPool, "block" => NetworkMessage::Block(Decodable::consensus_decode(&mut mem_d)?), - "headers" => NetworkMessage::Headers(Decodable::consensus_decode(&mut mem_d)?), + "headers" => + NetworkMessage::Headers(>>> + ::consensus_decode(&mut mem_d)?.0), "getaddr" => NetworkMessage::GetAddr, "ping" => NetworkMessage::Ping(Decodable::consensus_decode(&mut mem_d)?), "pong" => NetworkMessage::Pong(Decodable::consensus_decode(&mut mem_d)?),