580 lines
23 KiB
Swift
580 lines
23 KiB
Swift
// SPDX-FileCopyrightText: 2026 Kaede Akino
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
import Darwin
|
|
import Foundation
|
|
import Logging
|
|
import TermindApplication
|
|
import TermindDomain
|
|
|
|
/// A host-scoped group of physical SSH connections. Only local and dynamic rules fan out.
|
|
actor SSHForwardingConnectionPool: SSHForwardingConnectionProtocol {
|
|
typealias ConnectionFactory =
|
|
@Sendable (
|
|
UUID, @escaping SSHForwardingConnection.ConnectionHandler,
|
|
@escaping SSHForwardingConnection.FailureHandler
|
|
) -> SSHForwardingConnection
|
|
|
|
private final class Slot: @unchecked Sendable {
|
|
let id: UUID
|
|
let connection: SSHForwardingConnection
|
|
var rules: Set<PortForwardingRule.ID> = []
|
|
|
|
init(id: UUID, connection: SSHForwardingConnection) {
|
|
self.id = id
|
|
self.connection = connection
|
|
}
|
|
}
|
|
|
|
private struct RuleRegistration: Sendable {
|
|
let rule: PortForwardingRule
|
|
let generation: UUID
|
|
}
|
|
|
|
private var maximumConnections: Int
|
|
private let connectionHandler: SSHForwardingConnection.ConnectionHandler
|
|
private let failureHandler: SSHForwardingConnection.FailureHandler
|
|
private var accessCheck: HostOperationAccessCheck?
|
|
private let makeConnection: ConnectionFactory
|
|
private let listeners: SSHForwardingLocalListeners
|
|
private var slots: [UUID: Slot] = [:]
|
|
private var primaryID: UUID?
|
|
private var registrations: [PortForwardingRule.ID: RuleRegistration] = [:]
|
|
private var counts: [UUID: [PortForwardingRule.ID: Int]] = [:]
|
|
private var additionalRequest: SSHAdditionalConnectionRequestResolver?
|
|
private var opening: [UUID: Slot] = [:]
|
|
private var installing: [UUID: Slot] = [:]
|
|
private var openingTasks: [UUID: Task<Void, Never>] = [:]
|
|
private var draining: Set<UUID> = []
|
|
private var assignments: [UUID: Int] = [:]
|
|
private var checkingIdle: Set<UUID> = []
|
|
private var recheckIdle: Set<UUID> = []
|
|
private var retryTask: Task<Void, Never>?
|
|
private var stabilityTasks: [UUID: Task<Void, Never>] = [:]
|
|
private var retryAttempt = 0
|
|
private var roundRobin = 0
|
|
private var disconnecting = false
|
|
|
|
init(
|
|
maximumConnections: Int,
|
|
connectionHandler: @escaping SSHForwardingConnection.ConnectionHandler,
|
|
failureHandler: @escaping SSHForwardingConnection.FailureHandler,
|
|
logger: Logger = ApplicationLogger.disabled("termind.port-forwarding-pool"),
|
|
makeConnection: ConnectionFactory? = nil
|
|
) {
|
|
self.maximumConnections = Host.normalizedPortForwardingMaxConnections(maximumConnections)
|
|
self.connectionHandler = connectionHandler
|
|
self.failureHandler = failureHandler
|
|
self.makeConnection =
|
|
makeConnection ?? { _, connectionHandler, failureHandler in
|
|
SSHForwardingConnection(
|
|
connectionHandler: connectionHandler, failureHandler: failureHandler, logger: logger)
|
|
}
|
|
listeners = SSHForwardingLocalListeners { _, _, descriptor in
|
|
Darwin.close(descriptor)
|
|
}
|
|
listeners.setHandler { [weak self] ruleID, generation, descriptor in
|
|
guard let self else {
|
|
Darwin.close(descriptor)
|
|
return
|
|
}
|
|
Task { await self.assign(descriptor, ruleID: ruleID, generation: generation) }
|
|
}
|
|
}
|
|
|
|
func connect(_ request: SSHConnectionRequest, progress: SSHRouteProgressHandler?) async throws {
|
|
await disconnect()
|
|
accessCheck = request.configuration.accessCheck
|
|
disconnecting = false
|
|
let id = UUID()
|
|
let slot = makeSlot(id: id)
|
|
slots[id] = slot
|
|
primaryID = id
|
|
do {
|
|
try await slot.connection.connect(request, progress: progress)
|
|
} catch {
|
|
slots[id] = nil
|
|
primaryID = nil
|
|
throw error
|
|
}
|
|
}
|
|
|
|
func configureAdditionalConnections(_ resolver: @escaping SSHAdditionalConnectionRequestResolver) async {
|
|
additionalRequest = resolver
|
|
fillPool()
|
|
}
|
|
|
|
func connectionPoolStatus() async -> SSHForwardingPoolStatus {
|
|
SSHForwardingPoolStatus(connected: slots.count, target: desiredConnections)
|
|
}
|
|
|
|
func updateMaximumConnections(_ maximum: Int) async {
|
|
maximumConnections = Host.normalizedPortForwardingMaxConnections(maximum)
|
|
rebalanceSlots()
|
|
// Pending auxiliaries have no accepted flows and can be cancelled immediately.
|
|
let pending: [Slot]
|
|
if slots.count + opening.count + installing.count > desiredConnections {
|
|
pending = Array(opening.values) + Array(installing.values)
|
|
opening.removeAll()
|
|
installing.removeAll()
|
|
for task in openingTasks.values { task.cancel() }
|
|
openingTasks.removeAll()
|
|
} else {
|
|
pending = []
|
|
}
|
|
retryTask?.cancel()
|
|
retryTask = nil
|
|
publishPoolCounts()
|
|
for slot in pending { await slot.connection.disconnect() }
|
|
for id in Array(draining) { await retireIfIdle(id) }
|
|
fillPool()
|
|
}
|
|
|
|
func flowDistribution(for ruleID: PortForwardingRule.ID) async -> [Int] {
|
|
slots.keys.sorted { $0.uuidString < $1.uuidString }.map { counts[$0]?[ruleID] ?? 0 }
|
|
}
|
|
|
|
func start(_ rule: PortForwardingRule, generation: UUID) async throws {
|
|
try accessCheck?()
|
|
guard let primary = primarySlot else { throw TermindSSHError.disconnected }
|
|
registrations[rule.id] = RuleRegistration(rule: rule, generation: generation)
|
|
do {
|
|
if rule.kind == .remote {
|
|
try await primary.connection.start(rule, generation: generation)
|
|
primary.rules.insert(rule.id)
|
|
} else {
|
|
// Draining slots may be retained again if the limit increases before their flows finish.
|
|
for slot in slots.values {
|
|
try await slot.connection.registerShared(rule, generation: generation)
|
|
slot.rules.insert(rule.id)
|
|
}
|
|
try await listeners.start(rule, generation: generation)
|
|
fillPool()
|
|
}
|
|
} catch {
|
|
registrations[rule.id] = nil
|
|
for slot in slots.values where slot.rules.remove(rule.id) != nil {
|
|
await slot.connection.stop(rule.id)
|
|
}
|
|
await listeners.stop(rule.id)
|
|
throw error
|
|
}
|
|
}
|
|
|
|
func stop(_ id: PortForwardingRule.ID) async {
|
|
let registration = registrations.removeValue(forKey: id)
|
|
if registration?.rule.kind != .remote { await listeners.stop(id) }
|
|
for slot in slots.values where slot.rules.remove(id) != nil {
|
|
await slot.connection.stop(id)
|
|
}
|
|
for slotID in counts.keys { counts[slotID]?[id] = nil }
|
|
report(id, generation: registration?.generation)
|
|
if !needsPool { await retireAuxiliaries() }
|
|
}
|
|
|
|
func disconnect() async {
|
|
disconnecting = true
|
|
retryTask?.cancel()
|
|
retryTask = nil
|
|
for task in stabilityTasks.values { task.cancel() }
|
|
stabilityTasks.removeAll()
|
|
for task in openingTasks.values { task.cancel() }
|
|
openingTasks.removeAll()
|
|
draining.removeAll()
|
|
assignments.removeAll()
|
|
checkingIdle.removeAll()
|
|
recheckIdle.removeAll()
|
|
let pending = Array(opening.values) + Array(installing.values)
|
|
opening.removeAll()
|
|
installing.removeAll()
|
|
await listeners.disconnect()
|
|
let active = Array(slots.values)
|
|
slots.removeAll()
|
|
primaryID = nil
|
|
registrations.removeAll()
|
|
counts.removeAll()
|
|
additionalRequest = nil
|
|
await withTaskGroup(of: Void.self) { group in
|
|
for slot in active + pending { group.addTask { await slot.connection.disconnect() } }
|
|
}
|
|
}
|
|
|
|
private var primarySlot: Slot? { primaryID.flatMap { slots[$0] } }
|
|
private var needsPool: Bool { registrations.values.contains { $0.rule.kind != .remote } }
|
|
private var desiredConnections: Int { needsPool ? maximumConnections : 1 }
|
|
|
|
private func rebalanceSlots() {
|
|
let auxiliaries = slots.keys.filter { $0 != primaryID }.sorted { $0.uuidString < $1.uuidString }
|
|
let retained = Set(auxiliaries.prefix(max(0, desiredConnections - 1)))
|
|
draining = Set(auxiliaries).subtracting(retained)
|
|
}
|
|
|
|
private func retireIfIdle(_ id: UUID) async {
|
|
guard draining.contains(id), id != primaryID, assignments[id, default: 0] == 0,
|
|
let slot = slots[id]
|
|
else { return }
|
|
guard checkingIdle.insert(id).inserted else {
|
|
recheckIdle.insert(id)
|
|
return
|
|
}
|
|
let busy = await slot.connection.hasActiveFlows()
|
|
checkingIdle.remove(id)
|
|
let needsRecheck = recheckIdle.remove(id) != nil
|
|
if busy {
|
|
if needsRecheck { await retireIfIdle(id) }
|
|
return
|
|
}
|
|
// The worker query is authoritative: count callbacks may still be in flight.
|
|
guard draining.contains(id), id != primaryID, slots[id] === slot,
|
|
assignments[id, default: 0] == 0
|
|
else { return }
|
|
slots[id] = nil
|
|
counts[id] = nil
|
|
draining.remove(id)
|
|
stabilityTasks.removeValue(forKey: id)?.cancel()
|
|
publishPoolCounts()
|
|
await slot.connection.disconnect()
|
|
fillPool()
|
|
}
|
|
|
|
private func makeSlot(id: UUID) -> Slot {
|
|
let connection = makeConnection(
|
|
id,
|
|
{ [weak self] ruleID, generation, count in
|
|
Task { await self?.receivedCount(count, ruleID: ruleID, generation: generation, slotID: id) }
|
|
},
|
|
{ [weak self] message in
|
|
Task { await self?.failed(slotID: id, message: message) }
|
|
})
|
|
return Slot(id: id, connection: connection)
|
|
}
|
|
|
|
private func fillPool() {
|
|
guard !disconnecting, needsPool, additionalRequest != nil else { return }
|
|
let missing = desiredConnections - slots.count - opening.count - installing.count
|
|
guard missing > 0 else { return }
|
|
for _ in 0..<min(missing, max(0, 4 - opening.count - installing.count)) {
|
|
let id = UUID()
|
|
let slot = makeSlot(id: id)
|
|
opening[id] = slot
|
|
let resolver = additionalRequest!
|
|
openingTasks[id] = Task { [weak self] in
|
|
do {
|
|
let request = try await resolver()
|
|
try Task.checkCancellation()
|
|
try await slot.connection.connect(request, progress: nil)
|
|
try Task.checkCancellation()
|
|
await self?.opened(slot)
|
|
} catch {
|
|
await self?.openFailed(id: id, error: error)
|
|
await slot.connection.disconnect()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
private func opened(_ slot: Slot) async {
|
|
defer { openingTasks[slot.id] = nil }
|
|
guard opening.removeValue(forKey: slot.id) != nil, !disconnecting,
|
|
slots.count + installing.count < desiredConnections
|
|
else {
|
|
await slot.connection.disconnect()
|
|
return
|
|
}
|
|
installing[slot.id] = slot
|
|
do {
|
|
for registration in registrations.values where registration.rule.kind != .remote {
|
|
try await slot.connection.registerShared(registration.rule, generation: registration.generation)
|
|
slot.rules.insert(registration.rule.id)
|
|
}
|
|
guard installing.removeValue(forKey: slot.id) != nil, !disconnecting,
|
|
slots.count < desiredConnections
|
|
else {
|
|
await slot.connection.disconnect()
|
|
return
|
|
}
|
|
slots[slot.id] = slot
|
|
stabilityTasks[slot.id] = Task { [weak self] in
|
|
try? await Task.sleep(for: .seconds(60))
|
|
guard !Task.isCancelled else { return }
|
|
await self?.becameStable(slot.id)
|
|
}
|
|
publishPoolCounts()
|
|
fillPool()
|
|
} catch {
|
|
let wasInstalling = installing.removeValue(forKey: slot.id) != nil
|
|
await slot.connection.disconnect()
|
|
if wasInstalling { scheduleRetry(for: error) }
|
|
}
|
|
}
|
|
|
|
private func openFailed(id: UUID, error: any Error) {
|
|
openingTasks[id] = nil
|
|
guard opening.removeValue(forKey: id) != nil, !disconnecting else { return }
|
|
scheduleRetry(for: error)
|
|
}
|
|
|
|
private func scheduleRetry(for error: any Error) {
|
|
guard !Self.requiresManualRetry(error), retryTask == nil, needsPool,
|
|
!disconnecting, slots.count + opening.count + installing.count < desiredConnections
|
|
else { return }
|
|
let delays = [1, 2, 5, 10, 30]
|
|
let delay = delays[min(retryAttempt, delays.count - 1)]
|
|
retryAttempt += 1
|
|
retryTask = Task { [weak self] in
|
|
try? await Task.sleep(for: .seconds(delay))
|
|
guard !Task.isCancelled else { return }
|
|
await self?.retryReady()
|
|
}
|
|
}
|
|
|
|
private func retryReady() {
|
|
retryTask = nil
|
|
fillPool()
|
|
}
|
|
|
|
private func becameStable(_ id: UUID) {
|
|
guard slots[id] != nil else { return }
|
|
stabilityTasks[id] = nil
|
|
retryAttempt = 0
|
|
}
|
|
|
|
private func failed(slotID: UUID, message: String) async {
|
|
guard !disconnecting else { return }
|
|
if opening.removeValue(forKey: slotID) != nil || installing.removeValue(forKey: slotID) != nil {
|
|
scheduleRetry(for: TermindSSHError.disconnected)
|
|
return
|
|
}
|
|
guard slots.removeValue(forKey: slotID) != nil else { return }
|
|
let wasPrimary = primaryID == slotID
|
|
stabilityTasks.removeValue(forKey: slotID)?.cancel()
|
|
counts[slotID] = nil
|
|
draining.remove(slotID)
|
|
// A healthy draining connection can take over admission immediately.
|
|
// Counting it toward capacity without retaining it would suppress both
|
|
// new assignments and replacement opens until its old flows finish.
|
|
if !wasPrimary { rebalanceSlots() }
|
|
publishPoolCounts()
|
|
if wasPrimary, let promoted = slots.values.first {
|
|
primaryID = promoted.id
|
|
rebalanceSlots()
|
|
do {
|
|
for registration in registrations.values where registration.rule.kind == .remote {
|
|
try await promoted.connection.start(registration.rule, generation: registration.generation)
|
|
promoted.rules.insert(registration.rule.id)
|
|
}
|
|
fillPool()
|
|
for id in Array(draining) { await retireIfIdle(id) }
|
|
} catch {
|
|
failureHandler(error.localizedDescription)
|
|
}
|
|
} else if wasPrimary {
|
|
primaryID = nil
|
|
failureHandler(message)
|
|
} else {
|
|
scheduleRetry(for: TermindSSHError.disconnected)
|
|
}
|
|
}
|
|
|
|
private func retireAuxiliaries() async {
|
|
retryTask?.cancel()
|
|
retryTask = nil
|
|
for task in openingTasks.values { task.cancel() }
|
|
openingTasks.removeAll()
|
|
let pending = Array(opening.values) + Array(installing.values)
|
|
opening.removeAll()
|
|
installing.removeAll()
|
|
let primaryID = self.primaryID
|
|
let auxiliaries = slots.values.filter { $0.id != primaryID }
|
|
for slot in auxiliaries {
|
|
slots[slot.id] = nil
|
|
counts[slot.id] = nil
|
|
draining.remove(slot.id)
|
|
stabilityTasks.removeValue(forKey: slot.id)?.cancel()
|
|
}
|
|
await withTaskGroup(of: Void.self) { group in
|
|
for slot in auxiliaries + pending { group.addTask { await slot.connection.disconnect() } }
|
|
}
|
|
publishPoolCounts()
|
|
}
|
|
|
|
private func assign(_ descriptor: Int32, ruleID: UUID, generation: UUID) async {
|
|
do { try accessCheck?() } catch {
|
|
Darwin.close(descriptor)
|
|
return
|
|
}
|
|
guard let registration = registrations[ruleID], registration.generation == generation else {
|
|
Darwin.close(descriptor)
|
|
return
|
|
}
|
|
let eligible = slots.values.filter { $0.rules.contains(ruleID) && !draining.contains($0.id) }
|
|
guard !eligible.isEmpty else {
|
|
Darwin.close(descriptor)
|
|
return
|
|
}
|
|
let minimum = eligible.map { counts[$0.id]?[ruleID] ?? 0 }.min() ?? 0
|
|
let candidates = eligible.filter { counts[$0.id]?[ruleID] ?? 0 == minimum }.sorted {
|
|
$0.id.uuidString < $1.id.uuidString
|
|
}
|
|
let slot = candidates[roundRobin % candidates.count]
|
|
roundRobin &+= 1
|
|
assignments[slot.id, default: 0] += 1
|
|
let accepted = await slot.connection.acceptShared(descriptor, ruleID: ruleID, generation: generation)
|
|
assignments[slot.id, default: 1] -= 1
|
|
if assignments[slot.id] == 0 { assignments[slot.id] = nil }
|
|
if !accepted {
|
|
Darwin.close(descriptor)
|
|
}
|
|
await retireIfIdle(slot.id)
|
|
}
|
|
|
|
private func receivedCount(_ count: Int, ruleID: UUID, generation: UUID, slotID: UUID) async {
|
|
guard registrations[ruleID]?.generation == generation, slots[slotID] != nil else { return }
|
|
counts[slotID, default: [:]][ruleID] = count
|
|
report(ruleID, generation: generation)
|
|
if count == 0 { await retireIfIdle(slotID) }
|
|
}
|
|
|
|
private func report(_ ruleID: UUID, generation: UUID?) {
|
|
guard let generation else { return }
|
|
let total = counts.values.reduce(0) { $0 + ($1[ruleID] ?? 0) }
|
|
connectionHandler(ruleID, generation, total)
|
|
}
|
|
|
|
private func publishPoolCounts() {
|
|
for registration in registrations.values {
|
|
report(registration.rule.id, generation: registration.generation)
|
|
}
|
|
}
|
|
|
|
private static func requiresManualRetry(_ error: any Error) -> Bool {
|
|
if error is SSHInteractionRequired || error is CancellationError { return true }
|
|
guard let ssh = error as? TermindSSHError else { return false }
|
|
if ssh.isLocalAuthenticationRequirement || ssh.shouldRequestAuthentication || ssh.shouldRequestHostKey {
|
|
return true
|
|
}
|
|
switch ssh {
|
|
case .untrustedHost, .mismatchedHostKey, .hostKeyChanged: return true
|
|
default: return false
|
|
}
|
|
}
|
|
}
|
|
|
|
private nonisolated final class SSHForwardingLocalListeners: @unchecked Sendable {
|
|
typealias Handler = @Sendable (PortForwardingRule.ID, UUID, Int32) -> Void
|
|
|
|
private final class Entry: @unchecked Sendable {
|
|
let generation: UUID
|
|
let socket: Int32
|
|
var events: SSHSocketEvents?
|
|
|
|
init(generation: UUID, socket: Int32) {
|
|
self.generation = generation
|
|
self.socket = socket
|
|
}
|
|
}
|
|
|
|
private let queue = DispatchQueue(label: "com.akinokaede.termind.ssh.forwarding-listeners")
|
|
private var entries: [PortForwardingRule.ID: Entry] = [:]
|
|
private var handler: Handler
|
|
|
|
init(handler: @escaping Handler) { self.handler = handler }
|
|
|
|
func setHandler(_ handler: @escaping Handler) {
|
|
queue.sync { self.handler = handler }
|
|
}
|
|
|
|
func start(_ rule: PortForwardingRule, generation: UUID) async throws {
|
|
try await withCheckedThrowingContinuation { continuation in
|
|
queue.async { [self] in
|
|
do {
|
|
if let current = entries[rule.id] {
|
|
guard current.generation == generation else { throw CancellationError() }
|
|
continuation.resume()
|
|
return
|
|
}
|
|
let socket = try makeListener(host: rule.bindHost, port: rule.bindPort)
|
|
let entry = Entry(generation: generation, socket: socket)
|
|
entry.events = SSHSocketEvents(socket: socket, queue: queue) { [weak self, weak entry] in
|
|
guard let self, let entry else { return }
|
|
self.accept(ruleID: rule.id, entry: entry)
|
|
}
|
|
entries[rule.id] = entry
|
|
entry.events?.update(read: true, write: false)
|
|
continuation.resume()
|
|
} catch {
|
|
continuation.resume(throwing: error)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func stop(_ id: PortForwardingRule.ID) async {
|
|
await withCheckedContinuation { continuation in
|
|
queue.async { [self] in
|
|
guard let entry = entries.removeValue(forKey: id) else {
|
|
continuation.resume()
|
|
return
|
|
}
|
|
if let events = entry.events {
|
|
events.cancel {
|
|
Darwin.close(entry.socket)
|
|
continuation.resume()
|
|
}
|
|
} else {
|
|
Darwin.close(entry.socket)
|
|
continuation.resume()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func disconnect() async {
|
|
let ids = await withCheckedContinuation { continuation in
|
|
queue.async { continuation.resume(returning: Array(self.entries.keys)) }
|
|
}
|
|
for id in ids { await stop(id) }
|
|
}
|
|
|
|
private func accept(ruleID: UUID, entry: Entry) {
|
|
guard entries[ruleID] === entry else { return }
|
|
for _ in 0..<32 {
|
|
let descriptor = Darwin.accept(entry.socket, nil, nil)
|
|
guard descriptor >= 0 else { break }
|
|
handler(ruleID, entry.generation, descriptor)
|
|
}
|
|
}
|
|
|
|
private func makeListener(host: String, port: Int) throws -> Int32 {
|
|
var hints = addrinfo()
|
|
hints.ai_family = AF_UNSPEC
|
|
hints.ai_socktype = SOCK_STREAM
|
|
hints.ai_flags = AI_PASSIVE
|
|
var result: UnsafeMutablePointer<addrinfo>?
|
|
let resolved = getaddrinfo(host, String(port), &hints, &result)
|
|
guard resolved == 0, let first = result else {
|
|
throw TermindSSHError.socket(String(cString: gai_strerror(resolved)))
|
|
}
|
|
defer { freeaddrinfo(first) }
|
|
var cursor: UnsafeMutablePointer<addrinfo>? = first
|
|
while let info = cursor?.pointee {
|
|
let candidate = Darwin.socket(info.ai_family, info.ai_socktype, info.ai_protocol)
|
|
if candidate >= 0 {
|
|
var yes: Int32 = 1
|
|
setsockopt(candidate, SOL_SOCKET, SO_NOSIGPIPE, &yes, socklen_t(MemoryLayout<Int32>.size))
|
|
if setsockopt(candidate, SOL_SOCKET, SO_REUSEADDR, &yes, socklen_t(MemoryLayout<Int32>.size)) == 0,
|
|
Darwin.bind(candidate, info.ai_addr, info.ai_addrlen) == 0,
|
|
Darwin.listen(candidate, 128) == 0
|
|
{
|
|
let flags = fcntl(candidate, F_GETFL, 0)
|
|
if flags >= 0 { _ = fcntl(candidate, F_SETFL, flags | O_NONBLOCK) }
|
|
return candidate
|
|
}
|
|
Darwin.close(candidate)
|
|
}
|
|
cursor = info.ai_next
|
|
}
|
|
throw TermindSSHError.socket(String(cString: strerror(errno)))
|
|
}
|
|
}
|