receiver: wait for batched FDs instead of returning MismatchedCount

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: giuseppe

src/transport.rs

@@ -55,6 +55,7 @@ impl UnixSocketTransport {
                 stream,
                 buffer: Vec::new(),
                 fd_queue: VecDeque::new(),
+                pending_message: None,
             },
         )
     }
@@ -262,6 +263,8 @@ pub struct Receiver {
     stream: Arc<TokioUnixStream>,
     buffer: Vec<u8>,
     fd_queue: VecDeque<OwnedFd>,
+    /// A fully parsed JSON message waiting for its FDs to arrive.
+    pending_message: Option<(serde_json::Value, usize)>,
 }
 
 impl Receiver {
@@ -275,7 +278,18 @@ impl Receiver {
                 return Ok(message);
             }
 
-            self.read_more_data().await?;
+            match self.read_more_data().await {
+                Err(Error::ConnectionClosed) if self.pending_message.is_some() => {
+                    // Connection closed while waiting for FDs — per spec
+                    // Section 5, Step 4 this is a Mismatched Count error.
+                    let (_, fd_count) = self.pending_message.take().unwrap();
+                    return Err(Error::MismatchedCount {
+                        expected: fd_count,
+                        found: self.fd_queue.len(),
+                    });
+                }
+                other => other?,
+            }
         }
     }
 
@@ -300,7 +314,47 @@ impl Receiver {
         }
     }
 
