Adjust relay selector to support custom lists

8 files changed, 107 insertions, 36 deletions

diff --git a/ios/MullvadREST/Relay/RelaySelector.swift b/ios/MullvadREST/Relay/RelaySelector.swift
index 6fc016d2c7..b156c83fa4 100644
--- a/ios/MullvadREST/Relay/RelaySelector.swift
+++ b/ios/MullvadREST/Relay/RelaySelector.swift
@@ -150,24 +150,28 @@ public enum RelaySelector {
                 }
             }
 
-            switch constraints.location {
+            switch constraints.locations {
             case .any:
                 return true
             case let .only(relayConstraint):
-                switch relayConstraint {
-                case let .country(countryCode):
-                    return relayWithLocation.serverLocation.countryCode == countryCode &&
-                        relayWithLocation.relay.includeInCountry
+                for location in relayConstraint.locations {
+                    switch location {
+                    case let .country(countryCode):
+                        return relayWithLocation.serverLocation.countryCode == countryCode &&
+                            relayWithLocation.relay.includeInCountry
 
-                case let .city(countryCode, cityCode):
-                    return relayWithLocation.serverLocation.countryCode == countryCode &&
-                        relayWithLocation.serverLocation.cityCode == cityCode
+                    case let .city(countryCode, cityCode):
+                        return relayWithLocation.serverLocation.countryCode == countryCode &&
+                            relayWithLocation.serverLocation.cityCode == cityCode
 
-                case let .hostname(countryCode, cityCode, hostname):
-                    return relayWithLocation.serverLocation.countryCode == countryCode &&
-                        relayWithLocation.serverLocation.cityCode == cityCode &&
-                        relayWithLocation.relay.hostname == hostname
+                    case let .hostname(countryCode, cityCode, hostname):
+                        return relayWithLocation.serverLocation.countryCode == countryCode &&
+                            relayWithLocation.serverLocation.cityCode == cityCode &&
+                            relayWithLocation.relay.hostname == hostname
+                    }
                 }
+
+                return false
             }
         }.filter { relayWithLocation -> Bool in
             relayWithLocation.relay.active
diff --git a/ios/MullvadTypes/RelayConstraints.swift b/ios/MullvadTypes/RelayConstraints.swift
index 602e17b088..a756008e0b 100644
--- a/ios/MullvadTypes/RelayConstraints.swift
+++ b/ios/MullvadTypes/RelayConstraints.swift
@@ -21,32 +21,59 @@ public class RelayConstraintsUpdater: ConstraintsPropagation {
 }
 
 public struct RelayConstraints: Codable, Equatable, CustomDebugStringConvertible {
-    public var location: RelayConstraint<RelayLocation>
+    @available(*, deprecated, renamed: "locations")
+    private var location: RelayConstraint<RelayLocation> = .only(.country("se"))
 
     // Added in 2023.3
     public var port: RelayConstraint<UInt16>
     public var filter: RelayConstraint<RelayFilter>
 
+    // Added in 2024.1
+    public var locations: RelayConstraint<RelayLocations>
+
     public var debugDescription: String {
-        "RelayConstraints { location: \(location), port: \(port) }"
+        "RelayConstraints { locations: \(locations), port: \(port) }"
     }
 
     public init(
-        location: RelayConstraint<RelayLocation> = .only(.country("se")),
+        locations: RelayConstraint<RelayLocations> = .only(RelayLocations(locations: [.country("se")])),
         port: RelayConstraint<UInt16> = .any,
         filter: RelayConstraint<RelayFilter> = .any
     ) {
-        self.location = location
+        self.locations = locations
         self.port = port
         self.filter = filter
     }
 
     public init(from decoder: Decoder) throws {
         let container = try decoder.container(keyedBy: CodingKeys.self)
-        location = try container.decode(RelayConstraint<RelayLocation>.self, forKey: .location)
 
         // Added in 2023.3
         port = try container.decodeIfPresent(RelayConstraint<UInt16>.self, forKey: .port) ?? .any
         filter = try container.decodeIfPresent(RelayConstraint<RelayFilter>.self, forKey: .filter) ?? .any
+
+        // Added in 2024.1
+        locations = try container.decodeIfPresent(RelayConstraint<RelayLocations>.self, forKey: .locations)
+            ?? Self.migrateLocations(decoder: decoder)
+            ?? .only(RelayLocations(locations: [.country("se")]))
+    }
+}
+
+extension RelayConstraints {
+    private static func migrateLocations(decoder: Decoder) -> RelayConstraint<RelayLocations>? {
+        let container = try? decoder.container(keyedBy: CodingKeys.self)
+
+        guard
+            let location = try? container?.decodeIfPresent(RelayConstraint<RelayLocation>.self, forKey: .location)
+        else {
+            return nil
+        }
+
+        switch location {
+        case .any:
+            return .any
+        case let .only(location):
+            return .only(RelayLocations(locations: [location]))
+        }
     }
 }
diff --git a/ios/MullvadTypes/RelayLocation.swift b/ios/MullvadTypes/RelayLocation.swift
index b797e69d3c..d7dbb8d2a8 100644
--- a/ios/MullvadTypes/RelayLocation.swift
+++ b/ios/MullvadTypes/RelayLocation.swift
@@ -106,3 +106,13 @@ public enum RelayLocation: Codable, Hashable, CustomDebugStringConvertible {
         }
     }
 }
