Skip to content

Commit 42f9904

Browse files
author
Jon Petersson
committed
Add RelaySelectorWrapper tests
1 parent 8b647d4 commit 42f9904

File tree

15 files changed

+564
-283
lines changed

15 files changed

+564
-283
lines changed

‎ios/MullvadREST/Relay/MultihopDecisionFlow.swift

+116
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
//
2+
// MultihopDecisionFlow.swift
3+
// MullvadREST
4+
//
5+
// Created by Jon Petersson on 2024-06-14.
6+
// Copyright © 2024 Mullvad VPN AB. All rights reserved.
7+
//
8+
9+
import Foundation
10+
11+
protocol MultihopDecisionFlow {
12+
typealias RelayCandidate = RelayWithLocation<REST.ServerRelay>
13+
init(next: MultihopDecisionFlow?, relayPicker: RelayPicking)
14+
func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool
15+
func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays
16+
}
17+
18+
struct OneToOne: MultihopDecisionFlow {
19+
let next: MultihopDecisionFlow?
20+
let relayPicker: RelayPicking
21+
init(next: (any MultihopDecisionFlow)?, relayPicker: RelayPicking) {
22+
self.next = next
23+
self.relayPicker = relayPicker
24+
}
25+
26+
func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays {
27+
guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
28+
guard let next else {
29+
throw NoRelaysSatisfyingConstraintsError()
30+
}
31+
return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
32+
}
33+
34+
guard entryCandidates.first != exitCandidates.first else {
35+
throw NoRelaysSatisfyingConstraintsError()
36+
}
37+
38+
let entryMatch = try relayPicker.findBestMatch(from: entryCandidates)
39+
let exitMatch = try relayPicker.findBestMatch(from: exitCandidates)
40+
return SelectedRelays(entry: entryMatch, exit: exitMatch)
41+
}
42+
43+
func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool {
44+
entryCandidates.count == 1 && exitCandidates.count == 1
45+
}
46+
}
47+
48+
struct OneToMany: MultihopDecisionFlow {
49+
let next: MultihopDecisionFlow?
50+
let relayPicker: RelayPicking
51+
52+
init(next: (any MultihopDecisionFlow)?, relayPicker: RelayPicking) {
53+
self.next = next
54+
self.relayPicker = relayPicker
55+
}
56+
57+
func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays {
58+
guard let multihopPicker = relayPicker as? MultihopPicker else {
59+
fatalError("Could not cast picker to MultihopPicker")
60+
}
61+
62+
guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
63+
guard let next else {
64+
throw NoRelaysSatisfyingConstraintsError()
65+
}
66+
return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
67+
}
68+
69+
switch (entryCandidates.count, exitCandidates.count) {
70+
case let (1, count) where count > 1:
71+
let entryMatch = try multihopPicker.findBestMatch(from: entryCandidates)
72+
let exitMatch = try multihopPicker.exclude(relay: entryMatch, from: exitCandidates)
73+
return SelectedRelays(entry: entryMatch, exit: exitMatch)
74+
default:
75+
let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates)
76+
let entryMatch = try multihopPicker.exclude(relay: exitMatch, from: entryCandidates)
77+
return SelectedRelays(entry: entryMatch, exit: exitMatch)
78+
}
79+
}
80+
81+
func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool {
82+
(entryCandidates.count == 1 && exitCandidates.count > 1) ||
83+
(entryCandidates.count > 1 && exitCandidates.count == 1)
84+
}
85+
}
86+
87+
struct ManyToMany: MultihopDecisionFlow {
88+
let next: MultihopDecisionFlow?
89+
let relayPicker: RelayPicking
90+
91+
init(next: (any MultihopDecisionFlow)?, relayPicker: RelayPicking) {
92+
self.next = next
93+
self.relayPicker = relayPicker
94+
}
95+
96+
func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays {
97+
guard let multihopPicker = relayPicker as? MultihopPicker else {
98+
fatalError("Could not cast picker to MultihopPicker")
99+
}
100+
101+
guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
102+
guard let next else {
103+
throw NoRelaysSatisfyingConstraintsError()
104+
}
105+
return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
106+
}
107+
108+
let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates)
109+
let entryMatch = try multihopPicker.exclude(relay: exitMatch, from: entryCandidates)
110+
return SelectedRelays(entry: entryMatch, exit: exitMatch)
111+
}
112+
113+
func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool {
114+
entryCandidates.count > 1 && exitCandidates.count > 1
115+
}
116+
}

