receiver: wait for batched FDs instead of returning MismatchedCount

When the sender batches file descriptors across multiple sendmsg()
calls, the receiver may parse the complete JSON message before all
FDs have arrived.  Previously this returned a MismatchedCount error
immediately.

Buffer the parsed message in a pending_message field and let the
receive() loop continue calling read_more_data() until enough FDs
have been collected.

Assisted-by: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Giuseppe Scrivano <gscrivan@redhat.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 {
@@ -300,7 +303,34 @@ 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> = (0..fd_count)
+            .map(|_| fd_queue.pop_front().unwrap())
+            .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.
+        if let Some((value, fd_count)) = self.pending_message.take() {
+            if self.fd_queue.len() >= fd_count {
+                return Ok(Some(Self::build_message(
+                    &mut self.fd_queue,
+                    value,
+                    fd_count,
+                )?));
+            }
+            // Not enough FDs yet, put it back and wait for more.
+            self.pending_message = Some((value, fd_count));
+            return Ok(None);
+        }
+
         if self.buffer.is_empty() {
             return Ok(None);
         }
@@ -323,18 +353,23 @@ 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.
+                    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,85 @@ 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 {
+        if let Ok((stream, _)) = listener.accept().await {
+            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(())
+}