Add ability to specify our public addr for people to contact us back

2 files changed, 13 insertions, 3 deletions

diff --git a/src/message.rs b/src/message.rs
index bd54523..9ab14f9 100644
--- a/src/message.rs
+++ b/src/message.rs
@@ -1,3 +1,5 @@
+use std::net::IpAddr;
+
 use serde::{Deserialize, Serialize};
 
 pub type MessageKind = u32;
@@ -24,6 +26,7 @@ pub trait Message: Serialize + for<'de> Deserialize<'de> + Send + Sync {
 
 #[derive(Serialize, Deserialize)]
 pub(crate) struct HelloMessage {
+	pub server_addr: Option<IpAddr>,
 	pub server_port: u16,
 }
 
diff --git a/src/netapp.rs b/src/netapp.rs
index 6f6da5b..8397be9 100644
--- a/src/netapp.rs
+++ b/src/netapp.rs
@@ -1,6 +1,6 @@
 use std::any::Any;
 use std::collections::HashMap;
-use std::net::SocketAddr;
+use std::net::{IpAddr, SocketAddr};
 use std::pin::Pin;
 use std::sync::{Arc, RwLock};
 use std::time::Instant;
@@ -48,6 +48,8 @@ pub(crate) struct Handler {
 /// and RPS peering strategies take care of the most common use cases.
 pub struct NetApp {
 	pub listen_addr: SocketAddr,
+	pub public_addr: Option<IpAddr>,
+
 	pub netid: auth::Key,
 	pub pubkey: ed25519::PublicKey,
 	pub privkey: ed25519::SecretKey,
@@ -111,12 +113,14 @@ impl NetApp {
 	/// Creates a new instance of NetApp. No background process is
 	pub fn new(
 		listen_addr: SocketAddr,
+		public_addr: Option<IpAddr>,
 		netid: auth::Key,
 		privkey: ed25519::SecretKey,
 	) -> Arc<Self> {
 		let pubkey = privkey.public_key();
 		let netapp = Arc::new(Self {
 			listen_addr,
+			public_addr,
 			netid,
 			pubkey,
 			privkey,
@@ -311,7 +315,9 @@ impl NetApp {
 	fn handle_hello_message(&self, id: ed25519::PublicKey, msg: HelloMessage) {
 		if let Some(h) = self.on_connected_handler.load().as_ref() {
 			if let Some(c) = self.server_conns.read().unwrap().get(&id) {
-				let remote_addr = SocketAddr::new(c.remote_addr.ip(), msg.server_port);
+				let remote_ip = msg.server_addr
+					.unwrap_or(c.remote_addr.ip());
+				let remote_addr = SocketAddr::new(remote_ip, msg.server_port);
 				h(id, remote_addr, true);
 			}
 		}
@@ -357,9 +363,10 @@ impl NetApp {
 			h(conn.peer_pk, conn.remote_addr, false);
 		}
 
+		let server_addr = self.public_addr;
 		let server_port = self.listen_addr.port();
 		tokio::spawn(async move {
-			conn.request(HelloMessage { server_port }, PRIO_NORMAL)
+			conn.request(HelloMessage { server_addr, server_port }, PRIO_NORMAL)
 				.await
 				.log_err("Sending hello message");
 		});