external/secretive
1
0
Fork 0
mirror of https://github.com/maxgoedjen/secretive.git synced 2025-08-30 17:10:56 +00:00
Code Issues Projects Releases Wiki Activity

fix: partial data read issue

Browse Source
This commit is contained in:
yminghua 2025-08-14 15:00:26 -07:00
parent b5f2dfd153
commit 575cf7658f
2 changed files with 42 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
Sources/Packages/Sources/SecretAgentKit/Agent.swift
View File

@ -37,16 +37,24 @@ extension Agent {
/// - Boolean if data could be read
@discardableResult public func handle(reader: FileHandleReader, writer: FileHandleWriter) async -> Bool {
logger.debug("Agent handling new data")
let data = Data(reader.availableData)
guard data.count > 4 else { return false}
let requestTypeInt = data[4]
let newData = reader.availableData
// If client closed the connection, availableData will be empty
guard !newData.isEmpty else { return false }
guard let message = reader.appendAndParseMessage(from: newData) else {
return true // only return true if we received something
}
guard message.count >= 1 else { return false }
let requestTypeInt = message[0]
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 true
}
logger.debug("Agent handling request of type \(requestType.debugDescription)")
let subData = Data(data[5...])
let subData = Data(message.dropFirst())
let response = await handle(requestType: requestType, data: subData, reader: reader)
writer.write(response)
return true

30
Sources/Packages/Sources/SecretAgentKit/FileHandleProtocols.swift
View File

@ -9,6 +9,8 @@ public protocol FileHandleReader: Sendable {
var fileDescriptor: Int32 { get }
/// The process ID of the process coonnected to the other end of the FileHandle.
var pidOfConnectedProcess: Int32 { get }
/// Append data to a buffer and extract the next SSH message if available.
func appendAndParseMessage(from newData: Data) -> Data?
}
@ -20,6 +22,23 @@ public protocol FileHandleWriter: Sendable {
}
// SSHMessageBuffer for reassembly
final class SSHMessageBuffer {
private var buffer = Data()
func append(_ newData: Data) {
buffer.append(newData)
}
func nextMessage() -> Data? {
guard buffer.count >= 4 else { return nil }
let length = buffer.prefix(4).withUnsafeBytes { $0.load(as: UInt32.self).bigEndian }
guard length < 1_000_000 else { return nil }
guard buffer.count >= 4 + Int(length) else { return nil }
let message = buffer.subdata(in: 4..<4+Int(length))
buffer.removeSubrange(0..<4+Int(length))
return message
}
}
extension FileHandle: FileHandleReader, FileHandleWriter {
public var pidOfConnectedProcess: Int32 {
@ -29,4 +48,15 @@ extension FileHandle: FileHandleReader, FileHandleWriter {
return pidPointer.load(as: Int32.self)
}
private static var messageBuffers = [Int32: SSHMessageBuffer]()
private static let messageBuffersLock = NSLock()
public func appendAndParseMessage(from newData: Data) -> Data? {
let fd = self.fileDescriptor
FileHandle.messageBuffersLock.lock()
let buffer = FileHandle.messageBuffers[fd] ?? SSHMessageBuffer()
FileHandle.messageBuffers[fd] = buffer
FileHandle.messageBuffersLock.unlock()
buffer.append(newData)
return buffer.nextMessage()
}
}