Rewrite SocketController (#634)

* WIP

* Working

* Working

* Cleanup

Sources/Packages/Sources/SecretAgentKit/Agent.swift

@@ -11,7 +11,6 @@ public final class Agent: Sendable {
     private let witness: SigningWitness?
     private let publicKeyWriter = OpenSSHPublicKeyWriter()
     private let signatureWriter = OpenSSHSignatureWriter()
-    private let requestTracer = SigningRequestTracer()
     private let certificateHandler = OpenSSHCertificateHandler()
     private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "Agent")
 
@@ -34,28 +33,26 @@ extension Agent {
 
     /// Handles an incoming request.
     /// - Parameters:
-    ///   - reader: A ``FileHandleReader`` to read the content of the request.
-    ///   - writer: A ``FileHandleWriter`` to write the response to.
-    /// - Return value: 
-    ///   - Boolean if data could be read
-    @discardableResult public func handle(reader: FileHandleReader, writer: FileHandleWriter) async -> Bool {
+    ///   - data: The data to handle.
+    ///   - provenance: The origin of the request.
+    /// - Returns: A response data payload.
+    public func handle(data: Data, provenance: SigningRequestProvenance) async throws -> Data {
         logger.debug("Agent handling new data")
-        let data = Data(reader.availableData)
-        guard data.count > 4 else { return false}
+        guard data.count > 4 else {
+            throw AgentError.couldNotRead
+        }
         let requestTypeInt = data[4]
         guard let requestType = SSHAgent.RequestType(rawValue: requestTypeInt) else {
-            writer.write(SSHAgent.ResponseType.agentFailure.data.lengthAndData)
             logger.debug("Agent returned \(SSHAgent.ResponseType.agentFailure.debugDescription)")
-            return true
+            return SSHAgent.ResponseType.agentFailure.data.lengthAndData
         }
         logger.debug("Agent handling request of type \(requestType.debugDescription)")
         let subData = Data(data[5...])
-        let response = await handle(requestType: requestType, data: subData, reader: reader)
-        writer.write(response)
-        return true
+        let response = await handle(requestType: requestType, data: subData, provenance: provenance)
+        return response
     }
 
-    func handle(requestType: SSHAgent.RequestType, data: Data, reader: FileHandleReader) async -> Data {
+    private func handle(requestType: SSHAgent.RequestType, data: Data, provenance: SigningRequestProvenance) async -> Data {
         // Depending on the launch context (such as after macOS update), the agent may need to reload secrets before acting
         await reloadSecretsIfNeccessary()
         var response = Data()
@@ -66,7 +63,6 @@ extension Agent {
                 response.append(await identities())
                 logger.debug("Agent returned \(SSHAgent.ResponseType.agentIdentitiesAnswer.debugDescription)")
             case .signRequest:
-                let provenance = requestTracer.provenance(from: reader)
                 response.append(SSHAgent.ResponseType.agentSignResponse.data)
                 response.append(try await sign(data: data, provenance: provenance))
                 logger.debug("Agent returned \(SSHAgent.ResponseType.agentSignResponse.debugDescription)")
@@ -184,6 +180,7 @@ extension Agent {
 
     /// An error involving agent operations..
     enum AgentError: Error {
+        case couldNotRead
         case unhandledType
         case noMatchingKey
         case unsupportedKeyType

Sources/Packages/Sources/SecretAgentKit/SocketController.swift

@@ -1,23 +1,32 @@
 import Foundation
 import OSLog
+import SecretKit
 
 /// A controller that manages socket configuration and request dispatching.
-public final class SocketController {
+public struct SocketController {
 
-    /// The active FileHandle.
-    private var fileHandle: FileHandle?
-    /// The active SocketPort.
-    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: (@Sendable (FileHandleReader, FileHandleWriter) async -> Bool)?
-    /// Logger.
+    /// A stream of Sessions. Each session represents one connection to a class communicating with the socket. Multiple Sessions may be active simultaneously.
+    public let sessions: AsyncStream<Session>
+
+    /// A continuation to create new sessions.
+    private let sessionsContinuation: AsyncStream<Session>.Continuation
+
+    /// The active SocketPort. Must be retained to be kept valid.
+    private let port: SocketPort
+
+    /// The FileHandle for the main socket.
+    private let fileHandle: FileHandle
+    
+    /// Logger for the socket controller.
     private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "SocketController")
 
+    /// Tracer which determines who originates a socket connection.
+    private let requestTracer = SigningRequestTracer()
 
     /// Initializes a socket controller with a specified path.
     /// - Parameter path: The path to use as a socket.
     public init(path: String) {
+        (sessions, sessionsContinuation) = AsyncStream<Session>.makeStream()
         logger.debug("Socket controller setting up at \(path)")
         if let _ = try? FileManager.default.removeItem(atPath: path) {
             logger.debug("Socket controller removed existing socket")
@@ -25,25 +34,102 @@ public final class SocketController {
         let exists = FileManager.default.fileExists(atPath: path)
         assert(!exists)
         logger.debug("Socket controller path is clear")
-        port = socketPort(at: path)
-        configureSocket(at: path)
+        port = SocketPort(path: path)
+        fileHandle = FileHandle(fileDescriptor: port.socket, closeOnDealloc: true)
+        Task { [fileHandle, sessionsContinuation, logger] in
+            for await notification in NotificationCenter.default.notifications(named: .NSFileHandleConnectionAccepted) {
+                logger.debug("Socket controller accepted connection")
+                guard let new = notification.userInfo?[NSFileHandleNotificationFileHandleItem] as? FileHandle else { continue }
+                let session = Session(fileHandle: new)
+                sessionsContinuation.yield(session)
+                await fileHandle.acceptConnectionInBackgroundAndNotifyOnMainActor()
+            }
+        }
+        fileHandle.acceptConnectionInBackgroundAndNotify(forModes: [RunLoop.Mode.common])
         logger.debug("Socket listening at \(path)")
     }
 
-    /// Configures the socket and a corresponding FileHandle.
-    /// - Parameter path: The path to use as a socket.
-    func configureSocket(at path: String) {
-        guard let port = port else { return }
-        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.Mode.common])
+}
+
+extension SocketController {
+    
+    /// A session represents a connection that has been established between the two ends of the socket.
+    public struct Session: Sendable {
+        
+        /// Data received by the socket.
+        public let messages: AsyncStream<Data>
+
+        /// The provenance of the process that established the session.
+        public let provenance: SigningRequestProvenance
+        
+        /// A FileHandle used to communicate with the socket.
+        private let fileHandle: FileHandle
+
+        /// A continuation for issuing new messages.
+        private let messagesContinuation: AsyncStream<Data>.Continuation
+
+        /// A logger for the session.
+        private let logger = Logger(subsystem: "com.maxgoedjen.secretive.secretagent", category: "Session")
+        
+        /// Initializes a new Session.
+        /// - Parameter fileHandle: The FileHandle used to communicate with the socket.
+        init(fileHandle: FileHandle) {
+            self.fileHandle = fileHandle
+            provenance = SigningRequestTracer().provenance(from: fileHandle)
+            (messages, messagesContinuation) = AsyncStream.makeStream()
+            Task { [messagesContinuation, logger] in
+                await fileHandle.waitForDataInBackgroundAndNotifyOnMainActor()
+                for await _ in NotificationCenter.default.notifications(named: .NSFileHandleDataAvailable, object: fileHandle) {
+                    let data = fileHandle.availableData
+                    guard !data.isEmpty else {
+                        logger.debug("Socket controller received empty data, ending continuation.")
+                        messagesContinuation.finish()
+                        try fileHandle.close()
+                        return
+                    }
+                    messagesContinuation.yield(data)
+                    logger.debug("Socket controller yielded data.")
+                }
+            }
+        }
+        
+        /// Writes new data to the socket.
+        /// - Parameter data: The data to write.
+        public func write(_ data: Data) async throws {
+            try fileHandle.write(contentsOf: data)
+            await fileHandle.waitForDataInBackgroundAndNotifyOnMainActor()
+        }
+        
+        /// Closes the socket and cleans up resources.
+        public func close() throws {
+            logger.debug("Session closed.")
+            messagesContinuation.finish()
+            try fileHandle.close()
+        }
+
     }
 
-    /// Creates a SocketPort for a path.
-    /// - Parameter path: The path to use as a socket.
-    /// - Returns: A configured SocketPort.
-    func socketPort(at path: String) -> SocketPort {
+}
+
+private 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)
+    }
+
+}
+
+private extension SocketPort {
+
+    convenience init(path: String) {
         var addr = sockaddr_un()
         addr.sun_family = sa_family_t(AF_UNIX)
 
@@ -61,51 +147,7 @@ public final class SocketController {
             data = Data(bytes: pointer, count: MemoryLayout<sockaddr_un>.size)
         }
 
-        return SocketPort(protocolFamily: AF_UNIX, socketType: SOCK_STREAM, protocol: 0, address: data)!
-    }
-
-    /// 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")
-        guard let new = notification.userInfo?[NSFileHandleNotificationFileHandleItem] as? FileHandle else { return }
-        Task { [handler, fileHandle] in
-            _ = await handler?(new, new)
-            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")
-        guard let new = notification.object as? FileHandle else { return }
-        logger.debug("Socket controller received new file handle")
-        Task { [handler, logger = logger] in
-            if((await handler?(new, new)) == true) {
-                logger.debug("Socket controller handled data, wait for more data")
-                await new.waitForDataInBackgroundAndNotifyOnMainActor()
-            } 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)
+        self.init(protocolFamily: AF_UNIX, socketType: SOCK_STREAM, protocol: 0, address: data)!
     }
 
 }