Relay selector should use overridden IP addresses for relays

ios/MullvadREST/ApiHandlers/ServerRelaysResponse.swift

@@ -34,6 +34,19 @@ extension REST {
         public let ipv4AddrIn: IPv4Address
         public let weight: UInt64
         public let includeInCountry: Bool
+
+        public func override(ipv4AddrIn: IPv4Address?) -> Self {
+            return BridgeRelay(
+                hostname: hostname,
+                active: active,
+                owned: owned,
+                location: location,
+                provider: provider,
+                ipv4AddrIn: ipv4AddrIn ?? self.ipv4AddrIn,
+                weight: weight,
+                includeInCountry: includeInCountry
+            )
+        }
     }
 
     public struct ServerRelay: Codable, Equatable {
@@ -47,6 +60,21 @@ extension REST {
         public let ipv6AddrIn: IPv6Address
         public let publicKey: Data
         public let includeInCountry: Bool
+
+        public func override(ipv4AddrIn: IPv4Address?, ipv6AddrIn: IPv6Address?) -> Self {
+            return ServerRelay(
+                hostname: hostname,
+                active: active,
+                owned: owned,
+                location: location,
+                provider: provider,
+                weight: weight,
+                ipv4AddrIn: ipv4AddrIn ?? self.ipv4AddrIn,
+                ipv6AddrIn: ipv6AddrIn ?? self.ipv6AddrIn,
+                publicKey: publicKey,
+                includeInCountry: includeInCountry
+            )
+        }
     }
 
     public struct ServerWireguardTunnels: Codable, Equatable {

ios/MullvadREST/Relay/AnyRelay.swift

@@ -0,0 +1,29 @@
+//
+//  AnyRelay.swift
+//  MullvadREST
+//
+//  Created by Jon Petersson on 2024-01-31.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadTypes
+import Network
+
+public protocol AnyRelay {
+    var hostname: String { get }
+    var owned: Bool { get }
+    var location: String { get }
+    var provider: String { get }
+    var weight: UInt64 { get }
+    var active: Bool { get }
+    var includeInCountry: Bool { get }
+
+    func override(ipv4AddrIn: IPv4Address?, ipv6AddrIn: IPv6Address?) -> Self
+}
+
+extension REST.ServerRelay: AnyRelay {}
+extension REST.BridgeRelay: AnyRelay {
+    public func override(ipv4AddrIn: IPv4Address?, ipv6AddrIn: IPv6Address?) -> REST.BridgeRelay {
+        override(ipv4AddrIn: ipv4AddrIn)
+    }
+}

ios/MullvadREST/Relay/IPOverrideWrapper.swift

@@ -0,0 +1,72 @@
+//
+//  IPOverrideWrapper.swift
+//  MullvadREST
+//
+//  Created by Jon Petersson on 2024-02-05.
+//  Copyright © 2024 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadSettings
+import MullvadTypes
+
+public class IPOverrideWrapper: RelayCacheProtocol {
+    private let relayCache: RelayCacheProtocol
+    private let ipOverrideRepository: any IPOverrideRepositoryProtocol
+
+    public init(relayCache: RelayCacheProtocol, ipOverrideRepository: any IPOverrideRepositoryProtocol) {
+        self.relayCache = relayCache
+        self.ipOverrideRepository = ipOverrideRepository
+    }
+
+    public func read() throws -> CachedRelays {
+        let cache = try relayCache.read()
+        let relayResponse = apply(overrides: ipOverrideRepository.fetchAll(), to: cache.relays)
+
+        return CachedRelays(relays: relayResponse, updatedAt: cache.updatedAt)
+    }
+
+    public func write(record: CachedRelays) throws {
+        try relayCache.write(record: record)
+    }
+
+    private func apply(
+        overrides: [IPOverride],
+        to relayResponse: REST.ServerRelaysResponse
+    ) -> REST.ServerRelaysResponse {
+        let wireguard = relayResponse.wireguard
+        let bridge = relayResponse.bridge
+
+        let overridenWireguardRelays = wireguard.relays.map { relay in
+            return apply(overrides: overrides, to: relay)
+        }
+        let overridenBridgeRelays = bridge.relays.map { relay in
+            return apply(overrides: overrides, to: relay)
+        }
+
+        return REST.ServerRelaysResponse(
+            locations: relayResponse.locations,
+            wireguard: REST.ServerWireguardTunnels(
+                ipv4Gateway: wireguard.ipv4Gateway,
+                ipv6Gateway: wireguard.ipv6Gateway,
+                portRanges: wireguard.portRanges,
+                relays: overridenWireguardRelays
+            ),
+            bridge: REST.ServerBridges(
+                shadowsocks: bridge.shadowsocks,
+                relays: overridenBridgeRelays
+            )
+        )
+    }
+
+    private func apply<T: AnyRelay>(overrides: [IPOverride], to relay: T) -> T {
+        return overrides
+            .first { $0.hostname == relay.hostname }
+            .flatMap {
+                relay.override(
+                    ipv4AddrIn: $0.ipv4Address,
+                    ipv6AddrIn: $0.ipv6Address
+                )
+            }
+            ?? relay
+    }
+}

ios/MullvadREST/Relay/RelayCache.swift

@@ -14,6 +14,8 @@ public protocol RelayCacheProtocol {
     func write(record: CachedRelays) throws
 }
 
+/// - Warning: `RelayCache` should not be used directly. It should be used through `IPOverrideWrapper` to have
+/// ip overrides applied.
 public final class RelayCache: RelayCacheProtocol {
     private let fileCache: any FileCacheProtocol<CachedRelays>
 

ios/MullvadREST/Relay/RelaySelector.swift

@@ -306,19 +306,6 @@ public struct RelaySelectorResult: Codable, Equatable {
     public var location: Location
 }
 
-public protocol AnyRelay {
-    var hostname: String { get }
-    var owned: Bool { get }
-    var location: String { get }
-    var provider: String { get }
-    var weight: UInt64 { get }
-    var active: Bool { get }
-    var includeInCountry: Bool { get }
-}
-
-extension REST.ServerRelay: AnyRelay {}
-extension REST.BridgeRelay: AnyRelay {}
-
 private struct RelayWithLocation<T: AnyRelay> {
     let relay: T
     let serverLocation: Location

ios/MullvadSettings/IPOverride.swift

@@ -16,6 +16,10 @@ public struct RelayOverrides: Codable {
     }
 }
 
+public struct IPOverrideFormatError: LocalizedError {
+    public let errorDescription: String?
+}
+
 public struct IPOverride: Codable, Equatable {
     public let hostname: String
     public var ipv4Address: IPv4Address?
@@ -27,7 +31,7 @@ public struct IPOverride: Codable, Equatable {
         case ipv6Address = "ipv6_addr_in"
     }
 
-    init(hostname: String, ipv4Address: IPv4Address?, ipv6Address: IPv6Address?) throws {
+    public init(hostname: String, ipv4Address: IPv4Address?, ipv6Address: IPv6Address?) throws {
         self.hostname = hostname
         self.ipv4Address = ipv4Address
         self.ipv6Address = ipv6Address
@@ -49,7 +53,3 @@ public struct IPOverride: Codable, Equatable {
         }
     }
 }
-
-public struct IPOverrideFormatError: LocalizedError {
-    public let errorDescription: String?
-}

ios/MullvadSettings/IPOverrideRepository.swift

@@ -12,13 +12,13 @@ import MullvadLogging
 public protocol IPOverrideRepositoryProtocol {
     func add(_ overrides: [IPOverride])
     func fetchAll() -> [IPOverride]
-    func fetchByHostname(_ hostname: String) -> IPOverride?
     func deleteAll()
     func parse(data: Data) throws -> [IPOverride]
 }
 
 public class IPOverrideRepository: IPOverrideRepositoryProtocol {
     private let logger = Logger(label: "IPOverrideRepository")
+    private let readWriteLock = NSLock()
 
     public init() {}
 
@@ -54,13 +54,11 @@ public class IPOverrideRepository: IPOverrideRepositoryProtocol {
         return (try? readIpOverrides()) ?? []
     }
 
-    public func fetchByHostname(_ hostname: String) -> IPOverride? {
-        return fetchAll().first { $0.hostname == hostname }
-    }
-
     public func deleteAll() {
         do {
-            try SettingsManager.store.delete(key: .ipOverrides)
+            try readWriteLock.withLock {
+                try SettingsManager.store.delete(key: .ipOverrides)
+            }
         } catch {
             logger.error("Could not delete all overrides. \nError: \(error)")
         }
@@ -74,17 +72,20 @@ public class IPOverrideRepository: IPOverrideRepositoryProtocol {
     }
 
     private func readIpOverrides() throws -> [IPOverride] {
-        let parser = makeParser()
-        let data = try SettingsManager.store.read(key: .ipOverrides)
-
-        return try parser.parseUnversionedPayload(as: [IPOverride].self, from: data)
+        try readWriteLock.withLock {
+            let parser = makeParser()
+            let data = try SettingsManager.store.read(key: .ipOverrides)
+            return try parser.parseUnversionedPayload(as: [IPOverride].self, from: data)
+        }
     }
 
     private func writeIpOverrides(_ overrides: [IPOverride]) throws {
         let parser = makeParser()
         let data = try parser.produceUnversionedPayload(overrides)
 
-        try SettingsManager.store.write(data, for: .ipOverrides)
+        try readWriteLock.withLock {
+            try SettingsManager.store.write(data, for: .ipOverrides)
+        }
     }
 
     private func makeParser() -> SettingsParser {

ios/MullvadVPN.xcodeproj/project.pbxproj

@@ -486,6 +486,8 @@
 		7A3FD1B82AD54AE60042BEA6 /* TimeServerProxy.swift in Sources */ = {isa = PBXBuildFile; fileRef = 58BDEB9A2A98F58600F578F2 /* TimeServerProxy.swift */; };
 		7A42DEC92A05164100B209BE /* SettingsInputCell.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A42DEC82A05164100B209BE /* SettingsInputCell.swift */; };
 		7A516C2E2B6D357500BBD33D /* URL+Scoping.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A516C2D2B6D357500BBD33D /* URL+Scoping.swift */; };
+		7A516C3A2B7111A700BBD33D /* IPOverrideWrapper.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A516C392B7111A700BBD33D /* IPOverrideWrapper.swift */; };
+		7A516C3C2B712F0B00BBD33D /* IPOverrideWrapperTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A516C3B2B712F0B00BBD33D /* IPOverrideWrapperTests.swift */; };
 		7A5869952B32E9C700640D27 /* LinkButton.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A5869942B32E9C700640D27 /* LinkButton.swift */; };
 		7A5869972B32EA4500640D27 /* AppButton.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A5869962B32EA4500640D27 /* AppButton.swift */; };
 		7A58699B2B482FE200640D27 /* UITableViewCell+Disable.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7A58699A2B482FE200640D27 /* UITableViewCell+Disable.swift */; };
