From f91551a109284501a6759ce944a828cefbbf6151 Mon Sep 17 00:00:00 2001
From: Ryo Kawaguchi <ryo@kawagu.ch>
Date: Thu, 13 Jun 2024 16:20:58 +0900
Subject: [PATCH] Support override-endpoint with the unspecified address

---
 server/src/api/mod.rs | 38 ++++++++++++++++++++++++++------------
 shared/src/prompts.rs |  8 +++++++-
 shared/src/types.rs   | 17 ++++++++++++++++-
 3 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/server/src/api/mod.rs b/server/src/api/mod.rs
index da0a12c..43a2261 100644
--- a/server/src/api/mod.rs
+++ b/server/src/api/mod.rs
@@ -1,26 +1,40 @@
-use shared::Peer;
+use std::net::SocketAddr;
 
 use crate::Session;
+use shared::{Endpoint, Peer};
 
 pub mod admin;
 pub mod user;
 
-/// Inject the collected endpoints from the WG interface into a list of peers.
-/// This is essentially what adds NAT holepunching functionality. If a peer
-/// already has an endpoint specified (by calling the override-endpoint) API,
-/// the relatively recent wireguard endpoint will be added to the list of NAT
-/// candidates, so other peers have a better chance of connecting.
+/// Implements NAT traversal strategies.
+/// (1) NAT holepunching: Report the most recent wireguard endpoint as the peer's
+///     endpoint or add it to the list of NAT candidates if an override enpoint is
+///     specified. Note that NAT traversal does not always work e.g. if the peer is
+///     behind double NAT or address/port restricted cone NAT.
+/// (2) Unspecified endpoint IP: A peer may report an override endpoint with
+///     an unspecified IP. It typically indicates the peer does not have a fixed
+///     global IP, and it needs help from the innernet server to resolve it.
+///     Override the endpoint IP with what's most recently reported by wireguard.
 pub fn inject_endpoints(session: &Session, peers: &mut Vec<Peer>) {
     for peer in peers {
         let endpoints = session.context.endpoints.read();
         if let Some(wg_endpoint) = endpoints.get(&peer.public_key) {
-            if peer.contents.endpoint.is_none() {
-                peer.contents.endpoint = Some(wg_endpoint.to_owned().into());
+            let wg_endpoint_ip = wg_endpoint.ip();
+            let wg_endpoint: Endpoint = wg_endpoint.to_owned().into();
+            if let Some(endpoint) = &mut peer.contents.endpoint {
+                if endpoint.is_host_unspecified() {
+                    // (2) Unspecified endpoint host
+                    *endpoint = SocketAddr::new(wg_endpoint_ip, endpoint.port()).into();
+                } else if *endpoint != wg_endpoint {
+                    // (1) NAT holepunching
+                    // The peer already has an endpoint specified, but it might be stale.
+                    // If there is an endpoint reported from wireguard, we should add it
+                    // to the list of candidates so others can try to connect using it.
+                    peer.contents.candidates.push(wg_endpoint);
+                }
             } else {
-                // The peer already has an endpoint specified, but it might be stale.
-                // If there is an endpoint reported from wireguard, we should add it
-                // to the list of candidates so others can try to connect using it.
-                peer.contents.candidates.push(wg_endpoint.to_owned().into());
+                // (1) NAT holepunching
+                peer.contents.endpoint = Some(wg_endpoint);
             }
         }
     }
diff --git a/shared/src/prompts.rs b/shared/src/prompts.rs
index 238249a..f69a5bb 100644
--- a/shared/src/prompts.rs
+++ b/shared/src/prompts.rs
@@ -15,7 +15,7 @@ use std::{
     fmt::{Debug, Display},
     fs::{File, OpenOptions},
     io,
-    net::SocketAddr,
+    net::{IpAddr, Ipv6Addr, SocketAddr},
     str::FromStr,
     time::SystemTime,
 };
@@ -565,6 +565,12 @@ pub fn ask_endpoint(listen_port: u16) -> Result<Endpoint, Error> {
         .interact()?
     {
         publicip::get_any(Preference::Ipv4)
+    } else if Confirm::with_theme(&*THEME)
+        .wait_for_newline(true)
+        .with_prompt("You do not have a fixed global IP (use the unspecified address)?")
+        .interact()?
+    {
+        Some(IpAddr::V6(Ipv6Addr::UNSPECIFIED))
     } else {
         None
     };
diff --git a/shared/src/types.rs b/shared/src/types.rs
index cd8f156..c7388b5 100644
--- a/shared/src/types.rs
+++ b/shared/src/types.rs
@@ -146,6 +146,19 @@ impl Endpoint {
             )
         })
     }
+
+    /// Returns true if the endpoint host is unspecified e.g. 0.0.0.0
+    pub fn is_host_unspecified(&self) -> bool {
+        match self.host {
+            Host::Ipv4(ip) => ip.is_unspecified(),
+            Host::Ipv6(ip) => ip.is_unspecified(),
+            Host::Domain(_) => false,
+        }
+    }
+
+    pub fn port(&self) -> u16 {
+        self.port
+    }
 }
 
 #[derive(Deserialize, Serialize, Debug)]
@@ -437,7 +450,9 @@ pub struct ListenPortOpts {
 
 #[derive(Debug, Clone, PartialEq, Eq, Args)]
 pub struct OverrideEndpointOpts {
-    /// The listen port you'd like to set for the interface
+    /// The external endpoint that you'd like the innernet server to broadcast
+    /// to other peers. The IP address may be unspecified, in which case the
+    /// server will try to resolve it based on its most recent connection.
     #[clap(short, long)]
     pub endpoint: Option<Endpoint>,