fix(server): close connection on malformed WebSocket frames instead of stalling (#258)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Thomas Kosiewski <tk@coder.com>

Author: ThomasK33

lua/claudecode/server/client.lua

@@ -114,10 +114,33 @@ function M.process_data(client, data, on_message, on_close, on_error, auth_token
   end
 
   while #client.buffer >= 2 do -- Minimum frame size
-    local parsed_frame, bytes_consumed = frame.parse_frame(client.buffer)
+    -- Stop if a prior frame (or a TCP error) already initiated teardown. The
+    -- handle can still be open (Close-frame write / scheduled on_close/on_error
+    -- pending) and read is not stopped, so a later TCP segment could otherwise
+    -- re-enter process_data and dispatch frames against a closing/closed client.
+    if client.state == "closing" or client.state == "closed" then
+      break
+    end
+
+    local parsed_frame, bytes_consumed, close_code = frame.parse_frame(client.buffer)
 
     if not parsed_frame then
-      break
+      if close_code then
+        -- Fatal protocol violation: send the RFC 6455 Close frame, drop the
+        -- offending bytes (so they are not re-parsed forever, which previously
+        -- wedged the connection), and tear the connection down.
+        --
+        -- on_error runs the synchronous teardown (tcp.lua _disconnect_client ->
+        -- _remove_client -> tcp_handle:close()), which would cancel the pending
+        -- Close-frame write if it ran first. Passing it as close_client's on_done
+        -- runs it from the write callback, i.e. only after the status code has
+        -- been flushed, so the peer actually receives the 1002/1007/1009 close.
+        client.buffer = ""
+        M.close_client(client, close_code, "Protocol error", function()
+          on_error(client, "WebSocket protocol error in frame")
+        end)
+      end
+      break -- No close_code means the frame is incomplete; wait for more bytes.
     end
 
     -- Frame validation is now handled entirely within frame.parse_frame.
@@ -155,18 +178,35 @@ function M.process_data(client, data, on_message, on_close, on_error, auth_token
       vim.schedule(function()
         on_close(client, code, reason)
       end)
+      -- A CLOSE is terminal: drop any bytes that followed it (a frame after CLOSE
+      -- is a protocol violation) and stop iterating, so we neither echo a PONG
+      -- after our own Close nor dispatch on_message after on_close. Matches the
+      -- other termination branches.
+      client.buffer = ""
+      break
     elseif parsed_frame.opcode == frame.OPCODE.PING then
       local pong_frame = frame.create_pong_frame(parsed_frame.payload)
       client.tcp_handle:write(pong_frame)
     elseif parsed_frame.opcode == frame.OPCODE.PONG then
       client.last_pong = vim.loop.now()
     elseif parsed_frame.opcode == frame.OPCODE.CONTINUATION then
-      -- Continuation frame - for simplicity, we don't support fragmentation
-      on_error(client, "Fragmented messages not supported")
-      M.close_client(client, 1003, "Unsupported data")
+      -- Continuation frame - for simplicity, we don't support fragmentation.
+      -- Run on_error from close_client's write callback (so the Close frame is
+      -- flushed before teardown), then clear the buffer and break so we stop
+      -- iterating over a connection that is being torn down. See the fatal-frame
+      -- branch above for the rationale.
+      M.close_client(client, 1003, "Unsupported data", function()
+        on_error(client, "Fragmented messages not supported")
+      end)
+      client.buffer = ""
+      break
     else
-      on_error(client, "Unknown WebSocket opcode: " .. parsed_frame.opcode)
-      M.close_client(client, 1002, "Protocol error")
+      local opcode = parsed_frame.opcode
+      M.close_client(client, 1002, "Protocol error", function()
+        on_error(client, "Unknown WebSocket opcode: " .. opcode)
+      end)
+      client.buffer = ""
+      break
     end
   end
 end
@@ -204,26 +244,59 @@ end
 ---@param client WebSocketClient The client object
 ---@param code number|nil Close code (default: 1000)
 ---@param reason string|nil Close reason
-function M.close_client(client, code, reason)
+---@param on_done function|nil Optional callback run after the Close frame has been
+---  written (or once the connection is already gone). Protocol-violation paths use
+---  it to defer error/teardown until the status code has actually been flushed.
+function M.close_client(client, code, reason, on_done)
   if client.state == "closed" or client.state == "closing" then
+    if on_done then
+      vim.schedule(on_done)
+    end
     return
   end
 
   code = code or 1000
   reason = reason or ""
 
+  -- Defensive guard for callers that don't gate on client.state: if the
+  -- underlying handle is already closing (e.g. a TCP error path closed it),
+  -- writing a Close frame would target a dead handle and never reach the peer.
+  -- Just mark the client closed in that case. Mirrors the is_closing() check
+  -- in tcp.lua's _remove_client.
+  if client.tcp_handle:is_closing() then
+    client.state = "closed"
+    if on_done then
+      vim.schedule(on_done)
+    end
+    return
+  end
+
   if client.handshake_complete then
     local close_frame = frame.create_close_frame(code, reason)
     client.tcp_handle:write(close_frame, function()
       client.state = "closed"
-      client.tcp_handle:close()
+      if not client.tcp_handle:is_closing() then
+        client.tcp_handle:close()
+      end
+      -- Run any teardown only after the Close frame has been flushed, so a
+      -- caller's synchronous handle close cannot cancel the pending write. libuv
+      -- always invokes this callback (on success or error), so on_done is never
+      -- dropped.
+      if on_done then
+        on_done()
+      end
     end)
+    -- Mark as "closing" while the Close frame write/teardown is in flight. The
+    -- write callback transitions to "closed". Do not clobber the "closed" state
+    -- set synchronously on the not-handshake-complete branch below.
+    client.state = "closing"
   else
     client.state = "closed"
     client.tcp_handle:close()
+    if on_done then
+      vim.schedule(on_done)
+    end
   end
-
-  client.state = "closing"
 end
 
 ---Check if a client connection is alive

lua/claudecode/server/frame.lua

@@ -22,9 +22,15 @@ M.OPCODE = {
 ---@field payload string Frame payload data
 
 ---Parse a WebSocket frame from binary data
+---
+---Failure is disambiguated via the optional third return value:
+---  * incomplete input ("need more bytes") returns `nil, 0` with NO third value
+---  * a fatal protocol violation returns `nil, 0, <close_code>` where the close
+---    code is the RFC 6455 status to send in the Close frame (1002/1007/1009)
 ---@param data string The binary frame data
 ---@return WebSocketFrame|nil frame The parsed frame, or nil if incomplete/invalid
 ---@return number bytes_consumed Number of bytes consumed from input
+---@return number|nil close_code WebSocket close code for fatal protocol violations
 function M.parse_frame(data)
   if type(data) ~= "string" then
     return nil, 0
@@ -65,18 +71,18 @@ function M.parse_frame(data)
   }
 
   if not valid_opcodes[opcode] then
-    return nil, 0 -- Invalid opcode
+    return nil, 0, 1002 -- Invalid opcode (protocol error)
   end
 
   -- Check for reserved bits (must be 0)
   if rsv1 or rsv2 or rsv3 then
-    return nil, 0 -- Protocol error
+    return nil, 0, 1002 -- Reserved bits set (protocol error)
   end
 
   -- Control frames must have fin=1 and payload ≤ 125 (RFC 6455 Section 5.5)
   if opcode >= M.OPCODE.CLOSE then
     if not fin or payload_len > 125 then
-      return nil, 0 -- Protocol violation
+      return nil, 0, 1002 -- Control frame fragmented or oversized (protocol error)
     end
   end
 
@@ -103,13 +109,13 @@ function M.parse_frame(data)
 
     -- Prevent extremely large payloads (DOS protection)
     if actual_payload_len > 100 * 1024 * 1024 then -- 100MB limit
-      return nil, 0
+      return nil, 0, 1009 -- Message too big
     end
   end
 
   -- Additional payload length validation
   if actual_payload_len < 0 then
-    return nil, 0 -- Invalid negative length
+    return nil, 0, 1002 -- Invalid negative length (protocol error)
   end
 
   -- Read mask if present
@@ -138,19 +144,19 @@ function M.parse_frame(data)
 
   -- Validate text frame payload is valid UTF-8
   if opcode == M.OPCODE.TEXT and not utils.is_valid_utf8(payload) then
-    return nil, 0 -- Invalid UTF-8 in text frame
+    return nil, 0, 1007 -- Invalid UTF-8 in text frame (invalid payload data)
   end
 
   -- Basic validation for close frame payload
   if opcode == M.OPCODE.CLOSE and actual_payload_len > 0 then
     if actual_payload_len == 1 then
-      return nil, 0 -- Close frame with 1 byte payload is invalid
+      return nil, 0, 1002 -- Close frame with 1 byte payload is invalid (protocol error)
     end
     -- Allow most close codes for compatibility, only validate UTF-8 for reason text
     if actual_payload_len > 2 then
       local reason = payload:sub(3)
       if not utils.is_valid_utf8(reason) then
-        return nil, 0 -- Invalid UTF-8 in close reason
+        return nil, 0, 1007 -- Invalid UTF-8 in close reason (invalid payload data)
       end
     end
   end