Skip to content

Commit 5037919

Browse files
author
Jon Petersson
committed
Allow relay selector to filter DAITA enabled relays
1 parent a6c8f0c commit 5037919

12 files changed

+129
-19
lines changed

‎ios/MullvadREST/Relay/RelayPicking.swift

+8-3
Original file line numberDiff line numberDiff line change
@@ -42,22 +42,27 @@ struct SinglehopPicker: RelayPicking {
4242
let connectionAttemptCount: UInt
4343

4444
func pick() throws -> SelectedRelays {
45-
let candidates = try RelaySelector.WireGuard.findCandidates(
45+
let exitCandidates = try RelaySelector.WireGuard.findCandidates(
4646
by: constraints.exitLocations,
4747
in: relays,
4848
filterConstraint: constraints.filter,
4949
daita: daita
5050
)
5151

52-
if daita && candidates.isEmpty {
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
57+
5358
return try MultihopPicker(
5459
constraints: constraints,
5560
daita: daita,
5661
relays: relays,
5762
connectionAttemptCount: connectionAttemptCount
5863
).pick()
5964
} else {
60-
let match = try findBestMatch(from: candidates)
65+
let match = try findBestMatch(from: exitCandidates)
6166
return SelectedRelays(entry: nil, exit: match, retryAttempt: connectionAttemptCount)
6267
}
6368
}

‎ios/MullvadREST/Relay/RelaySelectorWrapper.swift

+6-1
Original file line numberDiff line numberDiff line change
@@ -13,9 +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
16+
private var daitaState: MultihopState = .off
1717
private var observer: MultihopObserverBlock!
1818

19+
// TODO: Remove, Jon
20+
func setDatita(state: MultihopState) {
21+
daitaState = state
22+
}
23+
1924
deinit {
2025
self.multihopUpdater.removeObserver(observer)
2126
}

‎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

+15-2
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")]))
@@ -200,6 +198,21 @@ class RelaySelectorTests: XCTestCase {
200198

201199
XCTAssertThrowsError(try pickRelay(by: constraints, in: sampleRelays, failedAttemptCount: 0))
202200
}
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+
}
203216
}
204217

205218
extension RelaySelectorTests {

‎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.setDatita(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.setDatita(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.setDatita(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.setDatita(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
}

‎ios/MullvadVPNTests/MullvadSettings/IPOverrideWrapperTests.swift

+2-1
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,8 @@ extension IPOverrideWrapperTests {
8383
ipv4AddrIn: .any,
8484
ipv6AddrIn: .any,
8585
publicKey: Data(),
86-
includeInCountry: true
86+
includeInCountry: true,
87+
daita: false
8788
)
8889
}
8990

‎ios/PacketTunnelCoreTests/AppMessageHandlerTests.swift

+2-1
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,8 @@ final class AppMessageHandlerTests: XCTestCase {
8585
let candidates = try RelaySelector.WireGuard.findCandidates(
8686
by: relayConstraints.exitLocations,
8787
in: ServerRelaysResponseStubs.sampleRelays,
88-
filterConstraint: relayConstraints.filter
88+
filterConstraint: relayConstraints.filter,
89+
daita: false
8990
)
9091

9192
let match = try RelaySelector.WireGuard.pickCandidate(

‎ios/PacketTunnelCoreTests/MultiHopPostQuantumKeyExchangingTests.swift

+4-2
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,8 @@ final class MultiHopPostQuantumKeyExchangingTests: XCTestCase {
2727
from: try RelaySelector.WireGuard.findCandidates(
2828
by: relayConstraints.exitLocations,
2929
in: ServerRelaysResponseStubs.sampleRelays,
30-
filterConstraint: relayConstraints.filter
30+
filterConstraint: relayConstraints.filter,
31+
daita: false
3132
),
3233
relays: ServerRelaysResponseStubs.sampleRelays,
3334
portConstraint: relayConstraints.port,
@@ -38,7 +39,8 @@ final class MultiHopPostQuantumKeyExchangingTests: XCTestCase {
3839
from: try RelaySelector.WireGuard.findCandidates(
3940
by: relayConstraints.entryLocations,
4041
in: ServerRelaysResponseStubs.sampleRelays,
41-
filterConstraint: relayConstraints.filter
42+
filterConstraint: relayConstraints.filter,
43+
daita: false
4244
),
4345
relays: ServerRelaysResponseStubs.sampleRelays,
4446
portConstraint: relayConstraints.port,

‎ios/PacketTunnelCoreTests/PostQuantumKeyExchangingPipelineTests.swift

+4-2
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,8 @@ final class PostQuantumKeyExchangingPipelineTests: XCTestCase {
2828
from: try RelaySelector.WireGuard.findCandidates(
2929
by: relayConstraints.exitLocations,
3030
in: ServerRelaysResponseStubs.sampleRelays,
31-
filterConstraint: relayConstraints.filter
31+
filterConstraint: relayConstraints.filter,
32+
daita: false
3233
),
3334
relays: ServerRelaysResponseStubs.sampleRelays,
3435
portConstraint: relayConstraints.port,
@@ -39,7 +40,8 @@ final class PostQuantumKeyExchangingPipelineTests: XCTestCase {
3940
from: try RelaySelector.WireGuard.findCandidates(
4041
by: relayConstraints.entryLocations,
4142
in: ServerRelaysResponseStubs.sampleRelays,
42-
filterConstraint: relayConstraints.filter
43+
filterConstraint: relayConstraints.filter,
44+
daita: false
4345
),
4446
relays: ServerRelaysResponseStubs.sampleRelays,
4547
portConstraint: relayConstraints.port,

‎ios/PacketTunnelCoreTests/SingleHopPostQuantumKeyExchangingTests.swift

+2-1
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,8 @@ final class SingleHopPostQuantumKeyExchangingTests: XCTestCase {
2424
let candidates = try RelaySelector.WireGuard.findCandidates(
2525
by: relayConstraints.exitLocations,
2626
in: ServerRelaysResponseStubs.sampleRelays,
27-
filterConstraint: relayConstraints.filter
27+
filterConstraint: relayConstraints.filter,
28+
daita: false
2829
)
2930

3031
let match = try RelaySelector.WireGuard.pickCandidate(

0 commit comments

Comments
 (0)