From 575cf7658fb4f0cb69d799f5c73eeef895386b5f Mon Sep 17 00:00:00 2001
From: yminghua <minghuay@andrew.cmu.edu>
Date: Thu, 14 Aug 2025 15:00:26 -0700
Subject: [PATCH] fix: partial data read issue

---
 .../Sources/SecretAgentKit/Agent.swift        | 16 +++++++---
 .../SecretAgentKit/FileHandleProtocols.swift  | 30 +++++++++++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/Sources/Packages/Sources/SecretAgentKit/Agent.swift b/Sources/Packages/Sources/SecretAgentKit/Agent.swift
index 7209635..625ac1e 100644
--- a/Sources/Packages/Sources/SecretAgentKit/Agent.swift
+++ b/Sources/Packages/Sources/SecretAgentKit/Agent.swift
@@ -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
diff --git a/Sources/Packages/Sources/SecretAgentKit/FileHandleProtocols.swift b/Sources/Packages/Sources/SecretAgentKit/FileHandleProtocols.swift
index 40a2840..4bca392 100644
--- a/Sources/Packages/Sources/SecretAgentKit/FileHandleProtocols.swift
+++ b/Sources/Packages/Sources/SecretAgentKit/FileHandleProtocols.swift
@@ -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()
+    }
 }