From 568a18d8bda316779010fd628b7391aa8c189046 Mon Sep 17 00:00:00 2001
From: Arkadiy Paronyan <arkady.paronyan@gmail.com>
Date: Mon, 31 Oct 2016 12:54:50 +0100
Subject: [PATCH] Prevent connecting to self (#2997)

---
 util/network/src/discovery.rs | 18 +++++++++++-------
 util/network/src/host.rs      |  6 +++---
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/util/network/src/discovery.rs b/util/network/src/discovery.rs
index 1f506ee89..ef5573d4d 100644
--- a/util/network/src/discovery.rs
+++ b/util/network/src/discovery.rs
@@ -130,7 +130,7 @@ impl Discovery {
 
 	/// Add a new node to discovery table. Pings the node.
 	pub fn add_node(&mut self, e: NodeEntry) {
-		if e.endpoint.is_allowed(self.allow_ips) {
+		if self.is_allowed(&e) {
 			let endpoint = e.endpoint.clone();
 			self.update_node(e);
 			self.ping(&endpoint);
@@ -146,7 +146,7 @@ impl Discovery {
 	/// Add a list of known nodes to the table.
 	pub fn init_node_list(&mut self, mut nodes: Vec<NodeEntry>) {
 		for n in nodes.drain(..) {
-			if n.endpoint.is_allowed(self.allow_ips) {
+			if self.is_allowed(&n) {
 				self.update_node(n);
 			}
 		}
@@ -399,6 +399,10 @@ impl Discovery {
 		Ok(())
 	}
 
+	fn is_allowed(&self, entry: &NodeEntry) -> bool {
+		entry.endpoint.is_allowed(self.allow_ips) && entry.id != self.id
+	}
+
 	fn on_ping(&mut self, rlp: &UntrustedRlp, node: &NodeId, from: &SocketAddr) -> Result<Option<TableUpdates>, NetworkError> {
 		trace!(target: "discovery", "Got Ping from {:?}", &from);
 		let source = try!(NodeEndpoint::from_rlp(&try!(rlp.at(1))));
@@ -409,7 +413,7 @@ impl Discovery {
 		let entry = NodeEntry { id: node.clone(), endpoint: source.clone() };
 		if !entry.endpoint.is_valid() {
 			debug!(target: "discovery", "Got bad address: {:?}", entry);
-		} else if !entry.endpoint.is_allowed(self.allow_ips) {
+		} else if !self.is_allowed(&entry) {
 			debug!(target: "discovery", "Address not allowed: {:?}", entry);
 		} else {
 			self.update_node(entry.clone());
@@ -484,15 +488,15 @@ impl Discovery {
 				debug!(target: "discovery", "Bad address: {:?}", endpoint);
 				continue;
 			}
-			if !endpoint.is_allowed(self.allow_ips) {
-				debug!(target: "discovery", "Address not allowed: {:?}", endpoint);
-				continue;
-			}
 			let node_id: NodeId = try!(r.val_at(3));
 			if node_id == self.id {
 				continue;
 			}
 			let entry = NodeEntry { id: node_id.clone(), endpoint: endpoint };
+			if !self.is_allowed(&entry) {
+				debug!(target: "discovery", "Address not allowed: {:?}", entry);
+				continue;
+			}
 			added.insert(node_id, entry.clone());
 			self.ping(&entry.endpoint);
 			self.update_node(entry);
diff --git a/util/network/src/host.rs b/util/network/src/host.rs
index b5e4ed914..b693d0abd 100644
--- a/util/network/src/host.rs
+++ b/util/network/src/host.rs
@@ -634,14 +634,14 @@ impl Host {
 	}
 
 	fn connect_peers(&self, io: &IoContext<NetworkIoMessage>) {
-		let (min_peers, mut pin, max_handshakes, allow_ips) = {
+		let (min_peers, mut pin, max_handshakes, allow_ips, self_id) = {
 			let info = self.info.read();
 			if info.capabilities.is_empty() {
 				return;
 			}
 			let config = &info.config;
 
-			(config.min_peers, config.non_reserved_mode == NonReservedPeerMode::Deny, config.max_handshakes as usize, config.allow_ips)
+			(config.min_peers, config.non_reserved_mode == NonReservedPeerMode::Deny, config.max_handshakes as usize, config.allow_ips, info.id().clone())
 		};
 
 		let session_count = self.session_count();
@@ -672,7 +672,7 @@ impl Host {
 
 		let max_handshakes_per_round = max_handshakes / 2;
 		let mut started: usize = 0;
-		for id in nodes.filter(|ref id| !self.have_session(id) && !self.connecting_to(id))
+		for id in nodes.filter(|id| !self.have_session(id) && !self.connecting_to(id) && *id != self_id)
 			.take(min(max_handshakes_per_round, max_handshakes - handshake_count)) {
 			self.connect_peer(&id, io);
 			started += 1;