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 { extension Agent {
public func handle(reader: FileHandleReader, writer: FileHandleWriter) { public func handle(reader: FileHandleReader, writer: FileHandleWriter) -> Bool {
Logger().debug("Agent handling new data") Logger().debug("Agent handling new data")
let data = reader.availableData let data = reader.availableData
guard !data.isEmpty else { return } guard !data.isEmpty else { return true }
let requestTypeInt = data[4] let requestTypeInt = data[4]
guard let requestType = SSHAgent.RequestType(rawValue: requestTypeInt) else { guard let requestType = SSHAgent.RequestType(rawValue: requestTypeInt) else {
writer.write(OpenSSHKeyWriter().lengthAndData(of: SSHAgent.ResponseType.agentFailure.data)) writer.write(OpenSSHKeyWriter().lengthAndData(of: SSHAgent.ResponseType.agentFailure.data))
Logger().debug("Agent returned \(SSHAgent.ResponseType.agentFailure.debugDescription)") Logger().debug("Agent returned \(SSHAgent.ResponseType.agentFailure.debugDescription)")
return return true
} }
Logger().debug("Agent handling request of type \(requestType.debugDescription)") Logger().debug("Agent handling request of type \(requestType.debugDescription)")
let subData = Data(data[5...]) 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) 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() var response = Data()
let shouldClose: Bool
do { do {
switch requestType { switch requestType {
case .requestIdentities: case .requestIdentities:
response.append(SSHAgent.ResponseType.agentIdentitiesAnswer.data) response.append(SSHAgent.ResponseType.agentIdentitiesAnswer.data)
response.append(identities()) response.append(identities())
Logger().debug("Agent returned \(SSHAgent.ResponseType.agentIdentitiesAnswer.debugDescription)") Logger().debug("Agent returned \(SSHAgent.ResponseType.agentIdentitiesAnswer.debugDescription)")
shouldClose = false
case .signRequest: case .signRequest:
let provenance = requestTracer.provenance(from: reader) let provenance = requestTracer.provenance(from: reader)
response.append(SSHAgent.ResponseType.agentSignResponse.data) response.append(SSHAgent.ResponseType.agentSignResponse.data)
response.append(try sign(data: data, provenance: provenance)) response.append(try sign(data: data, provenance: provenance))
Logger().debug("Agent returned \(SSHAgent.ResponseType.agentSignResponse.debugDescription)") Logger().debug("Agent returned \(SSHAgent.ResponseType.agentSignResponse.debugDescription)")
shouldClose = true
} }
} catch { } catch {
response.removeAll() response.removeAll()
response.append(SSHAgent.ResponseType.agentFailure.data) response.append(SSHAgent.ResponseType.agentFailure.data)
Logger().debug("Agent returned \(SSHAgent.ResponseType.agentFailure.debugDescription)") Logger().debug("Agent returned \(SSHAgent.ResponseType.agentFailure.debugDescription)")
shouldClose = true
} }
let full = OpenSSHKeyWriter().lengthAndData(of: response) let full = OpenSSHKeyWriter().lengthAndData(of: response)
return full return (full, shouldClose)
} }
} }

View File

@ -4,8 +4,9 @@ import OSLog
public class SocketController { public class SocketController {
private var fileHandle: FileHandle? private var fileHandle: FileHandle?
private var fileHandleBag: Set<FileHandle> = []
private var port: SocketPort? private var port: SocketPort?
public var handler: ((FileHandleReader, FileHandleWriter) -> Void)? public var handler: ((FileHandleReader, FileHandleWriter) -> Bool)?
public init(path: String) { public init(path: String) {
Logger().debug("Socket controller setting up at \(path)") Logger().debug("Socket controller setting up at \(path)")
@ -52,7 +53,7 @@ public class SocketController {
@objc func handleConnectionAccept(notification: Notification) { @objc func handleConnectionAccept(notification: Notification) {
Logger().debug("Socket controller accepted connection") Logger().debug("Socket controller accepted connection")
guard let new = notification.userInfo?[NSFileHandleNotificationFileHandleItem] as? FileHandle else { return } guard let new = notification.userInfo?[NSFileHandleNotificationFileHandleItem] as? FileHandle else { return }
handler?(new, new) _ = handler?(new, new)
new.waitForDataInBackgroundAndNotify() new.waitForDataInBackgroundAndNotify()
fileHandle?.acceptConnectionInBackgroundAndNotify(forModes: [RunLoop.current.currentMode!]) fileHandle?.acceptConnectionInBackgroundAndNotify(forModes: [RunLoop.current.currentMode!])
} }
@ -61,7 +62,15 @@ public class SocketController {
Logger().debug("Socket controller has new data available") Logger().debug("Socket controller has new data available")
guard let new = notification.object as? FileHandle else { return } guard let new = notification.object as? FileHandle else { return }
Logger().debug("Socket controller received new file handle") 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")
}
} }
} }