Merge branch 'main' into next-0.10

3 files changed, 67 insertions, 34 deletions

diff --git a/src/net/netapp.rs b/src/net/netapp.rs
index b1ad9db8..faa51a99 100644
--- a/src/net/netapp.rs
+++ b/src/net/netapp.rs
@@ -13,7 +13,7 @@ use sodiumoxide::crypto::sign::ed25519;
 
 use futures::stream::futures_unordered::FuturesUnordered;
 use futures::stream::StreamExt;
-use tokio::net::{TcpListener, TcpStream};
+use tokio::net::{TcpListener, TcpSocket, TcpStream};
 use tokio::select;
 use tokio::sync::{mpsc, watch};
 
@@ -38,6 +38,11 @@ pub(crate) type VersionTag = [u8; 16];
 /// Value of the Netapp version used in the version tag
 pub(crate) const NETAPP_VERSION_TAG: u64 = 0x6e65746170700005; // netapp 0x0005
 
+/// HelloMessage is sent by the client on a Netapp connection to indicate
+/// that they are also a server and ready to recieve incoming connections
+/// at the specified address and port. If the client doesn't know their
+/// public address, they don't need to specify it and we look at the
+/// remote address of the socket is used instead.
 #[derive(Serialize, Deserialize, Debug)]
 pub(crate) struct HelloMessage {
 	pub server_addr: Option<IpAddr>,
@@ -56,10 +61,8 @@ type OnDisconnectHandler = Box<dyn Fn(NodeID, bool) + Send + Sync>;
 /// 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 {
+	bind_outgoing_to: Option<IpAddr>,
 	listen_params: ArcSwapOption<ListenParams>,
 
 	/// Version tag, 8 bytes for netapp version, 8 bytes for app version
@@ -83,7 +86,7 @@ pub struct NetApp {
 
 struct ListenParams {
 	listen_addr: SocketAddr,
-	public_addr: Option<IpAddr>,
+	public_addr: Option<SocketAddr>,
 }
 
 impl NetApp {
@@ -92,13 +95,19 @@ impl NetApp {
 	/// using `.listen()`
 	///
 	/// Our Peer ID is the public key associated to the secret key given here.
-	pub fn new(app_version_tag: u64, netid: auth::Key, privkey: ed25519::SecretKey) -> Arc<Self> {
+	pub fn new(
+		app_version_tag: u64,
+		netid: auth::Key,
+		privkey: ed25519::SecretKey,
+		bind_outgoing_to: Option<IpAddr>,
+	) -> Arc<Self> {
 		let mut version_tag = [0u8; 16];
 		version_tag[0..8].copy_from_slice(&u64::to_be_bytes(NETAPP_VERSION_TAG)[..]);
 		version_tag[8..16].copy_from_slice(&u64::to_be_bytes(app_version_tag)[..]);
 
 		let id = privkey.public_key();
 		let netapp = Arc::new(Self {
+			bind_outgoing_to,
 			listen_params: ArcSwapOption::new(None),
 			version_tag,
 			netid,
@@ -180,7 +189,7 @@ impl NetApp {
 	pub async fn listen(
 		self: Arc<Self>,
 		listen_addr: SocketAddr,
-		public_addr: Option<IpAddr>,
+		public_addr: Option<SocketAddr>,
 		mut must_exit: watch::Receiver<bool>,
 	) {
 		let listen_params = ListenParams {
@@ -298,9 +307,20 @@ impl NetApp {
 			return Ok(());
 		}
 
-		let socket = TcpStream::connect(ip).await?;
+		let stream = match self.bind_outgoing_to {
+			Some(addr) => {
+				let socket = if addr.is_ipv4() {
+					TcpSocket::new_v4()?
+				} else {
+					TcpSocket::new_v6()?
+				};
+				socket.bind(SocketAddr::new(addr, 0))?;
+				socket.connect(ip).await?
+			}
+			None => TcpStream::connect(ip).await?,
+		};
 		info!("Connected to {}, negotiating handshake...", ip);
-		ClientConn::init(self, socket, id).await?;
+		ClientConn::init(self, stream, id).await?;
 		Ok(())
 	}
 
@@ -396,8 +416,11 @@ impl NetApp {
 		}
 
 		if let Some(lp) = self.listen_params.load_full() {
-			let server_addr = lp.public_addr;
-			let server_port = lp.listen_addr.port();
+			let server_addr = lp.public_addr.map(|x| x.ip());
+			let server_port = lp
+				.public_addr
+				.map(|x| x.port())
+				.unwrap_or(lp.listen_addr.port());
 			let hello_endpoint = self.hello_endpoint.load_full().unwrap();
 			tokio::spawn(async move {
 				hello_endpoint
diff --git a/src/net/peering.rs b/src/net/peering.rs
index 32199cf8..61882a18 100644
--- a/src/net/peering.rs
+++ b/src/net/peering.rs
@@ -80,6 +80,23 @@ impl PeerInfoInternal {
 			failed_pings: 0,
 		}
 	}
+	fn add_addr(&mut self, addr: SocketAddr) -> bool {
+		if !self.all_addrs.contains(&addr) {
+			self.all_addrs.push(addr);
+			// If we are learning a new address for this node,
+			// we want to retry connecting
+			self.state = match self.state {
+				PeerConnState::Trying(_) => PeerConnState::Trying(0),
+				PeerConnState::Waiting(_, _) | PeerConnState::Abandonned => {
+					PeerConnState::Waiting(0, Instant::now())
+				}
+				x @ (PeerConnState::Ourself | PeerConnState::Connected) => x,
+			};
+			true
+		} else {
+			false
+		}
+	}
 }
 
 /// Information that the full mesh peering strategy can return about the peers it knows of
@@ -147,23 +164,22 @@ struct KnownHosts {
 impl KnownHosts {
 	fn new() -> Self {
 		let list = HashMap::new();
-		let hash = Self::calculate_hash(&list);
+		let hash = Self::calculate_hash(vec![]);
 		Self { list, hash }
 	}
 	fn update_hash(&mut self) {
-		self.hash = Self::calculate_hash(&self.list);
+		self.hash = Self::calculate_hash(self.connected_peers_vec());
 	}
-	fn map_into_vec(input: &HashMap<NodeID, PeerInfoInternal>) -> Vec<(NodeID, SocketAddr)> {
-		let mut list = Vec::with_capacity(input.len());
-		for (id, peer) in input.iter() {
-			if peer.state == PeerConnState::Connected || peer.state == PeerConnState::Ourself {
+	fn connected_peers_vec(&self) -> Vec<(NodeID, SocketAddr)> {
+		let mut list = Vec::with_capacity(self.list.len());
+		for (id, peer) in self.list.iter() {
+			if peer.state.is_up() {
 				list.push((*id, peer.addr));
 			}
 		}
 		list
 	}
-	fn calculate_hash(input: &HashMap<NodeID, PeerInfoInternal>) -> hash::Digest {
-		let mut list = Self::map_into_vec(input);
+	fn calculate_hash(mut list: Vec<(NodeID, SocketAddr)>) -> hash::Digest {
 		list.sort();
 		let mut hash_state = hash::State::new();
 		for (id, addr) in list {
@@ -214,6 +230,7 @@ impl PeeringManager {
 				netapp.id,
 				PeerInfoInternal::new(addr, PeerConnState::Ourself),
 			);
+			known_hosts.update_hash();
 		}
 
 		// TODO for v0.10 / v1.0 : rename the endpoint (it will break compatibility)
@@ -234,13 +251,11 @@ impl PeeringManager {
 
 		let strat2 = strat.clone();
 		netapp.on_connected(move |id: NodeID, addr: SocketAddr, is_incoming: bool| {
-			let strat2 = strat2.clone();
 			strat2.on_connected(id, addr, is_incoming);
 		});
 
 		let strat2 = strat.clone();
 		netapp.on_disconnected(move |id: NodeID, is_incoming: bool| {
-			let strat2 = strat2.clone();
 			strat2.on_disconnected(id, is_incoming);
 		});
 
@@ -445,7 +460,7 @@ impl PeeringManager {
 	}
 
 	async fn exchange_peers(self: Arc<Self>, id: &NodeID) {
-		let peer_list = KnownHosts::map_into_vec(&self.known_hosts.read().unwrap().list);
+		let peer_list = self.known_hosts.read().unwrap().connected_peers_vec();
 		let pex_message = PeerListMessage { list: peer_list };
 		match self
 			.peer_list_endpoint
@@ -465,8 +480,7 @@ impl PeeringManager {
 		let mut changed = false;
 		for (id, addr) in list.iter() {
 			if let Some(kh) = known_hosts.list.get_mut(id) {
-				if !kh.all_addrs.contains(addr) {
-					kh.all_addrs.push(*addr);
+				if kh.add_addr(*addr) {
 					changed = true;
 				}
 			} else {
@@ -534,13 +548,11 @@ impl PeeringManager {
 		}
 	}
 
-	fn on_connected(self: Arc<Self>, id: NodeID, addr: SocketAddr, is_incoming: bool) {
+	fn on_connected(self: &Arc<Self>, id: NodeID, addr: SocketAddr, is_incoming: bool) {
 		let mut known_hosts = self.known_hosts.write().unwrap();
 		if is_incoming {
 			if let Some(host) = known_hosts.list.get_mut(&id) {
-				if !host.all_addrs.contains(&addr) {
-					host.all_addrs.push(addr);
-				}
+				host.add_addr(addr);
 			} else {
 				known_hosts.list.insert(id, self.new_peer(&id, addr));
 			}
@@ -553,9 +565,7 @@ impl PeeringManager {
 			if let Some(host) = known_hosts.list.get_mut(&id) {
 				host.state = PeerConnState::Connected;
 				host.addr = addr;
-				if !host.all_addrs.contains(&addr) {
-					host.all_addrs.push(addr);
-				}
+				host.add_addr(addr);
 			} else {
 				known_hosts
 					.list
@@ -566,7 +576,7 @@ impl PeeringManager {
 		self.update_public_peer_list(&known_hosts);
 	}
 
-	fn on_disconnected(self: Arc<Self>, id: NodeID, is_incoming: bool) {
+	fn on_disconnected(self: &Arc<Self>, id: NodeID, is_incoming: bool) {
 		if !is_incoming {
 			info!("Connection to {} was closed", hex::encode(&id[..8]));
 			let mut known_hosts = self.known_hosts.write().unwrap();
@@ -608,7 +618,7 @@ impl EndpointHandler<PeerListMessage> for PeeringManager {
 		_from: NodeID,
 	) -> PeerListMessage {
 		self.handle_peer_list(&peer_list.list[..]);
-		let peer_list = KnownHosts::map_into_vec(&self.known_hosts.read().unwrap().list);
+		let peer_list = self.known_hosts.read().unwrap().connected_peers_vec();
 		PeerListMessage { list: peer_list }
 	}
 }
diff --git a/src/net/test.rs b/src/net/test.rs
index c6259752..5a3f236d 100644
--- a/src/net/test.rs
+++ b/src/net/test.rs
@@ -102,7 +102,7 @@ fn run_netapp(
 	Arc<NetApp>,
 	Arc<PeeringManager>,
 ) {
-	let netapp = NetApp::new(0u64, netid, sk);
+	let netapp = NetApp::new(0u64, netid, sk, None);
 	let peering = PeeringManager::new(netapp.clone(), bootstrap_peers, None);
 
 	let peering2 = peering.clone();