1904 lines
80 KiB
Swift
1904 lines
80 KiB
Swift
// SPDX-FileCopyrightText: 2026 Kaede Akino
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
import AgentKit
|
|
import Foundation
|
|
import TermindApplication
|
|
import TermindDomain
|
|
|
|
import struct TermindDomain.Host
|
|
|
|
nonisolated public struct AgentHostTarget: Identifiable, Hashable, Sendable, Codable {
|
|
public struct ProtocolSupport: Hashable, Sendable, Codable {
|
|
public var name: HostTerminalProtocol
|
|
public var address: String
|
|
public var connectionState: String
|
|
|
|
public init(name: HostTerminalProtocol, address: String, connectionState: String) {
|
|
self.name = name
|
|
self.address = address
|
|
self.connectionState = connectionState
|
|
}
|
|
}
|
|
|
|
public var id: UUID
|
|
public var displayName: String
|
|
public var address: String
|
|
public var tags: [String]
|
|
public var connectionState: String
|
|
public var isTemporary: Bool
|
|
public var groupID: HostGroup.ID?
|
|
public var groupPath: String?
|
|
public var preferredTerminalProtocol: HostTerminalProtocol
|
|
public var protocols: [ProtocolSupport]
|
|
|
|
public init(
|
|
id: UUID, displayName: String, address: String, tags: [String], connectionState: String, isTemporary: Bool,
|
|
groupID: HostGroup.ID? = nil, groupPath: String? = nil,
|
|
preferredTerminalProtocol: HostTerminalProtocol = .ssh,
|
|
protocols: [ProtocolSupport]? = nil
|
|
) {
|
|
self.id = id
|
|
self.displayName = displayName
|
|
self.address = address
|
|
self.tags = tags
|
|
self.connectionState = connectionState
|
|
self.isTemporary = isTemporary
|
|
self.groupID = groupID
|
|
self.groupPath = groupPath
|
|
self.preferredTerminalProtocol = preferredTerminalProtocol
|
|
self.protocols =
|
|
protocols ?? [
|
|
ProtocolSupport(name: preferredTerminalProtocol, address: address, connectionState: connectionState)
|
|
]
|
|
}
|
|
}
|
|
|
|
nonisolated public struct AgentHostGroupTarget: Identifiable, Hashable, Sendable {
|
|
public var id: HostGroup.ID
|
|
public var name: String
|
|
public var path: String
|
|
public var children: [AgentHostGroupTarget]
|
|
|
|
public init(id: HostGroup.ID, name: String, path: String, children: [AgentHostGroupTarget]) {
|
|
self.id = id
|
|
self.name = name
|
|
self.path = path
|
|
self.children = children
|
|
}
|
|
}
|
|
|
|
nonisolated public struct AgentHostListing: Hashable, Sendable {
|
|
public var groups: [AgentHostGroupTarget]
|
|
public var hosts: [AgentHostTarget]
|
|
|
|
public init(groups: [AgentHostGroupTarget], hosts: [AgentHostTarget]) {
|
|
self.groups = groups
|
|
self.hosts = hosts
|
|
}
|
|
}
|
|
|
|
/// Main-actor bridge over the Vault. The Agent receives only public host
|
|
/// metadata; Keychain material is resolved directly into SSH configuration and
|
|
/// never serialized into a tool result.
|
|
@MainActor
|
|
public final class AgentHostResolver: Sendable {
|
|
private let store: HostLibraryModel
|
|
private var temporaryTargets: @MainActor @Sendable () -> [AgentHostTarget]
|
|
private var temporaryConfiguration: @MainActor @Sendable (UUID) throws -> SSHConnectionConfiguration?
|
|
|
|
public init(
|
|
store: HostLibraryModel,
|
|
temporaryTargets: @escaping @MainActor @Sendable () -> [AgentHostTarget] = { [] },
|
|
temporaryConfiguration: @escaping @MainActor @Sendable (UUID) throws -> SSHConnectionConfiguration? = { _ in nil
|
|
}
|
|
) {
|
|
self.store = store
|
|
self.temporaryTargets = temporaryTargets
|
|
self.temporaryConfiguration = temporaryConfiguration
|
|
}
|
|
|
|
public func configureTemporaryTargets(
|
|
targets: @escaping @MainActor @Sendable () -> [AgentHostTarget],
|
|
configuration: @escaping @MainActor @Sendable (UUID) throws -> SSHConnectionConfiguration?
|
|
) {
|
|
temporaryTargets = targets
|
|
temporaryConfiguration = configuration
|
|
}
|
|
|
|
public func listing() -> AgentHostListing {
|
|
let paths = groupPaths()
|
|
let saved = store.hosts.map { host in
|
|
let state = Self.stateName(store.status(of: host.id))
|
|
let preferred = host.preferredTerminalProtocol ?? .ssh
|
|
let protocols = host.supportedTerminalProtocols.map { connectionProtocol in
|
|
AgentHostTarget.ProtocolSupport(
|
|
name: connectionProtocol,
|
|
address: store.loginAddress(for: host, protocol: connectionProtocol),
|
|
connectionState: state)
|
|
}
|
|
return AgentHostTarget(
|
|
id: host.id, displayName: host.name,
|
|
address: store.loginAddress(for: host, protocol: preferred),
|
|
tags: host.tags, connectionState: state,
|
|
isTemporary: false, groupID: host.groupID,
|
|
groupPath: host.groupID.flatMap { paths[$0] },
|
|
preferredTerminalProtocol: preferred, protocols: protocols
|
|
)
|
|
}
|
|
return AgentHostListing(
|
|
groups: groupTree(paths: paths),
|
|
hosts: saved
|
|
+ temporaryTargets().map { target in
|
|
var target = target
|
|
target.groupID = nil
|
|
target.groupPath = nil
|
|
return target
|
|
}
|
|
)
|
|
}
|
|
|
|
/// Produces every valid path without trusting persisted parent links to be
|
|
/// acyclic. A broken parent simply ends the path; a cycle never recurses.
|
|
private func groupPaths() -> [HostGroup.ID: String] {
|
|
var result: [HostGroup.ID: String] = [:]
|
|
for group in store.groups {
|
|
var names: [String] = []
|
|
var current: HostGroup? = group
|
|
var visited: Set<HostGroup.ID> = []
|
|
while let node = current, visited.insert(node.id).inserted {
|
|
names.append(node.name)
|
|
current = node.parentID.flatMap(store.group)
|
|
}
|
|
result[group.id] = names.reversed().joined(separator: " / ")
|
|
}
|
|
return result
|
|
}
|
|
|
|
/// Roots include ordinary top-level groups and groups whose parent is
|
|
/// missing. Any cyclic component that has no root is appended afterwards,
|
|
/// exactly once, so corrupt metadata cannot make a group disappear.
|
|
private func groupTree(paths: [HostGroup.ID: String]) -> [AgentHostGroupTarget] {
|
|
let all = Dictionary(uniqueKeysWithValues: store.groups.map { ($0.id, $0) })
|
|
func ordered(_ groups: [HostGroup]) -> [HostGroup] {
|
|
groups.sorted {
|
|
if $0.sortIndex != $1.sortIndex { return $0.sortIndex < $1.sortIndex }
|
|
return $0.id.uuidString < $1.id.uuidString
|
|
}
|
|
}
|
|
var emitted: Set<HostGroup.ID> = []
|
|
func build(_ group: HostGroup, ancestors: Set<HostGroup.ID>) -> AgentHostGroupTarget {
|
|
emitted.insert(group.id)
|
|
let nextAncestors = ancestors.union([group.id])
|
|
let children = ordered(store.groups.filter { $0.parentID == group.id })
|
|
.filter { !nextAncestors.contains($0.id) }
|
|
.map { build($0, ancestors: nextAncestors) }
|
|
return AgentHostGroupTarget(
|
|
id: group.id, name: group.name, path: paths[group.id] ?? group.name,
|
|
children: children
|
|
)
|
|
}
|
|
let roots = ordered(
|
|
store.groups.filter { group in
|
|
guard let parentID = group.parentID else { return true }
|
|
return all[parentID] == nil
|
|
})
|
|
var result = roots.map { build($0, ancestors: []) }
|
|
for group in ordered(store.groups) where !emitted.contains(group.id) {
|
|
result.append(build(group, ancestors: []))
|
|
}
|
|
return result
|
|
}
|
|
|
|
public func configuration(for id: UUID) throws -> SSHConnectionConfiguration {
|
|
if let host = store.host(id) {
|
|
try store.requireHostAccess(host)
|
|
guard host.sshEnabled else { throw HostProtocolUnavailableError(useCase: .sshTerminal) }
|
|
return try store.sshConnectionConfiguration(for: host)
|
|
}
|
|
if let configuration = try temporaryConfiguration(id) { return configuration }
|
|
throw AgentToolError.invalidHost
|
|
}
|
|
|
|
public func resolvedConfiguration(for id: UUID) async throws -> SSHConnectionConfiguration {
|
|
if let host = store.host(id) { try await store.resolveHostAccess(host) }
|
|
return try configuration(for: id)
|
|
}
|
|
|
|
public func resolvedRouteFingerprint(for id: UUID) async throws -> String {
|
|
if let host = store.host(id) { try await store.resolveHostAccess(host) }
|
|
return try routeFingerprint(for: id)
|
|
}
|
|
|
|
public func canResolve(_ id: UUID) -> Bool {
|
|
store.host(id)?.sshEnabled == true
|
|
|| temporaryTargets().contains { target in
|
|
target.id == id && target.protocols.contains { $0.name == .ssh }
|
|
}
|
|
}
|
|
|
|
fileprivate struct InteractiveResolution: Sendable {
|
|
var configuration: SSHConnectionConfiguration
|
|
var submissions: [Host.ID: SSHAuthenticationSubmission]
|
|
var proxySubmissions: [Host.ID: SSHProxyAuthenticationSubmission]
|
|
var temporaryHostKeys: [Host.ID: SSHHostKeyIdentity]
|
|
}
|
|
|
|
/// Resolves every local authentication checkpoint in route order. The
|
|
/// interaction handler receives public metadata and returns secrets
|
|
/// directly to this main-actor bridge; none of it becomes tool content.
|
|
fileprivate func interactiveConfiguration(
|
|
for id: UUID,
|
|
invocationID: AgentToolInvocation.ID,
|
|
interaction: AgentSSHInteraction,
|
|
submissions initial: [Host.ID: SSHAuthenticationSubmission],
|
|
proxySubmissions initialProxySubmissions: [Host.ID: SSHProxyAuthenticationSubmission],
|
|
temporaryHostKeys initialTemporaryHostKeys: [Host.ID: SSHHostKeyIdentity],
|
|
forcedHostID initialForcedHostID: Host.ID? = nil,
|
|
forcedProxy initialForcedProxy: Bool = false,
|
|
failure initialFailure: String? = nil
|
|
) async throws -> InteractiveResolution {
|
|
guard let destination = store.host(id) else {
|
|
// Pane-backed targets already carry a complete configuration and
|
|
// cannot be saved back to the Vault from an Agent card.
|
|
if let configuration = try temporaryConfiguration(id) {
|
|
return InteractiveResolution(
|
|
configuration: configuration, submissions: initial,
|
|
proxySubmissions: initialProxySubmissions,
|
|
temporaryHostKeys: initialTemporaryHostKeys
|
|
)
|
|
}
|
|
throw AgentToolError.invalidHost
|
|
}
|
|
try await store.resolveHostAccess(destination)
|
|
guard destination.sshEnabled else {
|
|
throw HostProtocolUnavailableError(useCase: .sshTerminal)
|
|
}
|
|
|
|
var submissions = initial
|
|
var proxySubmissions = initialProxySubmissions
|
|
let temporaryHostKeys = initialTemporaryHostKeys
|
|
var forcedHostID = initialForcedHostID
|
|
var forcedProxy = initialForcedProxy
|
|
var failure = initialFailure
|
|
// Match Terminal and SFTP: resolve local credential requirements first,
|
|
// then begin the network handshake and host-key verification. In
|
|
// particular, a protected key must raise Passphrase even when the peer
|
|
// later fails before returning its SSH banner.
|
|
for node in store.route(for: destination) {
|
|
while proxyAuthenticationRequired(
|
|
for: node, submissions: proxySubmissions,
|
|
forcedHostID: forcedHostID, forcedProxy: forcedProxy
|
|
) {
|
|
guard let proxy = node.proxy else { break }
|
|
let existing = proxySubmissions[node.id]
|
|
let configuredIdentity = (existing?.identityID ?? proxy.identityID).flatMap(store.identity)
|
|
let selectedIdentity: Identity?
|
|
if let configuredIdentity {
|
|
switch store.passwordAvailability(for: configuredIdentity) {
|
|
case .available, .notConfigured: selectedIdentity = configuredIdentity
|
|
case .missing, .failed: selectedIdentity = nil
|
|
}
|
|
} else {
|
|
selectedIdentity = nil
|
|
}
|
|
let request = AgentSSHAuthenticationRequest(
|
|
invocationID: invocationID, hostID: node.id,
|
|
displayName: String(localized: "Proxy for \(node.name)"),
|
|
address: "\(proxy.hostname):\(proxy.port)", mode: .proxy(ownerHostID: node.id),
|
|
selectedIdentityID: selectedIdentity?.id,
|
|
username: existing?.username ?? proxy.username, failureMessage: failure)
|
|
switch await interaction.authenticate(request) {
|
|
case .cancelled, .keyboardInteractive, .securityKey, .securityKeyFailed:
|
|
throw AgentToolError.cancelled
|
|
case .submit(let answer):
|
|
let submission = SSHProxyAuthenticationSubmission(
|
|
username: answer.username, password: answer.password, identityID: answer.identityID)
|
|
if answer.save {
|
|
_ = try store.saveProxyAuthentication(
|
|
submission, for: node,
|
|
selectedIdentity: answer.identityID.flatMap(store.identity),
|
|
suppliedPassword: answer.password)
|
|
}
|
|
proxySubmissions[node.id] = submission
|
|
}
|
|
forcedHostID = nil
|
|
forcedProxy = false
|
|
failure = nil
|
|
}
|
|
while let requirement = authenticationRequirement(
|
|
for: node, submissions: submissions, forcedHostID: forcedHostID
|
|
) {
|
|
let existing = submissions[node.id]
|
|
let identity =
|
|
existing?.identityID.flatMap(store.identity)
|
|
?? store.identity(for: node)
|
|
let request = AgentSSHAuthenticationRequest(
|
|
invocationID: invocationID, hostID: node.id,
|
|
displayName: node.name, address: node.displayAddress,
|
|
mode: requirement.mode,
|
|
selectedIdentityID: existing?.identityID ?? identity?.id,
|
|
username: existing?.username
|
|
?? (node.isTailscaleManaged ? store.tailscaleUsername(for: node) : identity?.username) ?? "",
|
|
selectedKeyID: existing?.keyID ?? identity?.keyID,
|
|
failureMessage: failure
|
|
)
|
|
switch await interaction.authenticate(request) {
|
|
case .cancelled, .keyboardInteractive, .securityKey, .securityKeyFailed:
|
|
throw AgentToolError.cancelled
|
|
case .submit(let answer):
|
|
if case .passphrase = requirement.mode {
|
|
do {
|
|
submissions[node.id] = try store.passphraseSubmission(
|
|
for: request.journeyPrompt, existing: submissions[node.id],
|
|
passphrase: answer.passphrase, save: answer.save)
|
|
} catch {
|
|
failure = error.localizedDescription
|
|
continue
|
|
}
|
|
} else {
|
|
submissions[node.id] = try submission(from: answer, for: node)
|
|
}
|
|
}
|
|
forcedHostID = nil
|
|
forcedProxy = false
|
|
failure = nil
|
|
}
|
|
|
|
}
|
|
|
|
var retainedSubmissions = submissions
|
|
let attemptSubmissions = retainedSubmissions.takeAuthenticationSecrets()
|
|
return InteractiveResolution(
|
|
configuration: try store.sshConnectionConfiguration(
|
|
for: destination, submissions: attemptSubmissions,
|
|
proxySubmissions: proxySubmissions.takeProxyAuthenticationSecrets(),
|
|
temporaryHostKeys: temporaryHostKeys,
|
|
trustPolicy: .allowChallenge
|
|
),
|
|
submissions: retainedSubmissions, proxySubmissions: proxySubmissions,
|
|
temporaryHostKeys: temporaryHostKeys
|
|
)
|
|
}
|
|
|
|
private func proxyAuthenticationRequired(
|
|
for host: Host, submissions: [Host.ID: SSHProxyAuthenticationSubmission],
|
|
forcedHostID: Host.ID?, forcedProxy: Bool
|
|
) -> Bool {
|
|
guard let proxy = host.proxy else { return false }
|
|
if forcedProxy, forcedHostID == host.id { return true }
|
|
if let submission = submissions[host.id] {
|
|
return submission.requiresPasswordReentry
|
|
}
|
|
guard let identityID = proxy.identityID, let identity = store.identity(identityID) else {
|
|
return proxy.identityID != nil
|
|
}
|
|
switch store.passwordAvailability(for: identity) {
|
|
case .available: return false
|
|
case .notConfigured, .missing, .failed: return true
|
|
}
|
|
}
|
|
|
|
fileprivate func confirmHostKey(
|
|
_ challenge: SSHHostKeyChallenge,
|
|
invocationID: AgentToolInvocation.ID,
|
|
interaction: AgentSSHInteraction
|
|
) async -> AgentSSHHostKeyResponse {
|
|
let saved = store.host(challenge.hostID)
|
|
let temporary = temporaryTargets().first { $0.id == challenge.hostID }
|
|
let displayName = saved?.name ?? temporary?.displayName ?? challenge.observation.hostname
|
|
let address =
|
|
saved?.displayAddress ?? temporary?.address
|
|
?? (challenge.observation.port == 22
|
|
? challenge.observation.hostname
|
|
: "\(challenge.observation.hostname):\(challenge.observation.port)")
|
|
let response = await interaction.confirmHostKey(
|
|
AgentSSHHostKeyRequest(
|
|
invocationID: invocationID, hostID: challenge.hostID,
|
|
displayName: displayName, address: address,
|
|
prompt: SSHHostKeyPrompt(challenge: challenge),
|
|
allowsPersistentTrust: saved != nil
|
|
)
|
|
)
|
|
if response == .trustAndContinue, saved != nil {
|
|
store.trust(
|
|
KnownHostEntry(
|
|
hostname: challenge.observation.hostname,
|
|
port: challenge.observation.port,
|
|
algorithm: challenge.observation.algorithm,
|
|
fingerprint: challenge.observation.fingerprint,
|
|
addedAt: .now
|
|
)
|
|
)
|
|
}
|
|
return response
|
|
}
|
|
|
|
private struct AuthenticationRequirement {
|
|
var mode: AgentSSHAuthenticationRequest.Mode
|
|
}
|
|
|
|
private func authenticationRequirement(
|
|
for host: Host,
|
|
submissions: [Host.ID: SSHAuthenticationSubmission],
|
|
forcedHostID: Host.ID?
|
|
) -> AuthenticationRequirement? {
|
|
guard forcedHostID == nil || forcedHostID == host.id else { return nil }
|
|
if host.isTailscaleManaged {
|
|
let username = (submissions[host.id]?.username ?? store.savedTailscaleUsername(for: host))?
|
|
.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
return username?.isEmpty != false || forcedHostID == host.id
|
|
? AuthenticationRequirement(mode: .tailscaleUsername) : nil
|
|
}
|
|
if let submission = submissions[host.id] {
|
|
if let keyID = submission.keyID,
|
|
store.key(keyID).map({
|
|
(forcedHostID == host.id && $0.hasPassphrase) || store.requiresPassphraseEntry(for: $0)
|
|
}) == true,
|
|
submission.passphrase?.isEmpty != false
|
|
{
|
|
return AuthenticationRequirement(mode: .passphrase(keyID: keyID))
|
|
}
|
|
if submission.username.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty {
|
|
return AuthenticationRequirement(mode: .identity)
|
|
}
|
|
if forcedHostID == host.id {
|
|
return AuthenticationRequirement(mode: .identity)
|
|
}
|
|
return nil
|
|
}
|
|
guard let identity = store.identity(for: host) else {
|
|
return AuthenticationRequirement(mode: .identity)
|
|
}
|
|
if forcedHostID == host.id, let keyID = identity.keyID, store.key(keyID)?.hasPassphrase == true {
|
|
return AuthenticationRequirement(mode: .passphrase(keyID: keyID))
|
|
}
|
|
let resolution = store.sshCredentials(for: identity)
|
|
if let keyID = resolution.requiredPassphraseKeyID {
|
|
return AuthenticationRequirement(mode: .passphrase(keyID: keyID))
|
|
}
|
|
if identity.username.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
|
|
|| (identity.hasSecret && resolution.credentials.isEmpty)
|
|
{
|
|
return AuthenticationRequirement(mode: .identity)
|
|
}
|
|
if forcedHostID == host.id {
|
|
return AuthenticationRequirement(mode: .identity)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
private func submission(
|
|
from answer: AgentSSHAuthenticationAnswer, for host: Host
|
|
) throws -> SSHAuthenticationSubmission {
|
|
if host.isTailscaleManaged {
|
|
let username = answer.username.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
guard !username.isEmpty else { throw TermindSSHError.missingUsername }
|
|
if answer.save, !store.saveTailscaleUsername(username, hostID: host.id) {
|
|
throw TailscaleConfigurationError.unavailable
|
|
}
|
|
return SSHAuthenticationSubmission(username: username)
|
|
}
|
|
if let identityID = answer.identityID {
|
|
guard let identity = store.identity(identityID) else {
|
|
throw TermindSSHError.missingCredential
|
|
}
|
|
if answer.save {
|
|
var updated = host
|
|
updated.identityID = identity.id
|
|
store.upsert(updated)
|
|
}
|
|
return SSHAuthenticationSubmission(
|
|
username: identity.username,
|
|
passphrase: answer.passphrase, keyID: identity.keyID,
|
|
identityID: identity.id
|
|
)
|
|
}
|
|
|
|
let username = answer.username.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
guard !username.isEmpty else { throw TermindSSHError.missingUsername }
|
|
var identityID: Identity.ID?
|
|
if answer.save {
|
|
let name = answer.identityName?.trimmingCharacters(in: .whitespacesAndNewlines)
|
|
let identity = Identity(
|
|
name: name?.isEmpty == false ? name! : username,
|
|
username: username, keyID: answer.keyID
|
|
)
|
|
try store.upsert(
|
|
identity,
|
|
passwordUpdate: answer.password.map(IdentityPasswordUpdate.replace) ?? .remove
|
|
)
|
|
var updated = host
|
|
updated.identityID = identity.id
|
|
store.upsert(updated)
|
|
identityID = identity.id
|
|
}
|
|
return SSHAuthenticationSubmission(
|
|
username: username, password: answer.password,
|
|
passphrase: answer.passphrase, keyID: answer.keyID,
|
|
identityID: identityID
|
|
)
|
|
}
|
|
|
|
public func agentForwardingEnabled(for id: UUID) -> Bool { store.host(id)?.agentForwardingEnabled == true }
|
|
|
|
/// A stable description of *which machine, reached how, and as whom* — the
|
|
/// only question a pooled connection has to keep answering the same way.
|
|
///
|
|
/// Built from records the store already holds in memory, with deliberately
|
|
/// no Keychain read: this runs before every reuse, and resolving credentials
|
|
/// here would put private-key material on the hot path of every command to
|
|
/// no purpose. A rotated credential on an unchanged identity is therefore
|
|
/// not a mismatch — the live connection is still the same box and the same
|
|
/// login, and the new credential applies to the next connect.
|
|
public func routeFingerprint(for id: UUID) throws -> String {
|
|
guard let host = store.host(id) else {
|
|
// A pane-backed target has no Vault record to describe, and its
|
|
// configuration is rebuilt from the live pane on every connect. Tie
|
|
// the fingerprint to that pane still existing, so a closed workspace
|
|
// invalidates the Agent's connection to it.
|
|
guard temporaryTargets().contains(where: { $0.id == id }) else {
|
|
throw AgentToolError.invalidHost
|
|
}
|
|
return "temporary:\(id.uuidString)"
|
|
}
|
|
try store.requireHostAccess(host)
|
|
guard host.sshEnabled else { throw HostProtocolUnavailableError(useCase: .sshTerminal) }
|
|
let route = store.route(for: host)
|
|
return route.enumerated().map { index, node in
|
|
let identity = store.identity(for: node)
|
|
let proxy = node.proxy.map {
|
|
"\($0.kind.rawValue)://\($0.username)@\($0.hostname):\($0.port)"
|
|
}
|
|
let keys = store.knownHosts(for: node).map(\.fingerprint).sorted()
|
|
return [
|
|
node.id.uuidString, node.host, String(node.port),
|
|
identity?.id.uuidString ?? "", identity?.username ?? "",
|
|
proxy ?? "", keys.joined(separator: ","),
|
|
index == route.count - 1 ? node.charset.rawValue : "",
|
|
].joined(separator: "\u{1F}")
|
|
}.joined(separator: "\u{1E}")
|
|
}
|
|
|
|
private static func stateName(_ status: ConnectionStatus) -> String {
|
|
switch status {
|
|
case .disconnected: "disconnected"
|
|
case .connecting: "connecting"
|
|
case .connected: "connected"
|
|
case .reconnecting: "reconnecting"
|
|
case .failed: "failed"
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Keeps temporary credentials alive for one overlapping wave of Agent SSH
|
|
/// connection attempts. Exec and primary SFTP join one connection open, so every
|
|
/// overlapping caller observes the same ephemeral answer and generation.
|
|
public actor AgentSSHAuthenticationContext {
|
|
public struct Attempt: Sendable {
|
|
public var request: SSHConnectionRequest
|
|
public var generation: UInt64
|
|
|
|
public var configuration: SSHConnectionConfiguration { request.configuration }
|
|
|
|
public init(request: SSHConnectionRequest, generation: UInt64) {
|
|
self.request = request
|
|
self.generation = generation
|
|
}
|
|
|
|
}
|
|
|
|
private struct Pending {
|
|
var id: UUID
|
|
var generation: UInt64
|
|
var task: Task<AgentHostResolver.InteractiveResolution, any Error>
|
|
var waiters: [UUID: CheckedContinuation<AgentHostResolver.InteractiveResolution, any Error>]
|
|
}
|
|
|
|
private let resolver: AgentHostResolver
|
|
private let hostID: Host.ID
|
|
private let invocationID: AgentToolInvocation.ID
|
|
private let interaction: AgentSSHInteraction
|
|
private var submissions: [Host.ID: SSHAuthenticationSubmission] = [:]
|
|
private var proxySubmissions: [Host.ID: SSHProxyAuthenticationSubmission] = [:]
|
|
private var temporaryHostKeys: [Host.ID: SSHHostKeyIdentity] = [:]
|
|
private var cached: SSHConnectionConfiguration?
|
|
private var generation: UInt64 = 0
|
|
private var pending: Pending?
|
|
private var forcedHostID: Host.ID?
|
|
private var forcedProxy = false
|
|
private var failure: String?
|
|
private struct PendingHostKey {
|
|
var fingerprint: String
|
|
var task: Task<AgentSSHHostKeyResponse, Never>
|
|
var waiters: [UUID: CheckedContinuation<AgentSSHHostKeyResponse?, Never>]
|
|
}
|
|
private var pendingHostKeys: [Host.ID: PendingHostKey] = [:]
|
|
|
|
public init(
|
|
resolver: AgentHostResolver, hostID: Host.ID,
|
|
invocationID: AgentToolInvocation.ID,
|
|
interaction: AgentSSHInteraction
|
|
) {
|
|
self.resolver = resolver
|
|
self.hostID = hostID
|
|
self.invocationID = invocationID
|
|
self.interaction = interaction
|
|
}
|
|
|
|
public func configurationAttempt() async throws -> Attempt {
|
|
try Task.checkCancellation()
|
|
if let cached {
|
|
try cached.accessCheck?()
|
|
return makeAttempt(configuration: cached, generation: generation)
|
|
}
|
|
|
|
let pendingID: UUID
|
|
let pendingGeneration: UInt64
|
|
if let pending, pending.generation == generation {
|
|
pendingID = pending.id
|
|
pendingGeneration = pending.generation
|
|
} else {
|
|
let id = UUID()
|
|
let currentGeneration = generation
|
|
let resolver = resolver
|
|
let hostID = hostID
|
|
let invocationID = invocationID
|
|
let interaction = interaction
|
|
let submissions = submissions
|
|
let proxySubmissions = proxySubmissions
|
|
let temporaryHostKeys = temporaryHostKeys
|
|
let forcedHostID = forcedHostID
|
|
let forcedProxy = forcedProxy
|
|
let failure = failure
|
|
let task = Task {
|
|
try await resolver.interactiveConfiguration(
|
|
for: hostID, invocationID: invocationID, interaction: interaction,
|
|
submissions: submissions, proxySubmissions: proxySubmissions,
|
|
temporaryHostKeys: temporaryHostKeys,
|
|
forcedHostID: forcedHostID, forcedProxy: forcedProxy, failure: failure
|
|
)
|
|
}
|
|
pending = Pending(
|
|
id: id, generation: currentGeneration, task: task, waiters: [:]
|
|
)
|
|
pendingID = id
|
|
pendingGeneration = currentGeneration
|
|
Task { [weak self] in
|
|
do {
|
|
let resolution = try await task.value
|
|
await self?.finishConfigurationAttempt(
|
|
id: id, generation: currentGeneration, result: .success(resolution)
|
|
)
|
|
} catch {
|
|
await self?.finishConfigurationAttempt(
|
|
id: id, generation: currentGeneration, result: .failure(error)
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
let resolution = try await waitForConfigurationAttempt(
|
|
id: pendingID, generation: pendingGeneration
|
|
)
|
|
try Task.checkCancellation()
|
|
return makeAttempt(
|
|
configuration: resolution.configuration,
|
|
generation: pendingGeneration
|
|
)
|
|
}
|
|
|
|
/// Convenience for focused resolver tests that do not need to simulate a
|
|
/// rejected connection attempt.
|
|
public func configuration() async throws -> SSHConnectionConfiguration {
|
|
try await configurationAttempt().configuration
|
|
}
|
|
|
|
public var currentGeneration: UInt64 { generation }
|
|
|
|
public func canResolveTarget() async -> Bool {
|
|
await resolver.canResolve(hostID)
|
|
}
|
|
|
|
public func retry(after error: TermindSSHError, generation failedGeneration: UInt64) {
|
|
// Overlapping callers may report the same failed attempt. Only the
|
|
// first failure advances the generation; a late failure from the old
|
|
// generation must not erase credentials the user has just re-entered.
|
|
guard failedGeneration == generation else { return }
|
|
let target = error.routeHostID ?? hostID
|
|
if var retained = submissions[target] {
|
|
retained.clearSecrets()
|
|
submissions[target] = retained
|
|
}
|
|
generation &+= 1
|
|
cached?.clearSecrets()
|
|
cached = nil
|
|
cancelPendingConfigurationAttempt()
|
|
forcedHostID = target
|
|
forcedProxy = error.isProxyAuthentication
|
|
failure = error.localizedDescription
|
|
}
|
|
|
|
public func didAuthenticate(generation authenticatedGeneration: UInt64) {
|
|
guard authenticatedGeneration == generation else { return }
|
|
submissions.removeAll()
|
|
proxySubmissions.removeAll()
|
|
cached?.clearSecrets()
|
|
cached = nil
|
|
cancelPendingConfigurationAttempt()
|
|
}
|
|
|
|
private func waitForConfigurationAttempt(
|
|
id: UUID, generation: UInt64
|
|
) async throws -> AgentHostResolver.InteractiveResolution {
|
|
let waiterID = UUID()
|
|
return try await withTaskCancellationHandler {
|
|
try await withCheckedThrowingContinuation { continuation in
|
|
guard !Task.isCancelled else {
|
|
continuation.resume(throwing: CancellationError())
|
|
return
|
|
}
|
|
guard var active = pending,
|
|
active.id == id,
|
|
active.generation == generation
|
|
else {
|
|
continuation.resume(throwing: CancellationError())
|
|
return
|
|
}
|
|
active.waiters[waiterID] = continuation
|
|
pending = active
|
|
}
|
|
} onCancel: {
|
|
Task { [weak self] in
|
|
await self?.cancelConfigurationWaiter(
|
|
id: id, generation: generation, waiterID: waiterID
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
private func finishConfigurationAttempt(
|
|
id: UUID, generation completedGeneration: UInt64,
|
|
result: Result<AgentHostResolver.InteractiveResolution, any Error>
|
|
) {
|
|
guard let active = pending,
|
|
active.id == id,
|
|
active.generation == completedGeneration
|
|
else { return }
|
|
pending = nil
|
|
if case .success(let resolution) = result, generation == completedGeneration {
|
|
submissions = resolution.submissions
|
|
proxySubmissions = resolution.proxySubmissions
|
|
temporaryHostKeys = resolution.temporaryHostKeys
|
|
cached = resolution.configuration
|
|
forcedHostID = nil
|
|
forcedProxy = false
|
|
failure = nil
|
|
}
|
|
for continuation in active.waiters.values {
|
|
continuation.resume(with: result)
|
|
}
|
|
}
|
|
|
|
private func cancelConfigurationWaiter(
|
|
id: UUID, generation: UInt64, waiterID: UUID
|
|
) {
|
|
guard var active = pending,
|
|
active.id == id,
|
|
active.generation == generation,
|
|
let continuation = active.waiters.removeValue(forKey: waiterID)
|
|
else { return }
|
|
continuation.resume(throwing: CancellationError())
|
|
if active.waiters.isEmpty {
|
|
pending = nil
|
|
active.task.cancel()
|
|
} else {
|
|
pending = active
|
|
}
|
|
}
|
|
|
|
private func cancelPendingConfigurationAttempt() {
|
|
guard let active = pending else { return }
|
|
pending = nil
|
|
active.task.cancel()
|
|
for continuation in active.waiters.values {
|
|
continuation.resume(throwing: CancellationError())
|
|
}
|
|
}
|
|
|
|
private func makeAttempt(
|
|
configuration: SSHConnectionConfiguration, generation: UInt64
|
|
) -> Attempt {
|
|
Attempt(
|
|
request: SSHConnectionRequest(
|
|
configuration: configuration,
|
|
authorizeHostKey: { [weak self] challenge in
|
|
guard let self else { throw CancellationError() }
|
|
return await self.authorizeHostKey(
|
|
challenge, generation: generation
|
|
)
|
|
},
|
|
keyboardInteractive: { [weak self] challenge in
|
|
guard let self, await self.currentGeneration == generation else { throw CancellationError() }
|
|
let answers = try await self.interaction.respondToChallenge(
|
|
challenge, invocationID: self.invocationID)
|
|
guard await self.currentGeneration == generation else { throw CancellationError() }
|
|
return answers
|
|
},
|
|
securityKeySigning: { [weak self] request in
|
|
guard let self, await self.currentGeneration == generation else { throw CancellationError() }
|
|
let assertion = try await self.interaction.signSecurityKey(request, invocationID: self.invocationID)
|
|
guard await self.currentGeneration == generation else { throw CancellationError() }
|
|
return assertion
|
|
},
|
|
agentForwardingContextID: interaction.presentationContextID
|
|
),
|
|
generation: generation
|
|
)
|
|
}
|
|
|
|
private func authorizeHostKey(
|
|
_ challenge: SSHHostKeyChallenge, generation attemptedGeneration: UInt64
|
|
) async -> SSHHostKeyDecision {
|
|
guard attemptedGeneration == generation else { return .reject }
|
|
if temporaryHostKeys[challenge.hostID] == challenge.observation.identity {
|
|
return .accept
|
|
}
|
|
let pending: PendingHostKey
|
|
if let existing = pendingHostKeys[challenge.hostID],
|
|
existing.fingerprint == challenge.observation.fingerprint
|
|
{
|
|
pending = existing
|
|
} else {
|
|
let resolver = resolver
|
|
let invocationID = invocationID
|
|
let interaction = interaction
|
|
let task = Task {
|
|
await resolver.confirmHostKey(
|
|
challenge, invocationID: invocationID,
|
|
interaction: interaction
|
|
)
|
|
}
|
|
pending = PendingHostKey(
|
|
fingerprint: challenge.observation.fingerprint,
|
|
task: task, waiters: [:]
|
|
)
|
|
pendingHostKeys[challenge.hostID] = pending
|
|
Task { [weak self] in
|
|
let response = await task.value
|
|
await self?.finishHostKeyPrompt(
|
|
hostID: challenge.hostID,
|
|
fingerprint: challenge.observation.fingerprint,
|
|
response: response
|
|
)
|
|
}
|
|
}
|
|
let waiterID = UUID()
|
|
let response: AgentSSHHostKeyResponse? = await withTaskCancellationHandler {
|
|
await withCheckedContinuation {
|
|
(continuation: CheckedContinuation<AgentSSHHostKeyResponse?, Never>) in
|
|
guard var current = pendingHostKeys[challenge.hostID],
|
|
current.fingerprint == pending.fingerprint
|
|
else {
|
|
continuation.resume(returning: nil)
|
|
return
|
|
}
|
|
current.waiters[waiterID] = continuation
|
|
pendingHostKeys[challenge.hostID] = current
|
|
}
|
|
} onCancel: {
|
|
Task { [weak self] in
|
|
await self?.cancelHostKeyWaiter(
|
|
hostID: challenge.hostID,
|
|
fingerprint: challenge.observation.fingerprint,
|
|
waiterID: waiterID
|
|
)
|
|
}
|
|
}
|
|
guard attemptedGeneration == generation else { return .reject }
|
|
switch response {
|
|
case nil:
|
|
return .reject
|
|
case .deny:
|
|
return .reject
|
|
case .continueOnce, .trustAndContinue:
|
|
temporaryHostKeys[challenge.hostID] = challenge.observation.identity
|
|
return .accept
|
|
}
|
|
}
|
|
|
|
private func finishHostKeyPrompt(
|
|
hostID: Host.ID, fingerprint: String,
|
|
response: AgentSSHHostKeyResponse
|
|
) {
|
|
guard let pending = pendingHostKeys[hostID],
|
|
pending.fingerprint == fingerprint
|
|
else { return }
|
|
pendingHostKeys[hostID] = nil
|
|
for continuation in pending.waiters.values {
|
|
continuation.resume(returning: response)
|
|
}
|
|
}
|
|
|
|
private func cancelHostKeyWaiter(
|
|
hostID: Host.ID, fingerprint: String, waiterID: UUID
|
|
) {
|
|
guard var pending = pendingHostKeys[hostID],
|
|
pending.fingerprint == fingerprint,
|
|
let continuation = pending.waiters.removeValue(forKey: waiterID)
|
|
else { return }
|
|
continuation.resume(returning: nil)
|
|
if pending.waiters.isEmpty {
|
|
pending.task.cancel()
|
|
pendingHostKeys[hostID] = nil
|
|
} else {
|
|
pendingHostKeys[hostID] = pending
|
|
}
|
|
}
|
|
}
|
|
|
|
nonisolated public struct AgentSSHConnectionContext: Sendable {
|
|
public var authentication: AgentSSHAuthenticationContext
|
|
public var connection: any SSHConnectionProtocol
|
|
|
|
public init(authentication: AgentSSHAuthenticationContext, connection: any SSHConnectionProtocol) {
|
|
self.authentication = authentication
|
|
self.connection = connection
|
|
}
|
|
}
|
|
|
|
/// Shares one temporary authentication context among connection attempts that
|
|
/// overlap for the same destination. The entry disappears with the last lease,
|
|
/// so Custom passwords and passphrases never become a time-based credential
|
|
/// cache between otherwise unrelated tool calls.
|
|
public actor AgentSSHAuthenticationRegistry {
|
|
private struct Entry {
|
|
var token: UUID
|
|
var context: AgentSSHAuthenticationContext
|
|
var leases: Int
|
|
}
|
|
|
|
private let resolver: AgentHostResolver
|
|
private var entries: [Host.ID: Entry] = [:]
|
|
|
|
public init(resolver: AgentHostResolver) { self.resolver = resolver }
|
|
|
|
public func connection(
|
|
hostID: Host.ID, invocationID: AgentToolInvocation.ID,
|
|
interaction: AgentSSHInteraction, pool: AgentConnectionPool
|
|
) async throws -> any SSHConnectionProtocol {
|
|
try await withAuthenticationContext(
|
|
hostID: hostID, invocationID: invocationID, interaction: interaction
|
|
) { authentication in
|
|
try await self.openConnection(
|
|
hostID: hostID, authentication: authentication, pool: pool
|
|
)
|
|
}
|
|
}
|
|
|
|
public func connectionContext(
|
|
hostID: Host.ID, invocationID: AgentToolInvocation.ID,
|
|
interaction: AgentSSHInteraction, pool: AgentConnectionPool
|
|
) async throws -> AgentSSHConnectionContext {
|
|
try await withAuthenticationContext(
|
|
hostID: hostID, invocationID: invocationID, interaction: interaction
|
|
) { authentication in
|
|
let connection = try await self.openConnection(
|
|
hostID: hostID, authentication: authentication, pool: pool
|
|
)
|
|
return AgentSSHConnectionContext(
|
|
authentication: authentication, connection: connection
|
|
)
|
|
}
|
|
}
|
|
|
|
public func openConnection(
|
|
hostID: Host.ID, authentication: AgentSSHAuthenticationContext,
|
|
pool: AgentConnectionPool
|
|
) async throws -> any SSHConnectionProtocol {
|
|
try await openAuthenticated(
|
|
hostID: hostID, authentication: authentication, pool: pool
|
|
) { attempt in
|
|
if let attempt {
|
|
return try await pool.connection(hostID: hostID) { _ in attempt.request }
|
|
}
|
|
return try await pool.connection(hostID: hostID)
|
|
}
|
|
}
|
|
|
|
public func sftpChannel(
|
|
hostID: Host.ID, invocationID: AgentToolInvocation.ID,
|
|
interaction: AgentSSHInteraction, pool: AgentConnectionPool
|
|
) async throws -> any SFTPChannelProtocol {
|
|
try await withAuthenticationContext(
|
|
hostID: hostID, invocationID: invocationID, interaction: interaction
|
|
) { authentication in
|
|
try await self.openAuthenticated(
|
|
hostID: hostID, authentication: authentication, pool: pool
|
|
) { attempt in
|
|
if let attempt {
|
|
return try await pool.sftpChannel(hostID: hostID) { _ in attempt.request }
|
|
}
|
|
return try await pool.sftpChannel(hostID: hostID)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Runs one operation with an exclusive Agent SFTP lease. Authentication
|
|
/// may retry while acquiring the lease; the operation itself is never
|
|
/// replayed.
|
|
public func withSFTP<Value: Sendable>(
|
|
hostID: Host.ID, invocationID: AgentToolInvocation.ID,
|
|
interaction: AgentSSHInteraction, pool: AgentConnectionPool,
|
|
cancellation: TransferCancellationToken? = nil,
|
|
runID: UUID? = nil,
|
|
byteCount: Int64? = nil,
|
|
schedulesDataTransfer: Bool = false,
|
|
concurrencyPermit: SFTPConcurrencyPermit? = nil,
|
|
operation: @escaping @Sendable (any SFTPChannelProtocol) async throws -> Value
|
|
) async throws -> Value {
|
|
let ownedPermit: SFTPConcurrencyPermit?
|
|
if schedulesDataTransfer, concurrencyPermit == nil {
|
|
ownedPermit = try await pool.sftpConcurrencyScheduler.acquire(
|
|
anchorHostID: hostID,
|
|
hostIDs: [hostID],
|
|
source: .agentRun(runID ?? invocationID),
|
|
byteCount: byteCount
|
|
)
|
|
} else {
|
|
ownedPermit = nil
|
|
}
|
|
let permit = concurrencyPermit ?? ownedPermit
|
|
do {
|
|
let value = try await withAuthenticationContext(
|
|
hostID: hostID, invocationID: invocationID, interaction: interaction
|
|
) { authentication in
|
|
let lease = try await self.openAuthenticated(
|
|
hostID: hostID, authentication: authentication, pool: pool,
|
|
evictsPooledConnection: false, cancellation: cancellation
|
|
) { attempt in
|
|
let resolver: AgentConnectionPool.RequestResolving?
|
|
if let attempt {
|
|
let request = attempt.request
|
|
resolver = { _ in request }
|
|
} else {
|
|
resolver = nil
|
|
}
|
|
return try await pool.acquireSFTP(
|
|
hostID: hostID, cancellation: cancellation,
|
|
configuration: resolver,
|
|
mode: permit?.schedulingMode ?? .channelsOnly)
|
|
}
|
|
do {
|
|
let value = try await operation(lease.channel)
|
|
await pool.releaseSFTP(lease.id)
|
|
return value
|
|
} catch {
|
|
await pool.releaseSFTP(lease.id)
|
|
if let failure = error as? SFTPError {
|
|
switch failure {
|
|
case .notConnected, .connectionLost:
|
|
await pool.disconnect(hostID: hostID)
|
|
case .cancelled, .localFile, .operation:
|
|
break
|
|
}
|
|
}
|
|
throw error
|
|
}
|
|
}
|
|
await ownedPermit?.release()
|
|
return value
|
|
} catch {
|
|
await ownedPermit?.release()
|
|
throw error
|
|
}
|
|
}
|
|
|
|
private func openAuthenticated<Client: Sendable>(
|
|
hostID: Host.ID, authentication: AgentSSHAuthenticationContext,
|
|
pool: AgentConnectionPool,
|
|
evictsPooledConnection: Bool = true,
|
|
cancellation: TransferCancellationToken? = nil,
|
|
open: @Sendable (AgentSSHAuthenticationContext.Attempt?) async throws -> Client
|
|
) async throws -> Client {
|
|
var attempt: AgentSSHAuthenticationContext.Attempt?
|
|
if await pool.isConnected(hostID: hostID), !(await resolver.agentForwardingEnabled(for: hostID)) {
|
|
attempt = nil
|
|
} else if await authentication.canResolveTarget() {
|
|
attempt = try await authentication.configurationAttempt()
|
|
}
|
|
while true {
|
|
try Task.checkCancellation()
|
|
if cancellation?.isCancelled == true { throw SFTPError.cancelled }
|
|
let attemptedGeneration: UInt64
|
|
if let attempt {
|
|
attemptedGeneration = attempt.generation
|
|
} else {
|
|
attemptedGeneration = await authentication.currentGeneration
|
|
}
|
|
do {
|
|
let client = try await open(attempt)
|
|
await authentication.didAuthenticate(generation: attemptedGeneration)
|
|
return client
|
|
} catch let error as TermindSSHError where error.shouldRequestAuthentication {
|
|
if cancellation?.isCancelled == true { throw SFTPError.cancelled }
|
|
if evictsPooledConnection { await pool.disconnect(hostID: hostID) }
|
|
await authentication.retry(after: error, generation: attemptedGeneration)
|
|
attempt = try await authentication.configurationAttempt()
|
|
} catch let error as TermindSSHError where error.isLocalAuthenticationRequirement {
|
|
attempt = try await authentication.configurationAttempt()
|
|
}
|
|
}
|
|
}
|
|
|
|
fileprivate func withAuthenticationContext<Value: Sendable>(
|
|
hostID: Host.ID, invocationID: AgentToolInvocation.ID,
|
|
interaction: AgentSSHInteraction,
|
|
operation: @Sendable (AgentSSHAuthenticationContext) async throws -> Value
|
|
) async throws -> Value {
|
|
if await resolver.agentForwardingEnabled(for: hostID) {
|
|
return try await operation(
|
|
AgentSSHAuthenticationContext(
|
|
resolver: resolver, hostID: hostID, invocationID: invocationID, interaction: interaction))
|
|
}
|
|
let token: UUID
|
|
let context: AgentSSHAuthenticationContext
|
|
if var entry = entries[hostID] {
|
|
entry.leases += 1
|
|
entries[hostID] = entry
|
|
token = entry.token
|
|
context = entry.context
|
|
} else {
|
|
token = UUID()
|
|
context = AgentSSHAuthenticationContext(
|
|
resolver: resolver, hostID: hostID, invocationID: invocationID,
|
|
interaction: interaction
|
|
)
|
|
entries[hostID] = Entry(token: token, context: context, leases: 1)
|
|
}
|
|
defer { release(hostID: hostID, token: token) }
|
|
return try await operation(context)
|
|
}
|
|
|
|
private func release(hostID: Host.ID, token: UUID) {
|
|
guard var entry = entries[hostID], entry.token == token else { return }
|
|
entry.leases -= 1
|
|
if entry.leases == 0 {
|
|
entries[hostID] = nil
|
|
} else {
|
|
entries[hostID] = entry
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Keeps one channel-multiplexed SSH connection per Host alive across tool
|
|
/// calls, runs, and conversations. Exec and the primary SFTP subsystem share
|
|
/// that connection; SFTP topology may add silent auxiliary connections.
|
|
///
|
|
/// Before this existed, `ssh_exec` paid a full TCP connect, handshake, host-key
|
|
/// check, and signature for *every command* — and a jump-host route rebuilt
|
|
/// every hop each time. Keying by host rather than by run is what lets a
|
|
/// follow-up message reuse the connection; the idle sweep, not the end of a run,
|
|
/// is what gives it back.
|
|
///
|
|
public actor AgentConnectionPool {
|
|
public nonisolated let sftpConcurrencyScheduler: SFTPConcurrencyScheduler
|
|
|
|
public typealias ConfigurationResolving =
|
|
@Sendable (UUID) async throws -> SSHConnectionConfiguration
|
|
public typealias RequestResolving =
|
|
@Sendable (UUID) async throws -> SSHConnectionRequest
|
|
public typealias FingerprintResolving = @Sendable (UUID) async throws -> String
|
|
public typealias ConnectionOpening =
|
|
@Sendable (SSHConnectionRequest) async throws -> any SSHConnectionProtocol
|
|
public typealias SFTPOpening =
|
|
@Sendable (SSHConnectionRequest) async throws -> any SFTPClientProtocol
|
|
|
|
/// Long enough that a user reading output and asking a follow-up keeps the
|
|
/// connection, short enough that walking away does not leave an
|
|
/// authenticated connection open on a server indefinitely.
|
|
public static let defaultIdleTimeout = Duration.seconds(180)
|
|
|
|
private struct Entry {
|
|
var connection: any SSHConnectionProtocol
|
|
var fingerprint: String
|
|
var lastUsedAt: ContinuousClock.Instant
|
|
}
|
|
|
|
private struct ConnectionOpenKey: Hashable, Sendable {
|
|
var hostID: UUID
|
|
var fingerprint: String
|
|
}
|
|
|
|
private struct OpenedConnection: Sendable {
|
|
var value: any SSHConnectionProtocol
|
|
}
|
|
|
|
private struct InFlight {
|
|
var id: UUID
|
|
var task: Task<any SSHConnectionProtocol, any Error>
|
|
var waiters: [UUID: CheckedContinuation<OpenedConnection, any Error>]
|
|
}
|
|
|
|
private let configuration: ConfigurationResolving
|
|
private final class ForwardedConnection: @unchecked Sendable {
|
|
weak var value: (any SSHConnectionProtocol)?
|
|
init(_ value: any SSHConnectionProtocol) { self.value = value }
|
|
}
|
|
private var forwardedConnections: [UUID: [ForwardedConnection]] = [:]
|
|
private var forwardedOpenings: [UUID: [UUID: Task<any SSHConnectionProtocol, any Error>]] = [:]
|
|
private let forwardingEnabled: @Sendable (UUID) async -> Bool
|
|
private let routeFingerprint: FingerprintResolving
|
|
private let openConnection: ConnectionOpening
|
|
private let openSFTP: SFTPOpening
|
|
private let idleTimeout: Duration
|
|
private let sweeper: AgentIdleSweeper
|
|
private var entries: [UUID: Entry] = [:]
|
|
private var pending: [ConnectionOpenKey: InFlight] = [:]
|
|
private var activeExecLeases: [UUID: Int] = [:]
|
|
|
|
struct SFTPLease: Sendable {
|
|
fileprivate var id: UUID
|
|
var channel: any SFTPChannelProtocol
|
|
}
|
|
|
|
private enum SFTPCapability { case unknown, connections, channels }
|
|
|
|
private struct SFTPRecord {
|
|
var id: UUID
|
|
var channel: any SFTPChannelProtocol
|
|
var busy: Bool
|
|
var lastUsedAt: ContinuousClock.Instant
|
|
}
|
|
|
|
private struct SFTPState {
|
|
var fingerprint: String
|
|
var primaryID: UUID
|
|
var records: [UUID: SFTPRecord]
|
|
var capability = SFTPCapability.unknown
|
|
}
|
|
|
|
private struct ActiveSFTPLease {
|
|
var hostID: UUID
|
|
var recordID: UUID?
|
|
var channel: any SFTPChannelProtocol
|
|
var cancellation: TransferCancellationToken?
|
|
}
|
|
|
|
private struct SFTPProbe {
|
|
var id: UUID
|
|
var task: Task<any SFTPChannelProtocol, any Error>
|
|
}
|
|
|
|
private struct SFTPPrimaryOpen {
|
|
var id: UUID
|
|
var fingerprint: String
|
|
var connection: any SSHConnectionProtocol
|
|
var task: Task<any SFTPChannelProtocol, any Error>
|
|
}
|
|
|
|
private var sftpStates: [UUID: SFTPState] = [:]
|
|
private var activeSFTPLeases: [UUID: ActiveSFTPLease] = [:]
|
|
private var sftpProbes: [UUID: SFTPProbe] = [:]
|
|
private var sftpPrimaryOpens: [UUID: SFTPPrimaryOpen] = [:]
|
|
/// Acquires one exclusive Agent SFTP channel. Overlap prefers another SSH
|
|
/// connection; after silent connection setup fails, it opens a sibling
|
|
/// SFTP channel on the primary connection instead.
|
|
func acquireSFTP(
|
|
hostID: UUID, cancellation: TransferCancellationToken? = nil,
|
|
configuration override: RequestResolving? = nil,
|
|
mode: SSHChannelSchedulingMode = .connectionsPreferred
|
|
) async throws -> SFTPLease {
|
|
try Task.checkCancellation()
|
|
if cancellation?.isCancelled == true { throw SFTPError.cancelled }
|
|
|
|
let fingerprint = try await routeFingerprint(hostID)
|
|
if let state = sftpStates[hostID], state.fingerprint != fingerprint {
|
|
_ = await disconnectSFTPState(hostID: hostID)
|
|
}
|
|
|
|
let primary = try await sftpChannel(hostID: hostID, configuration: override)
|
|
let committedFingerprint = try await routeFingerprint(hostID)
|
|
if sftpStates[hostID] == nil {
|
|
let record = SFTPRecord(id: UUID(), channel: primary, busy: false, lastUsedAt: .now)
|
|
sftpStates[hostID] = SFTPState(
|
|
fingerprint: committedFingerprint, primaryID: record.id,
|
|
records: [record.id: record])
|
|
}
|
|
|
|
if let idle = sftpStates[hostID]?.records.values
|
|
.filter({ !$0.busy && (mode == .connectionsPreferred || $0.id == sftpStates[hostID]?.primaryID) })
|
|
.sorted(by: { $0.lastUsedAt < $1.lastUsedAt }).first
|
|
{
|
|
sftpStates[hostID]?.records[idle.id]?.busy = true
|
|
return registerSFTPLease(
|
|
hostID: hostID, recordID: idle.id, channel: idle.channel,
|
|
cancellation: cancellation)
|
|
}
|
|
|
|
guard let state = sftpStates[hostID], let primaryRecord = state.records[state.primaryID] else {
|
|
throw SFTPError.notConnected
|
|
}
|
|
|
|
if mode == .channelsOnly {
|
|
do {
|
|
let sibling = try await primaryRecord.channel.openSiblingChannel()
|
|
if cancellation?.isCancelled == true {
|
|
await sibling.disconnect()
|
|
throw SFTPError.cancelled
|
|
}
|
|
return registerSFTPLease(
|
|
hostID: hostID, recordID: nil, channel: sibling, cancellation: cancellation)
|
|
} catch {
|
|
try Task.checkCancellation()
|
|
if cancellation?.isCancelled == true { throw SFTPError.cancelled }
|
|
if error is CancellationError { throw error }
|
|
if let failure = error as? SFTPError, case .cancelled = failure { throw failure }
|
|
while sftpStates[hostID]?.records[primaryRecord.id]?.busy == true {
|
|
try await Task.sleep(for: .milliseconds(10))
|
|
try Task.checkCancellation()
|
|
if cancellation?.isCancelled == true { throw SFTPError.cancelled }
|
|
}
|
|
return try await acquireSFTP(
|
|
hostID: hostID, cancellation: cancellation,
|
|
configuration: override, mode: mode)
|
|
}
|
|
}
|
|
|
|
if state.capability == .unknown {
|
|
let probe: SFTPProbe
|
|
if let existing = sftpProbes[hostID] {
|
|
probe = existing
|
|
} else {
|
|
let id = UUID()
|
|
let task = Task { [configuration, openSFTP] () async throws -> any SFTPChannelProtocol in
|
|
let durable = try await configuration(hostID)
|
|
let request = SSHConnectionRequest(
|
|
configuration: durable, interactionPolicy: .forbidden
|
|
).enforcingInteractionPolicy()
|
|
return try await openSFTP(request)
|
|
}
|
|
probe = SFTPProbe(id: id, task: task)
|
|
sftpProbes[hostID] = probe
|
|
}
|
|
do {
|
|
let channel = try await probe.task.value
|
|
if sftpProbes[hostID]?.id == probe.id {
|
|
sftpProbes[hostID] = nil
|
|
let record = SFTPRecord(id: UUID(), channel: channel, busy: false, lastUsedAt: .now)
|
|
sftpStates[hostID]?.records[record.id] = record
|
|
sftpStates[hostID]?.capability = .connections
|
|
}
|
|
return try await acquireSFTP(
|
|
hostID: hostID, cancellation: cancellation,
|
|
configuration: override, mode: mode)
|
|
} catch {
|
|
if sftpProbes[hostID]?.id == probe.id {
|
|
sftpProbes[hostID] = nil
|
|
sftpStates[hostID]?.capability = .channels
|
|
}
|
|
if error is CancellationError { throw error }
|
|
if let failure = error as? SFTPError, case .cancelled = failure { throw failure }
|
|
}
|
|
} else if state.capability == .connections {
|
|
do {
|
|
let durable = try await configuration(hostID)
|
|
let request = SSHConnectionRequest(
|
|
configuration: durable, interactionPolicy: .forbidden
|
|
).enforcingInteractionPolicy()
|
|
let channel = try await openSFTP(request)
|
|
try Task.checkCancellation()
|
|
if cancellation?.isCancelled == true {
|
|
await channel.disconnect()
|
|
throw SFTPError.cancelled
|
|
}
|
|
let record = SFTPRecord(id: UUID(), channel: channel, busy: true, lastUsedAt: .now)
|
|
sftpStates[hostID]?.records[record.id] = record
|
|
sftpStates[hostID]?.capability = .connections
|
|
startSweeping()
|
|
return registerSFTPLease(
|
|
hostID: hostID, recordID: record.id, channel: channel,
|
|
cancellation: cancellation)
|
|
} catch {
|
|
if error is CancellationError { throw error }
|
|
if let failure = error as? SFTPError, case .cancelled = failure { throw failure }
|
|
sftpStates[hostID]?.capability = .channels
|
|
}
|
|
}
|
|
|
|
let sibling = try await primaryRecord.channel.openSiblingChannel()
|
|
if cancellation?.isCancelled == true {
|
|
await sibling.disconnect()
|
|
throw SFTPError.cancelled
|
|
}
|
|
return registerSFTPLease(
|
|
hostID: hostID, recordID: nil, channel: sibling, cancellation: cancellation)
|
|
}
|
|
|
|
func releaseSFTP(_ id: UUID) async {
|
|
guard let lease = activeSFTPLeases.removeValue(forKey: id) else { return }
|
|
if let recordID = lease.recordID, sftpStates[lease.hostID]?.records[recordID] != nil {
|
|
sftpStates[lease.hostID]?.records[recordID]?.busy = false
|
|
sftpStates[lease.hostID]?.records[recordID]?.lastUsedAt = .now
|
|
} else {
|
|
await lease.channel.disconnect()
|
|
}
|
|
}
|
|
|
|
private func registerSFTPLease(
|
|
hostID: UUID, recordID: UUID?, channel: any SFTPChannelProtocol,
|
|
cancellation: TransferCancellationToken?
|
|
) -> SFTPLease {
|
|
let id = UUID()
|
|
activeSFTPLeases[id] = ActiveSFTPLease(
|
|
hostID: hostID, recordID: recordID, channel: channel, cancellation: cancellation)
|
|
startSweeping()
|
|
return SFTPLease(id: id, channel: channel)
|
|
}
|
|
|
|
public init(
|
|
configuration: @escaping ConfigurationResolving,
|
|
routeFingerprint: @escaping FingerprintResolving,
|
|
forwardingEnabled: @escaping @Sendable (UUID) async -> Bool = { _ in false },
|
|
openConnection: @escaping ConnectionOpening,
|
|
openSFTP: @escaping SFTPOpening,
|
|
sftpConcurrencyScheduler: SFTPConcurrencyScheduler = SFTPConcurrencyScheduler(),
|
|
idleTimeout: Duration = AgentConnectionPool.defaultIdleTimeout,
|
|
sweepInterval: Duration? = .seconds(30)
|
|
) {
|
|
self.configuration = configuration
|
|
self.routeFingerprint = routeFingerprint
|
|
self.forwardingEnabled = forwardingEnabled
|
|
self.openConnection = openConnection
|
|
self.openSFTP = openSFTP
|
|
self.sftpConcurrencyScheduler = sftpConcurrencyScheduler
|
|
self.idleTimeout = idleTimeout
|
|
sweeper = AgentIdleSweeper(interval: sweepInterval)
|
|
}
|
|
|
|
/// Runs one command on the host's pooled connection, opening one if needed.
|
|
///
|
|
/// Retries exactly once only when the pooled connection was already gone
|
|
/// before a channel could open. A server-side channel rejection stays on
|
|
/// the one-connection policy and is surfaced; anything that failed after the
|
|
/// command started is also surfaced because replay could duplicate effects.
|
|
public func execute(
|
|
_ command: String, workingDirectory: String?, stdin: SSHCommandInput?,
|
|
timeout: TimeInterval, outputLimit: Int, hostID: UUID,
|
|
configuration override: RequestResolving? = nil,
|
|
onOutput: SSHExecOutputHandler? = nil
|
|
) async throws -> SSHExecResult {
|
|
do {
|
|
let connection = try await connection(hostID: hostID, configuration: override)
|
|
return try await withExecLease(hostID: hostID, connection: connection) {
|
|
try await connection.execute(
|
|
command, workingDirectory: workingDirectory, stdin: stdin,
|
|
timeout: timeout, outputLimit: outputLimit, onOutput: onOutput)
|
|
}
|
|
} catch SSHCommandExecutionError.connectionUnavailable {
|
|
await disconnect(hostID: hostID)
|
|
let connection = try await connection(hostID: hostID, configuration: override)
|
|
return try await withExecLease(hostID: hostID, connection: connection) {
|
|
try await connection.execute(
|
|
command, workingDirectory: workingDirectory, stdin: stdin,
|
|
timeout: timeout, outputLimit: outputLimit, onOutput: onOutput)
|
|
}
|
|
}
|
|
}
|
|
|
|
func withExecLease<Value: Sendable>(
|
|
hostID: UUID, connection: any SSHConnectionProtocol,
|
|
operation: @escaping @Sendable () async throws -> Value
|
|
) async throws -> Value {
|
|
if forwardedConnections[hostID]?.contains(where: { $0.value === connection }) == true {
|
|
defer {
|
|
forwardedConnections[hostID]?.removeAll { $0.value == nil || $0.value === connection }
|
|
if forwardedConnections[hostID]?.isEmpty == true { forwardedConnections[hostID] = nil }
|
|
}
|
|
do {
|
|
let result = try await operation()
|
|
await connection.disconnect()
|
|
return result
|
|
} catch {
|
|
await connection.disconnect()
|
|
throw error
|
|
}
|
|
}
|
|
guard let pooled = entries[hostID]?.connection, pooled === connection else {
|
|
throw SSHCommandExecutionError.connectionUnavailable(TermindSSHError.disconnected)
|
|
}
|
|
activeExecLeases[hostID, default: 0] += 1
|
|
defer {
|
|
let remaining = (activeExecLeases[hostID] ?? 1) - 1
|
|
activeExecLeases[hostID] = remaining == 0 ? nil : remaining
|
|
if var entry = entries[hostID] {
|
|
entry.lastUsedAt = .now
|
|
entries[hostID] = entry
|
|
}
|
|
}
|
|
return try await operation()
|
|
}
|
|
|
|
public func connection(
|
|
hostID: UUID, configuration override: RequestResolving? = nil
|
|
) async throws -> any SSHConnectionProtocol {
|
|
if await forwardingEnabled(hostID) {
|
|
var request: SSHConnectionRequest
|
|
if let override {
|
|
request = try await override(hostID)
|
|
} else {
|
|
request = SSHConnectionRequest(
|
|
configuration: try await configuration(hostID), interactionPolicy: .forbidden)
|
|
}
|
|
let openingID = UUID()
|
|
let isolated = request
|
|
let task = Task { [openConnection] in try await openConnection(isolated) }
|
|
forwardedOpenings[hostID, default: [:]][openingID] = task
|
|
do {
|
|
let connection = try await withTaskCancellationHandler {
|
|
try await task.value
|
|
} onCancel: {
|
|
task.cancel()
|
|
}
|
|
guard forwardedOpenings[hostID]?[openingID] != nil, !Task.isCancelled else {
|
|
await connection.disconnect()
|
|
throw CancellationError()
|
|
}
|
|
forwardedOpenings[hostID]?[openingID] = nil
|
|
if forwardedOpenings[hostID]?.isEmpty == true { forwardedOpenings[hostID] = nil }
|
|
forwardedConnections[hostID] = (forwardedConnections[hostID] ?? []).filter { $0.value != nil }
|
|
forwardedConnections[hostID, default: []].append(ForwardedConnection(connection))
|
|
return connection
|
|
} catch {
|
|
forwardedOpenings[hostID]?[openingID] = nil
|
|
if forwardedOpenings[hostID]?.isEmpty == true { forwardedOpenings[hostID] = nil }
|
|
throw error
|
|
}
|
|
}
|
|
return try await resolve(hostID: hostID, configuration: override)
|
|
}
|
|
|
|
public func sftpChannel(
|
|
hostID: UUID, configuration override: RequestResolving? = nil
|
|
) async throws -> any SFTPChannelProtocol {
|
|
let fingerprint = try await routeFingerprint(hostID)
|
|
if let state = sftpStates[hostID], state.fingerprint == fingerprint,
|
|
let primary = state.records[state.primaryID]
|
|
{
|
|
sftpStates[hostID]?.records[state.primaryID]?.lastUsedAt = .now
|
|
entries[hostID]?.lastUsedAt = .now
|
|
return primary.channel
|
|
}
|
|
if sftpStates[hostID] != nil { await disconnect(hostID: hostID) }
|
|
let connection = try await resolve(hostID: hostID, configuration: override)
|
|
let opening: SFTPPrimaryOpen
|
|
let ownsOpening: Bool
|
|
if let existing = sftpPrimaryOpens[hostID], existing.fingerprint == fingerprint,
|
|
existing.connection === connection
|
|
{
|
|
opening = existing
|
|
ownsOpening = false
|
|
} else {
|
|
let id = UUID()
|
|
let task = Task { try await connection.openSFTPChannel() }
|
|
opening = SFTPPrimaryOpen(
|
|
id: id, fingerprint: fingerprint, connection: connection, task: task)
|
|
sftpPrimaryOpens[hostID] = opening
|
|
ownsOpening = true
|
|
}
|
|
do {
|
|
let client = try await opening.task.value
|
|
if !ownsOpening {
|
|
while sftpPrimaryOpens[hostID]?.id == opening.id,
|
|
sftpStates[hostID] == nil
|
|
{
|
|
try Task.checkCancellation()
|
|
await Task.yield()
|
|
}
|
|
guard let state = sftpStates[hostID], state.fingerprint == fingerprint,
|
|
let primary = state.records[state.primaryID]
|
|
else { throw SFTPError.notConnected }
|
|
return primary.channel
|
|
}
|
|
let committedFingerprint = try await routeFingerprint(hostID)
|
|
guard committedFingerprint == fingerprint, entries[hostID]?.connection === connection else {
|
|
await client.disconnect()
|
|
await disconnect(hostID: hostID)
|
|
return try await sftpChannel(hostID: hostID, configuration: override)
|
|
}
|
|
let record = SFTPRecord(id: UUID(), channel: client, busy: false, lastUsedAt: .now)
|
|
sftpStates[hostID] = SFTPState(
|
|
fingerprint: fingerprint, primaryID: record.id,
|
|
records: [record.id: record])
|
|
sftpPrimaryOpens[hostID] = nil
|
|
if var entry = entries[hostID] {
|
|
entry.lastUsedAt = .now
|
|
entries[hostID] = entry
|
|
}
|
|
startSweeping()
|
|
return client
|
|
} catch {
|
|
if ownsOpening, sftpPrimaryOpens[hostID]?.id == opening.id {
|
|
sftpPrimaryOpens[hostID] = nil
|
|
}
|
|
throw error
|
|
}
|
|
}
|
|
|
|
private func disconnectSFTPState(hostID: UUID) async -> Bool {
|
|
sftpProbes.removeValue(forKey: hostID)?.task.cancel()
|
|
sftpPrimaryOpens.removeValue(forKey: hostID)?.task.cancel()
|
|
let leaseIDs = activeSFTPLeases.filter { $0.value.hostID == hostID }.map(\.key)
|
|
for id in leaseIDs { activeSFTPLeases[id]?.cancellation?.cancel() }
|
|
let leases = leaseIDs.compactMap { activeSFTPLeases.removeValue(forKey: $0) }
|
|
let state = sftpStates.removeValue(forKey: hostID)
|
|
var seen: Set<ObjectIdentifier> = []
|
|
for record in state?.records.values ?? [UUID: SFTPRecord]().values {
|
|
let identity = ObjectIdentifier(record.channel as AnyObject)
|
|
if seen.insert(identity).inserted { await record.channel.disconnect() }
|
|
}
|
|
for lease in leases where lease.recordID == nil { await lease.channel.disconnect() }
|
|
return state != nil || !leases.isEmpty
|
|
}
|
|
|
|
@discardableResult
|
|
public func disconnect(hostID: UUID) async -> Bool {
|
|
cancelConnectionOpenings(hostID: hostID)
|
|
let opening = forwardedOpenings.removeValue(forKey: hostID) ?? [:]
|
|
for task in opening.values { task.cancel() }
|
|
let forwarded = forwardedConnections.removeValue(forKey: hostID) ?? []
|
|
for connection in forwarded { await connection.value?.disconnect() }
|
|
let stateClosed = await disconnectSFTPState(hostID: hostID)
|
|
activeExecLeases[hostID] = nil
|
|
let entry = entries.removeValue(forKey: hostID)
|
|
stopSweepingIfEmpty()
|
|
await entry?.connection.disconnect()
|
|
return stateClosed || entry != nil || !forwarded.isEmpty
|
|
}
|
|
|
|
public func closeAll() async {
|
|
let hostIDs = Set(entries.keys)
|
|
.union(forwardedConnections.keys)
|
|
.union(forwardedOpenings.keys)
|
|
.union(sftpStates.keys)
|
|
.union(pending.keys.map(\.hostID))
|
|
.union(sftpPrimaryOpens.keys)
|
|
for hostID in hostIDs { _ = await disconnect(hostID: hostID) }
|
|
}
|
|
|
|
/// Closes everything untouched for longer than the idle timeout and reports
|
|
/// how many connections remain. Driven by a background task in the app and
|
|
/// called directly, with a synthetic instant, by tests.
|
|
@discardableResult
|
|
public func sweepIdle(now: ContinuousClock.Instant) async -> Int {
|
|
var expiredSFTP: [(UUID, SFTPRecord)] = []
|
|
for hostID in Array(sftpStates.keys) {
|
|
guard var state = sftpStates[hostID] else { continue }
|
|
for (id, record) in state.records
|
|
where id != state.primaryID && !record.busy && now - record.lastUsedAt >= idleTimeout {
|
|
state.records[id] = nil
|
|
expiredSFTP.append((hostID, record))
|
|
}
|
|
sftpStates[hostID] = state
|
|
}
|
|
let expiredHosts = entries.compactMap { hostID, entry -> UUID? in
|
|
guard activeExecLeases[hostID] == nil,
|
|
!activeSFTPLeases.values.contains(where: { $0.hostID == hostID })
|
|
else { return nil }
|
|
if let state = sftpStates[hostID], let primary = state.records[state.primaryID] {
|
|
guard state.records.count == 1, !primary.busy,
|
|
now - max(entry.lastUsedAt, primary.lastUsedAt) >= idleTimeout
|
|
else { return nil }
|
|
} else if now - entry.lastUsedAt < idleTimeout {
|
|
return nil
|
|
}
|
|
return hostID
|
|
}
|
|
for hostID in expiredHosts { _ = await disconnect(hostID: hostID) }
|
|
stopSweepingIfEmpty()
|
|
for (_, record) in expiredSFTP { await record.channel.disconnect() }
|
|
return openConnectionCount
|
|
}
|
|
|
|
public func isConnected(hostID: UUID) -> Bool { entries[hostID] != nil }
|
|
|
|
public var openConnectionCount: Int {
|
|
entries.count
|
|
+ sftpStates.values.reduce(0) { total, state in
|
|
total + max(0, state.records.count - 1)
|
|
}
|
|
}
|
|
|
|
private func resolve(
|
|
hostID: UUID, configuration override: RequestResolving?
|
|
) async throws -> any SSHConnectionProtocol {
|
|
let fingerprint: String
|
|
do {
|
|
fingerprint = try await routeFingerprint(hostID)
|
|
} catch {
|
|
// The host is gone — deleted from the Vault, or a temporary target
|
|
// whose pane closed. Do not keep a connection to it alive.
|
|
await disconnect(hostID: hostID)
|
|
throw error
|
|
}
|
|
if let entry = entries.removeValue(forKey: hostID) {
|
|
guard entry.fingerprint == fingerprint else {
|
|
// Nothing waits on a stale socket. Its worker also frees the
|
|
// libssh2 connection in `deinit`, so dropping the reference alone
|
|
// would be safe; this just makes it prompt.
|
|
Task { await entry.connection.disconnect() }
|
|
_ = await disconnectSFTPState(hostID: hostID)
|
|
return try await open(
|
|
hostID: hostID, fingerprint: fingerprint, configuration: override
|
|
)
|
|
}
|
|
entries[hostID] = Entry(
|
|
connection: entry.connection, fingerprint: fingerprint, lastUsedAt: .now
|
|
)
|
|
return entry.connection
|
|
}
|
|
return try await open(hostID: hostID, fingerprint: fingerprint, configuration: override)
|
|
}
|
|
|
|
private func open(
|
|
hostID: UUID, fingerprint: String, configuration override: RequestResolving?
|
|
) async throws -> any SSHConnectionProtocol {
|
|
// Without this join, two tool calls that both miss the cache would each
|
|
// build a connection and one of them would never be pooled — or closed.
|
|
let key = ConnectionOpenKey(hostID: hostID, fingerprint: fingerprint)
|
|
if let inFlight = pending[key] {
|
|
return try await waitForConnectionOpening(key: key, id: inFlight.id)
|
|
}
|
|
let id = UUID()
|
|
let task = Task { [configuration, openConnection] () async throws -> any SSHConnectionProtocol in
|
|
var request: SSHConnectionRequest
|
|
if let override {
|
|
request = try await override(hostID)
|
|
} else {
|
|
request = SSHConnectionRequest(
|
|
configuration: try await configuration(hostID)
|
|
)
|
|
}
|
|
request.permitsAgentForwarding = false
|
|
return try await openConnection(request)
|
|
}
|
|
pending[key] = InFlight(id: id, task: task, waiters: [:])
|
|
Task { [weak self] in
|
|
do {
|
|
let connection = try await task.value
|
|
await self?.finishConnectionOpening(
|
|
key: key, id: id, result: .success(connection)
|
|
)
|
|
} catch {
|
|
await self?.finishConnectionOpening(
|
|
key: key, id: id, result: .failure(error)
|
|
)
|
|
}
|
|
}
|
|
return try await waitForConnectionOpening(key: key, id: id)
|
|
}
|
|
|
|
private func waitForConnectionOpening(
|
|
key: ConnectionOpenKey, id: UUID
|
|
) async throws -> any SSHConnectionProtocol {
|
|
let waiterID = UUID()
|
|
let connection: OpenedConnection = try await withTaskCancellationHandler {
|
|
try await withCheckedThrowingContinuation {
|
|
(continuation: CheckedContinuation<OpenedConnection, any Error>) in
|
|
guard !Task.isCancelled else {
|
|
continuation.resume(throwing: CancellationError())
|
|
return
|
|
}
|
|
guard var opening = pending[key], opening.id == id else {
|
|
continuation.resume(throwing: CancellationError())
|
|
return
|
|
}
|
|
opening.waiters[waiterID] = continuation
|
|
pending[key] = opening
|
|
}
|
|
} onCancel: {
|
|
Task { [weak self] in
|
|
await self?.cancelConnectionOpeningWaiter(
|
|
key: key, id: id, waiterID: waiterID
|
|
)
|
|
}
|
|
}
|
|
try Task.checkCancellation()
|
|
return connection.value
|
|
}
|
|
|
|
private func finishConnectionOpening(
|
|
key: ConnectionOpenKey, id: UUID,
|
|
result initialResult: Result<any SSHConnectionProtocol, any Error>
|
|
) async {
|
|
guard let opening = pending[key], opening.id == id else {
|
|
if case .success(let connection) = initialResult {
|
|
await connection.disconnect()
|
|
}
|
|
return
|
|
}
|
|
|
|
let result: Result<OpenedConnection, any Error>
|
|
switch initialResult {
|
|
case .failure(let error):
|
|
result = .failure(error)
|
|
case .success(let connection):
|
|
do {
|
|
// Saving/binding an Identity while the connection prompt is open
|
|
// legitimately changes the route fingerprint. Bind the authenticated
|
|
// connection to that post-save route so the next command can reuse it.
|
|
let committedFingerprint = try await routeFingerprint(key.hostID)
|
|
guard pending[key]?.id == id else {
|
|
await connection.disconnect()
|
|
return
|
|
}
|
|
commit(
|
|
connection, hostID: key.hostID,
|
|
fingerprint: committedFingerprint
|
|
)
|
|
result = .success(OpenedConnection(value: connection))
|
|
} catch {
|
|
await connection.disconnect()
|
|
result = .failure(error)
|
|
}
|
|
}
|
|
|
|
guard let completed = pending.removeValue(forKey: key), completed.id == id else {
|
|
if case .success(let connection) = result {
|
|
await connection.value.disconnect()
|
|
}
|
|
return
|
|
}
|
|
for continuation in completed.waiters.values {
|
|
continuation.resume(with: result)
|
|
}
|
|
}
|
|
|
|
private func cancelConnectionOpeningWaiter(
|
|
key: ConnectionOpenKey, id: UUID, waiterID: UUID
|
|
) {
|
|
guard var opening = pending[key], opening.id == id,
|
|
let continuation = opening.waiters.removeValue(forKey: waiterID)
|
|
else { return }
|
|
continuation.resume(throwing: CancellationError())
|
|
if opening.waiters.isEmpty {
|
|
pending[key] = nil
|
|
opening.task.cancel()
|
|
} else {
|
|
pending[key] = opening
|
|
}
|
|
}
|
|
|
|
private func cancelConnectionOpenings(hostID: UUID) {
|
|
let keys = pending.keys.filter { $0.hostID == hostID }
|
|
for key in keys {
|
|
guard let opening = pending.removeValue(forKey: key) else { continue }
|
|
opening.task.cancel()
|
|
for continuation in opening.waiters.values {
|
|
continuation.resume(throwing: CancellationError())
|
|
}
|
|
}
|
|
}
|
|
|
|
private func commit(
|
|
_ connection: any SSHConnectionProtocol, hostID: UUID, fingerprint: String
|
|
) {
|
|
entries[hostID] = Entry(connection: connection, fingerprint: fingerprint, lastUsedAt: .now)
|
|
startSweeping()
|
|
}
|
|
|
|
private func startSweeping() {
|
|
sweeper.start { [weak self] in
|
|
guard let self else { return 0 }
|
|
return await self.sweepIdle(now: .now)
|
|
}
|
|
}
|
|
|
|
private func stopSweepingIfEmpty() {
|
|
guard entries.isEmpty, pending.isEmpty, sftpStates.isEmpty, activeSFTPLeases.isEmpty,
|
|
sftpProbes.isEmpty, sftpPrimaryOpens.isEmpty
|
|
else { return }
|
|
sweeper.stop()
|
|
}
|
|
}
|