Fast sync (#8884)

* State sync

* Importing state fixes

* Bugfixes

* Sync with proof

* Status reporting

* Unsafe sync mode

* Sync test

* Cleanup

* Apply suggestions from code review

Co-authored-by: cheme <emericchevalier.pro@gmail.com>
Co-authored-by: Pierre Krieger <pierre.krieger1708@gmail.com>

* set_genesis_storage

* Extract keys from range proof

* Detect iter completion

* Download and import bodies with fast sync

* Replaced meta updates tuple with a struct

* Fixed reverting finalized state

* Reverted timeout

* Typo

* Doc

* Doc

* Fixed light client test

* Fixed error handling

* Tweaks

* More UpdateMeta changes

* Rename convert_transaction

* Apply suggestions from code review

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Apply suggestions from code review

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

* Code review suggestions

* Fixed count handling

Co-authored-by: cheme <emericchevalier.pro@gmail.com>
Co-authored-by: Pierre Krieger <pierre.krieger1708@gmail.com>
Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

substrate/client/network/src/protocol/sync.rs

@@ -31,14 +31,16 @@
 
 use codec::Encode;
 use blocks::BlockCollection;
-use sp_blockchain::{Error as ClientError, Info as BlockchainInfo, HeaderMetadata};
+use state::StateSync;
+use sp_blockchain::{Error as ClientError, HeaderMetadata};
 use sp_consensus::{BlockOrigin, BlockStatus,
 	block_validation::{BlockAnnounceValidator, Validation},
 	import_queue::{IncomingBlock, BlockImportResult, BlockImportError}
 };
 use crate::protocol::message::{
-	self, BlockAnnounce, BlockAttributes, BlockRequest, BlockResponse, Roles,
+	self, BlockAnnounce, BlockAttributes, BlockRequest, BlockResponse,
 };
+use crate::schema::v1::{StateResponse, StateRequest};
 use either::Either;
 use extra_requests::ExtraRequests;
 use libp2p::PeerId;
@@ -59,6 +61,7 @@ use futures::{task::Poll, Future, stream::FuturesUnordered, FutureExt, StreamExt
 
 mod blocks;
 mod extra_requests;
+mod state;
 
 /// Maximum blocks to request in a single packet.
 const MAX_BLOCKS_TO_REQUEST: usize = 128;
@@ -84,6 +87,9 @@ const MAX_CONCURRENT_BLOCK_ANNOUNCE_VALIDATIONS: usize = 256;
 /// See [`MAX_CONCURRENT_BLOCK_ANNOUNCE_VALIDATIONS`] for more information.
 const MAX_CONCURRENT_BLOCK_ANNOUNCE_VALIDATIONS_PER_PEER: usize = 4;
 
+/// Pick the state to sync as the latest finalized number minus this.
+const STATE_SYNC_FINALITY_THRESHOLD: u32 = 8;
+
 /// We use a heuristic that with a high likelihood, by the time
 /// `MAJOR_SYNC_BLOCKS` have been imported we'll be on the same
 /// chain as (or at least closer to) the peer so we want to delay
@@ -183,11 +189,8 @@ pub struct ChainSync<B: BlockT> {
 	best_queued_number: NumberFor<B>,
 	/// The best block hash in our queue of blocks to import
 	best_queued_hash: B::Hash,
-	/// The role of this node, e.g. light or full
-	role: Roles,
-	/// What block attributes we require for this node, usually derived from
-	/// what role we are, but could be customized
-	required_block_attributes: message::BlockAttributes,
+	/// Current mode (full/light)
+	mode: SyncMode,
 	/// Any extra justification requests.
 	extra_justifications: ExtraRequests<B>,
 	/// A set of hashes of blocks that are being downloaded or have been
@@ -209,6 +212,11 @@ pub struct ChainSync<B: BlockT> {
 	>,
 	/// Stats per peer about the number of concurrent block announce validations.
 	block_announce_validation_per_peer_stats: HashMap<PeerId, usize>,
+	/// State sync in progress, if any.
+	state_sync: Option<StateSync<B>>,
+	/// Enable importing existing blocks. This is used used after the state download to
+	/// catch up to the latest state while re-importing blocks.
+	import_existing: bool,
 }
 
 /// All the data we have about a Peer that we are trying to sync with
@@ -281,6 +289,8 @@ pub enum PeerSyncState<B: BlockT> {
 	DownloadingStale(B::Hash),
 	/// Downloading justification for given block hash.
 	DownloadingJustification(B::Hash),
+	/// Downloading state.
+	DownloadingState,
 }
 
 impl<B: BlockT> PeerSyncState<B> {
@@ -298,6 +308,15 @@ pub enum SyncState {
 	Downloading
 }
 
+/// Reported state download progress.
+#[derive(Clone, Eq, PartialEq, Debug)]
+pub struct StateDownloadProgress {
+	/// Estimated download percentage.
+	pub percentage: u32,
+	/// Total state size in bytes downloaded so far.
+	pub size: u64,
+}
+
 /// Syncing status and statistics.
 #[derive(Clone)]
 pub struct Status<B: BlockT> {
@@ -309,6 +328,8 @@ pub struct Status<B: BlockT> {
 	pub num_peers: u32,
 	/// Number of blocks queued for import
 	pub queued_blocks: u32,
+	/// State sync status in progress, if any.
+	pub state_sync: Option<StateDownloadProgress>,
 }
 
 /// A peer did not behave as expected and should be reported.
@@ -344,6 +365,15 @@ impl<B: BlockT> OnBlockData<B> {
 	}
 }
 
+/// Result of [`ChainSync::on_state_data`].
+#[derive(Debug)]
+pub enum OnStateData<B: BlockT> {
+	/// The block and state that should be imported.
+	Import(BlockOrigin, IncomingBlock<B>),
+	/// A new state request needs to be made to the given peer.
+	Request(PeerId, StateRequest)
+}
+
 /// Result of [`ChainSync::poll_block_announce_validation`].
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum PollBlockAnnounceValidation<H> {
@@ -429,6 +459,20 @@ pub enum OnBlockJustification<B: BlockT> {
 	}
 }
 
+
+/// Operation mode.
+#[derive(Debug, PartialEq, Eq)]
+pub enum SyncMode {
+	// Sync headers only
+	Light,
+	// Sync headers and block bodies
+	Full,
+	// Sync headers and the last finalied state
+	LightState {
+		skip_proofs: bool
+	},
+}
+
 /// Result of [`ChainSync::has_slot_for_block_announce_validation`].
 enum HasSlotForBlockAnnounceValidation {
 	/// Yes, there is a slot for the block announce validation.
@@ -442,27 +486,19 @@ enum HasSlotForBlockAnnounceValidation {
 impl<B: BlockT> ChainSync<B> {
 	/// Create a new instance.
 	pub fn new(
-		role: Roles,
+		mode: SyncMode,
 		client: Arc<dyn crate::chain::Client<B>>,
-		info: &BlockchainInfo<B>,
 		block_announce_validator: Box<dyn BlockAnnounceValidator<B> + Send>,
 		max_parallel_downloads: u32,
-	) -> Self {
-		let mut required_block_attributes = BlockAttributes::HEADER | BlockAttributes::JUSTIFICATION;
-
-		if role.is_full() {
-			required_block_attributes |= BlockAttributes::BODY
-		}
-
-		ChainSync {
+	) -> Result<Self, ClientError> {
+		let mut sync = ChainSync {
 			client,
 			peers: HashMap::new(),
 			blocks: BlockCollection::new(),
-			best_queued_hash: info.best_hash,
-			best_queued_number: info.best_number,
+			best_queued_hash: Default::default(),
+			best_queued_number: Zero::zero(),
 			extra_justifications: ExtraRequests::new("justification"),
-			role,
-			required_block_attributes,
+			mode,
 			queue_blocks: Default::default(),
 			fork_targets: Default::default(),
 			pending_requests: Default::default(),
@@ -471,6 +507,27 @@ impl<B: BlockT> ChainSync<B> {
 			downloaded_blocks: 0,
 			block_announce_validation: Default::default(),
 			block_announce_validation_per_peer_stats: Default::default(),
+			state_sync: None,
+			import_existing: false,
+		};
+		sync.reset_sync_start_point()?;
+		Ok(sync)
+	}
+
+	fn required_block_attributes(&self) -> BlockAttributes {
+		match self.mode {
+			SyncMode::Full => BlockAttributes::HEADER | BlockAttributes::JUSTIFICATION | BlockAttributes::BODY,
+			SyncMode::Light => BlockAttributes::HEADER | BlockAttributes::JUSTIFICATION,
+			SyncMode::LightState { .. } =>
+				BlockAttributes::HEADER | BlockAttributes::JUSTIFICATION | BlockAttributes::BODY,
+		}
+	}
+
+	fn skip_execution(&self) -> bool {
+		match self.mode {
+			SyncMode::Full => false,
+			SyncMode::Light => true,
+			SyncMode::LightState { .. } => true,
 		}
 	}
 
@@ -502,6 +559,7 @@ impl<B: BlockT> ChainSync<B> {
 			best_seen_block: best_seen,
 			num_peers: self.peers.len() as u32,
 			queued_blocks: self.queue_blocks.len() as u32,
+			state_sync: self.state_sync.as_ref().map(|s| s.progress()),
 		}
 	}
 
@@ -607,7 +665,7 @@ impl<B: BlockT> ChainSync<B> {
 				);
 				self.peers.insert(who.clone(), PeerSync {
 					peer_id: who.clone(),
-					common_number: best_number,
+					common_number: std::cmp::min(self.best_queued_number, best_number),
 					best_hash,
 					best_number,
 					state: PeerSyncState::Available,
@@ -718,7 +776,7 @@ impl<B: BlockT> ChainSync<B> {
 
 	/// Get an iterator over all block requests of all peers.
 	pub fn block_requests(&mut self) -> impl Iterator<Item = (&PeerId, BlockRequest<B>)> + '_ {
-		if self.pending_requests.is_empty() {
+		if self.pending_requests.is_empty() || self.state_sync.is_some() {
 			return Either::Left(std::iter::empty())
 		}
 		if self.queue_blocks.len() > MAX_IMPORTING_BLOCKS {
@@ -726,10 +784,10 @@ impl<B: BlockT> ChainSync<B> {
 			return Either::Left(std::iter::empty())
 		}
 		let major_sync = self.status().state == SyncState::Downloading;
+		let attrs = self.required_block_attributes();
 		let blocks = &mut self.blocks;
-		let attrs = &self.required_block_attributes;
 		let fork_targets = &mut self.fork_targets;
-		let last_finalized = self.client.info().finalized_number;
+		let last_finalized = std::cmp::min(self.best_queued_number, self.client.info().finalized_number);
 		let best_queued = self.best_queued_number;
 		let client = &self.client;
 		let queue = &self.queue_blocks;
@@ -804,6 +862,28 @@ impl<B: BlockT> ChainSync<B> {
 		Either::Right(iter)
 	}
 
+	/// Get a state request, if any
+	pub fn state_request(&mut self) -> Option<(PeerId, StateRequest)> {
+		if let Some(sync) = &self.state_sync {
+			if sync.is_complete() {
+				return None;
+			}
+			if self.peers.iter().any(|(_, peer)| peer.state == PeerSyncState::DownloadingState) {
+				// Only one pending state request is allowed.
+				return None;
+			}
+			for (id, peer) in self.peers.iter_mut() {
+				if peer.state.is_available() && peer.common_number >= sync.target_block_num() {
+					trace!(target: "sync", "New StateRequest for {}", id);
+					peer.state = PeerSyncState::DownloadingState;
+					let request = sync.next_request();
+					return Some((id.clone(), request))
+				}
+			}
+		}
+		None
+	}
+
 	/// Handle a response from the remote to a block request that we made.
 	///
 	/// `request` must be the original request that triggered `response`.
@@ -848,7 +928,9 @@ impl<B: BlockT> ChainSync<B> {
 										justifications,
 										origin: block_data.origin,
 										allow_missing_state: true,
-										import_existing: false,
+										import_existing: self.import_existing,
+										skip_execution: self.skip_execution(),
+										state: None,
 									}
 								}).collect()
 						}
@@ -870,7 +952,9 @@ impl<B: BlockT> ChainSync<B> {
 									justifications,
 									origin: Some(who.clone()),
 									allow_missing_state: true,
-									import_existing: false,
+									import_existing: self.import_existing,
+									skip_execution: self.skip_execution(),
+									state: None,
 								}
 							}).collect()
 						}
@@ -963,10 +1047,11 @@ impl<B: BlockT> ChainSync<B> {
 								peer.state = PeerSyncState::Available;
 								Vec::new()
 							}
-						}
-
-						| PeerSyncState::Available
-						| PeerSyncState::DownloadingJustification(..) => Vec::new()
+						},
+						PeerSyncState::Available
+						| PeerSyncState::DownloadingJustification(..)
+						| PeerSyncState::DownloadingState
+							=> Vec::new()
 					}
 				} else {
 					// When request.is_none() this is a block announcement. Just accept blocks.
@@ -983,6 +1068,8 @@ impl<B: BlockT> ChainSync<B> {
 							origin: Some(who.clone()),
 							allow_missing_state: true,
 							import_existing: false,
+							skip_execution: true,
+							state: None,
 						}
 					}).collect()
 				}
@@ -994,6 +1081,60 @@ impl<B: BlockT> ChainSync<B> {
 		Ok(self.validate_and_queue_blocks(new_blocks))
 	}
 
+	/// Handle a response from the remote to a state request that we made.
+	///
+	/// Returns next request if any.
+	pub fn on_state_data(
+		&mut self,
+		who: &PeerId,
+		response: StateResponse,
+	) -> Result<OnStateData<B>, BadPeer> {
+		let import_result = if let Some(sync) = &mut self.state_sync {
+			debug!(
+				target: "sync",
+				"Importing state data from {} with {} keys, {} proof nodes.",
+				who,
+				response.entries.len(),
+				response.proof.len(),
+			);
+			sync.import(response)
+		} else {
+			debug!(target: "sync", "Ignored obsolete state response from {}", who);
+			return Err(BadPeer(who.clone(), rep::NOT_REQUESTED));
+		};
+
+		match import_result {
+			state::ImportResult::Import(hash, header, state) => {
+				let origin = if self.status().state != SyncState::Downloading {
+					BlockOrigin::NetworkBroadcast
+				} else {
+					BlockOrigin::NetworkInitialSync
+				};
+
+				let block = IncomingBlock {
+					hash,
+					header: Some(header),
+					body: None,
+					justifications: None,
+					origin: None,
+					allow_missing_state: true,
+					import_existing: true,
+					skip_execution: self.skip_execution(),
+					state: Some(state),
+				};
+				debug!(target: "sync", "State sync is complete. Import is queued");
+				Ok(OnStateData::Import(origin, block))
+			}
+			state::ImportResult::Continue(request) => {
+				Ok(OnStateData::Request(who.clone(), request))
+			}
+			state::ImportResult::BadResponse => {
+				debug!(target: "sync", "Bad state data received from {}", who);
+				Err(BadPeer(who.clone(), rep::BAD_BLOCK))
+			}
+		}
+	}
+
 	fn validate_and_queue_blocks(
 		&mut self,
 		mut new_blocks: Vec<IncomingBlock<B>>,
@@ -1048,7 +1189,7 @@ impl<B: BlockT> ChainSync<B> {
 			// We only request one justification at a time
 			let justification = if let Some(block) = response.blocks.into_iter().next() {
 				if hash != block.hash {
-					info!(
+					warn!(
 						target: "sync",
 						"💔 Invalid block justification provided by {}: requested: {:?} got: {:?}", who, hash, block.hash
 					);
@@ -1137,7 +1278,7 @@ impl<B: BlockT> ChainSync<B> {
 
 					if aux.bad_justification {
 						if let Some(ref peer) = who {
-							info!("💔 Sent block with bad justification to import");
+							warn!("💔 Sent block with bad justification to import");
 							output.push(Err(BadPeer(peer.clone(), rep::BAD_JUSTIFICATION)));
 						}
 					}
@@ -1145,6 +1286,17 @@ impl<B: BlockT> ChainSync<B> {
 					if let Some(peer) = who.and_then(|p| self.peers.get_mut(&p)) {
 						peer.update_common_number(number);
 					}
+					let state_sync_complete = self.state_sync.as_ref().map_or(false, |s| s.target() == hash);
+					if state_sync_complete {
+						info!(
+							target: "sync",
+							"State sync is complete ({} MiB), restarting block sync.",
+							self.state_sync.as_ref().map_or(0, |s| s.progress().size / (1024 * 1024)),
+						);
+						self.state_sync = None;
+						self.mode = SyncMode::Full;
+						output.extend(self.restart());
+					}
 				},
 				Err(BlockImportError::IncompleteHeader(who)) => {
 					if let Some(peer) = who {
@@ -1171,7 +1323,7 @@ impl<B: BlockT> ChainSync<B> {
 				},
 				Err(BlockImportError::BadBlock(who)) => {
 					if let Some(peer) = who {
-						info!(
+						warn!(
 							target: "sync",
 							"💔 Block {:?} received from peer {} has been blacklisted",
 							hash,
@@ -1189,6 +1341,7 @@ impl<B: BlockT> ChainSync<B> {
 				e @ Err(BlockImportError::UnknownParent) |
 				e @ Err(BlockImportError::Other(_)) => {
 					warn!(target: "sync", "💔 Error importing block {:?}: {:?}", hash, e);
+					self.state_sync = None;
 					output.extend(self.restart());
 				},
 				Err(BlockImportError::Cancelled) => {}
@@ -1214,6 +1367,29 @@ impl<B: BlockT> ChainSync<B> {
 			is_descendent_of(&**client, base, block)
 		});
 
+		if let SyncMode::LightState { skip_proofs } = &self.mode {
+			if self.state_sync.is_none()
+				&& !self.peers.is_empty()
+				&& self.queue_blocks.is_empty()
+			{
+				// Finalized a recent block.
+				let mut heads: Vec<_> = self.peers.iter().map(|(_, peer)| peer.best_number).collect();
+				heads.sort();
+				let median = heads[heads.len() / 2];
+				if number + STATE_SYNC_FINALITY_THRESHOLD.saturated_into() >= median {
+					if let Ok(Some(header)) = self.client.header(BlockId::hash(hash.clone())) {
+						log::debug!(
+							target: "sync",
+							"Starting state sync for #{} ({})",
+							number,
+							hash,
+						);
+						self.state_sync = Some(StateSync::new(self.client.clone(), header, *skip_proofs));
+					}
+				}
+			}
+		}
+
 		if let Err(err) = r {
 			warn!(
 				target: "sync",
@@ -1536,7 +1712,7 @@ impl<B: BlockT> ChainSync<B> {
 			return PollBlockAnnounceValidation::Nothing { is_best, who, announce }
 		}
 
-		let requires_additional_data = !self.role.is_light() || !known_parent;
+		let requires_additional_data = self.mode != SyncMode::Light || !known_parent;
 		if !requires_additional_data {
 			trace!(
 				target: "sync",
@@ -1595,6 +1771,8 @@ impl<B: BlockT> ChainSync<B> {
 					origin: block_data.origin,
 					allow_missing_state: true,
 					import_existing: false,
+					skip_execution: self.skip_execution(),
+					state: None,
 				}
 			}).collect();
 		if !blocks.is_empty() {
@@ -1611,9 +1789,9 @@ impl<B: BlockT> ChainSync<B> {
 		&'a mut self,
 	) -> impl Iterator<Item = Result<(PeerId, BlockRequest<B>), BadPeer>> + 'a {
 		self.blocks.clear();
-		let info = self.client.info();
-		self.best_queued_hash = info.best_hash;
-		self.best_queued_number = info.best_number;
+		if let Err(e) = self.reset_sync_start_point() {
+			warn!(target: "sync", "💔  Unable to restart sync. :{:?}", e);
+		}
 		self.pending_requests.set_all();
 		debug!(target:"sync", "Restarted with {} ({})", self.best_queued_number, self.best_queued_hash);
 		let old_peers = std::mem::take(&mut self.peers);
@@ -1624,7 +1802,7 @@ impl<B: BlockT> ChainSync<B> {
 			match p.state {
 				PeerSyncState::DownloadingJustification(_) => {
 					// We make sure our commmon number is at least something we have.
-					p.common_number = info.best_number;
+					p.common_number = self.best_queued_number;
 					self.peers.insert(id, p);
 					return None;
 				}
@@ -1640,6 +1818,38 @@ impl<B: BlockT> ChainSync<B> {
 		})
 	}
 
+	/// Find a block to start sync from. If we sync with state, that's the latest block we have state for.
+	fn reset_sync_start_point(&mut self) -> Result<(), ClientError> {
+		let info = self.client.info();
+		if matches!(self.mode, SyncMode::LightState {..}) && info.finalized_state.is_some() {
+			log::warn!(
+				target: "sync",
+				"Can't use fast sync mode with a partially synced database. Reverting to full sync mode."
+			);
+			self.mode = SyncMode::Full;
+		}
+		self.import_existing = false;
+		self.best_queued_hash = info.best_hash;
+		self.best_queued_number = info.best_number;
+		if self.mode == SyncMode::Full {
+			if self.client.block_status(&BlockId::hash(info.best_hash))? != BlockStatus::InChainWithState {
+				self.import_existing = true;
+				// Latest state is missing, start with the last finalized state or genesis instead.
+				if let Some((hash, number)) = info.finalized_state {
+					log::debug!(target: "sync", "Starting from finalized state #{}", number);
+					self.best_queued_hash = hash;
+					self.best_queued_number = number;
+				} else {
+					log::debug!(target: "sync", "Restarting from genesis");
+					self.best_queued_hash = Default::default();
+					self.best_queued_number = Zero::zero();
+				}
+			}
+		}
+		log::trace!(target: "sync", "Restarted sync at #{} ({:?})", self.best_queued_number, self.best_queued_hash);
+		Ok(())
+	}
+
 	/// What is the status of the block corresponding to the given hash?
 	fn block_status(&self, hash: &B::Hash) -> Result<BlockStatus, ClientError> {
 		if self.queue_blocks.contains(hash) {
@@ -1764,7 +1974,7 @@ fn peer_block_request<B: BlockT>(
 	id: &PeerId,
 	peer: &PeerSync<B>,
 	blocks: &mut BlockCollection<B>,
-	attrs: &message::BlockAttributes,
+	attrs: message::BlockAttributes,
 	max_parallel_downloads: u32,
 	finalized: NumberFor<B>,
 	best_num: NumberFor<B>,
@@ -1815,7 +2025,7 @@ fn fork_sync_request<B: BlockT>(
 	targets: &mut HashMap<B::Hash, ForkTarget<B>>,
 	best_num: NumberFor<B>,
 	finalized: NumberFor<B>,
-	attributes: &message::BlockAttributes,
+	attributes: message::BlockAttributes,
 	check_block: impl Fn(&B::Hash) -> BlockStatus,
 ) -> Option<(B::Hash, BlockRequest<B>)> {
 	targets.retain(|hash, r| {
@@ -1994,17 +2204,15 @@ mod test {
 		// internally we should process the response as the justification not being available.
 
 		let client = Arc::new(TestClientBuilder::new().build());
-		let info = client.info();
 		let block_announce_validator = Box::new(DefaultBlockAnnounceValidator);
 		let peer_id = PeerId::random();
 
 		let mut sync = ChainSync::new(
-			Roles::AUTHORITY,
+			SyncMode::Full,
 			client.clone(),
-			&info,
 			block_announce_validator,
 			1,
-		);
+		).unwrap();
 
 		let (a1_hash, a1_number) = {
 			let a1 = client.new_block(Default::default()).unwrap().build().unwrap().block;
@@ -2067,15 +2275,12 @@ mod test {
 	#[test]
 	fn restart_doesnt_affect_peers_downloading_finality_data() {
 		let mut client = Arc::new(TestClientBuilder::new().build());
-		let info = client.info();
-
 		let mut sync = ChainSync::new(
-			Roles::AUTHORITY,
+			SyncMode::Full,
 			client.clone(),
-			&info,
 			Box::new(DefaultBlockAnnounceValidator),
 			1,
-		);
+		).unwrap();
 
 		let peer_id1 = PeerId::random();
 		let peer_id2 = PeerId::random();
@@ -2242,15 +2447,13 @@ mod test {
 		sp_tracing::try_init_simple();
 
 		let mut client = Arc::new(TestClientBuilder::new().build());
-		let info = client.info();
 
 		let mut sync = ChainSync::new(
-			Roles::AUTHORITY,
+			SyncMode::Full,
 			client.clone(),
-			&info,
 			Box::new(DefaultBlockAnnounceValidator),
 			5,
-		);
+		).unwrap();
 
 		let peer_id1 = PeerId::random();
 		let peer_id2 = PeerId::random();
@@ -2359,12 +2562,11 @@ mod test {
 		let info = client.info();
 
 		let mut sync = ChainSync::new(
-			Roles::AUTHORITY,
+			SyncMode::Full,
 			client.clone(),
-			&info,
 			Box::new(DefaultBlockAnnounceValidator),
 			5,
-		);
+		).unwrap();
 
 		let peer_id1 = PeerId::random();
 		let peer_id2 = PeerId::random();
@@ -2481,12 +2683,11 @@ mod test {
 		let info = client.info();
 
 		let mut sync = ChainSync::new(
-			Roles::AUTHORITY,
+			SyncMode::Full,
 			client.clone(),
-			&info,
 			Box::new(DefaultBlockAnnounceValidator),
 			5,
-		);
+		).unwrap();
 
 		let finalized_block = blocks[MAX_BLOCKS_TO_LOOK_BACKWARDS as usize * 2 - 1].clone();
 		let just = (*b"TEST", Vec::new());
@@ -2592,15 +2793,12 @@ mod test {
 			.map(|_| build_block(&mut client, None, false))
 			.collect::<Vec<_>>();
 
-		let info = client.info();
-
 		let mut sync = ChainSync::new(
-			Roles::AUTHORITY,
+			SyncMode::Full,
 			client.clone(),
-			&info,
 			Box::new(DefaultBlockAnnounceValidator),
 			1,
-		);
+		).unwrap();
 
 		let peer_id1 = PeerId::random();
 		let common_block = blocks[1].clone();

substrate/client/network/src/protocol/sync/state.rs

@@ -0,0 +1,187 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2017-2021 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use std::sync::Arc;
+use codec::{Encode, Decode};
+use sp_runtime::traits::{Block as BlockT, Header, NumberFor};
+use sc_client_api::StorageProof;
+use crate::schema::v1::{StateRequest, StateResponse, StateEntry};
+use crate::chain::{Client, ImportedState};
+use super::StateDownloadProgress;
+
+/// State sync support.
+
+/// State sync state machine. Accumulates partial state data until it
+/// is ready to be imported.
+pub struct StateSync<B: BlockT> {
+	target_block: B::Hash,
+	target_header: B::Header,
+	target_root: B::Hash,
+	last_key: Vec<u8>,
+	state: Vec<(Vec<u8>, Vec<u8>)>,
+	complete: bool,
+	client: Arc<dyn Client<B>>,
+	imported_bytes: u64,
+	skip_proof: bool,
+}
+
+/// Import state chunk result.
+pub enum ImportResult<B: BlockT> {
+	/// State is complete and ready for import.
+	Import(B::Hash, B::Header, ImportedState<B>),
+	/// Continue dowloading.
+	Continue(StateRequest),
+	/// Bad state chunk.
+	BadResponse,
+}
+
+impl<B: BlockT> StateSync<B> {
+	///  Create a new instance.
+	pub fn new(client: Arc<dyn Client<B>>, target: B::Header, skip_proof: bool) -> Self {
+		StateSync {
+			client,
+			target_block: target.hash(),
+			target_root: target.state_root().clone(),
+			target_header: target,
+			last_key: Vec::default(),
+			state: Vec::default(),
+			complete: false,
+			imported_bytes: 0,
+			skip_proof,
+		}
+	}
+
+	///  Validate and import a state reponse.
+	pub fn import(&mut self, response: StateResponse) -> ImportResult<B> {
+		if response.entries.is_empty() && response.proof.is_empty() && !response.complete {
+			log::debug!(
+				target: "sync",
+				"Bad state response",
+			);
+			return ImportResult::BadResponse;
+		}
+		if !self.skip_proof && response.proof.is_empty() {
+			log::debug!(
+				target: "sync",
+				"Missing proof",
+			);
+			return ImportResult::BadResponse;
+		}
+		let complete = if !self.skip_proof {
+			log::debug!(
+				target: "sync",
+				"Importing state from {} trie nodes",
+				response.proof.len(),
+			);
+			let proof_size = response.proof.len() as u64;
+			let proof = match StorageProof::decode(&mut response.proof.as_ref()) {
+				Ok(proof) => proof,
+				Err(e) => {
+					log::debug!(target: "sync", "Error decoding proof: {:?}", e);
+					return ImportResult::BadResponse;
+				}
+			};
+			let (values, complete) = match self.client.verify_range_proof(
+				self.target_root,
+				proof,
+				&self.last_key
+			) {
+				Err(e) => {
+					log::debug!(
+						target: "sync",
+						"StateResponse failed proof verification: {:?}",
+						e,
+					);
+					return ImportResult::BadResponse;
+				},
+				Ok(values) => values,
+			};
+			log::debug!(target: "sync", "Imported with {} keys", values.len());
+
+			if let Some(last) = values.last().map(|(k, _)| k) {
+				self.last_key = last.clone();
+			}
+
+			for (key, value) in values {
+				self.imported_bytes += key.len() as u64;
+				self.state.push((key, value))
+			};
+			self.imported_bytes += proof_size;
+			complete
+		} else {
+			log::debug!(
+				target: "sync",
+				"Importing state from {:?} to {:?}",
+				response.entries.last().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
+				response.entries.first().map(|e| sp_core::hexdisplay::HexDisplay::from(&e.key)),
+			);
+
+			if let Some(e) = response.entries.last() {
+				self.last_key = e.key.clone();
+			}
+			for StateEntry { key, value } in response.entries {
+				self.imported_bytes += (key.len() + value.len()) as u64;
+				self.state.push((key, value))
+			}
+			response.complete
+		};
+		if complete {
+			self.complete = true;
+			ImportResult::Import(self.target_block.clone(), self.target_header.clone(), ImportedState {
+				block: self.target_block.clone(),
+				state: std::mem::take(&mut self.state)
+			})
+		} else {
+			ImportResult::Continue(self.next_request())
+		}
+	}
+
+	/// Produce next state request.
+	pub fn next_request(&self) -> StateRequest {
+		StateRequest {
+			block: self.target_block.encode(),
+			start: self.last_key.clone(),
+			no_proof: self.skip_proof,
+		}
+	}
+
+	/// Check if the state is complete.
+	pub fn is_complete(&self) -> bool {
+		self.complete
+	}
+
+	/// Returns target block number.
+	pub fn target_block_num(&self) -> NumberFor<B> {
+		self.target_header.number().clone()
+	}
+
+	/// Returns target block hash.
+	pub fn target(&self) -> B::Hash {
+		self.target_block.clone()
+	}
+
+	/// Returns state sync estimated progress.
+	pub fn progress(&self) -> StateDownloadProgress {
+		let percent_done = (*self.last_key.get(0).unwrap_or(&0u8) as u32) * 100 / 256;
+		StateDownloadProgress {
+			percentage: percent_done,
+			size: self.imported_bytes,
+		}
+	}
+}
+