Fix missing cities and relays in custom list

4 files changed, 98 insertions, 43 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt
index a71005a78a..3f138dee29 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/relaylist/RelayItemExtensions.kt
@@ -1,6 +1,10 @@
 package net.mullvad.mullvadvpn.relaylist
 
+import java.lang.IllegalArgumentException
+import net.mullvad.mullvadvpn.model.Constraint
 import net.mullvad.mullvadvpn.model.LocationConstraint
+import net.mullvad.mullvadvpn.model.Ownership
+import net.mullvad.mullvadvpn.model.Providers
 
 fun RelayItem.toLocationConstraint(): LocationConstraint {
     return when (this) {
@@ -24,3 +28,90 @@ fun RelayItem.descendants(): List<RelayItem> {
     val children = children()
     return children + children.flatMap { it.descendants() }
 }
+
+private fun RelayItem.hasOwnership(ownershipConstraint: Constraint<Ownership>): Boolean =
+    if (ownershipConstraint is Constraint.Only) {
+        when (this) {
+            is RelayItem.Country -> cities.any { it.hasOwnership(ownershipConstraint) }
+            is RelayItem.City -> relays.any { it.hasOwnership(ownershipConstraint) }
+            is RelayItem.Relay -> this.ownership == ownershipConstraint.value
+            is RelayItem.CustomList -> locations.any { it.hasOwnership(ownershipConstraint) }
+        }
+    } else {
+        true
+    }
+
+private fun RelayItem.hasProvider(providersConstraint: Constraint<Providers>): Boolean =
+    if (providersConstraint is Constraint.Only) {
+        when (this) {
+            is RelayItem.Country -> cities.any { it.hasProvider(providersConstraint) }
+            is RelayItem.City -> relays.any { it.hasProvider(providersConstraint) }
+            is RelayItem.Relay -> providersConstraint.value.providers.contains(providerName)
+            is RelayItem.CustomList -> locations.any { it.hasProvider(providersConstraint) }
+        }
+    } else {
+        true
+    }
+
+fun RelayItem.filterOnOwnershipAndProvider(
+    ownership: Constraint<Ownership>,
+    providers: Constraint<Providers>
+): RelayItem? =
+    when (this) {
+        is RelayItem.City -> filterOnOwnershipAndProvider(ownership, providers)
+        is RelayItem.Country -> filterOnOwnershipAndProvider(ownership, providers)
+        is RelayItem.CustomList -> filterOnOwnershipAndProvider(ownership, providers)
+        is RelayItem.Relay -> filterOnOwnershipAndProvider(ownership, providers)
+    }
+
+fun RelayItem.CustomList.filterOnOwnershipAndProvider(
+    ownership: Constraint<Ownership>,
+    providers: Constraint<Providers>
+): RelayItem.CustomList {
+    val newLocations =
+        locations.mapNotNull {
+            when (it) {
+                is RelayItem.City -> it.filterOnOwnershipAndProvider(ownership, providers)
+                is RelayItem.Country -> it.filterOnOwnershipAndProvider(ownership, providers)
+                is RelayItem.CustomList ->
+                    throw IllegalArgumentException("CustomList can't contain CustomList")
+                is RelayItem.Relay -> it.filterOnOwnershipAndProvider(ownership, providers)
+            }
+        }
+    return copy(locations = newLocations)
+}
+
+fun RelayItem.Country.filterOnOwnershipAndProvider(
+    ownership: Constraint<Ownership>,
+    providers: Constraint<Providers>
+): RelayItem.Country? {
+    val cities = cities.mapNotNull { it.filterOnOwnershipAndProvider(ownership, providers) }
+    return if (cities.isNotEmpty()) {
+        this.copy(cities = cities)
+    } else {
+        null
+    }
+}
+
+private fun RelayItem.City.filterOnOwnershipAndProvider(
+    ownership: Constraint<Ownership>,
+    providers: Constraint<Providers>
+): RelayItem.City? {
+    val relays = relays.mapNotNull { it.filterOnOwnershipAndProvider(ownership, providers) }
+    return if (relays.isNotEmpty()) {
+        this.copy(relays = relays)
+    } else {
+        null
+    }
+}
+
+private fun RelayItem.Relay.filterOnOwnershipAndProvider(
+    ownership: Constraint<Ownership>,
+    providers: Constraint<Providers>
+): RelayItem.Relay? {
+    return if (hasOwnership(ownership) && hasProvider(providers)) {
+        this
+    } else {
+        null
+    }
+}
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 8d2ea5f348..78b3732734 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
@@ -1,9 +1,7 @@
 package net.mullvad.mullvadvpn.relaylist
 
-import net.mullvad.mullvadvpn.model.Constraint
 import net.mullvad.mullvadvpn.model.GeographicLocationConstraint
 import net.mullvad.mullvadvpn.model.Ownership
-import net.mullvad.mullvadvpn.model.Providers
 import net.mullvad.mullvadvpn.model.Relay as DaemonRelay
 import net.mullvad.mullvadvpn.model.RelayList
 
@@ -175,39 +173,6 @@ private fun List<DaemonRelay>.filterValidRelays(): List<DaemonRelay> = filter {
     it.isWireguardRelay
 }
 
-fun List<RelayItem.Country>.filterOnOwnershipAndProviders(
-    ownership: Constraint<Ownership>,
-    providers: Constraint<Providers>
-): 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 -> it.ownership == ownership.value
-            }
-        }
-        .filter { relay ->
-            when (providers) {
-                is Constraint.Any -> true
-                is Constraint.Only -> providers.value.providers.contains(relay.providerName)
-            }
-        }
-
 /** Expand the parent(s), if any, for the current selected item */
 private fun List<RelayItem.Country>.expandItemForSelection(
     selectedItem: RelayItem?
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 bdeaac0d0c..b4197fe7b7 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,7 +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.filterOnOwnershipAndProvider
 import net.mullvad.mullvadvpn.relaylist.findItemForGeographicLocationConstraint
 import net.mullvad.mullvadvpn.relaylist.toRelayCountries
 import net.mullvad.mullvadvpn.relaylist.toRelayItemLists
@@ -41,11 +41,11 @@ class RelayListUseCase(
             val customLists =
                 settings?.customLists?.customLists?.toRelayItemLists(relayCountries) ?: emptyList()
             val relayCountriesFiltered =
-                relayCountries.filterOnOwnershipAndProviders(ownership, providers)
+                relayCountries.mapNotNull { it.filterOnOwnershipAndProvider(ownership, providers) }
             val selectedItem =
                 findSelectedRelayItem(
                     relaySettings = settings?.relaySettings,
-                    relayCountries = relayCountriesFiltered,
+                    relayCountries = relayCountries,
                     customLists = customLists,
                 )
             RelayList(
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 586bbf8bc1..1759055bda 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
@@ -19,6 +19,7 @@ import net.mullvad.mullvadvpn.model.Constraint
 import net.mullvad.mullvadvpn.model.Ownership
 import net.mullvad.mullvadvpn.relaylist.Provider
 import net.mullvad.mullvadvpn.relaylist.RelayItem
+import net.mullvad.mullvadvpn.relaylist.filterOnOwnershipAndProvider
 import net.mullvad.mullvadvpn.relaylist.filterOnSearchTerm
 import net.mullvad.mullvadvpn.relaylist.toLocationConstraint
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
@@ -66,11 +67,9 @@ class SelectLocationViewModel(
 
                 val filteredCustomLists =
                     customLists.filterOnSearchTerm(searchTerm).map { customList ->
-                        customList.copy(
-                            locations =
-                                customList.locations.filter { location ->
-                                    filteredRelayCountries.any { it.code == location.code }
-                                }
+                        customList.filterOnOwnershipAndProvider(
+                            selectedOwnership,
+                            selectedConstraintProviders
                         )
                     }