From 3924138c49c8edfbd7a78eb247d4a081e9084271 Mon Sep 17 00:00:00 2001 From: neequ57 Date: Wed, 26 Feb 2025 14:24:38 +0000 Subject: [PATCH] Fix geolocation feature after recent refactor --- .../src/routing_table/route_spec_store/mod.rs | 131 +++++++++--------- veilid-core/src/veilid_api/tests/fixtures.rs | 2 +- .../src/veilid_api/types/country_code.rs | 70 +++------- 3 files changed, 84 insertions(+), 119 deletions(-) diff --git a/veilid-core/src/routing_table/route_spec_store/mod.rs b/veilid-core/src/routing_table/route_spec_store/mod.rs index 8e0b1f23..623a228f 100644 --- a/veilid-core/src/routing_table/route_spec_store/mod.rs +++ b/veilid-core/src/routing_table/route_spec_store/mod.rs @@ -235,6 +235,11 @@ impl RouteSpecStore { .relay_node(RoutingDomain::PublicInternet) .map(|nr| nr.locked(rti)); + #[cfg(feature = "geolocation")] + let country_code_denylist = self + .config() + .with(|config| config.network.privacy.country_code_denylist.clone()); + // Get list of all nodes, and sort them for selection let cur_ts = Timestamp::now(); let filter = Box::new( @@ -277,78 +282,66 @@ impl RouteSpecStore { // Exclude nodes from blacklisted countries #[cfg(feature = "geolocation")] - { - let country_code_denylist = self - .unlocked_inner - .routing_table - .config - .get() - .network - .privacy - .country_code_denylist - .clone(); + if !country_code_denylist.is_empty() { + let geolocation_info = + sni.get_geolocation_info(RoutingDomain::PublicInternet); - if !country_code_denylist.is_empty() { - let geolocation_info = - sni.get_geolocation_info(RoutingDomain::PublicInternet); + // Since denylist is used, consider nodes with unknown countries to be automatically + // excluded as well + if geolocation_info.country_code().is_none() { + veilid_log!(self + debug "allocate_route_inner: skipping node {} from unknown country", + e.best_node_id() + ); + return false; + } + // Same thing applies to relays used by the node + if geolocation_info + .relay_country_codes() + .iter() + .any(Option::is_none) + { + veilid_log!(self + debug "allocate_route_inner: skipping node {} using relay from unknown country", + e.best_node_id() + ); + return false; + } - // Since denylist is used, consider nodes with unknown countries to be automatically - // excluded as well - if geolocation_info.country_code().is_none() { - veilid_log!(self - debug "allocate_route_inner: skipping node {} from unknown country", - e.best_node_id() - ); - return false; - } - // Same thing applies to relays used by the node - if geolocation_info - .relay_country_codes() - .iter() - .any(Option::is_none) - { - veilid_log!(self - debug "allocate_route_inner: skipping node {} using relay from unknown country", - e.best_node_id() - ); - return false; - } + // Ensure that node is not excluded + // Safe to unwrap here, checked above + if country_code_denylist.contains(&geolocation_info.country_code().unwrap()) + { + veilid_log!(self + debug "allocate_route_inner: skipping node {} from excluded country {}", + e.best_node_id(), + geolocation_info.country_code().unwrap() + ); + return false; + } - // Ensure that node is not excluded - // Safe to unwrap here, checked above - if country_code_denylist.contains(&geolocation_info.country_code().unwrap()) - { - veilid_log!(self - debug "allocate_route_inner: skipping node {} from excluded country {}", - e.best_node_id(), - geolocation_info.country_code().unwrap() - ); - return false; - } - - // Ensure that node relays are not excluded - // Safe to unwrap here, checked above - if geolocation_info - .relay_country_codes() - .iter() - .cloned() - .map(Option::unwrap) - .any(|cc| country_code_denylist.contains(&cc)) - { - veilid_log!(self - debug "allocate_route_inner: skipping node {} using relay from excluded country {:?}", - e.best_node_id(), - geolocation_info - .relay_country_codes() - .iter() - .cloned() - .map(Option::unwrap) - .filter(|cc| country_code_denylist.contains(&cc)) - .next() - .unwrap() - ); - return false; - } + // Ensure that node relays are not excluded + // Safe to unwrap here, checked above + if geolocation_info + .relay_country_codes() + .iter() + .cloned() + .map(Option::unwrap) + .any(|cc| country_code_denylist.contains(&cc)) + { + veilid_log!(self + debug "allocate_route_inner: skipping node {} using relay from excluded country {:?}", + e.best_node_id(), + geolocation_info + .relay_country_codes() + .iter() + .cloned() + .map(Option::unwrap) + .filter(|cc| country_code_denylist.contains(&cc)) + .next() + .unwrap() + ); + return false; } } diff --git a/veilid-core/src/veilid_api/tests/fixtures.rs b/veilid-core/src/veilid_api/tests/fixtures.rs index a71075d1..d2e1e454 100644 --- a/veilid-core/src/veilid_api/tests/fixtures.rs +++ b/veilid-core/src/veilid_api/tests/fixtures.rs @@ -239,7 +239,7 @@ pub fn fix_veilidconfiginner() -> VeilidConfigInner { }, #[cfg(feature = "geolocation")] privacy: VeilidConfigPrivacy { - country_code_denylist: vec![CountryCode([b'N', b'Z'])], + country_code_denylist: vec![CountryCode::from_str("NZ").unwrap()], }, #[cfg(feature = "virtual-network")] virtual_network: VeilidConfigVirtualNetwork { diff --git a/veilid-core/src/veilid_api/types/country_code.rs b/veilid-core/src/veilid_api/types/country_code.rs index db6f739b..164d4310 100644 --- a/veilid-core/src/veilid_api/types/country_code.rs +++ b/veilid-core/src/veilid_api/types/country_code.rs @@ -1,17 +1,12 @@ use super::*; -use std::hash::{Hash, Hasher}; -/// Two-letter country code. Case-insensitive when comparing. -#[derive(Copy, Default, Clone, Ord, Eq, Serialize, Deserialize, JsonSchema)] +/// Two-letter country code. Stored in upper case internally. +#[derive( + Copy, Default, Clone, Hash, PartialOrd, Ord, PartialEq, Eq, Serialize, Deserialize, JsonSchema, +)] #[serde(try_from = "String")] #[serde(into = "String")] -pub struct CountryCode(pub [u8; 2]); - -impl From<[u8; 2]> for CountryCode { - fn from(b: [u8; 2]) -> Self { - Self(b) - } -} +pub struct CountryCode([u8; 2]); impl From for String { fn from(u: CountryCode) -> Self { @@ -22,7 +17,18 @@ impl From for String { impl TryFrom<&[u8]> for CountryCode { type Error = VeilidAPIError; fn try_from(b: &[u8]) -> Result { - Ok(Self(b.try_into().map_err(VeilidAPIError::generic)?)) + let cc: [u8; 2] = b.try_into().map_err(VeilidAPIError::generic)?; + + if !cc[0].is_ascii_alphabetic() || !cc[1].is_ascii_alphabetic() { + return Err(VeilidAPIError::generic( + "country code must only contain alphabetic chars", + )); + } + + Ok(Self([ + cc[0].to_ascii_uppercase(), + cc[1].to_ascii_uppercase(), + ])) } } @@ -35,55 +41,21 @@ impl TryFrom for CountryCode { impl fmt::Display for CountryCode { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> { - write!(f, "{}", String::from_utf8_lossy(&self.0)) + // self.0 is guaranteed to be a valid ASCII string, checked in Self::try_from(&[u8]) + write!(f, "{}{}", self.0[0] as char, self.0[1] as char) } } impl fmt::Debug for CountryCode { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> { - write!(f, "{}", String::from_utf8_lossy(&self.0)) + ::fmt(self, f) } } impl FromStr for CountryCode { type Err = VeilidAPIError; fn from_str(s: &str) -> Result { - Ok(Self( - s.as_bytes().try_into().map_err(VeilidAPIError::generic)?, - )) - } -} - -impl Hash for CountryCode { - fn hash(&self, state: &mut H) { - let this = [ - self.0[0].to_ascii_uppercase(), - self.0[1].to_ascii_uppercase(), - ]; - - state.write(&this[..]); - } -} - -impl PartialEq for CountryCode { - fn eq(&self, other: &Self) -> bool { - self.0[0].to_ascii_uppercase() == other.0[0].to_ascii_uppercase() - && self.0[1].to_ascii_uppercase() == other.0[1].to_ascii_uppercase() - } -} - -impl PartialOrd for CountryCode { - fn partial_cmp(&self, other: &Self) -> Option { - let this = [ - self.0[0].to_ascii_uppercase(), - self.0[1].to_ascii_uppercase(), - ]; - let other = [ - other.0[0].to_ascii_uppercase(), - other.0[1].to_ascii_uppercase(), - ]; - - this.partial_cmp(&other) + Ok(Self::try_from(s.as_bytes())?) } }