@@ -560,6 +562,8 @@
 		7AD0AA1D2AD6A86700119E10 /* PacketTunnelActorProtocol.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AD0AA192AD69B6E00119E10 /* PacketTunnelActorProtocol.swift */; };
 		7AD0AA1F2AD6C8B900119E10 /* URLRequestProxyProtocol.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AD0AA1E2AD6C8B900119E10 /* URLRequestProxyProtocol.swift */; };
 		7AD0AA212AD6CB0000119E10 /* URLRequestProxyStub.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AD0AA202AD6CB0000119E10 /* URLRequestProxyStub.swift */; };
+		7ADCB2D82B6A6EB300C88F89 /* AnyRelay.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7ADCB2D72B6A6EB300C88F89 /* AnyRelay.swift */; };
+		7ADCB2DA2B6A730400C88F89 /* IPOverrideRepositoryStub.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7ADCB2D92B6A730400C88F89 /* IPOverrideRepositoryStub.swift */; };
 		7AE044BB2A935726003915D8 /* Routing.h in Headers */ = {isa = PBXBuildFile; fileRef = 7A88DCD02A8FABBE00D2FF0E /* Routing.h */; settings = {ATTRIBUTES = (Public, ); }; };
 		7AEF7F1A2AD00F52006FE45D /* AppMessageHandler.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AEF7F192AD00F52006FE45D /* AppMessageHandler.swift */; };
 		7AF10EB22ADE859200C090B9 /* AlertViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 7AF10EB12ADE859200C090B9 /* AlertViewController.swift */; };
