Return success or error when creating and updating a custom list

10 files changed, 146 insertions, 26 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemType.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemType.kt
deleted file mode 100644
index cdbd58b291..0000000000
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemType.kt
+++ /dev/null
@@ -1,7 +0,0 @@
-package net.mullvad.mullvadvpn.relaylist
-
-enum class RelayItemType {
-    Country,
-    City,
-    Relay,
-}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt
index 06e00a022a..882a3e42a4 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayListExtensions.kt
@@ -236,4 +236,32 @@ private fun List<RelayItem.Country>.expandItemForSelection(
     } ?: this
 }
 
-private const val MIN_SEARCH_LENGTH = 2
+@Suppress("NestedBlockDepth", "ReturnCount")
+fun RelayList.getGeographicLocationConstraintByCode(code: String): GeographicLocationConstraint? {
+    countries.forEach { country ->
+        val countryCode = country.code
+        if (country.code == code) {
+            return GeographicLocationConstraint.Country(countryCode)
+        }
+        country.cities.forEach { city ->
+            val cityCode = city.code
+            if (city.code == code) {
+                return GeographicLocationConstraint.City(countryCode, city.code)
+            }
+            city.relays.forEach { relay ->
+                if (relay.hostname == code) {
+                    return GeographicLocationConstraint.Hostname(
+                        countryCode,
+                        cityCode,
+                        relay.hostname
+                    )
+                }
+            }
+        }
+    }
+    return null
+}
+
+fun List<RelayItem.Country>.getRelayItemsByCodes(codes: List<String>): List<RelayItem> =
+    this.filter { codes.contains(it.code) } +
+        this.flatMap { it.descendants() }.filter { codes.contains(it.code) }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/repository/CustomListsRepository.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/repository/CustomListsRepository.kt
index 9660981688..f1a38871bd 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/repository/CustomListsRepository.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/repository/CustomListsRepository.kt
@@ -1,28 +1,80 @@
 package net.mullvad.mullvadvpn.repository
 
 import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.mapNotNull
 import net.mullvad.mullvadvpn.lib.ipc.Event
 import net.mullvad.mullvadvpn.lib.ipc.MessageHandler
 import net.mullvad.mullvadvpn.lib.ipc.Request
 import net.mullvad.mullvadvpn.lib.ipc.events
+import net.mullvad.mullvadvpn.model.CreateCustomListResult
 import net.mullvad.mullvadvpn.model.CustomList
+import net.mullvad.mullvadvpn.model.CustomListsError
+import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.UpdateCustomListResult
+import net.mullvad.mullvadvpn.relaylist.getGeographicLocationConstraintByCode
+import net.mullvad.mullvadvpn.ui.serviceconnection.RelayListListener
+import net.mullvad.mullvadvpn.util.firstOrNullWithTimeout
 
