From 8608c2eae475c3af926729671b0167b49e06b142 Mon Sep 17 00:00:00 2001
From: Pierre Krieger <pierre.krieger1708@gmail.com>
Date: Fri, 9 Apr 2021 17:22:47 +0200
Subject: [PATCH] Cap the warp sync proof by size, not by fragments (#8578)

* Cap the warp sync proof by size, not by fragments

* Add a final debug assert

* Check size after
---
 .../finality-grandpa-warp-sync/src/lib.rs     |  2 +-
 .../finality-grandpa-warp-sync/src/proof.rs   | 34 ++++++++++++-------
 2 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/substrate/client/finality-grandpa-warp-sync/src/lib.rs b/substrate/client/finality-grandpa-warp-sync/src/lib.rs
index 285a5fe736..dca6c2ad1b 100644
--- a/substrate/client/finality-grandpa-warp-sync/src/lib.rs
+++ b/substrate/client/finality-grandpa-warp-sync/src/lib.rs
@@ -66,7 +66,7 @@ pub fn generate_request_response_config(protocol_id: ProtocolId) -> RequestRespo
 	RequestResponseConfig {
 		name: generate_protocol_name(protocol_id).into(),
 		max_request_size: 32,
-		max_response_size: 16 * 1024 * 1024,
+		max_response_size: proof::MAX_WARP_SYNC_PROOF_SIZE as u64,
 		request_timeout: Duration::from_secs(10),
 		inbound_queue: None,
 	}
diff --git a/substrate/client/finality-grandpa-warp-sync/src/proof.rs b/substrate/client/finality-grandpa-warp-sync/src/proof.rs
index 08effcf1c2..26560c10fe 100644
--- a/substrate/client/finality-grandpa-warp-sync/src/proof.rs
+++ b/substrate/client/finality-grandpa-warp-sync/src/proof.rs
@@ -29,8 +29,8 @@ use sp_runtime::{
 
 use crate::HandleRequestError;
 
-/// The maximum number of authority set change proofs to include in a single warp sync proof.
-const MAX_CHANGES_PER_WARP_SYNC_PROOF: usize = 256;
+/// The maximum size in bytes of the `WarpSyncProof`.
+pub(super) const MAX_WARP_SYNC_PROOF_SIZE: usize = 16 * 1024 * 1024;
 
 /// A proof of an authority set change.
 #[derive(Decode, Encode)]
@@ -53,7 +53,7 @@ pub struct WarpSyncProof<Block: BlockT> {
 impl<Block: BlockT> WarpSyncProof<Block> {
 	/// Generates a warp sync proof starting at the given block. It will generate authority set
 	/// change proofs for all changes that happened from `begin` until the current authority set
-	/// (capped by MAX_CHANGES_PER_WARP_SYNC_PROOF).
+	/// (capped by MAX_WARP_SYNC_PROOF_SIZE).
 	pub fn generate<Backend>(
 		backend: &Backend,
 		begin: Block::Hash,
@@ -88,14 +88,10 @@ impl<Block: BlockT> WarpSyncProof<Block> {
 		}
 
 		let mut proofs = Vec::new();
+		let mut proofs_encoded_len = 0;
 		let mut proof_limit_reached = false;
 
 		for (_, last_block) in set_changes.iter_from(begin_number) {
-			if proofs.len() >= MAX_CHANGES_PER_WARP_SYNC_PROOF {
-				proof_limit_reached = true;
-				break;
-			}
-
 			let header = blockchain.header(BlockId::Number(*last_block))?.expect(
 				"header number comes from previously applied set changes; must exist in db; qed.",
 			);
@@ -120,10 +116,22 @@ impl<Block: BlockT> WarpSyncProof<Block> {
 
 			let justification = GrandpaJustification::<Block>::decode(&mut &justification[..])?;
 
-			proofs.push(WarpSyncFragment {
+			let proof = WarpSyncFragment {
 				header: header.clone(),
 				justification,
-			});
+			};
+			let proof_size = proof.encoded_size();
+
+			// Check for the limit. We remove some bytes from the maximum size, because we're only
+			// counting the size of the `WarpSyncFragment`s. The extra margin is here to leave
+			// room for rest of the data (the size of the `Vec` and the boolean).
+			if proofs_encoded_len + proof_size >= MAX_WARP_SYNC_PROOF_SIZE - 50 {
+				proof_limit_reached = true;
+				break;
+			}
+
+			proofs_encoded_len += proof_size;
+			proofs.push(proof);
 		}
 
 		let is_finished = if proof_limit_reached {
@@ -156,10 +164,12 @@ impl<Block: BlockT> WarpSyncProof<Block> {
 			true
 		};
 
-		Ok(WarpSyncProof {
+		let final_outcome = WarpSyncProof {
 			proofs,
 			is_finished,
-		})
+		};
+		debug_assert!(final_outcome.encoded_size() <= MAX_WARP_SYNC_PROOF_SIZE);
+		Ok(final_outcome)
 	}
 
 	/// Verifies the warp sync proof starting at the given set id and with the given authorities.