‎ios/MullvadREST/Relay/RelayPicking.swift

+107
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
//
2+
// RelaySelectorPicker.swift
3+
// MullvadREST
4+
//
5+
// Created by Jon Petersson on 2024-06-05.
6+
// Copyright © 2024 Mullvad VPN AB. All rights reserved.
7+
//
8+
9+
import MullvadSettings
10+
import MullvadTypes
11+
12+
protocol RelayPicking {
13+
var relays: REST.ServerRelaysResponse { get }
14+
var constraints: RelayConstraints { get }
15+
var connectionAttemptCount: UInt { get }
16+
func pick() throws -> SelectedRelays
17+
}
18+
19+
extension RelayPicking {
20+
func findBestMatch(
21+
from candidates: [RelayWithLocation<REST.ServerRelay>]
22+
) throws -> SelectedRelay {
23+
let match = try RelaySelector.WireGuard.pickCandidate(
24+
from: candidates,
25+
relays: relays,
26+
portConstraint: constraints.port,
27+
numberOfFailedAttempts: connectionAttemptCount
28+
)
29+
30+
return SelectedRelay(
31+
endpoint: match.endpoint,
32+
hostname: match.relay.hostname,
33+
location: match.location,
34+
retryAttempts: connectionAttemptCount
35+
)
36+
}
37+
}
38+
39+
struct SinglehopPicker: RelayPicking {
40+
let constraints: RelayConstraints
41+
let relays: REST.ServerRelaysResponse
42+
let connectionAttemptCount: UInt
43+
44+
func pick() throws -> SelectedRelays {
45+
let candidates = try RelaySelector.WireGuard.findCandidates(
46+
by: constraints.exitLocations,
47+
in: relays,
48+
filterConstraint: constraints.filter
49+
)
50+
51+
let match = try findBestMatch(from: candidates)
52+
53+
return SelectedRelays(entry: nil, exit: match)
54+
}
55+
}
56+
57+
struct MultihopPicker: RelayPicking {
58+
let constraints: RelayConstraints
59+
let relays: REST.ServerRelaysResponse
60+
let connectionAttemptCount: UInt
61+
62+
func pick() throws -> SelectedRelays {
63+
let entryCandidates = try RelaySelector.WireGuard.findCandidates(
64+
by: constraints.entryLocations,
65+
in: relays,
66+
filterConstraint: constraints.filter
67+
)
68+
69+
let exitCandidates = try RelaySelector.WireGuard.findCandidates(
70+
by: constraints.exitLocations,
71+
in: relays,
72+
filterConstraint: constraints.filter
73+
)
74+
75+
/*
76+
Relay selection is prioritised in the following order:
77+
1. Both entry and exit constraints match only a single relay. Both relays are selected.
78+
2. Either entry or exit constraint matches only a single relay and the other multiple relays. The single relays
79+
is selected and excluded from the list of multiple relays.
80+
3. Both entry and exit constraints match multiple relays. Exit relay is picked first and then excluded from
81+
the list of entry relays.
82+
*/
83+
let decisionFlow = OneToOne(
84+
next: OneToMany(
85+
next: ManyToMany(
86+
next: nil,
87+
relayPicker: self
88+
),
89+
relayPicker: self
90+
),
91+
relayPicker: self
92+
)
93+
94+
return try decisionFlow.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
95+
}
96+
97+
func exclude(
98+
relay: SelectedRelay,
99+
from candidates: [RelayWithLocation<REST.ServerRelay>]
100+
) throws -> SelectedRelay {
101+
let filteredCandidates = candidates.filter { relayWithLocation in
102+
relayWithLocation.relay.hostname != relay.hostname
103+
}
104+
105+
return try findBestMatch(from: filteredCandidates)
106+
}
107+
}

0 commit comments

Comments
 (0)