Do not filter locations for custom lists

With the exception of the select location screen list

9 files changed, 71 insertions, 30 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreen.kt
index 87ed88d263..9186e639c5 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/EditCustomListScreen.kt
@@ -45,6 +45,7 @@ import net.mullvad.mullvadvpn.compose.transitions.SlideInFromRightTransition
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
 import net.mullvad.mullvadvpn.lib.theme.Dimens
 import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.Ownership
 import net.mullvad.mullvadvpn.relaylist.RelayItem
 import net.mullvad.mullvadvpn.viewmodel.EditCustomListViewModel
 import org.koin.androidx.compose.koinViewModel
@@ -69,7 +70,9 @@ private fun PreviewEditCustomListScreen() {
                                     "hostname",
                                     "hostname",
                                     "hostname"
-                                )
+                                ),
+                                "Provider",
+                                Ownership.MullvadOwned
                             )
                         )
                 )
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/util/PreviewData.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/util/PreviewData.kt
index 61b563564c..3581d1d0b4 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/util/PreviewData.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/util/PreviewData.kt
@@ -1,6 +1,7 @@
 package net.mullvad.mullvadvpn.compose.util
 
 import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.Ownership
 import net.mullvad.mullvadvpn.relaylist.RelayItem
 
 fun generateRelayItemCountry(
@@ -65,7 +66,9 @@ fun generateRelayItemRelay(
                 hostname = hostName,
             ),
         locationName = "$cityCode $hostName",
-        active = active
+        active = active,
+        providerName = "Provider",
+        ownership = Ownership.MullvadOwned,
     )
 
 private fun String.generateCountryCode() = (take(1) + takeLast(1)).lowercase()
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt
index ce4be395b6..af4a0084d2 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItem.kt
@@ -3,6 +3,7 @@ package net.mullvad.mullvadvpn.relaylist
 import net.mullvad.mullvadvpn.model.CustomListName
 import net.mullvad.mullvadvpn.model.GeoIpLocation
 import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
