From 09fac449e1d2cf0d7e849cbb338962153050e238 Mon Sep 17 00:00:00 2001
From: Josh Wilson <joshuef@gmail.com>
Date: Wed, 17 Jan 2024 11:27:57 +0100
Subject: [PATCH] feat: add quic/tcp features and set quic as default

---
 Cargo.toml     |  6 ++++++
 src/local.rs   | 30 +++++++++++++++++++++---------
 src/service.rs | 12 +++++++++---
 3 files changed, 36 insertions(+), 12 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 267c0e1..94725fb 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,10 +7,16 @@ description = "A command-line application for installing, managing and operating
 license = "GPL-3.0"
 repository = "https://github.com/maidsafe/sn-node-manager"
 
+
 [[bin]]
 path="src/main.rs"
 name="safenode-manager"
 
+[features]
+default = ["quic"]
+quic = []
+tcp = []
+
 [dependencies]
 clap = { version = "4.4.6", features = ["derive", "env"]}
 colored = "2.0.4"
diff --git a/src/local.rs b/src/local.rs
index 5f3f52a..ff9a59d 100644
--- a/src/local.rs
+++ b/src/local.rs
@@ -273,15 +273,28 @@ pub async fn run_node(
         safenode_path: Some(launcher.get_safenode_path()),
     });
 
-    Ok(Multiaddr::from_str(&format!(
-        "/ip4/127.0.0.1/tcp/{port}/p2p/{peer_id}"
-    ))?)
+    Ok(build_multiaddr(port, peer_id))
 }
 
 ///
 /// Private Helpers
 ///
 
+/// Builds a MultiAddress taking into account any required ffeature set of the nodes.
+/// Default is a quic multiaddress
+fn build_multiaddr(port: u16, peer_id: PeerId) -> Multiaddr {
+    let addr = Multiaddr::from(std::net::Ipv4Addr::LOCALHOST);
+
+    #[cfg(feature = "tcp")]
+    let addr = addr.with(libp2p::multiaddr::Protocol::Tcp(port));
+    #[cfg(feature = "quic")]
+    let addr = addr
+        .with(libp2p::multiaddr::Protocol::Udp(port))
+        .with(libp2p::multiaddr::Protocol::QuicV1);
+
+    addr.with(libp2p::multiaddr::Protocol::P2p(peer_id))
+}
+
 #[cfg(target_os = "windows")]
 fn get_username() -> Result<String> {
     Ok(std::env::var("USERNAME")?)
@@ -376,8 +389,8 @@ mod tests {
         let peer_id = PeerId::from_str("12D3KooWS2tpXGGTmg2AHFiDh57yPQnat49YHnyqoggzXZWpqkCR")?;
         let port = 12000;
         let rpc_port = 13000;
-        let node_multiaddr =
-            Multiaddr::from_str(&format!("/ip4/127.0.0.1/tcp/{port}/p2p/{peer_id}"))?;
+
+        let node_multiaddr = build_multiaddr(port, peer_id);
 
         mock_launcher
             .expect_get_safenode_version()
@@ -452,8 +465,8 @@ mod tests {
     #[tokio::test]
     async fn run_node_should_launch_an_additional_node() -> Result<()> {
         let peer_id = PeerId::from_str("12D3KooWS2tpXGGTmg2AHFiDh57yPQnat49YHnyqoggzXZWpqkCR")?;
-        let genesis_peer_addr =
-            Multiaddr::from_str(&format!("/ip4/127.0.0.1/tcp/12000/p2p/{peer_id}"))?;
+
+        let genesis_peer_addr = build_multiaddr(12000, peer_id);
 
         let mut mock_launcher = MockLauncher::new();
         let mut node_registry = NodeRegistry {
@@ -480,8 +493,7 @@ mod tests {
         let peer_id = PeerId::from_str("12D3KooWS2tpXGGTmg2AHFiDh57yPQnat49YHnyqoggzXZWpqkCR")?;
         let port = 12001;
         let rpc_port = 13001;
-        let node_peer_addr =
-            Multiaddr::from_str(&format!("/ip4/127.0.0.1/tcp/{port}/p2p/{peer_id}"))?;
+        let node_peer_addr = build_multiaddr(port, peer_id);
 
         mock_launcher
             .expect_get_safenode_version()
diff --git a/src/service.rs b/src/service.rs
index 7489f55..8a81ffa 100644
--- a/src/service.rs
+++ b/src/service.rs
@@ -15,7 +15,12 @@ use service_manager::{
     ServiceUninstallCtx,
 };
 use std::ffi::OsString;
-use std::net::{SocketAddr, TcpListener};
+use std::net::SocketAddr;
+#[cfg(feature = "tcp")]
+use std::net::TcpListener as SocketBinder;
+#[cfg(not(feature = "tcp"))]
+use std::net::UdpSocket as SocketBinder;
+
 use std::path::PathBuf;
 use sysinfo::{Pid, System, SystemExt};
 
@@ -145,7 +150,7 @@ impl ServiceControl for NodeServiceManager {
     }
 
     fn is_port_free(&self, port: u16) -> bool {
-        TcpListener::bind(("127.0.0.1", port)).is_ok()
+        SocketBinder::bind(("127.0.0.1", port)).is_ok()
     }
 
     fn is_service_process_running(&self, pid: u32) -> bool {
@@ -156,7 +161,8 @@ impl ServiceControl for NodeServiceManager {
 
     fn get_available_port(&self) -> Result<u16> {
         let addr: SocketAddr = "127.0.0.1:0".parse().unwrap();
-        Ok(TcpListener::bind(addr)?.local_addr()?.port())
+
+        Ok(SocketBinder::bind(addr)?.local_addr()?.port())
     }
 
     fn install(&self, config: ServiceConfig) -> Result<()> {