fix BlockAttributes encoding (#6281)

substrate/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<Hash: Encode, Number: Encode>(
+	attributes: BlockAttributes,
+	from_block: message::FromBlock<Hash, Number>,
+	to_block: Option<Hash>,
+	direction: message::Direction,
+	max_blocks: Option<u32>,
+) -> 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),
+	}
+}

substrate/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<B: Block>(request: &Request<B>) -> usize {
 fn serialize_request<B: Block>(request: &Request<B>) -> Result<Vec<u8>, 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<B>>(
+				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::<Block>(&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)
+		);
+	}
 }

substrate/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<Self, Error> {
+		BlockAttributes::from_bits(encoded.to_be_bytes()[0])
+			.ok_or_else(|| Error::from("Invalid BlockAttributes"))
+	}
+}
+
 impl Encode for BlockAttributes {
 	fn encode_to<T: Output>(&self, dest: &mut T) {
 		dest.push_byte(self.bits())