Skip to content

Commit

Permalink
Add RelaySelectorWrapper tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Jon Petersson committed Jun 26, 2024
1 parent 091124e commit ce11a62
Show file tree
Hide file tree
Showing 13 changed files with 545 additions and 280 deletions.
116 changes: 116 additions & 0 deletions ios/MullvadREST/Relay/MultihopDecisionFlow.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
//
// MultihopDecisionFlow.swift
// MullvadREST
//
// Created by Jon Petersson on 2024-06-14.
// Copyright © 2024 Mullvad VPN AB. All rights reserved.
//

import Foundation

protocol MultihopDecisionFlow {
typealias RelayCandidate = RelayWithLocation<REST.ServerRelay>
init(next: MultihopDecisionFlow?, relayPicker: RelayPicking)
func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool
func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays
}

struct OneToOne: MultihopDecisionFlow {
let next: MultihopDecisionFlow?
let relayPicker: RelayPicking
init(next: (any MultihopDecisionFlow)?, relayPicker: RelayPicking) {
self.next = next
self.relayPicker = relayPicker
}

func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays {
guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
guard let next else {
throw NoRelaysSatisfyingConstraintsError()
}
return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
}

guard entryCandidates.first != exitCandidates.first else {
throw NoRelaysSatisfyingConstraintsError()
}

let entryMatch = try relayPicker.findBestMatch(from: entryCandidates)
let exitMatch = try relayPicker.findBestMatch(from: exitCandidates)
return SelectedRelays(entry: entryMatch, exit: exitMatch)
}

func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool {
entryCandidates.count == 1 && exitCandidates.count == 1
}
}

struct OneToMany: MultihopDecisionFlow {
let next: MultihopDecisionFlow?
let relayPicker: RelayPicking

init(next: (any MultihopDecisionFlow)?, relayPicker: RelayPicking) {
self.next = next
self.relayPicker = relayPicker
}

func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays {
guard let multihopPicker = relayPicker as? MultihopPicker else {
fatalError("Could not cast picker to MultihopPicker")
}

guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
guard let next else {
throw NoRelaysSatisfyingConstraintsError()
}
return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
}

switch (entryCandidates.count, exitCandidates.count) {
case let (1, count) where count > 1:
let entryMatch = try multihopPicker.findBestMatch(from: entryCandidates)
let exitMatch = try multihopPicker.exclude(relay: entryMatch, from: exitCandidates)
return SelectedRelays(entry: entryMatch, exit: exitMatch)
default:
let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates)
let entryMatch = try multihopPicker.exclude(relay: exitMatch, from: entryCandidates)
return SelectedRelays(entry: entryMatch, exit: exitMatch)
}
}

func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool {
(entryCandidates.count == 1 && exitCandidates.count > 1) ||
(entryCandidates.count > 1 && exitCandidates.count == 1)
}
}

struct ManyToMany: MultihopDecisionFlow {
let next: MultihopDecisionFlow?
let relayPicker: RelayPicking

init(next: (any MultihopDecisionFlow)?, relayPicker: RelayPicking) {
self.next = next
self.relayPicker = relayPicker
}

func pick(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) throws -> SelectedRelays {
guard let multihopPicker = relayPicker as? MultihopPicker else {
fatalError("Could not cast picker to MultihopPicker")
}

guard canHandle(entryCandidates: entryCandidates, exitCandidates: exitCandidates) else {
guard let next else {
throw NoRelaysSatisfyingConstraintsError()
}
return try next.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
}

let exitMatch = try multihopPicker.findBestMatch(from: exitCandidates)
let entryMatch = try multihopPicker.exclude(relay: exitMatch, from: entryCandidates)
return SelectedRelays(entry: entryMatch, exit: exitMatch)
}

func canHandle(entryCandidates: [RelayCandidate], exitCandidates: [RelayCandidate]) -> Bool {
entryCandidates.count > 1 && exitCandidates.count > 1
}
}
99 changes: 99 additions & 0 deletions ios/MullvadREST/Relay/RelayPicking.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
//
// RelaySelectorPicker.swift
// MullvadREST
//
// Created by Jon Petersson on 2024-06-05.
// Copyright © 2024 Mullvad VPN AB. All rights reserved.
//

import MullvadSettings
import MullvadTypes

protocol RelayPicking {
var relays: REST.ServerRelaysResponse { get }
var constraints: RelayConstraints { get }
var connectionAttemptCount: UInt { get }
func pick() throws -> SelectedRelays
}

extension RelayPicking {
func findBestMatch(
from candidates: [RelayWithLocation<REST.ServerRelay>]
) throws -> SelectedRelay {
let match = try RelaySelector.WireGuard.pickCandidate(
from: candidates,
relays: relays,
portConstraint: constraints.port,
numberOfFailedAttempts: connectionAttemptCount
)

return SelectedRelay(
endpoint: match.endpoint,
hostname: match.relay.hostname,
location: match.location,
retryAttempts: connectionAttemptCount
)
}
}

struct SinglehopPicker: RelayPicking {
let constraints: RelayConstraints
let relays: REST.ServerRelaysResponse
let connectionAttemptCount: UInt

func pick() throws -> SelectedRelays {
let candidates = try RelaySelector.WireGuard.findCandidates(
by: constraints.exitLocations,
in: relays,
filterConstraint: constraints.filter
)

let match = try findBestMatch(from: candidates)

return SelectedRelays(entry: nil, exit: match)
}
}

struct MultihopPicker: RelayPicking {
let constraints: RelayConstraints
let relays: REST.ServerRelaysResponse
let connectionAttemptCount: UInt

func pick() throws -> SelectedRelays {
let entryCandidates = try RelaySelector.WireGuard.findCandidates(
by: constraints.entryLocations,
in: relays,
filterConstraint: constraints.filter
)

let exitCandidates = try RelaySelector.WireGuard.findCandidates(
by: constraints.exitLocations,
in: relays,
filterConstraint: constraints.filter
)

let decisionFlow = OneToOne(
next: OneToMany(
next: ManyToMany(
next: nil,
relayPicker: self
),
relayPicker: self
),
relayPicker: self
)

return try decisionFlow.pick(entryCandidates: entryCandidates, exitCandidates: exitCandidates)
}

func exclude(
relay: SelectedRelay,
from candidates: [RelayWithLocation<REST.ServerRelay>]
) throws -> SelectedRelay {
let filteredCandidates = candidates.filter { relayWithLocation in
relayWithLocation.relay.hostname != relay.hostname
}

return try findBestMatch(from: filteredCandidates)
}
}
Loading

0 comments on commit ce11a62

Please sign in to comment.