From ed96b2a9e503d2aaebf4ad6c082a2dca7ceea9b3 Mon Sep 17 00:00:00 2001
From: magine <magine233@gmail.com>
Date: Mon, 21 Oct 2024 20:48:48 +0800
Subject: [PATCH] fix: set external ip to fix relay connection reservation

---
 Cargo.lock     |  2 +-
 Cargo.toml     |  2 +-
 src/lib.rs     |  4 +++-
 src/main.rs    | 11 +++++++++++
 src/p2p/mod.rs | 11 +++++++++++
 5 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index e577a17..a2fe68e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -665,7 +665,7 @@ dependencies = [
 
 [[package]]
 name = "dephy-pproxy"
-version = "0.4.0"
+version = "0.4.1"
 dependencies = [
  "async-trait",
  "bytes",
diff --git a/Cargo.toml b/Cargo.toml
index dc6d62e..cb5156e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "dephy-pproxy"
-version = "0.4.0"
+version = "0.4.1"
 edition = "2021"
 
 [dependencies]
diff --git a/src/lib.rs b/src/lib.rs
index 55ccdf1..90699c1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,4 +1,5 @@
 use std::collections::HashMap;
+use std::net::IpAddr;
 use std::net::SocketAddr;
 use std::sync::atomic::AtomicUsize;
 use std::sync::Arc;
@@ -118,9 +119,10 @@ impl PProxy {
         listen_addr: SocketAddr,
         proxy_addr: Option<SocketAddr>,
         access_server_endpoint: Option<reqwest::Url>,
+        external_ip: Option<IpAddr>,
     ) -> Result<(Self, PProxyHandle)> {
         let (command_tx, command_rx) = mpsc::channel(DEFAULT_CHANNEL_SIZE);
-        let swarm = crate::p2p::new_swarm(keypair, listen_addr)
+        let swarm = crate::p2p::new_swarm(keypair, listen_addr, external_ip)
             .map_err(|e| Error::Libp2pSwarmCreateError(e.to_string()))?;
         let stream_control = swarm.behaviour().stream.new_control();
         let access_client = access_server_endpoint.map(AccessClient::new);
diff --git a/src/main.rs b/src/main.rs
index aff11de..1b5f392 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -58,6 +58,13 @@ fn parse_args() -> Command {
                 .num_args(1)
                 .action(ArgAction::Set)
                 .help("Access server endpoint is used to verify if one peer can access another. If not set, all access is allowed."),
+        )
+        .arg(
+            Arg::new("EXTERNAL_IP")
+                .long("external-ip")
+                .num_args(1)
+                .action(ArgAction::Set)
+                .help("External ip for relay behaviour. If not set, no external ip will be assigned."),
         );
 
     let create_tunnel_server = Command::new("create_tunnel_server")
@@ -142,6 +149,9 @@ async fn serve(args: &ArgMatches) {
     let access_server_endpoint = args
         .get_one::<String>("ACCESS_SERVER_ENDPOINT")
         .map(|endpoint| Url::parse(endpoint).expect("Invalid access server endpoint"));
+    let external_ip = args
+        .get_one::<String>("EXTERNAL_IP")
+        .map(|addr| addr.parse().expect("Invalid external ip"));
 
     println!("server_addr: {}", server_addr);
     println!("commander_server_addr: {}", commander_server_addr);
@@ -151,6 +161,7 @@ async fn serve(args: &ArgMatches) {
         server_addr,
         proxy_addr,
         access_server_endpoint,
+        external_ip,
     )
     .expect("Create pproxy failed");
 
diff --git a/src/p2p/mod.rs b/src/p2p/mod.rs
index bab0d63..436b552 100644
--- a/src/p2p/mod.rs
+++ b/src/p2p/mod.rs
@@ -18,6 +18,7 @@ pub const PPROXY_PROTOCOL: StreamProtocol = StreamProtocol::new("/pproxy/1.0.0")
 pub(crate) fn new_swarm(
     keypair: Keypair,
     listen_addr: SocketAddr,
+    external_ip: Option<IpAddr>,
 ) -> std::result::Result<Swarm<PProxyNetworkBehaviour>, Box<dyn std::error::Error>> {
     let (ip_type, ip, port) = match listen_addr.ip() {
         IpAddr::V4(ip) => ("ip4", ip.to_string(), listen_addr.port()),
@@ -26,6 +27,12 @@ pub(crate) fn new_swarm(
 
     let listen_multiaddr = format!("/{ip_type}/{ip}/tcp/{port}").parse()?;
 
+    let external_multiaddr = match external_ip {
+        None => None,
+        Some(IpAddr::V4(ip)) => Some(format!("/ip4/{ip}/tcp/{port}").parse()?),
+        Some(IpAddr::V6(ip)) => Some(format!("/ip6/{ip}/tcp/{port}").parse()?),
+    };
+
     let mut swarm = libp2p::SwarmBuilder::with_existing_identity(keypair)
         .with_tokio()
         .with_tcp(
@@ -40,5 +47,9 @@ pub(crate) fn new_swarm(
 
     swarm.listen_on(listen_multiaddr)?;
 
+    if let Some(external_multiaddr) = external_multiaddr {
+        swarm.add_external_address(external_multiaddr);
+    }
+
     Ok(swarm)
 }