Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

wait for new packets on the agent socket after the handler is invoked… #267

Merged
merged 5 commits into from
Jan 31, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions Sources/Packages/Sources/SecretAgentKit/Agent.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
13 changes: 10 additions & 3 deletions Sources/Packages/Sources/SecretAgentKit/SocketController.swift
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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!])
}
Expand All @@ -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")
}
}

}