diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/message.rs | 3 | ||||
| -rw-r--r-- | src/netapp.rs | 13 |
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"); }); |