Handle concurrent requests to socket (#495)

* Socket updates

* Sendable.

* Update tests.

Sources/Packages/Sources/SecretAgentKit/Agent.swift

@@ -12,7 +12,7 @@ public class Agent {
     private let writer = OpenSSHKeyWriter()
     private let requestTracer = SigningRequestTracer()
     private let certificateHandler = OpenSSHCertificateHandler()
-    private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent.agent", category: "")
+    private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "Agent")
 
     /// Initializes an agent with a store list and a witness.
     /// - Parameters:
@@ -35,7 +35,7 @@ extension Agent {
     ///   - writer: A ``FileHandleWriter`` to write the response to.
     /// - Return value: 
     ///   - Boolean if data could be read
-    @discardableResult public func handle(reader: FileHandleReader, writer: FileHandleWriter) -> Bool {
+    @discardableResult @Sendable public func handle(reader: FileHandleReader, writer: FileHandleWriter) async -> Bool {
         logger.debug("Agent handling new data")
         let data = Data(reader.availableData)
         guard data.count > 4 else { return false}
@@ -47,12 +47,12 @@ extension Agent {
         }
         logger.debug("Agent handling request of type \(requestType.debugDescription)")
         let subData = Data(data[5...])
-        let response = handle(requestType: requestType, data: subData, reader: reader)
+        let response = await handle(requestType: requestType, data: subData, reader: reader)
         writer.write(response)
         return true
     }
 
-    func handle(requestType: SSHAgent.RequestType, data: Data, reader: FileHandleReader) -> Data {
+    func handle(requestType: SSHAgent.RequestType, data: Data, reader: FileHandleReader) async -> Data {
         // Depending on the launch context (such as after macOS update), the agent may need to reload secrets before acting
         reloadSecretsIfNeccessary()
         var response = Data()

Sources/Packages/Sources/SecretAgentKit/FileHandleProtocols.swift

@@ -1,7 +1,7 @@
 import Foundation
 
 /// Protocol abstraction of the reading aspects of FileHandle.
-public protocol FileHandleReader {
+public protocol FileHandleReader: Sendable {
 
     /// Gets data that is available for reading.
     var availableData: Data { get }
@@ -13,7 +13,7 @@ public protocol FileHandleReader {
 }
 
 /// Protocol abstraction of the writing aspects of FileHandle.
-public protocol FileHandleWriter {
+public protocol FileHandleWriter: Sendable {
 
     /// Writes data to the handle.
     func write(_ data: Data)

Sources/Packages/Sources/SecretAgentKit/SocketController.swift

@@ -10,22 +10,24 @@ public class SocketController {
     private var port: SocketPort?
     /// A handler that will be notified when a new read/write handle is available.
     /// False if no data could be read
-    public var handler: ((FileHandleReader, FileHandleWriter) -> Bool)?
+    public var handler: ((FileHandleReader, FileHandleWriter) async -> Bool)?
+    /// Logger.
+    private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "SocketController")
 
 
     /// Initializes a socket controller with a specified path.
     /// - Parameter path: The path to use as a socket.
     public init(path: String) {
-        Logger().debug("Socket controller setting up at \(path)")
+        logger.debug("Socket controller setting up at \(path)")
         if let _ = try? FileManager.default.removeItem(atPath: path) {
-            Logger().debug("Socket controller removed existing socket")
+            logger.debug("Socket controller removed existing socket")
         }
         let exists = FileManager.default.fileExists(atPath: path)
         assert(!exists)
-        Logger().debug("Socket controller path is clear")
+        logger.debug("Socket controller path is clear")
         port = socketPort(at: path)
         configureSocket(at: path)
-        Logger().debug("Socket listening at \(path)")
+        logger.debug("Socket listening at \(path)")
     }
 
     /// Configures the socket and a corresponding FileHandle.
@@ -35,7 +37,7 @@ public class SocketController {
         fileHandle = FileHandle(fileDescriptor: port.socket, closeOnDealloc: true)
         NotificationCenter.default.addObserver(self, selector: #selector(handleConnectionAccept(notification:)), name: .NSFileHandleConnectionAccepted, object: nil)
         NotificationCenter.default.addObserver(self, selector: #selector(handleConnectionDataAvailable(notification:)), name: .NSFileHandleDataAvailable, object: nil)
-        fileHandle?.acceptConnectionInBackgroundAndNotify(forModes: [RunLoop.current.currentMode!])
+        fileHandle?.acceptConnectionInBackgroundAndNotify(forModes: [RunLoop.Mode.common])
     }
 
     /// Creates a SocketPort for a path.
@@ -65,25 +67,45 @@ public class SocketController {
     /// Handles a new connection being accepted, invokes the handler, and prepares to accept new connections.
     /// - Parameter notification: A `Notification` that triggered the call.
     @objc func handleConnectionAccept(notification: Notification) {
-        Logger().debug("Socket controller accepted connection")
+        logger.debug("Socket controller accepted connection")
         guard let new = notification.userInfo?[NSFileHandleNotificationFileHandleItem] as? FileHandle else { return }
-        _ = handler?(new, new)
+        Task {
-        new.waitForDataInBackgroundAndNotify()
+            _ = await handler?(new, new)
-        fileHandle?.acceptConnectionInBackgroundAndNotify(forModes: [RunLoop.current.currentMode!])
+            await new.waitForDataInBackgroundAndNotifyOnMainActor()
+            await fileHandle?.acceptConnectionInBackgroundAndNotifyOnMainActor()
+        }
     }
 
     /// Handles a new connection providing data and invokes the handler callback.
     /// - Parameter notification: A `Notification` that triggered the call.
     @objc func handleConnectionDataAvailable(notification: Notification) {
-        Logger().debug("Socket controller has new data available")
+        logger.debug("Socket controller has new data available")
         guard let new = notification.object as? FileHandle else { return }
-        Logger().debug("Socket controller received new file handle")
+        logger.debug("Socket controller received new file handle")
-        if((handler?(new, new)) == true) {
+        Task {
-            Logger().debug("Socket controller handled data, wait for more data")
+            if((await handler?(new, new)) == true) {
-            new.waitForDataInBackgroundAndNotify()
+                logger.debug("Socket controller handled data, wait for more data")
-        } else {
+                await new.waitForDataInBackgroundAndNotifyOnMainActor()
-            Logger().debug("Socket controller called with empty data, socked closed")
+            } else {
+                logger.debug("Socket controller called with empty data, socked closed")
+            }
         }
     }
 
 }
+
+extension FileHandle {
+    
+    /// Ensures waitForDataInBackgroundAndNotify will be called on the main actor.
+    @MainActor func waitForDataInBackgroundAndNotifyOnMainActor() {
+        waitForDataInBackgroundAndNotify()
+    }
+
+
+    /// Ensures acceptConnectionInBackgroundAndNotify will be called on the main actor.
+    /// - Parameter modes: the runloop modes to use.
+    @MainActor func acceptConnectionInBackgroundAndNotifyOnMainActor(forModes modes: [RunLoop.Mode]? = [RunLoop.Mode.common]) {
+        acceptConnectionInBackgroundAndNotify(forModes: modes)
+    }
+
+}

Sources/Packages/Sources/SecretKit/OpenSSH/OpenSSHCertificateHandler.swift

@@ -5,7 +5,7 @@ import OSLog
 public class OpenSSHCertificateHandler {
 
     private let publicKeyFileStoreController = PublicKeyFileStoreController(homeDirectory: NSHomeDirectory())
-    private let logger = Logger()
+    private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "OpenSSHCertificateHandler")
     private let writer = OpenSSHKeyWriter()
     private var keyBlobsAndNames: [AnySecret: (Data, Data)] = [:]
 

Sources/Packages/Sources/SecretKit/PublicKeyStandinFileController.swift

@@ -4,7 +4,7 @@ import OSLog
 /// Controller responsible for writing public keys to disk, so that they're easily accessible by scripts.
 public class PublicKeyFileStoreController {
 
-    private let logger = Logger()
+    private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "PublicKeyFileStoreController")
     private let directory: String
     private let keyWriter = OpenSSHKeyWriter()
 

Sources/Packages/Tests/SecretAgentKitTests/AgentTests.swift

@@ -10,39 +10,39 @@ class AgentTests: XCTestCase {
 
     // MARK: Identity Listing
 
-    func testEmptyStores() {
+    func testEmptyStores() async {
         let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestIdentities)
         let agent = Agent(storeList: SecretStoreList())
-        agent.handle(reader: stubReader, writer: stubWriter)
+        await agent.handle(reader: stubReader, writer: stubWriter)
         XCTAssertEqual(stubWriter.data, Constants.Responses.requestIdentitiesEmpty)
     }
 
-    func testIdentitiesList() {
+    func testIdentitiesList() async {
         let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestIdentities)
         let list = storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
         let agent = Agent(storeList: list)
-        agent.handle(reader: stubReader, writer: stubWriter)
+        await agent.handle(reader: stubReader, writer: stubWriter)
         XCTAssertEqual(stubWriter.data, Constants.Responses.requestIdentitiesMultiple)
     }
 
     // MARK: Signatures
 
-    func testNoMatchingIdentities() {
+    func testNoMatchingIdentities() async {
         let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignatureWithNoneMatching)
         let list = storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
         let agent = Agent(storeList: list)
-        agent.handle(reader: stubReader, writer: stubWriter)
+        await agent.handle(reader: stubReader, writer: stubWriter)
 //        XCTAssertEqual(stubWriter.data, Constants.Responses.requestFailure)
     }
 
-    func testSignature() {
+    func testSignature() async {
         let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignature)
         let requestReader = OpenSSHReader(data: Constants.Requests.requestSignature[5...])
         _ = requestReader.readNextChunk()
         let dataToSign = requestReader.readNextChunk()
         let list = storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
         let agent = Agent(storeList: list)
-        agent.handle(reader: stubReader, writer: stubWriter)
+        await agent.handle(reader: stubReader, writer: stubWriter)
         let outer = OpenSSHReader(data: stubWriter.data[5...])
         let payload = outer.readNextChunk()
         let inner = OpenSSHReader(data: payload)
@@ -76,18 +76,18 @@ class AgentTests: XCTestCase {
 
     // MARK: Witness protocol
 
-    func testWitnessObjectionStopsRequest() {
+    func testWitnessObjectionStopsRequest() async {
         let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignature)
         let list = storeList(with: [Constants.Secrets.ecdsa256Secret])
         let witness = StubWitness(speakNow: { _,_  in
             return true
         }, witness: { _, _ in })
         let agent = Agent(storeList: list, witness: witness)
-        agent.handle(reader: stubReader, writer: stubWriter)
+        await agent.handle(reader: stubReader, writer: stubWriter)
         XCTAssertEqual(stubWriter.data, Constants.Responses.requestFailure)
     }
 
-    func testWitnessSignature() {
+    func testWitnessSignature() async {
         let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignature)
         let list = storeList(with: [Constants.Secrets.ecdsa256Secret])
         var witnessed = false
@@ -97,11 +97,11 @@ class AgentTests: XCTestCase {
             witnessed = true
         })
         let agent = Agent(storeList: list, witness: witness)
-        agent.handle(reader: stubReader, writer: stubWriter)
+        await agent.handle(reader: stubReader, writer: stubWriter)
         XCTAssertTrue(witnessed)
     }
 
-    func testRequestTracing() {
+    func testRequestTracing() async {
         let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignature)
         let list = storeList(with: [Constants.Secrets.ecdsa256Secret])
         var speakNowTrace: SigningRequestProvenance! = nil
@@ -113,7 +113,7 @@ class AgentTests: XCTestCase {
             witnessTrace = trace
         })
         let agent = Agent(storeList: list, witness: witness)
-        agent.handle(reader: stubReader, writer: stubWriter)
+        await agent.handle(reader: stubReader, writer: stubWriter)
         XCTAssertEqual(witnessTrace, speakNowTrace)
         XCTAssertEqual(witnessTrace.origin.displayName, "Finder")
         XCTAssertEqual(witnessTrace.origin.validSignature, true)
@@ -122,22 +122,22 @@ class AgentTests: XCTestCase {
 
     // MARK: Exception Handling
 
-    func testSignatureException() {
+    func testSignatureException() async {
         let stubReader = StubFileHandleReader(availableData: Constants.Requests.requestSignature)
         let list = storeList(with: [Constants.Secrets.ecdsa256Secret, Constants.Secrets.ecdsa384Secret])
         let store = list.stores.first?.base as! Stub.Store
         store.shouldThrow = true
         let agent = Agent(storeList: list)
-        agent.handle(reader: stubReader, writer: stubWriter)
+        await agent.handle(reader: stubReader, writer: stubWriter)
         XCTAssertEqual(stubWriter.data, Constants.Responses.requestFailure)
     }
 
     // MARK: Unsupported
 
-    func testUnhandledAdd() {
+    func testUnhandledAdd() async {
         let stubReader = StubFileHandleReader(availableData: Constants.Requests.addIdentity)
         let agent = Agent(storeList: SecretStoreList())
-        agent.handle(reader: stubReader, writer: stubWriter)
+        await agent.handle(reader: stubReader, writer: stubWriter)
         XCTAssertEqual(stubWriter.data, Constants.Responses.requestFailure)
     }
 

Sources/SecretAgent/AppDelegate.swift

@@ -27,9 +27,10 @@ class AppDelegate: NSObject, NSApplicationDelegate {
         return SocketController(path: path)
     }()
     private var updateSink: AnyCancellable?
+    private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "AppDelegate")
 
     func applicationDidFinishLaunching(_ aNotification: Notification) {
-        Logger().debug("SecretAgent finished launching")
+        logger.debug("SecretAgent finished launching")
         DispatchQueue.main.async {
             self.socketController.handler = self.agent.handle(reader:writer:)
         }

Sources/Secretive/Controllers/LaunchAgentController.swift

@@ -6,8 +6,10 @@ import SecretKit
 
 struct LaunchAgentController {
     
+    private let logger = Logger(subsystem: "com.maxgoedjen.secretive", category: "LaunchAgentController")
+
     func install(completion: (() -> Void)? = nil) {
-        Logger().debug("Installing agent")
+        logger.debug("Installing agent")
         _ = setEnabled(false)
         // This is definitely a bit of a "seems to work better" thing but:
         // Seems to more reliably hit if these are on separate runloops, otherwise it seems like it sometimes doesn't kill old
@@ -20,7 +22,7 @@ struct LaunchAgentController {
     }
 
     func forceLaunch(completion: ((Bool) -> Void)?) {
-        Logger().debug("Agent is not running, attempting to force launch")
+        logger.debug("Agent is not running, attempting to force launch")
         let url = Bundle.main.bundleURL.appendingPathComponent("Contents/Library/LoginItems/SecretAgent.app")
         let config = NSWorkspace.OpenConfiguration()
         config.activates = false
@@ -29,9 +31,9 @@ struct LaunchAgentController {
                 completion?(error == nil)
             }
             if let error = error {
-                Logger().error("Error force launching \(error.localizedDescription)")
+                logger.error("Error force launching \(error.localizedDescription)")
             } else {
-                Logger().debug("Agent force launched")
+                logger.debug("Agent force launched")
             }
         }
     }