Skip to main content

hyper_connector_mullvad/
location.rs

1use std::{fmt, str::FromStr};
2
3use crate::Error;
4
5/// A geographic restriction applied to discovered Mullvad relays.
6#[derive(Clone, Debug, Eq, PartialEq)]
7pub enum LocationFilter {
8    /// An ISO 3166-1 alpha-2 country code.
9    Country(String),
10    /// A continent containing one or more relay countries.
11    Continent(Continent),
12}
13
14/// A supported continent/region.
15#[derive(Clone, Copy, Debug, Eq, PartialEq)]
16pub enum Continent {
17    Africa,
18    Asia,
19    Europe,
20    NorthAmerica,
21    SouthAmerica,
22    Oceania,
23}
24
25impl LocationFilter {
26    /// Parse a country code, English country name, or continent name.
27    pub fn parse(value: &str) -> Result<Self, Error> {
28        value.parse()
29    }
30
31    /// Construct a country filter from a code or English country name.
32    pub fn country(value: &str) -> Result<Self, Error> {
33        country_code(value)
34            .map(Self::Country)
35            .ok_or_else(|| Error::InvalidLocationFilter(value.trim().to_owned()))
36    }
37
38    /// Construct a continent filter.
39    pub fn continent(value: &str) -> Result<Self, Error> {
40        continent(value)
41            .map(Self::Continent)
42            .ok_or_else(|| Error::InvalidLocationFilter(value.trim().to_owned()))
43    }
44
45    pub(crate) fn matches(&self, code: &str) -> bool {
46        match self {
47            Self::Country(wanted) => wanted == code,
48            Self::Continent(c) => country_continent(code) == Some(*c),
49        }
50    }
51}
52
53impl FromStr for LocationFilter {
54    type Err = Error;
55    fn from_str(value: &str) -> Result<Self, Self::Err> {
56        match (continent(value), country_code(value)) {
57            (Some(_), Some(_)) => Err(Error::InvalidLocationFilter(value.trim().to_owned())),
58            (Some(c), None) => Ok(Self::Continent(c)),
59            (None, Some(code)) => Ok(Self::Country(code)),
60            (None, None) => Err(Error::InvalidLocationFilter(value.trim().to_owned())),
61        }
62    }
63}
64
65impl fmt::Display for Continent {
66    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
67        write!(f, "{:?}", self)
68    }
69}
70
71fn normalized(value: &str) -> String {
72    value
73        .trim()
74        .to_ascii_lowercase()
75        .replace([' ', '-', '_'], "")
76}
77
78fn continent(value: &str) -> Option<Continent> {
79    match normalized(value).as_str() {
80        "africa" => Some(Continent::Africa),
81        "asia" => Some(Continent::Asia),
82        "europe" => Some(Continent::Europe),
83        "northamerica" => Some(Continent::NorthAmerica),
84        "southamerica" => Some(Continent::SouthAmerica),
85        "oceania" | "australia" => Some(Continent::Oceania),
86        _ => None,
87    }
88}
89
90fn country_code(value: &str) -> Option<String> {
91    let n = normalized(value);
92    if n.len() == 2 && n.bytes().all(|b| b.is_ascii_alphabetic()) {
93        return Some(n);
94    }
95    let code = match n.as_str() {
96        "sweden" => "se",
97        "unitedstates" | "usa" => "us",
98        "unitedkingdom" | "greatbritain" => "gb",
99        "australia" => "au",
100        "austria" => "at",
101        "belgium" => "be",
102        "brazil" => "br",
103        "bulgaria" => "bg",
104        "canada" => "ca",
105        "chile" => "cl",
106        "colombia" => "co",
107        "croatia" => "hr",
108        "czechia" | "czechrepublic" => "cz",
109        "denmark" => "dk",
110        "estonia" => "ee",
111        "finland" => "fi",
112        "france" => "fr",
113        "germany" => "de",
114        "greece" => "gr",
115        "hongkong" => "hk",
116        "hungary" => "hu",
117        "ireland" => "ie",
118        "israel" => "il",
119        "italy" => "it",
120        "japan" => "jp",
121        "latvia" => "lv",
122        "luxembourg" => "lu",
123        "mexico" => "mx",
124        "netherlands" => "nl",
125        "newzealand" => "nz",
126        "norway" => "no",
127        "poland" => "pl",
128        "portugal" => "pt",
129        "romania" => "ro",
130        "serbia" => "rs",
131        "singapore" => "sg",
132        "slovakia" => "sk",
133        "slovenia" => "si",
134        "southafrica" => "za",
135        "southkorea" | "korea" => "kr",
136        "spain" => "es",
137        "switzerland" => "ch",
138        "taiwan" => "tw",
139        "turkey" | "turkiye" => "tr",
140        "unitedarabemirates" => "ae",
141        _ => return None,
142    };
143    Some(code.into())
144}
145
146fn country_continent(code: &str) -> Option<Continent> {
147    const EUROPE: &[&str] = &[
148        "al", "at", "be", "bg", "ch", "cz", "de", "dk", "ee", "es", "fi", "fr", "gb", "gr", "hr",
149        "hu", "ie", "is", "it", "lt", "lu", "lv", "md", "mk", "mt", "nl", "no", "pl", "pt", "ro",
150        "rs", "se", "si", "sk", "ua",
151    ];
152    const ASIA: &[&str] = &[
153        "ae", "hk", "id", "il", "in", "jp", "kr", "my", "ph", "sg", "th", "tr", "tw", "vn",
154    ];
155    const NORTH_AMERICA: &[&str] = &["ca", "cr", "gt", "mx", "pa", "us"];
156    const SOUTH_AMERICA: &[&str] = &["ar", "bo", "br", "cl", "co", "ec", "pe", "py", "uy", "ve"];
157    const OCEANIA: &[&str] = &["au", "nz"];
158    const AFRICA: &[&str] = &["dz", "eg", "gh", "ke", "ma", "ng", "tn", "za"];
159    [
160        (EUROPE, Continent::Europe),
161        (ASIA, Continent::Asia),
162        (NORTH_AMERICA, Continent::NorthAmerica),
163        (SOUTH_AMERICA, Continent::SouthAmerica),
164        (OCEANIA, Continent::Oceania),
165        (AFRICA, Continent::Africa),
166    ]
167    .into_iter()
168    .find_map(|(codes, c)| codes.contains(&code).then_some(c))
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174    #[test]
175    fn parses_locations() {
176        assert_eq!(
177            LocationFilter::parse(" Sweden ").unwrap(),
178            LocationFilter::Country("se".into())
179        );
180        assert!(LocationFilter::parse("Europe").unwrap().matches("se"));
181        assert!(
182            LocationFilter::parse("North America")
183                .unwrap()
184                .matches("us")
185        );
186        assert!(LocationFilter::parse("Oceania").unwrap().matches("au"));
187        assert!(LocationFilter::parse("Australia").is_err());
188        assert!(LocationFilter::country("Australia").unwrap().matches("au"));
189        assert!(LocationFilter::parse("Atlantis").is_err());
190    }
191}