From 38804da23e2409989752caf36bce53fe32596b5d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= <jon.haggblad@gmail.com>
Date: Tue, 5 Nov 2024 15:42:47 +0100
Subject: [PATCH] Lookup gateway IP using vpn-api

---
 .../src/entries/gateway.rs                    | 13 ++++++++
 .../crates/nym-gateway-directory/src/error.rs |  3 ++
 .../src/gateway_client.rs                     | 30 ++++++++++++++++++-
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/nym-vpn-core/crates/nym-gateway-directory/src/entries/gateway.rs b/nym-vpn-core/crates/nym-gateway-directory/src/entries/gateway.rs
index 0915b30fe7..ad9559ffa6 100644
--- a/nym-vpn-core/crates/nym-gateway-directory/src/entries/gateway.rs
+++ b/nym-vpn-core/crates/nym-gateway-directory/src/entries/gateway.rs
@@ -64,6 +64,19 @@ impl Gateway {
         self.authenticator_address.is_some()
     }
 
+    pub fn host(&self) -> Option<&nym_topology::NetworkAddress> {
+        self.host.as_ref()
+    }
+
+    pub async fn lookup_ip(&self) -> Option<IpAddr> {
+        match self.host.clone()? {
+            nym_topology::NetworkAddress::IpAddr(ip) => Some(ip),
+            nym_topology::NetworkAddress::Hostname(hostname) => {
+                crate::helpers::try_resolve_hostname(&hostname).await.ok()
+            }
+        }
+    }
+
     pub fn clients_address_no_tls(&self) -> Option<String> {
         match (&self.host, &self.clients_ws_port) {
             (Some(host), Some(port)) => Some(format!("ws://{}:{}", host, port)),
diff --git a/nym-vpn-core/crates/nym-gateway-directory/src/error.rs b/nym-vpn-core/crates/nym-gateway-directory/src/error.rs
index bb18812582..3c5fecbc82 100644
--- a/nym-vpn-core/crates/nym-gateway-directory/src/error.rs
+++ b/nym-vpn-core/crates/nym-gateway-directory/src/error.rs
@@ -73,6 +73,9 @@ pub enum Error {
 
     #[error("gateway {0} doesn't have a description available")]
     NoGatewayDescriptionAvailable(String),
+
+    #[error("failed to lookup gateway ip for gateway {0}")]
+    FailedToLookupIp(String),
 }
 
 // Result type based on our error type
diff --git a/nym-vpn-core/crates/nym-gateway-directory/src/gateway_client.rs b/nym-vpn-core/crates/nym-gateway-directory/src/gateway_client.rs
index dbaadbae86..e4acc56bcb 100644
--- a/nym-vpn-core/crates/nym-gateway-directory/src/gateway_client.rs
+++ b/nym-vpn-core/crates/nym-gateway-directory/src/gateway_client.rs
@@ -173,7 +173,8 @@ impl GatewayClient {
         gateways.random_low_latency_gateway().await
     }
 
-    pub async fn lookup_gateway_ip(&self, gateway_identity: &str) -> Result<IpAddr> {
+    pub async fn lookup_gateway_ip_from_nym_api(&self, gateway_identity: &str) -> Result<IpAddr> {
+        info!("Fetching gateway ip from nym-api...");
         let mut ips = self
             .api_client
             .get_all_described_nodes()
@@ -262,6 +263,33 @@ impl GatewayClient {
             .map(GatewayList::into_vpn_gateways)
     }
 
+    pub async fn lookup_gateway_ip(&self, gateway_identity: &str) -> Result<IpAddr> {
+        if let Some(nym_vpn_api_client) = &self.nym_vpn_api_client {
+            info!("Fetching gateway ip from nym-vpn-api...");
+            let gateway = nym_vpn_api_client
+                .get_gateways(None)
+                .await?
+                .into_iter()
+                .find_map(|gw| {
+                    if gw.identity_key != gateway_identity {
+                        None
+                    } else {
+                        Gateway::try_from(gw)
+                            .inspect_err(|err| error!("Failed to parse gateway: {err}"))
+                            .ok()
+                    }
+                })
+                .ok_or_else(|| Error::RequestedGatewayIdNotFound(gateway_identity.to_string()))?;
+            gateway
+                .lookup_ip()
+                .await
+                .ok_or(Error::FailedToLookupIp(gateway_identity.to_string()))
+        } else {
+            warn!("OPERATING IN FALLBACK MODE WITHOUT NYM-VPN-API!");
+            self.lookup_gateway_ip_from_nym_api(gateway_identity).await
+        }
+    }
+
     pub async fn lookup_all_gateways(&self) -> Result<GatewayList> {
         if let Some(nym_vpn_api_client) = &self.nym_vpn_api_client {
             info!("Fetching all gateways from nym-vpn-api...");