+    /// Build a `MessageWithFds` by draining `fd_count` FDs from the queue.
+    fn build_message(
+        fd_queue: &mut VecDeque<OwnedFd>,
+        value: serde_json::Value,
+        fd_count: usize,
+    ) -> Result<MessageWithFds> {
+        let fds: Vec<OwnedFd> = fd_queue.drain(..fd_count).collect();
+        let message = JsonRpcMessage::from_json_value(value)?;
+        Ok(MessageWithFds::new(message, fds))
+    }
+
     fn try_parse_message(&mut self) -> Result<Option<MessageWithFds>> {
+        // Check if we have a pending message waiting for FDs.
+        // While a message is pending, all subsequent message parsing is
+        // blocked — even messages needing 0 FDs.  This preserves FIFO
+        // ordering on the Unix socket: FDs queued after the pending
+        // message's FDs belong to later messages and must not be
+        // consumed early.
+        if let Some((value, fd_count)) = self
+            .pending_message
+            .take_if(|(_, c)| self.fd_queue.len() >= *c)
+        {
+            return Ok(Some(Self::build_message(
+                &mut self.fd_queue,
+                value,
+                fd_count,
+            )?));
+        } else if let Some((_, fd_count)) = &self.pending_message {
+            // Not enough FDs yet.  Per the spec (Section 5, Step 4),
+            // if the buffer contains any non-whitespace byte the sender
+            // has started the next message before delivering all FDs for
+            // the current one — that is a fatal protocol violation.
+            if self.buffer.iter().any(|&b| !b.is_ascii_whitespace()) {
+                return Err(Error::MismatchedCount {
+                    expected: *fd_count,
+                    found: self.fd_queue.len(),
+                });
+            }
+            return Ok(None);
+        }
+
         if self.buffer.is_empty() {
             return Ok(None);
         }
@@ -323,18 +377,33 @@ impl Receiver {
                 let fd_count = get_fd_count(&value);
 
                 if fd_count > self.fd_queue.len() {
-                    return Err(Error::MismatchedCount {
-                        expected: fd_count,
-                        found: self.fd_queue.len(),
-                    });
+                    // FDs may arrive across multiple recvmsg() calls when the
+                    // sender batches them.  Buffer the parsed message and let
+                    // the receive() loop read more data.
+                    //
+                    // Per the spec (Section 5, Step 4), if the buffer already
+                    // contains non-whitespace bytes the sender has started the
+                    // next message before delivering all FDs — a fatal error.
+                    if self.buffer.iter().any(|&b| !b.is_ascii_whitespace()) {
+                        return Err(Error::MismatchedCount {
+                            expected: fd_count,
+                            found: self.fd_queue.len(),
+                        });
+                    }
+                    debug!(
+                        "Message expects {} FDs but only {} available, waiting for more",
+                        fd_count,
+                        self.fd_queue.len()
+                    );
+                    self.pending_message = Some((value, fd_count));
+                    return Ok(None);
                 }
 
-                let fds: Vec<OwnedFd> = (0..fd_count)
-                    .map(|_| self.fd_queue.pop_front().unwrap())
-                    .collect();
-
-                let message = JsonRpcMessage::from_json_value(value)?;
-                Ok(Some(MessageWithFds::new(message, fds)))
+                Ok(Some(Self::build_message(
+                    &mut self.fd_queue,
+                    value,
+                    fd_count,
+                )?))
             }
             Some(Err(e)) if e.is_eof() => {
                 // Incomplete JSON - need more data

tests/integration_tests.rs

@@ -2276,3 +2276,192 @@ async fn test_fd_batching_many_fds_small_batches() -> Result<()> {
     server_handle.abort();
     Ok(())
 }
+
+/// Test that the receiver correctly waits for batched FDs from the server.
+///
+/// When the server responds with many FDs using a small batch size, the
+/// receiver may parse the JSON message before all FDs have arrived.  The
+/// receiver must buffer the parsed message and keep reading until enough
+/// FDs are available, rather than returning a MismatchedCount error.
+#[tokio::test]
+async fn test_receiver_waits_for_batched_response_fds() -> Result<()> {
+    let temp_dir = TempDir::new().unwrap();
+    let socket_path = temp_dir.path().join("test_receiver_batch.sock");
+
+    let num_fds = 5;
+
+    let listener = tokio::net::UnixListener::bind(&socket_path).unwrap();
+
+    let server_handle = tokio::spawn(async move {
+        let (stream, _) = listener.accept().await.unwrap();
+        let transport = UnixSocketTransport::new(stream);
+        let (mut sender, mut receiver) = transport.split();
+
+        // Force small batches on the server side so the client's
+        // receiver sees FDs arriving across multiple recvmsg() calls.
+        sender.set_max_fds_per_sendmsg(NonZeroUsize::new(1).unwrap());
+
+        // Read the request.
+        let request = receiver.receive().await.unwrap();
+        assert!(request.file_descriptors.is_empty());
+
+        // Build a response with many FDs.
+        let mut fds: Vec<OwnedFd> = Vec::new();
+        for i in 0..num_fds {
+            let mut temp_file = NamedTempFile::new().unwrap();
+            write!(temp_file, "response file {i}").unwrap();
+            temp_file.flush().unwrap();
+            temp_file.seek(SeekFrom::Start(0)).unwrap();
+            fds.push(temp_file.into_file().into());
+        }
+
+        let response = jsonrpc_fdpass::JsonRpcResponse::success(
+            Value::String("here are your files".to_string()),
+            Value::Number(1.into()),
+        );
+        let msg = MessageWithFds::new(JsonRpcMessage::Response(response), fds);
+        sender.send(msg).await.unwrap();
+    });
+
+    // Client side: send request, receive response with batched FDs.
+    let stream = tokio::net::UnixStream::connect(&socket_path).await.unwrap();
+    let transport = UnixSocketTransport::new(stream);
+    let (mut sender, mut receiver) = transport.split();
+
+    let request = JsonRpcRequest::new("get_files".to_string(), None, Value::Number(1.into()));
+    sender
+        .send(MessageWithFds::new(
+            JsonRpcMessage::Request(request),
+            Vec::new(),
+        ))
+        .await?;
+
+    // This is the critical part: the receiver must wait for all FDs
+    // instead of failing with MismatchedCount.
+    let response = receiver.receive().await?;
+    assert_eq!(
+        response.file_descriptors.len(),
+        num_fds,
+        "Expected {num_fds} FDs in batched response"
+    );
+
+    // Verify FD contents are correct and in order.
+    for (i, fd) in response.file_descriptors.into_iter().enumerate() {
+        let mut file = File::from(fd);
+        let mut contents = String::new();
+        file.seek(SeekFrom::Start(0)).unwrap();
+        file.read_to_string(&mut contents).unwrap();
+        assert_eq!(contents, format!("response file {i}"));
+    }
+
+    server_handle.await.unwrap();
+    Ok(())
+}
+
+/// Test that the receiver returns MismatchedCount when the sender starts a new
+/// message before delivering all FDs for the current one (protocol violation).
+#[tokio::test]
+async fn test_receiver_errors_on_next_message_before_fds() -> Result<()> {
+    let temp_dir = TempDir::new().unwrap();
+    let socket_path = temp_dir.path().join("test_next_msg_violation.sock");
+
+    let listener = tokio::net::UnixListener::bind(&socket_path).unwrap();
+
+    let server_handle = tokio::spawn(async move {
+        let (stream, _) = listener.accept().await.unwrap();
+        let transport = UnixSocketTransport::new(stream);
+        let (_sender, mut receiver) = transport.split();
+
+        // The client will claim fds but send a second message before
+        // delivering them.  We expect a MismatchedCount error.
+        match receiver.receive().await {
+            Err(jsonrpc_fdpass::Error::MismatchedCount { expected, found }) => {
+                assert_eq!(expected, 2);
+                assert_eq!(found, 0);
+            }
+            Err(e) => panic!("Expected MismatchedCount, got: {e:?}"),
+            Ok(_) => panic!("Should have failed with MismatchedCount"),
+        }
+    });
+
+    // Connect and send a message claiming 2 FDs, then immediately send
+    // a second message without delivering any FDs.
+    let stream = tokio::net::UnixStream::connect(&socket_path).await.unwrap();
+
+    use tokio::io::AsyncWriteExt;
+    let mut stream = stream;
+
+    let first = serde_json::json!({
+        "jsonrpc": "2.0",
+        "method": "need_fds",
+        "params": {},
+        "id": 1,
+        "fds": 2
+    });
+    let second = serde_json::json!({
+        "jsonrpc": "2.0",
+        "method": "violation",
+        "params": {},
+        "id": 2
+    });
+
+    // Send both messages back-to-back without any FDs.
+    let mut payload = serde_json::to_vec(&first).unwrap();
+    payload.extend_from_slice(&serde_json::to_vec(&second).unwrap());
+    stream.write_all(&payload).await.unwrap();
+    stream.flush().await.unwrap();
+
+    server_handle.await.unwrap();
+    Ok(())
+}
+
+/// Test that the receiver returns MismatchedCount when the connection is
+/// closed while waiting for batched FDs.
+#[tokio::test]
+async fn test_receiver_errors_on_close_while_pending() -> Result<()> {
+    let temp_dir = TempDir::new().unwrap();
+    let socket_path = temp_dir.path().join("test_close_pending.sock");
+
+    let listener = tokio::net::UnixListener::bind(&socket_path).unwrap();
+
+    let server_handle = tokio::spawn(async move {
+        let (stream, _) = listener.accept().await.unwrap();
+        let transport = UnixSocketTransport::new(stream);
+        let (_sender, mut receiver) = transport.split();
+
+        match receiver.receive().await {
+            Err(jsonrpc_fdpass::Error::MismatchedCount { expected, found }) => {
+                assert_eq!(expected, 3);
+                assert_eq!(found, 0);
+            }
+            Err(e) => panic!("Expected MismatchedCount, got: {e:?}"),
+            Ok(_) => panic!("Should have failed with MismatchedCount"),
+        }
+    });
+
+    // Connect, send a message claiming 3 FDs, then drop the connection.
+    let stream = tokio::net::UnixStream::connect(&socket_path).await.unwrap();
+
+    use tokio::io::AsyncWriteExt;
+    let mut stream = stream;
+
+    let msg = serde_json::json!({
+        "jsonrpc": "2.0",
+        "method": "test",
+        "params": {},
+        "id": 1,
+        "fds": 3
+    });
+
+    stream
+        .write_all(&serde_json::to_vec(&msg).unwrap())
+        .await
+        .unwrap();
+    stream.flush().await.unwrap();
+
+    // Close the connection without sending any FDs.
+    drop(stream);
+
+    server_handle.await.unwrap();
+    Ok(())
+}