Refactor custom list implementation

5 files changed, 39 insertions, 180 deletions

diff --git a/mullvad-management-interface/src/client.rs b/mullvad-management-interface/src/client.rs
index e9acad1337..a1ddc5e39a 100644
--- a/mullvad-management-interface/src/client.rs
+++ b/mullvad-management-interface/src/client.rs
@@ -4,7 +4,7 @@ use crate::types;
 use futures::{Stream, StreamExt};
 use mullvad_types::{
     account::{AccountData, AccountToken, VoucherSubmission},
-    custom_list::{CustomList, CustomListLocationUpdate},
+    custom_list::{CustomList, Id},
     device::{Device, DeviceEvent, DeviceId, DeviceState, RemoveDeviceEvent},
     location::GeoIpLocation,
     relay_constraints::{BridgeSettings, BridgeState, ObfuscationSettings, RelaySettingsUpdate},
@@ -16,6 +16,7 @@ use mullvad_types::{
 };
 #[cfg(target_os = "windows")]
 use std::path::Path;
+use std::str::FromStr;
 #[cfg(target_os = "windows")]
 use talpid_types::split_tunnel::ExcludedProcess;
 use tonic::{Code, Status};
@@ -430,60 +431,27 @@ impl MullvadProxyClient {
         PublicKey::try_from(key).map_err(Error::InvalidResponse)
     }
 
-    pub async fn list_custom_lists(&mut self) -> Result<Vec<CustomList>> {
-        let result = self
-            .0
-            .list_custom_lists(())
-            .await
-            .map_err(map_custom_list_error)?
-            .into_inner()
-            .try_into()
-            .map_err(Error::InvalidResponse)?;
-        Ok(result)
-    }
-
-    pub async fn get_custom_list(&mut self, name: String) -> Result<CustomList> {
-        let result = self
+    pub async fn create_custom_list(&mut self, name: String) -> Result<Id> {
+        let id = self
             .0
-            .get_custom_list(name)
-            .await
-            .map_err(map_custom_list_error)?
-            .into_inner()
-            .try_into()
-            .map_err(Error::InvalidResponse)?;
-        Ok(result)
-    }
-
-    pub async fn create_custom_list(&mut self, name: String) -> Result<()> {
-        self.0
             .create_custom_list(name)
             .await
-            .map_err(map_custom_list_error)?;
-        Ok(())
-    }
-
-    pub async fn delete_custom_list(&mut self, name: String) -> Result<()> {
-        self.0
-            .delete_custom_list(name)
-            .await
-            .map_err(map_custom_list_error)?;
-        Ok(())
+            .map_err(map_custom_list_error)?
+            .into_inner();
+        Id::from_str(&id).map_err(|_| Error::CustomListListNotFound)
     }
 
-    pub async fn update_custom_list_location(
-        &mut self,
-        custom_list_update: CustomListLocationUpdate,
-    ) -> Result<()> {
+    pub async fn delete_custom_list(&mut self, id: String) -> Result<()> {
         self.0
-            .update_custom_list_location(types::CustomListLocationUpdate::from(custom_list_update))
+            .delete_custom_list(id)
             .await
             .map_err(map_custom_list_error)?;
         Ok(())
     }
 
-    pub async fn rename_custom_list(&mut self, name: String, new_name: String) -> Result<()> {
+    pub async fn update_custom_list(&mut self, custom_list: CustomList) -> Result<()> {
         self.0
-            .rename_custom_list(types::CustomListRename::from((name, new_name)))
+            .update_custom_list(types::CustomList::from(custom_list))
             .await
             .map_err(map_custom_list_error)?;
         Ok(())
@@ -605,26 +573,19 @@ fn map_location_error(status: Status) -> Error {
 fn map_custom_list_error(status: Status) -> Error {
     match status.code() {
         Code::NotFound => {
-            let details = status.details();
-            if details == crate::CUSTOM_LIST_LOCATION_NOT_FOUND_DETAILS {
-                Error::LocationNotFoundInCustomlist
-            } else if details == crate::CUSTOM_LIST_LIST_NOT_FOUND_DETAILS {
+            if status.details() == crate::CUSTOM_LIST_LIST_NOT_FOUND_DETAILS {
                 Error::CustomListListNotFound
             } else {
                 Error::Rpc(status)
             }
         }
         Code::AlreadyExists => {
-            let details = status.details();
-            if details == crate::CUSTOM_LIST_LOCATION_EXISTS_DETAILS {
-                Error::LocationExistsInCustomList
-            } else if details == crate::CUSTOM_LIST_LIST_EXISTS_DETAILS {
+            if status.details() == crate::CUSTOM_LIST_LIST_EXISTS_DETAILS {
                 Error::CustomListExists
             } else {
                 Error::Rpc(status)
             }
         }
-        Code::InvalidArgument => Error::CustomListCannotAddOrRemoveAny,
         _other => Error::Rpc(status),
     }
 }
diff --git a/mullvad-management-interface/src/lib.rs b/mullvad-management-interface/src/lib.rs
index b2ca7a4079..cf1a798878 100644
--- a/mullvad-management-interface/src/lib.rs
+++ b/mullvad-management-interface/src/lib.rs
@@ -26,9 +26,7 @@ use once_cell::sync::Lazy;
 static MULLVAD_MANAGEMENT_SOCKET_GROUP: Lazy<Option<String>> =
     Lazy::new(|| env::var("MULLVAD_MANAGEMENT_SOCKET_GROUP").ok());
 
-pub const CUSTOM_LIST_LOCATION_NOT_FOUND_DETAILS: &[u8] = b"custom_list_location_not_found";
 pub const CUSTOM_LIST_LIST_NOT_FOUND_DETAILS: &[u8] = b"custom_list_list_not_found";
-pub const CUSTOM_LIST_LOCATION_EXISTS_DETAILS: &[u8] = b"custom_list_location_exists";
 pub const CUSTOM_LIST_LIST_EXISTS_DETAILS: &[u8] = b"custom_list_list_exists";
 
 #[derive(err_derive::Error, Debug)]
@@ -100,9 +98,6 @@ pub enum Error {
     #[error(display = "A custom list with that name does not exist")]
     CustomListListNotFound,
 
-    #[error(display = "Can not add or remove 'any' to or from a custom list")]
-    CustomListCannotAddOrRemoveAny,
-
     #[error(display = "Location already exists in the custom list")]
     LocationExistsInCustomList,
 
diff --git a/mullvad-management-interface/src/types/conversions/custom_list.rs b/mullvad-management-interface/src/types/conversions/custom_list.rs
index 63bb73652b..62799168f7 100644
--- a/mullvad-management-interface/src/types/conversions/custom_list.rs
+++ b/mullvad-management-interface/src/types/conversions/custom_list.rs
@@ -1,30 +1,15 @@
-use crate::types::{proto, FromProtobufTypeError};
-use mullvad_types::{custom_list::Id, relay_constraints::GeographicLocationConstraint};
-use proto::RelayLocation;
-
-impl From<(String, String)> for proto::CustomListRename {
-    fn from(names: (String, String)) -> Self {
-        proto::CustomListRename {
-            name: names.0,
-            new_name: names.1,
-        }
-    }
-}
+use std::{collections::BTreeSet, str::FromStr};
 
-impl From<proto::CustomListRename> for (String, String) {
-    fn from(names: proto::CustomListRename) -> Self {
-        (names.name, names.new_name)
-    }
-}
+use crate::types::{proto, FromProtobufTypeError};
+use mullvad_types::{
+    custom_list::{CustomList, Id},
+    relay_constraints::GeographicLocationConstraint,
+};
 
-impl From<&mullvad_types::custom_list::CustomListsSettings> for proto::CustomListSettings {
-    fn from(settings: &mullvad_types::custom_list::CustomListsSettings) -> Self {
+impl From<mullvad_types::custom_list::CustomListsSettings> for proto::CustomListSettings {
+    fn from(settings: mullvad_types::custom_list::CustomListsSettings) -> Self {
         Self {
-            custom_lists: settings
-                .custom_lists
-                .iter()
-                .map(|custom_list| proto::CustomList::from(custom_list.clone()))
-                .collect(),
+            custom_lists: settings.into_iter().map(proto::CustomList::from).collect(),
         }
     }
 }
@@ -33,72 +18,13 @@ impl TryFrom<proto::CustomListSettings> for mullvad_types::custom_list::CustomLi
     type Error = FromProtobufTypeError;
 
     fn try_from(settings: proto::CustomListSettings) -> Result<Self, Self::Error> {
-        Ok(Self {
-            custom_lists: settings
+        Ok(Self::from(
+            settings
                 .custom_lists
                 .into_iter()
                 .map(mullvad_types::custom_list::CustomList::try_from)
-                .collect::<Result<Vec<_>, _>>()?,
-        })
-    }
-}
-
-impl From<mullvad_types::custom_list::CustomListLocationUpdate>
-    for proto::CustomListLocationUpdate
-{
-    fn from(custom_list: mullvad_types::custom_list::CustomListLocationUpdate) -> Self {
-        use mullvad_types::relay_constraints::Constraint;
-        match custom_list {
-            mullvad_types::custom_list::CustomListLocationUpdate::Add { name, location } => {
-                let location = match location {
-                    Constraint::Any => None,
-                    Constraint::Only(location) => Some(RelayLocation::from(location)),
-                };
-                Self {
-                    state: i32::from(proto::custom_list_location_update::State::Add),
-                    name,
-                    location,
-                }
-            }
-            mullvad_types::custom_list::CustomListLocationUpdate::Remove { name, location } => {
-                let location = match location {
-                    Constraint::Any => None,
-                    Constraint::Only(location) => Some(RelayLocation::from(location)),
-                };
-                Self {
-                    state: i32::from(proto::custom_list_location_update::State::Remove),
-                    name,
-                    location,
-                }
-            }
-        }
-    }
-}
-
-impl TryFrom<proto::CustomListLocationUpdate>
-    for mullvad_types::custom_list::CustomListLocationUpdate
-{
-    type Error = FromProtobufTypeError;
-
-    fn try_from(custom_list: proto::CustomListLocationUpdate) -> Result<Self, Self::Error> {
-        use mullvad_types::relay_constraints::Constraint;
-        let location: Constraint<GeographicLocationConstraint> =
-            Constraint::<GeographicLocationConstraint>::from(
-                custom_list
-                    .location
-                    .ok_or(FromProtobufTypeError::InvalidArgument("missing location"))?,
-            );
-        match proto::custom_list_location_update::State::try_from(custom_list.state) {
-            Ok(proto::custom_list_location_update::State::Add) => Ok(Self::Add {
-                name: custom_list.name,
-                location,
-            }),
-            Ok(proto::custom_list_location_update::State::Remove) => Ok(Self::Remove {
-                name: custom_list.name,
-                location,
-            }),
-            Err(_) => Err(FromProtobufTypeError::InvalidArgument("incorrect state")),
-        }
+                .collect::<Result<Vec<CustomList>, _>>()?,
+        ))
     }
 }
 
@@ -121,18 +47,14 @@ impl TryFrom<proto::CustomList> for mullvad_types::custom_list::CustomList {
     type Error = FromProtobufTypeError;
 
     fn try_from(custom_list: proto::CustomList) -> Result<Self, Self::Error> {
-        let locations: Result<Vec<GeographicLocationConstraint>, _> = custom_list
+        let locations = custom_list
             .locations
             .into_iter()
             .map(GeographicLocationConstraint::try_from)
-            .collect();
-        let locations = locations.map_err(|_| {
-            FromProtobufTypeError::InvalidArgument("Could not convert custom list from proto")
-        })?;
+            .collect::<Result<BTreeSet<_>, Self::Error>>()?;
         Ok(Self {
-            id: Id::try_from(custom_list.id.as_str()).map_err(|_| {
-                FromProtobufTypeError::InvalidArgument("Id could not be parsed to a uuid")
-            })?,
+            id: Id::from_str(&custom_list.id)
+                .map_err(|_| FromProtobufTypeError::InvalidArgument("Invalid list ID"))?,
             name: custom_list.name,
             locations,
         })
@@ -149,7 +71,7 @@ impl TryFrom<proto::RelayLocation> for GeographicLocationConstraint {
             relay_location.hostname.as_ref(),
         ) {
             ("", ..) => Err(FromProtobufTypeError::InvalidArgument(
-                "Relay location formatted incorrectly",
+                "Invalid geographic relay location",
             )),
             (_country, "", "") => Ok(GeographicLocationConstraint::Country(
                 relay_location.country,
@@ -174,27 +96,3 @@ impl TryFrom<proto::RelayLocation> for GeographicLocationConstraint {
         }
     }
 }
-
-impl From<Vec<mullvad_types::custom_list::CustomList>> for proto::CustomLists {
-    fn from(custom_lists: Vec<mullvad_types::custom_list::CustomList>) -> Self {
-        let custom_lists = custom_lists
-            .into_iter()
-            .map(proto::CustomList::from)
-            .collect();
-        proto::CustomLists { custom_lists }
-    }
-}
-
-impl TryFrom<proto::CustomLists> for Vec<mullvad_types::custom_list::CustomList> {
-    type Error = FromProtobufTypeError;
-
-    fn try_from(custom_lists: proto::CustomLists) -> Result<Self, Self::Error> {
-        let mut new_custom_lists = Vec::with_capacity(custom_lists.custom_lists.len());
-        for custom_list in custom_lists.custom_lists {
-            new_custom_lists.push(mullvad_types::custom_list::CustomList::try_from(
-                custom_list,
-            )?);
-        }
-        Ok(new_custom_lists)
-    }
-}
diff --git a/mullvad-management-interface/src/types/conversions/relay_constraints.rs b/mullvad-management-interface/src/types/conversions/relay_constraints.rs
index a3ec8af0da..b7b5f0060b 100644
--- a/mullvad-management-interface/src/types/conversions/relay_constraints.rs
+++ b/mullvad-management-interface/src/types/conversions/relay_constraints.rs
@@ -3,6 +3,7 @@ use mullvad_types::{
     custom_list::Id,
     relay_constraints::{Constraint, RelaySettingsUpdate},
 };
+use std::str::FromStr;
 use talpid_types::net::TunnelType;
 
 impl TryFrom<&proto::WireguardConstraints>
@@ -494,7 +495,9 @@ impl From<mullvad_types::relay_constraints::LocationConstraint> for proto::Locat
                 )),
             },
             LocationConstraint::CustomList { list_id } => Self {
-                r#type: Some(proto::location_constraint::Type::CustomList(list_id)),
+                r#type: Some(proto::location_constraint::Type::CustomList(
+                    list_id.to_string(),
+                )),
             },
         }
     }
@@ -521,7 +524,7 @@ impl TryFrom<proto::LocationConstraint>
             }
             Some(proto::location_constraint::Type::CustomList(list_id)) => {
                 let location = LocationConstraint::CustomList {
-                    list_id: Id::try_from(list_id.as_str()).map_err(|_| {
+                    list_id: Id::from_str(&list_id).map_err(|_| {
                         FromProtobufTypeError::InvalidArgument("Id could not be parsed to a uuid")
                     })?,
                 };
diff --git a/mullvad-management-interface/src/types/conversions/settings.rs b/mullvad-management-interface/src/types/conversions/settings.rs
index 5b8bf76d86..f123b41755 100644
--- a/mullvad-management-interface/src/types/conversions/settings.rs
+++ b/mullvad-management-interface/src/types/conversions/settings.rs
@@ -39,7 +39,9 @@ impl From<&mullvad_types::settings::Settings> for proto::Settings {
                 &settings.obfuscation_settings,
             )),
             split_tunnel,
-            custom_lists: Some(proto::CustomListSettings::from(&settings.custom_lists)),
+            custom_lists: Some(proto::CustomListSettings::from(
+                settings.custom_lists.clone(),
+            )),
         }
     }
 }