From 9fcb0dea9e49e6224e8e2fb07be0e33cf1f171ba Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Tue, 21 Jul 2026 17:43:28 +0200 Subject: [PATCH 01/15] feat(virtq): add stateful chain byte streams Read received payloads directly into caller owned storage and require paired completion of readable/writable chains. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/benches/common/mod.rs | 8 +- src/hyperlight_common/src/virtq/buffer.rs | 54 ++ .../src/virtq/concurrency.rs | 16 +- src/hyperlight_common/src/virtq/consumer.rs | 715 ++++++++++++++---- src/hyperlight_common/src/virtq/mod.rs | 105 +-- src/hyperlight_common/src/virtq/producer.rs | 198 +++-- 6 files changed, 785 insertions(+), 311 deletions(-) diff --git a/src/hyperlight_common/benches/common/mod.rs b/src/hyperlight_common/benches/common/mod.rs index cd68a8176..77ed2c78a 100644 --- a/src/hyperlight_common/benches/common/mod.rs +++ b/src/hyperlight_common/benches/common/mod.rs @@ -219,8 +219,8 @@ where let token = pair.producer.submit(chain).unwrap(); let (recv, reply) = pair.consumer.poll(payload.len()).unwrap().unwrap(); - black_box(recv.segments().segment_count()); - pair.consumer.complete(reply).unwrap(); + black_box(recv.len()); + pair.consumer.complete(recv, reply).unwrap(); let used = pair.producer.poll().unwrap().unwrap(); debug_assert_eq!(used.token(), token); @@ -245,13 +245,13 @@ where let token = pair.producer.submit(chain).unwrap(); let (recv, reply) = pair.consumer.poll(request.len()).unwrap().unwrap(); - black_box(recv.segments().segment_count()); + black_box(recv.len()); let ReplyChain::Writable(mut writable) = reply else { panic!("expected writable reply"); }; writable.write_all(response).unwrap(); - pair.consumer.complete(writable).unwrap(); + pair.consumer.complete(recv, writable).unwrap(); let used = pair.producer.poll().unwrap().unwrap(); debug_assert_eq!(used.token(), token); diff --git a/src/hyperlight_common/src/virtq/buffer.rs b/src/hyperlight_common/src/virtq/buffer.rs index 866a80ddd..8a5ca708f 100644 --- a/src/hyperlight_common/src/virtq/buffer.rs +++ b/src/hyperlight_common/src/virtq/buffer.rs @@ -174,6 +174,34 @@ impl Segments { self.0.iter() } + /// Split off an owned byte prefix without copying payload data. + /// + /// Returns `None` and leaves `self` unchanged when `len` exceeds the + /// remaining payload length. A split within a segment creates shared + /// [`Bytes`] slices backed by the same owner. + pub fn split_to(&mut self, len: usize) -> Option { + if len > self.len() { + return None; + } + + let mut prefix = SmallVec::<[Bytes; 4]>::new(); + let mut remaining = len; + + while remaining != 0 { + let mut segment = self.0.remove(0); + if segment.len() <= remaining { + remaining -= segment.len(); + prefix.push(segment); + } else { + prefix.push(segment.split_to(remaining)); + self.0.insert(0, segment); + remaining = 0; + } + } + + Some(Self(prefix)) + } + /// Borrow this payload as a [`Buf`] cursor. pub fn as_buf(&self) -> SegmentsBuf<'_> { SegmentsBuf::new(&self.0, self.len()) @@ -459,6 +487,32 @@ mod tests { assert_eq!(cursor.chunk(), b"world"); } + #[test] + fn segments_split_to_shares_boundary_segment() { + let boundary = Bytes::from(vec![b'd', b'e', b'f']); + let boundary_ptr = boundary.as_ptr(); + let mut segments = Segments::new([ + Bytes::from_static(b"abc"), + boundary, + Bytes::from_static(b"ghi"), + ]); + + let prefix = segments.split_to(5).unwrap(); + + assert_eq!(prefix.segment_count(), 2); + assert_eq!(prefix.to_bytes().as_ref(), b"abcde"); + assert_eq!(prefix.as_slice()[1].as_ptr(), boundary_ptr); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.to_bytes().as_ref(), b"fghi"); + assert_eq!( + segments.as_slice()[0].as_ptr(), + boundary_ptr.wrapping_add(2) + ); + + assert!(segments.split_to(5).is_none()); + assert_eq!(segments.to_bytes().as_ref(), b"fghi"); + } + #[test] fn segments_into_bytes_reuses_single_segment() { let segment = Bytes::from(vec![1, 2, 3, 4]); diff --git a/src/hyperlight_common/src/virtq/concurrency.rs b/src/hyperlight_common/src/virtq/concurrency.rs index 8592b9e66..33df1cc8c 100644 --- a/src/hyperlight_common/src/virtq/concurrency.rs +++ b/src/hyperlight_common/src/virtq/concurrency.rs @@ -343,12 +343,12 @@ fn virtq_ping_pong() { } thread::yield_now(); }; - assert_eq!(recv.to_bytes().as_ref(), b"ping"); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"ping"); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable reply"); }; wc.write_all(b"pong").unwrap(); - cons.complete(wc).unwrap(); + cons.complete(recv, wc).unwrap(); }); t_prod.join().unwrap(); @@ -390,9 +390,9 @@ fn virtq_ack_only() { } thread::yield_now(); }; - assert_eq!(recv.to_bytes().as_ref(), b"ping"); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"ping"); assert!(matches!(reply, ReplyChain::Ack(_))); - cons.complete(reply).unwrap(); + cons.complete(recv, reply).unwrap(); }); t_prod.join().unwrap(); @@ -458,7 +458,7 @@ fn virtq_out_of_order_completions() { } thread::yield_now(); }; - assert_eq!(recv1.to_bytes().as_ref(), b"first"); + assert_eq!(recv1.to_bytes().unwrap().as_ref(), b"first"); let (recv2, reply2) = loop { if let Some(r) = cons.poll(1024).unwrap() { @@ -466,17 +466,17 @@ fn virtq_out_of_order_completions() { } thread::yield_now(); }; - assert_eq!(recv2.to_bytes().as_ref(), b"second"); + assert_eq!(recv2.to_bytes().unwrap().as_ref(), b"second"); let ReplyChain::Writable(second) = reply2 else { panic!("expected writable reply"); }; - cons.complete(second).unwrap(); + cons.complete(recv2, second).unwrap(); let ReplyChain::Writable(first) = reply1 else { panic!("expected writable reply"); }; - cons.complete(first).unwrap(); + cons.complete(recv1, first).unwrap(); }); t_prod.join().unwrap(); diff --git a/src/hyperlight_common/src/virtq/consumer.rs b/src/hyperlight_common/src/virtq/consumer.rs index 7fa34e500..ffc50ef9d 100644 --- a/src/hyperlight_common/src/virtq/consumer.rs +++ b/src/hyperlight_common/src/virtq/consumer.rs @@ -2,6 +2,7 @@ // Copyright 2026 The Hyperlight Authors. use alloc::vec; +use core::fmt; use bytes::Bytes; use fixedbitset::FixedBitSet; @@ -9,46 +10,162 @@ use smallvec::SmallVec; use super::*; -type WritableElems = SmallVec<[BufferElement; 2]>; - -/// Payload received from the producer, safely copied out of shared memory. +/// Stateful reader over device-readable descriptors received from the producer. /// -/// Created by [`VirtqConsumer::poll`]. Device-readable segments are eagerly -/// copied during poll using [`MemOps::read`] (volatile on the host side), so -/// accessing data requires no unsafe code and no references into shared -/// memory. Segment boundaries are preserved in [`Segments`]. -#[derive(Debug, Clone)] -pub struct RecvChain { - token: Token, - segments: Segments, +/// Reads copy directly from shared memory into caller-provided final storage. +/// The chain must be returned together with its paired [`ReplyChain`] through +/// [`VirtqConsumer::complete`] before the descriptors can be reused. +#[must_use = "dropping without completing leaks the descriptor"] +pub struct RecvChain { + state: ChainState, } -impl RecvChain { +impl fmt::Debug for RecvChain { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RecvChain") + .field("token", &self.state.token) + .field("elems", &self.state.elems) + .field("len", &self.state.total) + .field("consumed", &self.state.position) + .field("desc_index", &self.state.desc_idx) + .field("desc_offset", &self.state.desc_off) + .finish() + } +} + +impl RecvChain { + fn new(mem: M, token: Token, elems: ChainElems, len: usize) -> Self { + Self { + state: ChainState::new(mem, token, elems, len), + } + } + /// The token identifying this chain. + #[inline] pub fn token(&self) -> Token { - self.token + self.state.token() + } + + /// Total readable payload length. + #[inline] + pub fn len(&self) -> usize { + self.state.total() } - /// The chain payload as ordered byte segments. - pub fn segments(&self) -> &Segments { - &self.segments + /// Whether this chain has no readable payload. + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 } - /// Consume the chain, taking ownership of the segments. - pub fn into_segments(self) -> Segments { - self.segments + /// Number of bytes consumed by the stateful reader. + #[inline] + pub fn consumed(&self) -> usize { + self.state.position() + } + + /// Number of bytes still available to the stateful reader. + #[inline] + pub fn remaining(&self) -> usize { + self.state.remaining() } - /// Return the chain payload as contiguous bytes. + /// Read bytes sequentially across descriptor boundaries. /// - /// Returns empty [`Bytes`] when the chain has no readable buffers. - pub fn to_bytes(&self) -> Bytes { - self.segments.to_bytes() + /// Returns the number of bytes copied, which may be smaller than `buf.len()` + /// at the end of the chain. If a later memory read fails, the cursor remains + /// advanced past any earlier chunks copied by the same call. + pub fn read(&mut self, buf: &mut [u8]) -> Result { + let len = buf.len().min(self.remaining()); + let mut dst = &mut buf[..len]; + let mut read = 0; + + while !dst.is_empty() { + let Some(elem) = self.state.current_elem() else { + break; + }; + + let desc_len = elem.len as usize; + let desc_offset = self.state.desc_offset(); + let len = (desc_len - desc_offset).min(dst.len()); + let (current, rest) = dst.split_at_mut(len); + + let addr = elem + .addr + .checked_add(desc_offset as u64) + .ok_or(VirtqError::MemoryReadError)?; + + self.state + .mem + .read(addr, current) + .map_err(|_| VirtqError::MemoryReadError)?; + + self.state.advance(len); + read += len; + dst = rest; + } + + Ok(read) } - /// Consume the chain and return the payload as contiguous bytes. - pub fn into_bytes(self) -> Bytes { - self.segments.into_bytes() + /// Read exactly `buf.len()` bytes or return an error. + #[inline] + pub fn read_exact(&mut self, buf: &mut [u8]) -> Result<&mut Self, VirtqError> { + if buf.len() > self.remaining() { + return Err(VirtqError::ReceiveTooShort { + requested: buf.len(), + remaining: self.remaining(), + }); + } + + let read = self.read(buf)?; + debug_assert_eq!(read, buf.len()); + Ok(self) + } + + /// Copy the complete payload into descriptor-preserving owned segments. + /// + /// This does not change the stateful read position. Each call takes a new + /// snapshot of shared memory; callers should validate and use the returned + /// owned value rather than reading the same untrusted payload again. + pub fn to_segments(&self) -> Result { + let mut segments = SmallVec::<[Bytes; 4]>::new(); + + for elem in &self.state.elems { + let mut buf = vec![0u8; elem.len as usize]; + self.state + .mem + .read(elem.addr, &mut buf) + .map_err(|_| VirtqError::MemoryReadError)?; + segments.push(Bytes::from(buf)); + } + + Ok(Segments::from_smallvec(segments)) + } + + /// Copy the complete payload directly into one contiguous allocation. + /// + /// This does not change the stateful read position. Each call takes a new + /// snapshot of shared memory; callers should validate and use the returned + /// owned value rather than reading the same untrusted payload again. + pub fn to_bytes(&self) -> Result { + if self.is_empty() { + return Ok(Bytes::new()); + } + + let mut buf = vec![0u8; self.len()]; + let mut offset = 0; + + for elem in &self.state.elems { + let end = offset + elem.len as usize; + self.state + .mem + .read(elem.addr, &mut buf[offset..end]) + .map_err(|_| VirtqError::MemoryReadError)?; + offset = end; + } + + Ok(Bytes::from(buf)) } } @@ -62,13 +179,15 @@ pub enum ReplyChain { /// Use the `write*` methods on [`WritableChain`] to fill the /// response buffer. Writable(WritableChain), - /// Ack-only reply (for chains with only readable buffers). No response buffer. - /// Just pass back to [`VirtqConsumer::complete`] to acknowledge. + /// Ack-only reply (for chains with only readable buffers). No response + /// buffer. Pass it back with the paired [`RecvChain`] through + /// [`VirtqConsumer::complete`] to acknowledge. Ack(AckChain), } impl ReplyChain { /// The token identifying this reply. + #[inline] pub fn token(&self) -> Token { match self { ReplyChain::Writable(wc) => wc.token(), @@ -77,9 +196,10 @@ impl ReplyChain { } /// Number of bytes written (0 for Ack). + #[inline] pub fn written(&self) -> usize { match self { - ReplyChain::Writable(wc) => wc.written, + ReplyChain::Writable(wc) => wc.written(), ReplyChain::Ack(_) => 0, } } @@ -103,48 +223,44 @@ impl ReplyChain { /// ```ignore /// if let ReplyChain::Writable(mut wc) = reply { /// wc.write_all(b"response data")?; -/// consumer.complete(wc)?; +/// consumer.complete(recv, wc)?; /// } /// ``` #[must_use = "dropping without completing leaks the descriptor"] pub struct WritableChain { - mem: M, - token: Token, - elems: WritableElems, - capacity: usize, - written: usize, + state: ChainState, } impl WritableChain { - fn new(mem: M, token: Token, elems: WritableElems) -> Self { + fn new(mem: M, token: Token, elems: ChainElems) -> Self { let capacity = elems.iter().map(|elem| elem.len as usize).sum(); Self { - mem, - token, - elems, - capacity, - written: 0, + state: ChainState::new(mem, token, elems, capacity), } } /// The token identifying this writable reply. + #[inline] pub fn token(&self) -> Token { - self.token + self.state.token() } /// Total reply capacity in bytes. + #[inline] pub fn capacity(&self) -> usize { - self.capacity + self.state.total() } /// Number of bytes written so far. + #[inline] pub fn written(&self) -> usize { - self.written + self.state.position() } /// Remaining reply capacity. + #[inline] pub fn remaining(&self) -> usize { - self.capacity() - self.written() + self.state.remaining() } /// Write bytes into writable buffers, returning how many were written. @@ -152,15 +268,39 @@ impl WritableChain { /// Appends at the current write position. If `buf` is larger than the /// remaining capacity, writes as many bytes as will fit (partial write). /// Segmentation is intentionally hidden; host-side writes must go through - /// [`MemOps::write`]. + /// [`MemOps::write`]. If a later memory write fails, the cursor and written + /// length retain any earlier chunks written by the same call. /// /// # Errors /// /// - [`VirtqError::MemoryWriteError`] - underlying MemOps write failed pub fn write(&mut self, buf: &[u8]) -> Result { - let written = write_elements(&self.mem, &self.elems, self.written, buf) - .map_err(|_| VirtqError::MemoryWriteError)?; - self.written += written; + let mut src = &buf[..buf.len().min(self.remaining())]; + let mut written = 0; + + while !src.is_empty() { + let Some(elem) = self.state.current_elem() else { + break; + }; + let desc_capacity = elem.len as usize; + let desc_offset = self.state.desc_offset(); + let len = (desc_capacity - desc_offset).min(src.len()); + + let addr = elem + .addr + .checked_add(desc_offset as u64) + .ok_or(VirtqError::MemoryWriteError)?; + + self.state + .mem + .write(addr, &src[..len]) + .map_err(|_| VirtqError::MemoryWriteError)?; + + self.state.advance(len); + written += len; + src = &src[len..]; + } + Ok(written) } @@ -170,6 +310,7 @@ impl WritableChain { /// /// - [`VirtqError::ReplyTooLarge`] - buf exceeds remaining capacity /// - [`VirtqError::MemoryWriteError`] - underlying MemOps write failed + #[inline] pub fn write_all(&mut self, buf: &[u8]) -> Result<&mut Self, VirtqError> { if buf.len() > self.remaining() { return Err(VirtqError::ReplyTooLarge); @@ -185,14 +326,15 @@ impl WritableChain { /// Previously written bytes in shared memory are not zeroed; the /// `written` count is simply reset to 0. pub fn rewind(&mut self) { - self.written = 0; + self.state.rewind(); } } /// An ack-only reply for chains with no writable buffers. /// -/// No response buffer - just pass back to [`VirtqConsumer::complete`] -/// to acknowledge processing and release the descriptor. +/// No response buffer - pass it back with the paired [`RecvChain`] through +/// [`VirtqConsumer::complete`] to acknowledge processing and release the descriptor. +/// /// This wrapper keeps ack replies as a must-use completion capability instead /// of exposing a bare token that could be accidentally ignored. #[must_use = "dropping without completing leaks the descriptor"] @@ -205,6 +347,7 @@ impl AckChain { Self { token } } + #[inline] pub fn token(&self) -> Token { self.token } @@ -221,29 +364,30 @@ impl AckChain { /// let mut consumer = VirtqConsumer::new(layout, mem, notifier); /// /// // Poll and process -/// while let Some((chain, reply)) = consumer.poll(MAX_RECV_LEN)? { -/// let data = chain.to_bytes(); +/// while let Some((recv, reply)) = consumer.poll(MAX_RECV_LEN)? { +/// let data = recv.to_bytes()?; /// match reply { /// ReplyChain::Writable(mut wc) => { /// let response = handle_request(data); /// wc.write_all(&response)?; -/// consumer.complete(wc)?; +/// consumer.complete(recv, wc)?; /// } /// ReplyChain::Ack(ack) => { -/// consumer.complete(ack)?; +/// consumer.complete(recv, ack)?; /// } /// } /// } /// /// // Or defer completions /// let mut pending = Vec::new(); -/// while let Some((chain, reply)) = consumer.poll(MAX_RECV_LEN)? { -/// pending.push((process(chain), reply)); +/// while let Some((recv, reply)) = consumer.poll(MAX_RECV_LEN)? { +/// let result = process(&recv); +/// pending.push((result, recv, reply)); /// } /// -/// for (result, reply) in pending { +/// for (result, recv, reply) in pending { /// // ... complete later ... -/// consumer.complete(reply)?; +/// consumer.complete(recv, reply)?; /// } /// ``` pub struct VirtqConsumer { @@ -275,29 +419,29 @@ impl VirtqConsumer { /// Poll for a single incoming chain from the driver. /// - /// Returns a [`RecvChain`] (copied data) and a [`ReplyChain`] (writable reply - /// capacity or ack token). Both are independent owned values with no borrow - /// on the consumer. + /// Returns a stateful [`RecvChain`] reader and a [`ReplyChain`] writable + /// reply or ack capability. Both are independent owned values with no + /// borrow on the consumer, but they must be returned together through + /// [`complete`](Self::complete). /// - /// On [`VirtqError::BadChain`], [`VirtqError::PayloadTooLarge`], and - /// [`VirtqError::MemoryReadError`] the descriptor is returned to the driver - /// (completed with zero length) before the error is propagated, so a - /// rejected chain does not leak. + /// On [`VirtqError::BadChain`] and [`VirtqError::PayloadTooLarge`] the + /// descriptor is returned to the driver (completed with zero length) before + /// the error is propagated, so a rejected chain does not leak. /// /// # Arguments /// - /// * `max_recv_len` - Maximum receive payload size to copy. Payloads larger + /// * `max_recv_len` - Maximum readable payload size. Payloads larger /// than this return [`VirtqError::PayloadTooLarge`]. /// /// # Errors /// /// - [`VirtqError::BadChain`] - Descriptor chain format not recognized /// - [`VirtqError::InvalidState`] - Descriptor ID collision (driver bug) - /// - [`VirtqError::MemoryReadError`] - Failed to read chain payload from shared memory + #[allow(clippy::type_complexity)] pub fn poll( &mut self, max_recv_len: usize, - ) -> Result)>, VirtqError> { + ) -> Result, ReplyChain)>, VirtqError> { let (id, chain) = match self.inner.poll_available() { Ok(x) => x, Err(RingError::WouldBlock) => return Ok(None), @@ -341,20 +485,19 @@ impl VirtqConsumer { )); } - // Copy chain payload from shared memory - let data = match self.read_elements(readables) { - Ok(d) => d, - Err(e) => return Err(self.abort_chain(id, e)), - }; - - let chain = RecvChain { + let chain = RecvChain::new( + self.inner.mem().clone(), token, - segments: data, - }; + readables.iter().copied().collect(), + recv_len, + ); let reply = if !writables.is_empty() { - let mem = self.inner.mem().clone(); - let writable = WritableChain::new(mem, token, writables.iter().copied().collect()); + let writable = WritableChain::new( + self.inner.mem().clone(), + token, + writables.iter().copied().collect(), + ); ReplyChain::Writable(writable) } else { let ack = AckChain::new(token); @@ -364,14 +507,29 @@ impl VirtqConsumer { Ok(Some((chain, reply))) } - /// Submit a reply/ack for a received chain back to the ring. + /// Submit both halves of a received chain back to the ring. + /// + /// Consuming the [`RecvChain`] prevents further reads once its descriptors + /// can be reused by the producer. `reply` accepts both [`WritableChain`] + /// (with written byte count) and [`AckChain`] (zero-length) through + /// [`ReplyChain`]. The two halves must have matching tokens. /// - /// Accepts both [`WritableChain`] (with written byte count) and - /// [`AckChain`] (zero-length) via the [`ReplyChain`] enum. - /// Clears the inflight slot and notifies the producer if event - /// suppression allows. - pub fn complete(&mut self, reply: impl Into>) -> Result<(), VirtqError> { + /// A mismatched pair returns [`VirtqError::InvalidState`] without returning + /// either descriptor. This fails closed: the descriptors remain in flight + /// because completing either could invalidate another still-live + /// [`RecvChain`]. + pub fn complete( + &mut self, + recv: impl Into>, + reply: impl Into>, + ) -> Result<(), VirtqError> { + let recv = recv.into(); let reply = reply.into(); + + if recv.token() != reply.token() { + return Err(VirtqError::InvalidState); + } + let id = reply.token().id; let written = u32::try_from(reply.written()).map_err(|_| VirtqError::ReplyTooLarge)?; @@ -462,63 +620,111 @@ impl VirtqConsumer { Ok(()) } - /// Read readable buffer elements from shared memory into `Bytes`. - fn read_elements(&self, elems: &[BufferElement]) -> Result { - let mut segments = SmallVec::<[Bytes; 4]>::new(); - - for elem in elems { - let mut buf = vec![0u8; elem.len as usize]; - self.inner - .mem() - .read(elem.addr, &mut buf) - .map_err(|_| VirtqError::MemoryReadError)?; - segments.push(Bytes::from(buf)); + /// Reset ring and inflight state to initial values. + /// + /// Fails while a polled chain has not yet been completed, preventing a + /// live [`RecvChain`] from reading descriptors after reset and reuse. + /// + /// # Errors + /// + /// - [`VirtqError::InvalidState`] - one or more chains are still in flight + pub fn reset(&mut self) -> Result<(), VirtqError> { + if self.inflight.ones().next().is_some() { + return Err(VirtqError::InvalidState); } - Ok(Segments::from_smallvec(segments)) - } - - /// Reset ring and inflight state to initial values. - pub fn reset(&mut self) { self.inner.reset(); self.inflight.clear(); + Ok(()) } } -fn write_elements( - mem: &M, - elems: &[BufferElement], - offset: usize, - buf: &[u8], -) -> Result { - let capacity: usize = elems.iter().map(|elem| elem.len as usize).sum(); - let mut src = &buf[..buf.len().min(capacity.saturating_sub(offset))]; - let mut written = 0; - let mut skip = offset; +type ChainElems = SmallVec<[BufferElement; 4]>; - for elem in elems { - if src.is_empty() { - break; - } +struct ChainState { + mem: M, + token: Token, + elems: ChainElems, + total: usize, + position: usize, + desc_idx: usize, + desc_off: usize, +} - let elem_len = elem.len as usize; - if skip >= elem_len { - skip -= elem_len; - continue; - } +impl ChainState { + fn new(mem: M, token: Token, elems: ChainElems, total: usize) -> Self { + let mut state = Self { + mem, + token, + elems, + total, + position: 0, + desc_idx: 0, + desc_off: 0, + }; + state.rewind(); + state + } - let elem_offset = skip; - skip = 0; - let n = (elem_len - elem_offset).min(src.len()); - let addr = elem.addr + elem_offset as u64; + #[inline] + fn token(&self) -> Token { + self.token + } - mem.write(addr, &src[..n])?; + #[inline] + fn total(&self) -> usize { + self.total + } + + #[inline] + fn position(&self) -> usize { + self.position + } + + #[inline] + fn remaining(&self) -> usize { + self.total - self.position + } + + #[inline(always)] + fn desc_len(&self) -> usize { + self.elems + .get(self.desc_idx) + .map(|elem| elem.len as usize) + .unwrap_or(0) + } - written += n; - src = &src[n..]; + #[inline(always)] + fn desc_offset(&self) -> usize { + self.desc_off } - Ok(written) + #[inline(always)] + fn current_elem(&self) -> Option { + self.elems.get(self.desc_idx).copied() + } + + #[inline(always)] + fn advance(&mut self, len: usize) { + debug_assert!(len <= self.desc_len() - self.desc_off); + self.desc_off += len; + self.position += len; + + while self.current_elem().is_some() && self.desc_off == self.desc_len() { + self.desc_idx += 1; + self.desc_off = 0; + } + } + + fn rewind(&mut self) { + self.position = 0; + self.desc_off = 0; + self.desc_idx = self + .elems + .iter() + .position(|elem| elem.len != 0) + .unwrap_or(self.elems.len()); + } } impl From> for ReplyChain { @@ -536,15 +742,59 @@ impl From for ReplyChain { #[cfg(test)] mod tests { use super::*; - use crate::virtq::ring::tests::{make_producer, make_ring}; + use crate::virtq::ring::tests::{TestMem, make_producer, make_ring}; use crate::virtq::test_utils::*; fn poll_data( - consumer: &mut VirtqConsumer, - ) -> (RecvChain, ReplyChain) { + consumer: &mut VirtqConsumer, + ) -> (RecvChain, ReplyChain) { consumer.poll(1024).unwrap().unwrap() } + #[derive(Clone)] + struct FailingPayloadReadMem { + inner: TestMem, + payload_addr: u64, + payload_len: usize, + } + + // SAFETY: All operations delegate to TestMem. Reads overlapping the + // configured payload range return an error before accessing memory. + unsafe impl MemOps for FailingPayloadReadMem { + type Error = (); + + fn read(&self, addr: u64, dst: &mut [u8]) -> Result<(), Self::Error> { + let read_end = addr.saturating_add(dst.len() as u64); + let payload_end = self.payload_addr.saturating_add(self.payload_len as u64); + if addr < payload_end && self.payload_addr < read_end { + return Err(()); + } + self.inner.read(addr, dst).map_err(|err| match err {}) + } + + fn write(&self, addr: u64, src: &[u8]) -> Result<(), Self::Error> { + self.inner.write(addr, src).map_err(|err| match err {}) + } + + fn load_acquire(&self, addr: u64) -> Result { + self.inner.load_acquire(addr).map_err(|err| match err {}) + } + + fn store_release(&self, addr: u64, val: u16) -> Result<(), Self::Error> { + self.inner + .store_release(addr, val) + .map_err(|err| match err {}) + } + + unsafe fn as_slice(&self, addr: u64, len: usize) -> Result<&[u8], Self::Error> { + unsafe { self.inner.as_slice(addr, len) }.map_err(|err| match err {}) + } + + unsafe fn as_mut_slice(&self, addr: u64, len: usize) -> Result<&mut [u8], Self::Error> { + unsafe { self.inner.as_mut_slice(addr, len) }.map_err(|err| match err {}) + } + } + #[test] fn test_write_only_recv_is_empty() { let ring = make_ring(16); @@ -554,12 +804,12 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_data(&mut consumer); - assert!(recv.to_bytes().is_empty()); + assert!(recv.to_bytes().unwrap().is_empty()); assert!(matches!(reply, ReplyChain::Writable(_))); if let ReplyChain::Writable(mut wc) = reply { wc.write_all(b"response").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } } @@ -573,10 +823,10 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_data(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello"); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello"); assert!(matches!(reply, ReplyChain::Ack(_))); - consumer.complete(reply).unwrap(); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -589,7 +839,7 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_data(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello world"); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello world"); if let ReplyChain::Writable(mut wc) = reply { assert_eq!(wc.capacity(), 64); @@ -598,12 +848,78 @@ mod tests { wc.write_all(b"response").unwrap(); assert_eq!(wc.written(), 8); assert_eq!(wc.remaining(), 56); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } else { panic!("expected Writable reply for recv+reply chain"); } } + #[test] + fn test_recv_reads_across_descriptor_boundaries() { + let ring = make_ring(16); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + let mut se = producer.chain().readable(4).readable(4).build().unwrap(); + se.write_all(b"abcdefgh").unwrap(); + producer.submit(se).unwrap(); + + let (mut recv, reply) = poll_data(&mut consumer); + assert_eq!(recv.len(), 8); + assert_eq!(recv.remaining(), 8); + let mut first = [0u8; 2]; + recv.read_exact(&mut first).unwrap(); + assert_eq!(&first, b"ab"); + assert_eq!(recv.consumed(), 2); + assert_eq!(recv.remaining(), 6); + + let mut second = [0u8; 3]; + recv.read_exact(&mut second).unwrap(); + assert_eq!(&second, b"cde"); + + let mut too_long = [0u8; 4]; + assert!(matches!( + recv.read_exact(&mut too_long), + Err(VirtqError::ReceiveTooShort { + requested: 4, + remaining: 3 + }) + )); + + let mut final_buf = [0u8; 4]; + assert_eq!(recv.read(&mut final_buf).unwrap(), 3); + assert_eq!(&final_buf[..3], b"fgh"); + assert_eq!(recv.read(&mut final_buf).unwrap(), 0); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"abcdefgh"); + + consumer.complete(recv, reply).unwrap(); + } + + #[test] + fn test_poll_defers_payload_reads() { + let ring = make_ring(16); + let mem = ring.mem(); + let mut ring_producer = make_producer(&ring); + let payload_addr = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + mem.write(payload_addr, b"data").unwrap(); + + let chain = BufferChainBuilder::new() + .readable(payload_addr, 4) + .build() + .unwrap(); + ring_producer.submit_available(&chain).unwrap(); + + let guarded_mem = FailingPayloadReadMem { + inner: mem, + payload_addr, + payload_len: 4, + }; + let mut consumer = VirtqConsumer::new(ring.layout(), guarded_mem, TestNotifier::new()); + + let (recv, reply) = consumer.poll(4).unwrap().unwrap(); + assert!(matches!(recv.to_bytes(), Err(VirtqError::MemoryReadError))); + consumer.complete(recv, reply).unwrap(); + } + #[test] fn test_writable_partial_write() { let ring = make_ring(16); @@ -612,13 +928,13 @@ mod tests { let se = producer.chain().writable(8).build().unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_data(&mut consumer); + let (recv, reply) = poll_data(&mut consumer); if let ReplyChain::Writable(mut wc) = reply { let n = wc.write(b"hello world!").unwrap(); assert_eq!(n, 8); assert_eq!(wc.remaining(), 0); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } else { panic!("expected Writable"); } @@ -679,8 +995,8 @@ mod tests { // A subsequent normal exchange still round-trips end to end. let se2 = producer.chain().writable(16).build().unwrap(); producer.submit(se2).unwrap(); - let (_recv, reply) = poll_data(&mut consumer); - consumer.complete(reply).unwrap(); + let (recv, reply) = poll_data(&mut consumer); + consumer.complete(recv, reply).unwrap(); assert!(producer.poll().unwrap().is_some()); } @@ -731,13 +1047,13 @@ mod tests { let se = producer.chain().writable(16).build().unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_data(&mut consumer); + let (recv, reply) = poll_data(&mut consumer); let ReplyChain::Writable(mut wc) = reply else { panic!("expected Writable"); }; wc.write_all(b"hello").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); let used = producer.poll().unwrap().unwrap(); assert_eq!(used.to_bytes().unwrap().as_ref(), b"hello"); @@ -751,7 +1067,7 @@ mod tests { let se = producer.chain().writable(16).build().unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_data(&mut consumer); + let (recv, reply) = poll_data(&mut consumer); if let ReplyChain::Writable(mut wc) = reply { wc.write_all(b"first").unwrap(); @@ -761,7 +1077,7 @@ mod tests { assert_eq!(wc.remaining(), 16); wc.write_all(b"second").unwrap(); assert_eq!(wc.written(), 6); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } else { panic!("expected Writable"); } @@ -783,7 +1099,7 @@ mod tests { let id = ring_producer.submit_available(&chain).unwrap(); let (recv, reply) = poll_data(&mut consumer); - assert!(recv.to_bytes().is_empty()); + assert!(recv.to_bytes().unwrap().is_empty()); let ReplyChain::Writable(mut wc) = reply else { panic!("expected Writable"); @@ -791,7 +1107,7 @@ mod tests { assert_eq!(wc.capacity(), 8); wc.write_all(b"abcdefgh").unwrap(); assert_eq!(wc.written(), 8); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); let mut first = [0u8; 4]; let mut second = [0u8; 4]; @@ -805,6 +1121,43 @@ mod tests { assert_eq!(used.len, 8); } + #[test] + fn test_writable_short_write_reports_contiguous_used_length() { + let ring = make_ring(16); + let mem = ring.mem(); + let mut ring_producer = make_producer(&ring); + let base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + mem.write(base, &[0xff; 8]).unwrap(); + + let chain = BufferChainBuilder::new() + .writable(base, 4) + .writable(base + 4, 4) + .build() + .unwrap(); + let id = ring_producer.submit_available(&chain).unwrap(); + + let mut consumer = VirtqConsumer::new(ring.layout(), mem.clone(), TestNotifier::new()); + let (recv, reply) = poll_data(&mut consumer); + let ReplyChain::Writable(mut writable) = reply else { + panic!("expected writable reply"); + }; + + writable.write_all(b"abc").unwrap(); + writable.write_all(b"def").unwrap(); + assert_eq!(writable.written(), 6); + assert_eq!(writable.remaining(), 2); + + consumer.complete(recv, writable).unwrap(); + + let mut contents = [0u8; 8]; + mem.read(base, &mut contents).unwrap(); + assert_eq!(&contents, b"abcdef\xff\xff"); + + let used = ring_producer.poll_used().unwrap(); + assert_eq!(used.id, id); + assert_eq!(used.len, 6); + } + #[test] fn test_multiple_pending_replies() { let ring = make_ring(16); @@ -815,16 +1168,43 @@ mod tests { let se2 = producer.chain().writable(16).build().unwrap(); producer.submit(se2).unwrap(); - let (_e1, c1) = poll_data(&mut consumer); - let (_e2, c2) = poll_data(&mut consumer); + let (e1, c1) = poll_data(&mut consumer); + let (e2, c2) = poll_data(&mut consumer); // Complete in reverse order - consumer.complete(c2).unwrap(); - consumer.complete(c1).unwrap(); + consumer.complete(e2, c2).unwrap(); + consumer.complete(e1, c1).unwrap(); } #[test] - fn test_recv_into_bytes() { + fn test_mismatched_completion_fails_closed() { + let ring = make_ring(16); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + let mut first = producer.chain().readable(1).build().unwrap(); + first.write_all(b"a").unwrap(); + producer.submit(first).unwrap(); + + let mut second = producer.chain().readable(1).build().unwrap(); + second.write_all(b"b").unwrap(); + producer.submit(second).unwrap(); + + let (recv1, reply1) = poll_data(&mut consumer); + let (recv2, reply2) = poll_data(&mut consumer); + + assert!(matches!( + consumer.complete(recv1, reply2), + Err(VirtqError::InvalidState) + )); + assert_eq!(consumer.inflight.count_ones(..), 2); + assert!(producer.poll().unwrap().is_none()); + assert!(matches!(consumer.reset(), Err(VirtqError::InvalidState))); + + drop((recv2, reply1)); + } + + #[test] + fn test_recv_to_bytes_preserves_reader_position() { let ring = make_ring(16); let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); @@ -832,10 +1212,14 @@ mod tests { se.write_all(b"abc").unwrap(); producer.submit(se).unwrap(); - let (recv, reply) = poll_data(&mut consumer); - let data = recv.into_bytes(); + let (mut recv, reply) = poll_data(&mut consumer); + let mut first = [0u8; 1]; + recv.read_exact(&mut first).unwrap(); + let data = recv.to_bytes().unwrap(); + assert_eq!(&first, b"a"); assert_eq!(data.as_ref(), b"abc"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.consumed(), 1); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -847,13 +1231,14 @@ mod tests { let se = producer.chain().writable(16).build().unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_data(&mut consumer); + let (recv, reply) = poll_data(&mut consumer); assert!(consumer.inflight.count_ones(..) > 0); + assert!(matches!(consumer.reset(), Err(VirtqError::InvalidState))); // Complete first so we do not leak - consumer.complete(reply).unwrap(); + consumer.complete(recv, reply).unwrap(); - consumer.reset(); + consumer.reset().unwrap(); assert_eq!(consumer.inflight.count_ones(..), 0); assert_eq!(consumer.inner.num_inflight(), 0); @@ -870,13 +1255,13 @@ mod tests { let se2 = producer.chain().writable(16).build().unwrap(); producer.submit(se2).unwrap(); - let (_e1, c1) = poll_data(&mut consumer); - let (_e2, c2) = poll_data(&mut consumer); + let (e1, c1) = poll_data(&mut consumer); + let (e2, c2) = poll_data(&mut consumer); // Complete both before reset - consumer.complete(c1).unwrap(); - consumer.complete(c2).unwrap(); + consumer.complete(e1, c1).unwrap(); + consumer.complete(e2, c2).unwrap(); - consumer.reset(); + consumer.reset().unwrap(); assert_eq!(consumer.inflight.count_ones(..), 0); assert_eq!(consumer.inner.num_inflight(), 0); diff --git a/src/hyperlight_common/src/virtq/mod.rs b/src/hyperlight_common/src/virtq/mod.rs index 02b6af120..467076de0 100644 --- a/src/hyperlight_common/src/virtq/mod.rs +++ b/src/hyperlight_common/src/virtq/mod.rs @@ -43,27 +43,28 @@ //! } //! //! // Consumer (device) side - receive a chain and reply/ack it -//! if let Some((chain, reply)) = consumer.poll(max_recv_len)? { -//! let request = chain.to_bytes(); +//! if let Some((recv, reply)) = consumer.poll(max_recv_len)? { +//! let request = recv.to_bytes()?; //! match reply { //! ReplyChain::Writable(mut wc) => { //! let response = handle(request); //! wc.write_all(&response)?; -//! consumer.complete(wc)?; +//! consumer.complete(recv, wc)?; //! } //! ReplyChain::Ack(ack) => { -//! consumer.complete(ack)?; +//! consumer.complete(recv, ack)?; //! } //! } //! } //! //! // Multiple pending completions (no borrow on consumer) //! let mut pending = Vec::new(); -//! while let Some((chain, reply)) = consumer.poll(max_recv_len)? { -//! pending.push((process(chain), reply)); +//! while let Some((recv, reply)) = consumer.poll(max_recv_len)? { +//! let result = process(&recv); +//! pending.push((result, recv, reply)); //! } -//! for (result, reply) in pending { -//! consumer.complete(reply)?; +//! for (result, recv, reply) in pending { +//! consumer.complete(recv, reply)?; //! } //! ``` //! @@ -191,6 +192,8 @@ pub enum VirtqError { BadChain, #[error("Payload data too large: received {recv} bytes, limit {limit} bytes")] PayloadTooLarge { recv: usize, limit: usize }, + #[error("Receive data too short: requested {requested} bytes, only {remaining} bytes remain")] + ReceiveTooShort { requested: usize, remaining: usize }, #[error("Reply data too large for allocated buffer")] ReplyTooLarge, #[error("Internal state error")] @@ -588,7 +591,7 @@ mod tests { fn poll_received( consumer: &mut VirtqConsumer, - ) -> (RecvChain, ReplyChain) { + ) -> (RecvChain, ReplyChain) { consumer.poll(1024).unwrap().unwrap() } @@ -617,8 +620,8 @@ mod tests { // Consumer sees all requests for _ in 0..3 { - let (_recv, reply) = poll_received(&mut consumer); - consumer.complete(reply).unwrap(); + let (recv, reply) = poll_received(&mut consumer); + consumer.complete(recv, reply).unwrap(); } // All completions available @@ -650,12 +653,12 @@ mod tests { // Consumer processes requests for _ in 0..3 { - let (_recv, reply) = poll_received(&mut consumer); + let (recv, reply) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable reply"); }; wc.write_all(b"used-data").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } // Producer can drain all responses @@ -736,19 +739,19 @@ mod tests { // Consumer sees all three entries let (recv1, reply1) = poll_received(&mut consumer); - assert_eq!(recv1.to_bytes().as_ref(), b"first-ent"); - consumer.complete(reply1).unwrap(); + assert_eq!(recv1.to_bytes().unwrap().as_ref(), b"first-ent"); + consumer.complete(recv1, reply1).unwrap(); let (recv2, reply2) = poll_received(&mut consumer); - assert_eq!(recv2.to_bytes().as_ref(), b"copy-ent"); - consumer.complete(reply2).unwrap(); + assert_eq!(recv2.to_bytes().unwrap().as_ref(), b"copy-ent"); + consumer.complete(recv2, reply2).unwrap(); - let (_recv3, reply3) = poll_received(&mut consumer); + let (recv3, reply3) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply3 else { panic!("expected writable reply"); }; wc.write_all(b"resp").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv3, wc).unwrap(); // Drain completions let _ = producer.poll().unwrap().unwrap(); @@ -771,14 +774,14 @@ mod tests { // Consumer sees the data let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), token); - assert_eq!(recv.to_bytes().as_ref(), b"hello"); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello"); // Write response let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable reply"); }; wc.write_all(b"world").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); let used = producer.poll().unwrap().unwrap(); assert_eq!(used.to_bytes().unwrap().as_ref(), b"world"); } @@ -794,14 +797,14 @@ mod tests { // Consumer receives and responds let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), token); - assert_eq!(recv.to_bytes().as_ref(), b"round-trip-recv"); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"round-trip-recv"); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable reply"); }; assert!(wc.capacity() >= 128); wc.write_all(b"round-trip-rsp").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); // Producer gets the reply let used = producer.poll().unwrap().unwrap(); @@ -816,8 +819,8 @@ mod tests { let token = send_readwrite(&mut producer, b"recv-data", 64); - let (_recv, reply) = poll_received(&mut consumer); - consumer.complete(reply).unwrap(); + let (recv, reply) = poll_received(&mut consumer); + consumer.complete(recv, reply).unwrap(); let used = producer.poll().unwrap().unwrap(); assert_eq!(used.token(), token); @@ -835,13 +838,13 @@ mod tests { // Poll and hold the reply let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), token); - assert_eq!(recv.to_bytes().as_ref(), b"deferred"); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"deferred"); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable reply"); }; wc.write_all(b"deferred-used").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); let used = producer.poll().unwrap().unwrap(); assert_eq!(used.token(), token); @@ -859,24 +862,24 @@ mod tests { // Poll both let (recv1, reply1) = poll_received(&mut consumer); assert_eq!(recv1.token(), tok1); - assert_eq!(recv1.to_bytes().as_ref(), b"first"); + assert_eq!(recv1.to_bytes().unwrap().as_ref(), b"first"); let (recv2, reply2) = poll_received(&mut consumer); assert_eq!(recv2.token(), tok2); - assert_eq!(recv2.to_bytes().as_ref(), b"second"); + assert_eq!(recv2.to_bytes().unwrap().as_ref(), b"second"); // Complete second first (out of order) let ReplyChain::Writable(mut wc2) = reply2 else { panic!("expected writable"); }; wc2.write_all(b"resp2").unwrap(); - consumer.complete(wc2).unwrap(); + consumer.complete(recv2, wc2).unwrap(); let ReplyChain::Writable(mut wc1) = reply1 else { panic!("expected writable"); }; wc1.write_all(b"resp1").unwrap(); - consumer.complete(wc1).unwrap(); + consumer.complete(recv1, wc1).unwrap(); let used1 = producer.poll().unwrap().unwrap(); let used2 = producer.poll().unwrap().unwrap(); @@ -924,8 +927,8 @@ mod tests { // Consumer acks all entries while let Some(result) = consumer.poll(1024).unwrap() { - let (_, reply) = result; - consumer.complete(reply).unwrap(); + let (recv, reply) = result; + consumer.complete(recv, reply).unwrap(); } // Reclaim should free ring slots without losing data @@ -945,12 +948,12 @@ mod tests { let tok = send_readwrite(&mut producer, b"request", 64); // Consumer processes and writes response - let (_, reply) = poll_received(&mut consumer); + let (recv, reply) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable"); }; wc.write_all(b"response-data").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); // Reclaim buffers the reply (doesn't discard it) let count = producer.reclaim().unwrap(); @@ -973,18 +976,18 @@ mod tests { let _tok_ro2 = send_readonly(&mut producer, b"log2"); // Consumer processes all 3 - let (_, reply1) = poll_received(&mut consumer); - consumer.complete(reply1).unwrap(); // ack RO + let (recv1, reply1) = poll_received(&mut consumer); + consumer.complete(recv1, reply1).unwrap(); // ack RO - let (_, reply2) = poll_received(&mut consumer); + let (recv2, reply2) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply2 else { panic!("expected writable"); }; wc.write_all(b"result").unwrap(); - consumer.complete(wc).unwrap(); // complete RW + consumer.complete(recv2, wc).unwrap(); // complete RW - let (_, reply3) = poll_received(&mut consumer); - consumer.complete(reply3).unwrap(); // ack RO + let (recv3, reply3) = poll_received(&mut consumer); + consumer.complete(recv3, reply3).unwrap(); // ack RO // Reclaim all 3 - RO completions are discarded, only RW is buffered let count = producer.reclaim().unwrap(); @@ -1008,15 +1011,15 @@ mod tests { send_readonly(&mut producer, b"x"); let tok_rw = send_readwrite(&mut producer, b"y", 64); - let (_, reply1) = poll_received(&mut consumer); - consumer.complete(reply1).unwrap(); + let (recv1, reply1) = poll_received(&mut consumer); + consumer.complete(recv1, reply1).unwrap(); - let (_, reply2) = poll_received(&mut consumer); + let (recv2, reply2) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply2 else { panic!("expected writable"); }; wc.write_all(b"reply").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv2, wc).unwrap(); // poll() consumes first recv directly from ring let used1 = producer.poll().unwrap().unwrap(); @@ -1041,8 +1044,8 @@ mod tests { // Submit and complete a ReadOnly recv let tok_old = send_readonly(&mut producer, b"log"); - let (_, reply) = poll_received(&mut consumer); - consumer.complete(reply).unwrap(); + let (recv, reply) = poll_received(&mut consumer); + consumer.complete(recv, reply).unwrap(); let count = producer.reclaim().unwrap(); assert_eq!(count, 1); @@ -1057,12 +1060,12 @@ mod tests { ); // Complete the ReadWrite recv - let (_, reply) = poll_received(&mut consumer); + let (recv, reply) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable"); }; wc.write_all(b"result").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); // Poll returns only the RW reply (RO was discarded by reclaim) let used = producer.poll().unwrap().unwrap(); @@ -1087,8 +1090,8 @@ mod tests { // Consumer acks all while let Some(result) = consumer.poll(1024).unwrap() { - let (_, reply) = result; - consumer.complete(reply).unwrap(); + let (recv, reply) = result; + consumer.complete(recv, reply).unwrap(); } // Reclaim frees ring slots; empty completions are discarded diff --git a/src/hyperlight_common/src/virtq/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index 04f59e0ab..0443c0ec6 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -739,10 +739,12 @@ pub struct SendChain { // public lifetime, so these `expect`s cannot fail. #[allow(clippy::expect_used)] impl SendChain { + #[inline(always)] fn chain(&self) -> &BufferChain { self.chain.as_ref().expect("SendChain missing BufferChain") } + #[inline(always)] fn chain_mut(&mut self) -> &mut BufferChain { self.chain.as_mut().expect("SendChain missing BufferChain") } @@ -762,22 +764,26 @@ impl SendChain { Inflight { token, chain } } - /// Number of producer-written readable segments in this chain. - pub fn segment_count(&self) -> usize { + /// Number of producer-written readable descriptors in this chain. + #[inline] + pub fn desc_count(&self) -> usize { self.chain().readables().len() } /// Total producer-written readable capacity in bytes. + #[inline] pub fn capacity(&self) -> usize { self.rd_capacity } /// Number of producer-written readable bytes written so far. + #[inline] pub fn written(&self) -> usize { self.rd_written } /// Remaining producer-written readable capacity. + #[inline] pub fn remaining(&self) -> usize { self.capacity() - self.written() } @@ -787,14 +793,15 @@ impl SendChain { /// Appends at the current aggregate write position and scatters across /// readable segments in chain order. Uses [`MemOps::write`] (volatile on /// host side). If `buf` is larger than the remaining capacity, writes as - /// many bytes as will fit. + /// many bytes as will fit. If a later memory write fails, the cursor and + /// written length retain any earlier chunks written by the same call. /// /// # Errors /// /// - [`VirtqError::NoPayloadSegment`] - no readable buffer allocated /// - [`VirtqError::MemoryWriteError`] - underlying write failed pub fn write(&mut self, buf: &[u8]) -> Result { - if self.segment_count() == 0 { + if self.desc_count() == 0 { return Err(VirtqError::NoPayloadSegment); } @@ -803,40 +810,34 @@ impl SendChain { let mut remaining = &buf[..buf.len().min(self.remaining())]; let mut written = 0; - let SendChain { - mem, - chain, - rd_caps, - .. - } = self; - - let readables = chain - .as_mut() - .expect("SendChain missing BufferChain") - .readables_mut(); - - for (readable, &cap) in readables.iter_mut().zip(rd_caps.iter()) { + for index in 0..self.rd_caps.len() { if remaining.is_empty() { break; } - let written_len = readable.len as usize; - let free = cap - written_len; - if free == 0 { + let cap = self.rd_caps[index]; + let elem = self.chain().readables()[index]; + let desc_off = elem.len as usize; + let len = (cap - desc_off).min(remaining.len()); + if len == 0 { continue; } - let n = free.min(remaining.len()); - let addr = readable.addr + written_len as u64; - mem.write(addr, &remaining[..n]) + let addr = elem + .addr + .checked_add(desc_off as u64) + .ok_or(VirtqError::MemoryWriteError)?; + + self.mem + .write(addr, &remaining[..len]) .map_err(|_| VirtqError::MemoryWriteError)?; - readable.len += n as u32; - written += n; - remaining = &remaining[n..]; + self.chain_mut().readables_mut()[index].len += len as u32; + self.rd_written += len; + written += len; + remaining = &remaining[len..]; } - self.rd_written += written; Ok(written) } @@ -851,8 +852,9 @@ impl SendChain { /// - [`VirtqError::PayloadTooLarge`] - buf exceeds remaining capacity /// - [`VirtqError::NoPayloadSegment`] - no readable buffer allocated /// - [`VirtqError::MemoryWriteError`] - underlying write failed + #[inline] pub fn write_all(&mut self, buf: &[u8]) -> Result<&mut Self, VirtqError> { - if self.segment_count() == 0 { + if self.desc_count() == 0 { return Err(VirtqError::NoPayloadSegment); } @@ -995,7 +997,7 @@ mod tests { fn poll_received( consumer: &mut VirtqConsumer, - ) -> (RecvChain, ReplyChain) { + ) -> (RecvChain, ReplyChain) { consumer.poll(1024).unwrap().unwrap() } @@ -1049,7 +1051,7 @@ mod tests { let (producer, _consumer, _notifier) = make_test_producer(&ring); let se = producer.chain().readable(16).writable(32).build().unwrap(); - assert_eq!(se.segment_count(), 1); + assert_eq!(se.desc_count(), 1); assert_eq!(se.capacity(), 16); } @@ -1072,11 +1074,38 @@ mod tests { let token = producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), token); - assert_eq!(recv.to_bytes().as_ref(), b"hello world"); - assert_eq!(recv.segments().segment_count(), 2); - assert_eq!(recv.segments().as_slice()[0].as_ref(), b"hello"); - assert_eq!(recv.segments().as_slice()[1].as_ref(), b" world"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello world"); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.as_slice()[0].as_ref(), b"hello"); + assert_eq!(segments.as_slice()[1].as_ref(), b" world"); + consumer.complete(recv, reply).unwrap(); + } + + #[test] + fn test_chain_multi_readable_appends_across_calls() { + let ring = make_ring(16); + let layout = ring.layout(); + let mem = ring.mem(); + let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + let pool = TestPool::new_with_max_alloc_len(pool_base, 0x8000, 4); + let notifier = TestNotifier::new(); + let mut producer = VirtqProducer::new(layout, mem.clone(), notifier.clone(), pool); + let mut consumer = VirtqConsumer::new(layout, mem, notifier); + + let mut send = producer.chain().readable(8).build().unwrap(); + send.write_all(b"abc").unwrap(); + send.write_all(b"def").unwrap(); + assert_eq!(send.written(), 6); + assert_eq!(send.remaining(), 2); + + producer.submit(send).unwrap(); + let (recv, reply) = poll_received(&mut consumer); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.as_slice()[0].as_ref(), b"abcd"); + assert_eq!(segments.as_slice()[1].as_ref(), b"ef"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1092,7 +1121,7 @@ mod tests { let mut se = producer.chain().readable(10).writable(32).build().unwrap(); - assert_eq!(se.segment_count(), 3); + assert_eq!(se.desc_count(), 3); assert_eq!(se.capacity(), 10); se.write_all(b"abcdefghij").unwrap(); @@ -1101,12 +1130,13 @@ mod tests { let token = producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), token); - assert_eq!(recv.to_bytes().as_ref(), b"abcdefghij"); - assert_eq!(recv.segments().segment_count(), 3); - assert_eq!(recv.segments().as_slice()[0].as_ref(), b"abcd"); - assert_eq!(recv.segments().as_slice()[1].as_ref(), b"efgh"); - assert_eq!(recv.segments().as_slice()[2].as_ref(), b"ij"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"abcdefghij"); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 3); + assert_eq!(segments.as_slice()[0].as_ref(), b"abcd"); + assert_eq!(segments.as_slice()[1].as_ref(), b"efgh"); + assert_eq!(segments.as_slice()[2].as_ref(), b"ij"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1134,14 +1164,14 @@ mod tests { let se = producer.chain().writable(10).build().unwrap(); let token = producer.submit(se).unwrap(); - let (_recv, reply) = poll_received(&mut consumer); + let (recv, reply) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable reply"); }; assert_eq!(wc.capacity(), 10); wc.write_all(b"abcdefghij").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); let used = producer.poll().unwrap().unwrap(); assert_eq!(used.token(), token); @@ -1173,8 +1203,8 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"headbody"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"headbody"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1196,9 +1226,10 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"headbody"); - assert_eq!(recv.segments().segment_count(), 2); - consumer.complete(reply).unwrap(); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.to_bytes().as_ref(), b"headbody"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1212,9 +1243,10 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"headbody"); - assert_eq!(recv.segments().segment_count(), 2); - consumer.complete(reply).unwrap(); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.to_bytes().as_ref(), b"headbody"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1225,14 +1257,14 @@ mod tests { let se = producer.chain().writable(5).writable(6).build().unwrap(); let token = producer.submit(se).unwrap(); - let (_recv, reply) = poll_received(&mut consumer); + let (recv, reply) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable reply"); }; assert_eq!(wc.capacity(), 11); wc.write_all(b"hello world").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); let used = producer.poll().unwrap().unwrap(); assert_eq!(used.token(), token); @@ -1251,13 +1283,13 @@ mod tests { let se = producer.chain().writable(5).writable(6).build().unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_received(&mut consumer); + let (recv, reply) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable reply"); }; wc.write_all(b"hello wo").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); let used = producer.poll().unwrap().unwrap(); let segments = used.segments().unwrap(); @@ -1275,8 +1307,8 @@ mod tests { let se = producer.chain().writable(5).writable(6).build().unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_received(&mut consumer); - consumer.complete(reply).unwrap(); + let (recv, reply) = poll_received(&mut consumer); + consumer.complete(recv, reply).unwrap(); let used = producer.poll().unwrap().unwrap(); let segments = used.segments().unwrap(); @@ -1329,8 +1361,8 @@ mod tests { let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), tok); - assert_eq!(recv.to_bytes().as_ref(), b"hello world"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello world"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1349,8 +1381,8 @@ mod tests { let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), tok); - assert_eq!(recv.to_bytes().as_ref(), b"hello world"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello world"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1365,8 +1397,8 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello wo"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello wo"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1384,8 +1416,8 @@ mod tests { let _tok = producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1404,8 +1436,8 @@ mod tests { let _tok = producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1431,7 +1463,7 @@ mod tests { let producer = VirtqProducer::new(layout, mem, notifier, pool); let mut se = producer.chain().readable(8).build().unwrap(); - assert_eq!(se.segment_count(), 2); + assert_eq!(se.desc_count(), 2); assert!(matches!( se.with_seg(2, |_| Ok::(0)), Err(VirtqError::NoPayloadSegment) @@ -1575,12 +1607,12 @@ mod tests { assert_eq!(notifier.notification_count(), initial_count + 1); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"first"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"first"); + consumer.complete(recv, reply).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"second"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"second"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1629,11 +1661,11 @@ mod tests { let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), token); - assert!(recv.to_bytes().is_empty()); + assert!(recv.to_bytes().unwrap().is_empty()); if let ReplyChain::Writable(mut wc) = reply { wc.write_all(b"filled-by-consumer").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } else { panic!("expected Writable"); } @@ -1655,9 +1687,9 @@ mod tests { let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), token); - assert_eq!(recv.to_bytes().as_ref(), b"fire-and-forget"); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"fire-and-forget"); assert!(matches!(reply, ReplyChain::Ack(_))); - consumer.complete(reply).unwrap(); + consumer.complete(recv, reply).unwrap(); let used = producer.poll().unwrap().unwrap(); assert!(matches!(used, UsedChain::Ack(t) if t == token)); @@ -1673,10 +1705,10 @@ mod tests { let token = producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"request data"); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"request data"); if let ReplyChain::Writable(mut wc) = reply { wc.write_all(b"response data").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } else { panic!("expected Writable"); } @@ -1702,10 +1734,10 @@ mod tests { se.write_all(b"request data").unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_received(&mut consumer); + let (recv, reply) = poll_received(&mut consumer); if let ReplyChain::Writable(mut wc) = reply { wc.write_all(b"response data").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } else { panic!("expected Writable"); } @@ -1763,8 +1795,8 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello"); + consumer.complete(recv, reply).unwrap(); let _ = producer.poll().unwrap().unwrap(); // Now reset From 824025ce1b6037994048baa1945f7bdf7f7ea6d5 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Wed, 22 Jul 2026 14:52:30 +0200 Subject: [PATCH 02/15] refactor(virtq): remove reset api and harden allocation rollback Make pool restoration transactional. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/virtq/buffer.rs | 47 +----- src/hyperlight_common/src/virtq/pool.rs | 79 ++------- src/hyperlight_common/src/virtq/producer.rs | 169 ++++++++------------ 3 files changed, 85 insertions(+), 210 deletions(-) diff --git a/src/hyperlight_common/src/virtq/buffer.rs b/src/hyperlight_common/src/virtq/buffer.rs index 8a5ca708f..7c696a4fc 100644 --- a/src/hyperlight_common/src/virtq/buffer.rs +++ b/src/hyperlight_common/src/virtq/buffer.rs @@ -53,9 +53,6 @@ pub trait BufferProvider { /// Free a previously allocated segment by start address. fn dealloc(&self, addr: u64) -> Result<(), AllocError>; - /// Reset the pool to initial state. - fn reset(&self) {} - /// Allocate scatter/gather segments for a logical payload of `total_len` bytes. fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { if total_len == 0 { @@ -101,9 +98,6 @@ impl BufferProvider for Rc { fn dealloc(&self, addr: u64) -> Result<(), AllocError> { (**self).dealloc(addr) } - fn reset(&self) { - (**self).reset() - } fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { (**self).alloc_sg(total_len) } @@ -119,9 +113,6 @@ impl BufferProvider for Arc { fn dealloc(&self, addr: u64) -> Result<(), AllocError> { (**self).dealloc(addr) } - fn reset(&self) { - (**self).reset() - } fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { (**self).alloc_sg(total_len) } @@ -337,16 +328,8 @@ impl AsRef<[u8]> for BufferOwner { } /// Pool-owned allocation that is returned to the pool on drop. -/// -/// Use [`into_raw`](Self::into_raw) to transfer ownership to a descriptor -/// state that will deallocate the raw [`Allocation`] through another path. #[derive(Debug)] pub struct OwnedAlloc { - inner: Option>, -} - -#[derive(Debug)] -struct Inner { pool: P, alloc: Allocation, } @@ -354,9 +337,7 @@ struct Inner { impl OwnedAlloc

{ /// Wrap an existing allocation with its owning pool. pub fn new(pool: P, alloc: Allocation) -> Self { - Self { - inner: Some(Inner { pool, alloc }), - } + Self { pool, alloc } } /// Allocate from `pool` and return an owning guard. @@ -366,35 +347,15 @@ impl OwnedAlloc

{ } /// The raw allocation currently owned by this guard. - // `inner` is `Some` for the whole lifetime of a live guard: it is only - // taken by `into_raw` which consumes `self` or on drop, so this access - // cannot fail. - #[allow(clippy::expect_used)] pub fn allocation(&self) -> Allocation { - self.inner - .as_ref() - .map(|inner| inner.alloc) - .expect("OwnedAlloc::allocation called after ownership transfer") - } - - /// Release ownership and return the raw allocation. - // `inner` is `Some` until ownership is released, and `into_raw` consumes - // `self`, so it can only ever observe `Some` here. - #[allow(clippy::expect_used)] - pub fn into_raw(mut self) -> Allocation { - self.inner - .take() - .map(|inner| inner.alloc) - .expect("OwnedAlloc::into_raw called after ownership transfer") + self.alloc } } impl Drop for OwnedAlloc

{ fn drop(&mut self) { - if let Some(Inner { pool, alloc }) = self.inner.take() { - let result = pool.dealloc(alloc.addr); - debug_assert!(result.is_ok(), "OwnedAlloc drop dealloc failed: {result:?}"); - } + let result = self.pool.dealloc(self.alloc.addr); + debug_assert!(result.is_ok(), "OwnedAlloc drop dealloc failed: {result:?}"); } } diff --git a/src/hyperlight_common/src/virtq/pool.rs b/src/hyperlight_common/src/virtq/pool.rs index af5b63425..99e3a9950 100644 --- a/src/hyperlight_common/src/virtq/pool.rs +++ b/src/hyperlight_common/src/virtq/pool.rs @@ -233,12 +233,6 @@ impl Slab { fn contains(&self, addr: u64) -> bool { self.range().contains(&addr) } - - fn reset(&mut self) { - self.used_slots.clear(); - self.run_starts.clear(); - self.last_free_run = None; - } } #[cfg(test)] @@ -391,12 +385,6 @@ impl BufferProvider for BufferPool { fn dealloc(&self, addr: u64) -> Result<(), AllocError> { self.inner.borrow_mut().dealloc_addr(addr) } - - fn reset(&self) { - let mut inner = self.inner.borrow_mut(); - inner.lower.reset(); - inner.upper.reset(); - } } impl BufferPool { @@ -544,26 +532,21 @@ impl RecycleList { /// Rebuild state so that exactly the addresses in `allocated` are marked /// live and every other slot is free. /// - /// On error the pool is left in an indeterminate state and should be - /// [`reset`](Self::reset) before reuse. + /// On error the pool is left unchanged. fn restore_allocated(&mut self, allocated: &[u64]) -> Result<(), AllocError> { - self.allocated.clear(); + let mut restored = FixedBitSet::with_capacity(self.count); for &addr in allocated { let slot = self.slot_of(addr)?; - if self.allocated.contains(slot) { + if restored.contains(slot) { return Err(AllocError::InvalidFree(addr, self.slot_size)); } - self.allocated.insert(slot); + restored.insert(slot); } + self.allocated = restored; self.rebuild_free(); Ok(()) } - fn reset(&mut self) { - self.allocated.clear(); - self.rebuild_free(); - } - /// Repopulate the free list with every slot whose allocated bit is clear. fn rebuild_free(&mut self) { self.free.clear(); @@ -619,7 +602,8 @@ impl RecyclePool { } /// Rebuild pool state so that every address in `allocated` is removed from - /// the free list, matching externally known inflight state. + /// the free list, matching externally known inflight state. Validation is + /// transactional: an error leaves the existing pool state unchanged. pub fn restore_allocated(&self, allocated: &[u64]) -> Result<(), AllocError> { self.inner.borrow_mut().restore_allocated(allocated) } @@ -674,10 +658,6 @@ impl BufferProvider for RecyclePool { fn dealloc(&self, addr: u64) -> Result<(), AllocError> { self.inner.borrow_mut().dealloc_addr(addr) } - - fn reset(&self) { - self.inner.borrow_mut().reset() - } } #[cfg(test)] @@ -859,40 +839,6 @@ mod tests { assert!(pool.inner.borrow().upper.contains(alloc2.addr)); } - #[test] - fn test_buffer_pool_reset_returns_to_initial_state() { - let pool = make_pool::<256, 4096>(0x20000); - - // Allocate from both tiers - let a1 = pool.inner.borrow_mut().alloc(128).unwrap(); - let a2 = pool.inner.borrow_mut().alloc(4096).unwrap(); - assert!(a1.len > 0); - assert!(a2.len > 0); - - pool.reset(); - - let inner = pool.inner.borrow(); - assert_eq!(inner.lower.free_bytes(), inner.lower.capacity()); - assert_eq!(inner.upper.free_bytes(), inner.upper.capacity()); - } - - #[test] - fn test_buffer_pool_reset_allows_reallocation() { - let pool = make_pool::<256, 4096>(0x20000); - - // Fill up some allocations - let mut allocs = Vec::new(); - for _ in 0..5 { - allocs.push(pool.inner.borrow_mut().alloc(256).unwrap()); - } - - pool.reset(); - - // Should be able to allocate as if fresh - let a = pool.inner.borrow_mut().alloc(256).unwrap(); - assert!(a.len > 0); - } - #[test] fn test_pool_dealloc_addr_routes_to_correct_tier() { let pool = make_pool::<256, 4096>(0x20000); @@ -970,8 +916,13 @@ mod tests { #[test] fn test_recycle_pool_restore_allocated_invalid_addr_returns_error() { let pool = make_recycle_pool(4, 4096); - let result = pool.restore_allocated(&[0xDEAD]); + pool.restore_allocated(&[0x80000]).unwrap(); + + let result = pool.restore_allocated(&[0x81000, 0xDEAD]); assert!(result.is_err()); + assert_eq!(pool.num_free(), 3); + assert_eq!(pool.allocation_len(0x80000).unwrap(), 4096); + assert!(pool.allocation_len(0x81000).is_err()); } #[test] @@ -1005,15 +956,13 @@ mod tests { } #[test] - fn test_recycle_pool_restore_allocated_resets_first() { + fn test_recycle_pool_restore_allocated_replaces_state() { let pool = make_recycle_pool(4, 4096); - // Allocate some slots let _ = pool.alloc(4096).unwrap(); let _ = pool.alloc(4096).unwrap(); assert_eq!(pool.num_free(), 2); - // restore_allocated resets then removes - so 4 - 1 = 3 pool.restore_allocated(&[0x80000]).unwrap(); assert_eq!(pool.num_free(), 3); } diff --git a/src/hyperlight_common/src/virtq/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index 0443c0ec6..2c9d7562b 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -149,7 +149,11 @@ where } } - fn dealloc_elems( + /// Retire allocations from a completed descriptor chain. + /// + /// Every element is attempted so one deallocation failure does not strand + /// later allocations; the first failure is returned after cleanup. + fn retire_elems( &self, elems: impl IntoIterator, ) -> Result<(), VirtqError> { @@ -289,39 +293,6 @@ where } Ok(()) } - - /// Reset ring, inflight, and pool state to initial values. - /// - /// # Safety - /// - /// No outstanding [`UsedChain::Data`] buffers, borrowed segment views, or - /// peer accesses to previously submitted descriptors may exist. Resetting - /// recycles the same backing addresses, so outstanding zero-copy buffers or - /// stale descriptor users could alias memory that is handed out again. - /// - /// TODO(virtq): find a way to allow guest to keep used chains across resets. - pub unsafe fn reset(&mut self) { - self.inflight.iter_mut().for_each(|slot| *slot = None); - self.pending.clear(); - self.inner.reset(); - self.pool.reset(); - } - - /// Replace the pool and reset ring, inflight, and pending state. - /// - /// # Safety - /// - /// No outstanding [`UsedChain::Data`] buffers, borrowed segment views, or - /// peer accesses to previously submitted descriptors may exist. The new pool - /// may manage the same shared-memory addresses as the old pool, so old - /// zero-copy buffers must not outlive this transition. - pub unsafe fn reset_with_pool(&mut self, pool: P) { - self.pending.clear(); - self.inflight.iter_mut().for_each(|slot| *slot = None); - self.inner.reset(); - self.pool = pool; - self.pool.reset(); - } } impl VirtqProducer @@ -393,7 +364,7 @@ where let written = used.len as usize; let Inflight { token, chain } = inf; - self.dealloc_elems(chain.readables().iter().copied())?; + self.retire_elems(chain.readables().iter().copied())?; let used = if chain.writables().is_empty() { UsedChain::Ack(token) @@ -426,14 +397,14 @@ where if remaining != 0 { let elems = owned.iter().map(|(elem, _)| *elem).chain(free); - self.dealloc_elems(elems)?; + self.retire_elems(elems)?; return Err(VirtqError::InvalidState); } for (elem, len) in &owned { if unsafe { self.inner.mem().as_slice(elem.addr, *len) }.is_err() { let elems = owned.iter().map(|(elem, _)| *elem).chain(free); - let _ = self.dealloc_elems(elems); + let _ = self.retire_elems(elems); return Err(VirtqError::MemoryReadError); } } @@ -456,7 +427,7 @@ where sgs.push(Bytes::from_owner(owner)); } - self.dealloc_elems(free)?; + self.retire_elems(free)?; Ok(Segments::from_smallvec(sgs)) } @@ -595,7 +566,8 @@ impl ChainBuilder { return Err(VirtqError::InvalidState); } - let mut rollback = Rollback::new(&self.pool); + let rd_capacity = self.rd_caps.iter().sum(); + let mut allocs = AllocTxn::new(&self.pool); let mut rd_caps = SmallVec::<[usize; 4]>::new(); let mut rd_elems = SmallVec::<[BufferElement; 4]>::new(); let mut wr_elems = SmallVec::<[BufferElement; 4]>::new(); @@ -604,7 +576,7 @@ impl ChainBuilder { // The buffer element lengths are initialized to zero and updated as the // `SendChain` writes. for &cap in &self.rd_caps { - let sgs = self.pool.alloc_sg(cap)?; + let sgs = allocs.alloc_sg(cap)?; let mut remaining = cap; for alloc in sgs { @@ -618,7 +590,6 @@ impl ChainBuilder { writable: false, }); remaining -= seg_cap; - rollback.allocs.push(alloc); } if remaining != 0 { @@ -630,7 +601,7 @@ impl ChainBuilder { // Writable buffer elements are initialized with their full capacity for the device to // write into. for &cap in &self.wr_caps { - let sgs = self.pool.alloc_sg(cap)?; + let sgs = allocs.alloc_sg(cap)?; for alloc in sgs { let len = checked_descriptor_len(alloc.len)?; wr_elems.push(BufferElement { @@ -638,7 +609,6 @@ impl ChainBuilder { len, writable: true, }); - rollback.allocs.push(alloc); } } @@ -647,43 +617,61 @@ impl ChainBuilder { .writables(wr_elems) .build()?; - rollback.release(); + allocs.commit(); Ok(SendChain { mem: self.mem, pool: self.pool, chain: Some(chain), rd_caps, - rd_capacity: self.rd_caps.iter().sum(), + rd_capacity, rd_written: 0, write_mode: WriteMode::Unset, }) } } -struct Rollback<'a, P: BufferProvider> { +/// Build-scoped allocation transaction. +/// +/// `SendChain` and `Inflight` intentionally retain lightweight descriptor +/// metadata instead of one allocation guard and cloned pool handle per +/// descriptor. While a valid `BufferChain` is being built, this transaction +/// provides aggregate RAII: it records every allocated address before the +/// caller can perform fallible validation and returns them all on drop. +/// [`commit`](Self::commit) disarms rollback once `SendChain` can take +/// responsibility for reclaiming the completed chain. +struct AllocTxn<'a, P: BufferProvider> { pool: &'a P, - allocs: SmallVec<[Allocation; 8]>, + addrs: SmallVec<[u64; 8]>, } -impl<'a, P: BufferProvider> Rollback<'a, P> { +impl<'a, P: BufferProvider> AllocTxn<'a, P> { fn new(pool: &'a P) -> Self { Self { pool, - allocs: SmallVec::new(), + addrs: SmallVec::new(), } } - fn release(mut self) { - self.allocs.clear(); + fn alloc_sg(&mut self, total_len: usize) -> Result, AllocError> { + let allocs = self.pool.alloc_sg(total_len)?; + self.addrs.extend(allocs.iter().map(|alloc| alloc.addr)); + Ok(allocs) + } + + fn commit(mut self) { + self.addrs.clear(); } } -impl Drop for Rollback<'_, P> { +impl Drop for AllocTxn<'_, P> { fn drop(&mut self) { - for alloc in self.allocs.drain(..) { - let result = self.pool.dealloc(alloc.addr); - debug_assert!(result.is_ok(), "rollback dealloc failed: {result:?}"); + for addr in self.addrs.drain(..) { + let result = self.pool.dealloc(addr); + debug_assert!( + result.is_ok(), + "allocation rollback dealloc failed: {result:?}" + ); } } } @@ -1544,6 +1532,30 @@ mod tests { assert!(tok.id < 16); } + #[cfg(target_pointer_width = "64")] + #[test] + fn test_chain_build_rolls_back_unrepresentable_allocations() { + let ring = make_ring(16); + let slot_size = u32::MAX as usize + 1; + let pool = RecyclePool::new(0, slot_size, slot_size).unwrap(); + let mem = ring.mem(); + let producer = VirtqProducer::new(ring.layout(), mem, TestNotifier::new(), pool.clone()); + + assert!(matches!( + producer.chain().readable(1).build(), + Err(VirtqError::PayloadTooLarge { recv, limit }) + if recv == slot_size && limit == u32::MAX as usize + )); + assert_eq!(pool.num_free(), 1); + + assert!(matches!( + producer.chain().writable(1).build(), + Err(VirtqError::PayloadTooLarge { recv, limit }) + if recv == slot_size && limit == u32::MAX as usize + )); + assert_eq!(pool.num_free(), 1); + } + #[test] fn test_submit_notifies() { let ring = make_ring(16); @@ -1784,51 +1796,4 @@ mod tests { assert_eq!(producer.inner.num_inflight(), 1); } - #[test] - fn test_virtq_producer_reset() { - let ring = make_ring(16); - let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); - - // Submit and complete a round trip - let mut se = producer.chain().readable(32).writable(64).build().unwrap(); - se.write_all(b"hello").unwrap(); - producer.submit(se).unwrap(); - - let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello"); - consumer.complete(recv, reply).unwrap(); - let _ = producer.poll().unwrap().unwrap(); - - // Now reset - // SAFETY: the used chain was dropped before reset and no peer can - // access the reset test ring concurrently. - unsafe { - producer.reset(); - } - - // All inflight slots should be cleared - assert_eq!(producer.inner.num_inflight(), 0); - // Ring state should be back to initial - assert_eq!(producer.inner.num_free(), producer.inner.len()); - } - - #[test] - fn test_virtq_producer_reset_clears_inflight() { - let ring = make_ring(16); - let (mut producer, _consumer, _notifier) = make_test_producer(&ring); - - // Submit without completing - let se = producer.chain().writable(64).build().unwrap(); - producer.submit(se).unwrap(); - - assert_eq!(producer.inner.num_inflight(), 1); - - // SAFETY: no peer can access the reset test ring concurrently. - unsafe { - producer.reset(); - } - - assert_eq!(producer.inner.num_inflight(), 0); - assert_eq!(producer.inner.num_free(), producer.inner.len()); - } } From d47e9d28014940cece95f170fd4a614e92aa326e Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Wed, 22 Jul 2026 17:53:39 +0200 Subject: [PATCH 03/15] feat(virtq): add tiered fixed slot allocation Split and rename the pool implementations, add explicit lower/upper SlotPool regions. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/benches/buffer_pool.rs | 38 +- src/hyperlight_common/benches/common/mod.rs | 23 +- src/hyperlight_common/benches/virtq_api.rs | 32 +- src/hyperlight_common/src/virtq/buffer.rs | 111 +- .../src/virtq/concurrency.rs | 10 +- src/hyperlight_common/src/virtq/pool.rs | 1345 ++--------------- src/hyperlight_common/src/virtq/pool/fuzz.rs | 369 +++++ src/hyperlight_common/src/virtq/pool/run.rs | 366 +++++ src/hyperlight_common/src/virtq/pool/slot.rs | 379 +++++ src/hyperlight_common/src/virtq/pool/tests.rs | 502 ++++++ src/hyperlight_common/src/virtq/producer.rs | 3 +- 11 files changed, 1797 insertions(+), 1381 deletions(-) create mode 100644 src/hyperlight_common/src/virtq/pool/fuzz.rs create mode 100644 src/hyperlight_common/src/virtq/pool/run.rs create mode 100644 src/hyperlight_common/src/virtq/pool/slot.rs create mode 100644 src/hyperlight_common/src/virtq/pool/tests.rs diff --git a/src/hyperlight_common/benches/buffer_pool.rs b/src/hyperlight_common/benches/buffer_pool.rs index 19ef5d464..4357b8992 100644 --- a/src/hyperlight_common/benches/buffer_pool.rs +++ b/src/hyperlight_common/benches/buffer_pool.rs @@ -4,12 +4,12 @@ use std::hint::black_box; use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; -use hyperlight_common::virtq::{BufferPool, BufferProvider, RecyclePool}; +use hyperlight_common::virtq::{BufferProvider, RunPool, SlotLayout, SlotPool}; // Helper to create a pool for benchmarking -fn make_pool(size: usize) -> BufferPool { +fn make_run_pool(size: usize) -> RunPool { let base = 0x10000; - BufferPool::::new(base, size).unwrap() + RunPool::::new(base, size).unwrap() } // Single allocation performance @@ -19,7 +19,7 @@ fn bench_alloc_single(c: &mut Criterion) { for size in [64, 128, 256, 512, 1024, 1500, 4096].iter() { group.throughput(Throughput::Elements(1)); group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, &size| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { let alloc = pool.alloc(black_box(size)).unwrap(); pool.dealloc(alloc.addr).unwrap(); @@ -36,7 +36,7 @@ fn bench_alloc_lifo(c: &mut Criterion) { for size in [256, 1500, 4096].iter() { group.throughput(Throughput::Elements(100)); group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, &size| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { for _ in 0..100 { let alloc = pool.alloc(black_box(size)).unwrap(); @@ -53,7 +53,7 @@ fn bench_alloc_fragmented(c: &mut Criterion) { let mut group = c.benchmark_group("alloc_fragmented"); group.bench_function("fragmented_256", |b| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); // Create fragmentation pattern: allocate many, free every other let mut allocations = Vec::new(); @@ -79,7 +79,7 @@ fn bench_free(c: &mut Criterion) { for size in [256, 1500, 4096].iter() { group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, &size| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { let alloc = pool.alloc(size).unwrap(); pool.dealloc(black_box(alloc.addr)).unwrap(); @@ -96,7 +96,7 @@ fn bench_free_list_reuse(c: &mut Criterion) { // With cursor optimization (LIFO) group.bench_function("lifo_pattern", |b| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { let alloc = pool.alloc(256).unwrap(); pool.dealloc(alloc.addr).unwrap(); @@ -107,7 +107,7 @@ fn bench_free_list_reuse(c: &mut Criterion) { // Without cursor benefit (FIFO-like) group.bench_function("fifo_pattern", |b| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); let mut queue = Vec::new(); // Pre-fill queue @@ -136,7 +136,7 @@ fn bench_segmented_payload(c: &mut Criterion) { BenchmarkId::from_parameter(payload_size), &payload_size, |b, &payload_size| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { let sgs = pool.alloc_sg(black_box(payload_size)).unwrap(); for sg in sgs { @@ -150,11 +150,12 @@ fn bench_segmented_payload(c: &mut Criterion) { group.finish(); } -fn bench_recycle_pool(c: &mut Criterion) { - let mut group = c.benchmark_group("recycle_pool"); +fn bench_slot_pool(c: &mut Criterion) { + let mut group = c.benchmark_group("slot_pool"); group.bench_function("alloc_dealloc_4096", |b| { - let pool = RecyclePool::new(0x80000, 4 * 1024 * 1024, 4096).unwrap(); + let layout = SlotLayout::new(0x80000, 4096, 1024); + let pool = SlotPool::new(layout).unwrap(); b.iter(|| { let alloc = pool.alloc(black_box(4096)).unwrap(); pool.dealloc(alloc.addr).unwrap(); @@ -162,7 +163,8 @@ fn bench_recycle_pool(c: &mut Criterion) { }); group.bench_function("alloc_dealloc_128", |b| { - let pool = RecyclePool::new(0x80000, 4 * 1024 * 1024, 256).unwrap(); + let layout = SlotLayout::new(0x80000, 256, 16 * 1024); + let pool = SlotPool::new(layout).unwrap(); b.iter(|| { let alloc = pool.alloc(black_box(128)).unwrap(); pool.dealloc(alloc.addr).unwrap(); @@ -170,7 +172,8 @@ fn bench_recycle_pool(c: &mut Criterion) { }); group.bench_function("alloc_dealloc_1500", |b| { - let pool = RecyclePool::new(0x80000, 4 * 1024 * 1024, 4096).unwrap(); + let layout = SlotLayout::new(0x80000, 4096, 1024); + let pool = SlotPool::new(layout).unwrap(); b.iter(|| { let alloc = pool.alloc(black_box(1500)).unwrap(); pool.dealloc(alloc.addr).unwrap(); @@ -178,7 +181,8 @@ fn bench_recycle_pool(c: &mut Criterion) { }); group.bench_function("alloc_sg_64k", |b| { - let pool = RecyclePool::new(0x80000, 4 * 1024 * 1024, 4096).unwrap(); + let layout = SlotLayout::new(0x80000, 4096, 1024); + let pool = SlotPool::new(layout).unwrap(); b.iter(|| { let sgs = pool.alloc_sg(black_box(64 * 1024)).unwrap(); for sg in sgs { @@ -198,7 +202,7 @@ criterion_group!( bench_free, bench_free_list_reuse, bench_segmented_payload, - bench_recycle_pool, + bench_slot_pool, ); criterion_main!(benches); diff --git a/src/hyperlight_common/benches/common/mod.rs b/src/hyperlight_common/benches/common/mod.rs index 77ed2c78a..e91a97d9a 100644 --- a/src/hyperlight_common/benches/common/mod.rs +++ b/src/hyperlight_common/benches/common/mod.rs @@ -14,15 +14,15 @@ use std::sync::atomic::{AtomicU16, AtomicUsize, Ordering}; use bytemuck::Pod; use hyperlight_common::virtq::{ - BufferPool, BufferProvider, Descriptor, Layout, MemOps, Notifier, QueueStats, RecyclePool, - ReplyChain, UsedChain, VirtqConsumer, VirtqProducer, + BufferProvider, Descriptor, Layout, MemOps, Notifier, QueueStats, ReplyChain, RunPool, + SlotLayout, SlotPool, UsedChain, VirtqConsumer, VirtqProducer, }; pub const LOWER_SLOT: usize = 256; pub const UPPER_SLOT: usize = 4096; pub const POOL_SIZE: usize = 8 * 1024 * 1024; -pub type RunBufferPool = BufferPool; +pub type BenchRunPool = RunPool; #[derive(Clone)] struct BenchMem { @@ -163,12 +163,12 @@ where BenchPair { producer, consumer } } -pub fn run_buffer_pool(base: u64, size: usize) -> RunBufferPool { - BufferPool::new(base, size).unwrap() +pub fn run_pool(base: u64, size: usize) -> BenchRunPool { + RunPool::new(base, size).unwrap() } -pub fn fragmented_run_buffer_pool(base: u64, size: usize, payload_size: usize) -> RunBufferPool { - let pool = run_buffer_pool(base, size); +pub fn fragmented_run_pool(base: u64, size: usize, payload_size: usize) -> BenchRunPool { + let pool = run_pool(base, size); let payload_slots = payload_size.div_ceil(UPPER_SLOT); let prefix_slots = 32; let suffix_slots = 32; @@ -184,12 +184,13 @@ pub fn fragmented_run_buffer_pool(base: u64, size: usize, payload_size: usize) - pool } -pub fn recycle_pool(base: u64, size: usize) -> RecyclePool { - RecyclePool::new(base, size, UPPER_SLOT).unwrap() +pub fn slot_pool(base: u64, size: usize) -> SlotPool { + let layout = SlotLayout::new(base, UPPER_SLOT, size / UPPER_SLOT); + SlotPool::new(layout).unwrap() } -pub fn fragmented_recycle_pool(base: u64, size: usize, payload_size: usize) -> RecyclePool { - let pool = recycle_pool(base, size); +pub fn fragmented_slot_pool(base: u64, size: usize, payload_size: usize) -> SlotPool { + let pool = slot_pool(base, size); let payload_slots = payload_size.div_ceil(UPPER_SLOT); let allocated: Vec<_> = (0..payload_slots * 2 + 16) .map(|_| pool.alloc(UPPER_SLOT).unwrap()) diff --git a/src/hyperlight_common/benches/virtq_api.rs b/src/hyperlight_common/benches/virtq_api.rs index 565ffde00..e17e71e01 100644 --- a/src/hyperlight_common/benches/virtq_api.rs +++ b/src/hyperlight_common/benches/virtq_api.rs @@ -17,10 +17,10 @@ fn bench_readonly_strategies(c: &mut Criterion) { group.throughput(Throughput::Bytes(size as u64)); group.bench_with_input( - BenchmarkId::new("buffer_pool_run", size), + BenchmarkId::new("run_pool", size), &payload, |b, payload| { - let mut pair = make_pair(128, run_buffer_pool); + let mut pair = make_pair(128, run_pool); b.iter(|| { let used = readonly_roundtrip(&mut pair, black_box(payload)); debug_assert!(matches!(used, UsedChain::Ack(_))); @@ -29,11 +29,11 @@ fn bench_readonly_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("buffer_pool_run_fragmented", size), + BenchmarkId::new("run_pool_fragmented", size), &payload, |b, payload| { let mut pair = make_pair(128, |base, pool_size| { - fragmented_run_buffer_pool(base, pool_size, payload.len()) + fragmented_run_pool(base, pool_size, payload.len()) }); b.iter(|| { let used = readonly_roundtrip(&mut pair, black_box(payload)); @@ -43,10 +43,10 @@ fn bench_readonly_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("recycle_pool_segmented", size), + BenchmarkId::new("slot_pool_segmented", size), &payload, |b, payload| { - let mut pair = make_pair(128, recycle_pool); + let mut pair = make_pair(128, slot_pool); b.iter(|| { let used = readonly_roundtrip(&mut pair, black_box(payload)); debug_assert!(matches!(used, UsedChain::Ack(_))); @@ -55,11 +55,11 @@ fn bench_readonly_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("recycle_pool_segmented_fragmented", size), + BenchmarkId::new("slot_pool_segmented_fragmented", size), &payload, |b, payload| { let mut pair = make_pair(128, |base, pool_size| { - fragmented_recycle_pool(base, pool_size, payload.len()) + fragmented_slot_pool(base, pool_size, payload.len()) }); b.iter(|| { let used = readonly_roundtrip(&mut pair, black_box(payload)); @@ -81,10 +81,10 @@ fn bench_readwrite_strategies(c: &mut Criterion) { group.throughput(Throughput::Bytes((request.len() + response.len()) as u64)); group.bench_with_input( - BenchmarkId::new("buffer_pool_run", size), + BenchmarkId::new("run_pool", size), &(request.clone(), response.clone()), |b, (request, response)| { - let mut pair = make_pair(128, run_buffer_pool); + let mut pair = make_pair(128, run_pool); b.iter(|| { let used = readwrite_roundtrip( &mut pair, @@ -97,11 +97,11 @@ fn bench_readwrite_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("buffer_pool_run_fragmented", size), + BenchmarkId::new("run_pool_fragmented", size), &(request.clone(), response.clone()), |b, (request, response)| { let mut pair = make_pair(128, |base, pool_size| { - fragmented_run_buffer_pool(base, pool_size, request.len()) + fragmented_run_pool(base, pool_size, request.len()) }); b.iter(|| { let used = readwrite_roundtrip( @@ -115,10 +115,10 @@ fn bench_readwrite_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("recycle_pool_segmented", size), + BenchmarkId::new("slot_pool_segmented", size), &(request.clone(), response.clone()), |b, (request, response)| { - let mut pair = make_pair(128, recycle_pool); + let mut pair = make_pair(128, slot_pool); b.iter(|| { let used = readwrite_roundtrip( &mut pair, @@ -131,11 +131,11 @@ fn bench_readwrite_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("recycle_pool_segmented_fragmented", size), + BenchmarkId::new("slot_pool_segmented_fragmented", size), &(request, response), |b, (request, response)| { let mut pair = make_pair(128, |base, pool_size| { - fragmented_recycle_pool(base, pool_size, request.len()) + fragmented_slot_pool(base, pool_size, request.len()) }); b.iter(|| { let used = readwrite_roundtrip( diff --git a/src/hyperlight_common/src/virtq/buffer.rs b/src/hyperlight_common/src/virtq/buffer.rs index 7c696a4fc..9f3265e9f 100644 --- a/src/hyperlight_common/src/virtq/buffer.rs +++ b/src/hyperlight_common/src/virtq/buffer.rs @@ -1,122 +1,15 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright 2026 The Hyperlight Authors. -//! Buffer allocation traits and shared types for virtqueue buffer management. +//! Owned and segmented virtqueue buffer representations. -use alloc::rc::Rc; -use alloc::sync::Arc; use alloc::vec::Vec; use bytes::{Buf, Bytes}; use smallvec::{SmallVec, smallvec}; -use thiserror::Error; use super::access::MemOps; - -#[derive(Debug, Error, Copy, Clone)] -pub enum AllocError { - #[error("Invalid region addr {0}")] - InvalidAlign(u64), - #[error("Invalid free addr {0} and size {1}")] - InvalidFree(u64, usize), - #[error("Invalid argument")] - InvalidArg, - #[error("Empty region")] - EmptyRegion, - #[error("No space available")] - NoSpace, - #[error("Requested size exceeds pool capacity")] - OutOfMemory, - #[error("Overflow")] - Overflow, -} - -/// Allocation result -#[derive(Debug, Clone, Copy)] -pub struct Allocation { - /// Starting address of the allocation - pub addr: u64, - /// Capacity of the allocation in bytes, rounded up to the allocator's slot size. - pub len: usize, -} - -/// Trait for buffer providers. -pub trait BufferProvider { - /// Preferred maximum size of one allocation segment. - fn max_alloc_len(&self) -> usize { - usize::MAX - } - - /// Allocate one buffer that can hold at least `len` bytes. - fn alloc(&self, len: usize) -> Result; - - /// Free a previously allocated segment by start address. - fn dealloc(&self, addr: u64) -> Result<(), AllocError>; - - /// Allocate scatter/gather segments for a logical payload of `total_len` bytes. - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - if total_len == 0 { - return Err(AllocError::InvalidArg); - } - - let seg_cap = self.max_alloc_len(); - if seg_cap == 0 { - return Err(AllocError::InvalidArg); - } - - let mut rem = total_len; - let mut sgs = SmallVec::<[Allocation; 4]>::new(); - - while rem > 0 { - let len = rem.min(seg_cap); - match self.alloc(len) { - Ok(alloc) => { - sgs.push(alloc); - rem -= len; - } - Err(err) => { - for sg in sgs { - let _res = self.dealloc(sg.addr); - debug_assert!(_res.is_ok(), "dealloc failed: {_res:?}"); - } - return Err(err); - } - } - } - - Ok(sgs) - } -} - -impl BufferProvider for Rc { - fn max_alloc_len(&self) -> usize { - (**self).max_alloc_len() - } - fn alloc(&self, len: usize) -> Result { - (**self).alloc(len) - } - fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - (**self).dealloc(addr) - } - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - (**self).alloc_sg(total_len) - } -} - -impl BufferProvider for Arc { - fn max_alloc_len(&self) -> usize { - (**self).max_alloc_len() - } - fn alloc(&self, len: usize) -> Result { - (**self).alloc(len) - } - fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - (**self).dealloc(addr) - } - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - (**self).alloc_sg(total_len) - } -} +use super::pool::{AllocError, Allocation, BufferProvider}; /// Ordered byte segments that make up one virtqueue payload. /// diff --git a/src/hyperlight_common/src/virtq/concurrency.rs b/src/hyperlight_common/src/virtq/concurrency.rs index 33df1cc8c..e91f24922 100644 --- a/src/hyperlight_common/src/virtq/concurrency.rs +++ b/src/hyperlight_common/src/virtq/concurrency.rs @@ -49,7 +49,7 @@ use loom::thread; use super::*; use crate::virtq::desc::Descriptor; -use crate::virtq::pool::BufferPoolSync; +use crate::virtq::pool::RunPoolSync; #[derive(Debug)] pub struct MemErr; @@ -316,7 +316,7 @@ fn virtq_ping_pong() { let pool_size = 0x10000; let mem = Arc::new(LoomMem::new(ring_base, 8, pool_base, pool_size)); - let pool = BufferPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); + let pool = RunPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); let notify = Arc::new(Notify::new()); let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); @@ -364,7 +364,7 @@ fn virtq_ack_only() { let pool_size = 0x10000; let mem = Arc::new(LoomMem::new(ring_base, 4, pool_base, pool_size)); - let pool = BufferPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); + let pool = RunPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); let notify = Arc::new(Notify::new()); let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); @@ -408,7 +408,7 @@ fn virtq_out_of_order_completions() { let pool_size = 0x10000; let mem = Arc::new(LoomMem::new(ring_base, 8, pool_base, pool_size)); - let pool = BufferPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); + let pool = RunPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); let notify = Arc::new(Notify::new()); let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); @@ -500,7 +500,7 @@ fn virtq_event_suppression_reconfig() { let pool_size = 0x10000; let mem = Arc::new(LoomMem::new(ring_base, 4, pool_base, pool_size)); - let pool = BufferPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); + let pool = RunPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); let notify = Arc::new(Notify::new()); let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); diff --git a/src/hyperlight_common/src/virtq/pool.rs b/src/hyperlight_common/src/virtq/pool.rs index 99e3a9950..a75fe684f 100644 --- a/src/hyperlight_common/src/virtq/pool.rs +++ b/src/hyperlight_common/src/virtq/pool.rs @@ -1,1287 +1,188 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright 2026 The Hyperlight Authors. -//! Buffer pool implementations for virtqueue buffer management. +//! Buffer pool APIs and implementations for virtqueue payloads. //! -//! This module provides concrete buffer allocators: -//! -//! - [`BufferPool`] - a two-tier run allocator for variable-sized allocations. -//! - [`RecyclePool`] - a single-tier fixed-slot free-list recycler for bounded -//! descriptor segments. -//! -//! All implement [`BufferProvider`] from the [`super::buffer`] module. -//! -//! # BufferPool design -//! -//! `BufferPool` is a variable-sized run allocator. -//! -//! # Two-tier layout -//! -//! [`BufferPool`] divides the underlying region into two slabs with different -//! slot sizes: -//! -//! - The lower tier (default `L = 256`) is intended for *smaller allocations* - -//! control messages, descriptor metadata, and other small structures. Small -//! allocations first try this tier. -//! - The upper tier (default `U = 4096`) uses page sized slots and is intended -//! for larger contiguous buffers. +//! - [`RunPool`] allocates variable-sized contiguous runs from two tiers. +//! - [`SlotPool`] recycles one or two tiers of fixed-size slots. use alloc::rc::Rc; -use core::cell::RefCell; +use alloc::sync::Arc; use core::ops::Deref; -use fixedbitset::FixedBitSet; use smallvec::SmallVec; +use thiserror::Error; -use super::buffer::{AllocError, Allocation, BufferProvider}; - -/// Wrapper asserting `Send` for an inner value that is only ever accessed from -/// a single thread. -/// -/// [`BufferPool`] and [`RecyclePool`] hold their state in an `Rc>`, -/// which is neither `Send` nor `Sync`. Their allocations are exposed as -/// zero-copy reply payloads through -/// [`Bytes::from_owner`](bytes::Bytes::from_owner), whose owner bound is -/// `Send + 'static`; this wrapper exists solely so the pools can satisfy that -/// bound. -/// -/// # Safety -/// -/// The `Send` assertion is only sound while the wrapped value - and every -/// `Bytes` handed out from it - stays on a single thread. Hyperlight guests are -/// single-threaded, so this holds for guest-side use. It is unsound to move a -/// pool (or a reply `Bytes`) to another thread, e.g. by using these pools with a -/// producer/consumer on the multi-threaded host. -#[derive(Debug)] -struct SendWrap(T); - -impl Clone for SendWrap { - fn clone(&self) -> Self { - Self(self.0.clone()) - } -} - -impl Deref for SendWrap { - type Target = T; - fn deref(&self) -> &T { - &self.0 - } -} - -#[derive(Debug, Clone)] -struct Slab { - base_addr: u64, - used_slots: FixedBitSet, - run_starts: FixedBitSet, - last_free_run: Option, -} - -impl Slab { - fn new(base_addr: u64, region_len: usize) -> Result { - let usable = region_len - (region_len % N); - let num_slots = usable / N; - let used_slots = FixedBitSet::with_capacity(num_slots); - let run_starts = FixedBitSet::with_capacity(num_slots); - - if !base_addr.is_multiple_of(N as u64) { - return Err(AllocError::InvalidAlign(base_addr)); - } - if num_slots == 0 { - return Err(AllocError::EmptyRegion); - } - - Ok(Self { - base_addr, - used_slots, - run_starts, - last_free_run: None, - }) - } - - fn addr_of(&self, slot_idx: usize) -> Option { - self.base_addr - .checked_add((slot_idx as u64).checked_mul(N as u64)?) - } - - fn slot_of(&self, addr: u64) -> usize { - let off = (addr - self.base_addr) as usize; - off / N - } - - fn checked_slot_of(&self, addr: u64, len: usize) -> Result { - if addr < self.base_addr { - return Err(AllocError::InvalidFree(addr, len)); - } - - let off = (addr - self.base_addr) as usize; - if !off.is_multiple_of(N) { - return Err(AllocError::InvalidFree(addr, len)); - } - - let slot = off / N; - if slot >= self.used_slots.len() { - return Err(AllocError::InvalidFree(addr, len)); - } - - Ok(slot) - } - - fn live_run_slots_at(&self, start: usize) -> Option { - if start >= self.used_slots.len() - || !self.used_slots.contains(start) - || !self.run_starts.contains(start) - { - return None; - } - - let mut end = start + 1; - while end < self.used_slots.len() - && self.used_slots.contains(end) - && !self.run_starts.contains(end) - { - end += 1; - } - - Some(end - start) - } - - fn maybe_invalidate_last_run(&mut self, alloc: Allocation) { - if let Some(run) = &self.last_free_run { - let new_end = alloc.addr + alloc.len as u64; - let run_end = run.addr + run.len as u64; +mod run; +mod slot; - if alloc.addr < run_end && run.addr < new_end { - self.last_free_run = None; - } - } +pub use run::RunPool; +#[cfg(all(test, loom))] +pub use run::RunPoolSync; +pub use slot::{SlotLayout, SlotPool}; + +/// Buffer allocation failure. +#[derive(Debug, Error, Copy, Clone)] +pub enum AllocError { + /// A region does not meet its required alignment. + #[error("Invalid region addr {0}")] + InvalidAlign(u64), + /// An address does not identify a live allocation. + #[error("Invalid free addr {0} and size {1}")] + InvalidFree(u64, usize), + /// An argument is zero or otherwise invalid. + #[error("Invalid argument")] + InvalidArg, + /// A region cannot hold any allocation. + #[error("Empty region")] + EmptyRegion, + /// No currently free allocation can satisfy the request. + #[error("No space available")] + NoSpace, + /// The request exceeds the pool's allocation capacity. + #[error("Requested size exceeds pool capacity")] + OutOfMemory, + /// Address or size arithmetic overflowed. + #[error("Overflow")] + Overflow, +} + +/// One allocation returned by a [`BufferProvider`]. +#[derive(Debug, Clone, Copy)] +pub struct Allocation { + /// Starting address of the allocation. + pub addr: u64, + /// Capacity in bytes, rounded up according to the provider's policy. + pub len: usize, +} + +/// Allocates and reclaims virtqueue payload buffers. +pub trait BufferProvider { + /// Preferred maximum size of one allocation segment. + fn max_alloc_len(&self) -> usize { + usize::MAX } - fn find_slots(&mut self, slots_num: usize) -> Option { - debug_assert!(slots_num > 0); + /// Allocate one buffer that can hold at least `len` bytes. + fn alloc(&self, len: usize) -> Result; - if let Some(alloc) = self.last_free_run - && alloc.len >= slots_num * N - { - let pos = self.slot_of(alloc.addr); - let _ = self.last_free_run.take(); - return Some(pos); - } - - let total = self.used_slots.len(); - self.used_slots.zeroes().find(|&next_free| { - let end = next_free + slots_num; - end <= total && self.used_slots.count_zeroes(next_free..end) == slots_num - }) - } + /// Free a previously allocated segment by start address. + fn dealloc(&self, addr: u64) -> Result<(), AllocError>; - fn alloc(&mut self, len: usize) -> Result { - if len == 0 { + /// Allocate scatter/gather segments for a logical payload of `total_len` bytes. + fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { + if total_len == 0 { return Err(AllocError::InvalidArg); } - let total = self.used_slots.len(); - let need_slots = len.div_ceil(N); - if need_slots > total { - return Err(AllocError::OutOfMemory); + let seg_cap = self.max_alloc_len(); + if seg_cap == 0 { + return Err(AllocError::InvalidArg); } - let idx = self.find_slots(need_slots).ok_or(AllocError::NoSpace)?; - self.used_slots.insert_range(idx..idx + need_slots); - self.run_starts.insert(idx); - let addr = self.addr_of(idx).ok_or(AllocError::Overflow)?; - - let alloc = Allocation { - addr, - len: need_slots * N, - }; - - self.maybe_invalidate_last_run(alloc); - Ok(alloc) - } - - fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { - let start = self.checked_slot_of(addr, 0)?; - let run_slots = self - .live_run_slots_at(start) - .ok_or(AllocError::InvalidFree(addr, 0))?; - self.dealloc_run(start, run_slots, addr) - } - - fn dealloc_run(&mut self, start: usize, run_slots: usize, addr: u64) -> Result<(), AllocError> { - let len = run_slots * N; - self.used_slots.remove_range(start..start + run_slots); - self.run_starts.set(start, false); - self.last_free_run = Some(Allocation { addr, len }); - Ok(()) - } - - fn allocation_len(&self, addr: u64) -> Result { - let start = self.checked_slot_of(addr, 0)?; - let run_slots = self - .live_run_slots_at(start) - .ok_or(AllocError::InvalidFree(addr, 0))?; - Ok(run_slots * N) - } - - fn capacity(&self) -> usize { - self.used_slots.len() * N - } - - fn range(&self) -> core::ops::Range { - self.base_addr..self.base_addr + self.capacity() as u64 - } - - fn contains(&self, addr: u64) -> bool { - self.range().contains(&addr) - } -} - -#[cfg(test)] -impl Slab { - fn free_bytes(&self) -> usize { - (self.used_slots.len() - self.used_slots.count_ones(..)) * N - } -} - -#[inline] -fn align_up(val: usize, align: usize) -> Result { - if align == 0 { - return Err(AllocError::InvalidArg); - } - - val.checked_next_multiple_of(align) - .ok_or(AllocError::Overflow) -} - -#[derive(Debug)] -struct Inner { - lower: Slab, - upper: Slab, -} - -// SAFETY: only sound for single-threaded (guest-side) access; see the -// type-level invariant on `SendWrap`. -unsafe impl Send for SendWrap>>> {} - -/// Two tier buffer pool with small and large slabs. -#[derive(Debug, Clone)] -pub struct BufferPool { - inner: SendWrap>>>, -} - -impl BufferPool { - /// Create a new buffer pool over a fixed region. - pub fn new(base_addr: u64, region_len: usize) -> Result { - let inner = Inner::::new(base_addr, region_len)?; - Ok(Self { - inner: SendWrap(Rc::new(RefCell::new(inner))), - }) - } -} - -impl BufferPool { - /// Upper slab slot size in bytes. - pub const fn upper_slot_size() -> usize { - 4096 - } - - /// Lower slab slot size in bytes. - pub const fn lower_slot_size() -> usize { - 256 - } -} - -#[cfg(all(test, loom))] -#[derive(Debug, Clone)] -pub struct BufferPoolSync { - inner: std::sync::Arc>>, -} - -#[cfg(all(test, loom))] -impl BufferPoolSync { - /// Create a new buffer pool over a fixed region. - pub fn new(base_addr: u64, region_len: usize) -> Result { - let inner = Inner::::new(base_addr, region_len)?; - Ok(Self { - inner: std::sync::Arc::new(std::sync::Mutex::new(inner)), - }) - } -} - -impl Inner { - /// Create a new buffer pool over a fixed region. - pub fn new(base_addr: u64, region_len: usize) -> Result { - const LOWER_FRACTION: usize = 8; - - let base = usize::try_from(base_addr).map_err(|_| AllocError::Overflow)?; - let region_end = base.checked_add(region_len).ok_or(AllocError::Overflow)?; - - let lower_base = align_up(base, L)?; - let usable = region_end - .checked_sub(lower_base) - .ok_or(AllocError::EmptyRegion)?; - - let lower_region = usable / LOWER_FRACTION; - let lower = Slab::::new(lower_base as u64, lower_region)?; - - let upper_base = lower_base - .checked_add(lower.capacity()) - .ok_or(AllocError::Overflow)?; + let mut rem = total_len; + let mut sgs = SmallVec::<[Allocation; 4]>::new(); - let upper_base = align_up(upper_base, U)?; - let upper_region = region_end - .checked_sub(upper_base) - .ok_or(AllocError::EmptyRegion)?; - - let upper = Slab::::new(upper_base as u64, upper_region)?; - Ok(Self { lower, upper }) - } - - /// Allocate at least `len` bytes. - pub fn alloc(&mut self, len: usize) -> Result { - if len <= L { - match self.lower.alloc(len) { - Ok(alloc) => return Ok(alloc), - Err(AllocError::NoSpace) => {} - Err(e) => return Err(e), + while rem > 0 { + let len = rem.min(seg_cap); + match self.alloc(len) { + Ok(alloc) => { + sgs.push(alloc); + rem -= len; + } + Err(err) => { + for sg in sgs { + let result = self.dealloc(sg.addr); + debug_assert!(result.is_ok(), "dealloc failed: {result:?}"); + } + return Err(err); + } } } - // fallback to upper slab - self.upper.alloc(len) - } - - /// Free a previously allocated block by its start address. - pub fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { - if self.lower.contains(addr) { - self.lower.dealloc_addr(addr) - } else { - self.upper.dealloc_addr(addr) - } - } - - /// Capacity of a live allocation by its start address. - pub fn allocation_len(&self, addr: u64) -> Result { - if self.lower.contains(addr) { - self.lower.allocation_len(addr) - } else { - self.upper.allocation_len(addr) - } + Ok(sgs) } } -impl BufferProvider for BufferPool { +impl BufferProvider for Rc { fn max_alloc_len(&self) -> usize { - U + (**self).max_alloc_len() } fn alloc(&self, len: usize) -> Result { - self.inner.borrow_mut().alloc(len) - } - - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - Ok(smallvec::smallvec![self.alloc(total_len)?]) + (**self).alloc(len) } fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) + (**self).dealloc(addr) } -} -impl BufferPool { - /// Free a previously allocated block by its start address. - pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) - } - - /// Capacity of a live allocation by its start address. - pub fn allocation_len(&self, addr: u64) -> Result { - self.inner.borrow().allocation_len(addr) + fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { + (**self).alloc_sg(total_len) } } -#[cfg(all(test, loom))] -impl BufferProvider for BufferPoolSync { +impl BufferProvider for Arc { fn max_alloc_len(&self) -> usize { - U + (**self).max_alloc_len() } fn alloc(&self, len: usize) -> Result { - self.inner.lock().expect("poisoned mutex").alloc(len) - } - - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - Ok(smallvec::smallvec![self.alloc(total_len)?]) + (**self).alloc(len) } fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - self.inner - .lock() - .expect("poisoned mutex") - .dealloc_addr(addr) - } -} - -/// Single-tier fixed-slot free list. -/// -/// Tracks a fixed set of equal-sized buffer slots. Allocation pops a free slot -/// and deallocation returns it, both O(1). A [`FixedBitSet`] records which slots -/// are currently allocated, so double frees and frees of unknown addresses are -/// rejected without scanning the free list. -struct RecycleList { - base_addr: u64, - slot_size: usize, - count: usize, - /// Free slot addresses, popped/pushed LIFO. - free: SmallVec<[u64; 64]>, - /// One bit per slot index; set means the slot is currently handed out. - allocated: FixedBitSet, -} - -// SAFETY: only sound for single-threaded (guest-side) access; see the -// type-level invariant on `SendWrap`. -unsafe impl Send for SendWrap>> {} - -impl RecycleList { - fn new(base_addr: u64, region_len: usize, slot_size: usize) -> Result { - if slot_size == 0 { - return Err(AllocError::InvalidArg); - } - - let count = region_len / slot_size; - if count == 0 { - return Err(AllocError::EmptyRegion); - } - - let mut free = SmallVec::with_capacity(count); - for i in 0..count { - free.push(base_addr + (i * slot_size) as u64); - } - - Ok(Self { - base_addr, - slot_size, - count, - free, - allocated: FixedBitSet::with_capacity(count), - }) - } - - fn end(&self) -> u64 { - self.base_addr + (self.count * self.slot_size) as u64 - } - - fn contains(&self, addr: u64) -> bool { - (self.base_addr..self.end()).contains(&addr) - } - - /// Validate that `addr` names a slot start within the region. - fn slot_of(&self, addr: u64) -> Result { - if !self.contains(addr) { - return Err(AllocError::InvalidFree(addr, 0)); - } - - let off = addr - self.base_addr; - if !off.is_multiple_of(self.slot_size as u64) { - return Err(AllocError::InvalidFree(addr, 0)); - } - - Ok((off / self.slot_size as u64) as usize) - } - - /// Validate that `addr` is a live (currently allocated) slot start. - fn live_slot_of(&self, addr: u64) -> Result { - let slot = self.slot_of(addr)?; - if !self.allocated.contains(slot) { - return Err(AllocError::InvalidFree(addr, 0)); - } - Ok(slot) - } - - fn alloc(&mut self, len: usize) -> Result { - if len == 0 { - return Err(AllocError::InvalidArg); - } - if len > self.slot_size { - return Err(AllocError::OutOfMemory); - } - - let addr = self.free.pop().ok_or(AllocError::NoSpace)?; - // Safety of the index: `addr` came from `free`, which only ever holds - // valid slot starts. - self.allocated - .insert(((addr - self.base_addr) / self.slot_size as u64) as usize); - - Ok(Allocation { - addr, - len: self.slot_size, - }) - } - - fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { - let slot = self.live_slot_of(addr)?; - self.allocated.set(slot, false); - self.free.push(addr); - Ok(()) - } - - fn allocation_len(&self, addr: u64) -> Result { - self.live_slot_of(addr)?; - Ok(self.slot_size) - } - - /// Rebuild state so that exactly the addresses in `allocated` are marked - /// live and every other slot is free. - /// - /// On error the pool is left unchanged. - fn restore_allocated(&mut self, allocated: &[u64]) -> Result<(), AllocError> { - let mut restored = FixedBitSet::with_capacity(self.count); - for &addr in allocated { - let slot = self.slot_of(addr)?; - if restored.contains(slot) { - return Err(AllocError::InvalidFree(addr, self.slot_size)); - } - restored.insert(slot); - } - self.allocated = restored; - self.rebuild_free(); - Ok(()) - } - - /// Repopulate the free list with every slot whose allocated bit is clear. - fn rebuild_free(&mut self) { - self.free.clear(); - for i in 0..self.count { - if !self.allocated.contains(i) { - self.free.push(self.base_addr + (i * self.slot_size) as u64); - } - } - } - - fn slot_addr(&self, index: usize) -> Option { - (index < self.count).then(|| self.base_addr + (index * self.slot_size) as u64) + (**self).dealloc(addr) } - fn num_free(&self) -> usize { - self.free.len() + fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { + (**self).alloc_sg(total_len) } } -/// A recycling buffer provider with fixed-size slots. +/// Wrapper asserting `Send` for an inner value that is only ever accessed from +/// a single thread. /// -/// Holds a fixed set of equal-sized buffer addresses in a free list. Alloc and -/// dealloc are O(1). It is intended for bounded scatter/gather descriptor -/// segments that are pre-allocated and recycled after use: -/// [`alloc_sg`](BufferProvider::alloc_sg) splits a logical payload into -/// `ceil(total_len / slot_size)` fixed-size segments. -#[derive(Clone)] -pub struct RecyclePool { - inner: SendWrap>>, -} - -impl RecyclePool { - /// Create a recycling pool of `slot_size`-byte slots over a fixed region. - /// - /// The base address is aligned up to `slot_size`; the slot count is based - /// on the remaining usable region after alignment. - pub fn new(base_addr: u64, region_len: usize, slot_size: usize) -> Result { - if slot_size == 0 { - return Err(AllocError::InvalidArg); - } - - let base = usize::try_from(base_addr).map_err(|_| AllocError::Overflow)?; - let region_end = base.checked_add(region_len).ok_or(AllocError::Overflow)?; - let aligned = align_up(base, slot_size)?; - let usable = region_end - .checked_sub(aligned) - .ok_or(AllocError::EmptyRegion)?; - let list = RecycleList::new(aligned as u64, usable, slot_size)?; - - Ok(Self { - inner: SendWrap(Rc::new(RefCell::new(list))), - }) - } - - /// Rebuild pool state so that every address in `allocated` is removed from - /// the free list, matching externally known inflight state. Validation is - /// transactional: an error leaves the existing pool state unchanged. - pub fn restore_allocated(&self, allocated: &[u64]) -> Result<(), AllocError> { - self.inner.borrow_mut().restore_allocated(allocated) - } - - /// Compute the address of slot `index`. - /// - /// Returns `None` if `index >= count`. - pub fn slot_addr(&self, index: usize) -> Option { - self.inner.borrow().slot_addr(index) - } - - /// Number of free slots. - pub fn num_free(&self) -> usize { - self.inner.borrow().num_free() - } - - /// Free a previously allocated slot by address. - pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) - } - - /// Capacity of a live allocation by its start address. - pub fn allocation_len(&self, addr: u64) -> Result { - self.inner.borrow().allocation_len(addr) - } - - /// Base address of the pool region. - pub fn base_addr(&self) -> u64 { - self.inner.borrow().base_addr - } - - /// Slot size in bytes. - pub fn slot_size(&self) -> usize { - self.inner.borrow().slot_size - } +/// [`RunPool`] and [`SlotPool`] hold their state in an `Rc>`, which +/// is neither `Send` nor `Sync`. Their allocations are exposed as zero-copy +/// reply payloads through [`Bytes::from_owner`](bytes::Bytes::from_owner), +/// whose owner bound is `Send + 'static`; this wrapper exists solely so the +/// pools can satisfy that bound. +/// +/// # Safety +/// +/// The `Send` assertion is only sound while the wrapped value and every +/// `Bytes` handed out from it stay on a single thread. Hyperlight guests are +/// single-threaded, so this holds for guest-side use. It is unsound to move a +/// pool or reply `Bytes` to another thread. +#[derive(Debug)] +struct SendWrap(T); - /// Number of slots in the pool. - pub fn count(&self) -> usize { - self.inner.borrow().count +impl Clone for SendWrap { + fn clone(&self) -> Self { + Self(self.0.clone()) } } -impl BufferProvider for RecyclePool { - fn max_alloc_len(&self) -> usize { - self.inner.borrow().slot_size - } - - fn alloc(&self, len: usize) -> Result { - self.inner.borrow_mut().alloc(len) - } +impl Deref for SendWrap { + type Target = T; - fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) + fn deref(&self) -> &T { + &self.0 } } -#[cfg(test)] -mod tests { - use super::*; - - fn make_pool(size: usize) -> BufferPool { - let base = align_up(0x10000, L.max(U)).unwrap() as u64; - BufferPool::::new(base, size).unwrap() - } - - fn make_recycle_pool(slot_count: usize, slot_size: usize) -> RecyclePool { - let base = 0x80000u64; - RecyclePool::new(base, slot_count * slot_size, slot_size).unwrap() - } - - #[test] - fn test_pool_new_success() { - let pool = BufferPool::<256, 4096>::new(0x10000, 1024 * 1024).unwrap(); - assert!(pool.inner.borrow().lower.capacity() > 0); - assert!(pool.inner.borrow().upper.capacity() > 0); - } - - #[test] - fn test_pool_alloc_small_to_lower() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(128).unwrap(); - - // Should come from lower slab - assert!(pool.inner.borrow().lower.contains(alloc.addr)); - assert_eq!(alloc.len, 256); - } - - #[test] - fn test_pool_alloc_large_to_upper() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(1500).unwrap(); - - // Should come from upper slab - assert!(pool.inner.borrow().upper.contains(alloc.addr)); - assert_eq!(alloc.len, 4096); - } - - #[test] - fn test_pool_alloc_fallback_to_upper() { - let pool = make_pool::<256, 4096>(1024 * 1024); - - // Fill lower slab completely - let mut allocations = Vec::new(); - while pool.inner.borrow().lower.free_bytes() > 0 { - allocations.push(pool.inner.borrow_mut().lower.alloc(256).unwrap()); - } - - // Small allocation should fallback to upper slab - let alloc = pool.alloc(128).unwrap(); - assert!(pool.inner.borrow().upper.contains(alloc.addr)); - } - - #[test] - fn test_pool_free_from_lower() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(128).unwrap(); - - let free_before = pool.inner.borrow().lower.free_bytes(); - pool.dealloc(alloc.addr).unwrap(); - assert_eq!( - pool.inner.borrow().lower.free_bytes(), - free_before + alloc.len - ); - } - - #[test] - fn test_pool_free_from_upper() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(1500).unwrap(); - - let free_before = pool.inner.borrow().upper.free_bytes(); - pool.dealloc(alloc.addr).unwrap(); - assert_eq!( - pool.inner.borrow().upper.free_bytes(), - free_before + alloc.len - ); - } - - #[test] - fn test_pool_stress_many_allocations() { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); - let mut allocations = Vec::new(); - - // Allocate many buffers - for i in 0..100 { - let size = if i % 2 == 0 { 128 } else { 1500 }; - allocations.push(pool.alloc(size).unwrap()); - } - - // Free half of them - for i in (0..100).step_by(2) { - pool.dealloc(allocations[i].addr).unwrap(); - } - - // Should be able to allocate again - for i in 0..50 { - let size = if i % 2 == 0 { 128 } else { 1500 }; - let _alloc = pool.alloc(size).unwrap(); - } - } - - #[test] - fn test_pool_mixed_workload() { - let pool = make_pool::<256, 4096>(2 * 1024 * 1024); - - // Simulate virtio-net workload - let desc_buf = pool.alloc(64).unwrap(); // Control message - let rx_buf1 = pool.alloc(1500).unwrap(); // MTU packet - let rx_buf2 = pool.alloc(1500).unwrap(); // MTU packet - let tx_buf = pool.alloc(4096).unwrap(); // Large buffer - - // Free and reallocate - pool.dealloc(rx_buf1.addr).unwrap(); - let rx_buf3 = pool.alloc(1500).unwrap(); - - // Should reuse freed buffer (LIFO) - assert_eq!(rx_buf3.addr, rx_buf1.addr); - - pool.dealloc(desc_buf.addr).unwrap(); - pool.dealloc(rx_buf2.addr).unwrap(); - pool.dealloc(rx_buf3.addr).unwrap(); - pool.dealloc(tx_buf.addr).unwrap(); - } - - #[test] - fn test_pool_zero_allocation_error() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let result = pool.alloc(0); - assert!(matches!(result, Err(AllocError::InvalidArg))); - } - - #[test] - fn test_pool_too_large_allocation() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let result = pool.alloc(2 * 1024 * 1024); // Larger than pool - assert!(matches!(result, Err(AllocError::OutOfMemory))); - } - - #[test] - fn test_align_up_helper() { - assert_eq!(align_up(0, 256).unwrap(), 0); - assert_eq!(align_up(1, 256).unwrap(), 256); - assert_eq!(align_up(256, 256).unwrap(), 256); - assert_eq!(align_up(257, 256).unwrap(), 512); - assert_eq!(align_up(511, 256).unwrap(), 512); - assert_eq!(align_up(512, 256).unwrap(), 512); - assert!(matches!(align_up(1, 0), Err(AllocError::InvalidArg))); - assert!(matches!( - align_up(usize::MAX, 256), - Err(AllocError::Overflow) - )); - } - - #[test] - fn test_recycle_pool_alignment_subtracts_padding() { - let pool = RecyclePool::new(0x80001, 8192, 4096).unwrap(); - - assert_eq!(pool.base_addr(), 0x81000); - assert_eq!(pool.count(), 1); - } - - // Edge case: allocation exactly at boundary - #[test] - fn test_pool_boundary_allocation() { - let pool = make_pool::<256, 4096>(1024 * 1024); - - // Allocate exactly at boundary - let alloc = pool.alloc(256).unwrap(); - assert!(pool.inner.borrow().lower.contains(alloc.addr)); - - // Allocate just over boundary - let alloc2 = pool.alloc(257).unwrap(); - assert!(pool.inner.borrow().upper.contains(alloc2.addr)); - } - - #[test] - fn test_pool_dealloc_addr_routes_to_correct_tier() { - let pool = make_pool::<256, 4096>(0x20000); - let lower = pool.alloc(128).unwrap(); - let upper = pool.alloc(1024).unwrap(); - - assert_eq!(pool.allocation_len(lower.addr).unwrap(), 256); - assert_eq!(pool.allocation_len(upper.addr).unwrap(), 4096); - - pool.dealloc_addr(lower.addr).unwrap(); - pool.dealloc_addr(upper.addr).unwrap(); - } - - #[test] - fn test_buffer_pool_alloc_sg_uses_one_contiguous_run() { - let pool = make_pool::<256, 4096>(0x20000); - let sgs = pool.alloc_sg(4096 * 2 + 1).unwrap(); - - assert_eq!(sgs.len(), 1); - assert_eq!(sgs[0].len, 4096 * 3); - - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); - } - } - - #[test] - fn test_buffer_pool_alloc_sg_large_run() { - let pool = make_pool::<256, 4096>(0x20000); - let sgs = pool.alloc_sg(8192).unwrap(); - - assert_eq!(sgs.len(), 1); - assert_eq!(sgs[0].len, 8192); - - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); - } - } - - #[test] - fn test_recycle_pool_alloc_sg_splits() { - let pool = make_recycle_pool(8, 4096); - let sgs = pool.alloc_sg(4096 * 2 + 1).unwrap(); - - assert_eq!(sgs.len(), 3); - assert_eq!(sgs[0].len, 4096); - assert_eq!(sgs[1].len, 4096); - assert_eq!(sgs[2].len, 4096); - - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); - } - } - - #[test] - fn test_recycle_pool_restore_allocated_removes_from_free_list() { - let pool = make_recycle_pool(4, 4096); - assert_eq!(pool.num_free(), 4); - - let addrs = [0x80000, 0x81000]; // slots 0 and 1 - pool.restore_allocated(&addrs).unwrap(); - assert_eq!(pool.num_free(), 2); - - // Allocating should only return the two remaining slots - let a1 = pool.alloc(4096).unwrap(); - let a2 = pool.alloc(4096).unwrap(); - assert!(pool.alloc(4096).is_err()); - - // The allocated addresses should be the non-restored ones - let mut got = [a1.addr, a2.addr]; - got.sort(); - assert_eq!(got, [0x82000, 0x83000]); - } - - #[test] - fn test_recycle_pool_restore_allocated_invalid_addr_returns_error() { - let pool = make_recycle_pool(4, 4096); - pool.restore_allocated(&[0x80000]).unwrap(); - - let result = pool.restore_allocated(&[0x81000, 0xDEAD]); - assert!(result.is_err()); - assert_eq!(pool.num_free(), 3); - assert_eq!(pool.allocation_len(0x80000).unwrap(), 4096); - assert!(pool.allocation_len(0x81000).is_err()); - } - - #[test] - fn test_recycle_pool_restore_allocated_then_dealloc_roundtrip() { - let pool = make_recycle_pool(4, 4096); - let addr = 0x81000u64; - - pool.restore_allocated(&[addr]).unwrap(); - assert_eq!(pool.num_free(), 3); - - // Dealloc the restored address - pool.dealloc(addr).unwrap(); - assert_eq!(pool.num_free(), 4); - } - - #[test] - fn test_recycle_pool_restore_allocated_all_slots() { - let pool = make_recycle_pool(4, 4096); - let addrs: Vec = (0..4).map(|i| 0x80000 + i * 4096).collect(); - - pool.restore_allocated(&addrs).unwrap(); - assert_eq!(pool.num_free(), 0); - assert!(pool.alloc(4096).is_err()); - } - - #[test] - fn test_recycle_pool_restore_allocated_empty_list_is_noop() { - let pool = make_recycle_pool(4, 4096); - pool.restore_allocated(&[]).unwrap(); - assert_eq!(pool.num_free(), 4); - } - - #[test] - fn test_recycle_pool_restore_allocated_replaces_state() { - let pool = make_recycle_pool(4, 4096); - - let _ = pool.alloc(4096).unwrap(); - let _ = pool.alloc(4096).unwrap(); - assert_eq!(pool.num_free(), 2); - - pool.restore_allocated(&[0x80000]).unwrap(); - assert_eq!(pool.num_free(), 3); - } - - #[test] - fn test_recycle_pool_dealloc_out_of_range() { - let pool = make_recycle_pool(4, 4096); - let _ = pool.alloc(4096).unwrap(); - - assert!(matches!( - pool.dealloc(0xDEAD), - Err(AllocError::InvalidFree(0xDEAD, 0)) - )); - } - - #[test] - fn test_recycle_pool_dealloc_misaligned() { - let pool = make_recycle_pool(4, 4096); - let _ = pool.alloc(4096).unwrap(); - - assert!(matches!( - pool.dealloc(0x80001), - Err(AllocError::InvalidFree(0x80001, 0)) - )); - } - - #[test] - fn test_recycle_pool_dealloc_double_free() { - let pool = make_recycle_pool(4, 4096); - let a = pool.alloc(4096).unwrap(); - pool.dealloc(a.addr).unwrap(); - - // Second dealloc should fail - address is already in the free list - assert!(matches!( - pool.dealloc(a.addr), - Err(AllocError::InvalidFree(_, _)) - )); - } - - #[test] - fn test_recycle_pool_alloc_sg_rolls_back_on_failure() { - let pool = make_recycle_pool(2, 4096); - - assert!(matches!(pool.alloc_sg(4096 * 3), Err(AllocError::NoSpace))); - assert_eq!(pool.num_free(), 2); - - let alloc = pool.alloc(4096).unwrap(); - assert_eq!(pool.num_free(), 1); - pool.dealloc(alloc.addr).unwrap(); - } - - #[test] - fn test_recycle_pool_dealloc_addr_and_allocation_len() { - let pool = make_recycle_pool(4, 4096); - let alloc = pool.alloc(4096).unwrap(); - - assert_eq!(pool.allocation_len(alloc.addr).unwrap(), 4096); - pool.dealloc_addr(alloc.addr).unwrap(); - assert!(matches!( - pool.allocation_len(alloc.addr), - Err(AllocError::InvalidFree(_, 0)) - )); - } - - #[test] - fn test_recycle_pool_random_order_dealloc() { - let pool = make_recycle_pool(8, 4096); - - let mut allocs: Vec = (0..8).map(|_| pool.alloc(4096).unwrap()).collect(); - assert_eq!(pool.num_free(), 0); - - // Dealloc in reverse order - allocs.reverse(); - for a in &allocs { - pool.dealloc(a.addr).unwrap(); - } - assert_eq!(pool.num_free(), 8); - - // All slots should be re-allocatable - let reallocs: Vec = (0..8).map(|_| pool.alloc(4096).unwrap()).collect(); - assert_eq!(pool.num_free(), 0); - - // Verify all addresses are distinct - let mut addrs: Vec = reallocs.iter().map(|a| a.addr).collect(); - addrs.sort(); - addrs.dedup(); - assert_eq!(addrs.len(), 8); - } - - #[test] - fn test_recycle_pool_interleaved_alloc_dealloc_order() { - let pool = make_recycle_pool(4, 4096); - - let a0 = pool.alloc(4096).unwrap(); - let a1 = pool.alloc(4096).unwrap(); - let a2 = pool.alloc(4096).unwrap(); - let a3 = pool.alloc(4096).unwrap(); - assert_eq!(pool.num_free(), 0); - - // Free middle slots first (out of allocation order) - pool.dealloc(a2.addr).unwrap(); - pool.dealloc(a0.addr).unwrap(); - assert_eq!(pool.num_free(), 2); - - // Re-alloc gets the out-of-order slots back (LIFO) - let b0 = pool.alloc(4096).unwrap(); - assert_eq!(b0.addr, a0.addr); - let b1 = pool.alloc(4096).unwrap(); - assert_eq!(b1.addr, a2.addr); - - // Free everything in yet another order - pool.dealloc(a1.addr).unwrap(); - pool.dealloc(b0.addr).unwrap(); - pool.dealloc(b1.addr).unwrap(); - pool.dealloc(a3.addr).unwrap(); - assert_eq!(pool.num_free(), 4); - - // All 4 original addresses should be available - let mut final_addrs: Vec = (0..4).map(|_| pool.alloc(4096).unwrap().addr).collect(); - final_addrs.sort(); - let expected: Vec = (0..4).map(|i| 0x80000 + i * 4096).collect(); - assert_eq!(final_addrs, expected); +#[inline] +fn align_up(val: usize, align: usize) -> Result { + if align == 0 { + return Err(AllocError::InvalidArg); } - #[test] - fn test_recycle_pool_dealloc_order_independent_of_alloc_order() { - let pool = make_recycle_pool(6, 256); - - // Allocate all - let allocs: Vec = (0..6).map(|_| pool.alloc(256).unwrap()).collect(); - - // Dealloc in scattered order: 4, 1, 5, 0, 3, 2 - let order = [4, 1, 5, 0, 3, 2]; - for &i in &order { - pool.dealloc(allocs[i].addr).unwrap(); - } - assert_eq!(pool.num_free(), 6); - - // Re-allocate all and verify we get back the full set - let mut realloc_addrs: Vec = (0..6).map(|_| pool.alloc(256).unwrap().addr).collect(); - realloc_addrs.sort(); - - let mut orig_addrs: Vec = allocs.iter().map(|a| a.addr).collect(); - orig_addrs.sort(); - - assert_eq!(realloc_addrs, orig_addrs); - } + val.checked_next_multiple_of(align) + .ok_or(AllocError::Overflow) } #[cfg(test)] -mod fuzz { - use quickcheck::{Arbitrary, Gen, QuickCheck}; +mod tests; - use super::*; - - const MAX_OPS: usize = 10; - const MAX_ALLOC_SIZE: usize = 8192; - - #[derive(Clone, Debug)] - enum Op { - Alloc(usize), - AllocSg(usize), - Dealloc(usize), - } - - impl Arbitrary for Op { - fn arbitrary(g: &mut Gen) -> Self { - match u8::arbitrary(g) % 3 { - 0 => Op::Alloc(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), - 1 => Op::AllocSg(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), - 2 => Op::Dealloc(usize::arbitrary(g)), - _ => unreachable!(), - } - } - } - - #[derive(Clone, Debug)] - struct Scenario { - pool_size: usize, - ops: Vec, - } - - impl Arbitrary for Scenario { - fn arbitrary(g: &mut Gen) -> Self { - let pool_size = (usize::arbitrary(g) % (4 * 1024 * 1024)) + (1024 * 1024); - let num_ops = usize::arbitrary(g) % MAX_OPS + 1; - let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); - - Scenario { pool_size, ops } - } - } - - fn run_scenario(s: Scenario) -> bool { - let base = align_up(0x10000, 4096).unwrap() as u64; - let pool = match BufferPool::<256, 4096>::new(base, s.pool_size) { - Ok(p) => p, - Err(_) => return true, - }; - - let mut allocations: Vec = Vec::new(); - - for op in &s.ops { - match op { - Op::Alloc(size) => match pool.alloc(*size) { - Ok(alloc) => { - assert!(alloc.len >= *size); - allocations.push(alloc); - } - Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} - Err(_) => { - return false; - } - }, - Op::AllocSg(size) => match pool.alloc_sg(*size) { - Ok(sgs) => { - let total: usize = sgs.iter().map(|sg| sg.len).sum(); - assert!(total >= *size); - allocations.extend(sgs); - } - Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} - Err(_) => { - return false; - } - }, - Op::Dealloc(idx) => { - if allocations.is_empty() { - continue; - } - - let idx = idx % allocations.len(); - let alloc = allocations.swap_remove(idx); - - match pool.dealloc(alloc.addr) { - Ok(_) => {} - Err(_) => return false, - } - } - } - - if check_pool_invariants(&pool, &allocations).is_err() { - return false; - } - } - - // Cleanup - for alloc in &allocations { - if pool.dealloc(alloc.addr).is_err() { - return false; - } - } - - check_pool_invariants(&pool, &allocations).is_ok() - } - - fn check_slab_invariants(slab: &Slab) -> Result<(), &'static str> { - let used = slab.used_slots.count_ones(..); - let free = slab.used_slots.count_zeroes(..); - if used + free != slab.used_slots.len() { - return Err("used + free != total slots"); - } - - let expected_free = free * N; - if slab.free_bytes() != expected_free { - return Err("free_bytes doesn't match bitmap"); - } - - if let Some(alloc) = slab.last_free_run { - if alloc.len == 0 || alloc.len % N != 0 { - return Err("last_free_run has invalid length"); - } - if !slab.contains(alloc.addr) { - return Err("last_free_run addr outside range"); - } - } - - Ok(()) - } - - fn check_pool_invariants( - pool: &BufferPool, - allocations: &[Allocation], - ) -> Result<(), &'static str> { - check_slab_invariants(&pool.inner.borrow().lower)?; - check_slab_invariants(&pool.inner.borrow().upper)?; - - if pool.inner.borrow().lower.range().end > pool.inner.borrow().upper.range().start { - return Err("lower and upper ranges overlap"); - } - - let mut seen = std::collections::HashSet::new(); - - for alloc in allocations { - if !pool.inner.borrow().lower.contains(alloc.addr) - && !pool.inner.borrow().upper.contains(alloc.addr) - { - return Err("allocation address outside pool ranges"); - } - - if alloc.len % L != 0 && alloc.len % U != 0 { - return Err("allocation length not aligned to any tier"); - } - - if !seen.insert(alloc.addr) { - return Err("duplicate allocation address in tracking"); - } - } - - Ok(()) - } - - #[test] - fn prop_allocator_invariants() { - #[cfg(miri)] - let tests = 10; - #[cfg(not(miri))] - let tests = 1000; - - QuickCheck::new() - .tests(tests) - .quickcheck(run_scenario as fn(Scenario) -> bool); - } -} +#[cfg(test)] +mod fuzz; diff --git a/src/hyperlight_common/src/virtq/pool/fuzz.rs b/src/hyperlight_common/src/virtq/pool/fuzz.rs new file mode 100644 index 000000000..22cd1c934 --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/fuzz.rs @@ -0,0 +1,369 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +use std::collections::{BTreeMap, HashSet}; + +use quickcheck::{Arbitrary, Gen, QuickCheck}; + +use super::run::Tier as RunTier; +use super::*; + +const MAX_OPS: usize = 10; +const MAX_ALLOC_SIZE: usize = 8192; +const MAX_TIER_SLOTS: usize = 16; +const LOWER_BASE: u64 = 0x80000; +const UPPER_BASE: u64 = 0x90000; +const LOWER_SLOT_SIZE: usize = 256; +const UPPER_SLOT_SIZE: usize = 4096; + +#[derive(Clone, Debug)] +enum Op { + Alloc(usize), + AllocSg(usize), + Dealloc(usize), +} + +impl Arbitrary for Op { + fn arbitrary(g: &mut Gen) -> Self { + match u8::arbitrary(g) % 3 { + 0 => Op::Alloc(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), + 1 => Op::AllocSg(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), + 2 => Op::Dealloc(usize::arbitrary(g)), + _ => unreachable!(), + } + } +} + +#[derive(Clone, Debug)] +struct RunScenario { + pool_size: usize, + ops: Vec, +} + +impl Arbitrary for RunScenario { + fn arbitrary(g: &mut Gen) -> Self { + let pool_size = (usize::arbitrary(g) % (4 * 1024 * 1024)) + (1024 * 1024); + let num_ops = usize::arbitrary(g) % MAX_OPS + 1; + let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); + + RunScenario { pool_size, ops } + } +} + +fn run_provider_ops(pool: &P, ops: &[Op], check: F) -> bool +where + P: BufferProvider, + F: Fn(&P, &[Allocation]) -> Result<(), &'static str>, +{ + let mut allocations: Vec = Vec::new(); + + for op in ops { + match op { + Op::Alloc(size) => match pool.alloc(*size) { + Ok(alloc) => { + if alloc.len < *size + || allocations + .iter() + .any(|existing| existing.addr == alloc.addr) + { + return false; + } + allocations.push(alloc); + } + Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} + Err(_) => return false, + }, + Op::AllocSg(size) => match pool.alloc_sg(*size) { + Ok(sgs) => { + let mut total = 0usize; + for sg in sgs { + let Some(next_total) = total.checked_add(sg.len) else { + return false; + }; + if allocations.iter().any(|existing| existing.addr == sg.addr) { + return false; + } + total = next_total; + allocations.push(sg); + } + if total < *size { + return false; + } + } + Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} + Err(_) => return false, + }, + Op::Dealloc(index) => { + if !allocations.is_empty() { + let index = index % allocations.len(); + if pool.dealloc(allocations[index].addr).is_err() { + return false; + } + allocations.swap_remove(index); + } + } + } + + if check(pool, &allocations).is_err() { + return false; + } + } + + while let Some(alloc) = allocations.pop() { + if pool.dealloc(alloc.addr).is_err() { + return false; + } + } + + check(pool, &allocations).is_ok() +} + +fn run_pool_scenario(scenario: RunScenario) -> bool { + let base = align_up(0x10000, 4096).unwrap() as u64; + let pool = match RunPool::<256, 4096>::new(base, scenario.pool_size) { + Ok(pool) => pool, + Err(_) => return true, + }; + + run_provider_ops(&pool, &scenario.ops, check_run_pool_invariants) +} + +fn check_run_tier_invariants( + tier: &RunTier, + allocations: &[Allocation], +) -> Result<(), &'static str> { + let mut expected_used = HashSet::new(); + let mut expected_starts = HashSet::new(); + + for alloc in allocations.iter().filter(|alloc| tier.contains(alloc.addr)) { + let offset = usize::try_from(alloc.addr - tier.base_addr) + .map_err(|_| "allocation offset overflows usize")?; + if alloc.len == 0 || !offset.is_multiple_of(N) || !alloc.len.is_multiple_of(N) { + return Err("allocation is not tier-aligned"); + } + + let start = offset / N; + let slots = alloc.len / N; + let end = start + .checked_add(slots) + .ok_or("allocation slot range overflow")?; + if end > tier.used_slots.len() || !expected_starts.insert(start) { + return Err("allocation run is invalid"); + } + for slot in start..end { + if !expected_used.insert(slot) { + return Err("allocation runs overlap"); + } + } + } + + for slot in 0..tier.used_slots.len() { + if tier.used_slots.contains(slot) != expected_used.contains(&slot) { + return Err("used bitmap does not match live allocations"); + } + if tier.run_starts.contains(slot) != expected_starts.contains(&slot) { + return Err("run-start bitmap does not match live allocations"); + } + } + + if tier.free_bytes() != (tier.used_slots.len() - expected_used.len()) * N { + return Err("free_bytes does not match live allocations"); + } + + if let Some(free_run) = tier.last_free_run { + if !tier.contains(free_run.addr) { + return Err("cached free-run address outside tier"); + } + let offset = usize::try_from(free_run.addr - tier.base_addr) + .map_err(|_| "cached free-run offset overflows usize")?; + if free_run.len == 0 || !offset.is_multiple_of(N) || !free_run.len.is_multiple_of(N) { + return Err("cached free run is not tier-aligned"); + } + + let start = offset / N; + let end = start + .checked_add(free_run.len / N) + .ok_or("cached free-run range overflow")?; + if end > tier.used_slots.len() || (start..end).any(|slot| tier.used_slots.contains(slot)) { + return Err("cached free run overlaps live allocations"); + } + } + + Ok(()) +} + +fn check_run_pool_invariants( + pool: &RunPool, + allocations: &[Allocation], +) -> Result<(), &'static str> { + let inner = pool.inner.borrow(); + if inner.lower.range().end > inner.upper.range().start { + return Err("lower and upper ranges overlap"); + } + + let mut seen = HashSet::new(); + for alloc in allocations { + let in_lower = inner.lower.contains(alloc.addr); + let in_upper = inner.upper.contains(alloc.addr); + if in_lower == in_upper { + return Err("allocation does not belong to exactly one tier"); + } + if !seen.insert(alloc.addr) { + return Err("duplicate allocation address in tracking"); + } + } + + check_run_tier_invariants(&inner.lower, allocations)?; + check_run_tier_invariants(&inner.upper, allocations) +} + +#[derive(Clone, Debug)] +struct SlotScenario { + tiered: bool, + lower_count: usize, + upper_count: usize, + ops: Vec, +} + +impl Arbitrary for SlotScenario { + fn arbitrary(g: &mut Gen) -> Self { + let tiered = bool::arbitrary(g); + let lower_count = usize::arbitrary(g) % MAX_TIER_SLOTS + 1; + let upper_count = usize::arbitrary(g) % MAX_TIER_SLOTS + 1; + let num_ops = usize::arbitrary(g) % MAX_OPS + 1; + let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); + + Self { + tiered, + lower_count, + upper_count, + ops, + } + } +} + +fn make_slot_pool(scenario: &SlotScenario) -> SlotPool { + if scenario.tiered { + let lower = SlotLayout::new(LOWER_BASE, LOWER_SLOT_SIZE, scenario.lower_count); + let upper = SlotLayout::new(UPPER_BASE, UPPER_SLOT_SIZE, scenario.upper_count); + SlotPool::new_tiered(lower, upper).unwrap() + } else { + let layout = SlotLayout::new(UPPER_BASE, UPPER_SLOT_SIZE, scenario.upper_count); + SlotPool::new(layout).unwrap() + } +} + +fn run_slot_pool_scenario(scenario: SlotScenario) -> bool { + let pool = make_slot_pool(&scenario); + run_provider_ops(&pool, &scenario.ops, check_slot_pool_invariants) +} + +fn layout_contains(layout: SlotLayout, addr: u64) -> bool { + let Ok(end) = layout.end_addr() else { + return false; + }; + (layout.base_addr..end).contains(&addr) +} + +fn slot_capacity(pool: &SlotPool, addr: u64) -> Option { + let (lower, upper) = pool.layouts(); + if let Some(lower) = lower + && layout_contains(lower, addr) + { + return Some(lower.slot_size); + } + layout_contains(upper, addr).then_some(upper.slot_size) +} + +fn check_slot_pool_invariants( + pool: &SlotPool, + allocations: &[Allocation], +) -> Result<(), &'static str> { + let mut expected_live = BTreeMap::new(); + for alloc in allocations { + if expected_live.insert(alloc.addr, alloc.len).is_some() { + return Err("duplicate allocation address in tracking"); + } + } + + let live = pool.live_addrs(); + let expected_addrs: Vec = expected_live.keys().copied().collect(); + if live != expected_addrs || live.windows(2).any(|pair| pair[0] >= pair[1]) { + return Err("live addresses are not unique and deterministic"); + } + if pool.num_free() + live.len() != pool.count() { + return Err("free + live != total slots"); + } + + let (lower, upper) = pool.layouts(); + let expected_base = lower.map_or(upper.base_addr, |layout| layout.base_addr); + if pool.base_addr() != expected_base || pool.slot_size() != upper.slot_size { + return Err("reported pool layout is inconsistent"); + } + + let mut expected_count = upper.slot_count; + if let Some(lower) = lower { + if lower.slot_size >= upper.slot_size + || lower.end_addr().map_err(|_| "lower layout overflow")? > upper.base_addr + { + return Err("tier layout is invalid"); + } + expected_count += lower.slot_count; + } + if pool.count() != expected_count || pool.slot_addr(pool.count()).is_some() { + return Err("reported slot count is inconsistent"); + } + + let mut seen = HashSet::new(); + for index in 0..pool.count() { + let Some(addr) = pool.slot_addr(index) else { + return Err("missing slot address"); + }; + if !seen.insert(addr) { + return Err("duplicate slot address"); + } + let Some(capacity) = slot_capacity(pool, addr) else { + return Err("slot address outside layout"); + }; + + match expected_live.get(&addr) { + Some(expected_capacity) => { + if *expected_capacity != capacity + || pool.allocation_len(addr).ok() != Some(capacity) + { + return Err("live slot capacity is inconsistent"); + } + } + None if pool.allocation_len(addr).is_ok() => { + return Err("free slot reported as live"); + } + None => {} + } + } + + Ok(()) +} + +#[test] +fn prop_run_pool_invariants() { + #[cfg(miri)] + let tests = 10; + #[cfg(not(miri))] + let tests = 1000; + + QuickCheck::new() + .tests(tests) + .quickcheck(run_pool_scenario as fn(RunScenario) -> bool); +} + +#[test] +fn prop_slot_pool_invariants() { + #[cfg(miri)] + let tests = 10; + #[cfg(not(miri))] + let tests = 1000; + + QuickCheck::new() + .tests(tests) + .quickcheck(run_slot_pool_scenario as fn(SlotScenario) -> bool); +} diff --git a/src/hyperlight_common/src/virtq/pool/run.rs b/src/hyperlight_common/src/virtq/pool/run.rs new file mode 100644 index 000000000..9ab69eadb --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/run.rs @@ -0,0 +1,366 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. +//! Variable-sized contiguous-run pool. +//! +//! [`RunPool`] partitions one backing region into lower and upper tiers with +//! compile-time slot sizes. The lower tier is carved from the first eighth of +//! the aligned usable region. Eligible requests try that tier first and fall +//! back to the upper tier only when no contiguous lower-tier run is available. +//! +//! Allocations are rounded to a tier's slot size and occupy contiguous runs, so +//! scatter/gather requests produce one allocation. Occupied-slot and run-start +//! bitmaps support reclaiming a complete run by its start address, while a +//! cached free run accelerates immediate reuse. This preserves contiguous +//! buffers but remains subject to fragmentation. + +use alloc::rc::Rc; +use core::cell::RefCell; + +use fixedbitset::FixedBitSet; +use smallvec::SmallVec; + +use super::{AllocError, Allocation, BufferProvider, SendWrap, align_up}; + +#[derive(Debug, Clone)] +pub(super) struct Tier { + pub(super) base_addr: u64, + pub(super) used_slots: FixedBitSet, + pub(super) run_starts: FixedBitSet, + pub(super) last_free_run: Option, +} + +impl Tier { + fn new(base_addr: u64, region_len: usize) -> Result { + let usable = region_len - (region_len % N); + let num_slots = usable / N; + let used_slots = FixedBitSet::with_capacity(num_slots); + let run_starts = FixedBitSet::with_capacity(num_slots); + + if !base_addr.is_multiple_of(N as u64) { + return Err(AllocError::InvalidAlign(base_addr)); + } + if num_slots == 0 { + return Err(AllocError::EmptyRegion); + } + + Ok(Self { + base_addr, + used_slots, + run_starts, + last_free_run: None, + }) + } + + fn addr_of(&self, slot_idx: usize) -> Option { + self.base_addr + .checked_add((slot_idx as u64).checked_mul(N as u64)?) + } + + fn slot_of(&self, addr: u64) -> usize { + let off = (addr - self.base_addr) as usize; + off / N + } + + fn checked_slot_of(&self, addr: u64, len: usize) -> Result { + if addr < self.base_addr { + return Err(AllocError::InvalidFree(addr, len)); + } + + let off = (addr - self.base_addr) as usize; + if !off.is_multiple_of(N) { + return Err(AllocError::InvalidFree(addr, len)); + } + + let slot = off / N; + if slot >= self.used_slots.len() { + return Err(AllocError::InvalidFree(addr, len)); + } + + Ok(slot) + } + + fn live_run_slots_at(&self, start: usize) -> Option { + if start >= self.used_slots.len() + || !self.used_slots.contains(start) + || !self.run_starts.contains(start) + { + return None; + } + + let mut end = start + 1; + while end < self.used_slots.len() + && self.used_slots.contains(end) + && !self.run_starts.contains(end) + { + end += 1; + } + + Some(end - start) + } + + fn maybe_invalidate_last_run(&mut self, alloc: Allocation) { + if let Some(run) = &self.last_free_run { + let new_end = alloc.addr + alloc.len as u64; + let run_end = run.addr + run.len as u64; + + if alloc.addr < run_end && run.addr < new_end { + self.last_free_run = None; + } + } + } + + fn find_slots(&mut self, slots_num: usize) -> Option { + debug_assert!(slots_num > 0); + + if let Some(alloc) = self.last_free_run + && alloc.len >= slots_num * N + { + let pos = self.slot_of(alloc.addr); + let _ = self.last_free_run.take(); + return Some(pos); + } + + let total = self.used_slots.len(); + self.used_slots.zeroes().find(|&next_free| { + let end = next_free + slots_num; + end <= total && self.used_slots.count_zeroes(next_free..end) == slots_num + }) + } + + pub(super) fn alloc(&mut self, len: usize) -> Result { + if len == 0 { + return Err(AllocError::InvalidArg); + } + + let total = self.used_slots.len(); + let need_slots = len.div_ceil(N); + if need_slots > total { + return Err(AllocError::OutOfMemory); + } + + let idx = self.find_slots(need_slots).ok_or(AllocError::NoSpace)?; + self.used_slots.insert_range(idx..idx + need_slots); + self.run_starts.insert(idx); + let addr = self.addr_of(idx).ok_or(AllocError::Overflow)?; + + let alloc = Allocation { + addr, + len: need_slots * N, + }; + + self.maybe_invalidate_last_run(alloc); + Ok(alloc) + } + + fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { + let start = self.checked_slot_of(addr, 0)?; + let run_slots = self + .live_run_slots_at(start) + .ok_or(AllocError::InvalidFree(addr, 0))?; + self.dealloc_run(start, run_slots, addr) + } + + fn dealloc_run(&mut self, start: usize, run_slots: usize, addr: u64) -> Result<(), AllocError> { + let len = run_slots * N; + self.used_slots.remove_range(start..start + run_slots); + self.run_starts.set(start, false); + self.last_free_run = Some(Allocation { addr, len }); + Ok(()) + } + + fn allocation_len(&self, addr: u64) -> Result { + let start = self.checked_slot_of(addr, 0)?; + let run_slots = self + .live_run_slots_at(start) + .ok_or(AllocError::InvalidFree(addr, 0))?; + Ok(run_slots * N) + } + + pub(super) fn capacity(&self) -> usize { + self.used_slots.len() * N + } + + pub(super) fn range(&self) -> core::ops::Range { + self.base_addr..self.base_addr + self.capacity() as u64 + } + + pub(super) fn contains(&self, addr: u64) -> bool { + self.range().contains(&addr) + } +} + +#[cfg(test)] +impl Tier { + pub(super) fn free_bytes(&self) -> usize { + (self.used_slots.len() - self.used_slots.count_ones(..)) * N + } +} + +#[derive(Debug)] +pub(super) struct Inner { + pub(super) lower: Tier, + pub(super) upper: Tier, +} + +// SAFETY: only sound for single-threaded (guest-side) access; see the +// type-level invariant on `SendWrap`. +unsafe impl Send for SendWrap>>> {} + +/// Two-tier pool for variable-sized contiguous runs. +#[derive(Debug, Clone)] +pub struct RunPool { + pub(super) inner: SendWrap>>>, +} + +impl RunPool { + /// Create a new run pool over a fixed region. + pub fn new(base_addr: u64, region_len: usize) -> Result { + let inner = Inner::::new(base_addr, region_len)?; + Ok(Self { + inner: SendWrap(Rc::new(RefCell::new(inner))), + }) + } +} + +impl RunPool { + /// Upper tier slot size in bytes. + pub const fn upper_slot_size() -> usize { + 4096 + } + + /// Lower tier slot size in bytes. + pub const fn lower_slot_size() -> usize { + 256 + } +} + +#[cfg(all(test, loom))] +#[derive(Debug, Clone)] +pub struct RunPoolSync { + inner: std::sync::Arc>>, +} + +#[cfg(all(test, loom))] +impl RunPoolSync { + /// Create a new synchronized run pool over a fixed region. + pub fn new(base_addr: u64, region_len: usize) -> Result { + let inner = Inner::::new(base_addr, region_len)?; + Ok(Self { + inner: std::sync::Arc::new(std::sync::Mutex::new(inner)), + }) + } +} + +impl Inner { + /// Create new run-pool state over a fixed region. + pub fn new(base_addr: u64, region_len: usize) -> Result { + const LOWER_FRACTION: usize = 8; + + let base = usize::try_from(base_addr).map_err(|_| AllocError::Overflow)?; + let region_end = base.checked_add(region_len).ok_or(AllocError::Overflow)?; + + let lower_base = align_up(base, L)?; + let usable = region_end + .checked_sub(lower_base) + .ok_or(AllocError::EmptyRegion)?; + + let lower_region = usable / LOWER_FRACTION; + let lower = Tier::::new(lower_base as u64, lower_region)?; + + let upper_base = lower_base + .checked_add(lower.capacity()) + .ok_or(AllocError::Overflow)?; + + let upper_base = align_up(upper_base, U)?; + let upper_region = region_end + .checked_sub(upper_base) + .ok_or(AllocError::EmptyRegion)?; + + let upper = Tier::::new(upper_base as u64, upper_region)?; + Ok(Self { lower, upper }) + } + + /// Allocate at least `len` bytes. + pub fn alloc(&mut self, len: usize) -> Result { + if len <= L { + match self.lower.alloc(len) { + Ok(alloc) => return Ok(alloc), + Err(AllocError::NoSpace) => {} + Err(e) => return Err(e), + } + } + + // Fall back to the upper tier. + self.upper.alloc(len) + } + + /// Free a previously allocated block by its start address. + pub fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { + if self.lower.contains(addr) { + self.lower.dealloc_addr(addr) + } else { + self.upper.dealloc_addr(addr) + } + } + + /// Capacity of a live allocation by its start address. + pub fn allocation_len(&self, addr: u64) -> Result { + if self.lower.contains(addr) { + self.lower.allocation_len(addr) + } else { + self.upper.allocation_len(addr) + } + } +} + +impl BufferProvider for RunPool { + fn max_alloc_len(&self) -> usize { + U + } + + fn alloc(&self, len: usize) -> Result { + self.inner.borrow_mut().alloc(len) + } + + fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { + Ok(smallvec::smallvec![self.alloc(total_len)?]) + } + + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } +} + +impl RunPool { + /// Free a previously allocated block by its start address. + pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } + + /// Capacity of a live allocation by its start address. + pub fn allocation_len(&self, addr: u64) -> Result { + self.inner.borrow().allocation_len(addr) + } +} + +#[cfg(all(test, loom))] +impl BufferProvider for RunPoolSync { + fn max_alloc_len(&self) -> usize { + U + } + + fn alloc(&self, len: usize) -> Result { + self.inner.lock().expect("poisoned mutex").alloc(len) + } + + fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { + Ok(smallvec::smallvec![self.alloc(total_len)?]) + } + + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { + self.inner + .lock() + .expect("poisoned mutex") + .dealloc_addr(addr) + } +} diff --git a/src/hyperlight_common/src/virtq/pool/slot.rs b/src/hyperlight_common/src/virtq/pool/slot.rs new file mode 100644 index 000000000..8e8a50535 --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/slot.rs @@ -0,0 +1,379 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! Fixed-slot pool with optional lower and required upper tiers. +//! +//! [`SlotPool`] manages one or two non-overlapping [`SlotLayout`]s. Each tier +//! contains independent, equal-sized slots tracked by a free list and an +//! allocation bitmap. Eligible requests try the lower tier first and fall back +//! to the upper tier only when the lower tier has no free slot. +//! +//! Slots need not be contiguous, so scatter/gather allocation splits a logical +//! buffer at the upper-tier slot size and may place an eligible final segment +//! in the lower tier. [`SlotPool::live_addrs`] reports ownership in deterministic +//! lower-then-upper index order without mutating pool state. + +use alloc::rc::Rc; +use alloc::vec::Vec; +use core::cell::RefCell; + +use fixedbitset::FixedBitSet; +use smallvec::SmallVec; + +use super::{AllocError, Allocation, BufferProvider, SendWrap}; + +/// Exact memory layout for one [`SlotPool`] tier. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SlotLayout { + /// Start of the first slot. + pub base_addr: u64, + /// Capacity of each slot. + pub slot_size: usize, + /// Number of slots. + pub slot_count: usize, +} + +impl SlotLayout { + /// Describe exact fixed-slot placement. + pub const fn new(base_addr: u64, slot_size: usize, slot_count: usize) -> Self { + Self { + base_addr, + slot_size, + slot_count, + } + } + + /// Total bytes occupied by the slots. + pub fn byte_len(self) -> Result { + self.slot_size + .checked_mul(self.slot_count) + .ok_or(AllocError::Overflow) + } + + /// Exclusive end address. + pub fn end_addr(self) -> Result { + self.base_addr + .checked_add(u64::try_from(self.byte_len()?).map_err(|_| AllocError::Overflow)?) + .ok_or(AllocError::Overflow) + } +} + +/// Single-tier fixed-slot free list. +/// +/// Tracks a fixed set of equal-sized buffer slots. Allocation pops a free slot +/// and deallocation returns it, both O(1). A [`FixedBitSet`] records which slots +/// are currently allocated, so double frees and frees of unknown addresses are +/// rejected without scanning the free list. +struct Tier { + /// Start of this tier's backing memory. + base_addr: u64, + /// Capacity of this slot. + slot_size: usize, + /// Number of slots in this tier. + count: usize, + /// Free slot addresses, popped/pushed LIFO. + free: SmallVec<[u64; 64]>, + /// One bit per slot index; set means the slot is currently handed out. + allocated: FixedBitSet, +} + +// SAFETY: only sound for single-threaded (guest-side) access; see the +// type-level invariant on `SendWrap`. +unsafe impl Send for SendWrap>> {} + +impl Tier { + fn from_layout(layout: SlotLayout) -> Result { + if layout.slot_size == 0 { + return Err(AllocError::InvalidArg); + } + + if layout.slot_count == 0 { + return Err(AllocError::EmptyRegion); + } + + layout.end_addr()?; + + let mut free = SmallVec::with_capacity(layout.slot_count); + for i in 0..layout.slot_count { + free.push(layout.base_addr + (i * layout.slot_size) as u64); + } + + Ok(Self { + base_addr: layout.base_addr, + slot_size: layout.slot_size, + count: layout.slot_count, + free, + allocated: FixedBitSet::with_capacity(layout.slot_count), + }) + } + + fn end(&self) -> u64 { + self.base_addr + (self.count * self.slot_size) as u64 + } + + fn contains(&self, addr: u64) -> bool { + (self.base_addr..self.end()).contains(&addr) + } + + /// Validate that `addr` names a slot start within the region. + fn slot_of(&self, addr: u64) -> Result { + if !self.contains(addr) { + return Err(AllocError::InvalidFree(addr, 0)); + } + + let off = addr - self.base_addr; + if !off.is_multiple_of(self.slot_size as u64) { + return Err(AllocError::InvalidFree(addr, 0)); + } + + Ok((off / self.slot_size as u64) as usize) + } + + /// Validate that `addr` is a live (currently allocated) slot start. + fn live_slot_of(&self, addr: u64) -> Result { + let slot = self.slot_of(addr)?; + if !self.allocated.contains(slot) { + return Err(AllocError::InvalidFree(addr, 0)); + } + Ok(slot) + } + + fn alloc(&mut self, len: usize) -> Result { + if len == 0 { + return Err(AllocError::InvalidArg); + } + if len > self.slot_size { + return Err(AllocError::OutOfMemory); + } + + let addr = self.free.pop().ok_or(AllocError::NoSpace)?; + // Safety of the index: `addr` came from `free`, which only ever holds + // valid slot starts. + self.allocated + .insert(((addr - self.base_addr) / self.slot_size as u64) as usize); + + Ok(Allocation { + addr, + len: self.slot_size, + }) + } + + fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { + let slot = self.live_slot_of(addr)?; + self.allocated.set(slot, false); + self.free.push(addr); + Ok(()) + } + + fn allocation_len(&self, addr: u64) -> Result { + self.live_slot_of(addr)?; + Ok(self.slot_size) + } + + fn slot_addr(&self, index: usize) -> Option { + (index < self.count).then(|| self.base_addr + (index * self.slot_size) as u64) + } + + fn num_free(&self) -> usize { + self.free.len() + } + + fn append_live_addrs(&self, addrs: &mut Vec) { + addrs.extend( + self.allocated + .ones() + .map(|slot| self.base_addr + (slot * self.slot_size) as u64), + ); + } + + fn layout(&self) -> SlotLayout { + SlotLayout::new(self.base_addr, self.slot_size, self.count) + } +} + +struct Inner { + lower: Option, + upper: Tier, +} + +impl Inner { + fn new(lower: Option, upper: SlotLayout) -> Result { + let lower = lower.map(Tier::from_layout).transpose()?; + let upper = Tier::from_layout(upper)?; + + if let Some(lower) = &lower + && (lower.slot_size >= upper.slot_size || lower.end() > upper.base_addr) + { + return Err(AllocError::InvalidArg); + } + + Ok(Self { lower, upper }) + } + + fn max_alloc_len(&self) -> usize { + self.upper.slot_size + } + + fn alloc(&mut self, len: usize) -> Result { + if let Some(lower) = &mut self.lower + && len <= lower.slot_size + { + match lower.alloc(len) { + Ok(alloc) => return Ok(alloc), + Err(AllocError::NoSpace) => {} + Err(err) => return Err(err), + } + } + + self.upper.alloc(len) + } + + fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { + if let Some(lower) = &mut self.lower + && lower.contains(addr) + { + return lower.dealloc_addr(addr); + } + self.upper.dealloc_addr(addr) + } + + fn allocation_len(&self, addr: u64) -> Result { + if let Some(lower) = &self.lower + && lower.contains(addr) + { + return lower.allocation_len(addr); + } + self.upper.allocation_len(addr) + } + + fn slot_addr(&self, index: usize) -> Option { + if let Some(lower) = &self.lower { + if index < lower.count { + return lower.slot_addr(index); + } + return self.upper.slot_addr(index - lower.count); + } + self.upper.slot_addr(index) + } + + fn live_addrs(&self) -> Vec { + let mut addrs = Vec::with_capacity(self.count() - self.num_free()); + if let Some(lower) = &self.lower { + lower.append_live_addrs(&mut addrs); + } + self.upper.append_live_addrs(&mut addrs); + addrs + } + + fn base_addr(&self) -> u64 { + self.lower + .as_ref() + .map_or(self.upper.base_addr, |lower| lower.base_addr) + } + + fn count(&self) -> usize { + self.lower.as_ref().map_or(0, |lower| lower.count) + self.upper.count + } + + fn num_free(&self) -> usize { + self.lower.as_ref().map_or(0, Tier::num_free) + self.upper.num_free() + } + + fn layouts(&self) -> (Option, SlotLayout) { + (self.lower.as_ref().map(Tier::layout), self.upper.layout()) + } +} + +/// A buffer pool with one or two fixed-slot tiers. +/// +/// Allocation and deallocation are O(1) per slot. Eligible allocations first +/// try the optional lower tier and fall back to the required upper tier when +/// the lower tier is full. [`alloc_sg`](BufferProvider::alloc_sg) splits logical +/// payloads into bounded descriptor segments. +#[derive(Clone)] +pub struct SlotPool { + inner: SendWrap>>, +} + +impl SlotPool { + /// Create a single-tier recycling pool from exact slot placement. + pub fn new(layout: SlotLayout) -> Result { + Self::from_layouts(None, layout) + } + + /// Create a two-tier recycling pool from exact lower and upper layouts. + /// + /// The lower layout must precede the upper layout without overlap, and its + /// slot size must be strictly smaller. + pub fn new_tiered(lower: SlotLayout, upper: SlotLayout) -> Result { + Self::from_layouts(Some(lower), upper) + } + + fn from_layouts(lower: Option, upper: SlotLayout) -> Result { + let inner = Inner::new(lower, upper)?; + Ok(Self { + inner: SendWrap(Rc::new(RefCell::new(inner))), + }) + } + + /// Return every live slot address in deterministic tier and index order. + pub fn live_addrs(&self) -> Vec { + self.inner.borrow().live_addrs() + } + + /// Return the lower and upper tier layouts. + pub fn layouts(&self) -> (Option, SlotLayout) { + self.inner.borrow().layouts() + } + + /// Compute the address of slot `index`, with lower-tier slots first. + /// + /// Returns `None` if `index >= count`. + pub fn slot_addr(&self, index: usize) -> Option { + self.inner.borrow().slot_addr(index) + } + + /// Total number of free slots across all tiers. + pub fn num_free(&self) -> usize { + self.inner.borrow().num_free() + } + + /// Free a previously allocated slot by address. + pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } + + /// Capacity of a live allocation by its start address. + pub fn allocation_len(&self, addr: u64) -> Result { + self.inner.borrow().allocation_len(addr) + } + + /// Base address of the first managed tier. + pub fn base_addr(&self) -> u64 { + self.inner.borrow().base_addr() + } + + /// Maximum slot size in bytes. + pub fn slot_size(&self) -> usize { + self.inner.borrow().max_alloc_len() + } + + /// Total number of slots across all tiers. + pub fn count(&self) -> usize { + self.inner.borrow().count() + } +} + +impl BufferProvider for SlotPool { + fn max_alloc_len(&self) -> usize { + self.inner.borrow().max_alloc_len() + } + + fn alloc(&self, len: usize) -> Result { + self.inner.borrow_mut().alloc(len) + } + + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } +} diff --git a/src/hyperlight_common/src/virtq/pool/tests.rs b/src/hyperlight_common/src/virtq/pool/tests.rs new file mode 100644 index 000000000..9090f9fcd --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/tests.rs @@ -0,0 +1,502 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +use super::*; + +fn make_run_pool(size: usize) -> RunPool { + let base = align_up(0x10000, L.max(U)).unwrap() as u64; + RunPool::::new(base, size).unwrap() +} + +fn make_slot_pool(slot_count: usize, slot_size: usize) -> SlotPool { + let layout = SlotLayout::new(0x80000, slot_size, slot_count); + SlotPool::new(layout).unwrap() +} + +fn make_tiered_slot_pool(lower_count: usize, upper_count: usize) -> SlotPool { + let lower = SlotLayout::new(0x80000, 256, lower_count); + let upper = SlotLayout::new(0x90000, 4096, upper_count); + SlotPool::new_tiered(lower, upper).unwrap() +} + +#[test] +fn test_run_pool_new_success() { + let pool = RunPool::<256, 4096>::new(0x10000, 1024 * 1024).unwrap(); + assert!(pool.inner.borrow().lower.capacity() > 0); + assert!(pool.inner.borrow().upper.capacity() > 0); +} + +#[test] +fn test_run_pool_alloc_small_to_lower() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(128).unwrap(); + + // Should come from the lower tier. + assert!(pool.inner.borrow().lower.contains(alloc.addr)); + assert_eq!(alloc.len, 256); +} + +#[test] +fn test_run_pool_alloc_large_to_upper() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(1500).unwrap(); + + // Should come from the upper tier. + assert!(pool.inner.borrow().upper.contains(alloc.addr)); + assert_eq!(alloc.len, 4096); +} + +#[test] +fn test_run_pool_alloc_fallback_to_upper() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + + // Fill the lower tier completely. + let mut allocations = Vec::new(); + while pool.inner.borrow().lower.free_bytes() > 0 { + allocations.push(pool.inner.borrow_mut().lower.alloc(256).unwrap()); + } + + // Small allocation should fall back to the upper tier. + let alloc = pool.alloc(128).unwrap(); + assert!(pool.inner.borrow().upper.contains(alloc.addr)); +} + +#[test] +fn test_run_pool_free_from_lower() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(128).unwrap(); + + let free_before = pool.inner.borrow().lower.free_bytes(); + pool.dealloc(alloc.addr).unwrap(); + assert_eq!( + pool.inner.borrow().lower.free_bytes(), + free_before + alloc.len + ); +} + +#[test] +fn test_run_pool_free_from_upper() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(1500).unwrap(); + + let free_before = pool.inner.borrow().upper.free_bytes(); + pool.dealloc(alloc.addr).unwrap(); + assert_eq!( + pool.inner.borrow().upper.free_bytes(), + free_before + alloc.len + ); +} + +#[test] +fn test_run_pool_stress_many_allocations() { + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); + let mut allocations = Vec::new(); + + // Allocate many buffers + for i in 0..100 { + let size = if i % 2 == 0 { 128 } else { 1500 }; + allocations.push(pool.alloc(size).unwrap()); + } + + // Free half of them + for i in (0..100).step_by(2) { + pool.dealloc(allocations[i].addr).unwrap(); + } + + // Should be able to allocate again + for i in 0..50 { + let size = if i % 2 == 0 { 128 } else { 1500 }; + let _alloc = pool.alloc(size).unwrap(); + } +} + +#[test] +fn test_run_pool_mixed_workload() { + let pool = make_run_pool::<256, 4096>(2 * 1024 * 1024); + + // Simulate virtio-net workload + let desc_buf = pool.alloc(64).unwrap(); // Control message + let rx_buf1 = pool.alloc(1500).unwrap(); // MTU packet + let rx_buf2 = pool.alloc(1500).unwrap(); // MTU packet + let tx_buf = pool.alloc(4096).unwrap(); // Large buffer + + // Free and reallocate + pool.dealloc(rx_buf1.addr).unwrap(); + let rx_buf3 = pool.alloc(1500).unwrap(); + + // Should reuse freed buffer (LIFO) + assert_eq!(rx_buf3.addr, rx_buf1.addr); + + pool.dealloc(desc_buf.addr).unwrap(); + pool.dealloc(rx_buf2.addr).unwrap(); + pool.dealloc(rx_buf3.addr).unwrap(); + pool.dealloc(tx_buf.addr).unwrap(); +} + +#[test] +fn test_run_pool_zero_allocation_error() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let result = pool.alloc(0); + assert!(matches!(result, Err(AllocError::InvalidArg))); +} + +#[test] +fn test_run_pool_too_large_allocation() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let result = pool.alloc(2 * 1024 * 1024); // Larger than pool + assert!(matches!(result, Err(AllocError::OutOfMemory))); +} + +#[test] +fn test_align_up_helper() { + assert_eq!(align_up(0, 256).unwrap(), 0); + assert_eq!(align_up(1, 256).unwrap(), 256); + assert_eq!(align_up(256, 256).unwrap(), 256); + assert_eq!(align_up(257, 256).unwrap(), 512); + assert_eq!(align_up(511, 256).unwrap(), 512); + assert_eq!(align_up(512, 256).unwrap(), 512); + assert!(matches!(align_up(1, 0), Err(AllocError::InvalidArg))); + assert!(matches!( + align_up(usize::MAX, 256), + Err(AllocError::Overflow) + )); +} + +#[test] +fn test_slot_pool_preserves_exact_base() { + let layout = SlotLayout::new(0x80001, 4096, 2); + let pool = SlotPool::new(layout).unwrap(); + + assert_eq!(pool.base_addr(), 0x80001); + assert_eq!(pool.count(), 2); + assert_eq!(pool.slot_addr(0), Some(0x80001)); + assert_eq!(pool.slot_addr(1), Some(0x81001)); +} + +#[test] +fn test_tiered_slot_pool_reports_layouts() { + let lower = SlotLayout::new(0x80001, 0x100, 2); + let upper = SlotLayout::new(0x90001, 0x1000, 2); + let pool = SlotPool::new_tiered(lower, upper).unwrap(); + + let (lower, upper) = pool.layouts(); + assert_eq!(lower, Some(SlotLayout::new(0x80001, 0x100, 2))); + assert_eq!(upper, SlotLayout::new(0x90001, 0x1000, 2)); + assert_eq!(pool.base_addr(), 0x80001); + assert_eq!(pool.slot_size(), 0x1000); + assert_eq!(pool.count(), 4); + assert_eq!(pool.slot_addr(0), Some(0x80001)); + assert_eq!(pool.slot_addr(1), Some(0x80101)); + assert_eq!(pool.slot_addr(2), Some(0x90001)); + assert_eq!(pool.slot_addr(3), Some(0x91001)); + assert_eq!(pool.slot_addr(4), None); +} + +#[test] +fn test_tiered_slot_pool_rejects_invalid_layout() { + let lower = SlotLayout::new(0x80000, 0x100, 32); + let overlapping_upper = SlotLayout::new(0x81000, 0x1000, 2); + let overlapping = SlotPool::new_tiered(lower, overlapping_upper); + assert!(matches!(overlapping, Err(AllocError::InvalidArg))); + + let lower = SlotLayout::new(0x80000, 0x1000, 2); + let smaller_upper = SlotLayout::new(0x90000, 0x100, 32); + let reversed_sizes = SlotPool::new_tiered(lower, smaller_upper); + assert!(matches!(reversed_sizes, Err(AllocError::InvalidArg))); +} + +#[test] +fn test_tiered_slot_pool_routes_by_size() { + let pool = make_tiered_slot_pool(2, 2); + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(257).unwrap(); + + assert_eq!(lower.len, 256); + assert!((0x80000..0x80200).contains(&lower.addr)); + assert_eq!(upper.len, 4096); + assert!((0x90000..0x92000).contains(&upper.addr)); + assert_eq!(pool.allocation_len(lower.addr).unwrap(), 256); + assert_eq!(pool.allocation_len(upper.addr).unwrap(), 4096); +} + +#[test] +fn test_tiered_slot_pool_lower_falls_back_when_full() { + let pool = make_tiered_slot_pool(1, 2); + + let lower = pool.alloc(128).unwrap(); + let fallback = pool.alloc(128).unwrap(); + + assert_eq!(lower.len, 256); + assert_eq!(fallback.len, 4096); + assert!((0x90000..0x92000).contains(&fallback.addr)); +} + +#[test] +fn test_tiered_slot_pool_does_not_mask_lower_errors() { + let pool = make_tiered_slot_pool(1, 1); + + assert!(matches!(pool.alloc(0), Err(AllocError::InvalidArg))); + assert!(matches!(pool.alloc(4097), Err(AllocError::OutOfMemory))); + assert_eq!(pool.num_free(), 2); +} + +#[test] +fn test_tiered_slot_pool_alloc_sg_uses_both_tiers() { + let pool = make_tiered_slot_pool(1, 2); + let sgs = pool.alloc_sg(4096 + 128).unwrap(); + + assert_eq!(sgs.len(), 2); + assert_eq!(sgs[0].len, 4096); + assert_eq!(sgs[1].len, 256); + assert!((0x90000..0x92000).contains(&sgs[0].addr)); + assert!((0x80000..0x80100).contains(&sgs[1].addr)); + + for sg in sgs { + pool.dealloc(sg.addr).unwrap(); + } + assert_eq!(pool.num_free(), 3); +} + +#[test] +fn test_tiered_slot_pool_dealloc_routes_by_region() { + let pool = make_tiered_slot_pool(1, 1); + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(1024).unwrap(); + + pool.dealloc(lower.addr).unwrap(); + pool.dealloc(upper.addr).unwrap(); + assert_eq!(pool.num_free(), 2); + assert!(matches!( + pool.dealloc(lower.addr), + Err(AllocError::InvalidFree(_, _)) + )); + assert!(matches!( + pool.dealloc(0x88000), + Err(AllocError::InvalidFree(_, _)) + )); +} + +// Edge case: allocation exactly at boundary +#[test] +fn test_run_pool_boundary_allocation() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + + // Allocate exactly at boundary + let alloc = pool.alloc(256).unwrap(); + assert!(pool.inner.borrow().lower.contains(alloc.addr)); + + // Allocate just over boundary + let alloc2 = pool.alloc(257).unwrap(); + assert!(pool.inner.borrow().upper.contains(alloc2.addr)); +} + +#[test] +fn test_run_pool_dealloc_addr_routes_to_correct_tier() { + let pool = make_run_pool::<256, 4096>(0x20000); + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(1024).unwrap(); + + assert_eq!(pool.allocation_len(lower.addr).unwrap(), 256); + assert_eq!(pool.allocation_len(upper.addr).unwrap(), 4096); + + pool.dealloc_addr(lower.addr).unwrap(); + pool.dealloc_addr(upper.addr).unwrap(); +} + +#[test] +fn test_run_pool_alloc_sg_uses_one_contiguous_run() { + let pool = make_run_pool::<256, 4096>(0x20000); + let sgs = pool.alloc_sg(4096 * 2 + 1).unwrap(); + + assert_eq!(sgs.len(), 1); + assert_eq!(sgs[0].len, 4096 * 3); + + for sg in sgs { + pool.dealloc(sg.addr).unwrap(); + } +} + +#[test] +fn test_run_pool_alloc_sg_large_run() { + let pool = make_run_pool::<256, 4096>(0x20000); + let sgs = pool.alloc_sg(8192).unwrap(); + + assert_eq!(sgs.len(), 1); + assert_eq!(sgs[0].len, 8192); + + for sg in sgs { + pool.dealloc(sg.addr).unwrap(); + } +} + +#[test] +fn test_slot_pool_alloc_sg_splits() { + let pool = make_slot_pool(8, 4096); + let sgs = pool.alloc_sg(4096 * 2 + 1).unwrap(); + + assert_eq!(sgs.len(), 3); + assert_eq!(sgs[0].len, 4096); + assert_eq!(sgs[1].len, 4096); + assert_eq!(sgs[2].len, 4096); + + for sg in sgs { + pool.dealloc(sg.addr).unwrap(); + } +} + +#[test] +fn test_tiered_slot_pool_live_addrs_are_deterministic() { + let pool = make_tiered_slot_pool(2, 2); + let lower_high = pool.alloc(128).unwrap(); + let upper_high = pool.alloc(1024).unwrap(); + let lower_low = pool.alloc(128).unwrap(); + + assert_eq!( + pool.live_addrs(), + vec![lower_low.addr, lower_high.addr, upper_high.addr] + ); +} + +#[test] +fn test_slot_pool_dealloc_out_of_range() { + let pool = make_slot_pool(4, 4096); + let _ = pool.alloc(4096).unwrap(); + + assert!(matches!( + pool.dealloc(0xDEAD), + Err(AllocError::InvalidFree(0xDEAD, 0)) + )); +} + +#[test] +fn test_slot_pool_dealloc_misaligned() { + let pool = make_slot_pool(4, 4096); + let _ = pool.alloc(4096).unwrap(); + + assert!(matches!( + pool.dealloc(0x80001), + Err(AllocError::InvalidFree(0x80001, 0)) + )); +} + +#[test] +fn test_slot_pool_dealloc_double_free() { + let pool = make_slot_pool(4, 4096); + let a = pool.alloc(4096).unwrap(); + pool.dealloc(a.addr).unwrap(); + + // Second dealloc should fail - address is already in the free list + assert!(matches!( + pool.dealloc(a.addr), + Err(AllocError::InvalidFree(_, _)) + )); +} + +#[test] +fn test_slot_pool_alloc_sg_rolls_back_on_failure() { + let pool = make_slot_pool(2, 4096); + + assert!(matches!(pool.alloc_sg(4096 * 3), Err(AllocError::NoSpace))); + assert_eq!(pool.num_free(), 2); + + let alloc = pool.alloc(4096).unwrap(); + assert_eq!(pool.num_free(), 1); + pool.dealloc(alloc.addr).unwrap(); +} + +#[test] +fn test_slot_pool_dealloc_addr_and_allocation_len() { + let pool = make_slot_pool(4, 4096); + let alloc = pool.alloc(4096).unwrap(); + + assert_eq!(pool.allocation_len(alloc.addr).unwrap(), 4096); + pool.dealloc_addr(alloc.addr).unwrap(); + assert!(matches!( + pool.allocation_len(alloc.addr), + Err(AllocError::InvalidFree(_, 0)) + )); +} + +#[test] +fn test_slot_pool_random_order_dealloc() { + let pool = make_slot_pool(8, 4096); + + let mut allocs: Vec = (0..8).map(|_| pool.alloc(4096).unwrap()).collect(); + assert_eq!(pool.num_free(), 0); + + // Dealloc in reverse order + allocs.reverse(); + for a in &allocs { + pool.dealloc(a.addr).unwrap(); + } + assert_eq!(pool.num_free(), 8); + + // All slots should be re-allocatable + let reallocs: Vec = (0..8).map(|_| pool.alloc(4096).unwrap()).collect(); + assert_eq!(pool.num_free(), 0); + + // Verify all addresses are distinct + let mut addrs: Vec = reallocs.iter().map(|a| a.addr).collect(); + addrs.sort(); + addrs.dedup(); + assert_eq!(addrs.len(), 8); +} + +#[test] +fn test_slot_pool_interleaved_alloc_dealloc_order() { + let pool = make_slot_pool(4, 4096); + + let a0 = pool.alloc(4096).unwrap(); + let a1 = pool.alloc(4096).unwrap(); + let a2 = pool.alloc(4096).unwrap(); + let a3 = pool.alloc(4096).unwrap(); + assert_eq!(pool.num_free(), 0); + + // Free middle slots first (out of allocation order) + pool.dealloc(a2.addr).unwrap(); + pool.dealloc(a0.addr).unwrap(); + assert_eq!(pool.num_free(), 2); + + // Re-alloc gets the out-of-order slots back (LIFO) + let b0 = pool.alloc(4096).unwrap(); + assert_eq!(b0.addr, a0.addr); + let b1 = pool.alloc(4096).unwrap(); + assert_eq!(b1.addr, a2.addr); + + // Free everything in yet another order + pool.dealloc(a1.addr).unwrap(); + pool.dealloc(b0.addr).unwrap(); + pool.dealloc(b1.addr).unwrap(); + pool.dealloc(a3.addr).unwrap(); + assert_eq!(pool.num_free(), 4); + + // All 4 original addresses should be available + let mut final_addrs: Vec = (0..4).map(|_| pool.alloc(4096).unwrap().addr).collect(); + final_addrs.sort(); + let expected: Vec = (0..4).map(|i| 0x80000 + i * 4096).collect(); + assert_eq!(final_addrs, expected); +} + +#[test] +fn test_slot_pool_dealloc_order_independent_of_alloc_order() { + let pool = make_slot_pool(6, 256); + + // Allocate all + let allocs: Vec = (0..6).map(|_| pool.alloc(256).unwrap()).collect(); + + // Dealloc in scattered order: 4, 1, 5, 0, 3, 2 + let order = [4, 1, 5, 0, 3, 2]; + for &i in &order { + pool.dealloc(allocs[i].addr).unwrap(); + } + assert_eq!(pool.num_free(), 6); + + // Re-allocate all and verify we get back the full set + let mut realloc_addrs: Vec = (0..6).map(|_| pool.alloc(256).unwrap().addr).collect(); + realloc_addrs.sort(); + + let mut orig_addrs: Vec = allocs.iter().map(|a| a.addr).collect(); + orig_addrs.sort(); + + assert_eq!(realloc_addrs, orig_addrs); +} diff --git a/src/hyperlight_common/src/virtq/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index 2c9d7562b..c3090b968 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -1537,7 +1537,8 @@ mod tests { fn test_chain_build_rolls_back_unrepresentable_allocations() { let ring = make_ring(16); let slot_size = u32::MAX as usize + 1; - let pool = RecyclePool::new(0, slot_size, slot_size).unwrap(); + let layout = SlotLayout::new(0, slot_size, 1); + let pool = SlotPool::new(layout).unwrap(); let mem = ring.mem(); let producer = VirtqProducer::new(ring.layout(), mem, TestNotifier::new(), pool.clone()); From dc73880add9a110f8084413927ca113dac89ec19 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Thu, 23 Jul 2026 14:55:22 +0200 Subject: [PATCH 04/15] feat(virtq): add canonical packed ring images Define deterministic producer and consumer reset behavior. Validate canonical events, descriptor chains, IDs, buffer policy, and unused descriptors. Signed-off-by: Tomasz Andrzejak --- fuzz/README.md | 2 +- fuzz/fuzz_targets/virtq_packed_ring.rs | 81 ++- src/hyperlight_common/src/virtq/consumer.rs | 3 +- src/hyperlight_common/src/virtq/desc.rs | 15 + src/hyperlight_common/src/virtq/event.rs | 9 + src/hyperlight_common/src/virtq/producer.rs | 1 - src/hyperlight_common/src/virtq/ring.rs | 488 +++++--------- .../src/virtq/ring/canonical.rs | 611 ++++++++++++++++++ src/hyperlight_common/src/virtq/ring/fuzz.rs | 204 ++++++ 9 files changed, 1076 insertions(+), 338 deletions(-) create mode 100644 src/hyperlight_common/src/virtq/ring/canonical.rs create mode 100644 src/hyperlight_common/src/virtq/ring/fuzz.rs diff --git a/fuzz/README.md b/fuzz/README.md index b08611786..2144ad6c1 100644 --- a/fuzz/README.md +++ b/fuzz/README.md @@ -10,7 +10,7 @@ which evaluates to the following command `cargo +nightly fuzz run fuzz_host_prin As per Microsoft's Offensive Research & Security Engineering (MORSE) team, all host exposed functions that receive or interact with guest data must be continuously fuzzed for, at least, 500 million fuzz test cases without any crashes. Because `cargo-fuzz` doesn't support setting a maximum number of iterations; instead, we use the `--max_total_time` flag to set a maximum time to run the fuzzer. We have a GitHub action (acting like a CRON job) that runs the fuzzers for 24 hours every week. -Currently, we fuzz the parameters and return type to a hardcoded `PrintOutput` guest function, the `HostPrint` host function, and the packed virtqueue ring parser. We plan to add more fuzzers in the future. +Currently, we fuzz the parameters and return type to a hardcoded `PrintOutput` guest function, the `HostPrint` host function, the packed virtqueue ring parser, and canonical ring image validation. We plan to add more fuzzers in the future. ## On Failure diff --git a/fuzz/fuzz_targets/virtq_packed_ring.rs b/fuzz/fuzz_targets/virtq_packed_ring.rs index b69dc877b..750c01f9f 100644 --- a/fuzz/fuzz_targets/virtq_packed_ring.rs +++ b/fuzz/fuzz_targets/virtq_packed_ring.rs @@ -8,7 +8,8 @@ use std::num::NonZeroU16; use std::ops::Range; use std::rc::Rc; -use hyperlight_common::virtq::{Descriptor, Layout, MemOps, RingConsumer}; +use hyperlight_common::virtq::canonical::validate_canon_image; +use hyperlight_common::virtq::{Descriptor, Layout, MemOps, RingConsumer, RingError}; use libfuzzer_sys::{Corpus, fuzz_target}; const DEFAULT_QUEUE_SIZE: usize = 16; @@ -30,6 +31,7 @@ struct FuzzDesc { #[derive(Clone, Debug)] struct FuzzCase { queue_size: usize, + avail_descs: usize, driver_event_off_wrap: u16, driver_event_flags: u16, written_len: u32, @@ -112,9 +114,9 @@ unsafe impl MemOps for FuzzMem { } } -fn write_driver_event(mem: &FuzzMem, layout: Layout, off_wrap: u16, flags: u16) -> Result<(), ()> { +fn write_event(mem: &FuzzMem, addr: u64, off_wrap: u16, flags: u16) -> Result<(), ()> { mem.write( - layout.drv_evt_addr(), + addr, &[ (off_wrap & 0xff) as u8, (off_wrap >> 8) as u8, @@ -150,7 +152,8 @@ fn parse_case(data: &[u8]) -> Option { let raw_queue_size = read_u16(0); let queue_size = normalize_queue_size(raw_queue_size); - let desc_count = usize::from(read_u16(2)).min(MAX_DESCS).min(queue_size); + let avail_descs = usize::from(read_u16(2)); + let desc_count = avail_descs.min(MAX_DESCS).min(queue_size); let driver_event_off_wrap = read_u16(4); let driver_event_flags = read_u16(6); @@ -177,6 +180,7 @@ fn parse_case(data: &[u8]) -> Option { Some(FuzzCase { queue_size, + avail_descs, driver_event_off_wrap, driver_event_flags, written_len, @@ -194,6 +198,60 @@ fn normalize_queue_size(raw: u16) -> usize { raw.min(MAX_QUEUE_SIZE) } +fn fuzz_canon_image( + mem: &FuzzMem, + layout: Layout, + case: &FuzzCase, + payload_base: u64, +) -> Result<(), ()> { + write_event( + mem, + layout.drv_evt_addr(), + case.driver_event_off_wrap, + case.driver_event_flags, + )?; + let _ = validate_canon_image(mem, layout, case.avail_descs, |_, _| true); + + write_event(mem, layout.drv_evt_addr(), 0, 0)?; + let canon = validate_canon_image(mem, layout, case.avail_descs, |_, _| true); + + let payload_end = payload_base + PAYLOAD_SIZE as u64; + let _ = validate_canon_image(mem, layout, case.avail_descs, |_, elem| { + elem.addr >= payload_base + && elem + .addr + .checked_add(u64::from(elem.len)) + .is_some_and(|end| end <= payload_end) + }); + + if let Ok(chains) = canon { + let mut consumer = RingConsumer::new(layout, mem.clone()); + for expected in chains { + let Ok((id, actual)) = consumer.poll_available() else { + panic!("canonical image was rejected by the ring consumer"); + }; + assert_eq!(id, expected.id()); + assert_eq!(actual.elems().len(), expected.buffers().elems().len()); + for (actual, expected) in actual.elems().iter().zip(expected.buffers().elems()) { + assert_eq!(actual.addr, expected.addr); + assert_eq!(actual.len, expected.len); + assert_eq!(actual.writable, expected.writable); + } + } + assert!(matches!( + consumer.poll_available(), + Err(RingError::WouldBlock) + )); + } + + write_event( + mem, + layout.drv_evt_addr(), + case.driver_event_off_wrap, + case.driver_event_flags, + ) +} + fn run_case(case: FuzzCase) -> Corpus { let Some(num_descs) = NonZeroU16::new(case.queue_size as u16) else { return Corpus::Reject; @@ -206,17 +264,6 @@ fn run_case(case: FuzzCase) -> Corpus { Err(_) => return Corpus::Reject, }; - if write_driver_event( - &mem, - layout, - case.driver_event_off_wrap, - case.driver_event_flags, - ) - .is_err() - { - return Corpus::Reject; - } - let payload_base = BASE_ADDR + ring_size as u64; for (idx, fuzz_desc) in case.descs.iter().enumerate() { let payload_offset = fuzz_desc.addr_offset as usize % PAYLOAD_SIZE; @@ -232,6 +279,10 @@ fn run_case(case: FuzzCase) -> Corpus { } } + if fuzz_canon_image(&mem, layout, &case, payload_base).is_err() { + return Corpus::Reject; + } + let mut consumer = RingConsumer::new(layout, mem); for _ in 0..case.poll_count { let Ok((id, _chain)) = consumer.poll_available() else { diff --git a/src/hyperlight_common/src/virtq/consumer.rs b/src/hyperlight_common/src/virtq/consumer.rs index ffc50ef9d..40cbdfb19 100644 --- a/src/hyperlight_common/src/virtq/consumer.rs +++ b/src/hyperlight_common/src/virtq/consumer.rs @@ -628,12 +628,13 @@ impl VirtqConsumer { /// # Errors /// /// - [`VirtqError::InvalidState`] - one or more chains are still in flight + /// - [`VirtqError::RingError`] - device-event normalization failed pub fn reset(&mut self) -> Result<(), VirtqError> { if self.inflight.ones().next().is_some() { return Err(VirtqError::InvalidState); } - self.inner.reset(); + self.inner.reset()?; self.inflight.clear(); Ok(()) } diff --git a/src/hyperlight_common/src/virtq/desc.rs b/src/hyperlight_common/src/virtq/desc.rs index bc14af310..ae6054f68 100644 --- a/src/hyperlight_common/src/virtq/desc.rs +++ b/src/hyperlight_common/src/virtq/desc.rs @@ -225,6 +225,16 @@ impl DescTable { Some(self.base_addr + (idx as u64 * Descriptor::SIZE as u64)) } + /// Clear all descriptors in the table by writing zeroed descriptors to memory. + pub fn clear(&self, mem: &M) -> Result<(), M::Error> { + let zeroed = Descriptor::zeroed(); + for idx in 0..self.len { + let addr = self.base_addr + (idx as u64 * Descriptor::SIZE as u64); + zeroed.write_release(mem, addr)?; + } + Ok(()) + } + /// Get number of descriptors in table pub fn len(&self) -> usize { self.len @@ -235,6 +245,11 @@ impl DescTable { self.len == 0 } + /// Get the base address of the descriptor table in shared memory + pub fn base_addr(&self) -> u64 { + self.base_addr + } + pub const fn default_len() -> usize { Self::DEFAULT_LEN } diff --git a/src/hyperlight_common/src/virtq/event.rs b/src/hyperlight_common/src/virtq/event.rs index 649beab8a..c6af0c701 100644 --- a/src/hyperlight_common/src/virtq/event.rs +++ b/src/hyperlight_common/src/virtq/event.rs @@ -110,6 +110,15 @@ impl EventSuppression { }) } + /// Clear an `EventSuppression` to the canonical enabled state. + /// + /// # Invariant + /// + /// The caller must ensure that `addr` is a valid pointer to an `EventSuppression`. + pub fn clear(mem: &M, addr: u64) -> Result<(), M::Error> { + Self::new(0, EventFlags::ENABLE).write_release(mem, addr) + } + /// Write an `EventSuppression` to a raw pointer with release semantics. /// /// # Invariant diff --git a/src/hyperlight_common/src/virtq/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index c3090b968..f7ed2576a 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -1796,5 +1796,4 @@ mod tests { )); assert_eq!(producer.inner.num_inflight(), 1); } - } diff --git a/src/hyperlight_common/src/virtq/ring.rs b/src/hyperlight_common/src/virtq/ring.rs index 2d38271f8..56fa7a415 100644 --- a/src/hyperlight_common/src/virtq/ring.rs +++ b/src/hyperlight_common/src/virtq/ring.rs @@ -61,6 +61,8 @@ //! - **DESC**: Notify only when a specific descriptor index is reached //! ``` +pub mod canonical; + use core::fmt; use core::marker::PhantomData; use core::sync::atomic::{Ordering, fence}; @@ -154,6 +156,10 @@ pub enum RingError { InvalidState, #[error("Invalid memory layout")] InvalidLayout, + /// A backend memory operation failed. + /// + /// A failed write may have partially modified shared memory. After a write + /// error, retry reset or discard the endpoint before reuse. #[error("Backend memory error while {op} at address 0x{addr:x}, len {len}")] MemError { /// Memory operation that failed. @@ -892,45 +898,38 @@ impl RingProducer { should_notify_evt(&self.mem, self.dev_evt_addr, self.len() as u16, old, new) } - /// Reset to initial state matching a freshly zeroed ring. - pub fn reset(&mut self) { + /// Reset producer state and its shared ring image to the canonical empty state. + /// + /// The peer must not access the ring during this operation. This clears + /// every descriptor and sets the driver event to `ENABLE`. The consumer + /// separately owns the device event. This low-level operation does not + /// reclaim payload allocations or reconcile higher-level in-flight tracking. + /// + /// # Errors + /// + /// Returns [`RingError::MemError`] if descriptor or event normalization + /// cannot be written to shared memory. Local bookkeeping remains unchanged + /// on error. + pub fn reset(&mut self) -> Result<(), RingError> { + let table_addr = self.desc_table.base_addr(); let size = self.desc_table.len(); + + self.desc_table + .clear(&self.mem) + .map_err(|_| RingError::mem_err(MemOp::WriteDesc, table_addr))?; + + EventSuppression::clear(&self.mem, self.drv_evt_addr) + .map_err(|_| RingError::mem_err(MemOp::WriteEvent, self.drv_evt_addr))?; + self.avail_cursor.reset(); self.used_cursor.reset(); + self.num_free = size; self.id_free.clear(); self.id_free.extend(0..size as u16); self.id_num.iter_mut().for_each(|n| *n = 0); self.event_flags_shadow = EventFlags::ENABLE; - } - - /// Reset the ring to the "N slots submitted, none completed" state. - /// - /// `ids` contains the descriptor IDs that are in-flight. - /// Sets cursors, counters, and `id_num` accordingly. The chain lengths are all set to 1. - pub fn reset_prefilled(&mut self, ids: &[u16]) { - let size = self.desc_table.len(); - let count = ids.len(); - assert!(count <= size); - - let wrapped = count >= size; - self.avail_cursor.head = if wrapped { 0 } else { count as u16 }; - self.avail_cursor.wrap = !wrapped; - - self.used_cursor.head = 0; - self.used_cursor.wrap = true; - - self.id_num.iter_mut().for_each(|n| *n = 0); - for &id in ids { - assert!((id as usize) < size); - assert_eq!(self.id_num[id as usize], 0); - self.id_num[id as usize] = 1; - } - - self.num_free = size - count; - self.id_free.clear(); - self.id_free - .extend((0..size as u16).filter(|id| self.id_num[*id as usize] == 0)); + Ok(()) } } @@ -1309,14 +1308,27 @@ impl RingConsumer { should_notify_evt(&self.mem, self.drv_evt_addr, self.len() as u16, old, new) } - /// Reset to initial state matching a freshly zeroed ring. - /// Does not reallocate internal buffers. - pub fn reset(&mut self) { + /// Reset consumer state and normalize its event-suppression structure. + /// + /// The peer must not access the ring during this operation. Descriptor + /// contents remain producer-owned. This lets a fresh consumer adopt a + /// canonical prefill. A higher-level caller must first rule out outstanding + /// descriptor views. + /// + /// # Errors + /// + /// Returns [`RingError::MemError`] if the device event cannot be normalized + /// in shared memory. Local bookkeeping remains unchanged on error. + pub fn reset(&mut self) -> Result<(), RingError> { + EventSuppression::clear(&self.mem, self.dev_evt_addr) + .map_err(|_| RingError::mem_err(MemOp::WriteEvent, self.dev_evt_addr))?; + self.avail_cursor.reset(); self.used_cursor.reset(); self.id_num.iter_mut().for_each(|n| *n = 0); self.num_inflight = 0; self.event_flags_shadow = EventFlags::ENABLE; + Ok(()) } } @@ -1389,10 +1401,11 @@ impl From<&Descriptor> for BufferElement { #[cfg(test)] pub(crate) mod tests { use alloc::sync::Arc; + use alloc::vec::Vec; use core::cell::UnsafeCell; use core::num::NonZeroU16; use core::ptr; - use core::sync::atomic::{AtomicU16, Ordering}; + use core::sync::atomic::{AtomicU16, AtomicUsize, Ordering}; use bytemuck::{Pod, Zeroable}; @@ -1502,6 +1515,77 @@ pub(crate) mod tests { } } + #[derive(Clone)] + struct FailingWriteMem { + inner: TestMem, + fail_at: Arc, + writes: Arc, + } + + impl FailingWriteMem { + fn new(inner: TestMem) -> Self { + Self { + inner, + fail_at: Arc::new(AtomicUsize::new(usize::MAX)), + writes: Arc::new(AtomicUsize::new(0)), + } + } + + fn fail_at(&self, write: usize) { + self.writes.store(0, Ordering::Relaxed); + self.fail_at.store(write, Ordering::Relaxed); + } + + fn allow_writes(&self) { + self.writes.store(0, Ordering::Relaxed); + self.fail_at.store(usize::MAX, Ordering::Relaxed); + } + + fn check_write(&self) -> Result<(), ()> { + let write = self.writes.fetch_add(1, Ordering::Relaxed); + if write == self.fail_at.load(Ordering::Relaxed) { + Err(()) + } else { + Ok(()) + } + } + } + + // SAFETY: FailingWriteMem delegates to TestMem and only injects errors + // before writes. + unsafe impl MemOps for FailingWriteMem { + type Error = (); + + fn read(&self, addr: u64, dst: &mut [u8]) -> Result<(), Self::Error> { + self.inner.read(addr, dst).unwrap(); + Ok(()) + } + + fn write(&self, addr: u64, src: &[u8]) -> Result<(), Self::Error> { + self.check_write()?; + self.inner.write(addr, src).unwrap(); + Ok(()) + } + + fn load_acquire(&self, addr: u64) -> Result { + Ok(self.inner.load_acquire(addr).unwrap()) + } + + fn store_release(&self, addr: u64, val: u16) -> Result<(), Self::Error> { + self.check_write()?; + self.inner.store_release(addr, val).unwrap(); + Ok(()) + } + + unsafe fn as_slice(&self, addr: u64, len: usize) -> Result<&[u8], Self::Error> { + Ok(unsafe { self.inner.as_slice(addr, len) }.unwrap()) + } + + unsafe fn as_mut_slice(&self, addr: u64, len: usize) -> Result<&mut [u8], Self::Error> { + Ok(unsafe { self.inner.as_mut_slice(addr, len) }.unwrap()) + } + } + /// Owns the descriptor table and event suppression structures pub struct OwnedRing { mem: TestMem, @@ -3214,7 +3298,7 @@ pub(crate) mod tests { used.submit_one(0x1000, 64, false).unwrap(); used.submit_one(0x2000, 128, true).unwrap(); - used.reset(); + used.reset().unwrap(); assert_eq!(used.avail_cursor, fresh.avail_cursor); assert_eq!(used.used_cursor, fresh.used_cursor); @@ -3235,7 +3319,7 @@ pub(crate) mod tests { } assert_eq!(producer.num_free, 4); - producer.reset(); + producer.reset().unwrap(); assert_eq!(producer.num_free, 8); assert_eq!(producer.id_free.len(), 8); @@ -3245,6 +3329,50 @@ pub(crate) mod tests { } } + #[test] + fn test_ring_producer_failed_reset_preserves_local_state() { + let ring = make_ring(4); + let mem = FailingWriteMem::new(ring.mem()); + let mut producer = RingProducer::new(ring.layout(), mem.clone()); + + producer.submit_one(0x1000, 64, false).unwrap(); + producer.submit_one(0x2000, 128, true).unwrap(); + + let avail_cursor = producer.avail_cursor; + let used_cursor = producer.used_cursor; + let num_free = producer.num_free; + let id_free = producer.id_free.clone(); + let id_num = producer.id_num.clone(); + let event_flags_shadow = producer.event_flags_shadow; + + mem.fail_at(1); + assert!(matches!( + producer.reset(), + Err(RingError::MemError { + op: MemOp::WriteDesc, + .. + }) + )); + + assert_eq!(producer.avail_cursor, avail_cursor); + assert_eq!(producer.used_cursor, used_cursor); + assert_eq!(producer.num_free, num_free); + assert_eq!(producer.id_free, id_free); + assert_eq!(producer.id_num, id_num); + assert_eq!(producer.event_flags_shadow, event_flags_shadow); + + mem.allow_writes(); + producer.reset().unwrap(); + + let fresh = RingProducer::new(ring.layout(), mem); + assert_eq!(producer.avail_cursor, fresh.avail_cursor); + assert_eq!(producer.used_cursor, fresh.used_cursor); + assert_eq!(producer.num_free, fresh.num_free); + assert_eq!(producer.id_free, fresh.id_free); + assert_eq!(producer.id_num, fresh.id_num); + assert_eq!(producer.event_flags_shadow, fresh.event_flags_shadow); + } + #[test] fn test_ring_consumer_reset_matches_new() { let ring = make_ring(8); @@ -3260,7 +3388,7 @@ pub(crate) mod tests { let (id, _chain) = used.poll_available().unwrap(); used.submit_used(id, 64).unwrap(); - used.reset(); + used.reset().unwrap(); assert_eq!(used.avail_cursor, fresh.avail_cursor); assert_eq!(used.used_cursor, fresh.used_cursor); @@ -3282,101 +3410,10 @@ pub(crate) mod tests { let _ = consumer.poll_available().unwrap(); assert_eq!(consumer.num_inflight, 2); - consumer.reset(); + consumer.reset().unwrap(); assert_eq!(consumer.num_inflight, 0); } - #[test] - fn test_reset_prefilled_sets_cursors() { - let ring = make_ring(8); - let mut producer = make_producer(&ring); - let ids: Vec = (0..8).collect(); - producer.reset_prefilled(&ids); - - // avail wrapped once (all 8 slots submitted) - assert_eq!(producer.avail_cursor.head(), 0); - assert!(!producer.avail_cursor.wrap()); - // used cursor at initial position - assert_eq!(producer.used_cursor.head(), 0); - assert!(producer.used_cursor.wrap()); - } - - #[test] - fn test_reset_prefilled_all_ids_inflight() { - let ring = make_ring(8); - let mut producer = make_producer(&ring); - let ids: Vec = (0..8).collect(); - producer.reset_prefilled(&ids); - - assert_eq!(producer.num_free, 0); - assert!(producer.id_free.is_empty()); - assert!(producer.id_num.iter().all(|&n| n == 1)); - } - - #[test] - fn test_reset_prefilled_partial() { - let ring = make_ring(8); - let mut producer = make_producer(&ring); - producer.reset_prefilled(&[5, 6, 7, 3]); - - // avail cursor at position 4, no wrap - assert_eq!(producer.avail_cursor.head(), 4); - assert!(producer.avail_cursor.wrap()); - // used cursor at initial position - assert_eq!(producer.used_cursor.head(), 0); - assert!(producer.used_cursor.wrap()); - - assert_eq!(producer.num_free, 4); - assert_eq!(producer.id_free.len(), 4); - for &id in &[0, 1, 2, 4] { - assert!(producer.id_free.contains(&id)); - } - // Only the specified IDs are in-flight - for &id in &[5, 6, 7, 3] { - assert_eq!(producer.id_num[id as usize], 1); - } - for &id in &[0, 1, 2, 4] { - assert_eq!(producer.id_num[id as usize], 0); - } - } - - #[test] - fn test_reset_prefilled_partial_then_submit() { - let ring = make_ring(8); - let mut producer = make_producer(&ring); - producer.reset_prefilled(&[4, 5, 6, 7]); - - let id = producer.submit_one(0x8000, 128, false).unwrap(); - - assert!([0, 1, 2, 3].contains(&id)); - assert_eq!(producer.num_free, 3); - assert_eq!(producer.id_num[id as usize], 1); - } - - #[test] - fn test_reset_prefilled_then_poll_used() { - let ring = make_ring(4); - let mut producer = make_producer(&ring); - - // Simulate host prefill: LIFO assigns IDs 3, 2, 1, 0 - for i in 0..4u64 { - producer.submit_one(0x1000 + i * 4096, 4096, true).unwrap(); - } - - // Consumer marks one as used - let mut consumer = make_consumer(&ring); - let (id, _chain) = consumer.poll_available().unwrap(); - consumer.submit_used(id, 64).unwrap(); - - // Fresh producer restores via reset_prefilled with all IDs - let mut restored = make_producer(&ring); - restored.reset_prefilled(&[0, 1, 2, 3]); - - // poll_used should discover the consumed descriptor - let used = restored.poll_used().unwrap(); - assert_eq!(used.id, id); - } - #[test] fn test_desc_table_read_after_submit() { let ring = make_ring(8); @@ -4210,193 +4247,4 @@ mod virtio_villain { } #[cfg(test)] -mod fuzz { - use quickcheck::{Arbitrary, Gen, QuickCheck}; - - use super::tests::{OwnedRing, make_consumer, make_producer}; - use super::*; - - const MAX_RING: usize = 64; - const MAX_OPS: usize = 128; - const MAX_CHAIN_LEN: usize = 8; - - #[allow(clippy::large_enum_variant)] - #[derive(Clone, Debug)] - enum Op { - /// submit one chain - Submit(BufferChain), - /// poll up to N chains - PollAvail(u8), - /// driver reclaims up to N completions - PollUsed(u8), - /// complete one previously polled chain - CompleteOne, - } - - impl Arbitrary for Op { - fn arbitrary(g: &mut Gen) -> Self { - let choice = u8::arbitrary(g) % 4; - match choice { - 0 => Op::Submit(BufferChain::arbitrary(g)), - 1 => Op::PollAvail(u8::arbitrary(g) % 8 + 1), - 2 => Op::PollUsed(u8::arbitrary(g) % 8 + 1), - 3 => Op::CompleteOne, - _ => unreachable!(), - } - } - } - - #[derive(Clone, Debug)] - struct Scenario { - table_size: usize, - ops: Vec, - } - - impl Arbitrary for Scenario { - fn arbitrary(g: &mut Gen) -> Self { - let table_size = (usize::arbitrary(g) % MAX_RING + 1).next_power_of_two(); - let num_ops = usize::arbitrary(g) % MAX_OPS + 1; - - let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); - Scenario { table_size, ops } - } - } - - impl Arbitrary for BufferElement { - fn arbitrary(g: &mut Gen) -> Self { - let addr = u64::arbitrary(g); - let len = u32::arbitrary(g); - let writable = bool::arbitrary(g); - - BufferElement { - addr, - len, - writable, - } - } - } - - impl Arbitrary for BufferChain { - fn arbitrary(g: &mut Gen) -> Self { - let chain_len = usize::arbitrary(g) % MAX_CHAIN_LEN + 1; - - let mut elems = vec![BufferElement::zeroed(); chain_len]; - let mut readables = 0; - let mut writables = 0; - - for _ in 0..chain_len { - let elem = BufferElement::arbitrary(g); - if elem.writable { - elems[chain_len - 1 - writables] = elem; - writables += 1; - } else { - elems[readables] = elem; - readables += 1; - } - } - - BufferChain { - elems: elems.into(), - split: readables, - } - } - } - - fn run_scenario(s: Scenario) -> bool { - let ring = OwnedRing::new(s.table_size); - let mut producer = make_producer(&ring); - let mut consumer = make_consumer(&ring); - - // Order logs - let mut dev_order: Vec = Vec::new(); - let mut drv_order: Vec = Vec::new(); - - // Device-tracked polled-but-not-completed IDs - let mut dev_ready: Vec<(u16, u32)> = Vec::new(); - - for op in &s.ops { - match op { - Op::Submit(chain) => { - // Submit only if space; otherwise skip - let _ = producer.submit_available(chain); - } - Op::PollAvail(n) => { - for _ in 0..*n { - if let Ok((id, chain)) = consumer.poll_available() { - dev_ready.push((id, chain.len() as u32)); - } else { - break; - } - } - } - Op::PollUsed(n) => { - for _ in 0..*n { - match producer.poll_used() { - Ok(u) => { - drv_order.push(u.id); - if producer.id_num[u.id as usize] != 0 { - return false; - } - if !producer.id_free.contains(&u.id) { - return false; - } - } - Err(RingError::WouldBlock) => break, - Err(_) => return false, - } - } - } - Op::CompleteOne => { - if let Some((id, len)) = dev_ready.pop() { - if consumer.submit_used(id, len).is_err() { - return false; - } - - dev_order.push(id); - } - } - } - - // assert invariants after each op - let outstanding: u16 = producer.id_num.iter().copied().sum(); - if outstanding as usize + producer.num_free != ring.len() { - return false; - } - - for id in producer.id_free.iter() { - if producer.id_num[*id as usize] != 0 { - return false; - } - } - } - - // Drain remaining completions and reclaims - while let Some((id, len)) = dev_ready.pop() { - if consumer.submit_used(id, len).is_err() { - return false; - } - } - - loop { - match producer.poll_used() { - Ok(u) => drv_order.push(u.id), - Err(RingError::WouldBlock) => break, - Err(_) => return false, - } - } - - true - } - - #[test] - fn prop_interleaved_with_order_verification() { - #[cfg(miri)] - let tests = 1; - #[cfg(not(miri))] - let tests = 100; - - QuickCheck::new() - .tests(tests) - .quickcheck(run_scenario as fn(Scenario) -> bool); - } -} +mod fuzz; diff --git a/src/hyperlight_common/src/virtq/ring/canonical.rs b/src/hyperlight_common/src/virtq/ring/canonical.rs new file mode 100644 index 000000000..684ecf68b --- /dev/null +++ b/src/hyperlight_common/src/virtq/ring/canonical.rs @@ -0,0 +1,611 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + */ + +//! Canonical packed virtqueue images. +//! +//! A canonical image starts at the initial wrap round. Available descriptors +//! occupy a complete prefix, unused descriptors are zero, and both event +//! structures are enabled at offset zero. This form can be restored as bytes +//! and validated before either peer resumes. +//! +//! Ring resets normalize the structures owned by each peer. Buffer range +//! policy remains with the integration through the validator callback. + +use alloc::vec::Vec; + +use bytemuck::Zeroable; +use fixedbitset::FixedBitSet; +use smallvec::SmallVec; +use thiserror::Error; + +use super::super::desc::{DescFlags, DescTable, Descriptor}; +use super::super::event::{EventFlags, EventSuppression}; +use super::super::{Layout, MemOps}; +use super::{BufferChain, BufferElement, MemOp, RingError}; + +/// Why a descriptor does not belong to a canonical packed-ring image. +#[derive(Error, Debug, Copy, Clone, PartialEq, Eq)] +pub enum DescError { + /// An unused descriptor was not completely zeroed. + #[error("unused descriptor is not zeroed")] + ExpectedZero, + /// The raw descriptor contains unsupported or reserved flag bits. + #[error("descriptor contains unknown flags")] + UnknownFlags, + /// The descriptor is not available in the initial packed-ring wrap round. + #[error("descriptor is not initially available")] + NotAvailable, + /// Indirect descriptor tables are unsupported. + #[error("indirect descriptor is unsupported")] + Indirect, + /// A chain's NEXT flag extends beyond the available descriptor prefix. + #[error("chain continues beyond the available descriptor prefix")] + ChainContinues, + /// A chain ID is outside the descriptor-table bounds. + #[error("descriptor ID is out of range")] + IdOutOfRange, + /// A tail descriptor does not carry its head descriptor's ID. + #[error("descriptor ID differs within a chain")] + IdMismatch, + /// Two available chains use the same descriptor ID. + #[error("descriptor ID is already used by another chain")] + DuplicateId, + /// A readable descriptor follows a writable descriptor. + #[error("readable descriptor follows a writable descriptor")] + ReadableAfterWritable, +} + +/// Validation failure for a canonical packed-ring image. +#[derive(Error, Debug)] +pub enum ImageError { + /// Reading the shared ring image failed. + #[error(transparent)] + Ring(#[from] RingError), + /// The caller supplied an impossible available-descriptor prefix length. + #[error("available descriptor count {available} exceeds ring capacity {capacity}")] + DescCount { + /// Number of descriptors expected to be available. + available: usize, + /// Descriptor-table capacity. + capacity: usize, + }, + /// An event-suppression structure is not the canonical enabled value. + #[error("event suppression at address 0x{addr:x} is not canonical")] + Event { + /// Address of the invalid event-suppression structure. + addr: u64, + }, + /// A descriptor violates the canonical packed-ring structure. + #[error("descriptor {index} is not canonical: {reason}")] + Desc { + /// Descriptor-table index. + index: u16, + /// Structural validation failure. + reason: DescError, + }, + /// The caller rejected a descriptor's payload range or attributes. + #[error("descriptor {index} buffer at 0x{addr:x} with length {len} was rejected")] + Buffer { + /// Descriptor-table index. + index: u16, + /// Buffer address from the descriptor. + addr: u64, + /// Buffer length from the descriptor. + len: u32, + }, +} + +impl ImageError { + fn desc_count(available: usize, capacity: usize) -> Self { + Self::DescCount { + available, + capacity, + } + } + + fn event(addr: u64) -> Self { + Self::Event { addr } + } + + fn desc(index: u16, reason: DescError) -> Self { + Self::Desc { index, reason } + } + + fn buffer(index: u16, addr: u64, len: u32) -> Self { + Self::Buffer { index, addr, len } + } +} + +/// One available descriptor chain from a validated canonical ring image. +#[derive(Debug, Clone)] +pub struct CanonChain { + id: u16, + inner: BufferChain, +} + +impl CanonChain { + fn new(id: u16, chain: BufferChain) -> Self { + Self { id, inner: chain } + } + + /// Descriptor ID shared by every buffer in the chain. + pub fn id(&self) -> u16 { + self.id + } + + /// Validated buffers in descriptor order. + pub fn buffers(&self) -> &BufferChain { + &self.inner + } + + /// Consume the image metadata and return its buffer chain. + pub fn into_buffers(self) -> BufferChain { + self.inner + } +} + +/// Validate a packed ring while neither peer can modify it. +/// +/// The first `avail_descs` descriptors must form complete available chains +/// beginning at descriptor zero. Every remaining descriptor must be zeroed, +/// and both event-suppression structures must be the canonical enabled value. +/// `validate_buf` supplies integration-specific address, length, and +/// direction bounds without embedding them in the ring implementation. +/// +/// The returned chains preserve descriptor IDs and chain boundaries for +/// cross-checking against producer and pool ownership. +/// +/// # Errors +/// +/// Returns [`ImageError`] if event state is not normalized, descriptor +/// structure is malformed, an unused descriptor is not zero, a buffer is +/// rejected by `validate_buf`, or shared memory cannot be read. +pub fn validate_canon_image( + mem: &M, + layout: Layout, + avail_descs: usize, + mut validate_buf: F, +) -> Result, ImageError> +where + M: MemOps, + F: FnMut(u16, BufferElement) -> bool, +{ + let cap = layout.desc_table_len() as usize; + if avail_descs > cap { + return Err(ImageError::desc_count(avail_descs, cap)); + } + + let canon_evt = EventSuppression::new(0, EventFlags::ENABLE); + for addr in [layout.drv_evt_addr(), layout.dev_evt_addr()] { + let evt = mem + .read_val::(addr) + .map_err(|_| RingError::mem_err(MemOp::ReadEvent, addr))?; + + if evt != canon_evt { + return Err(ImageError::event(addr)); + } + } + + // SAFETY: `Layout` validates the table base, alignment, and descriptor count. + let table = unsafe { DescTable::from_raw_parts(layout.desc_table_addr(), cap) }; + + let mut seen_ids = FixedBitSet::with_capacity(cap); + let mut chains = Vec::new(); + let mut pos = 0usize; + + while pos < avail_descs { + let head_idx = u16::try_from(pos).map_err(|_| RingError::InvalidState)?; + let (head, _) = read_canon_avail_desc(mem, &table, head_idx)?; + let id_idx = head.id as usize; + if id_idx >= cap { + return Err(ImageError::desc(head_idx, DescError::IdOutOfRange)); + } + + if seen_ids.contains(id_idx) { + return Err(ImageError::desc(head_idx, DescError::DuplicateId)); + } + + seen_ids.insert(id_idx); + + let mut elems = SmallVec::<[BufferElement; 16]>::new(); + let mut split = 0usize; + + loop { + let idx = u16::try_from(pos).map_err(|_| RingError::InvalidState)?; + let (desc, flags) = read_canon_avail_desc(mem, &table, idx)?; + if desc.id != head.id { + return Err(ImageError::desc(idx, DescError::IdMismatch)); + } + + let elem = BufferElement::from(&desc); + if !elem.writable && split != elems.len() { + return Err(ImageError::desc(idx, DescError::ReadableAfterWritable)); + } + + split += usize::from(!elem.writable); + + if !validate_buf(idx, elem) { + return Err(ImageError::buffer(idx, elem.addr, elem.len)); + } + + elems.push(elem); + pos += 1; + + if !flags.contains(DescFlags::NEXT) { + break; + } + if pos >= avail_descs { + return Err(ImageError::desc(idx, DescError::ChainContinues)); + } + } + + let canon = CanonChain::new(head.id, BufferChain { elems, split }); + chains.push(canon); + } + + let empty = Descriptor::zeroed(); + for pos in avail_descs..cap { + let idx = u16::try_from(pos).map_err(|_| RingError::InvalidState)?; + let addr = table.desc_addr(idx).ok_or(RingError::InvalidState)?; + + let desc = mem + .read_val::(addr) + .map_err(|_| RingError::mem_err(MemOp::ReadDesc, addr))?; + + if desc != empty { + return Err(ImageError::desc(idx, DescError::ExpectedZero)); + } + } + + Ok(chains) +} + +fn read_canon_avail_desc( + mem: &M, + table: &DescTable, + idx: u16, +) -> Result<(Descriptor, DescFlags), ImageError> { + let addr = table.desc_addr(idx).ok_or(RingError::InvalidState)?; + let desc = mem + .read_val::(addr) + .map_err(|_| RingError::mem_err(MemOp::ReadDesc, addr))?; + + let flags = DescFlags::from_bits(desc.flags) + .ok_or_else(|| ImageError::desc(idx, DescError::UnknownFlags))?; + + if flags.contains(DescFlags::INDIRECT) { + return Err(ImageError::desc(idx, DescError::Indirect)); + } + if !flags.is_avail(true) { + return Err(ImageError::desc(idx, DescError::NotAvailable)); + } + + Ok((desc, flags)) +} + +#[cfg(test)] +mod tests { + use super::super::BufferChainBuilder; + use super::super::tests::{OwnedRing, make_consumer, make_producer, make_ring}; + use super::*; + + fn writable_chain(base: u64, lengths: &[u32]) -> BufferChain { + BufferChainBuilder::new() + .writables(lengths.iter().scan(base, |addr, &len| { + let element = BufferElement { + addr: *addr, + len, + writable: true, + }; + *addr += len as u64; + Some(element) + })) + .build() + .unwrap() + } + + fn validate_all(ring: &OwnedRing, avail_descs: usize) -> Result, ImageError> { + validate_canon_image(&ring.mem(), ring.layout(), avail_descs, |_, _| true) + } + + #[test] + fn canon_reset_normalizes_empty_image() { + let ring = make_ring(8); + let mut producer = make_producer(&ring); + let mut consumer = make_consumer(&ring); + + producer.submit_one(0x1000, 64, true).unwrap(); + producer.enable_used_notifications_desc(3, false).unwrap(); + consumer.enable_avail_notifications_desc(5, false).unwrap(); + + producer.reset().unwrap(); + consumer.reset().unwrap(); + + assert!(validate_all(&ring, 0).unwrap().is_empty()); + assert_eq!( + ring.read_driver_event(), + EventSuppression::new(0, EventFlags::ENABLE) + ); + assert_eq!( + ring.read_device_event(), + EventSuppression::new(0, EventFlags::ENABLE) + ); + for index in 0..ring.len() as u16 { + assert_eq!(ring.read_desc(index), Descriptor::zeroed()); + } + } + + #[test] + fn canon_multi_desc_refill_is_visible_to_fresh_consumer() { + let ring = make_ring(8); + let mut producer = make_producer(&ring); + let mut consumer = make_consumer(&ring); + + producer.reset().unwrap(); + consumer.reset().unwrap(); + + let chains = [ + writable_chain(0x1000, &[64, 128, 256]), + writable_chain(0x2000, &[512, 1024]), + writable_chain(0x4000, &[64, 64, 64]), + ]; + let mut expected = Vec::new(); + for chain in &chains { + let id = producer.submit_available(chain).unwrap(); + expected.push((id, chain.len())); + } + + assert_eq!(producer.num_free(), 0); + assert_eq!(producer.avail_cursor().head(), 0); + assert!(!producer.avail_cursor().wrap()); + + let image = validate_canon_image(&ring.mem(), ring.layout(), ring.len(), |_, elem| { + elem.writable + && elem.len > 0 + && elem + .addr + .checked_add(elem.len as u64) + .is_some_and(|end| end <= 0x5000) + }) + .unwrap(); + + assert_eq!(image.len(), expected.len()); + for (validated, (id, len)) in image.iter().zip(&expected) { + assert_eq!(validated.id(), *id); + assert_eq!(validated.buffers().len(), *len); + assert!(validated.buffers().elems().iter().all(|elem| elem.writable)); + assert_eq!(producer.id_num[*id as usize] as usize, *len); + } + + let mut fresh = make_consumer(&ring); + for (expected_id, expected_len) in expected { + let (id, chain) = fresh.poll_available().unwrap(); + assert_eq!(id, expected_id); + assert_eq!(chain.len(), expected_len); + } + assert!(matches!(fresh.poll_available(), Err(RingError::WouldBlock))); + } + + #[test] + fn canon_image_rejects_invalid_ids() { + let duplicate_ring = make_ring(4); + let mut producer = make_producer(&duplicate_ring); + producer.submit_one(0x1000, 64, true).unwrap(); + producer.submit_one(0x2000, 64, true).unwrap(); + + let head_id = duplicate_ring.read_desc(0).id; + let mut duplicate = duplicate_ring.read_desc(1); + duplicate.id = head_id; + duplicate_ring.write_desc(1, duplicate); + + assert!(matches!( + validate_all(&duplicate_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::DuplicateId, + }) + )); + + let mismatch_ring = make_ring(4); + let mut producer = make_producer(&mismatch_ring); + producer + .submit_available(&writable_chain(0x3000, &[64, 64])) + .unwrap(); + + let mut tail = mismatch_ring.read_desc(1); + tail.id = tail.id.wrapping_sub(1); + mismatch_ring.write_desc(1, tail); + + assert!(matches!( + validate_all(&mismatch_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::IdMismatch, + }) + )); + + let range_ring = make_ring(4); + let mut producer = make_producer(&range_ring); + producer.submit_one(0x4000, 64, true).unwrap(); + + let mut out_of_range = range_ring.read_desc(0); + out_of_range.id = range_ring.len() as u16; + range_ring.write_desc(0, out_of_range); + + assert!(matches!( + validate_all(&range_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::IdOutOfRange, + }) + )); + } + + #[test] + fn canon_image_rejects_invalid_flags_and_wrap_state() { + let unknown_ring = make_ring(4); + let mut producer = make_producer(&unknown_ring); + producer.submit_one(0x1000, 64, true).unwrap(); + + let mut unknown = unknown_ring.read_desc(0); + unknown.flags |= 1 << 3; + unknown_ring.write_desc(0, unknown); + + assert!(matches!( + validate_all(&unknown_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::UnknownFlags, + }) + )); + + let used_ring = make_ring(4); + let mut producer = make_producer(&used_ring); + producer.submit_one(0x2000, 64, true).unwrap(); + + let mut used = used_ring.read_desc(0); + used.mark_used(true); + used_ring.write_desc(0, used); + + assert!(matches!( + validate_all(&used_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::NotAvailable, + }) + )); + + let indirect_ring = make_ring(4); + let mut producer = make_producer(&indirect_ring); + producer.submit_one(0x3000, 64, true).unwrap(); + + let mut indirect = indirect_ring.read_desc(0); + indirect.flags |= DescFlags::INDIRECT.bits(); + indirect_ring.write_desc(0, indirect); + + assert!(matches!( + validate_all(&indirect_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::Indirect, + }) + )); + } + + #[test] + fn canon_image_rejects_malformed_chain_and_unused_desc() { + let chain_ring = make_ring(4); + let mut producer = make_producer(&chain_ring); + producer + .submit_available(&writable_chain(0x1000, &[64, 64])) + .unwrap(); + + let mut tail = chain_ring.read_desc(1); + tail.flags |= DescFlags::NEXT.bits(); + chain_ring.write_desc(1, tail); + + assert!(matches!( + validate_all(&chain_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::ChainContinues, + }) + )); + + let unused_ring = make_ring(4); + let mut producer = make_producer(&unused_ring); + producer.submit_one(0x2000, 64, true).unwrap(); + unused_ring.write_desc(3, Descriptor::new(0x3000, 64, 0, DescFlags::empty())); + + assert!(matches!( + validate_all(&unused_ring, 1), + Err(ImageError::Desc { + index: 3, + reason: DescError::ExpectedZero, + }) + )); + + let direction_ring = make_ring(4); + let mut producer = make_producer(&direction_ring); + producer + .submit_available(&writable_chain(0x4000, &[64, 64])) + .unwrap(); + + let mut readable_tail = direction_ring.read_desc(1); + readable_tail.flags &= !DescFlags::WRITE.bits(); + direction_ring.write_desc(1, readable_tail); + + assert!(matches!( + validate_all(&direction_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::ReadableAfterWritable, + }) + )); + } + + #[test] + fn canon_image_rejects_noncanon_evt_and_buffer_bounds() { + let evt_ring = make_ring(4); + evt_ring + .mem() + .write_val( + evt_ring.layout().drv_evt_addr(), + EventSuppression::new(1, EventFlags::ENABLE), + ) + .unwrap(); + + match validate_all(&evt_ring, 0) { + Err(ImageError::Event { addr }) => { + let expected = evt_ring.layout().drv_evt_addr(); + assert_eq!(addr, expected); + } + other => unreachable!("unexpected result: {other:?}"), + } + + let bounds = make_ring(4); + let mut producer = make_producer(&bounds); + producer.submit_one(u64::MAX - 15, 32, true).unwrap(); + + let res = validate_canon_image(&bounds.mem(), bounds.layout(), 1, |_, element| { + element + .addr + .checked_add(element.len as u64) + .is_some_and(|end| end <= 0x8000) + }); + + match res { + Err(ImageError::Buffer { index, addr, len }) => { + assert_eq!(index, 0); + assert_eq!(addr, u64::MAX - 15); + assert_eq!(len, 32); + } + other => unreachable!("unexpected result: {other:?}"), + } + } + + #[test] + fn canon_image_rejects_avail_count_over_capacity() { + let ring = make_ring(4); + assert!(matches!( + validate_all(&ring, 5), + Err(ImageError::DescCount { + available: 5, + capacity: 4, + }) + )); + } +} diff --git a/src/hyperlight_common/src/virtq/ring/fuzz.rs b/src/hyperlight_common/src/virtq/ring/fuzz.rs new file mode 100644 index 000000000..cef6d12e0 --- /dev/null +++ b/src/hyperlight_common/src/virtq/ring/fuzz.rs @@ -0,0 +1,204 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use quickcheck::{Arbitrary, Gen, QuickCheck}; + +use super::tests::{OwnedRing, make_consumer, make_producer}; +use super::*; + +const MAX_RING: usize = 64; +const MAX_OPS: usize = 128; +const MAX_CHAIN_LEN: usize = 8; + +#[allow(clippy::large_enum_variant)] +#[derive(Clone, Debug)] +enum Op { + /// submit one chain + Submit(BufferChain), + /// poll up to N chains + PollAvail(u8), + /// driver reclaims up to N completions + PollUsed(u8), + /// complete one previously polled chain + CompleteOne, +} + +impl Arbitrary for Op { + fn arbitrary(g: &mut Gen) -> Self { + let choice = u8::arbitrary(g) % 4; + match choice { + 0 => Op::Submit(BufferChain::arbitrary(g)), + 1 => Op::PollAvail(u8::arbitrary(g) % 8 + 1), + 2 => Op::PollUsed(u8::arbitrary(g) % 8 + 1), + 3 => Op::CompleteOne, + _ => unreachable!(), + } + } +} + +#[derive(Clone, Debug)] +struct Scenario { + table_size: usize, + ops: Vec, +} + +impl Arbitrary for Scenario { + fn arbitrary(g: &mut Gen) -> Self { + let table_size = (usize::arbitrary(g) % MAX_RING + 1).next_power_of_two(); + let num_ops = usize::arbitrary(g) % MAX_OPS + 1; + + let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); + Scenario { table_size, ops } + } +} + +impl Arbitrary for BufferElement { + fn arbitrary(g: &mut Gen) -> Self { + let addr = u64::arbitrary(g); + let len = u32::arbitrary(g); + let writable = bool::arbitrary(g); + + BufferElement { + addr, + len, + writable, + } + } +} + +impl Arbitrary for BufferChain { + fn arbitrary(g: &mut Gen) -> Self { + let chain_len = usize::arbitrary(g) % MAX_CHAIN_LEN + 1; + + let mut elems = vec![BufferElement::zeroed(); chain_len]; + let mut readables = 0; + let mut writables = 0; + + for _ in 0..chain_len { + let elem = BufferElement::arbitrary(g); + if elem.writable { + elems[chain_len - 1 - writables] = elem; + writables += 1; + } else { + elems[readables] = elem; + readables += 1; + } + } + + BufferChain { + elems: elems.into(), + split: readables, + } + } +} + +fn run_scenario(s: Scenario) -> bool { + let ring = OwnedRing::new(s.table_size); + let mut producer = make_producer(&ring); + let mut consumer = make_consumer(&ring); + + // Order logs + let mut dev_order: Vec = Vec::new(); + let mut drv_order: Vec = Vec::new(); + + // Device-tracked polled-but-not-completed IDs + let mut dev_ready: Vec<(u16, u32)> = Vec::new(); + + for op in &s.ops { + match op { + Op::Submit(chain) => { + // Submit only if space; otherwise skip + let _ = producer.submit_available(chain); + } + Op::PollAvail(n) => { + for _ in 0..*n { + if let Ok((id, chain)) = consumer.poll_available() { + dev_ready.push((id, chain.len() as u32)); + } else { + break; + } + } + } + Op::PollUsed(n) => { + for _ in 0..*n { + match producer.poll_used() { + Ok(u) => { + drv_order.push(u.id); + if producer.id_num[u.id as usize] != 0 { + return false; + } + if !producer.id_free.contains(&u.id) { + return false; + } + } + Err(RingError::WouldBlock) => break, + Err(_) => return false, + } + } + } + Op::CompleteOne => { + if let Some((id, len)) = dev_ready.pop() { + if consumer.submit_used(id, len).is_err() { + return false; + } + + dev_order.push(id); + } + } + } + + // assert invariants after each op + let outstanding: u16 = producer.id_num.iter().copied().sum(); + if outstanding as usize + producer.num_free != ring.len() { + return false; + } + + for id in producer.id_free.iter() { + if producer.id_num[*id as usize] != 0 { + return false; + } + } + } + + // Drain remaining completions and reclaims + while let Some((id, len)) = dev_ready.pop() { + if consumer.submit_used(id, len).is_err() { + return false; + } + } + + loop { + match producer.poll_used() { + Ok(u) => drv_order.push(u.id), + Err(RingError::WouldBlock) => break, + Err(_) => return false, + } + } + + true +} + +#[test] +fn prop_interleaved_with_order_verification() { + #[cfg(miri)] + let tests = 1; + #[cfg(not(miri))] + let tests = 100; + + QuickCheck::new() + .tests(tests) + .quickcheck(run_scenario as fn(Scenario) -> bool); +} From 0047aff657c7ec45d4210cb9c4295a88bd4f3129 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Fri, 24 Jul 2026 16:56:12 +0200 Subject: [PATCH 05/15] refactor(layout): model scratch-top metadata Represent scratch bookkeeping with one repr(C) layout. Derive offsets and assert the host/guest ABI at compile time. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/layout.rs | 54 +++++++++++++++++++++++++---- 1 file changed, 47 insertions(+), 7 deletions(-) diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index 28a03eb9e..3ab626c44 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright 2025 The Hyperlight Authors. +use core::mem::{offset_of, size_of}; + #[cfg_attr(target_arch = "x86_64", path = "arch/amd64/layout.rs")] #[cfg_attr(target_arch = "aarch64", path = "arch/aarch64/layout.rs")] mod arch; @@ -9,13 +11,51 @@ pub use arch::{ SCRATCH_TOP_GPA, SCRATCH_TOP_GVA, SNAPSHOT_PT_GVA_MAX, SNAPSHOT_PT_GVA_MIN, io_page, }; -// offsets down from the top of scratch memory for various things -pub const SCRATCH_TOP_SIZE_OFFSET: u64 = 0x08; -pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = 0x10; -pub const SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET: u64 = 0x18; -pub const SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET: u64 = 0x20; -pub const SCRATCH_TOP_LIBC_RNG_SEED_OFFSET: u64 = 0x28; -pub const SCRATCH_TOP_EXN_STACK_OFFSET: u64 = 0x30; +const EXN_STACK_ALIGNMENT: usize = 16; + +// Fields are listed in ascending-address order. Public offsets are measured +// down from the top of scratch memory. +#[repr(C)] +struct ScratchTopMetadata { + /// Keep the exception stack pointer aligned 16 bytes aligned. + _alignment_padding: [u8; 8], + /// Seed request for libc's pseudorandom number generator. + libc_rng_seed: u64, + /// Generation of the snapshot backing the sandbox. + snapshot_generation: u64, + /// GPA of the snapshot page-table copy in scratch memory. + snapshot_pt_gpa_base: u64, + /// Next GPA available to the dynamic scratch allocator. + allocator: u64, + /// Size of the scratch region in bytes. + scratch_size: u64, +} + +const fn scratch_top_offset(field_offset: usize) -> u64 { + (size_of::() - field_offset) as u64 +} + +pub const SCRATCH_TOP_SIZE_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, scratch_size)); +pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, allocator)); +pub const SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, snapshot_pt_gpa_base)); +pub const SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, snapshot_generation)); +pub const SCRATCH_TOP_LIBC_RNG_SEED_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, libc_rng_seed)); +pub const SCRATCH_TOP_EXN_STACK_OFFSET: u64 = size_of::() as u64; + +const _: () = { + assert!(size_of::().is_multiple_of(EXN_STACK_ALIGNMENT)); + assert!(SCRATCH_TOP_SIZE_OFFSET == 0x08); + assert!(SCRATCH_TOP_ALLOCATOR_OFFSET == 0x10); + assert!(SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET == 0x18); + assert!(SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET == 0x20); + assert!(SCRATCH_TOP_LIBC_RNG_SEED_OFFSET == 0x28); + assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x30); +}; pub fn scratch_base_gpa(size: usize) -> u64 { (SCRATCH_TOP_GPA - size + 1) as u64 From 7fa5e9e63ab813e3830d64aa0a96cf9456ae2944 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Fri, 24 Jul 2026 17:08:54 +0200 Subject: [PATCH 06/15] fix(guest): correct scratch allocator boundary Use the first GPA of the reserved pages as an exclusive limit. Accept allocations ending at the limit and reject address overflow. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/layout.rs | 7 +++++++ .../src/arch/aarch64/prim_alloc.rs | 9 ++------- .../src/arch/amd64/prim_alloc.rs | 12 ++++-------- src/hyperlight_guest/src/prim_alloc.rs | 17 +++++++++++++++++ 4 files changed, 30 insertions(+), 15 deletions(-) diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index 3ab626c44..dd5bdf963 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -12,6 +12,8 @@ pub use arch::{ }; const EXN_STACK_ALIGNMENT: usize = 16; +/// Pages reserved for the exception stack and scratch-top metadata. +pub const SCRATCH_TOP_RESERVED_PAGES: usize = 2; // Fields are listed in ascending-address order. Public offsets are measured // down from the top of scratch memory. @@ -57,6 +59,11 @@ const _: () = { assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x30); }; +/// Exclusive upper GPA boundary for dynamic scratch allocations. +pub const fn scratch_allocator_limit_gpa() -> u64 { + (SCRATCH_TOP_GPA + 1 - SCRATCH_TOP_RESERVED_PAGES * crate::vmem::PAGE_SIZE) as u64 +} + pub fn scratch_base_gpa(size: usize) -> u64 { (SCRATCH_TOP_GPA - size + 1) as u64 } diff --git a/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs b/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs index 8b5110bb2..9455a7353 100644 --- a/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs +++ b/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs @@ -22,13 +22,8 @@ pub unsafe fn alloc_phys_pages(n: u64) -> u64 { prev_base = out(reg) prev_base, ); } - // Set aside two pages at the top of the scratch region for the - // exception stack, shared state, etc - let max_avail = layout::SCRATCH_TOP_GPA - vmem::PAGE_SIZE * 2; - if prev_base - .checked_add(nbytes) - .is_none_or(|xx| xx >= max_avail as u64) - { + let limit = layout::scratch_allocator_limit_gpa(); + if super::allocation_exceeds_limit(prev_base, nbytes, limit) { unsafe { crate::exit::abort_with_code_and_message( &[ErrorCode::MallocFailed as u8], diff --git a/src/hyperlight_guest/src/arch/amd64/prim_alloc.rs b/src/hyperlight_guest/src/arch/amd64/prim_alloc.rs index 7dbf879df..d56df7aa4 100644 --- a/src/hyperlight_guest/src/arch/amd64/prim_alloc.rs +++ b/src/hyperlight_guest/src/arch/amd64/prim_alloc.rs @@ -2,6 +2,7 @@ // Copyright 2025 The Hyperlight Authors. use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; +use hyperlight_common::{layout, vmem}; // There are no notable architecture-specific safety considerations // here, and the general conditions are documented in the @@ -9,7 +10,7 @@ use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; #[allow(clippy::missing_safety_doc)] pub unsafe fn alloc_phys_pages(n: u64) -> u64 { let addr = crate::layout::allocator_gva(); - let nbytes = n * hyperlight_common::vmem::PAGE_SIZE as u64; + let nbytes = n * vmem::PAGE_SIZE as u64; let mut x = nbytes; unsafe { core::arch::asm!( @@ -18,13 +19,8 @@ pub unsafe fn alloc_phys_pages(n: u64) -> u64 { x = inout(reg) x ); } - // Set aside two pages at the top of the scratch region for the - // exception stack, shared state, etc - let max_avail = - hyperlight_common::layout::SCRATCH_TOP_GPA - hyperlight_common::vmem::PAGE_SIZE * 2; - if x.checked_add(nbytes) - .is_none_or(|xx| xx >= max_avail as u64) - { + let limit = layout::scratch_allocator_limit_gpa(); + if super::allocation_exceeds_limit(x, nbytes, limit) { unsafe { crate::exit::abort_with_code_and_message( &[ErrorCode::MallocFailed as u8], diff --git a/src/hyperlight_guest/src/prim_alloc.rs b/src/hyperlight_guest/src/prim_alloc.rs index dc206af18..988dad347 100644 --- a/src/hyperlight_guest/src/prim_alloc.rs +++ b/src/hyperlight_guest/src/prim_alloc.rs @@ -22,3 +22,20 @@ mod arch; /// latter cannot be perfectly satisfied due to the lack of per-byte /// atomic memcpy in the host. pub use arch::alloc_phys_pages; + +#[inline] +fn allocation_exceeds_limit(base: u64, len: u64, limit: u64) -> bool { + base.checked_add(len).is_none_or(|end| end > limit) +} + +#[cfg(test)] +mod tests { + use super::allocation_exceeds_limit; + + #[test] + fn allocation_may_end_at_limit() { + assert!(!allocation_exceeds_limit(0x1000, 0x2000, 0x3000)); + assert!(allocation_exceeds_limit(0x1000, 0x2001, 0x3000)); + assert!(allocation_exceeds_limit(u64::MAX, 1, u64::MAX)); + } +} From 78412edc038b1611f57832899e46423f0f395b16 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Tue, 28 Jul 2026 00:20:39 +0200 Subject: [PATCH 07/15] feat(virtq): define virtq transport metadata Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/layout.rs | 42 ++++++++++++++++++++++++++++- src/hyperlight_guest/src/layout.rs | 41 ++++++++++++++++++++++------ 2 files changed, 74 insertions(+), 9 deletions(-) diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index dd5bdf963..9c4e1b53e 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -21,6 +21,22 @@ pub const SCRATCH_TOP_RESERVED_PAGES: usize = 2; struct ScratchTopMetadata { /// Keep the exception stack pointer aligned 16 bytes aligned. _alignment_padding: [u8; 8], + /// Number of pages reserved for the H2G pool. + h2g_pool_pages: u64, + /// Guest-published GPA of the H2G pool. + h2g_pool_gpa: u64, + /// Host-published GPA of the H2G ring. + h2g_ring_gpa: u64, + /// Host-published H2G descriptor count. + h2g_queue_depth: u64, + /// Number of pages reserved for the G2H pool. + g2h_pool_pages: u64, + /// Guest-published GPA of the G2H pool. + g2h_pool_gpa: u64, + /// Host-published GPA of the G2H ring. + g2h_ring_gpa: u64, + /// Host-published G2H descriptor count. + g2h_queue_depth: u64, /// Seed request for libc's pseudorandom number generator. libc_rng_seed: u64, /// Generation of the snapshot backing the sandbox. @@ -37,6 +53,22 @@ const fn scratch_top_offset(field_offset: usize) -> u64 { (size_of::() - field_offset) as u64 } +pub const SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_queue_depth)); +pub const SCRATCH_TOP_G2H_RING_GPA_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_ring_gpa)); +pub const SCRATCH_TOP_G2H_POOL_GPA_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_gpa)); +pub const SCRATCH_TOP_G2H_POOL_PAGES_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_pages)); +pub const SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_queue_depth)); +pub const SCRATCH_TOP_H2G_RING_GPA_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_ring_gpa)); +pub const SCRATCH_TOP_H2G_POOL_GPA_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_gpa)); +pub const SCRATCH_TOP_H2G_POOL_PAGES_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_pages)); pub const SCRATCH_TOP_SIZE_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, scratch_size)); pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = @@ -56,7 +88,15 @@ const _: () = { assert!(SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET == 0x18); assert!(SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET == 0x20); assert!(SCRATCH_TOP_LIBC_RNG_SEED_OFFSET == 0x28); - assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x30); + assert!(SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET == 0x30); + assert!(SCRATCH_TOP_G2H_RING_GPA_OFFSET == 0x38); + assert!(SCRATCH_TOP_G2H_POOL_GPA_OFFSET == 0x40); + assert!(SCRATCH_TOP_G2H_POOL_PAGES_OFFSET == 0x48); + assert!(SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET == 0x50); + assert!(SCRATCH_TOP_H2G_RING_GPA_OFFSET == 0x58); + assert!(SCRATCH_TOP_H2G_POOL_GPA_OFFSET == 0x60); + assert!(SCRATCH_TOP_H2G_POOL_PAGES_OFFSET == 0x68); + assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x70); }; /// Exclusive upper GPA boundary for dynamic scratch allocations. diff --git a/src/hyperlight_guest/src/layout.rs b/src/hyperlight_guest/src/layout.rs index ecb8f9d43..95dbc652d 100644 --- a/src/hyperlight_guest/src/layout.rs +++ b/src/hyperlight_guest/src/layout.rs @@ -6,21 +6,46 @@ mod arch; pub use arch::{MAIN_STACK_LIMIT_GVA, MAIN_STACK_TOP_GVA}; + +fn scratch_top_gva(offset: u64) -> *mut u64 { + (hyperlight_common::layout::SCRATCH_TOP_GVA as u64 - offset + 1) as *mut u64 +} + pub fn scratch_size_gva() -> *mut u64 { - use hyperlight_common::layout::{SCRATCH_TOP_GVA, SCRATCH_TOP_SIZE_OFFSET}; - (SCRATCH_TOP_GVA as u64 - SCRATCH_TOP_SIZE_OFFSET + 1) as *mut u64 + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_SIZE_OFFSET) } pub fn allocator_gva() -> *mut u64 { - use hyperlight_common::layout::{SCRATCH_TOP_ALLOCATOR_OFFSET, SCRATCH_TOP_GVA}; - (SCRATCH_TOP_GVA as u64 - SCRATCH_TOP_ALLOCATOR_OFFSET + 1) as *mut u64 + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_ALLOCATOR_OFFSET) } pub fn snapshot_pt_gpa_base_gva() -> *mut u64 { - use hyperlight_common::layout::{SCRATCH_TOP_GVA, SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET}; - (SCRATCH_TOP_GVA as u64 - SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET + 1) as *mut u64 + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET) } pub fn snapshot_generation_gva() -> *mut u64 { - use hyperlight_common::layout::{SCRATCH_TOP_GVA, SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET}; - (SCRATCH_TOP_GVA as u64 - SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET + 1) as *mut u64 + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET) +} +pub fn g2h_queue_depth_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET) +} +pub fn g2h_ring_gpa_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_RING_GPA_OFFSET) +} +pub fn h2g_ring_gpa_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_RING_GPA_OFFSET) +} +pub fn g2h_pool_gpa_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_POOL_GPA_OFFSET) +} +pub fn h2g_pool_gpa_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_POOL_GPA_OFFSET) +} +pub fn g2h_pool_pages_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_POOL_PAGES_OFFSET) +} +pub fn h2g_queue_depth_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET) +} +pub fn h2g_pool_pages_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_POOL_PAGES_OFFSET) } pub fn libc_rng_seed_gva() -> *mut u64 { use hyperlight_common::layout::{SCRATCH_TOP_GVA, SCRATCH_TOP_LIBC_RNG_SEED_OFFSET}; From 4ff0ad346e8d47cd893967e11b107ae6d3199b15 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Tue, 28 Jul 2026 19:39:07 +0200 Subject: [PATCH 08/15] feat(virtq): implement host side memory access Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/virtq/consumer.rs | 30 ++- src/hyperlight_host/src/mem/mod.rs | 2 + src/hyperlight_host/src/mem/shared_mem.rs | 170 +++++++++++++++++ src/hyperlight_host/src/mem/virtq_mem.rs | 196 ++++++++++++++++++++ 4 files changed, 389 insertions(+), 9 deletions(-) create mode 100644 src/hyperlight_host/src/mem/virtq_mem.rs diff --git a/src/hyperlight_common/src/virtq/consumer.rs b/src/hyperlight_common/src/virtq/consumer.rs index 40cbdfb19..a063d8ae0 100644 --- a/src/hyperlight_common/src/virtq/consumer.rs +++ b/src/hyperlight_common/src/virtq/consumer.rs @@ -392,6 +392,7 @@ impl AckChain { /// ``` pub struct VirtqConsumer { inner: RingConsumer, + mem: M, notifier: N, inflight: FixedBitSet, next_token: u32, @@ -406,11 +407,17 @@ impl VirtqConsumer { /// * `mem` - Memory ops implementation for reading/writing to shared memory /// * `notifier` - Callback for notifying the driver about replies pub fn new(layout: Layout, mem: M, notifier: N) -> Self { - let inner = RingConsumer::new(layout, mem); + Self::new_split(layout, mem.clone(), mem, notifier) + } + + /// Create a consumer with separate ring and buffer memory accessors. + pub fn new_split(layout: Layout, ring_mem: M, buf_mem: M, notifier: N) -> Self { + let inner = RingConsumer::new(layout, ring_mem); let inflight = FixedBitSet::with_capacity(inner.len()); Self { inner, + mem: buf_mem, notifier, inflight, next_token: 0, @@ -486,18 +493,16 @@ impl VirtqConsumer { } let chain = RecvChain::new( - self.inner.mem().clone(), + self.mem.clone(), token, readables.iter().copied().collect(), recv_len, ); let reply = if !writables.is_empty() { - let writable = WritableChain::new( - self.inner.mem().clone(), - token, - writables.iter().copied().collect(), - ); + let mem = self.mem.clone(); + let elems = writables.iter().copied().collect(); + let writable = WritableChain::new(mem, token, elems); ReplyChain::Writable(writable) } else { let ack = AckChain::new(token); @@ -909,12 +914,19 @@ mod tests { .unwrap(); ring_producer.submit_available(&chain).unwrap(); - let guarded_mem = FailingPayloadReadMem { + let ring_mem = FailingPayloadReadMem { + inner: mem.clone(), + payload_addr, + payload_len: 0, + }; + let mem = FailingPayloadReadMem { inner: mem, payload_addr, payload_len: 4, }; - let mut consumer = VirtqConsumer::new(ring.layout(), guarded_mem, TestNotifier::new()); + + let mut consumer = + VirtqConsumer::new_split(ring.layout(), ring_mem, mem, TestNotifier::new()); let (recv, reply) = consumer.poll(4).unwrap().unwrap(); assert!(matches!(recv.to_bytes(), Err(VirtqError::MemoryReadError))); diff --git a/src/hyperlight_host/src/mem/mod.rs b/src/hyperlight_host/src/mem/mod.rs index 96e784acd..7bb110c28 100644 --- a/src/hyperlight_host/src/mem/mod.rs +++ b/src/hyperlight_host/src/mem/mod.rs @@ -25,3 +25,5 @@ pub mod shared_mem; /// Utilities for writing shared memory tests #[cfg(all(test, not(miri)))] // uses proptest which isn't miri-compatible pub(crate) mod shared_mem_tests; +#[allow(dead_code)] +pub(crate) mod virtq_mem; diff --git a/src/hyperlight_host/src/mem/shared_mem.rs b/src/hyperlight_host/src/mem/shared_mem.rs index 0f843f96a..0d1eb01e5 100644 --- a/src/hyperlight_host/src/mem/shared_mem.rs +++ b/src/hyperlight_host/src/mem/shared_mem.rs @@ -7,6 +7,7 @@ use std::io::Error; use std::mem::{align_of, size_of}; #[cfg(unix)] use std::ptr::null_mut; +use std::sync::atomic::Ordering; use std::sync::{Arc, RwLock}; use bytemuck::Pod; @@ -98,6 +99,18 @@ pub enum SharedMemoryError { #[error("Cannot access a value with size {0} at offset {1} in memory of size {2}")] Bounds(usize, usize, usize), + /// An atomic access was not aligned for its atomic type. + #[error("Atomic access at offset {0} is not aligned to {1} bytes")] + AtomicAccessUnaligned(usize, usize), + + /// An atomic load used an unsupported ordering. + #[error("Invalid atomic load ordering: {0:?}")] + InvalidAtomicLoadOrdering(Ordering), + + /// An atomic store used an unsupported ordering. + #[error("Invalid atomic store ordering: {0:?}")] + InvalidAtomicStoreOrdering(Ordering), + /// When creating a memory with contents from a file, metadata for /// that file could not be read #[error("Could not access metadata for file: {0}")] @@ -175,6 +188,46 @@ macro_rules! bounds_check { }; } +mod atomic_access { + pub trait Sealed {} +} + +/// An integer atomic supported by [`HostSharedMemory`] atomic operations. +/// +/// This trait is sealed and implemented for the standard signed and unsigned +/// integer atomic types. +#[allow(private_bounds)] +pub trait AtomicAccess: atomic_access::Sealed { + /// The integer stored by this atomic type. + type Value: Copy; + + /// Load the atomic value with `ordering`. + #[doc(hidden)] + fn load(&self, ordering: Ordering) -> Self::Value; + + /// Store `value` with `ordering`. + #[doc(hidden)] + fn store(&self, value: Self::Value, ordering: Ordering); +} + +macro_rules! impl_atomic_access { + ($atomic:ty, $value:ty) => { + impl atomic_access::Sealed for $atomic {} + + impl AtomicAccess for $atomic { + type Value = $value; + + fn load(&self, ordering: Ordering) -> Self::Value { + <$atomic>::load(self, ordering) + } + + fn store(&self, value: Self::Value, ordering: Ordering) { + <$atomic>::store(self, value, ordering); + } + } + }; +} + /// generates a reader function for the given type macro_rules! generate_reader { ($fname:ident, $ty:ty) => { @@ -205,6 +258,17 @@ macro_rules! generate_writer { }; } +impl_atomic_access!(std::sync::atomic::AtomicI8, i8); +impl_atomic_access!(std::sync::atomic::AtomicI16, i16); +impl_atomic_access!(std::sync::atomic::AtomicI32, i32); +impl_atomic_access!(std::sync::atomic::AtomicI64, i64); +impl_atomic_access!(std::sync::atomic::AtomicIsize, isize); +impl_atomic_access!(std::sync::atomic::AtomicU8, u8); +impl_atomic_access!(std::sync::atomic::AtomicU16, u16); +impl_atomic_access!(std::sync::atomic::AtomicU32, u32); +impl_atomic_access!(std::sync::atomic::AtomicU64, u64); +impl_atomic_access!(std::sync::atomic::AtomicUsize, usize); + /// A representation of a host mapping of a shared memory region, /// which will be released when this structure is Drop'd. This is not /// individually Clone (since it holds ownership of the mapping), or @@ -1214,6 +1278,62 @@ impl HostSharedMemory { self.copy_from_slice(bytemuck::bytes_of(&data), offset) } + /// Load an integer atomic at `offset` with `ordering`. + pub fn load_atomic( + &self, + offset: usize, + ordering: Ordering, + ) -> Result { + if matches!(ordering, Ordering::Release | Ordering::AcqRel) { + return Err(SharedMemoryError::InvalidAtomicLoadOrdering(ordering)); + } + + bounds_check!(offset, size_of::(), self.mem_size()); + let ptr = self.base_ptr().wrapping_add(offset); + if !(ptr as usize).is_multiple_of(align_of::()) { + return Err(SharedMemoryError::AtomicAccessUnaligned( + offset, + align_of::(), + )); + } + + let _guard = self.lock.try_read()?; + + // SAFETY: The bounds and alignment checks cover an A within the mapping. + // AtomicAccess is sealed to integer atomics, whose bit patterns are valid. + let atomic = unsafe { &*ptr.cast::() }; + Ok(atomic.load(ordering)) + } + + /// Store an integer atomic at `offset` with `ordering`. + pub fn store_atomic( + &self, + offset: usize, + value: A::Value, + ordering: Ordering, + ) -> Result<()> { + if matches!(ordering, Ordering::Acquire | Ordering::AcqRel) { + return Err(SharedMemoryError::InvalidAtomicStoreOrdering(ordering)); + } + + bounds_check!(offset, size_of::(), self.mem_size()); + let ptr = self.base_ptr().wrapping_add(offset); + if !(ptr as usize).is_multiple_of(align_of::()) { + return Err(SharedMemoryError::AtomicAccessUnaligned( + offset, + align_of::(), + )); + } + + let _guard = self.lock.try_read()?; + + // SAFETY: The bounds and alignment checks cover an A within the mapping. + // AtomicAccess is sealed to integer atomics, whose bit patterns are valid. + let atomic = unsafe { &*ptr.cast::() }; + atomic.store(value, ordering); + Ok(()) + } + /// Copy the contents of the slice into the sandbox at the /// specified offset pub fn copy_to_slice(&self, slice: &mut [u8], offset: usize) -> Result<()> { @@ -1884,6 +2004,8 @@ impl PartialEq for ReadonlySharedMemory { #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicI32, AtomicU16, Ordering}; + #[cfg(not(miri))] use proptest::prelude::*; @@ -1952,6 +2074,54 @@ mod tests { assert!(hshm.fill(0, 1, usize::MAX).is_err()); } + #[test] + fn atomic_access() { + let mem_size = page_size::get(); + let eshm = ExclusiveSharedMemory::new(mem_size).unwrap(); + let (hshm, _) = eshm.build(); + + hshm.store_atomic::(0, 0x1234, Ordering::Release) + .unwrap(); + assert_eq!( + hshm.load_atomic::(0, Ordering::Acquire).unwrap(), + 0x1234 + ); + + hshm.store_atomic::(4, -42, Ordering::SeqCst) + .unwrap(); + assert_eq!( + hshm.load_atomic::(4, Ordering::SeqCst).unwrap(), + -42 + ); + + assert!(hshm.load_atomic::(1, Ordering::Relaxed).is_err()); + assert!( + hshm.load_atomic::(mem_size - 1, Ordering::Relaxed) + .is_err() + ); + assert!(hshm.load_atomic::(0, Ordering::Release).is_err()); + assert!( + hshm.store_atomic::(0, 0, Ordering::Acquire) + .is_err() + ); + } + + #[test] + fn atomic_access_observes_exclusivity() { + let eshm = ExclusiveSharedMemory::new(page_size::get()).unwrap(); + let (mut hshm, _) = eshm.build(); + let other = hshm.clone(); + + hshm.with_exclusivity(|_| { + assert!( + other + .load_atomic::(0, Ordering::Relaxed) + .is_err() + ); + }) + .unwrap(); + } + #[test] fn copy_into_from() -> Result<()> { let mem_size: usize = page_size::get(); diff --git a/src/hyperlight_host/src/mem/virtq_mem.rs b/src/hyperlight_host/src/mem/virtq_mem.rs new file mode 100644 index 000000000..f0e2d447a --- /dev/null +++ b/src/hyperlight_host/src/mem/virtq_mem.rs @@ -0,0 +1,196 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! Host [`MemOps`] access to a bounded scratch region. +//! +//! Every operation uses [`HostSharedMemory`]'s checked API and acquires its +//! lifecycle read lock. This preserves exclusive-memory coordination but makes +//! descriptor traversal pay for one lock acquisition per field access. + +use core::mem::size_of; +use core::ops::Range; +use core::sync::atomic::{AtomicU16, Ordering}; + +use hyperlight_common::layout::scratch_base_gva; +use hyperlight_common::virtq::MemOps; + +use super::shared_mem::{HostSharedMemory, SharedMemory}; +use crate::{HyperlightError, Result, new_error}; + +/// Host virtqueue memory access confined to one scratch GVA range. +/// +/// Accepted guest virtual addresses are translated relative to +/// `scratch_base_gva` and delegated to `scratch_mem`. Separate instances +/// confine ring metadata and payload pools independently. Clones share the +/// backing mapping and lifecycle lock while retaining the same range. +#[derive(Clone)] +pub(crate) struct HostMemOps { + /// Shared scratch mapping used for checked memory operations. + scratch_mem: HostSharedMemory, + /// Guest virtual address corresponding to offset zero in `scratch_mem`. + scratch_base_gva: u64, + /// End-exclusive guest virtual address range accepted by this accessor. + region: Range, +} + +impl HostMemOps { + /// Create a memory accessor for `region`. + pub(crate) fn new(scratch: &HostSharedMemory, region: Range) -> Result { + let scratch_size = scratch.mem_size(); + let scratch_base_gva = scratch_base_gva(scratch_size); + + let scratch_end = u64::try_from(scratch_size) + .ok() + .and_then(|size| scratch_base_gva.checked_add(size)); + + if scratch_end.is_none_or(|end| region.end > end) + || region.start >= region.end + || region.start < scratch_base_gva + { + return Err(new_error!( + "region [{:#x}, {:#x}) is outside scratch at {:#x} with size {}", + region.start, + region.end, + scratch_base_gva, + scratch_size + )); + } + + Ok(Self { + scratch_mem: scratch.clone(), + scratch_base_gva, + region, + }) + } + + fn to_offset(&self, addr: u64, len: usize) -> Result { + let out_of_bounds = || { + new_error!( + "address {:#x} with length {} is outside region [{:#x}, {:#x})", + addr, + len, + self.region.start, + self.region.end + ) + }; + + let access_end = u64::try_from(len) + .ok() + .and_then(|len| addr.checked_add(len)); + + if addr < self.region.start || access_end.is_none_or(|end| end > self.region.end) { + return Err(out_of_bounds()); + } + + addr.checked_sub(self.scratch_base_gva) + .and_then(|offset| usize::try_from(offset).ok()) + .ok_or_else(out_of_bounds) + } +} + +// TODO: Hold one HostSharedMemory read guard across a virtq transaction. +// Descriptor metadata requires several reads and writes, so locking every +// operation scales with chain length and dominates the cached metadata path. + +// SAFETY: HostMemOps rejects accesses outside its assigned region. The backing +// HostSharedMemory keeps the mapping alive, bounds-checks each operation, and +// coordinates every byte and atomic access with exclusive memory operations. +unsafe impl MemOps for HostMemOps { + type Error = HyperlightError; + + fn read(&self, addr: u64, dst: &mut [u8]) -> Result<()> { + let offset = self.to_offset(addr, dst.len())?; + Ok(self.scratch_mem.copy_to_slice(dst, offset)?) + } + + fn write(&self, addr: u64, src: &[u8]) -> Result<()> { + let offset = self.to_offset(addr, src.len())?; + Ok(self.scratch_mem.copy_from_slice(src, offset)?) + } + + fn load_acquire(&self, addr: u64) -> Result { + let offset = self.to_offset(addr, size_of::())?; + Ok(self + .scratch_mem + .load_atomic::(offset, Ordering::Acquire)?) + } + + fn store_release(&self, addr: u64, val: u16) -> Result<()> { + let offset = self.to_offset(addr, size_of::())?; + Ok(self + .scratch_mem + .store_atomic::(offset, val, Ordering::Release)?) + } + + unsafe fn as_slice(&self, _addr: u64, _len: usize) -> Result<&[u8]> { + Err(new_error!("as_slice/as_mut_slice not supported on host")) + } + + #[allow(clippy::mut_from_ref)] + unsafe fn as_mut_slice(&self, _addr: u64, _len: usize) -> Result<&mut [u8]> { + Err(new_error!("as_slice/as_mut_slice not supported on host")) + } +} + +#[cfg(test)] +mod tests { + use hyperlight_common::virtq::MemOps; + + use super::*; + use crate::mem::shared_mem::ExclusiveSharedMemory; + + const SCRATCH_SIZE: usize = 0x4000; + + fn scratch_base() -> u64 { + scratch_base_gva(SCRATCH_SIZE) + } + + fn region() -> Range { + let scratch_base = scratch_base(); + scratch_base + 0x1000..scratch_base + 0x2000 + } + + fn host_mem_ops() -> HostMemOps { + let scratch = ExclusiveSharedMemory::new(SCRATCH_SIZE).unwrap(); + let (scratch, _) = scratch.build(); + HostMemOps::new(&scratch, region()).unwrap() + } + + #[test] + fn accesses_only_assigned_region() { + let mem = host_mem_ops(); + let region = region(); + + mem.write(region.start, &[1, 2, 3, 4]).unwrap(); + let mut bytes = [0; 4]; + mem.read(region.start, &mut bytes).unwrap(); + assert_eq!(bytes, [1, 2, 3, 4]); + + assert!(mem.read(region.start - 1, &mut [0]).is_err()); + assert!(mem.write(region.end - 1, &[1, 2]).is_err()); + assert!(mem.read(region.end, &mut [0]).is_err()); + assert!(mem.read(u64::MAX, &mut [0]).is_err()); + } + + #[test] + fn atomics_use_shared_memory_checks() { + let mem = host_mem_ops(); + let region = region(); + + mem.store_release(region.start, 0x1234).unwrap(); + assert_eq!(mem.load_acquire(region.start).unwrap(), 0x1234); + assert!(mem.load_acquire(region.start + 1).is_err()); + assert!(mem.load_acquire(region.end - 1).is_err()); + } + + #[test] + fn rejects_regions_outside_scratch() { + let scratch = ExclusiveSharedMemory::new(SCRATCH_SIZE).unwrap(); + let (scratch, _) = scratch.build(); + let scratch_base = scratch_base(); + let scratch_end = scratch_base + SCRATCH_SIZE as u64; + + assert!(HostMemOps::new(&scratch, scratch_base - 1..scratch_base).is_err()); + assert!(HostMemOps::new(&scratch, scratch_end - 1..scratch_end + 1).is_err()); + } +} From 7e70dd4f794b64f082363f2412580feccc0ceb19 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Wed, 29 Jul 2026 10:51:59 +0200 Subject: [PATCH 09/15] feat(virtq): configure transport geometry Define directional queue depths, buffer sizes, and pool page counts. Account for guest allocated rings and pools in minimum scratch calculations. Publish the transport contract through scratch-top metadata. Signed-off-by: Tomasz Andrzejak --- CHANGELOG.md | 2 + .../src/arch/aarch64/layout.rs | 8 +- .../src/arch/amd64/layout.rs | 10 +- src/hyperlight_common/src/layout.rs | 76 +++++- src/hyperlight_common/src/virtq/mod.rs | 7 + src/hyperlight_common/src/virtq/pool/slot.rs | 32 ++- src/hyperlight_common/src/virtq/pool/tests.rs | 18 ++ src/hyperlight_guest/src/layout.rs | 9 +- src/hyperlight_host/src/mem/layout.rs | 123 ++++++++- src/hyperlight_host/src/sandbox/config.rs | 239 +++++++++++++++++- .../src/sandbox/initialized_multi_use.rs | 4 + 11 files changed, 492 insertions(+), 36 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a575be40..27a12d094 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ## [Prerelease] - Unreleased ### Added +* Add per-direction virtqueue configuration and account its allocations in + scratch sizing. ### Changed * Expose C guest `ByteChunks` values as pointer and length arrays. diff --git a/src/hyperlight_common/src/arch/aarch64/layout.rs b/src/hyperlight_common/src/arch/aarch64/layout.rs index eb5913faf..d466455e8 100644 --- a/src/hyperlight_common/src/arch/aarch64/layout.rs +++ b/src/hyperlight_common/src/arch/aarch64/layout.rs @@ -15,7 +15,9 @@ pub const fn io_page() -> Option<(crate::vmem::PhysAddr, crate::vmem::VirtAddr)> Some((IO_PAGE_GPA, IO_PAGE_GVA)) } -pub fn min_scratch_size(input_data_size: usize, output_data_size: usize) -> usize { - (input_data_size + output_data_size).next_multiple_of(crate::vmem::PAGE_SIZE) - + 12 * crate::vmem::PAGE_SIZE +pub(super) fn min_scratch_size(input_data_size: usize, output_data_size: usize) -> Option { + input_data_size + .checked_add(output_data_size)? + .checked_next_multiple_of(crate::vmem::PAGE_SIZE)? + .checked_add(12 * crate::vmem::PAGE_SIZE) } diff --git a/src/hyperlight_common/src/arch/amd64/layout.rs b/src/hyperlight_common/src/arch/amd64/layout.rs index e0e1d4aac..d6c0c9895 100644 --- a/src/hyperlight_common/src/arch/amd64/layout.rs +++ b/src/hyperlight_common/src/arch/amd64/layout.rs @@ -28,8 +28,10 @@ pub fn io_page() -> Option<(u64, u64)> { /// - A page for the smallest possible non-exception stack /// - (up to) 3 pages for mapping that /// - Two pages for the exception stack and metadata -/// - A page-aligned amount of memory for I/O buffers (for now) -pub fn min_scratch_size(input_data_size: usize, output_data_size: usize) -> usize { - (input_data_size + output_data_size).next_multiple_of(crate::vmem::PAGE_SIZE) - + 12 * crate::vmem::PAGE_SIZE +/// - A page-aligned amount of memory for I/O buffers +pub(super) fn min_scratch_size(input_data_size: usize, output_data_size: usize) -> Option { + input_data_size + .checked_add(output_data_size)? + .checked_next_multiple_of(crate::vmem::PAGE_SIZE)? + .checked_add(12 * crate::vmem::PAGE_SIZE) } diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index 9c4e1b53e..59f3e7f9b 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -11,6 +11,8 @@ pub use arch::{ SCRATCH_TOP_GPA, SCRATCH_TOP_GVA, SNAPSHOT_PT_GVA_MAX, SNAPSHOT_PT_GVA_MIN, io_page, }; +use crate::virtq; + const EXN_STACK_ALIGNMENT: usize = 16; /// Pages reserved for the exception stack and scratch-top metadata. pub const SCRATCH_TOP_RESERVED_PAGES: usize = 2; @@ -21,19 +23,23 @@ pub const SCRATCH_TOP_RESERVED_PAGES: usize = 2; struct ScratchTopMetadata { /// Keep the exception stack pointer aligned 16 bytes aligned. _alignment_padding: [u8; 8], + /// Host-published capacity of each H2G buffer. + h2g_buffer_size: u64, /// Number of pages reserved for the H2G pool. h2g_pool_pages: u64, /// Guest-published GPA of the H2G pool. h2g_pool_gpa: u64, - /// Host-published GPA of the H2G ring. + /// Guest-published GPA of the H2G ring. h2g_ring_gpa: u64, /// Host-published H2G descriptor count. h2g_queue_depth: u64, + /// Host-published capacity of each G2H upper-tier buffer. + g2h_buffer_size: u64, /// Number of pages reserved for the G2H pool. g2h_pool_pages: u64, /// Guest-published GPA of the G2H pool. g2h_pool_gpa: u64, - /// Host-published GPA of the G2H ring. + /// Guest-published GPA of the G2H ring. g2h_ring_gpa: u64, /// Host-published G2H descriptor count. g2h_queue_depth: u64, @@ -61,6 +67,8 @@ pub const SCRATCH_TOP_G2H_POOL_GPA_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_gpa)); pub const SCRATCH_TOP_G2H_POOL_PAGES_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_pages)); +pub const SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_buffer_size)); pub const SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_queue_depth)); pub const SCRATCH_TOP_H2G_RING_GPA_OFFSET: u64 = @@ -69,6 +77,8 @@ pub const SCRATCH_TOP_H2G_POOL_GPA_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_gpa)); pub const SCRATCH_TOP_H2G_POOL_PAGES_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_pages)); +pub const SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_buffer_size)); pub const SCRATCH_TOP_SIZE_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, scratch_size)); pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = @@ -92,11 +102,13 @@ const _: () = { assert!(SCRATCH_TOP_G2H_RING_GPA_OFFSET == 0x38); assert!(SCRATCH_TOP_G2H_POOL_GPA_OFFSET == 0x40); assert!(SCRATCH_TOP_G2H_POOL_PAGES_OFFSET == 0x48); - assert!(SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET == 0x50); - assert!(SCRATCH_TOP_H2G_RING_GPA_OFFSET == 0x58); - assert!(SCRATCH_TOP_H2G_POOL_GPA_OFFSET == 0x60); - assert!(SCRATCH_TOP_H2G_POOL_PAGES_OFFSET == 0x68); - assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x70); + assert!(SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET == 0x50); + assert!(SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET == 0x58); + assert!(SCRATCH_TOP_H2G_RING_GPA_OFFSET == 0x60); + assert!(SCRATCH_TOP_H2G_POOL_GPA_OFFSET == 0x68); + assert!(SCRATCH_TOP_H2G_POOL_PAGES_OFFSET == 0x70); + assert!(SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET == 0x78); + assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x80); }; /// Exclusive upper GPA boundary for dynamic scratch allocations. @@ -112,4 +124,52 @@ pub fn scratch_base_gva(size: usize) -> u64 { } /// Compute the minimum scratch region size needed for a sandbox. -pub use arch::min_scratch_size; +/// +/// The transport allowance contains one page-backed ring arena and both +/// page-backed buffer pools. The result saturates at [`usize::MAX`]. +pub fn min_scratch_size( + input_data_size: usize, + output_data_size: usize, + g2h_queue_depth: usize, + h2g_queue_depth: usize, + g2h_pool_pages: usize, + h2g_pool_pages: usize, +) -> usize { + let size = arch::min_scratch_size(input_data_size, output_data_size).and_then(|fixed| { + let h2g_ring_offset = virtq::Layout::query_size(g2h_queue_depth) + .checked_next_multiple_of(virtq::Descriptor::ALIGN)?; + + let ring_pages = h2g_ring_offset + .checked_add(virtq::Layout::query_size(h2g_queue_depth))? + .checked_next_multiple_of(crate::vmem::PAGE_SIZE)?; + + let pool_size = g2h_pool_pages + .checked_add(h2g_pool_pages)? + .checked_mul(crate::vmem::PAGE_SIZE)?; + + fixed.checked_add(ring_pages)?.checked_add(pool_size) + }); + + size.unwrap_or(usize::MAX) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn minimum_scratch_includes_ring_arena_and_pools() { + let fixed = arch::min_scratch_size(0, 0).unwrap(); + let transport_pages = 1 + 8 + 4; + + assert_eq!( + fixed + transport_pages * crate::vmem::PAGE_SIZE, + min_scratch_size(0, 0, 64, 32, 8, 4) + ); + } + + #[test] + fn minimum_scratch_saturates_on_overflow() { + assert_eq!(usize::MAX, min_scratch_size(0, 0, 64, 32, usize::MAX, 4)); + } +} diff --git a/src/hyperlight_common/src/virtq/mod.rs b/src/hyperlight_common/src/virtq/mod.rs index 467076de0..da13bcd65 100644 --- a/src/hyperlight_common/src/virtq/mod.rs +++ b/src/hyperlight_common/src/virtq/mod.rs @@ -172,6 +172,13 @@ pub use producer::*; pub use ring::*; use thiserror::Error; +/// Capacity of each fixed G2H lower-tier slot. +pub const G2H_LOWER_SLOT_SIZE: usize = 256; +/// Number of G2H lower-tier slots occupying the first pool page. +pub const G2H_LOWER_SLOT_COUNT: usize = crate::vmem::PAGE_SIZE / G2H_LOWER_SLOT_SIZE; + +const _: () = assert!(G2H_LOWER_SLOT_COUNT * G2H_LOWER_SLOT_SIZE == crate::vmem::PAGE_SIZE); + /// A trait for notifying the consumer about virtqueue events. pub trait Notifier { fn notify(&self, stats: QueueStats); diff --git a/src/hyperlight_common/src/virtq/pool/slot.rs b/src/hyperlight_common/src/virtq/pool/slot.rs index 8e8a50535..ea2801970 100644 --- a/src/hyperlight_common/src/virtq/pool/slot.rs +++ b/src/hyperlight_common/src/virtq/pool/slot.rs @@ -201,13 +201,34 @@ impl Inner { let lower = lower.map(Tier::from_layout).transpose()?; let upper = Tier::from_layout(upper)?; - if let Some(lower) = &lower - && (lower.slot_size >= upper.slot_size || lower.end() > upper.base_addr) - { + let Some(lower) = lower else { + return Ok(Self { lower: None, upper }); + }; + + if lower.slot_size > upper.slot_size || lower.end() > upper.base_addr { return Err(AllocError::InvalidArg); } - Ok(Self { lower, upper }) + if lower.slot_size == upper.slot_size { + if lower.end() != upper.base_addr { + return Err(AllocError::InvalidArg); + } + + let count = lower + .count + .checked_add(upper.count) + .ok_or(AllocError::Overflow)?; + let layout = SlotLayout::new(lower.base_addr, lower.slot_size, count); + return Ok(Self { + lower: None, + upper: Tier::from_layout(layout)?, + }); + } + + Ok(Self { + lower: Some(lower), + upper, + }) } fn max_alloc_len(&self) -> usize { @@ -304,7 +325,8 @@ impl SlotPool { /// Create a two-tier recycling pool from exact lower and upper layouts. /// /// The lower layout must precede the upper layout without overlap, and its - /// slot size must be strictly smaller. + /// slot size must not exceed the upper slot size. Adjacent equal-sized + /// layouts form one tier. pub fn new_tiered(lower: SlotLayout, upper: SlotLayout) -> Result { Self::from_layouts(Some(lower), upper) } diff --git a/src/hyperlight_common/src/virtq/pool/tests.rs b/src/hyperlight_common/src/virtq/pool/tests.rs index 9090f9fcd..031759f85 100644 --- a/src/hyperlight_common/src/virtq/pool/tests.rs +++ b/src/hyperlight_common/src/virtq/pool/tests.rs @@ -192,6 +192,19 @@ fn test_tiered_slot_pool_reports_layouts() { assert_eq!(pool.slot_addr(4), None); } +#[test] +fn test_tiered_slot_pool_combines_contiguous_equal_sized_layouts() { + let lower = SlotLayout::new(0x80000, 0x100, 2); + let upper = SlotLayout::new(0x80200, 0x100, 3); + let pool = SlotPool::new_tiered(lower, upper).unwrap(); + + assert_eq!(pool.layouts(), (None, SlotLayout::new(0x80000, 0x100, 5))); + assert_eq!(pool.base_addr(), 0x80000); + assert_eq!(pool.slot_size(), 0x100); + assert_eq!(pool.count(), 5); + assert_eq!(pool.slot_addr(4), Some(0x80400)); +} + #[test] fn test_tiered_slot_pool_rejects_invalid_layout() { let lower = SlotLayout::new(0x80000, 0x100, 32); @@ -199,6 +212,11 @@ fn test_tiered_slot_pool_rejects_invalid_layout() { let overlapping = SlotPool::new_tiered(lower, overlapping_upper); assert!(matches!(overlapping, Err(AllocError::InvalidArg))); + let lower = SlotLayout::new(0x80000, 0x100, 2); + let separated_upper = SlotLayout::new(0x80300, 0x100, 2); + let separated = SlotPool::new_tiered(lower, separated_upper); + assert!(matches!(separated, Err(AllocError::InvalidArg))); + let lower = SlotLayout::new(0x80000, 0x1000, 2); let smaller_upper = SlotLayout::new(0x90000, 0x100, 32); let reversed_sizes = SlotPool::new_tiered(lower, smaller_upper); diff --git a/src/hyperlight_guest/src/layout.rs b/src/hyperlight_guest/src/layout.rs index 95dbc652d..6d403768c 100644 --- a/src/hyperlight_guest/src/layout.rs +++ b/src/hyperlight_guest/src/layout.rs @@ -41,14 +41,19 @@ pub fn h2g_pool_gpa_gva() -> *mut u64 { pub fn g2h_pool_pages_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_POOL_PAGES_OFFSET) } +pub fn g2h_buffer_size_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET) +} pub fn h2g_queue_depth_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET) } pub fn h2g_pool_pages_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_POOL_PAGES_OFFSET) } +pub fn h2g_buffer_size_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET) +} pub fn libc_rng_seed_gva() -> *mut u64 { - use hyperlight_common::layout::{SCRATCH_TOP_GVA, SCRATCH_TOP_LIBC_RNG_SEED_OFFSET}; - (SCRATCH_TOP_GVA as u64 - SCRATCH_TOP_LIBC_RNG_SEED_OFFSET + 1) as *mut u64 + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_LIBC_RNG_SEED_OFFSET) } pub use arch::{scratch_base_gpa, scratch_base_gva}; diff --git a/src/hyperlight_host/src/mem/layout.rs b/src/hyperlight_host/src/mem/layout.rs index 22b0abc37..a40d74548 100644 --- a/src/hyperlight_host/src/mem/layout.rs +++ b/src/hyperlight_host/src/mem/layout.rs @@ -250,6 +250,18 @@ pub(crate) struct SandboxMemoryLayout { init_data_permissions: Option, /// The size of the scratch region in physical memory. scratch_size: usize, + /// Number of descriptors in the G2H virtqueue. + g2h_queue_depth: usize, + /// Number of descriptors in the H2G virtqueue. + h2g_queue_depth: usize, + /// Capacity of each G2H upper-tier buffer. + g2h_buffer_size: usize, + /// Capacity of each H2G buffer. + h2g_buffer_size: usize, + /// Number of pages in the G2H buffer pool. + g2h_pool_pages: usize, + /// Number of pages in the H2G buffer pool. + h2g_pool_pages: usize, /// Size of the primary guest memory region at `BASE_ADDRESS` /// (code, PEB, heap, init data). For a snapshot-backed layout /// this is also the guest-visible prefix of the host snapshot @@ -284,6 +296,12 @@ impl Debug for SandboxMemoryLayout { &format_args!("{:#x}", self.output_data_size), ) .field("Scratch Size", &format_args!("{:#x}", self.scratch_size)) + .field("G2H Queue Depth", &self.g2h_queue_depth) + .field("H2G Queue Depth", &self.h2g_queue_depth) + .field("G2H Buffer Size", &self.g2h_buffer_size) + .field("H2G Buffer Size", &self.h2g_buffer_size) + .field("G2H Pool Pages", &self.g2h_pool_pages) + .field("H2G Pool Pages", &self.h2g_pool_pages) .field("Snapshot Size", &format_args!("{:#x}", self.snapshot_size)) .field("PT Size", &format_args!("{:#x}", self.pt_size.unwrap_or(0))) .field( @@ -334,8 +352,20 @@ impl SandboxMemoryLayout { } let input_data_size = cfg.get_input_data_size(); let output_data_size = cfg.get_output_data_size(); - let min_scratch_size = - hyperlight_common::layout::min_scratch_size(input_data_size, output_data_size); + let g2h_queue_depth = cfg.get_g2h_queue_depth(); + let h2g_queue_depth = cfg.get_h2g_queue_depth(); + let g2h_buffer_size = cfg.get_g2h_buffer_size(); + let h2g_buffer_size = cfg.get_h2g_buffer_size(); + let g2h_pool_pages = cfg.get_g2h_pool_pages(); + let h2g_pool_pages = cfg.get_h2g_pool_pages(); + let min_scratch_size = hyperlight_common::layout::min_scratch_size( + input_data_size, + output_data_size, + g2h_queue_depth, + h2g_queue_depth, + g2h_pool_pages, + h2g_pool_pages, + ); if scratch_size < min_scratch_size { return Err(MemoryRequestTooSmall(scratch_size, min_scratch_size)); } @@ -349,6 +379,12 @@ impl SandboxMemoryLayout { init_data_permissions, pt_size: None, scratch_size, + g2h_queue_depth, + h2g_queue_depth, + g2h_buffer_size, + h2g_buffer_size, + g2h_pool_pages, + h2g_pool_pages, snapshot_size: 0, }; ret.set_snapshot_size(ret.get_memory_size()?); @@ -383,6 +419,36 @@ impl SandboxMemoryLayout { self.scratch_size } + #[allow(dead_code)] + pub(crate) fn get_g2h_queue_depth(&self) -> usize { + self.g2h_queue_depth + } + + #[allow(dead_code)] + pub(crate) fn get_h2g_queue_depth(&self) -> usize { + self.h2g_queue_depth + } + + #[allow(dead_code)] + pub(crate) fn get_g2h_buffer_size(&self) -> usize { + self.g2h_buffer_size + } + + #[allow(dead_code)] + pub(crate) fn get_h2g_buffer_size(&self) -> usize { + self.h2g_buffer_size + } + + #[allow(dead_code)] + pub(crate) fn get_g2h_pool_pages(&self) -> usize { + self.g2h_pool_pages + } + + #[allow(dead_code)] + pub(crate) fn get_h2g_pool_pages(&self) -> usize { + self.h2g_pool_pages + } + /// Guest-visible prefix size of the snapshot blob. pub(crate) fn snapshot_size(&self) -> usize { self.snapshot_size @@ -408,8 +474,12 @@ impl SandboxMemoryLayout { let min_fixed_scratch = hyperlight_common::layout::min_scratch_size( self.input_data_size, self.output_data_size, + self.g2h_queue_depth, + self.h2g_queue_depth, + self.g2h_pool_pages, + self.h2g_pool_pages, ); - let min_scratch = min_fixed_scratch + size; + let min_scratch = min_fixed_scratch.saturating_add(size); if self.scratch_size < min_scratch { return Err(MemoryRequestTooSmall(self.scratch_size, min_scratch)); } @@ -679,8 +749,7 @@ impl SandboxMemoryLayout { + self.get_pt_base_scratch_offset() as u64 } - /// First GPA of the scratch region the host has not used for - /// something else. + /// First GPA available to the guest scratch allocator. pub(crate) fn get_first_free_scratch_gpa(&self) -> u64 { self.get_pt_base_gpa() + self.pt_size.unwrap_or(0) as u64 } @@ -740,6 +809,34 @@ mod tests { ); } + #[test] + fn transport_memory_is_part_of_minimum_scratch_size() { + let mut cfg = SandboxConfiguration::default(); + let minimum = hyperlight_common::layout::min_scratch_size( + cfg.get_input_data_size(), + cfg.get_output_data_size(), + cfg.get_g2h_queue_depth(), + cfg.get_h2g_queue_depth(), + cfg.get_g2h_pool_pages(), + cfg.get_h2g_pool_pages(), + ); + cfg.set_scratch_size(minimum); + let mut layout = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap(); + + assert!(matches!( + layout.set_pt_size(PAGE_SIZE), + Err(MemoryRequestTooSmall(..)) + )); + } + + #[test] + fn transport_minimum_rejects_capacity_overflow() { + let mut cfg = SandboxConfiguration::default(); + cfg.set_g2h_pool_pages(usize::MAX); + let layout = SandboxMemoryLayout::new(cfg, 4096, 0, None); + assert!(matches!(layout, Err(MemoryRequestTooSmall(_, usize::MAX)))); + } + #[test] fn test_max_memory_sandbox() { let mut cfg = SandboxConfiguration::default(); @@ -811,7 +908,7 @@ mod tests { cfg.set_input_data_size(0x2000); cfg.set_output_data_size(0x2000); cfg.set_heap_size(0x2000); - cfg.set_scratch_size(0x10000); + cfg.set_scratch_size(0x20000); let layout = SandboxMemoryLayout::new(cfg, 0x1000, 0, None).unwrap(); pin_eq!(layout.guest_code_offset(), 0); @@ -821,7 +918,7 @@ mod tests { pin_eq!(layout.init_data_offset(), 0x4000); pin_eq!(layout.get_memory_size().unwrap(), 0x4000); - pin_eq!(layout.get_scratch_size(), 0x10000); + pin_eq!(layout.get_scratch_size(), 0x20000); pin_eq!(layout.get_pt_size(), 0); pin_eq!(layout.get_input_data_buffer_scratch_host_offset(), 0); @@ -840,11 +937,11 @@ mod tests { // `SCRATCH_TOP` pins above, these fix the absolute addresses. pin_eq!( layout.get_input_data_buffer_gva() - - hyperlight_common::layout::scratch_base_gva(0x10000), + - hyperlight_common::layout::scratch_base_gva(0x20000), 0 ); pin_eq!( - layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x10000), + layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x20000), 0x4000 ); // pt_size is zero here, so the first free scratch GPA equals @@ -860,7 +957,7 @@ mod tests { cfg.set_input_data_size(0x4000); cfg.set_output_data_size(0x2000); cfg.set_heap_size(0x5000); - cfg.set_scratch_size(0x20000); + cfg.set_scratch_size(0x30000); let layout = SandboxMemoryLayout::new(cfg, 0x3000, 0, None).unwrap(); pin_eq!(layout.guest_code_offset(), 0); @@ -873,7 +970,7 @@ mod tests { 0x9000_usize.next_multiple_of(page_size::get()) ); - pin_eq!(layout.get_scratch_size(), 0x20000); + pin_eq!(layout.get_scratch_size(), 0x30000); pin_eq!(layout.get_pt_size(), 0); pin_eq!(layout.get_input_data_buffer_scratch_host_offset(), 0); @@ -887,11 +984,11 @@ mod tests { pin_eq!( layout.get_input_data_buffer_gva() - - hyperlight_common::layout::scratch_base_gva(0x20000), + - hyperlight_common::layout::scratch_base_gva(0x30000), 0 ); pin_eq!( - layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x20000), + layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x30000), 0x6000 ); pin_eq!( diff --git a/src/hyperlight_host/src/sandbox/config.rs b/src/hyperlight_host/src/sandbox/config.rs index 1b701c3d4..12705b455 100644 --- a/src/hyperlight_host/src/sandbox/config.rs +++ b/src/hyperlight_host/src/sandbox/config.rs @@ -4,6 +4,8 @@ use std::cmp::max; use std::time::Duration; +use hyperlight_common::virtq::G2H_LOWER_SLOT_SIZE; +use hyperlight_common::vmem::PAGE_SIZE; #[cfg(target_os = "linux")] use libc::c_int; use tracing::{Span, instrument}; @@ -73,6 +75,18 @@ pub struct SandboxConfiguration { interrupt_vcpu_sigrtmin_offset: u8, /// How much writable memory to offer the guest scratch_size: usize, + /// Number of descriptors in the G2H virtqueue. + g2h_queue_depth: usize, + /// Number of descriptors in the H2G virtqueue. + h2g_queue_depth: usize, + /// Capacity of each G2H upper-tier buffer. + g2h_buffer_size: usize, + /// Capacity of each H2G buffer. + h2g_buffer_size: usize, + /// Number of pages in the G2H buffer pool. + g2h_pool_pages: usize, + /// Number of pages in the H2G buffer pool. + h2g_pool_pages: usize, /// Declared guest MSRs, stored inline to keep this type `Copy`. #[cfg(target_arch = "x86_64")] guest_msrs: [u32; Self::MAX_GUEST_MSRS], @@ -97,7 +111,27 @@ impl SandboxConfiguration { /// The default heap size of a hyperlight sandbox pub const DEFAULT_HEAP_SIZE: u64 = 131072; /// The default size of the scratch region - pub const DEFAULT_SCRATCH_SIZE: usize = 0x48000; + pub const DEFAULT_SCRATCH_SIZE: usize = 0x55000; + /// The default G2H virtqueue descriptor count. + pub const DEFAULT_G2H_QUEUE_DEPTH: usize = 64; + /// The default H2G virtqueue descriptor count. + pub const DEFAULT_H2G_QUEUE_DEPTH: usize = 32; + /// The default G2H upper-tier buffer size. + pub const DEFAULT_G2H_BUFFER_SIZE: usize = PAGE_SIZE; + /// The default H2G buffer size. + pub const DEFAULT_H2G_BUFFER_SIZE: usize = PAGE_SIZE; + /// The default total number of G2H pool pages. + pub const DEFAULT_G2H_POOL_PAGES: usize = 8; + /// The default total number of H2G pool pages. + pub const DEFAULT_H2G_POOL_PAGES: usize = 4; + /// The minimum G2H virtqueue descriptor count. + const MIN_QUEUE_DEPTH: usize = 2; + /// The maximum G2H virtqueue descriptor count. + const MAX_QUEUE_DEPTH: usize = 32_768; + /// The minimum configured transport buffer size. + const MIN_BUFFER_SIZE: usize = G2H_LOWER_SLOT_SIZE; + /// The maximum configured transport buffer size. + const MAX_BUFFER_SIZE: usize = u32::MAX as usize; /// Maximum number of distinct guest MSRs that can be declared. /// KVM supports at most 16 MSR filter ranges. Each index may require its /// own range, so 16 is the portable limit across backends. @@ -122,6 +156,12 @@ impl SandboxConfiguration { output_data_size: max(output_data_size, Self::MIN_OUTPUT_SIZE), heap_size_override: heap_size_override.unwrap_or(0), scratch_size, + g2h_queue_depth: Self::DEFAULT_G2H_QUEUE_DEPTH, + h2g_queue_depth: Self::DEFAULT_H2G_QUEUE_DEPTH, + g2h_buffer_size: Self::DEFAULT_G2H_BUFFER_SIZE, + h2g_buffer_size: Self::DEFAULT_H2G_BUFFER_SIZE, + g2h_pool_pages: Self::DEFAULT_G2H_POOL_PAGES, + h2g_pool_pages: Self::DEFAULT_H2G_POOL_PAGES, interrupt_retry_delay, interrupt_vcpu_sigrtmin_offset, #[cfg(gdb)] @@ -286,6 +326,98 @@ impl SandboxConfiguration { self.scratch_size = scratch_size; } + /// Get the G2H virtqueue descriptor count. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_g2h_queue_depth(&self) -> usize { + self.g2h_queue_depth + } + + /// Set the G2H virtqueue descriptor count. + /// + /// Values are rounded up to a power of two in `2..=32768`. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_g2h_queue_depth(&mut self, depth: usize) { + self.g2h_queue_depth = Self::normalize_queue_depth(depth); + } + + /// Get the H2G virtqueue descriptor count. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_h2g_queue_depth(&self) -> usize { + self.h2g_queue_depth + } + + /// Set the H2G virtqueue descriptor count. + /// + /// Values are rounded up to a power of two in `2..=32768`. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_h2g_queue_depth(&mut self, depth: usize) { + self.h2g_queue_depth = Self::normalize_queue_depth(depth); + } + + /// Get the capacity of each G2H upper-tier buffer. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_g2h_buffer_size(&self) -> usize { + self.g2h_buffer_size + } + + /// Set the capacity of each G2H upper-tier buffer. + /// + /// Values are clamped to `256..=u32::MAX`. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_g2h_buffer_size(&mut self, size: usize) { + self.g2h_buffer_size = size.clamp(Self::MIN_BUFFER_SIZE, Self::MAX_BUFFER_SIZE); + self.g2h_pool_pages = max( + self.g2h_pool_pages, + Self::min_g2h_pool_pages(self.g2h_buffer_size), + ); + } + + /// Get the capacity of each H2G buffer. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_h2g_buffer_size(&self) -> usize { + self.h2g_buffer_size + } + + /// Set the capacity of each H2G buffer. + /// + /// Values are clamped to `256..=u32::MAX`. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_h2g_buffer_size(&mut self, size: usize) { + self.h2g_buffer_size = size.clamp(Self::MIN_BUFFER_SIZE, Self::MAX_BUFFER_SIZE); + self.h2g_pool_pages = max( + self.h2g_pool_pages, + Self::min_h2g_pool_pages(self.h2g_buffer_size), + ); + } + + /// Get the total number of G2H pool pages. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_g2h_pool_pages(&self) -> usize { + self.g2h_pool_pages + } + + /// Set the total number of G2H pool pages. + /// + /// The pool contains one lower-tier page and at least one upper buffer. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_g2h_pool_pages(&mut self, pages: usize) { + self.g2h_pool_pages = max(pages, Self::min_g2h_pool_pages(self.g2h_buffer_size)); + } + + /// Get the total number of H2G pool pages. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_h2g_pool_pages(&self) -> usize { + self.h2g_pool_pages + } + + /// Set the total number of H2G pool pages. + /// + /// The pool contains at least one H2G buffer. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_h2g_pool_pages(&mut self, pages: usize) { + self.h2g_pool_pages = max(pages, Self::min_h2g_pool_pages(self.h2g_buffer_size)); + } + #[cfg(crashdump)] #[instrument(skip_all, parent = Span::current(), level= "Trace")] pub(crate) fn get_guest_core_dump(&self) -> bool { @@ -310,6 +442,20 @@ impl SandboxConfiguration { self.heap_size_override_opt() .unwrap_or(Self::DEFAULT_HEAP_SIZE) } + + fn normalize_queue_depth(depth: usize) -> usize { + depth + .clamp(Self::MIN_QUEUE_DEPTH, Self::MAX_QUEUE_DEPTH) + .next_power_of_two() + } + + fn min_g2h_pool_pages(buffer_size: usize) -> usize { + 1 + Self::min_h2g_pool_pages(buffer_size) + } + + fn min_h2g_pool_pages(buffer_size: usize) -> usize { + buffer_size.div_ceil(PAGE_SIZE) + } } impl Default for SandboxConfiguration { @@ -334,6 +480,8 @@ impl Default for SandboxConfiguration { mod tests { #[cfg(target_arch = "x86_64")] use super::GuestMsrError; + use hyperlight_common::vmem::PAGE_SIZE; + use super::SandboxConfiguration; #[test] @@ -422,6 +570,30 @@ mod tests { assert_eq!(0x40000, cfg.scratch_size); assert_eq!(INPUT_DATA_SIZE_OVERRIDE, cfg.input_data_size); assert_eq!(OUTPUT_DATA_SIZE_OVERRIDE, cfg.output_data_size); + assert_eq!( + SandboxConfiguration::DEFAULT_G2H_QUEUE_DEPTH, + cfg.get_g2h_queue_depth() + ); + assert_eq!( + SandboxConfiguration::DEFAULT_H2G_QUEUE_DEPTH, + cfg.get_h2g_queue_depth() + ); + assert_eq!( + SandboxConfiguration::DEFAULT_G2H_BUFFER_SIZE, + cfg.get_g2h_buffer_size() + ); + assert_eq!( + SandboxConfiguration::DEFAULT_H2G_BUFFER_SIZE, + cfg.get_h2g_buffer_size() + ); + assert_eq!( + SandboxConfiguration::DEFAULT_G2H_POOL_PAGES, + cfg.get_g2h_pool_pages() + ); + assert_eq!( + SandboxConfiguration::DEFAULT_H2G_POOL_PAGES, + cfg.get_h2g_pool_pages() + ); } #[test] @@ -449,6 +621,71 @@ mod tests { assert_eq!(SandboxConfiguration::MIN_OUTPUT_SIZE, cfg.output_data_size); } + #[test] + fn queue_depths_are_normalized() { + let mut cfg = SandboxConfiguration::default(); + for (depth, expected) in [ + (0, 2), + (1, 2), + (2, 2), + (3, 4), + (32_767, 32_768), + (32_768, 32_768), + (32_769, 32_768), + (usize::MAX, 32_768), + ] { + cfg.set_g2h_queue_depth(depth); + cfg.set_h2g_queue_depth(depth); + assert_eq!(expected, cfg.get_g2h_queue_depth()); + assert_eq!(expected, cfg.get_h2g_queue_depth()); + } + } + + #[test] + fn buffer_sizes_are_normalized_without_page_rounding() { + let mut cfg = SandboxConfiguration::default(); + + cfg.set_g2h_buffer_size(0); + cfg.set_h2g_buffer_size(0); + assert_eq!(256, cfg.get_g2h_buffer_size()); + assert_eq!(256, cfg.get_h2g_buffer_size()); + + cfg.set_g2h_buffer_size(3000); + cfg.set_h2g_buffer_size(3001); + assert_eq!(3000, cfg.get_g2h_buffer_size()); + assert_eq!(3001, cfg.get_h2g_buffer_size()); + + cfg.set_g2h_buffer_size(usize::MAX); + cfg.set_h2g_buffer_size(usize::MAX); + assert_eq!(u32::MAX as usize, cfg.get_g2h_buffer_size()); + assert_eq!(u32::MAX as usize, cfg.get_h2g_buffer_size()); + } + + #[test] + fn pool_page_counts_are_normalized() { + let mut cfg = SandboxConfiguration::default(); + + cfg.set_g2h_pool_pages(0); + cfg.set_h2g_pool_pages(0); + assert_eq!(2, cfg.get_g2h_pool_pages()); + assert_eq!(1, cfg.get_h2g_pool_pages()); + + cfg.set_g2h_buffer_size(PAGE_SIZE + 1); + cfg.set_h2g_buffer_size(PAGE_SIZE + 1); + assert_eq!(3, cfg.get_g2h_pool_pages()); + assert_eq!(2, cfg.get_h2g_pool_pages()); + + cfg.set_g2h_pool_pages(2); + cfg.set_h2g_pool_pages(1); + assert_eq!(3, cfg.get_g2h_pool_pages()); + assert_eq!(2, cfg.get_h2g_pool_pages()); + + cfg.set_g2h_pool_pages(4); + cfg.set_h2g_pool_pages(3); + assert_eq!(4, cfg.get_g2h_pool_pages()); + assert_eq!(3, cfg.get_h2g_pool_pages()); + } + mod proptests { use proptest::prelude::*; diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index 2f7c9c1b4..7b317afdc 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -1370,6 +1370,10 @@ mod tests { hyperlight_common::layout::min_scratch_size( defaults.get_input_data_size(), defaults.get_output_data_size(), + defaults.get_g2h_queue_depth(), + defaults.get_h2g_queue_depth(), + defaults.get_g2h_pool_pages(), + defaults.get_h2g_pool_pages(), ) } + 0x10000 + 0x10000; From fa9ab9c099689523c061f3af4ccb6c79d3c7a14b Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Thu, 30 Jul 2026 17:58:00 +0200 Subject: [PATCH 10/15] feat(virtq): initialize runtime transport This patch adds guest owned G2H and H2G rings and pools during guest initialization, and prefill H2G receive capacity. The guest then publishes their gpas through scratch metadata. The patch is also validates allocation order, scratch ownership, and canonical ring images before installing either host consumer. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/layout.rs | 269 +++++++-- src/hyperlight_common/src/virtq/mod.rs | 17 + src/hyperlight_guest/src/error.rs | 10 + src/hyperlight_guest/src/layout.rs | 13 +- src/hyperlight_guest/src/lib.rs | 1 + src/hyperlight_guest/src/transport/context.rs | 139 +++++ src/hyperlight_guest/src/transport/mem.rs | 135 +++++ src/hyperlight_guest/src/transport/mod.rs | 62 ++ src/hyperlight_guest_bin/src/lib.rs | 4 + src/hyperlight_guest_bin/src/transport.rs | 90 +++ src/hyperlight_host/src/mem/layout.rs | 81 ++- src/hyperlight_host/src/mem/mgr.rs | 84 +++ src/hyperlight_host/src/mem/mod.rs | 2 + src/hyperlight_host/src/mem/virtq.rs | 568 ++++++++++++++++++ .../src/sandbox/initialized_multi_use.rs | 8 + .../src/sandbox/uninitialized_evolve.rs | 8 + 16 files changed, 1435 insertions(+), 56 deletions(-) create mode 100644 src/hyperlight_guest/src/transport/context.rs create mode 100644 src/hyperlight_guest/src/transport/mem.rs create mode 100644 src/hyperlight_guest/src/transport/mod.rs create mode 100644 src/hyperlight_guest_bin/src/transport.rs create mode 100644 src/hyperlight_host/src/mem/virtq.rs diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index 59f3e7f9b..f0fd03be1 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -2,6 +2,7 @@ // Copyright 2025 The Hyperlight Authors. use core::mem::{offset_of, size_of}; +use core::num::{NonZeroU16, NonZeroUsize}; #[cfg_attr(target_arch = "x86_64", path = "arch/amd64/layout.rs")] #[cfg_attr(target_arch = "aarch64", path = "arch/aarch64/layout.rs")] @@ -21,28 +22,22 @@ pub const SCRATCH_TOP_RESERVED_PAGES: usize = 2; // down from the top of scratch memory. #[repr(C)] struct ScratchTopMetadata { - /// Keep the exception stack pointer aligned 16 bytes aligned. - _alignment_padding: [u8; 8], + /// Padding that keeps the exception stack 16-byte aligned. + _alignment_padding: [u8; 16], /// Host-published capacity of each H2G buffer. h2g_buffer_size: u64, /// Number of pages reserved for the H2G pool. h2g_pool_pages: u64, - /// Guest-published GPA of the H2G pool. - h2g_pool_gpa: u64, - /// Guest-published GPA of the H2G ring. - h2g_ring_gpa: u64, /// Host-published H2G descriptor count. h2g_queue_depth: u64, /// Host-published capacity of each G2H upper-tier buffer. g2h_buffer_size: u64, /// Number of pages reserved for the G2H pool. g2h_pool_pages: u64, - /// Guest-published GPA of the G2H pool. - g2h_pool_gpa: u64, - /// Guest-published GPA of the G2H ring. - g2h_ring_gpa: u64, /// Host-published G2H descriptor count. g2h_queue_depth: u64, + /// Host-published GPA of the fixed transport arena. + transport_arena_gpa: u64, /// Seed request for libc's pseudorandom number generator. libc_rng_seed: u64, /// Generation of the snapshot backing the sandbox. @@ -61,24 +56,18 @@ const fn scratch_top_offset(field_offset: usize) -> u64 { pub const SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_queue_depth)); -pub const SCRATCH_TOP_G2H_RING_GPA_OFFSET: u64 = - scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_ring_gpa)); -pub const SCRATCH_TOP_G2H_POOL_GPA_OFFSET: u64 = - scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_gpa)); pub const SCRATCH_TOP_G2H_POOL_PAGES_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_pages)); pub const SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_buffer_size)); pub const SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_queue_depth)); -pub const SCRATCH_TOP_H2G_RING_GPA_OFFSET: u64 = - scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_ring_gpa)); -pub const SCRATCH_TOP_H2G_POOL_GPA_OFFSET: u64 = - scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_gpa)); pub const SCRATCH_TOP_H2G_POOL_PAGES_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_pages)); pub const SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_buffer_size)); +pub const SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, transport_arena_gpa)); pub const SCRATCH_TOP_SIZE_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, scratch_size)); pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = @@ -98,17 +87,14 @@ const _: () = { assert!(SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET == 0x18); assert!(SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET == 0x20); assert!(SCRATCH_TOP_LIBC_RNG_SEED_OFFSET == 0x28); - assert!(SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET == 0x30); - assert!(SCRATCH_TOP_G2H_RING_GPA_OFFSET == 0x38); - assert!(SCRATCH_TOP_G2H_POOL_GPA_OFFSET == 0x40); - assert!(SCRATCH_TOP_G2H_POOL_PAGES_OFFSET == 0x48); - assert!(SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET == 0x50); - assert!(SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET == 0x58); - assert!(SCRATCH_TOP_H2G_RING_GPA_OFFSET == 0x60); - assert!(SCRATCH_TOP_H2G_POOL_GPA_OFFSET == 0x68); - assert!(SCRATCH_TOP_H2G_POOL_PAGES_OFFSET == 0x70); - assert!(SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET == 0x78); - assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x80); + assert!(SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET == 0x30); + assert!(SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET == 0x38); + assert!(SCRATCH_TOP_G2H_POOL_PAGES_OFFSET == 0x40); + assert!(SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET == 0x48); + assert!(SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET == 0x50); + assert!(SCRATCH_TOP_H2G_POOL_PAGES_OFFSET == 0x58); + assert!(SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET == 0x60); + assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x70); }; /// Exclusive upper GPA boundary for dynamic scratch allocations. @@ -125,7 +111,7 @@ pub fn scratch_base_gva(size: usize) -> u64 { /// Compute the minimum scratch region size needed for a sandbox. /// -/// The transport allowance contains one page-backed ring arena and both +/// The fixed transport prefix contains one page-backed ring arena and both /// page-backed buffer pools. The result saturates at [`usize::MAX`]. pub fn min_scratch_size( input_data_size: usize, @@ -136,27 +122,231 @@ pub fn min_scratch_size( h2g_pool_pages: usize, ) -> usize { let size = arch::min_scratch_size(input_data_size, output_data_size).and_then(|fixed| { - let h2g_ring_offset = virtq::Layout::query_size(g2h_queue_depth) + let g2h = QueueDims::new(g2h_queue_depth, g2h_pool_pages)?; + let h2g = QueueDims::new(h2g_queue_depth, h2g_pool_pages)?; + + let transport_len = TransportArena::checked_query_size(g2h, h2g)?; + fixed.checked_add(transport_len) + }); + + size.unwrap_or(usize::MAX) +} + +/// Validated address independent dimensions for one transport queue. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct QueueDims { + depth: NonZeroU16, + pool_pages: NonZeroUsize, +} + +impl QueueDims { + /// Validate one queue descriptor count and pool page count. + pub fn new(depth: usize, pool_pages: usize) -> Option { + let depth = u16::try_from(depth).ok()?; + let depth = NonZeroU16::new(depth)?; + + if !depth.get().is_power_of_two() { + return None; + } + + let pool_pages = NonZeroUsize::new(pool_pages)?; + Some(Self { depth, pool_pages }) + } + + /// Number of descriptors in the queue. + pub const fn depth(&self) -> NonZeroU16 { + self.depth + } + + /// Number of pages in the queue's buffer pool. + pub const fn pool_pages(&self) -> NonZeroUsize { + self.pool_pages + } + + /// Compute the ring length, returning `None` on arithmetic overflow. + pub fn checked_ring_len(&self) -> Option { + virtq::Layout::checked_query_size(usize::from(self.depth.get())) + } + + /// Compute the pool length, returning `None` on arithmetic overflow. + pub fn checked_pool_len(&self) -> Option { + self.pool_pages.get().checked_mul(crate::vmem::PAGE_SIZE) + } +} + +/// Addresses of both rings and pools in one fixed transport arena. +/// +/// The G2H ring begins at the arena base. The H2G ring is descriptor aligned. +/// Both pools are page aligned. +/// +/// ```text +/// +----------+------------+----------+-----+----------+----------+ +/// | G2H ring | align pad | H2G ring | pad | G2H pool | H2G pool | +/// +----------+------------+----------+-----+----------+----------+ +/// ``` +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TransportArena { + /// Address of the G2H ring and base of the arena. + g2h_ring_addr: u64, + /// Address of the H2G ring. + h2g_ring_addr: u64, + /// Address of the G2H pool. + g2h_pool_addr: u64, + /// Address of the H2G pool. + h2g_pool_addr: u64, + /// Page-aligned length occupied by both rings. + ring_span_len: usize, + /// Total page-aligned arena length. + len: usize, +} + +impl TransportArena { + /// Derive one transport arena from its base address and queue dimensions. + pub fn new(base_addr: u64, g2h: QueueDims, h2g: QueueDims) -> Option { + if !base_addr.is_multiple_of(crate::vmem::PAGE_SIZE as u64) { + return None; + } + + let h2g_ring_offset = g2h + .checked_ring_len()? .checked_next_multiple_of(virtq::Descriptor::ALIGN)?; - let ring_pages = h2g_ring_offset - .checked_add(virtq::Layout::query_size(h2g_queue_depth))? + let g2h_pool_offset = h2g_ring_offset + .checked_add(h2g.checked_ring_len()?)? .checked_next_multiple_of(crate::vmem::PAGE_SIZE)?; - let pool_size = g2h_pool_pages - .checked_add(h2g_pool_pages)? - .checked_mul(crate::vmem::PAGE_SIZE)?; + let g2h_pool_len = g2h.checked_pool_len()?; + let h2g_pool_offset = g2h_pool_offset.checked_add(g2h_pool_len)?; - fixed.checked_add(ring_pages)?.checked_add(pool_size) - }); + let h2g_pool_len = h2g.checked_pool_len()?; + let len = h2g_pool_offset.checked_add(h2g_pool_len)?; - size.unwrap_or(usize::MAX) + let addr = |offset: usize| base_addr.checked_add(u64::try_from(offset).ok()?); + let _end_addr = addr(len)?; + + Some(Self { + g2h_ring_addr: base_addr, + h2g_ring_addr: addr(h2g_ring_offset)?, + g2h_pool_addr: addr(g2h_pool_offset)?, + h2g_pool_addr: addr(h2g_pool_offset)?, + ring_span_len: g2h_pool_offset, + len, + }) + } + + /// Compute the total arena size without assigning an address. + pub fn checked_query_size(g2h: QueueDims, h2g: QueueDims) -> Option { + Some(Self::new(0, g2h, h2g)?.len) + } + + /// Base address of the arena. + pub const fn base_addr(&self) -> u64 { + self.g2h_ring_addr + } + + /// Address of the G2H ring. + pub const fn g2h_ring_addr(&self) -> u64 { + self.g2h_ring_addr + } + + /// Address of the H2G ring. + pub const fn h2g_ring_addr(&self) -> u64 { + self.h2g_ring_addr + } + + /// Address of the G2H pool. + pub const fn g2h_pool_addr(&self) -> u64 { + self.g2h_pool_addr + } + + /// Address of the H2G pool. + pub const fn h2g_pool_addr(&self) -> u64 { + self.h2g_pool_addr + } + + /// Page-aligned length occupied by both rings. + pub const fn ring_span_len(&self) -> usize { + self.ring_span_len + } + + /// Total page-aligned arena length. + pub const fn size(&self) -> usize { + self.len + } + + /// Exclusive end address of the arena. + pub const fn end_addr(&self) -> u64 { + self.g2h_ring_addr + self.len as u64 + } + + /// Convert the arena's absolute addresses into offsets from the arena base. + pub fn to_offsets(&self) -> (usize, usize, usize, usize) { + // Already validated by `TransportArena::new`. + let to_offset = |addr| usize::try_from(addr - self.g2h_ring_addr).unwrap(); + + ( + to_offset(self.h2g_ring_addr), + to_offset(self.g2h_pool_addr), + to_offset(self.h2g_pool_addr), + self.len, + ) + } } #[cfg(test)] mod tests { use super::*; + #[test] + fn transport_arena_derives_aligned_regions() { + let base = 0x1_0000; + let g2h = QueueDims::new(64, 8).unwrap(); + let h2g = QueueDims::new(32, 4).unwrap(); + let arena = TransportArena::new(base, g2h, h2g).unwrap(); + + assert_eq!(arena.g2h_ring_addr(), base); + assert!( + arena + .h2g_ring_addr() + .is_multiple_of(virtq::Descriptor::ALIGN as u64) + ); + assert!( + arena + .g2h_pool_addr() + .is_multiple_of(crate::vmem::PAGE_SIZE as u64) + ); + assert_eq!( + arena.h2g_pool_addr(), + base + 9 * crate::vmem::PAGE_SIZE as u64 + ); + assert_eq!(arena.end_addr(), base + 13 * crate::vmem::PAGE_SIZE as u64); + assert_eq!( + arena.to_offsets(), + ( + 0x410, + crate::vmem::PAGE_SIZE, + 9 * crate::vmem::PAGE_SIZE, + 13 * crate::vmem::PAGE_SIZE, + ) + ); + assert_eq!(arena.ring_span_len(), crate::vmem::PAGE_SIZE); + assert_eq!(arena.size(), 13 * crate::vmem::PAGE_SIZE); + assert_eq!( + TransportArena::checked_query_size(g2h, h2g), + Some(arena.size()) + ); + assert_eq!(TransportArena::new(base + 1, g2h, h2g), None); + assert_eq!(QueueDims::new(3, 8), None); + assert_eq!(QueueDims::new(64, 0), None); + assert_eq!(QueueDims::new(usize::MAX, 8), None); + let oversized = QueueDims::new(64, usize::MAX).unwrap(); + assert_eq!(TransportArena::new(base, oversized, h2g), None); + assert_eq!( + TransportArena::new(u64::MAX - crate::vmem::PAGE_SIZE as u64 + 1, g2h, h2g,), + None + ); + } + #[test] fn minimum_scratch_includes_ring_arena_and_pools() { let fixed = arch::min_scratch_size(0, 0).unwrap(); @@ -171,5 +361,6 @@ mod tests { #[test] fn minimum_scratch_saturates_on_overflow() { assert_eq!(usize::MAX, min_scratch_size(0, 0, 64, 32, usize::MAX, 4)); + assert_eq!(usize::MAX, min_scratch_size(0, 0, usize::MAX, 32, 8, 4)); } } diff --git a/src/hyperlight_common/src/virtq/mod.rs b/src/hyperlight_common/src/virtq/mod.rs index da13bcd65..906c3f14f 100644 --- a/src/hyperlight_common/src/virtq/mod.rs +++ b/src/hyperlight_common/src/virtq/mod.rs @@ -269,6 +269,11 @@ const fn align_up(val: usize, align: usize) -> usize { val.next_multiple_of(align) } +#[inline] +const fn align_up_checked(val: usize, align: usize) -> Option { + val.checked_next_multiple_of(align) +} + impl Layout { /// Create a Layout from a base address and number of descriptors. /// @@ -345,6 +350,18 @@ impl Layout { dev_evt_offset + event_size } + + /// Calculate the ring size, returning `None` on arithmetic overflow. + pub fn checked_query_size(num_descs: usize) -> Option { + let desc_size = num_descs.checked_mul(Descriptor::SIZE)?; + let event_size = EventSuppression::SIZE; + let align = EventSuppression::ALIGN; + + let drv_evt_offset = align_up_checked(desc_size, align)?; + let dev_evt_offset = align_up_checked(drv_evt_offset.checked_add(event_size)?, align)?; + + dev_evt_offset.checked_add(event_size) + } } /// Statistics about the current virtqueue state. diff --git a/src/hyperlight_guest/src/error.rs b/src/hyperlight_guest/src/error.rs index 58eb31de3..a6014fde6 100644 --- a/src/hyperlight_guest/src/error.rs +++ b/src/hyperlight_guest/src/error.rs @@ -6,6 +6,7 @@ use alloc::string::{String, ToString as _}; pub use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; use hyperlight_common::func::Error as FuncError; +use hyperlight_common::virtq::VirtqError; use {anyhow, serde_json}; pub type Result = core::result::Result; @@ -67,6 +68,15 @@ impl From for HyperlightGuestError { } } +impl From for HyperlightGuestError { + fn from(error: VirtqError) -> Self { + Self { + kind: ErrorCode::GuestError, + message: format!("virtq: {error}"), + } + } +} + /// Extension trait to add context to `Option` and `Result` types in guest code, /// converting them to `Result`. /// diff --git a/src/hyperlight_guest/src/layout.rs b/src/hyperlight_guest/src/layout.rs index 6d403768c..e689a0c4f 100644 --- a/src/hyperlight_guest/src/layout.rs +++ b/src/hyperlight_guest/src/layout.rs @@ -26,17 +26,8 @@ pub fn snapshot_generation_gva() -> *mut u64 { pub fn g2h_queue_depth_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET) } -pub fn g2h_ring_gpa_gva() -> *mut u64 { - scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_RING_GPA_OFFSET) -} -pub fn h2g_ring_gpa_gva() -> *mut u64 { - scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_RING_GPA_OFFSET) -} -pub fn g2h_pool_gpa_gva() -> *mut u64 { - scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_POOL_GPA_OFFSET) -} -pub fn h2g_pool_gpa_gva() -> *mut u64 { - scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_POOL_GPA_OFFSET) +pub fn transport_arena_gpa_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET) } pub fn g2h_pool_pages_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_POOL_PAGES_OFFSET) diff --git a/src/hyperlight_guest/src/lib.rs b/src/hyperlight_guest/src/lib.rs index 4f793e944..aa723d784 100644 --- a/src/hyperlight_guest/src/lib.rs +++ b/src/hyperlight_guest/src/lib.rs @@ -12,6 +12,7 @@ pub mod error; pub mod exit; pub mod layout; pub mod prim_alloc; +pub mod transport; pub mod types; pub mod guest_handle { diff --git a/src/hyperlight_guest/src/transport/context.rs b/src/hyperlight_guest/src/transport/context.rs new file mode 100644 index 000000000..5a9e11d11 --- /dev/null +++ b/src/hyperlight_guest/src/transport/context.rs @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! Guest virtqueue context. + +use core::result; + +use hyperlight_common::virtq::{ + AllocError, G2H_LOWER_SLOT_COUNT, G2H_LOWER_SLOT_SIZE, Layout, Notifier, QueueStats, + SlotLayout, SlotPool, VirtqProducer, +}; + +use super::GuestMemOps; +use crate::error::{GuestErrorContext, Result}; + +/// Guest-side notifier for polled transport operation. +#[derive(Clone, Copy)] +pub struct GuestNotifier; + +impl Notifier for GuestNotifier { + fn notify(&self, _stats: QueueStats) {} +} + +/// Type alias for the guest-side G2H producer. +pub type G2hProducer = VirtqProducer; + +/// Type alias for the guest-side H2G producer. +pub type H2gProducer = VirtqProducer; + +/// Configuration for one queue passed to [`GuestContext::new`]. +pub struct QueueConfig { + /// Ring descriptor layout in shared memory. + pub layout: Layout, + /// Base GVA of the buffer pool region. + pub pool_gva: u64, + /// Number of pages in the buffer pool. + pub pool_pages: usize, + /// Size of each upper-tier buffer. + pub buffer_size: usize, +} + +/// Virtqueue runtime state for guest-host communication. +pub struct GuestContext { + /// Guest-to-host driver. + _g2h_producer: G2hProducer, + /// Host-to-guest driver. + h2g_producer: H2gProducer, + /// Size of each prefilled H2G buffer. + h2g_slot_size: usize, +} + +impl GuestContext { + /// Create a new context with G2H and H2G queues. + pub fn new(g2h: QueueConfig, h2g: QueueConfig) -> Result { + Self::with_mem(g2h, h2g, GuestMemOps::for_scratch()) + } + + /// Create a new context with memory access provided. + fn with_mem(g2h: QueueConfig, h2g: QueueConfig, mem: GuestMemOps) -> Result { + let g2h_pool = g2h_pool(g2h.pool_gva, g2h.pool_pages, g2h.buffer_size) + .with_context(|| "failed to create G2H pool")?; + let g2h_producer = VirtqProducer::new(g2h.layout, mem, GuestNotifier, g2h_pool); + + let h2g_pool = h2g_pool(h2g.pool_gva, h2g.pool_pages, h2g.buffer_size) + .with_context(|| "failed to create H2G slot pool")?; + let h2g_producer = VirtqProducer::new(h2g.layout, mem, GuestNotifier, h2g_pool); + + let mut ctx = Self { + _g2h_producer: g2h_producer, + h2g_producer, + h2g_slot_size: h2g.buffer_size, + }; + + ctx.prefill_h2g().expect("H2G initial prefill failed"); + Ok(ctx) + } + + /// Pre-fill H2G with writable buffers until its ring or pool is full. + fn prefill_h2g(&mut self) -> Result<()> { + let mut batch = self.h2g_producer.batch(); + + loop { + let chain = match batch.chain().writable(self.h2g_slot_size).build() { + Ok(chain) => chain, + Err(error) if error.is_transient() => { + batch.finish()?; + return Ok(()); + } + Err(error) => return Err(error.into()), + }; + + match batch.submit(chain) { + Ok(_) => {} + Err(error) if error.is_transient() => { + batch.finish()?; + return Ok(()); + } + Err(error) => return Err(error.into()), + } + } + } +} + +fn pool_len(pages: usize) -> result::Result { + pages + .checked_mul(hyperlight_common::vmem::PAGE_SIZE) + .ok_or(AllocError::Overflow) +} + +/// Build the uniform H2G pool. +/// +/// Every preposted receive buffer has the configured size so the host sees one +/// predictable capacity for guest calls. +fn h2g_pool(base: u64, pages: usize, buffer_size: usize) -> result::Result { + let count = pool_len(pages)? / buffer_size; + SlotPool::new(SlotLayout::new(base, buffer_size, count)) +} + +/// Build the tiered G2H pool. +/// +/// One page of 256-byte slots serves small control and log messages without +/// consuming configured-size slots. Complete slots in the remaining pages form +/// the upper tier. +fn g2h_pool(base: u64, pages: usize, upper_size: usize) -> result::Result { + let pool_len = pool_len(pages)?; + let lower_len = G2H_LOWER_SLOT_COUNT + .checked_mul(G2H_LOWER_SLOT_SIZE) + .ok_or(AllocError::Overflow)?; + + let upper_len = pool_len + .checked_sub(lower_len) + .ok_or(AllocError::EmptyRegion)?; + + let upper_count = upper_len / upper_size; + + let lower = SlotLayout::new(base, G2H_LOWER_SLOT_SIZE, G2H_LOWER_SLOT_COUNT); + let upper = SlotLayout::new(lower.end_addr()?, upper_size, upper_count); + SlotPool::new_tiered(lower, upper) +} diff --git a/src/hyperlight_guest/src/transport/mem.rs b/src/hyperlight_guest/src/transport/mem.rs new file mode 100644 index 000000000..f2f9e4e16 --- /dev/null +++ b/src/hyperlight_guest/src/transport/mem.rs @@ -0,0 +1,135 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! Guest-side [`MemOps`] implementation for virtqueue access. + +use core::mem::{align_of, size_of}; +use core::sync::atomic::{AtomicU16, Ordering}; + +use hyperlight_common::virtq::MemOps; + +use crate::layout; + +/// Guest-side memory accessor for GVA-valued virtqueue addresses. +#[derive(Clone, Copy, Debug)] +pub struct GuestMemOps { + scratch_gva: u64, + scratch_end: u64, +} + +/// Invalid guest virtqueue memory access. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct GuestMemError; + +impl GuestMemOps { + pub(super) fn for_scratch() -> Self { + let scratch_len = unsafe { layout::scratch_size_gva().read_volatile() }; + // SAFETY: Generic initialization keeps the scratch GVA range mapped. + unsafe { Self::from_raw_parts(layout::scratch_base_gva(), scratch_len) } + } + + /// Create an accessor for a scratch virtual address range. + /// + /// # Safety + /// + /// The range must remain mapped for this value's lifetime. Peer access must + /// follow virtqueue descriptor ownership. + pub unsafe fn from_raw_parts(scratch_gva: u64, scratch_len: u64) -> Self { + let scratch_end = scratch_gva + .checked_add(scratch_len) + .expect("scratch end overflow"); + + Self { + scratch_gva, + scratch_end, + } + } + + fn ptr(&self, addr: u64, len: usize) -> Result<*mut u8, GuestMemError> { + let end = addr.checked_add(len as u64).ok_or(GuestMemError)?; + if addr < self.scratch_gva || end > self.scratch_end { + return Err(GuestMemError); + } + Ok(addr as *mut u8) + } + + fn atomic(&self, addr: u64) -> Result<&AtomicU16, GuestMemError> { + let ptr = self.ptr(addr, size_of::())?; + if !(ptr as usize).is_multiple_of(align_of::()) { + return Err(GuestMemError); + } + // SAFETY: `ptr` is inside the live scratch mapping and is aligned. + Ok(unsafe { &*ptr.cast::() }) + } +} + +// SAFETY: Every address is restricted to the scratch mapping. Payload +// references rely on descriptor ownership, and ring flags use aligned atomics. +unsafe impl MemOps for GuestMemOps { + type Error = GuestMemError; + + fn read(&self, addr: u64, dst: &mut [u8]) -> Result<(), Self::Error> { + let src = self.ptr(addr, dst.len())?; + // SAFETY: `src` covers `dst.len()` initialized scratch bytes. + unsafe { src.copy_to_nonoverlapping(dst.as_mut_ptr(), dst.len()) }; + Ok(()) + } + + fn write(&self, addr: u64, src: &[u8]) -> Result<(), Self::Error> { + let dst = self.ptr(addr, src.len())?; + // SAFETY: `dst` covers `src.len()` scratch bytes. + unsafe { src.as_ptr().copy_to_nonoverlapping(dst, src.len()) }; + Ok(()) + } + + fn load_acquire(&self, addr: u64) -> Result { + Ok(self.atomic(addr)?.load(Ordering::Acquire)) + } + + fn store_release(&self, addr: u64, val: u16) -> Result<(), Self::Error> { + self.atomic(addr)?.store(val, Ordering::Release); + Ok(()) + } + + unsafe fn as_slice(&self, addr: u64, len: usize) -> Result<&[u8], Self::Error> { + let ptr = self.ptr(addr, len)?; + // SAFETY: The caller upholds descriptor ownership for this range. + Ok(unsafe { core::slice::from_raw_parts(ptr, len) }) + } + + #[allow(clippy::mut_from_ref)] + unsafe fn as_mut_slice(&self, addr: u64, len: usize) -> Result<&mut [u8], Self::Error> { + let ptr = self.ptr(addr, len)?; + // SAFETY: The caller upholds exclusive descriptor ownership. + Ok(unsafe { core::slice::from_raw_parts_mut(ptr, len) }) + } +} + +#[cfg(test)] +mod tests { + use alloc::vec; + use core::mem::size_of; + + use hyperlight_common::virtq::MemOps; + + use super::*; + + #[test] + fn guest_mem_access_is_bounded_by_scratch() { + const LEN: usize = 0x4000; + let mut backing = vec![0u64; LEN / size_of::()]; + let base = backing.as_mut_ptr() as usize as u64; + let mem = unsafe { GuestMemOps::from_raw_parts(base, LEN as u64) }; + + mem.write(base, &[1, 2, 3, 4]).unwrap(); + let mut bytes = [0; 4]; + mem.read(base, &mut bytes).unwrap(); + assert_eq!(bytes, [1, 2, 3, 4]); + + mem.store_release(base, 0x1234).unwrap(); + assert_eq!(mem.load_acquire(base).unwrap(), 0x1234); + + assert!(mem.write(base + LEN as u64 - 1, &[1, 2]).is_err()); + assert!(mem.load_acquire(base + 1).is_err()); + } +} diff --git a/src/hyperlight_guest/src/transport/mod.rs b/src/hyperlight_guest/src/transport/mod.rs new file mode 100644 index 000000000..6a7101e7a --- /dev/null +++ b/src/hyperlight_guest/src/transport/mod.rs @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! Guest transport context and memory access. +//! +//! Global context is installed once via [`set_global_context`] and accessed via [`with_context`]. + +pub mod context; +pub mod mem; + +use core::cell::RefCell; +use core::sync::atomic::{AtomicU8, Ordering}; + +pub use context::{GuestContext, QueueConfig}; +pub use mem::GuestMemOps; + +const UNINITIALIZED: u8 = 0; +const INITIALIZED: u8 = 1; + +static INIT_STATE: AtomicU8 = AtomicU8::new(UNINITIALIZED); +static GLOBAL_CONTEXT: SyncWrap>> = SyncWrap(RefCell::new(None)); + +struct SyncWrap(T); + +// SAFETY: Hyperlight guests have one vCPU and serialize guest entry. +unsafe impl Sync for SyncWrap {} + +/// Whether the virtqueue context is installed. +pub fn is_initialized() -> bool { + INIT_STATE.load(Ordering::Acquire) == INITIALIZED +} + +/// Run a closure with the global virtqueue context. +/// +/// # Panics +/// +/// Panics if the context is uninitialized or already borrowed. +pub fn with_context(f: impl FnOnce(&mut GuestContext) -> R) -> R { + assert!(is_initialized(), "transport context not initialized"); + let mut context = GLOBAL_CONTEXT.0.borrow_mut(); + f(context.as_mut().expect("transport context missing")) +} + +/// Install the global transport context. +/// +/// # Panics +/// +/// Panics if a context was already installed. +pub fn set_global_context(context: GuestContext) { + assert!( + INIT_STATE + .compare_exchange( + UNINITIALIZED, + INITIALIZED, + Ordering::SeqCst, + Ordering::SeqCst, + ) + .is_ok(), + "virtqueue context already initialized" + ); + *GLOBAL_CONTEXT.0.borrow_mut() = Some(context); +} diff --git a/src/hyperlight_guest_bin/src/lib.rs b/src/hyperlight_guest_bin/src/lib.rs index 1bd765797..398dc77a6 100644 --- a/src/hyperlight_guest_bin/src/lib.rs +++ b/src/hyperlight_guest_bin/src/lib.rs @@ -39,6 +39,7 @@ pub mod guest_logger; pub mod host_comm; pub mod memory; pub mod paging; +pub mod transport; /// Bridge between picolibc's POSIX expectations and the Hyperlight host. /// cbindgen:ignore @@ -291,6 +292,9 @@ pub(crate) extern "C" fn generic_init( registration(); } + // Prepare transport before guest code starts. + transport::initialize(); + unsafe { hyperlight_main(); } diff --git a/src/hyperlight_guest_bin/src/transport.rs b/src/hyperlight_guest_bin/src/transport.rs new file mode 100644 index 000000000..a920df7a9 --- /dev/null +++ b/src/hyperlight_guest_bin/src/transport.rs @@ -0,0 +1,90 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! Guest virtqueue initialization. + +use hyperlight_common::layout::{QueueDims, TransportArena}; +use hyperlight_common::virtq::Layout; +use hyperlight_guest::transport::{GuestContext, QueueConfig}; +use hyperlight_guest::{layout, transport as guest_transport}; + +use crate::paging::phys_to_virt; + +/// Initialize the guest transport queues in host-assigned scratch regions. +pub(crate) fn initialize() { + // The host writes normalized transport dimensions and the arena base before entry. + // SAFETY: Generic initialization has mapped writable scratch metadata. + let transport_arena_gpa = unsafe { layout::transport_arena_gpa_gva().read_volatile() }; + + let (depth, pages, g2h_bufsz) = read_published_g2h(); + let g2h = QueueDims::new(depth, pages).expect("invalid G2H queue dimensions"); + + let (depth, pages, h2g_bufsz) = read_published_h2g(); + let h2g = QueueDims::new(depth, pages).expect("invalid H2G queue dimensions"); + + assert!(g2h_bufsz > 0 && h2g_bufsz > 0); + + let arena = TransportArena::new(transport_arena_gpa, g2h, h2g).expect("invalid virtq arena"); + let g2h_pages = g2h.pool_pages().get(); + let h2g_pages = h2g.pool_pages().get(); + + let g2h_ring_gva = scratch_gva(arena.g2h_ring_addr()); + let h2g_ring_gva = scratch_gva(arena.h2g_ring_addr()); + let g2h_pool_gva = scratch_gva(arena.g2h_pool_addr()); + let h2g_pool_gva = scratch_gva(arena.h2g_pool_addr()); + + let g2h_layout = + unsafe { Layout::from_base(g2h_ring_gva, g2h.depth()) }.expect("G2H layout is invalid"); + let h2g_layout = + unsafe { Layout::from_base(h2g_ring_gva, h2g.depth()) }.expect("H2G layout is invalid"); + + // Build the queues and prefill H2G before exposing either queue to the host. + let context = GuestContext::new( + QueueConfig { + layout: g2h_layout, + pool_gva: g2h_pool_gva, + pool_pages: g2h_pages, + buffer_size: g2h_bufsz, + }, + QueueConfig { + layout: h2g_layout, + pool_gva: h2g_pool_gva, + pool_pages: h2g_pages, + buffer_size: h2g_bufsz, + }, + ) + .expect("failed to create guest context"); + + guest_transport::set_global_context(context); +} + +fn scratch_gva(gpa: u64) -> u64 { + let ptr = phys_to_virt(gpa).expect("transport GPA is outside scratch"); + u64::try_from(ptr as usize).expect("transport GVA exceeds u64") +} + +fn read_published_g2h() -> (usize, usize, usize) { + // SAFETY: Generic initialization has mapped writable scratch metadata. + let depth_raw = unsafe { layout::g2h_queue_depth_gva().read_volatile() }; + let pages_raw = unsafe { layout::g2h_pool_pages_gva().read_volatile() }; + let bufsz_raw = unsafe { layout::g2h_buffer_size_gva().read_volatile() }; + + let depth = usize::try_from(depth_raw).expect("G2H queue depth exceeds usize"); + let pages = usize::try_from(pages_raw).expect("G2H pool page count exceeds usize"); + let bufsz = usize::try_from(bufsz_raw).expect("G2H buffer size exceeds usize"); + + (depth, pages, bufsz) +} + +fn read_published_h2g() -> (usize, usize, usize) { + // SAFETY: Generic initialization has mapped writable scratch metadata. + let depth_raw = unsafe { layout::h2g_queue_depth_gva().read_volatile() }; + let pages_raw = unsafe { layout::h2g_pool_pages_gva().read_volatile() }; + let bufsz_raw = unsafe { layout::h2g_buffer_size_gva().read_volatile() }; + + let depth = usize::try_from(depth_raw).expect("H2G queue depth exceeds usize"); + let pages = usize::try_from(pages_raw).expect("H2G pool page count exceeds usize"); + let bufsz = usize::try_from(bufsz_raw).expect("H2G buffer size exceeds usize"); + + (depth, pages, bufsz) +} diff --git a/src/hyperlight_host/src/mem/layout.rs b/src/hyperlight_host/src/mem/layout.rs index a40d74548..cbc97e13b 100644 --- a/src/hyperlight_host/src/mem/layout.rs +++ b/src/hyperlight_host/src/mem/layout.rs @@ -35,13 +35,17 @@ //! There is also a scratch region at the top of physical memory, //! which is mostly laid out as a large undifferentiated blob of //! memory, although at present the snapshot process specially -//! privileges the statically allocated input and output data regions: +//! privileges fixed input, output, and transport regions: //! //! +-------------------------------------------+ (top of physical memory) //! | Exception Stack, Metadata | //! +-------------------------------------------+ (1 page below) //! | Scratch Memory | //! +-------------------------------------------+ +//! | Guest Page Tables | +//! +-------------------------------------------+ +//! | Transport Arena | +//! +-------------------------------------------+ //! | Output Data | //! +-------------------------------------------+ //! | Input Data | @@ -50,6 +54,7 @@ use std::fmt::Debug; use std::mem::size_of; +use hyperlight_common::layout::TransportArena; use hyperlight_common::mem::HyperlightPEB; use hyperlight_common::vmem::PAGE_SIZE; use tracing::{Span, instrument}; @@ -350,6 +355,11 @@ impl SandboxMemoryLayout { if scratch_size > Self::MAX_MEMORY_SIZE { return Err(MemoryRequestTooBig(scratch_size, Self::MAX_MEMORY_SIZE)); } + if !scratch_size.is_multiple_of(PAGE_SIZE) { + return Err(new_error!( + "scratch size {scratch_size} must be a multiple of {PAGE_SIZE}" + )); + } let input_data_size = cfg.get_input_data_size(); let output_data_size = cfg.get_output_data_size(); let g2h_queue_depth = cfg.get_g2h_queue_depth(); @@ -449,6 +459,16 @@ impl SandboxMemoryLayout { self.h2g_pool_pages } + pub(crate) fn get_g2h_queue_dims(&self) -> hyperlight_common::layout::QueueDims { + hyperlight_common::layout::QueueDims::new(self.g2h_queue_depth, self.g2h_pool_pages) + .expect("validated G2H queue dimensions") + } + + pub(crate) fn get_h2g_queue_dims(&self) -> hyperlight_common::layout::QueueDims { + hyperlight_common::layout::QueueDims::new(self.h2g_queue_depth, self.h2g_pool_pages) + .expect("validated H2G queue dimensions") + } + /// Guest-visible prefix size of the snapshot blob. pub(crate) fn snapshot_size(&self) -> usize { self.snapshot_size @@ -740,7 +760,7 @@ impl SandboxMemoryLayout { /// Offset from the beginning of the scratch region to the location /// where page tables are eagerly copied on restore. pub(crate) fn get_pt_base_scratch_offset(&self) -> usize { - (self.input_data_size + self.output_data_size).next_multiple_of(PAGE_SIZE) + self.get_virtq_base_scratch_offset() + self.get_transport_arena().size() } /// Base GPA to which the page tables are eagerly copied on restore. @@ -754,6 +774,24 @@ impl SandboxMemoryLayout { self.get_pt_base_gpa() + self.pt_size.unwrap_or(0) as u64 } + fn get_virtq_base_scratch_offset(&self) -> usize { + (self.input_data_size + self.output_data_size) + .next_multiple_of(hyperlight_common::vmem::PAGE_SIZE) + } + + /// Exact transport placement in the fixed scratch prefix. + pub(crate) fn get_transport_arena(&self) -> TransportArena { + let base_gpa = hyperlight_common::layout::scratch_base_gpa(self.scratch_size) + + self.get_virtq_base_scratch_offset() as u64; + + TransportArena::new( + base_gpa, + self.get_g2h_queue_dims(), + self.get_h2g_queue_dims(), + ) + .expect("validated virtqueue arena dimensions") + } + /// Total size of guest memory in `self`'s memory layout. fn get_unaligned_memory_size(&self) -> usize { self.init_data_offset() + self.init_data_size @@ -837,6 +875,21 @@ mod tests { assert!(matches!(layout, Err(MemoryRequestTooSmall(_, usize::MAX)))); } + #[test] + fn rejects_unaligned_scratch_size() { + let mut cfg = SandboxConfiguration::default(); + cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 1); + + let error = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap_err(); + assert_eq!( + error.to_string(), + format!( + "scratch size {} must be a multiple of {PAGE_SIZE}", + SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 1 + ) + ); + } + #[test] fn test_max_memory_sandbox() { let mut cfg = SandboxConfiguration::default(); @@ -923,7 +976,15 @@ mod tests { pin_eq!(layout.get_input_data_buffer_scratch_host_offset(), 0); pin_eq!(layout.get_output_data_buffer_scratch_host_offset(), 0x2000); - pin_eq!(layout.get_pt_base_scratch_offset(), 0x4000); + pin_eq!(layout.get_pt_base_scratch_offset(), 0x11000); + + let arena = layout.get_transport_arena(); + let scratch_base_gpa = hyperlight_common::layout::scratch_base_gpa(0x20000); + pin_eq!(arena.g2h_ring_addr() - scratch_base_gpa, 0x4000); + pin_eq!(arena.h2g_ring_addr() - scratch_base_gpa, 0x4410); + pin_eq!(arena.g2h_pool_addr() - scratch_base_gpa, 0x5000); + pin_eq!(arena.h2g_pool_addr() - scratch_base_gpa, 0xd000); + pin_eq!(arena.end_addr() - scratch_base_gpa, 0x11000); // The output buffer sits one input buffer past the input // buffer in the guest's scratch view. @@ -942,7 +1003,7 @@ mod tests { ); pin_eq!( layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x20000), - 0x4000 + 0x11000 ); // pt_size is zero here, so the first free scratch GPA equals // the page table base. @@ -975,7 +1036,15 @@ mod tests { pin_eq!(layout.get_input_data_buffer_scratch_host_offset(), 0); pin_eq!(layout.get_output_data_buffer_scratch_host_offset(), 0x4000); - pin_eq!(layout.get_pt_base_scratch_offset(), 0x6000); + pin_eq!(layout.get_pt_base_scratch_offset(), 0x13000); + + let arena = layout.get_transport_arena(); + let scratch_base_gpa = hyperlight_common::layout::scratch_base_gpa(0x30000); + pin_eq!(arena.g2h_ring_addr() - scratch_base_gpa, 0x6000); + pin_eq!(arena.h2g_ring_addr() - scratch_base_gpa, 0x6410); + pin_eq!(arena.g2h_pool_addr() - scratch_base_gpa, 0x7000); + pin_eq!(arena.h2g_pool_addr() - scratch_base_gpa, 0xf000); + pin_eq!(arena.end_addr() - scratch_base_gpa, 0x13000); pin_eq!( layout.get_output_data_buffer_gva() - layout.get_input_data_buffer_gva(), @@ -989,7 +1058,7 @@ mod tests { ); pin_eq!( layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x30000), - 0x6000 + 0x13000 ); pin_eq!( layout.get_first_free_scratch_gpa(), diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs index 520ade313..11efee1ed 100644 --- a/src/hyperlight_host/src/mem/mgr.rs +++ b/src/hyperlight_host/src/mem/mgr.rs @@ -17,6 +17,7 @@ use super::layout::SandboxMemoryLayout; use super::shared_mem::{ ExclusiveSharedMemory, GuestSharedMemory, HostSharedMemory, ReadonlySharedMemory, SharedMemory, }; +use super::virtq::{self, G2hConsumer, H2gConsumer}; use crate::hypervisor::regs::CommonSpecialRegisters; use crate::mem::memory_region::MemoryRegion; #[cfg(crashdump)] @@ -117,6 +118,7 @@ impl ReadonlySharedMemory { } } pub(crate) use unused_hack::SnapshotSharedMemory; + /// A struct that is responsible for laying out and managing the memory /// for a given `Sandbox`. pub(crate) struct SandboxMemoryManager { @@ -141,6 +143,26 @@ pub(crate) struct SandboxMemoryManager { /// restored snapshot's own generation number so the guest-visible /// counter tracks which snapshot the sandbox is a clone of. pub(crate) snapshot_count: u64, + /// G2H consumer bound to the current scratch mapping. + pub(crate) g2h_consumer: Option, + /// H2G consumer bound to the current scratch mapping. + pub(crate) h2g_consumer: Option, +} + +impl Clone for SandboxMemoryManager { + fn clone(&self) -> Self { + Self { + shared_mem: self.shared_mem.clone(), + scratch_mem: self.scratch_mem.clone(), + layout: self.layout, + next_action: self.next_action, + original_entrypoint: self.original_entrypoint, + abort_buffer: self.abort_buffer.clone(), + snapshot_count: self.snapshot_count, + g2h_consumer: None, + h2g_consumer: None, + } + } } /// Buffer for building guest page tables during snapshot creation. @@ -276,6 +298,8 @@ where original_entrypoint: 0, abort_buffer: Vec::new(), snapshot_count: 0, + g2h_consumer: None, + h2g_consumer: None, } } @@ -358,6 +382,8 @@ impl SandboxMemoryManager { original_entrypoint: self.original_entrypoint, abort_buffer: self.abort_buffer, snapshot_count: self.snapshot_count, + g2h_consumer: None, + h2g_consumer: None, }; let guest_mgr = SandboxMemoryManager { shared_mem: gshm, @@ -367,6 +393,8 @@ impl SandboxMemoryManager { original_entrypoint: self.original_entrypoint, abort_buffer: Vec::new(), // Guest doesn't need abort buffer snapshot_count: self.snapshot_count, + g2h_consumer: None, + h2g_consumer: None, }; host_mgr.update_scratch_bookkeeping()?; Ok((host_mgr, guest_mgr)) @@ -374,6 +402,27 @@ impl SandboxMemoryManager { } impl SandboxMemoryManager { + /// Attach host consumers to a guest-produced initial transport image. + /// + /// Before guest initialization, the host publishes queue dimensions and the + /// transport arena GPA. The guest derives and initializes every fixed region + /// without consuming dynamic scratch. + /// + /// This method runs after the initialization VM exit. It checks the + /// published arena against the host layout, derives bounded GVA views, + /// and validates each directional ring before exposing either consumer. + /// Fresh sandboxes and pre-initialization restores use this path. + pub(crate) fn attach_virtq(&mut self) -> Result<()> { + if self.g2h_consumer.is_some() || self.h2g_consumer.is_some() { + return Err(new_error!("virtqueue consumers are already attached")); + } + + let (g2h, h2g) = virtq::attach(&self.layout, &self.scratch_mem)?; + self.g2h_consumer = Some(g2h); + self.h2g_consumer = Some(h2g); + Ok(()) + } + /// Reads a host function call from memory #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] pub(crate) fn get_host_function_call(&mut self) -> Result { @@ -473,6 +522,9 @@ impl SandboxMemoryManager { Option>, Option, )> { + self.g2h_consumer = None; + self.h2g_consumer = None; + let gsnapshot = if *snapshot.memory() == self.shared_mem { // If the snapshot memory is already the correct memory, // which is readonly, don't bother with restoring it, @@ -561,6 +613,38 @@ impl SandboxMemoryManager { self.snapshot_count, )?; + // Record the G2H and H2G queue depths, pool page counts, and buffer sizes. + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET, + u64::try_from(self.layout.get_g2h_queue_depth())?, + )?; + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_G2H_POOL_PAGES_OFFSET, + u64::try_from(self.layout.get_g2h_pool_pages())?, + )?; + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET, + u64::try_from(self.layout.get_g2h_buffer_size())?, + )?; + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET, + u64::try_from(self.layout.get_h2g_queue_depth())?, + )?; + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_H2G_POOL_PAGES_OFFSET, + u64::try_from(self.layout.get_h2g_pool_pages())?, + )?; + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET, + u64::try_from(self.layout.get_h2g_buffer_size())?, + )?; + + let transport_arena = self.layout.get_transport_arena(); + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET, + transport_arena.base_addr(), + )?; + // Initialise the guest input and output data buffers in // scratch memory. TODO: remove the need for this. self.scratch_mem.write::( diff --git a/src/hyperlight_host/src/mem/mod.rs b/src/hyperlight_host/src/mem/mod.rs index 7bb110c28..693dc2af7 100644 --- a/src/hyperlight_host/src/mem/mod.rs +++ b/src/hyperlight_host/src/mem/mod.rs @@ -25,5 +25,7 @@ pub mod shared_mem; /// Utilities for writing shared memory tests #[cfg(all(test, not(miri)))] // uses proptest which isn't miri-compatible pub(crate) mod shared_mem_tests; +/// Host virtqueue attachment and validation. +pub(crate) mod virtq; #[allow(dead_code)] pub(crate) mod virtq_mem; diff --git a/src/hyperlight_host/src/mem/virtq.rs b/src/hyperlight_host/src/mem/virtq.rs new file mode 100644 index 000000000..01b641dab --- /dev/null +++ b/src/hyperlight_host/src/mem/virtq.rs @@ -0,0 +1,568 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! Host virtqueue attachment. +//! +//! The host publishes one transport arena address in scratch-top metadata. Guest +//! initialization builds both queues in those fixed regions. This module +//! validates the complete initial image before returning either consumer. + +use core::ops::Range; + +use hyperlight_common::layout::{QueueDims, TransportArena}; +use hyperlight_common::virtq::canonical::validate_canon_image; +use hyperlight_common::virtq::{ + Layout as VirtqLayout, MemOps, Notifier, QueueStats, VirtqConsumer, +}; + +use super::layout::{BaseGpaRegion, SandboxMemoryLayout}; +use super::shared_mem::{HostSharedMemory, SharedMemory}; +use super::virtq_mem::HostMemOps; +use crate::{Result, new_error}; + +/// Host-side G2H virtqueue consumer. +pub(crate) type G2hConsumer = VirtqConsumer; +/// Host-side H2G virtqueue consumer. +pub(crate) type H2gConsumer = VirtqConsumer; + +/// No-op notifier for polled host transport. +#[derive(Clone, Copy)] +pub(crate) struct HostNotifier; + +impl Notifier for HostNotifier { + fn notify(&self, _stats: QueueStats) {} +} + +/// Build both host consumers from a guest-produced initial transport image. +/// +/// The consumers are returned only after the host-assigned arena and both +/// directional ring images have passed validation. +pub(crate) fn attach( + layout: &SandboxMemoryLayout, + scratch_mem: &HostSharedMemory, +) -> Result<(G2hConsumer, H2gConsumer)> { + let validator = Validator::new(layout)?; + let arena_gpa = read_published_arena_gpa(scratch_mem)?; + let regions = validator.validate_published_arena(arena_gpa)?; + + let g2h_ring_mem = HostMemOps::new(scratch_mem, regions.g2h_ring.clone())?; + let g2h_pool_mem = HostMemOps::new(scratch_mem, regions.g2h_pool)?; + let g2h_layout = validator.validate_g2h(&g2h_ring_mem, regions.g2h_ring)?; + + let h2g_ring_mem = HostMemOps::new(scratch_mem, regions.h2g_ring.clone())?; + let h2g_pool_mem = HostMemOps::new(scratch_mem, regions.h2g_pool.clone())?; + let h2g_layout = validator.validate_h2g(&h2g_ring_mem, regions.h2g_ring, regions.h2g_pool)?; + + Ok(( + VirtqConsumer::new_split(g2h_layout, g2h_ring_mem, g2h_pool_mem, HostNotifier), + VirtqConsumer::new_split(h2g_layout, h2g_ring_mem, h2g_pool_mem, HostNotifier), + )) +} + +/// Bounded GVA regions derived from validated transport GPAs. +struct GvaRegions { + g2h_ring: Range, + h2g_ring: Range, + g2h_pool: Range, + h2g_pool: Range, +} + +#[derive(Clone, Copy)] +struct QueueConfig { + /// Address-independent queue dimensions. + dims: QueueDims, + /// Size of the ring image in bytes including event suppressions. + ring_len: usize, + /// Size of the buffer pool in bytes. + pool_len: usize, + /// Size of each buffer in the pool in bytes. + buffer_size: usize, +} + +impl QueueConfig { + fn new(dims: QueueDims, buffer_size: usize) -> Result { + let ring_len = dims + .checked_ring_len() + .ok_or_else(|| new_error!("ring size overflow"))?; + let pool_len = dims + .checked_pool_len() + .ok_or_else(|| new_error!("pool size overflow"))?; + + if buffer_size == 0 { + return Err(new_error!("buffer size is zero")); + } + + Ok(Self { + dims, + ring_len, + pool_len, + buffer_size, + }) + } +} + +/// Host-owned transport dimensions. +#[derive(Clone, Copy)] +struct Config { + /// Host-requested G2H configuration. + g2h: QueueConfig, + /// Host-requested H2G configuration. + h2g: QueueConfig, + /// Fixed host-assigned transport arena. + arena: TransportArena, + /// Number of one-descriptor chains posted before the H2G ring or pool fills. + h2g_prefill_chains: usize, +} + +impl Config { + /// Compute the host transport configuration from the memory layout. + fn from_layout(layout: &SandboxMemoryLayout) -> Result { + let g2h = QueueConfig::new(layout.get_g2h_queue_dims(), layout.get_g2h_buffer_size())?; + + let h2g = QueueConfig::new(layout.get_h2g_queue_dims(), layout.get_h2g_buffer_size())?; + + let h2g_prefill_chains = + usize::from(h2g.dims.depth().get()).min(h2g.pool_len / h2g.buffer_size); + let arena = layout.get_transport_arena(); + + Ok(Self { + g2h, + h2g, + arena, + h2g_prefill_chains, + }) + } +} + +/// Validates one initial transport image against one host layout. +struct Validator<'a> { + config: Config, + layout: &'a SandboxMemoryLayout, +} + +impl<'a> Validator<'a> { + fn new(layout: &'a SandboxMemoryLayout) -> Result { + Ok(Self { + config: Config::from_layout(layout)?, + layout, + }) + } + + /// Validate the initial G2H queue and return its layout. + fn validate_g2h(&self, mem: &M, ring: Range) -> Result { + // SAFETY: `ring` spans the configured image and `mem` keeps that image + // valid for the duration of validation. + let layout = unsafe { VirtqLayout::from_base(ring.start, self.config.g2h.dims.depth()) } + .map_err(|error| new_error!("invalid G2H ring layout: {error}"))?; + + validate_canon_image(mem, layout, 0, |_, _| false) + .map_err(|error| new_error!("invalid canonical G2H image: {error}"))?; + + Ok(layout) + } + + /// Validate the initial H2G queue and return its layout. + /// + /// Every available chain contains one configured size writable descriptor. + /// Descriptors must name distinct, slot-aligned ranges inside the H2G pool. + fn validate_h2g( + &self, + mem: &M, + ring: Range, + pool: Range, + ) -> Result { + // SAFETY: `ring` spans the configured image and `mem` keeps that image + // valid for the duration of validation. + let layout = unsafe { VirtqLayout::from_base(ring.start, self.config.h2g.dims.depth()) } + .map_err(|error| new_error!("invalid H2G ring layout: {error}"))?; + + let bufsz = self.config.h2g.buffer_size; + let prefill = self.config.h2g_prefill_chains; + + if prefill == 0 { + return Err(new_error!("H2G pool has no complete buffers")); + } + + // Record the accepted descriptor ranges to detect overlaps. + let mut accepted: Vec> = Vec::with_capacity(prefill); + + let image = validate_canon_image(mem, layout, prefill, |_, elem| { + let Ok(bufsz_u64) = u64::try_from(bufsz) else { + return false; + }; + + // all descriptors must be writable and match the configured buffer size + if !elem.writable || usize::try_from(elem.len).ok() != Some(bufsz) { + return false; + } + + let Some(offset) = elem.addr.checked_sub(pool.start) else { + return false; + }; + let Some(end) = elem.addr.checked_add(u64::from(elem.len)) else { + return false; + }; + + // all descriptors must be slot-aligned and remain inside the pool + if !offset.is_multiple_of(bufsz_u64) || end > pool.end { + return false; + } + + let buf = elem.addr..end; + + // all descriptors must name distinct ranges + if accepted + .iter() + .any(|other| buf.start < other.end && other.start < buf.end) + { + return false; + } + + accepted.push(buf); + true + }) + .map_err(|error| new_error!("invalid canonical H2G image: {error}"))?; + + // compare the number of accepted chains to the expected prefill count + if image.len() != prefill { + return Err(new_error!("invalid initial H2G chains")); + } + + Ok(layout) + } + + /// Validate the published arena and return its GVA regions. + fn validate_published_arena(&self, arena_gpa: u64) -> Result { + if arena_gpa != self.config.arena.base_addr() { + return Err(new_error!("published transport arena is invalid")); + } + + self.resolve_gva_regions() + } + + /// Translate validated transport GPAs into the GVA ranges used by descriptors. + fn resolve_gva_regions(&self) -> Result { + let to_gva = |gpa| { + let resolved = self + .layout + .resolve_gpa(gpa, &[]) + .ok_or_else(|| new_error!("GPA {gpa:#x} is outside scratch"))?; + + if !matches!(resolved.base, BaseGpaRegion::Scratch(())) { + return Err(new_error!("GPA {gpa:#x} is outside scratch")); + } + + hyperlight_common::layout::scratch_base_gva(self.layout.get_scratch_size()) + .checked_add(u64::try_from(resolved.offset)?) + .ok_or_else(|| new_error!("GPA {gpa:#x} to GVA translation overflow")) + }; + + let ( + g2h_ring_addr, + h2g_ring_addr, + g2h_pool_addr, + h2g_pool_addr, + g2h_ring_len, + h2g_ring_len, + g2h_pool_len, + h2g_pool_len, + ) = ( + self.config.arena.g2h_ring_addr(), + self.config.arena.h2g_ring_addr(), + self.config.arena.g2h_pool_addr(), + self.config.arena.h2g_pool_addr(), + self.config.g2h.ring_len, + self.config.h2g.ring_len, + self.config.g2h.pool_len, + self.config.h2g.pool_len, + ); + + Ok(GvaRegions { + g2h_ring: checked_region(to_gva(g2h_ring_addr)?, g2h_ring_len, "G2H ring")?, + h2g_ring: checked_region(to_gva(h2g_ring_addr)?, h2g_ring_len, "H2G ring")?, + g2h_pool: checked_region(to_gva(g2h_pool_addr)?, g2h_pool_len, "G2H pool")?, + h2g_pool: checked_region(to_gva(h2g_pool_addr)?, h2g_pool_len, "H2G pool")?, + }) + } +} + +/// Read the transport arena GPA from scratch-top metadata. +fn read_published_arena_gpa(scratch_mem: &HostSharedMemory) -> Result { + let offset = hyperlight_common::layout::SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET as usize; + Ok(scratch_mem.read::(scratch_mem.mem_size() - offset)?) +} + +#[cfg(test)] +fn write_published_arena_gpa(scratch_mem: &HostSharedMemory, arena_gpa: u64) -> Result<()> { + let offset = hyperlight_common::layout::SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET as usize; + Ok(scratch_mem.write::(scratch_mem.mem_size() - offset, arena_gpa)?) +} + +fn checked_region(start: u64, len: usize, tag: &str) -> Result> { + let end = start + .checked_add(u64::try_from(len)?) + .ok_or_else(|| new_error!("{tag} GVA range overflow"))?; + + Ok(start..end) +} + +#[cfg(test)] +mod tests { + use core::num::NonZeroU16; + + use hyperlight_common::virtq::{ + DescFlags, Descriptor, MemOps, SlotLayout, SlotPool, VirtqProducer, + }; + use hyperlight_common::vmem; + + use super::*; + use crate::mem::shared_mem::ExclusiveSharedMemory; + use crate::sandbox::SandboxConfiguration; + + const SCRATCH_SIZE: usize = 0x20_000; + const G2H_DEPTH: u16 = 16; + const H2G_DEPTH: u16 = 8; + const G2H_POOL_PAGES: usize = 3; + const H2G_POOL_PAGES: usize = 2; + const H2G_BUFFER_SIZE: usize = 3000; + + fn memory_layout() -> SandboxMemoryLayout { + let mut config = SandboxConfiguration::default(); + config.set_scratch_size(SCRATCH_SIZE); + config.set_g2h_queue_depth(G2H_DEPTH as usize); + config.set_h2g_queue_depth(H2G_DEPTH as usize); + config.set_h2g_buffer_size(H2G_BUFFER_SIZE); + config.set_g2h_pool_pages(G2H_POOL_PAGES); + config.set_h2g_pool_pages(H2G_POOL_PAGES); + SandboxMemoryLayout::new(config, 4096, 0, None).unwrap() + } + + fn attach_config() -> Config { + Config::from_layout(&memory_layout()).unwrap() + } + + fn host_scratch() -> HostSharedMemory { + let scratch = ExclusiveSharedMemory::new(SCRATCH_SIZE).unwrap(); + scratch.build().0 + } + + fn validate_published(arena_gpa: u64, config: Config) -> Result { + let layout = memory_layout(); + Validator { + config, + layout: &layout, + } + .validate_published_arena(arena_gpa) + } + + struct PreparedVirtq { + g2h_mem: HostMemOps, + h2g_mem: HostMemOps, + g2h_ring: Range, + h2g_ring: Range, + g2h_pool: Range, + h2g_pool: Range, + g2h_layout: VirtqLayout, + h2g_layout: VirtqLayout, + } + + fn prepared_virtq() -> PreparedVirtq { + let scratch = ExclusiveSharedMemory::new(SCRATCH_SIZE).unwrap(); + let (scratch, _) = scratch.build(); + + let layout = memory_layout(); + let config = Config::from_layout(&layout).unwrap(); + let scratch_base_gpa = hyperlight_common::layout::scratch_base_gpa(SCRATCH_SIZE); + let scratch_base_gva = hyperlight_common::layout::scratch_base_gva(SCRATCH_SIZE); + let to_gva = |gpa| scratch_base_gva + (gpa - scratch_base_gpa); + + let ring_base = to_gva(config.arena.g2h_ring_addr()); + let h2g_base = to_gva(config.arena.h2g_ring_addr()); + let g2h_pool_base = to_gva(config.arena.g2h_pool_addr()); + let g2h_pool_end = g2h_pool_base + (G2H_POOL_PAGES * vmem::PAGE_SIZE) as u64; + let h2g_pool_base = to_gva(config.arena.h2g_pool_addr()); + let h2g_pool_end = h2g_pool_base + (H2G_POOL_PAGES * vmem::PAGE_SIZE) as u64; + + // SAFETY: The scratch mapping covers both ring layouts. + let g2h_layout = unsafe { + VirtqLayout::from_base(ring_base, NonZeroU16::new(G2H_DEPTH).unwrap()).unwrap() + }; + // SAFETY: The scratch mapping covers both ring layouts. + let h2g_layout = unsafe { + VirtqLayout::from_base(h2g_base, NonZeroU16::new(H2G_DEPTH).unwrap()).unwrap() + }; + + let mem = HostMemOps::new(&scratch, ring_base..h2g_pool_end).unwrap(); + let h2g_prefill_chains = (H2G_POOL_PAGES * vmem::PAGE_SIZE) / H2G_BUFFER_SIZE; + + let h2g_pool = SlotPool::new(SlotLayout::new( + h2g_pool_base, + H2G_BUFFER_SIZE, + h2g_prefill_chains, + )) + .unwrap(); + + let mut h2g = VirtqProducer::new(h2g_layout, mem, HostNotifier, h2g_pool.clone()); + let mut batch = h2g.batch(); + + for _ in 0..h2g_pool.num_free() { + let chain = batch.chain().writable(H2G_BUFFER_SIZE).build().unwrap(); + batch.submit(chain).unwrap(); + } + + batch.finish().unwrap(); + write_published_arena_gpa(&scratch, config.arena.base_addr()).unwrap(); + + let g2h_ring = ring_base..ring_base + VirtqLayout::query_size(G2H_DEPTH as usize) as u64; + let h2g_ring = h2g_base..h2g_base + VirtqLayout::query_size(H2G_DEPTH as usize) as u64; + let g2h_pool = g2h_pool_base..g2h_pool_end; + let h2g_pool = h2g_pool_base..h2g_pool_end; + let g2h_mem = HostMemOps::new(&scratch, g2h_ring.clone()).unwrap(); + let h2g_mem = HostMemOps::new(&scratch, h2g_ring.clone()).unwrap(); + + PreparedVirtq { + g2h_mem, + h2g_mem, + g2h_ring, + h2g_ring, + g2h_pool, + h2g_pool, + g2h_layout, + h2g_layout, + } + } + + fn validate(prepared: &PreparedVirtq) -> Result<()> { + let layout = memory_layout(); + let validator = Validator::new(&layout)?; + + validator.validate_g2h(&prepared.g2h_mem, prepared.g2h_ring.clone())?; + validator.validate_h2g( + &prepared.h2g_mem, + prepared.h2g_ring.clone(), + prepared.h2g_pool.clone(), + )?; + Ok(()) + } + + fn read_desc(mem: &HostMemOps, layout: VirtqLayout, index: u16) -> Descriptor { + mem.read_val(layout.desc_table_addr() + u64::from(index) * Descriptor::SIZE as u64) + .unwrap() + } + + fn write_desc(mem: &HostMemOps, layout: VirtqLayout, index: u16, desc: Descriptor) { + mem.write_val( + layout.desc_table_addr() + u64::from(index) * Descriptor::SIZE as u64, + desc, + ) + .unwrap(); + } + + #[test] + fn validates_host_placed_regions() { + let config = attach_config(); + let regions = validate_published(config.arena.base_addr(), config).unwrap(); + + assert_eq!( + regions.g2h_ring.end - regions.g2h_ring.start, + config.g2h.ring_len as u64 + ); + assert_eq!( + regions.h2g_ring.end - regions.h2g_ring.start, + config.h2g.ring_len as u64 + ); + assert_eq!( + regions.g2h_pool.end - regions.g2h_pool.start, + config.g2h.pool_len as u64 + ); + assert_eq!( + regions.h2g_pool.end - regions.h2g_pool.start, + config.h2g.pool_len as u64 + ); + } + + #[test] + fn rejects_invalid_published_regions() { + let config = attach_config(); + let arena_gpa = config.arena.base_addr() + 1; + assert!(validate_published(arena_gpa, config).is_err()); + } + + #[test] + fn rejects_published_region_overflow() { + let config = attach_config(); + assert!(validate_published(u64::MAX, config).is_err()); + } + + #[test] + fn rejects_untranslatable_or_overflowing_gva_regions() { + let config = attach_config(); + let arena_gpa = config.arena.base_addr(); + let invalid = hyperlight_common::layout::scratch_base_gpa(SCRATCH_SIZE) - 1; + assert!(validate_published(invalid, config).is_err()); + + let mut config = config; + config.g2h.ring_len = usize::MAX; + assert!(validate_published(arena_gpa, config).is_err()); + } + + #[test] + fn validates_initial_virtq_images() { + validate(&prepared_virtq()).unwrap(); + } + + #[test] + fn rejects_h2g_descriptors_outside_pool() { + let prepared = prepared_virtq(); + let mut desc = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 0); + desc.addr = prepared.g2h_pool.start; + write_desc(&prepared.h2g_mem, prepared.h2g_layout, 0, desc); + assert!(validate(&prepared).is_err()); + } + + #[test] + fn rejects_nonzero_g2h_descriptors() { + let prepared = prepared_virtq(); + let mut desc = read_desc(&prepared.g2h_mem, prepared.g2h_layout, 0); + desc.addr = prepared.g2h_pool.start; + write_desc(&prepared.g2h_mem, prepared.g2h_layout, 0, desc); + assert!(validate(&prepared).is_err()); + } + + #[test] + fn rejects_readable_h2g_descriptor() { + let prepared = prepared_virtq(); + let mut desc = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 0); + desc.flags &= !DescFlags::WRITE.bits(); + write_desc(&prepared.h2g_mem, prepared.h2g_layout, 0, desc); + assert!(validate(&prepared).is_err()); + } + + #[test] + fn rejects_invalid_h2g_size() { + let prepared = prepared_virtq(); + let mut desc = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 0); + desc.len -= 1; + write_desc(&prepared.h2g_mem, prepared.h2g_layout, 0, desc); + assert!(validate(&prepared).is_err()); + } + + #[test] + fn rejects_misaligned_h2g_descriptor() { + let prepared = prepared_virtq(); + let mut desc = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 0); + desc.addr += 1; + write_desc(&prepared.h2g_mem, prepared.h2g_layout, 0, desc); + assert!(validate(&prepared).is_err()); + } + + #[test] + fn rejects_overlapping_h2g_descriptors() { + let prepared = prepared_virtq(); + let first = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 0); + let mut second = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 1); + second.addr = first.addr; + write_desc(&prepared.h2g_mem, prepared.h2g_layout, 1, second); + assert!(validate(&prepared).is_err()); + } +} diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index 7b317afdc..298d5d83e 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -247,6 +247,10 @@ impl MultiUseSandbox { let mgr = crate::mem::mgr::SandboxMemoryManager::from_snapshot(&snapshot)?; let (mut hshm, gshm) = mgr.build()?; + let attach_virtq = matches!( + snapshot.next_action(), + super::snapshot::NextAction::Initialise(_) + ); let page_size = u32::try_from(page_size::get())? as usize; @@ -321,6 +325,10 @@ impl MultiUseSandbox { })?; } + if attach_virtq { + hshm.attach_virtq()?; + } + let sbox = MultiUseSandbox::from_uninit(host_funcs, hshm, vm); Ok(sbox) } diff --git a/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs index 63ae19547..718d7e5c4 100644 --- a/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs +++ b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs @@ -22,6 +22,10 @@ use crate::{MultiUseSandbox, Result, UninitializedSandbox}; #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")] pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result { let (mut hshm, gshm) = u_sbox.mgr.build()?; + let attach_virtq = matches!( + hshm.next_action, + crate::sandbox::snapshot::NextAction::Initialise(_) + ); // Get the host page size. Narrowed to u32 because the guest ABI // passes it via a 32-bit register (rdx), but widened back to usize @@ -85,6 +89,10 @@ pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result Date: Fri, 31 Jul 2026 14:22:21 +0200 Subject: [PATCH 11/15] feat(snapshot): preserve canonical virtq state Persist validated G2H and H2G ring images in running snapshots. Restore fixed transport allocations and install fresh host consumers before sandbox execution. Signed-off-by: Tomasz Andrzejak --- CHANGELOG.md | 2 + docs/snapshot-oci-format.md | 8 +- docs/snapshot-versioning.md | 5 +- src/hyperlight_host/src/mem/mgr.rs | 90 +++++---- src/hyperlight_host/src/mem/virtq.rs | 173 +++++++++++++++++- src/hyperlight_host/src/mem/virtq_mem.rs | 72 +++++++- .../src/sandbox/initialized_multi_use.rs | 58 +++++- .../src/sandbox/snapshot/file/config.rs | 85 ++++++++- .../src/sandbox/snapshot/file/media_types.rs | 6 +- .../src/sandbox/snapshot/file/mod.rs | 29 ++- .../src/sandbox/snapshot/file_tests.rs | 56 +++++- .../src/sandbox/snapshot/mod.rs | 20 ++ .../src/sandbox/snapshot/tripwires.rs | 4 +- .../tests/snapshot_goldens/goldens_version.rs | 2 +- 14 files changed, 545 insertions(+), 65 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 27a12d094..3feea9cbf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). * Expose C guest `ByteChunks` values as pointer and length arrays. * Return typed `hl_ReturnValue` objects from C guest functions through `hl_result_from_*` constructors. +* Place virtqueue rings and pools in host-owned scratch before page tables. + Snapshot ABI 3 rejects snapshots created with earlier layouts. ### Removed diff --git a/docs/snapshot-oci-format.md b/docs/snapshot-oci-format.md index 971b3c868..e77b892b9 100644 --- a/docs/snapshot-oci-format.md +++ b/docs/snapshot-oci-format.md @@ -31,11 +31,11 @@ Three blob kinds per tag: * **manifest** (`application/vnd.oci.image.manifest.v1+json`). Tiny JSON pointer record selected via `index.json`. References one config and one layer by digest. -* **config** (`application/vnd.hyperlight.snapshot.config.v1+json`). The +* **config** (`application/vnd.hyperlight.snapshot.config.v2+json`). The snapshot descriptor: arch, hypervisor, CPU vendor, ABI version, - resume address and captured registers, memory layout, registered - host functions, snapshot generation counter. Loaded eagerly and - fully parsed. + resume address and captured registers, memory and transport layout, + registered host functions, snapshot generation counter. Loaded + eagerly and fully parsed. * **layer / memory** (`application/vnd.hyperlight.snapshot.memory.v1`). The raw guest memory image, exactly `memory_size` bytes. mmap'd on restore. diff --git a/docs/snapshot-versioning.md b/docs/snapshot-versioning.md index ddf2f8ce8..783813bba 100644 --- a/docs/snapshot-versioning.md +++ b/docs/snapshot-versioning.md @@ -23,8 +23,8 @@ A snapshot carries three independently evolvable version markers: `MT_SNAPSHOT_CURRENT`. This is the on-wire format of the snapshot blob: framing, section ordering, alignment, dirty/zero-page elision, anything about how the bytes are packed inside the OCI layer. -* **Config schema**, `MT_CONFIG_V1` - (`application/vnd.hyperlight.snapshot.config.v1+json`), aliased as +* **Config schema**, `MT_CONFIG_V2` + (`application/vnd.hyperlight.snapshot.config.v2+json`), aliased as `MT_CONFIG_CURRENT`. This is the JSON shape of the config blob: field names, types, required vs optional, the descriptors the loader needs in order to reconstruct the sandbox (memory sizes, buffer @@ -367,4 +367,3 @@ major: * The loader accepts the old `abi_version` (Option 2 step 4), so the old golden loads. * Register the host functions the old golden's checks call. - diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs index 11efee1ed..06654236f 100644 --- a/src/hyperlight_host/src/mem/mgr.rs +++ b/src/hyperlight_host/src/mem/mgr.rs @@ -307,37 +307,6 @@ where pub(crate) fn get_abort_buffer_mut(&mut self) -> &mut Vec { &mut self.abort_buffer } - - /// Create a snapshot with the given mapped regions - #[allow(clippy::too_many_arguments)] - pub(crate) fn snapshot( - &mut self, - mapped_regions: Vec, - root_pt_gpas: &[u64], - rsp_gva: u64, - sregs: CommonSpecialRegisters, - #[cfg(target_arch = "x86_64")] msrs: Vec, - next_action: NextAction, - host_functions: HostFunctionDetails, - ) -> Result { - self.snapshot_count += 1; - Snapshot::new( - &mut self.shared_mem, - &mut self.scratch_mem, - self.layout, - crate::mem::exe::LoadInfo::dummy(), - mapped_regions, - root_pt_gpas, - rsp_gva, - sregs, - #[cfg(target_arch = "x86_64")] - msrs, - next_action, - self.original_entrypoint, - self.snapshot_count, - host_functions, - ) - } } impl SandboxMemoryManager { @@ -402,6 +371,44 @@ impl SandboxMemoryManager { } impl SandboxMemoryManager { + /// Create a snapshot with the given mapped regions. + #[allow(clippy::too_many_arguments)] + pub(crate) fn snapshot( + &mut self, + mapped_regions: Vec, + root_pt_gpas: &[u64], + rsp_gva: u64, + sregs: CommonSpecialRegisters, + #[cfg(target_arch = "x86_64")] msrs: Vec, + next_action: NextAction, + host_functions: HostFunctionDetails, + ) -> Result { + let virtq = match (&self.g2h_consumer, &self.h2g_consumer) { + (Some(_), Some(_)) => Some(virtq::snapshot(&self.layout, &self.scratch_mem)?), + (None, None) => None, + _ => return Err(new_error!("virtqueue consumer ownership is incomplete")), + }; + + self.snapshot_count += 1; + Snapshot::new( + &mut self.shared_mem, + &mut self.scratch_mem, + self.layout, + crate::mem::exe::LoadInfo::dummy(), + mapped_regions, + root_pt_gpas, + rsp_gva, + sregs, + #[cfg(target_arch = "x86_64")] + msrs, + next_action, + self.original_entrypoint, + self.snapshot_count, + host_functions, + virtq, + ) + } + /// Attach host consumers to a guest-produced initial transport image. /// /// Before guest initialization, the host publishes queue dimensions and the @@ -423,6 +430,22 @@ impl SandboxMemoryManager { Ok(()) } + /// Restore a captured canonical transport image against this scratch mapping. + pub(crate) fn restore_virtq(&mut self, snapshot: Option<&virtq::VirtqSnapshot>) -> Result<()> { + let Some(snapshot) = snapshot else { + return Ok(()); + }; + + if self.g2h_consumer.is_some() || self.h2g_consumer.is_some() { + return Err(new_error!("virtqueue consumers are already attached")); + } + + let (g2h, h2g) = virtq::restore(&self.layout, &self.scratch_mem, snapshot)?; + self.g2h_consumer = Some(g2h); + self.h2g_consumer = Some(h2g); + Ok(()) + } + /// Reads a host function call from memory #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] pub(crate) fn get_host_function_call(&mut self) -> Result { @@ -522,6 +545,10 @@ impl SandboxMemoryManager { Option>, Option, )> { + if let Some(virtq) = snapshot.virtq() { + virtq.preflight(snapshot.layout())?; + } + self.g2h_consumer = None; self.h2g_consumer = None; @@ -567,6 +594,7 @@ impl SandboxMemoryManager { self.original_entrypoint = snapshot.original_entrypoint(); self.update_scratch_bookkeeping()?; + self.restore_virtq(snapshot.virtq())?; Ok((gsnapshot, gscratch)) } diff --git a/src/hyperlight_host/src/mem/virtq.rs b/src/hyperlight_host/src/mem/virtq.rs index 01b641dab..59543e4c6 100644 --- a/src/hyperlight_host/src/mem/virtq.rs +++ b/src/hyperlight_host/src/mem/virtq.rs @@ -17,7 +17,7 @@ use hyperlight_common::virtq::{ use super::layout::{BaseGpaRegion, SandboxMemoryLayout}; use super::shared_mem::{HostSharedMemory, SharedMemory}; -use super::virtq_mem::HostMemOps; +use super::virtq_mem::{HostMemOps, ImageMem}; use crate::{Result, new_error}; /// Host-side G2H virtqueue consumer. @@ -59,6 +59,45 @@ pub(crate) fn attach( )) } +/// Capture the canonical transport state omitted from the main memory snapshot. +pub(crate) fn snapshot( + layout: &SandboxMemoryLayout, + scratch_mem: &HostSharedMemory, +) -> Result { + let validator = Validator::new(layout)?; + + let arena_gpa = read_published_arena_gpa(scratch_mem)?; + let regions = validator.validate_published_arena(arena_gpa)?; + + let g2h_mem = HostMemOps::new(scratch_mem, regions.g2h_ring.clone())?; + validator.validate_g2h(&g2h_mem, regions.g2h_ring.clone())?; + + let h2g_mem = HostMemOps::new(scratch_mem, regions.h2g_ring.clone())?; + validator.validate_h2g(&h2g_mem, regions.h2g_ring.clone(), regions.h2g_pool.clone())?; + + // The vCPU is stopped, so the ring images and snapshotted guest producer + // bookkeeping describe the same instant. + Ok(VirtqSnapshot { + scratch_size: layout.get_scratch_size(), + g2h_ring: read_ring(scratch_mem, regions.g2h_ring)?, + h2g_ring: read_ring(scratch_mem, regions.h2g_ring)?, + }) +} + +/// Restore one captured canonical transport image and return fresh consumers. +pub(crate) fn restore( + layout: &SandboxMemoryLayout, + scratch_mem: &HostSharedMemory, + snapshot: &VirtqSnapshot, +) -> Result<(G2hConsumer, H2gConsumer)> { + let regions = Validator::new(layout)?.validate_snapshot(snapshot)?; + + write_published_arena_gpa(scratch_mem, layout.get_transport_arena().base_addr())?; + write_ring(scratch_mem, regions.g2h_ring, &snapshot.g2h_ring)?; + write_ring(scratch_mem, regions.h2g_ring, &snapshot.h2g_ring)?; + attach(layout, scratch_mem) +} + /// Bounded GVA regions derived from validated transport GPAs. struct GvaRegions { g2h_ring: Range, @@ -134,7 +173,22 @@ impl Config { } } -/// Validates one initial transport image against one host layout. +/// Canonical in-memory transport state excluded from ordinary snapshot pages. +#[derive(Debug, PartialEq, Eq)] +pub(crate) struct VirtqSnapshot { + scratch_size: usize, + g2h_ring: Vec, + h2g_ring: Vec, +} + +impl VirtqSnapshot { + /// Validate every captured field before mutating restored scratch. + pub(crate) fn preflight(&self, layout: &SandboxMemoryLayout) -> Result<()> { + Validator::new(layout)?.validate_snapshot(self).map(|_| ()) + } +} + +/// Validates live and captured transport images against one host layout. struct Validator<'a> { config: Config, layout: &'a SandboxMemoryLayout, @@ -240,6 +294,28 @@ impl<'a> Validator<'a> { self.resolve_gva_regions() } + fn validate_snapshot(&self, snapshot: &VirtqSnapshot) -> Result { + if snapshot.scratch_size != self.layout.get_scratch_size() { + return Err(new_error!( + "virtqueue snapshot scratch size {} does not match layout size {}", + snapshot.scratch_size, + self.layout.get_scratch_size() + )); + } + + let regions = self.resolve_gva_regions()?; + validate_ring_len("G2H", &snapshot.g2h_ring, self.config.g2h.ring_len)?; + validate_ring_len("H2G", &snapshot.h2g_ring, self.config.h2g.ring_len)?; + + let g2h_mem = ImageMem::new(regions.g2h_ring.start, &snapshot.g2h_ring); + self.validate_g2h(&g2h_mem, regions.g2h_ring.clone())?; + + let h2g_mem = ImageMem::new(regions.h2g_ring.start, &snapshot.h2g_ring); + self.validate_h2g(&h2g_mem, regions.h2g_ring.clone(), regions.h2g_pool.clone())?; + + Ok(regions) + } + /// Translate validated transport GPAs into the GVA ranges used by descriptors. fn resolve_gva_regions(&self) -> Result { let to_gva = |gpa| { @@ -292,12 +368,41 @@ fn read_published_arena_gpa(scratch_mem: &HostSharedMemory) -> Result { Ok(scratch_mem.read::(scratch_mem.mem_size() - offset)?) } -#[cfg(test)] fn write_published_arena_gpa(scratch_mem: &HostSharedMemory, arena_gpa: u64) -> Result<()> { let offset = hyperlight_common::layout::SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET as usize; Ok(scratch_mem.write::(scratch_mem.mem_size() - offset, arena_gpa)?) } +fn read_ring(scratch_mem: &HostSharedMemory, ring: Range) -> Result> { + let len = usize::try_from( + ring.end + .checked_sub(ring.start) + .ok_or_else(|| new_error!("invalid ring range"))?, + )?; + + let mem = HostMemOps::new(scratch_mem, ring.clone())?; + let mut bytes = vec![0; len]; + mem.read(ring.start, &mut bytes)?; + + Ok(bytes) +} + +fn write_ring(scratch_mem: &HostSharedMemory, ring: Range, bytes: &[u8]) -> Result<()> { + validate_ring_len("restored", bytes, usize::try_from(ring.end - ring.start)?)?; + let mem = HostMemOps::new(scratch_mem, ring.clone())?; + mem.write(ring.start, bytes) +} + +fn validate_ring_len(direction: &str, bytes: &[u8], expected: usize) -> Result<()> { + if bytes.len() != expected { + return Err(new_error!( + "{direction} snapshot ring length {} and expected length {expected}", + bytes.len() + )); + } + Ok(()) +} + fn checked_region(start: u64, len: usize, tag: &str) -> Result> { let end = start .checked_add(u64::try_from(len)?) @@ -356,6 +461,7 @@ mod tests { } struct PreparedVirtq { + scratch: HostSharedMemory, g2h_mem: HostMemOps, h2g_mem: HostMemOps, g2h_ring: Range, @@ -421,6 +527,7 @@ mod tests { let h2g_mem = HostMemOps::new(&scratch, h2g_ring.clone()).unwrap(); PreparedVirtq { + scratch, g2h_mem, h2g_mem, g2h_ring, @@ -511,6 +618,66 @@ mod tests { validate(&prepared_virtq()).unwrap(); } + #[test] + fn snapshots_and_restores_canonical_image() { + let prepared = prepared_virtq(); + let layout = memory_layout(); + let stale_pool = [0xa5; 16]; + let pool_mem = HostMemOps::new(&prepared.scratch, prepared.h2g_pool.clone()).unwrap(); + pool_mem + .write(prepared.h2g_pool.start, &stale_pool) + .unwrap(); + + let captured = snapshot(&layout, &prepared.scratch).unwrap(); + let restored = host_scratch(); + let allocator = layout.get_first_free_scratch_gpa(); + let allocator_offset = + restored.mem_size() - hyperlight_common::layout::SCRATCH_TOP_ALLOCATOR_OFFSET as usize; + restored.write::(allocator_offset, allocator).unwrap(); + + restore(&layout, &restored, &captured).unwrap(); + let restored_snapshot = snapshot(&layout, &restored).unwrap(); + let restored_pool = HostMemOps::new(&restored, prepared.h2g_pool.clone()).unwrap(); + let mut pool_bytes = [0; 16]; + restored_pool + .read(prepared.h2g_pool.start, &mut pool_bytes) + .unwrap(); + + assert_eq!(restored_snapshot, captured); + assert_eq!(restored.read::(allocator_offset).unwrap(), allocator); + assert_eq!(pool_bytes, [0; 16]); + } + + #[test] + fn rejects_corrupt_snapshot_ring_before_restore() { + let prepared = prepared_virtq(); + let layout = memory_layout(); + let mut snapshot = snapshot(&layout, &prepared.scratch).unwrap(); + snapshot.h2g_ring.fill(0); + let restored = host_scratch(); + + assert!(restore(&layout, &restored, &snapshot).is_err()); + assert_eq!(read_published_arena_gpa(&restored).unwrap(), 0); + } + + #[test] + fn restores_with_grown_page_tables() { + let prepared = prepared_virtq(); + let layout = memory_layout(); + let snapshot = snapshot(&layout, &prepared.scratch).unwrap(); + let mut grown_layout = layout; + grown_layout + .set_pt_size(layout.get_pt_size() + vmem::PAGE_SIZE) + .unwrap(); + let restored = host_scratch(); + + restore(&grown_layout, &restored, &snapshot).unwrap(); + assert_eq!( + read_published_arena_gpa(&restored).unwrap(), + grown_layout.get_transport_arena().base_addr() + ); + } + #[test] fn rejects_h2g_descriptors_outside_pool() { let prepared = prepared_virtq(); diff --git a/src/hyperlight_host/src/mem/virtq_mem.rs b/src/hyperlight_host/src/mem/virtq_mem.rs index f0e2d447a..944eee1b7 100644 --- a/src/hyperlight_host/src/mem/virtq_mem.rs +++ b/src/hyperlight_host/src/mem/virtq_mem.rs @@ -1,11 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright 2026 The Hyperlight Authors. -//! Host [`MemOps`] access to a bounded scratch region. +//! Host [`MemOps`] implementations for live scratch and captured ring images. //! -//! Every operation uses [`HostSharedMemory`]'s checked API and acquires its -//! lifecycle read lock. This preserves exclusive-memory coordination but makes -//! descriptor traversal pay for one lock acquisition per field access. +//! Live scratch operations use [`HostSharedMemory`]'s checked API and acquire +//! its lifecycle read lock. This preserves exclusive-memory coordination but +//! makes descriptor traversal pay for one lock acquisition per field access. use core::mem::size_of; use core::ops::Range; @@ -132,6 +132,70 @@ unsafe impl MemOps for HostMemOps { } } +/// Read-only [`MemOps`] view over a captured ring image. +/// +/// Snapshot preflight must validate captured bytes before writing them into +/// restored scratch. This view maps the image to its captured ring GVA, letting +/// the same directional validators handle snapshots and live [`HostMemOps`]. +pub(super) struct ImageMem<'a> { + base: u64, + bytes: &'a [u8], +} + +impl<'a> ImageMem<'a> { + pub(super) fn new(base: u64, bytes: &'a [u8]) -> Self { + Self { base, bytes } + } + + fn offset(&self, addr: u64, len: usize) -> Result { + let out_of_bounds = || new_error!("image memory access is out of bounds"); + // VirtqLayout uses absolute GVAs, while the captured image starts at index zero. + let offset = addr.checked_sub(self.base).ok_or_else(&out_of_bounds)?; + let offset = usize::try_from(offset).map_err(|_| out_of_bounds())?; + let end = offset.checked_add(len).ok_or_else(&out_of_bounds)?; + + (end <= self.bytes.len()) + .then_some(offset) + .ok_or_else(out_of_bounds) + } +} + +// SAFETY: ImageMem provides immutable access only within `bytes`. Write +// operations fail, and the backing slice outlives every returned shared slice. +unsafe impl MemOps for ImageMem<'_> { + type Error = HyperlightError; + + fn read(&self, addr: u64, dst: &mut [u8]) -> Result<()> { + let offset = self.offset(addr, dst.len())?; + dst.copy_from_slice(&self.bytes[offset..offset + dst.len()]); + Ok(()) + } + + fn load_acquire(&self, addr: u64) -> Result { + let mut bytes = [0; size_of::()]; + self.read(addr, &mut bytes)?; + Ok(u16::from_ne_bytes(bytes)) + } + + unsafe fn as_slice(&self, addr: u64, len: usize) -> Result<&[u8]> { + let offset = self.offset(addr, len)?; + Ok(&self.bytes[offset..offset + len]) + } + + fn write(&self, _addr: u64, _src: &[u8]) -> Result<()> { + Err(new_error!("image memory is read-only")) + } + + fn store_release(&self, _addr: u64, _val: u16) -> Result<()> { + Err(new_error!("image memory is read-only")) + } + + #[allow(clippy::mut_from_ref)] + unsafe fn as_mut_slice(&self, _addr: u64, _len: usize) -> Result<&mut [u8]> { + Err(new_error!("image memory is read-only")) + } +} + #[cfg(test)] mod tests { use hyperlight_common::virtq::MemOps; diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index 298d5d83e..bf57b176c 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -166,10 +166,9 @@ impl MultiUseSandbox { /// /// An optional [`SandboxConfiguration`](crate::sandbox::SandboxConfiguration) /// can be supplied to override runtime settings such as timeouts and - /// interrupt behavior. Memory layout fields - /// (`input_data_size`, `output_data_size`, `heap_size`, `scratch_size`) - /// are always taken from the snapshot. Any values supplied in - /// `config` for those fields are ignored. On x86_64 the `config` must + /// interrupt behavior. Memory layout fields and transport geometry are + /// always taken from the snapshot. Any values supplied in `config` for + /// those fields are ignored. On x86_64 the `config` must /// declare every guest MSR the snapshot was taken with (see /// [`SandboxConfiguration::guest_msrs`](crate::sandbox::SandboxConfiguration::guest_msrs)), /// or the load fails with an MSR mismatch. @@ -243,6 +242,12 @@ impl MultiUseSandbox { config.set_output_data_size(snapshot.layout().output_data_size()); config.set_heap_size(snapshot.layout().heap_size() as u64); config.set_scratch_size(snapshot.layout().get_scratch_size()); + config.set_g2h_queue_depth(snapshot.layout().get_g2h_queue_depth()); + config.set_h2g_queue_depth(snapshot.layout().get_h2g_queue_depth()); + config.set_g2h_buffer_size(snapshot.layout().get_g2h_buffer_size()); + config.set_h2g_buffer_size(snapshot.layout().get_h2g_buffer_size()); + config.set_g2h_pool_pages(snapshot.layout().get_g2h_pool_pages()); + config.set_h2g_pool_pages(snapshot.layout().get_h2g_pool_pages()); let load_info = snapshot.load_info(); let mgr = crate::mem::mgr::SandboxMemoryManager::from_snapshot(&snapshot)?; @@ -327,6 +332,8 @@ impl MultiUseSandbox { if attach_virtq { hshm.attach_virtq()?; + } else { + hshm.restore_virtq(snapshot.virtq())?; } let sbox = MultiUseSandbox::from_uninit(host_funcs, hshm, vm); @@ -1150,6 +1157,36 @@ fn warn_on_layout_override( caller.get_scratch_size() as u64, snapshot.get_scratch_size() as u64, ), + ( + "g2h_queue_depth", + caller.get_g2h_queue_depth() as u64, + snapshot.get_g2h_queue_depth() as u64, + ), + ( + "h2g_queue_depth", + caller.get_h2g_queue_depth() as u64, + snapshot.get_h2g_queue_depth() as u64, + ), + ( + "g2h_buffer_size", + caller.get_g2h_buffer_size() as u64, + snapshot.get_g2h_buffer_size() as u64, + ), + ( + "h2g_buffer_size", + caller.get_h2g_buffer_size() as u64, + snapshot.get_h2g_buffer_size() as u64, + ), + ( + "g2h_pool_pages", + caller.get_g2h_pool_pages() as u64, + snapshot.get_g2h_pool_pages() as u64, + ), + ( + "h2g_pool_pages", + caller.get_h2g_pool_pages() as u64, + snapshot.get_h2g_pool_pages() as u64, + ), ]; for (name, supplied, snap) in mismatches { if supplied != snap { @@ -1199,6 +1236,11 @@ mod tests { assert!(SandboxStatus::Unrecoverable.is_unrecoverable()); } + fn assert_virtq_attached(sbox: &MultiUseSandbox) { + assert!(sbox.mem_mgr.g2h_consumer.is_some()); + assert!(sbox.mem_mgr.h2g_consumer.is_some()); + } + #[test] fn poison() { let mut sbox = SandboxBuilder::from_file(simple_guest_as_pathbuf()) @@ -1736,6 +1778,7 @@ mod tests { let snapshot = sandbox.snapshot().unwrap(); sandbox2.restore(snapshot).unwrap(); + assert_virtq_attached(&sandbox2); assert_eq!(sandbox2.call::("GetStatic", ()).unwrap(), 42); } @@ -4646,7 +4689,9 @@ mod tests { let mut sbox = make_sandbox(); sbox.call::("AddToStatic", 11i32).unwrap(); let snapshot = sbox.snapshot().unwrap(); + assert!(snapshot.virtq().is_some()); let mut sbox2 = SandboxBuilder::from_snapshot(snapshot).build().unwrap(); + super::assert_virtq_attached(&sbox2); assert_eq!(sbox2.call::("GetStatic", ()).unwrap(), 11); let echoed: String = sbox2.call("Echo", "hi".to_string()).unwrap(); assert_eq!(echoed, "hi"); @@ -4658,6 +4703,7 @@ mod tests { let snap = Snapshot::from_env(GuestBinary::FilePath(path), SandboxConfiguration::default()) .unwrap(); + assert!(snap.virtq().is_none()); let mut sbox = SandboxBuilder::from_snapshot(Arc::new(snap)) .build() .unwrap(); @@ -4679,6 +4725,8 @@ mod tests { let mut b = SandboxBuilder::from_snapshot(snapshot.clone()) .build() .unwrap(); + super::assert_virtq_attached(&a); + super::assert_virtq_attached(&b); assert_eq!(a.call::("GetStatic", ()).unwrap(), 3); assert_eq!(b.call::("GetStatic", ()).unwrap(), 3); @@ -4688,6 +4736,8 @@ mod tests { a.restore(snapshot.clone()).unwrap(); b.restore(snapshot).unwrap(); + super::assert_virtq_attached(&a); + super::assert_virtq_attached(&b); assert_eq!(a.call::("GetStatic", ()).unwrap(), 3); assert_eq!(b.call::("GetStatic", ()).unwrap(), 3); } diff --git a/src/hyperlight_host/src/sandbox/snapshot/file/config.rs b/src/hyperlight_host/src/sandbox/snapshot/file/config.rs index 4faedf8c1..3a704f218 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/config.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file/config.rs @@ -153,7 +153,7 @@ impl CpuVendor { /// Top-level Hyperlight snapshot config JSON. Lives at /// `blobs/sha256/` with media type -/// `application/vnd.hyperlight.snapshot.config.v1+json`. +/// `application/vnd.hyperlight.snapshot.config.v2+json`. /// /// In OCI terms this is the "image config" blob that the manifest's /// `config` descriptor points to. It describes the accompanying @@ -214,6 +214,12 @@ pub(super) struct MemoryLayout { /// Memory region flag bits. `None` means default permissions. pub(super) init_data_permissions: Option, pub(super) scratch_size: usize, + pub(super) g2h_queue_depth: usize, + pub(super) h2g_queue_depth: usize, + pub(super) g2h_buffer_size: usize, + pub(super) h2g_buffer_size: usize, + pub(super) g2h_pool_pages: usize, + pub(super) h2g_pool_pages: usize, pub(super) snapshot_size: usize, pub(super) pt_size: Option, } @@ -471,6 +477,10 @@ impl OciSnapshotConfig { ("code_size", self.layout.code_size), ("init_data_size", self.layout.init_data_size), ("scratch_size", self.layout.scratch_size), + ("g2h_buffer_size", self.layout.g2h_buffer_size), + ("h2g_buffer_size", self.layout.h2g_buffer_size), + ("g2h_pool_pages", self.layout.g2h_pool_pages), + ("h2g_pool_pages", self.layout.h2g_pool_pages), ] { if value > max_region { return Err(crate::new_error!( @@ -482,6 +492,55 @@ impl OciSnapshotConfig { } } + let mut transport = crate::sandbox::SandboxConfiguration::default(); + transport.set_g2h_queue_depth(self.layout.g2h_queue_depth); + transport.set_h2g_queue_depth(self.layout.h2g_queue_depth); + transport.set_g2h_buffer_size(self.layout.g2h_buffer_size); + transport.set_h2g_buffer_size(self.layout.h2g_buffer_size); + transport.set_g2h_pool_pages(self.layout.g2h_pool_pages); + transport.set_h2g_pool_pages(self.layout.h2g_pool_pages); + + for (name, saved, normalized) in [ + ( + "g2h_queue_depth", + self.layout.g2h_queue_depth, + transport.get_g2h_queue_depth(), + ), + ( + "h2g_queue_depth", + self.layout.h2g_queue_depth, + transport.get_h2g_queue_depth(), + ), + ( + "g2h_buffer_size", + self.layout.g2h_buffer_size, + transport.get_g2h_buffer_size(), + ), + ( + "h2g_buffer_size", + self.layout.h2g_buffer_size, + transport.get_h2g_buffer_size(), + ), + ( + "g2h_pool_pages", + self.layout.g2h_pool_pages, + transport.get_g2h_pool_pages(), + ), + ( + "h2g_pool_pages", + self.layout.h2g_pool_pages, + transport.get_h2g_pool_pages(), + ), + ] { + if saved != normalized { + return Err(crate::new_error!( + "snapshot layout field {} ({}) is not a valid transport value", + name, + saved + )); + } + } + // The saved dispatch entrypoint must be in the executable code // region. Code occupies the page-rounded prefix of the snapshot. let code_lo = SandboxMemoryLayout::BASE_ADDRESS as u64; @@ -782,6 +841,12 @@ mod tests { init_data_size: 0, init_data_permissions: None, scratch_size: 0, + g2h_queue_depth: 64, + h2g_queue_depth: 32, + g2h_buffer_size: PAGE_SIZE, + h2g_buffer_size: PAGE_SIZE, + g2h_pool_pages: 8, + h2g_pool_pages: 4, snapshot_size: PAGE_SIZE, pt_size: None, }, @@ -848,7 +913,7 @@ mod schema_pin { const PINNED_CALL: &str = r#"{ "hyperlight_version": "x.y.z", "arch": "x86_64", - "abi_version": 1, + "abi_version": 3, "hypervisor": "mshv", "cpu_vendor": "intel", "stack_top_gva": 3735928559, @@ -1014,6 +1079,12 @@ mod schema_pin { "init_data_size": 5, "init_data_permissions": null, "scratch_size": 8, + "g2h_queue_depth": 64, + "h2g_queue_depth": 32, + "g2h_buffer_size": 4096, + "h2g_buffer_size": 4096, + "g2h_pool_pages": 8, + "h2g_pool_pages": 4, "snapshot_size": 9, "pt_size": null }, @@ -1034,7 +1105,7 @@ mod schema_pin { const PINNED_CALL: &str = r#"{ "hyperlight_version": "x.y.z", "arch": "aarch64", - "abi_version": 1, + "abi_version": 3, "hypervisor": "mshv", "cpu_vendor": "intel", "stack_top_gva": 3735928559, @@ -1056,6 +1127,12 @@ mod schema_pin { "init_data_size": 5, "init_data_permissions": null, "scratch_size": 8, + "g2h_queue_depth": 64, + "h2g_queue_depth": 32, + "g2h_buffer_size": 4096, + "h2g_buffer_size": 4096, + "g2h_pool_pages": 8, + "h2g_pool_pages": 4, "snapshot_size": 9, "pt_size": null }, @@ -1094,7 +1171,7 @@ mod schema_pin { assert_eq!( actual_value, pinned_value, "Snapshot config JSON schema changed. If the change can break \ - existing snapshots on disk, bump `MT_CONFIG_V1` in \ + existing snapshots on disk, bump `MT_CONFIG_CURRENT` in \ `super::media_types` and follow `docs/snapshot-versioning.md`. \ Either way, paste the actual output below into the matching \ `PINNED_*`.\n\nactual:\n{actual}" diff --git a/src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs b/src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs index 0f664edbc..8ec23eed2 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs @@ -6,7 +6,9 @@ // docs/snapshot-versioning.md for how to add a version. pub(in crate::sandbox::snapshot) const MT_CONFIG_V1: &str = "application/vnd.hyperlight.snapshot.config.v1+json"; -pub(in crate::sandbox::snapshot) const MT_CONFIG_CURRENT: &str = MT_CONFIG_V1; +pub(in crate::sandbox::snapshot) const MT_CONFIG_V2: &str = + "application/vnd.hyperlight.snapshot.config.v2+json"; +pub(in crate::sandbox::snapshot) const MT_CONFIG_CURRENT: &str = MT_CONFIG_V2; pub(in crate::sandbox::snapshot) const MT_SNAPSHOT_V1: &str = "application/vnd.hyperlight.snapshot.memory.v1"; pub(in crate::sandbox::snapshot) const MT_SNAPSHOT_CURRENT: &str = MT_SNAPSHOT_V1; @@ -14,7 +16,7 @@ pub(in crate::sandbox::snapshot) const MT_SNAPSHOT_CURRENT: &str = MT_SNAPSHOT_V /// ABI version for the snapshot memory blob. Bumped when the /// host-guest contract for the snapshot bytes changes. See /// docs/snapshot-versioning.md. -pub(in crate::sandbox::snapshot) const SNAPSHOT_ABI_VERSION: u32 = 2; +pub(in crate::sandbox::snapshot) const SNAPSHOT_ABI_VERSION: u32 = 3; /// OCI standard annotation key for a manifest's tag inside an image /// index. Set on the manifest descriptor in `index.json`, not on the diff --git a/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs b/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs index 0331628de..75769464a 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs @@ -26,7 +26,8 @@ use self::media_types::{ ANNOTATION_ARCH, ANNOTATION_CPU, ANNOTATION_HYPERVISOR, ANNOTATION_REF_NAME, }; pub(super) use self::media_types::{ - MT_CONFIG_CURRENT, MT_CONFIG_V1, MT_SNAPSHOT_CURRENT, MT_SNAPSHOT_V1, SNAPSHOT_ABI_VERSION, + MT_CONFIG_CURRENT, MT_CONFIG_V1, MT_CONFIG_V2, MT_SNAPSHOT_CURRENT, MT_SNAPSHOT_V1, + SNAPSHOT_ABI_VERSION, }; use self::reference::{OciDigest, OciReference, OciTag}; use super::{NextAction, Snapshot}; @@ -609,6 +610,12 @@ impl Snapshot { init_data_size: l.init_data_size(), init_data_permissions: l.init_data_permissions().map(|f| f.bits()), scratch_size: l.get_scratch_size(), + g2h_queue_depth: l.get_g2h_queue_depth(), + h2g_queue_depth: l.get_h2g_queue_depth(), + g2h_buffer_size: l.get_g2h_buffer_size(), + h2g_buffer_size: l.get_h2g_buffer_size(), + g2h_pool_pages: l.get_g2h_pool_pages(), + h2g_pool_pages: l.get_h2g_pool_pages(), snapshot_size: l.snapshot_size(), pt_size: l.pt_size(), }, @@ -731,16 +738,21 @@ impl Snapshot { // digest. let manifest = load_manifest(path, &blobs_dir, reference, verify_blobs)?; let cfg_desc = manifest.config(); - // Loader dispatch on config media type. A future v2 lands - // as a new arm that converts to the in-memory current shape. + // Loader dispatch on config media type. let cfg_media = cfg_desc.media_type().to_string(); match cfg_media.as_str() { - MT_CONFIG_V1 => {} + MT_CONFIG_V2 => {} + MT_CONFIG_V1 => { + return Err(crate::new_error!( + "snapshot config v1 is incompatible with snapshot ABI {}", + SNAPSHOT_ABI_VERSION + )); + } other => { return Err(crate::new_error!( "unexpected config media type {:?} (supported: {:?})", other, - MT_CONFIG_V1 + MT_CONFIG_V2 )); } } @@ -800,6 +812,12 @@ impl Snapshot { sbox_cfg.set_output_data_size(cfg.layout.output_data_size); sbox_cfg.set_heap_size(cfg.layout.heap_size as u64); sbox_cfg.set_scratch_size(cfg.layout.scratch_size); + sbox_cfg.set_g2h_queue_depth(cfg.layout.g2h_queue_depth); + sbox_cfg.set_h2g_queue_depth(cfg.layout.h2g_queue_depth); + sbox_cfg.set_g2h_buffer_size(cfg.layout.g2h_buffer_size); + sbox_cfg.set_h2g_buffer_size(cfg.layout.h2g_buffer_size); + sbox_cfg.set_g2h_pool_pages(cfg.layout.g2h_pool_pages); + sbox_cfg.set_h2g_pool_pages(cfg.layout.h2g_pool_pages); let init_data_perms = match cfg.layout.init_data_permissions { None => None, Some(bits) => Some(MemoryRegionFlags::from_bits(bits).ok_or_else(|| { @@ -893,6 +911,7 @@ impl Snapshot { original_entrypoint: cfg.original_entrypoint_addr, snapshot_generation, host_functions, + virtq: None, }) } } diff --git a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs b/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs index 8d7572955..ee28108b2 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs @@ -1830,6 +1830,20 @@ fn unknown_config_media_type_rejected() { assert_err_contains(err, "config media type"); } +#[test] +fn config_v1_rejected() { + let (_dir, path) = save_for_mutation(); + rewrite_manifest(&path, |m| { + m["config"]["mediaType"] = + Value::from("application/vnd.hyperlight.snapshot.config.v1+json"); + }); + let err = unwrap_err_snapshot(Snapshot::checked_load( + &path, + OciTag::new("latest").unwrap(), + )); + assert_err_contains(err, "incompatible with snapshot ABI 3"); +} + #[test] fn empty_layers_rejected() { let (_dir, path) = save_for_mutation(); @@ -2296,7 +2310,7 @@ fn manifest_uses_correct_config_and_layer_media_types() { serde_json::from_slice(&std::fs::read(manifest_path(&path)).unwrap()).unwrap(); assert_eq!( manifest["config"]["mediaType"].as_str().unwrap(), - "application/vnd.hyperlight.snapshot.config.v1+json" + "application/vnd.hyperlight.snapshot.config.v2+json" ); assert_eq!(manifest["layers"].as_array().unwrap().len(), 1); assert_eq!( @@ -2308,7 +2322,7 @@ fn manifest_uses_correct_config_and_layer_media_types() { // that falls back to `config.mediaType` sees the same value. assert_eq!( manifest["artifactType"].as_str().unwrap(), - "application/vnd.hyperlight.snapshot.config.v1+json" + "application/vnd.hyperlight.snapshot.config.v2+json" ); } @@ -2827,6 +2841,44 @@ fn persisted_non_default_layout_loads_and_runs() { ); } +#[test] +fn round_trip_preserves_transport_layout() { + use crate::sandbox::SandboxConfiguration; + + let mut cfg = SandboxConfiguration::default(); + cfg.set_scratch_size(512 * 1024); + cfg.set_heap_size(512 * 1024); + cfg.set_g2h_queue_depth(128); + cfg.set_h2g_queue_depth(16); + cfg.set_g2h_buffer_size(8192); + cfg.set_h2g_buffer_size(2048); + cfg.set_g2h_pool_pages(16); + cfg.set_h2g_pool_pages(6); + + let mut sbox = + UninitializedSandbox::new(GuestBinary::FilePath(simple_guest_as_pathbuf()), Some(cfg)) + .unwrap() + .evolve() + .unwrap(); + let snapshot = sbox.snapshot().unwrap(); + let expected = snapshot.layout().get_transport_arena(); + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("layout"); + snapshot + .save(&path, &OciTag::new("latest").unwrap()) + .unwrap(); + let loaded = Snapshot::checked_load(&path, OciTag::new("latest").unwrap()).unwrap(); + + assert_eq!(loaded.layout().get_g2h_queue_depth(), 128); + assert_eq!(loaded.layout().get_h2g_queue_depth(), 16); + assert_eq!(loaded.layout().get_g2h_buffer_size(), 8192); + assert_eq!(loaded.layout().get_h2g_buffer_size(), 2048); + assert_eq!(loaded.layout().get_g2h_pool_pages(), 16); + assert_eq!(loaded.layout().get_h2g_pool_pages(), 6); + assert_eq!(loaded.layout().get_transport_arena(), expected); +} + #[test] fn snapshot_config_records_entrypoint_and_sregs() { let snap = create_snapshot(); diff --git a/src/hyperlight_host/src/sandbox/snapshot/mod.rs b/src/hyperlight_host/src/sandbox/snapshot/mod.rs index 4a5484497..bef2bbf7c 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/mod.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/mod.rs @@ -26,6 +26,7 @@ use crate::mem::layout::SandboxMemoryLayout; use crate::mem::memory_region::{GuestMemoryRegion, MemoryRegion, MemoryRegionFlags}; use crate::mem::mgr::{GuestPageTableBuffer, SnapshotSharedMemory}; use crate::mem::shared_mem::{ReadonlySharedMemory, SharedMemory}; +use crate::mem::virtq::VirtqSnapshot; use crate::sandbox::SandboxConfiguration; use crate::sandbox::uninitialized::{GuestBinary, GuestEnvironment}; @@ -110,6 +111,12 @@ pub struct Snapshot { /// `HostFunctions` set that is missing required functions or /// has mismatched signatures. host_functions: HostFunctionDetails, + + /// Canonical in-memory virtqueue state omitted from ordinary snapshot pages. + /// + /// File snapshot persistence is deferred while stack communication remains + /// active. + virtq: Option, } impl core::convert::AsRef for Snapshot { fn as_ref(&self) -> &Self { @@ -393,6 +400,7 @@ impl Snapshot { host_functions: HostFunctionDetails { host_functions: None, }, + virtq: None, }) } @@ -419,6 +427,7 @@ impl Snapshot { original_entrypoint: u64, snapshot_generation: u64, host_functions: HostFunctionDetails, + virtq: Option, ) -> Result { let mut phys_seen = HashMap::::new(); let scratch_gva = scratch_base_gva(layout.get_scratch_size()); @@ -568,6 +577,10 @@ impl Snapshot { debug_assert!(guest_visible_size.is_multiple_of(page_size::get())); layout.set_snapshot_size(guest_visible_size); + if let Some(virtq) = &virtq { + virtq.preflight(&layout)?; + } + Ok(Self { layout, memory: ReadonlySharedMemory::from_bytes(&memory, guest_visible_size)?, @@ -580,6 +593,7 @@ impl Snapshot { original_entrypoint, snapshot_generation, host_functions, + virtq, }) } @@ -630,6 +644,10 @@ impl Snapshot { self.next_action } + pub(crate) fn virtq(&self) -> Option<&VirtqSnapshot> { + self.virtq.as_ref() + } + /// Guest virtual address of the guest binary's ELF entry point, /// preserved across the `Initialise` -> `Call` transition. Used /// to fill `AT_ENTRY` in guest core dumps. 0 if unknown. @@ -779,6 +797,7 @@ mod tests { 0, 1, HostFunctionDetails::default(), + None, ) .unwrap(); @@ -799,6 +818,7 @@ mod tests { 0, 2, HostFunctionDetails::default(), + None, ) .unwrap(); diff --git a/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs b/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs index c6dde9df1..7c7103884 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs @@ -15,8 +15,8 @@ use super::file::{ MT_CONFIG_CURRENT, MT_SNAPSHOT_CURRENT, OCI_LAYOUT_VERSION, SNAPSHOT_ABI_VERSION, }; -const EXPECTED_ABI_VERSION: u32 = 2; -const EXPECTED_MT_CONFIG: &str = "application/vnd.hyperlight.snapshot.config.v1+json"; +const EXPECTED_ABI_VERSION: u32 = 3; +const EXPECTED_MT_CONFIG: &str = "application/vnd.hyperlight.snapshot.config.v2+json"; const EXPECTED_MT_SNAPSHOT: &str = "application/vnd.hyperlight.snapshot.memory.v1"; const EXPECTED_OCI_LAYOUT_VERSION: &str = "1.0.0"; diff --git a/src/hyperlight_host/tests/snapshot_goldens/goldens_version.rs b/src/hyperlight_host/tests/snapshot_goldens/goldens_version.rs index 9752c0ffd..fc3c1f3b5 100644 --- a/src/hyperlight_host/tests/snapshot_goldens/goldens_version.rs +++ b/src/hyperlight_host/tests/snapshot_goldens/goldens_version.rs @@ -8,7 +8,7 @@ //! publish. See `docs/snapshot-versioning.md`. /// Goldens version, a `vMAJOR.MINOR` string. -pub(crate) const GOLDENS_VERSION: &str = "v2.0"; +pub(crate) const GOLDENS_VERSION: &str = "v3.0"; /// Old majors kept loadable through a compatibility path, verified /// alongside `GOLDENS_VERSION`. A backwards-compatible break (Option 2) From 46f92694f7e45540cc3d9cf23974b549c68c2e2c Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Mon, 10 Aug 2026 19:11:41 +0200 Subject: [PATCH 12/15] feat(virtq): optimize inflight bookkeeping memory Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/virtq/mod.rs | 2 + src/hyperlight_common/src/virtq/producer.rs | 251 +++++++++++++++++++- 2 files changed, 246 insertions(+), 7 deletions(-) diff --git a/src/hyperlight_common/src/virtq/mod.rs b/src/hyperlight_common/src/virtq/mod.rs index 906c3f14f..66479e4dd 100644 --- a/src/hyperlight_common/src/virtq/mod.rs +++ b/src/hyperlight_common/src/virtq/mod.rs @@ -195,6 +195,8 @@ pub enum VirtqError { Backpressure, #[error("Allocation exceeds pool capacity")] OutOfMemory, + #[error("Failed to allocate virtqueue bookkeeping")] + BookkeepingAllocation, #[error("Invalid chain received")] BadChain, #[error("Payload data too large: received {recv} bytes, limit {limit} bytes")] diff --git a/src/hyperlight_common/src/virtq/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index f7ed2576a..831f38c81 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -2,6 +2,7 @@ // Copyright 2026 The Hyperlight Authors. use alloc::collections::VecDeque; +use alloc::vec; use alloc::vec::Vec; use bytes::Bytes; @@ -80,6 +81,86 @@ pub(crate) struct Inflight { chain: BufferChain, } +/// Compact in-flight chains with constant-time descriptor-ID lookup. +/// +/// Descriptor IDs span the full ring, but live chains are normally bounded by +/// the much smaller buffer pool. `by_id` maps each descriptor ID to a packed +/// `live` index. Removal uses `swap_remove` and repairs the moved entry's map. +struct InflightTable { + by_id: Vec, + live: Vec, +} + +impl InflightTable { + const VACANT: u16 = u16::MAX; + + fn new(ring_len: usize) -> Self { + Self { + by_id: vec![Self::VACANT; ring_len], + live: Vec::new(), + } + } + + fn try_reserve_one(&mut self) -> Result<(), VirtqError> { + if self.live.len() > self.by_id.len() { + return Err(VirtqError::InvalidState); + } + + if self.live.len() == self.by_id.len() { + return Err(VirtqError::Backpressure); + } + + // Producers with one live chain should not pay for four large inline + // chain records. + let result = if self.live.capacity() == 0 { + self.live.try_reserve_exact(1) + } else { + self.live.try_reserve(1) + }; + + result.map_err(|_| VirtqError::BookkeepingAllocation) + } + + fn contains(&self, id: u16) -> bool { + self.by_id + .get(id as usize) + .is_some_and(|slot| *slot != Self::VACANT) + } + + fn insert(&mut self, inflight: Inflight) { + let id = inflight.token.id; + debug_assert!(!self.contains(id)); + debug_assert!(self.live.len() < Self::VACANT as usize); + + let slot = self.live.len() as u16; + self.live.push(inflight); + self.by_id[id as usize] = slot; + } + + fn remove(&mut self, id: u16) -> Option { + let slot = self.by_id.get_mut(id as usize)?; + if *slot == Self::VACANT { + return None; + } + + let index = usize::from(*slot); + *slot = Self::VACANT; + + let removed = self.live.swap_remove(index); + if let Some(moved) = self.live.get(index) { + self.by_id[moved.token.id as usize] = index as u16; + } + + Some(removed) + } + + fn pop(&mut self) -> Option { + let inflight = self.live.pop()?; + self.by_id[inflight.token.id as usize] = Self::VACANT; + Some(inflight) + } +} + /// A high-level virtqueue producer (driver side). /// /// The producer sends chains to the consumer (device), and receives used chains. @@ -117,7 +198,7 @@ pub struct VirtqProducer { notifier: N, pool: P, next_token: u32, - inflight: Vec>, + inflight: InflightTable, pending: VecDeque, } @@ -138,14 +219,15 @@ where pub fn new(layout: Layout, mem: M, notifier: N, pool: P) -> Self { let inner = RingProducer::new(layout, mem); let ring_len = inner.len(); + let inflight = InflightTable::new(ring_len); Self { inner, pool, notifier, + inflight, next_token: 0, - inflight: (0..ring_len).map(|_| None).collect(), - pending: VecDeque::with_capacity(ring_len), + pending: VecDeque::new(), } } @@ -210,17 +292,19 @@ where } fn publish(&mut self, send: SendChain) -> Result { + self.inflight.try_reserve_one()?; + let token_id = self.next_token; let id = self.inner.submit_available(send.chain())?; let token = Token { seq: token_id, id }; // A free descriptor id must never already be tracked as inflight. - if self.inflight[id as usize].is_some() { + if self.inflight.contains(id) { return Err(VirtqError::InvalidState); } let inf = send.into_inflight(token); - self.inflight[id as usize] = Some(inf); + self.inflight.insert(inf); self.next_token = self.next_token.wrapping_add(1); Ok(token) @@ -264,6 +348,44 @@ where self.inner.num_free() } + /// Number of submitted descriptors not yet polled as used. + #[inline] + pub fn num_inflight(&self) -> usize { + self.inner.num_inflight() + } + + /// Reset a stopped producer and release transport-owned allocations. + /// + /// The peer must not access the ring until its consumer is reset. Buffered + /// writable completions are guest-owned and make this operation fail. + /// Owner-backed payloads already returned to callers are not tracked as + /// in-flight and remain allocated. + pub fn reset(&mut self) -> Result<(), VirtqError> { + if !self.pending.is_empty() { + return Err(VirtqError::InvalidState); + } + + self.inner.reset()?; + self.next_token = 0; + + let mut maybe_err = None; + + // Drain all in-flight chains and retire their allocations. This is a best-effort + while let Some(inflight) = self.inflight.pop() { + let ret = self.retire_elems(inflight.chain.elems().iter().copied()); + if let Err(err) = ret + && maybe_err.is_none() + { + maybe_err = Some(err); + } + } + + match maybe_err { + Some(error) => Err(error), + None => Ok(()), + } + } + /// Configure event suppression for used buffer notifications. /// /// This controls when the device (consumer) signals us about completed buffers: @@ -340,6 +462,9 @@ where while let Some(chain) = self.poll_ring()? { if matches!(chain, UsedChain::Data(_, _)) { debug_assert!(self.pending.len() < self.inner.len()); + self.pending + .try_reserve(1) + .map_err(|_| VirtqError::BookkeepingAllocation)?; self.pending.push_back(chain); } count += 1; @@ -357,8 +482,7 @@ where let inf = self .inflight - .get_mut(used.id as usize) - .and_then(Option::take) + .remove(used.id) .ok_or(VirtqError::InvalidState)?; let written = used.len as usize; @@ -989,6 +1113,119 @@ mod tests { consumer.poll(1024).unwrap().unwrap() } + fn inflight(seq: u32, id: u16) -> Inflight { + let chain = BufferChainBuilder::new() + .readable(0x1000 + u64::from(id) * 0x10, 8) + .build() + .unwrap(); + Inflight { + token: Token { seq, id }, + chain, + } + } + + #[test] + fn inflight_table_repairs_moved_entry_after_removal() { + let mut table = InflightTable::new(16); + for (seq, id) in [(0, 3), (1, 7), (2, 5)] { + table.try_reserve_one().unwrap(); + table.insert(inflight(seq, id)); + } + + assert_eq!(table.remove(7).unwrap().token.seq, 1); + assert!(!table.contains(7)); + assert_eq!(table.remove(5).unwrap().token.seq, 2); + assert_eq!(table.remove(3).unwrap().token.seq, 0); + assert!(table.live.is_empty()); + assert!(table.remove(7).is_none()); + } + + #[test] + fn producer_bookkeeping_starts_compact_and_lazy() { + let ring = make_ring(64); + let (producer, _consumer, _notifier) = make_test_producer(&ring); + + assert_eq!(producer.inflight.by_id.len(), ring.len()); + assert!(producer.inflight.live.is_empty()); + assert_eq!(producer.inflight.live.capacity(), 0); + assert_eq!(producer.pending.capacity(), 0); + } + + #[test] + fn full_ring_still_reports_backpressure() { + let ring = make_ring(4); + let (mut producer, _consumer, _notifier) = make_test_producer(&ring); + + for _ in 0..ring.len() { + let chain = producer.chain().readable(1).build().unwrap(); + producer.submit(chain).unwrap(); + } + + let chain = producer.chain().readable(1).build().unwrap(); + assert!(matches!( + producer.submit(chain), + Err(VirtqError::Backpressure) + )); + } + + #[test] + fn reset_reclaims_inflight_slots_and_reuses_ring() { + let ring = make_ring(8); + let mem = ring.mem(); + let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + let pool = SlotPool::new(SlotLayout::new(pool_base, 64, ring.len())).unwrap(); + let notifier = TestNotifier::new(); + let mut producer = VirtqProducer::new(ring.layout(), mem, notifier, pool.clone()); + + for _ in 0..ring.len() { + let chain = producer.chain().writable(64).build().unwrap(); + producer.submit(chain).unwrap(); + } + + assert_eq!(pool.num_free(), 0); + + producer.reset().unwrap(); + + assert_eq!(producer.num_inflight(), 0); + assert_eq!(producer.num_free(), ring.len()); + assert_eq!(pool.num_free(), ring.len()); + assert!(producer.inflight.live.is_empty()); + assert!( + producer + .inflight + .by_id + .iter() + .all(|slot| *slot == InflightTable::VACANT) + ); + + for _ in 0..ring.len() { + let chain = producer.chain().writable(64).build().unwrap(); + producer.submit(chain).unwrap(); + } + } + + #[test] + fn reset_rejects_buffered_writable_completion() { + let ring = make_ring(8); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + let chain = producer.chain().writable(64).build().unwrap(); + producer.submit(chain).unwrap(); + + let (recv, reply) = poll_received(&mut consumer); + let ReplyChain::Writable(mut reply) = reply else { + panic!("expected writable reply"); + }; + + reply.write_all(b"retained").unwrap(); + consumer.complete(recv, reply).unwrap(); + producer.reclaim().unwrap(); + + assert!(matches!(producer.reset(), Err(VirtqError::InvalidState))); + + drop(producer.poll().unwrap().unwrap()); + producer.reset().unwrap(); + } + #[derive(Clone)] struct NoDirectSliceMem(TestMem); From 854dc5793530320ecf9fa290fec0aed5aec60df2 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Thu, 13 Aug 2026 16:29:56 +0200 Subject: [PATCH 13/15] feat(virtq): use grouped allocation for virtq chains Allocate logical regions atomically Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/benches/buffer_pool.rs | 14 +- src/hyperlight_common/src/virtq/buffer.rs | 2 +- src/hyperlight_common/src/virtq/mod.rs | 43 +-- src/hyperlight_common/src/virtq/pool.rs | 75 ++---- src/hyperlight_common/src/virtq/pool/fuzz.rs | 33 ++- src/hyperlight_common/src/virtq/pool/run.rs | 76 +++++- src/hyperlight_common/src/virtq/pool/slot.rs | 195 ++++++++++++-- src/hyperlight_common/src/virtq/pool/tests.rs | 202 +++++++++++--- src/hyperlight_common/src/virtq/producer.rs | 253 +++++++++--------- src/hyperlight_common/src/virtq/ring.rs | 20 ++ 10 files changed, 628 insertions(+), 285 deletions(-) diff --git a/src/hyperlight_common/benches/buffer_pool.rs b/src/hyperlight_common/benches/buffer_pool.rs index 4357b8992..50f76e014 100644 --- a/src/hyperlight_common/benches/buffer_pool.rs +++ b/src/hyperlight_common/benches/buffer_pool.rs @@ -138,9 +138,9 @@ fn bench_segmented_payload(c: &mut Criterion) { |b, &payload_size| { let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { - let sgs = pool.alloc_sg(black_box(payload_size)).unwrap(); - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); + let regions = pool.alloc_regions([black_box(payload_size)]).unwrap(); + for alloc in regions.into_iter().flatten() { + pool.dealloc(alloc.addr).unwrap(); } }); }, @@ -180,13 +180,13 @@ fn bench_slot_pool(c: &mut Criterion) { }); }); - group.bench_function("alloc_sg_64k", |b| { + group.bench_function("alloc_regions_64k", |b| { let layout = SlotLayout::new(0x80000, 4096, 1024); let pool = SlotPool::new(layout).unwrap(); b.iter(|| { - let sgs = pool.alloc_sg(black_box(64 * 1024)).unwrap(); - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); + let regions = pool.alloc_regions([black_box(64 * 1024)]).unwrap(); + for alloc in regions.into_iter().flatten() { + pool.dealloc(alloc.addr).unwrap(); } }); }); diff --git a/src/hyperlight_common/src/virtq/buffer.rs b/src/hyperlight_common/src/virtq/buffer.rs index 9f3265e9f..4ca31a3a0 100644 --- a/src/hyperlight_common/src/virtq/buffer.rs +++ b/src/hyperlight_common/src/virtq/buffer.rs @@ -207,7 +207,7 @@ pub struct BufferOwner { impl AsRef<[u8]> for BufferOwner { fn as_ref(&self) -> &[u8] { let alloc = self.alloc.allocation(); - let len = self.written.min(alloc.len); + let len = self.written.min(alloc.len as usize); // Safety: BufferOwner keeps both the pool allocation and the M alive, // so the memory region is valid. match unsafe { self.mem.as_slice(alloc.addr, len) } { diff --git a/src/hyperlight_common/src/virtq/mod.rs b/src/hyperlight_common/src/virtq/mod.rs index 66479e4dd..1ee43d538 100644 --- a/src/hyperlight_common/src/virtq/mod.rs +++ b/src/hyperlight_common/src/virtq/mod.rs @@ -408,7 +408,7 @@ impl From for Allocation { fn from(value: BufferElement) -> Self { Allocation { addr: value.addr, - len: value.len as usize, + len: value.len, } } } @@ -509,7 +509,6 @@ pub(crate) mod test_utils { base: u64, next: Arc, size: usize, - max_alloc_len: usize, allocations: Arc>>, } @@ -519,33 +518,23 @@ pub(crate) mod test_utils { base, next: Arc::new(AtomicU64::new(base)), size, - max_alloc_len: usize::MAX, - allocations: Arc::new(Mutex::new(BTreeMap::new())), - } - } - - pub(crate) fn new_with_max_alloc_len(base: u64, size: usize, max_alloc_len: usize) -> Self { - Self { - base, - next: Arc::new(AtomicU64::new(base)), - size, - max_alloc_len, allocations: Arc::new(Mutex::new(BTreeMap::new())), } } } impl BufferProvider for TestPool { - fn max_alloc_len(&self) -> usize { - self.max_alloc_len + fn preferred_segment_len(&self) -> usize { + u32::MAX as usize } fn alloc(&self, len: usize) -> Result { if len == 0 { return Err(AllocError::InvalidArg); } + let len = u32::try_from(len).map_err(|_| AllocError::OutOfMemory)?; - let addr = self.next.fetch_add(len as u64, Ordering::Relaxed); + let addr = self.next.fetch_add(u64::from(len), Ordering::Relaxed); let end = addr + len as u64; if end > self.base + self.size as u64 { return Err(AllocError::NoSpace); @@ -553,10 +542,30 @@ pub(crate) mod test_utils { self.allocations .lock() .expect("poisoned mutex") - .insert(addr, len); + .insert(addr, len as usize); + Ok(Allocation { addr, len }) } + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + let mut regions = Regions::new(); + for len in lengths { + match self.alloc(len) { + Ok(alloc) => regions.push(Allocations::from_iter([alloc])), + Err(error) => return Err(error), + } + } + + if regions.is_empty() { + return Err(AllocError::InvalidArg); + } + + Ok(regions) + } + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { self.allocations .lock() diff --git a/src/hyperlight_common/src/virtq/pool.rs b/src/hyperlight_common/src/virtq/pool.rs index a75fe684f..ea23dde5e 100644 --- a/src/hyperlight_common/src/virtq/pool.rs +++ b/src/hyperlight_common/src/virtq/pool.rs @@ -51,16 +51,20 @@ pub enum AllocError { pub struct Allocation { /// Starting address of the allocation. pub addr: u64, - /// Capacity in bytes, rounded up according to the provider's policy. - pub len: usize, + /// Nonzero descriptor-safe capacity in bytes. + pub len: u32, } +/// Ordered nonoverlapping allocations that back one logical region. +pub type Allocations = SmallVec<[Allocation; 4]>; + +/// Ordered allocation groups, one for each requested logical region. +pub type Regions = SmallVec<[Allocations; 4]>; + /// Allocates and reclaims virtqueue payload buffers. pub trait BufferProvider { - /// Preferred maximum size of one allocation segment. - fn max_alloc_len(&self) -> usize { - usize::MAX - } + /// Preferred nonzero descriptor-safe size of one bulk allocation segment. + fn preferred_segment_len(&self) -> usize; /// Allocate one buffer that can hold at least `len` bytes. fn alloc(&self, len: usize) -> Result; @@ -68,44 +72,15 @@ pub trait BufferProvider { /// Free a previously allocated segment by start address. fn dealloc(&self, addr: u64) -> Result<(), AllocError>; - /// Allocate scatter/gather segments for a logical payload of `total_len` bytes. - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - if total_len == 0 { - return Err(AllocError::InvalidArg); - } - - let seg_cap = self.max_alloc_len(); - if seg_cap == 0 { - return Err(AllocError::InvalidArg); - } - - let mut rem = total_len; - let mut sgs = SmallVec::<[Allocation; 4]>::new(); - - while rem > 0 { - let len = rem.min(seg_cap); - match self.alloc(len) { - Ok(alloc) => { - sgs.push(alloc); - rem -= len; - } - Err(err) => { - for sg in sgs { - let result = self.dealloc(sg.addr); - debug_assert!(result.is_ok(), "dealloc failed: {result:?}"); - } - return Err(err); - } - } - } - - Ok(sgs) - } + /// Allocate independent logical regions in input order. + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator; } impl BufferProvider for Rc { - fn max_alloc_len(&self) -> usize { - (**self).max_alloc_len() + fn preferred_segment_len(&self) -> usize { + (**self).preferred_segment_len() } fn alloc(&self, len: usize) -> Result { @@ -116,14 +91,17 @@ impl BufferProvider for Rc { (**self).dealloc(addr) } - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - (**self).alloc_sg(total_len) + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + (**self).alloc_regions(lengths) } } impl BufferProvider for Arc { - fn max_alloc_len(&self) -> usize { - (**self).max_alloc_len() + fn preferred_segment_len(&self) -> usize { + (**self).preferred_segment_len() } fn alloc(&self, len: usize) -> Result { @@ -134,8 +112,11 @@ impl BufferProvider for Arc { (**self).dealloc(addr) } - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - (**self).alloc_sg(total_len) + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + (**self).alloc_regions(lengths) } } diff --git a/src/hyperlight_common/src/virtq/pool/fuzz.rs b/src/hyperlight_common/src/virtq/pool/fuzz.rs index 22cd1c934..6eed13b69 100644 --- a/src/hyperlight_common/src/virtq/pool/fuzz.rs +++ b/src/hyperlight_common/src/virtq/pool/fuzz.rs @@ -19,7 +19,7 @@ const UPPER_SLOT_SIZE: usize = 4096; #[derive(Clone, Debug)] enum Op { Alloc(usize), - AllocSg(usize), + AllocRegions(usize), Dealloc(usize), } @@ -27,7 +27,7 @@ impl Arbitrary for Op { fn arbitrary(g: &mut Gen) -> Self { match u8::arbitrary(g) % 3 { 0 => Op::Alloc(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), - 1 => Op::AllocSg(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), + 1 => Op::AllocRegions(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), 2 => Op::Dealloc(usize::arbitrary(g)), _ => unreachable!(), } @@ -61,7 +61,7 @@ where match op { Op::Alloc(size) => match pool.alloc(*size) { Ok(alloc) => { - if alloc.len < *size + if (alloc.len as usize) < *size || allocations .iter() .any(|existing| existing.addr == alloc.addr) @@ -73,18 +73,21 @@ where Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} Err(_) => return false, }, - Op::AllocSg(size) => match pool.alloc_sg(*size) { - Ok(sgs) => { + Op::AllocRegions(size) => match pool.alloc_regions([*size]) { + Ok(regions) => { let mut total = 0usize; - for sg in sgs { - let Some(next_total) = total.checked_add(sg.len) else { + for allocation in regions.into_iter().flatten() { + let Some(next_total) = total.checked_add(allocation.len as usize) else { return false; }; - if allocations.iter().any(|existing| existing.addr == sg.addr) { + if allocations + .iter() + .any(|existing| existing.addr == allocation.addr) + { return false; } total = next_total; - allocations.push(sg); + allocations.push(allocation); } if total < *size { return false; @@ -138,12 +141,13 @@ fn check_run_tier_invariants( for alloc in allocations.iter().filter(|alloc| tier.contains(alloc.addr)) { let offset = usize::try_from(alloc.addr - tier.base_addr) .map_err(|_| "allocation offset overflows usize")?; - if alloc.len == 0 || !offset.is_multiple_of(N) || !alloc.len.is_multiple_of(N) { + let len = alloc.len as usize; + if len == 0 || !offset.is_multiple_of(N) || !len.is_multiple_of(N) { return Err("allocation is not tier-aligned"); } let start = offset / N; - let slots = alloc.len / N; + let slots = len / N; let end = start .checked_add(slots) .ok_or("allocation slot range overflow")?; @@ -176,13 +180,14 @@ fn check_run_tier_invariants( } let offset = usize::try_from(free_run.addr - tier.base_addr) .map_err(|_| "cached free-run offset overflows usize")?; - if free_run.len == 0 || !offset.is_multiple_of(N) || !free_run.len.is_multiple_of(N) { + let len = free_run.len as usize; + if len == 0 || !offset.is_multiple_of(N) || !len.is_multiple_of(N) { return Err("cached free run is not tier-aligned"); } let start = offset / N; let end = start - .checked_add(free_run.len / N) + .checked_add(len / N) .ok_or("cached free-run range overflow")?; if end > tier.used_slots.len() || (start..end).any(|slot| tier.used_slots.contains(slot)) { return Err("cached free run overlaps live allocations"); @@ -328,7 +333,7 @@ fn check_slot_pool_invariants( match expected_live.get(&addr) { Some(expected_capacity) => { - if *expected_capacity != capacity + if *expected_capacity as usize != capacity || pool.allocation_len(addr).ok() != Some(capacity) { return Err("live slot capacity is inconsistent"); diff --git a/src/hyperlight_common/src/virtq/pool/run.rs b/src/hyperlight_common/src/virtq/pool/run.rs index 9ab69eadb..5e1e53d81 100644 --- a/src/hyperlight_common/src/virtq/pool/run.rs +++ b/src/hyperlight_common/src/virtq/pool/run.rs @@ -17,9 +17,8 @@ use alloc::rc::Rc; use core::cell::RefCell; use fixedbitset::FixedBitSet; -use smallvec::SmallVec; -use super::{AllocError, Allocation, BufferProvider, SendWrap, align_up}; +use super::{AllocError, Allocation, Allocations, BufferProvider, Regions, SendWrap, align_up}; #[derive(Debug, Clone)] pub(super) struct Tier { @@ -113,7 +112,7 @@ impl Tier { debug_assert!(slots_num > 0); if let Some(alloc) = self.last_free_run - && alloc.len >= slots_num * N + && alloc.len as usize >= slots_num * N { let pos = self.slot_of(alloc.addr); let _ = self.last_free_run.take(); @@ -138,6 +137,11 @@ impl Tier { return Err(AllocError::OutOfMemory); } + let alloc_len = need_slots + .checked_mul(N) + .and_then(|len| u32::try_from(len).ok()) + .ok_or(AllocError::OutOfMemory)?; + let idx = self.find_slots(need_slots).ok_or(AllocError::NoSpace)?; self.used_slots.insert_range(idx..idx + need_slots); self.run_starts.insert(idx); @@ -145,7 +149,7 @@ impl Tier { let alloc = Allocation { addr, - len: need_slots * N, + len: alloc_len, }; self.maybe_invalidate_last_run(alloc); @@ -161,7 +165,7 @@ impl Tier { } fn dealloc_run(&mut self, start: usize, run_slots: usize, addr: u64) -> Result<(), AllocError> { - let len = run_slots * N; + let len = u32::try_from(run_slots * N).map_err(|_| AllocError::Overflow)?; self.used_slots.remove_range(start..start + run_slots); self.run_starts.set(start, false); self.last_free_run = Some(Allocation { addr, len }); @@ -311,19 +315,59 @@ impl Inner { self.upper.allocation_len(addr) } } + + /// Allocate independent logical regions in order. + fn alloc_regions(&mut self, lengths: I) -> Result + where + I: IntoIterator, + { + let mut regions = Regions::new(); + + for len in lengths { + match self.alloc(len) { + Ok(allocation) => { + let mut allocations = Allocations::new(); + allocations.push(allocation); + regions.push(allocations); + } + Err(error) => { + self.rollback_regions(®ions); + return Err(error); + } + } + } + + if regions.is_empty() { + return Err(AllocError::InvalidArg); + } + + Ok(regions) + } + + fn rollback_regions(&mut self, regions: &Regions) { + for allocations in regions { + for allocation in allocations { + let result = self.dealloc_addr(allocation.addr); + debug_assert!(result.is_ok(), "dealloc failed: {result:?}"); + } + } + } } impl BufferProvider for RunPool { - fn max_alloc_len(&self) -> usize { - U + fn preferred_segment_len(&self) -> usize { + U.min(u32::MAX as usize) } fn alloc(&self, len: usize) -> Result { self.inner.borrow_mut().alloc(len) } - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - Ok(smallvec::smallvec![self.alloc(total_len)?]) + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + self.inner.borrow_mut().alloc_regions(lengths) } fn dealloc(&self, addr: u64) -> Result<(), AllocError> { @@ -345,16 +389,22 @@ impl RunPool { #[cfg(all(test, loom))] impl BufferProvider for RunPoolSync { - fn max_alloc_len(&self) -> usize { - U + fn preferred_segment_len(&self) -> usize { + U.min(u32::MAX as usize) } fn alloc(&self, len: usize) -> Result { self.inner.lock().expect("poisoned mutex").alloc(len) } - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - Ok(smallvec::smallvec![self.alloc(total_len)?]) + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + self.inner + .lock() + .expect("poisoned mutex") + .alloc_regions(lengths) } fn dealloc(&self, addr: u64) -> Result<(), AllocError> { diff --git a/src/hyperlight_common/src/virtq/pool/slot.rs b/src/hyperlight_common/src/virtq/pool/slot.rs index ea2801970..b103abcea 100644 --- a/src/hyperlight_common/src/virtq/pool/slot.rs +++ b/src/hyperlight_common/src/virtq/pool/slot.rs @@ -20,14 +20,14 @@ use core::cell::RefCell; use fixedbitset::FixedBitSet; use smallvec::SmallVec; -use super::{AllocError, Allocation, BufferProvider, SendWrap}; +use super::{AllocError, Allocation, Allocations, BufferProvider, Regions, SendWrap}; /// Exact memory layout for one [`SlotPool`] tier. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct SlotLayout { /// Start of the first slot. pub base_addr: u64, - /// Capacity of each slot. + /// Capacity of each slot. Must fit in [`Allocation::len`]. pub slot_size: usize, /// Number of slots. pub slot_count: usize, @@ -68,7 +68,7 @@ struct Tier { /// Start of this tier's backing memory. base_addr: u64, /// Capacity of this slot. - slot_size: usize, + slot_size: u32, /// Number of slots in this tier. count: usize, /// Free slot addresses, popped/pushed LIFO. @@ -86,6 +86,7 @@ impl Tier { if layout.slot_size == 0 { return Err(AllocError::InvalidArg); } + let slot_size = u32::try_from(layout.slot_size).map_err(|_| AllocError::InvalidArg)?; if layout.slot_count == 0 { return Err(AllocError::EmptyRegion); @@ -100,7 +101,7 @@ impl Tier { Ok(Self { base_addr: layout.base_addr, - slot_size: layout.slot_size, + slot_size, count: layout.slot_count, free, allocated: FixedBitSet::with_capacity(layout.slot_count), @@ -108,7 +109,7 @@ impl Tier { } fn end(&self) -> u64 { - self.base_addr + (self.count * self.slot_size) as u64 + self.base_addr + self.count as u64 * u64::from(self.slot_size) } fn contains(&self, addr: u64) -> bool { @@ -122,11 +123,11 @@ impl Tier { } let off = addr - self.base_addr; - if !off.is_multiple_of(self.slot_size as u64) { + if !off.is_multiple_of(u64::from(self.slot_size)) { return Err(AllocError::InvalidFree(addr, 0)); } - Ok((off / self.slot_size as u64) as usize) + Ok((off / u64::from(self.slot_size)) as usize) } /// Validate that `addr` is a live (currently allocated) slot start. @@ -142,7 +143,7 @@ impl Tier { if len == 0 { return Err(AllocError::InvalidArg); } - if len > self.slot_size { + if len > self.slot_size as usize { return Err(AllocError::OutOfMemory); } @@ -150,7 +151,7 @@ impl Tier { // Safety of the index: `addr` came from `free`, which only ever holds // valid slot starts. self.allocated - .insert(((addr - self.base_addr) / self.slot_size as u64) as usize); + .insert(((addr - self.base_addr) / u64::from(self.slot_size)) as usize); Ok(Allocation { addr, @@ -167,11 +168,11 @@ impl Tier { fn allocation_len(&self, addr: u64) -> Result { self.live_slot_of(addr)?; - Ok(self.slot_size) + Ok(self.slot_size as usize) } fn slot_addr(&self, index: usize) -> Option { - (index < self.count).then(|| self.base_addr + (index * self.slot_size) as u64) + (index < self.count).then(|| self.base_addr + (index * self.slot_size as usize) as u64) } fn num_free(&self) -> usize { @@ -182,12 +183,12 @@ impl Tier { addrs.extend( self.allocated .ones() - .map(|slot| self.base_addr + (slot * self.slot_size) as u64), + .map(|slot| self.base_addr + (slot * self.slot_size as usize) as u64), ); } fn layout(&self) -> SlotLayout { - SlotLayout::new(self.base_addr, self.slot_size, self.count) + SlotLayout::new(self.base_addr, self.slot_size as usize, self.count) } } @@ -218,7 +219,9 @@ impl Inner { .count .checked_add(upper.count) .ok_or(AllocError::Overflow)?; - let layout = SlotLayout::new(lower.base_addr, lower.slot_size, count); + + let layout = SlotLayout::new(lower.base_addr, lower.slot_size as usize, count); + return Ok(Self { lower: None, upper: Tier::from_layout(layout)?, @@ -232,12 +235,12 @@ impl Inner { } fn max_alloc_len(&self) -> usize { - self.upper.slot_size + self.upper.slot_size as usize } fn alloc(&mut self, len: usize) -> Result { if let Some(lower) = &mut self.lower - && len <= lower.slot_size + && len <= lower.slot_size as usize { match lower.alloc(len) { Ok(alloc) => return Ok(alloc), @@ -249,6 +252,101 @@ impl Inner { self.upper.alloc(len) } + fn alloc_counts( + &self, + lengths: impl IntoIterator, + ) -> Result<(usize, usize), AllocError> { + let lower_size = self.lower.as_ref().map(|lower| lower.slot_size as usize); + let upper_size = self.upper.slot_size as usize; + let free_lower = self.lower.as_ref().map_or(0, Tier::num_free); + let free_upper = self.upper.num_free(); + + let mut alloc_count = 0usize; + let mut lower_count = 0usize; + + for len in lengths { + if len == 0 { + return Err(AllocError::InvalidArg); + } + + alloc_count = alloc_count + .checked_add(len.div_ceil(upper_size)) + .ok_or(AllocError::Overflow)?; + + let tail_len = len % upper_size; + if tail_len != 0 + && lower_size.is_some_and(|size| tail_len <= size) + && lower_count < free_lower + { + lower_count += 1; + } + + if alloc_count - lower_count > free_upper { + return Err(AllocError::NoSpace); + } + } + + Ok((alloc_count, alloc_count - lower_count)) + } + + fn max_alloc( + &self, + lengths: impl IntoIterator, + alloc_limit: usize, + ) -> Result { + let (used, upper_used) = self.alloc_counts(lengths)?; + let remaining = alloc_limit.checked_sub(used).ok_or(AllocError::NoSpace)?; + if remaining == 0 { + return Err(AllocError::NoSpace); + } + + let free_upper = self.upper.num_free() - upper_used; + let upper_count = remaining.min(free_upper); + + let len = upper_count + .checked_mul(self.upper.slot_size as usize) + .ok_or(AllocError::Overflow)?; + + if len == 0 { + return Err(AllocError::NoSpace); + } + + Ok(len) + } + + fn alloc_regions(&mut self, lengths: I) -> Result + where + I: IntoIterator, + { + let lengths = SmallVec::<[usize; 4]>::from_iter(lengths); + self.alloc_counts(lengths.iter().copied())?; + + let mut regions = Regions::with_capacity(lengths.len()); + let slot_size = self.max_alloc_len(); + + for total_len in lengths { + let mut allocs = Allocations::new(); + let mut remaining = total_len; + + while remaining > 0 { + let len = remaining.min(slot_size); + + // `alloc_counts` preflights the complete request. + #[allow(clippy::expect_used)] + allocs.push(self.alloc(len).expect("plan validated upstream")); + + remaining -= len; + } + regions.push(allocs); + } + + if regions.is_empty() { + return Err(AllocError::InvalidArg); + } + + Ok(regions) + } + fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { if let Some(lower) = &mut self.lower && lower.contains(addr) @@ -278,7 +376,7 @@ impl Inner { } fn live_addrs(&self) -> Vec { - let mut addrs = Vec::with_capacity(self.count() - self.num_free()); + let mut addrs = Vec::with_capacity(self.num_live()); if let Some(lower) = &self.lower { lower.append_live_addrs(&mut addrs); } @@ -300,6 +398,10 @@ impl Inner { self.lower.as_ref().map_or(0, Tier::num_free) + self.upper.num_free() } + fn num_live(&self) -> usize { + self.count() - self.num_free() + } + fn layouts(&self) -> (Option, SlotLayout) { (self.lower.as_ref().map(Tier::layout), self.upper.layout()) } @@ -309,8 +411,8 @@ impl Inner { /// /// Allocation and deallocation are O(1) per slot. Eligible allocations first /// try the optional lower tier and fall back to the required upper tier when -/// the lower tier is full. [`alloc_sg`](BufferProvider::alloc_sg) splits logical -/// payloads into bounded descriptor segments. +/// the lower tier is full. [`alloc_regions`](BufferProvider::alloc_regions) +/// splits logical payloads into bounded descriptor segments. #[derive(Clone)] pub struct SlotPool { inner: SendWrap>>, @@ -360,6 +462,21 @@ impl SlotPool { self.inner.borrow().num_free() } + /// Total number of currently allocated slots across all tiers. + pub fn num_live(&self) -> usize { + self.inner.borrow().num_live() + } + + /// Total number of free slots in the lower tier. + pub fn num_free_lower(&self) -> usize { + self.inner.borrow().lower.as_ref().map_or(0, Tier::num_free) + } + + /// Total number of free slots in the upper tier. + pub fn num_free_upper(&self) -> usize { + self.inner.borrow().upper.num_free() + } + /// Free a previously allocated slot by address. pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { self.inner.borrow_mut().dealloc_addr(addr) @@ -380,14 +497,45 @@ impl SlotPool { self.inner.borrow().max_alloc_len() } + /// Slot size in bytes for the lower tier, if present. + pub fn lower_slot_size(&self) -> Option { + self.inner + .borrow() + .lower + .as_ref() + .map(|lower| lower.slot_size as usize) + } + + /// Maximum slot size in bytes for the upper tier. + pub fn upper_slot_size(&self) -> usize { + self.inner.borrow().upper.slot_size as usize + } + /// Total number of slots across all tiers. pub fn count(&self) -> usize { self.inner.borrow().count() } + + /// Maximum upper-tier region length after reserving `lengths`. + /// + /// The returned region uses only upper-tier slots. It and the reserved + /// regions use at most `alloc_limit` allocations in total. This query does + /// not mutate the pool. + /// + /// # Errors + /// + /// Returns an error when a reserved region is invalid or unavailable, no + /// additional allocation fits, or capacity arithmetic overflows. + pub fn max_alloc(&self, lengths: I, alloc_limit: usize) -> Result + where + I: IntoIterator, + { + self.inner.borrow().max_alloc(lengths, alloc_limit) + } } impl BufferProvider for SlotPool { - fn max_alloc_len(&self) -> usize { + fn preferred_segment_len(&self) -> usize { self.inner.borrow().max_alloc_len() } @@ -395,6 +543,13 @@ impl BufferProvider for SlotPool { self.inner.borrow_mut().alloc(len) } + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + self.inner.borrow_mut().alloc_regions(lengths) + } + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { self.inner.borrow_mut().dealloc_addr(addr) } diff --git a/src/hyperlight_common/src/virtq/pool/tests.rs b/src/hyperlight_common/src/virtq/pool/tests.rs index 031759f85..cf4761eb7 100644 --- a/src/hyperlight_common/src/virtq/pool/tests.rs +++ b/src/hyperlight_common/src/virtq/pool/tests.rs @@ -19,6 +19,13 @@ fn make_tiered_slot_pool(lower_count: usize, upper_count: usize) -> SlotPool { SlotPool::new_tiered(lower, upper).unwrap() } +fn alloc_exact_regions( + pool: &impl BufferProvider, + lengths: impl IntoIterator, +) -> Result { + pool.alloc_regions(lengths) +} + #[test] fn test_run_pool_new_success() { let pool = RunPool::<256, 4096>::new(0x10000, 1024 * 1024).unwrap(); @@ -70,7 +77,7 @@ fn test_run_pool_free_from_lower() { pool.dealloc(alloc.addr).unwrap(); assert_eq!( pool.inner.borrow().lower.free_bytes(), - free_before + alloc.len + free_before + alloc.len as usize ); } @@ -83,7 +90,7 @@ fn test_run_pool_free_from_upper() { pool.dealloc(alloc.addr).unwrap(); assert_eq!( pool.inner.borrow().upper.free_bytes(), - free_before + alloc.len + free_before + alloc.len as usize ); } @@ -202,6 +209,8 @@ fn test_tiered_slot_pool_combines_contiguous_equal_sized_layouts() { assert_eq!(pool.base_addr(), 0x80000); assert_eq!(pool.slot_size(), 0x100); assert_eq!(pool.count(), 5); + assert_eq!(pool.num_free_lower(), 0); + assert_eq!(pool.num_free_upper(), 5); assert_eq!(pool.slot_addr(4), Some(0x80400)); } @@ -223,6 +232,13 @@ fn test_tiered_slot_pool_rejects_invalid_layout() { assert!(matches!(reversed_sizes, Err(AllocError::InvalidArg))); } +#[cfg(target_pointer_width = "64")] +#[test] +fn test_slot_pool_rejects_unrepresentable_slot_size() { + let layout = SlotLayout::new(0x80000, u32::MAX as usize + 1, 1); + assert!(matches!(SlotPool::new(layout), Err(AllocError::InvalidArg))); +} + #[test] fn test_tiered_slot_pool_routes_by_size() { let pool = make_tiered_slot_pool(2, 2); @@ -259,22 +275,120 @@ fn test_tiered_slot_pool_does_not_mask_lower_errors() { } #[test] -fn test_tiered_slot_pool_alloc_sg_uses_both_tiers() { - let pool = make_tiered_slot_pool(1, 2); - let sgs = pool.alloc_sg(4096 + 128).unwrap(); +fn test_tiered_slot_pool_reports_free_tier_counts() { + let pool = make_tiered_slot_pool(2, 3); + + assert_eq!(pool.num_free_lower(), 2); + assert_eq!(pool.num_free_upper(), 3); - assert_eq!(sgs.len(), 2); - assert_eq!(sgs[0].len, 4096); - assert_eq!(sgs[1].len, 256); - assert!((0x90000..0x92000).contains(&sgs[0].addr)); - assert!((0x80000..0x80100).contains(&sgs[1].addr)); + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(4096).unwrap(); + assert_eq!(pool.num_free_lower(), 1); + assert_eq!(pool.num_free_upper(), 2); - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); + pool.dealloc(lower.addr).unwrap(); + pool.dealloc(upper.addr).unwrap(); + assert_eq!(pool.num_free_lower(), 2); + assert_eq!(pool.num_free_upper(), 3); +} + +#[test] +fn test_tiered_slot_pool_region_uses_both_tiers() { + let pool = make_tiered_slot_pool(1, 2); + let regions = alloc_exact_regions(&pool, [4096 + 128]).unwrap(); + let allocations = ®ions[0]; + + assert_eq!(regions.len(), 1); + assert_eq!(allocations.len(), 2); + assert_eq!(allocations[0].len, 4096); + assert_eq!(allocations[1].len, 256); + assert!((0x90000..0x92000).contains(&allocations[0].addr)); + assert!((0x80000..0x80100).contains(&allocations[1].addr)); + + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); } assert_eq!(pool.num_free(), 3); } +#[test] +fn test_tiered_slot_pool_allocates_regions_in_order() { + let pool = make_tiered_slot_pool(1, 2); + let regions = alloc_exact_regions(&pool, [128, 128]).unwrap(); + + assert_eq!(regions.len(), 2); + assert_eq!(regions[0].len(), 1); + assert_eq!(regions[1].len(), 1); + assert!((0x80000..0x80100).contains(®ions[0][0].addr)); + assert!((0x90000..0x92000).contains(®ions[1][0].addr)); + + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); + } +} + +#[test] +fn test_slot_pool_max_alloc_reserves_regions_and_honors_limit() { + let pool = make_tiered_slot_pool(2, 3); + + let max = pool.max_alloc([128], 3).unwrap(); + assert_eq!(max, 2 * 4096); + assert_eq!(pool.num_free(), 5); + + let regions = pool.alloc_regions([128, max]).unwrap(); + assert_eq!(regions.iter().map(Allocations::len).sum::(), 3); + + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); + } +} + +#[test] +fn test_slot_pool_max_alloc_requires_one_remaining_allocation() { + let pool = make_tiered_slot_pool(1, 1); + + assert!(matches!(pool.max_alloc([128], 1), Err(AllocError::NoSpace))); + assert_eq!(pool.num_free(), 2); +} + +#[test] +fn test_slot_pool_max_alloc_ignores_remaining_lower_slot() { + let pool = make_tiered_slot_pool(1, 1); + + assert!(matches!( + pool.max_alloc([4096], 2), + Err(AllocError::NoSpace) + )); + assert_eq!(pool.num_free(), 2); +} + +#[test] +fn test_slot_pool_rejects_invalid_or_unavailable_regions() { + let pool = make_tiered_slot_pool(1, 1); + + assert!(matches!( + pool.alloc_regions([0]), + Err(AllocError::InvalidArg) + )); + assert!(matches!( + pool.alloc_regions([]), + Err(AllocError::InvalidArg) + )); + assert!(matches!( + pool.alloc_regions([128, 0]), + Err(AllocError::InvalidArg) + )); + assert!(matches!( + pool.alloc_regions([4096 + 257]), + Err(AllocError::NoSpace) + )); + assert!(matches!( + pool.alloc_regions([128, 4096, 1]), + Err(AllocError::NoSpace) + )); + assert_eq!(pool.num_free(), 2); +} + #[test] fn test_tiered_slot_pool_dealloc_routes_by_region() { let pool = make_tiered_slot_pool(1, 1); @@ -322,53 +436,62 @@ fn test_run_pool_dealloc_addr_routes_to_correct_tier() { } #[test] -fn test_run_pool_alloc_sg_uses_one_contiguous_run() { +fn test_run_pool_region_uses_one_contiguous_run() { let pool = make_run_pool::<256, 4096>(0x20000); - let sgs = pool.alloc_sg(4096 * 2 + 1).unwrap(); + let regions = alloc_exact_regions(&pool, [4096 * 2 + 1]).unwrap(); - assert_eq!(sgs.len(), 1); - assert_eq!(sgs[0].len, 4096 * 3); + assert_eq!(regions.len(), 1); + assert_eq!(regions[0].len(), 1); + assert_eq!(regions[0][0].len, 4096 * 3); - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); } } #[test] -fn test_run_pool_alloc_sg_large_run() { +fn test_run_pool_allocates_each_region_as_one_run() { let pool = make_run_pool::<256, 4096>(0x20000); - let sgs = pool.alloc_sg(8192).unwrap(); + let regions = alloc_exact_regions(&pool, [8192, 128]).unwrap(); - assert_eq!(sgs.len(), 1); - assert_eq!(sgs[0].len, 8192); + assert_eq!(regions.len(), 2); + assert_eq!(regions[0].len(), 1); + assert_eq!(regions[1].len(), 1); + assert_eq!(regions[0][0].len, 8192); + assert_eq!(regions[1][0].len, 256); - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); } } #[test] -fn test_slot_pool_alloc_sg_splits() { +fn test_slot_pool_region_splits() { let pool = make_slot_pool(8, 4096); - let sgs = pool.alloc_sg(4096 * 2 + 1).unwrap(); + let regions = alloc_exact_regions(&pool, [4096 * 2 + 1]).unwrap(); + let allocations = ®ions[0]; - assert_eq!(sgs.len(), 3); - assert_eq!(sgs[0].len, 4096); - assert_eq!(sgs[1].len, 4096); - assert_eq!(sgs[2].len, 4096); + assert_eq!(regions.len(), 1); + assert_eq!(allocations.len(), 3); + assert_eq!(allocations[0].len, 4096); + assert_eq!(allocations[1].len, 4096); + assert_eq!(allocations[2].len, 4096); - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); } } #[test] fn test_tiered_slot_pool_live_addrs_are_deterministic() { let pool = make_tiered_slot_pool(2, 2); + assert_eq!(pool.num_live(), 0); + let lower_high = pool.alloc(128).unwrap(); let upper_high = pool.alloc(1024).unwrap(); let lower_low = pool.alloc(128).unwrap(); + assert_eq!(pool.num_live(), 3); assert_eq!( pool.live_addrs(), vec![lower_low.addr, lower_high.addr, upper_high.addr] @@ -411,10 +534,21 @@ fn test_slot_pool_dealloc_double_free() { } #[test] -fn test_slot_pool_alloc_sg_rolls_back_on_failure() { +fn test_slot_pool_alloc_regions_preflights_sequence() { let pool = make_slot_pool(2, 4096); - assert!(matches!(pool.alloc_sg(4096 * 3), Err(AllocError::NoSpace))); + assert!(matches!( + pool.alloc_regions([]), + Err(AllocError::InvalidArg) + )); + assert!(matches!( + pool.alloc_regions([4096, 0]), + Err(AllocError::InvalidArg) + )); + assert!(matches!( + pool.alloc_regions([4096, 4096, 1]), + Err(AllocError::NoSpace) + )); assert_eq!(pool.num_free(), 2); let alloc = pool.alloc(4096).unwrap(); diff --git a/src/hyperlight_common/src/virtq/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index 831f38c81..2b78cdb10 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -262,6 +262,11 @@ where ChainBuilder::new(self.inner.mem().clone(), self.pool.clone()) } + /// Preferred size of one bulk payload segment. + pub fn preferred_segment_len(&self) -> usize { + self.pool.preferred_segment_len() + } + /// Begin a batch of submissions. /// /// Chains submitted through the returned [`SubmitBatch`] are published to @@ -539,7 +544,7 @@ where self.pool.clone(), Allocation { addr: elem.addr, - len: elem.len as usize, + len: elem.len, }, ); let mem = self.inner.mem().clone(); @@ -684,65 +689,70 @@ impl ChainBuilder { /// # Errors /// /// - [`VirtqError::InvalidState`] - No buffers requested - /// - [`VirtqError::Alloc`] - Pool exhausted + /// - [`VirtqError::Alloc`] - Buffer allocation failed pub fn build(self) -> Result, VirtqError> { if self.rd_caps.is_empty() && self.wr_caps.is_empty() { return Err(VirtqError::InvalidState); } - let rd_capacity = self.rd_caps.iter().sum(); - let mut allocs = AllocTxn::new(&self.pool); + let rd_capacity = self.rd_caps.iter().try_fold(0usize, |total, &cap| { + total.checked_add(cap).ok_or(AllocError::Overflow) + })?; + + let lengths = self.rd_caps.iter().chain(&self.wr_caps).copied(); + let regions = self.pool.alloc_regions(lengths)?; + + debug_assert_eq!(regions.len(), self.rd_caps.len() + self.wr_caps.len()); + + let mut regions = regions.into_iter(); let mut rd_caps = SmallVec::<[usize; 4]>::new(); let mut rd_elems = SmallVec::<[BufferElement; 4]>::new(); let mut wr_elems = SmallVec::<[BufferElement; 4]>::new(); - // Allocate readable buffers, splitting into multiple descriptors if needed. // The buffer element lengths are initialized to zero and updated as the // `SendChain` writes. - for &cap in &self.rd_caps { - let sgs = allocs.alloc_sg(cap)?; + for (&cap, allocs) in Iterator::zip(self.rd_caps.iter(), regions.by_ref()) { let mut remaining = cap; - for alloc in sgs { - let _ = checked_descriptor_len(alloc.len)?; - let seg_cap = remaining.min(alloc.len); + for alloc in allocs { + debug_assert_ne!(remaining, 0); + + let seg_cap = remaining.min(alloc.len as usize); + let elem = BufferElement::readable(alloc.addr); rd_caps.push(seg_cap); - rd_elems.push(BufferElement { - addr: alloc.addr, - len: 0, - writable: false, - }); + rd_elems.push(elem); + remaining -= seg_cap; } - if remaining != 0 { - return Err(VirtqError::InvalidState); - } + // The sum of the allocation lengths must equal the requested capacity. + debug_assert_eq!(remaining, 0); } - // Allocate writable buffers, with the same caveat about splitting as readable buffers. // Writable buffer elements are initialized with their full capacity for the device to // write into. - for &cap in &self.wr_caps { - let sgs = allocs.alloc_sg(cap)?; - for alloc in sgs { - let len = checked_descriptor_len(alloc.len)?; - wr_elems.push(BufferElement { - addr: alloc.addr, - len, - writable: true, - }); + for (&cap, allocs) in Iterator::zip(self.wr_caps.iter(), regions.by_ref()) { + let mut remaining = cap; + + for alloc in allocs { + debug_assert_ne!(remaining, 0); + let elem = BufferElement::writable(alloc.addr, alloc.len); + + wr_elems.push(elem); + remaining = remaining.saturating_sub(alloc.len as usize); } + debug_assert_eq!(remaining, 0); } + // All requested readable and writable buffers must have been allocated. + debug_assert!(regions.next().is_none()); + let chain = BufferChainBuilder::new() .readables(rd_elems) .writables(wr_elems) .build()?; - allocs.commit(); - Ok(SendChain { mem: self.mem, pool: self.pool, @@ -755,51 +765,6 @@ impl ChainBuilder { } } -/// Build-scoped allocation transaction. -/// -/// `SendChain` and `Inflight` intentionally retain lightweight descriptor -/// metadata instead of one allocation guard and cloned pool handle per -/// descriptor. While a valid `BufferChain` is being built, this transaction -/// provides aggregate RAII: it records every allocated address before the -/// caller can perform fallible validation and returns them all on drop. -/// [`commit`](Self::commit) disarms rollback once `SendChain` can take -/// responsibility for reclaiming the completed chain. -struct AllocTxn<'a, P: BufferProvider> { - pool: &'a P, - addrs: SmallVec<[u64; 8]>, -} - -impl<'a, P: BufferProvider> AllocTxn<'a, P> { - fn new(pool: &'a P) -> Self { - Self { - pool, - addrs: SmallVec::new(), - } - } - - fn alloc_sg(&mut self, total_len: usize) -> Result, AllocError> { - let allocs = self.pool.alloc_sg(total_len)?; - self.addrs.extend(allocs.iter().map(|alloc| alloc.addr)); - Ok(allocs) - } - - fn commit(mut self) { - self.addrs.clear(); - } -} - -impl Drop for AllocTxn<'_, P> { - fn drop(&mut self) { - for addr in self.addrs.drain(..) { - let result = self.pool.dealloc(addr); - debug_assert!( - result.is_ok(), - "allocation rollback dealloc failed: {result:?}" - ); - } - } -} - /// Tracks which write API a [`SendChain`] payload uses, so the two paths are /// not mixed. /// @@ -876,12 +841,24 @@ impl SendChain { Inflight { token, chain } } - /// Number of producer-written readable descriptors in this chain. + /// Total number of descriptors in this chain. #[inline] pub fn desc_count(&self) -> usize { + self.chain().len() + } + + /// Number of readable descriptors in this chain. + #[inline] + pub fn rd_desc_count(&self) -> usize { self.chain().readables().len() } + /// Number of writable descriptors in this chain. + #[inline] + pub fn wr_desc_count(&self) -> usize { + self.chain().writables().len() + } + /// Total producer-written readable capacity in bytes. #[inline] pub fn capacity(&self) -> usize { @@ -913,7 +890,7 @@ impl SendChain { /// - [`VirtqError::NoPayloadSegment`] - no readable buffer allocated /// - [`VirtqError::MemoryWriteError`] - underlying write failed pub fn write(&mut self, buf: &[u8]) -> Result { - if self.desc_count() == 0 { + if self.rd_desc_count() == 0 { return Err(VirtqError::NoPayloadSegment); } @@ -966,7 +943,7 @@ impl SendChain { /// - [`VirtqError::MemoryWriteError`] - underlying write failed #[inline] pub fn write_all(&mut self, buf: &[u8]) -> Result<&mut Self, VirtqError> { - if self.desc_count() == 0 { + if self.rd_desc_count() == 0 { return Err(VirtqError::NoPayloadSegment); } @@ -1104,7 +1081,7 @@ fn checked_descriptor_len(len: usize) -> Result { #[cfg(test)] mod tests { use super::*; - use crate::virtq::ring::tests::{TestMem, make_consumer, make_ring}; + use crate::virtq::ring::tests::{OwnedRing, TestMem, make_consumer, make_ring}; use crate::virtq::test_utils::*; fn poll_received( @@ -1113,6 +1090,26 @@ mod tests { consumer.poll(1024).unwrap().unwrap() } + fn make_slot_producer( + ring: &OwnedRing, + slot_size: usize, + ) -> ( + VirtqProducer, + VirtqConsumer, + ) { + let mem = ring.mem(); + let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + + let lower = SlotLayout::new(pool_base, slot_size / 2, ring.len()); + let upper = SlotLayout::new(lower.end_addr().unwrap(), slot_size, ring.len()); + let pool = SlotPool::new_tiered(lower, upper).unwrap(); + + let notifier = TestNotifier::new(); + let producer = VirtqProducer::new(ring.layout(), mem.clone(), notifier.clone(), pool); + let consumer = VirtqConsumer::new(ring.layout(), mem, notifier); + (producer, consumer) + } + fn inflight(seq: u32, id: u16) -> Inflight { let chain = BufferChainBuilder::new() .readable(0x1000 + u64::from(id) * 0x10, 8) @@ -1276,7 +1273,9 @@ mod tests { let (producer, _consumer, _notifier) = make_test_producer(&ring); let se = producer.chain().readable(16).writable(32).build().unwrap(); - assert_eq!(se.desc_count(), 1); + assert_eq!(se.desc_count(), 2); + assert_eq!(se.rd_desc_count(), 1); + assert_eq!(se.wr_desc_count(), 1); assert_eq!(se.capacity(), 16); } @@ -1308,15 +1307,47 @@ mod tests { } #[test] - fn test_chain_multi_readable_appends_across_calls() { + fn test_chain_independent_readables_preserve_pool_tiers() { let ring = make_ring(16); let layout = ring.layout(); let mem = ring.mem(); - let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; - let pool = TestPool::new_with_max_alloc_len(pool_base, 0x8000, 4); + + let lower = SlotLayout::new( + mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100, + 256, + 1, + ); + + let upper = SlotLayout::new(lower.end_addr().unwrap(), 4096, 1); + let pool = SlotPool::new_tiered(lower, upper).unwrap(); let notifier = TestNotifier::new(); - let mut producer = VirtqProducer::new(layout, mem.clone(), notifier.clone(), pool); - let mut consumer = VirtqConsumer::new(layout, mem, notifier); + let producer = VirtqProducer::new(layout, mem, notifier, pool.clone()); + + let send = producer + .chain() + .readable(128) + .readable(4096) + .build() + .unwrap(); + + let readables = send.chain().readables(); + + assert_eq!(readables.len(), 2); + assert_eq!(readables[0].addr, lower.base_addr); + assert_eq!(readables[1].addr, upper.base_addr); + assert_eq!(pool.num_free_lower(), 0); + assert_eq!(pool.num_free_upper(), 0); + + drop(send); + + assert_eq!(pool.num_free_lower(), 1); + assert_eq!(pool.num_free_upper(), 1); + } + + #[test] + fn test_chain_multi_readable_appends_across_calls() { + let ring = make_ring(16); + let (mut producer, mut consumer) = make_slot_producer(&ring, 4); let mut send = producer.chain().readable(8).build().unwrap(); send.write_all(b"abc").unwrap(); @@ -1336,17 +1367,11 @@ mod tests { #[test] fn test_chain_readable_splits_logical_capacity() { let ring = make_ring(16); - let layout = ring.layout(); - let mem = ring.mem(); - let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; - let pool = TestPool::new_with_max_alloc_len(pool_base, 0x8000, 4); - let notifier = TestNotifier::new(); - let mut producer = VirtqProducer::new(layout, mem.clone(), notifier.clone(), pool); - let mut consumer = VirtqConsumer::new(layout, mem, notifier); + let (mut producer, mut consumer) = make_slot_producer(&ring, 4); let mut se = producer.chain().readable(10).writable(32).build().unwrap(); - assert_eq!(se.desc_count(), 3); + assert_eq!(se.rd_desc_count(), 3); assert_eq!(se.capacity(), 10); se.write_all(b"abcdefghij").unwrap(); @@ -1378,13 +1403,7 @@ mod tests { #[test] fn test_chain_writable_splits_logical_capacity() { let ring = make_ring(16); - let layout = ring.layout(); - let mem = ring.mem(); - let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; - let pool = TestPool::new_with_max_alloc_len(pool_base, 0x8000, 4); - let notifier = TestNotifier::new(); - let mut producer = VirtqProducer::new(layout, mem.clone(), notifier.clone(), pool); - let mut consumer = VirtqConsumer::new(layout, mem, notifier); + let (mut producer, mut consumer) = make_slot_producer(&ring, 4); let se = producer.chain().writable(10).build().unwrap(); let token = producer.submit(se).unwrap(); @@ -1680,15 +1699,10 @@ mod tests { #[test] fn test_send_chain_single_segment_writer_rejects_auto_split_chain() { let ring = make_ring(16); - let layout = ring.layout(); - let mem = ring.mem(); - let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; - let pool = TestPool::new_with_max_alloc_len(pool_base, 0x8000, 4); - let notifier = TestNotifier::new(); - let producer = VirtqProducer::new(layout, mem, notifier, pool); + let (producer, _consumer) = make_slot_producer(&ring, 4); let mut se = producer.chain().readable(8).build().unwrap(); - assert_eq!(se.desc_count(), 2); + assert_eq!(se.rd_desc_count(), 2); assert!(matches!( se.with_seg(2, |_| Ok::(0)), Err(VirtqError::NoPayloadSegment) @@ -1769,31 +1783,6 @@ mod tests { assert!(tok.id < 16); } - #[cfg(target_pointer_width = "64")] - #[test] - fn test_chain_build_rolls_back_unrepresentable_allocations() { - let ring = make_ring(16); - let slot_size = u32::MAX as usize + 1; - let layout = SlotLayout::new(0, slot_size, 1); - let pool = SlotPool::new(layout).unwrap(); - let mem = ring.mem(); - let producer = VirtqProducer::new(ring.layout(), mem, TestNotifier::new(), pool.clone()); - - assert!(matches!( - producer.chain().readable(1).build(), - Err(VirtqError::PayloadTooLarge { recv, limit }) - if recv == slot_size && limit == u32::MAX as usize - )); - assert_eq!(pool.num_free(), 1); - - assert!(matches!( - producer.chain().writable(1).build(), - Err(VirtqError::PayloadTooLarge { recv, limit }) - if recv == slot_size && limit == u32::MAX as usize - )); - assert_eq!(pool.num_free(), 1); - } - #[test] fn test_submit_notifies() { let ring = make_ring(16); diff --git a/src/hyperlight_common/src/virtq/ring.rs b/src/hyperlight_common/src/virtq/ring.rs index 56fa7a415..65557812f 100644 --- a/src/hyperlight_common/src/virtq/ring.rs +++ b/src/hyperlight_common/src/virtq/ring.rs @@ -90,6 +90,26 @@ pub struct BufferElement { pub writable: bool, } +impl BufferElement { + /// Create a readable buffer element + pub fn readable(addr: u64) -> Self { + Self { + addr, + len: 0, + writable: false, + } + } + + /// Create a writable buffer element + pub fn writable(addr: u64, len: u32) -> Self { + Self { + addr, + len, + writable: true, + } + } +} + /// A buffer returned from the ring after being used by the device. /// /// When the device completes processing a buffer chain, it returns this From 70de7f8db171407862ec8e78756befe84e3938db Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Tue, 11 Aug 2026 19:06:55 +0200 Subject: [PATCH 14/15] refactor(virtq) adjust virtio terminology Use queue_size instead of queue_depth Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/layout.rs | 44 ++++++------- src/hyperlight_guest/src/layout.rs | 8 +-- src/hyperlight_guest_bin/src/transport.rs | 24 ++++---- src/hyperlight_host/src/mem/layout.rs | 40 ++++++------ src/hyperlight_host/src/mem/mgr.rs | 10 +-- src/hyperlight_host/src/mem/virtq.rs | 10 +-- src/hyperlight_host/src/sandbox/config.rs | 61 +++++++++---------- .../src/sandbox/initialized_multi_use.rs | 20 +++--- .../src/sandbox/snapshot/file/config.rs | 32 +++++----- .../src/sandbox/snapshot/file/mod.rs | 8 +-- .../src/sandbox/snapshot/file_tests.rs | 8 +-- 11 files changed, 132 insertions(+), 133 deletions(-) diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index f0fd03be1..3e8dfd6c8 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -29,13 +29,13 @@ struct ScratchTopMetadata { /// Number of pages reserved for the H2G pool. h2g_pool_pages: u64, /// Host-published H2G descriptor count. - h2g_queue_depth: u64, + h2g_queue_size: u64, /// Host-published capacity of each G2H upper-tier buffer. g2h_buffer_size: u64, /// Number of pages reserved for the G2H pool. g2h_pool_pages: u64, /// Host-published G2H descriptor count. - g2h_queue_depth: u64, + g2h_queue_size: u64, /// Host-published GPA of the fixed transport arena. transport_arena_gpa: u64, /// Seed request for libc's pseudorandom number generator. @@ -54,14 +54,14 @@ const fn scratch_top_offset(field_offset: usize) -> u64 { (size_of::() - field_offset) as u64 } -pub const SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET: u64 = - scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_queue_depth)); +pub const SCRATCH_TOP_G2H_QUEUE_SIZE_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_queue_size)); pub const SCRATCH_TOP_G2H_POOL_PAGES_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_pages)); pub const SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_buffer_size)); -pub const SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET: u64 = - scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_queue_depth)); +pub const SCRATCH_TOP_H2G_QUEUE_SIZE_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_queue_size)); pub const SCRATCH_TOP_H2G_POOL_PAGES_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_pages)); pub const SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET: u64 = @@ -88,10 +88,10 @@ const _: () = { assert!(SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET == 0x20); assert!(SCRATCH_TOP_LIBC_RNG_SEED_OFFSET == 0x28); assert!(SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET == 0x30); - assert!(SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET == 0x38); + assert!(SCRATCH_TOP_G2H_QUEUE_SIZE_OFFSET == 0x38); assert!(SCRATCH_TOP_G2H_POOL_PAGES_OFFSET == 0x40); assert!(SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET == 0x48); - assert!(SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET == 0x50); + assert!(SCRATCH_TOP_H2G_QUEUE_SIZE_OFFSET == 0x50); assert!(SCRATCH_TOP_H2G_POOL_PAGES_OFFSET == 0x58); assert!(SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET == 0x60); assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x70); @@ -116,14 +116,14 @@ pub fn scratch_base_gva(size: usize) -> u64 { pub fn min_scratch_size( input_data_size: usize, output_data_size: usize, - g2h_queue_depth: usize, - h2g_queue_depth: usize, + g2h_queue_size: usize, + h2g_queue_size: usize, g2h_pool_pages: usize, h2g_pool_pages: usize, ) -> usize { let size = arch::min_scratch_size(input_data_size, output_data_size).and_then(|fixed| { - let g2h = QueueDims::new(g2h_queue_depth, g2h_pool_pages)?; - let h2g = QueueDims::new(h2g_queue_depth, h2g_pool_pages)?; + let g2h = QueueDims::new(g2h_queue_size, g2h_pool_pages)?; + let h2g = QueueDims::new(h2g_queue_size, h2g_pool_pages)?; let transport_len = TransportArena::checked_query_size(g2h, h2g)?; fixed.checked_add(transport_len) @@ -135,27 +135,27 @@ pub fn min_scratch_size( /// Validated address independent dimensions for one transport queue. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct QueueDims { - depth: NonZeroU16, + size: NonZeroU16, pool_pages: NonZeroUsize, } impl QueueDims { /// Validate one queue descriptor count and pool page count. - pub fn new(depth: usize, pool_pages: usize) -> Option { - let depth = u16::try_from(depth).ok()?; - let depth = NonZeroU16::new(depth)?; + pub fn new(size: usize, pool_pages: usize) -> Option { + let size = u16::try_from(size).ok()?; + let size = NonZeroU16::new(size)?; - if !depth.get().is_power_of_two() { + if !size.get().is_power_of_two() { return None; } let pool_pages = NonZeroUsize::new(pool_pages)?; - Some(Self { depth, pool_pages }) + Some(Self { size, pool_pages }) } /// Number of descriptors in the queue. - pub const fn depth(&self) -> NonZeroU16 { - self.depth + pub const fn size(&self) -> NonZeroU16 { + self.size } /// Number of pages in the queue's buffer pool. @@ -165,7 +165,7 @@ impl QueueDims { /// Compute the ring length, returning `None` on arithmetic overflow. pub fn checked_ring_len(&self) -> Option { - virtq::Layout::checked_query_size(usize::from(self.depth.get())) + virtq::Layout::checked_query_size(usize::from(self.size.get())) } /// Compute the pool length, returning `None` on arithmetic overflow. @@ -281,7 +281,7 @@ impl TransportArena { /// Convert the arena's absolute addresses into offsets from the arena base. pub fn to_offsets(&self) -> (usize, usize, usize, usize) { - // Already validated by `TransportArena::new`. + #[allow(clippy::unwrap_used)] // `new` proves every stored offset fits in `usize`. let to_offset = |addr| usize::try_from(addr - self.g2h_ring_addr).unwrap(); ( diff --git a/src/hyperlight_guest/src/layout.rs b/src/hyperlight_guest/src/layout.rs index e689a0c4f..3650150d2 100644 --- a/src/hyperlight_guest/src/layout.rs +++ b/src/hyperlight_guest/src/layout.rs @@ -23,8 +23,8 @@ pub fn snapshot_pt_gpa_base_gva() -> *mut u64 { pub fn snapshot_generation_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET) } -pub fn g2h_queue_depth_gva() -> *mut u64 { - scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET) +pub fn g2h_queue_size_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_QUEUE_SIZE_OFFSET) } pub fn transport_arena_gpa_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET) @@ -35,8 +35,8 @@ pub fn g2h_pool_pages_gva() -> *mut u64 { pub fn g2h_buffer_size_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET) } -pub fn h2g_queue_depth_gva() -> *mut u64 { - scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET) +pub fn h2g_queue_size_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_QUEUE_SIZE_OFFSET) } pub fn h2g_pool_pages_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_POOL_PAGES_OFFSET) diff --git a/src/hyperlight_guest_bin/src/transport.rs b/src/hyperlight_guest_bin/src/transport.rs index a920df7a9..0d66c01f9 100644 --- a/src/hyperlight_guest_bin/src/transport.rs +++ b/src/hyperlight_guest_bin/src/transport.rs @@ -16,11 +16,11 @@ pub(crate) fn initialize() { // SAFETY: Generic initialization has mapped writable scratch metadata. let transport_arena_gpa = unsafe { layout::transport_arena_gpa_gva().read_volatile() }; - let (depth, pages, g2h_bufsz) = read_published_g2h(); - let g2h = QueueDims::new(depth, pages).expect("invalid G2H queue dimensions"); + let (size, pages, g2h_bufsz) = read_published_g2h(); + let g2h = QueueDims::new(size, pages).expect("invalid G2H queue dimensions"); - let (depth, pages, h2g_bufsz) = read_published_h2g(); - let h2g = QueueDims::new(depth, pages).expect("invalid H2G queue dimensions"); + let (size, pages, h2g_bufsz) = read_published_h2g(); + let h2g = QueueDims::new(size, pages).expect("invalid H2G queue dimensions"); assert!(g2h_bufsz > 0 && h2g_bufsz > 0); @@ -34,9 +34,9 @@ pub(crate) fn initialize() { let h2g_pool_gva = scratch_gva(arena.h2g_pool_addr()); let g2h_layout = - unsafe { Layout::from_base(g2h_ring_gva, g2h.depth()) }.expect("G2H layout is invalid"); + unsafe { Layout::from_base(g2h_ring_gva, g2h.size()) }.expect("G2H layout is invalid"); let h2g_layout = - unsafe { Layout::from_base(h2g_ring_gva, h2g.depth()) }.expect("H2G layout is invalid"); + unsafe { Layout::from_base(h2g_ring_gva, h2g.size()) }.expect("H2G layout is invalid"); // Build the queues and prefill H2G before exposing either queue to the host. let context = GuestContext::new( @@ -65,26 +65,26 @@ fn scratch_gva(gpa: u64) -> u64 { fn read_published_g2h() -> (usize, usize, usize) { // SAFETY: Generic initialization has mapped writable scratch metadata. - let depth_raw = unsafe { layout::g2h_queue_depth_gva().read_volatile() }; + let size_raw = unsafe { layout::g2h_queue_size_gva().read_volatile() }; let pages_raw = unsafe { layout::g2h_pool_pages_gva().read_volatile() }; let bufsz_raw = unsafe { layout::g2h_buffer_size_gva().read_volatile() }; - let depth = usize::try_from(depth_raw).expect("G2H queue depth exceeds usize"); + let size = usize::try_from(size_raw).expect("G2H queue size exceeds usize"); let pages = usize::try_from(pages_raw).expect("G2H pool page count exceeds usize"); let bufsz = usize::try_from(bufsz_raw).expect("G2H buffer size exceeds usize"); - (depth, pages, bufsz) + (size, pages, bufsz) } fn read_published_h2g() -> (usize, usize, usize) { // SAFETY: Generic initialization has mapped writable scratch metadata. - let depth_raw = unsafe { layout::h2g_queue_depth_gva().read_volatile() }; + let size_raw = unsafe { layout::h2g_queue_size_gva().read_volatile() }; let pages_raw = unsafe { layout::h2g_pool_pages_gva().read_volatile() }; let bufsz_raw = unsafe { layout::h2g_buffer_size_gva().read_volatile() }; - let depth = usize::try_from(depth_raw).expect("H2G queue depth exceeds usize"); + let size = usize::try_from(size_raw).expect("H2G queue size exceeds usize"); let pages = usize::try_from(pages_raw).expect("H2G pool page count exceeds usize"); let bufsz = usize::try_from(bufsz_raw).expect("H2G buffer size exceeds usize"); - (depth, pages, bufsz) + (size, pages, bufsz) } diff --git a/src/hyperlight_host/src/mem/layout.rs b/src/hyperlight_host/src/mem/layout.rs index cbc97e13b..bcd144521 100644 --- a/src/hyperlight_host/src/mem/layout.rs +++ b/src/hyperlight_host/src/mem/layout.rs @@ -256,9 +256,9 @@ pub(crate) struct SandboxMemoryLayout { /// The size of the scratch region in physical memory. scratch_size: usize, /// Number of descriptors in the G2H virtqueue. - g2h_queue_depth: usize, + g2h_queue_size: usize, /// Number of descriptors in the H2G virtqueue. - h2g_queue_depth: usize, + h2g_queue_size: usize, /// Capacity of each G2H upper-tier buffer. g2h_buffer_size: usize, /// Capacity of each H2G buffer. @@ -301,8 +301,8 @@ impl Debug for SandboxMemoryLayout { &format_args!("{:#x}", self.output_data_size), ) .field("Scratch Size", &format_args!("{:#x}", self.scratch_size)) - .field("G2H Queue Depth", &self.g2h_queue_depth) - .field("H2G Queue Depth", &self.h2g_queue_depth) + .field("G2H Queue Size", &self.g2h_queue_size) + .field("H2G Queue Size", &self.h2g_queue_size) .field("G2H Buffer Size", &self.g2h_buffer_size) .field("H2G Buffer Size", &self.h2g_buffer_size) .field("G2H Pool Pages", &self.g2h_pool_pages) @@ -362,8 +362,8 @@ impl SandboxMemoryLayout { } let input_data_size = cfg.get_input_data_size(); let output_data_size = cfg.get_output_data_size(); - let g2h_queue_depth = cfg.get_g2h_queue_depth(); - let h2g_queue_depth = cfg.get_h2g_queue_depth(); + let g2h_queue_size = cfg.get_g2h_queue_size(); + let h2g_queue_size = cfg.get_h2g_queue_size(); let g2h_buffer_size = cfg.get_g2h_buffer_size(); let h2g_buffer_size = cfg.get_h2g_buffer_size(); let g2h_pool_pages = cfg.get_g2h_pool_pages(); @@ -371,8 +371,8 @@ impl SandboxMemoryLayout { let min_scratch_size = hyperlight_common::layout::min_scratch_size( input_data_size, output_data_size, - g2h_queue_depth, - h2g_queue_depth, + g2h_queue_size, + h2g_queue_size, g2h_pool_pages, h2g_pool_pages, ); @@ -389,8 +389,8 @@ impl SandboxMemoryLayout { init_data_permissions, pt_size: None, scratch_size, - g2h_queue_depth, - h2g_queue_depth, + g2h_queue_size, + h2g_queue_size, g2h_buffer_size, h2g_buffer_size, g2h_pool_pages, @@ -430,13 +430,13 @@ impl SandboxMemoryLayout { } #[allow(dead_code)] - pub(crate) fn get_g2h_queue_depth(&self) -> usize { - self.g2h_queue_depth + pub(crate) fn get_g2h_queue_size(&self) -> usize { + self.g2h_queue_size } #[allow(dead_code)] - pub(crate) fn get_h2g_queue_depth(&self) -> usize { - self.h2g_queue_depth + pub(crate) fn get_h2g_queue_size(&self) -> usize { + self.h2g_queue_size } #[allow(dead_code)] @@ -460,12 +460,12 @@ impl SandboxMemoryLayout { } pub(crate) fn get_g2h_queue_dims(&self) -> hyperlight_common::layout::QueueDims { - hyperlight_common::layout::QueueDims::new(self.g2h_queue_depth, self.g2h_pool_pages) + hyperlight_common::layout::QueueDims::new(self.g2h_queue_size, self.g2h_pool_pages) .expect("validated G2H queue dimensions") } pub(crate) fn get_h2g_queue_dims(&self) -> hyperlight_common::layout::QueueDims { - hyperlight_common::layout::QueueDims::new(self.h2g_queue_depth, self.h2g_pool_pages) + hyperlight_common::layout::QueueDims::new(self.h2g_queue_size, self.h2g_pool_pages) .expect("validated H2G queue dimensions") } @@ -494,8 +494,8 @@ impl SandboxMemoryLayout { let min_fixed_scratch = hyperlight_common::layout::min_scratch_size( self.input_data_size, self.output_data_size, - self.g2h_queue_depth, - self.h2g_queue_depth, + self.g2h_queue_size, + self.h2g_queue_size, self.g2h_pool_pages, self.h2g_pool_pages, ); @@ -853,8 +853,8 @@ mod tests { let minimum = hyperlight_common::layout::min_scratch_size( cfg.get_input_data_size(), cfg.get_output_data_size(), - cfg.get_g2h_queue_depth(), - cfg.get_h2g_queue_depth(), + cfg.get_g2h_queue_size(), + cfg.get_h2g_queue_size(), cfg.get_g2h_pool_pages(), cfg.get_h2g_pool_pages(), ); diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs index 06654236f..30ec8a442 100644 --- a/src/hyperlight_host/src/mem/mgr.rs +++ b/src/hyperlight_host/src/mem/mgr.rs @@ -641,10 +641,10 @@ impl SandboxMemoryManager { self.snapshot_count, )?; - // Record the G2H and H2G queue depths, pool page counts, and buffer sizes. + // Record the G2H and H2G queue sizes, pool page counts, and buffer sizes. self.update_scratch_bookkeeping_item( - SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET, - u64::try_from(self.layout.get_g2h_queue_depth())?, + SCRATCH_TOP_G2H_QUEUE_SIZE_OFFSET, + u64::try_from(self.layout.get_g2h_queue_size())?, )?; self.update_scratch_bookkeeping_item( SCRATCH_TOP_G2H_POOL_PAGES_OFFSET, @@ -655,8 +655,8 @@ impl SandboxMemoryManager { u64::try_from(self.layout.get_g2h_buffer_size())?, )?; self.update_scratch_bookkeeping_item( - SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET, - u64::try_from(self.layout.get_h2g_queue_depth())?, + SCRATCH_TOP_H2G_QUEUE_SIZE_OFFSET, + u64::try_from(self.layout.get_h2g_queue_size())?, )?; self.update_scratch_bookkeeping_item( SCRATCH_TOP_H2G_POOL_PAGES_OFFSET, diff --git a/src/hyperlight_host/src/mem/virtq.rs b/src/hyperlight_host/src/mem/virtq.rs index 59543e4c6..2d46a41a0 100644 --- a/src/hyperlight_host/src/mem/virtq.rs +++ b/src/hyperlight_host/src/mem/virtq.rs @@ -161,7 +161,7 @@ impl Config { let h2g = QueueConfig::new(layout.get_h2g_queue_dims(), layout.get_h2g_buffer_size())?; let h2g_prefill_chains = - usize::from(h2g.dims.depth().get()).min(h2g.pool_len / h2g.buffer_size); + usize::from(h2g.dims.size().get()).min(h2g.pool_len / h2g.buffer_size); let arena = layout.get_transport_arena(); Ok(Self { @@ -206,7 +206,7 @@ impl<'a> Validator<'a> { fn validate_g2h(&self, mem: &M, ring: Range) -> Result { // SAFETY: `ring` spans the configured image and `mem` keeps that image // valid for the duration of validation. - let layout = unsafe { VirtqLayout::from_base(ring.start, self.config.g2h.dims.depth()) } + let layout = unsafe { VirtqLayout::from_base(ring.start, self.config.g2h.dims.size()) } .map_err(|error| new_error!("invalid G2H ring layout: {error}"))?; validate_canon_image(mem, layout, 0, |_, _| false) @@ -227,7 +227,7 @@ impl<'a> Validator<'a> { ) -> Result { // SAFETY: `ring` spans the configured image and `mem` keeps that image // valid for the duration of validation. - let layout = unsafe { VirtqLayout::from_base(ring.start, self.config.h2g.dims.depth()) } + let layout = unsafe { VirtqLayout::from_base(ring.start, self.config.h2g.dims.size()) } .map_err(|error| new_error!("invalid H2G ring layout: {error}"))?; let bufsz = self.config.h2g.buffer_size; @@ -434,8 +434,8 @@ mod tests { fn memory_layout() -> SandboxMemoryLayout { let mut config = SandboxConfiguration::default(); config.set_scratch_size(SCRATCH_SIZE); - config.set_g2h_queue_depth(G2H_DEPTH as usize); - config.set_h2g_queue_depth(H2G_DEPTH as usize); + config.set_g2h_queue_size(G2H_DEPTH as usize); + config.set_h2g_queue_size(H2G_DEPTH as usize); config.set_h2g_buffer_size(H2G_BUFFER_SIZE); config.set_g2h_pool_pages(G2H_POOL_PAGES); config.set_h2g_pool_pages(H2G_POOL_PAGES); diff --git a/src/hyperlight_host/src/sandbox/config.rs b/src/hyperlight_host/src/sandbox/config.rs index 12705b455..d9ce766c3 100644 --- a/src/hyperlight_host/src/sandbox/config.rs +++ b/src/hyperlight_host/src/sandbox/config.rs @@ -76,9 +76,9 @@ pub struct SandboxConfiguration { /// How much writable memory to offer the guest scratch_size: usize, /// Number of descriptors in the G2H virtqueue. - g2h_queue_depth: usize, + g2h_queue_size: usize, /// Number of descriptors in the H2G virtqueue. - h2g_queue_depth: usize, + h2g_queue_size: usize, /// Capacity of each G2H upper-tier buffer. g2h_buffer_size: usize, /// Capacity of each H2G buffer. @@ -113,9 +113,9 @@ impl SandboxConfiguration { /// The default size of the scratch region pub const DEFAULT_SCRATCH_SIZE: usize = 0x55000; /// The default G2H virtqueue descriptor count. - pub const DEFAULT_G2H_QUEUE_DEPTH: usize = 64; + pub const DEFAULT_G2H_QUEUE_SIZE: usize = 64; /// The default H2G virtqueue descriptor count. - pub const DEFAULT_H2G_QUEUE_DEPTH: usize = 32; + pub const DEFAULT_H2G_QUEUE_SIZE: usize = 32; /// The default G2H upper-tier buffer size. pub const DEFAULT_G2H_BUFFER_SIZE: usize = PAGE_SIZE; /// The default H2G buffer size. @@ -125,9 +125,9 @@ impl SandboxConfiguration { /// The default total number of H2G pool pages. pub const DEFAULT_H2G_POOL_PAGES: usize = 4; /// The minimum G2H virtqueue descriptor count. - const MIN_QUEUE_DEPTH: usize = 2; + const MIN_QUEUE_SIZE: usize = 2; /// The maximum G2H virtqueue descriptor count. - const MAX_QUEUE_DEPTH: usize = 32_768; + const MAX_QUEUE_SIZE: usize = 32_768; /// The minimum configured transport buffer size. const MIN_BUFFER_SIZE: usize = G2H_LOWER_SLOT_SIZE; /// The maximum configured transport buffer size. @@ -156,8 +156,8 @@ impl SandboxConfiguration { output_data_size: max(output_data_size, Self::MIN_OUTPUT_SIZE), heap_size_override: heap_size_override.unwrap_or(0), scratch_size, - g2h_queue_depth: Self::DEFAULT_G2H_QUEUE_DEPTH, - h2g_queue_depth: Self::DEFAULT_H2G_QUEUE_DEPTH, + g2h_queue_size: Self::DEFAULT_G2H_QUEUE_SIZE, + h2g_queue_size: Self::DEFAULT_H2G_QUEUE_SIZE, g2h_buffer_size: Self::DEFAULT_G2H_BUFFER_SIZE, h2g_buffer_size: Self::DEFAULT_H2G_BUFFER_SIZE, g2h_pool_pages: Self::DEFAULT_G2H_POOL_PAGES, @@ -328,30 +328,30 @@ impl SandboxConfiguration { /// Get the G2H virtqueue descriptor count. #[instrument(skip_all, parent = Span::current(), level= "Trace")] - pub fn get_g2h_queue_depth(&self) -> usize { - self.g2h_queue_depth + pub fn get_g2h_queue_size(&self) -> usize { + self.g2h_queue_size } /// Set the G2H virtqueue descriptor count. /// /// Values are rounded up to a power of two in `2..=32768`. #[instrument(skip_all, parent = Span::current(), level= "Trace")] - pub fn set_g2h_queue_depth(&mut self, depth: usize) { - self.g2h_queue_depth = Self::normalize_queue_depth(depth); + pub fn set_g2h_queue_size(&mut self, size: usize) { + self.g2h_queue_size = Self::normalize_queue_size(size); } /// Get the H2G virtqueue descriptor count. #[instrument(skip_all, parent = Span::current(), level= "Trace")] - pub fn get_h2g_queue_depth(&self) -> usize { - self.h2g_queue_depth + pub fn get_h2g_queue_size(&self) -> usize { + self.h2g_queue_size } /// Set the H2G virtqueue descriptor count. /// /// Values are rounded up to a power of two in `2..=32768`. #[instrument(skip_all, parent = Span::current(), level= "Trace")] - pub fn set_h2g_queue_depth(&mut self, depth: usize) { - self.h2g_queue_depth = Self::normalize_queue_depth(depth); + pub fn set_h2g_queue_size(&mut self, size: usize) { + self.h2g_queue_size = Self::normalize_queue_size(size); } /// Get the capacity of each G2H upper-tier buffer. @@ -443,9 +443,8 @@ impl SandboxConfiguration { .unwrap_or(Self::DEFAULT_HEAP_SIZE) } - fn normalize_queue_depth(depth: usize) -> usize { - depth - .clamp(Self::MIN_QUEUE_DEPTH, Self::MAX_QUEUE_DEPTH) + fn normalize_queue_size(size: usize) -> usize { + size.clamp(Self::MIN_QUEUE_SIZE, Self::MAX_QUEUE_SIZE) .next_power_of_two() } @@ -478,10 +477,10 @@ impl Default for SandboxConfiguration { #[cfg(test)] mod tests { - #[cfg(target_arch = "x86_64")] - use super::GuestMsrError; use hyperlight_common::vmem::PAGE_SIZE; + #[cfg(target_arch = "x86_64")] + use super::GuestMsrError; use super::SandboxConfiguration; #[test] @@ -571,12 +570,12 @@ mod tests { assert_eq!(INPUT_DATA_SIZE_OVERRIDE, cfg.input_data_size); assert_eq!(OUTPUT_DATA_SIZE_OVERRIDE, cfg.output_data_size); assert_eq!( - SandboxConfiguration::DEFAULT_G2H_QUEUE_DEPTH, - cfg.get_g2h_queue_depth() + SandboxConfiguration::DEFAULT_G2H_QUEUE_SIZE, + cfg.get_g2h_queue_size() ); assert_eq!( - SandboxConfiguration::DEFAULT_H2G_QUEUE_DEPTH, - cfg.get_h2g_queue_depth() + SandboxConfiguration::DEFAULT_H2G_QUEUE_SIZE, + cfg.get_h2g_queue_size() ); assert_eq!( SandboxConfiguration::DEFAULT_G2H_BUFFER_SIZE, @@ -622,9 +621,9 @@ mod tests { } #[test] - fn queue_depths_are_normalized() { + fn queue_sizes_are_normalized() { let mut cfg = SandboxConfiguration::default(); - for (depth, expected) in [ + for (size, expected) in [ (0, 2), (1, 2), (2, 2), @@ -634,10 +633,10 @@ mod tests { (32_769, 32_768), (usize::MAX, 32_768), ] { - cfg.set_g2h_queue_depth(depth); - cfg.set_h2g_queue_depth(depth); - assert_eq!(expected, cfg.get_g2h_queue_depth()); - assert_eq!(expected, cfg.get_h2g_queue_depth()); + cfg.set_g2h_queue_size(size); + cfg.set_h2g_queue_size(size); + assert_eq!(expected, cfg.get_g2h_queue_size()); + assert_eq!(expected, cfg.get_h2g_queue_size()); } } diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index bf57b176c..446b66b74 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -242,8 +242,8 @@ impl MultiUseSandbox { config.set_output_data_size(snapshot.layout().output_data_size()); config.set_heap_size(snapshot.layout().heap_size() as u64); config.set_scratch_size(snapshot.layout().get_scratch_size()); - config.set_g2h_queue_depth(snapshot.layout().get_g2h_queue_depth()); - config.set_h2g_queue_depth(snapshot.layout().get_h2g_queue_depth()); + config.set_g2h_queue_size(snapshot.layout().get_g2h_queue_size()); + config.set_h2g_queue_size(snapshot.layout().get_h2g_queue_size()); config.set_g2h_buffer_size(snapshot.layout().get_g2h_buffer_size()); config.set_h2g_buffer_size(snapshot.layout().get_h2g_buffer_size()); config.set_g2h_pool_pages(snapshot.layout().get_g2h_pool_pages()); @@ -1158,14 +1158,14 @@ fn warn_on_layout_override( snapshot.get_scratch_size() as u64, ), ( - "g2h_queue_depth", - caller.get_g2h_queue_depth() as u64, - snapshot.get_g2h_queue_depth() as u64, + "g2h_queue_size", + caller.get_g2h_queue_size() as u64, + snapshot.get_g2h_queue_size() as u64, ), ( - "h2g_queue_depth", - caller.get_h2g_queue_depth() as u64, - snapshot.get_h2g_queue_depth() as u64, + "h2g_queue_size", + caller.get_h2g_queue_size() as u64, + snapshot.get_h2g_queue_size() as u64, ), ( "g2h_buffer_size", @@ -1420,8 +1420,8 @@ mod tests { hyperlight_common::layout::min_scratch_size( defaults.get_input_data_size(), defaults.get_output_data_size(), - defaults.get_g2h_queue_depth(), - defaults.get_h2g_queue_depth(), + defaults.get_g2h_queue_size(), + defaults.get_h2g_queue_size(), defaults.get_g2h_pool_pages(), defaults.get_h2g_pool_pages(), ) diff --git a/src/hyperlight_host/src/sandbox/snapshot/file/config.rs b/src/hyperlight_host/src/sandbox/snapshot/file/config.rs index 3a704f218..b50fca035 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/config.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file/config.rs @@ -214,8 +214,8 @@ pub(super) struct MemoryLayout { /// Memory region flag bits. `None` means default permissions. pub(super) init_data_permissions: Option, pub(super) scratch_size: usize, - pub(super) g2h_queue_depth: usize, - pub(super) h2g_queue_depth: usize, + pub(super) g2h_queue_size: usize, + pub(super) h2g_queue_size: usize, pub(super) g2h_buffer_size: usize, pub(super) h2g_buffer_size: usize, pub(super) g2h_pool_pages: usize, @@ -493,8 +493,8 @@ impl OciSnapshotConfig { } let mut transport = crate::sandbox::SandboxConfiguration::default(); - transport.set_g2h_queue_depth(self.layout.g2h_queue_depth); - transport.set_h2g_queue_depth(self.layout.h2g_queue_depth); + transport.set_g2h_queue_size(self.layout.g2h_queue_size); + transport.set_h2g_queue_size(self.layout.h2g_queue_size); transport.set_g2h_buffer_size(self.layout.g2h_buffer_size); transport.set_h2g_buffer_size(self.layout.h2g_buffer_size); transport.set_g2h_pool_pages(self.layout.g2h_pool_pages); @@ -502,14 +502,14 @@ impl OciSnapshotConfig { for (name, saved, normalized) in [ ( - "g2h_queue_depth", - self.layout.g2h_queue_depth, - transport.get_g2h_queue_depth(), + "g2h_queue_size", + self.layout.g2h_queue_size, + transport.get_g2h_queue_size(), ), ( - "h2g_queue_depth", - self.layout.h2g_queue_depth, - transport.get_h2g_queue_depth(), + "h2g_queue_size", + self.layout.h2g_queue_size, + transport.get_h2g_queue_size(), ), ( "g2h_buffer_size", @@ -841,8 +841,8 @@ mod tests { init_data_size: 0, init_data_permissions: None, scratch_size: 0, - g2h_queue_depth: 64, - h2g_queue_depth: 32, + g2h_queue_size: 64, + h2g_queue_size: 32, g2h_buffer_size: PAGE_SIZE, h2g_buffer_size: PAGE_SIZE, g2h_pool_pages: 8, @@ -1079,8 +1079,8 @@ mod schema_pin { "init_data_size": 5, "init_data_permissions": null, "scratch_size": 8, - "g2h_queue_depth": 64, - "h2g_queue_depth": 32, + "g2h_queue_size": 64, + "h2g_queue_size": 32, "g2h_buffer_size": 4096, "h2g_buffer_size": 4096, "g2h_pool_pages": 8, @@ -1127,8 +1127,8 @@ mod schema_pin { "init_data_size": 5, "init_data_permissions": null, "scratch_size": 8, - "g2h_queue_depth": 64, - "h2g_queue_depth": 32, + "g2h_queue_size": 64, + "h2g_queue_size": 32, "g2h_buffer_size": 4096, "h2g_buffer_size": 4096, "g2h_pool_pages": 8, diff --git a/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs b/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs index 75769464a..84fd33723 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs @@ -610,8 +610,8 @@ impl Snapshot { init_data_size: l.init_data_size(), init_data_permissions: l.init_data_permissions().map(|f| f.bits()), scratch_size: l.get_scratch_size(), - g2h_queue_depth: l.get_g2h_queue_depth(), - h2g_queue_depth: l.get_h2g_queue_depth(), + g2h_queue_size: l.get_g2h_queue_size(), + h2g_queue_size: l.get_h2g_queue_size(), g2h_buffer_size: l.get_g2h_buffer_size(), h2g_buffer_size: l.get_h2g_buffer_size(), g2h_pool_pages: l.get_g2h_pool_pages(), @@ -812,8 +812,8 @@ impl Snapshot { sbox_cfg.set_output_data_size(cfg.layout.output_data_size); sbox_cfg.set_heap_size(cfg.layout.heap_size as u64); sbox_cfg.set_scratch_size(cfg.layout.scratch_size); - sbox_cfg.set_g2h_queue_depth(cfg.layout.g2h_queue_depth); - sbox_cfg.set_h2g_queue_depth(cfg.layout.h2g_queue_depth); + sbox_cfg.set_g2h_queue_size(cfg.layout.g2h_queue_size); + sbox_cfg.set_h2g_queue_size(cfg.layout.h2g_queue_size); sbox_cfg.set_g2h_buffer_size(cfg.layout.g2h_buffer_size); sbox_cfg.set_h2g_buffer_size(cfg.layout.h2g_buffer_size); sbox_cfg.set_g2h_pool_pages(cfg.layout.g2h_pool_pages); diff --git a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs b/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs index ee28108b2..0cd4a5616 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs @@ -2848,8 +2848,8 @@ fn round_trip_preserves_transport_layout() { let mut cfg = SandboxConfiguration::default(); cfg.set_scratch_size(512 * 1024); cfg.set_heap_size(512 * 1024); - cfg.set_g2h_queue_depth(128); - cfg.set_h2g_queue_depth(16); + cfg.set_g2h_queue_size(128); + cfg.set_h2g_queue_size(16); cfg.set_g2h_buffer_size(8192); cfg.set_h2g_buffer_size(2048); cfg.set_g2h_pool_pages(16); @@ -2870,8 +2870,8 @@ fn round_trip_preserves_transport_layout() { .unwrap(); let loaded = Snapshot::checked_load(&path, OciTag::new("latest").unwrap()).unwrap(); - assert_eq!(loaded.layout().get_g2h_queue_depth(), 128); - assert_eq!(loaded.layout().get_h2g_queue_depth(), 16); + assert_eq!(loaded.layout().get_g2h_queue_size(), 128); + assert_eq!(loaded.layout().get_h2g_queue_size(), 16); assert_eq!(loaded.layout().get_g2h_buffer_size(), 8192); assert_eq!(loaded.layout().get_h2g_buffer_size(), 2048); assert_eq!(loaded.layout().get_g2h_pool_pages(), 16); From 7d903d0e49676ad847461d05a0e9f74d28bc3bfb Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Wed, 2 Sep 2026 18:25:23 +0200 Subject: [PATCH 15/15] fix: size foundation memory tests Signed-off-by: Tomasz Andrzejak --- .../src/sandbox/initialized_multi_use.rs | 40 +++++++++---------- .../src/sandbox/snapshot/file_tests.rs | 4 +- .../src/sandbox/uninitialized.rs | 6 ++- src/hyperlight_host/tests/integration_test.rs | 15 +++++-- 4 files changed, 38 insertions(+), 27 deletions(-) diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index 446b66b74..e825b8cfe 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -1406,15 +1406,12 @@ mod tests { assert_eq!(res, 0); } - // Tests to ensure that many (1000) function calls can be made in a call context with a small stack (24K) and heap(32K). - // This test effectively ensures that the stack is being properly reset after each call and we are not leaking memory in the Guest. + // Checks that 1,000 calls work with constrained guest memory. + // This catches guest stack reset and heap leaks. #[test] fn test_with_small_stack_and_heap() { - const HEAP_SIZE: u64 = 32 * 1024; - // min_scratch_size already includes 1 page (4k on most - // platforms) of guest stack, so add 20k more to get 24k - // total, and then add some more for the eagerly-copied page - // tables on amd64 + const HEAP_SIZE: u64 = 128 * 1024; + // Leave headroom for legacy transport and eagerly copied page tables. let scratch_size = { let defaults = SandboxConfiguration::default(); hyperlight_common::layout::min_scratch_size( @@ -1425,8 +1422,7 @@ mod tests { defaults.get_g2h_pool_pages(), defaults.get_h2g_pool_pages(), ) - } + 0x10000 - + 0x10000; + } + 0x40000; let mut sbox1 = SandboxBuilder::from_file(simple_guest_as_pathbuf()) .heap_size(HEAP_SIZE) @@ -2147,7 +2143,7 @@ mod tests { #[test] fn snapshot_restore_recovers_oom_with_larger_heap() { let mut source_cfg = SandboxConfiguration::default(); - source_cfg.set_heap_size(0x20_000); + source_cfg.set_heap_size(0x40_000); let path = simple_guest_as_pathbuf(); let mut source = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(source_cfg)) .unwrap() @@ -2156,7 +2152,7 @@ mod tests { let snapshot = source.snapshot().unwrap(); let mut target_cfg = SandboxConfiguration::default(); - target_cfg.set_heap_size(0x8000); + target_cfg.set_heap_size(0x20_000); let path = simple_guest_as_pathbuf(); let mut target = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(target_cfg)) .unwrap() @@ -2177,7 +2173,7 @@ mod tests { #[test] fn snapshot_restore_applies_smaller_heap_limit() { let mut source_cfg = SandboxConfiguration::default(); - source_cfg.set_heap_size(0x8000); + source_cfg.set_heap_size(0x20_000); let path = simple_guest_as_pathbuf(); let mut source = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(source_cfg)) .unwrap() @@ -2186,7 +2182,7 @@ mod tests { let snapshot = source.snapshot().unwrap(); let mut target_cfg = SandboxConfiguration::default(); - target_cfg.set_heap_size(0x20_000); + target_cfg.set_heap_size(0x80_000); let path = simple_guest_as_pathbuf(); let mut target = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(target_cfg)) .unwrap() @@ -2194,18 +2190,20 @@ mod tests { .unwrap(); assert_eq!( - target.call::("CallMalloc", 0x10_000i32).unwrap(), - 0x10_000 + target.call::("CallMalloc", 0x30_000i32).unwrap(), + 0x30_000 ); target.restore(snapshot).unwrap(); - assert_eq!(target.mem_mgr.layout.heap_size(), 0x8000); - assert!(target.call::("CallMalloc", 0x10_000i32).is_err()); + assert_eq!(target.mem_mgr.layout.heap_size(), 0x20_000); + assert!(target.call::("CallMalloc", 0x30_000i32).is_err()); assert!(target.status().is_poisoned()); } #[test] fn snapshot_restore_applies_smaller_io_limits() { let mut source_cfg = SandboxConfiguration::default(); + source_cfg.set_heap_size(0x40_000); + source_cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 256 * 1024); source_cfg.set_input_data_size(0x2000); source_cfg.set_output_data_size(0x2000); let path = simple_guest_as_pathbuf(); @@ -2216,6 +2214,8 @@ mod tests { let snapshot = source.snapshot().unwrap(); let mut target_cfg = SandboxConfiguration::default(); + target_cfg.set_heap_size(0x40_000); + target_cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 256 * 1024); target_cfg.set_input_data_size(0x8000); target_cfg.set_output_data_size(0x8000); let path = simple_guest_as_pathbuf(); @@ -2242,7 +2242,7 @@ mod tests { let mut small_cfg = SandboxConfiguration::default(); small_cfg.set_input_data_size(0x2000); small_cfg.set_output_data_size(0x2000); - small_cfg.set_heap_size(0x8000); + small_cfg.set_heap_size(0x20_000); let path = simple_guest_as_pathbuf(); let mut small = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(small_cfg)) .unwrap() @@ -2272,7 +2272,7 @@ mod tests { target.restore(small_snapshot.clone()).unwrap(); assert_eq!(target.call::("GetStatic", ()).unwrap(), 11); - assert_eq!(target.mem_mgr.layout.heap_size(), 0x8000); + assert_eq!(target.mem_mgr.layout.heap_size(), 0x20_000); target.restore(large_snapshot).unwrap(); assert_eq!(target.call::("GetStatic", ()).unwrap(), 22); @@ -2280,7 +2280,7 @@ mod tests { target.restore(small_snapshot).unwrap(); assert_eq!(target.call::("GetStatic", ()).unwrap(), 11); - assert_eq!(target.mem_mgr.layout.heap_size(), 0x8000); + assert_eq!(target.mem_mgr.layout.heap_size(), 0x20_000); } #[test] diff --git a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs b/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs index 0cd4a5616..f9900ef92 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs @@ -2781,7 +2781,9 @@ fn round_trip_preserves_stack_top_gva() { #[test] fn round_trip_preserves_non_default_scratch_size() { - let custom_scratch: usize = 256 * 1024; + use crate::sandbox::SandboxConfiguration; + + let custom_scratch = SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 64 * 1024; let mut sbox = SandboxBuilder::from_file(simple_guest_as_pathbuf()) .scratch_size(custom_scratch) .build() diff --git a/src/hyperlight_host/src/sandbox/uninitialized.rs b/src/hyperlight_host/src/sandbox/uninitialized.rs index c1bfd1088..756078e02 100644 --- a/src/hyperlight_host/src/sandbox/uninitialized.rs +++ b/src/hyperlight_host/src/sandbox/uninitialized.rs @@ -1163,6 +1163,7 @@ mod tests { { let mut cfg = SandboxConfiguration::default(); cfg.set_heap_size(16 * 1024 * 1024); // 16MB heap + cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 256 * 1024); let env = GuestEnvironment::new(GuestBinary::FilePath(binary_path.clone()), None); @@ -1185,7 +1186,7 @@ mod tests { // Test 3: Create snapshot with custom scratch size { let mut cfg = SandboxConfiguration::default(); - cfg.set_scratch_size(256 * 1024); // 256KB scratch + cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 64 * 1024); let env = GuestEnvironment::new(GuestBinary::FilePath(binary_path.clone()), None); @@ -1208,6 +1209,7 @@ mod tests { // Test 4: Create snapshot with custom input/output buffer sizes { let mut cfg = SandboxConfiguration::default(); + cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 128 * 1024); cfg.set_input_data_size(64 * 1024); // 64KB input cfg.set_output_data_size(64 * 1024); // 64KB output @@ -1233,7 +1235,7 @@ mod tests { { let mut cfg = SandboxConfiguration::default(); cfg.set_heap_size(32 * 1024 * 1024); // 32MB heap - cfg.set_scratch_size(256 * 1024 * 2); // 512KB scratch (256KB will be input/output) + cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 1024 * 1024); cfg.set_input_data_size(128 * 1024); // 128KB input cfg.set_output_data_size(128 * 1024); // 128KB output diff --git a/src/hyperlight_host/tests/integration_test.rs b/src/hyperlight_host/tests/integration_test.rs index 64cd780de..e0f2aff33 100644 --- a/src/hyperlight_host/tests/integration_test.rs +++ b/src/hyperlight_host/tests/integration_test.rs @@ -523,8 +523,8 @@ fn guest_malloc_abort() { }); // allocate a vector (on heap) that is bigger than the heap - let heap_size = 0x8000; - let size_to_allocate = 0x10000; + let heap_size = 128 * 1024; + let size_to_allocate = 256 * 1024; assert!( size_to_allocate > heap_size, "precondition: size_to_allocate ({size_to_allocate}) must be > heap_size ({heap_size})" @@ -601,7 +601,7 @@ fn corrupt_output_back_pointer_rejected() { #[test] fn guest_panic_no_alloc() { - let heap_size = 0x8000; + let heap_size = 128 * 1024; let configure = |builder: SandboxBuilder| builder.heap_size(heap_size); with_rust_sandbox_from(configure, |mut sbox| { @@ -612,10 +612,15 @@ fn guest_panic_no_alloc() { ) .unwrap_err(); + // Legacy transport may report its own allocation failure. assert!( matches!( &res, - HyperlightError::GuestAborted(code, msg) if *code == ErrorCode::UnknownError as u8 && msg.contains("memory allocation of ") && msg.contains("bytes failed") + HyperlightError::GuestAborted(code, msg) + if (*code == ErrorCode::UnknownError as u8 + && msg.contains("memory allocation of ") + && msg.contains("bytes failed")) + || *code == ErrorCode::MallocFailed as u8 ), "unexpected error: {res:?}" ); @@ -1664,6 +1669,8 @@ fn fill_heap_and_cause_exception() { let err = result.unwrap_err(); match &err { + // Legacy transport may report its own allocation failure. + HyperlightError::GuestAborted(code, _) if *code == ErrorCode::MallocFailed as u8 => {} HyperlightError::GuestAborted(code, message) => { assert_eq!(*code, ErrorCode::GuestError as u8, "Full error: {:?}", err);