Add incoming connection to node table

This commit is contained in:
arkpar 2016-02-15 19:54:27 +01:00
parent 0e1e80477a
commit 4b9c7f7517
4 changed files with 31 additions and 1 deletions

View File

@ -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));

View File

@ -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() {

View File

@ -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);
}

View File

@ -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)) {