diff --git a/CHANGELOG.md b/CHANGELOG.md index 1cc80f8fe..3a3f330d0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,11 @@ 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. +* Shared virtqueue framing with a 12-byte `MsgHeader` and external byte values. +* Producer batch completion without notification and segmented payload + extraction without flattening. ### Changed * `Snapshot::save` now writes the guest memory blob sparsely, skipping all-zero @@ -17,10 +22,19 @@ 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. +* Virtqueue producers use concrete `SlotPool` allocation and `BufferLease` + ownership. `BufferMap` supplies complete owners exposing initialized bytes. +* `ChainBuilder::build()` allocates readable and writable requests. + `writable_avail()` reserves available upper-tier slots within the descriptor + budget. It may add zero slots to a nonempty chain. ### Removed +* `RunPool` and the run-specific `AllocError::InvalidAlign` variant. ### Fixed +* Use a 16 KiB-aligned default scratch size for Apple Silicon compatibility. ## [v0.17.0] - 2026-08-27 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..d8afad453 100644 --- a/docs/snapshot-versioning.md +++ b/docs/snapshot-versioning.md @@ -14,17 +14,17 @@ A snapshot carries three independently evolvable version markers: [src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs](../src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs)). This is what the host reads back from a snapshot: the `OutBAction` and `VmAction` port numbers, the input and output buffer stack - format, the offset and size of each memory region (including the - `HyperlightPEB` size), and the calling convention for guest function - entry. A change to any of these breaks older snapshots unless the - loader adds a compat path. + format, the virtqueue transport layout, the offset and size of each + memory region (including the `HyperlightPEB` size), and the calling + convention for guest function entry. A change to any of these breaks + older snapshots unless the loader adds a compat path. * **Snapshot blob encoding**, `MT_SNAPSHOT_V1` (`application/vnd.hyperlight.snapshot.memory.v1`), aliased as `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/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/guest_trace.rs b/fuzz/fuzz_targets/guest_trace.rs index 9cb05050f..a76828fd4 100644 --- a/fuzz/fuzz_targets/guest_trace.rs +++ b/fuzz/fuzz_targets/guest_trace.rs @@ -53,9 +53,7 @@ impl<'a> Arbitrary<'a> for FuzzInput { // Any unexpected errors from the guest should be reported. fuzz_target!( init: { - // In local tests, 256 KiB seemed sufficient for deep recursion let mu_sbox = SandboxBuilder::from_file(simple_guest_for_fuzzing_as_pathbuf()) - .scratch_size(256 * 1024) .build() .unwrap(); 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/benches/buffer_pool.rs b/src/hyperlight_common/benches/buffer_pool.rs index 19ef5d464..10fe44a68 100644 --- a/src/hyperlight_common/benches/buffer_pool.rs +++ b/src/hyperlight_common/benches/buffer_pool.rs @@ -4,143 +4,25 @@ 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::{SlotLayout, SlotPool}; -// Helper to create a pool for benchmarking -fn make_pool(size: usize) -> BufferPool { - let base = 0x10000; - BufferPool::::new(base, size).unwrap() -} - -// Single allocation performance -fn bench_alloc_single(c: &mut Criterion) { - let mut group = c.benchmark_group("alloc_single"); - - 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); - b.iter(|| { - let alloc = pool.alloc(black_box(size)).unwrap(); - pool.dealloc(alloc.addr).unwrap(); - }); - }); - } - group.finish(); -} - -// LIFO recycling -fn bench_alloc_lifo(c: &mut Criterion) { - let mut group = c.benchmark_group("alloc_lifo"); - - 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); - b.iter(|| { - for _ in 0..100 { - let alloc = pool.alloc(black_box(size)).unwrap(); - pool.dealloc(alloc.addr).unwrap(); - } - }); - }); - } - group.finish(); -} - -// Fragmented allocation worst case -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); - - // Create fragmentation pattern: allocate many, free every other - let mut allocations = Vec::new(); - for _ in 0..100 { - allocations.push(pool.alloc(128).unwrap()); - } - for i in (0..100).step_by(2) { - pool.dealloc(allocations[i].addr).unwrap(); - } - - b.iter(|| { - let alloc = pool.alloc(black_box(256)).unwrap(); - pool.dealloc(alloc.addr).unwrap(); - }); - }); - - group.finish(); -} - -// Free performance -fn bench_free(c: &mut Criterion) { - let mut group = c.benchmark_group("free"); - - 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); - b.iter(|| { - let alloc = pool.alloc(size).unwrap(); - pool.dealloc(black_box(alloc.addr)).unwrap(); - }); - }); - } - - group.finish(); -} - -// Free-list reuse -fn bench_free_list_reuse(c: &mut Criterion) { - let mut group = c.benchmark_group("free_list_reuse"); - - // With cursor optimization (LIFO) - group.bench_function("lifo_pattern", |b| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); - b.iter(|| { - let alloc = pool.alloc(256).unwrap(); - pool.dealloc(alloc.addr).unwrap(); - let alloc2 = pool.alloc(black_box(256)).unwrap(); - pool.dealloc(alloc2.addr).unwrap(); - }); - }); - - // Without cursor benefit (FIFO-like) - group.bench_function("fifo_pattern", |b| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); - let mut queue = Vec::new(); - - // Pre-fill queue - for _ in 0..10 { - queue.push(pool.alloc(256).unwrap()); - } - - b.iter(|| { - // FIFO: free oldest, allocate new - let old = queue.remove(0); - pool.dealloc(old.addr).unwrap(); - queue.push(pool.alloc(black_box(256)).unwrap()); - }); - }); - - group.finish(); -} - -// Segmented logical payload allocation +// Raw segmented allocation without producer bookkeeping. fn bench_segmented_payload(c: &mut Criterion) { - let mut group = c.benchmark_group("segmented_payload"); + let mut group = c.benchmark_group("payload_allocation"); for payload_size in [8 * 1024usize, 64 * 1024, 256 * 1024] { group.throughput(Throughput::Bytes(payload_size as u64)); group.bench_with_input( - BenchmarkId::from_parameter(payload_size), + BenchmarkId::new("slot_pool_segmented", payload_size), &payload_size, |b, &payload_size| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = SlotPool::new(SlotLayout::new(0x80000, 4096, 1024)).unwrap(); b.iter(|| { - let sgs = pool.alloc_sg(black_box(payload_size)).unwrap(); - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); + let allocations: Vec<_> = (0..black_box(payload_size).div_ceil(4096)) + .map(|_| pool.alloc(4096).unwrap()) + .collect(); + for allocation in allocations.into_iter().rev() { + pool.dealloc(allocation.addr).unwrap(); } }); }, @@ -150,11 +32,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 +45,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,35 +54,17 @@ 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(); }); }); - group.bench_function("alloc_sg_64k", |b| { - let pool = RecyclePool::new(0x80000, 4 * 1024 * 1024, 4096).unwrap(); - b.iter(|| { - let sgs = pool.alloc_sg(black_box(64 * 1024)).unwrap(); - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); - } - }); - }); - group.finish(); } -criterion_group!( - benches, - bench_alloc_single, - bench_alloc_lifo, - bench_alloc_fragmented, - bench_free, - bench_free_list_reuse, - bench_segmented_payload, - bench_recycle_pool, -); +criterion_group!(benches, bench_segmented_payload, 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 cd68a8176..f9a948aeb 100644 --- a/src/hyperlight_common/benches/common/mod.rs +++ b/src/hyperlight_common/benches/common/mod.rs @@ -4,6 +4,7 @@ //! Shared harness for the `virtq_api` benchmarks: an in-memory [`MemOps`] //! backend, a counting [`Notifier`], producer/consumer pair construction, pool //! factories, and request/response round-trip drivers. +//! Reply owners copy payloads to keep native buffer leases thread-local. use std::cell::UnsafeCell; use std::hint::black_box; @@ -14,16 +15,13 @@ 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, + BufferLease, BufferMap, Descriptor, Layout, MemOps, Notifier, QueueStats, ReplyChain, + 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; - #[derive(Clone)] struct BenchMem { inner: Arc, @@ -112,6 +110,29 @@ unsafe impl MemOps for BenchMem { } } +impl BufferMap for BenchMem { + type Mapping = Vec; + + unsafe fn map_buffer( + &self, + lease: BufferLease, + written: usize, + ) -> Result { + let allocation = lease.allocation(); + assert!(written <= allocation.len as usize); + + let offset = allocation.addr.checked_sub(self.base_addr()).unwrap() as usize; + // SAFETY: Benchmark storage is never resized. + let size = unsafe { &*self.inner.storage.get() }.len(); + assert!(offset.checked_add(allocation.len as usize).unwrap() <= size); + + let mut bytes = vec![0; written]; + self.read(allocation.addr, &mut bytes)?; + + Ok(bytes) + } +} + #[derive(Clone)] struct BenchNotifier { count: Arc, @@ -132,8 +153,8 @@ impl Notifier for BenchNotifier { } /// A producer/consumer pair sharing one in-memory ring and pool. -pub struct BenchPair

