Enable support for creating a custom list with locations

3 files changed, 39 insertions, 10 deletions

diff --git a/mullvad-daemon/src/custom_list.rs b/mullvad-daemon/src/custom_list.rs
index ba662a2dac..06f5926a2c 100644
--- a/mullvad-daemon/src/custom_list.rs
+++ b/mullvad-daemon/src/custom_list.rs
@@ -1,19 +1,27 @@
 use crate::{Daemon, Error};
 use mullvad_relay_selector::SelectorConfig;
+use mullvad_types::relay_constraints::GeographicLocationConstraint;
 use mullvad_types::{
     constraints::Constraint,
     custom_list::{CustomList, Id},
     relay_constraints::{BridgeState, LocationConstraint, RelaySettings, ResolvedBridgeSettings},
 };
+use std::collections::BTreeSet;
 use talpid_types::net::TunnelType;
 
 impl Daemon {
     /// Create a new custom list.
     ///
     /// Returns an error if the name is not unique.
-    pub async fn create_custom_list(&mut self, name: String) -> Result<Id, crate::Error> {
-        let new_list = CustomList::new(name).map_err(crate::Error::CustomListError)?;
-        let id = new_list.id;
+    pub async fn create_custom_list(
+        &mut self,
+        name: String,
+        locations: BTreeSet<GeographicLocationConstraint>,
+    ) -> Result<Id, crate::Error> {
+        let mut new_list = CustomList::new(name).map_err(crate::Error::CustomListError)?;
+        new_list.append(locations);
+
+        let id = new_list.id();
 
         self.settings
             .try_update(|settings| settings.custom_lists.add(new_list))
@@ -57,7 +65,7 @@ impl Daemon {
     /// - there is no existing list with the same ID,
     /// - or the existing list has a different name.
     pub async fn update_custom_list(&mut self, new_list: CustomList) -> Result<(), Error> {
-        let list_id = new_list.id;
+        let list_id = new_list.id();
         let settings_changed = self
             .settings
             .try_update(|settings| settings.custom_lists.update(new_list))
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index f73ab6927c..5550ef3a8a 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -45,6 +45,7 @@ use mullvad_encrypted_dns_proxy::state::EncryptedDnsProxyState;
 use mullvad_relay_selector::{RelaySelector, SelectorConfig};
 #[cfg(target_os = "android")]
 use mullvad_types::account::{PlayPurchase, PlayPurchasePaymentToken};
+use mullvad_types::relay_constraints::GeographicLocationConstraint;
 #[cfg(any(windows, target_os = "android", target_os = "macos"))]
 use mullvad_types::settings::SplitApp;
 #[cfg(daita)]
@@ -70,6 +71,7 @@ use mullvad_types::{
 };
 use relay_list::{RelayListUpdater, RelayListUpdaterHandle, RELAYS_FILENAME};
 use settings::SettingsPersister;
+use std::collections::BTreeSet;
 #[cfg(any(windows, target_os = "android", target_os = "macos"))]
 use std::collections::HashSet;
 #[cfg(target_os = "android")]
@@ -310,7 +312,11 @@ pub enum DaemonCommand {
     /// Return a public key of the currently set wireguard private key, if there is one
     GetWireguardKey(ResponseTx<Option<PublicKey>, Error>),
     /// Create custom list
-    CreateCustomList(ResponseTx<mullvad_types::custom_list::Id, Error>, String),
+    CreateCustomList(
+        ResponseTx<mullvad_types::custom_list::Id, Error>,
+        String,
+        BTreeSet<GeographicLocationConstraint>,
+    ),
     /// Delete custom list
     DeleteCustomList(ResponseTx<(), Error>, mullvad_types::custom_list::Id),
     /// Update a custom list with a given id
@@ -1442,7 +1448,9 @@ impl Daemon {
             ResetSettings(tx) => self.on_reset_settings(tx).await,
             RotateWireguardKey(tx) => self.on_rotate_wireguard_key(tx),
             GetWireguardKey(tx) => self.on_get_wireguard_key(tx).await,
-            CreateCustomList(tx, name) => self.on_create_custom_list(tx, name).await,
+            CreateCustomList(tx, name, locations) => {
+                self.on_create_custom_list(tx, name, locations).await
+            }
             DeleteCustomList(tx, id) => self.on_delete_custom_list(tx, id).await,
             UpdateCustomList(tx, update) => self.on_update_custom_list(tx, update).await,
             ClearCustomLists(tx) => self.on_clear_custom_lists(tx).await,
@@ -2865,8 +2873,9 @@ impl Daemon {
         &mut self,
         tx: ResponseTx<mullvad_types::custom_list::Id, Error>,
         name: String,
+        locations: BTreeSet<GeographicLocationConstraint>,
     ) {
-        let result = self.create_custom_list(name).await;
+        let result = self.create_custom_list(name, locations).await;
         Self::oneshot_send(tx, result, "create_custom_list response");
     }
 
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index 5695683fcf..aef1d17c45 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -4,10 +4,12 @@ use futures::{
     StreamExt,
 };
 use mullvad_api::{rest::Error as RestError, StatusCode};
+use mullvad_management_interface::types::FromProtobufTypeError;
 use mullvad_management_interface::{
     types::{self, daemon_event, management_service_server::ManagementService},
     Code, Request, Response, ServerJoinHandle, Status,
 };
+use mullvad_types::relay_constraints::GeographicLocationConstraint;
 use mullvad_types::{
     account::AccountNumber,
     relay_constraints::{
@@ -20,6 +22,7 @@ use mullvad_types::{
     version,
     wireguard::{RotationInterval, RotationIntervalError},
 };
+use std::collections::BTreeSet;
 use std::{
     path::Path,
     str::FromStr,
@@ -663,13 +666,22 @@ impl ManagementService for ManagementServiceImpl {
     // Custom lists
     //
 
-    async fn create_custom_list(&self, request: Request<String>) -> ServiceResult<String> {
+    async fn create_custom_list(
+        &self,
+        request: Request<types::NewCustomList>,
+    ) -> ServiceResult<String> {
         log::debug!("create_custom_list");
+        let request = request.into_inner();
+        let locations = request
+            .locations
+            .into_iter()
+            .map(GeographicLocationConstraint::try_from)
+            .collect::<Result<BTreeSet<_>, FromProtobufTypeError>>()?;
         let (tx, rx) = oneshot::channel();
-        self.send_command_to_daemon(DaemonCommand::CreateCustomList(tx, request.into_inner()))?;
+        self.send_command_to_daemon(DaemonCommand::CreateCustomList(tx, request.name, locations))?;
         self.wait_for_result(rx)
             .await?
-            .map(|response| Response::new(response.to_string()))
+            .map(|id| Response::new(id.to_string()))
             .map_err(map_daemon_error)
     }