From a27268a9f4e6fe04bbc49dc4909e9b78b98954ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Drwi=C4=99ga?= <tomusdrw@users.noreply.github.com>
Date: Tue, 28 Aug 2018 11:55:39 +0200
Subject: [PATCH] Make sure to ban invalid transactions. (#615)

---
 substrate/extrinsic-pool/src/pool.rs    | 25 +++++++++++++++++++++
 substrate/extrinsic-pool/src/rotator.rs | 29 +++++++++++++++----------
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/substrate/extrinsic-pool/src/pool.rs b/substrate/extrinsic-pool/src/pool.rs
index 100fa700e4c5f..ea9ccfebd02cd 100644
--- a/substrate/extrinsic-pool/src/pool.rs
+++ b/substrate/extrinsic-pool/src/pool.rs
@@ -262,9 +262,17 @@ impl<B: ChainApi> Pool<B> {
 	pub fn remove(&self, hashes: &[B::Hash], is_valid: bool) -> Vec<Option<Arc<VerifiedFor<B>>>> {
 		let mut pool = self.pool.write();
 		let mut results = Vec::with_capacity(hashes.len());
+
+		// temporarily ban invalid transactions
+		if !is_valid {
+			debug!(target: "transaction-pool", "Banning invalid transactions: {:?}", hashes);
+			self.rotator.ban(&time::Instant::now(), hashes);
+		}
+
 		for hash in hashes {
 			results.push(pool.remove(hash, is_valid));
 		}
+
 		results
 	}
 
@@ -578,4 +586,21 @@ pub mod tests {
 		let pending: Vec<_> = pool.cull_and_get_pending(&BlockId::number(0), |p| p.map(|a| (*a.sender(), a.original.transfer.nonce)).collect()).unwrap();
 		assert_eq!(pending, vec![(Alice.to_raw_public().into(), 209), (Alice.to_raw_public().into(), 210)]);
 	}
+
+	#[test]
+	fn should_ban_invalid_transactions() {
+		let pool = pool();
+		let uxt = uxt(Alice, 209);
+		let hash = *pool.submit_one(&BlockId::number(0), uxt.clone()).unwrap().hash();
+		pool.remove(&[hash], true);
+		pool.submit_one(&BlockId::number(0), uxt.clone()).unwrap();
+
+		// when
+		pool.remove(&[hash], false);
+		let pending: Vec<AccountId> = pool.cull_and_get_pending(&BlockId::number(0), |p| p.map(|a| *a.sender()).collect()).unwrap();
+		assert_eq!(pending, vec![]);
+
+		// then
+		pool.submit_one(&BlockId::number(0), uxt.clone()).unwrap_err();
+	}
 }
diff --git a/substrate/extrinsic-pool/src/rotator.rs b/substrate/extrinsic-pool/src/rotator.rs
index cacf079879506..93acce9e9130c 100644
--- a/substrate/extrinsic-pool/src/rotator.rs
+++ b/substrate/extrinsic-pool/src/rotator.rs
@@ -58,6 +58,23 @@ impl<Hash: hash::Hash + Eq + Clone> PoolRotator<Hash> {
 		self.banned_until.read().contains_key(hash)
 	}
 
+	/// Bans given set of hashes.
+	pub fn ban(&self, now: &Instant, hashes: &[Hash]) {
+		let mut banned = self.banned_until.write();
+
+		for hash in hashes {
+			banned.insert(hash.clone(), *now + self.ban_time);
+		}
+
+		if banned.len() > 2 * EXPECTED_SIZE {
+			while banned.len() > EXPECTED_SIZE {
+				if let Some(key) = banned.keys().next().cloned() {
+					banned.remove(&key);
+				}
+			}
+		}
+	}
+
 	/// Bans extrinsic if it's stale.
 	///
 	/// Returns `true` if extrinsic is stale and got banned.
@@ -69,17 +86,7 @@ impl<Hash: hash::Hash + Eq + Clone> PoolRotator<Hash> {
 			return false;
 		}
 
-		let mut banned = self.banned_until.write();
-		banned.insert(xt.verified.hash().clone(), *now + self.ban_time);
-
-		if banned.len() > 2 * EXPECTED_SIZE {
-			while banned.len() > EXPECTED_SIZE {
-				if let Some(key) = banned.keys().next().cloned() {
-					banned.remove(&key);
-				}
-			}
-		}
-
+		self.ban(now, &[xt.verified.hash().clone()]);
 		true
 	}