@@ -1666,6 +1670,8 @@
 		7A3FD1B42AD4465A0042BEA6 /* AppMessageHandlerTests.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = AppMessageHandlerTests.swift; sourceTree = "<group>"; };
 		7A42DEC82A05164100B209BE /* SettingsInputCell.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SettingsInputCell.swift; sourceTree = "<group>"; };
 		7A516C2D2B6D357500BBD33D /* URL+Scoping.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "URL+Scoping.swift"; sourceTree = "<group>"; };
+		7A516C392B7111A700BBD33D /* IPOverrideWrapper.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IPOverrideWrapper.swift; sourceTree = "<group>"; };
+		7A516C3B2B712F0B00BBD33D /* IPOverrideWrapperTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IPOverrideWrapperTests.swift; sourceTree = "<group>"; };
 		7A5869942B32E9C700640D27 /* LinkButton.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LinkButton.swift; sourceTree = "<group>"; };
 		7A5869962B32EA4500640D27 /* AppButton.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppButton.swift; sourceTree = "<group>"; };
 		7A58699A2B482FE200640D27 /* UITableViewCell+Disable.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = "UITableViewCell+Disable.swift"; sourceTree = "<group>"; };
@@ -1732,6 +1738,8 @@
 		7AD0AA1B2AD6A63F00119E10 /* PacketTunnelActorStub.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = PacketTunnelActorStub.swift; sourceTree = "<group>"; };
 		7AD0AA1E2AD6C8B900119E10 /* URLRequestProxyProtocol.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = URLRequestProxyProtocol.swift; sourceTree = "<group>"; };
 		7AD0AA202AD6CB0000119E10 /* URLRequestProxyStub.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = URLRequestProxyStub.swift; sourceTree = "<group>"; };
+		7ADCB2D72B6A6EB300C88F89 /* AnyRelay.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AnyRelay.swift; sourceTree = "<group>"; };
+		7ADCB2D92B6A730400C88F89 /* IPOverrideRepositoryStub.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IPOverrideRepositoryStub.swift; sourceTree = "<group>"; };
 		7AEF7F192AD00F52006FE45D /* AppMessageHandler.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = AppMessageHandler.swift; sourceTree = "<group>"; };
 		7AF10EB12ADE859200C090B9 /* AlertViewController.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = AlertViewController.swift; sourceTree = "<group>"; };
 		7AF10EB32ADE85BC00C090B9 /* RelayFilterCoordinator.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = RelayFilterCoordinator.swift; sourceTree = "<group>"; };
@@ -2745,8 +2753,10 @@
 				58B0A2A4238EE67E00BC001D /* Info.plist */,
 				A9B6AC192ADE8FBB00F7802A /* InMemorySettingsStore.swift */,
 				F07BF2572A26112D00042943 /* InputTextFormatterTests.swift */,
