From db8916a48e2bfc9ae9c18c3fa617f7302432c685 Mon Sep 17 00:00:00 2001 From: Svyatoslav Nikolsky Date: Tue, 16 Jun 2020 23:51:45 +0300 Subject: [PATCH] fix BlockAttributes encoding (#6281) --- client/network/src/block_requests.rs | 49 +++++++++++++++------- client/network/src/light_client_handler.rs | 35 ++++++++++++---- client/network/src/protocol/message.rs | 14 +++++++ 3 files changed, 74 insertions(+), 24 deletions(-) diff --git a/client/network/src/block_requests.rs b/client/network/src/block_requests.rs index ae5a3a0b4e894..6d698a7300126 100644 --- a/client/network/src/block_requests.rs +++ b/client/network/src/block_requests.rs @@ -277,21 +277,13 @@ where return SendRequestOutcome::NotConnected; }; - let protobuf_rq = schema::v1::BlockRequest { - fields: u32::from_be_bytes([req.fields.bits(), 0, 0, 0]), - from_block: match req.from { - message::FromBlock::Hash(h) => - Some(schema::v1::block_request::FromBlock::Hash(h.encode())), - message::FromBlock::Number(n) => - Some(schema::v1::block_request::FromBlock::Number(n.encode())), - }, - to_block: req.to.map(|h| h.encode()).unwrap_or_default(), - direction: match req.direction { - message::Direction::Ascending => schema::v1::Direction::Ascending as i32, - message::Direction::Descending => schema::v1::Direction::Descending as i32, - }, - max_blocks: req.max.unwrap_or(0), - }; + let protobuf_rq = build_protobuf_block_request( + req.fields, + req.from.clone(), + req.to.clone(), + req.direction, + req.max, + ); let mut buf = Vec::with_capacity(protobuf_rq.encoded_len()); if let Err(err) = protobuf_rq.encode(&mut buf) { @@ -386,7 +378,7 @@ where return Err(io::Error::new(io::ErrorKind::Other, msg).into()) }; - let attributes = BlockAttributes::decode(&mut request.fields.to_be_bytes().as_ref())?; + let attributes = BlockAttributes::from_be_u32(request.fields)?; let get_header = attributes.contains(BlockAttributes::HEADER); let get_body = attributes.contains(BlockAttributes::BODY); let get_justification = attributes.contains(BlockAttributes::JUSTIFICATION); @@ -826,3 +818,28 @@ where }.boxed() } } + +/// Build protobuf block request message. +pub(crate) fn build_protobuf_block_request( + attributes: BlockAttributes, + from_block: message::FromBlock, + to_block: Option, + direction: message::Direction, + max_blocks: Option, +) -> schema::v1::BlockRequest { + schema::v1::BlockRequest { + fields: attributes.to_be_u32(), + from_block: match from_block { + message::FromBlock::Hash(h) => + Some(schema::v1::block_request::FromBlock::Hash(h.encode())), + message::FromBlock::Number(n) => + Some(schema::v1::block_request::FromBlock::Number(n.encode())), + }, + to_block: to_block.map(|h| h.encode()).unwrap_or_default(), + direction: match direction { + message::Direction::Ascending => schema::v1::Direction::Ascending as i32, + message::Direction::Descending => schema::v1::Direction::Descending as i32, + }, + max_blocks: max_blocks.unwrap_or(0), + } +} diff --git a/client/network/src/light_client_handler.rs b/client/network/src/light_client_handler.rs index 236ae817474ab..ab6bea8761b95 100644 --- a/client/network/src/light_client_handler.rs +++ b/client/network/src/light_client_handler.rs @@ -27,9 +27,10 @@ use bytes::Bytes; use codec::{self, Encode, Decode}; use crate::{ + block_requests::build_protobuf_block_request, chain::Client, config::ProtocolId, - protocol::message::BlockAttributes, + protocol::message::{BlockAttributes, Direction, FromBlock}, schema, }; use futures::{channel::oneshot, future::BoxFuture, prelude::*, stream::FuturesUnordered}; @@ -1062,13 +1063,13 @@ fn retries(request: &Request) -> usize { fn serialize_request(request: &Request) -> Result, prost::EncodeError> { let request = match request { Request::Body { request, .. } => { - let rq = schema::v1::BlockRequest { - fields: u32::from(BlockAttributes::BODY.bits()), - from_block: Some(schema::v1::block_request::FromBlock::Hash(request.header.hash().encode())), - to_block: Vec::new(), - direction: schema::v1::Direction::Ascending as i32, - max_blocks: 1, - }; + let rq = build_protobuf_block_request::<_, NumberFor>( + BlockAttributes::BODY, + FromBlock::Hash(request.header.hash()), + None, + Direction::Ascending, + Some(1), + ); let mut buf = Vec::with_capacity(rq.encoded_len()); rq.encode(&mut buf)?; return Ok(buf); @@ -2036,4 +2037,22 @@ mod tests { assert_eq!(vec![(100, 2)], task::block_on(chan.1).unwrap().unwrap()); // ^--- from `DummyFetchChecker::check_changes_proof` } + + #[test] + fn body_request_fields_encoded_properly() { + let (sender, _) = oneshot::channel(); + let serialized_request = serialize_request::(&Request::Body { + request: RemoteBodyRequest { + header: dummy_header(), + retry_count: None, + }, + sender, + }).unwrap(); + let deserialized_request = schema::v1::BlockRequest::decode(&serialized_request[..]).unwrap(); + assert!( + BlockAttributes::from_be_u32(deserialized_request.fields) + .unwrap() + .contains(BlockAttributes::BODY) + ); + } } diff --git a/client/network/src/protocol/message.rs b/client/network/src/protocol/message.rs index bb2253b733864..a7fbb92387cf6 100644 --- a/client/network/src/protocol/message.rs +++ b/client/network/src/protocol/message.rs @@ -87,6 +87,20 @@ bitflags! { } } +impl BlockAttributes { + /// Encodes attributes as big endian u32, compatible with SCALE-encoding (i.e the + /// significant byte has zero index). + pub fn to_be_u32(&self) -> u32 { + u32::from_be_bytes([self.bits(), 0, 0, 0]) + } + + /// Decodes attributes, encoded with the `encode_to_be_u32()` call. + pub fn from_be_u32(encoded: u32) -> Result { + BlockAttributes::from_bits(encoded.to_be_bytes()[0]) + .ok_or_else(|| Error::from("Invalid BlockAttributes")) + } +} + impl Encode for BlockAttributes { fn encode_to(&self, dest: &mut T) { dest.push_byte(self.bits())