Migrating RelayCache to REST relays

7 files changed, 91 insertions, 138 deletions

diff --git a/ios/MullvadVPN/AutomaticKeyRotationManager.swift b/ios/MullvadVPN/AutomaticKeyRotationManager.swift
index 67df0cd020..9d6b8677a0 100644
--- a/ios/MullvadVPN/AutomaticKeyRotationManager.swift
+++ b/ios/MullvadVPN/AutomaticKeyRotationManager.swift
@@ -177,8 +177,13 @@ class AutomaticKeyRotationManager {
             self.dispatchQueue.async {
                 let updateResult = result.mapError { (error) -> Error in
                     return .rest(error)
-                }.flatMap { (addresses) -> Result<TunnelSettings, Error> in
-                    self.updateTunnelSettings(privateKey: newPrivateKey, addresses: addresses)
+                }.flatMap { (response) -> Result<TunnelSettings, Error> in
+                    let addresses = WireguardAssociatedAddresses(
+                        ipv4Address: response.ipv4Address,
+                        ipv6Address: response.ipv6Address
+                    )
+
+                    return self.updateTunnelSettings(privateKey: newPrivateKey, addresses: addresses)
                 }
                 completionHandler(updateResult)
             }
diff --git a/ios/MullvadVPN/MullvadRest.swift b/ios/MullvadVPN/MullvadRest.swift
index 7d587f52d6..90cf44fc20 100644
--- a/ios/MullvadVPN/MullvadRest.swift
+++ b/ios/MullvadVPN/MullvadRest.swift
@@ -452,35 +452,36 @@ struct AccountResponse: Decodable, RestResponse {
     let expires: Date
 }
 
-struct ServerLocation: Decodable {
+struct ServerLocation: Codable {
     let country: String
     let city: String
     let latitude: Double
     let longitude: Double
 }
 
-struct ServerRelay: Decodable {
+struct ServerRelay: Codable {
     let hostname: String
     let active: Bool
     let owned: Bool
     let location: String
     let provider: String
-    let ipv4AddrIn: IPv4Address
     let weight: Int32
+    let ipv4AddrIn: IPv4Address
+    let ipv6AddrIn: IPv6Address
+    let publicKey: Data
     let includeInCountry: Bool
 }
 
-struct ServerWireguardTunnel: Decodable {
+struct ServerWireguardTunnels: Codable {
     let ipv4Gateway: IPv4Address
     let ipv6Gateway: IPv6Address
-    let publicKey: Data
     let portRanges: [ClosedRange<UInt16>]
     let relays: [ServerRelay]
 }
 
-struct ServerRelaysResponse: Decodable, RestResponse {
+struct ServerRelaysResponse: Codable, RestResponse {
     let locations: [String: ServerLocation]
-    let wireguard: [ServerWireguardTunnel]
+    let wireguard: ServerWireguardTunnels
 }
 
 struct PushWireguardKeyRequest: Encodable, RestPayload {
diff --git a/ios/MullvadVPN/RelayCache.swift b/ios/MullvadVPN/RelayCache.swift
index 89c7a4c480..e72c8e9106 100644
--- a/ios/MullvadVPN/RelayCache.swift
+++ b/ios/MullvadVPN/RelayCache.swift
@@ -20,7 +20,7 @@ enum RelayCacheError: ChainedError {
     case writeCache(Error)
     case encodeCache(Error)
     case decodeCache(Error)
-    case rpc(MullvadRpc.Error)
+    case rest(RestError)
 
     var errorDescription: String? {
         switch self {
@@ -36,14 +36,14 @@ enum RelayCacheError: ChainedError {
             return "Decode pre-bundled relays error"
         case .writeCache:
             return "Write cache error"
-        case .rpc:
-            return "RPC error"
+        case .rest:
+            return "REST error"
         }
     }
 }
 
 protocol RelayCacheObserver: class {
-    func relayCache(_ relayCache: RelayCache, didUpdateCachedRelayList cachedRelayList: CachedRelayList)
+    func relayCache(_ relayCache: RelayCache, didUpdateCachedRelays cachedRelays: CachedRelays)
 }
 
 private class AnyRelayCacheObserver: WeakObserverBox, RelayCacheObserver {
@@ -56,8 +56,8 @@ private class AnyRelayCacheObserver: WeakObserverBox, RelayCacheObserver {
         self.inner = inner
     }
 
-    func relayCache(_ relayCache: RelayCache, didUpdateCachedRelayList cachedRelayList: CachedRelayList) {
-        inner?.relayCache(relayCache, didUpdateCachedRelayList: cachedRelayList)
+    func relayCache(_ relayCache: RelayCache, didUpdateCachedRelays cachedRelays: CachedRelays) {
+        inner?.relayCache(relayCache, didUpdateCachedRelays: cachedRelays)
     }
 
     static func == (lhs: AnyRelayCacheObserver, rhs: AnyRelayCacheObserver) -> Bool {
@@ -155,10 +155,10 @@ class RelayCache {
     }
 
     /// Read the relay cache from disk
-    func read(completionHandler: @escaping (Result<CachedRelayList, RelayCacheError>) -> Void) {
+    func read(completionHandler: @escaping (Result<CachedRelays, RelayCacheError>) -> Void) {
         dispatchQueue.async {
             let result = Self.read(cacheFileURL: self.cacheFileURL)
-                .flatMapError { (error) -> Result<CachedRelayList, RelayCacheError> in
+                .flatMapError { (error) -> Result<CachedRelays, RelayCacheError> in
                     if case .readCache(let ioError as CocoaError) = error, ioError.code == .fileReadNoSuchFile {
                         return Self.readPrebundledRelays(fileURL: Self.preBundledRelaysFileURL)
                     } else {
@@ -201,19 +201,19 @@ class RelayCache {
 
     private func downloadRelays() {
         let taskResult = makeDownloadTask { (result) in
-            let result = result.flatMap { (relayList) -> Result<CachedRelayList, RelayCacheError> in
-                let cachedRelayList = CachedRelayList(relayList: relayList, updatedAt: Date())
+            let result = result.flatMap { (relays) -> Result<CachedRelays, RelayCacheError> in
+                let cachedRelays = CachedRelays(relays: relays, updatedAt: Date())
 
-                return Self.write(cacheFileURL: self.cacheFileURL, record: cachedRelayList)
-                    .map { cachedRelayList }
+                return Self.write(cacheFileURL: self.cacheFileURL, record: cachedRelays)
+                    .map { cachedRelays }
             }
 
             switch result {
-            case .success(let cachedRelayList):
-                os_log(.default, "Downloaded %d relays", cachedRelayList.relayList.numRelays)
+            case .success(let cachedRelays):
+                os_log(.default, "Downloaded %d relays", cachedRelays.relays.wireguard.relays.count)
 
                 self.observerList.forEach { (observer) in
-                    observer.relayCache(self, didUpdateCachedRelayList: cachedRelayList)
+                    observer.relayCache(self, didUpdateCachedRelays: cachedRelays)
                 }
 
             case .failure(let error):
@@ -248,52 +248,19 @@ class RelayCache {
         self.timerSource = timerSource
     }
 
-    private func makeDownloadTask(completionHandler: @escaping (Result<RelayList, RelayCacheError>) -> Void) -> Result<URLSessionDataTask, RestError> {
+    private func makeDownloadTask(completionHandler: @escaping (Result<ServerRelaysResponse, RelayCacheError>) -> Void) -> Result<URLSessionDataTask, RestError> {
         return rest.getRelays().dataTask(payload: EmptyPayload()) { (result) in
-            let result = result
-                .map(Self.filterRelayList)
-                .mapError { RelayCacheError.rpc($0) }
-
             self.dispatchQueue.async {
-                completionHandler(result)
+                completionHandler(result.mapError { RelayCacheError.rest($0) })
             }
         }
     }
 
     // MARK: - Private class methods
 
-    /// Filters the given `RelayList` removing empty leaf nodes, relays without Wireguard tunnels or
-    /// Wireguard tunnels without any available ports.
-    private class func filterRelayList(_ relayList: RelayList) -> RelayList {
-        let filteredCountries = relayList.countries
-            .map { (country) -> RelayList.Country in
-                var filteredCountry = country
-
-                filteredCountry.cities = country.cities.map { (city) -> RelayList.City in
-                    var filteredCity = city
-
-                    filteredCity.relays = city.relays
-                        .map { (relay) -> RelayList.Relay in
-                            var filteredRelay = relay
-
-                            // filter out tunnels without ports
-                            filteredRelay.tunnels?.wireguard = relay.tunnels?.wireguard?
-                                .filter { !$0.portRanges.isEmpty }
-
-                            return filteredRelay
-                    }.filter { $0.tunnels?.wireguard.flatMap { !$0.isEmpty } ?? false }
-
-                    return filteredCity
-                }.filter { !$0.relays.isEmpty }
-
-                return filteredCountry
-        }.filter { !$0.cities.isEmpty }
-
-        return RelayList(countries: filteredCountries)
-    }
     /// Safely read the cache file from disk using file coordinator
-    private class func read(cacheFileURL: URL) -> Result<CachedRelayList, RelayCacheError> {
-        var result: Result<CachedRelayList, RelayCacheError>?
+    private class func read(cacheFileURL: URL) -> Result<CachedRelays, RelayCacheError> {
+        var result: Result<CachedRelays, RelayCacheError>?
         let fileCoordinator = NSFileCoordinator(filePresenter: nil)
 
         let accessor = { (fileURLForReading: URL) -> Void in
@@ -301,7 +268,7 @@ class RelayCache {
             result = Result { try Data(contentsOf: fileURLForReading) }
                 .mapError { RelayCacheError.readCache($0) }
                 .flatMap { (data) in
-                    Result { try JSONDecoder().decode(CachedRelayList.self, from: data) }
+                    Result { try JSONDecoder().decode(CachedRelays.self, from: data) }
                         .mapError { RelayCacheError.decodeCache($0) }
                 }
         }
@@ -319,15 +286,15 @@ class RelayCache {
         return result!
     }
 
-    private class func readPrebundledRelays(fileURL: URL) -> Result<CachedRelayList, RelayCacheError> {
+    private class func readPrebundledRelays(fileURL: URL) -> Result<CachedRelays, RelayCacheError> {
         return Result { try Data(contentsOf: fileURL) }
             .mapError { RelayCacheError.readPrebundledRelays($0) }
-            .flatMap { (data) -> Result<CachedRelayList, RelayCacheError> in
-                return Result { try MullvadRpc.makeJSONDecoder().decode(RelayList.self, from: data) }
+            .flatMap { (data) -> Result<CachedRelays, RelayCacheError> in
+                return Result { try MullvadRest.makeJSONDecoder().decode(ServerRelaysResponse.self, from: data) }
                     .mapError { RelayCacheError.decodePrebundledRelays($0) }
-                    .map { (relayList) -> CachedRelayList in
-                        return CachedRelayList(
-                            relayList: Self.filterRelayList(relayList),
+                    .map { (relays) -> CachedRelays in
+                        return CachedRelays(
+                            relays: relays,
                             updatedAt: Date(timeIntervalSince1970: 0)
                         )
                 }
@@ -335,7 +302,7 @@ class RelayCache {
     }
 
     /// Safely write the cache file on disk using file coordinator
-    private class func write(cacheFileURL: URL, record: CachedRelayList) -> Result<(), RelayCacheError> {
+    private class func write(cacheFileURL: URL, record: CachedRelays) -> Result<(), RelayCacheError> {
         var result: Result<(), RelayCacheError>?
         let fileCoordinator = NSFileCoordinator(filePresenter: nil)
 
@@ -393,9 +360,9 @@ class RelayCache {
 }
 
 /// A struct that represents the relay cache on disk
-struct CachedRelayList: Codable {
+struct CachedRelays: Codable {
     /// The relay list stored within the cache entry
-    var relayList: RelayList
+    var relays: ServerRelaysResponse
 
     /// The date when this cache was last updated
     var updatedAt: Date
diff --git a/ios/MullvadVPN/RelayList.swift b/ios/MullvadVPN/RelayList.swift
index cf9fe7edd7..d4384696d7 100644
--- a/ios/MullvadVPN/RelayList.swift
+++ b/ios/MullvadVPN/RelayList.swift
@@ -27,9 +27,11 @@ struct RelayList: Codable {
 
     struct Relay: Codable {
         var hostname: String
+        var provider: String
         var ipv4AddrIn: IPv4Address
         var includeInCountry: Bool
         var active: Bool
+        var owned: Bool
         var weight: Int32
         var tunnels: Tunnels?
     }
diff --git a/ios/MullvadVPN/RelaySelector.swift b/ios/MullvadVPN/RelaySelector.swift
index 716a62fc9a..1ee17a31ac 100644
--- a/ios/MullvadVPN/RelaySelector.swift
+++ b/ios/MullvadVPN/RelaySelector.swift
@@ -10,57 +10,51 @@ import Foundation
 import Network
 
 struct RelaySelectorResult {
-    var relay: RelayList.Relay
-    var tunnel: RelayList.WireguardTunnel
     var endpoint: MullvadEndpoint
+    var relay: ServerRelay
     var location: Location
 }
 
 private struct RelayWithLocation {
-    var relay: RelayList.Relay
+    var relay: ServerRelay
     var location: Location
 }
 
 struct RelaySelector {
 
-    private let relayList: RelayList
+    private let relays: ServerRelaysResponse
 
-    init(relayList: RelayList) {
-        self.relayList = relayList
+    init(relays: ServerRelaysResponse) {
+        self.relays = relays
     }
 
     func evaluate(with constraints: RelayConstraints) -> RelaySelectorResult? {
-        let relays = Self.applyConstraints(constraints, relays: Self.parseRelayList(self.relayList))
-        let totalWeight = relays.reduce(0) { $0 + $1.relay.weight }
+        let filteredRelays = Self.applyConstraints(constraints, relays: Self.parseRelaysResponse(self.relays))
+        let totalWeight = filteredRelays.reduce(0) { $0 + $1.relay.weight }
 
         guard totalWeight > 0 else { return nil }
         guard var i = (0...totalWeight).randomElement() else { return nil }
 
-        let relayWithLocation = relays.first { (relayWithLocation) -> Bool in
+        let relayWithLocation = filteredRelays.first { (relayWithLocation) -> Bool in
             i -= relayWithLocation.relay.weight
             return i <= 0
         }.unsafelyUnwrapped
 
-        guard let tunnel = relayWithLocation.relay.tunnels?.wireguard?.randomElement() else {
-            return nil
-        }
-
-        guard let port = tunnel.portRanges.randomElement()?.randomElement() else {
+        guard let port = relays.wireguard.portRanges.randomElement()?.randomElement() else {
             return nil
         }
 
         let endpoint = MullvadEndpoint(
             ipv4Relay: IPv4Endpoint(ip: relayWithLocation.relay.ipv4AddrIn, port: port),
-            ipv6Relay: nil,
-            ipv4Gateway: tunnel.ipv4Gateway,
-            ipv6Gateway: tunnel.ipv6Gateway,
-            publicKey: tunnel.publicKey
+            ipv6Relay: IPv6Endpoint(ip: relayWithLocation.relay.ipv6AddrIn, port: port),
+            ipv4Gateway: relays.wireguard.ipv4Gateway,
+            ipv6Gateway: relays.wireguard.ipv6Gateway,
+            publicKey: relayWithLocation.relay.publicKey
         )
 
         return RelaySelectorResult(
-            relay: relayWithLocation.relay,
-            tunnel: tunnel,
             endpoint: endpoint,
+            relay: relayWithLocation.relay,
             location: relayWithLocation.location
         )
     }
@@ -87,45 +81,29 @@ struct RelaySelector {
                         relayWithLocation.relay.hostname == hostname
                 }
             }
-        }.map({ (relayWithLocation) -> RelayWithLocation in
-            var filteredRelay = relayWithLocation
-            let wireguardTunnels = filteredRelay.relay.tunnels?.wireguard?
-                .filter { !$0.portRanges.isEmpty }
-
-            filteredRelay.relay.tunnels?.wireguard = wireguardTunnels
-
-            return filteredRelay
-        }).filter { (relayWithLocation) -> Bool in
-            guard let wireguardTunnels = relayWithLocation.relay.tunnels?.wireguard else { return false }
-
-            return relayWithLocation.relay.active && !wireguardTunnels.isEmpty
+        }.filter { (relayWithLocation) -> Bool in
+            return relayWithLocation.relay.active
         }
     }
 
-    private static func parseRelayList(_ relayList: RelayList) -> [RelayWithLocation] {
-        var relays = [RelayWithLocation]()
+    private static func parseRelaysResponse(_ response: ServerRelaysResponse) -> [RelayWithLocation] {
+        return response.wireguard.relays.compactMap { (serverRelay) -> RelayWithLocation? in
+            guard let serverLocation = response.locations[serverRelay.location] else { return nil }
 
-        for country in relayList.countries {
-            for city in country.cities {
-                for relay in city.relays {
-                    let location = Location(
-                        country: country.name,
-                        countryCode: country.code,
-                        city: city.name,
-                        cityCode: city.code,
-                        latitude: city.latitude,
-                        longitude: city.longitude
-                    )
-                    let relayWithLocation = RelayWithLocation(
-                        relay: relay,
-                        location: location
-                    )
-                    relays.append(relayWithLocation)
-                }
-            }
-        }
+            let locationComponents = serverRelay.location.split(separator: "-")
+            guard locationComponents.count > 1 else { return nil }
+            
+            let location = Location(
+                country: serverLocation.country,
+                countryCode: String(locationComponents[0]),
+                city: serverLocation.city,
+                cityCode: String(locationComponents[1]),
+                latitude: serverLocation.latitude,
+                longitude: serverLocation.longitude
+            )
 
-        return relays
+            return RelayWithLocation(relay: serverRelay, location: location)
+        }
     }
 
 }
diff --git a/ios/MullvadVPN/SelectLocationController.swift b/ios/MullvadVPN/SelectLocationController.swift
index c1c8b7d938..26aae03a95 100644
--- a/ios/MullvadVPN/SelectLocationController.swift
+++ b/ios/MullvadVPN/SelectLocationController.swift
@@ -99,12 +99,12 @@ class SelectLocationController: UITableViewController, RelayCacheObserver {
 
     // MARK: - RelayCacheObserver
 
-    func relayCache(_ relayCache: RelayCache, didUpdateCachedRelayList cachedRelayList: CachedRelayList) {
-        self.didReceiveCachedRelays(cachedRelayList: cachedRelayList) { (result) in
+    func relayCache(_ relayCache: RelayCache, didUpdateCachedRelays cachedRelays: CachedRelays) {
+        self.didReceiveCachedRelays(cachedRelays) { (result) in
             DispatchQueue.main.async {
                 switch result {
-                case .success(let (cachedRelayList, relayConstraints)):
-                    self.didReceive(relayList: cachedRelayList.relayList, relayConstraints: relayConstraints)
+                case .success(let (cachedRelays, relayConstraints)):
+                    self.didReceiveCachedRelays(cachedRelays, relayConstraints: relayConstraints)
 
                 case .failure(let error):
                     error.logChain()
@@ -119,8 +119,8 @@ class SelectLocationController: UITableViewController, RelayCacheObserver {
         fetchRelays { (result) in
             DispatchQueue.main.async {
                 switch result {
-                case .success(let (cachedRelayList, relayConstraints)):
-                    self.didReceive(relayList: cachedRelayList.relayList, relayConstraints: relayConstraints)
+                case .success(let (cachedRelays, relayConstraints)):
+                    self.didReceiveCachedRelays(cachedRelays, relayConstraints: relayConstraints)
 
                 case .failure(let error):
                     error.logChain()
@@ -129,11 +129,11 @@ class SelectLocationController: UITableViewController, RelayCacheObserver {
         }
     }
 
-    private func fetchRelays(completionHandler: @escaping (Result<(CachedRelayList, RelayConstraints), Error>) -> Void) {
+    private func fetchRelays(completionHandler: @escaping (Result<(CachedRelays, RelayConstraints), Error>) -> Void) {
         RelayCache.shared.read { (result) in
             switch result {
-            case .success(let cachedRelayList):
-                self.didReceiveCachedRelays(cachedRelayList: cachedRelayList, completionHandler: completionHandler)
+            case .success(let cachedRelays):
+                self.didReceiveCachedRelays(cachedRelays, completionHandler: completionHandler)
 
             case .failure(let error):
                 completionHandler(.failure(.loadRelayList(error)))
@@ -141,17 +141,17 @@ class SelectLocationController: UITableViewController, RelayCacheObserver {
         }
     }
 
-    private func didReceiveCachedRelays(cachedRelayList: CachedRelayList, completionHandler: @escaping (Result<(CachedRelayList, RelayConstraints), Error>) -> Void) {
+    private func didReceiveCachedRelays(_ cachedRelays: CachedRelays, completionHandler: @escaping (Result<(CachedRelays, RelayConstraints), Error>) -> Void) {
         TunnelManager.shared.getRelayConstraints { (result) in
             let result = result
-                .map { (cachedRelayList, $0) }
+                .map { (cachedRelays, $0) }
                 .mapError { Error.getRelayConstraints($0) }
 
             completionHandler(result)
         }
     }
 
-    private func didReceive(relayList: RelayList, relayConstraints: RelayConstraints) {
+    private func didReceiveCachedRelays(_ cachedRelays: CachedRelays, relayConstraints: RelayConstraints) {
         self.relayList = relayList.sorted()
         self.relayConstraints = relayConstraints
 
diff --git a/ios/PacketTunnel/PacketTunnelProvider.swift b/ios/PacketTunnel/PacketTunnelProvider.swift
index d9e40ed496..5c5da30b15 100644
--- a/ios/PacketTunnel/PacketTunnelProvider.swift
+++ b/ios/PacketTunnel/PacketTunnelProvider.swift
@@ -459,7 +459,7 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
         RelayCache.shared.read { (result) in
             switch result {
             case .success(let cachedRelayList):
-                let relaySelector = RelaySelector(relayList: cachedRelayList.relayList)
+                let relaySelector = RelaySelector(relays: cachedRelayList.relays)
 
                 if let selectorResult = relaySelector.evaluate(with: relayConstraints) {
                     completionHandler(.success(selectorResult))