Skip to content

Commit c65edb4

Browse files
author
Jon Petersson
committed
Allow relay selector to filter DAITA enabled relays
1 parent 91e0194 commit c65edb4

17 files changed

+168
-26
lines changed

‎ios/MullvadREST/ApiHandlers/ServerRelaysResponse.swift

+4-1
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ extension REST {
3434
public let ipv4AddrIn: IPv4Address
3535
public let weight: UInt64
3636
public let includeInCountry: Bool
37+
public var daita: Bool? = nil
3738

3839
public func override(ipv4AddrIn: IPv4Address?) -> Self {
3940
return BridgeRelay(
@@ -60,6 +61,7 @@ extension REST {
6061
public let ipv6AddrIn: IPv6Address
6162
public let publicKey: Data
6263
public let includeInCountry: Bool
64+
public let daita: Bool?
6365

6466
public func override(ipv4AddrIn: IPv4Address?, ipv6AddrIn: IPv6Address?) -> Self {
6567
return ServerRelay(
@@ -72,7 +74,8 @@ extension REST {
7274
ipv4AddrIn: ipv4AddrIn ?? self.ipv4AddrIn,
7375
ipv6AddrIn: ipv6AddrIn ?? self.ipv6AddrIn,
7476
publicKey: publicKey,
75-
includeInCountry: includeInCountry
77+
includeInCountry: includeInCountry,
78+
daita: daita
7679
)
7780
}
7881
}

‎ios/MullvadREST/Relay/AnyRelay.swift

+1
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ public protocol AnyRelay {
1717
var weight: UInt64 { get }
1818
var active: Bool { get }
1919
var includeInCountry: Bool { get }
20+
var daita: Bool? { get }
2021

2122
func override(ipv4AddrIn: IPv4Address?, ipv6AddrIn: IPv6Address?) -> Self
2223
}

‎ios/MullvadREST/Relay/RelayPicking.swift

+24-6
Original file line numberDiff line numberDiff line change
@@ -37,38 +37,56 @@ extension RelayPicking {
3737

3838
struct SinglehopPicker: RelayPicking {
3939
let constraints: RelayConstraints
40+
let daita: Bool
4041
let relays: REST.ServerRelaysResponse
4142
let connectionAttemptCount: UInt
4243

4344
func pick() throws -> SelectedRelays {
44-
let candidates = try RelaySelector.WireGuard.findCandidates(
45+
let exitCandidates = try RelaySelector.WireGuard.findCandidates(
4546
by: constraints.exitLocations,
4647
in: relays,
47-
filterConstraint: constraints.filter
48+
filterConstraint: constraints.filter,
49+
daita: daita
4850
)
4951

50-
let match = try findBestMatch(from: candidates)
52+
// If DAITA is enabled and no supported relays are found, we should try to find the nearest
53+
// available relay that supports DAITA and use it as entry in a multihop selection.
54+
if daita && exitCandidates.isEmpty {
55+
var constraints = constraints
56+
constraints.entryLocations = .any
5157

52-
return SelectedRelays(entry: nil, exit: match, retryAttempt: connectionAttemptCount)
58+
return try MultihopPicker(
59+
constraints: constraints,
60+
daita: daita,
61+
relays: relays,
62+
connectionAttemptCount: connectionAttemptCount
63+
).pick()
64+
} else {
65+
let match = try findBestMatch(from: exitCandidates)
66+
return SelectedRelays(entry: nil, exit: match, retryAttempt: connectionAttemptCount)
67+
}
5368
}
5469
}
5570

5671
struct MultihopPicker: RelayPicking {
5772
let constraints: RelayConstraints
73+
let daita: Bool
5874
let relays: REST.ServerRelaysResponse
5975
let connectionAttemptCount: UInt
6076

6177
func pick() throws -> SelectedRelays {
6278
let entryCandidates = try RelaySelector.WireGuard.findCandidates(
6379
by: constraints.entryLocations,
6480
in: relays,
65-
filterConstraint: constraints.filter
81+
filterConstraint: constraints.filter,
82+
daita: daita
6683
)
6784

6885
let exitCandidates = try RelaySelector.WireGuard.findCandidates(
6986
by: constraints.exitLocations,
7087
in: relays,
71-
filterConstraint: constraints.filter
88+
filterConstraint: constraints.filter,
89+
daita: false
7290
)
7391

7492
/*

‎ios/MullvadREST/Relay/RelaySelector+Shadowsocks.swift

+1
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ extension RelaySelector {
4646
let filteredRelays = applyConstraints(
4747
location,
4848
filterConstraint: filter,
49+
daita: false,
4950
relays: mappedBridges
5051
)
5152
guard filteredRelays.isEmpty == false else { return relay(from: relaysResponse) }

‎ios/MullvadREST/Relay/RelaySelector+Wireguard.swift

+3-1
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,15 @@ extension RelaySelector {
1414
public static func findCandidates(
1515
by relayConstraint: RelayConstraint<UserSelectedRelays>,
1616
in relays: REST.ServerRelaysResponse,
17-
filterConstraint: RelayConstraint<RelayFilter>
17+
filterConstraint: RelayConstraint<RelayFilter>,
18+
daita: Bool
1819
) throws -> [RelayWithLocation<REST.ServerRelay>] {
1920
let mappedRelays = mapRelays(relays: relays.wireguard.relays, locations: relays.locations)
2021

2122
return applyConstraints(
2223
relayConstraint,
2324
filterConstraint: filterConstraint,
25+
daita: daita,
2426
relays: mappedRelays
2527
)
2628
}

‎ios/MullvadREST/Relay/RelaySelector.swift

+7-1
Original file line numberDiff line numberDiff line change
@@ -135,14 +135,20 @@ public enum RelaySelector {
135135
static func applyConstraints<T: AnyRelay>(
136136
_ relayConstraint: RelayConstraint<UserSelectedRelays>,
137137
filterConstraint: RelayConstraint<RelayFilter>,
138+
daita: Bool,
138139
relays: [RelayWithLocation<T>]
139140
) -> [RelayWithLocation<T>] {
140-
// Filter on active status, filter, and location.
141+
// Filter on active status, daita support, filter constraint and location constraint.
141142
let filteredRelays = relays.filter { relayWithLocation -> Bool in
142143
guard relayWithLocation.relay.active else {
143144
return false
144145
}
145146

147+
let relaySupportsDaita = (relayWithLocation.relay.daita ?? false) == true
148+
if daita && !relaySupportsDaita {
149+
return false
150+
}
151+
146152
switch filterConstraint {
147153
case .any:
148154
break

‎ios/MullvadREST/Relay/RelaySelectorWrapper.swift

+8
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,14 @@ public final class RelaySelectorWrapper: RelaySelectorProtocol {
1313
let relayCache: RelayCacheProtocol
1414
let multihopUpdater: MultihopUpdater
1515
private var multihopState: MultihopState = .off
16+
private var daitaState: MultihopState = .on
1617
private var observer: MultihopObserverBlock!
1718

19+
// TODO: Remove, Jon
20+
func setDaita(state: MultihopState) {
21+
daitaState = state
22+
}
23+
1824
deinit {
1925
self.multihopUpdater.removeObserver(observer)
2026
}
@@ -39,12 +45,14 @@ public final class RelaySelectorWrapper: RelaySelectorProtocol {
3945
case .off:
4046
return try SinglehopPicker(
4147
constraints: constraints,
48+
daita: daitaState == .on,
4249
relays: relays,
4350
connectionAttemptCount: connectionAttemptCount
4451
).pick()
4552
case .on:
4653
return try MultihopPicker(
4754
constraints: constraints,
55+
daita: daitaState == .on,
4856
relays: relays,
4957
connectionAttemptCount: connectionAttemptCount
5058
).pick()

‎ios/MullvadVPNTests/MullvadREST/ApiHandlers/ServerRelaysResponse+Stubs.swift

+12-6
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,8 @@ enum ServerRelaysResponseStubs {
8585
ipv4AddrIn: .loopback,
8686
ipv6AddrIn: .loopback,
8787
publicKey: PrivateKey().publicKey.rawValue,
88-
includeInCountry: true
88+
includeInCountry: true,
89+
daita: true
8990
),
9091
REST.ServerRelay(
9192
hostname: "se10-wireguard",
@@ -97,7 +98,8 @@ enum ServerRelaysResponseStubs {
9798
ipv4AddrIn: .loopback,
9899
ipv6AddrIn: .loopback,
99100
publicKey: PrivateKey().publicKey.rawValue,
100-
includeInCountry: true
101+
includeInCountry: true,
102+
daita: false
101103
),
102104
REST.ServerRelay(
103105
hostname: "se2-wireguard",
@@ -109,7 +111,8 @@ enum ServerRelaysResponseStubs {
109111
ipv4AddrIn: .loopback,
110112
ipv6AddrIn: .loopback,
111113
publicKey: PrivateKey().publicKey.rawValue,
112-
includeInCountry: true
114+
includeInCountry: true,
115+
daita: false
113116
),
114117
REST.ServerRelay(
115118
hostname: "se6-wireguard",
@@ -121,7 +124,8 @@ enum ServerRelaysResponseStubs {
121124
ipv4AddrIn: .loopback,
122125
ipv6AddrIn: .loopback,
123126
publicKey: PrivateKey().publicKey.rawValue,
124-
includeInCountry: true
127+
includeInCountry: true,
128+
daita: false
125129
),
126130
REST.ServerRelay(
127131
hostname: "us-dal-wg-001",
@@ -133,7 +137,8 @@ enum ServerRelaysResponseStubs {
133137
ipv4AddrIn: .loopback,
134138
ipv6AddrIn: .loopback,
135139
publicKey: PrivateKey().publicKey.rawValue,
136-
includeInCountry: true
140+
includeInCountry: true,
141+
daita: true
137142
),
138143
REST.ServerRelay(
139144
hostname: "us-nyc-wg-301",
@@ -145,7 +150,8 @@ enum ServerRelaysResponseStubs {
145150
ipv4AddrIn: .loopback,
146151
ipv6AddrIn: .loopback,
147152
publicKey: PrivateKey().publicKey.rawValue,
148-
includeInCountry: true
153+
includeInCountry: true,
154+
daita: true
149155
),
150156
]
151157
),

‎ios/MullvadVPNTests/MullvadREST/Relay/MultihopDecisionFlowTests.swift

+1
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,7 @@ extension MultihopDecisionFlowTests {
119119

120120
return MultihopPicker(
121121
constraints: constraints,
122+
daita: false,
122123
relays: sampleRelays,
123124
connectionAttemptCount: 0
124125
)

‎ios/MullvadVPNTests/MullvadREST/Relay/RelayPickingTests.swift

+3
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ class RelayPickingTests: XCTestCase {
2323

2424
let picker = SinglehopPicker(
2525
constraints: constraints,
26+
daita: false,
2627
relays: sampleRelays,
2728
connectionAttemptCount: 0
2829
)
@@ -41,6 +42,7 @@ class RelayPickingTests: XCTestCase {
4142

4243
let picker = MultihopPicker(
4344
constraints: constraints,
45+
daita: false,
4446
relays: sampleRelays,
4547
connectionAttemptCount: 0
4648
)
@@ -59,6 +61,7 @@ class RelayPickingTests: XCTestCase {
5961

6062
let picker = MultihopPicker(
6163
constraints: constraints,
64+
daita: false,
6265
relays: sampleRelays,
6366
connectionAttemptCount: 0
6467
)

‎ios/MullvadVPNTests/MullvadREST/Relay/RelaySelectorTests.swift

+20-4
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,6 @@ private let defaultPort: UInt16 = 53
1717
class RelaySelectorTests: XCTestCase {
1818
let sampleRelays = ServerRelaysResponseStubs.sampleRelays
1919

20-
// MARK: - single-Hop tests
21-
2220
func testCountryConstraint() throws {
2321
let constraints = RelayConstraints(
2422
exitLocations: .only(UserSelectedRelays(locations: [.country("es")]))
@@ -74,6 +72,7 @@ class RelaySelectorTests: XCTestCase {
7472
let constrainedLocations = RelaySelector.applyConstraints(
7573
constraints.exitLocations,
7674
filterConstraint: constraints.filter,
75+
daita: false,
7776
relays: relayWithLocations
7877
)
7978

@@ -199,18 +198,35 @@ class RelaySelectorTests: XCTestCase {
199198

200199
XCTAssertThrowsError(try pickRelay(by: constraints, in: sampleRelays, failedAttemptCount: 0))
201200
}
201+
202+
func testRelayWithDaita() throws {
203+
let hasDaitaConstraints = RelayConstraints(
204+
exitLocations: .only(UserSelectedRelays(locations: [.country("es")]))
205+
)
206+
207+
let noDaitaConstraints = RelayConstraints(
208+
exitLocations: .only(UserSelectedRelays(locations: [.country("se")]))
209+
)
210+
211+
XCTAssertNoThrow(try pickRelay(by: hasDaitaConstraints, in: sampleRelays, failedAttemptCount: 0, daita: true))
212+
XCTAssertThrowsError(
213+
try pickRelay(by: noDaitaConstraints, in: sampleRelays, failedAttemptCount: 0, daita: true)
214+
)
215+
}
202216
}
203217

204218
extension RelaySelectorTests {
205219
private func pickRelay(
206220
by constraints: RelayConstraints,
207221
in relays: REST.ServerRelaysResponse,
208-
failedAttemptCount: UInt
222+
failedAttemptCount: UInt,
223+
daita: Bool = false
209224
) throws -> RelaySelectorMatch {
210225
let candidates = try RelaySelector.WireGuard.findCandidates(
211226
by: constraints.exitLocations,
212227
in: relays,
213-
filterConstraint: constraints.filter
228+
filterConstraint: constraints.filter,
229+
daita: daita
214230
)
215231

216232
return try RelaySelector.WireGuard.pickCandidate(

‎ios/MullvadVPNTests/MullvadREST/Relay/RelaySelectorWrapperTests.swift

+70
Original file line numberDiff line numberDiff line change
@@ -52,4 +52,74 @@ class RelaySelectorWrapperTests: XCTestCase {
5252
let selectedRelays = try wrapper.selectRelays(with: RelayConstraints(), connectionAttemptCount: 0)
5353
XCTAssertNotNil(selectedRelays.entry)
5454
}
55+
56+
func testCanSelectRelayWithMultihopOnAndDaitaOn() throws {
57+
let wrapper = RelaySelectorWrapper(
58+
relayCache: relayCache,
59+
multihopUpdater: multihopUpdater
60+
)
61+
62+
multihopStateListener.onNewMultihop?(.on)
63+
wrapper.setDaita(state: .on)
64+
65+
let constraints = RelayConstraints(
66+
entryLocations: .only(UserSelectedRelays(locations: [.country("es")])), // Relay with DAITA.
67+
exitLocations: .only(UserSelectedRelays(locations: [.country("us")]))
68+
)
69+
70+
XCTAssertNoThrow(try wrapper.selectRelays(with: constraints, connectionAttemptCount: 0))
71+
}
72+
73+
func testCannotSelectRelayWithMultihopOnAndDaitaOn() throws {
74+
let wrapper = RelaySelectorWrapper(
75+
relayCache: relayCache,
76+
multihopUpdater: multihopUpdater
77+
)
78+
79+
multihopStateListener.onNewMultihop?(.on)
80+
wrapper.setDaita(state: .on)
81+
82+
let constraints = RelayConstraints(
83+
entryLocations: .only(UserSelectedRelays(locations: [.country("se")])), // Relay without DAITA.
84+
exitLocations: .only(UserSelectedRelays(locations: [.country("us")]))
85+
)
86+
87+
XCTAssertThrowsError(try wrapper.selectRelays(with: constraints, connectionAttemptCount: 0))
88+
}
89+
90+
func testCanSelectRelayWithMultihopOffAndDaitaOn() throws {
91+
let wrapper = RelaySelectorWrapper(
92+
relayCache: relayCache,
93+
multihopUpdater: multihopUpdater
94+
)
95+
96+
multihopStateListener.onNewMultihop?(.off)
97+
wrapper.setDaita(state: .on)
98+
99+
let constraints = RelayConstraints(
100+
exitLocations: .only(UserSelectedRelays(locations: [.country("es")])) // Relay with DAITA.
101+
)
102+
103+
let selectedRelays = try wrapper.selectRelays(with: constraints, connectionAttemptCount: 0)
104+
XCTAssertNil(selectedRelays.entry)
105+
}
106+
107+
// If DAITA is enabled and no supported relays are found, we should try to find the nearest
108+
// available relay that supports DAITA and use it as entry in a multihop selection.
109+
func testCanSelectRelayWithMultihopOffAndDaitaOnThroughMultihop() throws {
110+
let wrapper = RelaySelectorWrapper(
111+
relayCache: relayCache,
112+
multihopUpdater: multihopUpdater
113+
)
114+
115+
multihopStateListener.onNewMultihop?(.off)
116+
wrapper.setDaita(state: .on)
117+
118+
let constraints = RelayConstraints(
119+
exitLocations: .only(UserSelectedRelays(locations: [.country("se")])) // Relay without DAITA.
120+
)
121+
122+
let selectedRelays = try wrapper.selectRelays(with: constraints, connectionAttemptCount: 0)
123+
XCTAssertNotNil(selectedRelays.entry)
124+
}
55125
}

0 commit comments

Comments
 (0)