+
+public struct RelayLocations: Codable, Equatable {
+    public let locations: [RelayLocation]
+    public let customListId: UUID?
+
+    public init(locations: [RelayLocation], customListId: UUID? = nil) {
+        self.locations = locations
+        self.customListId = customListId
+    }
+}
diff --git a/ios/MullvadVPN/Coordinators/SelectLocationCoordinator.swift b/ios/MullvadVPN/Coordinators/SelectLocationCoordinator.swift
index 936c6ce4aa..dcaf47347d 100644
--- a/ios/MullvadVPN/Coordinators/SelectLocationCoordinator.swift
+++ b/ios/MullvadVPN/Coordinators/SelectLocationCoordinator.swift
@@ -58,7 +58,10 @@ class SelectLocationCoordinator: Coordinator, Presentable, Presenting, RelayCach
             guard let self else { return }
 
             var relayConstraints = tunnelManager.settings.relayConstraints
-            relayConstraints.location = .only(relay)
+            relayConstraints.locations = .only(RelayLocations(
+                locations: [relay],
+                customListId: nil
+            ))
 
             tunnelManager.updateSettings([.relayConstraints(relayConstraints)]) {
                 self.tunnelManager.startTunnel()
@@ -98,7 +101,8 @@ class SelectLocationCoordinator: Coordinator, Presentable, Presenting, RelayCach
             selectLocationViewController.setCachedRelays(cachedRelays, filter: relayFilter)
         }
 
-        selectLocationViewController.relayLocation = tunnelManager.settings.relayConstraints.location.value
+        selectLocationViewController.relayLocation =
+            tunnelManager.settings.relayConstraints.locations.value?.locations.first
 
         navigationController.pushViewController(selectLocationViewController, animated: false)
     }
