Also filter gateways on location

common/client-core/src/client/topology_control/geo_aware_provider.rs

@@ -1,14 +1,14 @@
 use std::{collections::HashMap, fmt};
 
 use log::{debug, error, info};
-use nym_explorer_api_requests::PrettyDetailedMixNodeBond;
+use nym_explorer_api_requests::{PrettyDetailedMixNodeBond, PrettyDetailedGatewayBond};
 use nym_network_defaults::var_names::EXPLORER_API;
 use nym_topology::{
     nym_topology_from_detailed,
     provider_trait::{async_trait, TopologyProvider},
     NymTopology,
 };
-use nym_validator_client::client::MixId;
+use nym_validator_client::client::{MixId, IdentityKey};
 use rand::{prelude::SliceRandom, thread_rng};
 use serde::{Deserialize, Serialize};
 use url::Url;
@@ -47,6 +47,25 @@ async fn fetch_mixnodes_from_explorer_api() -> Option<Vec<PrettyDetailedMixNodeB
         .ok()
 }
 
+// TODO: create a explorer-api-client
+async fn fetch_gateways_from_explorer_api() -> Option<Vec<PrettyDetailedGatewayBond>> {
+    let explorer_api_url = std::env::var(EXPLORER_API).ok()?;
+    let explorer_api_url = Url::parse(&explorer_api_url)
+        .ok()?
+        .join("v1/gateways")
+        .ok()?;
+
+    debug!("Fetching: {}", explorer_api_url);
+    reqwest_client()?
+        .get(explorer_api_url)
+        .send()
+        .await
+        .ok()?
+        .json::<Vec<PrettyDetailedGatewayBond>>()
+        .await
+        .ok()
+}
+
 #[derive(Copy, Clone, Hash, PartialEq, Eq, Serialize, Deserialize, Debug)]
 pub enum CountryGroup {
     Europe,
@@ -213,6 +232,22 @@ fn group_mixnodes_by_country_code(
         })
 }
 
+fn group_gateways_by_country_code(
+    gateways: Vec<PrettyDetailedGatewayBond>,
+) -> HashMap<CountryGroup, Vec<IdentityKey>> {
+    gateways
+        .into_iter()
+        .fold(HashMap::<CountryGroup, Vec<IdentityKey>>::new(), |mut acc, g| {
+            if let Some(ref location) = g.location {
+                let country_code = location.two_letter_iso_country_code.clone();
+                let group_code = CountryGroup::new(country_code.as_str());
+                let gateways = acc.entry(group_code).or_insert_with(Vec::new);
+                gateways.push(g.gateway.identity_key)
+            }
+            acc
+        })
+}
+
 fn log_mixnode_distribution(mixnodes: &HashMap<CountryGroup, Vec<MixId>>) {
     let mixnode_distribution = mixnodes
         .iter()
@@ -222,6 +257,15 @@ fn log_mixnode_distribution(mixnodes: &HashMap<CountryGroup, Vec<MixId>>) {
     debug!("Mixnode distribution - {}", mixnode_distribution);
 }
 
+fn log_gateway_distribution(gateways: &HashMap<CountryGroup, Vec<IdentityKey>>) {
+    let gateway_distribution = gateways
+        .iter()
+        .map(|(k, v)| format!("{}: {}", k, v.len()))
+        .collect::<Vec<_>>()
+        .join(", ");
+    debug!("Gateway distribution - {}", gateway_distribution);
+}
+
 fn check_layer_integrity(topology: NymTopology) -> Result<(), ()> {
     let mixes = topology.mixes();
     if mixes.keys().len() < 3 {
@@ -294,6 +338,12 @@ impl GeoAwareTopologyProvider {
             return None;
         };
 
+        debug!("Fetching gateways from explorer-api...");
+        let Some(gateways_from_explorer_api) = fetch_gateways_from_explorer_api().await else {
+            error!("failed to get mixnodes from explorer-api");
+            return None;
+        };
+
         // Partition mixnodes_from_explorer_api according to the value of
         // two_letter_iso_country_code.
         // NOTE: we construct the full distribution here, but only use the one we're interested in.
@@ -302,16 +352,29 @@ impl GeoAwareTopologyProvider {
         let mixnode_distribution = group_mixnodes_by_country_code(mixnodes_from_explorer_api);
         log_mixnode_distribution(&mixnode_distribution);
 
+        let gateway_distribution = group_gateways_by_country_code(gateways_from_explorer_api);
+        log_gateway_distribution(&gateway_distribution);
+
         let Some(filtered_mixnode_ids) = mixnode_distribution.get(&self.filter_on) else {
             error!("no mixnodes found for: {}", self.filter_on);
             return None;
         };
 
+        let Some(filtered_gateway_ids) = gateway_distribution.get(&self.filter_on) else {
+            error!("no gateways found for: {}", self.filter_on);
+            return None;
+        };
+
         let mixnodes = mixnodes
             .into_iter()
             .filter(|m| filtered_mixnode_ids.contains(&m.mix_id()))
             .collect::<Vec<_>>();
 
+        let gateways = gateways
+            .into_iter()
+            .filter(|g| filtered_gateway_ids.contains(g.identity()))
+            .collect::<Vec<_>>();
+
         let topology = nym_topology_from_detailed(mixnodes, gateways)
             .filter_system_version(&self.client_version);