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 copyWith(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 copyWith(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/RelaySelector.swift

@@ -7,10 +7,21 @@
 //
 
 import Foundation
+import MullvadSettings
 import MullvadTypes
+import Network
 
 private let defaultPort: UInt16 = 53
 
+///
+/// Tool for selecting a relay depending on multiple factors and constraints, such as provider filtering,
+/// custom port, distance etc.
+///
+/// - Important: Any public function returning a relay (ie. ``REST.ServerRelay`` or ``REST.BridgeRelay``)
+/// must apply any existing IP overrides before returning.
+///
+/// - Seealso: ``RelaySelector.applyIpOverrides()``
+///
 public enum RelaySelector {
     /**
      Returns random shadowsocks TCP bridge, otherwise `nil` if there are no shadowdsocks bridges.
@@ -19,15 +30,6 @@ public enum RelaySelector {
         relays.bridge.shadowsocks.filter { $0.protocol == "tcp" }.randomElement()
     }
 
-    /// Return a random Shadowsocks bridge relay, or `nil` if no relay were found.
-    ///
-    /// Non `active` relays are filtered out.
-    /// - Parameter relays: The list of relays to randomly select from.
-    /// - Returns: A Shadowsocks relay or `nil` if no active relay were found.
-    public static func shadowsocksRelay(from relaysResponse: REST.ServerRelaysResponse) -> REST.BridgeRelay? {
-        relaysResponse.bridge.relays.filter { $0.active }.randomElement()
-    }
-
     /// Returns the closest Shadowsocks relay using the given `constraints`, or a random relay if `constraints` were
     /// unsatisfiable.
     ///
@@ -41,7 +43,11 @@ public enum RelaySelector {
     ) -> REST.BridgeRelay? {
         let mappedBridges = mapRelays(relays: relaysResponse.bridge.relays, locations: relaysResponse.locations)
         let filteredRelays = applyConstraints(constraints, relays: mappedBridges)
-        guard filteredRelays.isEmpty == false else { return shadowsocksRelay(from: relaysResponse) }
+
+        guard filteredRelays.isEmpty == false else {
+            let relay = shadowsocksRelay(from: relaysResponse)
+            return relay.flatMap { applyIpOverrides(to: $0) }
+        }
 
         // Compute the midpoint location from all the filtered relays
         // Take *either* the first five relays, OR the relays below maximum bridge distance
@@ -77,7 +83,8 @@ public enum RelaySelector {
             UInt64(1 + greatestDistance - relay.distance)
         })
 
-        return randomRelay?.relay ?? filteredRelays.randomElement()?.relay
+        let relayToReturn = randomRelay?.relay ?? filteredRelays.randomElement()?.relay
+        return relayToReturn.flatMap { applyIpOverrides(to: $0) }
     }
 
     /**
@@ -97,10 +104,15 @@ public enum RelaySelector {
             numberOfFailedAttempts: numberOfFailedAttempts
         )
 
-        guard let relayWithLocation = pickRandomRelayByWeight(relays: filteredRelays), let port else {
+        guard var relayWithLocation = pickRandomRelayByWeight(relays: filteredRelays), let port else {
             throw NoRelaysSatisfyingConstraintsError()
         }
 
+        relayWithLocation = RelayWithLocation(
+            relay: applyIpOverrides(to: relayWithLocation.relay),
+            serverLocation: relayWithLocation.serverLocation
+        )
+
         let endpoint = MullvadEndpoint(
             ipv4Relay: IPv4Endpoint(
                 ip: relayWithLocation.relay.ipv4AddrIn,
@@ -135,6 +147,15 @@ public enum RelaySelector {
         }
     }
 
+    /// Return a random Shadowsocks bridge relay, or `nil` if no relay were found.
+    ///
+    /// Non `active` relays are filtered out.
+    /// - Parameter relays: The list of relays to randomly select from.
+    /// - Returns: A Shadowsocks relay or `nil` if no active relay were found.
+    private static func shadowsocksRelay(from relaysResponse: REST.ServerRelaysResponse) -> REST.BridgeRelay? {
+        relaysResponse.bridge.relays.filter { $0.active }.randomElement()
+    }
+
     /// Produce a list of `RelayWithLocation` items satisfying the given constraints
     private static func applyConstraints<T: AnyRelay>(
         _ constraints: RelayConstraints,
@@ -157,16 +178,16 @@ public enum RelaySelector {
                 switch relayConstraint {
                 case let .country(countryCode):
                     return relayWithLocation.serverLocation.countryCode == countryCode &&
-                        relayWithLocation.relay.includeInCountry
+                    relayWithLocation.relay.includeInCountry
 
                 case let .city(countryCode, cityCode):
                     return relayWithLocation.serverLocation.countryCode == countryCode &&
-                        relayWithLocation.serverLocation.cityCode == cityCode
+                    relayWithLocation.serverLocation.cityCode == cityCode
 
                 case let .hostname(countryCode, cityCode, hostname):
                     return relayWithLocation.serverLocation.countryCode == countryCode &&
-                        relayWithLocation.serverLocation.cityCode == cityCode &&
-                        relayWithLocation.relay.hostname == hostname
+                    relayWithLocation.serverLocation.cityCode == cityCode &&
+                    relayWithLocation.relay.hostname == hostname
                 }
             }
         }.filter { relayWithLocation -> Bool in
@@ -194,8 +215,23 @@ public enum RelaySelector {
         }
     }
 
+    private static func applyIpOverrides<T: AnyRelay>(to relay: T) -> T {
+        let overrides = IPOverrideRepository().fetchAll()
+
+        if let override = overrides.first(where: { host in
+            host.hostname == relay.hostname
+        }) {
+            return relay.copyWith(
+                ipv4AddrIn: override.ipv4Address,
+                ipv6AddrIn: override.ipv6Address
+            )
+        }
+
+        return relay
+    }
+
     private static func pickRandomRelayByWeight<T: AnyRelay>(relays: [RelayWithLocation<T>])
-        -> RelayWithLocation<T>? {
+    -> RelayWithLocation<T>? {
         rouletteSelection(relays: relays, weightFunction: { relayWithLocation in relayWithLocation.relay.weight })
     }
 
@@ -314,10 +350,16 @@ public protocol AnyRelay {
     var weight: UInt64 { get }
     var active: Bool { get }
     var includeInCountry: Bool { get }
+
+    func copyWith(ipv4AddrIn: IPv4Address?, ipv6AddrIn: IPv6Address?) -> Self
 }
 
 extension REST.ServerRelay: AnyRelay {}
-extension REST.BridgeRelay: AnyRelay {}
+extension REST.BridgeRelay: AnyRelay {
+    public func copyWith(ipv4AddrIn: IPv4Address?, ipv6AddrIn: IPv6Address?) -> REST.BridgeRelay {
+        copyWith(ipv4AddrIn: ipv4AddrIn)
+    }
+}
 
 private struct RelayWithLocation<T: AnyRelay> {
     let relay: T