fix(entropy): cap per-request entropy allocation to 64 KiB

Overlapping descriptors within a single chain can cause buffer.len() to
exceed the distinct guest memory backing the request. Without a bound,
handle_one() allocates a vec proportional to the inflated length, which
can reach ~4 GiB from a 17 MiB guest.

Introduce MAX_ENTROPY_BYTES (64 KiB) and clamp the allocation in
handle_one() to that limit.  Legitimate requests are unaffected since a
256-entry descriptor chain with typical page-sized buffers fits well
within the cap.

Add tests covering the capped path, the large inflated buffer path, and
the pass-through for small requests.

Signed-off-by: Nikita Kalyazin <kalyazin@amazon.com>

Author: kalyazin

CHANGELOG.md

@@ -6,6 +6,16 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to
 [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.14.4]
+
+### Fixed
+
+- [#5762](https://github.com/firecracker-microvm/firecracker/pull/5762): Cap
+  virtio-rng per-request entropy to 64 KiB. Previously, a guest could construct
+  a descriptor chain that caused Firecracker to allocate more host memory than
+  the guest actually provided, potentially leading to excessive host memory
+  consumption.
+
 ## [1.14.3]
 
 ### Fixed

src/vmm/src/devices/virtio/rng/device.rs

@@ -28,6 +28,14 @@ use crate::vstate::memory::GuestMemoryMmap;
 
 pub const ENTROPY_DEV_ID: &str = "rng";
 
+/// Maximum number of bytes `handle_one()` will serve per request.
+///
+/// Overlapping descriptors within a single chain can cause `buffer.len()` to
+/// exceed the amount of distinct guest memory actually backing the request.
+/// Capping the per-request allocation to 64 KiB keeps host memory usage
+/// bounded regardless of how the descriptor chain is constructed.
+const MAX_ENTROPY_BYTES: u32 = 64 * 1024;
+
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
 pub enum EntropyError {
     /// Error while handling an Event file descriptor: {0}
@@ -119,14 +127,19 @@ impl Entropy {
             return Ok(0);
         }
 
-        let mut rand_bytes = vec![0; self.buffer.len() as usize];
+        // Cap the number of bytes we actually generate so that the host-side
+        // allocation stays bounded even when buffer.len() is inflated by
+        // overlapping descriptors in the chain.
+        let len = std::cmp::min(self.buffer.len(), MAX_ENTROPY_BYTES);
+
+        let mut rand_bytes = vec![0; len as usize];
         rand::fill(&mut rand_bytes).inspect_err(|_| {
             METRICS.host_rng_fails.inc();
         })?;
 
-        // It is ok to unwrap here. We are writing `iovec.len()` bytes at offset 0.
+        // It is ok to unwrap here. We are writing `len` bytes at offset 0.
         self.buffer.write_all_volatile_at(&rand_bytes, 0).unwrap();
-        Ok(self.buffer.len())
+        Ok(len)
     }
 
     fn process_entropy_queue(&mut self) -> Result<(), InvalidAvailIdx> {
@@ -611,4 +624,125 @@ mod tests {
         // The rate limiter event should have processed the pending buffer as well
         assert_eq!(METRICS.entropy_bytes.count(), entropy_bytes + 128);
     }
+
+    /// Verify that handle_one() caps the host allocation to MAX_ENTROPY_BYTES
+    /// when overlapping descriptors inflate buffer.len() beyond the limit.
+    #[test]
+    fn test_handle_one_caps_overlapping_descriptors() {
+        use crate::devices::virtio::queue::VIRTQ_DESC_F_NEXT;
+        use crate::devices::virtio::test_utils::VirtQueue;
+        use crate::test_utils::single_region_mem;
+        use crate::vstate::memory::GuestAddress;
+
+        // 32 descriptors × 4 KiB = 128 KiB claimed, which exceeds MAX_ENTROPY_BYTES (64 KiB).
+        const N_DESC: u16 = 32;
+        const CHUNK: u32 = 4096;
+
+        let mem = single_region_mem(0x20000);
+        let vq = VirtQueue::new(GuestAddress(0), &mem, 256);
+        let mut queue = vq.create_queue();
+
+        let target: u64 = 0x10000;
+        for i in 0..N_DESC {
+            let flags = VIRTQ_DESC_F_WRITE | if i < N_DESC - 1 { VIRTQ_DESC_F_NEXT } else { 0 };
+            vq.dtable[i as usize].set(target, CHUNK, flags, i + 1);
+        }
+        vq.avail.ring[0].set(0);
+        vq.avail.idx.set(1);
+
+        let head = queue.pop().unwrap().unwrap();
+        // SAFETY: `mem` is a valid guest memory region and `head` is a descriptor chain
+        // obtained from the virtqueue backed by that memory.
+        let buf = unsafe { IoVecBufferMut::<256>::from_descriptor_chain(&mem, head).unwrap() };
+        // buffer.len() is inflated well past the cap.
+        assert_eq!(buf.len(), u32::from(N_DESC) * CHUNK); // 128 KiB
+
+        let mut dev = default_entropy();
+        dev.buffer = buf;
+        let bytes = dev.handle_one().unwrap();
+
+        assert_eq!(
+            bytes,
+            MAX_ENTROPY_BYTES,
+            "handle_one() must cap at MAX_ENTROPY_BYTES ({MAX_ENTROPY_BYTES}), got {bytes} for \
+             inflated buffer.len() = {}",
+            u32::from(N_DESC) * CHUNK
+        );
+    }
+
+    /// Verify that handle_one() caps a large inflated buffer (~4 GiB from
+    /// 255 overlapping descriptors) to MAX_ENTROPY_BYTES.
+    #[test]
+    fn test_handle_one_caps_large_inflated_buffer() {
+        use crate::devices::virtio::queue::VIRTQ_DESC_F_NEXT;
+        use crate::devices::virtio::test_utils::VirtQueue;
+        use crate::test_utils::single_region_mem;
+        use crate::vstate::memory::GuestAddress;
+
+        const N_DESC: u16 = 255;
+        const CHUNK: u32 = 16 * 1024 * 1024; // 16 MiB
+        const TOTAL: u64 = (N_DESC as u64) * (CHUNK as u64); // ~4 GiB
+
+        let mem = single_region_mem((CHUNK as usize) + 0x100000);
+        let vq = VirtQueue::new(GuestAddress(0), &mem, 256);
+        let mut queue = vq.create_queue();
+
+        let target: u64 = 0x80000;
+        for i in 0..N_DESC {
+            let flags = VIRTQ_DESC_F_WRITE | if i < N_DESC - 1 { VIRTQ_DESC_F_NEXT } else { 0 };
+            vq.dtable[i as usize].set(target, CHUNK, flags, i + 1);
+        }
+        vq.avail.ring[0].set(0);
+        vq.avail.idx.set(1);
+
+        let head = queue.pop().unwrap().unwrap();
+        // SAFETY: `mem` is a valid guest memory region and `head` is a descriptor chain
+        // obtained from the virtqueue backed by that memory.
+        let buf = unsafe { IoVecBufferMut::<256>::from_descriptor_chain(&mem, head).unwrap() };
+        assert_eq!(buf.len() as u64, TOTAL);
+
+        let mut dev = default_entropy();
+        dev.buffer = buf;
+        let bytes = dev.handle_one().unwrap();
+
+        assert_eq!(
+            bytes, MAX_ENTROPY_BYTES,
+            "handle_one() must cap at MAX_ENTROPY_BYTES, not allocate {} bytes",
+            TOTAL
+        );
+    }
+
+    /// Verify that a request within MAX_ENTROPY_BYTES is served in full
+    /// (the cap does not truncate legitimate small requests).
+    #[test]
+    fn test_handle_one_serves_small_request_in_full() {
+        use crate::devices::virtio::test_utils::VirtQueue;
+        use crate::test_utils::single_region_mem;
+        use crate::vstate::memory::GuestAddress;
+
+        const SIZE: u32 = 256;
+
+        let mem = single_region_mem(0x20000);
+        let vq = VirtQueue::new(GuestAddress(0), &mem, 256);
+        let mut queue = vq.create_queue();
+
+        vq.dtable[0].set(0x10000, SIZE, VIRTQ_DESC_F_WRITE, 0);
+        vq.avail.ring[0].set(0);
+        vq.avail.idx.set(1);
+
+        let head = queue.pop().unwrap().unwrap();
+        // SAFETY: `mem` is a valid guest memory region and `head` is a descriptor chain
+        // obtained from the virtqueue backed by that memory.
+        let buf = unsafe { IoVecBufferMut::<256>::from_descriptor_chain(&mem, head).unwrap() };
+        assert_eq!(buf.len(), SIZE);
+
+        let mut dev = default_entropy();
+        dev.buffer = buf;
+        let bytes = dev.handle_one().unwrap();
+
+        assert_eq!(
+            bytes, SIZE,
+            "small request ({SIZE} bytes) should be served in full, got {bytes}"
+        );
+    }
 }