Merge pull request #85 from LockInTime/t3code/transport-hardening

fix: back off failing accepts and correlate responses by id

Author: yashranaway

apps/headless/Sources/HeadlessProtocol/Protocol.swift

@@ -445,6 +445,11 @@ public struct CommandResponse: Codable, Equatable, Sendable {
     public let result: JSONValue?
     public let error: CommandError?
 
+    /// Used only when the host could not read the request well enough to know
+    /// its id. Clients treat it as "this reply is about your request even
+    /// though it is not correlated", so nothing else may use it.
+    public static let unknownRequestIdentifier = "unknown"
+
     public static func success(id: String, result: JSONValue = .object([:])) -> CommandResponse {
         CommandResponse(id: id, version: headlessProtocolVersion, ok: true, result: result, error: nil)
     }

apps/headless/Sources/HeadlessProtocol/Transport.swift

@@ -15,6 +15,7 @@ public enum LocalTransportError: Error, CustomStringConvertible {
     case connectionClosed
     case messageTooLarge
     case alreadyRunning
+    case mismatchedResponse
 
     public var description: String {
         switch self {
@@ -26,6 +27,7 @@ public enum LocalTransportError: Error, CustomStringConvertible {
         case .connectionClosed: return "Headless host closed the connection"
         case .messageTooLarge: return "Headless host message exceeded the size limit"
         case .alreadyRunning: return "Another Headless host is already using the local socket"
+        case .mismatchedResponse: return "Headless host replied to a different request"
         }
     }
 }
@@ -97,7 +99,19 @@ public final class LocalSocketClient {
 
         try writeAll(try ProtocolCodec.encodeLine(request), to: fd)
         let responseData = try readLine(from: fd)
-        return try ProtocolCodec.decodeLine(CommandResponse.self, from: responseData)
+        let response = try ProtocolCodec.decodeLine(CommandResponse.self, from: responseData)
+        // One request, one response, one connection — so a mismatched id means
+        // this reply belongs to something else. Correlating by convention was
+        // enough only while nothing ever got it wrong.
+        //
+        // `unknownRequestIdentifier` is the documented exception: the host uses
+        // it only when it could not read the request at all (peer rejected,
+        // unreadable frame), and those replies still carry the reason the
+        // caller needs to see.
+        guard response.id == request.id || response.id == CommandResponse.unknownRequestIdentifier else {
+            throw LocalTransportError.mismatchedResponse
+        }
+        return response
     }
 }
 
@@ -118,6 +132,8 @@ public final class LocalSocketServer: @unchecked Sendable {
     private let stateLock = NSLock()
     private var listeningDescriptor: Int32 = -1
     private var running = false
+    /// Roughly 30 seconds of backed-off retries before the listener gives up.
+    static let maximumAcceptFailures = 64
 
     public init(socketPath: String = LocalRuntime.socketURL.path) {
         self.socketPath = socketPath
@@ -175,13 +191,27 @@ public final class LocalSocketServer: @unchecked Sendable {
     }
 
     private func acceptLoop(handler: @escaping Handler) {
+        // A persistent accept() failure — a descriptor limit is the realistic
+        // one — used to spin this loop at full speed forever. Back off instead,
+        // and give up rather than pretend to serve a socket we cannot accept
+        // on: a host that exits is recoverable, a host that burns a core while
+        // silently refusing every agent is not.
+        var consecutiveFailures = 0
         while isRunning {
             let client = systemAccept(currentDescriptor)
             if client < 0 {
                 if !isRunning { return }
                 if errno == EINTR { continue }
+                consecutiveFailures += 1
+                if consecutiveFailures >= Self.maximumAcceptFailures {
+                    stop()
+                    return
+                }
+                let backoff = min(0.05 * Double(consecutiveFailures), 1.0)
+                Thread.sleep(forTimeInterval: backoff)
                 continue
             }
+            consecutiveFailures = 0
             clientQueue.async { [weak self] in
                 guard let self else { systemClose(client); return }
                 #if canImport(Darwin)
@@ -195,16 +225,23 @@ public final class LocalSocketServer: @unchecked Sendable {
     }
 
     private func handleClient(_ fd: Int32, handler: Handler) {
+        // Echo the request id as soon as it is known so a failure reply is
+        // still correlated. Only a request the host could not read at all
+        // falls back to the unknown-id sentinel.
+        var identifier = CommandResponse.unknownRequestIdentifier
         do {
             try configureNoSigPipe(fd: fd)
             guard try peerUserID(fd: fd) == currentUserID() else {
-                let response = CommandResponse.failure(id: "unknown", code: "PEER_DENIED", message: "Socket peer user is not authorized.")
+                let response = CommandResponse.failure(
+                    id: identifier, code: "PEER_DENIED", message: "Socket peer user is not authorized."
+                )
                 try writeAll(try ProtocolCodec.encodeLine(response), to: fd)
                 return
             }
             try configureTimeout(fd: fd, seconds: 5)
             let data = try readLine(from: fd)
             let request = try ProtocolCodec.decodeLine(CommandRequest.self, from: data)
+            identifier = request.id
             try request.validate()
             try configureTimeout(fd: fd, seconds: 125)
             // `shutdown` only signals the host's main loop and does not mutate
@@ -233,7 +270,7 @@ public final class LocalSocketServer: @unchecked Sendable {
             try writeAll(payload, to: fd)
         } catch {
             let response = CommandResponse.failure(
-                id: "unknown",
+                id: identifier,
                 code: "INVALID_REQUEST",
                 message: String(describing: error)
             )

apps/headless/Tests/HeadlessProtocolTests/ProtocolTests.swift

@@ -771,6 +771,39 @@ struct ProtocolTests {
         try expect(response.result == .object(["pong": .bool(true)]), "socket response should decode")
     }
 
+    static func rejectsMismatchedResponseIdentifier() throws {
+        try LocalRuntime.preparePrivateDirectory()
+        let socketPath = LocalRuntime.directoryURL
+            .appendingPathComponent("test-\(UUID().uuidString).sock").path
+        let server = LocalSocketServer(socketPath: socketPath)
+        // A host that answers with someone else's id is answering the wrong
+        // question. One request per connection means the client can say so.
+        try server.start { _ in
+            CommandResponse.success(id: "a-different-request", result: .object(["pong": .bool(true)]))
+        }
+        defer { server.stop() }
+        try expectThrows("a mismatched response identifier should be rejected") {
+            _ = try LocalSocketClient(socketPath: socketPath)
+                .send(CommandRequest(id: "ping-correlated", command: .ping), timeout: 2)
+        }
+        // The unknown-id sentinel stays usable, because a host that could not
+        // read the request still has to be able to explain why.
+        let sentinelPath = LocalRuntime.directoryURL
+            .appendingPathComponent("test-\(UUID().uuidString).sock").path
+        let sentinelServer = LocalSocketServer(socketPath: sentinelPath)
+        try sentinelServer.start { _ in
+            CommandResponse.failure(
+                id: CommandResponse.unknownRequestIdentifier,
+                code: "INVALID_REQUEST", message: "unreadable"
+            )
+        }
+        defer { sentinelServer.stop() }
+        let sentinel = try LocalSocketClient(socketPath: sentinelPath)
+            .send(CommandRequest(id: "ping-sentinel", command: .ping), timeout: 2)
+        try expect(!sentinel.ok, "the sentinel reply should still reach the caller")
+        try expect(sentinel.error?.code == "INVALID_REQUEST", "the sentinel reply should keep its reason")
+    }
+
     static func liveSocketCannotBeReplaced() throws {
         try LocalRuntime.preparePrivateDirectory()
         let socketPath = LocalRuntime.directoryURL
@@ -855,6 +888,7 @@ struct ProtocolTests {
             ("diagnostic services", diagnosticServices),
             ("diagnostic CLI", diagnosticCLI),
             ("local socket round-trip", localSocketRoundTrip),
+            ("response identifier correlation", rejectsMismatchedResponseIdentifier),
             ("live socket replacement protection", liveSocketCannotBeReplaced),
             ("private socket directory", serverRejectsSocketOutsidePrivateDirectory),
             ("shutdown bypasses busy request", shutdownBypassesBusyRequest),

docs/roadmap/improvements-backlog.md

@@ -66,7 +66,9 @@ richer answer and is tracked separately (§G3).
 
 **A4. Accept-loop error spin.** ([#15](https://github.com/LockInTime/headless/issues/15)) All `accept()` errors are swallowed with
 `continue` (`HP/Transport.swift:180-184`); persistent EMFILE becomes a hot
-loop. Add backoff + a fatal threshold.
+loop. ~~Add backoff + a fatal threshold.~~ **Done:** failures back off from
+50 ms to 1 s and the listener stops after 64 consecutive failures rather than
+burning a core while silently refusing every agent.
 
 **A5. `@eN` refs silently invalidated by every snapshot.** ([#16](https://github.com/LockInTime/headless/issues/16)) The `current` ref
 map is reset on each `snapshot()` (`HP/AgentRuntime.swift:376`), so a
@@ -96,8 +98,12 @@ containing `--json`, tabs, double spaces.
 
 **A7. Client never verifies response `id`.** ([#18](https://github.com/LockInTime/headless/issues/18)) Failure paths return
 `id:"unknown"` (`HP/Transport.swift:201,222`); `LocalSocketClient.send`
-doesn't check correlation. Echo the request id everywhere and assert
-client-side.
+doesn't check correlation. ~~Echo the request id everywhere and assert
+client-side.~~ **Done:** the host echoes the id as soon as it can decode one,
+so validation failures are correlated too, and the client rejects any other
+id. `CommandResponse.unknownRequestIdentifier` is the one documented
+exception, for replies where the host could not read the request at all —
+those still have to reach the caller with their reason.
 
 **A8. CDP O(n²) buffering.** ([#19](https://github.com/LockInTime/headless/issues/19)) `receiveText` rescans the whole buffer and
 `removeFirst`s per 8 KiB read (`LinuxHost/CDP.swift:229-256`); a 30 MB