Implement FromStr on GeographicLocationConstraint

2 files changed, 53 insertions, 1 deletions

diff --git a/mullvad-cli/src/cmds/relay_constraints.rs b/mullvad-cli/src/cmds/relay_constraints.rs
index 97555997fc..72b3bdb88e 100644
--- a/mullvad-cli/src/cmds/relay_constraints.rs
+++ b/mullvad-cli/src/cmds/relay_constraints.rs
@@ -27,7 +27,6 @@ impl From<LocationArgs> for Constraint<GeographicLocationConstraint> {
             (country, Some(city), Some(hostname)) => {
                 GeographicLocationConstraint::Hostname(country, city, hostname)
             }
-
             _ => unreachable!("invalid location arguments"),
         })
     }
diff --git a/mullvad-types/src/relay_constraints.rs b/mullvad-types/src/relay_constraints.rs
index 4be7e25a4b..e6eaa7ca55 100644
--- a/mullvad-types/src/relay_constraints.rs
+++ b/mullvad-types/src/relay_constraints.rs
@@ -174,6 +174,12 @@ pub enum GeographicLocationConstraint {
     Hostname(CountryCode, CityCode, Hostname),
 }
 
+#[derive(thiserror::Error, Debug)]
+#[error("Failed to parse {input} into a geographic location constraint")]
+pub struct ParseGeoLocationError {
+    input: String,
+}
+
 impl GeographicLocationConstraint {
     /// Create a new [`GeographicLocationConstraint`] given a country.
     pub fn country(country: impl Into<String>) -> Self {
@@ -227,6 +233,27 @@ impl Match<Relay> for GeographicLocationConstraint {
     }
 }
 
+impl FromStr for GeographicLocationConstraint {
+    type Err = ParseGeoLocationError;
+
+    // TODO: Implement for country and city as well?
+    fn from_str(input: &str) -> Result<Self, Self::Err> {
+        // A host name, such as "se-got-wg-101" maps to
+        // Country: se
+        // City: got
+        // hostname: se-got-wg-101
+        let x = input.split("-").collect::<Vec<_>>();
+        match x[..] {
+            [country] => Ok(GeographicLocationConstraint::country(country)),
+            [country, city] => Ok(GeographicLocationConstraint::city(country, city)),
+            [country, city, ..] => Ok(GeographicLocationConstraint::hostname(country, city, input)),
+            _ => Err(ParseGeoLocationError {
+                input: input.to_string(),
+            }),
+        }
+    }
+}
+
 /// Limits the set of servers to choose based on ownership.
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Deserialize, Serialize)]
 #[cfg_attr(feature = "clap", derive(clap::ValueEnum))]
@@ -677,3 +704,29 @@ impl RelayOverride {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parse_hostname() {
+        // Parse a country
+        assert_eq!(
+            "se".parse::<GeographicLocationConstraint>().unwrap(),
+            GeographicLocationConstraint::country("se")
+        );
+        // Parse a city
+        assert_eq!(
+            "se-got".parse::<GeographicLocationConstraint>().unwrap(),
+            GeographicLocationConstraint::city("se", "got")
+        );
+        // Parse a hostname
+        assert_eq!(
+            "se-got-wg-101"
+                .parse::<GeographicLocationConstraint>()
+                .unwrap(),
+            GeographicLocationConstraint::hostname("se", "got", "se-got-wg-101")
+        );
+    }
+}