[codex] Harden core reliability

ListenBar/Sources/Services/PortScanning/PortScannerService.swift

@@ -1,3 +1,5 @@
+import Darwin
+import Dispatch
 import Foundation
 
 enum PortScannerError: LocalizedError, Equatable {
@@ -21,7 +23,19 @@ enum PortScannerError: LocalizedError, Equatable {
     }
 }
 
+struct PortScannerProcessResult: Equatable, Sendable {
+    let terminationStatus: Int32
+    let standardOutput: Data
+    let standardError: Data
+}
+
 enum PortScannerService {
+    private static let processIOQueue = DispatchQueue(
+        label: "top.ygsgdbd.ListenBar.PortScannerService.IO",
+        qos: .userInitiated,
+        attributes: .concurrent,
+    )
+
     static let lsofPath = "/usr/sbin/lsof"
     static let lsofArguments = [
         "-nP",
@@ -36,39 +50,116 @@ enum PortScannerService {
     ]
 
     static func scanListeningPorts() async throws -> [PortEntry] {
-        try await Task.detached(priority: .userInitiated) {
-            let process = Process()
-            process.executableURL = URL(fileURLWithPath: lsofPath)
-            process.arguments = lsofArguments
+        let result = try await executeProcess(
+            executableURL: URL(fileURLWithPath: lsofPath),
+            arguments: lsofArguments,
+        )
+        return try interpretLsofResult(result)
+    }
+
+    static func executeProcess(
+        executableURL: URL,
+        arguments: [String],
+    ) async throws -> PortScannerProcessResult {
+        let process = Process()
+        process.executableURL = executableURL
+        process.arguments = arguments
+
+        let standardOutput = Pipe()
+        let standardError = Pipe()
+        process.standardOutput = standardOutput
+        process.standardError = standardError
+
+        let outputDescriptor = try duplicateDescriptor(
+            standardOutput.fileHandleForReading.fileDescriptor,
+        )
+        let errorDescriptor: Int32
+        do {
+            errorDescriptor = try duplicateDescriptor(
+                standardError.fileHandleForReading.fileDescriptor,
+            )
+        } catch {
+            Darwin.close(outputDescriptor)
+            throw error
+        }
 
-            let standardOutput = Pipe()
-            let standardError = Pipe()
-            process.standardOutput = standardOutput
-            process.standardError = standardError
+        async let outputData = readData(from: outputDescriptor)
+        async let errorData = readData(from: errorDescriptor)
 
-            try process.run()
-            process.waitUntilExit()
+        do {
+            let terminationStatus = try await runAndWaitForTermination(process)
+            let (standardOutputData, standardErrorData) = try await (outputData, errorData)
+            return PortScannerProcessResult(
+                terminationStatus: terminationStatus,
+                standardOutput: standardOutputData,
+                standardError: standardErrorData,
+            )
+        } catch {
+            try? standardOutput.fileHandleForWriting.close()
+            try? standardError.fileHandleForWriting.close()
+            _ = try? await (outputData, errorData)
+            throw error
+        }
+    }
 
-            let outputData = standardOutput.fileHandleForReading.readDataToEndOfFile()
-            let errorData = standardError.fileHandleForReading.readDataToEndOfFile()
+    static func interpretLsofResult(_ result: PortScannerProcessResult) throws -> [PortEntry] {
+        guard result.terminationStatus == 0 else {
+            let message = String(data: result.standardError, encoding: .utf8)?
+                .trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
+            throw PortScannerError.lsofFailed(
+                status: result.terminationStatus,
+                message: message,
+            )
+        }
 
-            guard let output = String(data: outputData, encoding: .utf8) else {
-                throw PortScannerError.lsofOutputUnreadable
+        guard let output = String(data: result.standardOutput, encoding: .utf8) else {
+            throw PortScannerError.lsofOutputUnreadable
+        }
+
+        return parseLsofFieldOutput(output)
+    }
+
+    private static func duplicateDescriptor(_ descriptor: Int32) throws -> Int32 {
+        let duplicate = Darwin.dup(descriptor)
+        guard duplicate >= 0 else {
+            throw POSIXError(POSIXErrorCode(rawValue: errno) ?? .EIO)
+        }
+        return duplicate
+    }
+
+    private static func readData(from descriptor: Int32) async throws -> Data {
+        try await withCheckedThrowingContinuation { continuation in
+            let state = DispatchIOReadState(continuation: continuation)
+            let channel = DispatchIO(
+                type: .stream,
+                fileDescriptor: descriptor,
+                queue: processIOQueue,
+            ) { _ in
+                Darwin.close(descriptor)
+            }
+            channel.setLimit(lowWater: 64 * 1024)
+            channel.read(offset: 0, length: Int.max, queue: processIOQueue) { done, chunk, error in
+                if done {
+                    channel.close()
+                }
+                state.receive(chunk: chunk, done: done, error: error)
             }
+        }
+    }
 
-            let ports = parseLsofFieldOutput(output)
-            if ports.isEmpty, process.terminationStatus != 0 {
-                let message = String(data: errorData, encoding: .utf8)?
-                    .trimmingCharacters(in: .whitespacesAndNewlines) ?? ""
-                throw PortScannerError.lsofFailed(
-                    status: process.terminationStatus,
-                    message: message,
-                )
+    private static func runAndWaitForTermination(_ process: Process) async throws -> Int32 {
+        try await withCheckedThrowingContinuation { continuation in
+            process.terminationHandler = { terminatedProcess in
+                continuation.resume(returning: terminatedProcess.terminationStatus)
             }
 
-            return ports
+            do {
+                try process.run()
+            } catch {
+                process.terminationHandler = nil
+                continuation.resume(throwing: error)
+            }
         }
-        .value
     }
 
     static func parseLsofFieldOutput(_ output: String) -> [PortEntry] {
@@ -186,3 +277,28 @@ private struct FileRecord {
     var networkProtocol: NetworkProtocol?
     var name: String?
 }
+
+private final class DispatchIOReadState: @unchecked Sendable {
+    private let continuation: CheckedContinuation<Data, Error>
+    private var data = Data()
+
+    init(continuation: CheckedContinuation<Data, Error>) {
+        self.continuation = continuation
+    }
+
+    // DispatchIO guarantees that handlers for one read operation are not reentrant.
+    func receive(chunk: DispatchData?, done: Bool, error: Int32) {
+        if let chunk {
+            chunk.enumerateBytes { buffer, _, _ in
+                data.append(contentsOf: buffer)
+            }
+        }
+
+        guard done else { return }
+        if error == 0 {
+            continuation.resume(returning: data)
+        } else {
+            continuation.resume(throwing: POSIXError(POSIXErrorCode(rawValue: error) ?? .EIO))
+        }
+    }
+}

ListenBar/Sources/Services/System/LaunchAtLoginService.swift

@@ -10,11 +10,13 @@ struct LaunchAtLoginServiceEnvironment {
     var runLaunchctl: ([String]) throws -> Void
     var userID: () -> uid_t
     var fileManager: FileManager
+    var replaceItemAt: (URL, URL) throws -> Void
 }
 
 enum LaunchAtLoginService {
     private static let launchAgentIdentifier = "top.ygsgdbd.ListenBar"
     private static let launchAgentPlistName = "\(launchAgentIdentifier).plist"
+    private static let launchAgentStagingDirectoryName = ".\(launchAgentIdentifier).staging"
 
     private static var liveEnvironment: LaunchAtLoginServiceEnvironment {
         LaunchAtLoginServiceEnvironment(
@@ -28,6 +30,12 @@ enum LaunchAtLoginService {
             runLaunchctl: runLaunchctl,
             userID: { getuid() },
             fileManager: .default,
+            replaceItemAt: { originalItemURL, newItemURL in
+                _ = try FileManager.default.replaceItemAt(
+                    originalItemURL,
+                    withItemAt: newItemURL,
+                )
+            },
         )
     }
 
@@ -166,6 +174,13 @@ private extension LaunchAtLoginService {
             at: environment.plistURL.deletingLastPathComponent(),
             withIntermediateDirectories: true,
         )
+        let stagingDirectoryURL = environment.plistURL.deletingLastPathComponent()
+            .appendingPathComponent(launchAgentStagingDirectoryName, isDirectory: true)
+        let stagingPlistURL = stagingDirectoryURL.appendingPathComponent(launchAgentPlistName)
+        try environment.fileManager.createDirectory(
+            at: stagingDirectoryURL,
+            withIntermediateDirectories: true,
+        )
 
         let plist: [String: Any] = [
             "Label": launchAgentIdentifier,
@@ -177,19 +192,92 @@ private extension LaunchAtLoginService {
             format: .xml,
             options: 0,
         )
-        try data.write(to: environment.plistURL, options: .atomic)
+        try data.write(to: stagingPlistURL, options: .atomic)
 
         let domain = "gui/\(environment.userID())"
-        try? environment.runLaunchctl(["bootout", domain, environment.plistURL.path])
-        try environment.runLaunchctl(["bootstrap", domain, environment.plistURL.path])
+        let serviceTarget = "\(domain)/\(launchAgentIdentifier)"
+        try? environment.runLaunchctl(["bootout", serviceTarget])
+        do {
+            try environment.runLaunchctl(["bootstrap", domain, stagingPlistURL.path])
+            try commitFallbackLaunchAgent(
+                stagingPlistURL: stagingPlistURL,
+                environment: environment,
+            )
+        } catch {
+            cleanUpFailedFallbackInstallation(
+                stagingDirectoryURL: stagingDirectoryURL,
+                serviceTarget: serviceTarget,
+                environment: environment,
+            )
+            throw error
+        }
+
+        removeStagingDirectory(
+            stagingDirectoryURL,
+            fileManager: environment.fileManager,
+        )
     }
 
     static func removeFallbackLaunchAgent(environment: LaunchAtLoginServiceEnvironment) {
-        let domain = "gui/\(environment.userID())"
-        try? environment.runLaunchctl(["bootout", domain, environment.plistURL.path])
+        let serviceTarget = "gui/\(environment.userID())/\(launchAgentIdentifier)"
+        try? environment.runLaunchctl(["bootout", serviceTarget])
         try? environment.fileManager.removeItem(at: environment.plistURL)
     }
 
+    static func commitFallbackLaunchAgent(
+        stagingPlistURL: URL,
+        environment: LaunchAtLoginServiceEnvironment,
+    ) throws {
+        if environment.fileManager.fileExists(atPath: environment.plistURL.path) {
+            try environment.replaceItemAt(
+                environment.plistURL,
+                stagingPlistURL,
+            )
+        } else {
+            try environment.fileManager.moveItem(
+                at: stagingPlistURL,
+                to: environment.plistURL,
+            )
+        }
+    }
+
+    static func cleanUpFailedFallbackInstallation(
+        stagingDirectoryURL: URL,
+        serviceTarget: String,
+        environment: LaunchAtLoginServiceEnvironment,
+    ) {
+        do {
+            try environment.runLaunchctl(["bootout", serviceTarget])
+        } catch {
+            print("Failed to clean up ListenBar LaunchAgent service: \(error.localizedDescription)")
+        }
+
+        removeStagingDirectory(
+            stagingDirectoryURL,
+            fileManager: environment.fileManager,
+        )
+
+        if environment.fileManager.fileExists(atPath: environment.plistURL.path) {
+            do {
+                try environment.fileManager.removeItem(at: environment.plistURL)
+            } catch {
+                print("Failed to remove ListenBar LaunchAgent plist after rollback: \(error.localizedDescription)")
+            }
+        }
+    }
+
+    static func removeStagingDirectory(_ url: URL, fileManager: FileManager) {
+        guard fileManager.fileExists(atPath: url.path) else {
+            return
+        }
+
+        do {
+            try fileManager.removeItem(at: url)
+        } catch {
+            print("Failed to remove ListenBar LaunchAgent staging files: \(error.localizedDescription)")
+        }
+    }
+
     static func runLaunchctl(arguments: [String]) throws {
         let process = Process()
         process.executableURL = URL(fileURLWithPath: "/bin/launchctl")