diff --git a/SecretAgentKit/Agent.swift b/SecretAgentKit/Agent.swift index 3e69037..5e4d259 100644 --- a/SecretAgentKit/Agent.swift +++ b/SecretAgentKit/Agent.swift @@ -21,43 +21,48 @@ public class Agent { extension Agent { - public func handle(reader: FileHandleReader, writer: FileHandleWriter) { + public func handle(reader: FileHandleReader, writer: FileHandleWriter) -> Bool { Logger().debug("Agent handling new data") let data = reader.availableData - guard !data.isEmpty else { return } + guard !data.isEmpty else { return true } let requestTypeInt = data[4] guard let requestType = SSHAgent.RequestType(rawValue: requestTypeInt) else { writer.write(OpenSSHKeyWriter().lengthAndData(of: SSHAgent.ResponseType.agentFailure.data)) Logger().debug("Agent returned \(SSHAgent.ResponseType.agentFailure.debugDescription)") - return + return true } 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, shouldClose) = handle(requestType: requestType, data: subData, reader: reader) writer.write(response) + return shouldClose } - func handle(requestType: SSHAgent.RequestType, data: Data, reader: FileHandleReader) -> Data { + func handle(requestType: SSHAgent.RequestType, data: Data, reader: FileHandleReader) -> (Data, Bool) { var response = Data() + let shouldClose: Bool do { switch requestType { case .requestIdentities: response.append(SSHAgent.ResponseType.agentIdentitiesAnswer.data) response.append(identities()) Logger().debug("Agent returned \(SSHAgent.ResponseType.agentIdentitiesAnswer.debugDescription)") + shouldClose = false case .signRequest: let provenance = requestTracer.provenance(from: reader) response.append(SSHAgent.ResponseType.agentSignResponse.data) response.append(try sign(data: data, provenance: provenance)) Logger().debug("Agent returned \(SSHAgent.ResponseType.agentSignResponse.debugDescription)") + shouldClose = true } } catch { response.removeAll() response.append(SSHAgent.ResponseType.agentFailure.data) Logger().debug("Agent returned \(SSHAgent.ResponseType.agentFailure.debugDescription)") + shouldClose = true } let full = OpenSSHKeyWriter().lengthAndData(of: response) - return full + return (full, shouldClose) } } diff --git a/SecretAgentKit/SocketController.swift b/SecretAgentKit/SocketController.swift index 0c916c3..4fd3b0f 100644 --- a/SecretAgentKit/SocketController.swift +++ b/SecretAgentKit/SocketController.swift @@ -4,8 +4,9 @@ import OSLog public class SocketController { private var fileHandle: FileHandle? + private var fileHandleBag: Set = [] private var port: SocketPort? - public var handler: ((FileHandleReader, FileHandleWriter) -> Void)? + public var handler: ((FileHandleReader, FileHandleWriter) -> Bool)? public init(path: String) { Logger().debug("Socket controller setting up at \(path)") @@ -52,7 +53,7 @@ public class SocketController { @objc func handleConnectionAccept(notification: Notification) { Logger().debug("Socket controller accepted connection") guard let new = notification.userInfo?[NSFileHandleNotificationFileHandleItem] as? FileHandle else { return } - handler?(new, new) + _ = handler?(new, new) new.waitForDataInBackgroundAndNotify() fileHandle?.acceptConnectionInBackgroundAndNotify(forModes: [RunLoop.current.currentMode!]) } @@ -61,7 +62,15 @@ public class SocketController { 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") - handler?(new, new) + let shouldClose = handler?(new, new) ?? false + if shouldClose { + fileHandleBag.remove(new) + Logger().debug("Closing handle") + } else if !fileHandleBag.contains(new) { + fileHandleBag.insert(new) + new.waitForDataInBackgroundAndNotify() + Logger().debug("Keeping handle open") + } } }