use std::collections::HashMap; use std::net::{IpAddr, SocketAddr}; use std::sync::{Arc, RwLock}; use log::{debug, info, error}; use arc_swap::ArcSwapOption; use async_trait::async_trait; use serde::{Deserialize, Serialize}; use sodiumoxide::crypto::auth; use sodiumoxide::crypto::sign::ed25519; use tokio::net::{TcpListener, TcpStream}; use crate::client::*; use crate::server::*; use crate::endpoint::*; use crate::error::*; use crate::proto::*; use crate::util::*; #[derive(Serialize, Deserialize)] pub(crate) struct HelloMessage { pub server_addr: Option, pub server_port: u16, } impl Message for HelloMessage { type Response = (); } type OnConnectHandler = Box; type OnDisconnectHandler = Box; /// NetApp is the main class that handles incoming and outgoing connections. /// /// NetApp can be used in a stand-alone fashion or together with a peering strategy. /// If using it alone, you will want to set `on_connect` and `on_disconnect` events /// in order to manage information about the current peer list. /// /// It is generally not necessary to use NetApp stand-alone, as the provided full mesh /// and RPS peering strategies take care of the most common use cases. pub struct NetApp { listen_params: ArcSwapOption, /// Network secret key pub netid: auth::Key, /// Our peer ID pub id: NodeID, /// Private key associated with our peer ID pub privkey: ed25519::SecretKey, pub(crate) server_conns: RwLock>>, pub(crate) client_conns: RwLock>>, pub(crate) endpoints: RwLock>, hello_endpoint: ArcSwapOption>, on_connected_handler: ArcSwapOption, on_disconnected_handler: ArcSwapOption, } struct ListenParams { listen_addr: SocketAddr, public_addr: Option, } impl NetApp { /// Creates a new instance of NetApp, which can serve either as a full p2p node, /// or just as a passive client. To upgrade to a full p2p node, spawn a listener /// using `.listen()` /// /// Our Peer ID is the public key associated to the secret key given here. pub fn new(netid: auth::Key, privkey: ed25519::SecretKey) -> Arc { let id = privkey.public_key(); let netapp = Arc::new(Self { listen_params: ArcSwapOption::new(None), netid, id, privkey, server_conns: RwLock::new(HashMap::new()), client_conns: RwLock::new(HashMap::new()), endpoints: RwLock::new(HashMap::new()), hello_endpoint: ArcSwapOption::new(None), on_connected_handler: ArcSwapOption::new(None), on_disconnected_handler: ArcSwapOption::new(None), }); netapp .hello_endpoint .swap(Some(netapp.endpoint("__netapp/netapp.rs/Hello".into()))); netapp .hello_endpoint .load_full() .unwrap() .set_handler(netapp.clone()); netapp } /// Set the handler to be called when a new connection (incoming or outgoing) has /// been successfully established. Do not set this if using a peering strategy, /// as the peering strategy will need to set this itself. pub fn on_connected(&self, handler: F) where F: Fn(NodeID, SocketAddr, bool) + Sized + Send + Sync + 'static, { self.on_connected_handler .store(Some(Arc::new(Box::new(handler)))); } /// Set the handler to be called when an existing connection (incoming or outgoing) has /// been closed by either party. Do not set this if using a peering strategy, /// as the peering strategy will need to set this itself. pub fn on_disconnected(&self, handler: F) where F: Fn(NodeID, bool) + Sized + Send + Sync + 'static, { self.on_disconnected_handler .store(Some(Arc::new(Box::new(handler)))); } pub fn endpoint(self: &Arc, name: String) -> Arc> where M: Message + 'static, H: EndpointHandler + 'static, { let endpoint = Arc::new(Endpoint::::new(self.clone(), name.clone())); let endpoint_arc = EndpointArc(endpoint.clone()); if self .endpoints .write() .unwrap() .insert(name.clone(), Box::new(endpoint_arc)) .is_some() { panic!("Redefining endpoint: {}", name); }; endpoint } /// Main listening process for our app. This future runs during the whole /// run time of our application. /// If this is not called, the NetApp instance remains a passive client. pub async fn listen(self: Arc, listen_addr: SocketAddr, public_addr: Option) { let listen_params = ListenParams { listen_addr, public_addr, }; if self.listen_params.swap(Some(Arc::new(listen_params))).is_some() { error!("Trying to listen on NetApp but we're already listening!"); } let listener = TcpListener::bind(listen_addr).await.unwrap(); info!("Listening on {}", listen_addr); loop { // The second item contains the IP and port of the new connection. let (socket, _) = listener.accept().await.unwrap(); info!( "Incoming connection from {}, negotiating handshake...", match socket.peer_addr() { Ok(x) => format!("{}", x), Err(e) => format!("", e), } ); let self2 = self.clone(); tokio::spawn(async move { ServerConn::run(self2, socket) .await .log_err("ServerConn::run"); }); } } /// Attempt to connect to a peer, given by its ip:port and its public key. /// The public key will be checked during the secret handshake process. /// This function returns once the connection has been established and a /// successfull handshake was made. At this point we can send messages to /// the other node with `Netapp::request` pub async fn try_connect(self: Arc, ip: SocketAddr, id: NodeID) -> Result<(), Error> { // Don't connect to ourself, we don't care // but pretend we did if id == self.id { tokio::spawn(async move { if let Some(h) = self.on_connected_handler.load().as_ref() { h(id, ip, false); } }); return Ok(()); } // Don't connect if already connected if self.client_conns.read().unwrap().contains_key(&id) { return Ok(()); } let socket = TcpStream::connect(ip).await?; info!("Connected to {}, negotiating handshake...", ip); ClientConn::init(self, socket, id).await?; Ok(()) } /// Close the outgoing connection we have to a node specified by its public key, /// if such a connection is currently open. pub fn disconnect(self: &Arc, id: &NodeID) { // If id is ourself, we're not supposed to have a connection open if *id != self.id { let conn = self.client_conns.write().unwrap().remove(id); if let Some(c) = conn { debug!( "Closing connection to {} ({})", hex::encode(c.peer_id), c.remote_addr ); c.close(); } else { return; } } // call on_disconnected_handler immediately, since the connection // was removed // (if id == self.id, we pretend we disconnected) let id = *id; let self2 = self.clone(); tokio::spawn(async move { if let Some(h) = self2.on_disconnected_handler.load().as_ref() { h(id, false); } }); } /// Close the incoming connection from a certain client to us, /// if such a connection is currently open. pub fn server_disconnect(self: &Arc, id: &NodeID) { let conn = self.server_conns.read().unwrap().get(id).cloned(); if let Some(c) = conn { debug!( "Closing incoming connection from {} ({})", hex::encode(c.peer_id), c.remote_addr ); c.close(); } } // Called from conn.rs when an incoming connection is successfully established // Registers the connection in our list of connections // Do not yet call the on_connected handler, because we don't know if the remote // has an actual IP address and port we can call them back on. // We will know this when they send a Hello message, which is handled below. pub(crate) fn connected_as_server(&self, id: NodeID, conn: Arc) { info!("Accepted connection from {}", hex::encode(id)); self.server_conns.write().unwrap().insert(id, conn); } // Handle hello message from a client. This message is used for them to tell us // that they are listening on a certain port number on which we can call them back. // At this point we know they are a full network member, and not just a client, // and we call the on_connected handler so that the peering strategy knows // we have a new potential peer // Called from conn.rs when an incoming connection is closed. // We deregister the connection from server_conns and call the // handler registered by on_disconnected pub(crate) fn disconnected_as_server(&self, id: &NodeID, conn: Arc) { info!("Connection from {} closed", hex::encode(id)); let mut conn_list = self.server_conns.write().unwrap(); if let Some(c) = conn_list.get(id) { if Arc::ptr_eq(c, &conn) { conn_list.remove(id); drop(conn_list); if let Some(h) = self.on_disconnected_handler.load().as_ref() { h(conn.peer_id, true); } } } } // Called from conn.rs when an outgoinc connection is successfully established. // The connection is registered in self.client_conns, and the // on_connected handler is called. // // Since we are ourself listening, we send them a Hello message so that // they know on which port to call us back. (TODO: don't do this if we are // just a simple client and not a full p2p node) pub(crate) fn connected_as_client(&self, id: NodeID, conn: Arc) { info!("Connection established to {}", hex::encode(id)); { let old_c_opt = self.client_conns.write().unwrap().insert(id, conn.clone()); if let Some(old_c) = old_c_opt { tokio::spawn(async move { old_c.close() }); } } if let Some(h) = self.on_connected_handler.load().as_ref() { h(conn.peer_id, conn.remote_addr, false); } if let Some(lp) = self.listen_params.load_full() { let server_addr = lp.public_addr; let server_port = lp.listen_addr.port(); let hello_endpoint = self.hello_endpoint.load_full().unwrap(); tokio::spawn(async move { hello_endpoint .call( &conn.peer_id, HelloMessage { server_addr, server_port, }, PRIO_NORMAL, ) .await .log_err("Sending hello message"); }); } } // Called from conn.rs when an outgoinc connection is closed. // The connection is removed from conn_list, and the on_disconnected handler // is called. pub(crate) fn disconnected_as_client(&self, id: &NodeID, conn: Arc) { info!("Connection to {} closed", hex::encode(id)); let mut conn_list = self.client_conns.write().unwrap(); if let Some(c) = conn_list.get(id) { if Arc::ptr_eq(c, &conn) { conn_list.remove(id); drop(conn_list); if let Some(h) = self.on_disconnected_handler.load().as_ref() { h(conn.peer_id, false); } } } // else case: happens if connection was removed in .disconnect() // in which case on_disconnected_handler was already called } } #[async_trait] impl EndpointHandler for NetApp { async fn handle(self: &Arc, msg: HelloMessage, from: NodeID) { if let Some(h) = self.on_connected_handler.load().as_ref() { if let Some(c) = self.server_conns.read().unwrap().get(&from) { let remote_ip = msg.server_addr.unwrap_or_else(|| c.remote_addr.ip()); let remote_addr = SocketAddr::new(remote_ip, msg.server_port); h(from, remote_addr, true); } } } }