Dynamically allocate receive buffer to avoid unbounded committed memory growth

Signed-off-by: Bob Weinand <bob.weinand@datadoghq.com>

Author: bwoebi

datadog-ipc-macros/src/lib.rs

@@ -319,17 +319,16 @@ fn gen_serve_fn(
                     return;
                 }
             };
-            let mut buf = vec![0u8; datadog_ipc::max_message_size() + datadog_ipc::HANDLE_SUFFIX_SIZE];
             loop {
-                let (n, fds) = match datadog_ipc::recv_raw_async(&async_fd, &mut buf).await {
+                let (buf, fds) = match datadog_ipc::recv_raw_async(&async_fd).await {
                     Ok(x) => x,
                     Err(e) => {
                         ::tracing::trace!("IPC serve: recv (connection closed?): {e}");
                         break;
                     }
                 };
                 let Ok((discriminant, mut req)) =
-                    datadog_ipc::codec::decode::<#enum_name>(&buf[..n])
+                    datadog_ipc::codec::decode::<#enum_name>(&buf)
                 else {
                     ::tracing::warn!("IPC serve: failed to decode request");
                     break;

datadog-ipc/src/platform/unix/sockets/mod.rs

@@ -348,12 +348,23 @@ pub type AsyncConn = AsyncFd<OwnedFd>;
 
 /// Async receive on a Tokio `AsyncFd`-wrapped IPC connection.
 ///
+/// Allocates a buffer sized to `max_message_size()` per call and returns only the received
+/// bytes (truncated), so no large buffer is held between receives.
+///
 /// Used by the server dispatch loop (generated by `#[service]` macro).
-pub async fn recv_raw_async(fd: &AsyncConn, buf: &mut [u8]) -> io::Result<(usize, Vec<OwnedFd>)> {
+pub async fn recv_raw_async(fd: &AsyncConn) -> io::Result<(Vec<u8>, Vec<OwnedFd>)> {
     loop {
         let mut guard = fd.readable().await?;
-        match guard.try_io(|inner| recvmsg_raw(inner.as_raw_fd(), buf, MsgFlags::empty())) {
-            Ok(result) => return result,
+        // SAFETY: recvmsg writes exactly the first n bytes; we truncate to n before returning,
+        // so no uninitialized bytes are ever exposed to the caller.
+        let mut buf = Vec::with_capacity(max_message_size());
+        unsafe { buf.set_len(max_message_size()) };
+        match guard.try_io(|inner| recvmsg_raw(inner.as_raw_fd(), &mut buf, MsgFlags::empty())) {
+            Ok(Ok((n, fds))) => {
+                buf.truncate(n);
+                return Ok((buf, fds));
+            }
+            Ok(Err(e)) => return Err(e),
             Err(_would_block) => continue,
         }
     }

datadog-ipc/src/platform/windows/sockets.rs

@@ -713,14 +713,19 @@ impl SeqpacketConn {
 
 /// Async receive on a Windows named pipe IPC connection.
 ///
-/// Calls `block_in_place` with a direct blocking `ReadFile` into the caller-supplied buffer,
+/// Calls `block_in_place` with a direct blocking `ReadFile` into a caller-owned buffer,
 /// bypassing mio's 4 KB internal read buffer and correctly handling messages of any size.
-pub async fn recv_raw_async(
-    conn: &AsyncConn,
-    buf: &mut [u8],
-) -> io::Result<(usize, Vec<OwnedHandle>)> {
+pub async fn recv_raw_async(conn: &AsyncConn) -> io::Result<(Vec<u8>, Vec<OwnedHandle>)> {
     let h = conn.handle.as_raw_handle() as SysHANDLE;
-    tokio::task::block_in_place(|| pipe_read(h, buf, true))
+    tokio::task::block_in_place(|| {
+        // SAFETY: ReadFile writes exactly the first n bytes; we truncate to n before returning,
+        // so no uninitialized bytes are ever exposed to the caller.
+        let mut buf = Vec::with_capacity(max_message_size() + HANDLE_SUFFIX_SIZE);
+        unsafe { buf.set_len(max_message_size() + HANDLE_SUFFIX_SIZE) };
+        let (n, handles) = pipe_read(h, &mut buf, true)?;
+        buf.truncate(n);
+        Ok((buf, handles))
+    })
 }
 
 /// Async send on a Windows named pipe IPC connection.