+import net.mullvad.mullvadvpn.model.Ownership
 
 sealed interface RelayItem {
     val name: String
@@ -66,6 +67,8 @@ sealed interface RelayItem {
         override val locationName: String,
         override val active: Boolean,
         val location: GeographicLocationConstraint.Hostname,
+        val providerName: String,
+        val ownership: Ownership,
     ) : RelayItem {
         override val code = name
         override val hasChildren = false
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
index 30e4146245..e469aec118 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayList.kt
@@ -2,6 +2,7 @@ package net.mullvad.mullvadvpn.relaylist
 
 data class RelayList(
     val customLists: List<RelayItem.CustomList>,
-    val country: List<RelayItem.Country>,
+    val allCountries: List<RelayItem.Country>,
+    val filteredCountries: List<RelayItem.Country>,
     val selectedItem: RelayItem?,
 )
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 882a3e42a4..8d2ea5f348 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
@@ -12,10 +12,7 @@ import net.mullvad.mullvadvpn.model.RelayList
  * filtered out and also relays that do not fit the ownership and provider list So are also cities
  * that only contains non-wireguard relays Countries, cities and relays are ordered by name
  */
-fun RelayList.toRelayCountries(
-    ownership: Constraint<Ownership>,
-    providers: Constraint<Providers>
-): List<RelayItem.Country> {
+fun RelayList.toRelayCountries(): List<RelayItem.Country> {
     val relayCountries =
         this.countries
             .map { country ->
@@ -33,8 +30,7 @@ fun RelayList.toRelayCountries(
                             relays = relays
                         )
 
-                    val validCityRelays =
-                        city.relays.filterValidRelays(ownership = ownership, providers = providers)
+                    val validCityRelays = city.relays.filterValidRelays()
 
                     for (relay in validCityRelays) {
                         relays.add(
@@ -47,7 +43,10 @@ fun RelayList.toRelayCountries(
                                         relay.hostname
                                     ),
                                 locationName = "${city.name} (${relay.hostname})",
-                                active = relay.active
+                                active = relay.active,
+                                providerName = relay.provider,
+                                ownership =
+                                    if (relay.owned) Ownership.MullvadOwned else Ownership.Rented
                             )
                         )
                     }
@@ -172,25 +171,40 @@ fun List<RelayItem.Country>.filterOnSearchTerm(
     }
 }
 
-private fun List<DaemonRelay>.filterValidRelays(
+private fun List<DaemonRelay>.filterValidRelays(): List<DaemonRelay> = filter {
+    it.isWireguardRelay
+}
+
+fun List<RelayItem.Country>.filterOnOwnershipAndProviders(
     ownership: Constraint<Ownership>,
     providers: Constraint<Providers>
-): List<DaemonRelay> =
-    filter { it.isWireguardRelay }
-        .filter {
+): List<RelayItem.Country> {
+    return map { country ->
+            val cities =
+                country.cities.map { city ->
+                    val relays =
+                        city.relays.filterRelayByOwnershipAndProviders(ownership, providers)
+                    city.copy(relays = relays)
+                }
+            country.copy(cities = cities.filter { it.relays.isNotEmpty() })
+        }
+        .filter { it.cities.isNotEmpty() }
+}
+
+private fun List<RelayItem.Relay>.filterRelayByOwnershipAndProviders(
+    ownership: Constraint<Ownership>,
+    providers: Constraint<Providers>
+): List<RelayItem.Relay> =
+    filter {
             when (ownership) {
                 is Constraint.Any -> true
-                is Constraint.Only ->
-                    when (ownership.value) {
-                        Ownership.MullvadOwned -> it.owned
-                        Ownership.Rented -> !it.owned
-                    }
+                is Constraint.Only -> it.ownership == ownership.value
             }
         }
         .filter { relay ->
             when (providers) {
                 is Constraint.Any -> true
-                is Constraint.Only -> providers.value.providers.contains(relay.provider)
+                is Constraint.Only -> providers.value.providers.contains(relay.providerName)
             }
         }
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/RelayListUseCase.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/RelayListUseCase.kt
index 4957818283..bdeaac0d0c 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/RelayListUseCase.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/RelayListUseCase.kt
@@ -9,6 +9,7 @@ import net.mullvad.mullvadvpn.model.RelaySettings
 import net.mullvad.mullvadvpn.model.WireguardConstraints
 import net.mullvad.mullvadvpn.relaylist.RelayItem
 import net.mullvad.mullvadvpn.relaylist.RelayList
+import net.mullvad.mullvadvpn.relaylist.filterOnOwnershipAndProviders
 import net.mullvad.mullvadvpn.relaylist.findItemForGeographicLocationConstraint
 import net.mullvad.mullvadvpn.relaylist.toRelayCountries
 import net.mullvad.mullvadvpn.relaylist.toRelayItemLists
@@ -36,22 +37,29 @@ class RelayListUseCase(
                 settings?.relaySettings?.relayConstraints()?.ownership ?: Constraint.Any()
             val providers =
                 settings?.relaySettings?.relayConstraints()?.providers ?: Constraint.Any()
-            val relayCountries =
-                relayList.toRelayCountries(ownership = ownership, providers = providers)
+            val relayCountries = relayList.toRelayCountries()
             val customLists =
                 settings?.customLists?.customLists?.toRelayItemLists(relayCountries) ?: emptyList()
+            val relayCountriesFiltered =
+                relayCountries.filterOnOwnershipAndProviders(ownership, providers)
             val selectedItem =
                 findSelectedRelayItem(
                     relaySettings = settings?.relaySettings,
-                    relayCountries = relayCountries,
+                    relayCountries = relayCountriesFiltered,
                     customLists = customLists,
                 )
-            RelayList(customLists, relayCountries, selectedItem)
+            RelayList(
+                customLists = customLists,
+                allCountries = relayCountries,
+                filteredCountries = relayCountriesFiltered,
+                selectedItem = selectedItem
+            )
         }
 
     fun selectedRelayItem(): Flow<RelayItem?> = relayListWithSelection().map { it.selectedItem }
 
-    fun relayList(): Flow<List<RelayItem.Country>> = relayListWithSelection().map { it.country }
+    fun fullRelayList(): Flow<List<RelayItem.Country>> =
+        relayListWithSelection().map { it.allCountries }
 
     fun customLists(): Flow<List<RelayItem.CustomList>> =
         relayListWithSelection().map { it.customLists }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/customlists/CustomListActionUseCase.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/customlists/CustomListActionUseCase.kt
index 8d722325d6..16c86c0d59 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/customlists/CustomListActionUseCase.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/usecase/customlists/CustomListActionUseCase.kt
@@ -53,7 +53,7 @@ class CustomListActionUseCase(
                     )
                     val locationNames =
                         relayListUseCase
-                            .relayList()
+                            .fullRelayList()
                             .firstOrNull()
                             ?.getRelayItemsByCodes(action.locations)
                             ?.map { it.name }
@@ -61,7 +61,7 @@ class CustomListActionUseCase(
                         CustomListResult.Created(
                             id = result.id,
                             name = action.name,
-                            locationName = locationNames?.first(),
+                            locationName = locationNames?.firstOrNull(),
                             undo = action.not(result.id)
                         )
                     )
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt
index 5efba5321e..cdbcebbb83 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/CustomListLocationsViewModel.kt
@@ -37,7 +37,7 @@ class CustomListLocationsViewModel(
     private val _searchTerm = MutableStateFlow(EMPTY_SEARCH_TERM)
 
     val uiState =
-        combine(relayListUseCase.relayList(), _searchTerm, _selectedLocations) {
+        combine(relayListUseCase.fullRelayList(), _searchTerm, _selectedLocations) {
                 relayCountries,
                 searchTerm,
                 selectedLocations ->
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt
index 15df90be9e..586bbf8bc1 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/viewmodel/SelectLocationViewModel.kt
@@ -35,6 +35,7 @@ class SelectLocationViewModel(
 ) : ViewModel() {
     private val _searchTerm = MutableStateFlow(EMPTY_SEARCH_TERM)
 
+    @Suppress("DestructuringDeclarationWithTooManyEntries")
     val uiState =
         combine(
                 relayListUseCase.relayListWithSelection(),
@@ -43,7 +44,7 @@ class SelectLocationViewModel(
                 relayListFilterUseCase.availableProviders(),
                 relayListFilterUseCase.selectedProviders(),
             ) {
-                (customLists, relayCountries, selectedItem),
+                (customLists, _, relayCountries, selectedItem),
                 searchTerm,
                 selectedOwnership,
                 allProviders,
@@ -63,7 +64,15 @@ class SelectLocationViewModel(
                 val filteredRelayCountries =
                     relayCountries.filterOnSearchTerm(searchTerm, selectedItem)
 
-                val filteredCustomLists = customLists.filterOnSearchTerm(searchTerm)
+                val filteredCustomLists =
+                    customLists.filterOnSearchTerm(searchTerm).map { customList ->
+                        customList.copy(
+                            locations =
+                                customList.locations.filter { location ->
+                                    filteredRelayCountries.any { it.code == location.code }
+                                }
+                        )
+                    }
 
                 SelectLocationUiState.Content(
                     searchTerm = searchTerm,