From df3b40b63754c4dfe560ee0a2b9155bb6667c15b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Silva?= <andre.beat@gmail.com>
Date: Wed, 27 Jun 2018 19:04:59 +0100
Subject: [PATCH] aura: handle immediate transitions

---
 ethcore/src/engines/authority_round/mod.rs | 26 ++++++++++++++++------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/ethcore/src/engines/authority_round/mod.rs b/ethcore/src/engines/authority_round/mod.rs
index 9b9eb80f3de..fa8822a89ac 100644
--- a/ethcore/src/engines/authority_round/mod.rs
+++ b/ethcore/src/engines/authority_round/mod.rs
@@ -916,17 +916,17 @@ impl Engine<EthereumMachine> for AuthorityRound {
 		}
 
 		let active_set;
-		let validators = if self.immediate_transitions {
-			&*self.validators
+		let (validators, set_number) = if self.immediate_transitions {
+			(&*self.validators, header.number())
 		} else {
 			match self.epoch_set(header) {
 				Err(err) => {
 					warn!(target: "engine", "Unable to generate seal: {}", err);
 					return Seal::None;
 				},
-				Ok((set, _)) => {
+				Ok((set, set_number)) => {
 					active_set = set;
-					&active_set
+					(&active_set as &_, set_number)
 				},
 			}
 		};
@@ -1095,7 +1095,12 @@ impl Engine<EthereumMachine> for AuthorityRound {
 
 		match verify_timestamp(&self.step.inner, header_step(header, self.empty_steps_transition)?) {
 			Err(BlockError::InvalidSeal) => {
-				let (_, set_number) = self.epoch_set(header)?;
+				let set_number = if self.immediate_transitions {
+					header.number()
+				} else {
+					self.epoch_set(header)?.1
+				};
+
 				self.validators.report_benign(header.author(), set_number, header.number());
 				Err(BlockError::InvalidSeal.into())
 			}
@@ -1109,7 +1114,14 @@ impl Engine<EthereumMachine> for AuthorityRound {
 		let step = header_step(header, self.empty_steps_transition)?;
 		let parent_step = header_step(parent, self.empty_steps_transition)?;
 
-		let (validators, set_number) = self.epoch_set(header)?;
+		let active_set;
+		let (validators, set_number) = if self.immediate_transitions {
+			(&*self.validators, header.number())
+		} else {
+			let (set, set_number) = self.epoch_set(header)?;
+			active_set = set;
+			(&active_set as &_, set_number)
+		};
 
 		// Ensure header is from the step after parent.
 		if step == parent_step
@@ -1136,7 +1148,7 @@ impl Engine<EthereumMachine> for AuthorityRound {
 							format!("empty step proof for invalid parent hash: {:?}", empty_step.parent_hash)))?;
 					}
 
-					if !empty_step.verify(&validators).unwrap_or(false) {
+					if !empty_step.verify(validators).unwrap_or(false) {
 						Err(EngineError::InsufficientProof(
 							format!("invalid empty step proof: {:?}", empty_step)))?;
 					}