refactor(connections): cap command output, precise timeout, log malformed source (#1254)

Author: datlechin

TablePro/Core/Storage/ConnectionStorage.swift

@@ -814,7 +814,7 @@ private struct StoredConnection: Codable {
         sshTunnelModeJson = try container.decodeIfPresent(Data.self, forKey: .sshTunnelModeJson)
         cloudflareTunnelModeJson = try container.decodeIfPresent(Data.self, forKey: .cloudflareTunnelModeJson)
         additionalFields = try container.decodeIfPresent([String: String].self, forKey: .additionalFields)
-        passwordSource = try? container.decodeIfPresent(PasswordSource.self, forKey: .passwordSource)
+        passwordSource = PasswordSource.resilientlyDecoded(from: container, forKey: .passwordSource)
         localOnly = try container.decodeIfPresent(Bool.self, forKey: .localOnly) ?? false
         isSample = try container.decodeIfPresent(Bool.self, forKey: .isSample) ?? false
         isFavorite = try container.decodeIfPresent(Bool.self, forKey: .isFavorite) ?? false

TablePro/Core/Utilities/Connection/PasswordSourceResolver.swift

@@ -13,13 +13,15 @@ enum PasswordSourceResolver {
     private static let logger = Logger(subsystem: "com.TablePro", category: "PasswordSourceResolver")
 
     private static let commandTimeoutSeconds: UInt64 = 30
+    private static let maxOutputBytes = 1_048_576
 
     enum ResolutionError: LocalizedError {
         case fileNotFound(path: String)
         case fileUnreadable(path: String)
         case environmentVariableNotSet(name: String)
         case commandFailed(exitCode: Int32, stderr: String)
         case commandTimedOut
+        case outputTooLarge
         case emptyPassword
 
         var errorDescription: String? {
@@ -45,6 +47,8 @@ enum PasswordSourceResolver {
                 return String(format: String(localized: "Password command failed (exit %d): %@"), exitCode, message)
             case .commandTimedOut:
                 return String(localized: "Password command timed out after 30 seconds")
+            case .outputTooLarge:
+                return String(localized: "Password command produced too much output")
             case .emptyPassword:
                 return String(localized: "The password source produced an empty password")
             }
@@ -94,11 +98,15 @@ enum PasswordSourceResolver {
             process.standardOutput = stdoutPipe
             process.standardError = stderrPipe
 
-            let stdoutCollector = PipeDataCollector()
-            let stderrCollector = PipeDataCollector()
+            let stdoutCollector = PipeDataCollector(maxBytes: maxOutputBytes)
+            let stderrCollector = PipeDataCollector(maxBytes: maxOutputBytes)
             stdoutPipe.fileHandleForReading.readabilityHandler = { handle in
                 let chunk = handle.availableData
-                if !chunk.isEmpty { stdoutCollector.append(chunk) }
+                guard !chunk.isEmpty else { return }
+                stdoutCollector.append(chunk)
+                if stdoutCollector.overflowed, process.isRunning {
+                    process.terminate()
+                }
             }
             stderrPipe.fileHandleForReading.readabilityHandler = { handle in
                 let chunk = handle.availableData
@@ -107,9 +115,13 @@ enum PasswordSourceResolver {
 
             try process.run()
 
+            let didTimeout = AtomicFlag()
             let timeoutTask = Task.detached {
                 try await Task.sleep(nanoseconds: timeoutSeconds * 1_000_000_000)
-                if process.isRunning { process.terminate() }
+                if process.isRunning {
+                    didTimeout.set()
+                    process.terminate()
+                }
             }
 
             process.waitUntilExit()
@@ -118,7 +130,10 @@ enum PasswordSourceResolver {
             stdoutPipe.fileHandleForReading.readabilityHandler = nil
             stderrPipe.fileHandleForReading.readabilityHandler = nil
 
-            if process.terminationReason == .uncaughtSignal {
+            if stdoutCollector.overflowed {
+                throw ResolutionError.outputTooLarge
+            }
+            if didTimeout.isSet {
                 throw ResolutionError.commandTimedOut
             }
             if process.terminationStatus != 0 {
@@ -167,12 +182,34 @@ enum PasswordSourceResolver {
 
 private final class PipeDataCollector: @unchecked Sendable {
     private let lock = NSLock()
+    private let maxBytes: Int
     private var data = Data()
+    private var didOverflow = false
+
+    init(maxBytes: Int) {
+        self.maxBytes = maxBytes
+    }
 
     func append(_ chunk: Data) {
         lock.lock()
-        data.append(chunk)
-        lock.unlock()
+        defer { lock.unlock() }
+        let remaining = maxBytes - data.count
+        guard remaining > 0 else {
+            didOverflow = true
+            return
+        }
+        if chunk.count > remaining {
+            data.append(chunk.prefix(remaining))
+            didOverflow = true
+        } else {
+            data.append(chunk)
+        }
+    }
+
+    var overflowed: Bool {
+        lock.lock()
+        defer { lock.unlock() }
+        return didOverflow
     }
 
     var string: String {
@@ -181,3 +218,20 @@ private final class PipeDataCollector: @unchecked Sendable {
         return String(data: data, encoding: .utf8) ?? ""
     }
 }
+
+private final class AtomicFlag: @unchecked Sendable {
+    private let lock = NSLock()
+    private var value = false
+
+    func set() {
+        lock.lock()
+        value = true
+        lock.unlock()
+    }
+
+    var isSet: Bool {
+        lock.lock()
+        defer { lock.unlock() }
+        return value
+    }
+}

TablePro/Models/Connection/DatabaseConnection.swift

@@ -553,7 +553,7 @@ extension DatabaseConnection: Codable {
         localOnly = try container.decodeIfPresent(Bool.self, forKey: .localOnly) ?? false
         isSample = try container.decodeIfPresent(Bool.self, forKey: .isSample) ?? false
         isFavorite = try container.decodeIfPresent(Bool.self, forKey: .isFavorite) ?? false
-        passwordSource = try? container.decodeIfPresent(PasswordSource.self, forKey: .passwordSource)
+        passwordSource = PasswordSource.resilientlyDecoded(from: container, forKey: .passwordSource)
         cloudflareTunnelMode = try container.decodeIfPresent(CloudflareTunnelMode.self, forKey: .cloudflareTunnelMode) ?? .disabled
 
         // Migrate from legacy fields if sshTunnelMode is not present

TablePro/Models/Connection/PasswordSource.swift

@@ -4,6 +4,7 @@
 //
 
 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.
@@ -12,6 +13,8 @@ enum PasswordSource: Codable, Hashable, Sendable {
     case env(variable: String)
     case command(shell: String)
 
+    private static let logger = Logger(subsystem: "com.TablePro", category: "PasswordSource")
+
     private enum CodingKeys: String, CodingKey {
         case kind, path, variable, shell
     }
@@ -53,4 +56,18 @@ enum PasswordSource: Codable, Hashable, Sendable {
             try container.encode(shell, forKey: .shell)
         }
     }
+
+    /// 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<Key>(
+        from container: KeyedDecodingContainer<Key>,
+        forKey key: Key
+    ) -> PasswordSource? {
+        do {
+            return try container.decodeIfPresent(PasswordSource.self, forKey: key)
+        } catch {
+            logger.warning("Ignoring malformed passwordSource in a connection")
+            return nil
+        }
+    }
 }

TableProTests/Core/Utilities/PasswordSourceResolverTests.swift

@@ -7,7 +7,7 @@ import Foundation
 import Testing
 @testable import TablePro
 
-@Suite("PasswordSourceResolver")
+@Suite("PasswordSourceResolver", .serialized)
 struct PasswordSourceResolverTests {
     @Test("Reads a password from a file")
     func fileHappyPath() async throws {
@@ -154,6 +154,24 @@ struct PasswordSourceResolverTests {
         }
     }
 
+    @Test("Throws when command output exceeds the size cap")
+    func commandOutputTooLarge() async {
+        do {
+            _ = try await PasswordSourceResolver.resolveCommand(
+                shell: "head -c 2000000 /dev/zero | tr '\\0' 'a'",
+                timeoutSeconds: 30
+            )
+            Issue.record("Expected resolveCommand to reject oversized output")
+        } catch let error as PasswordSourceResolver.ResolutionError {
+            guard case .outputTooLarge = error else {
+                Issue.record("Expected outputTooLarge, got \(error)")
+                return
+            }
+        } catch {
+            Issue.record("Unexpected error: \(error)")
+        }
+    }
+
     @Test("Times out a slow command")
     func commandTimesOut() async {
         do {