From b2803493f160dc4c68190446fbc0d76f228c081c Mon Sep 17 00:00:00 2001
From: Ljubisa Isakovic <93769705+sh3ll3x3c@users.noreply.github.com>
Date: Thu, 16 Nov 2023 19:31:26 +0100
Subject: [PATCH] added dns address parsing for autonat servers (#348)

* added dns address parsing for autonat servers

* changed how autonat addresses are reformed

---------

Co-authored-by: momosh <momcilo@ethernal.tech>
---
 src/network/p2p/client.rs | 51 ++++++++++++++++++++++-----------------
 1 file changed, 29 insertions(+), 22 deletions(-)

diff --git a/src/network/p2p/client.rs b/src/network/p2p/client.rs
index 3c9820d65..57f327e01 100644
--- a/src/network/p2p/client.rs
+++ b/src/network/p2p/client.rs
@@ -12,7 +12,7 @@ use libp2p::{
 };
 use std::time::{Duration, Instant};
 use tokio::sync::{mpsc, oneshot};
-use tracing::{debug, trace, warn};
+use tracing::{debug, trace};
 
 use super::DHTPutSuccess;
 
@@ -143,27 +143,8 @@ impl Client {
 				.context("Error dialing bootstrap peer")?;
 			self.add_address(peer, addr.clone()).await?;
 
-			// Re-form the bootstrap Udp protocol part of the multiaddress for the autonat bootstrap TCP listener
-			// check for UDP protocol parts
-			if let Some(Protocol::Udp(port)) = addr
-				.iter()
-				.find(|protocol| matches!(protocol, Protocol::Udp(_)))
-			{
-				// take the IP Protocol part
-				if let Some(ip_part) = addr
-					.iter()
-					.find(|protocol| matches!(protocol, Protocol::Ip4(_)))
-				{
-					// crate new address for the same IP with TCP Protocol
-					let addr = Multiaddr::empty()
-						.with(ip_part)
-						.with(Protocol::Tcp(port + 1));
-					// add this address as Autonat server
-					self.add_autonat_server(peer, addr).await?;
-				} else {
-					warn!("Unable to re-form autonat server multi-address.")
-				}
-			}
+			// add proper address for auto-nat servers
+			self.add_autonat_server(peer, autonat_address(addr)).await?;
 		}
 
 		self.command_sender
@@ -465,3 +446,29 @@ pub enum Command {
 		response_sender: oneshot::Sender<Result<()>>,
 	},
 }
+
+/// This utility function takes the Multiaddress as parameter and searches for
+/// the Protocol::Udp(_) part of it, and replaces it with a TCP variant, while
+/// shifting the port number by 1 up.
+///
+/// Used to setup a proper Multiaddress for AutoNat servers on TCP.
+fn autonat_address(addr: Multiaddr) -> Multiaddr {
+	let mut stacks = addr.iter().collect::<Vec<_>>();
+	// search for the first occurrence of Protocol::Udp, to replace it with Tcp variant
+	stacks.iter_mut().find_map(|protocol| {
+		if let Protocol::Udp(port) = protocol {
+			// replace the UDP variant, moving the port number 1 forward
+			*protocol = Protocol::Tcp(*port + 1);
+			Some(protocol)
+		} else {
+			None
+		}
+	});
+
+	let mut addr = Multiaddr::empty();
+	for stack in stacks {
+		addr.push(stack)
+	}
+
+	addr
+}