From 1b4e029b26437ab0414cc974124972e65bc71195 Mon Sep 17 00:00:00 2001
From: ljedrz <ljedrz@gmail.com>
Date: Tue, 18 Jan 2022 15:20:47 +0100
Subject: [PATCH] feat: enable UPnP-powered automatic port forwarding

Signed-off-by: ljedrz <ljedrz@gmail.com>
---
 Cargo.toml            |  3 +++
 src/network/server.rs | 24 ++++++++++++++++++++++++
 src/node.rs           |  3 +++
 3 files changed, 30 insertions(+)

diff --git a/Cargo.toml b/Cargo.toml
index ab10ce507f..bb26304fa2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -71,6 +71,9 @@ tui = "0.17.0"
   [dependencies.hex]
   version = "0.4"
 
+  [dependencies.igd]
+  version = "0.12"
+
   [dependencies.jsonrpsee]
   version = "0.9"
   optional = true
diff --git a/src/network/server.rs b/src/network/server.rs
index 3e66460c13..14a26acf92 100644
--- a/src/network/server.rs
+++ b/src/network/server.rs
@@ -33,6 +33,7 @@ use crate::rpc::{context::RpcContext, initialize_rpc_server};
 use tokio::sync::RwLock;
 
 use anyhow::Result;
+use igd::{AddPortError, PortMappingProtocol};
 use std::{net::SocketAddr, sync::Arc, time::Duration};
 use tokio::{net::TcpListener, sync::oneshot, task};
 
@@ -61,6 +62,29 @@ impl<N: Network, E: Environment> Server<N, E> {
     ///
     #[inline]
     pub async fn initialize(node: &Node, address: Option<Address<N>>, pool_ip: Option<SocketAddr>) -> Result<Self> {
+        // Use UPnP to detect the gateway, if the upnp option is enabled.
+        if node.upnp {
+            let gateway_search_opts = igd::SearchOptions {
+                timeout: Some(Duration::from_secs(1)),
+                ..Default::default()
+            };
+            let gateway = igd::search_gateway(gateway_search_opts)
+                .map_err(|e| warn!("Can't obtain the node's gateway details: {}; perhaps it's not UPnP-enabled?", e))
+                .ok();
+
+            if let Some(ref gateway) = gateway {
+                if let SocketAddr::V4(internal_addr) = node.node {
+                    let external_port = internal_addr.port();
+
+                    match gateway.add_port(PortMappingProtocol::TCP, external_port, internal_addr, 0, "snarkOS") {
+                        Err(AddPortError::PortInUse) => debug!("Port {} is already forwarded", external_port),
+                        Err(e) => error!("Can't map external port {} to address {}: {}", external_port, internal_addr, e),
+                        Ok(_) => info!("Enabled port forwarding via UPnP"),
+                    }
+                }
+            }
+        }
+
         // Initialize a new TCP listener at the given IP.
         let (local_ip, listener) = match TcpListener::bind(node.node).await {
             Ok(listener) => (listener.local_addr().expect("Failed to fetch the local IP"), listener),
diff --git a/src/node.rs b/src/node.rs
index 0a3c9a4eee..0edb18163d 100644
--- a/src/node.rs
+++ b/src/node.rs
@@ -55,6 +55,9 @@ pub struct Node {
     /// Specify the IP address and port for the node server.
     #[structopt(parse(try_from_str), default_value = "0.0.0.0:4132", long = "node")]
     pub node: SocketAddr,
+    /// If the flag is set, the node will attempt to use UPnP to open its listening port.
+    #[structopt(long = "upnp")]
+    pub upnp: bool,
     /// Specify the IP address and port for the RPC server.
     #[structopt(parse(try_from_str), default_value = "0.0.0.0:3032", long = "rpc")]
     pub rpc: SocketAddr,