diff --git a/Sources/Packages/Sources/SecretAgentKit/Agent.swift b/Sources/Packages/Sources/SecretAgentKit/Agent.swift index 68fae7c..15c3545 100644 --- a/Sources/Packages/Sources/SecretAgentKit/Agent.swift +++ b/Sources/Packages/Sources/SecretAgentKit/Agent.swift @@ -30,20 +30,23 @@ extension Agent { /// - Parameters: /// - reader: A ``FileHandleReader`` to read the content of the request. /// - writer: A ``FileHandleWriter`` to write the response to. - public func handle(reader: FileHandleReader, writer: FileHandleWriter) { + /// - Return value: + /// - Boolean if data could be read + public func handle(reader: FileHandleReader, writer: FileHandleWriter) -> Bool { Logger().debug("Agent handling new data") let data = Data(reader.availableData) - guard data.count > 4 else { return } + guard data.count > 4 else { return false} 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) writer.write(response) + return true } func handle(requestType: SSHAgent.RequestType, data: Data, reader: FileHandleReader) -> Data { diff --git a/Sources/Packages/Sources/SecretAgentKit/SocketController.swift b/Sources/Packages/Sources/SecretAgentKit/SocketController.swift index 5334dbf..3143d99 100644 --- a/Sources/Packages/Sources/SecretAgentKit/SocketController.swift +++ b/Sources/Packages/Sources/SecretAgentKit/SocketController.swift @@ -9,7 +9,9 @@ public class SocketController { /// The active SocketPort. private var port: SocketPort? /// A handler that will be notified when a new read/write handle is available. - public var handler: ((FileHandleReader, FileHandleWriter) -> Void)? + /// False if no data could be read + public var handler: ((FileHandleReader, FileHandleWriter) -> Bool)? + /// Initializes a socket controller with a specified path. /// - Parameter path: The path to use as a socket. @@ -65,7 +67,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!]) } @@ -76,7 +78,12 @@ 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) + if((handler?(new, new)) == true) { + Logger().debug("Socket controller handled data, wait for more data") + new.waitForDataInBackgroundAndNotify() + } else { + Logger().debug("Socket controller called with empty data, socked closed") + } } }