Return success or error when creating and updating a custom list

4 files changed, 94 insertions, 16 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() }
+}