From 609f0387ca7e04cf1406f10029115f7c09e3a3d9 Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Fri, 11 Dec 2020 15:53:59 +0100
Subject: [PATCH] Add ability to specify our public addr for people to contact
 us back

---
 Makefile             |  4 +++-
 examples/basalt.rs   |  6 +++++-
 examples/fullmesh.rs |  6 +++++-
 src/message.rs       |  3 +++
 src/netapp.rs        | 13 ++++++++++---
 5 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/Makefile b/Makefile
index 0ec01f0..53921a9 100644
--- a/Makefile
+++ b/Makefile
@@ -1,6 +1,8 @@
 all:
 	cargo build
-	RUST_LOG=netapp=debug cargo run --example fullmesh -- -n 3242ce79e05e8b6a0e43441fbd140a906e13f335f298ae3a52f29784abbab500 -p 6c304114a0e1018bbe60502a34d33f4f439f370856c3333dda2726da01eb93a4894b7ef7249a71f11d342b69702f1beb7c93ec95fbcf122ad1eca583bb0629e7
+	cargo build --example fullmesh
+	cargo build --example basalt
+	#RUST_LOG=netapp=debug cargo run --example fullmesh -- -n 3242ce79e05e8b6a0e43441fbd140a906e13f335f298ae3a52f29784abbab500 -p 6c304114a0e1018bbe60502a34d33f4f439f370856c3333dda2726da01eb93a4894b7ef7249a71f11d342b69702f1beb7c93ec95fbcf122ad1eca583bb0629e7
 
 docker_basalt:
 	cargo build --release --example basalt
diff --git a/examples/basalt.rs b/examples/basalt.rs
index eaf056b..e1ac2e4 100644
--- a/examples/basalt.rs
+++ b/examples/basalt.rs
@@ -31,6 +31,9 @@ pub struct Opt {
 	#[structopt(long = "listen-addr", short = "l", default_value = "127.0.0.1:1980")]
 	listen_addr: String,
 
+	#[structopt(long = "public-addr", short = "a")]
+	public_addr: Option<String>,
+
 	#[structopt(long = "view-size", short = "v", default_value = "100")]
 	view_size: usize,
 
@@ -83,7 +86,8 @@ async fn main() {
 	info!("KYEV PK {}", hex::encode(&privkey.public_key()));
 
 	let listen_addr = opt.listen_addr.parse().unwrap();
-	let netapp = NetApp::new(listen_addr, netid, privkey);
+	let public_addr = opt.public_addr.map(|x| x.parse().unwrap());
+	let netapp = NetApp::new(listen_addr, public_addr, netid, privkey);
 
 	let mut bootstrap_peers = vec![];
 	for peer in opt.bootstrap_peers.iter() {
diff --git a/examples/fullmesh.rs b/examples/fullmesh.rs
index 5addcea..88784c1 100644
--- a/examples/fullmesh.rs
+++ b/examples/fullmesh.rs
@@ -25,6 +25,9 @@ pub struct Opt {
 
 	#[structopt(long = "listen-addr", short = "l", default_value = "127.0.0.1:1980")]
 	listen_addr: String,
+
+	#[structopt(long = "public-addr", short = "a")]
+	public_addr: Option<String>,
 }
 
 #[tokio::main]
@@ -63,7 +66,8 @@ async fn main() {
 	info!("Node public key: {}", hex::encode(&privkey.public_key()));
 
 	let listen_addr = opt.listen_addr.parse().unwrap();
-	let netapp = NetApp::new(listen_addr, netid, privkey);
+	let public_addr = opt.public_addr.map(|x| x.parse().unwrap());
+	let netapp = NetApp::new(listen_addr, public_addr, netid, privkey);
 
 	let mut bootstrap_peers = vec![];
 	for peer in opt.bootstrap_peers.iter() {
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");
 		});