This commit is contained in:
Max Goedjen 2021-11-30 22:25:27 -08:00
parent 6bb9fd376f
commit b144868e8b
No known key found for this signature in database
GPG Key ID: E58C21DD77B9B8E8
2 changed files with 23 additions and 9 deletions

View File

@ -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)
}
}

View File

@ -4,8 +4,9 @@ import OSLog
public class SocketController {
private var fileHandle: FileHandle?
private var fileHandleBag: Set<FileHandle> = []
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")
}
}
}