From bdfd5405c044e7c7ed5c0a9754697517dc70bcce Mon Sep 17 00:00:00 2001
From: Ignotus Peverell <igno.peverell@protonmail.com>
Date: Thu, 12 Jul 2018 17:06:52 +0100
Subject: [PATCH] Leverage sync state to accept txhashset (#1251)

Fixes #1246
---
 p2p/src/peer.rs                | 4 ++++
 p2p/src/peers.rs               | 4 ++++
 p2p/src/protocol.rs            | 3 +++
 p2p/src/serv.rs                | 4 ++++
 p2p/src/types.rs               | 6 ++++++
 servers/src/common/adapters.rs | 7 +++++--
 6 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/p2p/src/peer.rs b/p2p/src/peer.rs
index 4de79322..dad6b899 100644
--- a/p2p/src/peer.rs
+++ b/p2p/src/peer.rs
@@ -428,6 +428,10 @@ impl ChainAdapter for TrackingAdapter {
 		self.adapter.txhashset_read(h)
 	}
 
+	fn txhashset_receive_ready(&self) -> bool {
+		self.adapter.txhashset_receive_ready()
+	}
+
 	fn txhashset_write(
 		&self,
 		h: Hash,
diff --git a/p2p/src/peers.rs b/p2p/src/peers.rs
index fa1c730a..f8d63ddc 100644
--- a/p2p/src/peers.rs
+++ b/p2p/src/peers.rs
@@ -611,6 +611,10 @@ impl ChainAdapter for Peers {
 		self.adapter.txhashset_read(h)
 	}
 
+	fn txhashset_receive_ready(&self) -> bool {
+		self.adapter.txhashset_receive_ready()
+	}
+
 	fn txhashset_write(
 		&self,
 		h: Hash,
diff --git a/p2p/src/protocol.rs b/p2p/src/protocol.rs
index 7099f63e..b5e3aa0a 100644
--- a/p2p/src/protocol.rs
+++ b/p2p/src/protocol.rs
@@ -233,6 +233,9 @@ impl MessageHandler for Protocol {
 					sm_arch.hash,
 					sm_arch.height,
 				);
+				if !self.adapter.txhashset_receive_ready() {
+					return Err(Error::BadMessage);
+				}
 
 				let mut tmp = env::temp_dir();
 				tmp.push("txhashset.zip");
diff --git a/p2p/src/serv.rs b/p2p/src/serv.rs
index ca391454..57ecb0db 100644
--- a/p2p/src/serv.rs
+++ b/p2p/src/serv.rs
@@ -250,6 +250,10 @@ impl ChainAdapter for DummyAdapter {
 		unimplemented!()
 	}
 
+	fn txhashset_receive_ready(&self) -> bool {
+		false
+	}
+
 	fn txhashset_write(
 		&self,
 		_h: Hash,
diff --git a/p2p/src/types.rs b/p2p/src/types.rs
index 5d9f7a96..e66aeaa0 100644
--- a/p2p/src/types.rs
+++ b/p2p/src/types.rs
@@ -256,6 +256,12 @@ pub trait ChainAdapter: Sync + Send {
 	/// at the provided block hash.
 	fn txhashset_read(&self, h: Hash) -> Option<TxHashSetRead>;
 
+	/// Whether the node is ready to accept a new txhashset. If this isn't the
+	/// case, the archive is provided without being requested and likely an
+	/// attack attempt. This should be checked *before* downloading the whole
+	/// state data.
+	fn txhashset_receive_ready(&self) -> bool;
+
 	/// Writes a reading view on a txhashset state that's been provided to us.
 	/// If we're willing to accept that new state, the data stream will be
 	/// read as a zip file, unzipped and the resulting state files should be
diff --git a/servers/src/common/adapters.rs b/servers/src/common/adapters.rs
index c4580d93..6713aec1 100644
--- a/servers/src/common/adapters.rs
+++ b/servers/src/common/adapters.rs
@@ -24,7 +24,7 @@ use std::thread;
 use std::time::Instant;
 
 use chain::{self, ChainAdapter, Options, Tip};
-use common::types::{ChainValidationMode, ServerConfig, SyncState};
+use common::types::{ChainValidationMode, ServerConfig, SyncState, SyncStatus};
 use core::core;
 use core::core::block::BlockHeader;
 use core::core::hash::{Hash, Hashed};
@@ -320,6 +320,10 @@ impl p2p::ChainAdapter for NetToChainAdapter {
 		}
 	}
 
+	fn txhashset_receive_ready(&self) -> bool {
+		self.sync_state.status() == SyncStatus::TxHashsetDownload
+	}
+
 	/// Writes a reading view on a txhashset state that's been provided to us.
 	/// If we're willing to accept that new state, the data stream will be
 	/// read as a zip file, unzipped and the resulting state files should be
@@ -330,7 +334,6 @@ impl p2p::ChainAdapter for NetToChainAdapter {
 		txhashset_data: File,
 		_peer_addr: SocketAddr,
 	) -> bool {
-		// TODO check whether we should accept any txhashset now
 		if let Err(e) =
 			w(&self.chain).txhashset_write(h, txhashset_data, self.sync_state.as_ref())
 		{