Whitelist relays in an end to end test framwork config

mullvad-cli/src/cmds/relay.rs

@@ -958,7 +958,7 @@ pub async fn resolve_location_constraint(
     } else {
         // The Constraint was not a relay, assuming it to be a location
         let location_constraint: Constraint<GeographicLocationConstraint> =
-            Constraint::from(location_constraint_args);
+            Constraint::try_from(location_constraint_args)?;
 
         // If the location constraint was not "any", then validate the country/city
         if let Constraint::Only(constraint) = &location_constraint {

mullvad-cli/src/cmds/relay_constraints.rs

@@ -15,38 +15,51 @@ pub struct LocationArgs {
     pub hostname: Option<Hostname>,
 }
 
-impl From<LocationArgs> for Constraint<GeographicLocationConstraint> {
-    fn from(value: LocationArgs) -> Self {
-        if value.country.eq_ignore_ascii_case("any") {
-            return Constraint::Any;
+#[derive(thiserror::Error, Debug)]
+pub enum Error {
+    #[error("Failed to parse location constraint from input: TODO")]
+    Parse,
+}
+
+impl TryFrom<LocationArgs> for GeographicLocationConstraint {
+    type Error = Error;
+
+    fn try_from(value: LocationArgs) -> Result<Self, Self::Error> {
+        match (value.country, value.city, value.hostname) {
+            (country, None, None) => Ok(GeographicLocationConstraint::Country(country)),
+            (country, Some(city), None) => Ok(GeographicLocationConstraint::City(country, city)),
+            (country, Some(city), Some(hostname)) => Ok(GeographicLocationConstraint::Hostname(
+                country, city, hostname,
+            )),
+            _ => Err(Error::Parse),
+            //_ => unreachable!("invalid location arguments"),
         }
+    }
+}
 
-        Constraint::Only(match (value.country, value.city, value.hostname) {
-            (country, None, None) => GeographicLocationConstraint::Country(country),
-            (country, Some(city), None) => GeographicLocationConstraint::City(country, city),
-            (country, Some(city), Some(hostname)) => {
-                GeographicLocationConstraint::Hostname(country, city, hostname)
-            }
+impl TryFrom<LocationArgs> for LocationConstraint {
+    type Error = Error;
 
-            _ => unreachable!("invalid location arguments"),
-        })
+    fn try_from(value: LocationArgs) -> Result<Self, Self::Error> {
+        GeographicLocationConstraint::try_from(value).map(LocationConstraint::from)
     }
 }
 
-impl From<LocationArgs> for Constraint<LocationConstraint> {
-    fn from(value: LocationArgs) -> Self {
+impl TryFrom<LocationArgs> for Constraint<GeographicLocationConstraint> {
+    type Error = Error;
+
+    fn try_from(value: LocationArgs) -> Result<Self, Self::Error> {
         if value.country.eq_ignore_ascii_case("any") {
-            return Constraint::Any;
+            return Ok(Constraint::Any);
         }
+        GeographicLocationConstraint::try_from(value).map(Constraint::Only)
+    }
+}
+
+impl TryFrom<LocationArgs> for Constraint<LocationConstraint> {
+    type Error = Error;
 
-        let location = match (value.country, value.city, value.hostname) {
-            (country, None, None) => GeographicLocationConstraint::Country(country),
-            (country, Some(city), None) => GeographicLocationConstraint::City(country, city),
-            (country, Some(city), Some(hostname)) => {
-                GeographicLocationConstraint::Hostname(country, city, hostname)
-            }
-            _ => unreachable!("invalid location arguments"),
-        };
-        Constraint::Only(LocationConstraint::Location(location))
+    fn try_from(value: LocationArgs) -> Result<Self, Self::Error> {
+        LocationConstraint::try_from(value).map(Constraint::Only)
     }
 }

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")
+        );
+    }
+}

test/test-manager/Cargo.toml

@@ -32,6 +32,7 @@ async-trait = { workspace = true }
 uuid = "1.3"
 dirs = "5.0.1"
 scopeguard = "1.2"
+glob = "0.3"
 
 serde = { workspace = true }
 serde_json = { workspace = true }

test/test-manager/src/config/error.rs

@@ -0,0 +1,15 @@
+use std::io;
+
+#[derive(thiserror::Error, Debug)]
+pub enum Error {
+    #[error("Could not find config dir")]
+    FindConfigDir,
+    #[error("Could not create config dir")]
+    CreateConfigDir(#[source] io::Error),
+    #[error("Failed to read config")]
+    Read(#[source] io::Error),
+    #[error("Failed to parse config")]
+    InvalidConfig(#[from] serde_json::Error),
+    #[error("Failed to write config")]
+    Write(#[source] io::Error),
+}

test/test-manager/src/config/io.rs

@@ -0,0 +1,80 @@
+//! See [ConfigFile].
+
+use std::io;
+use std::ops::Deref;
+use std::path::{Path, PathBuf};
+
+use super::{Config, Error};
+
+/// On-disk representation of [Config].
+pub struct ConfigFile {
+    path: PathBuf,
+    config: Config,
+}
+
+impl ConfigFile {
+    /// Make config changes and save them to disk
+    pub async fn edit(&mut self, edit: impl FnOnce(&mut Config)) -> Result<(), Error> {
+        Self::ensure_config_dir().await?;
+        edit(&mut self.config);
+        self.config_save().await
+    }
+
+    /// Make config changes and save them to disk
+    pub async fn load_or_default() -> Result<Self, Error> {
+        let path = Self::get_config_path()?;
+        let config = Self::config_load_or_default(&path).await?;
+        let config_file = Self { path, config };
+        Ok(config_file)
+    }
+
+    async fn config_load_or_default<P: AsRef<Path>>(path: P) -> Result<Config, Error> {
+        Self::config_load(path).await.or_else(|error| match error {
+            Error::Read(ref io_err) if io_err.kind() == io::ErrorKind::NotFound => {
+                log::trace!("Failed to read config file");
+                Ok(Config::default())
+            }
+            error => Err(error),
+        })
+    }
+
+    async fn config_load<P: AsRef<Path>>(path: P) -> Result<Config, Error> {
+        let data = tokio::fs::read(path).await.map_err(Error::Read)?;
+        serde_json::from_slice(&data).map_err(Error::InvalidConfig)
+    }
+
+    async fn config_save(&self) -> Result<(), Error> {
+        let data = serde_json::to_vec_pretty(&self.config).unwrap();
+        tokio::fs::write(&self.path, &data)
+            .await
+            .map_err(Error::Write)
+    }
+
+    /// Get configuration file path
+    fn get_config_path() -> Result<PathBuf, Error> {
+        Ok(Self::get_config_dir()?.join("config.json"))
+    }
+
+    /// Get configuration file directory
+    fn get_config_dir() -> Result<PathBuf, Error> {
+        let dir = dirs::config_dir()
+            .ok_or(Error::FindConfigDir)?
+            .join("mullvad-test");
+        Ok(dir)
+    }
+
+    /// Create configuration file directory if it does not exist
+    async fn ensure_config_dir() -> Result<(), Error> {
+        tokio::fs::create_dir_all(Self::get_config_dir()?)
+            .await
+            .map_err(Error::CreateConfigDir)
+    }
+}
+
+impl Deref for ConfigFile {
+    type Target = Config;
+
+    fn deref(&self) -> &Self::Target {
+        &self.config
+    }
+}