-class CustomListsRepository(private val messageHandler: MessageHandler) {
-    suspend fun createCustomList(name: String): String? {
+class CustomListsRepository(
+    private val messageHandler: MessageHandler,
+    private val settingsRepository: SettingsRepository,
+    private val relayListListener: RelayListListener
+) {
+    suspend fun createCustomList(name: String): CreateCustomListResult {
         val result = messageHandler.trySendRequest(Request.CreateCustomList(name))
 
         return if (result) {
-            messageHandler.events<Event.CreateCustomListResult>().first().listId
+            messageHandler.events<Event.CreateCustomListResultEvent>().first().result
         } else {
-            null
+            CreateCustomListResult.Error(CustomListsError.OtherError)
         }
     }
 
-    fun deleteCustomList(id: String) {
-        messageHandler.trySendRequest(Request.DeleteCustomList(id))
+    fun deleteCustomList(id: String) = messageHandler.trySendRequest(Request.DeleteCustomList(id))
+
+    private suspend fun updateCustomList(customList: CustomList): UpdateCustomListResult {
+        val result = messageHandler.trySendRequest(Request.UpdateCustomList(customList))
+
+        return if (result) {
+            messageHandler.events<Event.UpdateCustomListResultEvent>().first().result
+        } else {
+            UpdateCustomListResult.Error(CustomListsError.OtherError)
+        }
     }
 
-    fun updateCustomList(customList: CustomList) {
-        messageHandler.trySendRequest(Request.UpdateCustomList(customList))
+    suspend fun updateCustomListLocationsFromCodes(
+        id: String,
+        locationCodes: List<String>
+    ): UpdateCustomListResult =
+        updateCustomListLocations(
+            id = id,
+            locations =
+                ArrayList(locationCodes.mapNotNull { getGeographicLocationConstraintByCode(it) })
+        )
+
+    suspend fun updateCustomListName(id: String, name: String): UpdateCustomListResult =
+        getCustomListById(id)?.let { updateCustomList(it.copy(name = name)) }
+            ?: UpdateCustomListResult.Error(CustomListsError.OtherError)
+
+    private suspend fun updateCustomListLocations(
+        id: String,
+        locations: ArrayList<GeographicLocationConstraint>
+    ): UpdateCustomListResult =
+        awaitCustomListById(id)?.let { updateCustomList(it.copy(locations = locations)) }
+            ?: UpdateCustomListResult.Error(CustomListsError.OtherError)
+
+    private suspend fun awaitCustomListById(id: String): CustomList? =
+        settingsRepository.settingsUpdates
+            .mapNotNull { settings -> settings?.customLists?.customLists?.find { it.id == id } }
+            .firstOrNullWithTimeout(GET_CUSTOM_LIST_TIMEOUT_MS)
+
+    fun getCustomListById(id: String): CustomList? =
+        settingsRepository.settingsUpdates.value?.customLists?.customLists?.find { it.id == id }
+
+    private fun getGeographicLocationConstraintByCode(code: String): GeographicLocationConstraint? =
+        relayListListener.relayListEvents.value.getGeographicLocationConstraintByCode(code)
+
+    companion object {
+        private const val GET_CUSTOM_LIST_TIMEOUT_MS = 5000L
     }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
index 843c7c2930..b3a8727df9 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/FlowUtils.kt
@@ -8,6 +8,7 @@ import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.flow.catch
+import kotlinx.coroutines.flow.firstOrNull
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.retryWhen
@@ -149,3 +150,7 @@ suspend inline fun <T> Flow<T>.retryWithExponentialBackOff(
         }
 
 class ExceptionWrapper(val item: Any) : Throwable()
+
+suspend fun <T> Flow<T>.firstOrNullWithTimeout(timeMillis: Long): T? {
+    return withTimeoutOrNull(timeMillis) { firstOrNull() }
+}
diff --git a/android/lib/ipc/src/main/kotlin/net/mullvad/mullvadvpn/lib/ipc/Event.kt b/android/lib/ipc/src/main/kotlin/net/mullvad/mullvadvpn/lib/ipc/Event.kt
index 1136ae8c55..cce2ab1f87 100644
--- a/android/lib/ipc/src/main/kotlin/net/mullvad/mullvadvpn/lib/ipc/Event.kt
+++ b/android/lib/ipc/src/main/kotlin/net/mullvad/mullvadvpn/lib/ipc/Event.kt
@@ -5,6 +5,7 @@ import kotlinx.parcelize.Parcelize
 import net.mullvad.mullvadvpn.model.AccountCreationResult
 import net.mullvad.mullvadvpn.model.AccountExpiry
 import net.mullvad.mullvadvpn.model.AccountHistory
+import net.mullvad.mullvadvpn.model.CreateCustomListResult
 import net.mullvad.mullvadvpn.model.DeviceListEvent
 import net.mullvad.mullvadvpn.model.DeviceState
 import net.mullvad.mullvadvpn.model.LoginResult
@@ -14,6 +15,7 @@ import net.mullvad.mullvadvpn.model.RelayList
 import net.mullvad.mullvadvpn.model.RemoveDeviceResult
 import net.mullvad.mullvadvpn.model.Settings
 import net.mullvad.mullvadvpn.model.TunnelState
+import net.mullvad.mullvadvpn.model.UpdateCustomListResult
 
 // Events that can be sent from the service
 sealed class Event : Message.EventMessage() {
@@ -65,7 +67,9 @@ sealed class Event : Message.EventMessage() {
 
     @Parcelize object VpnPermissionRequest : Event()
 
-    @Parcelize data class CreateCustomListResult(val listId: String) : Event()
+    @Parcelize data class CreateCustomListResultEvent(val result: CreateCustomListResult) : Event()
+
+    @Parcelize data class UpdateCustomListResultEvent(val result: UpdateCustomListResult) : Event()
 
     companion object {
         private const val MESSAGE_KEY = "event"
diff --git a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/model/CreateCustomListResult.kt b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/model/CreateCustomListResult.kt
new file mode 100644
index 0000000000..73eaa209c8
--- /dev/null
+++ b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/model/CreateCustomListResult.kt
@@ -0,0 +1,10 @@
+package net.mullvad.mullvadvpn.model
+
+import android.os.Parcelable
+import kotlinx.parcelize.Parcelize
+
+sealed class CreateCustomListResult : Parcelable {
+    @Parcelize data class Ok(val id: String) : CreateCustomListResult()
+
+    @Parcelize data class Error(val error: CustomListsError) : CreateCustomListResult()
+}
diff --git a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/model/CustomListsError.kt b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/model/CustomListsError.kt
new file mode 100644
index 0000000000..83806af4f7
--- /dev/null
+++ b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/model/CustomListsError.kt
@@ -0,0 +1,6 @@
+package net.mullvad.mullvadvpn.model
+
+enum class CustomListsError {
+    CustomListExists,
+    OtherError
+}
diff --git a/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/model/UpdateCustomListResult.kt b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/model/UpdateCustomListResult.kt
new file mode 100644
index 0000000000..ebfe9e8cd6
--- /dev/null
+++ b/android/lib/model/src/main/kotlin/net/mullvad/mullvadvpn/model/UpdateCustomListResult.kt
@@ -0,0 +1,10 @@
+package net.mullvad.mullvadvpn.model
+
+import android.os.Parcelable
+import kotlinx.parcelize.Parcelize
+
+sealed class UpdateCustomListResult : Parcelable {
+    @Parcelize data object Ok : UpdateCustomListResult()
+
+    @Parcelize data class Error(val error: CustomListsError) : UpdateCustomListResult()
+}
diff --git a/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt b/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt
index d09287fbf2..f99d36c679 100644
--- a/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt
+++ b/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/MullvadDaemon.kt
@@ -5,6 +5,7 @@ import kotlinx.coroutines.flow.asSharedFlow
 import net.mullvad.mullvadvpn.lib.endpoint.ApiEndpoint
 import net.mullvad.mullvadvpn.lib.endpoint.ApiEndpointConfiguration
 import net.mullvad.mullvadvpn.model.AppVersionInfo
+import net.mullvad.mullvadvpn.model.CreateCustomListResult
 import net.mullvad.mullvadvpn.model.CustomList
 import net.mullvad.mullvadvpn.model.Device
 import net.mullvad.mullvadvpn.model.DeviceEvent
@@ -24,6 +25,7 @@ import net.mullvad.mullvadvpn.model.RemoveDeviceEvent
 import net.mullvad.mullvadvpn.model.RemoveDeviceResult
 import net.mullvad.mullvadvpn.model.Settings
 import net.mullvad.mullvadvpn.model.TunnelState
+import net.mullvad.mullvadvpn.model.UpdateCustomListResult
 import net.mullvad.mullvadvpn.model.VoucherSubmissionResult
 import net.mullvad.talpid.util.EventNotifier
 
@@ -190,17 +192,15 @@ class MullvadDaemon(
         setQuantumResistantTunnel(daemonInterfaceAddress, quantumResistant)
     }
 
-    fun createCustomList(name: String): String {
-        return createCustomList(daemonInterfaceAddress, name)
-    }
+    fun createCustomList(name: String): CreateCustomListResult =
+        createCustomList(daemonInterfaceAddress, name)
 
     fun deleteCustomList(id: String) {
         deleteCustomList(daemonInterfaceAddress, id)
     }
 
-    fun updateCustomList(customList: CustomList) {
+    fun updateCustomList(customList: CustomList): UpdateCustomListResult =
         updateCustomList(daemonInterfaceAddress, customList)
-    }
 
     fun onDestroy() {
         onSettingsChange.unsubscribeAll()
@@ -311,11 +311,17 @@ class MullvadDaemon(
 
     // Used by JNI
 
-    private external fun createCustomList(daemonInterfaceAddress: Long, name: String): String
+    private external fun createCustomList(
+        daemonInterfaceAddress: Long,
+        name: String
+    ): CreateCustomListResult
 
     private external fun deleteCustomList(daemonInterfaceAddress: Long, id: String)
 
-    private external fun updateCustomList(daemonInterfaceAddress: Long, customList: CustomList)
+    private external fun updateCustomList(
+        daemonInterfaceAddress: Long,
+        customList: CustomList
+    ): UpdateCustomListResult
 
     @Suppress("unused")
     private fun notifyAppVersionInfoEvent(appVersionInfo: AppVersionInfo) {
diff --git a/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/CustomLists.kt b/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/CustomLists.kt
index d80bcf04ff..39702398c7 100644
--- a/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/CustomLists.kt
+++ b/android/service/src/main/kotlin/net/mullvad/mullvadvpn/service/endpoint/CustomLists.kt
@@ -9,6 +9,7 @@ import kotlinx.coroutines.flow.filterIsInstance
 import kotlinx.coroutines.launch
 import net.mullvad.mullvadvpn.lib.ipc.Event
 import net.mullvad.mullvadvpn.lib.ipc.Request
+import net.mullvad.mullvadvpn.model.CustomList
 
 class CustomLists(
     private val endpoint: ServiceEndpoint,
@@ -34,13 +35,18 @@ class CustomLists(
         scope.launch {
             endpoint.dispatcher.parsedMessages
                 .filterIsInstance<Request.UpdateCustomList>()
-                .collect { daemon.await().updateCustomList(it.customList) }
+                .collect { updateCustomList(it.customList) }
         }
     }
 
     private suspend fun createCustomList(name: String) {
         val result = daemon.await().createCustomList(name)
-        endpoint.sendEvent(Event.CreateCustomListResult(result))
+        endpoint.sendEvent(Event.CreateCustomListResultEvent(result))
+    }
+
+    private suspend fun updateCustomList(customList: CustomList) {
+        val result = daemon.await().updateCustomList(customList)
+        endpoint.sendEvent(Event.UpdateCustomListResultEvent(result))
     }
 
     fun onDestroy() {