{ - producer: VirtqProducer, +pub struct BenchPair { + producer: VirtqProducer, consumer: VirtqConsumer, } @@ -143,10 +164,7 @@ fn align_up(value: usize, align: usize) -> usize { /// Build a [`BenchPair`] with `descs` ring descriptors and a pool built by /// `make_pool`. -pub fn make_pair

(descs: usize, make_pool: impl FnOnce(u64, usize) -> P) -> BenchPair

-where - P: BufferProvider + Clone, -{ +pub fn make_pair(descs: usize, make_pool: impl FnOnce(u64, usize) -> SlotPool) -> BenchPair { let ring_size = Layout::query_size(descs); let mem = BenchMem::new(ring_size + POOL_SIZE + 0x20000); let ring_base = align_up(mem.base_addr() as usize, Descriptor::ALIGN) as u64; @@ -163,33 +181,13 @@ where BenchPair { producer, consumer } } -pub fn run_buffer_pool(base: u64, size: usize) -> RunBufferPool { - BufferPool::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); - let payload_slots = payload_size.div_ceil(UPPER_SLOT); - let prefix_slots = 32; - let suffix_slots = 32; - - let allocated: Vec<_> = (0..prefix_slots + payload_slots + suffix_slots) - .map(|_| pool.alloc(UPPER_SLOT).unwrap()) - .collect(); - - for alloc in &allocated[prefix_slots..prefix_slots + payload_slots] { - pool.dealloc(alloc.addr).unwrap(); - } - - 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()) @@ -204,10 +202,7 @@ pub fn fragmented_recycle_pool(base: u64, size: usize, payload_size: usize) -> R /// Drive one read-only (fire-and-forget) chain through submit, consume, ack, and /// poll, returning the producer-observed used chain. -pub fn readonly_roundtrip

(pair: &mut BenchPair

, payload: &[u8]) -> UsedChain -where - P: BufferProvider + Clone + Send + 'static, -{ +pub fn readonly_roundtrip(pair: &mut BenchPair, payload: &[u8]) -> UsedChain { let mut chain = pair .producer .chain() @@ -219,8 +214,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); @@ -229,10 +224,7 @@ where /// Drive one request/response chain through submit, consume, write reply, /// complete, and poll, returning the producer-observed used chain. -pub fn readwrite_roundtrip

(pair: &mut BenchPair

, request: &[u8], response: &[u8]) -> UsedChain -where - P: BufferProvider + Clone + Send + 'static, -{ +pub fn readwrite_roundtrip(pair: &mut BenchPair, request: &[u8], response: &[u8]) -> UsedChain { let mut chain = pair .producer .chain() @@ -245,13 +237,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/benches/virtq_api.rs b/src/hyperlight_common/benches/virtq_api.rs index 565ffde00..c29b1edd0 100644 --- a/src/hyperlight_common/benches/virtq_api.rs +++ b/src/hyperlight_common/benches/virtq_api.rs @@ -9,18 +9,18 @@ use hyperlight_common::virtq::UsedChain; mod common; use common::*; -fn bench_readonly_strategies(c: &mut Criterion) { - let mut group = c.benchmark_group("virtq_readonly_allocator_strategy"); +fn bench_readonly(c: &mut Criterion) { + let mut group = c.benchmark_group("virtq_readonly"); for size in [8 * 1024usize, 64 * 1024, 256 * 1024] { let payload = vec![0xA5u8; size]; group.throughput(Throughput::Bytes(size as u64)); group.bench_with_input( - BenchmarkId::new("buffer_pool_run", size), + BenchmarkId::new("slot_pool_segmented", size), &payload, |b, payload| { - let mut pair = make_pair(128, run_buffer_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(_))); @@ -29,37 +29,11 @@ fn bench_readonly_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("buffer_pool_run_fragmented", size), + BenchmarkId::new("slot_pool_segmented_fragmented", size), &payload, |b, payload| { let mut pair = make_pair(128, |base, pool_size| { - fragmented_run_buffer_pool(base, pool_size, payload.len()) - }); - b.iter(|| { - let used = readonly_roundtrip(&mut pair, black_box(payload)); - debug_assert!(matches!(used, UsedChain::Ack(_))); - }); - }, - ); - - group.bench_with_input( - BenchmarkId::new("recycle_pool_segmented", size), - &payload, - |b, payload| { - let mut pair = make_pair(128, recycle_pool); - b.iter(|| { - let used = readonly_roundtrip(&mut pair, black_box(payload)); - debug_assert!(matches!(used, UsedChain::Ack(_))); - }); - }, - ); - - group.bench_with_input( - BenchmarkId::new("recycle_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)); @@ -72,8 +46,8 @@ fn bench_readonly_strategies(c: &mut Criterion) { group.finish(); } -fn bench_readwrite_strategies(c: &mut Criterion) { - let mut group = c.benchmark_group("virtq_readwrite_allocator_strategy"); +fn bench_readwrite(c: &mut Criterion) { + let mut group = c.benchmark_group("virtq_readwrite"); for size in [8 * 1024usize, 64 * 1024, 256 * 1024] { let request = vec![0x11u8; size]; @@ -81,44 +55,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), - &(request.clone(), response.clone()), - |b, (request, response)| { - let mut pair = make_pair(128, run_buffer_pool); - b.iter(|| { - let used = readwrite_roundtrip( - &mut pair, - black_box(request.as_slice()), - black_box(response.as_slice()), - ); - black_box(used.segments().unwrap().segment_count()); - }); - }, - ); - - group.bench_with_input( - BenchmarkId::new("buffer_pool_run_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()) - }); - b.iter(|| { - let used = readwrite_roundtrip( - &mut pair, - black_box(request.as_slice()), - black_box(response.as_slice()), - ); - black_box(used.segments().unwrap().segment_count()); - }); - }, - ); - - 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 +71,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( @@ -152,10 +92,6 @@ fn bench_readwrite_strategies(c: &mut Criterion) { group.finish(); } -criterion_group!( - benches, - bench_readonly_strategies, - bench_readwrite_strategies, -); +criterion_group!(benches, bench_readonly, bench_readwrite,); criterion_main!(benches); 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 28a03eb9e..1666ba6be 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -1,6 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // 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")] mod arch; @@ -9,13 +12,95 @@ 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; +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; + +// Fields are listed in ascending-address order. Public offsets are measured +// down from the top of scratch memory. +#[repr(C)] +struct ScratchTopMetadata { + /// 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, + /// Host-published H2G descriptor count. + 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_size: 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. + 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_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_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 = + 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 = + 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_TRANSPORT_ARENA_GPA_OFFSET == 0x30); + 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_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); +}; + +/// 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 @@ -25,4 +110,254 @@ pub fn scratch_base_gva(size: usize) -> u64 { } /// Compute the minimum scratch region size needed for a sandbox. -pub use arch::min_scratch_size; +/// +/// `transport_len` includes both rings and buffer pools. +/// The result saturates at [`usize::MAX`]. +pub fn min_scratch_size( + input_data_size: usize, + output_data_size: usize, + transport_len: usize, +) -> usize { + arch::min_scratch_size(input_data_size, output_data_size) + .and_then(|fixed| fixed.checked_add(transport_len)) + .unwrap_or(usize::MAX) +} + +/// Validated address independent dimensions for one transport queue. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct QueueDims { + size: NonZeroU16, + pool_pages: NonZeroUsize, +} + +impl QueueDims { + /// Validate queue dimensions and their byte lengths. + pub fn new(size: usize, pool_pages: usize) -> Option { + let size = u16::try_from(size).ok()?; + let size = NonZeroU16::new(size)?; + + if !size.get().is_power_of_two() { + return None; + } + + let pool_pages = NonZeroUsize::new(pool_pages)?; + pool_pages.get().checked_mul(crate::vmem::PAGE_SIZE)?; + Some(Self { size, pool_pages }) + } + + /// Number of descriptors in the queue. + pub const fn size(&self) -> NonZeroU16 { + self.size + } + + /// Number of pages in the queue's buffer pool. + pub const fn pool_pages(&self) -> NonZeroUsize { + self.pool_pages + } + + /// Ring length in bytes, including event suppressions. + pub fn ring_len(&self) -> usize { + virtq::Layout::query_size(usize::from(self.size.get())) + } + + /// Pool length in bytes. + pub fn pool_len(&self) -> usize { + self.pool_pages.get() * 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 + .ring_len() + .checked_next_multiple_of(virtq::Descriptor::ALIGN)?; + + let g2h_pool_offset = h2g_ring_offset + .checked_add(h2g.ring_len())? + .checked_next_multiple_of(crate::vmem::PAGE_SIZE)?; + + let g2h_pool_len = g2h.pool_len(); + let h2g_pool_offset = g2h_pool_offset.checked_add(g2h_pool_len)?; + + let h2g_pool_len = h2g.pool_len(); + let len = h2g_pool_offset.checked_add(h2g_pool_len)?; + + 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, + }) + } + + /// 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) { + #[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(); + + ( + 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 queue_dims_validate_byte_lengths() { + let max_pages = usize::MAX / crate::vmem::PAGE_SIZE; + let dims = QueueDims::new(32768, max_pages).unwrap(); + + assert_eq!(dims.ring_len(), 0x80008); + assert_eq!(dims.pool_len(), max_pages * crate::vmem::PAGE_SIZE); + + for (size, pages) in [ + (0, 1), + (3, 1), + (usize::MAX, 1), + (64, 0), + (64, max_pages + 1), + ] { + assert_eq!(QueueDims::new(size, pages), None); + } + } + + #[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::new(base + 1, g2h, h2g), None); + + let oversized = QueueDims::new(64, usize::MAX / crate::vmem::PAGE_SIZE).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(); + let transport_len = (1 + 8 + 4) * crate::vmem::PAGE_SIZE; + + assert_eq!(fixed + transport_len, min_scratch_size(0, 0, transport_len)); + } + + #[test] + fn minimum_scratch_saturates_on_overflow() { + assert_eq!(usize::MAX, min_scratch_size(0, 0, usize::MAX)); + assert_eq!(usize::MAX, min_scratch_size(usize::MAX, 1, 0)); + } +} diff --git a/src/hyperlight_common/src/lib.rs b/src/hyperlight_common/src/lib.rs index 7450e4800..dcf791fec 100644 --- a/src/hyperlight_common/src/lib.rs +++ b/src/hyperlight_common/src/lib.rs @@ -33,6 +33,10 @@ pub mod outb; /// cbindgen:ignore pub mod resource; +/// Shared guest and host transport protocol. +// cbindgen:ignore +pub mod transport; + /// cbindgen:ignore pub mod func; diff --git a/src/hyperlight_common/src/transport.rs b/src/hyperlight_common/src/transport.rs new file mode 100644 index 000000000..2f8eac43a --- /dev/null +++ b/src/hyperlight_common/src/transport.rs @@ -0,0 +1,450 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! Shared guest and host transport protocol. +//! +//! Every logical message starts with this fixed header. It enables message type +//! discrimination, request/response correlation, and payload length validation. + +use alloc::vec::Vec; + +use anyhow::Result; +pub use bytes::Buf; +use bytes::Bytes; + +use crate::flatbuffer_wrappers::ExternalValueSink; + +/// Length of a FlatBuffer size prefix. +pub const SIZE_PREFIX_LEN: usize = core::mem::size_of::(); + +/// Message types for the virtqueue wire protocol. +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MsgKind { + /// A function call request (FunctionCall payload follows). + Request = 0x01, + /// A function call response (FunctionCallResult payload follows). + Response = 0x02, + /// A stream data chunk. + StreamChunk = 0x03, + /// End-of-stream marker. + StreamEnd = 0x04, + /// Cancel a pending request. + Cancel = 0x05, + /// A guest log message (GuestLogData payload follows). + Log = 0x06, + /// Internal request to prepare canonical transport state for snapshotting. + SnapshotCheckpoint = 0x07, +} + +impl TryFrom for MsgKind { + type Error = u8; + + fn try_from(value: u8) -> Result { + match value { + 0x01 => Ok(Self::Request), + 0x02 => Ok(Self::Response), + 0x03 => Ok(Self::StreamChunk), + 0x04 => Ok(Self::StreamEnd), + 0x05 => Ok(Self::Cancel), + 0x06 => Ok(Self::Log), + 0x07 => Ok(Self::SnapshotCheckpoint), + other => Err(other), + } + } +} + +/// Wire header for all virtqueue messages. +#[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable)] +#[repr(C)] +pub struct MsgHeader { + /// Discriminates the message type. + pub kind: u8, + /// Keep the header aligned to four bytes. + reserved: [u8; 3], + /// Caller-assigned correlation ID. Responses echo the request's ID. + pub cid: u32, + /// Total number of payload bytes in this logical message. + pub payload_len: u32, +} + +impl MsgHeader { + pub const SIZE: usize = core::mem::size_of::(); + + /// Create a message header. + pub const fn new(kind: MsgKind, cid: u32, payload_len: u32) -> Self { + Self { + kind: kind as u8, + reserved: [0; 3], + cid, + payload_len, + } + } + + /// Parse the kind field into a [`MsgKind`] enum. + pub fn msg_kind(&self) -> Result { + MsgKind::try_from(self.kind) + } + + /// Return the wire representation. + pub fn as_bytes(&self) -> &[u8] { + bytemuck::bytes_of(self) + } + + /// Parse and validate a wire header. + pub fn from_bytes(bytes: &[u8]) -> Option { + if bytes.len() != Self::SIZE { + return None; + } + + let header: Self = bytemuck::pod_read_unaligned(bytes); + (header.reserved == [0; 3] && header.msg_kind().is_ok()).then_some(header) + } +} + +/// Borrowed wire message split into transport-ready chunks. +#[derive(Debug)] +pub struct EncodedMessage<'a> { + header: MsgHeader, + control: &'a [u8], + externals: ExternalValues<'a>, + total_len: usize, +} + +impl<'a> EncodedMessage<'a> { + /// Build a message, returning `None` if its payload exceeds the wire field. + pub fn new( + kind: MsgKind, + cid: u32, + control: &'a [u8], + externals: ExternalValues<'a>, + ) -> Option { + let payload_len = control.len().checked_add(externals.total_len())?; + let payload_len = u32::try_from(payload_len).ok()?; + let total_len = MsgHeader::SIZE.checked_add(payload_len as usize)?; + + Some(Self { + header: MsgHeader::new(kind, cid, payload_len), + control, + externals, + total_len, + }) + } + + // Build a snapshot checkpoint message with no payload. + pub fn new_snapshot_cp() -> Self { + let total_len = MsgHeader::SIZE; + let externals = ExternalValues::new(); + + Self { + header: MsgHeader::new(MsgKind::SnapshotCheckpoint, 0, 0), + control: &[], + externals, + total_len, + } + } + + /// Borrow the complete wire message as a zero-copy byte cursor. + pub fn as_buf(&self) -> impl Buf + '_ { + EncodedMessageBuf::new( + self.header.as_bytes(), + self.control, + &self.externals.chunks, + self.total_len, + ) + } + + /// Iterate over the complete wire message in transmission order. + pub fn chunks(&self) -> impl Iterator + '_ { + core::iter::once(self.header.as_bytes()) + .chain(core::iter::once(self.control)) + .chain(self.externals.chunks()) + } + + /// Iterate over external transport chunks in wire order. + pub fn external_chunks(&self) -> impl Iterator + '_ { + self.externals.chunks() + } + + /// Message header. + pub const fn header(&self) -> &MsgHeader { + &self.header + } + + /// Size-prefixed FlatBuffer control data. + pub const fn control(&self) -> &[u8] { + self.control + } + + /// Total external byte-stream length. + pub const fn external_len(&self) -> usize { + self.payload_len() - self.control.len() + } + + /// Length of the header and control prefix before external bytes. + pub const fn prefix_len(&self) -> usize { + MsgHeader::SIZE + self.control.len() + } + + /// Logical payload length after the header. + pub const fn payload_len(&self) -> usize { + self.header.payload_len as usize + } + + /// Total wire length of all chunks. + pub const fn total_len(&self) -> usize { + self.total_len + } +} + +/// Borrowed [`Buf`] cursor over an [`EncodedMessage`]. +/// +/// Advancing the cursor does not mutate the message or copy its chunks. +struct EncodedMessageBuf<'a> { + header: &'a [u8], + control: &'a [u8], + externals: &'a [&'a [u8]], + index: usize, + offset: usize, + remaining: usize, +} + +impl<'a> EncodedMessageBuf<'a> { + fn new( + header: &'a [u8], + control: &'a [u8], + externals: &'a [&'a [u8]], + remaining: usize, + ) -> Self { + let mut this = Self { + header, + control, + externals, + index: 0, + offset: 0, + remaining, + }; + + this.skip_empty_chunks(); + this + } + + fn current(&self) -> Option<&[u8]> { + match self.index { + 0 => Some(self.header), + 1 => Some(self.control), + index => self.externals.get(index - 2).copied(), + } + } + + fn skip_empty_chunks(&mut self) { + while self + .current() + .is_some_and(|chunk| self.offset >= chunk.len()) + { + self.index += 1; + self.offset = 0; + } + } +} + +impl Buf for EncodedMessageBuf<'_> { + fn remaining(&self) -> usize { + self.remaining + } + + fn chunk(&self) -> &[u8] { + if self.remaining == 0 { + return &[]; + } + + #[allow(clippy::expect_used)] // `remaining` is derived from the chunks. + let chunk = self.current().expect("message length mismatch"); + &chunk[self.offset..] + } + + fn advance(&mut self, cnt: usize) { + assert!(cnt <= self.remaining, "cannot advance past remaining bytes"); + + self.remaining -= cnt; + let mut cnt = cnt; + + while cnt != 0 { + #[allow(clippy::expect_used)] // `remaining` advances with `index`. + let chunk = self.current().expect("message length mismatch"); + let advanced = cnt.min(chunk.len() - self.offset); + + self.offset += advanced; + cnt -= advanced; + self.skip_empty_chunks(); + } + } +} + +/// Borrowed external values collected while encoding a FlatBuffer. +#[derive(Debug, Default)] +pub struct ExternalValues<'a> { + chunks: Vec<&'a [u8]>, + total_len: usize, +} + +impl<'a> ExternalValues<'a> { + /// Create an empty collection. + pub fn new() -> Self { + Self::default() + } + + /// Iterate over transport chunks in wire order. + fn chunks(&self) -> impl Iterator + '_ { + self.chunks.iter().copied() + } + + /// Total byte length of all collected values. + pub const fn total_len(&self) -> usize { + self.total_len + } +} + +impl<'a> ExternalValueSink<'a> for ExternalValues<'a> { + fn push_bytes(&mut self, value: &'a [u8]) -> Result<()> { + if value.is_empty() { + return Ok(()); + } + + self.total_len = self + .total_len + .checked_add(value.len()) + .ok_or_else(|| anyhow::anyhow!("external value length overflow"))?; + + self.chunks.push(value); + Ok(()) + } + + fn push_chunks(&mut self, value: &'a [Bytes]) -> Result<()> { + let total_len = value + .iter() + .try_fold(self.total_len, |len, chunk| len.checked_add(chunk.len())) + .ok_or_else(|| anyhow::anyhow!("external value length overflow"))?; + + let chunks = value + .iter() + .map(Bytes::as_ref) + .filter(|chunk| !chunk.is_empty()); + + self.chunks.extend(chunks); + self.total_len = total_len; + Ok(()) + } +} + +/// Decode a FlatBuffer size prefix. +pub fn size_prefix_payload_len(prefix: &[u8]) -> Option { + // TODO: this is flatbuffer-specific and should be moved probably somewhere else. + let prefix = <[u8; SIZE_PREFIX_LEN]>::try_from(prefix).ok()?; + usize::try_from(u32::from_le_bytes(prefix)).ok() +} + +/// Add the FlatBuffer size prefix to a payload length. +pub const fn size_prefixed_len(payload_len: usize) -> Option { + // TODO: this is flatbuffer-specific and should be moved probably somewhere else. + SIZE_PREFIX_LEN.checked_add(payload_len) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::flatbuffer_wrappers::ExternalValueSink; + + #[test] + fn header_contains_framing_fields() { + let header = MsgHeader::new(MsgKind::Response, 0x1234_5678, 4096); + + assert_eq!(MsgHeader::SIZE, 12); + assert_eq!(header.msg_kind(), Ok(MsgKind::Response)); + assert_eq!(header.cid, 0x1234_5678); + assert_eq!(header.payload_len, 4096); + assert_eq!(header.reserved, [0; 3]); + } + + #[test] + fn rejects_invalid_wire_headers() { + let header = MsgHeader::new(MsgKind::Request, 1, 4); + let mut bytes = [0; MsgHeader::SIZE]; + bytes.copy_from_slice(header.as_bytes()); + + bytes[1] = 1; + assert_eq!(MsgHeader::from_bytes(&bytes), None); + + bytes[1] = 0; + bytes[0] = u8::MAX; + assert_eq!(MsgHeader::from_bytes(&bytes), None); + + bytes[0] = MsgKind::Request as u8; + assert_eq!(MsgHeader::from_bytes(&bytes[..MsgHeader::SIZE - 1]), None); + } + + #[test] + fn encoded_message_yields_wire_chunks_in_order() { + let chunks = [ + bytes::Bytes::from_static(b"ef"), + bytes::Bytes::from_static(b"gh"), + ]; + let mut external_values = ExternalValues::new(); + external_values.push_bytes(b"cd").unwrap(); + external_values.push_chunks(&chunks).unwrap(); + + let message = EncodedMessage::new(MsgKind::Request, 7, b"ab", external_values).unwrap(); + let visited: Vec<_> = message.chunks().map(<[u8]>::to_vec).collect(); + + assert_eq!(message.total_len(), MsgHeader::SIZE + 8); + assert_eq!(message.prefix_len(), MsgHeader::SIZE + 2); + assert_eq!(message.payload_len(), 8); + assert_eq!(visited[1..], [b"ab", b"cd", b"ef", b"gh"]); + } + + #[test] + fn encoded_message_buf_skips_empty_chunks() { + let mut external_values = ExternalValues::new(); + external_values.chunks.push(&[]); + external_values.push_bytes(b"ab").unwrap(); + + let message = EncodedMessage::new(MsgKind::Request, 7, &[], external_values).unwrap(); + let expected = message.chunks().flatten().copied().collect::>(); + let mut cursor = message.as_buf(); + let mut actual = vec![0; cursor.remaining()]; + + cursor.copy_to_slice(&mut actual); + + assert_eq!(actual, expected); + assert!(!cursor.has_remaining()); + } + + #[test] + fn encoded_message_rejects_length_overflow() { + let external_values = ExternalValues { + chunks: Vec::new(), + total_len: usize::MAX, + }; + + assert!(EncodedMessage::new(MsgKind::Request, 7, b"x", external_values).is_none()); + + let mut external_values = ExternalValues { + chunks: Vec::new(), + total_len: usize::MAX, + }; + assert!(external_values.push_bytes(b"x").is_err()); + assert!(external_values.chunks.is_empty()); + + let chunks = [Bytes::from_static(b"x")]; + assert!(external_values.push_chunks(&chunks).is_err()); + assert!(external_values.chunks.is_empty()); + } + + #[test] + fn size_prefix_helpers_validate_length() { + assert_eq!(size_prefix_payload_len(&4u32.to_le_bytes()), Some(4)); + assert_eq!(size_prefix_payload_len(&[0; 3]), None); + assert_eq!(size_prefixed_len(4), Some(SIZE_PREFIX_LEN + 4)); + assert_eq!(size_prefixed_len(usize::MAX), None); + } +} diff --git a/src/hyperlight_common/src/virtq/access.rs b/src/hyperlight_common/src/virtq/access.rs index 8a52fd321..1cc6b7929 100644 --- a/src/hyperlight_common/src/virtq/access.rs +++ b/src/hyperlight_common/src/virtq/access.rs @@ -11,6 +11,8 @@ use alloc::sync::Arc; use bytemuck::Pod; +use super::BufferLease; + /// Backend-provided memory access for virtqueue. /// /// # Safety @@ -29,20 +31,20 @@ use bytemuck::Pod; pub unsafe trait MemOps { type Error; - /// Read bytes from physical memory. + /// Read bytes at a backend address. /// /// Used for reading buffer contents pointed to by descriptors. /// /// # Arguments /// - /// * `addr` - Guest physical address to read from + /// * `addr` - Address in the backend's memory model /// * `dst` - Destination buffer to fill /// /// Implementations must return an error if `addr` cannot be read for /// at least `dst.len()` bytes. fn read(&self, addr: u64, dst: &mut [u8]) -> Result<(), Self::Error>; - /// Write bytes to physical memory. + /// Write bytes at a backend address. /// /// # Arguments /// @@ -74,8 +76,6 @@ pub unsafe trait MemOps { /// - The memory region is not concurrently modified for the lifetime of /// the returned slice. Caller must uphold this via protocol-level /// synchronisation, e.g. descriptor ownership transfer. - /// - /// See also [`BufferOwner`]: super::BufferOwner unsafe fn as_slice(&self, addr: u64, len: usize) -> Result<&[u8], Self::Error>; /// Get a direct mutable slice into shared memory. @@ -114,6 +114,31 @@ pub unsafe trait MemOps { } } +/// Owned immutable views of completed queue buffers. +pub trait BufferMap: MemOps { + /// A complete owner exposing exactly the initialized prefix. + /// + /// Its bytes must stay valid at the same address until the mapping drops, + /// including when the mapping is moved. A borrowed view must retain its + /// lease and release the view before returning the slot. + type Mapping: AsRef<[u8]> + Send + 'static; + + /// Retain a view of the first `written` bytes of an allocation. + /// + /// Implementations must release the lease on error. + /// + /// # Safety + /// + /// The first `written` bytes of the leased allocation are initialized, + /// and `written` does not exceed its capacity. No peer may write or reuse + /// the allocation while the lease survives. + unsafe fn map_buffer( + &self, + lease: BufferLease, + written: usize, + ) -> Result; +} + // SAFETY: Arc delegates all memory operations to the wrapped backend, preserving // that backend's MemOps contract. unsafe impl MemOps for Arc { @@ -136,11 +161,13 @@ unsafe impl MemOps for Arc { } unsafe fn as_slice(&self, addr: u64, len: usize) -> Result<&[u8], Self::Error> { + // SAFETY: The caller supplies the wrapped backend's slice preconditions. unsafe { (**self).as_slice(addr, len) } } #[allow(clippy::mut_from_ref)] unsafe fn as_mut_slice(&self, addr: u64, len: usize) -> Result<&mut [u8], Self::Error> { + // SAFETY: The caller supplies the wrapped backend's exclusive-access preconditions. unsafe { (**self).as_mut_slice(addr, len) } } } diff --git a/src/hyperlight_common/src/virtq/buffer.rs b/src/hyperlight_common/src/virtq/buffer.rs index 866a80ddd..3eb0d10b7 100644 --- a/src/hyperlight_common/src/virtq/buffer.rs +++ b/src/hyperlight_common/src/virtq/buffer.rs @@ -1,131 +1,14 @@ // 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>; - - /// 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 { - 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 reset(&self) { - (**self).reset() - } - 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 reset(&self) { - (**self).reset() - } - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - (**self).alloc_sg(total_len) - } -} +use super::{Allocation, SlotPool}; /// Ordered byte segments that make up one virtqueue payload. /// @@ -174,6 +57,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()) @@ -203,6 +114,11 @@ impl Segments { } } + /// Consume this payload without flattening its segments. + pub fn into_chunks(self) -> Vec { + self.0.into_vec() + } + fn collect(&self, sgs: &[Bytes], len: usize) -> Bytes { let mut out = Vec::with_capacity(len); out.extend(sgs.iter().flat_map(|seg| seg.iter().copied())); @@ -277,95 +193,31 @@ impl Buf for SegmentsBuf<'_> { } } -/// The owner of a mapped buffer, ensuring its lifetime. -/// -/// Holds an [`OwnedAlloc`] and provides direct access to the underlying -/// shared memory via [`MemOps::as_slice`]. Implements `AsRef<[u8]>` so it -/// can be used with [`Bytes::from_owner`](bytes::Bytes::from_owner) for -/// zero-copy `Bytes` backed by shared memory. -/// -/// When dropped, the allocation is returned to the pool. -#[derive(Debug)] -pub struct BufferOwner { - pub(crate) mem: M, - pub(crate) alloc: OwnedAlloc

, - pub(crate) written: usize, -} - -impl AsRef<[u8]> for BufferOwner { - fn as_ref(&self) -> &[u8] { - let alloc = self.alloc.allocation(); - let len = self.written.min(alloc.len); - // 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) } { - Ok(slice) => slice, - Err(_) => { - debug_assert!(false, "BufferOwner direct slice failed"); - &[] - } - } - } -} - -/// 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>, +/// An exclusively owned buffer allocation returned to its pool on drop. +pub struct BufferLease { + /// The pool that allocated the buffer. + pool: SlotPool, + /// The buffer's start address and full allocation capacity. + allocation: Allocation, } -#[derive(Debug)] -struct Inner { - pool: P, - alloc: Allocation, -} - -impl OwnedAlloc

{ - /// Wrap an existing allocation with its owning pool. - pub fn new(pool: P, alloc: Allocation) -> Self { - Self { - inner: Some(Inner { pool, alloc }), - } - } - - /// Allocate from `pool` and return an owning guard. - pub fn allocate(pool: P, len: usize) -> Result { - let alloc = pool.alloc(len)?; - Ok(Self::new(pool, alloc)) +impl BufferLease { + /// Create a new buffer lease from a pool and allocation. + pub fn new(pool: SlotPool, allocation: Allocation) -> Self { + Self { pool, allocation } } - /// 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)] + /// The buffer's start address and full allocation capacity. 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.allocation } } -impl Drop for OwnedAlloc

{ +impl Drop for BufferLease { 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:?}"); + if let Err(error) = self.pool.dealloc(self.allocation.addr) { + log::error!("Failed to release a virtqueue buffer: {error}"); + debug_assert!(false, "BufferLease deallocation failed: {error}"); } } } @@ -375,6 +227,25 @@ mod tests { use bytes::Buf; use super::*; + use crate::virtq::{SlotLayout, SlotPool}; + + #[test] + fn lease_returns_slot_on_drop() { + let pool = SlotPool::new(SlotLayout::new(0, 4, 1)).unwrap(); + let allocation = pool.alloc(4).unwrap(); + let lease = BufferLease::new(pool.clone(), allocation); + + assert_eq!(lease.allocation().addr, allocation.addr); + assert_eq!(lease.allocation().len, 4); + assert_eq!(pool.num_free(), 0); + + drop(lease); + assert_eq!(pool.num_free(), 1); + + let reused = pool.alloc(4).unwrap(); + assert_eq!(reused.addr, allocation.addr); + pool.dealloc(reused.addr).unwrap(); + } #[test] fn segments_cursor_advances_across_segments() { @@ -459,6 +330,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]); @@ -469,4 +366,18 @@ mod tests { assert_eq!(collected.as_ptr(), ptr); assert_eq!(collected.as_ref(), &[1, 2, 3, 4]); } + + #[test] + fn segments_into_chunks_preserves_segment_storage() { + let first = Bytes::from(vec![1, 2]); + let second = Bytes::from(vec![3, 4]); + let first_ptr = first.as_ptr(); + let second_ptr = second.as_ptr(); + + let chunks = Segments::new([first, second]).into_chunks(); + + assert_eq!(chunks.len(), 2); + assert_eq!(chunks[0].as_ptr(), first_ptr); + assert_eq!(chunks[1].as_ptr(), second_ptr); + } } diff --git a/src/hyperlight_common/src/virtq/concurrency.rs b/src/hyperlight_common/src/virtq/concurrency.rs index 8592b9e66..60d392ef1 100644 --- a/src/hyperlight_common/src/virtq/concurrency.rs +++ b/src/hyperlight_common/src/virtq/concurrency.rs @@ -41,6 +41,7 @@ use alloc::collections::BTreeMap; use alloc::sync::Arc; use alloc::vec; +use core::mem::ManuallyDrop; use core::num::NonZeroU16; use bytemuck::Zeroable; @@ -49,7 +50,6 @@ use loom::thread; use super::*; use crate::virtq::desc::Descriptor; -use crate::virtq::pool::BufferPoolSync; #[derive(Debug)] pub struct MemErr; @@ -289,6 +289,76 @@ unsafe impl MemOps for LoomMem { } } +pub struct LoomMapping { + creator: thread::ThreadId, + owner: ManuallyDrop, +} + +struct LoomBufferOwner { + view: loom::cell::ConstPtr>, + offset: usize, + written: usize, + _mem: Arc, + _lease: BufferLease, +} + +// SAFETY: The view is immutable. Drop checks the creator thread before +// destroying the owner containing the Rc-backed lease. +unsafe impl Send for LoomMapping {} + +impl AsRef<[u8]> for LoomMapping { + fn as_ref(&self) -> &[u8] { + // SAFETY: Construction checks this initialized range. The owned read + // guard excludes writes while the backing is borrowed. + self.owner.view.with(|buf| unsafe { + &(&*buf)[self.owner.offset..self.owner.offset + self.owner.written] + }) + } +} + +impl Drop for LoomMapping { + fn drop(&mut self) { + assert_eq!( + self.creator, + thread::current().id(), + "mapping dropped on another thread" + ); + // SAFETY: The creator thread releases the read guard before its lease. + unsafe { ManuallyDrop::drop(&mut self.owner) }; + } +} + +impl BufferMap for Arc { + type Mapping = LoomMapping; + + unsafe fn map_buffer( + &self, + lease: BufferLease, + written: usize, + ) -> Result { + let allocation = lease.allocation(); + let (info, offset) = self.region(allocation.addr).ok_or(MemErr)?; + + if !matches!(info.kind, RegionKind::Pool) + || written > allocation.len as usize + || offset.checked_add(allocation.len as usize).ok_or(MemErr)? > info.size + { + return Err(MemErr); + } + + Ok(LoomMapping { + creator: thread::current().id(), + owner: ManuallyDrop::new(LoomBufferOwner { + view: self.pool.get(), + offset, + written, + _mem: self.clone(), + _lease: lease, + }), + }) + } +} + #[derive(Debug)] pub struct Notify { kicks: AtomicUsize, @@ -316,13 +386,13 @@ 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 notify = Arc::new(Notify::new()); - let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); let mut cons = VirtqConsumer::new(mem.layout(), mem.clone(), notify.clone()); let t_prod = thread::spawn(move || { + let pool = SlotPool::new(SlotLayout::new(pool_base, 256, pool_size / 256)).unwrap(); + let mut prod = VirtqProducer::new(mem.layout(), mem, notify, pool); let mut se = prod.chain().readable(4).writable(32).build().unwrap(); se.write_all(b"ping").unwrap(); let tok = prod.submit(se).unwrap(); @@ -343,12 +413,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(); @@ -364,13 +434,13 @@ 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 notify = Arc::new(Notify::new()); - let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); let mut cons = VirtqConsumer::new(mem.layout(), mem.clone(), notify.clone()); let t_prod = thread::spawn(move || { + let pool = SlotPool::new(SlotLayout::new(pool_base, 256, pool_size / 256)).unwrap(); + let mut prod = VirtqProducer::new(mem.layout(), mem, notify, pool); let mut se = prod.chain().readable(4).build().unwrap(); se.write_all(b"ping").unwrap(); let tok = prod.submit(se).unwrap(); @@ -390,9 +460,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(); @@ -408,44 +478,48 @@ 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 notify = Arc::new(Notify::new()); - let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); let mut cons = VirtqConsumer::new(mem.layout(), mem.clone(), notify.clone()); let submitted = Arc::new(AtomicUsize::new(0)); let submitted_for_consumer = submitted.clone(); - let t_prod = thread::spawn(move || { - let mut first = prod.chain().readable(5).writable(8).build().unwrap(); - first.write_all(b"first").unwrap(); - let tok1 = prod.submit(first).unwrap(); - - let mut second = prod.chain().readable(6).writable(8).build().unwrap(); - second.write_all(b"second").unwrap(); - let tok2 = prod.submit(second).unwrap(); - submitted.store(1, Ordering::Release); - - let mut got_first = false; - let mut got_second = false; - while !(got_first && got_second) { - if let Some(r) = prod.poll().unwrap() { - let token = r.token(); - let bytes = r.to_bytes().unwrap(); - if token == tok1 { - assert!(bytes.is_empty()); - got_first = true; - } else if token == tok2 { - assert!(bytes.is_empty()); - got_second = true; + // Pool construction exceeds Loom's small default stack in this test. + let t_prod = thread::Builder::new() + .stack_size(256 * 1024) + .spawn(move || { + let pool = SlotPool::new(SlotLayout::new(pool_base, 256, pool_size / 256)).unwrap(); + let mut prod = VirtqProducer::new(mem.layout(), mem, notify, pool); + let mut first = prod.chain().readable(5).writable(8).build().unwrap(); + first.write_all(b"first").unwrap(); + let tok1 = prod.submit(first).unwrap(); + + let mut second = prod.chain().readable(6).writable(8).build().unwrap(); + second.write_all(b"second").unwrap(); + let tok2 = prod.submit(second).unwrap(); + submitted.store(1, Ordering::Release); + + let mut got_first = false; + let mut got_second = false; + while !(got_first && got_second) { + if let Some(r) = prod.poll().unwrap() { + let token = r.token(); + let bytes = r.to_bytes().unwrap(); + if token == tok1 { + assert!(bytes.is_empty()); + got_first = true; + } else if token == tok2 { + assert!(bytes.is_empty()); + got_second = true; + } else { + panic!("unexpected token"); + } } else { - panic!("unexpected token"); + thread::yield_now(); } - } else { - thread::yield_now(); } - } - }); + }) + .unwrap(); let t_cons = thread::spawn(move || { while submitted_for_consumer.load(Ordering::Acquire) == 0 { @@ -458,7 +532,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 +540,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(); @@ -500,10 +574,8 @@ 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 notify = Arc::new(Notify::new()); - let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); let mut cons = VirtqConsumer::new(mem.layout(), mem.clone(), notify.clone()); // Descriptor-mode suppression writes the `off_wrap` field that the @@ -516,12 +588,15 @@ fn virtq_event_suppression_reconfig() { }); let t_prod = thread::spawn(move || { + let pool = SlotPool::new(SlotLayout::new(pool_base, 256, pool_size / 256)).unwrap(); + let mut prod = VirtqProducer::new(mem.layout(), mem, notify, pool); let mut se = prod.chain().readable(4).build().unwrap(); se.write_all(b"ping").unwrap(); prod.submit(se).unwrap(); + t_cons.join().unwrap(); + prod.reset().unwrap(); }); - t_cons.join().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..f2380fdc1 100644 --- a/src/hyperlight_common/src/virtq/consumer.rs +++ b/src/hyperlight_common/src/virtq/consumer.rs @@ -2,53 +2,169 @@ // Copyright 2026 The Hyperlight Authors. use alloc::vec; +use core::fmt; use bytes::Bytes; -use fixedbitset::FixedBitSet; 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 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 { +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 + } + + /// 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() + } + + /// Read bytes sequentially across descriptor boundaries. + /// + /// 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, taking ownership of the segments. - pub fn into_segments(self) -> Segments { - self.segments + /// 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) } - /// Return the chain payload as contiguous bytes. + /// Copy the complete payload into descriptor-preserving owned segments. /// - /// Returns empty [`Bytes`] when the chain has no readable buffers. - pub fn to_bytes(&self) -> Bytes { - self.segments.to_bytes() + /// 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)) } - /// Consume the chain and return the payload as contiguous bytes. - pub fn into_bytes(self) -> Bytes { - self.segments.into_bytes() + /// 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 +178,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 +195,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 +222,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 +267,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 +309,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 +325,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 +346,7 @@ impl AckChain { Self { token } } + #[inline] pub fn token(&self) -> Token { self.token } @@ -221,35 +363,36 @@ 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 { inner: RingConsumer, + mem: M, notifier: N, - inflight: FixedBitSet, next_token: u32, } @@ -262,42 +405,46 @@ 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); - let inflight = FixedBitSet::with_capacity(inner.len()); + 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); Self { inner, + mem: buf_mem, notifier, - inflight, next_token: 0, } } /// 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 + /// - [`VirtqError::RingError`] - Invalid ring state or memory access failure + #[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), @@ -314,17 +461,6 @@ impl VirtqConsumer { .iter() .fold(0usize, |acc, elem| acc.saturating_add(elem.len as usize)); - // Reserve the inflight slot - let id_idx = id as usize; - if id_idx >= self.inflight.len() { - return Err(VirtqError::InvalidState); - } - - if self.inflight.contains(id_idx) { - return Err(VirtqError::InvalidState); - } - - self.inflight.insert(id_idx); let token = Token { seq: self.next_token, id, @@ -341,20 +477,17 @@ 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.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 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); @@ -364,24 +497,31 @@ 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. /// - /// 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> { + /// 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. + /// + /// 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(); - let id = reply.token().id; - let written = u32::try_from(reply.written()).map_err(|_| VirtqError::ReplyTooLarge)?; - let id_idx = id as usize; - let slot_set = id_idx < self.inflight.len() && self.inflight.contains(id_idx); - if !slot_set { + if recv.token() != reply.token() { return Err(VirtqError::InvalidState); } - self.inflight.set(id_idx, false); + let id = reply.token().id; + let written = u32::try_from(reply.written()).map_err(|_| VirtqError::ReplyTooLarge)?; if self.inner.submit_used_with_notify(id, written)? { self.notifier.notify(QueueStats { @@ -398,11 +538,6 @@ impl VirtqConsumer { /// The ring's `poll_available` removes the descriptor from the available /// ring before [`poll`](Self::poll) validates the chain. fn abort_chain(&mut self, id: u16, err: VirtqError) -> VirtqError { - let id_idx = id as usize; - if id_idx < self.inflight.len() { - self.inflight.set(id_idx, false); - } - // Best effort: failing to return the descriptor means the ring is // already in an unrecoverable state, so surface the original error. if let Ok(true) = self.inner.submit_used_with_notify(id, 0) { @@ -462,63 +597,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 + /// - [`VirtqError::RingError`] - device-event normalization failed + pub fn reset(&mut self) -> Result<(), VirtqError> { + if self.inner.num_inflight() != 0 { + return Err(VirtqError::InvalidState); } - Ok(Segments::from_smallvec(segments)) + self.inner.reset()?; + Ok(()) } +} - /// Reset ring and inflight state to initial values. - pub fn reset(&mut self) { - self.inner.reset(); - self.inflight.clear(); - } +type ChainElems = SmallVec<[BufferElement; 4]>; + +struct ChainState { + mem: M, + token: Token, + elems: ChainElems, + total: usize, + position: usize, + desc_idx: usize, + desc_off: usize, } -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; +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 + } - for elem in elems { - if src.is_empty() { - break; - } + #[inline] + fn token(&self) -> Token { + self.token + } - let elem_len = elem.len as usize; - if skip >= elem_len { - skip -= elem_len; - continue; - } + #[inline] + fn total(&self) -> usize { + self.total + } + + #[inline] + fn position(&self) -> usize { + self.position + } - 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 remaining(&self) -> usize { + self.total - self.position + } - mem.write(addr, &src[..n])?; + #[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 +719,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,13 +781,14 @@ 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(); } + producer.reset().unwrap(); } #[test] @@ -573,23 +801,24 @@ 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(); + producer.reset().unwrap(); } #[test] fn test_readwrite_round_trip() { let ring = make_ring(16); - let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + let (mut producer, mut consumer, _notifier) = make_test_producer_with_slot_size(&ring, 64); let mut se = producer.chain().readable(32).writable(64).build().unwrap(); se.write_all(b"hello world").unwrap(); 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,36 +827,112 @@ 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"); } + producer.reset().unwrap(); } #[test] - fn test_writable_partial_write() { + 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(); + producer.reset().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 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_split(ring.layout(), ring_mem, 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); + let (mut producer, mut consumer, _notifier) = make_test_producer_with_slot_size(&ring, 8); + 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"); } + producer.reset().unwrap(); } #[test] fn test_writable_write_all_too_large() { let ring = make_ring(16); - let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + let (mut producer, mut consumer, _notifier) = make_test_producer_with_slot_size(&ring, 4); let se = producer.chain().writable(4).build().unwrap(); producer.submit(se).unwrap(); @@ -639,6 +944,7 @@ mod tests { } else { panic!("expected Writable"); } + producer.reset().unwrap(); } #[test] @@ -654,6 +960,7 @@ mod tests { consumer.poll(4), Err(VirtqError::PayloadTooLarge { recv: 8, limit: 4 }) )); + producer.reset().unwrap(); } #[test] @@ -679,8 +986,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()); } @@ -698,7 +1005,6 @@ mod tests { consumer.poll(1024), Err(VirtqError::RingError(RingError::BadChain)) )); - assert_eq!(consumer.inflight.count_ones(..), 0); assert_eq!(consumer.inner.num_inflight(), 0); } @@ -720,7 +1026,6 @@ mod tests { consumer.poll(1024), Err(VirtqError::RingError(RingError::BadChain)) )); - assert_eq!(consumer.inflight.count_ones(..), 0); assert_eq!(consumer.inner.num_inflight(), 0); } @@ -731,13 +1036,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"); @@ -746,12 +1051,12 @@ mod tests { #[test] fn test_writable_rewind() { let ring = make_ring(16); - let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + let (mut producer, mut consumer, _notifier) = make_test_producer_with_slot_size(&ring, 16); 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,10 +1066,11 @@ 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"); } + producer.reset().unwrap(); } #[test] @@ -783,7 +1089,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 +1097,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 +1111,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 +1158,45 @@ 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(); + producer.reset().unwrap(); + } + + #[test] + 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.inner.num_inflight(), 2); + assert!(producer.poll().unwrap().is_none()); + assert!(matches!(consumer.reset(), Err(VirtqError::InvalidState))); + + drop((recv2, reply1)); + producer.reset().unwrap(); } #[test] - fn test_recv_into_bytes() { + 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 +1204,15 @@ 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(); + producer.reset().unwrap(); } #[test] @@ -847,16 +1224,17 @@ mod tests { let se = producer.chain().writable(16).build().unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_data(&mut consumer); - assert!(consumer.inflight.count_ones(..) > 0); + let (recv, reply) = poll_data(&mut consumer); + assert!(consumer.inner.num_inflight() > 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); + producer.reset().unwrap(); } #[test] @@ -870,15 +1248,38 @@ 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); + producer.reset().unwrap(); + } + + #[test] + fn failed_completion_keeps_reset_blocked() { + use crate::virtq::ring::tests::FaultMem; + + let ring = make_ring(4); + let (mut producer, _, _) = make_test_producer(&ring); + let mem = FaultMem::new(ring.mem()); + let mut consumer = VirtqConsumer::new(ring.layout(), mem.clone(), TestNotifier::new()); + let chain = producer.chain().writable(8).build().unwrap(); + producer.submit(chain).unwrap(); + let (recv, reply) = consumer.poll(0).unwrap().unwrap(); + + mem.fail_write_at(0); + assert!(matches!( + consumer.complete(recv, reply), + Err(VirtqError::RingError(RingError::MemError { .. })) + )); + mem.allow_writes(); + assert_eq!(consumer.inner.num_inflight(), 1); + assert!(matches!(consumer.reset(), Err(VirtqError::InvalidState))); + producer.reset().unwrap(); } } 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/mod.rs b/src/hyperlight_common/src/virtq/mod.rs index 02b6af120..443e42488 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)?; //! } //! ``` //! @@ -151,7 +152,6 @@ mod buffer; mod consumer; mod desc; mod event; -pub mod msg; mod pool; mod producer; mod ring; @@ -171,6 +171,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); @@ -187,10 +194,14 @@ pub enum VirtqError { Backpressure, #[error("Allocation exceeds pool capacity")] OutOfMemory, + #[error("Failed to allocate virtqueue bookkeeping")] + Bookkeeping, #[error("Invalid chain received")] 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")] @@ -225,6 +236,7 @@ impl From for VirtqError { match e { AllocError::NoSpace => Self::Backpressure, AllocError::OutOfMemory => Self::OutOfMemory, + AllocError::Bookkeeping => Self::Bookkeeping, other => Self::Alloc(other), } } @@ -379,7 +391,7 @@ impl From for Allocation { fn from(value: BufferElement) -> Self { Allocation { addr: value.addr, - len: value.len as usize, + len: value.len, } } } @@ -442,10 +454,8 @@ const _: () = { /// Shared test utilities for virtqueue tests. #[cfg(test)] pub(crate) mod test_utils { - use alloc::collections::BTreeMap; use alloc::sync::Arc; - use core::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; - use std::sync::Mutex; + use core::sync::atomic::{AtomicUsize, Ordering}; use super::*; use crate::virtq::ring::tests::{OwnedRing, TestMem}; @@ -474,84 +484,28 @@ pub(crate) mod test_utils { } } - /// Simple test buffer pool that allocates from a range. - #[derive(Clone)] - pub(crate) struct TestPool { - base: u64, - next: Arc, - size: usize, - max_alloc_len: usize, - allocations: Arc>>, - } - - impl TestPool { - pub(crate) fn new(base: u64, size: usize) -> Self { - Self { - 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 alloc(&self, len: usize) -> Result { - if len == 0 { - return Err(AllocError::InvalidArg); - } - - let addr = self.next.fetch_add(len as u64, Ordering::Relaxed); - let end = addr + len as u64; - if end > self.base + self.size as u64 { - return Err(AllocError::NoSpace); - } - self.allocations - .lock() - .expect("poisoned mutex") - .insert(addr, len); - Ok(Allocation { addr, len }) - } - - fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - self.allocations - .lock() - .expect("poisoned mutex") - .remove(&addr) - .map(|_| ()) - .ok_or(AllocError::InvalidFree(addr, 0)) - } - } - - type TestProducer = VirtqProducer; + type TestProducer = VirtqProducer; type TestConsumer = VirtqConsumer; /// Create test infrastructure: a producer, consumer, and notifier backed /// by the supplied [`OwnedRing`]. pub(crate) fn make_test_producer( ring: &OwnedRing, + ) -> (TestProducer, TestConsumer, TestNotifier) { + make_test_producer_with_slot_size(ring, 128) + } + + pub(crate) fn make_test_producer_with_slot_size( + ring: &OwnedRing, + slot_size: usize, ) -> (TestProducer, TestConsumer, TestNotifier) { let layout = ring.layout(); let mem = ring.mem(); // Pool needs to be in memory accessible via mem - use memory after ring layout let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; - let pool = TestPool::new(pool_base, 0x8000); + let pool = + SlotPool::new(SlotLayout::new(pool_base, slot_size, 0x8000 / slot_size)).unwrap(); let notifier = TestNotifier::new(); let producer = VirtqProducer::new(layout, mem.clone(), notifier.clone(), pool); @@ -572,7 +526,7 @@ mod tests { /// Helper: build and submit a readable+writable chain using the chain() builder. fn send_readwrite( - producer: &mut VirtqProducer, + producer: &mut VirtqProducer, entry_data: &[u8], used_cap: usize, ) -> Token { @@ -588,7 +542,7 @@ mod tests { fn poll_received( consumer: &mut VirtqConsumer, - ) -> (RecvChain, ReplyChain) { + ) -> (RecvChain, ReplyChain) { consumer.poll(1024).unwrap().unwrap() } @@ -604,6 +558,7 @@ mod tests { let (recv, _reply) = poll_received(&mut consumer); assert_eq!(recv.token(), token); + producer.reset().unwrap(); } #[test] @@ -617,8 +572,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 +605,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 @@ -694,7 +649,7 @@ mod tests { 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(pool_base, 0x8000); + let pool = SlotPool::new(SlotLayout::new(pool_base, 128, 0x8000 / 128)).unwrap(); let notifier = CtxNotifier { last_num_free: Arc::new(AtomicUsize::new(0)), last_num_inflight: Arc::new(AtomicUsize::new(0)), @@ -708,6 +663,7 @@ mod tests { producer.submit(se).unwrap(); assert_eq!(notifier.count.load(Ordering::Relaxed), 1); assert!(notifier.last_num_inflight.load(Ordering::Relaxed) > 0); + producer.reset().unwrap(); } #[test] @@ -736,19 +692,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 +727,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 +750,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 +772,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 +791,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 +815,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(); @@ -894,7 +850,7 @@ mod tests { /// Helper: submit a read-only chain (readable data, no writable reply). fn send_readonly( - producer: &mut VirtqProducer, + producer: &mut VirtqProducer, entry_data: &[u8], ) -> Token { let mut se = producer.chain().readable(entry_data.len()).build().unwrap(); @@ -912,11 +868,9 @@ mod tests { send_readonly(&mut producer, b"b"); send_readonly(&mut producer, b"c"); send_readonly(&mut producer, b"d"); + assert_eq!(producer.num_inflight(), 4); - // Ring is now full - next submit should fail with Backpressure - let mut se = producer.chain().readable(1).build().unwrap(); - se.write_all(b"e").unwrap(); - let res = producer.submit(se); + let res = producer.chain().readable(1).build(); assert!( matches!(res, Err(VirtqError::Backpressure)), "expected Backpressure from full ring" @@ -924,16 +878,18 @@ 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 let count = producer.reclaim().unwrap(); assert_eq!(count, 4, "expected 4 reclaimed entries"); + assert_eq!(producer.num_inflight(), 0); // Ring should have space now send_readonly(&mut producer, b"e"); + producer.reset().unwrap(); } #[test] @@ -945,12 +901,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 +929,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 +964,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 +997,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 +1013,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 +1043,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/msg.rs b/src/hyperlight_common/src/virtq/msg.rs deleted file mode 100644 index bac9da90c..000000000 --- a/src/hyperlight_common/src/virtq/msg.rs +++ /dev/null @@ -1,107 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// Copyright 2026 The Hyperlight Authors. - -//! Wire format header for all virtqueue messages. -//! -//! Every payload on both the G2H and H2G queues starts with this -//! fixed 8-byte header, enabling message type discrimination and -//! request/response correlation. - -use bitflags::bitflags; - -/// Message types for the virtqueue wire protocol. -#[repr(u8)] -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum MsgKind { - /// A function call request (FunctionCall payload follows). - Request = 0x01, - /// A function call response (FunctionCallResult payload follows). - Response = 0x02, - /// A stream data chunk. - StreamChunk = 0x03, - /// End-of-stream marker. - StreamEnd = 0x04, - /// Cancel a pending request. - Cancel = 0x05, - /// A guest log message (GuestLogData payload follows). - Log = 0x06, -} - -impl TryFrom for MsgKind { - type Error = u8; - - fn try_from(value: u8) -> Result { - match value { - 0x01 => Ok(Self::Request), - 0x02 => Ok(Self::Response), - 0x03 => Ok(Self::StreamChunk), - 0x04 => Ok(Self::StreamEnd), - 0x05 => Ok(Self::Cancel), - 0x06 => Ok(Self::Log), - other => Err(other), - } - } -} - -bitflags! { - #[repr(transparent)] - #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] - pub struct MsgFlags: u8 { - /// More descriptors follow for this message. - const MORE = 1 << 0; - } -} - -/// Wire header for all virtqueue messages -#[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)] -#[repr(C)] -pub struct VirtqMsgHeader { - /// Discriminates the message type. - pub kind: u8, - /// Per-message flags (see [`MsgFlags`]). - pub flags: u8, - /// Caller-assigned correlation ID. Responses echo the request's ID. - pub req_id: u16, - /// Byte length of the payload following this header in this descriptor. - pub payload_len: u32, -} - -impl VirtqMsgHeader { - pub const SIZE: usize = core::mem::size_of::(); - - /// Create a new message header with no flags set. - pub const fn new(kind: MsgKind, req_id: u16, payload_len: u32) -> Self { - Self { - kind: kind as u8, - flags: 0, - req_id, - payload_len, - } - } - - /// Create a new header with flags. - pub const fn with_flags(kind: MsgKind, flags: MsgFlags, req_id: u16, payload_len: u32) -> Self { - Self { - kind: kind as u8, - flags: flags.bits(), - req_id, - payload_len, - } - } - - /// Parse the kind field into a [`MsgKind`] enum. - pub fn msg_kind(&self) -> Result { - MsgKind::try_from(self.kind) - } - - /// Interpret the raw flags field as [`MsgFlags`]. - pub fn msg_flags(&self) -> MsgFlags { - MsgFlags::from_bits_truncate(self.flags) - } - - /// Returns true if [`MsgFlags::MORE`] is set, indicating more - /// descriptors follow for this message. - pub const fn has_more(&self) -> bool { - self.flags & MsgFlags::MORE.bits() != 0 - } -} diff --git a/src/hyperlight_common/src/virtq/pool.rs b/src/hyperlight_common/src/virtq/pool.rs index af5b63425..0422e5393 100644 --- a/src/hyperlight_common/src/virtq/pool.rs +++ b/src/hyperlight_common/src/virtq/pool.rs @@ -1,1338 +1,50 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright 2026 The Hyperlight Authors. -//! Buffer pool implementations for virtqueue buffer management. -//! -//! 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. - -use alloc::rc::Rc; -use core::cell::RefCell; -use core::ops::Deref; - -use fixedbitset::FixedBitSet; -use smallvec::SmallVec; - -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; - - 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 - }) - } - - 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) - } - - 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) - } - - fn reset(&mut self) { - self.used_slots.clear(); - self.run_starts.clear(); - self.last_free_run = None; - } +//! Fixed-slot allocation for virtqueue payloads. + +use thiserror::Error; + +mod slot; + +pub use slot::{SlotLayout, SlotPool}; + +/// Buffer allocation failure. +#[derive(Debug, Error, Copy, Clone)] +pub enum AllocError { + /// 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, + /// Allocation bookkeeping could not be reserved. + #[error("Failed to allocate buffer bookkeeping")] + Bookkeeping, + /// Address or size arithmetic overflowed. + #[error("Overflow")] + Overflow, +} + +/// One pool allocation. +#[derive(Debug, Clone, Copy)] +pub struct Allocation { + /// Starting address of the allocation. + pub addr: u64, + /// Nonzero descriptor-safe capacity in bytes. + pub len: u32, } #[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 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), - } - } - - // 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) - } - } -} - -impl BufferProvider for BufferPool { - 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) - } - - fn reset(&self) { - let mut inner = self.inner.borrow_mut(); - inner.lower.reset(); - inner.upper.reset(); - } -} - -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) - } -} - -#[cfg(all(test, loom))] -impl BufferProvider for BufferPoolSync { - 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) - } -} - -/// 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 in an indeterminate state and should be - /// [`reset`](Self::reset) before reuse. - fn restore_allocated(&mut self, allocated: &[u64]) -> Result<(), AllocError> { - self.allocated.clear(); - for &addr in allocated { - let slot = self.slot_of(addr)?; - if self.allocated.contains(slot) { - return Err(AllocError::InvalidFree(addr, self.slot_size)); - } - self.allocated.insert(slot); - } - 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(); - 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) - } - - fn num_free(&self) -> usize { - self.free.len() - } -} - -/// A recycling buffer provider with fixed-size slots. -/// -/// 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. - 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 - } - - /// Number of slots in the pool. - pub fn count(&self) -> usize { - self.inner.borrow().count - } -} - -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) - } - - fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) - } - - fn reset(&self) { - self.inner.borrow_mut().reset() - } -} +mod tests; #[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_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); - 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); - let result = pool.restore_allocated(&[0xDEAD]); - assert!(result.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_resets_first() { - 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); - } - - #[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); - } - - #[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); - } -} - -#[cfg(test)] -mod fuzz { - use quickcheck::{Arbitrary, Gen, QuickCheck}; - - 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); - } -} +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..af156b31f --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/fuzz.rs @@ -0,0 +1,235 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +use std::collections::{BTreeMap, HashSet}; + +use quickcheck::{Arbitrary, Gen, QuickCheck}; + +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), + Dealloc(usize), +} + +impl Arbitrary for Op { + fn arbitrary(g: &mut Gen) -> Self { + match u8::arbitrary(g) % 2 { + 0 => Op::Alloc(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), + 1 => Op::Dealloc(usize::arbitrary(g)), + _ => unreachable!(), + } + } +} + +#[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); + let mut allocations: Vec = Vec::new(); + + for op in &scenario.ops { + match op { + Op::Alloc(size) => match pool.alloc(*size) { + Ok(allocation) => { + if (allocation.len as usize) < *size + || allocations + .iter() + .any(|existing| existing.addr == allocation.addr) + { + return false; + } + allocations.push(allocation); + } + 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_slot_pool_invariants(&pool, &allocations).is_err() { + return false; + } + } + + while let Some(alloc) = allocations.pop() { + if pool.dealloc(alloc.addr).is_err() { + return false; + } + } + + check_slot_pool_invariants(&pool, &allocations).is_ok() +} + +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"); + } + } + + // Live addresses must match the order. Free plus live must cover every slot. + 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"); + } + + // Free-slot enumeration must match the reported count and be strictly ordered. + let mut free = Vec::new(); + pool.for_each_free(|allocation| free.push(allocation)); + + if free.len() != pool.num_free() { + return Err("free-slot visitation is inconsistent"); + } + + if free.windows(2).any(|pair| pair[0].addr >= pair[1].addr) { + return Err("free-slot visitation is inconsistent"); + } + + // Free slots cannot also be live and must report their tier's full capacity. + if free.iter().any(|allocation| { + expected_live.contains_key(&allocation.addr) + || slot_capacity(pool, allocation.addr) != Some(allocation.len as usize) + }) { + return Err("free-slot visitation is inconsistent"); + } + + // Reported geometry must agree with the stored tier layouts. + 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"); + } + + // Distinct tiers need increasing slot sizes and ordered, non-overlapping ranges. + 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"); + } + + // Indexed slots must be unique and inside a tier. Only live slots may report allocation lengths. + 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 as usize != 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_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/slot.rs b/src/hyperlight_common/src/virtq/pool/slot.rs new file mode 100644 index 000000000..361e40ffc --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/slot.rs @@ -0,0 +1,446 @@ +// 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. +//! +//! Each allocation occupies one slot. [`SlotPool::live_addrs`] reports ownership +//! in deterministic index order. + +use alloc::rc::Rc; +use alloc::vec::Vec; +use core::cell::RefCell; + +use fixedbitset::FixedBitSet; +use smallvec::SmallVec; + +use super::{AllocError, Allocation}; + +/// 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. Must fit in [`Allocation::len`]. + 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: u32, + /// 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, +} + +impl Tier { + fn from_layout(layout: SlotLayout) -> Result { + 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); + } + + 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, + count: layout.slot_count, + free, + allocated: FixedBitSet::with_capacity(layout.slot_count), + }) + } + + fn end(&self) -> u64 { + self.base_addr + self.count as u64 * u64::from(self.slot_size) + } + + 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(u64::from(self.slot_size)) { + return Err(AllocError::InvalidFree(addr, 0)); + } + + Ok((off / u64::from(self.slot_size)) 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 as usize { + 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) / u64::from(self.slot_size)) 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 as usize) + } + + fn slot_addr(&self, index: usize) -> Option { + (index < self.count).then(|| self.base_addr + (index * self.slot_size as usize) 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 usize) as u64), + ); + } + + fn for_each_free(&self, f: &mut impl FnMut(Allocation)) { + for slot in self.allocated.zeroes() { + f(Allocation { + addr: self.base_addr + slot as u64 * u64::from(self.slot_size), + len: self.slot_size, + }); + } + } + + fn layout(&self) -> SlotLayout { + SlotLayout::new(self.base_addr, self.slot_size as usize, 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)?; + + 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); + } + + 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 as usize, count); + + return Ok(Self { + lower: None, + upper: Tier::from_layout(layout)?, + }); + } + + Ok(Self { + lower: Some(lower), + upper, + }) + } + + fn max_alloc_len(&self) -> usize { + self.upper.slot_size as usize + } + + fn alloc(&mut self, len: usize) -> Result { + if let Some(lower) = &mut self.lower + && len <= lower.slot_size as usize + { + 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.num_live()); + 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 num_live(&self) -> usize { + self.count() - self.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. +#[derive(Clone)] +pub struct SlotPool { + inner: Rc>, +} + +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 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) + } + + fn from_layouts(lower: Option, upper: SlotLayout) -> Result { + let inner = Inner::new(lower, upper)?; + Ok(Self { + inner: 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() + } + + /// Visit every free slot in lower-then-upper index order. + /// + /// The callback must not allocate or free slots in this pool. + pub fn for_each_free(&self, mut f: impl FnMut(Allocation)) { + let inner = self.inner.borrow(); + if let Some(lower) = &inner.lower { + lower.for_each_free(&mut f); + } + inner.upper.for_each_free(&mut f); + } + + /// 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() + } + + /// 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(&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() + } + + /// 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() + } + + /// Allocate one slot holding at least `len` bytes. + pub fn alloc(&self, len: usize) -> Result { + self.inner.borrow_mut().alloc(len) + } + + #[cfg(test)] + pub(crate) fn strong_count(&self) -> usize { + Rc::strong_count(&self.inner) + } +} 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..e06e2afaa --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/tests.rs @@ -0,0 +1,333 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +use super::*; + +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_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_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.num_free_lower(), 0); + assert_eq!(pool.num_free_upper(), 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); + 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, 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); + 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); + 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_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); + + 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); + + 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_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(_, _)) + )); +} + +#[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] + ); +} + +#[test] +fn free_slots_include_full_capacities_and_preserve_allocation_order() { + let pool = make_tiered_slot_pool(2, 2); + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(1024).unwrap(); + let mut free = Vec::new(); + pool.for_each_free(|allocation| free.push((allocation.addr, allocation.len))); + assert_eq!(free, [(0x80000, 256), (0x90000, 4096)]); + assert_eq!(pool.live_addrs(), [lower.addr, upper.addr]); + + pool.dealloc(lower.addr).unwrap(); + let repeated = pool.alloc(128).unwrap(); + assert_eq!(repeated.addr, lower.addr); + pool.dealloc(repeated.addr).unwrap(); + pool.dealloc(upper.addr).unwrap(); + free.clear(); + pool.for_each_free(|allocation| free.push((allocation.addr, allocation.len))); + assert_eq!( + free, + [ + (0x80000, 256), + (0x80100, 256), + (0x90000, 4096), + (0x91000, 4096) + ] + ); +} + +#[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_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(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 04f59e0ab..145ac83a4 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -2,7 +2,9 @@ // Copyright 2026 The Hyperlight Authors. use alloc::collections::VecDeque; +use alloc::vec; use alloc::vec::Vec; +use core::mem::ManuallyDrop; use bytes::Bytes; use smallvec::SmallVec; @@ -13,9 +15,8 @@ use super::*; /// /// Read-only chains are returned as [`Ack`](Self::Ack). Chains with a writable /// buffer complete as [`Data`](Self::Data), even when the device wrote zero -/// bytes. Non-empty segments in [`Data`](Self::Data) are backed by -/// shared-memory pool allocations that are returned when the last clone is -/// dropped. +/// bytes. Non-empty segments use backend-owned mappings. Borrowed mappings +/// keep their pool slots allocated until the last [`Bytes`] owner drops. #[derive(Debug)] pub enum UsedChain { /// Acknowledgement for a read-only/fire-and-forget chain. @@ -69,15 +70,91 @@ impl UsedChain { } } -/// Allocation tracking for an in-flight descriptor chain. -/// -/// Descriptor lengths have already been published to the ring, so in-flight -/// state only needs the completion token and allocation ownership for later -/// reclaim. -#[derive(Debug)] -pub(crate) struct Inflight { +struct Inflight { token: Token, - chain: BufferChain, + // Automatic cleanup could recycle buffers still accessible to the peer. + // Only completion or a stopped reset permits releasing this owner. + chain: ManuallyDrop>, +} + +/// 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::Bookkeeping) + } + + 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). @@ -87,11 +164,15 @@ pub(crate) struct Inflight { /// /// # Threading /// -/// The producer is intended for single-threaded, guest-side use. Reply payloads -/// are exposed as zero-copy [`Bytes`] via [`Bytes::from_owner`](bytes::Bytes::from_owner), -/// which requires the owning pool to be `Send`. Do not move the producer -/// or its replies across threads, and do not instantiate it on the multi-threaded -/// host with those pools. +/// The producer and its pool use single-threaded allocation state. +/// [`BufferMap`] supplies the complete `Send` owner for returned [`Bytes`]. +/// +/// # Cleanup +/// +/// Dropping a producer does not stop its peer, so inflight owners are +/// deliberately leaked to avoid freeing buffers the peer may still use. +/// Drain completions for all submitted chains, or stop the peer and call +/// [`reset`](Self::reset), before dropping the producer. /// /// # Example /// @@ -112,20 +193,19 @@ pub(crate) struct Inflight { /// } /// } /// ``` -pub struct VirtqProducer { +pub struct VirtqProducer { inner: RingProducer, notifier: N, - pool: P, + pool: SlotPool, next_token: u32, - inflight: Vec>, + inflight: InflightTable, pending: VecDeque, } -impl VirtqProducer +impl VirtqProducer where M: MemOps + Clone, N: Notifier, - P: BufferProvider + Clone, { /// Create a new virtqueue producer. /// @@ -135,45 +215,46 @@ where /// * `mem` - Memory operations implementation for reading/writing to shared memory /// * `notifier` - Callback for notifying the device (consumer) about new chains /// * `pool` - Buffer allocator for chain payload and reply data - pub fn new(layout: Layout, mem: M, notifier: N, pool: P) -> Self { + pub fn new(layout: Layout, mem: M, notifier: N, pool: SlotPool) -> 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(), } } - fn dealloc_elems( - &self, - elems: impl IntoIterator, - ) -> Result<(), VirtqError> { - let mut first_err = None; - for elem in elems { - if let Err(err) = self.pool.dealloc(elem.addr) - && first_err.is_none() - { - first_err = Some(VirtqError::Alloc(err)); - } - } - - if let Some(err) = first_err { - return Err(err); - } + /// Borrow the pool used for new chains. + pub fn pool(&self) -> &SlotPool { + &self.pool + } - Ok(()) + /// Borrow the backend used for ring access and new chains. + pub fn memory(&self) -> &M { + self.inner.mem() } /// Begin building a descriptor chain for submission. /// - /// Returns a [`ChainBuilder`] that allocates buffers from the pool. - pub fn chain(&self) -> ChainBuilder { - ChainBuilder::new(self.inner.mem().clone(), self.pool.clone()) + /// The builder captures the current free-descriptor budget. + /// Submission checks capacity again. + pub fn chain(&self) -> ChainBuilder { + ChainBuilder::new( + self.inner.mem().clone(), + self.pool.clone(), + self.inner.num_free(), + ) + } + + /// Preferred size of one bulk payload segment. + pub fn preferred_segment_len(&self) -> usize { + self.pool.slot_size() } /// Begin a batch of submissions. @@ -182,7 +263,7 @@ where /// the ring immediately, but the consumer is notified at most once when /// [`SubmitBatch::finish`] is called. This mirrors the virtio pattern of /// adding multiple buffers and then kicking the queue once. - pub fn batch(&mut self) -> SubmitBatch<'_, M, N, P> { + pub fn batch(&mut self) -> SubmitBatch<'_, M, N> { SubmitBatch::new(self) } @@ -195,30 +276,53 @@ where /// /// # Errors /// - /// - [`VirtqError::PayloadTooLarge`] - written exceeds readable buffer capacity - /// - [`VirtqError::RingError`] - ring is full - /// - [`VirtqError::InvalidState`] - descriptor ID collision - pub fn submit(&mut self, chain: SendChain) -> Result { + /// * [`VirtqError::Backpressure`] - ring or in-flight tracking is full + /// * [`VirtqError::RingError`] - publication or notification memory access failed + /// * [`VirtqError::InvalidState`] - descriptor ID collision + pub fn submit(&mut self, chain: SendChain) -> Result { let cursor_before = self.inner.avail_cursor(); let token = self.publish(chain)?; self.notify_since(cursor_before)?; Ok(token) } - fn publish(&mut self, send: SendChain) -> Result { + fn publish(&mut self, send: SendChain) -> Result { + self.inflight.try_reserve_one()?; + + if send.desc_count() > self.inner.num_free() { + return Err(VirtqError::Backpressure); + } let token_id = self.next_token; - let id = self.inner.submit_available(send.chain())?; + let id = self.inner.next_id()?; 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); + let mut descriptors = send.owned.descriptors(); + let mut builder = BufferChainBuilder::new(); + + builder.reserve_exact(send.desc_count()); + let chain = builder + .readables(descriptors.by_ref().take(send.rd_desc_count())) + .writables(descriptors) + .build()?; + + let published = self.inner.submit_available(&chain); + + // A failed write can still publish descriptors. Keep their allocations + // until a stopped reset, including when submission reports an error. + let inf = Inflight { + token, + chain: ManuallyDrop::new(send.owned), + }; + self.inflight.insert(inf); self.next_token = self.next_token.wrapping_add(1); + let published_id = published?; + debug_assert_eq!(published_id, id); Ok(token) } @@ -260,6 +364,43 @@ 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; + + while let Some(inflight) = self.inflight.pop() { + let ret = ManuallyDrop::into_inner(inflight.chain).release(); + if let Err(err) = ret + && maybe_err.is_none() + { + maybe_err = Some(err); + } + } + + match maybe_err { + Some(error) => Err(error.into()), + None => Ok(()), + } + } + /// Configure event suppression for used buffer notifications. /// /// This controls when the device (consumer) signals us about completed buffers: @@ -289,46 +430,12 @@ 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 +impl VirtqProducer where - M: MemOps + Clone + Send + 'static, + M: BufferMap + Clone, N: Notifier, - P: BufferProvider + Clone + Send + 'static, { /// Poll for a single used chain from the device. /// @@ -338,10 +445,8 @@ where /// Returns `Ok(Some(used))` if a used chain is available, `Ok(None)` if no /// used chains are ready (would block), or an error if the device misbehaved. /// - /// Data used chains contain zero-copy [`Bytes`] backed by the shared-memory - /// allocation via [`BufferOwner`]. The pool allocation is held alive as long - /// as any `Bytes` clone exists, and is returned to the pool when the last - /// clone is dropped. + /// Data used chains contain [`Bytes`] owned by the [`BufferMap`] backend. + /// Borrowed views retain their pool allocations until the last clone drops. /// /// # Errors /// @@ -369,6 +474,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::Bookkeeping)?; self.pending.push_back(chain); } count += 1; @@ -386,81 +494,23 @@ 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; let Inflight { token, chain } = inf; - self.dealloc_elems(chain.readables().iter().copied())?; - - let used = if chain.writables().is_empty() { + let chain = ManuallyDrop::into_inner(chain); + let used = if chain.readable == chain.buffers.len() { + chain.release()?; UsedChain::Ack(token) } else { - UsedChain::Data(token, self.recv_segments(chain.writables(), written)?) + UsedChain::Data(token, chain.into_segments(written)?) }; Ok(Some(used)) } - fn recv_segments( - &self, - writables: &[BufferElement], - written: usize, - ) -> Result { - let mut owned = SmallVec::<[(BufferElement, usize); 4]>::new(); - let mut free = SmallVec::<[BufferElement; 4]>::new(); - let mut remaining = written; - - for &alloc in writables { - if remaining == 0 { - free.push(alloc); - continue; - } - - let len = remaining.min(alloc.len as usize); - owned.push((alloc, len)); - remaining -= len; - } - - if remaining != 0 { - let elems = owned.iter().map(|(elem, _)| *elem).chain(free); - self.dealloc_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); - return Err(VirtqError::MemoryReadError); - } - } - - let mut sgs = SmallVec::<[Bytes; 4]>::new(); - for (elem, written) in owned { - let alloc = OwnedAlloc::new( - self.pool.clone(), - Allocation { - addr: elem.addr, - len: elem.len as usize, - }, - ); - let mem = self.inner.mem().clone(); - let owner = BufferOwner { - alloc, - mem, - written, - }; - sgs.push(Bytes::from_owner(owner)); - } - - self.dealloc_elems(free)?; - - Ok(Segments::from_smallvec(sgs)) - } - /// Drain all available used chains, calling the provided closure for each. /// /// This is a convenience method that repeatedly calls [`poll`](Self::poll) @@ -489,21 +539,20 @@ where /// A scoped batch of producer submissions. /// /// Submissions are published immediately, while notification is delayed until -/// [`finish`](Self::finish). `finish` is explicit because the event-suppression -/// check can fail; dropping a batch does not notify. -#[must_use = "call finish to notify the consumer about batched submissions"] -pub struct SubmitBatch<'a, M, N, P> { - producer: &'a mut VirtqProducer, +/// [`finish`](Self::finish). [`finish_without_notify`](Self::finish_without_notify) +/// supports protocols whose peer is already scheduled to inspect the queue. +#[must_use = "finish the batch explicitly"] +pub struct SubmitBatch<'a, M, N> { + producer: &'a mut VirtqProducer, notify_from: Option, } -impl<'a, M, N, P> SubmitBatch<'a, M, N, P> +impl<'a, M, N> SubmitBatch<'a, M, N> where M: MemOps + Clone, N: Notifier, - P: BufferProvider + Clone, { - fn new(producer: &'a mut VirtqProducer) -> Self { + fn new(producer: &'a mut VirtqProducer) -> Self { Self { producer, notify_from: None, @@ -511,12 +560,12 @@ where } /// Begin building a descriptor chain for this batch. - pub fn chain(&self) -> ChainBuilder { + pub fn chain(&self) -> ChainBuilder { self.producer.chain() } /// Publish a chain as part of this batch without notifying yet. - pub fn submit(&mut self, chain: SendChain) -> Result { + pub fn submit(&mut self, chain: SendChain) -> Result { let cursor_before = self.producer.inner.avail_cursor(); let token = self.producer.publish(chain)?; @@ -537,6 +586,12 @@ where self.producer.notify_since(notify_from) } + + /// Finish the batch without notifying the consumer. + /// + /// Use this only when another protocol event guarantees that the consumer + /// will inspect the published descriptors. + pub fn finish_without_notify(self) {} } /// Builder for configuring a descriptor chain's buffer layout. @@ -544,20 +599,24 @@ where /// If dropped without building, no resources are leaked (allocations are /// deferred to [`build`](Self::build)). #[must_use = "call .build() to create a SendChain"] -pub struct ChainBuilder { +pub struct ChainBuilder { mem: M, - pool: P, + pool: SlotPool, rd_caps: SmallVec<[usize; 4]>, wr_caps: SmallVec<[usize; 4]>, + writable_avail: bool, + max_descs: usize, } -impl ChainBuilder { - fn new(mem: M, pool: P) -> Self { +impl ChainBuilder { + fn new(mem: M, pool: SlotPool, max_descs: usize) -> Self { Self { mem, pool, rd_caps: SmallVec::new(), wr_caps: SmallVec::new(), + writable_avail: false, + max_descs, } } @@ -584,110 +643,124 @@ impl ChainBuilder { self } + /// Request available upper-tier buffers within the ring budget. + /// + /// [`build`](Self::build) allocates these after all explicit requests. + /// This may add zero buffers. The complete chain must be nonempty. + pub fn writable_avail(mut self) -> Self { + self.writable_avail = true; + self + } + /// Allocate buffers and return a [`SendChain`] for writing. /// /// # Errors /// - /// - [`VirtqError::InvalidState`] - No buffers requested - /// - [`VirtqError::Alloc`] - Pool exhausted - pub fn build(self) -> Result, VirtqError> { - if self.rd_caps.is_empty() && self.wr_caps.is_empty() { + /// * [`VirtqError::InvalidState`] - no buffers requested + /// * [`VirtqError::Backpressure`] - insufficient pool slots or ring descriptors + /// * [`VirtqError::Bookkeeping`] - buffer record storage allocation failed + /// * [`VirtqError::Alloc`] - zero-length request or buffer allocation failed + pub fn build(self) -> Result, VirtqError> { + if self.rd_caps.is_empty() && self.wr_caps.is_empty() && !self.writable_avail { return Err(VirtqError::InvalidState); } - let mut rollback = Rollback::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(); + // Count explicit descriptors against the captured ring budget. + let slot_size = self.pool.slot_size(); + let rd_capacity = self.rd_caps.iter().try_fold(0usize, |total, &cap| { + total.checked_add(cap).ok_or(AllocError::Overflow) + })?; - // 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 = self.pool.alloc_sg(cap)?; - let mut remaining = cap; + let mut caps = self.rd_caps.iter().chain(&self.wr_caps); + let desc_count = caps.try_fold(0usize, |total, &cap| { + if cap == 0 { + return Err(AllocError::InvalidArg); + } + total + .checked_add(cap.div_ceil(slot_size)) + .ok_or(AllocError::Overflow) + })?; + + let remaining_descs = self + .max_descs + .checked_sub(desc_count) + .ok_or(VirtqError::Backpressure)?; + + // Reserve explicit records before taking slots. OwnedChain handles rollback. + let mut buffers = SmallVec::new(); + buffers + .try_reserve_exact(desc_count) + .map_err(|_| VirtqError::Bookkeeping)?; + + let mut owned = OwnedChain { + mem: self.mem, + pool: self.pool, + buffers, + readable: 0, + }; - for alloc in sgs { - let _ = checked_descriptor_len(alloc.len)?; - let seg_cap = remaining.min(alloc.len); + // Allocate readable regions before writable ones, splitting each at the upper slot size. + let regions = self + .rd_caps + .iter() + .map(|&len| (len, false)) + .chain(self.wr_caps.iter().map(|&len| (len, true))); + + for (total_len, wr) in regions { + let mut remaining = total_len; + + while remaining > 0 { + let len = remaining.min(slot_size); + let alloc = owned.pool.alloc(len)?; + let capacity = if wr { alloc.len } else { len as u32 }; - rd_caps.push(seg_cap); - rd_elems.push(BufferElement { + let ent = BufferEntry { addr: alloc.addr, - len: 0, - writable: false, - }); - remaining -= seg_cap; - rollback.allocs.push(alloc); - } + capacity, + written: 0, + }; + + owned.buffers.push(ent); + owned.readable += usize::from(!wr); - if remaining != 0 { - return Err(VirtqError::InvalidState); + remaining -= len; } } - // 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 = self.pool.alloc_sg(cap)?; - for alloc in sgs { - let len = checked_descriptor_len(alloc.len)?; - wr_elems.push(BufferElement { + // Use spare descriptors for upper-tier slots left after the explicit allocations. + if self.writable_avail { + let extra = owned.pool.num_free_upper().min(remaining_descs); + owned + .buffers + .try_reserve_exact(extra) + .map_err(|_| VirtqError::Bookkeeping)?; + + for _ in 0..extra { + let alloc = owned.pool.alloc(slot_size)?; + let ent = BufferEntry { addr: alloc.addr, - len, - writable: true, - }); - rollback.allocs.push(alloc); + capacity: alloc.len, + written: 0, + }; + + owned.buffers.push(ent); } } - let chain = BufferChainBuilder::new() - .readables(rd_elems) - .writables(wr_elems) - .build()?; - - rollback.release(); + // An availability-only request can leave no buffers to publish. + if owned.buffers.is_empty() { + return Err(VirtqError::Backpressure); + } Ok(SendChain { - mem: self.mem, - pool: self.pool, - chain: Some(chain), - rd_caps, - rd_capacity: self.rd_caps.iter().sum(), + owned, + rd_capacity, rd_written: 0, write_mode: WriteMode::Unset, }) } } -struct Rollback<'a, P: BufferProvider> { - pool: &'a P, - allocs: SmallVec<[Allocation; 8]>, -} - -impl<'a, P: BufferProvider> Rollback<'a, P> { - fn new(pool: &'a P) -> Self { - Self { - pool, - allocs: SmallVec::new(), - } - } - - fn release(mut self) { - self.allocs.clear(); - } -} - -impl Drop for Rollback<'_, 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:?}"); - } - } -} - /// Tracks which write API a [`SendChain`] payload uses, so the two paths are /// not mixed. /// @@ -724,29 +797,14 @@ enum WriteMode { /// /// If dropped without submitting, allocated buffers are returned to the pool. #[must_use = "dropping without submitting deallocates the buffers"] -pub struct SendChain { - mem: M, - pool: P, - chain: Option, - rd_caps: SmallVec<[usize; 4]>, +pub struct SendChain { + owned: OwnedChain, rd_capacity: usize, rd_written: usize, write_mode: WriteMode, } -// `chain` is wrapped in `Option` only so `into_inflight` can `take()` it -// without moving out of this `Drop` type; it stays `Some` for a chain's whole -// public lifetime, so these `expect`s cannot fail. -#[allow(clippy::expect_used)] -impl SendChain { - fn chain(&self) -> &BufferChain { - self.chain.as_ref().expect("SendChain missing BufferChain") - } - - fn chain_mut(&mut self) -> &mut BufferChain { - self.chain.as_mut().expect("SendChain missing BufferChain") - } - +impl SendChain { /// Record that this chain uses `mode`, asserting it is not mixed with the /// other write path. fn note_write_mode(&mut self, mode: WriteMode) { @@ -757,27 +815,38 @@ impl SendChain { self.write_mode = mode; } - fn into_inflight(mut self, token: Token) -> Inflight { - let chain = self.chain.take().expect("SendChain missing BufferChain"); - Inflight { token, chain } + /// Total number of descriptors in this chain. + #[inline] + pub fn desc_count(&self) -> usize { + self.owned.buffers.len() } - /// Number of producer-written readable segments in this chain. - pub fn segment_count(&self) -> usize { - self.chain().readables().len() + /// Number of readable descriptors in this chain. + #[inline] + pub fn rd_desc_count(&self) -> usize { + self.owned.readable + } + + /// Number of writable descriptors in this chain. + #[inline] + pub fn wr_desc_count(&self) -> usize { + self.desc_count() - self.rd_desc_count() } /// 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 +856,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.rd_desc_count() == 0 { return Err(VirtqError::NoPayloadSegment); } @@ -803,40 +873,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 buffer in &mut self.owned.buffers[..self.owned.readable] { if remaining.is_empty() { break; } - let written_len = readable.len as usize; - let free = cap - written_len; - if free == 0 { + let cap = buffer.capacity as usize; + let desc_off = buffer.written 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 = buffer + .addr + .checked_add(desc_off as u64) + .ok_or(VirtqError::MemoryWriteError)?; + + self.owned + .mem + .write(addr, &remaining[..len]) .map_err(|_| VirtqError::MemoryWriteError)?; - readable.len += n as u32; - written += n; - remaining = &remaining[n..]; + buffer.written += len as u32; + self.rd_written += len; + written += len; + remaining = &remaining[len..]; } - self.rd_written += written; Ok(written) } @@ -851,8 +915,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.rd_desc_count() == 0 { return Err(VirtqError::NoPayloadSegment); } @@ -881,11 +946,11 @@ impl SendChain { pub fn write_seg(&mut self, index: usize, buf: &[u8]) -> Result<&mut Self, VirtqError> { self.note_write_mode(WriteMode::Direct); - let cap = *self - .rd_caps - .get(index) + let buffer = self.owned.buffers[..self.owned.readable] + .get_mut(index) .ok_or(VirtqError::NoPayloadSegment)?; + let cap = buffer.capacity as usize; if buf.len() > cap { return Err(VirtqError::PayloadTooLarge { recv: buf.len(), @@ -893,19 +958,13 @@ impl SendChain { }); } - let addr = self - .chain() - .readables() - .get(index) - .ok_or(VirtqError::NoPayloadSegment)? - .addr; - - self.mem - .write(addr, buf) + self.owned + .mem + .write(buffer.addr, buf) .map_err(|_| VirtqError::MemoryWriteError)?; - let previous = self.chain().readables()[index].len as usize; - self.chain_mut().readables_mut()[index].len = checked_descriptor_len(buf.len())?; + let previous = buffer.written as usize; + buffer.written = checked_descriptor_len(buf.len())?; self.rd_written = self.rd_written - previous + buf.len(); Ok(self) } @@ -930,21 +989,17 @@ impl SendChain { { self.note_write_mode(WriteMode::Direct); - let cap = *self - .rd_caps - .get(index) + let buffer = self.owned.buffers[..self.owned.readable] + .get_mut(index) .ok_or_else(|| E::from(VirtqError::NoPayloadSegment))?; - let addr = self - .chain() - .readables() - .get(index) - .ok_or_else(|| E::from(VirtqError::NoPayloadSegment))? - .addr; + let cap = buffer.capacity as usize; + // SAFETY: This unpublished chain owns the allocation exclusively. let buf = unsafe { - self.mem - .as_mut_slice(addr, cap) + self.owned + .mem + .as_mut_slice(buffer.addr, cap) .map_err(|_| E::from(VirtqError::MemoryWriteError))? }; @@ -956,80 +1011,784 @@ impl SendChain { })); } - let previous = self.chain().readables()[index].len as usize; - // SAFETY: index was validated by the earlier get() call, so the readable element exists. - self.chain_mut().readables_mut()[index].len = - checked_descriptor_len(written).map_err(E::from)?; + let previous = buffer.written as usize; + buffer.written = checked_descriptor_len(written).map_err(E::from)?; self.rd_written = self.rd_written - previous + written; Ok(self) } } -impl Drop for SendChain { - fn drop(&mut self) { - if let Some(chain) = self.chain.take() { - for elem in chain.elems() { - let result = self.pool.dealloc(elem.addr); - debug_assert!(result.is_ok(), "SendChain drop dealloc failed: {result:?}"); - } - } +/// Tracks a pool slot's capacity and initialized length for writes and reply mappings. +/// +/// [`BufferElement::len`] carries only the length published to the peer. +#[derive(Debug)] +struct BufferEntry { + /// Buffer base address used for memory access and release through the pool. + addr: u64, + /// Usable bytes, + capacity: u32, + /// Initialized prefix length + written: u32, +} + +/// Buffer ownership passed from [`SendChain`] to [`Inflight`] at publication. +/// +/// Keeps the original memory backend and pool for reply mappings and slot release. +struct OwnedChain { + /// Original backend for writing payloads and mapping completed replies. + mem: M, + /// Shared allocator. + pool: SlotPool, + /// Allocation records in descriptor order. + buffers: SmallVec<[BufferEntry; 4]>, + /// Number of leading device readable buffers. + readable: usize, +} + +impl OwnedChain { + /// One [`BufferElement`] per direct descriptor, in chain order. + /// + /// The ring adds the descriptor IDs and flags that link the buffers. + fn descriptors(&self) -> impl ExactSizeIterator + Clone + '_ { + self.buffers.iter().enumerate().map(|(i, buf)| { + let wr = i >= self.readable; + let len = if wr { buf.capacity } else { buf.written }; + + BufferElement { + addr: buf.addr, + len, + writable: wr, + } + }) + } + + fn release(mut self) -> Result<(), AllocError> { + self.release_all() + } + + fn release_all(&mut self) -> Result<(), AllocError> { + let mut maybe_err = None; + while let Some(buf) = self.buffers.pop() { + if let Err(error) = self.pool.dealloc(buf.addr) + && maybe_err.is_none() + { + maybe_err = Some(error); + } + } + + self.readable = 0; + maybe_err.map_or(Ok(()), Err) + } +} + +impl OwnedChain { + fn into_segments(mut self, written: usize) -> Result { + let mut remaining = written; + let mut nonempty = 0; + + for buf in &mut self.buffers[self.readable..] { + let len = remaining.min(buf.capacity as usize); + buf.written = len as u32; + + nonempty += usize::from(len != 0); + remaining -= len; + } + + if remaining != 0 { + self.release()?; + return Err(VirtqError::InvalidState); + } + + let mut segments = SmallVec::<[Bytes; 4]>::new(); + segments + .try_reserve_exact(nonempty) + .map_err(|_| VirtqError::Bookkeeping)?; + + while let Some(buf) = self.buffers.last() { + if self.buffers.len() <= self.readable || buf.written == 0 { + let addr = buf.addr; + self.buffers.pop(); + self.pool.dealloc(addr)?; + + continue; + } + + let alloc = Allocation { + addr: buf.addr, + len: buf.capacity, + }; + + let written = buf.written as usize; + let lease = BufferLease::new(self.pool.clone(), alloc); + self.buffers.pop(); + + // SAFETY: Completion returns exclusive ownership of the initialized + // prefix. The mapper owns the slot until its borrowed view drops. + let mapping = unsafe { self.mem.map_buffer(lease, written) } + .map_err(|_| VirtqError::MemoryReadError)?; + + segments.push(Bytes::from_owner(mapping)); + } + + segments.reverse(); + Ok(Segments::from_smallvec(segments)) + } +} + +impl Drop for OwnedChain { + fn drop(&mut self) { + // best effort: if the pool deallocation fails, we can't do much about it here + if let Err(error) = self.release_all() { + log::error!("Failed to release virtqueue buffers: {error}"); + debug_assert!(false, "OwnedChain deallocation failed: {error}"); + } + } +} + +fn checked_descriptor_len(len: usize) -> Result { + if len > u32::MAX as usize { + return Err(VirtqError::PayloadTooLarge { + recv: len, + limit: u32::MAX as usize, + }); + } + Ok(len as u32) +} + +#[cfg(test)] +mod tests { + use alloc::rc::Rc; + use alloc::sync::Arc; + use core::sync::atomic::Ordering; + + use super::*; + use crate::virtq::ring::tests::{FaultMem, OwnedRing, TestMem, make_consumer, make_ring}; + use crate::virtq::test_utils::*; + + fn poll_received( + consumer: &mut VirtqConsumer, + ) -> (RecvChain, ReplyChain) { + consumer.poll(1024).unwrap().unwrap() + } + + #[derive(Clone)] + struct CopyingMem<'a>(&'a Rc); + + // SAFETY: All bounded operations delegate to TestMem. + unsafe impl MemOps for CopyingMem<'_> { + type Error = core::convert::Infallible; + + fn read(&self, addr: u64, dst: &mut [u8]) -> Result<(), Self::Error> { + self.0.read(addr, dst) + } + + fn write(&self, addr: u64, src: &[u8]) -> Result<(), Self::Error> { + self.0.write(addr, src) + } + + fn load_acquire(&self, addr: u64) -> Result { + self.0.load_acquire(addr) + } + + fn store_release(&self, addr: u64, val: u16) -> Result<(), Self::Error> { + self.0.store_release(addr, val) + } + + unsafe fn as_slice(&self, addr: u64, len: usize) -> Result<&[u8], Self::Error> { + // SAFETY: The caller supplies TestMem's slice preconditions. + unsafe { self.0.as_slice(addr, len) } + } + + unsafe fn as_mut_slice(&self, addr: u64, len: usize) -> Result<&mut [u8], Self::Error> { + // SAFETY: The caller supplies exclusive access to this range. + unsafe { self.0.as_mut_slice(addr, len) } + } + } + + impl BufferMap for CopyingMem<'_> { + type Mapping = Vec; + + unsafe fn map_buffer( + &self, + lease: BufferLease, + written: usize, + ) -> Result { + assert!(written <= lease.allocation().len as usize); + let mut bytes = vec![0; written]; + self.read(lease.allocation().addr, &mut bytes)?; + Ok(bytes) + } + } + + fn make_virtq_pair( + 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 make_chain( + lower_count: usize, + upper_count: usize, + slot_size: usize, + max_descs: usize, + ) -> ChainBuilder { + let mem = TestMem::new(lower_count * 256 + upper_count * slot_size); + let lower = SlotLayout::new(mem.base_addr(), 256, lower_count); + let upper = SlotLayout::new(lower.end_addr().unwrap(), slot_size, upper_count); + + let pool = if lower_count == 0 { + SlotPool::new(upper) + } else { + SlotPool::new_tiered(lower, upper) + } + .unwrap(); + + ChainBuilder::new(mem, pool, max_descs) + } + + fn inflight(seq: u32, id: u16) -> Inflight { + let mem = TestMem::new(8); + let pool = SlotPool::new(SlotLayout::new(mem.base_addr(), 8, 1)).unwrap(); + let chain = ChainBuilder::new(mem, pool, 1).readable(8).build().unwrap(); + Inflight { + token: Token { seq, id }, + chain: ManuallyDrop::new(chain.owned), + } + } + + #[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)); + } + + let removed = table.remove(7).unwrap(); + assert_eq!(removed.token.seq, 1); + + ManuallyDrop::into_inner(removed.chain).release().unwrap(); + assert!(!table.contains(7)); + + for (id, seq) in [(5, 2), (3, 0)] { + let removed = table.remove(id).unwrap(); + assert_eq!(removed.token.seq, seq); + ManuallyDrop::into_inner(removed.chain).release().unwrap(); + } + 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(); + } + + assert!(matches!( + producer.chain().readable(1).build(), + Err(VirtqError::Backpressure) + )); + producer.reset().unwrap(); + } + + #[test] + fn submission_rechecks_capacity_after_reservation() { + let ring = make_ring(2); + let (mut producer, mut consumer) = make_virtq_pair(&ring, 64); + let pool = producer.pool.clone(); + + let chain = producer.chain().writable(64).build().unwrap(); + for _ in 0..ring.len() { + let other = producer.chain().writable(32).build().unwrap(); + producer.submit(other).unwrap(); + } + + assert_eq!(pool.num_live(), 3); + assert!(matches!( + producer.submit(chain), + Err(VirtqError::Backpressure) + )); + assert_eq!(pool.num_live(), 2); + + producer.reset().unwrap(); + consumer.reset().unwrap(); + + assert_eq!(pool.num_live(), 0); + } + + #[test] + fn publication_failure_keeps_allocations_until_stopped_reset() { + for failed_write in 0..4 { + let ring = make_ring(4); + let orig_mem = FaultMem::new(ring.mem()); + let orig_gen = Arc::downgrade(&orig_mem.0); + + let base = ring.mem().base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + let pool = SlotPool::new(SlotLayout::new(base, 64, 4)).unwrap(); + let notif = TestNotifier::new(); + + let source = VirtqProducer::new(ring.layout(), orig_mem, notif.clone(), pool.clone()); + + let chain = source.chain().writable(64).build().unwrap(); + drop(source); + + let mem = FaultMem::new(ring.mem()); + let mut producer = + VirtqProducer::new(ring.layout(), mem.clone(), notif.clone(), pool.clone()); + + mem.fail_write_at(failed_write); + + let res = producer.submit(chain); + assert!(matches!( + res, + Err(VirtqError::RingError(RingError::MemError { .. })) + )); + assert_eq!(pool.num_live(), 1); + assert_eq!(producer.inflight.live.len(), 1); + assert!(orig_gen.upgrade().is_some()); + + mem.allow_writes(); + producer.reset().unwrap(); + + assert_eq!(pool.num_live(), 0); + assert!(orig_gen.upgrade().is_none()); + } + } + + #[test] + fn chain_clones_pool_only_for_independent_owners() { + let ring = make_ring(8); + let mem = ring.mem(); + let base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + let pool = SlotPool::new(SlotLayout::new(base, 64, 8)).unwrap(); + + let notif = TestNotifier::new(); + let mapping = FaultMem::new(mem.clone()); + + let mut producer = + VirtqProducer::new(ring.layout(), mapping.clone(), notif.clone(), pool.clone()); + + let mut consumer = VirtqConsumer::new(ring.layout(), mem, TestNotifier::new()); + + assert_eq!(pool.strong_count(), 2); + + let mut chain = producer + .chain() + .readable(192) + .writable(192) + .build() + .unwrap(); + + chain.write_all(b"request").unwrap(); + + assert_eq!(pool.num_live(), 6); + assert_eq!(pool.strong_count(), 3); + assert!(chain.owned.buffers.spilled()); + + let buffers = chain.owned.buffers.as_ptr(); + producer.submit(chain).unwrap(); + + assert_eq!(pool.strong_count(), 3); + assert_eq!(producer.inflight.live[0].chain.buffers.as_ptr(), buffers); + + let (recv, reply) = poll_received(&mut consumer); + let ReplyChain::Writable(mut reply) = reply else { + panic!("expected writable reply"); + }; + reply.write_all(&[0xa5; 70]).unwrap(); + consumer.complete(recv, reply).unwrap(); + + let segments = producer.poll().unwrap().unwrap().into_segments().unwrap(); + + assert_eq!(segments.segment_count(), 2); + assert_eq!(mapping.0.map_calls.load(Ordering::Relaxed), 2); + assert_eq!(pool.strong_count(), 4); + assert_eq!(pool.num_live(), 2); + + let retained = segments.as_slice()[0].slice(1..); + let cloned = retained.clone(); + + drop(segments); + + assert_eq!(pool.num_live(), 1); + assert_eq!(pool.strong_count(), 3); + + producer.reset().unwrap(); + consumer.reset().unwrap(); + + drop(producer); + + assert_eq!(pool.strong_count(), 2); + assert_eq!(retained.as_ref(), &[0xa5; 63]); + + drop(retained); + + assert_eq!(pool.num_live(), 1); + + drop(cloned); + + assert_eq!(pool.num_live(), 0); + assert_eq!(pool.strong_count(), 1); + assert_eq!(mapping.0.map_calls.load(Ordering::Relaxed), 2); + } + + #[test] + fn copied_mapping_crosses_threads_without_its_pool_or_borrowed_backend() { + let ring = make_ring(4); + let mem = Rc::new(ring.mem()); + let base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + let pool = SlotPool::new(SlotLayout::new(base, 64, 1)).unwrap(); + let notif = TestNotifier::new(); + + let mut producer = + VirtqProducer::new(ring.layout(), CopyingMem(&mem), notif.clone(), pool.clone()); + let mut consumer = VirtqConsumer::new(ring.layout(), ring.mem(), notif); + + 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"copied").unwrap(); + consumer.complete(recv, reply).unwrap(); + + let bytes = producer.poll().unwrap().unwrap().into_bytes().unwrap(); + + assert_eq!(pool.num_live(), 0); + assert_eq!(pool.strong_count(), 2); + + let mut reused = producer.chain().readable(64).build().unwrap(); + reused.write_all(b"reused").unwrap(); + + std::thread::spawn(move || assert_eq!(bytes.as_ref(), b"copied")) + .join() + .unwrap(); + + assert_eq!(pool.num_live(), 1); + assert_eq!(pool.strong_count(), 3); + + drop(reused); + + assert_eq!(pool.num_live(), 0); + } + + #[test] + fn completion_uses_its_original_memory_and_allocating_pool() { + let ring = make_ring(4); + let mem = FaultMem::new(ring.mem()); + let orig_gen = Arc::downgrade(&mem.0); + let base = ring.mem().base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + + let notifier = TestNotifier::new(); + let original = SlotPool::new(SlotLayout::new(base, 64, 4)).unwrap(); + let source = VirtqProducer::new(ring.layout(), mem, TestNotifier::new(), original.clone()); + + let chain = source.chain().writable(64).build().unwrap(); + drop(source); + + assert!(orig_gen.upgrade().is_some()); + + let replacement = + SlotPool::new(SlotLayout::new(original.base_addr() + 0x1000, 64, 4)).unwrap(); + + let replacement_mem = FaultMem::new(ring.mem()); + replacement_mem.fail_mapping_at(0); + + let mut producer = VirtqProducer::new( + ring.layout(), + replacement_mem.clone(), + notifier.clone(), + replacement.clone(), + ); + let mut consumer = VirtqConsumer::new(ring.layout(), ring.mem(), TestNotifier::new()); + + producer.submit(chain).unwrap(); + + assert!(orig_gen.upgrade().is_some()); + + 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(); + + let data = producer.poll().unwrap().unwrap().into_bytes().unwrap(); + + assert_eq!(data.as_ref(), b"retained"); + assert_eq!(original.num_live(), 1); + assert_eq!(replacement.num_live(), 0); + assert_eq!(replacement_mem.0.map_calls.load(Ordering::Relaxed), 0); + + let map_calls = orig_gen + .upgrade() + .unwrap() + .map_calls + .load(Ordering::Relaxed); + assert_eq!(map_calls, 1); + + drop(producer); + + assert_eq!(data.as_ref(), b"retained"); + + drop(data); + + assert_eq!(original.num_live(), 0); + assert!(orig_gen.upgrade().is_none()); + } + + #[test] + fn mapping_failure_releases_returned_and_unmapped_slots() { + for fail_at in 0..2 { + let ring = make_ring(8); + let mem = FaultMem::new(ring.mem()); + mem.fail_mapping_at(fail_at); + + let base = ring.mem().base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + let pool = SlotPool::new(SlotLayout::new(base, 64, 8)).unwrap(); + let notif = TestNotifier::new(); + let mut producer = + VirtqProducer::new(ring.layout(), mem.clone(), notif.clone(), pool.clone()); + + let mut consumer = make_consumer(&ring); + + let chain = producer + .chain() + .readable(192) + .writable(256) + .build() + .unwrap(); + + let mut addresses: Vec<_> = chain + .owned + .buffers + .iter() + .map(|buffer| buffer.addr) + .collect(); + + producer.submit(chain).unwrap(); + + let (id, _) = consumer.poll_available().unwrap(); + // TestMem's zeroed backing supplies initialized bytes for the mapped prefix. + consumer.submit_used(id, 70).unwrap(); + + assert!(matches!(producer.poll(), Err(VirtqError::MemoryReadError))); + assert_eq!(mem.0.map_calls.load(Ordering::Relaxed), fail_at + 1); + assert_eq!(pool.num_live(), 0); + assert_eq!(pool.strong_count(), 2); + assert_eq!(producer.num_inflight(), 0); + + if fail_at == 1 { + // The mapper releases the failed lease before completed owners unwind. + addresses.swap(3, 4); + } + + let repeated = producer + .chain() + .readable(192) + .writable(256) + .build() + .unwrap(); + + assert_eq!( + repeated + .owned + .buffers + .iter() + .map(|buffer| buffer.addr) + .collect::>(), + addresses + ); + } + } + + #[test] + fn empty_and_malformed_completions_skip_mapping_and_release_every_slot() { + let ring = make_ring(4); + let mem = FaultMem::new(ring.mem()); + mem.fail_mapping_at(0); + + let base = ring.mem().base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + let pool = SlotPool::new(SlotLayout::new(base, 64, 4)).unwrap(); + let notif = TestNotifier::new(); + + let mut producer = + VirtqProducer::new(ring.layout(), mem.clone(), notif.clone(), pool.clone()); + + let mut consumer = make_consumer(&ring); + + for (writable, written) in [(false, 0), (true, 0), (true, 129)] { + let builder = producer.chain().readable(64); + let builder = if writable { + builder.writable(128) + } else { + builder + }; + + let token = producer.submit(builder.build().unwrap()).unwrap(); + let (id, _) = consumer.poll_available().unwrap(); + consumer.submit_used(id, written).unwrap(); + + match producer.poll() { + Ok(Some(UsedChain::Ack(returned))) if !writable => assert_eq!(returned, token), + Ok(Some(UsedChain::Data(returned, segments))) if writable && written == 0 => { + assert_eq!(returned, token); + assert_eq!(segments.segment_count(), 0); + } + Err(VirtqError::InvalidState) if written == 129 => {} + other => panic!("unexpected completion: {other:?}"), + } + + assert_eq!(mem.0.map_calls.load(Ordering::Relaxed), 0); + assert_eq!(pool.num_live(), 0); + assert_eq!(pool.strong_count(), 2); + assert_eq!(producer.num_inflight(), 0); + } + } + + #[test] + fn cancelling_a_chain_preserves_slot_order_after_a_partial_write_error() { + let storage = TestMem::new(16); + let base = storage.base_addr(); + let mem = FaultMem::new(storage); + let t1 = SlotLayout::new(base, 2, 2); + let t2 = SlotLayout::new(base + 4, 4, 3); + let pool = SlotPool::new_tiered(t1, t2).unwrap(); + + let mut chain = ChainBuilder::new(mem.clone(), pool.clone(), 8) + .readable(9) + .writable_avail() + .build() + .unwrap(); + + let addresses: Vec<_> = chain + .owned + .buffers + .iter() + .map(|buffer| buffer.addr) + .collect(); + + mem.fail_write_at(1); + assert!(matches!( + chain.write_all(b"abcdefghi"), + Err(VirtqError::MemoryWriteError) + )); + assert_eq!(chain.written(), 4); + + drop(chain); + + assert_eq!(pool.num_live(), 0); + assert_eq!(pool.strong_count(), 1); + + mem.allow_writes(); + let repeated = ChainBuilder::new(mem, pool.clone(), 8) + .readable(9) + .writable_avail() + .build() + .unwrap(); + + assert_eq!( + repeated + .owned + .buffers + .iter() + .map(|buffer| buffer.addr) + .collect::>(), + addresses + ); } -} -fn checked_descriptor_len(len: usize) -> Result { - if len > u32::MAX as usize { - return Err(VirtqError::PayloadTooLarge { - recv: len, - limit: u32::MAX as usize, - }); - } - Ok(len as u32) -} + #[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()); -#[cfg(test)] -mod tests { - use super::*; - use crate::virtq::ring::tests::{TestMem, make_consumer, make_ring}; - use crate::virtq::test_utils::*; + for _ in 0..ring.len() { + let chain = producer.chain().writable(64).build().unwrap(); + producer.submit(chain).unwrap(); + } - fn poll_received( - consumer: &mut VirtqConsumer, - ) -> (RecvChain, ReplyChain) { - consumer.poll(1024).unwrap().unwrap() - } + assert_eq!(pool.num_free(), 0); - #[derive(Clone)] - struct NoDirectSliceMem(TestMem); + producer.reset().unwrap(); - // SAFETY: Delegates all non-slice memory operations to TestMem. Direct - // slices are intentionally unsupported to exercise producer error handling. - unsafe impl MemOps for NoDirectSliceMem { - type Error = (); + 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) + ); - fn read(&self, addr: u64, dst: &mut [u8]) -> Result<(), Self::Error> { - self.0.read(addr, dst).map_err(|e| match e {}) + for _ in 0..ring.len() { + let chain = producer.chain().writable(64).build().unwrap(); + producer.submit(chain).unwrap(); } + producer.reset().unwrap(); + } - fn write(&self, addr: u64, src: &[u8]) -> Result<(), Self::Error> { - self.0.write(addr, src).map_err(|e| match e {}) - } + #[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(); - fn load_acquire(&self, addr: u64) -> Result { - self.0.load_acquire(addr).map_err(|e| match e {}) - } + let (recv, reply) = poll_received(&mut consumer); + let ReplyChain::Writable(mut reply) = reply else { + panic!("expected writable reply"); + }; - fn store_release(&self, addr: u64, val: u16) -> Result<(), Self::Error> { - self.0.store_release(addr, val).map_err(|e| match e {}) - } + reply.write_all(b"retained").unwrap(); + consumer.complete(recv, reply).unwrap(); + producer.reclaim().unwrap(); - unsafe fn as_slice(&self, _addr: u64, _len: usize) -> Result<&[u8], Self::Error> { - Err(()) - } + assert!(matches!(producer.reset(), Err(VirtqError::InvalidState))); - unsafe fn as_mut_slice(&self, _addr: u64, _len: usize) -> Result<&mut [u8], Self::Error> { - Err(()) - } + drop(producer.poll().unwrap().unwrap()); + producer.reset().unwrap(); } #[test] @@ -1049,7 +1808,9 @@ 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(), 2); + assert_eq!(se.rd_desc_count(), 1); + assert_eq!(se.wr_desc_count(), 1); assert_eq!(se.capacity(), 16); } @@ -1065,48 +1826,427 @@ mod tests { .writable(32) .build() .unwrap(); - se.write_all(b"hello world").unwrap(); + se.write_all(b"hello world").unwrap(); assert_eq!(se.written(), 11); 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(); + producer.reset().unwrap(); } #[test] - fn test_chain_readable_splits_logical_capacity() { + 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.owned.descriptors().collect::>(); + + 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 chain_avail_uses_pool_availability_at_build() { + for held_count in [0, 2, 3] { + let builder = make_chain(2, 3, 4096, 3); + let pool = builder.pool.clone(); + let builder = builder.writable_avail().readable(128); + + assert_eq!(pool.num_live(), 0); + + let held: Vec<_> = (0..held_count).map(|_| pool.alloc(4096).unwrap()).collect(); + + let chain = builder.build().unwrap(); + + let writable = (3 - held_count).min(2); + assert_eq!(chain.rd_desc_count(), 1); + assert_eq!(chain.wr_desc_count(), writable); + assert_eq!(chain.desc_count(), 1 + writable); + assert_eq!(chain.owned.buffers[0].capacity, 128); + assert!( + chain.owned.buffers[1..] + .iter() + .all(|buf| buf.capacity == 4096) + ); + assert_eq!(pool.num_free_lower(), 1); + assert_eq!(pool.num_free_upper(), 3 - held_count - writable); + assert_eq!(pool.num_live(), held_count + chain.desc_count()); + + drop(chain); + + assert_eq!(pool.num_live(), held_count); + + for allocation in held.into_iter().rev() { + pool.dealloc(allocation.addr).unwrap(); + } + + assert_eq!(pool.num_live(), 0); + } + } + + #[test] + fn chain_avail_respects_mandatory_descriptor_budget() { + for max_descs in 0..3 { + let builder = make_chain(2, 3, 4096, max_descs); + let pool = builder.pool.clone(); + + let result = builder + .readable(4096) + .readable(128) + .writable_avail() + .build(); + + if max_descs == 2 { + let chain = result.unwrap(); + + assert_eq!(chain.rd_desc_count(), 2); + assert_eq!(chain.wr_desc_count(), 0); + + drop(chain); + } else { + assert!(matches!(result, Err(VirtqError::Backpressure))); + } + + assert_eq!(pool.num_live(), 0); + + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(4096).unwrap(); + + assert_eq!(lower.addr, pool.slot_addr(1).unwrap()); + assert_eq!(upper.addr, pool.slot_addr(4).unwrap()); + + pool.dealloc(lower.addr).unwrap(); + pool.dealloc(upper.addr).unwrap(); + } + } + + #[test] + fn chain_avail_preserves_explicit_writable_capacity_when_full() { + for (upper_count, max_descs) in [(2, 2), (1, 3)] { + let builder = make_chain(1, upper_count, 4096, max_descs); + let pool = builder.pool.clone(); + + let chain = builder + .writable_avail() + .writable(128) + .readable(128) + .build() + .unwrap(); + + assert_eq!(chain.rd_desc_count(), 1); + assert_eq!(chain.wr_desc_count(), 1); + assert_eq!(chain.owned.buffers[0].capacity, 128); + assert_eq!(chain.owned.buffers[1].capacity, 4096); + assert!(chain.owned.buffers.iter().all(|buf| buf.written == 0)); + assert_eq!(pool.num_free_upper(), upper_count - 1); + + drop(chain); + + assert_eq!(pool.num_live(), 0); + } + } + + #[test] + fn chain_avail_cannot_build_an_empty_chain() { + for max_descs in [0, 2] { + let builder = make_chain(1, 1, 4096, max_descs); + let pool = builder.pool.clone(); + let held = (max_descs != 0).then(|| pool.alloc(4096).unwrap()); + + assert!(matches!( + builder.writable_avail().build(), + Err(VirtqError::Backpressure) + )); + assert_eq!(pool.num_free_lower(), 1); + assert_eq!(pool.num_live(), usize::from(held.is_some())); + + if let Some(allocation) = held { + pool.dealloc(allocation.addr).unwrap(); + } + } + } + + #[test] + fn chain_avail_only_uses_upper_slots() { + let builder = make_chain(4, 2, 4096, 6); + let pool = builder.pool.clone(); + + let chain = builder.writable_avail().writable_avail().build().unwrap(); + + assert_eq!(chain.rd_desc_count(), 0); + assert_eq!(chain.wr_desc_count(), 2); + assert_eq!(pool.num_free_lower(), 4); + assert_eq!(pool.num_free_upper(), 0); + + drop(chain); + + assert_eq!(pool.num_live(), 0); + } + + #[test] + fn chain_preserves_region_order_with_odd_slot_sizes() { + let builder = make_chain(1, 4, 3001, 5); + let pool = builder.pool.clone(); + + // The final short readable must fall back to an upper slot. + let chain = builder + .readable(128) + .readable(6003) + .writable_avail() + .build() + .unwrap(); + + assert_eq!(chain.rd_desc_count(), 4); + + let actual_offsets = chain + .owned + .buffers + .iter() + .map(|buf| { + ( + buf.addr, + buf.capacity, + pool.allocation_len(buf.addr).unwrap(), + ) + }) + .collect::>(); + + let expected_offsets = [ + (0, 128, 256), + (4, 3001, 3001), + (3, 3001, 3001), + (2, 1, 3001), + (1, 3001, 3001), + ] + .map(|(slot, capacity, allocated)| (pool.slot_addr(slot).unwrap(), capacity, allocated)); + + assert_eq!(actual_offsets, expected_offsets); + + drop(chain); + + assert_eq!(pool.num_live(), 0); + } + + #[test] + fn chain_failure_releases_only_its_own_slots() { + let builder = make_chain(2, 3, 4096, 4); + let pool = builder.pool.clone(); + let held = pool.alloc(4096).unwrap(); + + assert!(matches!( + builder.readable(128).writable(4096 * 3).build(), + Err(VirtqError::Backpressure) + )); + assert_eq!(pool.live_addrs(), [held.addr]); + + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(4096).unwrap(); + let next_upper = pool.alloc(4096).unwrap(); + + assert_eq!(lower.addr, pool.slot_addr(1).unwrap()); + assert_eq!(upper.addr, pool.slot_addr(3).unwrap()); + assert_eq!(next_upper.addr, pool.slot_addr(2).unwrap()); + + pool.dealloc(next_upper.addr).unwrap(); + pool.dealloc(upper.addr).unwrap(); + pool.dealloc(lower.addr).unwrap(); + pool.dealloc(held.addr).unwrap(); + } + + #[test] + fn chain_rejects_invalid_request_sequences() { + for (readable, writable) in [(vec![128, 0], vec![]), (vec![128], vec![4096, 0])] { + let mut builder = make_chain(1, 1, 4096, 4); + let pool = builder.pool.clone(); + + for cap in readable { + builder = builder.readable(cap); + } + for cap in writable { + builder = builder.writable(cap); + } + + assert!(matches!( + builder.build(), + Err(VirtqError::Alloc(AllocError::InvalidArg)) + )); + assert_eq!(pool.num_live(), 0); + } + + let builder = make_chain(1, 1, 4096, 4); + let pool = builder.pool.clone(); + + assert!(matches!( + builder.readable(usize::MAX).readable(1).build(), + Err(VirtqError::Alloc(AllocError::Overflow)) + )); + assert_eq!(pool.num_live(), 0); + } + + #[test] + fn chain_keeps_four_records_inline_and_one_pool_handle() { + for count in [4, 5] { + let builder = make_chain(0, count, 4096, count); + let pool = builder.pool.clone(); + + let chain = builder + .readable((count - 1) * 4096) + .writable_avail() + .build() + .unwrap(); + + assert_eq!(chain.desc_count(), count); + assert_eq!(chain.owned.buffers.spilled(), count > 4); + assert_eq!(pool.strong_count(), 2); + + drop(chain); + + assert_eq!(pool.num_live(), 0); + } + } + + #[test] + fn test_chain_multi_readable_appends_across_calls() { + let ring = make_ring(16); + let (mut producer, mut consumer) = make_virtq_pair(&ring, 4); + + 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(); + producer.reset().unwrap(); + } + + #[test] + fn test_chain_readable_splits_logical_capacity() { + let expected_lengths = [4, 4, 4]; + let ring = make_ring(16); + let (mut producer, mut consumer, _) = make_test_producer_with_slot_size(&ring, 4); + + let mut se = producer.chain().readable(10).writable(32).build().unwrap(); + let readables = &se.owned.buffers[..se.rd_desc_count()]; + + assert_eq!(se.rd_desc_count(), 3); + assert_eq!(se.capacity(), 10); + + let caps = readables.iter().map(|buf| buf.capacity).collect::>(); + assert_eq!(caps, [4, 4, 2]); + + let lengths = readables + .iter() + .map(|buf| producer.pool.allocation_len(buf.addr).unwrap()) + .collect::>(); + assert_eq!(lengths, expected_lengths); + + se.write_all(b"abcdefghij").unwrap(); + assert_eq!(se.written(), 10); + + let token = producer.submit(se).unwrap(); + let (recv, reply) = poll_received(&mut consumer); + + assert_eq!(recv.token(), token); + 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(); + producer.reset().unwrap(); + + assert_eq!(producer.pool.num_live(), 0); + } + + #[test] + fn test_chain_readable_splits_logical_capacity_tiered() { + let expected_lengths = [4, 4, 2]; + let ring = make_ring(16); + let (mut producer, mut consumer) = make_virtq_pair(&ring, 4); let mut se = producer.chain().readable(10).writable(32).build().unwrap(); + let readables = &se.owned.buffers[..se.rd_desc_count()]; - assert_eq!(se.segment_count(), 3); + assert_eq!(se.rd_desc_count(), 3); assert_eq!(se.capacity(), 10); + let caps = readables.iter().map(|buf| buf.capacity).collect::>(); + assert_eq!(caps, [4, 4, 2]); + + let lengths = readables + .iter() + .map(|buf| producer.pool.allocation_len(buf.addr).unwrap()) + .collect::>(); + assert_eq!(lengths, expected_lengths); + se.write_all(b"abcdefghij").unwrap(); assert_eq!(se.written(), 10); 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(); + producer.reset().unwrap(); + + assert_eq!(producer.pool.num_live(), 0); } #[test] @@ -1123,25 +2263,19 @@ 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_virtq_pair(&ring, 4); 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 +2307,9 @@ 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(); + producer.reset().unwrap(); } #[test] @@ -1196,9 +2331,11 @@ 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(); + producer.reset().unwrap(); } #[test] @@ -1212,58 +2349,60 @@ 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(); + producer.reset().unwrap(); } #[test] fn test_chain_multi_writable_used_returns_segments() { let ring = make_ring(16); - let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + let (mut producer, mut consumer, _notifier) = make_test_producer_with_slot_size(&ring, 6); 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); + assert_eq!(wc.capacity(), 12); 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); let segments = used.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"); + assert_eq!(segments.as_slice()[0].as_ref(), b"hello "); + assert_eq!(segments.as_slice()[1].as_ref(), b"world"); assert_eq!(segments.to_bytes().as_ref(), b"hello world"); } #[test] fn test_chain_multi_writable_short_used_truncates_last_segment() { let ring = make_ring(16); - let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + let (mut producer, mut consumer, _notifier) = make_test_producer_with_slot_size(&ring, 6); 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(); 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" wo"); + assert_eq!(segments.as_slice()[0].as_ref(), b"hello "); + assert_eq!(segments.as_slice()[1].as_ref(), b"wo"); assert_eq!(segments.to_bytes().as_ref(), b"hello wo"); } @@ -1275,8 +2414,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 +2468,9 @@ 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(); + producer.reset().unwrap(); } #[test] @@ -1349,8 +2489,9 @@ 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(); + producer.reset().unwrap(); } #[test] @@ -1365,8 +2506,9 @@ 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(); + producer.reset().unwrap(); } #[test] @@ -1384,8 +2526,9 @@ 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(); + producer.reset().unwrap(); } #[test] @@ -1404,8 +2547,9 @@ 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(); + producer.reset().unwrap(); } #[test] @@ -1423,15 +2567,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_virtq_pair(&ring, 4); let mut se = producer.chain().readable(8).build().unwrap(); - assert_eq!(se.segment_count(), 2); + assert_eq!(se.rd_desc_count(), 2); assert!(matches!( se.with_seg(2, |_| Ok::(0)), Err(VirtqError::NoPayloadSegment) @@ -1494,6 +2633,7 @@ mod tests { let se = producer.chain().readable(64).writable(128).build().unwrap(); let tok = producer.submit(se).unwrap(); assert!(tok.id < 16); + producer.reset().unwrap(); } #[test] @@ -1510,6 +2650,7 @@ mod tests { let se = producer.chain().readable(64).writable(128).build().unwrap(); let tok = producer.submit(se).unwrap(); assert!(tok.id < 16); + producer.reset().unwrap(); } #[test] @@ -1524,6 +2665,7 @@ mod tests { producer.submit(se).unwrap(); assert!(notifier.notification_count() > initial_count); + producer.reset().unwrap(); } #[test] @@ -1538,6 +2680,7 @@ mod tests { producer.submit(se).unwrap(); assert!(notifier.notification_count() > initial_count); + producer.reset().unwrap(); } #[test] @@ -1551,6 +2694,7 @@ mod tests { producer.submit(se).unwrap(); assert!(notifier.notification_count() > initial_count); + producer.reset().unwrap(); } #[test] @@ -1575,12 +2719,13 @@ 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(); + producer.reset().unwrap(); } #[test] @@ -1607,6 +2752,7 @@ mod tests { assert!(batch.finish().unwrap()); assert_eq!(notifier.notification_count(), 1); + producer.reset().unwrap(); } #[test] @@ -1619,6 +2765,26 @@ mod tests { assert_eq!(notifier.notification_count(), 0); } + #[test] + fn test_batch_can_finish_without_notification() { + let ring = make_ring(16); + let (mut producer, mut consumer, notifier) = make_test_producer(&ring); + + let mut batch = producer.batch(); + let mut chain = batch.chain().readable(4).build().unwrap(); + chain.write_all(b"data").unwrap(); + batch.submit(chain).unwrap(); + batch.finish_without_notify(); + + assert_eq!(notifier.notification_count(), 0); + + let (recv, reply) = poll_received(&mut consumer); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"data"); + + consumer.complete(recv, reply).unwrap(); + producer.reset().unwrap(); + } + #[test] fn test_write_only_round_trip() { let ring = make_ring(16); @@ -1629,11 +2795,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 +2821,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 +2839,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"); } @@ -1687,36 +2853,38 @@ mod tests { } #[test] - fn test_poll_used_requires_direct_slice() { + fn test_poll_used_requires_owned_mapping() { let ring = make_ring(16); let layout = ring.layout(); let test_mem = ring.mem(); let pool_base = test_mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; - let pool = TestPool::new(pool_base, 0x8000); + let pool = SlotPool::new(SlotLayout::new(pool_base, 128, 0x8000 / 128)).unwrap(); let notifier = TestNotifier::new(); - let mem = NoDirectSliceMem(test_mem); - let mut producer = VirtqProducer::new(layout, mem.clone(), notifier.clone(), pool); + let mem = FaultMem::new(test_mem); + mem.deny_views(); + let mut producer = VirtqProducer::new(layout, mem.clone(), notifier.clone(), pool.clone()); let mut consumer = VirtqConsumer::new(layout, mem, notifier); let mut se = producer.chain().readable(64).writable(128).build().unwrap(); 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"); } assert!(matches!(producer.poll(), Err(VirtqError::MemoryReadError))); + assert_eq!(pool.num_live(), 0); } #[test] fn test_villain_used_len_exceeding_writable_capacity_is_rejected_and_released() { let ring = make_ring(16); - let (mut producer, _consumer, _notifier) = make_test_producer(&ring); + let (mut producer, _consumer, _notifier) = make_test_producer_with_slot_size(&ring, 4); let mut ring_consumer = make_consumer(&ring); let se = producer.chain().writable(4).build().unwrap(); @@ -1731,6 +2899,7 @@ mod tests { let se = producer.chain().writable(4).build().unwrap(); producer.submit(se).unwrap(); assert_eq!(producer.inner.num_inflight(), 1); + producer.reset().unwrap(); } #[test] @@ -1750,53 +2919,6 @@ mod tests { Err(VirtqError::RingError(RingError::InvalidState)) )); 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().as_ref(), b"hello"); - consumer.complete(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()); + producer.reset().unwrap(); } } diff --git a/src/hyperlight_common/src/virtq/ring.rs b/src/hyperlight_common/src/virtq/ring.rs index 2d38271f8..eaf180d96 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}; @@ -88,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 @@ -154,6 +176,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. @@ -188,15 +214,21 @@ pub struct Writable; /// Upholds invariants: at least one buffer must be present in the chain, /// and readable buffers must be added before writable buffers. /// -/// The builder stores up to 16 buffer elements inline to avoid allocation for +/// The builder stores up to four buffer elements inline to avoid allocation for /// common small chains. Larger chains are still supported and spill to the heap. #[derive(Debug, Default)] pub struct BufferChainBuilder { - elems: SmallVec<[BufferElement; 16]>, + elems: SmallVec<[BufferElement; 4]>, split: usize, marker: PhantomData, } +impl BufferChainBuilder { + pub(super) fn reserve_exact(&mut self, additional: usize) { + self.elems.reserve_exact(additional); + } +} + impl BufferChainBuilder { /// Create a new builder in the [`Readable`] state. pub fn new() -> Self { @@ -331,7 +363,7 @@ impl BufferChainBuilder { #[derive(Debug, Clone)] pub struct BufferChain { /// All buffer elements (readable followed by writable) - elems: SmallVec<[BufferElement; 16]>, + elems: SmallVec<[BufferElement; 4]>, /// Split index between readable and writable buffers split: usize, } @@ -347,11 +379,6 @@ impl BufferChain { &self.elems[..self.split] } - /// Get mutable readable buffers in chain. - pub(crate) fn readables_mut(&mut self) -> &mut [BufferElement] { - &mut self.elems[..self.split] - } - /// Get writable buffers in chain pub fn writables(&self) -> &[BufferElement] { &self.elems[self.split..] @@ -737,6 +764,15 @@ impl RingProducer { Ok(UsedBuffer { id, len: desc.len }) } + /// Get the next available descriptor ID without consuming it. + pub fn next_id(&self) -> Result { + let id = *self.id_free.last().ok_or(RingError::OutOfMemory)?; + if self.id_num[id as usize] != 0 { + return Err(RingError::InvalidState); + } + Ok(id) + } + /// Get number of free descriptors in the ring. #[inline] pub fn num_free(&self) -> usize { @@ -892,45 +928,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(()) } } @@ -1031,7 +1060,7 @@ impl RingConsumer { } // Build chain (head + tails), tracking readable/writable split inline. - let mut elements = SmallVec::<[BufferElement; 16]>::new(); + let mut elements = SmallVec::<[BufferElement; 4]>::new(); let mut pos = self.avail_cursor; let mut chain_len: u16 = 1; @@ -1309,14 +1338,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,16 +1431,18 @@ 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::ptr::{self, NonNull}; + use core::sync::atomic::{AtomicBool, AtomicU16, AtomicUsize, Ordering}; use bytemuck::{Pod, Zeroable}; use super::super::align_up; use super::*; use crate::virtq::event::EventSuppression; + use crate::virtq::{BufferLease, BufferMap}; /// Test MemOps implementation that maintains pointer provenance. /// @@ -1448,6 +1492,72 @@ pub(crate) mod tests { } } + pub struct TestMapping { + creator: std::thread::ThreadId, + owner: core::mem::ManuallyDrop, + } + + struct TestBufferOwner { + data: NonNull<[u8]>, + _mem: TestMem, + _lease: BufferLease, + } + + // SAFETY: The view is immutable. Drop checks the creator thread before + // destroying the owner containing the Rc-backed lease. + unsafe impl Send for TestMapping {} + + impl AsRef<[u8]> for TestMapping { + fn as_ref(&self) -> &[u8] { + // SAFETY: The mapping owns the backing and its initialized immutable prefix. + unsafe { self.owner.data.as_ref() } + } + } + + impl Drop for TestMapping { + fn drop(&mut self) { + assert_eq!( + self.creator, + std::thread::current().id(), + "mapping dropped on another thread" + ); + // SAFETY: The creator thread releases the mapping before its lease. + unsafe { core::mem::ManuallyDrop::drop(&mut self.owner) }; + } + } + + impl BufferMap for TestMem { + type Mapping = TestMapping; + + unsafe fn map_buffer( + &self, + lease: BufferLease, + written: usize, + ) -> Result { + let alloc = lease.allocation(); + assert!(written <= alloc.len as usize); + + let offset = alloc.addr.checked_sub(self.base_addr()).unwrap() as usize; + + // SAFETY: Test storage is never resized. + let size = unsafe { &*self.inner.storage.get() }.len(); + assert!(offset.checked_add(alloc.len as usize).unwrap() <= size); + + // SAFETY: The caller owns the allocation and excludes writes. The + // cloned TestMem retains its stable backing after this borrow. + let data = NonNull::from(unsafe { self.as_slice(alloc.addr, written)? }); + + Ok(TestMapping { + creator: std::thread::current().id(), + owner: core::mem::ManuallyDrop::new(TestBufferOwner { + data, + _mem: self.clone(), + _lease: lease, + }), + }) + } + } + // SAFETY: TestMem translates addresses into its owned backing storage. Unit // tests construct layouts within that storage and avoid concurrent access. unsafe impl MemOps for TestMem { @@ -1502,6 +1612,145 @@ pub(crate) mod tests { } } + /// Shared fault injection over real test memory and buffer owners. + #[derive(Clone)] + pub(crate) struct FaultMem(pub Arc); + + pub(crate) struct FaultState { + pub mem: TestMem, + pub writes: AtomicUsize, + pub fail_write_at: AtomicUsize, + pub fail_mapping_at: AtomicUsize, + pub map_calls: AtomicUsize, + pub deny_views: AtomicBool, + } + + impl FaultMem { + pub(crate) fn new(mem: TestMem) -> Self { + Self(Arc::new(FaultState { + mem, + fail_write_at: AtomicUsize::new(usize::MAX), + writes: AtomicUsize::new(0), + fail_mapping_at: AtomicUsize::new(usize::MAX), + map_calls: AtomicUsize::new(0), + deny_views: AtomicBool::new(false), + })) + } + + /// Reset the write count and fail at this zero-based index. + /// Byte writes and release stores share the count. + pub(crate) fn fail_write_at(&self, write: usize) { + self.0.writes.store(0, Ordering::Relaxed); + self.0.fail_write_at.store(write, Ordering::Relaxed); + } + + pub(crate) fn allow_writes(&self) { + self.fail_write_at(usize::MAX); + } + + /// Fail the zero-based mapping attempt after resetting its counter. + pub(crate) fn fail_mapping_at(&self, mapping: usize) { + self.0.map_calls.store(0, Ordering::Relaxed); + self.0.fail_mapping_at.store(mapping, Ordering::Relaxed); + } + + /// Reject borrowed slices and owned mappings. + pub(crate) fn deny_views(&self) { + self.0.deny_views.store(true, Ordering::Relaxed); + } + + fn check_write(&self) -> Result<(), ()> { + let write = self.0.writes.fetch_add(1, Ordering::Relaxed); + if write == self.0.fail_write_at.load(Ordering::Relaxed) { + Err(()) + } else { + Ok(()) + } + } + } + + // SAFETY: Successful operations delegate to TestMem with the same address + // and ownership preconditions. Injected failures perform no memory access. + unsafe impl MemOps for FaultMem { + type Error = (); + + fn read(&self, addr: u64, dst: &mut [u8]) -> Result<(), Self::Error> { + self.0.mem.read(addr, dst).map_err(|never| match never {}) + } + + fn write(&self, addr: u64, src: &[u8]) -> Result<(), Self::Error> { + self.check_write()?; + self.0.mem.write(addr, src).map_err(|never| match never {}) + } + + fn load_acquire(&self, addr: u64) -> Result { + self.0 + .mem + .load_acquire(addr) + .map_err(|never| match never {}) + } + + fn store_release(&self, addr: u64, val: u16) -> Result<(), Self::Error> { + self.check_write()?; + self.0 + .mem + .store_release(addr, val) + .map_err(|never| match never {}) + } + + unsafe fn as_slice(&self, addr: u64, len: usize) -> Result<&[u8], Self::Error> { + if self.0.deny_views.load(Ordering::Relaxed) { + return Err(()); + } + // SAFETY: The caller supplies TestMem's immutable slice preconditions. + unsafe { self.0.mem.as_slice(addr, len) }.map_err(|never| match never {}) + } + + unsafe fn as_mut_slice(&self, addr: u64, len: usize) -> Result<&mut [u8], Self::Error> { + if self.0.deny_views.load(Ordering::Relaxed) { + return Err(()); + } + // SAFETY: The caller supplies exclusive access to this range. + unsafe { self.0.mem.as_mut_slice(addr, len) }.map_err(|never| match never {}) + } + } + + /// Retains the backend generation for lifetime assertions. + pub(crate) struct TrackedMapping { + data: TestMapping, + _mem: FaultMem, + } + + impl AsRef<[u8]> for TrackedMapping { + fn as_ref(&self) -> &[u8] { + self.data.as_ref() + } + } + + impl BufferMap for FaultMem { + type Mapping = TrackedMapping; + + unsafe fn map_buffer( + &self, + lease: BufferLease, + written: usize, + ) -> Result { + let call = self.0.map_calls.fetch_add(1, Ordering::Relaxed); + if self.0.deny_views.load(Ordering::Relaxed) + || call == self.0.fail_mapping_at.load(Ordering::Relaxed) + { + return Err(()); + } + // SAFETY: The caller supplies TestMem's mapping preconditions. + let data = unsafe { BufferMap::map_buffer(&self.0.mem, lease, written) } + .map_err(|never| match never {})?; + Ok(TrackedMapping { + data, + _mem: self.clone(), + }) + } + } + /// Owns the descriptor table and event suppression structures pub struct OwnedRing { mem: TestMem, @@ -3214,7 +3463,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 +3484,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 +3494,50 @@ pub(crate) mod tests { } } + #[test] + fn test_ring_producer_failed_reset_preserves_local_state() { + let ring = make_ring(4); + let mem = FaultMem::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_write_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 +3553,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 +3575,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 +4412,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..cb56f8da3 --- /dev/null +++ b/src/hyperlight_common/src/virtq/ring/canonical.rs @@ -0,0 +1,598 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! 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; 4]>::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..878b84b80 --- /dev/null +++ b/src/hyperlight_common/src/virtq/ring/fuzz.rs @@ -0,0 +1,191 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +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); +} diff --git a/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs b/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs index 8b5110bb2..bc3147f73 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 prev_base.checked_add(nbytes).is_none_or(|end| end > 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..bc3bf6f97 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 x.checked_add(nbytes).is_none_or(|end| end > limit) { unsafe { crate::exit::abort_with_code_and_message( &[ErrorCode::MallocFailed as u8], 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 ecb8f9d43..3650150d2 100644 --- a/src/hyperlight_guest/src/layout.rs +++ b/src/hyperlight_guest/src/layout.rs @@ -6,24 +6,45 @@ 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_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) +} +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_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) +} +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_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..48c24218f --- /dev/null +++ b/src/hyperlight_guest/src/transport/context.rs @@ -0,0 +1,144 @@ +// 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. +/// +/// Each slot becomes one independent preposted receive buffer. +fn h2g_pool(base: u64, pages: usize, buffer_size: usize) -> result::Result { + if buffer_size == 0 { + return Err(AllocError::InvalidArg); + } + 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 { + if upper_size == 0 { + return Err(AllocError::InvalidArg); + } + 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..0d66c01f9 --- /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 (size, pages, g2h_bufsz) = read_published_g2h(); + let g2h = QueueDims::new(size, pages).expect("invalid G2H 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); + + 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.size()) }.expect("G2H layout is invalid"); + let h2g_layout = + 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( + 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 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 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"); + + (size, pages, bufsz) +} + +fn read_published_h2g() -> (usize, usize, usize) { + // SAFETY: Generic initialization has mapped writable scratch metadata. + 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 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"); + + (size, pages, bufsz) +} diff --git a/src/hyperlight_host/src/mem/layout.rs b/src/hyperlight_host/src/mem/layout.rs index 22b0abc37..11ea68cbe 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::{QueueDims, TransportArena}; use hyperlight_common::mem::HyperlightPEB; use hyperlight_common::vmem::PAGE_SIZE; use tracing::{Span, instrument}; @@ -250,6 +255,16 @@ pub(crate) struct SandboxMemoryLayout { init_data_permissions: Option, /// The size of the scratch region in physical memory. scratch_size: usize, + /// G2H ring and buffer pool dimensions. + g2h_dims: QueueDims, + /// H2G ring and buffer pool dimensions. + h2g_dims: QueueDims, + /// Capacity of each G2H upper-tier buffer. + g2h_buffer_size: usize, + /// Capacity of each H2G buffer. + h2g_buffer_size: usize, + /// Fixed ring and pool placement within scratch. + transport_arena: TransportArena, /// 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 +299,12 @@ impl Debug for SandboxMemoryLayout { &format_args!("{:#x}", self.output_data_size), ) .field("Scratch Size", &format_args!("{:#x}", self.scratch_size)) + .field("G2H Queue Size", &self.get_g2h_queue_size()) + .field("H2G Queue Size", &self.get_h2g_queue_size()) + .field("G2H Buffer Size", &self.g2h_buffer_size) + .field("H2G Buffer Size", &self.h2g_buffer_size) + .field("G2H Pool Pages", &self.get_g2h_pool_pages()) + .field("H2G Pool Pages", &self.get_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( @@ -332,10 +353,38 @@ 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 min_scratch_size = - hyperlight_common::layout::min_scratch_size(input_data_size, output_data_size); + 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(); + let h2g_pool_pages = cfg.get_h2g_pool_pages(); + + let g2h_dims = QueueDims::new(g2h_queue_size, g2h_pool_pages) + .ok_or(MemoryRequestTooSmall(scratch_size, usize::MAX))?; + let h2g_dims = QueueDims::new(h2g_queue_size, h2g_pool_pages) + .ok_or(MemoryRequestTooSmall(scratch_size, usize::MAX))?; + let io_len = input_data_size + .checked_add(output_data_size) + .and_then(|len| len.checked_next_multiple_of(PAGE_SIZE)) + .ok_or(MemoryRequestTooSmall(scratch_size, usize::MAX))?; + let arena_base_gpa = hyperlight_common::layout::scratch_base_gpa(scratch_size) + .checked_add(io_len as u64) + .ok_or(MemoryRequestTooSmall(scratch_size, usize::MAX))?; + let transport_arena = TransportArena::new(arena_base_gpa, g2h_dims, h2g_dims) + .ok_or(MemoryRequestTooSmall(scratch_size, usize::MAX))?; + let min_scratch_size = hyperlight_common::layout::min_scratch_size( + input_data_size, + output_data_size, + transport_arena.size(), + ); if scratch_size < min_scratch_size { return Err(MemoryRequestTooSmall(scratch_size, min_scratch_size)); } @@ -349,6 +398,11 @@ impl SandboxMemoryLayout { init_data_permissions, pt_size: None, scratch_size, + g2h_dims, + h2g_dims, + g2h_buffer_size, + h2g_buffer_size, + transport_arena, snapshot_size: 0, }; ret.set_snapshot_size(ret.get_memory_size()?); @@ -383,6 +437,44 @@ impl SandboxMemoryLayout { self.scratch_size } + #[allow(dead_code)] + pub(crate) fn get_g2h_queue_size(&self) -> usize { + usize::from(self.g2h_dims.size().get()) + } + + #[allow(dead_code)] + pub(crate) fn get_h2g_queue_size(&self) -> usize { + usize::from(self.h2g_dims.size().get()) + } + + #[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_dims.pool_pages().get() + } + + #[allow(dead_code)] + pub(crate) fn get_h2g_pool_pages(&self) -> usize { + self.h2g_dims.pool_pages().get() + } + + pub(crate) fn get_g2h_queue_dims(&self) -> QueueDims { + self.g2h_dims + } + + pub(crate) fn get_h2g_queue_dims(&self) -> QueueDims { + self.h2g_dims + } + /// Guest-visible prefix size of the snapshot blob. pub(crate) fn snapshot_size(&self) -> usize { self.snapshot_size @@ -408,8 +500,9 @@ impl SandboxMemoryLayout { let min_fixed_scratch = hyperlight_common::layout::min_scratch_size( self.input_data_size, self.output_data_size, + self.transport_arena.size(), ); - 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)); } @@ -670,21 +763,25 @@ 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_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(self.scratch_size)) + as usize } /// Base GPA to which the page tables are eagerly copied on restore. pub(crate) fn get_pt_base_gpa(&self) -> u64 { - hyperlight_common::layout::scratch_base_gpa(self.scratch_size) - + self.get_pt_base_scratch_offset() as u64 + self.transport_arena.end_addr() } - /// 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 } + /// Exact transport placement in the fixed scratch prefix. + pub(crate) fn get_transport_arena(&self) -> TransportArena { + self.transport_arena + } + /// 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 @@ -740,6 +837,88 @@ mod tests { ); } + #[test] + fn transport_arena_after_unaligned_io_buffers() { + let mut cfg = SandboxConfiguration::default(); + cfg.set_input_data_size(0x4001); + cfg.set_output_data_size(0x2001); + let mut layout = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap(); + let arena = layout.get_transport_arena(); + let scratch_base = hyperlight_common::layout::scratch_base_gpa(layout.get_scratch_size()); + + assert_eq!(arena.base_addr(), scratch_base + 0x7000); + assert_eq!(layout.get_pt_base_gpa(), arena.end_addr()); + + layout.set_pt_size(PAGE_SIZE).unwrap(); + + assert_eq!(layout.get_transport_arena(), arena); + assert_eq!( + layout.get_first_free_scratch_gpa(), + arena.end_addr() + PAGE_SIZE as u64 + ); + } + + #[test] + fn transport_memory_is_part_of_minimum_scratch_size() { + let mut cfg = SandboxConfiguration::default(); + let layout = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap(); + let minimum = hyperlight_common::layout::min_scratch_size( + cfg.get_input_data_size(), + cfg.get_output_data_size(), + layout.get_transport_arena().size(), + ); + 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() { + for (g2h_pages, h2g_pages) in [ + (usize::MAX, 4), + (8, usize::MAX), + (usize::MAX / PAGE_SIZE, 1), + ] { + let mut cfg = SandboxConfiguration::default(); + cfg.set_g2h_pool_pages(g2h_pages); + cfg.set_h2g_pool_pages(h2g_pages); + + let layout = SandboxMemoryLayout::new(cfg, 4096, 0, None); + assert!(matches!(layout, Err(MemoryRequestTooSmall(_, usize::MAX)))); + } + } + + #[test] + fn transport_minimum_rejects_io_overflow() { + for input_size in [usize::MAX, usize::MAX - 0x2000, usize::MAX - 0x5000 + 1] { + let mut cfg = SandboxConfiguration::default(); + cfg.set_input_data_size(input_size); + cfg.set_output_data_size(0x2000); + + let layout = SandboxMemoryLayout::new(cfg, 4096, 0, None); + 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(); @@ -811,7 +990,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,12 +1000,20 @@ 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); 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. @@ -840,12 +1027,12 @@ 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), - 0x4000 + layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x20000), + 0x11000 ); // pt_size is zero here, so the first free scratch GPA equals // the page table base. @@ -860,7 +1047,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,12 +1060,20 @@ 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); 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(), @@ -887,12 +1082,12 @@ 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), - 0x6000 + layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x30000), + 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..30ec8a442 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, } } @@ -283,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 { @@ -358,6 +351,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 +362,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 +371,81 @@ 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 + /// 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(()) + } + + /// 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 { @@ -473,6 +545,13 @@ impl SandboxMemoryManager { Option>, Option, )> { + if let Some(virtq) = snapshot.virtq() { + virtq.preflight(snapshot.layout())?; + } + + 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, @@ -515,6 +594,7 @@ impl SandboxMemoryManager { self.original_entrypoint = snapshot.original_entrypoint(); self.update_scratch_bookkeeping()?; + self.restore_virtq(snapshot.virtq())?; Ok((gsnapshot, gscratch)) } @@ -561,6 +641,38 @@ impl SandboxMemoryManager { self.snapshot_count, )?; + // Record the G2H and H2G queue sizes, pool page counts, and buffer sizes. + self.update_scratch_bookkeeping_item( + 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, + 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_SIZE_OFFSET, + u64::try_from(self.layout.get_h2g_queue_size())?, + )?; + 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 96e784acd..693dc2af7 100644 --- a/src/hyperlight_host/src/mem/mod.rs +++ b/src/hyperlight_host/src/mem/mod.rs @@ -25,3 +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/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.rs b/src/hyperlight_host/src/mem/virtq.rs new file mode 100644 index 000000000..99ecd82dc --- /dev/null +++ b/src/hyperlight_host/src/mem/virtq.rs @@ -0,0 +1,649 @@ +// 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::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, ImageMem}; +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 { 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), + )) +} + +/// Capture the canonical transport state omitted from the main memory snapshot. +pub(crate) fn snapshot( + layout: &SandboxMemoryLayout, + scratch_mem: &HostSharedMemory, +) -> Result { + let validator = Validator { 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 { 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, + h2g_ring: Range, + g2h_pool: Range, + h2g_pool: Range, +} + +/// 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 { layout }.validate_snapshot(self).map(|_| ()) + } +} + +/// Validates live and captured transport images against one host layout. +struct Validator<'a> { + layout: &'a SandboxMemoryLayout, +} + +impl Validator<'_> { + /// Validate the initial G2H queue and return its layout. + fn validate_g2h(&self, mem: &M, ring: Range) -> Result { + let dims = self.layout.get_g2h_queue_dims(); + + // 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, dims.size()) } + .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 { + let dims = self.layout.get_h2g_queue_dims(); + + // 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, dims.size()) } + .map_err(|error| new_error!("invalid H2G ring layout: {error}"))?; + + // SandboxConfiguration guarantees a nonzero buffer size. + let bufsz = self.layout.get_h2g_buffer_size(); + let prefill = usize::from(dims.size().get()).min(dims.pool_len() / bufsz); + + 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.layout.get_transport_arena().base_addr() { + return Err(new_error!("published transport arena is invalid")); + } + + 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.layout.get_g2h_queue_dims().ring_len(), + )?; + validate_ring_len( + "H2G", + &snapshot.h2g_ring, + self.layout.get_h2g_queue_dims().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) + } + + fn scratch_gva(&self, gpa: u64) -> Result { + 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")) + } + + /// Translate validated transport GPAs into the GVA ranges used by descriptors. + fn resolve_gva_regions(&self) -> Result { + let arena = self.layout.get_transport_arena(); + let g2h = self.layout.get_g2h_queue_dims(); + let h2g = self.layout.get_h2g_queue_dims(); + + Ok(GvaRegions { + g2h_ring: checked_region(self.scratch_gva(arena.g2h_ring_addr())?, g2h.ring_len())?, + h2g_ring: checked_region(self.scratch_gva(arena.h2g_ring_addr())?, h2g.ring_len())?, + g2h_pool: checked_region(self.scratch_gva(arena.g2h_pool_addr())?, g2h.pool_len())?, + h2g_pool: checked_region(self.scratch_gva(arena.h2g_pool_addr())?, h2g.pool_len())?, + }) + } +} + +/// 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)?) +} + +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) -> Result> { + let end = start + .checked_add(u64::try_from(len)?) + .ok_or_else(|| new_error!("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_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); + SandboxMemoryLayout::new(config, 4096, 0, None).unwrap() + } + + fn host_scratch() -> HostSharedMemory { + let scratch = ExclusiveSharedMemory::new(SCRATCH_SIZE).unwrap(); + scratch.build().0 + } + + struct PreparedVirtq { + scratch: HostSharedMemory, + 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 arena = layout.get_transport_arena(); + 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(arena.g2h_ring_addr()); + let h2g_base = to_gva(arena.h2g_ring_addr()); + let g2h_pool_base = to_gva(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(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, 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 { + scratch, + 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 { layout: &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 layout = memory_layout(); + let validator = Validator { layout: &layout }; + let regions = validator + .validate_published_arena(layout.get_transport_arena().base_addr()) + .unwrap(); + + assert_eq!( + regions.g2h_ring.end - regions.g2h_ring.start, + layout.get_g2h_queue_dims().ring_len() as u64 + ); + assert_eq!( + regions.h2g_ring.end - regions.h2g_ring.start, + layout.get_h2g_queue_dims().ring_len() as u64 + ); + assert_eq!( + regions.g2h_pool.end - regions.g2h_pool.start, + layout.get_g2h_queue_dims().pool_len() as u64 + ); + assert_eq!( + regions.h2g_pool.end - regions.h2g_pool.start, + layout.get_h2g_queue_dims().pool_len() as u64 + ); + } + + #[test] + fn rejects_invalid_published_regions() { + let layout = memory_layout(); + let validator = Validator { layout: &layout }; + let arena_gpa = layout.get_transport_arena().base_addr() + 1; + assert!(validator.validate_published_arena(arena_gpa).is_err()); + } + + #[test] + fn rejects_published_region_overflow() { + let layout = memory_layout(); + let validator = Validator { layout: &layout }; + assert!(validator.validate_published_arena(u64::MAX).is_err()); + } + + #[test] + fn rejects_untranslatable_or_overflowing_gva_regions() { + let layout = memory_layout(); + let validator = Validator { layout: &layout }; + let invalid = hyperlight_common::layout::scratch_base_gpa(SCRATCH_SIZE) - 1; + assert!(validator.scratch_gva(invalid).is_err()); + + let arena_gva = validator + .scratch_gva(layout.get_transport_arena().base_addr()) + .unwrap(); + assert!(checked_region(arena_gva, usize::MAX).is_err()); + } + + #[test] + fn validates_initial_virtq_images() { + 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(); + 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/mem/virtq_mem.rs b/src/hyperlight_host/src/mem/virtq_mem.rs new file mode 100644 index 000000000..944eee1b7 --- /dev/null +++ b/src/hyperlight_host/src/mem/virtq_mem.rs @@ -0,0 +1,260 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright 2026 The Hyperlight Authors. + +//! Host [`MemOps`] implementations for live scratch and captured ring images. +//! +//! 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; +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")) + } +} + +/// 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; + + 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()); + } +} diff --git a/src/hyperlight_host/src/sandbox/config.rs b/src/hyperlight_host/src/sandbox/config.rs index ceb6107f9..83e8f42c7 100644 --- a/src/hyperlight_host/src/sandbox/config.rs +++ b/src/hyperlight_host/src/sandbox/config.rs @@ -5,6 +5,8 @@ use std::cmp::max; use std::time::Duration; use hyperlight_common::log_level::GuestLogFilter; +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}; @@ -82,6 +84,18 @@ pub struct SandboxConfiguration { /// Stored as the guest ABI's numeric log-filter value, with `u64::MAX` /// representing an unset value, to keep this `#[repr(C)]` struct FFI-safe. max_guest_log_level: u64, + /// Number of descriptors in the G2H virtqueue. + g2h_queue_size: usize, + /// Number of descriptors in the H2G virtqueue. + h2g_queue_size: 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], @@ -105,8 +119,29 @@ impl SandboxConfiguration { pub const INTERRUPT_VCPU_SIGRTMIN_OFFSET: u8 = 0; /// 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; + // TODO: Reassess scratch sizing when virtqueues replace the input/output regions. + /// The default scratch size, aligned to 16 KiB for macOS hosts. + pub const DEFAULT_SCRATCH_SIZE: usize = 0x58000; + /// The default G2H virtqueue descriptor count. + pub const DEFAULT_G2H_QUEUE_SIZE: usize = 64; + /// The default H2G virtqueue descriptor count. + 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. + 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_SIZE: usize = 2; + /// The maximum G2H virtqueue descriptor count. + 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. + 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. @@ -133,6 +168,12 @@ impl SandboxConfiguration { heap_size_override: heap_size_override.unwrap_or(0), scratch_size, max_guest_log_level: Self::MAX_GUEST_LOG_LEVEL_UNSET, + 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, + h2g_pool_pages: Self::DEFAULT_H2G_POOL_PAGES, interrupt_retry_delay, interrupt_vcpu_sigrtmin_offset, #[cfg(gdb)] @@ -316,6 +357,98 @@ impl SandboxConfiguration { } } + /// Get the G2H virtqueue descriptor count. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + 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_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_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_size(&mut self, size: usize) { + self.h2g_queue_size = Self::normalize_queue_size(size); + } + + /// 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 { @@ -340,6 +473,19 @@ impl SandboxConfiguration { self.heap_size_override_opt() .unwrap_or(Self::DEFAULT_HEAP_SIZE) } + + fn normalize_queue_size(size: usize) -> usize { + size.clamp(Self::MIN_QUEUE_SIZE, Self::MAX_QUEUE_SIZE) + .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 { @@ -362,6 +508,7 @@ impl Default for SandboxConfiguration { #[cfg(test)] mod tests { + use hyperlight_common::vmem::PAGE_SIZE; use tracing_core::LevelFilter; #[cfg(target_arch = "x86_64")] @@ -386,6 +533,12 @@ mod tests { } } + #[test] + fn default_scratch_size_supports_16k_pages() { + let cfg = SandboxConfiguration::default(); + assert!(cfg.get_scratch_size().is_multiple_of(16 * 1024)); + } + #[test] #[cfg(target_arch = "x86_64")] fn guest_msrs_reports_overflow() { @@ -472,6 +625,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_SIZE, + cfg.get_g2h_queue_size() + ); + assert_eq!( + SandboxConfiguration::DEFAULT_H2G_QUEUE_SIZE, + cfg.get_h2g_queue_size() + ); + 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] @@ -499,6 +676,71 @@ mod tests { assert_eq!(SandboxConfiguration::MIN_OUTPUT_SIZE, cfg.output_data_size); } + #[test] + fn queue_sizes_are_normalized() { + let mut cfg = SandboxConfiguration::default(); + for (size, 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_size(size); + cfg.set_h2g_queue_size(size); + assert_eq!(expected, cfg.get_g2h_queue_size()); + assert_eq!(expected, cfg.get_h2g_queue_size()); + } + } + + #[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 55b6908db..50f2b2aef 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. @@ -251,11 +250,21 @@ 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_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()); + config.set_h2g_pool_pages(snapshot.layout().get_h2g_pool_pages()); let max_guest_log_level = config.get_max_guest_log_level(); let load_info = snapshot.load_info(); 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; @@ -345,6 +354,12 @@ impl MultiUseSandbox { })?; } + if attach_virtq { + hshm.attach_virtq()?; + } else { + hshm.restore_virtq(snapshot.virtq())?; + } + let sbox = MultiUseSandbox::from_uninit(host_funcs, hshm, vm); Ok(sbox) } @@ -1166,6 +1181,36 @@ fn warn_on_layout_override( caller.get_scratch_size() as u64, snapshot.get_scratch_size() as u64, ), + ( + "g2h_queue_size", + caller.get_g2h_queue_size() as u64, + snapshot.get_g2h_queue_size() as u64, + ), + ( + "h2g_queue_size", + caller.get_h2g_queue_size() as u64, + snapshot.get_h2g_queue_size() 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 { @@ -1188,8 +1233,6 @@ mod tests { use hyperlight_testing::sandbox_sizes::{LARGE_HEAP_SIZE, MEDIUM_HEAP_SIZE, SMALL_HEAP_SIZE}; use hyperlight_testing::{c_simple_guest_as_pathbuf, simple_guest_as_pathbuf}; - #[cfg(any(target_arch = "x86_64", feature = "trace_guest"))] - use crate::MultiUseSandbox; use crate::func::host_functions::Registerable; #[cfg(not(gdb))] use crate::hypervisor::hyperlight_vm::test_support::VmOperation; @@ -1198,7 +1241,8 @@ mod tests { use crate::sandbox::SandboxConfiguration; use crate::sandbox::uninitialized::{GuestBlob, GuestEnvironment}; use crate::{ - GuestBinary, HyperlightError, Result, SandboxBuilder, SandboxStatus, UninitializedSandbox, + GuestBinary, HyperlightError, MultiUseSandbox, Result, SandboxBuilder, SandboxStatus, + UninitializedSandbox, }; #[test] @@ -1216,6 +1260,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()) @@ -1381,23 +1430,23 @@ 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(); + let layout = + crate::mem::layout::SandboxMemoryLayout::new(defaults, 0, 0, None).unwrap(); hyperlight_common::layout::min_scratch_size( defaults.get_input_data_size(), defaults.get_output_data_size(), + layout.get_transport_arena().size(), ) - } + 0x10000 - + 0x10000; + .next_multiple_of(page_size::get()) + } + 0x40000; let mut sbox1 = SandboxBuilder::from_file(simple_guest_as_pathbuf()) .heap_size(HEAP_SIZE) @@ -1749,6 +1798,7 @@ mod tests { let snapshot = sandbox.snapshot().unwrap(); sandbox2.restore(snapshot).unwrap(); + assert_virtq_attached(&sandbox2); assert_eq!(sandbox2.call::("GetStatic", ()).unwrap(), 42); } @@ -2117,7 +2167,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() @@ -2126,7 +2176,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() @@ -2147,7 +2197,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() @@ -2156,7 +2206,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() @@ -2164,18 +2214,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(); @@ -2186,6 +2238,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(); @@ -2212,7 +2266,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() @@ -2242,7 +2296,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); @@ -2250,7 +2304,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] @@ -4659,7 +4713,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"); @@ -4671,6 +4727,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(); @@ -4767,6 +4824,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); @@ -4776,6 +4835,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..b50fca035 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_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, + 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_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); + transport.set_h2g_pool_pages(self.layout.h2g_pool_pages); + + for (name, saved, normalized) in [ + ( + "g2h_queue_size", + self.layout.g2h_queue_size, + transport.get_g2h_queue_size(), + ), + ( + "h2g_queue_size", + self.layout.h2g_queue_size, + transport.get_h2g_queue_size(), + ), + ( + "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_size: 64, + h2g_queue_size: 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_size": 64, + "h2g_queue_size": 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_size": 64, + "h2g_queue_size": 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..84fd33723 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_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(), + 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_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); + 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..f9900ef92 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" ); } @@ -2767,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() @@ -2827,6 +2843,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_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); + 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_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); + 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/src/sandbox/uninitialized.rs b/src/hyperlight_host/src/sandbox/uninitialized.rs index ec8eb14bf..5a73b261d 100644 --- a/src/hyperlight_host/src/sandbox/uninitialized.rs +++ b/src/hyperlight_host/src/sandbox/uninitialized.rs @@ -1161,6 +1161,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); @@ -1183,7 +1184,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); @@ -1206,6 +1207,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 @@ -1231,7 +1233,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/src/sandbox/uninitialized_evolve.rs b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs index f04a7ad60..73bb19235 100644 --- a/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs +++ b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs @@ -23,6 +23,10 @@ use crate::{MultiUseSandbox, Result, UninitializedSandbox}; pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result { let max_guest_log_level = u_sbox.config.get_max_guest_log_level(); 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 @@ -86,6 +90,10 @@ pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result 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); 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)