Relay selector should use overridden IP addresses for relays

5 files changed, 131 insertions, 13 deletions

diff --git a/ios/MullvadREST/ApiHandlers/ServerRelaysResponse.swift b/ios/MullvadREST/ApiHandlers/ServerRelaysResponse.swift
index 63f0822e63..a575da03cc 100644
--- a/ios/MullvadREST/ApiHandlers/ServerRelaysResponse.swift
+++ b/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 {
diff --git a/ios/MullvadREST/Relay/AnyRelay.swift b/ios/MullvadREST/Relay/AnyRelay.swift
new file mode 100644
index 0000000000..6c3c49aa55
--- /dev/null
+++ b/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)
+    }
+}
diff --git a/ios/MullvadREST/Relay/IPOverrideWrapper.swift b/ios/MullvadREST/Relay/IPOverrideWrapper.swift
new file mode 100644
index 0000000000..531112cf54
--- /dev/null
+++ b/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
+    }
+}
diff --git a/ios/MullvadREST/Relay/RelayCache.swift b/ios/MullvadREST/Relay/RelayCache.swift
index 6cc9ddc616..47ea44c6f5 100644
--- a/ios/MullvadREST/Relay/RelayCache.swift
+++ b/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>
 
diff --git a/ios/MullvadREST/Relay/RelaySelector.swift b/ios/MullvadREST/Relay/RelaySelector.swift
index bc0378c217..6fc016d2c7 100644
--- a/ios/MullvadREST/Relay/RelaySelector.swift
+++ b/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