From 4b9c7f7517eb79e95d7c9326cca734508614b1e0 Mon Sep 17 00:00:00 2001
From: arkpar <arkady.paronyan@gmail.com>
Date: Mon, 15 Feb 2016 19:54:27 +0100
Subject: [PATCH] Add incoming connection to node table

---
 util/src/network/connection.rs | 11 +++++++++++
 util/src/network/host.rs       |  9 +++++++++
 util/src/network/node_table.rs |  5 ++++-
 util/src/network/session.rs    |  7 +++++++
 4 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/util/src/network/connection.rs b/util/src/network/connection.rs
index 242e8935e..9e9304ca6 100644
--- a/util/src/network/connection.rs
+++ b/util/src/network/connection.rs
@@ -16,6 +16,7 @@
 
 use std::sync::Arc;
 use std::collections::VecDeque;
+use std::net::SocketAddr;
 use mio::{Handler, Token, EventSet, EventLoop, PollOpt, TryRead, TryWrite};
 use mio::tcp::*;
 use hash::*;
@@ -169,6 +170,11 @@ impl Connection {
 		self.token = token;
 	}
 
+	/// Get remote peer address
+	pub fn remote_addr(&self) -> io::Result<SocketAddr> {
+		self.socket.peer_addr()
+	}
+
 	/// Register this connection with the IO event loop.
 	pub fn register_socket<Host: Handler>(&self, reg: Token, event_loop: &mut EventLoop<Host>) -> io::Result<()> {
 		trace!(target: "net", "connection register; token={:?}", reg);
@@ -253,6 +259,11 @@ impl EncryptedConnection {
 		self.connection.set_token(token);
 	}
 
+	/// Get remote peer address
+	pub fn remote_addr(&self) -> io::Result<SocketAddr> {
+		self.connection.remote_addr()
+	}
+
 	/// Create an encrypted connection out of the handshake. Consumes a handshake object.
 	pub fn new(mut handshake: Handshake) -> Result<EncryptedConnection, UtilError> {
 		let shared = try!(crypto::ecdh::agree(handshake.ecdhe.secret(), &handshake.remote_public));
diff --git a/util/src/network/host.rs b/util/src/network/host.rs
index fb3bcee62..5f0bf19b9 100644
--- a/util/src/network/host.rs
+++ b/util/src/network/host.rs
@@ -663,6 +663,7 @@ impl<Message> Host<Message> where Message: Send + Sync + Clone {
 			}
 		}
 		let h = Arc::try_unwrap(h).ok().unwrap().into_inner().unwrap();
+		let originated = h.originated;
 		let mut session = match Session::new(h, &self.info.read().unwrap()) {
 			Ok(s) => s,
 			Err(e) => {
@@ -674,6 +675,14 @@ impl<Message> Host<Message> where Message: Send + Sync + Clone {
 			session.set_token(session_token);
 			io.update_registration(session_token).expect("Error updating session registration");
 			self.stats.inc_sessions();
+			if !originated {
+				// Add it no node table
+				if let Ok(address) = session.remote_addr() {
+					let entry = NodeEntry { id: session.id().clone(), endpoint: NodeEndpoint { address: address, udp_port: address.port() } };
+					self.nodes.write().unwrap().add_node(Node::new(entry.id.clone(), entry.endpoint.clone()));
+					self.discovery.lock().unwrap().add_node(entry);
+				}
+			}
 			Arc::new(Mutex::new(session))
 		});
 		if result.is_none() {
diff --git a/util/src/network/node_table.rs b/util/src/network/node_table.rs
index e5c47cafb..40cc14743 100644
--- a/util/src/network/node_table.rs
+++ b/util/src/network/node_table.rs
@@ -200,7 +200,10 @@ impl NodeTable {
 	}
 
 	/// Add a node to table
-	pub fn add_node(&mut self, node: Node) {
+	pub fn add_node(&mut self, mut node: Node) {
+		// preserve failure counter
+		let failures = self.nodes.get(&node.id).map_or(0, |n| n.failures);
+		node.failures = failures;
 		self.nodes.insert(node.id.clone(), node);
 	}
 
diff --git a/util/src/network/session.rs b/util/src/network/session.rs
index 5cda00567..f08fef385 100644
--- a/util/src/network/session.rs
+++ b/util/src/network/session.rs
@@ -14,6 +14,8 @@
 // You should have received a copy of the GNU General Public License
 // along with Parity.  If not, see <http://www.gnu.org/licenses/>.
 
+use std::net::SocketAddr;
+use std::io;
 use mio::*;
 use hash::*;
 use rlp::*;
@@ -144,6 +146,11 @@ impl Session {
 		self.connection.set_token(token);
 	}
 
+	/// Get remote peer address
+	pub fn remote_addr(&self) -> io::Result<SocketAddr> {
+		self.connection.remote_addr()
+	}
+
 	/// Readable IO handler. Returns packet data if available.
 	pub fn readable<Message>(&mut self, io: &IoContext<Message>, host: &HostInfo) -> Result<SessionData, UtilError>  where Message: Send + Sync + Clone {
 		match try!(self.connection.readable(io)) {