From c6d41ce712b994450d9faa42d1ecc6f18230b871 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 b2435bef85..53eec1c7dd 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -68,6 +68,9 @@ version = "2.0"
 [dependencies.hex]
 version = "0.4"
 
+[dependencies.igd]
+version = "0.12"
+
 [dependencies.hyper]
 version = "0.14"
 features = [ "http1", "runtime", "server", "tcp" ]
diff --git a/src/network/server.rs b/src/network/server.rs
index 4643f1a51f..e8387a424f 100644
--- a/src/network/server.rs
+++ b/src/network/server.rs
@@ -29,6 +29,7 @@ use snarkos_storage::{storage::rocksdb::RocksDB, LedgerState};
 use snarkvm::prelude::*;
 
 use anyhow::Result;
+use igd::{AddPortError, PortMappingProtocol};
 use std::{net::SocketAddr, sync::Arc, time::Duration};
 use tokio::{
     net::TcpListener,
@@ -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 4036d6045f..075700f236 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,