From fcd5a0f8753ad02439dad6e3a30bcb697a9852e3 Mon Sep 17 00:00:00 2001
From: Kasey <klhuizinga@gmail.com>
Date: Mon, 16 Sep 2024 13:31:01 -0400
Subject: [PATCH] adjust arguments to take `SocketAddrV4` and `SocketAddrV6`
 rather than a port

---
 src/main.rs | 108 +++++++++++++++++++++++++---------------------------
 1 file changed, 51 insertions(+), 57 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 6ebbea2..1c4a6af 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -9,7 +9,7 @@ use iroh_net::{
 };
 use std::{
     io,
-    net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs},
+    net::{SocketAddr, SocketAddrV4, SocketAddrV6, ToSocketAddrs},
     str::FromStr,
 };
 use tokio::{
@@ -72,9 +72,19 @@ pub enum Commands {
 
 #[derive(Parser, Debug)]
 pub struct CommonArgs {
-    /// The port to use for the magicsocket. Random by default.
-    #[clap(long, default_value_t = 0)]
-    pub magic_port: u16,
+    /// The IPv4 address that magicsocket will listen on.
+    ///
+    /// If None, defaults to a random free port, but it can be useful to specify a fixed
+    /// port, e.g. to configure a firewall rule.
+    #[clap(long, default_value = None)]
+    pub magic_ipv4_addr: Option<SocketAddrV4>,
+
+    /// The IPv6 address that magicsocket will listen on.
+    ///
+    /// If None, defaults to a random free port, but it can be useful to specify a fixed
+    /// port, e.g. to configure a firewall rule.
+    #[clap(long, default_value = None)]
+    pub magic_ipv6_addr: Option<SocketAddrV6>,
 
     /// A custom ALPN to use for the magicsocket.
     ///
@@ -258,21 +268,16 @@ async fn forward_bidi(
 
 async fn listen_stdio(args: ListenArgs) -> anyhow::Result<()> {
     let secret_key = get_or_create_secret()?;
-    let endpoint = Endpoint::builder()
+    let mut builder = Endpoint::builder()
         .alpns(vec![args.common.alpn()?])
-        .secret_key(secret_key)
-        .bind_addr_v4(SocketAddrV4::new(
-            Ipv4Addr::UNSPECIFIED,
-            args.common.magic_port,
-        ))
-        .bind_addr_v6(SocketAddrV6::new(
-            Ipv6Addr::UNSPECIFIED,
-            args.common.magic_port + 1,
-            0,
-            0,
-        ))
-        .bind()
-        .await?;
+        .secret_key(secret_key);
+    if let Some(addr) = args.common.magic_ipv4_addr {
+        builder = builder.bind_addr_v4(addr);
+    }
+    if let Some(addr) = args.common.magic_ipv6_addr {
+        builder = builder.bind_addr_v6(addr);
+    }
+    let endpoint = builder.bind().await?;
     // wait for the endpoint to figure out its address before making a ticket
     while endpoint.home_relay().is_none() {
         tokio::time::sleep(std::time::Duration::from_millis(100)).await;
@@ -330,20 +335,17 @@ async fn listen_stdio(args: ListenArgs) -> anyhow::Result<()> {
 
 async fn connect_stdio(args: ConnectArgs) -> anyhow::Result<()> {
     let secret_key = get_or_create_secret()?;
-    let endpoint = Endpoint::builder()
+    let mut builder = Endpoint::builder()
         .secret_key(secret_key)
-        .alpns(vec![])
-        .bind_addr_v4(SocketAddrV4::new(
-            Ipv4Addr::UNSPECIFIED,
-            args.common.magic_port,
-        ))
-        .bind_addr_v6(SocketAddrV6::new(
-            Ipv6Addr::UNSPECIFIED,
-            args.common.magic_port + 1,
-            0,
-            0,
-        ))
-        .bind()
+        .alpns(vec![]);
+   
+    if let Some(addr) = args.common.magic_ipv4_addr {
+        builder = builder.bind_addr_v4(addr);
+    }
+    if let Some(addr) = args.common.magic_ipv6_addr {
+        builder = builder.bind_addr_v6(addr);
+    }
+    let endpoint = builder.bind()
         .await?;
     let addr = args.ticket.node_addr();
     let remote_node_id = addr.node_id;
@@ -373,20 +375,16 @@ async fn connect_tcp(args: ConnectTcpArgs) -> anyhow::Result<()> {
         .to_socket_addrs()
         .context(format!("invalid host string {}", args.addr))?;
     let secret_key = get_or_create_secret()?;
-    let endpoint = Endpoint::builder()
+    let mut builder = Endpoint::builder()
         .alpns(vec![])
-        .secret_key(secret_key)
-        .bind_addr_v4(SocketAddrV4::new(
-            Ipv4Addr::UNSPECIFIED,
-            args.common.magic_port,
-        ))
-        .bind_addr_v6(SocketAddrV6::new(
-            Ipv6Addr::UNSPECIFIED,
-            args.common.magic_port + 1,
-            0,
-            0,
-        ))
-        .bind()
+        .secret_key(secret_key);
+    if let Some(addr) = args.common.magic_ipv4_addr {
+        builder = builder.bind_addr_v4(addr);
+    }
+    if let Some(addr) = args.common.magic_ipv6_addr {
+        builder = builder.bind_addr_v6(addr);
+    }
+        let endpoint = builder.bind()
         .await
         .context("unable to bind magicsock")?;
     tracing::info!("tcp listening on {:?}", addrs);
@@ -459,20 +457,16 @@ async fn listen_tcp(args: ListenTcpArgs) -> anyhow::Result<()> {
         Err(e) => anyhow::bail!("invalid host string {}: {}", args.host, e),
     };
     let secret_key = get_or_create_secret()?;
-    let endpoint = Endpoint::builder()
+    let mut builder = Endpoint::builder()
         .alpns(vec![args.common.alpn()?])
-        .secret_key(secret_key)
-        .bind_addr_v4(SocketAddrV4::new(
-            Ipv4Addr::UNSPECIFIED,
-            args.common.magic_port,
-        ))
-        .bind_addr_v6(SocketAddrV6::new(
-            Ipv6Addr::UNSPECIFIED,
-            args.common.magic_port + 1,
-            0,
-            0,
-        ))
-        .bind()
+        .secret_key(secret_key);
+    if let Some(addr) = args.common.magic_ipv4_addr {
+        builder = builder.bind_addr_v4(addr);
+    }
+    if let Some(addr) = args.common.magic_ipv6_addr {
+        builder = builder.bind_addr_v6(addr);
+    }
+    let endpoint = builder.bind()
         .await?;
     // wait for the endpoint to figure out its address before making a ticket
     while endpoint.home_relay().is_none() {