diff --git a/ios/MullvadVPNTests/MigrationManagerTests.swift b/ios/MullvadVPNTests/MigrationManagerTests.swift
index c5f693ad01..bfe7213628 100644
--- a/ios/MullvadVPNTests/MigrationManagerTests.swift
+++ b/ios/MullvadVPNTests/MigrationManagerTests.swift
@@ -121,7 +121,10 @@ final class MigrationManagerTests: XCTestCase {
 
     func testSuccessfulMigrationFromV2ToLatest() throws {
         var settingsV2 = TunnelSettingsV2()
-        let osakaRelayConstraints: RelayConstraints = .init(location: .only(.city("jp", "osa")))
+        let osakaRelayConstraints = RelayConstraints(
+            locations: .only(RelayLocations(locations: [.city("jp", "osa")]))
+        )
+
         settingsV2.relayConstraints = osakaRelayConstraints
 
         try migrateToLatest(settingsV2, version: .v2)
@@ -132,7 +135,10 @@ final class MigrationManagerTests: XCTestCase {
 
     func testSuccessfulMigrationFromV1ToLatest() throws {
         var settingsV1 = TunnelSettingsV1()
-        let osakaRelayConstraints: RelayConstraints = .init(location: .only(.city("jp", "osa")))
+        let osakaRelayConstraints = RelayConstraints(
+            locations: .only(RelayLocations(locations: [.city("jp", "osa")]))
+        )
+
         settingsV1.relayConstraints = osakaRelayConstraints
 
         try migrateToLatest(settingsV1, version: .v1)
diff --git a/ios/MullvadVPNTests/MullvadSettings/TunnelSettingsUpdateTests.swift b/ios/MullvadVPNTests/MullvadSettings/TunnelSettingsUpdateTests.swift
index 3e4860889e..1a89f822a2 100644
--- a/ios/MullvadVPNTests/MullvadSettings/TunnelSettingsUpdateTests.swift
+++ b/ios/MullvadVPNTests/MullvadSettings/TunnelSettingsUpdateTests.swift
@@ -48,7 +48,7 @@ final class TunnelSettingsUpdateTests: XCTestCase {
 
         // When:
         let relayConstraints = RelayConstraints(
-            location: .only(.country("zz")),
+            locations: .only(RelayLocations(locations: [.country("zz")])),
             port: .only(9999),
             filter: .only(.init(ownership: .rented, providers: .only(["foo", "bar"])))
         )
diff --git a/ios/MullvadVPNTests/RelaySelectorTests.swift b/ios/MullvadVPNTests/RelaySelectorTests.swift
index 68bfadfd4c..03ff9983d9 100644
--- a/ios/MullvadVPNTests/RelaySelectorTests.swift
+++ b/ios/MullvadVPNTests/RelaySelectorTests.swift
@@ -18,7 +18,9 @@ class RelaySelectorTests: XCTestCase {
     let sampleRelays = ServerRelaysResponseStubs.sampleRelays
 
     func testCountryConstraint() throws {
-        let constraints = RelayConstraints(location: .only(.country("es")))
+        let constraints = RelayConstraints(
+            locations: .only(RelayLocations(locations: [.country("es")]))
+        )
 
         let result = try RelaySelector.evaluate(
             relays: sampleRelays,
@@ -30,7 +32,10 @@ class RelaySelectorTests: XCTestCase {
     }
 
     func testCityConstraint() throws {
-        let constraints = RelayConstraints(location: .only(.city("se", "got")))
+        let constraints = RelayConstraints(
+            locations: .only(RelayLocations(locations: [.city("se", "got")]))
+        )
+
         let result = try RelaySelector.evaluate(
             relays: sampleRelays,
             constraints: constraints,
@@ -41,7 +46,9 @@ class RelaySelectorTests: XCTestCase {
     }
 
     func testHostnameConstraint() throws {
-        let constraints = RelayConstraints(location: .only(.hostname("se", "sto", "se6-wireguard")))
+        let constraints = RelayConstraints(
+            locations: .only(RelayLocations(locations: [.hostname("se", "sto", "se6-wireguard")]))
+        )
 
         let result = try RelaySelector.evaluate(
             relays: sampleRelays,
@@ -53,7 +60,10 @@ class RelaySelectorTests: XCTestCase {
     }
 
     func testSpecificPortConstraint() throws {
-        let constraints = RelayConstraints(location: .only(.hostname("se", "sto", "se6-wireguard")), port: .only(1))
+        let constraints = RelayConstraints(
+            locations: .only(RelayLocations(locations: [.hostname("se", "sto", "se6-wireguard")])),
+            port: .only(1)
+        )
 
         let result = try RelaySelector.evaluate(
             relays: sampleRelays,
@@ -65,7 +75,9 @@ class RelaySelectorTests: XCTestCase {
     }
 
     func testRandomPortSelectionWithFailedAttempts() throws {
-        let constraints = RelayConstraints(location: .only(.hostname("se", "sto", "se6-wireguard")))
+        let constraints = RelayConstraints(
+            locations: .only(RelayLocations(locations: [.hostname("se", "sto", "se6-wireguard")]))
+        )
         let allPorts = portRanges.flatMap { $0 }
 
         var result = try RelaySelector.evaluate(
@@ -89,7 +101,9 @@ class RelaySelectorTests: XCTestCase {
     }
 
     func testClosestShadowsocksRelay() throws {
-        let constraints = RelayConstraints(location: .only(.city("se", "sto")))
+        let constraints = RelayConstraints(
+            locations: .only(RelayLocations(locations: [.city("se", "sto")]))
+        )
 
         let selectedRelay = RelaySelector.closestShadowsocksRelayConstrained(by: constraints, in: sampleRelays)
 
@@ -97,7 +111,9 @@ class RelaySelectorTests: XCTestCase {
     }
 
     func testClosestShadowsocksRelayIsRandomWhenNoContraintsAreSatisfied() throws {
-        let constraints = RelayConstraints(location: .only(.country("INVALID COUNTRY")))
+        let constraints = RelayConstraints(
+            locations: .only(RelayLocations(locations: [.country("INVALID COUNTRY")]))
+        )
 
         let selectedRelay = try XCTUnwrap(RelaySelector.closestShadowsocksRelayConstrained(
             by: constraints,
@@ -109,8 +125,9 @@ class RelaySelectorTests: XCTestCase {
 
     func testRelayFilterConstraintWithOwnedOwnership() throws {
         let filter = RelayFilter(ownership: .owned, providers: .any)
+
         let constraints = RelayConstraints(
-            location: .only(.hostname("se", "sto", "se6-wireguard")),
+            locations: .only(RelayLocations(locations: [.hostname("se", "sto", "se6-wireguard")])),
             filter: .only(filter)
         )
 
@@ -125,8 +142,9 @@ class RelaySelectorTests: XCTestCase {
 
     func testRelayFilterConstraintWithRentedOwnership() throws {
         let filter = RelayFilter(ownership: .rented, providers: .any)
+
         let constraints = RelayConstraints(
-            location: .only(.hostname("se", "sto", "se6-wireguard")),
+            locations: .only(RelayLocations(locations: [.hostname("se", "sto", "se6-wireguard")])),
             filter: .only(filter)
         )
 
@@ -141,10 +159,10 @@ class RelaySelectorTests: XCTestCase {
 
     func testRelayFilterConstraintWithCorrectProvider() throws {
         let provider = "31173"
-
         let filter = RelayFilter(ownership: .any, providers: .only([provider]))
+
         let constraints = RelayConstraints(
-            location: .only(.hostname("se", "sto", "se6-wireguard")),
+            locations: .only(RelayLocations(locations: [.hostname("se", "sto", "se6-wireguard")])),
             filter: .only(filter)
         )
 
@@ -159,10 +177,10 @@ class RelaySelectorTests: XCTestCase {
 
     func testRelayFilterConstraintWithIncorrectProvider() throws {
         let provider = "DataPacket"
-
         let filter = RelayFilter(ownership: .any, providers: .only([provider]))
+
         let constraints = RelayConstraints(
-            location: .only(.hostname("se", "sto", "se6-wireguard")),
+            locations: .only(RelayLocations(locations: [.hostname("se", "sto", "se6-wireguard")])),
             filter: .only(filter)
         )
 
diff --git a/ios/PacketTunnelCoreTests/AppMessageHandlerTests.swift b/ios/PacketTunnelCoreTests/AppMessageHandlerTests.swift
index 21291f6334..94dcbcf500 100644
--- a/ios/PacketTunnelCoreTests/AppMessageHandlerTests.swift
+++ b/ios/PacketTunnelCoreTests/AppMessageHandlerTests.swift
@@ -77,7 +77,9 @@ final class AppMessageHandlerTests: XCTestCase {
         let actor = PacketTunnelActorStub(reconnectExpectation: reconnectExpectation)
         let appMessageHandler = createAppMessageHandler(actor: actor)
 
-        let relayConstraints = RelayConstraints(location: .only(.hostname("se", "sto", "se6-wireguard")))
+        let relayConstraints = RelayConstraints(
+            locations: .only(RelayLocations(locations: [.hostname("se", "sto", "se6-wireguard")]))
+        )
         let selectorResult = try XCTUnwrap(try? RelaySelector.evaluate(
             relays: ServerRelaysResponseStubs.sampleRelays,
             constraints: relayConstraints,