From 5f37c936595426b367a6f04dbb38c1f3593b8244 Mon Sep 17 00:00:00 2001
From: Robert Habermeier <rphmeier@gmail.com>
Date: Fri, 9 Dec 2016 15:04:54 +0100
Subject: [PATCH] les: use negotiated protocol version

---
 ethcore/light/src/net/context.rs |  7 +++++++
 ethcore/light/src/net/error.rs   |  8 ++++++++
 ethcore/light/src/net/mod.rs     | 22 +++++++++++++++++-----
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/ethcore/light/src/net/context.rs b/ethcore/light/src/net/context.rs
index f2d5ab907..c05e69b0f 100644
--- a/ethcore/light/src/net/context.rs
+++ b/ethcore/light/src/net/context.rs
@@ -38,6 +38,9 @@ pub trait IoContext {
 
     /// Disable a peer -- this is a disconnect + a time-out.
     fn disable_peer(&self, peer: PeerId);
+
+    /// Get a peer's protocol version.
+    fn protocol_version(&self, peer: PeerId) -> Option<u8>;
 }
 
 impl<'a> IoContext for NetworkContext<'a> {
@@ -60,6 +63,10 @@ impl<'a> IoContext for NetworkContext<'a> {
     fn disable_peer(&self, peer: PeerId) {
         NetworkContext::disable_peer(self, peer);
     }
+
+    fn protocol_version(&self, peer: PeerId) -> Option<u8> {
+        self.protocol_version(self.subprotocol_name(), peer)
+    }
 }
 
 /// Context for a protocol event.
diff --git a/ethcore/light/src/net/error.rs b/ethcore/light/src/net/error.rs
index c9ba85a42..01a15928b 100644
--- a/ethcore/light/src/net/error.rs
+++ b/ethcore/light/src/net/error.rs
@@ -58,6 +58,10 @@ pub enum Error {
 	UnsolicitedResponse,
 	/// Not a server.
 	NotServer,
+	/// Unsupported protocol version.
+	UnsupportedProtocolVersion(u8),
+	/// Bad protocol version.
+	BadProtocolVersion,
 }
 
 impl Error {
@@ -73,6 +77,8 @@ impl Error {
 			Error::UnknownPeer => Punishment::Disconnect,
 			Error::UnsolicitedResponse => Punishment::Disable,
 			Error::NotServer => Punishment::Disable,
+			Error::UnsupportedProtocolVersion(_) => Punishment::Disable,
+			Error::BadProtocolVersion => Punishment::Disable,
 		}
 	}
 }
@@ -101,6 +107,8 @@ impl fmt::Display for Error {
 			Error::UnknownPeer => write!(f, "Unknown peer"),
 			Error::UnsolicitedResponse => write!(f, "Peer provided unsolicited data"),
 			Error::NotServer => write!(f, "Peer not a server."),
+			Error::UnsupportedProtocolVersion(pv) => write!(f, "Unsupported protocol version: {}", pv),	
+			Error::BadProtocolVersion => write!(f, "Bad protocol version in handshake"),
 		}
 	}
 }
\ No newline at end of file
diff --git a/ethcore/light/src/net/mod.rs b/ethcore/light/src/net/mod.rs
index 1eb89a92f..481740a48 100644
--- a/ethcore/light/src/net/mod.rs
+++ b/ethcore/light/src/net/mod.rs
@@ -23,7 +23,7 @@ use ethcore::transaction::SignedTransaction;
 use ethcore::receipt::Receipt;
 
 use io::TimerToken;
-use network::{NetworkProtocolHandler, NetworkContext, NetworkError, PeerId};
+use network::{NetworkProtocolHandler, NetworkContext, PeerId};
 use rlp::{RlpStream, Stream, UntrustedRlp, View};
 use util::hash::H256;
 use util::{Bytes, Mutex, RwLock, U256};
@@ -111,6 +111,7 @@ pub struct ReqId(usize);
 struct PendingPeer {
 	sent_head: H256,
 	last_update: SteadyTime,
+	proto_version: u8,
 }
 
 // data about each peer.
@@ -121,6 +122,7 @@ struct Peer {
 	remote_flow: Option<(Buffer, FlowParams)>,
 	sent_head: H256, // last head we've given them.
 	last_update: SteadyTime,
+	proto_version: u8,
 }
 
 impl Peer {
@@ -507,17 +509,21 @@ impl LightProtocol {
 	}
 
 	// send status to a peer.
-	fn send_status(&self, peer: PeerId, io: &IoContext) -> Result<PendingPeer, NetworkError> {
-		let chain_info = self.provider.chain_info();
+	fn send_status(&self, peer: PeerId, io: &IoContext) -> Result<PendingPeer, Error> {
+		let proto_version = try!(io.protocol_version(peer).ok_or(Error::WrongNetwork));
 
-		// TODO: could update capabilities here.
+		if PROTOCOL_VERSIONS.iter().find(|x| **x == proto_version).is_none() {
+			return Err(Error::UnsupportedProtocolVersion(proto_version));
+		}
+		
+		let chain_info = self.provider.chain_info();
 
 		let status = Status {
 			head_td: chain_info.total_difficulty,
 			head_hash: chain_info.best_block_hash,
 			head_num: chain_info.best_block_number,
 			genesis_hash: chain_info.genesis_hash,
-			protocol_version: MAX_PROTOCOL_VERSION as u32,
+			protocol_version: proto_version as u32, // match peer proto version
 			network_id: self.network_id,
 			last_head: None,
 		};
@@ -530,6 +536,7 @@ impl LightProtocol {
 		Ok(PendingPeer {
 			sent_head: chain_info.best_block_hash,
 			last_update: SteadyTime::now(),
+			proto_version: proto_version,			
 		})
 	}
 
@@ -550,6 +557,10 @@ impl LightProtocol {
 			return Err(Error::WrongNetwork);
 		}
 
+		if Some(status.protocol_version as u8) != io.protocol_version(*peer) {
+			return Err(Error::BadProtocolVersion);
+		}
+
 		let remote_flow = flow_params.map(|params| (params.create_buffer(), params));
 
 		self.peers.write().insert(*peer, Mutex::new(Peer {
@@ -559,6 +570,7 @@ impl LightProtocol {
 			remote_flow: remote_flow,
 			sent_head: pending.sent_head,
 			last_update: pending.last_update,
+			proto_version: pending.proto_version,
 		}));
 
 		for handler in &self.handlers {