+				7ADCB2D92B6A730400C88F89 /* IPOverrideRepositoryStub.swift */,
 				7A5869C22B5820CE00640D27 /* IPOverrideRepositoryTests.swift */,
 				7AB4CCB82B69097E006037F5 /* IPOverrideTests.swift */,
+				7A516C3B2B712F0B00BBD33D /* IPOverrideWrapperTests.swift */,
 				A9B6AC172ADE8F4300F7802A /* MigrationManagerTests.swift */,
 				58C3FA652A38549D006A450A /* MockFileCache.swift */,
 				F09D04B42AE93CB6003D4F89 /* OutgoingConnectionProxy+Stub.swift */,
@@ -3461,8 +3471,10 @@
 		F0DC779F2B2222D20087F09D /* Relay */ = {
 			isa = PBXGroup;
 			children = (
+				7ADCB2D72B6A6EB300C88F89 /* AnyRelay.swift */,
 				585DA87626B024A600B8C587 /* CachedRelays.swift */,
 				F0DDE4272B220A15006B57A7 /* Haversine.swift */,
+				7A516C392B7111A700BBD33D /* IPOverrideWrapper.swift */,
 				F0DDE4292B220A15006B57A7 /* Midpoint.swift */,
 				5820675A26E6576800655B05 /* RelayCache.swift */,
 				F0DDE4282B220A15006B57A7 /* RelaySelector.swift */,
@@ -4438,6 +4450,7 @@
 				06799AF228F98E4800ACD94E /* RESTAccessTokenManager.swift in Sources */,
 				A90763B12B2857D50045ADF0 /* Socks5Endpoint.swift in Sources */,
 				06799AF328F98E4800ACD94E /* RESTAuthenticationProxy.swift in Sources */,
+				7A516C3A2B7111A700BBD33D /* IPOverrideWrapper.swift in Sources */,
 				F0DDE4142B220458006B57A7 /* ShadowSocksProxy.swift in Sources */,
 				A90763B62B2857D50045ADF0 /* Socks5ConnectNegotiation.swift in Sources */,
 				F06045E62B231EB700B2D37A /* URLSessionTransport.swift in Sources */,
@@ -4459,6 +4472,7 @@
 				A90763B42B2857D50045ADF0 /* NWConnection+Extensions.swift in Sources */,
 				F06045EA2B23217E00B2D37A /* ShadowsocksTransport.swift in Sources */,
 				06799AFC28F98EE300ACD94E /* AddressCache.swift in Sources */,
+				7ADCB2D82B6A6EB300C88F89 /* AnyRelay.swift in Sources */,
 				06799AF028F98E4800ACD94E /* REST.swift in Sources */,
 				06799ADF28F98E4800ACD94E /* RESTDevicesProxy.swift in Sources */,
 				06799ADA28F98E4800ACD94E /* RESTResponseHandler.swift in Sources */,
@@ -4574,6 +4588,7 @@
 				7A6F2FA52AFA3CB2006D0856 /* AccountExpiryTests.swift in Sources */,
 				A9A5FA082ACB05160083449F /* StorePaymentBlockObserver.swift in Sources */,
 				A9E0317C2ACBFC7E0095D843 /* TunnelStore+Stubs.swift in Sources */,
+				7A516C3C2B712F0B00BBD33D /* IPOverrideWrapperTests.swift in Sources */,
 				A9A5FA092ACB05160083449F /* SendStoreReceiptOperation.swift in Sources */,
 				A9A5FA0A2ACB05160083449F /* StorePaymentEvent.swift in Sources */,
 				A9A5FA0B2ACB05160083449F /* StorePaymentManager.swift in Sources */,
@@ -4624,6 +4639,7 @@
 				A9A5FA2F2ACB05160083449F /* FixedWidthIntegerArithmeticsTests.swift in Sources */,
 				A9A5FA302ACB05160083449F /* InputTextFormatterTests.swift in Sources */,
 				F0B0E6972AFE6E7E001DC66B /* XCTest+Async.swift in Sources */,
+				7ADCB2DA2B6A730400C88F89 /* IPOverrideRepositoryStub.swift in Sources */,
 				A9A5FA312ACB05160083449F /* MockFileCache.swift in Sources */,
 				A9A5FA322ACB05160083449F /* RelayCacheTests.swift in Sources */,
 				A9A5FA332ACB05160083449F /* RelaySelectorTests.swift in Sources */,

ios/MullvadVPN/AppDelegate.swift

@@ -45,6 +45,7 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
     private(set) var accessMethodRepository = AccessMethodRepository()
     private(set) var shadowsocksLoader: ShadowsocksLoaderProtocol!
     private(set) var configuredTransportProvider: ProxyConfigurationTransportProvider!
+    private(set) var ipOverrideRepository = IPOverrideRepository()
 
     // MARK: - Application lifecycle
 
@@ -66,8 +67,16 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
 
         setUpProxies(containerURL: containerURL)
 
-        let relayCache = RelayCache(cacheDirectory: containerURL)
-        relayCacheTracker = RelayCacheTracker(relayCache: relayCache, application: application, apiProxy: apiProxy)
+        let ipOverrideWrapper = IPOverrideWrapper(
+            relayCache: RelayCache(cacheDirectory: containerURL),
+            ipOverrideRepository: ipOverrideRepository
+        )
+
+        relayCacheTracker = RelayCacheTracker(
+            relayCache: ipOverrideWrapper,
+            application: application,
+            apiProxy: apiProxy
+        )
 
         addressCacheTracker = AddressCacheTracker(application: application, apiProxy: apiProxy, store: addressCache)
 
@@ -93,7 +102,7 @@ class AppDelegate: UIResponder, UIApplicationDelegate, UNUserNotificationCenterD
 
         shadowsocksLoader = ShadowsocksLoader(
             shadowsocksCache: shadowsocksCache,
-            relayCache: relayCache,
+            relayCache: ipOverrideWrapper,
             constraintsUpdater: constraintsUpdater
         )