diff --git a/CHANGELOG.md b/CHANGELOG.md index 0546852eb..6bd599cea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ## [Prerelease] - Unreleased ### Added +* Add per-direction virtqueue configuration and account its allocations in + scratch sizing. ### Changed * **Breaking:** Guest MSR state is now saved and restored across snapshots. @@ -13,10 +15,26 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). resets to a clean default. On KVM the guest may only read or write declared MSRs, on MSHV and WHP this is not enforced. by @ludfjig in https://github.com/hyperlight-dev/hyperlight/pull/991 * **Breaking:** Filesystem paths are now represented using `PathBuf`. `GuestBinary::FilePath` now stores a `PathBuf` instead of a `String`, and `MultiUseSandbox::generate_crashdump_to_dir` accepts `Into` instead of `Into`. Callers passing a `String` to `GuestBinary::FilePath` must convert it using `.into()`. +* Expose C guest `ByteChunks` values as pointer and length arrays. +* Place virtqueue rings and pools in host-owned scratch before page tables. + Snapshot ABI 2 rejects snapshots created with earlier layouts. +* Require guest logs and all host and guest function calls to use virtqueues. +* Keep registered Rust guest return values typed until transport encoding so + external byte results avoid intermediate FlatBuffer copies. +* Store canonical virtqueue rings in versioned OCI transport layers. Config v2 + rejects snapshots without transport state. +* Running snapshots checkpoint dirty virtqueues before capture. Ordinary calls + keep their deferred result path. +* Reject snapshot capture while guest-owned transport buffers are retained. +* Use the reclaimed stack pages to raise the default G2H and H2G pools to 12 + and 8 pages. ### Removed +* Remove legacy stack I/O, its `GuestHandle` methods, and its sandbox + configuration options. ### Fixed +* Keep sandboxes usable after an H2G request exceeds available virtqueue capacity. * Fix symbol resolution in guest core dumps for sandboxes created from snapshots by @ludfjig in https://github.com/hyperlight-dev/hyperlight/pull/1618 * Reject malformed OCI snapshot metadata and non-regular artifact files during load. diff --git a/Justfile b/Justfile index f69d88c41..aaf1b4de9 100644 --- a/Justfile +++ b/Justfile @@ -240,7 +240,7 @@ test-loom: # runs tests that requires being run separately, for example due to global state test-isolated target=default-target features="" : {{ cargo-cmd }} test {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" { "dev" } else { target } }} {{ target-triple-flag }} -p hyperlight-host --lib -- sandbox::uninitialized::tests::test_log_trace --exact --ignored - {{ cargo-cmd }} test {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" { "dev" } else { target } }} {{ target-triple-flag }} -p hyperlight-host --lib -- sandbox::outb::tests::test_log_outb_log --exact --ignored + {{ cargo-cmd }} test {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" { "dev" } else { target } }} {{ target-triple-flag }} -p hyperlight-host --lib -- sandbox::outb::tests::test_log_emit_guest_log --exact --ignored {{ cargo-cmd }} test {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" { "dev" } else { target } }} {{ target-triple-flag }} -p hyperlight-host --test integration_test -- log_message --exact --ignored @# CPU vendor check, gated to known CI runner hardware {{ cargo-cmd }} test {{ if features =="" {''} else if features=="no-default-features" {"--no-default-features" } else {"--no-default-features -F " + features } }} --profile={{ if target == "debug" { "dev" } else { target } }} {{ target-triple-flag }} -p hyperlight-host --lib -- sandbox::snapshot::file::config::tests::cpu_vendor_current_is_recognized --exact --ignored @@ -524,7 +524,7 @@ coverage-run hypervisor="kvm": ensure-cargo-llvm-cov # isolated tests (require running separately due to global state) cargo +nightly test -p hyperlight-host --lib -- sandbox::uninitialized::tests::test_log_trace --exact --ignored - cargo +nightly test -p hyperlight-host --lib -- sandbox::outb::tests::test_log_outb_log --exact --ignored + cargo +nightly test -p hyperlight-host --lib -- sandbox::outb::tests::test_log_emit_guest_log --exact --ignored cargo +nightly test -p hyperlight-host --test integration_test -- log_message --exact --ignored cargo +nightly test -p hyperlight-host --no-default-features -F function_call_metrics,{{ if hypervisor == "mshv3" { "mshv3" } else { "kvm" } }} --lib -- metrics::tests::test_metrics_are_emitted --exact diff --git a/docs/README.md b/docs/README.md index d99173756..9682b9269 100644 --- a/docs/README.md +++ b/docs/README.md @@ -35,6 +35,7 @@ This project is composed internally of several components, depicted in the below * [Security guidance for developers](./security-guidance-for-developers.md) * [Paging Development Notes](./paging-development-notes.md) +* [Virtqueue host and guest communication](./virtio-host-guest-communication.md) * [How to debug a Hyperlight guest](./how-to-debug-a-hyperlight-guest.md) * [How to use Flatbuffers in Hyperlight](./how-to-use-flatbuffers.md) * [How to make a Hyperlight release](./how-to-make-releases.md) diff --git a/docs/paging-development-notes.md b/docs/paging-development-notes.md index da08f6f24..7d2bbbdcd 100644 --- a/docs/paging-development-notes.md +++ b/docs/paging-development-notes.md @@ -139,13 +139,10 @@ calls, i.e. there may be no calls in flight at the time of snapshotting. This is not enforced, but odd things may happen if it is violated. -Buffer management between the host and guest is needed to pass call -arguments and return values. Ideally, buffers would be dynamically -allocated from the scratch region as needed. - -Currently, I/O buffers are statically allocated at the bottom of the -scratch region. This is a stopgap pending improved -physical allocation and buffer management. +Host and guest calls use two virtqueues in a fixed transport arena at +the bottom of scratch. The arena contains both rings and their +fixed-slot buffer pools. Copied page tables follow the arena. Dynamic +scratch allocations begin after the copied page tables. The minimum scratch size is calculated by `min_scratch_size()` in the architecture-specific layout modules under `hyperlight_common`; see @@ -177,4 +174,3 @@ paging) and enables PAE. The guest is always entered in long mode. Hyperlight unconditionally uses 48-bit virtual addresses. Hyperlight presently only uses addresses in the lower (ttbr0) half of the address range. - diff --git a/docs/snapshot-oci-format.md b/docs/snapshot-oci-format.md index 971b3c868..c66a2d2b8 100644 --- a/docs/snapshot-oci-format.md +++ b/docs/snapshot-oci-format.md @@ -24,21 +24,28 @@ path/ Hyperlight config JSON raw memory bytes (`memory_size` bytes) + canonical virtqueue rings ``` -Three blob kinds per tag: +Four 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 + two layers by digest. +* **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, and 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. +* **layer / transport** + (`application/vnd.hyperlight.snapshot.transport.v1`). A bounded + binary image of the canonical G2H and H2G rings. + +The runtime queue protocol and canonical checkpoint are described in +[Virtqueue host and guest communication](./virtio-host-guest-communication.md). Blob filenames are the sha256 of the blob bytes, so identical blobs across tags are stored once. @@ -55,8 +62,8 @@ A single saved `Snapshot` consists of exactly: config blob for tooling visibility, * one **manifest** blob (referenced by that index entry), * one **config** blob (referenced by the manifest's `config` field), -* one **layer** blob (the only entry in the manifest's `layers` - array, holding the raw memory image). +* one memory **layer** blob, +* one transport **layer** blob. Saving two snapshots under different tags into the same `path` produces two index entries and two manifests. Configs and layers are @@ -98,12 +105,12 @@ podman), `go-containerregistry` (crane), and `regclient`. ## Read semantics `Snapshot::load(path, reference)` reads a snapshot. It does not check -the manifest, config, or snapshot blobs against their sha256 digests. +the manifest, config, memory, or transport blobs against their sha256 digests. `reference` is an [`OciReference`], either a tag that matches the `org.opencontainers.image.ref.name` annotation or the manifest digest returned by `save`. `Snapshot::checked_load` adds the digest -check on those three blobs, catching accidental corruption on disk. +check on all four blobs, catching accidental corruption on disk. Both run every other check (OCI structure, descriptor sizes, schema versions, arch / hypervisor / CPU vendor / ABI tags, layout bounds, entrypoint bounds). The caller is responsible for trusting the source. diff --git a/docs/snapshot-versioning.md b/docs/snapshot-versioning.md index f855c1576..604575227 100644 --- a/docs/snapshot-versioning.md +++ b/docs/snapshot-versioning.md @@ -7,30 +7,34 @@ existing snapshots loadable, or while rejecting them with a clear error. ## What is versioned -A snapshot carries three independently evolvable version markers: +A snapshot carries four independently evolvable version markers: * **Memory blob ABI**, `SNAPSHOT_ABI_VERSION` (a `u32` inside the config blob, defined in [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. + and `VmAction` port numbers, 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 +* **Transport blob encoding**, `MT_TRANSPORT_V1` + (`application/vnd.hyperlight.snapshot.transport.v1`), aliased as + `MT_TRANSPORT_CURRENT`. This is the binary encoding of canonical + virtqueue state stored outside the memory layer. +* **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 sizes, `abi_version`, `hyperlight_version`, etc.). Renaming a field, changing its type, or adding a required field is a schema change and - bumps this constant. + bumps this constant. Version 2 requires a transport layer. The `OCI_LAYOUT_VERSION` constant is pinned by the OCI image-layout spec at `1.0.0`. @@ -382,4 +386,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/docs/virtio-host-guest-communication.md b/docs/virtio-host-guest-communication.md new file mode 100644 index 000000000..4cf4db82f --- /dev/null +++ b/docs/virtio-host-guest-communication.md @@ -0,0 +1,443 @@ +# Virtqueue host and guest communication + +Hyperlight transports typed function calls over two shared memory VIRTIO +packed virtqueues. It uses the packed ring layout and ownership rules, but it +is not a discoverable VIRTIO device. Queue configuration, arena placement, and +notification behavior are part of the Hyperlight ABI. + +This document describes the runtime transport, snapshot checkpoint, retention +mailbox, and placement constraints. + +## Architecture + +The guest is the driver (producer) for both queues. The host is the device +(consumer) for both queues. + +```text + Guest Host + + G2H producer === G2H packed ring and buffer pool ===> G2H consumer + + H2G producer === H2G packed ring and buffer pool ===> H2G consumer +``` + +Producer ownership describes who publishes descriptors. It does not always +describe the direction in which payload bytes move. + +* **G2H** carries guest requests, guest function results, and logs. Guest + readable descriptors carry bytes to the host. A guest call to a host + function also includes writable descriptors in the same chain for the host + response. +* **H2G** carries host requests to guest functions and internal control + requests. The guest preposts writable buffers. The host fills and completes + them before entering the VM. + +Two queues keep directional validation and capacity independent. H2G always +contains uniform preposted receive buffers. G2H supports readable messages and +optional writable response capacity. + +## Transport arena + +Both rings, the checkpoint mailbox, and both pools occupy one fixed prefix of +guest scratch memory. + +```text + scratch base + | + v + +----------+-----+----------+-----+-----+-----+----------+----------+ + | G2H ring | pad | H2G ring | pad | mbx | pad | G2H pool | H2G pool | + +----------+-----+----------+-----+-----+-----+----------+----------+ +``` + +The host derives this layout from `SandboxConfiguration`. Ring starts follow +packed ring alignment rules. The mailbox is `u64` aligned. Pools are page +aligned. + +The default layout is: + +| Region | Default size or capacity | +|---|---:| +| G2H ring | 64 descriptors | +| H2G ring | 32 descriptors | +| Mailbox | one `u64` | +| G2H pool | 12 pages | +| H2G pool | 8 pages | +| Arena | 21 pages total | + +Offsets after the G2H ring depend on configured queue sizes and pool pages. +`TransportArena` addresses are GPAs. The guest converts them to scratch GVAs +when constructing rings and pools. Descriptor buffer addresses are GVAs. + +The configured upper buffer size is 4 KiB by default. The G2H pool uses two +slot tiers: + +* The first page contains sixteen 256 byte slots for control messages and + logs. +* Complete configured size slots occupy the remaining pages. + +The lower tier is a memory efficiency optimization. Most control messages, +scalar function arguments, and scalar results fit in a small slot. Giving each +of them a full upper slot would waste most of that slot and reduce the number +of concurrent allocations the pool can hold. + +G2H senders allocate the header and control prefix separately when the external +byte stream aligns to the upper slot size. A small prefix uses a lower slot +while the external payload fills complete upper slots. Unaligned streams stay +combined to avoid adding a descriptor. + +The H2G pool contains uniform configured size slots. The same tier selection +does not fit its preposted receive model. The guest publishes writable buffers +before it knows the size of the next host written payload. Uniform slots let +the host calculate how many buffers it needs without negotiating a size class +or searching the ring. + +Queue sizes, upper buffer sizes, and pool page counts are configurable when +the sandbox is created. + +### Initialization + +The host writes the normalized queue sizes, pool page counts, buffer sizes, +and arena GPA into fixed metadata at the top of scratch. It creates both +consumers at cursor zero without reading uninitialized ring contents. + +On the first VM entry, the guest: + +1. Reads the published configuration. +2. Reconstructs `TransportArena`. +3. Converts each transport GPA into its scratch GVA. +4. Creates both packed ring producers and slot pools. +5. Prefills H2G with one writable descriptor per available H2G slot, bounded + by queue size. +6. Publishes the resulting `GuestContext`. + +The host consumers observe the descriptors after guest initialization. + +## Wire format + +Every logical message has this byte layout: + +```text + +----------------+-----------------------------+---------------------+ + | MsgHeader | size-prefixed FlatBuffer | external byte data | + | 12 bytes | control data | zero or more values | + +----------------+-----------------------------+---------------------+ +``` + +`MsgHeader` contains: + +* `kind: u8` +* three reserved zero bytes +* `cid: u32` +* `payload_len: u32` + +`payload_len` covers the control data and all external bytes. RPC correlation +IDs are nonzero. Responses echo the request ID. Logs and snapshot checkpoints +use ID zero. + +The active message kinds are: + +* `Request` +* `Response` +* `Log` +* `SnapshotCheckpoint` + +The FlatBuffer holds the typed function call or result and the lengths of +external byte values. External bytes follow it in the same logical message. +A logical message may span several descriptors or several H2G receive buffers. + +### External byte values + +`ByteChunks` values stay outside the FlatBuffer. The FlatBuffer contains the +total logical value length and whether the value is chunked. The encoder can +then reference the caller's byte slices directly without first copying them into +one contiguous FlatBuffer. + +On the guest, completed shared memory allocations can become +`Bytes::from_owner` values. `ByteChunks` can therefore map transport storage +directly and keep its pool slots allocated until the final `Bytes` owner +drops. `VecBytes` deliberately copies into one contiguous `Vec`. The host +also copies every G2H external value before passing it to host code because +guest writable scratch is untrusted. External values remove intermediate +serialization copies. They do not guarantee that every direction is +end-to-end zero copy. + +C guest function parameters expose `ByteChunks` as a borrowed +`hl_ByteChunks` array. Each `hl_ByteChunk` contains a pointer and length. The +descriptor array is allocated, but its payload pointers reference the +underlying `Bytes` directly. The view is valid until the guest function +returns. `hl_get_host_return_value_as_ByteChunks` returns an owning view that +must be released with `hl_free_byte_chunks`. Chunk arrays produced by C are +copied by `hl_result_from_ByteChunks`. + +The wire format does not preserve the sender's `Vec` boundaries. It +records one total length, not each source chunk length. The receiver sees the +logical byte sequence split where it intersects transport buffers: + +```text + sender chunks: [------][----------][----] + logical byte stream: [------------------------] + transport buffers: [--][--------][--------][--------] + receiver chunks: [--][--------][--------][--------] +``` + +H2G chunking follows the preposted H2G slot size. G2H responses returned to +the guest follow the G2H writable slot size. The message header and FlatBuffer +can consume part of the first slot. The final slot can also be partial. + +## Host calls a guest function + +```text + Host H2G Guest + | | | + | encode Request(cid) | | + | fill posted buffers ----+---------------------->| + | complete buffers | poll and decode | + | | run guest call | + | | | + |<----------------------- G2H Response(cid) ------| + | poll after guest halt | +``` + +The complete flow is: + +1. The host encodes a `FunctionCall` and external values. +2. The host polls enough H2G receive buffers for the complete message. +3. The host writes the message and completes each buffer. +4. The host enters the VM. +5. The guest polls completed H2G buffers, reconstructs the message, and + invokes the registered guest function. +6. The guest submits a G2H `Response` with the same correlation ID. +7. The guest refills H2G and halts without notifying for the deferred response. +8. The host polls G2H, decodes the result, and completes the chain. + +An H2G request containing external bytes must leave one posted buffer +available. This reserve allows a later control call to release retained guest +values. + +## Guest calls a host function + +```text + Guest G2H Host + | | | + | Request(cid) | | + | readable request -------+---------------------->| + | writable reply buffers | copy and decode | + | OUT notification | run host function | + | | | + |<---------------- same chain completed ----------| + | poll and decode Response(cid) | +``` + +The complete flow is: + +1. The guest encodes a `FunctionCall`. The G2H producer allocates its readable + regions and reserves writable response capacity. +2. The guest submits one G2H chain. Its readable region contains the request. + Its writable region reserves the response. +3. The guest notifies the host through `OutBAction::VirtqNotify`. +4. The host polls G2H and copies all request data out of guest writable scratch. +5. The host invokes the registered host function. +6. The host writes a `Response` into the writable region and completes the + same chain. +7. The VM resumes. The guest polls the completion and checks its correlation + ID. + +The host never retains references into guest scratch. It verifies framing, +copies control and external data into host owned values, then invokes host +code. + +Logs use readable G2H chains without writable response capacity. The host +drains and acknowledges them during the same VM exit. + +## Buffer ownership + +Guest `SlotPool` instances own all transport buffers. Pool clones share one +allocation bitmap with each producer. + +```text + Free -> allocated -> published -> completed -> owner-backed Bytes -> Free +``` + +Some stages are skipped by one-way messages. Final ownership matters for +external `ByteChunks`: + +* H2G `ByteChunks` can retain host written receive slots after a guest function + returns. +* G2H host responses can become owner-backed guest `Bytes`. +* `VecBytes` values copy into a contiguous `Vec`. +* Multiple `Bytes` clones or slices backed by one owner keep one slot live. +* The slot returns to the pool when the final owner drops. + +Producer reset releases allocations still owned by queue bookkeeping. After +both producers reset and before H2G prefill, every live pool slot belongs to +guest retained `Bytes`. + +### Trust boundary + +The host treats guest rings, descriptors, headers, FlatBuffers, and payload +lengths as untrusted. + +* Ring and pool access use separately bounded memory views. +* H2G descriptors must be writable, single buffer chains of the configured + size before the host writes to them. +* G2H control and external values are copied into host owned storage before + host code receives them. +* Canonical snapshot validation checks descriptor structure, addresses, + lengths, alignment, pool bounds, uniqueness, and overlap. + +## Snapshot checkpoint + +The transport arena lives in scratch and is not captured as ordinary guest +memory. Guest producer and pool bookkeeping is normal guest state, while ring +and pool bytes live in scratch. Snapshot capture needs a canonical transport +state. + +`MultiUseSandbox` tracks whether queue traffic occurred after the last +canonical boundary. A cached or clean snapshot needs no VM entry. A dirty +snapshot uses this flow: + +```text + Host Guest + | | + | mailbox = u64::MAX | + | H2G SnapshotCheckpoint ------------>| + | enter VM | + | | reclaim completed G2H work + | | reset G2H producer + | | reset H2G producer + | | count live pool slots + | | publish mailbox count + | | prefill H2G + |<------------------------------------| halt + | reset both consumers | + | read mailbox | + | validate and capture rings | +``` + +The canonical state is: + +* G2H is empty at cursor zero. +* H2G starts at cursor zero with one writable descriptor per complete free + slot, bounded by queue size. +* Guest producer and pool bookkeeping matches the rings. +* Driver and device event suppression is normalized. +* Host consumers start at cursor zero. + +The snapshot stores normal guest memory plus the two canonical ring images. +The OCI representation places ring images in the +[transport layer](./snapshot-oci-format.md). Pool payload bytes, the mailbox, +and host consumer cursors are not stored. + +### Restore + +Restore validates the persisted queue configuration, scratch size, ring +lengths, canonical descriptor structure, H2G slot alignment, pool bounds, and +descriptor overlap before exposing either queue. + +It writes the arena GPA metadata and both ring images into fresh scratch, then +attaches new host consumers at cursor zero. Normal guest memory restores the +matching producer and pool bookkeeping. Restore does not need a preparatory VM +entry. + +## Retention mailbox + +The mailbox is one `u64` in the ring to pool alignment gap. It is outside both +rings and pools. Both sides derive its address from trusted arena geometry. +The host accesses it before VM entry and after guest halt. + +The mailbox avoids a G2H checkpoint response. G2H can remain empty in the +canonical image even when retained G2H slots reduce available capacity. + +Before a dirty checkpoint, the host writes `u64::MAX` as a pending marker. +After producer reset, the guest writes: + +```text +g2h_pool.num_live() + h2g_pool.num_live() +``` + +The host reads the value after a successful guest halt and after resetting +both consumers. + +* `u64::MAX` is a fatal incomplete checkpoint. +* Zero permits snapshot capture. +* A nonzero count rejects capture without poisoning the sandbox. + +A nonzero rejection leaves the queues usable and keeps transport dirty. +Guest code can release retained values and retry the snapshot. + +The count only answers whether retained slots exist. It does not contain pool +identity, addresses, or initialized lengths. Retained pool payloads cannot be +restored because pool bytes are absent from the snapshot. + +## Future guest allocated pools and retained snapshots + +Transport pools can leave the fixed arena and use guest allocated scratch. +The rings and mailbox remain at fixed host assigned addresses. At startup, the +guest allocates each complete pool with `alloc_phys_pages`. It allocates fresh +pools when the snapshot generation changes. + +The host accepts descriptor payloads anywhere in guest allocator scratch. It +validates complete ranges, writable H2G buffers, uniqueness, and overlap. Ring +access remains restricted to the fixed arena. + +Pool GVAs are transient and cannot back retained `Bytes` directly. Before +constructing owner backed `Bytes`, the guest maps the buffer's physical pages +at a stable GVA in a reserved alias region. The `Bytes` pointer uses that +alias. The final owner unmaps the alias before returning the slot to its pool. + +Stable aliases make retained payloads ordinary snapshot mappings. Snapshot +capture copies each mapped physical page into snapshot memory while preserving +its alias GVA. Multiple aliases to one physical page share one copied page. +Owner construction clears the unused slot tail. Checkpointing clears free +slots in pools with retained owners, so captured pages contain retained bytes +and zeros. + +Pool backing belongs to one snapshot generation. Retained owners keep the old +pool metadata and stable aliases. After restore, the host enters the guest +without an H2G request. The guest resets both producers, allocates fresh pools, +prefills H2G, and returns before the host uses the restored queues. + +## Placement and relocation limitations + +Arena placement is host owned. `SandboxMemoryLayout` places it at the scratch +base, publishes the GPA, and requires the guest to reconstruct that exact +layout. Host attachment rejects any published arena base that differs from the +configured address. The guest cannot choose placement around its other scratch +allocations. + +The transport also stores absolute guest virtual addresses in descriptors, +pool owners, and guest producer state. Snapshot restore depends on the same +scratch size, queue geometry, and transport addresses. + +### Retained virtual addresses + +Pool relocation cannot transparently change a retained buffer's GVA. +`Bytes::from_owner` stores an absolute data pointer. Its clones and slices can +exist anywhere in guest state. Unsafe Rust and C guests can also retain raw +pointers derived from a live value. The host cannot discover and rebase every +such pointer during restore. + +Wrapping `Bytes` does not solve this because the wrapped `Bytes` still contains +an absolute pointer. A relocatable value would need to replace `Bytes` with an +arena relative handle that resolves its address on every access and does not +promise a stable borrowed slice. That would be a different guest API and would +not constrain pointers created by unsafe code. + +Snapshots containing retained transport values must restore each pool at the +same GVA. The GPA or host backing may move only if page tables and host memory +access preserve that GVA. Restore must fail if it cannot reserve or recreate +the original virtual range. + +Transport capacity is also fixed when the sandbox is created. Runtime queue +resize and VIRTIO feature negotiation are not supported. + +## Source map + +* Shared framing: [`src/hyperlight_common/src/transport.rs`](../src/hyperlight_common/src/transport.rs) +* Packed rings and pools: [`src/hyperlight_common/src/virtq`](../src/hyperlight_common/src/virtq) +* Arena layout: [`src/hyperlight_common/src/layout.rs`](../src/hyperlight_common/src/layout.rs) +* Guest transport: [`src/hyperlight_guest/src/transport`](../src/hyperlight_guest/src/transport) +* Guest initialization: [`src/hyperlight_guest_bin/src/transport.rs`](../src/hyperlight_guest_bin/src/transport.rs) +* Host runtime transport: [`src/hyperlight_host/src/mem/mgr.rs`](../src/hyperlight_host/src/mem/mgr.rs) +* Host validation and snapshots: [`src/hyperlight_host/src/mem/virtq`](../src/hyperlight_host/src/mem/virtq) 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 1ef633804..18802c164 100644 --- a/fuzz/fuzz_targets/guest_trace.rs +++ b/fuzz/fuzz_targets/guest_trace.rs @@ -68,9 +68,7 @@ impl<'a> Arbitrary<'a> for FuzzInput { // Any unexpected errors from the guest should be reported. fuzz_target!( init: { - let mut cfg = SandboxConfiguration::default(); - // In local tests, 256 KiB seemed sufficient for deep recursion - cfg.set_scratch_size(256 * 1024); + let cfg = SandboxConfiguration::default(); let path = simple_guest_for_fuzzing_as_pathbuf(); let u_sbox = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(cfg)).unwrap(); diff --git a/fuzz/fuzz_targets/host_call.rs b/fuzz/fuzz_targets/host_call.rs index 71f9891e4..894a750aa 100644 --- a/fuzz/fuzz_targets/host_call.rs +++ b/fuzz/fuzz_targets/host_call.rs @@ -33,9 +33,10 @@ static SANDBOX: OnceLock> = OnceLock::new(); fuzz_target!( init: { let mut cfg = SandboxConfiguration::default(); - cfg.set_output_data_size(64 * 1024); // 64 KB output buffer - cfg.set_input_data_size(64 * 1024); // 64 KB input buffer - cfg.set_scratch_size(512 * 1024); // large scratch region to contain those buffers, any data copies, etc. + cfg.set_heap_size(512 * 1024); + cfg.set_g2h_pool_pages(16); + cfg.set_h2g_pool_pages(16); + cfg.set_scratch_size(512 * 1024); let u_sbox = UninitializedSandbox::new( GuestBinary::FilePath(simple_guest_for_fuzzing_as_pathbuf()), Some(cfg) @@ -57,6 +58,7 @@ fuzz_target!( // to call with. HyperlightError::HostFunctionNotFound(_) => {} HyperlightError::GuestError(ErrorCode::HostFunctionError, msg) if msg == format!("HostFunction {} was not found", host_func_name) => {} + HyperlightError::GuestError(ErrorCode::HostFunctionError, msg) if msg == "Host response exceeds virtqueue capacity" => {} HyperlightError::UnexpectedNoOfArguments(_, _) => {}, HyperlightError::GuestError(ErrorCode::HostFunctionError, msg) if msg.contains("The number of arguments to the function is wrong") => {} HyperlightError::ParameterValueConversionFailure(_, _) => {}, diff --git a/fuzz/fuzz_targets/virtq_packed_ring.rs b/fuzz/fuzz_targets/virtq_packed_ring.rs index bcde1bfc4..fa91696ea 100644 --- a/fuzz/fuzz_targets/virtq_packed_ring.rs +++ b/fuzz/fuzz_targets/virtq_packed_ring.rs @@ -21,7 +21,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; @@ -43,6 +44,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, @@ -125,9 +127,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, @@ -163,7 +165,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); @@ -190,6 +193,7 @@ fn parse_case(data: &[u8]) -> Option { Some(FuzzCase { queue_size, + avail_descs, driver_event_off_wrap, driver_event_flags, written_len, @@ -207,6 +211,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; @@ -219,17 +277,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; @@ -245,6 +292,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 31c0bb566..8c4cb1625 100644 --- a/src/hyperlight_common/benches/buffer_pool.rs +++ b/src/hyperlight_common/benches/buffer_pool.rs @@ -17,12 +17,12 @@ limitations under the License. use std::hint::black_box; use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; -use hyperlight_common::virtq::{BufferPool, BufferProvider, RecyclePool}; +use hyperlight_common::virtq::{BufferProvider, RunPool, SlotLayout, SlotPool}; // Helper to create a pool for benchmarking -fn make_pool(size: usize) -> BufferPool { +fn make_run_pool(size: usize) -> RunPool { let base = 0x10000; - BufferPool::::new(base, size).unwrap() + RunPool::::new(base, size).unwrap() } // Single allocation performance @@ -32,7 +32,7 @@ fn bench_alloc_single(c: &mut Criterion) { for size in [64, 128, 256, 512, 1024, 1500, 4096].iter() { group.throughput(Throughput::Elements(1)); group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, &size| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { let alloc = pool.alloc(black_box(size)).unwrap(); pool.dealloc(alloc.addr).unwrap(); @@ -49,7 +49,7 @@ fn bench_alloc_lifo(c: &mut Criterion) { for size in [256, 1500, 4096].iter() { group.throughput(Throughput::Elements(100)); group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, &size| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { for _ in 0..100 { let alloc = pool.alloc(black_box(size)).unwrap(); @@ -66,7 +66,7 @@ fn bench_alloc_fragmented(c: &mut Criterion) { let mut group = c.benchmark_group("alloc_fragmented"); group.bench_function("fragmented_256", |b| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); // Create fragmentation pattern: allocate many, free every other let mut allocations = Vec::new(); @@ -92,7 +92,7 @@ fn bench_free(c: &mut Criterion) { for size in [256, 1500, 4096].iter() { group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, &size| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { let alloc = pool.alloc(size).unwrap(); pool.dealloc(black_box(alloc.addr)).unwrap(); @@ -109,7 +109,7 @@ fn bench_free_list_reuse(c: &mut Criterion) { // With cursor optimization (LIFO) group.bench_function("lifo_pattern", |b| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { let alloc = pool.alloc(256).unwrap(); pool.dealloc(alloc.addr).unwrap(); @@ -120,7 +120,7 @@ fn bench_free_list_reuse(c: &mut Criterion) { // Without cursor benefit (FIFO-like) group.bench_function("fifo_pattern", |b| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); let mut queue = Vec::new(); // Pre-fill queue @@ -149,11 +149,11 @@ fn bench_segmented_payload(c: &mut Criterion) { BenchmarkId::from_parameter(payload_size), &payload_size, |b, &payload_size| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { - let sgs = pool.alloc_sg(black_box(payload_size)).unwrap(); - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); + let regions = pool.alloc_regions([black_box(payload_size)]).unwrap(); + for alloc in regions.into_iter().flatten() { + pool.dealloc(alloc.addr).unwrap(); } }); }, @@ -163,11 +163,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(); @@ -175,7 +176,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(); @@ -183,19 +185,21 @@ 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(); + group.bench_function("alloc_regions_64k", |b| { + let layout = SlotLayout::new(0x80000, 4096, 1024); + let pool = SlotPool::new(layout).unwrap(); b.iter(|| { - let sgs = pool.alloc_sg(black_box(64 * 1024)).unwrap(); - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); + let regions = pool.alloc_regions([black_box(64 * 1024)]).unwrap(); + for alloc in regions.into_iter().flatten() { + pool.dealloc(alloc.addr).unwrap(); } }); }); @@ -211,7 +215,7 @@ criterion_group!( bench_free, bench_free_list_reuse, bench_segmented_payload, - bench_recycle_pool, + bench_slot_pool, ); criterion_main!(benches); diff --git a/src/hyperlight_common/benches/common/mod.rs b/src/hyperlight_common/benches/common/mod.rs index bd0930fcd..0102ec8ff 100644 --- a/src/hyperlight_common/benches/common/mod.rs +++ b/src/hyperlight_common/benches/common/mod.rs @@ -27,15 +27,15 @@ use std::sync::atomic::{AtomicU16, AtomicUsize, Ordering}; use bytemuck::Pod; use hyperlight_common::virtq::{ - BufferPool, BufferProvider, Descriptor, Layout, MemOps, Notifier, QueueStats, RecyclePool, - ReplyChain, UsedChain, VirtqConsumer, VirtqProducer, + BufferProvider, Descriptor, Layout, MemOps, Notifier, QueueStats, ReplyChain, RunPool, + SlotLayout, SlotPool, UsedChain, VirtqConsumer, VirtqProducer, }; pub const LOWER_SLOT: usize = 256; pub const UPPER_SLOT: usize = 4096; pub const POOL_SIZE: usize = 8 * 1024 * 1024; -pub type RunBufferPool = BufferPool; +pub type BenchRunPool = RunPool; #[derive(Clone)] struct BenchMem { @@ -176,12 +176,12 @@ where BenchPair { producer, consumer } } -pub fn run_buffer_pool(base: u64, size: usize) -> RunBufferPool { - BufferPool::new(base, size).unwrap() +pub fn run_pool(base: u64, size: usize) -> BenchRunPool { + RunPool::new(base, size).unwrap() } -pub fn fragmented_run_buffer_pool(base: u64, size: usize, payload_size: usize) -> RunBufferPool { - let pool = run_buffer_pool(base, size); +pub fn fragmented_run_pool(base: u64, size: usize, payload_size: usize) -> BenchRunPool { + let pool = run_pool(base, size); let payload_slots = payload_size.div_ceil(UPPER_SLOT); let prefix_slots = 32; let suffix_slots = 32; @@ -197,12 +197,13 @@ pub fn fragmented_run_buffer_pool(base: u64, size: usize, payload_size: usize) - pool } -pub fn recycle_pool(base: u64, size: usize) -> RecyclePool { - RecyclePool::new(base, size, UPPER_SLOT).unwrap() +pub fn slot_pool(base: u64, size: usize) -> SlotPool { + let layout = SlotLayout::new(base, UPPER_SLOT, size / UPPER_SLOT); + SlotPool::new(layout).unwrap() } -pub fn fragmented_recycle_pool(base: u64, size: usize, payload_size: usize) -> RecyclePool { - let pool = recycle_pool(base, size); +pub fn fragmented_slot_pool(base: u64, size: usize, payload_size: usize) -> SlotPool { + let pool = slot_pool(base, size); let payload_slots = payload_size.div_ceil(UPPER_SLOT); let allocated: Vec<_> = (0..payload_slots * 2 + 16) .map(|_| pool.alloc(UPPER_SLOT).unwrap()) @@ -232,8 +233,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); @@ -258,13 +259,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 be70bc5f3..b265773fd 100644 --- a/src/hyperlight_common/benches/virtq_api.rs +++ b/src/hyperlight_common/benches/virtq_api.rs @@ -30,10 +30,10 @@ fn bench_readonly_strategies(c: &mut Criterion) { group.throughput(Throughput::Bytes(size as u64)); group.bench_with_input( - BenchmarkId::new("buffer_pool_run", size), + BenchmarkId::new("run_pool", size), &payload, |b, payload| { - let mut pair = make_pair(128, run_buffer_pool); + let mut pair = make_pair(128, run_pool); b.iter(|| { let used = readonly_roundtrip(&mut pair, black_box(payload)); debug_assert!(matches!(used, UsedChain::Ack(_))); @@ -42,11 +42,11 @@ fn bench_readonly_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("buffer_pool_run_fragmented", size), + BenchmarkId::new("run_pool_fragmented", size), &payload, |b, payload| { let mut pair = make_pair(128, |base, pool_size| { - fragmented_run_buffer_pool(base, pool_size, payload.len()) + fragmented_run_pool(base, pool_size, payload.len()) }); b.iter(|| { let used = readonly_roundtrip(&mut pair, black_box(payload)); @@ -56,10 +56,10 @@ fn bench_readonly_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("recycle_pool_segmented", size), + BenchmarkId::new("slot_pool_segmented", size), &payload, |b, payload| { - let mut pair = make_pair(128, recycle_pool); + let mut pair = make_pair(128, slot_pool); b.iter(|| { let used = readonly_roundtrip(&mut pair, black_box(payload)); debug_assert!(matches!(used, UsedChain::Ack(_))); @@ -68,11 +68,11 @@ fn bench_readonly_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("recycle_pool_segmented_fragmented", size), + BenchmarkId::new("slot_pool_segmented_fragmented", size), &payload, |b, payload| { let mut pair = make_pair(128, |base, pool_size| { - fragmented_recycle_pool(base, pool_size, payload.len()) + fragmented_slot_pool(base, pool_size, payload.len()) }); b.iter(|| { let used = readonly_roundtrip(&mut pair, black_box(payload)); @@ -94,10 +94,10 @@ fn bench_readwrite_strategies(c: &mut Criterion) { group.throughput(Throughput::Bytes((request.len() + response.len()) as u64)); group.bench_with_input( - BenchmarkId::new("buffer_pool_run", size), + BenchmarkId::new("run_pool", size), &(request.clone(), response.clone()), |b, (request, response)| { - let mut pair = make_pair(128, run_buffer_pool); + let mut pair = make_pair(128, run_pool); b.iter(|| { let used = readwrite_roundtrip( &mut pair, @@ -110,11 +110,11 @@ fn bench_readwrite_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("buffer_pool_run_fragmented", size), + BenchmarkId::new("run_pool_fragmented", size), &(request.clone(), response.clone()), |b, (request, response)| { let mut pair = make_pair(128, |base, pool_size| { - fragmented_run_buffer_pool(base, pool_size, request.len()) + fragmented_run_pool(base, pool_size, request.len()) }); b.iter(|| { let used = readwrite_roundtrip( @@ -128,10 +128,10 @@ fn bench_readwrite_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("recycle_pool_segmented", size), + BenchmarkId::new("slot_pool_segmented", size), &(request.clone(), response.clone()), |b, (request, response)| { - let mut pair = make_pair(128, recycle_pool); + let mut pair = make_pair(128, slot_pool); b.iter(|| { let used = readwrite_roundtrip( &mut pair, @@ -144,11 +144,11 @@ fn bench_readwrite_strategies(c: &mut Criterion) { ); group.bench_with_input( - BenchmarkId::new("recycle_pool_segmented_fragmented", size), + BenchmarkId::new("slot_pool_segmented_fragmented", size), &(request, response), |b, (request, response)| { let mut pair = make_pair(128, |base, pool_size| { - fragmented_recycle_pool(base, pool_size, request.len()) + fragmented_slot_pool(base, pool_size, request.len()) }); b.iter(|| { let used = readwrite_roundtrip( diff --git a/src/hyperlight_common/src/arch/aarch64/layout.rs b/src/hyperlight_common/src/arch/aarch64/layout.rs index cb32cfe8e..75629bbc4 100644 --- a/src/hyperlight_common/src/arch/aarch64/layout.rs +++ b/src/hyperlight_common/src/arch/aarch64/layout.rs @@ -28,7 +28,6 @@ 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() -> Option { + 12usize.checked_mul(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 4237caf51..db570e38b 100644 --- a/src/hyperlight_common/src/arch/amd64/layout.rs +++ b/src/hyperlight_common/src/arch/amd64/layout.rs @@ -41,8 +41,6 @@ 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 +pub(super) fn min_scratch_size() -> Option { + 12usize.checked_mul(crate::vmem::PAGE_SIZE) } diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/codec.rs b/src/hyperlight_common/src/flatbuffer_wrappers/codec.rs new file mode 100644 index 000000000..155518b7c --- /dev/null +++ b/src/hyperlight_common/src/flatbuffer_wrappers/codec.rs @@ -0,0 +1,45 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use alloc::vec::Vec; + +use anyhow::Result; +use bytes::Bytes; + +/// Receives external byte values while their FlatBuffer metadata is encoded. +/// +/// Values are delivered in their logical order without flattening chunked +/// values. The value lifetime allows a sink to retain references until the +/// control buffer has been written, without copying payload bytes. +pub trait ExternalValueSink<'a> { + /// Receive one contiguous byte value. + fn push_bytes(&mut self, value: &'a [u8]) -> Result<()>; + + /// Receive one chunk-preserving byte value. + fn push_chunks(&mut self, value: &'a [Bytes]) -> Result<()>; +} + +/// Supplies complete external byte values while a FlatBuffer is decoded. +pub trait ExternalValueSource { + /// Take the next external value as contiguous bytes. + fn take_bytes(&mut self, length: usize) -> Result>; + + /// Take the next external value as owned byte chunks. + fn take_chunks(&mut self, length: usize) -> Result>; + + /// Finish decoding and reject any unused external values. + fn finish(&mut self) -> Result<()>; +} diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs b/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs index 056ced8e0..e9de2e6e3 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs @@ -17,18 +17,20 @@ limitations under the License. use alloc::string::{String, ToString}; use alloc::vec::Vec; -use anyhow::{Error, Result, bail}; +use anyhow::{Result, bail}; use flatbuffers::{FlatBufferBuilder, WIPOffset, size_prefixed_root}; #[cfg(feature = "tracing")] use tracing::{Span, instrument}; -use super::function_types::{ParameterValue, ReturnType}; +use super::codec::{ExternalValueSink, ExternalValueSource}; +use super::function_types::{ParameterValue, ReturnType, decode_parameter_value}; +use super::util::try_byte_chunks_len; use crate::flatbuffers::hyperlight::generated::{ FunctionCall as FbFunctionCall, FunctionCallArgs as FbFunctionCallArgs, FunctionCallType as FbFunctionCallType, Parameter, ParameterArgs, - ParameterValue as FbParameterValue, hlbool, hlboolArgs, hldouble, hldoubleArgs, hlfloat, - hlfloatArgs, hlint, hlintArgs, hllong, hllongArgs, hlstring, hlstringArgs, hluint, hluintArgs, - hlulong, hlulongArgs, hlvecbytes, hlvecbytesArgs, + ParameterValue as FbParameterValue, hlbool, hlboolArgs, hldouble, hldoubleArgs, + hlexternalbytes, hlexternalbytesArgs, hlfloat, hlfloatArgs, hlint, hlintArgs, hllong, + hllongArgs, hlstring, hlstringArgs, hluint, hluintArgs, hlulong, hlulongArgs, }; /// The type of function call. @@ -73,14 +75,15 @@ impl FunctionCall { self.function_call_type.clone() } - /// Encodes self into the given builder and returns the encoded data. - /// - /// # Notes - /// - /// The builder should not be reused after a call to encode, since this function - /// does not reset the state of the builder. If you want to reuse the builder, - /// you'll need to reset it first. - pub fn encode<'a>(&self, builder: &'a mut FlatBufferBuilder) -> &'a [u8] { + /// Encode control data and collect byte parameters as external values. + pub fn encode<'a, 'b, S>( + &'a self, + builder: &'b mut FlatBufferBuilder, + external_values: &mut S, + ) -> Result<&'b [u8]> + where + S: ExternalValueSink<'a> + ?Sized, + { let function_name = builder.create_string(&self.function_name); let function_call_type = match self.function_call_type { @@ -91,110 +94,147 @@ impl FunctionCall { let expected_return_type = self.expected_return_type.into(); let parameters = match &self.parameters { - Some(p) if !p.is_empty() => { - let parameter_offsets: Vec> = p + Some(parameters) if !parameters.is_empty() => { + let parameter_offsets: Vec> = parameters .iter() - .map(|param| match param { - ParameterValue::Int(i) => { - let hlint = hlint::create(builder, &hlintArgs { value: *i }); - Parameter::create( - builder, - &ParameterArgs { - value_type: FbParameterValue::hlint, - value: Some(hlint.as_union_value()), - }, - ) - } - ParameterValue::UInt(ui) => { - let hluint = hluint::create(builder, &hluintArgs { value: *ui }); - Parameter::create( - builder, - &ParameterArgs { - value_type: FbParameterValue::hluint, - value: Some(hluint.as_union_value()), - }, - ) - } - ParameterValue::Long(l) => { - let hllong = hllong::create(builder, &hllongArgs { value: *l }); - Parameter::create( - builder, - &ParameterArgs { - value_type: FbParameterValue::hllong, - value: Some(hllong.as_union_value()), - }, - ) - } - ParameterValue::ULong(ul) => { - let hlulong = hlulong::create(builder, &hlulongArgs { value: *ul }); - Parameter::create( - builder, - &ParameterArgs { - value_type: FbParameterValue::hlulong, - value: Some(hlulong.as_union_value()), - }, - ) - } - ParameterValue::Float(f) => { - let hlfloat = hlfloat::create(builder, &hlfloatArgs { value: *f }); - Parameter::create( - builder, - &ParameterArgs { - value_type: FbParameterValue::hlfloat, - value: Some(hlfloat.as_union_value()), - }, - ) - } - ParameterValue::Double(d) => { - let hldouble = hldouble::create(builder, &hldoubleArgs { value: *d }); - Parameter::create( - builder, - &ParameterArgs { - value_type: FbParameterValue::hldouble, - value: Some(hldouble.as_union_value()), - }, - ) - } - ParameterValue::Bool(b) => { - let hlbool = hlbool::create(builder, &hlboolArgs { value: *b }); - Parameter::create( - builder, - &ParameterArgs { - value_type: FbParameterValue::hlbool, - value: Some(hlbool.as_union_value()), - }, - ) - } - ParameterValue::String(s) => { - let val = builder.create_string(s.as_str()); - let hlstring = - hlstring::create(builder, &hlstringArgs { value: Some(val) }); - Parameter::create( - builder, - &ParameterArgs { - value_type: FbParameterValue::hlstring, - value: Some(hlstring.as_union_value()), - }, - ) - } - ParameterValue::VecBytes(v) => { - let vec_bytes = builder.create_vector(v); - let hlvecbytes = hlvecbytes::create( - builder, - &hlvecbytesArgs { - value: Some(vec_bytes), - }, - ); - Parameter::create( - builder, - &ParameterArgs { - value_type: FbParameterValue::hlvecbytes, - value: Some(hlvecbytes.as_union_value()), - }, - ) - } + .map(|parameter| -> Result> { + let parameter = match parameter { + ParameterValue::Int(value) => { + let value = hlint::create(builder, &hlintArgs { value: *value }); + Parameter::create( + builder, + &ParameterArgs { + value_type: FbParameterValue::hlint, + value: Some(value.as_union_value()), + }, + ) + } + ParameterValue::UInt(value) => { + let value = hluint::create(builder, &hluintArgs { value: *value }); + Parameter::create( + builder, + &ParameterArgs { + value_type: FbParameterValue::hluint, + value: Some(value.as_union_value()), + }, + ) + } + ParameterValue::Long(value) => { + let value = hllong::create(builder, &hllongArgs { value: *value }); + Parameter::create( + builder, + &ParameterArgs { + value_type: FbParameterValue::hllong, + value: Some(value.as_union_value()), + }, + ) + } + ParameterValue::ULong(value) => { + let value = + hlulong::create(builder, &hlulongArgs { value: *value }); + Parameter::create( + builder, + &ParameterArgs { + value_type: FbParameterValue::hlulong, + value: Some(value.as_union_value()), + }, + ) + } + ParameterValue::Float(value) => { + let value = + hlfloat::create(builder, &hlfloatArgs { value: *value }); + Parameter::create( + builder, + &ParameterArgs { + value_type: FbParameterValue::hlfloat, + value: Some(value.as_union_value()), + }, + ) + } + ParameterValue::Double(value) => { + let value = + hldouble::create(builder, &hldoubleArgs { value: *value }); + Parameter::create( + builder, + &ParameterArgs { + value_type: FbParameterValue::hldouble, + value: Some(value.as_union_value()), + }, + ) + } + ParameterValue::Bool(value) => { + let value = hlbool::create(builder, &hlboolArgs { value: *value }); + Parameter::create( + builder, + &ParameterArgs { + value_type: FbParameterValue::hlbool, + value: Some(value.as_union_value()), + }, + ) + } + ParameterValue::String(value) => { + let value = builder.create_string(value.as_str()); + let value = + hlstring::create(builder, &hlstringArgs { value: Some(value) }); + Parameter::create( + builder, + &ParameterArgs { + value_type: FbParameterValue::hlstring, + value: Some(value.as_union_value()), + }, + ) + } + ParameterValue::VecBytes(value) => { + let length = u64::try_from(value.len()).map_err(|_| { + anyhow::anyhow!( + "External VecBytes parameter length does not fit in u64" + ) + })?; + external_values.push_bytes(value)?; + let value = hlexternalbytes::create( + builder, + &hlexternalbytesArgs { + length, + chunked: false, + }, + ); + Parameter::create( + builder, + &ParameterArgs { + value_type: FbParameterValue::hlexternalbytes, + value: Some(value.as_union_value()), + }, + ) + } + ParameterValue::ByteChunks(value) => { + let length = try_byte_chunks_len(value).ok_or_else(|| { + anyhow::anyhow!("External ByteChunks parameter length overflow") + })?; + let length = u64::try_from(length).map_err(|_| { + anyhow::anyhow!( + "External ByteChunks parameter length does not fit in u64" + ) + })?; + external_values.push_chunks(value)?; + let value = hlexternalbytes::create( + builder, + &hlexternalbytesArgs { + length, + chunked: true, + }, + ); + Parameter::create( + builder, + &ParameterArgs { + value_type: FbParameterValue::hlexternalbytes, + value: Some(value.as_union_value()), + }, + ) + } + }; + Ok(parameter) }) - .collect(); + .collect::>>()?; Some(builder.create_vector(¶meter_offsets)) } _ => None, @@ -210,38 +250,14 @@ impl FunctionCall { }, ); builder.finish_size_prefixed(function_call, None); - builder.finished_data() + Ok(builder.finished_data()) } -} - -#[cfg_attr(feature = "tracing", instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace"))] -pub fn validate_guest_function_call_buffer(function_call_buffer: &[u8]) -> Result<()> { - let guest_function_call_fb = size_prefixed_root::(function_call_buffer) - .map_err(|e| anyhow::anyhow!("Error reading function call buffer: {:?}", e))?; - match guest_function_call_fb.function_call_type() { - FbFunctionCallType::guest => Ok(()), - other => { - bail!("Invalid function call type: {:?}", other); - } - } -} - -#[cfg_attr(feature = "tracing", instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace"))] -pub fn validate_host_function_call_buffer(function_call_buffer: &[u8]) -> Result<()> { - let host_function_call_fb = size_prefixed_root::(function_call_buffer) - .map_err(|e| anyhow::anyhow!("Error reading function call buffer: {:?}", e))?; - match host_function_call_fb.function_call_type() { - FbFunctionCallType::host => Ok(()), - other => { - bail!("Invalid function call type: {:?}", other); - } - } -} -impl TryFrom<&[u8]> for FunctionCall { - type Error = Error; - #[cfg_attr(feature = "tracing", instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace"))] - fn try_from(value: &[u8]) -> Result { + /// Decode control data and consume external byte parameters. + pub fn decode(value: &[u8], external_values: &mut S) -> Result + where + S: ExternalValueSource + ?Sized, + { let function_call_fb = size_prefixed_root::(value) .map_err(|e| anyhow::anyhow!("Error reading function call buffer: {:?}", e))?; let function_name = function_call_fb.function_name(); @@ -256,13 +272,15 @@ impl TryFrom<&[u8]> for FunctionCall { let parameters = function_call_fb .parameters() - .map(|v| { - v.iter() - .map(|p| p.try_into()) + .map(|parameters| { + parameters + .iter() + .map(|parameter| decode_parameter_value(parameter, external_values)) .collect::>>() }) .transpose()?; + external_values.finish()?; Ok(Self { function_name: function_name.to_string(), parameters, @@ -274,14 +292,79 @@ impl TryFrom<&[u8]> for FunctionCall { #[cfg(test)] mod tests { + use alloc::collections::VecDeque; use alloc::vec; use super::*; - use crate::flatbuffer_wrappers::function_types::ReturnType; + use crate::flatbuffer_wrappers::function_types::{Bytes, ReturnType}; + + #[derive(Debug, Clone, PartialEq)] + enum TestExternalValue { + VecBytes(Vec), + ByteChunks(Vec), + } + + #[derive(Default)] + struct TestExternalValues { + values: VecDeque, + } + + impl TestExternalValues { + fn from_values(values: impl IntoIterator) -> Self { + Self { + values: values.into_iter().collect(), + } + } + } + + impl<'a> ExternalValueSink<'a> for TestExternalValues { + fn push_bytes(&mut self, value: &'a [u8]) -> Result<()> { + self.values + .push_back(TestExternalValue::VecBytes(value.to_vec())); + Ok(()) + } + + fn push_chunks(&mut self, value: &'a [Bytes]) -> Result<()> { + self.values + .push_back(TestExternalValue::ByteChunks(value.to_vec())); + Ok(()) + } + } + + impl ExternalValueSource for TestExternalValues { + fn take_bytes(&mut self, _length: usize) -> Result> { + match self.values.pop_front() { + Some(TestExternalValue::VecBytes(value)) => Ok(value), + Some(TestExternalValue::ByteChunks(_)) => { + anyhow::bail!("Expected external VecBytes value") + } + None => anyhow::bail!("Missing external VecBytes value"), + } + } + + fn take_chunks(&mut self, _length: usize) -> Result> { + match self.values.pop_front() { + Some(TestExternalValue::ByteChunks(value)) => Ok(value), + Some(TestExternalValue::VecBytes(_)) => { + anyhow::bail!("Expected external ByteChunks value") + } + None => anyhow::bail!("Missing external ByteChunks value"), + } + } + + fn finish(&mut self) -> Result<()> { + if self.values.is_empty() { + Ok(()) + } else { + anyhow::bail!("Unused external values") + } + } + } #[test] fn read_from_flatbuffer() -> Result<()> { let mut builder = FlatBufferBuilder::new(); + let mut external_values = TestExternalValues::default(); let test_data = FunctionCall::new( "PrintTwelveArgs".to_string(), Some(vec![ @@ -301,9 +384,9 @@ mod tests { FunctionCallType::Guest, ReturnType::Int, ) - .encode(&mut builder); + .encode(&mut builder, &mut external_values)?; - let function_call = FunctionCall::try_from(test_data)?; + let function_call = FunctionCall::decode(test_data, &mut external_values)?; assert_eq!(function_call.function_name, "PrintTwelveArgs"); assert!(function_call.parameters.is_some()); let parameters = function_call.parameters.unwrap(); @@ -327,4 +410,122 @@ mod tests { Ok(()) } + + #[test] + fn byte_parameters_round_trip_in_external_value_order() { + let mut builder = FlatBufferBuilder::new(); + let expected_parameters = vec![ + ParameterValue::Int(7), + ParameterValue::UInt(8), + ParameterValue::Long(-9), + ParameterValue::ULong(10), + ParameterValue::Float(1.25), + ParameterValue::Double(2.5), + ParameterValue::Bool(true), + ParameterValue::VecBytes(vec![0xa5; 4096]), + ParameterValue::String("middle".to_string()), + ParameterValue::ByteChunks(vec![ + Bytes::from_static(b"chunk one"), + Bytes::from_static(b" and two"), + ]), + ParameterValue::VecBytes(Vec::new()), + ParameterValue::ByteChunks(Vec::new()), + ]; + let call = FunctionCall::new( + "external_bytes".to_string(), + Some(expected_parameters.clone()), + FunctionCallType::Host, + ReturnType::ByteChunks, + ); + let mut external_values = TestExternalValues::default(); + let encoded = call.encode(&mut builder, &mut external_values).unwrap(); + + assert!(encoded.len() < 4096); + assert_eq!( + external_values.values, + VecDeque::from([ + TestExternalValue::VecBytes(vec![0xa5; 4096]), + TestExternalValue::ByteChunks(vec![ + Bytes::from_static(b"chunk one"), + Bytes::from_static(b" and two"), + ]), + TestExternalValue::VecBytes(Vec::new()), + TestExternalValue::ByteChunks(Vec::new()), + ]) + ); + + let encoded_call = size_prefixed_root::(encoded).unwrap(); + let encoded_parameters = encoded_call.parameters().unwrap(); + for (index, length, chunked) in [ + (7, 4096, false), + (9, 17, true), + (10, 0, false), + (11, 0, true), + ] { + let parameter = encoded_parameters.get(index); + assert_eq!(parameter.value_type(), FbParameterValue::hlexternalbytes); + let metadata = parameter.value_as_hlexternalbytes().unwrap(); + assert_eq!(metadata.length(), length); + assert_eq!(metadata.chunked(), chunked); + } + + let decoded = FunctionCall::decode(encoded, &mut external_values).unwrap(); + assert_eq!(decoded.function_name, "external_bytes"); + assert_eq!(decoded.parameters, Some(expected_parameters)); + assert_eq!(decoded.function_call_type(), FunctionCallType::Host); + assert_eq!(decoded.expected_return_type, ReturnType::ByteChunks); + assert!(external_values.values.is_empty()); + } + + #[test] + fn scalar_call_uses_no_external_values() { + let call = FunctionCall::new( + "scalars".to_string(), + Some(vec![ + ParameterValue::Int(42), + ParameterValue::String("value".to_string()), + ]), + FunctionCallType::Guest, + ReturnType::Bool, + ); + let mut builder = FlatBufferBuilder::new(); + let mut external_values = TestExternalValues::default(); + let encoded = call.encode(&mut builder, &mut external_values).unwrap(); + + assert!(external_values.values.is_empty()); + let decoded = FunctionCall::decode(encoded, &mut external_values).unwrap(); + assert_eq!(decoded.function_name, "scalars"); + } + + #[test] + fn external_parameter_decoder_rejects_invalid_value_sequences() { + let mut builder = FlatBufferBuilder::new(); + let call = FunctionCall::new( + "external_bytes".to_string(), + Some(vec![ParameterValue::VecBytes(vec![1, 2, 3])]), + FunctionCallType::Guest, + ReturnType::Void, + ); + let mut encoded_values = TestExternalValues::default(); + let encoded = call.encode(&mut builder, &mut encoded_values).unwrap(); + + let mut missing = TestExternalValues::default(); + assert!(FunctionCall::decode(encoded, &mut missing).is_err()); + + let mut wrong_type = + TestExternalValues::from_values([TestExternalValue::ByteChunks(vec![ + Bytes::from_static(b"123"), + ])]); + assert!(FunctionCall::decode(encoded, &mut wrong_type).is_err()); + + let mut wrong_length = + TestExternalValues::from_values([TestExternalValue::VecBytes(vec![1, 2])]); + assert!(FunctionCall::decode(encoded, &mut wrong_length).is_err()); + + let mut extra = TestExternalValues::from_values([ + TestExternalValue::VecBytes(vec![1, 2, 3]), + TestExternalValue::VecBytes(Vec::new()), + ]); + assert!(FunctionCall::decode(encoded, &mut extra).is_err()); + } } diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs b/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs index 42c7ff823..1a7ab969b 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs @@ -18,32 +18,38 @@ use alloc::string::{String, ToString}; use alloc::vec::Vec; use anyhow::{Error, Result, anyhow, bail}; +pub use bytes::Bytes; use flatbuffers::size_prefixed_root; #[cfg(feature = "tracing")] use tracing::{Span, instrument}; +use super::codec::{ExternalValueSink, ExternalValueSource}; use super::guest_error::GuestError; +#[cfg(feature = "fuzzing")] +use super::util::arbitrary_byte_chunks; +use super::util::try_byte_chunks_len; use crate::flatbuffers::hyperlight::generated::{ FunctionCallResult as FbFunctionCallResult, FunctionCallResultArgs as FbFunctionCallResultArgs, FunctionCallResultType, Parameter, ParameterType as FbParameterType, ParameterValue as FbParameterValue, ReturnType as FbReturnType, ReturnValue as FbReturnValue, - ReturnValueBox, ReturnValueBoxArgs, hlbool, hlboolArgs, hldouble, hldoubleArgs, hlfloat, - hlfloatArgs, hlint, hlintArgs, hllong, hllongArgs, hlsizeprefixedbuffer, - hlsizeprefixedbufferArgs, hlstring, hlstringArgs, hluint, hluintArgs, hlulong, hlulongArgs, - hlvoid, hlvoidArgs, + ReturnValueBox, ReturnValueBoxArgs, hlbool, hlboolArgs, hldouble, hldoubleArgs, + hlexternalbytes, hlexternalbytesArgs, hlfloat, hlfloatArgs, hlint, hlintArgs, hllong, + hllongArgs, hlstring, hlstringArgs, hluint, hluintArgs, hlulong, hlulongArgs, hlvoid, + hlvoidArgs, }; pub struct FunctionCallResult(core::result::Result); impl FunctionCallResult { - /// Encodes self into the given builder and returns the encoded data. - /// - /// # Notes - /// - /// The builder should not be reused after a call to encode, since this function - /// does not reset the state of the builder. If you want to reuse the builder, - /// you'll need to reset it first. - pub fn encode<'a>(&self, builder: &'a mut flatbuffers::FlatBufferBuilder) -> &'a [u8] { + /// Encode control data and collect a byte return as an external value. + pub fn encode<'a, 'b, S>( + &'a self, + builder: &'b mut flatbuffers::FlatBufferBuilder, + external_values: &mut S, + ) -> Result<&'b [u8]> + where + S: ExternalValueSink<'a> + ?Sized, + { match &self.0 { Ok(rv) => { // Encode ReturnValue as ReturnValueBox @@ -82,18 +88,36 @@ impl FunctionCallResult { (Some(off.as_union_value()), FbReturnValue::hlstring) } ReturnValue::VecBytes(v) => { - let val = builder.create_vector(v); - let off = hlsizeprefixedbuffer::create( + let length = u64::try_from(v.len()) + .map_err(|_| anyhow!("External VecBytes length does not fit in u64"))?; + + external_values.push_bytes(v)?; + let off = hlexternalbytes::create( builder, - &hlsizeprefixedbufferArgs { - value: Some(val), - size: v.len() as i32, + &hlexternalbytesArgs { + length, + chunked: false, }, ); - ( - Some(off.as_union_value()), - FbReturnValue::hlsizeprefixedbuffer, - ) + (Some(off.as_union_value()), FbReturnValue::hlexternalbytes) + } + ReturnValue::ByteChunks(v) => { + let length = try_byte_chunks_len(v) + .ok_or_else(|| anyhow!("External ByteChunks length overflow"))?; + + let length = u64::try_from(length).map_err(|_| { + anyhow!("External ByteChunks length does not fit in u64") + })?; + + external_values.push_chunks(v)?; + let off = hlexternalbytes::create( + builder, + &hlexternalbytesArgs { + length, + chunked: true, + }, + ); + (Some(off.as_union_value()), FbReturnValue::hlexternalbytes) } ReturnValue::Void(()) => { let off = hlvoid::create(builder, &hlvoidArgs {}); @@ -110,7 +134,7 @@ impl FunctionCallResult { }, ); builder.finish_size_prefixed(fcr, None); - builder.finished_data() + Ok(builder.finished_data()) } Err(ge) => { // Encode GuestError @@ -131,10 +155,11 @@ impl FunctionCallResult { }, ); builder.finish_size_prefixed(fcr, None); - builder.finished_data() + Ok(builder.finished_data()) } } } + pub fn new(value: core::result::Result) -> Self { FunctionCallResult(value) } @@ -142,24 +167,23 @@ impl FunctionCallResult { pub fn into_inner(self) -> core::result::Result { self.0 } -} -impl TryFrom<&[u8]> for FunctionCallResult { - type Error = Error; - - fn try_from(value: &[u8]) -> Result { + /// Decode control data and consume an external byte return. + pub fn decode(value: &[u8], external_values: &mut S) -> Result + where + S: ExternalValueSource + ?Sized, + { let function_call_result_fb = size_prefixed_root::(value) .map_err(|e| anyhow!("Failed to get FunctionCallResult from bytes: {:?}", e))?; - match function_call_result_fb.result_type() { + let result = match function_call_result_fb.result_type() { FunctionCallResultType::ReturnValueBox => { let boxed = function_call_result_fb .result_as_return_value_box() .ok_or_else(|| { anyhow!("Failed to get ReturnValueBox from function call result") })?; - let return_value = ReturnValue::try_from(boxed)?; - Ok(FunctionCallResult(Ok(return_value))) + Ok(decode_return_value(boxed, external_values)?) } FunctionCallResultType::GuestError => { let guest_error_table = function_call_result_fb @@ -170,15 +194,15 @@ impl TryFrom<&[u8]> for FunctionCallResult { .message() .map(|s| s.to_string()) .unwrap_or_default(); - Ok(FunctionCallResult(Err(GuestError::new( - code.into(), - message, - )))) + Err(GuestError::new(code.into(), message)) } other => { bail!("Unexpected function call result type: {:?}", other) } - } + }; + + external_values.finish()?; + Ok(FunctionCallResult(result)) } } @@ -204,6 +228,13 @@ pub enum ParameterValue { Bool(bool), /// `Vec` VecBytes(Vec), + /// One complete chunk-preserving byte value. + /// + /// Chunk boundaries are not framing and are not guaranteed to survive + /// transport. + ByteChunks( + #[cfg_attr(feature = "fuzzing", arbitrary(with = arbitrary_byte_chunks))] Vec, + ), } /// Supported parameter types for function calling. @@ -228,6 +259,8 @@ pub enum ParameterType { Bool, /// `Vec` VecBytes, + /// One complete chunk-preserving byte value. + ByteChunks, } /// Supported return types with values from function calling. @@ -253,6 +286,11 @@ pub enum ReturnValue { Void(()), /// `Vec` VecBytes(Vec), + /// One complete chunk-preserving byte value. + /// + /// Chunk boundaries are not framing and are not guaranteed to survive + /// transport. + ByteChunks(Vec), } /// Supported return types from function calling. @@ -281,6 +319,99 @@ pub enum ReturnType { Void, /// `Vec` VecBytes, + /// One complete chunk-preserving byte value. + ByteChunks, +} + +enum DecodedExternalBytes { + VecBytes(Vec), + ByteChunks(Vec), +} + +fn decode_external_bytes( + metadata: hlexternalbytes<'_>, + externals: &mut S, +) -> Result +where + S: ExternalValueSource + ?Sized, +{ + // The length delimits this value in the ordered external payload stream. + let length = usize::try_from(metadata.length()).map_err(|_| { + anyhow!( + "External byte length {} does not fit in usize", + metadata.length() + ) + })?; + + // `chunked` selects the logical API type. Sources define chunk boundaries. + if metadata.chunked() { + let value = externals.take_chunks(length)?; + let actual = try_byte_chunks_len(&value) + .ok_or_else(|| anyhow!("External ByteChunks length overflow"))?; + + if actual != length { + bail!("External ByteChunks length mismatch: declared {length}, received {actual}",); + } + Ok(DecodedExternalBytes::ByteChunks(value)) + } else { + let value = externals.take_bytes(length)?; + let value_len = value.len(); + + if value_len != length { + bail!("External VecBytes length mismatch: declared {length}, received {value_len}",); + } + Ok(DecodedExternalBytes::VecBytes(value)) + } +} + +pub(crate) fn decode_parameter_value( + param: Parameter<'_>, + externals: &mut S, +) -> Result +where + S: ExternalValueSource + ?Sized, +{ + match param.value_type() { + FbParameterValue::hlexternalbytes => { + let Some(metadata) = param.value_as_hlexternalbytes() else { + bail!("External byte parameter metadata is missing"); + }; + + match decode_external_bytes(metadata, externals)? { + DecodedExternalBytes::VecBytes(value) => Ok(ParameterValue::VecBytes(value)), + DecodedExternalBytes::ByteChunks(value) => Ok(ParameterValue::ByteChunks(value)), + } + } + FbParameterValue::hlvecbytes | FbParameterValue::hlbytechunks => { + bail!("Embedded byte parameters are not supported") + } + _ => param.try_into(), + } +} + +fn decode_return_value( + return_value: ReturnValueBox<'_>, + externals: &mut S, +) -> Result +where + S: ExternalValueSource + ?Sized, +{ + match return_value.value_type() { + FbReturnValue::hlexternalbytes => { + let Some(metadata) = return_value.value_as_hlexternalbytes() else { + bail!("External byte parameter metadata is missing"); + }; + + match decode_external_bytes(metadata, externals)? { + DecodedExternalBytes::VecBytes(value) => Ok(ReturnValue::VecBytes(value)), + DecodedExternalBytes::ByteChunks(value) => Ok(ReturnValue::ByteChunks(value)), + } + } + FbReturnValue::hlsizeprefixedbuffer | FbReturnValue::hlsizeprefixedbytechunks => { + bail!("Embedded byte returns are not supported") + } + _ => return_value.try_into(), + } } impl From<&ParameterValue> for ParameterType { @@ -296,6 +427,7 @@ impl From<&ParameterValue> for ParameterType { ParameterValue::String(_) => ParameterType::String, ParameterValue::Bool(_) => ParameterType::Bool, ParameterValue::VecBytes(_) => ParameterType::VecBytes, + ParameterValue::ByteChunks(_) => ParameterType::ByteChunks, } } } @@ -331,9 +463,12 @@ impl TryFrom> for ParameterValue { FbParameterValue::hlstring => param.value_as_hlstring().map(|hlstring| { ParameterValue::String(hlstring.value().unwrap_or_default().to_string()) }), - FbParameterValue::hlvecbytes => param.value_as_hlvecbytes().map(|hlvecbytes| { - ParameterValue::VecBytes(hlvecbytes.value().unwrap_or_default().bytes().to_vec()) - }), + FbParameterValue::hlvecbytes | FbParameterValue::hlbytechunks => { + bail!("Embedded byte parameters are not supported") + } + FbParameterValue::hlexternalbytes => { + bail!("External byte parameter requires an external value source") + } other => { bail!("Unexpected flatbuffer parameter value type: {:?}", other); } @@ -355,6 +490,7 @@ impl From for FbParameterType { ParameterType::String => FbParameterType::hlstring, ParameterType::Bool => FbParameterType::hlbool, ParameterType::VecBytes => FbParameterType::hlvecbytes, + ParameterType::ByteChunks => FbParameterType::hlbytechunks, } } } @@ -373,6 +509,7 @@ impl From for FbReturnType { ReturnType::Bool => FbReturnType::hlbool, ReturnType::Void => FbReturnType::hlvoid, ReturnType::VecBytes => FbReturnType::hlsizeprefixedbuffer, + ReturnType::ByteChunks => FbReturnType::hlbytechunks, } } } @@ -391,6 +528,7 @@ impl TryFrom for ParameterType { FbParameterType::hlstring => Ok(ParameterType::String), FbParameterType::hlbool => Ok(ParameterType::Bool), FbParameterType::hlvecbytes => Ok(ParameterType::VecBytes), + FbParameterType::hlbytechunks => Ok(ParameterType::ByteChunks), _ => { bail!("Unexpected flatbuffer parameter type: {:?}", value) } @@ -413,6 +551,7 @@ impl TryFrom for ReturnType { FbReturnType::hlbool => Ok(ReturnType::Bool), FbReturnType::hlvoid => Ok(ReturnType::Void), FbReturnType::hlsizeprefixedbuffer => Ok(ReturnType::VecBytes), + FbReturnType::hlbytechunks => Ok(ReturnType::ByteChunks), _ => { bail!("Unexpected flatbuffer return type: {:?}", value) } @@ -537,6 +676,17 @@ impl TryFrom for Vec { } } +impl TryFrom for Vec { + type Error = Error; + + fn try_from(value: ParameterValue) -> Result { + match value { + ParameterValue::ByteChunks(v) => Ok(v), + _ => bail!("Unexpected parameter value type: {:?}", value), + } + } +} + impl TryFrom for i32 { type Error = Error; #[cfg_attr(feature = "tracing", instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace"))] @@ -654,6 +804,17 @@ impl TryFrom for Vec { } } +impl TryFrom for Vec { + type Error = Error; + + fn try_from(value: ReturnValue) -> Result { + match value { + ReturnValue::ByteChunks(v) => Ok(v), + _ => bail!("Unexpected return value type: {:?}", value), + } + } +} + impl TryFrom for () { type Error = Error; #[cfg_attr(feature = "tracing", instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace"))] @@ -722,14 +883,11 @@ impl TryFrom> for ReturnValue { Ok(ReturnValue::String(hlstring.unwrap_or("".to_string()))) } FbReturnValue::hlvoid => Ok(ReturnValue::Void(())), - FbReturnValue::hlsizeprefixedbuffer => { - let hlvecbytes = match return_value_box.value_as_hlsizeprefixedbuffer() { - Some(hlvecbytes) => hlvecbytes - .value() - .map(|val| val.iter().collect::>()), - None => None, - }; - Ok(ReturnValue::VecBytes(hlvecbytes.unwrap_or(Vec::new()))) + FbReturnValue::hlsizeprefixedbuffer | FbReturnValue::hlsizeprefixedbytechunks => { + bail!("Embedded byte returns are not supported") + } + FbReturnValue::hlexternalbytes => { + bail!("External byte return requires an external value source") } other => { bail!("Unexpected flatbuffer return value type: {:?}", other) @@ -738,233 +896,92 @@ impl TryFrom> for ReturnValue { } } -impl TryFrom<&ReturnValue> for Vec { - type Error = Error; - #[cfg_attr(feature = "tracing", instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace"))] - fn try_from(value: &ReturnValue) -> Result> { - let mut builder = flatbuffers::FlatBufferBuilder::new(); - let result_bytes = match value { - ReturnValue::Int(i) => { - let hlint_off = hlint::create(&mut builder, &hlintArgs { value: *i }); - let rv_box = ReturnValueBox::create( - &mut builder, - &ReturnValueBoxArgs { - value: Some(hlint_off.as_union_value()), - value_type: FbReturnValue::hlint, - }, - ); - let fcr = FbFunctionCallResult::create( - &mut builder, - &FbFunctionCallResultArgs { - result: Some(rv_box.as_union_value()), - result_type: FunctionCallResultType::ReturnValueBox, - }, - ); - builder.finish_size_prefixed(fcr, None); - builder.finished_data().to_vec() - } - ReturnValue::UInt(ui) => { - let off = hluint::create(&mut builder, &hluintArgs { value: *ui }); - let rv_box = ReturnValueBox::create( - &mut builder, - &ReturnValueBoxArgs { - value: Some(off.as_union_value()), - value_type: FbReturnValue::hluint, - }, - ); - let fcr = FbFunctionCallResult::create( - &mut builder, - &FbFunctionCallResultArgs { - result: Some(rv_box.as_union_value()), - result_type: FunctionCallResultType::ReturnValueBox, - }, - ); - builder.finish_size_prefixed(fcr, None); - builder.finished_data().to_vec() - } - ReturnValue::Long(l) => { - let off = hllong::create(&mut builder, &hllongArgs { value: *l }); - let rv_box = ReturnValueBox::create( - &mut builder, - &ReturnValueBoxArgs { - value: Some(off.as_union_value()), - value_type: FbReturnValue::hllong, - }, - ); - let fcr = FbFunctionCallResult::create( - &mut builder, - &FbFunctionCallResultArgs { - result: Some(rv_box.as_union_value()), - result_type: FunctionCallResultType::ReturnValueBox, - }, - ); - builder.finish_size_prefixed(fcr, None); - builder.finished_data().to_vec() - } - ReturnValue::ULong(ul) => { - let off = hlulong::create(&mut builder, &hlulongArgs { value: *ul }); - let rv_box = ReturnValueBox::create( - &mut builder, - &ReturnValueBoxArgs { - value: Some(off.as_union_value()), - value_type: FbReturnValue::hlulong, - }, - ); - let fcr = FbFunctionCallResult::create( - &mut builder, - &FbFunctionCallResultArgs { - result: Some(rv_box.as_union_value()), - result_type: FunctionCallResultType::ReturnValueBox, - }, - ); - builder.finish_size_prefixed(fcr, None); - builder.finished_data().to_vec() - } - ReturnValue::Float(f) => { - let off = hlfloat::create(&mut builder, &hlfloatArgs { value: *f }); - let rv_box = ReturnValueBox::create( - &mut builder, - &ReturnValueBoxArgs { - value: Some(off.as_union_value()), - value_type: FbReturnValue::hlfloat, - }, - ); - let fcr = FbFunctionCallResult::create( - &mut builder, - &FbFunctionCallResultArgs { - result: Some(rv_box.as_union_value()), - result_type: FunctionCallResultType::ReturnValueBox, - }, - ); - builder.finish_size_prefixed(fcr, None); - builder.finished_data().to_vec() - } - ReturnValue::Double(d) => { - let off = hldouble::create(&mut builder, &hldoubleArgs { value: *d }); - let rv_box = ReturnValueBox::create( - &mut builder, - &ReturnValueBoxArgs { - value: Some(off.as_union_value()), - value_type: FbReturnValue::hldouble, - }, - ); - let fcr = FbFunctionCallResult::create( - &mut builder, - &FbFunctionCallResultArgs { - result: Some(rv_box.as_union_value()), - result_type: FunctionCallResultType::ReturnValueBox, - }, - ); - builder.finish_size_prefixed(fcr, None); - builder.finished_data().to_vec() - } - ReturnValue::Bool(b) => { - let off = hlbool::create(&mut builder, &hlboolArgs { value: *b }); - let rv_box = ReturnValueBox::create( - &mut builder, - &ReturnValueBoxArgs { - value: Some(off.as_union_value()), - value_type: FbReturnValue::hlbool, - }, - ); - let fcr = FbFunctionCallResult::create( - &mut builder, - &FbFunctionCallResultArgs { - result: Some(rv_box.as_union_value()), - result_type: FunctionCallResultType::ReturnValueBox, - }, - ); - builder.finish_size_prefixed(fcr, None); - builder.finished_data().to_vec() - } - ReturnValue::String(s) => { - let off = { - let val = builder.create_string(s.as_str()); - hlstring::create(&mut builder, &hlstringArgs { value: Some(val) }) - }; - let rv_box = ReturnValueBox::create( - &mut builder, - &ReturnValueBoxArgs { - value: Some(off.as_union_value()), - value_type: FbReturnValue::hlstring, - }, - ); - let fcr = FbFunctionCallResult::create( - &mut builder, - &FbFunctionCallResultArgs { - result: Some(rv_box.as_union_value()), - result_type: FunctionCallResultType::ReturnValueBox, - }, - ); - builder.finish_size_prefixed(fcr, None); - builder.finished_data().to_vec() - } - ReturnValue::VecBytes(v) => { - let off = { - let val = builder.create_vector(v.as_slice()); - hlsizeprefixedbuffer::create( - &mut builder, - &hlsizeprefixedbufferArgs { - value: Some(val), - size: v.len() as i32, - }, - ) - }; - let rv_box = ReturnValueBox::create( - &mut builder, - &ReturnValueBoxArgs { - value: Some(off.as_union_value()), - value_type: FbReturnValue::hlsizeprefixedbuffer, - }, - ); - let fcr = FbFunctionCallResult::create( - &mut builder, - &FbFunctionCallResultArgs { - result: Some(rv_box.as_union_value()), - result_type: FunctionCallResultType::ReturnValueBox, - }, - ); - builder.finish_size_prefixed(fcr, None); - builder.finished_data().to_vec() - } - ReturnValue::Void(()) => { - let off = hlvoid::create(&mut builder, &hlvoidArgs {}); - let rv_box = ReturnValueBox::create( - &mut builder, - &ReturnValueBoxArgs { - value: Some(off.as_union_value()), - value_type: FbReturnValue::hlvoid, - }, - ); - let fcr = FbFunctionCallResult::create( - &mut builder, - &FbFunctionCallResultArgs { - result: Some(rv_box.as_union_value()), - result_type: FunctionCallResultType::ReturnValueBox, - }, - ); - builder.finish_size_prefixed(fcr, None); - builder.finished_data().to_vec() - } - }; - - Ok(result_bytes) - } -} - #[cfg(test)] mod tests { + use alloc::collections::VecDeque; + use alloc::vec; + use flatbuffers::FlatBufferBuilder; use super::super::guest_error::ErrorCode; use super::*; + use crate::flatbuffers::hyperlight::generated::{hlexternalbytes, hlexternalbytesArgs}; + + #[derive(Debug, Clone, PartialEq)] + enum TestExternalValue { + VecBytes(Vec), + ByteChunks(Vec), + } + + #[derive(Default)] + struct TestExternalValues { + values: VecDeque, + } + + impl TestExternalValues { + fn from_values(values: impl IntoIterator) -> Self { + Self { + values: values.into_iter().collect(), + } + } + } + + impl<'a> ExternalValueSink<'a> for TestExternalValues { + fn push_bytes(&mut self, value: &'a [u8]) -> Result<()> { + self.values + .push_back(TestExternalValue::VecBytes(value.to_vec())); + Ok(()) + } + + fn push_chunks(&mut self, value: &'a [Bytes]) -> Result<()> { + self.values + .push_back(TestExternalValue::ByteChunks(value.to_vec())); + Ok(()) + } + } + + impl ExternalValueSource for TestExternalValues { + fn take_bytes(&mut self, _length: usize) -> Result> { + match self.values.pop_front() { + Some(TestExternalValue::VecBytes(value)) => Ok(value), + Some(TestExternalValue::ByteChunks(_)) => { + anyhow::bail!("Expected external VecBytes value") + } + None => anyhow::bail!("Missing external VecBytes value"), + } + } + + fn take_chunks(&mut self, _length: usize) -> Result> { + match self.values.pop_front() { + Some(TestExternalValue::ByteChunks(value)) => Ok(value), + Some(TestExternalValue::VecBytes(_)) => { + anyhow::bail!("Expected external ByteChunks value") + } + None => anyhow::bail!("Missing external ByteChunks value"), + } + } + + fn finish(&mut self) -> Result<()> { + if self.values.is_empty() { + Ok(()) + } else { + anyhow::bail!("Unused external values") + } + } + } #[test] fn encode_success_result() { let mut builder = FlatBufferBuilder::new(); - let test_data = FunctionCallResult::new(Ok(ReturnValue::Int(42))).encode(&mut builder); + let mut external_values = TestExternalValues::default(); + + let test_data = FunctionCallResult::new(Ok(ReturnValue::Int(42))) + .encode(&mut builder, &mut external_values) + .unwrap(); + + let function_call_result = + FunctionCallResult::decode(test_data, &mut external_values).unwrap(); - let function_call_result = FunctionCallResult::try_from(test_data).unwrap(); let result = function_call_result.into_inner().unwrap(); assert_eq!(result, ReturnValue::Int(42)); } @@ -976,11 +993,124 @@ mod tests { ErrorCode::GuestFunctionNotFound, "Function not found".to_string(), ); - let test_data = FunctionCallResult::new(Err(test_error.clone())).encode(&mut builder); - let function_call_result = FunctionCallResult::try_from(test_data).unwrap(); + let mut external_values = TestExternalValues::default(); + let test_data = FunctionCallResult::new(Err(test_error.clone())) + .encode(&mut builder, &mut external_values) + .unwrap(); + + let function_call_result = + FunctionCallResult::decode(test_data, &mut external_values).unwrap(); + let error = function_call_result.into_inner().unwrap_err(); assert_eq!(error.code, test_error.code); assert_eq!(error.message, test_error.message); } + + #[test] + fn external_bytes_marks_chunked_values_only() { + fn round_trip(chunked: bool) -> bool { + let mut builder = FlatBufferBuilder::new(); + let value = hlexternalbytes::create( + &mut builder, + &hlexternalbytesArgs { + length: 42, + chunked, + }, + ); + builder.finish(value, None); + let value = flatbuffers::root::(builder.finished_data()).unwrap(); + + assert_eq!(value.length(), 42); + value.chunked() + } + + assert!(!round_trip(false)); + assert!(round_trip(true)); + } + + #[test] + fn byte_returns_round_trip_as_external_values() { + for expected in [ + ReturnValue::VecBytes(vec![0xa5; 4096]), + ReturnValue::ByteChunks(vec![ + Bytes::from_static(b"chunk one"), + Bytes::from_static(b" and two"), + ]), + ReturnValue::VecBytes(Vec::new()), + ReturnValue::ByteChunks(Vec::new()), + ] { + let mut builder = FlatBufferBuilder::new(); + let mut external_values = TestExternalValues::default(); + + let encoded = FunctionCallResult::new(Ok(expected.clone())) + .encode(&mut builder, &mut external_values) + .unwrap(); + + assert!(encoded.len() < 4096); + let encoded_result = size_prefixed_root::(encoded).unwrap(); + let return_value = encoded_result.result_as_return_value_box().unwrap(); + + assert_eq!(return_value.value_type(), FbReturnValue::hlexternalbytes); + + let metadata = return_value.value_as_hlexternalbytes().unwrap(); + let (length, chunked) = match &expected { + ReturnValue::VecBytes(value) => (value.len(), false), + ReturnValue::ByteChunks(value) => (try_byte_chunks_len(value).unwrap(), true), + _ => unreachable!(), + }; + + assert_eq!(metadata.length(), length as u64); + assert_eq!(metadata.chunked(), chunked); + + let decoded = FunctionCallResult::decode(encoded, &mut external_values) + .unwrap() + .into_inner() + .unwrap(); + + assert_eq!(decoded, expected); + assert!(external_values.values.is_empty()); + } + } + + #[test] + fn external_return_decoder_rejects_invalid_values() { + let mut builder = FlatBufferBuilder::new(); + let mut encoded_values = TestExternalValues::default(); + let encoded = + FunctionCallResult::new(Ok(ReturnValue::ByteChunks(vec![Bytes::from_static( + b"123", + )]))) + .encode(&mut builder, &mut encoded_values) + .unwrap(); + + let mut missing = TestExternalValues::default(); + assert!(FunctionCallResult::decode(encoded, &mut missing).is_err()); + + let mut wrong_type = + TestExternalValues::from_values([TestExternalValue::VecBytes(vec![1, 2, 3])]); + + assert!(FunctionCallResult::decode(encoded, &mut wrong_type).is_err()); + + let mut wrong_length = + TestExternalValues::from_values([TestExternalValue::ByteChunks(vec![ + Bytes::from_static(b"12"), + ])]); + + assert!(FunctionCallResult::decode(encoded, &mut wrong_length).is_err()); + } + + #[test] + fn external_result_decoder_rejects_unused_values() { + let result = FunctionCallResult::new(Ok(ReturnValue::Int(42))); + let mut builder = FlatBufferBuilder::new(); + let mut external_values = TestExternalValues::default(); + let encoded = result.encode(&mut builder, &mut external_values).unwrap(); + assert!(external_values.values.is_empty()); + + external_values + .values + .push_back(TestExternalValue::VecBytes(Vec::new())); + assert!(FunctionCallResult::decode(encoded, &mut external_values).is_err()); + } } diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/guest_log_level.rs b/src/hyperlight_common/src/flatbuffer_wrappers/guest_log_level.rs index deec57070..d3d1db431 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/guest_log_level.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/guest_log_level.rs @@ -83,7 +83,7 @@ impl From<&LogLevel> for FbLogLevel { } impl From<&LogLevel> for Level { - // There is a test (sandbox::outb::tests::test_log_outb_log) which emits trace record as logs + // There is a test (sandbox::outb::tests::test_log_emit_guest_log) which emits trace record as logs // which causes a panic when this function is instrumented as the logger is contained in refcell and // instrumentation ends up causing a double mutborrow. So this is not instrumented. //TODO: instrument this once we fix the test diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs b/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs index 57b320de4..d013cb78c 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs @@ -14,6 +14,7 @@ See the License for the specific language governing permissions and limitations under the License. */ +mod codec; pub mod function_call; pub mod function_types; pub mod guest_error; @@ -29,3 +30,5 @@ pub mod host_function_definition; /// cbindgen:ignore pub mod host_function_details; pub mod util; + +pub use codec::{ExternalValueSink, ExternalValueSource}; diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/util.rs b/src/hyperlight_common/src/flatbuffer_wrappers/util.rs index 2ee32c9a3..7e0ac007d 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/util.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/util.rs @@ -14,226 +14,14 @@ See the License for the specific language governing permissions and limitations under the License. */ +use alloc::vec; use alloc::vec::Vec; -use flatbuffers::FlatBufferBuilder; +use bytes::Bytes; use crate::flatbuffer_wrappers::function_types::ParameterValue; -use crate::flatbuffers::hyperlight::generated::{ - FunctionCallResult as FbFunctionCallResult, FunctionCallResultArgs as FbFunctionCallResultArgs, - FunctionCallResultType as FbFunctionCallResultType, ReturnValue as FbReturnValue, - ReturnValueBox, ReturnValueBoxArgs, hlbool as Fbhlbool, hlboolArgs as FbhlboolArgs, - hldouble as Fbhldouble, hldoubleArgs as FbhldoubleArgs, hlfloat as Fbhlfloat, - hlfloatArgs as FbhlfloatArgs, hlint as Fbhlint, hlintArgs as FbhlintArgs, hllong as Fbhllong, - hllongArgs as FbhllongArgs, hlsizeprefixedbuffer as Fbhlsizeprefixedbuffer, - hlsizeprefixedbufferArgs as FbhlsizeprefixedbufferArgs, hlstring as Fbhlstring, - hlstringArgs as FbhlstringArgs, hluint as Fbhluint, hluintArgs as FbhluintArgs, - hlulong as Fbhlulong, hlulongArgs as FbhlulongArgs, hlvoid as Fbhlvoid, - hlvoidArgs as FbhlvoidArgs, -}; - -/// Flatbuffer-encodes the given value -pub fn get_flatbuffer_result(val: T) -> Vec { - let mut builder = FlatBufferBuilder::new(); - let res = T::serialize(&val, &mut builder); - let result_offset = FbFunctionCallResult::create(&mut builder, &res); - - builder.finish_size_prefixed(result_offset, None); - - builder.finished_data().to_vec() -} - -pub trait FlatbufferSerializable { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs; -} - -// Implementations for basic types below - -impl FlatbufferSerializable for () { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { - let void_off = Fbhlvoid::create(builder, &FbhlvoidArgs {}); - let rv_box = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value_type: FbReturnValue::hlvoid, - value: Some(void_off.as_union_value()), - }, - ); - FbFunctionCallResultArgs { - result_type: FbFunctionCallResultType::ReturnValueBox, - result: Some(rv_box.as_union_value()), - } - } -} - -impl FlatbufferSerializable for &str { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { - let string_offset = builder.create_string(self); - let str_off = Fbhlstring::create( - builder, - &FbhlstringArgs { - value: Some(string_offset), - }, - ); - let rv_box = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value_type: FbReturnValue::hlstring, - value: Some(str_off.as_union_value()), - }, - ); - FbFunctionCallResultArgs { - result_type: FbFunctionCallResultType::ReturnValueBox, - result: Some(rv_box.as_union_value()), - } - } -} - -impl FlatbufferSerializable for &[u8] { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { - let vec_off = builder.create_vector(self); - let buf_off = Fbhlsizeprefixedbuffer::create( - builder, - &FbhlsizeprefixedbufferArgs { - size: self.len() as i32, - value: Some(vec_off), - }, - ); - let rv_box = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value_type: FbReturnValue::hlsizeprefixedbuffer, - value: Some(buf_off.as_union_value()), - }, - ); - FbFunctionCallResultArgs { - result_type: FbFunctionCallResultType::ReturnValueBox, - result: Some(rv_box.as_union_value()), - } - } -} - -impl FlatbufferSerializable for f32 { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { - let off = Fbhlfloat::create(builder, &FbhlfloatArgs { value: *self }); - let rv_box = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value_type: FbReturnValue::hlfloat, - value: Some(off.as_union_value()), - }, - ); - FbFunctionCallResultArgs { - result_type: FbFunctionCallResultType::ReturnValueBox, - result: Some(rv_box.as_union_value()), - } - } -} - -impl FlatbufferSerializable for f64 { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { - let off = Fbhldouble::create(builder, &FbhldoubleArgs { value: *self }); - let rv_box = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value_type: FbReturnValue::hldouble, - value: Some(off.as_union_value()), - }, - ); - FbFunctionCallResultArgs { - result_type: FbFunctionCallResultType::ReturnValueBox, - result: Some(rv_box.as_union_value()), - } - } -} - -impl FlatbufferSerializable for i32 { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { - let off = Fbhlint::create(builder, &FbhlintArgs { value: *self }); - let rv_box = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value_type: FbReturnValue::hlint, - value: Some(off.as_union_value()), - }, - ); - FbFunctionCallResultArgs { - result_type: FbFunctionCallResultType::ReturnValueBox, - result: Some(rv_box.as_union_value()), - } - } -} - -impl FlatbufferSerializable for i64 { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { - let off = Fbhllong::create(builder, &FbhllongArgs { value: *self }); - let rv_box = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value_type: FbReturnValue::hllong, - value: Some(off.as_union_value()), - }, - ); - FbFunctionCallResultArgs { - result_type: FbFunctionCallResultType::ReturnValueBox, - result: Some(rv_box.as_union_value()), - } - } -} - -impl FlatbufferSerializable for u32 { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { - let off = Fbhluint::create(builder, &FbhluintArgs { value: *self }); - let rv_box = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value_type: FbReturnValue::hluint, - value: Some(off.as_union_value()), - }, - ); - FbFunctionCallResultArgs { - result_type: FbFunctionCallResultType::ReturnValueBox, - result: Some(rv_box.as_union_value()), - } - } -} - -impl FlatbufferSerializable for u64 { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { - let off = Fbhlulong::create(builder, &FbhlulongArgs { value: *self }); - let rv_box = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value_type: FbReturnValue::hlulong, - value: Some(off.as_union_value()), - }, - ); - FbFunctionCallResultArgs { - result_type: FbFunctionCallResultType::ReturnValueBox, - result: Some(rv_box.as_union_value()), - } - } -} - -impl FlatbufferSerializable for bool { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { - let off = Fbhlbool::create(builder, &FbhlboolArgs { value: *self }); - let rv_box = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value_type: FbReturnValue::hlbool, - value: Some(off.as_union_value()), - }, - ); - FbFunctionCallResultArgs { - result_type: FbFunctionCallResultType::ReturnValueBox, - result: Some(rv_box.as_union_value()), - } - } -} -/// Estimates the required buffer capacity for encoding a FunctionCall with the given parameters. -/// This helps avoid reallocation during FlatBuffer encoding when passing large slices and strings. +/// Estimate the control-buffer capacity for encoding a function call. /// /// The function aims to be lightweight and fast and run in O(1) as long as the number of parameters is limited /// (which it is since hyperlight only currently supports up to 12). @@ -244,7 +32,7 @@ impl FlatbufferSerializable for bool { /// /// The estimations are numbers used are empirically derived based on the tests below and vaguely based /// on https://flatbuffers.dev/internals/ and https://github.com/dvidelabs/flatcc/blob/f064cefb2034d1e7407407ce32a6085c322212a7/doc/binary-format.md#flatbuffers-binary-format -#[inline] // allow cross-crate inlining (for hyperlight-host calls) +#[inline] pub fn estimate_flatbuffer_capacity(function_name: &str, args: &[ParameterValue]) -> usize { let mut estimated_capacity = 20; @@ -259,7 +47,7 @@ pub fn estimate_flatbuffer_capacity(function_name: &str, args: &[ParameterValue] estimated_capacity += 16; // Base parameter structure estimated_capacity += match arg { ParameterValue::String(s) => s.len() + 20, - ParameterValue::VecBytes(v) => v.len() + 20, + ParameterValue::VecBytes(_) | ParameterValue::ByteChunks(_) => 20, ParameterValue::Int(_) | ParameterValue::UInt(_) => 16, ParameterValue::Long(_) | ParameterValue::ULong(_) => 20, ParameterValue::Float(_) => 16, @@ -272,15 +60,73 @@ pub fn estimate_flatbuffer_capacity(function_name: &str, args: &[ParameterValue] estimated_capacity.next_power_of_two() } +/// Wrap contiguous bytes as one chunk without copying. +pub fn byte_chunks_from_bytes(value: Bytes) -> Vec { + if value.is_empty() { + Vec::new() + } else { + vec![value] + } +} + +/// Wrap a contiguous vector as one chunk without copying. +pub fn byte_chunks_from_vec(value: Vec) -> Vec { + byte_chunks_from_bytes(Bytes::from(value)) +} + +/// Return the complete logical length of a chunked byte value. +pub(crate) fn byte_chunks_len(value: &[Bytes]) -> usize { + value.iter().map(Bytes::len).sum() +} + +/// Return the complete logical length, or `None` if the sum overflows. +pub(crate) fn try_byte_chunks_len(value: &[Bytes]) -> Option { + value + .iter() + .try_fold(0usize, |length, chunk| length.checked_add(chunk.len())) +} + +/// Materialize byte chunks as contiguous [`Bytes`]. +/// +/// This is O(1) for zero or one chunk and copies once for multiple chunks. +pub fn byte_chunks_to_bytes(value: &[Bytes]) -> Bytes { + match value { + [] => Bytes::new(), + [chunk] => chunk.clone(), + chunks => Bytes::from(byte_chunks_to_vec(chunks)), + } +} + +/// Materialize byte chunks as one contiguous vector. +/// +/// This always allocates a new vector and copies the complete logical value. +pub fn byte_chunks_to_vec(value: &[Bytes]) -> Vec { + let mut output = Vec::with_capacity(byte_chunks_len(value)); + for chunk in value { + output.extend_from_slice(chunk); + } + output +} + +#[cfg(feature = "fuzzing")] +pub(crate) fn arbitrary_byte_chunks( + input: &mut arbitrary::Unstructured<'_>, +) -> arbitrary::Result> { + as arbitrary::Arbitrary>::arbitrary(input).map(byte_chunks_from_vec) +} + #[cfg(test)] mod tests { use alloc::string::ToString; use alloc::vec; use alloc::vec::Vec; + use flatbuffers::FlatBufferBuilder; + use super::*; use crate::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType}; use crate::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType}; + use crate::transport::ExternalValues; /// Helper function to check that estimation is within reasonable bounds (±25%) fn assert_estimation_accuracy( @@ -299,7 +145,8 @@ mod tests { ); // Important that this FlatBufferBuilder is created with capacity 0 so it grows to its needed capacity let mut builder = FlatBufferBuilder::new(); - let _buffer = fc.encode(&mut builder); + let mut external_values = ExternalValues::new(); + let _buffer = fc.encode(&mut builder, &mut external_values).unwrap(); let actual = builder.collapse().0.capacity(); let lower_bound = (actual as f64 * 0.75) as usize; @@ -328,6 +175,17 @@ mod tests { ); } + #[test] + fn capacity_ignores_external_byte_payload_length() { + let small = [ParameterValue::VecBytes(vec![0])]; + let large = [ParameterValue::VecBytes(vec![0; 1024 * 1024])]; + + assert_eq!( + estimate_flatbuffer_capacity("call", &small), + estimate_flatbuffer_capacity("call", &large) + ); + } + #[test] fn test_estimate_single_int_parameter() { assert_estimation_accuracy( diff --git a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlbytechunks_generated.rs b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlbytechunks_generated.rs new file mode 100644 index 000000000..d4f1940fd --- /dev/null +++ b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlbytechunks_generated.rs @@ -0,0 +1,124 @@ +// automatically generated by the FlatBuffers compiler, do not modify +// @generated +extern crate alloc; +extern crate flatbuffers; +use alloc::boxed::Box; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; +use core::cmp::Ordering; +use core::mem; + +use self::flatbuffers::{EndianScalar, Follow}; +use super::*; +pub enum hlbytechunksOffset {} +#[derive(Copy, Clone, PartialEq)] + +pub struct hlbytechunks<'a> { + pub _tab: flatbuffers::Table<'a>, +} + +impl<'a> flatbuffers::Follow<'a> for hlbytechunks<'a> { + type Inner = hlbytechunks<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { + _tab: unsafe { flatbuffers::Table::new(buf, loc) }, + } + } +} + +impl<'a> hlbytechunks<'a> { + pub const VT_VALUE: flatbuffers::VOffsetT = 4; + + #[inline] + pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self { + hlbytechunks { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr, A>, + args: &'args hlbytechunksArgs<'args>, + ) -> flatbuffers::WIPOffset> { + let mut builder = hlbytechunksBuilder::new(_fbb); + if let Some(x) = args.value { + builder.add_value(x); + } + builder.finish() + } + + #[inline] + pub fn value(&self) -> Option> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { + self._tab + .get::>>( + hlbytechunks::VT_VALUE, + None, + ) + } + } +} + +impl flatbuffers::Verifiable for hlbytechunks<'_> { + #[inline] + fn run_verifier( + v: &mut flatbuffers::Verifier, + pos: usize, + ) -> Result<(), flatbuffers::InvalidFlatbuffer> { + use self::flatbuffers::Verifiable; + v.visit_table(pos)? + .visit_field::>>( + "value", + Self::VT_VALUE, + false, + )? + .finish(); + Ok(()) + } +} +pub struct hlbytechunksArgs<'a> { + pub value: Option>>, +} +impl<'a> Default for hlbytechunksArgs<'a> { + #[inline] + fn default() -> Self { + hlbytechunksArgs { value: None } + } +} + +pub struct hlbytechunksBuilder<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> { + fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a, A>, + start_: flatbuffers::WIPOffset, +} +impl<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> hlbytechunksBuilder<'a, 'b, A> { + #[inline] + pub fn add_value(&mut self, value: flatbuffers::WIPOffset>) { + self.fbb_ + .push_slot_always::>(hlbytechunks::VT_VALUE, value); + } + #[inline] + pub fn new( + _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a, A>, + ) -> hlbytechunksBuilder<'a, 'b, A> { + let start = _fbb.start_table(); + hlbytechunksBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + flatbuffers::WIPOffset::new(o.value()) + } +} + +impl core::fmt::Debug for hlbytechunks<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut ds = f.debug_struct("hlbytechunks"); + ds.field("value", &self.value()); + ds.finish() + } +} diff --git a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlexternalbytes_generated.rs b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlexternalbytes_generated.rs new file mode 100644 index 000000000..d984c1bb8 --- /dev/null +++ b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlexternalbytes_generated.rs @@ -0,0 +1,140 @@ +// automatically generated by the FlatBuffers compiler, do not modify +// @generated +extern crate alloc; +extern crate flatbuffers; +use alloc::boxed::Box; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; +use core::cmp::Ordering; +use core::mem; + +use self::flatbuffers::{EndianScalar, Follow}; +use super::*; +pub enum hlexternalbytesOffset {} +#[derive(Copy, Clone, PartialEq)] + +pub struct hlexternalbytes<'a> { + pub _tab: flatbuffers::Table<'a>, +} + +impl<'a> flatbuffers::Follow<'a> for hlexternalbytes<'a> { + type Inner = hlexternalbytes<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { + _tab: unsafe { flatbuffers::Table::new(buf, loc) }, + } + } +} + +impl<'a> hlexternalbytes<'a> { + pub const VT_LENGTH: flatbuffers::VOffsetT = 4; + pub const VT_CHUNKED: flatbuffers::VOffsetT = 6; + + #[inline] + pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self { + hlexternalbytes { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr, A>, + args: &'args hlexternalbytesArgs, + ) -> flatbuffers::WIPOffset> { + let mut builder = hlexternalbytesBuilder::new(_fbb); + builder.add_length(args.length); + builder.add_chunked(args.chunked); + builder.finish() + } + + #[inline] + pub fn length(&self) -> u64 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { + self._tab + .get::(hlexternalbytes::VT_LENGTH, Some(0)) + .unwrap() + } + } + #[inline] + pub fn chunked(&self) -> bool { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { + self._tab + .get::(hlexternalbytes::VT_CHUNKED, Some(false)) + .unwrap() + } + } +} + +impl flatbuffers::Verifiable for hlexternalbytes<'_> { + #[inline] + fn run_verifier( + v: &mut flatbuffers::Verifier, + pos: usize, + ) -> Result<(), flatbuffers::InvalidFlatbuffer> { + use self::flatbuffers::Verifiable; + v.visit_table(pos)? + .visit_field::("length", Self::VT_LENGTH, false)? + .visit_field::("chunked", Self::VT_CHUNKED, false)? + .finish(); + Ok(()) + } +} +pub struct hlexternalbytesArgs { + pub length: u64, + pub chunked: bool, +} +impl<'a> Default for hlexternalbytesArgs { + #[inline] + fn default() -> Self { + hlexternalbytesArgs { + length: 0, + chunked: false, + } + } +} + +pub struct hlexternalbytesBuilder<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> { + fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a, A>, + start_: flatbuffers::WIPOffset, +} +impl<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> hlexternalbytesBuilder<'a, 'b, A> { + #[inline] + pub fn add_length(&mut self, length: u64) { + self.fbb_ + .push_slot::(hlexternalbytes::VT_LENGTH, length, 0); + } + #[inline] + pub fn add_chunked(&mut self, chunked: bool) { + self.fbb_ + .push_slot::(hlexternalbytes::VT_CHUNKED, chunked, false); + } + #[inline] + pub fn new( + _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a, A>, + ) -> hlexternalbytesBuilder<'a, 'b, A> { + let start = _fbb.start_table(); + hlexternalbytesBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + flatbuffers::WIPOffset::new(o.value()) + } +} + +impl core::fmt::Debug for hlexternalbytes<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut ds = f.debug_struct("hlexternalbytes"); + ds.field("length", &self.length()); + ds.field("chunked", &self.chunked()); + ds.finish() + } +} diff --git a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlsizeprefixedbytechunks_generated.rs b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlsizeprefixedbytechunks_generated.rs new file mode 100644 index 000000000..54661a89f --- /dev/null +++ b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlsizeprefixedbytechunks_generated.rs @@ -0,0 +1,150 @@ +// automatically generated by the FlatBuffers compiler, do not modify +// @generated +extern crate alloc; +extern crate flatbuffers; +use alloc::boxed::Box; +use alloc::string::{String, ToString}; +use alloc::vec::Vec; +use core::cmp::Ordering; +use core::mem; + +use self::flatbuffers::{EndianScalar, Follow}; +use super::*; +pub enum hlsizeprefixedbytechunksOffset {} +#[derive(Copy, Clone, PartialEq)] + +pub struct hlsizeprefixedbytechunks<'a> { + pub _tab: flatbuffers::Table<'a>, +} + +impl<'a> flatbuffers::Follow<'a> for hlsizeprefixedbytechunks<'a> { + type Inner = hlsizeprefixedbytechunks<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { + _tab: unsafe { flatbuffers::Table::new(buf, loc) }, + } + } +} + +impl<'a> hlsizeprefixedbytechunks<'a> { + pub const VT_SIZE: flatbuffers::VOffsetT = 4; + pub const VT_VALUE: flatbuffers::VOffsetT = 6; + + #[inline] + pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self { + hlsizeprefixedbytechunks { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut flatbuffers::FlatBufferBuilder<'bldr, A>, + args: &'args hlsizeprefixedbytechunksArgs<'args>, + ) -> flatbuffers::WIPOffset> { + let mut builder = hlsizeprefixedbytechunksBuilder::new(_fbb); + if let Some(x) = args.value { + builder.add_value(x); + } + builder.add_size(args.size); + builder.finish() + } + + #[inline] + pub fn size(&self) -> i32 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { + self._tab + .get::(hlsizeprefixedbytechunks::VT_SIZE, Some(0)) + .unwrap() + } + } + #[inline] + pub fn value(&self) -> Option> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { + self._tab + .get::>>( + hlsizeprefixedbytechunks::VT_VALUE, + None, + ) + } + } +} + +impl flatbuffers::Verifiable for hlsizeprefixedbytechunks<'_> { + #[inline] + fn run_verifier( + v: &mut flatbuffers::Verifier, + pos: usize, + ) -> Result<(), flatbuffers::InvalidFlatbuffer> { + use self::flatbuffers::Verifiable; + v.visit_table(pos)? + .visit_field::("size", Self::VT_SIZE, false)? + .visit_field::>>( + "value", + Self::VT_VALUE, + false, + )? + .finish(); + Ok(()) + } +} +pub struct hlsizeprefixedbytechunksArgs<'a> { + pub size: i32, + pub value: Option>>, +} +impl<'a> Default for hlsizeprefixedbytechunksArgs<'a> { + #[inline] + fn default() -> Self { + hlsizeprefixedbytechunksArgs { + size: 0, + value: None, + } + } +} + +pub struct hlsizeprefixedbytechunksBuilder<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> { + fbb_: &'b mut flatbuffers::FlatBufferBuilder<'a, A>, + start_: flatbuffers::WIPOffset, +} +impl<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> hlsizeprefixedbytechunksBuilder<'a, 'b, A> { + #[inline] + pub fn add_size(&mut self, size: i32) { + self.fbb_ + .push_slot::(hlsizeprefixedbytechunks::VT_SIZE, size, 0); + } + #[inline] + pub fn add_value(&mut self, value: flatbuffers::WIPOffset>) { + self.fbb_.push_slot_always::>( + hlsizeprefixedbytechunks::VT_VALUE, + value, + ); + } + #[inline] + pub fn new( + _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a, A>, + ) -> hlsizeprefixedbytechunksBuilder<'a, 'b, A> { + let start = _fbb.start_table(); + hlsizeprefixedbytechunksBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + flatbuffers::WIPOffset::new(o.value()) + } +} + +impl core::fmt::Debug for hlsizeprefixedbytechunks<'_> { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let mut ds = f.debug_struct("hlsizeprefixedbytechunks"); + ds.field("size", &self.size()); + ds.field("value", &self.value()); + ds.finish() + } +} diff --git a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_generated.rs b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_generated.rs index b0e803ec5..33f7c9819 100644 --- a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_generated.rs +++ b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_generated.rs @@ -198,6 +198,34 @@ impl<'a> Parameter<'a> { None } } + + #[inline] + #[allow(non_snake_case)] + pub fn value_as_hlexternalbytes(&self) -> Option> { + if self.value_type() == ParameterValue::hlexternalbytes { + let u = self.value(); + // Safety: + // Created from a valid Table for this object + // Which contains a valid union in this slot + Some(unsafe { hlexternalbytes::init_from_table(u) }) + } else { + None + } + } + + #[inline] + #[allow(non_snake_case)] + pub fn value_as_hlbytechunks(&self) -> Option> { + if self.value_type() == ParameterValue::hlbytechunks { + let u = self.value(); + // Safety: + // Created from a valid Table for this object + // Which contains a valid union in this slot + Some(unsafe { hlbytechunks::init_from_table(u) }) + } else { + None + } + } } impl flatbuffers::Verifiable for Parameter<'_> { @@ -260,6 +288,16 @@ impl flatbuffers::Verifiable for Parameter<'_> { "ParameterValue::hlvecbytes", pos, ), + ParameterValue::hlexternalbytes => v + .verify_union_variant::>( + "ParameterValue::hlexternalbytes", + pos, + ), + ParameterValue::hlbytechunks => v + .verify_union_variant::>( + "ParameterValue::hlbytechunks", + pos, + ), _ => Ok(()), }, )? @@ -410,6 +448,26 @@ impl core::fmt::Debug for Parameter<'_> { ) } } + ParameterValue::hlexternalbytes => { + if let Some(x) = self.value_as_hlexternalbytes() { + ds.field("value", &x) + } else { + ds.field( + "value", + &"InvalidFlatbuffer: Union discriminant does not match value.", + ) + } + } + ParameterValue::hlbytechunks => { + if let Some(x) = self.value_as_hlbytechunks() { + ds.field("value", &x) + } else { + ds.field( + "value", + &"InvalidFlatbuffer: Union discriminant does not match value.", + ) + } + } _ => { let x: Option<()> = None; ds.field("value", &x) diff --git a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_type_generated.rs b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_type_generated.rs index cf46560b1..cd280599a 100644 --- a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_type_generated.rs +++ b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_type_generated.rs @@ -19,13 +19,13 @@ pub const ENUM_MIN_PARAMETER_TYPE: u8 = 0; since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021." )] -pub const ENUM_MAX_PARAMETER_TYPE: u8 = 8; +pub const ENUM_MAX_PARAMETER_TYPE: u8 = 9; #[deprecated( since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021." )] #[allow(non_camel_case_types)] -pub const ENUM_VALUES_PARAMETER_TYPE: [ParameterType; 9] = [ +pub const ENUM_VALUES_PARAMETER_TYPE: [ParameterType; 10] = [ ParameterType::hlint, ParameterType::hluint, ParameterType::hllong, @@ -35,6 +35,7 @@ pub const ENUM_VALUES_PARAMETER_TYPE: [ParameterType; 9] = [ ParameterType::hlstring, ParameterType::hlbool, ParameterType::hlvecbytes, + ParameterType::hlbytechunks, ]; #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] @@ -51,9 +52,10 @@ impl ParameterType { pub const hlstring: Self = Self(6); pub const hlbool: Self = Self(7); pub const hlvecbytes: Self = Self(8); + pub const hlbytechunks: Self = Self(9); pub const ENUM_MIN: u8 = 0; - pub const ENUM_MAX: u8 = 8; + pub const ENUM_MAX: u8 = 9; pub const ENUM_VALUES: &'static [Self] = &[ Self::hlint, Self::hluint, @@ -64,6 +66,7 @@ impl ParameterType { Self::hlstring, Self::hlbool, Self::hlvecbytes, + Self::hlbytechunks, ]; /// Returns the variant's name or "" if unknown. pub fn variant_name(self) -> Option<&'static str> { @@ -77,6 +80,7 @@ impl ParameterType { Self::hlstring => Some("hlstring"), Self::hlbool => Some("hlbool"), Self::hlvecbytes => Some("hlvecbytes"), + Self::hlbytechunks => Some("hlbytechunks"), _ => None, } } diff --git a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_value_generated.rs b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_value_generated.rs index 8113df5fc..5ddab887f 100644 --- a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_value_generated.rs +++ b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/parameter_value_generated.rs @@ -19,13 +19,13 @@ pub const ENUM_MIN_PARAMETER_VALUE: u8 = 0; since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021." )] -pub const ENUM_MAX_PARAMETER_VALUE: u8 = 9; +pub const ENUM_MAX_PARAMETER_VALUE: u8 = 11; #[deprecated( since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021." )] #[allow(non_camel_case_types)] -pub const ENUM_VALUES_PARAMETER_VALUE: [ParameterValue; 10] = [ +pub const ENUM_VALUES_PARAMETER_VALUE: [ParameterValue; 12] = [ ParameterValue::NONE, ParameterValue::hlint, ParameterValue::hluint, @@ -36,6 +36,8 @@ pub const ENUM_VALUES_PARAMETER_VALUE: [ParameterValue; 10] = [ ParameterValue::hlstring, ParameterValue::hlbool, ParameterValue::hlvecbytes, + ParameterValue::hlexternalbytes, + ParameterValue::hlbytechunks, ]; #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] @@ -53,9 +55,11 @@ impl ParameterValue { pub const hlstring: Self = Self(7); pub const hlbool: Self = Self(8); pub const hlvecbytes: Self = Self(9); + pub const hlexternalbytes: Self = Self(10); + pub const hlbytechunks: Self = Self(11); pub const ENUM_MIN: u8 = 0; - pub const ENUM_MAX: u8 = 9; + pub const ENUM_MAX: u8 = 11; pub const ENUM_VALUES: &'static [Self] = &[ Self::NONE, Self::hlint, @@ -67,6 +71,8 @@ impl ParameterValue { Self::hlstring, Self::hlbool, Self::hlvecbytes, + Self::hlexternalbytes, + Self::hlbytechunks, ]; /// Returns the variant's name or "" if unknown. pub fn variant_name(self) -> Option<&'static str> { @@ -81,6 +87,8 @@ impl ParameterValue { Self::hlstring => Some("hlstring"), Self::hlbool => Some("hlbool"), Self::hlvecbytes => Some("hlvecbytes"), + Self::hlexternalbytes => Some("hlexternalbytes"), + Self::hlbytechunks => Some("hlbytechunks"), _ => None, } } diff --git a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_type_generated.rs b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_type_generated.rs index 913b1fe78..06ea93d9c 100644 --- a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_type_generated.rs +++ b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_type_generated.rs @@ -19,13 +19,13 @@ pub const ENUM_MIN_RETURN_TYPE: u8 = 0; since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021." )] -pub const ENUM_MAX_RETURN_TYPE: u8 = 9; +pub const ENUM_MAX_RETURN_TYPE: u8 = 10; #[deprecated( since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021." )] #[allow(non_camel_case_types)] -pub const ENUM_VALUES_RETURN_TYPE: [ReturnType; 10] = [ +pub const ENUM_VALUES_RETURN_TYPE: [ReturnType; 11] = [ ReturnType::hlint, ReturnType::hluint, ReturnType::hllong, @@ -36,6 +36,7 @@ pub const ENUM_VALUES_RETURN_TYPE: [ReturnType; 10] = [ ReturnType::hlbool, ReturnType::hlvoid, ReturnType::hlsizeprefixedbuffer, + ReturnType::hlbytechunks, ]; #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] @@ -53,9 +54,10 @@ impl ReturnType { pub const hlbool: Self = Self(7); pub const hlvoid: Self = Self(8); pub const hlsizeprefixedbuffer: Self = Self(9); + pub const hlbytechunks: Self = Self(10); pub const ENUM_MIN: u8 = 0; - pub const ENUM_MAX: u8 = 9; + pub const ENUM_MAX: u8 = 10; pub const ENUM_VALUES: &'static [Self] = &[ Self::hlint, Self::hluint, @@ -67,6 +69,7 @@ impl ReturnType { Self::hlbool, Self::hlvoid, Self::hlsizeprefixedbuffer, + Self::hlbytechunks, ]; /// Returns the variant's name or "" if unknown. pub fn variant_name(self) -> Option<&'static str> { @@ -81,6 +84,7 @@ impl ReturnType { Self::hlbool => Some("hlbool"), Self::hlvoid => Some("hlvoid"), Self::hlsizeprefixedbuffer => Some("hlsizeprefixedbuffer"), + Self::hlbytechunks => Some("hlbytechunks"), _ => None, } } diff --git a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_value_box_generated.rs b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_value_box_generated.rs index cecd8b6c1..854879a0e 100644 --- a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_value_box_generated.rs +++ b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_value_box_generated.rs @@ -212,6 +212,34 @@ impl<'a> ReturnValueBox<'a> { None } } + + #[inline] + #[allow(non_snake_case)] + pub fn value_as_hlexternalbytes(&self) -> Option> { + if self.value_type() == ReturnValue::hlexternalbytes { + let u = self.value(); + // Safety: + // Created from a valid Table for this object + // Which contains a valid union in this slot + Some(unsafe { hlexternalbytes::init_from_table(u) }) + } else { + None + } + } + + #[inline] + #[allow(non_snake_case)] + pub fn value_as_hlsizeprefixedbytechunks(&self) -> Option> { + if self.value_type() == ReturnValue::hlsizeprefixedbytechunks { + let u = self.value(); + // Safety: + // Created from a valid Table for this object + // Which contains a valid union in this slot + Some(unsafe { hlsizeprefixedbytechunks::init_from_table(u) }) + } else { + None + } + } } impl flatbuffers::Verifiable for ReturnValueBox<'_> { @@ -222,67 +250,24 @@ impl flatbuffers::Verifiable for ReturnValueBox<'_> { ) -> Result<(), flatbuffers::InvalidFlatbuffer> { use self::flatbuffers::Verifiable; v.visit_table(pos)? - .visit_union::( - "value_type", - Self::VT_VALUE_TYPE, - "value", - Self::VT_VALUE, - true, - |key, v, pos| match key { - ReturnValue::hlint => v - .verify_union_variant::>( - "ReturnValue::hlint", - pos, - ), - ReturnValue::hluint => v - .verify_union_variant::>( - "ReturnValue::hluint", - pos, - ), - ReturnValue::hllong => v - .verify_union_variant::>( - "ReturnValue::hllong", - pos, - ), - ReturnValue::hlulong => v - .verify_union_variant::>( - "ReturnValue::hlulong", - pos, - ), - ReturnValue::hlfloat => v - .verify_union_variant::>( - "ReturnValue::hlfloat", - pos, - ), - ReturnValue::hldouble => v - .verify_union_variant::>( - "ReturnValue::hldouble", - pos, - ), - ReturnValue::hlstring => v - .verify_union_variant::>( - "ReturnValue::hlstring", - pos, - ), - ReturnValue::hlbool => v - .verify_union_variant::>( - "ReturnValue::hlbool", - pos, - ), - ReturnValue::hlvoid => v - .verify_union_variant::>( - "ReturnValue::hlvoid", - pos, - ), - ReturnValue::hlsizeprefixedbuffer => v - .verify_union_variant::>( - "ReturnValue::hlsizeprefixedbuffer", - pos, - ), - _ => Ok(()), - }, - )? - .finish(); + .visit_union::("value_type", Self::VT_VALUE_TYPE, "value", Self::VT_VALUE, true, |key, v, pos| { + match key { + ReturnValue::hlint => v.verify_union_variant::>("ReturnValue::hlint", pos), + ReturnValue::hluint => v.verify_union_variant::>("ReturnValue::hluint", pos), + ReturnValue::hllong => v.verify_union_variant::>("ReturnValue::hllong", pos), + ReturnValue::hlulong => v.verify_union_variant::>("ReturnValue::hlulong", pos), + ReturnValue::hlfloat => v.verify_union_variant::>("ReturnValue::hlfloat", pos), + ReturnValue::hldouble => v.verify_union_variant::>("ReturnValue::hldouble", pos), + ReturnValue::hlstring => v.verify_union_variant::>("ReturnValue::hlstring", pos), + ReturnValue::hlbool => v.verify_union_variant::>("ReturnValue::hlbool", pos), + ReturnValue::hlvoid => v.verify_union_variant::>("ReturnValue::hlvoid", pos), + ReturnValue::hlsizeprefixedbuffer => v.verify_union_variant::>("ReturnValue::hlsizeprefixedbuffer", pos), + ReturnValue::hlexternalbytes => v.verify_union_variant::>("ReturnValue::hlexternalbytes", pos), + ReturnValue::hlsizeprefixedbytechunks => v.verify_union_variant::>("ReturnValue::hlsizeprefixedbytechunks", pos), + _ => Ok(()), + } + })? + .finish(); Ok(()) } } @@ -441,6 +426,26 @@ impl core::fmt::Debug for ReturnValueBox<'_> { ) } } + ReturnValue::hlexternalbytes => { + if let Some(x) = self.value_as_hlexternalbytes() { + ds.field("value", &x) + } else { + ds.field( + "value", + &"InvalidFlatbuffer: Union discriminant does not match value.", + ) + } + } + ReturnValue::hlsizeprefixedbytechunks => { + if let Some(x) = self.value_as_hlsizeprefixedbytechunks() { + ds.field("value", &x) + } else { + ds.field( + "value", + &"InvalidFlatbuffer: Union discriminant does not match value.", + ) + } + } _ => { let x: Option<()> = None; ds.field("value", &x) diff --git a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_value_generated.rs b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_value_generated.rs index d13c73623..6a6e619f6 100644 --- a/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_value_generated.rs +++ b/src/hyperlight_common/src/flatbuffers/hyperlight/generated/return_value_generated.rs @@ -19,13 +19,13 @@ pub const ENUM_MIN_RETURN_VALUE: u8 = 0; since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021." )] -pub const ENUM_MAX_RETURN_VALUE: u8 = 10; +pub const ENUM_MAX_RETURN_VALUE: u8 = 12; #[deprecated( since = "2.0.0", note = "Use associated constants instead. This will no longer be generated in 2021." )] #[allow(non_camel_case_types)] -pub const ENUM_VALUES_RETURN_VALUE: [ReturnValue; 11] = [ +pub const ENUM_VALUES_RETURN_VALUE: [ReturnValue; 13] = [ ReturnValue::NONE, ReturnValue::hlint, ReturnValue::hluint, @@ -37,6 +37,8 @@ pub const ENUM_VALUES_RETURN_VALUE: [ReturnValue; 11] = [ ReturnValue::hlbool, ReturnValue::hlvoid, ReturnValue::hlsizeprefixedbuffer, + ReturnValue::hlexternalbytes, + ReturnValue::hlsizeprefixedbytechunks, ]; #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)] @@ -55,9 +57,11 @@ impl ReturnValue { pub const hlbool: Self = Self(8); pub const hlvoid: Self = Self(9); pub const hlsizeprefixedbuffer: Self = Self(10); + pub const hlexternalbytes: Self = Self(11); + pub const hlsizeprefixedbytechunks: Self = Self(12); pub const ENUM_MIN: u8 = 0; - pub const ENUM_MAX: u8 = 10; + pub const ENUM_MAX: u8 = 12; pub const ENUM_VALUES: &'static [Self] = &[ Self::NONE, Self::hlint, @@ -70,6 +74,8 @@ impl ReturnValue { Self::hlbool, Self::hlvoid, Self::hlsizeprefixedbuffer, + Self::hlexternalbytes, + Self::hlsizeprefixedbytechunks, ]; /// Returns the variant's name or "" if unknown. pub fn variant_name(self) -> Option<&'static str> { @@ -85,6 +91,8 @@ impl ReturnValue { Self::hlbool => Some("hlbool"), Self::hlvoid => Some("hlvoid"), Self::hlsizeprefixedbuffer => Some("hlsizeprefixedbuffer"), + Self::hlexternalbytes => Some("hlexternalbytes"), + Self::hlsizeprefixedbytechunks => Some("hlsizeprefixedbytechunks"), _ => None, } } diff --git a/src/hyperlight_common/src/flatbuffers/mod.rs b/src/hyperlight_common/src/flatbuffers/mod.rs index 184260572..6e8f1125b 100644 --- a/src/hyperlight_common/src/flatbuffers/mod.rs +++ b/src/hyperlight_common/src/flatbuffers/mod.rs @@ -40,8 +40,14 @@ pub mod hyperlight { pub use self::hlbool_generated::*; mod hlvecbytes_generated; pub use self::hlvecbytes_generated::*; + mod hlbytechunks_generated; + pub use self::hlbytechunks_generated::*; + mod hlexternalbytes_generated; + pub use self::hlexternalbytes_generated::*; mod hlsizeprefixedbuffer_generated; pub use self::hlsizeprefixedbuffer_generated::*; + mod hlsizeprefixedbytechunks_generated; + pub use self::hlsizeprefixedbytechunks_generated::*; mod hlvoid_generated; pub use self::hlvoid_generated::*; mod guest_error_generated; diff --git a/src/hyperlight_common/src/func/mod.rs b/src/hyperlight_common/src/func/mod.rs index 5556f6fe7..9f55119dd 100644 --- a/src/hyperlight_common/src/func/mod.rs +++ b/src/hyperlight_common/src/func/mod.rs @@ -39,6 +39,8 @@ pub use functions::Function; pub use param_type::{ParameterTuple, SupportedParameterType}; pub use ret_type::{ResultType, SupportedReturnType}; +/// Re-export for chunk-preserving byte values +pub use crate::flatbuffer_wrappers::function_types::Bytes; /// Re-export for `ParameterValue` enum pub use crate::flatbuffer_wrappers::function_types::ParameterValue; /// Re-export for `ReturnType` enum diff --git a/src/hyperlight_common/src/func/param_type.rs b/src/hyperlight_common/src/func/param_type.rs index b4db004d8..07ae343d5 100644 --- a/src/hyperlight_common/src/func/param_type.rs +++ b/src/hyperlight_common/src/func/param_type.rs @@ -20,7 +20,7 @@ use alloc::vec::Vec; use super::error::Error; use super::utils::for_each_tuple; -use crate::flatbuffer_wrappers::function_types::{ParameterType, ParameterValue}; +use crate::flatbuffer_wrappers::function_types::{Bytes, ParameterType, ParameterValue}; /// This is a marker trait that is used to indicate that a type is a /// valid Hyperlight parameter type. @@ -50,6 +50,7 @@ macro_rules! for_each_param_type { $macro!(f64, Double); $macro!(bool, Bool); $macro!(Vec, VecBytes); + $macro!(Vec, ByteChunks); }; } @@ -135,3 +136,22 @@ macro_rules! impl_param_tuple { } for_each_tuple!(impl_param_tuple); + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn byte_chunks_parameter_round_trips_without_copying() { + let chunks = vec![Bytes::from_static(b"hello"), Bytes::from_static(b" world")]; + let first_chunk = chunks[0].as_ptr(); + let value = as SupportedParameterType>::into_value(chunks); + let chunks = as SupportedParameterType>::from_value(value).unwrap(); + + assert_eq!(chunks[0].as_ptr(), first_chunk); + assert_eq!( + chunks, + [Bytes::from_static(b"hello"), Bytes::from_static(b" world")] + ); + } +} diff --git a/src/hyperlight_common/src/func/ret_type.rs b/src/hyperlight_common/src/func/ret_type.rs index 69e0f1a13..fa8b64646 100644 --- a/src/hyperlight_common/src/func/ret_type.rs +++ b/src/hyperlight_common/src/func/ret_type.rs @@ -21,6 +21,10 @@ use super::error::Error; use crate::flatbuffer_wrappers::function_types::{ReturnType, ReturnValue}; /// This is a marker trait that is used to indicate that a type is a valid Hyperlight return type. +/// +/// `Vec` and `Vec` are distinct logical return types. `Vec` +/// carries contiguous bytes, while `Vec` preserves local chunk +/// ownership. Chunk boundaries are not application framing. pub trait SupportedReturnType: Sized + Clone + Send + Sync + 'static { /// The return type of the supported return value const TYPE: ReturnType; @@ -46,6 +50,7 @@ macro_rules! for_each_return_type { $macro!(f64, Double); $macro!(bool, Bool); $macro!(Vec, VecBytes); + $macro!(Vec<$crate::func::Bytes>, ByteChunks); }; } @@ -105,3 +110,23 @@ where } for_each_return_type!(impl_supported_return_type); + +#[cfg(test)] +mod tests { + use super::*; + use crate::flatbuffer_wrappers::function_types::Bytes; + + #[test] + fn byte_chunks_return_round_trips_without_copying() { + let chunks = vec![Bytes::from_static(b"hello"), Bytes::from_static(b" world")]; + let first_chunk = chunks[0].as_ptr(); + let value = as SupportedReturnType>::into_value(chunks); + let chunks = as SupportedReturnType>::from_value(value).unwrap(); + + assert_eq!(chunks[0].as_ptr(), first_chunk); + assert_eq!( + chunks, + [Bytes::from_static(b"hello"), Bytes::from_static(b" world")] + ); + } +} diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index bf25a2e0c..9da4d53b5 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -14,6 +14,9 @@ See the License for the specific language governing permissions and limitations under the License. */ +use core::mem::{align_of, 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; @@ -22,12 +25,90 @@ 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_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. + _reserved: u64, + /// 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, + /// 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_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_TRANSPORT_ARENA_GPA_OFFSET == 0x28); + assert!(SCRATCH_TOP_G2H_QUEUE_SIZE_OFFSET == 0x30); + assert!(SCRATCH_TOP_G2H_POOL_PAGES_OFFSET == 0x38); + assert!(SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET == 0x40); + assert!(SCRATCH_TOP_H2G_QUEUE_SIZE_OFFSET == 0x48); + assert!(SCRATCH_TOP_H2G_POOL_PAGES_OFFSET == 0x50); + assert!(SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET == 0x58); + assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x60); +}; + +/// 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 @@ -37,4 +118,273 @@ pub fn scratch_base_gva(size: usize) -> u64 { } /// Compute the minimum scratch region size needed for a sandbox. -pub use arch::min_scratch_size; +/// +/// The fixed transport prefix contains one page-backed ring arena and both +/// page-backed buffer pools. The result saturates at [`usize::MAX`]. +pub fn min_scratch_size( + g2h_queue_size: usize, + h2g_queue_size: usize, + g2h_pool_pages: usize, + h2g_pool_pages: usize, +) -> usize { + let size = arch::min_scratch_size().and_then(|fixed| { + let g2h = QueueDims::new(g2h_queue_size, g2h_pool_pages)?; + let h2g = QueueDims::new(h2g_queue_size, h2g_pool_pages)?; + + let transport_len = TransportArena::checked_query_size(g2h, h2g)?; + fixed.checked_add(transport_len) + }); + + size.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 one queue descriptor count and pool page count. + 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)?; + + virtq::Layout::checked_query_size(usize::from(size.get()))?; + + 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. + pub const fn ring_len(&self) -> usize { + virtq::Layout::query_size(self.size.get() as usize) + } + + /// Buffer pool length in bytes. + pub const fn pool_len(&self) -> usize { + self.pool_pages.get() * crate::vmem::PAGE_SIZE + } +} + +/// Addresses of both rings, the checkpoint mailbox, and pools in one fixed arena. +/// +/// The G2H ring begins at the arena base. The H2G ring is descriptor aligned. +/// The mailbox is `u64` aligned. Both pools are page aligned. +/// +/// ```text +/// +----------+-----+----------+-----+-----+-----+----------+----------+ +/// | G2H ring | pad | H2G ring | pad | mbx | 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 snapshot checkpoint mailbox. + mbx_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 and the mailbox. + 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 mbx_offset = h2g_ring_offset + .checked_add(h2g.ring_len())? + .checked_next_multiple_of(align_of::())?; + + let g2h_pool_offset = mbx_offset + .checked_add(size_of::())? + .checked_next_multiple_of(crate::vmem::PAGE_SIZE)?; + + let h2g_pool_offset = g2h_pool_offset.checked_add(g2h.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)?, + mbx_addr: addr(mbx_offset)?, + g2h_pool_addr: addr(g2h_pool_offset)?, + h2g_pool_addr: addr(h2g_pool_offset)?, + ring_span_len: g2h_pool_offset, + len, + }) + } + + /// Compute the total arena size without assigning an address. + pub fn checked_query_size(g2h: QueueDims, h2g: QueueDims) -> Option { + Some(Self::new(0, g2h, h2g)?.len) + } + + /// Base address of the arena. + pub const fn base_addr(&self) -> u64 { + self.g2h_ring_addr + } + + /// Address of the G2H ring. + pub const fn g2h_ring_addr(&self) -> u64 { + self.g2h_ring_addr + } + + /// Address of the H2G ring. + pub const fn h2g_ring_addr(&self) -> u64 { + self.h2g_ring_addr + } + + /// Address of the snapshot checkpoint mailbox. + pub const fn mbx_addr(&self) -> u64 { + self.mbx_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 and the mailbox. + 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, 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.mbx_addr), + to_offset(self.g2h_pool_addr), + to_offset(self.h2g_pool_addr), + self.len, + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn transport_arena_derives_aligned_regions() { + let base = 0x1_0000; + let g2h = QueueDims::new(64, 8).unwrap(); + let h2g = QueueDims::new(32, 4).unwrap(); + let arena = TransportArena::new(base, g2h, h2g).unwrap(); + + assert_eq!(g2h.ring_len(), virtq::Layout::query_size(64)); + assert_eq!(g2h.pool_len(), 8 * crate::vmem::PAGE_SIZE); + assert_eq!(arena.g2h_ring_addr(), base); + assert!( + arena + .h2g_ring_addr() + .is_multiple_of(virtq::Descriptor::ALIGN as u64) + ); + assert!(arena.mbx_addr().is_multiple_of(align_of::() 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, + 0x618, + crate::vmem::PAGE_SIZE, + 9 * crate::vmem::PAGE_SIZE, + 13 * crate::vmem::PAGE_SIZE, + ) + ); + assert_eq!(arena.ring_span_len(), crate::vmem::PAGE_SIZE); + assert_eq!(arena.size(), 13 * crate::vmem::PAGE_SIZE); + assert_eq!( + TransportArena::checked_query_size(g2h, h2g), + Some(arena.size()) + ); + assert_eq!(TransportArena::new(base + 1, g2h, h2g), None); + assert_eq!(QueueDims::new(3, 8), None); + assert_eq!(QueueDims::new(64, 0), None); + assert_eq!(QueueDims::new(usize::MAX, 8), None); + assert_eq!(QueueDims::new(64, usize::MAX), 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().unwrap(); + let transport_pages = 1 + 8 + 4; + + assert_eq!( + fixed + transport_pages * crate::vmem::PAGE_SIZE, + min_scratch_size(64, 32, 8, 4) + ); + } + + #[test] + fn minimum_scratch_saturates_on_overflow() { + assert_eq!(usize::MAX, min_scratch_size(64, 32, usize::MAX, 4)); + assert_eq!(usize::MAX, min_scratch_size(usize::MAX, 32, 8, 4)); + } +} diff --git a/src/hyperlight_common/src/lib.rs b/src/hyperlight_common/src/lib.rs index 6e12d8cd4..41b58fb03 100644 --- a/src/hyperlight_common/src/lib.rs +++ b/src/hyperlight_common/src/lib.rs @@ -42,6 +42,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/mem.rs b/src/hyperlight_common/src/mem.rs index 46798b7ee..653056ae8 100644 --- a/src/hyperlight_common/src/mem.rs +++ b/src/hyperlight_common/src/mem.rs @@ -31,8 +31,6 @@ pub struct GuestMemoryRegion { #[derive(Debug, Clone, Copy, PartialEq, bytemuck::Pod, bytemuck::Zeroable)] #[repr(C)] pub struct HyperlightPEB { - pub input_stack: GuestMemoryRegion, - pub output_stack: GuestMemoryRegion, pub init_data: GuestMemoryRegion, pub guest_heap: GuestMemoryRegion, } @@ -44,22 +42,14 @@ mod tests { #[test] fn peb_round_trip() { let peb = HyperlightPEB { - input_stack: GuestMemoryRegion { + init_data: GuestMemoryRegion { size: 0x1111, ptr: 0x2222, }, - output_stack: GuestMemoryRegion { + guest_heap: GuestMemoryRegion { size: 0x3333, ptr: 0x4444, }, - init_data: GuestMemoryRegion { - size: 0x5555, - ptr: 0x6666, - }, - guest_heap: GuestMemoryRegion { - size: 0x7777, - ptr: 0x8888, - }, }; let bytes = bytemuck::bytes_of(&peb); let peb2 = *bytemuck::from_bytes::(bytes); diff --git a/src/hyperlight_common/src/outb.rs b/src/hyperlight_common/src/outb.rs index 3bfb99848..d9a2386a1 100644 --- a/src/hyperlight_common/src/outb.rs +++ b/src/hyperlight_common/src/outb.rs @@ -87,16 +87,13 @@ impl TryFrom for Exception { /// Supported actions when issuing an OUTB actions by Hyperlight. /// These are handled by the sandbox-level outb dispatcher. -/// - Log: for logging, -/// - CallFunction: makes a call to a host function, /// - Abort: aborts the execution of the guest, /// - DebugPrint: prints a message to the host /// - TraceBatch: reports a batch of spans and events from the guest /// - TraceMemoryAlloc: records memory allocation events /// - TraceMemoryFree: records memory deallocation events +/// - VirtqNotify: reports newly available virtqueue work pub enum OutBAction { - Log = 99, - CallFunction = 101, Abort = 102, DebugPrint = 103, #[cfg(feature = "trace_guest")] @@ -105,6 +102,7 @@ pub enum OutBAction { TraceMemoryAlloc = 105, #[cfg(feature = "mem_profile")] TraceMemoryFree = 106, + VirtqNotify = 109, } /// IO-port actions intercepted at the hypervisor level (in `run_vcpu`) @@ -127,8 +125,6 @@ impl TryFrom for OutBAction { type Error = anyhow::Error; fn try_from(val: u16) -> anyhow::Result { match val { - 99 => Ok(OutBAction::Log), - 101 => Ok(OutBAction::CallFunction), 102 => Ok(OutBAction::Abort), 103 => Ok(OutBAction::DebugPrint), #[cfg(feature = "trace_guest")] @@ -137,6 +133,7 @@ impl TryFrom for OutBAction { 105 => Ok(OutBAction::TraceMemoryAlloc), #[cfg(feature = "mem_profile")] 106 => Ok(OutBAction::TraceMemoryFree), + 109 => Ok(OutBAction::VirtqNotify), _ => Err(anyhow::anyhow!("Invalid OutBAction value: {}", val)), } } @@ -152,3 +149,14 @@ impl TryFrom for VmAction { } } } + +#[cfg(test)] +mod tests { + use super::OutBAction; + + #[test] + fn rejects_legacy_stack_actions() { + assert!(OutBAction::try_from(99).is_err()); + assert!(OutBAction::try_from(101).is_err()); + } +} diff --git a/src/hyperlight_common/src/transport.rs b/src/hyperlight_common/src/transport.rs new file mode 100644 index 000000000..7e587738d --- /dev/null +++ b/src/hyperlight_common/src/transport.rs @@ -0,0 +1,463 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +//! 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/buffer.rs b/src/hyperlight_common/src/virtq/buffer.rs index 0d1ca4d6a..2ba8839bc 100644 --- a/src/hyperlight_common/src/virtq/buffer.rs +++ b/src/hyperlight_common/src/virtq/buffer.rs @@ -14,131 +14,15 @@ See the License for the specific language governing permissions and limitations under the License. */ -//! 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::pool::{AllocError, Allocation, BufferProvider}; /// Ordered byte segments that make up one virtqueue payload. /// @@ -187,6 +71,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()) @@ -216,6 +128,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())); @@ -308,7 +225,7 @@ pub struct BufferOwner { impl AsRef<[u8]> for BufferOwner { fn as_ref(&self) -> &[u8] { let alloc = self.alloc.allocation(); - let len = self.written.min(alloc.len); + let len = self.written.min(alloc.len as usize); // Safety: BufferOwner keeps both the pool allocation and the M alive, // so the memory region is valid. match unsafe { self.mem.as_slice(alloc.addr, len) } { @@ -322,16 +239,8 @@ impl AsRef<[u8]> for BufferOwner { } /// Pool-owned allocation that is returned to the pool on drop. -/// -/// Use [`into_raw`](Self::into_raw) to transfer ownership to a descriptor -/// state that will deallocate the raw [`Allocation`] through another path. #[derive(Debug)] pub struct OwnedAlloc { - inner: Option>, -} - -#[derive(Debug)] -struct Inner { pool: P, alloc: Allocation, } @@ -339,9 +248,7 @@ struct Inner { impl OwnedAlloc

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

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

{ fn drop(&mut self) { - if let Some(Inner { pool, alloc }) = self.inner.take() { - let result = pool.dealloc(alloc.addr); - debug_assert!(result.is_ok(), "OwnedAlloc drop dealloc failed: {result:?}"); - } + let result = self.pool.dealloc(self.alloc.addr); + debug_assert!(result.is_ok(), "OwnedAlloc drop dealloc failed: {result:?}"); } } @@ -472,6 +359,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]); @@ -482,4 +395,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 ad58b05cb..c4a08aa99 100644 --- a/src/hyperlight_common/src/virtq/concurrency.rs +++ b/src/hyperlight_common/src/virtq/concurrency.rs @@ -62,7 +62,7 @@ use loom::thread; use super::*; use crate::virtq::desc::Descriptor; -use crate::virtq::pool::BufferPoolSync; +use crate::virtq::pool::RunPoolSync; #[derive(Debug)] pub struct MemErr; @@ -329,7 +329,7 @@ fn virtq_ping_pong() { let pool_size = 0x10000; let mem = Arc::new(LoomMem::new(ring_base, 8, pool_base, pool_size)); - let pool = BufferPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); + let pool = RunPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); let notify = Arc::new(Notify::new()); let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); @@ -356,12 +356,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(); @@ -377,7 +377,7 @@ fn virtq_ack_only() { let pool_size = 0x10000; let mem = Arc::new(LoomMem::new(ring_base, 4, pool_base, pool_size)); - let pool = BufferPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); + let pool = RunPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); let notify = Arc::new(Notify::new()); let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); @@ -403,9 +403,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(); @@ -421,7 +421,7 @@ fn virtq_out_of_order_completions() { let pool_size = 0x10000; let mem = Arc::new(LoomMem::new(ring_base, 8, pool_base, pool_size)); - let pool = BufferPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); + let pool = RunPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); let notify = Arc::new(Notify::new()); let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); @@ -471,7 +471,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() { @@ -479,17 +479,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(); @@ -513,7 +513,7 @@ fn virtq_event_suppression_reconfig() { let pool_size = 0x10000; let mem = Arc::new(LoomMem::new(ring_base, 4, pool_base, pool_size)); - let pool = BufferPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); + let pool = RunPoolSync::<256, 4096>::new(pool_base, pool_size).unwrap(); let notify = Arc::new(Notify::new()); let mut prod = VirtqProducer::new(mem.layout(), mem.clone(), notify.clone(), pool); diff --git a/src/hyperlight_common/src/virtq/consumer.rs b/src/hyperlight_common/src/virtq/consumer.rs index d4fe0ff45..48d77b401 100644 --- a/src/hyperlight_common/src/virtq/consumer.rs +++ b/src/hyperlight_common/src/virtq/consumer.rs @@ -15,6 +15,8 @@ limitations under the License. */ use alloc::vec; +use alloc::vec::Vec; +use core::fmt; use bytes::Bytes; use fixedbitset::FixedBitSet; @@ -22,46 +24,162 @@ use smallvec::SmallVec; use super::*; -type WritableElems = SmallVec<[BufferElement; 2]>; - -/// Payload received from the producer, safely copied out of shared memory. +/// Stateful reader over device-readable descriptors received from the producer. /// -/// Created by [`VirtqConsumer::poll`]. Device-readable segments are eagerly -/// copied during poll using [`MemOps::read`] (volatile on the host side), so -/// accessing data requires no unsafe code and no references into shared -/// memory. Segment boundaries are preserved in [`Segments`]. -#[derive(Debug, Clone)] -pub struct RecvChain { - token: Token, - segments: Segments, +/// Reads copy directly from shared memory into caller-provided final storage. +/// The chain must be returned together with its paired [`ReplyChain`] through +/// [`VirtqConsumer::complete`] before the descriptors can be reused. +#[must_use = "dropping without completing leaks the descriptor"] +pub struct RecvChain { + state: ChainState, +} + +impl 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() + } + + /// Whether this chain has no readable payload. + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 } - /// The chain payload as ordered byte segments. - pub fn segments(&self) -> &Segments { - &self.segments + /// Number of bytes consumed by the stateful reader. + #[inline] + pub fn consumed(&self) -> usize { + self.state.position() } - /// Consume the chain, taking ownership of the segments. - pub fn into_segments(self) -> Segments { - self.segments + /// Number of bytes still available to the stateful reader. + #[inline] + pub fn remaining(&self) -> usize { + self.state.remaining() } - /// Return the chain payload as contiguous bytes. + /// Read bytes sequentially across descriptor boundaries. /// - /// Returns empty [`Bytes`] when the chain has no readable buffers. - pub fn to_bytes(&self) -> Bytes { - self.segments.to_bytes() + /// Returns the number of bytes copied, which may be smaller than `buf.len()` + /// at the end of the chain. If a later memory read fails, the cursor remains + /// advanced past any earlier chunks copied by the same call. + pub fn read(&mut self, buf: &mut [u8]) -> Result { + let len = buf.len().min(self.remaining()); + let mut dst = &mut buf[..len]; + let mut read = 0; + + while !dst.is_empty() { + let Some(elem) = self.state.current_elem() else { + break; + }; + + let desc_len = elem.len as usize; + let desc_offset = self.state.desc_offset(); + let len = (desc_len - desc_offset).min(dst.len()); + let (current, rest) = dst.split_at_mut(len); + + let addr = elem + .addr + .checked_add(desc_offset as u64) + .ok_or(VirtqError::MemoryReadError)?; + + self.state + .mem + .read(addr, current) + .map_err(|_| VirtqError::MemoryReadError)?; + + self.state.advance(len); + read += len; + dst = rest; + } + + Ok(read) + } + + /// Read exactly `buf.len()` bytes or return an error. + #[inline] + pub fn read_exact(&mut self, buf: &mut [u8]) -> Result<&mut Self, VirtqError> { + if buf.len() > self.remaining() { + return Err(VirtqError::ReceiveTooShort { + requested: buf.len(), + remaining: self.remaining(), + }); + } + + let read = self.read(buf)?; + debug_assert_eq!(read, buf.len()); + Ok(self) + } + + /// Copy the complete payload into descriptor-preserving owned segments. + /// + /// This does not change the stateful read position. Each call takes a new + /// snapshot of shared memory; callers should validate and use the returned + /// owned value rather than reading the same untrusted payload again. + pub fn to_segments(&self) -> Result { + let mut segments = SmallVec::<[Bytes; 4]>::new(); + + for elem in &self.state.elems { + let mut buf = vec![0u8; elem.len as usize]; + self.state + .mem + .read(elem.addr, &mut buf) + .map_err(|_| VirtqError::MemoryReadError)?; + segments.push(Bytes::from(buf)); + } + + Ok(Segments::from_smallvec(segments)) } - /// 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)) } } @@ -75,13 +193,21 @@ 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), } +/// One polled chain and its matching completion capability. +pub type PolledChain = (RecvChain, ReplyChain); + +/// An exact batch returned by [`VirtqConsumer::poll_exact`]. +pub type PolledChains = Vec>; + impl ReplyChain { /// The token identifying this reply. + #[inline] pub fn token(&self) -> Token { match self { ReplyChain::Writable(wc) => wc.token(), @@ -90,9 +216,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, } } @@ -116,48 +243,50 @@ 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 writable descriptors in this chain. + #[inline] + pub fn desc_count(&self) -> usize { + self.state.elems.len() } /// 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. @@ -165,15 +294,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) } @@ -183,6 +336,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); @@ -198,14 +352,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"] @@ -218,6 +373,7 @@ impl AckChain { Self { token } } + #[inline] pub fn token(&self) -> Token { self.token } @@ -234,33 +390,35 @@ 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, @@ -275,11 +433,17 @@ impl VirtqConsumer { /// * `mem` - Memory ops implementation for reading/writing to shared memory /// * `notifier` - Callback for notifying the driver about replies pub fn new(layout: Layout, mem: M, notifier: N) -> Self { - let inner = RingConsumer::new(layout, mem); + Self::new_split(layout, mem.clone(), mem, notifier) + } + + /// Create a consumer with separate ring and buffer memory accessors. + pub fn new_split(layout: Layout, ring_mem: M, buf_mem: M, notifier: N) -> Self { + let inner = RingConsumer::new(layout, ring_mem); let inflight = FixedBitSet::with_capacity(inner.len()); Self { inner, + mem: buf_mem, notifier, inflight, next_token: 0, @@ -288,29 +452,25 @@ impl VirtqConsumer { /// Poll for a single incoming chain from the driver. /// - /// Returns a [`RecvChain`] (copied data) and a [`ReplyChain`] (writable reply - /// capacity or ack token). Both are independent owned values with no borrow - /// on the consumer. + /// Returns a stateful [`RecvChain`] reader and a [`ReplyChain`] writable + /// reply or ack capability. Both are independent owned values with no + /// borrow on the consumer, but they must be returned together through + /// [`complete`](Self::complete). /// - /// On [`VirtqError::BadChain`], [`VirtqError::PayloadTooLarge`], and - /// [`VirtqError::MemoryReadError`] the descriptor is returned to the driver - /// (completed with zero length) before the error is propagated, so a - /// rejected chain does not leak. + /// On [`VirtqError::BadChain`] and [`VirtqError::PayloadTooLarge`] the + /// descriptor is returned to the driver (completed with zero length) before + /// the error is propagated, so a rejected chain does not leak. /// /// # Arguments /// - /// * `max_recv_len` - Maximum receive payload size to copy. Payloads larger + /// * `max_recv_len` - Maximum readable payload size. Payloads larger /// than this return [`VirtqError::PayloadTooLarge`]. /// /// # Errors /// /// - [`VirtqError::BadChain`] - Descriptor chain format not recognized /// - [`VirtqError::InvalidState`] - Descriptor ID collision (driver bug) - /// - [`VirtqError::MemoryReadError`] - Failed to read chain payload from shared memory - pub fn poll( - &mut self, - max_recv_len: usize, - ) -> Result)>, VirtqError> { + pub fn poll(&mut self, max_recv_len: usize) -> Result>, VirtqError> { let (id, chain) = match self.inner.poll_available() { Ok(x) => x, Err(RingError::WouldBlock) => return Ok(None), @@ -354,20 +514,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); @@ -377,14 +534,103 @@ impl VirtqConsumer { Ok(Some((chain, reply))) } - /// Submit a reply/ack for a received chain back to the ring. + /// Poll exactly `count` chains without consuming a partial batch. /// - /// 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> { + /// Returns `None` and restores the consumer's local state when fewer than + /// `count` chains are available. Each returned chain must be completed + /// through [`complete`](Self::complete). + /// + /// # Arguments + /// + /// * `count` - Exact number of chains to poll. Zero returns an empty batch. + /// * `max_recv_len` - Maximum readable payload size accepted for each chain + /// independently. + pub fn poll_exact( + &mut self, + count: usize, + max_recv_len: usize, + ) -> Result>, VirtqError> { + self.poll_exact_with_spare(count, 0, max_recv_len) + } + + /// Poll exactly `count` chains while leaving `spare` chains available. + /// + /// Returns `None` and restores the consumer's local state when fewer than + /// `count + spare` chains are available. The spare chains are inspected + /// without completing them and remain available for a later poll. + /// + /// # Arguments + /// + /// * `count` - Exact number of chains to poll. Zero returns an empty batch. + /// * `spare` - Number of chains to leave available for later inspection. + /// * `max_recv_len` - Maximum readable payload size accepted for each chain + /// independently. + pub fn poll_exact_with_spare( + &mut self, + count: usize, + spare: usize, + max_recv_len: usize, + ) -> Result>, VirtqError> { + let Some(total) = count.checked_add(spare) else { + return Ok(None); + }; + + // Every chain consumes at least one descriptor. + if total > self.inner.num_free() { + return Ok(None); + } + + // Polling changes only local bookkeeping until a chain is completed. + let cp = self.inner.poll_checkpoint(); + let next_token = self.next_token; + + let mut spare_cp = None; + let mut polled = PolledChains::with_capacity(total); + + while polled.len() < total { + if spare != 0 && polled.len() == count { + spare_cp = Some((self.inner.poll_checkpoint(), self.next_token)); + } + + if let Some(chain) = self.poll(max_recv_len)? { + polled.push(chain); + continue; + } + + self.rollback_polled(cp, next_token, polled)?; + return Ok(None); + } + + if let Some((checkpoint, next_token)) = spare_cp { + self.rollback_polled(checkpoint, next_token, polled.drain(count..))?; + } + + Ok(Some(polled)) + } + + /// Submit both halves of a received chain back to the ring. + /// + /// Consuming the [`RecvChain`] prevents further reads once its descriptors + /// can be reused by the producer. `reply` accepts both [`WritableChain`] + /// (with written byte count) and [`AckChain`] (zero-length) through + /// [`ReplyChain`]. The two halves must have matching tokens. + /// + /// A mismatched pair returns [`VirtqError::InvalidState`] without returning + /// either descriptor. This fails closed: the descriptors remain in flight + /// because completing either could invalidate another still-live + /// [`RecvChain`]. + pub fn complete( + &mut self, + recv: impl Into>, + reply: impl Into>, + ) -> Result<(), VirtqError> { + let recv = recv.into(); let reply = reply.into(); + + if recv.token() != reply.token() { + return Err(VirtqError::InvalidState); + } + let id = reply.token().id; let written = u32::try_from(reply.written()).map_err(|_| VirtqError::ReplyTooLarge)?; @@ -475,63 +721,135 @@ 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.inflight.ones().next().is_some() { + return Err(VirtqError::InvalidState); } - Ok(Segments::from_smallvec(segments)) + self.inner.reset()?; + self.inflight.clear(); + self.next_token = 0; + Ok(()) } - /// Reset ring and inflight state to initial values. - pub fn reset(&mut self) { - self.inner.reset(); - self.inflight.clear(); + fn rollback_polled( + &mut self, + checkpoint: Checkpoint, + next_token: u32, + polled: impl IntoIterator>, + ) -> Result<(), VirtqError> { + // No chain handle may survive when its descriptor becomes pollable again. + let ids = polled + .into_iter() + .map(|(recv, _)| recv.token().id) + .collect::>(); + + self.inner.rollback_polls(checkpoint, &ids)?; + self.next_token = next_token; + + for id in ids { + self.inflight.set(id as usize, false); + } + + Ok(()) } } -fn write_elements( - mem: &M, - elems: &[BufferElement], - offset: usize, - buf: &[u8], -) -> Result { - let capacity: usize = elems.iter().map(|elem| elem.len as usize).sum(); - let mut src = &buf[..buf.len().min(capacity.saturating_sub(offset))]; - let mut written = 0; - let mut skip = offset; +type ChainElems = SmallVec<[BufferElement; 4]>; - for elem in elems { - if src.is_empty() { - break; - } +struct ChainState { + mem: M, + token: Token, + elems: ChainElems, + total: usize, + position: usize, + desc_idx: usize, + desc_off: usize, +} - let elem_len = elem.len as usize; - if skip >= elem_len { - skip -= elem_len; - continue; - } +impl ChainState { + fn new(mem: M, token: Token, elems: ChainElems, total: usize) -> Self { + let mut state = Self { + mem, + token, + elems, + total, + position: 0, + desc_idx: 0, + desc_off: 0, + }; + state.rewind(); + state + } + + #[inline] + fn token(&self) -> Token { + self.token + } + + #[inline] + fn total(&self) -> usize { + self.total + } + + #[inline] + fn position(&self) -> usize { + self.position + } + + #[inline] + fn remaining(&self) -> usize { + self.total - self.position + } + + #[inline(always)] + fn desc_len(&self) -> usize { + self.elems + .get(self.desc_idx) + .map(|elem| elem.len as usize) + .unwrap_or(0) + } + + #[inline(always)] + fn desc_offset(&self) -> usize { + self.desc_off + } - let elem_offset = skip; - skip = 0; - let n = (elem_len - elem_offset).min(src.len()); - let addr = elem.addr + elem_offset as u64; + #[inline(always)] + fn current_elem(&self) -> Option { + self.elems.get(self.desc_idx).copied() + } - mem.write(addr, &src[..n])?; + #[inline(always)] + fn advance(&mut self, len: usize) { + debug_assert!(len <= self.desc_len() - self.desc_off); + self.desc_off += len; + self.position += len; - written += n; - src = &src[n..]; + while self.current_elem().is_some() && self.desc_off == self.desc_len() { + self.desc_idx += 1; + self.desc_off = 0; + } } - Ok(written) + 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 { @@ -549,15 +867,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); @@ -567,12 +929,12 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_data(&mut consumer); - assert!(recv.to_bytes().is_empty()); + assert!(recv.to_bytes().unwrap().is_empty()); assert!(matches!(reply, ReplyChain::Writable(_))); if let ReplyChain::Writable(mut wc) = reply { wc.write_all(b"response").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } } @@ -586,10 +948,10 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_data(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello"); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello"); assert!(matches!(reply, ReplyChain::Ack(_))); - consumer.complete(reply).unwrap(); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -602,7 +964,7 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_data(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello world"); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello world"); if let ReplyChain::Writable(mut wc) = reply { assert_eq!(wc.capacity(), 64); @@ -611,12 +973,85 @@ mod tests { wc.write_all(b"response").unwrap(); assert_eq!(wc.written(), 8); assert_eq!(wc.remaining(), 56); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } else { panic!("expected Writable reply for recv+reply chain"); } } + #[test] + fn test_recv_reads_across_descriptor_boundaries() { + let ring = make_ring(16); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + let mut se = producer.chain().readable(4).readable(4).build().unwrap(); + se.write_all(b"abcdefgh").unwrap(); + producer.submit(se).unwrap(); + + let (mut recv, reply) = poll_data(&mut consumer); + assert_eq!(recv.len(), 8); + assert_eq!(recv.remaining(), 8); + let mut first = [0u8; 2]; + recv.read_exact(&mut first).unwrap(); + assert_eq!(&first, b"ab"); + assert_eq!(recv.consumed(), 2); + assert_eq!(recv.remaining(), 6); + + let mut second = [0u8; 3]; + recv.read_exact(&mut second).unwrap(); + assert_eq!(&second, b"cde"); + + let mut too_long = [0u8; 4]; + assert!(matches!( + recv.read_exact(&mut too_long), + Err(VirtqError::ReceiveTooShort { + requested: 4, + remaining: 3 + }) + )); + + let mut final_buf = [0u8; 4]; + assert_eq!(recv.read(&mut final_buf).unwrap(), 3); + assert_eq!(&final_buf[..3], b"fgh"); + assert_eq!(recv.read(&mut final_buf).unwrap(), 0); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"abcdefgh"); + + consumer.complete(recv, reply).unwrap(); + } + + #[test] + fn test_poll_defers_payload_reads() { + let ring = make_ring(16); + let mem = ring.mem(); + let mut ring_producer = make_producer(&ring); + let payload_addr = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + mem.write(payload_addr, b"data").unwrap(); + + let chain = BufferChainBuilder::new() + .readable(payload_addr, 4) + .build() + .unwrap(); + ring_producer.submit_available(&chain).unwrap(); + + let 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); @@ -625,13 +1060,13 @@ mod tests { let se = producer.chain().writable(8).build().unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_data(&mut consumer); + let (recv, reply) = poll_data(&mut consumer); if let ReplyChain::Writable(mut wc) = reply { let n = wc.write(b"hello world!").unwrap(); assert_eq!(n, 8); assert_eq!(wc.remaining(), 0); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } else { panic!("expected Writable"); } @@ -654,6 +1089,155 @@ mod tests { } } + #[test] + fn test_poll_exact_rolls_back_partial_batch() { + let ring = make_ring(16); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + for _ in 0..2 { + let chain = producer.chain().writable(16).build().unwrap(); + producer.submit(chain).unwrap(); + } + + let cursor = consumer.avail_cursor(); + assert!(consumer.poll_exact(3, 0).unwrap().is_none()); + assert_eq!(consumer.avail_cursor(), cursor); + assert_eq!(consumer.inflight.count_ones(..), 0); + assert_eq!(consumer.inner.num_inflight(), 0); + assert_eq!(consumer.next_token, 0); + assert!(producer.poll().unwrap().is_none()); + + let reserved = consumer.poll_exact(2, 0).unwrap().unwrap(); + for (recv, reply) in reserved { + consumer.complete(recv, reply).unwrap(); + } + assert!(producer.poll().unwrap().is_some()); + assert!(producer.poll().unwrap().is_some()); + } + + #[test] + fn test_poll_exact_with_spare_leaves_spare_available() { + let ring = make_ring(16); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + for _ in 0..3 { + let chain = producer.chain().writable(16).build().unwrap(); + producer.submit(chain).unwrap(); + } + + let polled = consumer.poll_exact_with_spare(2, 1, 0).unwrap().unwrap(); + assert_eq!(consumer.avail_cursor().head(), 2); + assert_eq!(consumer.inflight.count_ones(..), 2); + assert_eq!(consumer.next_token, 2); + + for (recv, reply) in polled { + consumer.complete(recv, reply).unwrap(); + } + + let (recv, reply) = consumer.poll(0).unwrap().unwrap(); + assert_eq!(recv.token().seq, 2); + consumer.complete(recv, reply).unwrap(); + } + + #[test] + fn test_poll_exact_with_spare_rolls_back_requested_chains() { + let ring = make_ring(16); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + for _ in 0..2 { + let chain = producer.chain().writable(16).build().unwrap(); + producer.submit(chain).unwrap(); + } + + let cursor = consumer.avail_cursor(); + assert!(consumer.poll_exact_with_spare(2, 1, 0).unwrap().is_none()); + assert_eq!(consumer.avail_cursor(), cursor); + assert_eq!(consumer.inflight.count_ones(..), 0); + assert_eq!(consumer.inner.num_inflight(), 0); + assert_eq!(consumer.next_token, 0); + + let polled = consumer.poll_exact(2, 0).unwrap().unwrap(); + for (recv, reply) in polled { + consumer.complete(recv, reply).unwrap(); + } + } + + #[test] + fn test_poll_exact_preserves_existing_inflight_chain() { + let ring = make_ring(16); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + for _ in 0..3 { + let chain = producer.chain().writable(16).build().unwrap(); + producer.submit(chain).unwrap(); + } + + let existing = consumer.poll(0).unwrap().unwrap(); + let cursor = consumer.avail_cursor(); + assert!(consumer.poll_exact(3, 0).unwrap().is_none()); + assert_eq!(consumer.avail_cursor(), cursor); + assert_eq!(consumer.inflight.count_ones(..), 1); + assert_eq!(consumer.inner.num_inflight(), 1); + + let reserved = consumer.poll_exact(2, 0).unwrap().unwrap(); + consumer.complete(existing.0, existing.1).unwrap(); + for (recv, reply) in reserved { + consumer.complete(recv, reply).unwrap(); + } + } + + #[test] + fn test_poll_exact_rolls_back_multi_descriptor_chain() { + let ring = make_ring(16); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + let chain = producer.chain().writable(8).writable(8).build().unwrap(); + producer.submit(chain).unwrap(); + + let cursor = consumer.avail_cursor(); + assert!(consumer.poll_exact(2, 0).unwrap().is_none()); + assert_eq!(consumer.avail_cursor(), cursor); + assert_eq!(consumer.inner.num_inflight(), 0); + + let mut reserved = consumer.poll_exact(1, 0).unwrap().unwrap(); + let (recv, reply) = reserved.pop().unwrap(); + let ReplyChain::Writable(writable) = &reply else { + panic!("expected writable chain"); + }; + assert_eq!(writable.desc_count(), 2); + consumer.complete(recv, reply).unwrap(); + } + + #[test] + fn test_poll_exact_rolls_back_across_wrap() { + let ring = make_ring(4); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + for _ in 0..3 { + let chain = producer.chain().writable(16).build().unwrap(); + producer.submit(chain).unwrap(); + } + let reserved = consumer.poll_exact(3, 0).unwrap().unwrap(); + for (recv, reply) in reserved { + consumer.complete(recv, reply).unwrap(); + } + for _ in 0..3 { + assert!(producer.poll().unwrap().is_some()); + } + + let chain = producer.chain().writable(16).build().unwrap(); + producer.submit(chain).unwrap(); + + let cursor = consumer.avail_cursor(); + assert_eq!(cursor.head(), 3); + assert!(consumer.poll_exact(2, 0).unwrap().is_none()); + assert_eq!(consumer.avail_cursor(), cursor); + + let mut reserved = consumer.poll_exact(1, 0).unwrap().unwrap(); + let (recv, reply) = reserved.pop().unwrap(); + consumer.complete(recv, reply).unwrap(); + } + #[test] fn test_poll_too_large_returns_payload_error() { let ring = make_ring(16); @@ -692,8 +1276,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()); } @@ -744,13 +1328,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"); @@ -764,7 +1348,7 @@ mod tests { let se = producer.chain().writable(16).build().unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_data(&mut consumer); + let (recv, reply) = poll_data(&mut consumer); if let ReplyChain::Writable(mut wc) = reply { wc.write_all(b"first").unwrap(); @@ -774,7 +1358,7 @@ mod tests { assert_eq!(wc.remaining(), 16); wc.write_all(b"second").unwrap(); assert_eq!(wc.written(), 6); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } else { panic!("expected Writable"); } @@ -796,7 +1380,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"); @@ -804,7 +1388,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]; @@ -818,6 +1402,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); @@ -828,16 +1449,43 @@ mod tests { let se2 = producer.chain().writable(16).build().unwrap(); producer.submit(se2).unwrap(); - let (_e1, c1) = poll_data(&mut consumer); - let (_e2, c2) = poll_data(&mut consumer); + let (e1, c1) = poll_data(&mut consumer); + let (e2, c2) = poll_data(&mut consumer); // Complete in reverse order - consumer.complete(c2).unwrap(); - consumer.complete(c1).unwrap(); + consumer.complete(e2, c2).unwrap(); + consumer.complete(e1, c1).unwrap(); + } + + #[test] + fn test_mismatched_completion_fails_closed() { + let ring = make_ring(16); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + let mut first = producer.chain().readable(1).build().unwrap(); + first.write_all(b"a").unwrap(); + producer.submit(first).unwrap(); + + let mut second = producer.chain().readable(1).build().unwrap(); + second.write_all(b"b").unwrap(); + producer.submit(second).unwrap(); + + let (recv1, reply1) = poll_data(&mut consumer); + let (recv2, reply2) = poll_data(&mut consumer); + + assert!(matches!( + consumer.complete(recv1, reply2), + Err(VirtqError::InvalidState) + )); + assert_eq!(consumer.inflight.count_ones(..), 2); + assert!(producer.poll().unwrap().is_none()); + assert!(matches!(consumer.reset(), Err(VirtqError::InvalidState))); + + drop((recv2, reply1)); } #[test] - fn test_recv_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); @@ -845,10 +1493,14 @@ mod tests { se.write_all(b"abc").unwrap(); producer.submit(se).unwrap(); - let (recv, reply) = poll_data(&mut consumer); - let data = recv.into_bytes(); + let (mut recv, reply) = poll_data(&mut consumer); + let mut first = [0u8; 1]; + recv.read_exact(&mut first).unwrap(); + let data = recv.to_bytes().unwrap(); + assert_eq!(&first, b"a"); assert_eq!(data.as_ref(), b"abc"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.consumed(), 1); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -860,13 +1512,14 @@ mod tests { let se = producer.chain().writable(16).build().unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_data(&mut consumer); + let (recv, reply) = poll_data(&mut consumer); assert!(consumer.inflight.count_ones(..) > 0); + assert!(matches!(consumer.reset(), Err(VirtqError::InvalidState))); // Complete first so we do not leak - consumer.complete(reply).unwrap(); + consumer.complete(recv, reply).unwrap(); - consumer.reset(); + consumer.reset().unwrap(); assert_eq!(consumer.inflight.count_ones(..), 0); assert_eq!(consumer.inner.num_inflight(), 0); @@ -883,15 +1536,16 @@ 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); + assert_eq!(consumer.next_token, 0); } } diff --git a/src/hyperlight_common/src/virtq/desc.rs b/src/hyperlight_common/src/virtq/desc.rs index 9383f0acc..b67366468 100644 --- a/src/hyperlight_common/src/virtq/desc.rs +++ b/src/hyperlight_common/src/virtq/desc.rs @@ -238,6 +238,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 @@ -248,6 +258,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 46fce8dc9..234279bd4 100644 --- a/src/hyperlight_common/src/virtq/event.rs +++ b/src/hyperlight_common/src/virtq/event.rs @@ -123,6 +123,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 e676e7cc7..50a6be855 100644 --- a/src/hyperlight_common/src/virtq/mod.rs +++ b/src/hyperlight_common/src/virtq/mod.rs @@ -56,27 +56,28 @@ limitations under the License. //! } //! //! // 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)?; //! } //! ``` //! @@ -164,7 +165,6 @@ mod buffer; mod consumer; mod desc; mod event; -pub mod msg; mod pool; mod producer; mod ring; @@ -184,6 +184,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); @@ -200,10 +207,14 @@ pub enum VirtqError { Backpressure, #[error("Allocation exceeds pool capacity")] OutOfMemory, + #[error("Failed to allocate virtqueue bookkeeping")] + BookkeepingAllocation, #[error("Invalid chain received")] BadChain, #[error("Payload data too large: received {recv} bytes, limit {limit} bytes")] 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")] @@ -272,6 +283,11 @@ const fn align_up(val: usize, align: usize) -> usize { val.next_multiple_of(align) } +#[inline] +const fn align_up_checked(val: usize, align: usize) -> Option { + val.checked_next_multiple_of(align) +} + impl Layout { /// Create a Layout from a base address and number of descriptors. /// @@ -348,6 +364,18 @@ impl Layout { dev_evt_offset + event_size } + + /// Calculate the ring size, returning `None` on arithmetic overflow. + pub fn checked_query_size(num_descs: usize) -> Option { + let desc_size = num_descs.checked_mul(Descriptor::SIZE)?; + let event_size = EventSuppression::SIZE; + let align = EventSuppression::ALIGN; + + let drv_evt_offset = align_up_checked(desc_size, align)?; + let dev_evt_offset = align_up_checked(drv_evt_offset.checked_add(event_size)?, align)?; + + dev_evt_offset.checked_add(event_size) + } } /// Statistics about the current virtqueue state. @@ -392,7 +420,7 @@ impl From for Allocation { fn from(value: BufferElement) -> Self { Allocation { addr: value.addr, - len: value.len as usize, + len: value.len, } } } @@ -493,7 +521,6 @@ pub(crate) mod test_utils { base: u64, next: Arc, size: usize, - max_alloc_len: usize, allocations: Arc>>, } @@ -503,33 +530,23 @@ pub(crate) mod test_utils { base, next: Arc::new(AtomicU64::new(base)), size, - max_alloc_len: usize::MAX, - allocations: Arc::new(Mutex::new(BTreeMap::new())), - } - } - - pub(crate) fn new_with_max_alloc_len(base: u64, size: usize, max_alloc_len: usize) -> Self { - Self { - base, - next: Arc::new(AtomicU64::new(base)), - size, - max_alloc_len, allocations: Arc::new(Mutex::new(BTreeMap::new())), } } } impl BufferProvider for TestPool { - fn max_alloc_len(&self) -> usize { - self.max_alloc_len + fn preferred_segment_len(&self) -> usize { + u32::MAX as usize } fn alloc(&self, len: usize) -> Result { if len == 0 { return Err(AllocError::InvalidArg); } + let len = u32::try_from(len).map_err(|_| AllocError::OutOfMemory)?; - let addr = self.next.fetch_add(len as u64, Ordering::Relaxed); + let addr = self.next.fetch_add(u64::from(len), Ordering::Relaxed); let end = addr + len as u64; if end > self.base + self.size as u64 { return Err(AllocError::NoSpace); @@ -537,10 +554,30 @@ pub(crate) mod test_utils { self.allocations .lock() .expect("poisoned mutex") - .insert(addr, len); + .insert(addr, len as usize); + Ok(Allocation { addr, len }) } + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + let mut regions = Regions::new(); + for len in lengths { + match self.alloc(len) { + Ok(alloc) => regions.push(Allocations::from_iter([alloc])), + Err(error) => return Err(error), + } + } + + if regions.is_empty() { + return Err(AllocError::InvalidArg); + } + + Ok(regions) + } + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { self.allocations .lock() @@ -601,7 +638,7 @@ mod tests { fn poll_received( consumer: &mut VirtqConsumer, - ) -> (RecvChain, ReplyChain) { + ) -> (RecvChain, ReplyChain) { consumer.poll(1024).unwrap().unwrap() } @@ -630,8 +667,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 @@ -663,12 +700,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 @@ -749,19 +786,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(); @@ -784,14 +821,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"); } @@ -807,14 +844,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(); @@ -829,8 +866,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); @@ -848,13 +885,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); @@ -872,24 +909,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(); @@ -925,6 +962,7 @@ 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(); @@ -937,13 +975,14 @@ 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"); @@ -958,12 +997,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(); @@ -986,18 +1025,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(); @@ -1021,15 +1060,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(); @@ -1054,8 +1093,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); @@ -1070,12 +1109,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(); @@ -1100,8 +1139,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 090c2eb5b..000000000 --- a/src/hyperlight_common/src/virtq/msg.rs +++ /dev/null @@ -1,120 +0,0 @@ -/* -Copyright 2026 The Hyperlight Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -//! 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 3f0acc963..78e2f4468 100644 --- a/src/hyperlight_common/src/virtq/pool.rs +++ b/src/hyperlight_common/src/virtq/pool.rs @@ -13,1339 +13,170 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. */ -//! Buffer pool implementations for virtqueue buffer management. +//! Buffer pool APIs and implementations for virtqueue payloads. //! -//! This module provides concrete buffer allocators: -//! -//! - [`BufferPool`] - a two-tier run allocator for variable-sized allocations. -//! - [`RecyclePool`] - a single-tier fixed-slot free-list recycler for bounded -//! descriptor segments. -//! -//! All implement [`BufferProvider`] from the [`super::buffer`] module. -//! -//! # BufferPool design -//! -//! `BufferPool` is a variable-sized run allocator. -//! -//! # Two-tier layout -//! -//! [`BufferPool`] divides the underlying region into two slabs with different -//! slot sizes: -//! -//! - The lower tier (default `L = 256`) is intended for *smaller allocations* - -//! control messages, descriptor metadata, and other small structures. Small -//! allocations first try this tier. -//! - The upper tier (default `U = 4096`) uses page sized slots and is intended -//! for larger contiguous buffers. +//! - [`RunPool`] allocates variable-sized contiguous runs from two tiers. +//! - [`SlotPool`] recycles one or two tiers of fixed-size slots. use alloc::rc::Rc; -use core::cell::RefCell; +use alloc::sync::Arc; use core::ops::Deref; -use fixedbitset::FixedBitSet; use smallvec::SmallVec; +use thiserror::Error; -use super::buffer::{AllocError, Allocation, BufferProvider}; - -/// Wrapper asserting `Send` for an inner value that is only ever accessed from -/// a single thread. -/// -/// [`BufferPool`] and [`RecyclePool`] hold their state in an `Rc>`, -/// which is neither `Send` nor `Sync`. Their allocations are exposed as -/// zero-copy reply payloads through -/// [`Bytes::from_owner`](bytes::Bytes::from_owner), whose owner bound is -/// `Send + 'static`; this wrapper exists solely so the pools can satisfy that -/// bound. -/// -/// # Safety -/// -/// The `Send` assertion is only sound while the wrapped value - and every -/// `Bytes` handed out from it - stays on a single thread. Hyperlight guests are -/// single-threaded, so this holds for guest-side use. It is unsound to move a -/// pool (or a reply `Bytes`) to another thread, e.g. by using these pools with a -/// producer/consumer on the multi-threaded host. -#[derive(Debug)] -struct SendWrap(T); - -impl Clone for SendWrap { - fn clone(&self) -> Self { - Self(self.0.clone()) - } -} - -impl Deref for SendWrap { - type Target = T; - fn deref(&self) -> &T { - &self.0 - } -} - -#[derive(Debug, Clone)] -struct Slab { - base_addr: u64, - used_slots: FixedBitSet, - run_starts: FixedBitSet, - last_free_run: Option, -} - -impl Slab { - fn new(base_addr: u64, region_len: usize) -> Result { - let usable = region_len - (region_len % N); - let num_slots = usable / N; - let used_slots = FixedBitSet::with_capacity(num_slots); - let run_starts = FixedBitSet::with_capacity(num_slots); - - if !base_addr.is_multiple_of(N as u64) { - return Err(AllocError::InvalidAlign(base_addr)); - } - if num_slots == 0 { - return Err(AllocError::EmptyRegion); - } - - Ok(Self { - base_addr, - used_slots, - run_starts, - last_free_run: None, - }) - } - - fn addr_of(&self, slot_idx: usize) -> Option { - self.base_addr - .checked_add((slot_idx as u64).checked_mul(N as u64)?) - } - - fn slot_of(&self, addr: u64) -> usize { - let off = (addr - self.base_addr) as usize; - off / N - } - - fn checked_slot_of(&self, addr: u64, len: usize) -> Result { - if addr < self.base_addr { - return Err(AllocError::InvalidFree(addr, len)); - } - - let off = (addr - self.base_addr) as usize; - if !off.is_multiple_of(N) { - return Err(AllocError::InvalidFree(addr, len)); - } - - let slot = off / N; - if slot >= self.used_slots.len() { - return Err(AllocError::InvalidFree(addr, len)); - } - - Ok(slot) - } - - fn live_run_slots_at(&self, start: usize) -> Option { - if start >= self.used_slots.len() - || !self.used_slots.contains(start) - || !self.run_starts.contains(start) - { - return None; - } - - let mut end = start + 1; - while end < self.used_slots.len() - && self.used_slots.contains(end) - && !self.run_starts.contains(end) - { - end += 1; - } - - Some(end - start) - } - - fn maybe_invalidate_last_run(&mut self, alloc: Allocation) { - if let Some(run) = &self.last_free_run { - let new_end = alloc.addr + alloc.len as u64; - let run_end = run.addr + run.len as u64; - - 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; - } -} - -#[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>>, -} +mod run; +mod slot; +pub use run::RunPool; #[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 +pub use run::RunPoolSync; +pub use slot::{SlotLayout, SlotPool}; + +/// Buffer allocation failure. +#[derive(Debug, Error, Copy, Clone)] +pub enum AllocError { + /// A region does not meet its required alignment. + #[error("Invalid region addr {0}")] + InvalidAlign(u64), + /// An address does not identify a live allocation. + #[error("Invalid free addr {0} and size {1}")] + InvalidFree(u64, usize), + /// An argument is zero or otherwise invalid. + #[error("Invalid argument")] + InvalidArg, + /// A region cannot hold any allocation. + #[error("Empty region")] + EmptyRegion, + /// No currently free allocation can satisfy the request. + #[error("No space available")] + NoSpace, + /// The request exceeds the pool's allocation capacity. + #[error("Requested size exceeds pool capacity")] + OutOfMemory, + /// Address or size arithmetic overflowed. + #[error("Overflow")] + Overflow, +} + +/// One allocation returned by a [`BufferProvider`]. +#[derive(Debug, Clone, Copy)] +pub struct Allocation { + /// Starting address of the allocation. + pub addr: u64, + /// Nonzero descriptor-safe capacity in bytes. + pub len: u32, +} + +/// Ordered nonoverlapping allocations that back one logical region. +pub type Allocations = SmallVec<[Allocation; 4]>; + +/// Ordered allocation groups, one for each requested logical region. +pub type Regions = SmallVec<[Allocations; 4]>; + +/// Allocates and reclaims virtqueue payload buffers. +pub trait BufferProvider { + /// Preferred nonzero descriptor-safe size of one bulk allocation segment. + fn preferred_segment_len(&self) -> usize; + + /// Allocate one buffer that can hold at least `len` bytes. + fn alloc(&self, len: usize) -> Result; + + /// Free a previously allocated segment by start address. + fn dealloc(&self, addr: u64) -> Result<(), AllocError>; + + /// Allocate independent logical regions in input order. + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator; +} + +impl BufferProvider for Rc { + fn preferred_segment_len(&self) -> usize { + (**self).preferred_segment_len() } fn alloc(&self, len: usize) -> Result { - self.inner.borrow_mut().alloc(len) - } - - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - Ok(smallvec::smallvec![self.alloc(total_len)?]) + (**self).alloc(len) } fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) - } - - 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) + (**self).dealloc(addr) } - /// Capacity of a live allocation by its start address. - pub fn allocation_len(&self, addr: u64) -> Result { - self.inner.borrow().allocation_len(addr) + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + (**self).alloc_regions(lengths) } } -#[cfg(all(test, loom))] -impl BufferProvider for BufferPoolSync { - fn max_alloc_len(&self) -> usize { - U +impl BufferProvider for Arc { + fn preferred_segment_len(&self) -> usize { + (**self).preferred_segment_len() } fn alloc(&self, len: usize) -> Result { - self.inner.lock().expect("poisoned mutex").alloc(len) - } - - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - Ok(smallvec::smallvec![self.alloc(total_len)?]) + (**self).alloc(len) } fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - self.inner - .lock() - .expect("poisoned mutex") - .dealloc_addr(addr) - } -} - -/// Single-tier fixed-slot free list. -/// -/// Tracks a fixed set of equal-sized buffer slots. Allocation pops a free slot -/// and deallocation returns it, both O(1). A [`FixedBitSet`] records which slots -/// are currently allocated, so double frees and frees of unknown addresses are -/// rejected without scanning the free list. -struct RecycleList { - base_addr: u64, - slot_size: usize, - count: usize, - /// Free slot addresses, popped/pushed LIFO. - free: SmallVec<[u64; 64]>, - /// One bit per slot index; set means the slot is currently handed out. - allocated: FixedBitSet, -} - -// SAFETY: only sound for single-threaded (guest-side) access; see the -// type-level invariant on `SendWrap`. -unsafe impl Send for SendWrap>> {} - -impl RecycleList { - fn new(base_addr: u64, region_len: usize, slot_size: usize) -> Result { - if slot_size == 0 { - return Err(AllocError::InvalidArg); - } - - let count = region_len / slot_size; - if count == 0 { - return Err(AllocError::EmptyRegion); - } - - let mut free = SmallVec::with_capacity(count); - for i in 0..count { - free.push(base_addr + (i * slot_size) as u64); - } - - Ok(Self { - base_addr, - slot_size, - count, - free, - allocated: FixedBitSet::with_capacity(count), - }) - } - - fn end(&self) -> u64 { - self.base_addr + (self.count * self.slot_size) as u64 - } - - fn contains(&self, addr: u64) -> bool { - (self.base_addr..self.end()).contains(&addr) - } - - /// Validate that `addr` names a slot start within the region. - fn slot_of(&self, addr: u64) -> Result { - if !self.contains(addr) { - return Err(AllocError::InvalidFree(addr, 0)); - } - - let off = addr - self.base_addr; - if !off.is_multiple_of(self.slot_size as u64) { - return Err(AllocError::InvalidFree(addr, 0)); - } - - Ok((off / self.slot_size as u64) as usize) - } - - /// Validate that `addr` is a live (currently allocated) slot start. - fn live_slot_of(&self, addr: u64) -> Result { - let slot = self.slot_of(addr)?; - if !self.allocated.contains(slot) { - return Err(AllocError::InvalidFree(addr, 0)); - } - Ok(slot) - } - - fn alloc(&mut self, len: usize) -> Result { - if len == 0 { - return Err(AllocError::InvalidArg); - } - if len > self.slot_size { - return Err(AllocError::OutOfMemory); - } - - let addr = self.free.pop().ok_or(AllocError::NoSpace)?; - // Safety of the index: `addr` came from `free`, which only ever holds - // valid slot starts. - self.allocated - .insert(((addr - self.base_addr) / self.slot_size as u64) as usize); - - Ok(Allocation { - addr, - len: self.slot_size, - }) - } - - fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { - let slot = self.live_slot_of(addr)?; - self.allocated.set(slot, false); - self.free.push(addr); - Ok(()) - } - - fn allocation_len(&self, addr: u64) -> Result { - self.live_slot_of(addr)?; - Ok(self.slot_size) - } - - /// Rebuild state so that exactly the addresses in `allocated` are marked - /// live and every other slot is free. - /// - /// On error the pool is left 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); - } - } + (**self).dealloc(addr) } - fn slot_addr(&self, index: usize) -> Option { - (index < self.count).then(|| self.base_addr + (index * self.slot_size) as u64) - } - - fn num_free(&self) -> usize { - self.free.len() + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + (**self).alloc_regions(lengths) } } -/// A recycling buffer provider with fixed-size slots. +/// Wrapper asserting `Send` for an inner value that is only ever accessed from +/// a single thread. /// -/// Holds a fixed set of equal-sized buffer addresses in a free list. Alloc and -/// dealloc are O(1). It is intended for bounded scatter/gather descriptor -/// segments that are pre-allocated and recycled after use: -/// [`alloc_sg`](BufferProvider::alloc_sg) splits a logical payload into -/// `ceil(total_len / slot_size)` fixed-size segments. -#[derive(Clone)] -pub struct RecyclePool { - inner: SendWrap>>, -} - -impl RecyclePool { - /// Create a recycling pool of `slot_size`-byte slots over a fixed region. - /// - /// The base address is aligned up to `slot_size`; the slot count is based - /// on the remaining usable region after alignment. - pub fn new(base_addr: u64, region_len: usize, slot_size: usize) -> Result { - if slot_size == 0 { - return Err(AllocError::InvalidArg); - } - - let base = usize::try_from(base_addr).map_err(|_| AllocError::Overflow)?; - let region_end = base.checked_add(region_len).ok_or(AllocError::Overflow)?; - let aligned = align_up(base, slot_size)?; - let usable = region_end - .checked_sub(aligned) - .ok_or(AllocError::EmptyRegion)?; - let list = RecycleList::new(aligned as u64, usable, slot_size)?; - - Ok(Self { - inner: SendWrap(Rc::new(RefCell::new(list))), - }) - } - - /// Rebuild pool state so that every address in `allocated` is removed from - /// the free list, matching externally known inflight state. - pub fn restore_allocated(&self, allocated: &[u64]) -> Result<(), AllocError> { - self.inner.borrow_mut().restore_allocated(allocated) - } - - /// Compute the address of slot `index`. - /// - /// Returns `None` if `index >= count`. - pub fn slot_addr(&self, index: usize) -> Option { - self.inner.borrow().slot_addr(index) - } - - /// Number of free slots. - pub fn num_free(&self) -> usize { - self.inner.borrow().num_free() - } - - /// Free a previously allocated slot by address. - pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) - } - - /// Capacity of a live allocation by its start address. - pub fn allocation_len(&self, addr: u64) -> Result { - self.inner.borrow().allocation_len(addr) - } - - /// Base address of the pool region. - pub fn base_addr(&self) -> u64 { - self.inner.borrow().base_addr - } - - /// Slot size in bytes. - pub fn slot_size(&self) -> usize { - self.inner.borrow().slot_size - } +/// [`RunPool`] and [`SlotPool`] hold their state in an `Rc>`, which +/// is neither `Send` nor `Sync`. Their allocations are exposed as zero-copy +/// reply payloads through [`Bytes::from_owner`](bytes::Bytes::from_owner), +/// whose owner bound is `Send + 'static`; this wrapper exists solely so the +/// pools can satisfy that bound. +/// +/// # Safety +/// +/// The `Send` assertion is only sound while the wrapped value and every +/// `Bytes` handed out from it stay on a single thread. Hyperlight guests are +/// single-threaded, so this holds for guest-side use. It is unsound to move a +/// pool or reply `Bytes` to another thread. +#[derive(Debug)] +struct SendWrap(T); - /// Number of slots in the pool. - pub fn count(&self) -> usize { - self.inner.borrow().count +impl Clone for SendWrap { + fn clone(&self) -> Self { + Self(self.0.clone()) } } -impl BufferProvider for RecyclePool { - fn max_alloc_len(&self) -> usize { - self.inner.borrow().slot_size - } - - fn alloc(&self, len: usize) -> Result { - self.inner.borrow_mut().alloc(len) - } - - fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) - } +impl Deref for SendWrap { + type Target = T; - fn reset(&self) { - self.inner.borrow_mut().reset() + fn deref(&self) -> &T { + &self.0 } } -#[cfg(test)] -mod tests { - use super::*; - - fn make_pool(size: usize) -> BufferPool { - let base = align_up(0x10000, L.max(U)).unwrap() as u64; - BufferPool::::new(base, size).unwrap() - } - - fn make_recycle_pool(slot_count: usize, slot_size: usize) -> RecyclePool { - let base = 0x80000u64; - RecyclePool::new(base, slot_count * slot_size, slot_size).unwrap() - } - - #[test] - fn test_pool_new_success() { - let pool = BufferPool::<256, 4096>::new(0x10000, 1024 * 1024).unwrap(); - assert!(pool.inner.borrow().lower.capacity() > 0); - assert!(pool.inner.borrow().upper.capacity() > 0); - } - - #[test] - fn test_pool_alloc_small_to_lower() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(128).unwrap(); - - // Should come from lower slab - assert!(pool.inner.borrow().lower.contains(alloc.addr)); - assert_eq!(alloc.len, 256); - } - - #[test] - fn test_pool_alloc_large_to_upper() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(1500).unwrap(); - - // Should come from upper slab - assert!(pool.inner.borrow().upper.contains(alloc.addr)); - assert_eq!(alloc.len, 4096); - } - - #[test] - fn test_pool_alloc_fallback_to_upper() { - let pool = make_pool::<256, 4096>(1024 * 1024); - - // Fill lower slab completely - let mut allocations = Vec::new(); - while pool.inner.borrow().lower.free_bytes() > 0 { - allocations.push(pool.inner.borrow_mut().lower.alloc(256).unwrap()); - } - - // Small allocation should fallback to upper slab - let alloc = pool.alloc(128).unwrap(); - assert!(pool.inner.borrow().upper.contains(alloc.addr)); - } - - #[test] - fn test_pool_free_from_lower() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(128).unwrap(); - - let free_before = pool.inner.borrow().lower.free_bytes(); - pool.dealloc(alloc.addr).unwrap(); - assert_eq!( - pool.inner.borrow().lower.free_bytes(), - free_before + alloc.len - ); - } - - #[test] - fn test_pool_free_from_upper() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(1500).unwrap(); - - let free_before = pool.inner.borrow().upper.free_bytes(); - pool.dealloc(alloc.addr).unwrap(); - assert_eq!( - pool.inner.borrow().upper.free_bytes(), - free_before + alloc.len - ); - } - - #[test] - fn test_pool_stress_many_allocations() { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); - let mut allocations = Vec::new(); - - // Allocate many buffers - for i in 0..100 { - let size = if i % 2 == 0 { 128 } else { 1500 }; - allocations.push(pool.alloc(size).unwrap()); - } - - // Free half of them - for i in (0..100).step_by(2) { - pool.dealloc(allocations[i].addr).unwrap(); - } - - // Should be able to allocate again - for i in 0..50 { - let size = if i % 2 == 0 { 128 } else { 1500 }; - let _alloc = pool.alloc(size).unwrap(); - } - } - - #[test] - fn test_pool_mixed_workload() { - let pool = make_pool::<256, 4096>(2 * 1024 * 1024); - - // Simulate virtio-net workload - let desc_buf = pool.alloc(64).unwrap(); // Control message - let rx_buf1 = pool.alloc(1500).unwrap(); // MTU packet - let rx_buf2 = pool.alloc(1500).unwrap(); // MTU packet - let tx_buf = pool.alloc(4096).unwrap(); // Large buffer - - // Free and reallocate - pool.dealloc(rx_buf1.addr).unwrap(); - let rx_buf3 = pool.alloc(1500).unwrap(); - - // Should reuse freed buffer (LIFO) - assert_eq!(rx_buf3.addr, rx_buf1.addr); - - pool.dealloc(desc_buf.addr).unwrap(); - pool.dealloc(rx_buf2.addr).unwrap(); - pool.dealloc(rx_buf3.addr).unwrap(); - pool.dealloc(tx_buf.addr).unwrap(); - } - - #[test] - fn test_pool_zero_allocation_error() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let result = pool.alloc(0); - assert!(matches!(result, Err(AllocError::InvalidArg))); - } - - #[test] - fn test_pool_too_large_allocation() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let result = pool.alloc(2 * 1024 * 1024); // Larger than pool - assert!(matches!(result, Err(AllocError::OutOfMemory))); - } - - #[test] - fn test_align_up_helper() { - assert_eq!(align_up(0, 256).unwrap(), 0); - assert_eq!(align_up(1, 256).unwrap(), 256); - assert_eq!(align_up(256, 256).unwrap(), 256); - assert_eq!(align_up(257, 256).unwrap(), 512); - assert_eq!(align_up(511, 256).unwrap(), 512); - assert_eq!(align_up(512, 256).unwrap(), 512); - assert!(matches!(align_up(1, 0), Err(AllocError::InvalidArg))); - assert!(matches!( - align_up(usize::MAX, 256), - Err(AllocError::Overflow) - )); - } - - #[test] - fn test_recycle_pool_alignment_subtracts_padding() { - let pool = RecyclePool::new(0x80001, 8192, 4096).unwrap(); - - assert_eq!(pool.base_addr(), 0x81000); - assert_eq!(pool.count(), 1); - } - - // Edge case: allocation exactly at boundary - #[test] - fn test_pool_boundary_allocation() { - let pool = make_pool::<256, 4096>(1024 * 1024); - - // Allocate exactly at boundary - let alloc = pool.alloc(256).unwrap(); - assert!(pool.inner.borrow().lower.contains(alloc.addr)); - - // Allocate just over boundary - let alloc2 = pool.alloc(257).unwrap(); - assert!(pool.inner.borrow().upper.contains(alloc2.addr)); - } - - #[test] - fn test_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); +#[inline] +fn align_up(val: usize, align: usize) -> Result { + if align == 0 { + return Err(AllocError::InvalidArg); } - #[test] - fn test_recycle_pool_dealloc_order_independent_of_alloc_order() { - let pool = make_recycle_pool(6, 256); - - // Allocate all - let allocs: Vec = (0..6).map(|_| pool.alloc(256).unwrap()).collect(); - - // Dealloc in scattered order: 4, 1, 5, 0, 3, 2 - let order = [4, 1, 5, 0, 3, 2]; - for &i in &order { - pool.dealloc(allocs[i].addr).unwrap(); - } - assert_eq!(pool.num_free(), 6); - - // Re-allocate all and verify we get back the full set - let mut realloc_addrs: Vec = (0..6).map(|_| pool.alloc(256).unwrap().addr).collect(); - realloc_addrs.sort(); - - let mut orig_addrs: Vec = allocs.iter().map(|a| a.addr).collect(); - orig_addrs.sort(); - - assert_eq!(realloc_addrs, orig_addrs); - } + val.checked_next_multiple_of(align) + .ok_or(AllocError::Overflow) } #[cfg(test)] -mod fuzz { - use quickcheck::{Arbitrary, Gen, QuickCheck}; - - 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(); +mod tests; - for op in &s.ops { - match op { - Op::Alloc(size) => match pool.alloc(*size) { - Ok(alloc) => { - assert!(alloc.len >= *size); - allocations.push(alloc); - } - Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} - Err(_) => { - return false; - } - }, - Op::AllocSg(size) => match pool.alloc_sg(*size) { - Ok(sgs) => { - let total: usize = sgs.iter().map(|sg| sg.len).sum(); - assert!(total >= *size); - allocations.extend(sgs); - } - Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} - Err(_) => { - return false; - } - }, - Op::Dealloc(idx) => { - if allocations.is_empty() { - continue; - } - - let idx = idx % allocations.len(); - let alloc = allocations.swap_remove(idx); - - match pool.dealloc(alloc.addr) { - Ok(_) => {} - Err(_) => return false, - } - } - } - - if check_pool_invariants(&pool, &allocations).is_err() { - return false; - } - } - - // Cleanup - for alloc in &allocations { - if pool.dealloc(alloc.addr).is_err() { - return false; - } - } - - check_pool_invariants(&pool, &allocations).is_ok() - } - - fn check_slab_invariants(slab: &Slab) -> Result<(), &'static str> { - let used = slab.used_slots.count_ones(..); - let free = slab.used_slots.count_zeroes(..); - if used + free != slab.used_slots.len() { - return Err("used + free != total slots"); - } - - let expected_free = free * N; - if slab.free_bytes() != expected_free { - return Err("free_bytes doesn't match bitmap"); - } - - if let Some(alloc) = slab.last_free_run { - if alloc.len == 0 || alloc.len % N != 0 { - return Err("last_free_run has invalid length"); - } - if !slab.contains(alloc.addr) { - return Err("last_free_run addr outside range"); - } - } - - Ok(()) - } - - fn check_pool_invariants( - pool: &BufferPool, - allocations: &[Allocation], - ) -> Result<(), &'static str> { - check_slab_invariants(&pool.inner.borrow().lower)?; - check_slab_invariants(&pool.inner.borrow().upper)?; - - if pool.inner.borrow().lower.range().end > pool.inner.borrow().upper.range().start { - return Err("lower and upper ranges overlap"); - } - - let mut seen = std::collections::HashSet::new(); - - for alloc in allocations { - if !pool.inner.borrow().lower.contains(alloc.addr) - && !pool.inner.borrow().upper.contains(alloc.addr) - { - return Err("allocation address outside pool ranges"); - } - - if alloc.len % L != 0 && alloc.len % U != 0 { - return Err("allocation length not aligned to any tier"); - } - - if !seen.insert(alloc.addr) { - return Err("duplicate allocation address in tracking"); - } - } - - Ok(()) - } - - #[test] - fn prop_allocator_invariants() { - #[cfg(miri)] - let tests = 10; - #[cfg(not(miri))] - let tests = 1000; - - QuickCheck::new() - .tests(tests) - .quickcheck(run_scenario as fn(Scenario) -> bool); - } -} +#[cfg(test)] +mod fuzz; diff --git a/src/hyperlight_common/src/virtq/pool/fuzz.rs b/src/hyperlight_common/src/virtq/pool/fuzz.rs new file mode 100644 index 000000000..a2e1c6dee --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/fuzz.rs @@ -0,0 +1,387 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use std::collections::{BTreeMap, HashSet}; + +use quickcheck::{Arbitrary, Gen, QuickCheck}; + +use super::run::Tier as RunTier; +use super::*; + +const MAX_OPS: usize = 10; +const MAX_ALLOC_SIZE: usize = 8192; +const MAX_TIER_SLOTS: usize = 16; +const LOWER_BASE: u64 = 0x80000; +const UPPER_BASE: u64 = 0x90000; +const LOWER_SLOT_SIZE: usize = 256; +const UPPER_SLOT_SIZE: usize = 4096; + +#[derive(Clone, Debug)] +enum Op { + Alloc(usize), + AllocRegions(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::AllocRegions(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), + 2 => Op::Dealloc(usize::arbitrary(g)), + _ => unreachable!(), + } + } +} + +#[derive(Clone, Debug)] +struct RunScenario { + pool_size: usize, + ops: Vec, +} + +impl Arbitrary for RunScenario { + fn arbitrary(g: &mut Gen) -> Self { + let pool_size = (usize::arbitrary(g) % (4 * 1024 * 1024)) + (1024 * 1024); + let num_ops = usize::arbitrary(g) % MAX_OPS + 1; + let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); + + RunScenario { pool_size, ops } + } +} + +fn run_provider_ops(pool: &P, ops: &[Op], check: F) -> bool +where + P: BufferProvider, + F: Fn(&P, &[Allocation]) -> Result<(), &'static str>, +{ + let mut allocations: Vec = Vec::new(); + + for op in ops { + match op { + Op::Alloc(size) => match pool.alloc(*size) { + Ok(alloc) => { + if (alloc.len as usize) < *size + || allocations + .iter() + .any(|existing| existing.addr == alloc.addr) + { + return false; + } + allocations.push(alloc); + } + Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} + Err(_) => return false, + }, + Op::AllocRegions(size) => match pool.alloc_regions([*size]) { + Ok(regions) => { + let mut total = 0usize; + for allocation in regions.into_iter().flatten() { + let Some(next_total) = total.checked_add(allocation.len as usize) else { + return false; + }; + if allocations + .iter() + .any(|existing| existing.addr == allocation.addr) + { + return false; + } + total = next_total; + allocations.push(allocation); + } + if total < *size { + return false; + } + } + Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} + Err(_) => return false, + }, + Op::Dealloc(index) => { + if !allocations.is_empty() { + let index = index % allocations.len(); + if pool.dealloc(allocations[index].addr).is_err() { + return false; + } + allocations.swap_remove(index); + } + } + } + + if check(pool, &allocations).is_err() { + return false; + } + } + + while let Some(alloc) = allocations.pop() { + if pool.dealloc(alloc.addr).is_err() { + return false; + } + } + + check(pool, &allocations).is_ok() +} + +fn run_pool_scenario(scenario: RunScenario) -> bool { + let base = align_up(0x10000, 4096).unwrap() as u64; + let pool = match RunPool::<256, 4096>::new(base, scenario.pool_size) { + Ok(pool) => pool, + Err(_) => return true, + }; + + run_provider_ops(&pool, &scenario.ops, check_run_pool_invariants) +} + +fn check_run_tier_invariants( + tier: &RunTier, + allocations: &[Allocation], +) -> Result<(), &'static str> { + let mut expected_used = HashSet::new(); + let mut expected_starts = HashSet::new(); + + for alloc in allocations.iter().filter(|alloc| tier.contains(alloc.addr)) { + let offset = usize::try_from(alloc.addr - tier.base_addr) + .map_err(|_| "allocation offset overflows usize")?; + let len = alloc.len as usize; + if len == 0 || !offset.is_multiple_of(N) || !len.is_multiple_of(N) { + return Err("allocation is not tier-aligned"); + } + + let start = offset / N; + let slots = len / N; + let end = start + .checked_add(slots) + .ok_or("allocation slot range overflow")?; + if end > tier.used_slots.len() || !expected_starts.insert(start) { + return Err("allocation run is invalid"); + } + for slot in start..end { + if !expected_used.insert(slot) { + return Err("allocation runs overlap"); + } + } + } + + for slot in 0..tier.used_slots.len() { + if tier.used_slots.contains(slot) != expected_used.contains(&slot) { + return Err("used bitmap does not match live allocations"); + } + if tier.run_starts.contains(slot) != expected_starts.contains(&slot) { + return Err("run-start bitmap does not match live allocations"); + } + } + + if tier.free_bytes() != (tier.used_slots.len() - expected_used.len()) * N { + return Err("free_bytes does not match live allocations"); + } + + if let Some(free_run) = tier.last_free_run { + if !tier.contains(free_run.addr) { + return Err("cached free-run address outside tier"); + } + let offset = usize::try_from(free_run.addr - tier.base_addr) + .map_err(|_| "cached free-run offset overflows usize")?; + let len = free_run.len as usize; + if len == 0 || !offset.is_multiple_of(N) || !len.is_multiple_of(N) { + return Err("cached free run is not tier-aligned"); + } + + let start = offset / N; + let end = start + .checked_add(len / N) + .ok_or("cached free-run range overflow")?; + if end > tier.used_slots.len() || (start..end).any(|slot| tier.used_slots.contains(slot)) { + return Err("cached free run overlaps live allocations"); + } + } + + Ok(()) +} + +fn check_run_pool_invariants( + pool: &RunPool, + allocations: &[Allocation], +) -> Result<(), &'static str> { + let inner = pool.inner.borrow(); + if inner.lower.range().end > inner.upper.range().start { + return Err("lower and upper ranges overlap"); + } + + let mut seen = HashSet::new(); + for alloc in allocations { + let in_lower = inner.lower.contains(alloc.addr); + let in_upper = inner.upper.contains(alloc.addr); + if in_lower == in_upper { + return Err("allocation does not belong to exactly one tier"); + } + if !seen.insert(alloc.addr) { + return Err("duplicate allocation address in tracking"); + } + } + + check_run_tier_invariants(&inner.lower, allocations)?; + check_run_tier_invariants(&inner.upper, allocations) +} + +#[derive(Clone, Debug)] +struct SlotScenario { + tiered: bool, + lower_count: usize, + upper_count: usize, + ops: Vec, +} + +impl Arbitrary for SlotScenario { + fn arbitrary(g: &mut Gen) -> Self { + let tiered = bool::arbitrary(g); + let lower_count = usize::arbitrary(g) % MAX_TIER_SLOTS + 1; + let upper_count = usize::arbitrary(g) % MAX_TIER_SLOTS + 1; + let num_ops = usize::arbitrary(g) % MAX_OPS + 1; + let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); + + Self { + tiered, + lower_count, + upper_count, + ops, + } + } +} + +fn make_slot_pool(scenario: &SlotScenario) -> SlotPool { + if scenario.tiered { + let lower = SlotLayout::new(LOWER_BASE, LOWER_SLOT_SIZE, scenario.lower_count); + let upper = SlotLayout::new(UPPER_BASE, UPPER_SLOT_SIZE, scenario.upper_count); + SlotPool::new_tiered(lower, upper).unwrap() + } else { + let layout = SlotLayout::new(UPPER_BASE, UPPER_SLOT_SIZE, scenario.upper_count); + SlotPool::new(layout).unwrap() + } +} + +fn run_slot_pool_scenario(scenario: SlotScenario) -> bool { + let pool = make_slot_pool(&scenario); + run_provider_ops(&pool, &scenario.ops, check_slot_pool_invariants) +} + +fn layout_contains(layout: SlotLayout, addr: u64) -> bool { + let Ok(end) = layout.end_addr() else { + return false; + }; + (layout.base_addr..end).contains(&addr) +} + +fn slot_capacity(pool: &SlotPool, addr: u64) -> Option { + let (lower, upper) = pool.layouts(); + if let Some(lower) = lower + && layout_contains(lower, addr) + { + return Some(lower.slot_size); + } + layout_contains(upper, addr).then_some(upper.slot_size) +} + +fn check_slot_pool_invariants( + pool: &SlotPool, + allocations: &[Allocation], +) -> Result<(), &'static str> { + let mut expected_live = BTreeMap::new(); + for alloc in allocations { + if expected_live.insert(alloc.addr, alloc.len).is_some() { + return Err("duplicate allocation address in tracking"); + } + } + + let live = pool.live_addrs(); + let expected_addrs: Vec = expected_live.keys().copied().collect(); + if live != expected_addrs || live.windows(2).any(|pair| pair[0] >= pair[1]) { + return Err("live addresses are not unique and deterministic"); + } + if pool.num_free() + live.len() != pool.count() { + return Err("free + live != total slots"); + } + + let (lower, upper) = pool.layouts(); + let expected_base = lower.map_or(upper.base_addr, |layout| layout.base_addr); + if pool.base_addr() != expected_base || pool.slot_size() != upper.slot_size { + return Err("reported pool layout is inconsistent"); + } + + let mut expected_count = upper.slot_count; + if let Some(lower) = lower { + if lower.slot_size >= upper.slot_size + || lower.end_addr().map_err(|_| "lower layout overflow")? > upper.base_addr + { + return Err("tier layout is invalid"); + } + expected_count += lower.slot_count; + } + if pool.count() != expected_count || pool.slot_addr(pool.count()).is_some() { + return Err("reported slot count is inconsistent"); + } + + let mut seen = HashSet::new(); + for index in 0..pool.count() { + let Some(addr) = pool.slot_addr(index) else { + return Err("missing slot address"); + }; + if !seen.insert(addr) { + return Err("duplicate slot address"); + } + let Some(capacity) = slot_capacity(pool, addr) else { + return Err("slot address outside layout"); + }; + + match expected_live.get(&addr) { + Some(expected_capacity) => { + if *expected_capacity 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_run_pool_invariants() { + #[cfg(miri)] + let tests = 10; + #[cfg(not(miri))] + let tests = 1000; + + QuickCheck::new() + .tests(tests) + .quickcheck(run_pool_scenario as fn(RunScenario) -> bool); +} + +#[test] +fn prop_slot_pool_invariants() { + #[cfg(miri)] + let tests = 10; + #[cfg(not(miri))] + let tests = 1000; + + QuickCheck::new() + .tests(tests) + .quickcheck(run_slot_pool_scenario as fn(SlotScenario) -> bool); +} diff --git a/src/hyperlight_common/src/virtq/pool/run.rs b/src/hyperlight_common/src/virtq/pool/run.rs new file mode 100644 index 000000000..4cebb1455 --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/run.rs @@ -0,0 +1,429 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ +//! Variable-sized contiguous-run pool. +//! +//! [`RunPool`] partitions one backing region into lower and upper tiers with +//! compile-time slot sizes. The lower tier is carved from the first eighth of +//! the aligned usable region. Eligible requests try that tier first and fall +//! back to the upper tier only when no contiguous lower-tier run is available. +//! +//! Allocations are rounded to a tier's slot size and occupy contiguous runs, so +//! scatter/gather requests produce one allocation. Occupied-slot and run-start +//! bitmaps support reclaiming a complete run by its start address, while a +//! cached free run accelerates immediate reuse. This preserves contiguous +//! buffers but remains subject to fragmentation. + +use alloc::rc::Rc; +use core::cell::RefCell; + +use fixedbitset::FixedBitSet; + +use super::{AllocError, Allocation, Allocations, BufferProvider, Regions, SendWrap, align_up}; + +#[derive(Debug, Clone)] +pub(super) struct Tier { + pub(super) base_addr: u64, + pub(super) used_slots: FixedBitSet, + pub(super) run_starts: FixedBitSet, + pub(super) last_free_run: Option, +} + +impl Tier { + fn new(base_addr: u64, region_len: usize) -> Result { + let usable = region_len - (region_len % N); + let num_slots = usable / N; + let used_slots = FixedBitSet::with_capacity(num_slots); + let run_starts = FixedBitSet::with_capacity(num_slots); + + if !base_addr.is_multiple_of(N as u64) { + return Err(AllocError::InvalidAlign(base_addr)); + } + if num_slots == 0 { + return Err(AllocError::EmptyRegion); + } + + Ok(Self { + base_addr, + used_slots, + run_starts, + last_free_run: None, + }) + } + + fn addr_of(&self, slot_idx: usize) -> Option { + self.base_addr + .checked_add((slot_idx as u64).checked_mul(N as u64)?) + } + + fn slot_of(&self, addr: u64) -> usize { + let off = (addr - self.base_addr) as usize; + off / N + } + + fn checked_slot_of(&self, addr: u64, len: usize) -> Result { + if addr < self.base_addr { + return Err(AllocError::InvalidFree(addr, len)); + } + + let off = (addr - self.base_addr) as usize; + if !off.is_multiple_of(N) { + return Err(AllocError::InvalidFree(addr, len)); + } + + let slot = off / N; + if slot >= self.used_slots.len() { + return Err(AllocError::InvalidFree(addr, len)); + } + + Ok(slot) + } + + fn live_run_slots_at(&self, start: usize) -> Option { + if start >= self.used_slots.len() + || !self.used_slots.contains(start) + || !self.run_starts.contains(start) + { + return None; + } + + let mut end = start + 1; + while end < self.used_slots.len() + && self.used_slots.contains(end) + && !self.run_starts.contains(end) + { + end += 1; + } + + Some(end - start) + } + + fn maybe_invalidate_last_run(&mut self, alloc: Allocation) { + if let Some(run) = &self.last_free_run { + let new_end = alloc.addr + alloc.len as u64; + let run_end = run.addr + run.len as u64; + + if alloc.addr < run_end && run.addr < new_end { + self.last_free_run = None; + } + } + } + + fn find_slots(&mut self, slots_num: usize) -> Option { + debug_assert!(slots_num > 0); + + if let Some(alloc) = self.last_free_run + && alloc.len as usize >= slots_num * N + { + let pos = self.slot_of(alloc.addr); + let _ = self.last_free_run.take(); + return Some(pos); + } + + let total = self.used_slots.len(); + self.used_slots.zeroes().find(|&next_free| { + let end = next_free + slots_num; + end <= total && self.used_slots.count_zeroes(next_free..end) == slots_num + }) + } + + pub(super) fn alloc(&mut self, len: usize) -> Result { + if len == 0 { + return Err(AllocError::InvalidArg); + } + + let total = self.used_slots.len(); + let need_slots = len.div_ceil(N); + if need_slots > total { + return Err(AllocError::OutOfMemory); + } + + let alloc_len = need_slots + .checked_mul(N) + .and_then(|len| u32::try_from(len).ok()) + .ok_or(AllocError::OutOfMemory)?; + + let idx = self.find_slots(need_slots).ok_or(AllocError::NoSpace)?; + self.used_slots.insert_range(idx..idx + need_slots); + self.run_starts.insert(idx); + let addr = self.addr_of(idx).ok_or(AllocError::Overflow)?; + + let alloc = Allocation { + addr, + len: alloc_len, + }; + + 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 = u32::try_from(run_slots * N).map_err(|_| AllocError::Overflow)?; + self.used_slots.remove_range(start..start + run_slots); + self.run_starts.set(start, false); + self.last_free_run = Some(Allocation { addr, len }); + Ok(()) + } + + fn allocation_len(&self, addr: u64) -> Result { + let start = self.checked_slot_of(addr, 0)?; + let run_slots = self + .live_run_slots_at(start) + .ok_or(AllocError::InvalidFree(addr, 0))?; + Ok(run_slots * N) + } + + pub(super) fn capacity(&self) -> usize { + self.used_slots.len() * N + } + + pub(super) fn range(&self) -> core::ops::Range { + self.base_addr..self.base_addr + self.capacity() as u64 + } + + pub(super) fn contains(&self, addr: u64) -> bool { + self.range().contains(&addr) + } +} + +#[cfg(test)] +impl Tier { + pub(super) fn free_bytes(&self) -> usize { + (self.used_slots.len() - self.used_slots.count_ones(..)) * N + } +} + +#[derive(Debug)] +pub(super) struct Inner { + pub(super) lower: Tier, + pub(super) upper: Tier, +} + +// SAFETY: only sound for single-threaded (guest-side) access; see the +// type-level invariant on `SendWrap`. +unsafe impl Send for SendWrap>>> {} + +/// Two-tier pool for variable-sized contiguous runs. +#[derive(Debug, Clone)] +pub struct RunPool { + pub(super) inner: SendWrap>>>, +} + +impl RunPool { + /// Create a new run pool over a fixed region. + pub fn new(base_addr: u64, region_len: usize) -> Result { + let inner = Inner::::new(base_addr, region_len)?; + Ok(Self { + inner: SendWrap(Rc::new(RefCell::new(inner))), + }) + } +} + +impl RunPool { + /// Upper tier slot size in bytes. + pub const fn upper_slot_size() -> usize { + 4096 + } + + /// Lower tier slot size in bytes. + pub const fn lower_slot_size() -> usize { + 256 + } +} + +#[cfg(all(test, loom))] +#[derive(Debug, Clone)] +pub struct RunPoolSync { + inner: std::sync::Arc>>, +} + +#[cfg(all(test, loom))] +impl RunPoolSync { + /// Create a new synchronized run pool over a fixed region. + pub fn new(base_addr: u64, region_len: usize) -> Result { + let inner = Inner::::new(base_addr, region_len)?; + Ok(Self { + inner: std::sync::Arc::new(std::sync::Mutex::new(inner)), + }) + } +} + +impl Inner { + /// Create new run-pool state over a fixed region. + pub fn new(base_addr: u64, region_len: usize) -> Result { + const LOWER_FRACTION: usize = 8; + + let base = usize::try_from(base_addr).map_err(|_| AllocError::Overflow)?; + let region_end = base.checked_add(region_len).ok_or(AllocError::Overflow)?; + + let lower_base = align_up(base, L)?; + let usable = region_end + .checked_sub(lower_base) + .ok_or(AllocError::EmptyRegion)?; + + let lower_region = usable / LOWER_FRACTION; + let lower = Tier::::new(lower_base as u64, lower_region)?; + + let upper_base = lower_base + .checked_add(lower.capacity()) + .ok_or(AllocError::Overflow)?; + + let upper_base = align_up(upper_base, U)?; + let upper_region = region_end + .checked_sub(upper_base) + .ok_or(AllocError::EmptyRegion)?; + + let upper = Tier::::new(upper_base as u64, upper_region)?; + Ok(Self { lower, upper }) + } + + /// Allocate at least `len` bytes. + pub fn alloc(&mut self, len: usize) -> Result { + if len <= L { + match self.lower.alloc(len) { + Ok(alloc) => return Ok(alloc), + Err(AllocError::NoSpace) => {} + Err(e) => return Err(e), + } + } + + // Fall back to the upper tier. + self.upper.alloc(len) + } + + /// Free a previously allocated block by its start address. + pub fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { + if self.lower.contains(addr) { + self.lower.dealloc_addr(addr) + } else { + self.upper.dealloc_addr(addr) + } + } + + /// Capacity of a live allocation by its start address. + pub fn allocation_len(&self, addr: u64) -> Result { + if self.lower.contains(addr) { + self.lower.allocation_len(addr) + } else { + self.upper.allocation_len(addr) + } + } + + /// Allocate independent logical regions in order. + fn alloc_regions(&mut self, lengths: I) -> Result + where + I: IntoIterator, + { + let mut regions = Regions::new(); + + for len in lengths { + match self.alloc(len) { + Ok(allocation) => { + let mut allocations = Allocations::new(); + allocations.push(allocation); + regions.push(allocations); + } + Err(error) => { + self.rollback_regions(®ions); + return Err(error); + } + } + } + + if regions.is_empty() { + return Err(AllocError::InvalidArg); + } + + Ok(regions) + } + + fn rollback_regions(&mut self, regions: &Regions) { + for allocations in regions { + for allocation in allocations { + let result = self.dealloc_addr(allocation.addr); + debug_assert!(result.is_ok(), "dealloc failed: {result:?}"); + } + } + } +} + +impl BufferProvider for RunPool { + fn preferred_segment_len(&self) -> usize { + U.min(u32::MAX as usize) + } + + fn alloc(&self, len: usize) -> Result { + self.inner.borrow_mut().alloc(len) + } + + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + self.inner.borrow_mut().alloc_regions(lengths) + } + + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } +} + +impl RunPool { + /// Free a previously allocated block by its start address. + pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } + + /// Capacity of a live allocation by its start address. + pub fn allocation_len(&self, addr: u64) -> Result { + self.inner.borrow().allocation_len(addr) + } +} + +#[cfg(all(test, loom))] +impl BufferProvider for RunPoolSync { + fn preferred_segment_len(&self) -> usize { + U.min(u32::MAX as usize) + } + + fn alloc(&self, len: usize) -> Result { + self.inner.lock().expect("poisoned mutex").alloc(len) + } + + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + self.inner + .lock() + .expect("poisoned mutex") + .alloc_regions(lengths) + } + + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { + self.inner + .lock() + .expect("poisoned mutex") + .dealloc_addr(addr) + } +} diff --git a/src/hyperlight_common/src/virtq/pool/slot.rs b/src/hyperlight_common/src/virtq/pool/slot.rs new file mode 100644 index 000000000..96d472a04 --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/slot.rs @@ -0,0 +1,564 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ +//! Fixed-slot pool with optional lower and required upper tiers. +//! +//! [`SlotPool`] manages one or two non-overlapping [`SlotLayout`]s. Each tier +//! contains independent, equal-sized slots tracked by a free list and an +//! allocation bitmap. Eligible requests try the lower tier first and fall back +//! to the upper tier only when the lower tier has no free slot. +//! +//! Slots need not be contiguous, so scatter/gather allocation splits a logical +//! buffer at the upper-tier slot size and may place an eligible final segment +//! in the lower tier. [`SlotPool::live_addrs`] reports ownership in deterministic +//! lower-then-upper index order without mutating pool state. + +use alloc::rc::Rc; +use alloc::vec::Vec; +use core::cell::RefCell; + +use fixedbitset::FixedBitSet; +use smallvec::SmallVec; + +use super::{AllocError, Allocation, Allocations, BufferProvider, Regions, SendWrap}; + +/// Exact memory layout for one [`SlotPool`] tier. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SlotLayout { + /// Start of the first slot. + pub base_addr: u64, + /// Capacity of each slot. 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, +} + +// SAFETY: only sound for single-threaded (guest-side) access; see the +// type-level invariant on `SendWrap`. +unsafe impl Send for SendWrap>> {} + +impl Tier { + fn from_layout(layout: SlotLayout) -> Result { + if layout.slot_size == 0 { + return Err(AllocError::InvalidArg); + } + 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 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 alloc_counts( + &self, + lengths: impl IntoIterator, + ) -> Result<(usize, usize), AllocError> { + let lower_size = self.lower.as_ref().map(|lower| lower.slot_size as usize); + let upper_size = self.upper.slot_size as usize; + let free_lower = self.lower.as_ref().map_or(0, Tier::num_free); + let free_upper = self.upper.num_free(); + + let mut alloc_count = 0usize; + let mut lower_count = 0usize; + + for len in lengths { + if len == 0 { + return Err(AllocError::InvalidArg); + } + + alloc_count = alloc_count + .checked_add(len.div_ceil(upper_size)) + .ok_or(AllocError::Overflow)?; + + let tail_len = len % upper_size; + if tail_len != 0 + && lower_size.is_some_and(|size| tail_len <= size) + && lower_count < free_lower + { + lower_count += 1; + } + + if alloc_count - lower_count > free_upper { + return Err(AllocError::NoSpace); + } + } + + Ok((alloc_count, alloc_count - lower_count)) + } + + fn max_alloc( + &self, + lengths: impl IntoIterator, + alloc_limit: usize, + ) -> Result { + let (used, upper_used) = self.alloc_counts(lengths)?; + let remaining = alloc_limit.checked_sub(used).ok_or(AllocError::NoSpace)?; + if remaining == 0 { + return Err(AllocError::NoSpace); + } + + let free_upper = self.upper.num_free() - upper_used; + let upper_count = remaining.min(free_upper); + + let len = upper_count + .checked_mul(self.upper.slot_size as usize) + .ok_or(AllocError::Overflow)?; + + if len == 0 { + return Err(AllocError::NoSpace); + } + + Ok(len) + } + + fn alloc_regions(&mut self, lengths: I) -> Result + where + I: IntoIterator, + { + let lengths = SmallVec::<[usize; 4]>::from_iter(lengths); + self.alloc_counts(lengths.iter().copied())?; + let mut regions = Regions::with_capacity(lengths.len()); + let slot_size = self.max_alloc_len(); + + for total_len in lengths { + let mut allocs = Allocations::new(); + let mut remaining = total_len; + + while remaining > 0 { + let len = remaining.min(slot_size); + allocs.push(self.alloc(len).expect("plan validated upstream")); + + remaining -= len; + } + regions.push(allocs); + } + + if regions.is_empty() { + return Err(AllocError::InvalidArg); + } + + Ok(regions) + } + + fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { + if let Some(lower) = &mut self.lower + && lower.contains(addr) + { + 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. [`alloc_regions`](BufferProvider::alloc_regions) +/// splits logical payloads into bounded descriptor segments. +#[derive(Clone)] +pub struct SlotPool { + inner: SendWrap>>, +} + +impl SlotPool { + /// Create a single-tier recycling pool from exact slot placement. + pub fn new(layout: SlotLayout) -> Result { + Self::from_layouts(None, layout) + } + + /// Create a two-tier recycling pool from exact lower and upper layouts. + /// + /// The lower layout must precede the upper layout without overlap, and its + /// slot size must 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: SendWrap(Rc::new(RefCell::new(inner))), + }) + } + + /// Return every live slot address in deterministic tier and index order. + pub fn live_addrs(&self) -> Vec { + self.inner.borrow().live_addrs() + } + + /// Return the lower and upper tier layouts. + pub fn layouts(&self) -> (Option, SlotLayout) { + self.inner.borrow().layouts() + } + + /// Compute the address of slot `index`, with lower-tier slots first. + /// + /// Returns `None` if `index >= count`. + pub fn slot_addr(&self, index: usize) -> Option { + self.inner.borrow().slot_addr(index) + } + + /// Total number of free slots across all tiers. + pub fn num_free(&self) -> usize { + self.inner.borrow().num_free() + } + + /// Total number of currently allocated slots across all tiers. + pub fn num_live(&self) -> usize { + self.inner.borrow().num_live() + } + + /// Total number of free slots in the lower tier. + pub fn num_free_lower(&self) -> usize { + self.inner.borrow().lower.as_ref().map_or(0, Tier::num_free) + } + + /// Total number of free slots in the upper tier. + pub fn num_free_upper(&self) -> usize { + self.inner.borrow().upper.num_free() + } + + /// Free a previously allocated slot by address. + pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } + + /// 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() + } + + /// Maximum upper-tier region length after reserving `lengths`. + /// + /// The returned region uses only upper-tier slots. It and the reserved + /// regions use at most `alloc_limit` allocations in total. This query does + /// not mutate the pool. + /// + /// # Errors + /// + /// Returns an error when a reserved region is invalid or unavailable, no + /// additional allocation fits, or capacity arithmetic overflows. + pub fn max_alloc(&self, lengths: I, alloc_limit: usize) -> Result + where + I: IntoIterator, + { + self.inner.borrow().max_alloc(lengths, alloc_limit) + } +} + +impl BufferProvider for SlotPool { + fn preferred_segment_len(&self) -> usize { + self.inner.borrow().max_alloc_len() + } + + fn alloc(&self, len: usize) -> Result { + self.inner.borrow_mut().alloc(len) + } + + fn alloc_regions(&self, lengths: I) -> Result + where + I: IntoIterator, + { + self.inner.borrow_mut().alloc_regions(lengths) + } + + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } +} diff --git a/src/hyperlight_common/src/virtq/pool/tests.rs b/src/hyperlight_common/src/virtq/pool/tests.rs new file mode 100644 index 000000000..af5829649 --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/tests.rs @@ -0,0 +1,667 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use super::*; + +fn make_run_pool(size: usize) -> RunPool { + let base = align_up(0x10000, L.max(U)).unwrap() as u64; + RunPool::::new(base, size).unwrap() +} + +fn make_slot_pool(slot_count: usize, slot_size: usize) -> SlotPool { + let layout = SlotLayout::new(0x80000, slot_size, slot_count); + SlotPool::new(layout).unwrap() +} + +fn make_tiered_slot_pool(lower_count: usize, upper_count: usize) -> SlotPool { + let lower = SlotLayout::new(0x80000, 256, lower_count); + let upper = SlotLayout::new(0x90000, 4096, upper_count); + SlotPool::new_tiered(lower, upper).unwrap() +} + +fn alloc_exact_regions( + pool: &impl BufferProvider, + lengths: impl IntoIterator, +) -> Result { + pool.alloc_regions(lengths) +} + +#[test] +fn test_run_pool_new_success() { + let pool = RunPool::<256, 4096>::new(0x10000, 1024 * 1024).unwrap(); + assert!(pool.inner.borrow().lower.capacity() > 0); + assert!(pool.inner.borrow().upper.capacity() > 0); +} + +#[test] +fn test_run_pool_alloc_small_to_lower() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(128).unwrap(); + + // Should come from the lower tier. + assert!(pool.inner.borrow().lower.contains(alloc.addr)); + assert_eq!(alloc.len, 256); +} + +#[test] +fn test_run_pool_alloc_large_to_upper() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(1500).unwrap(); + + // Should come from the upper tier. + assert!(pool.inner.borrow().upper.contains(alloc.addr)); + assert_eq!(alloc.len, 4096); +} + +#[test] +fn test_run_pool_alloc_fallback_to_upper() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + + // Fill the lower tier completely. + let mut allocations = Vec::new(); + while pool.inner.borrow().lower.free_bytes() > 0 { + allocations.push(pool.inner.borrow_mut().lower.alloc(256).unwrap()); + } + + // Small allocation should fall back to the upper tier. + let alloc = pool.alloc(128).unwrap(); + assert!(pool.inner.borrow().upper.contains(alloc.addr)); +} + +#[test] +fn test_run_pool_free_from_lower() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(128).unwrap(); + + let free_before = pool.inner.borrow().lower.free_bytes(); + pool.dealloc(alloc.addr).unwrap(); + assert_eq!( + pool.inner.borrow().lower.free_bytes(), + free_before + alloc.len as usize + ); +} + +#[test] +fn test_run_pool_free_from_upper() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(1500).unwrap(); + + let free_before = pool.inner.borrow().upper.free_bytes(); + pool.dealloc(alloc.addr).unwrap(); + assert_eq!( + pool.inner.borrow().upper.free_bytes(), + free_before + alloc.len as usize + ); +} + +#[test] +fn test_run_pool_stress_many_allocations() { + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); + let mut allocations = Vec::new(); + + // Allocate many buffers + for i in 0..100 { + let size = if i % 2 == 0 { 128 } else { 1500 }; + allocations.push(pool.alloc(size).unwrap()); + } + + // Free half of them + for i in (0..100).step_by(2) { + pool.dealloc(allocations[i].addr).unwrap(); + } + + // Should be able to allocate again + for i in 0..50 { + let size = if i % 2 == 0 { 128 } else { 1500 }; + let _alloc = pool.alloc(size).unwrap(); + } +} + +#[test] +fn test_run_pool_mixed_workload() { + let pool = make_run_pool::<256, 4096>(2 * 1024 * 1024); + + // Simulate virtio-net workload + let desc_buf = pool.alloc(64).unwrap(); // Control message + let rx_buf1 = pool.alloc(1500).unwrap(); // MTU packet + let rx_buf2 = pool.alloc(1500).unwrap(); // MTU packet + let tx_buf = pool.alloc(4096).unwrap(); // Large buffer + + // Free and reallocate + pool.dealloc(rx_buf1.addr).unwrap(); + let rx_buf3 = pool.alloc(1500).unwrap(); + + // Should reuse freed buffer (LIFO) + assert_eq!(rx_buf3.addr, rx_buf1.addr); + + pool.dealloc(desc_buf.addr).unwrap(); + pool.dealloc(rx_buf2.addr).unwrap(); + pool.dealloc(rx_buf3.addr).unwrap(); + pool.dealloc(tx_buf.addr).unwrap(); +} + +#[test] +fn test_run_pool_zero_allocation_error() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let result = pool.alloc(0); + assert!(matches!(result, Err(AllocError::InvalidArg))); +} + +#[test] +fn test_run_pool_too_large_allocation() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let result = pool.alloc(2 * 1024 * 1024); // Larger than pool + assert!(matches!(result, Err(AllocError::OutOfMemory))); +} + +#[test] +fn test_align_up_helper() { + assert_eq!(align_up(0, 256).unwrap(), 0); + assert_eq!(align_up(1, 256).unwrap(), 256); + assert_eq!(align_up(256, 256).unwrap(), 256); + assert_eq!(align_up(257, 256).unwrap(), 512); + assert_eq!(align_up(511, 256).unwrap(), 512); + assert_eq!(align_up(512, 256).unwrap(), 512); + assert!(matches!(align_up(1, 0), Err(AllocError::InvalidArg))); + assert!(matches!( + align_up(usize::MAX, 256), + Err(AllocError::Overflow) + )); +} + +#[test] +fn test_slot_pool_preserves_exact_base() { + let layout = SlotLayout::new(0x80001, 4096, 2); + let pool = SlotPool::new(layout).unwrap(); + + assert_eq!(pool.base_addr(), 0x80001); + assert_eq!(pool.count(), 2); + assert_eq!(pool.slot_addr(0), Some(0x80001)); + assert_eq!(pool.slot_addr(1), Some(0x81001)); +} + +#[test] +fn test_tiered_slot_pool_reports_layouts() { + let lower = SlotLayout::new(0x80001, 0x100, 2); + let upper = SlotLayout::new(0x90001, 0x1000, 2); + let pool = SlotPool::new_tiered(lower, upper).unwrap(); + + let (lower, upper) = pool.layouts(); + assert_eq!(lower, Some(SlotLayout::new(0x80001, 0x100, 2))); + assert_eq!(upper, SlotLayout::new(0x90001, 0x1000, 2)); + assert_eq!(pool.base_addr(), 0x80001); + assert_eq!(pool.slot_size(), 0x1000); + assert_eq!(pool.count(), 4); + assert_eq!(pool.slot_addr(0), Some(0x80001)); + assert_eq!(pool.slot_addr(1), Some(0x80101)); + assert_eq!(pool.slot_addr(2), Some(0x90001)); + assert_eq!(pool.slot_addr(3), Some(0x91001)); + assert_eq!(pool.slot_addr(4), None); +} + +#[test] +fn test_tiered_slot_pool_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_region_uses_both_tiers() { + let pool = make_tiered_slot_pool(1, 2); + let regions = alloc_exact_regions(&pool, [4096 + 128]).unwrap(); + let allocations = ®ions[0]; + + assert_eq!(regions.len(), 1); + assert_eq!(allocations.len(), 2); + assert_eq!(allocations[0].len, 4096); + assert_eq!(allocations[1].len, 256); + assert!((0x90000..0x92000).contains(&allocations[0].addr)); + assert!((0x80000..0x80100).contains(&allocations[1].addr)); + + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); + } + assert_eq!(pool.num_free(), 3); +} + +#[test] +fn test_tiered_slot_pool_allocates_regions_in_order() { + let pool = make_tiered_slot_pool(1, 2); + let regions = alloc_exact_regions(&pool, [128, 128]).unwrap(); + + assert_eq!(regions.len(), 2); + assert_eq!(regions[0].len(), 1); + assert_eq!(regions[1].len(), 1); + assert!((0x80000..0x80100).contains(®ions[0][0].addr)); + assert!((0x90000..0x92000).contains(®ions[1][0].addr)); + + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); + } +} + +#[test] +fn test_slot_pool_max_alloc_reserves_regions_and_honors_limit() { + let pool = make_tiered_slot_pool(2, 3); + + let max = pool.max_alloc([128], 3).unwrap(); + assert_eq!(max, 2 * 4096); + assert_eq!(pool.num_free(), 5); + + let regions = pool.alloc_regions([128, max]).unwrap(); + assert_eq!(regions.iter().map(Allocations::len).sum::(), 3); + + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); + } +} + +#[test] +fn test_slot_pool_max_alloc_requires_one_remaining_allocation() { + let pool = make_tiered_slot_pool(1, 1); + + assert!(matches!(pool.max_alloc([128], 1), Err(AllocError::NoSpace))); + assert_eq!(pool.num_free(), 2); +} + +#[test] +fn test_slot_pool_max_alloc_ignores_remaining_lower_slot() { + let pool = make_tiered_slot_pool(1, 1); + + assert!(matches!( + pool.max_alloc([4096], 2), + Err(AllocError::NoSpace) + )); + assert_eq!(pool.num_free(), 2); +} + +#[test] +fn test_slot_pool_rejects_invalid_or_unavailable_regions() { + let pool = make_tiered_slot_pool(1, 1); + + assert!(matches!( + pool.alloc_regions([0]), + Err(AllocError::InvalidArg) + )); + assert!(matches!( + pool.alloc_regions([]), + Err(AllocError::InvalidArg) + )); + assert!(matches!( + pool.alloc_regions([128, 0]), + Err(AllocError::InvalidArg) + )); + assert!(matches!( + pool.alloc_regions([4096 + 257]), + Err(AllocError::NoSpace) + )); + assert!(matches!( + pool.alloc_regions([128, 4096, 1]), + Err(AllocError::NoSpace) + )); + assert_eq!(pool.num_free(), 2); +} + +#[test] +fn test_tiered_slot_pool_dealloc_routes_by_region() { + let pool = make_tiered_slot_pool(1, 1); + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(1024).unwrap(); + + pool.dealloc(lower.addr).unwrap(); + pool.dealloc(upper.addr).unwrap(); + assert_eq!(pool.num_free(), 2); + assert!(matches!( + pool.dealloc(lower.addr), + Err(AllocError::InvalidFree(_, _)) + )); + assert!(matches!( + pool.dealloc(0x88000), + Err(AllocError::InvalidFree(_, _)) + )); +} + +// Edge case: allocation exactly at boundary +#[test] +fn test_run_pool_boundary_allocation() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + + // Allocate exactly at boundary + let alloc = pool.alloc(256).unwrap(); + assert!(pool.inner.borrow().lower.contains(alloc.addr)); + + // Allocate just over boundary + let alloc2 = pool.alloc(257).unwrap(); + assert!(pool.inner.borrow().upper.contains(alloc2.addr)); +} + +#[test] +fn test_run_pool_dealloc_addr_routes_to_correct_tier() { + let pool = make_run_pool::<256, 4096>(0x20000); + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(1024).unwrap(); + + assert_eq!(pool.allocation_len(lower.addr).unwrap(), 256); + assert_eq!(pool.allocation_len(upper.addr).unwrap(), 4096); + + pool.dealloc_addr(lower.addr).unwrap(); + pool.dealloc_addr(upper.addr).unwrap(); +} + +#[test] +fn test_run_pool_region_uses_one_contiguous_run() { + let pool = make_run_pool::<256, 4096>(0x20000); + let regions = alloc_exact_regions(&pool, [4096 * 2 + 1]).unwrap(); + + assert_eq!(regions.len(), 1); + assert_eq!(regions[0].len(), 1); + assert_eq!(regions[0][0].len, 4096 * 3); + + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); + } +} + +#[test] +fn test_run_pool_allocates_each_region_as_one_run() { + let pool = make_run_pool::<256, 4096>(0x20000); + let regions = alloc_exact_regions(&pool, [8192, 128]).unwrap(); + + assert_eq!(regions.len(), 2); + assert_eq!(regions[0].len(), 1); + assert_eq!(regions[1].len(), 1); + assert_eq!(regions[0][0].len, 8192); + assert_eq!(regions[1][0].len, 256); + + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); + } +} + +#[test] +fn test_slot_pool_region_splits() { + let pool = make_slot_pool(8, 4096); + let regions = alloc_exact_regions(&pool, [4096 * 2 + 1]).unwrap(); + let allocations = ®ions[0]; + + assert_eq!(regions.len(), 1); + assert_eq!(allocations.len(), 3); + assert_eq!(allocations[0].len, 4096); + assert_eq!(allocations[1].len, 4096); + assert_eq!(allocations[2].len, 4096); + + for allocation in regions.into_iter().flatten() { + pool.dealloc(allocation.addr).unwrap(); + } +} + +#[test] +fn test_tiered_slot_pool_live_addrs_are_deterministic() { + let pool = make_tiered_slot_pool(2, 2); + assert_eq!(pool.num_live(), 0); + + let lower_high = pool.alloc(128).unwrap(); + let upper_high = pool.alloc(1024).unwrap(); + let lower_low = pool.alloc(128).unwrap(); + + assert_eq!(pool.num_live(), 3); + assert_eq!( + pool.live_addrs(), + vec![lower_low.addr, lower_high.addr, upper_high.addr] + ); +} + +#[test] +fn test_slot_pool_dealloc_out_of_range() { + let pool = make_slot_pool(4, 4096); + let _ = pool.alloc(4096).unwrap(); + + assert!(matches!( + pool.dealloc(0xDEAD), + Err(AllocError::InvalidFree(0xDEAD, 0)) + )); +} + +#[test] +fn test_slot_pool_dealloc_misaligned() { + let pool = make_slot_pool(4, 4096); + let _ = pool.alloc(4096).unwrap(); + + assert!(matches!( + pool.dealloc(0x80001), + Err(AllocError::InvalidFree(0x80001, 0)) + )); +} + +#[test] +fn test_slot_pool_dealloc_double_free() { + let pool = make_slot_pool(4, 4096); + let a = pool.alloc(4096).unwrap(); + pool.dealloc(a.addr).unwrap(); + + // Second dealloc should fail - address is already in the free list + assert!(matches!( + pool.dealloc(a.addr), + Err(AllocError::InvalidFree(_, _)) + )); +} + +#[test] +fn test_slot_pool_alloc_regions_preflights_sequence() { + let pool = make_slot_pool(2, 4096); + + assert!(matches!( + pool.alloc_regions([]), + Err(AllocError::InvalidArg) + )); + assert!(matches!( + pool.alloc_regions([4096, 0]), + Err(AllocError::InvalidArg) + )); + assert!(matches!( + pool.alloc_regions([4096, 4096, 1]), + Err(AllocError::NoSpace) + )); + assert_eq!(pool.num_free(), 2); + + let alloc = pool.alloc(4096).unwrap(); + assert_eq!(pool.num_free(), 1); + pool.dealloc(alloc.addr).unwrap(); +} + +#[test] +fn test_slot_pool_dealloc_addr_and_allocation_len() { + let pool = make_slot_pool(4, 4096); + let alloc = pool.alloc(4096).unwrap(); + + assert_eq!(pool.allocation_len(alloc.addr).unwrap(), 4096); + pool.dealloc_addr(alloc.addr).unwrap(); + assert!(matches!( + pool.allocation_len(alloc.addr), + Err(AllocError::InvalidFree(_, 0)) + )); +} + +#[test] +fn test_slot_pool_random_order_dealloc() { + let pool = make_slot_pool(8, 4096); + + let mut allocs: Vec = (0..8).map(|_| pool.alloc(4096).unwrap()).collect(); + assert_eq!(pool.num_free(), 0); + + // Dealloc in reverse order + allocs.reverse(); + for a in &allocs { + pool.dealloc(a.addr).unwrap(); + } + assert_eq!(pool.num_free(), 8); + + // All slots should be re-allocatable + let reallocs: Vec = (0..8).map(|_| pool.alloc(4096).unwrap()).collect(); + assert_eq!(pool.num_free(), 0); + + // Verify all addresses are distinct + let mut addrs: Vec = reallocs.iter().map(|a| a.addr).collect(); + addrs.sort(); + addrs.dedup(); + assert_eq!(addrs.len(), 8); +} + +#[test] +fn test_slot_pool_interleaved_alloc_dealloc_order() { + let pool = make_slot_pool(4, 4096); + + let a0 = pool.alloc(4096).unwrap(); + let a1 = pool.alloc(4096).unwrap(); + let a2 = pool.alloc(4096).unwrap(); + let a3 = pool.alloc(4096).unwrap(); + assert_eq!(pool.num_free(), 0); + + // Free middle slots first (out of allocation order) + pool.dealloc(a2.addr).unwrap(); + pool.dealloc(a0.addr).unwrap(); + assert_eq!(pool.num_free(), 2); + + // Re-alloc gets the out-of-order slots back (LIFO) + let b0 = pool.alloc(4096).unwrap(); + assert_eq!(b0.addr, a0.addr); + let b1 = pool.alloc(4096).unwrap(); + assert_eq!(b1.addr, a2.addr); + + // Free everything in yet another order + pool.dealloc(a1.addr).unwrap(); + pool.dealloc(b0.addr).unwrap(); + pool.dealloc(b1.addr).unwrap(); + pool.dealloc(a3.addr).unwrap(); + assert_eq!(pool.num_free(), 4); + + // All 4 original addresses should be available + let mut final_addrs: Vec = (0..4).map(|_| pool.alloc(4096).unwrap().addr).collect(); + final_addrs.sort(); + let expected: Vec = (0..4).map(|i| 0x80000 + i * 4096).collect(); + assert_eq!(final_addrs, expected); +} + +#[test] +fn test_slot_pool_dealloc_order_independent_of_alloc_order() { + let pool = make_slot_pool(6, 256); + + // Allocate all + let allocs: Vec = (0..6).map(|_| pool.alloc(256).unwrap()).collect(); + + // Dealloc in scattered order: 4, 1, 5, 0, 3, 2 + let order = [4, 1, 5, 0, 3, 2]; + for &i in &order { + pool.dealloc(allocs[i].addr).unwrap(); + } + assert_eq!(pool.num_free(), 6); + + // Re-allocate all and verify we get back the full set + let mut realloc_addrs: Vec = (0..6).map(|_| pool.alloc(256).unwrap().addr).collect(); + realloc_addrs.sort(); + + let mut orig_addrs: Vec = allocs.iter().map(|a| a.addr).collect(); + orig_addrs.sort(); + + assert_eq!(realloc_addrs, orig_addrs); +} diff --git a/src/hyperlight_common/src/virtq/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index 1d3750e6c..8d5cba28e 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -15,6 +15,7 @@ limitations under the License. */ use alloc::collections::VecDeque; +use alloc::vec; use alloc::vec::Vec; use bytes::Bytes; @@ -93,6 +94,86 @@ pub(crate) struct Inflight { chain: BufferChain, } +/// Compact in-flight chains with constant-time descriptor-ID lookup. +/// +/// Descriptor IDs span the full ring, but live chains are normally bounded by +/// the much smaller buffer pool. `by_id` maps each descriptor ID to a packed +/// `live` index. Removal uses `swap_remove` and repairs the moved entry's map. +struct InflightTable { + by_id: Vec, + live: Vec, +} + +impl InflightTable { + const VACANT: u16 = u16::MAX; + + fn new(ring_len: usize) -> Self { + Self { + by_id: vec![Self::VACANT; ring_len], + live: Vec::new(), + } + } + + fn try_reserve_one(&mut self) -> Result<(), VirtqError> { + if self.live.len() > self.by_id.len() { + return Err(VirtqError::InvalidState); + } + + if self.live.len() == self.by_id.len() { + return Err(VirtqError::Backpressure); + } + + // Producers with one live chain should not pay for four large inline + // chain records. + let result = if self.live.capacity() == 0 { + self.live.try_reserve_exact(1) + } else { + self.live.try_reserve(1) + }; + + result.map_err(|_| VirtqError::BookkeepingAllocation) + } + + fn contains(&self, id: u16) -> bool { + self.by_id + .get(id as usize) + .is_some_and(|slot| *slot != Self::VACANT) + } + + fn insert(&mut self, inflight: Inflight) { + let id = inflight.token.id; + debug_assert!(!self.contains(id)); + debug_assert!(self.live.len() < Self::VACANT as usize); + + let slot = self.live.len() as u16; + self.live.push(inflight); + self.by_id[id as usize] = slot; + } + + fn remove(&mut self, id: u16) -> Option { + let slot = self.by_id.get_mut(id as usize)?; + if *slot == Self::VACANT { + return None; + } + + let index = usize::from(*slot); + *slot = Self::VACANT; + + let removed = self.live.swap_remove(index); + if let Some(moved) = self.live.get(index) { + self.by_id[moved.token.id as usize] = index as u16; + } + + Some(removed) + } + + fn pop(&mut self) -> Option { + let inflight = self.live.pop()?; + self.by_id[inflight.token.id as usize] = Self::VACANT; + Some(inflight) + } +} + /// A high-level virtqueue producer (driver side). /// /// The producer sends chains to the consumer (device), and receives used chains. @@ -130,7 +211,7 @@ pub struct VirtqProducer { notifier: N, pool: P, next_token: u32, - inflight: Vec>, + inflight: InflightTable, pending: VecDeque, } @@ -151,18 +232,23 @@ where pub fn new(layout: Layout, mem: M, notifier: N, pool: P) -> Self { let inner = RingProducer::new(layout, mem); let ring_len = inner.len(); + let inflight = InflightTable::new(ring_len); Self { inner, pool, notifier, + inflight, next_token: 0, - inflight: (0..ring_len).map(|_| None).collect(), - pending: VecDeque::with_capacity(ring_len), + pending: VecDeque::new(), } } - fn dealloc_elems( + /// Retire allocations from a completed descriptor chain. + /// + /// Every element is attempted so one deallocation failure does not strand + /// later allocations; the first failure is returned after cleanup. + fn retire_elems( &self, elems: impl IntoIterator, ) -> Result<(), VirtqError> { @@ -189,6 +275,11 @@ where ChainBuilder::new(self.inner.mem().clone(), self.pool.clone()) } + /// Preferred size of one bulk payload segment. + pub fn preferred_segment_len(&self) -> usize { + self.pool.preferred_segment_len() + } + /// Begin a batch of submissions. /// /// Chains submitted through the returned [`SubmitBatch`] are published to @@ -219,17 +310,19 @@ where } fn publish(&mut self, send: SendChain) -> Result { + self.inflight.try_reserve_one()?; + let token_id = self.next_token; let id = self.inner.submit_available(send.chain())?; let token = Token { seq: token_id, id }; // A free descriptor id must never already be tracked as inflight. - if self.inflight[id as usize].is_some() { + if self.inflight.contains(id) { return Err(VirtqError::InvalidState); } let inf = send.into_inflight(token); - self.inflight[id as usize] = Some(inf); + self.inflight.insert(inf); self.next_token = self.next_token.wrapping_add(1); Ok(token) @@ -273,6 +366,44 @@ where self.inner.num_free() } + /// Number of submitted descriptors not yet polled as used. + #[inline] + pub fn num_inflight(&self) -> usize { + self.inner.num_inflight() + } + + /// Reset a stopped producer and release transport-owned allocations. + /// + /// The peer must not access the ring until its consumer is reset. Buffered + /// writable completions are guest-owned and make this operation fail. + /// Owner-backed payloads already returned to callers are not tracked as + /// in-flight and remain allocated. + pub fn reset(&mut self) -> Result<(), VirtqError> { + if !self.pending.is_empty() { + return Err(VirtqError::InvalidState); + } + + self.inner.reset()?; + self.next_token = 0; + + let mut maybe_err = None; + + // Drain all in-flight chains and retire their allocations. This is a best-effort + while let Some(inflight) = self.inflight.pop() { + let ret = self.retire_elems(inflight.chain.elems().iter().copied()); + if let Err(err) = ret + && maybe_err.is_none() + { + maybe_err = Some(err); + } + } + + match maybe_err { + Some(error) => Err(error), + None => Ok(()), + } + } + /// Configure event suppression for used buffer notifications. /// /// This controls when the device (consumer) signals us about completed buffers: @@ -302,39 +433,6 @@ where } Ok(()) } - - /// Reset ring, inflight, and pool state to initial values. - /// - /// # Safety - /// - /// No outstanding [`UsedChain::Data`] buffers, borrowed segment views, or - /// peer accesses to previously submitted descriptors may exist. Resetting - /// recycles the same backing addresses, so outstanding zero-copy buffers or - /// stale descriptor users could alias memory that is handed out again. - /// - /// TODO(virtq): find a way to allow guest to keep used chains across resets. - pub unsafe fn reset(&mut self) { - self.inflight.iter_mut().for_each(|slot| *slot = None); - self.pending.clear(); - self.inner.reset(); - self.pool.reset(); - } - - /// Replace the pool and reset ring, inflight, and pending state. - /// - /// # Safety - /// - /// No outstanding [`UsedChain::Data`] buffers, borrowed segment views, or - /// peer accesses to previously submitted descriptors may exist. The new pool - /// may manage the same shared-memory addresses as the old pool, so old - /// zero-copy buffers must not outlive this transition. - pub unsafe fn reset_with_pool(&mut self, pool: P) { - self.pending.clear(); - self.inflight.iter_mut().for_each(|slot| *slot = None); - self.inner.reset(); - self.pool = pool; - self.pool.reset(); - } } impl VirtqProducer @@ -382,6 +480,9 @@ where while let Some(chain) = self.poll_ring()? { if matches!(chain, UsedChain::Data(_, _)) { debug_assert!(self.pending.len() < self.inner.len()); + self.pending + .try_reserve(1) + .map_err(|_| VirtqError::BookkeepingAllocation)?; self.pending.push_back(chain); } count += 1; @@ -399,14 +500,13 @@ 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())?; + self.retire_elems(chain.readables().iter().copied())?; let used = if chain.writables().is_empty() { UsedChain::Ack(token) @@ -439,14 +539,14 @@ where if remaining != 0 { let elems = owned.iter().map(|(elem, _)| *elem).chain(free); - self.dealloc_elems(elems)?; + self.retire_elems(elems)?; return Err(VirtqError::InvalidState); } for (elem, len) in &owned { if unsafe { self.inner.mem().as_slice(elem.addr, *len) }.is_err() { let elems = owned.iter().map(|(elem, _)| *elem).chain(free); - let _ = self.dealloc_elems(elems); + let _ = self.retire_elems(elems); return Err(VirtqError::MemoryReadError); } } @@ -457,7 +557,7 @@ where self.pool.clone(), Allocation { addr: elem.addr, - len: elem.len as usize, + len: elem.len, }, ); let mem = self.inner.mem().clone(); @@ -469,7 +569,7 @@ where sgs.push(Bytes::from_owner(owner)); } - self.dealloc_elems(free)?; + self.retire_elems(free)?; Ok(Segments::from_smallvec(sgs)) } @@ -502,9 +602,9 @@ 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"] +/// [`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, P> { producer: &'a mut VirtqProducer, notify_from: Option, @@ -550,6 +650,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. @@ -602,105 +708,83 @@ impl ChainBuilder { /// # Errors /// /// - [`VirtqError::InvalidState`] - No buffers requested - /// - [`VirtqError::Alloc`] - Pool exhausted + /// - [`VirtqError::Alloc`] - Buffer allocation failed pub fn build(self) -> Result, VirtqError> { if self.rd_caps.is_empty() && self.wr_caps.is_empty() { return Err(VirtqError::InvalidState); } - let mut rollback = Rollback::new(&self.pool); + let rd_capacity = self.rd_caps.iter().try_fold(0usize, |total, &cap| { + total.checked_add(cap).ok_or(AllocError::Overflow) + })?; + + let lengths = self.rd_caps.iter().chain(&self.wr_caps).copied(); + let regions = self.pool.alloc_regions(lengths)?; + + debug_assert_eq!(regions.len(), self.rd_caps.len() + self.wr_caps.len()); + + let mut regions = regions.into_iter(); let mut rd_caps = SmallVec::<[usize; 4]>::new(); let mut rd_elems = SmallVec::<[BufferElement; 4]>::new(); let mut wr_elems = SmallVec::<[BufferElement; 4]>::new(); - // Allocate readable buffers, splitting into multiple descriptors if needed. // The buffer element lengths are initialized to zero and updated as the // `SendChain` writes. - for &cap in &self.rd_caps { - let sgs = self.pool.alloc_sg(cap)?; + for (&cap, allocs) in Iterator::zip(self.rd_caps.iter(), regions.by_ref()) { let mut remaining = cap; - for alloc in sgs { - let _ = checked_descriptor_len(alloc.len)?; - let seg_cap = remaining.min(alloc.len); + for alloc in allocs { + debug_assert_ne!(remaining, 0); + + let seg_cap = remaining.min(alloc.len as usize); + let elem = BufferElement::readable(alloc.addr); rd_caps.push(seg_cap); - rd_elems.push(BufferElement { - addr: alloc.addr, - len: 0, - writable: false, - }); + rd_elems.push(elem); + remaining -= seg_cap; - rollback.allocs.push(alloc); } - if remaining != 0 { - return Err(VirtqError::InvalidState); - } + // The sum of the allocation lengths must equal the requested capacity. + debug_assert_eq!(remaining, 0); } - // Allocate writable buffers, with the same caveat about splitting as readable buffers. // Writable buffer elements are initialized with their full capacity for the device to // write into. - for &cap in &self.wr_caps { - let sgs = self.pool.alloc_sg(cap)?; - for alloc in sgs { - let len = checked_descriptor_len(alloc.len)?; - wr_elems.push(BufferElement { - addr: alloc.addr, - len, - writable: true, - }); - rollback.allocs.push(alloc); + for (&cap, allocs) in Iterator::zip(self.wr_caps.iter(), regions.by_ref()) { + let mut remaining = cap; + + for alloc in allocs { + debug_assert_ne!(remaining, 0); + let elem = BufferElement::writable(alloc.addr, alloc.len); + + wr_elems.push(elem); + remaining = remaining.saturating_sub(alloc.len as usize); } + debug_assert_eq!(remaining, 0); } + // All requested readable and writable buffers must have been allocated. + debug_assert!(regions.next().is_none()); + let chain = BufferChainBuilder::new() .readables(rd_elems) .writables(wr_elems) - .build()?; - - rollback.release(); + .build() + .expect("validated regions produce a nonempty chain"); Ok(SendChain { mem: self.mem, pool: self.pool, chain: Some(chain), rd_caps, - rd_capacity: self.rd_caps.iter().sum(), + rd_capacity, rd_written: 0, write_mode: WriteMode::Unset, }) } } -struct Rollback<'a, P: BufferProvider> { - 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. /// @@ -752,10 +836,12 @@ pub struct SendChain { // public lifetime, so these `expect`s cannot fail. #[allow(clippy::expect_used)] impl SendChain { + #[inline(always)] fn chain(&self) -> &BufferChain { self.chain.as_ref().expect("SendChain missing BufferChain") } + #[inline(always)] fn chain_mut(&mut self) -> &mut BufferChain { self.chain.as_mut().expect("SendChain missing BufferChain") } @@ -775,22 +861,38 @@ impl SendChain { Inflight { token, chain } } - /// Number of producer-written readable segments in this chain. - pub fn segment_count(&self) -> usize { + /// Total number of descriptors in this chain. + #[inline] + pub fn desc_count(&self) -> usize { + self.chain().len() + } + + /// Number of readable descriptors in this chain. + #[inline] + pub fn rd_desc_count(&self) -> usize { self.chain().readables().len() } + /// Number of writable descriptors in this chain. + #[inline] + pub fn wr_desc_count(&self) -> usize { + self.chain().writables().len() + } + /// Total producer-written readable capacity in bytes. + #[inline] pub fn capacity(&self) -> usize { 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() } @@ -800,14 +902,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); } @@ -816,40 +919,34 @@ impl SendChain { let mut remaining = &buf[..buf.len().min(self.remaining())]; let mut written = 0; - let SendChain { - mem, - chain, - rd_caps, - .. - } = self; - - let readables = chain - .as_mut() - .expect("SendChain missing BufferChain") - .readables_mut(); - - for (readable, &cap) in readables.iter_mut().zip(rd_caps.iter()) { + for index in 0..self.rd_caps.len() { if remaining.is_empty() { break; } - let written_len = readable.len as usize; - let free = cap - written_len; - if free == 0 { + let cap = self.rd_caps[index]; + let elem = self.chain().readables()[index]; + let desc_off = elem.len as usize; + let len = (cap - desc_off).min(remaining.len()); + if len == 0 { continue; } - let n = free.min(remaining.len()); - let addr = readable.addr + written_len as u64; - mem.write(addr, &remaining[..n]) + let addr = elem + .addr + .checked_add(desc_off as u64) + .ok_or(VirtqError::MemoryWriteError)?; + + self.mem + .write(addr, &remaining[..len]) .map_err(|_| VirtqError::MemoryWriteError)?; - readable.len += n as u32; - written += n; - remaining = &remaining[n..]; + self.chain_mut().readables_mut()[index].len += len as u32; + self.rd_written += len; + written += len; + remaining = &remaining[len..]; } - self.rd_written += written; Ok(written) } @@ -864,8 +961,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); } @@ -1003,15 +1101,148 @@ fn checked_descriptor_len(len: usize) -> Result { #[cfg(test)] mod tests { use super::*; - use crate::virtq::ring::tests::{TestMem, make_consumer, make_ring}; + use crate::virtq::ring::tests::{OwnedRing, TestMem, make_consumer, make_ring}; use crate::virtq::test_utils::*; fn poll_received( consumer: &mut VirtqConsumer, - ) -> (RecvChain, ReplyChain) { + ) -> (RecvChain, ReplyChain) { consumer.poll(1024).unwrap().unwrap() } + fn make_slot_producer( + ring: &OwnedRing, + slot_size: usize, + ) -> ( + VirtqProducer, + VirtqConsumer, + ) { + let mem = ring.mem(); + let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + + let lower = SlotLayout::new(pool_base, slot_size / 2, ring.len()); + let upper = SlotLayout::new(lower.end_addr().unwrap(), slot_size, ring.len()); + let pool = SlotPool::new_tiered(lower, upper).unwrap(); + + let notifier = TestNotifier::new(); + let producer = VirtqProducer::new(ring.layout(), mem.clone(), notifier.clone(), pool); + let consumer = VirtqConsumer::new(ring.layout(), mem, notifier); + (producer, consumer) + } + + fn inflight(seq: u32, id: u16) -> Inflight { + let chain = BufferChainBuilder::new() + .readable(0x1000 + u64::from(id) * 0x10, 8) + .build() + .unwrap(); + Inflight { + token: Token { seq, id }, + chain, + } + } + + #[test] + fn inflight_table_repairs_moved_entry_after_removal() { + let mut table = InflightTable::new(16); + for (seq, id) in [(0, 3), (1, 7), (2, 5)] { + table.try_reserve_one().unwrap(); + table.insert(inflight(seq, id)); + } + + assert_eq!(table.remove(7).unwrap().token.seq, 1); + assert!(!table.contains(7)); + assert_eq!(table.remove(5).unwrap().token.seq, 2); + assert_eq!(table.remove(3).unwrap().token.seq, 0); + assert!(table.live.is_empty()); + assert!(table.remove(7).is_none()); + } + + #[test] + fn producer_bookkeeping_starts_compact_and_lazy() { + let ring = make_ring(64); + let (producer, _consumer, _notifier) = make_test_producer(&ring); + + assert_eq!(producer.inflight.by_id.len(), ring.len()); + assert!(producer.inflight.live.is_empty()); + assert_eq!(producer.inflight.live.capacity(), 0); + assert_eq!(producer.pending.capacity(), 0); + } + + #[test] + fn full_ring_still_reports_backpressure() { + let ring = make_ring(4); + let (mut producer, _consumer, _notifier) = make_test_producer(&ring); + + for _ in 0..ring.len() { + let chain = producer.chain().readable(1).build().unwrap(); + producer.submit(chain).unwrap(); + } + + let chain = producer.chain().readable(1).build().unwrap(); + assert!(matches!( + producer.submit(chain), + Err(VirtqError::Backpressure) + )); + } + + #[test] + fn reset_reclaims_inflight_slots_and_reuses_ring() { + let ring = make_ring(8); + let mem = ring.mem(); + let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + let pool = SlotPool::new(SlotLayout::new(pool_base, 64, ring.len())).unwrap(); + let notifier = TestNotifier::new(); + let mut producer = VirtqProducer::new(ring.layout(), mem, notifier, pool.clone()); + + for _ in 0..ring.len() { + let chain = producer.chain().writable(64).build().unwrap(); + producer.submit(chain).unwrap(); + } + + assert_eq!(pool.num_free(), 0); + + producer.reset().unwrap(); + + assert_eq!(producer.num_inflight(), 0); + assert_eq!(producer.num_free(), ring.len()); + assert_eq!(pool.num_free(), ring.len()); + assert!(producer.inflight.live.is_empty()); + assert!( + producer + .inflight + .by_id + .iter() + .all(|slot| *slot == InflightTable::VACANT) + ); + + for _ in 0..ring.len() { + let chain = producer.chain().writable(64).build().unwrap(); + producer.submit(chain).unwrap(); + } + } + + #[test] + fn reset_rejects_buffered_writable_completion() { + let ring = make_ring(8); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + let chain = producer.chain().writable(64).build().unwrap(); + producer.submit(chain).unwrap(); + + let (recv, reply) = poll_received(&mut consumer); + let ReplyChain::Writable(mut reply) = reply else { + panic!("expected writable reply"); + }; + + reply.write_all(b"retained").unwrap(); + consumer.complete(recv, reply).unwrap(); + producer.reclaim().unwrap(); + + assert!(matches!(producer.reset(), Err(VirtqError::InvalidState))); + + drop(producer.poll().unwrap().unwrap()); + producer.reset().unwrap(); + } + #[derive(Clone)] struct NoDirectSliceMem(TestMem); @@ -1062,7 +1293,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); } @@ -1085,27 +1318,80 @@ mod tests { let token = producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), token); - assert_eq!(recv.to_bytes().as_ref(), b"hello world"); - assert_eq!(recv.segments().segment_count(), 2); - assert_eq!(recv.segments().as_slice()[0].as_ref(), b"hello"); - assert_eq!(recv.segments().as_slice()[1].as_ref(), b" world"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello world"); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.as_slice()[0].as_ref(), b"hello"); + assert_eq!(segments.as_slice()[1].as_ref(), b" world"); + consumer.complete(recv, reply).unwrap(); } #[test] - fn test_chain_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.chain().readables(); + + assert_eq!(readables.len(), 2); + assert_eq!(readables[0].addr, lower.base_addr); + assert_eq!(readables[1].addr, upper.base_addr); + assert_eq!(pool.num_free_lower(), 0); + assert_eq!(pool.num_free_upper(), 0); + + drop(send); + + assert_eq!(pool.num_free_lower(), 1); + assert_eq!(pool.num_free_upper(), 1); + } + + #[test] + fn test_chain_multi_readable_appends_across_calls() { + let ring = make_ring(16); + let (mut producer, mut consumer) = make_slot_producer(&ring, 4); + + let mut send = producer.chain().readable(8).build().unwrap(); + send.write_all(b"abc").unwrap(); + send.write_all(b"def").unwrap(); + assert_eq!(send.written(), 6); + assert_eq!(send.remaining(), 2); + + producer.submit(send).unwrap(); + let (recv, reply) = poll_received(&mut consumer); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.as_slice()[0].as_ref(), b"abcd"); + assert_eq!(segments.as_slice()[1].as_ref(), b"ef"); + consumer.complete(recv, reply).unwrap(); + } + + #[test] + fn test_chain_readable_splits_logical_capacity() { + let ring = make_ring(16); + let (mut producer, mut consumer) = make_slot_producer(&ring, 4); let mut se = producer.chain().readable(10).writable(32).build().unwrap(); - assert_eq!(se.segment_count(), 3); + assert_eq!(se.rd_desc_count(), 3); assert_eq!(se.capacity(), 10); se.write_all(b"abcdefghij").unwrap(); @@ -1114,12 +1400,13 @@ mod tests { let token = producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), token); - assert_eq!(recv.to_bytes().as_ref(), b"abcdefghij"); - assert_eq!(recv.segments().segment_count(), 3); - assert_eq!(recv.segments().as_slice()[0].as_ref(), b"abcd"); - assert_eq!(recv.segments().as_slice()[1].as_ref(), b"efgh"); - assert_eq!(recv.segments().as_slice()[2].as_ref(), b"ij"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"abcdefghij"); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 3); + assert_eq!(segments.as_slice()[0].as_ref(), b"abcd"); + assert_eq!(segments.as_slice()[1].as_ref(), b"efgh"); + assert_eq!(segments.as_slice()[2].as_ref(), b"ij"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1136,25 +1423,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_slot_producer(&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); @@ -1186,8 +1467,8 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"headbody"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"headbody"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1209,9 +1490,10 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"headbody"); - assert_eq!(recv.segments().segment_count(), 2); - consumer.complete(reply).unwrap(); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.to_bytes().as_ref(), b"headbody"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1225,9 +1507,10 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"headbody"); - assert_eq!(recv.segments().segment_count(), 2); - consumer.complete(reply).unwrap(); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.to_bytes().as_ref(), b"headbody"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1238,14 +1521,14 @@ mod tests { let se = producer.chain().writable(5).writable(6).build().unwrap(); let token = producer.submit(se).unwrap(); - let (_recv, reply) = poll_received(&mut consumer); + let (recv, reply) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable reply"); }; assert_eq!(wc.capacity(), 11); wc.write_all(b"hello world").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); let used = producer.poll().unwrap().unwrap(); assert_eq!(used.token(), token); @@ -1264,13 +1547,13 @@ mod tests { let se = producer.chain().writable(5).writable(6).build().unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_received(&mut consumer); + let (recv, reply) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable reply"); }; wc.write_all(b"hello wo").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); let used = producer.poll().unwrap().unwrap(); let segments = used.segments().unwrap(); @@ -1288,8 +1571,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(); @@ -1342,8 +1625,8 @@ mod tests { let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), tok); - assert_eq!(recv.to_bytes().as_ref(), b"hello world"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello world"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1362,8 +1645,8 @@ mod tests { let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), tok); - assert_eq!(recv.to_bytes().as_ref(), b"hello world"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello world"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1378,8 +1661,8 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello wo"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello wo"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1397,8 +1680,8 @@ mod tests { let _tok = producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1417,8 +1700,8 @@ mod tests { let _tok = producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1436,15 +1719,10 @@ mod tests { #[test] fn test_send_chain_single_segment_writer_rejects_auto_split_chain() { let ring = make_ring(16); - let layout = ring.layout(); - let mem = ring.mem(); - let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; - let pool = TestPool::new_with_max_alloc_len(pool_base, 0x8000, 4); - let notifier = TestNotifier::new(); - let producer = VirtqProducer::new(layout, mem, notifier, pool); + let (producer, _consumer) = make_slot_producer(&ring, 4); let mut se = producer.chain().readable(8).build().unwrap(); - assert_eq!(se.segment_count(), 2); + assert_eq!(se.rd_desc_count(), 2); assert!(matches!( se.with_seg(2, |_| Ok::(0)), Err(VirtqError::NoPayloadSegment) @@ -1588,12 +1866,12 @@ mod tests { assert_eq!(notifier.notification_count(), initial_count + 1); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"first"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"first"); + consumer.complete(recv, reply).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"second"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"second"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1632,6 +1910,23 @@ 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(); + } + #[test] fn test_write_only_round_trip() { let ring = make_ring(16); @@ -1642,11 +1937,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"); } @@ -1668,9 +1963,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)); @@ -1686,10 +1981,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"); } @@ -1715,10 +2010,10 @@ mod tests { se.write_all(b"request data").unwrap(); producer.submit(se).unwrap(); - let (_recv, reply) = poll_received(&mut consumer); + let (recv, reply) = poll_received(&mut consumer); if let ReplyChain::Writable(mut wc) = reply { wc.write_all(b"response data").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } else { panic!("expected Writable"); } @@ -1764,52 +2059,4 @@ mod tests { )); assert_eq!(producer.inner.num_inflight(), 1); } - - #[test] - fn test_virtq_producer_reset() { - let ring = make_ring(16); - let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); - - // Submit and complete a round trip - let mut se = producer.chain().readable(32).writable(64).build().unwrap(); - se.write_all(b"hello").unwrap(); - producer.submit(se).unwrap(); - - let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().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()); - } } diff --git a/src/hyperlight_common/src/virtq/ring.rs b/src/hyperlight_common/src/virtq/ring.rs index 060c70765..0a01160b4 100644 --- a/src/hyperlight_common/src/virtq/ring.rs +++ b/src/hyperlight_common/src/virtq/ring.rs @@ -74,6 +74,8 @@ limitations under the License. //! - **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}; @@ -101,6 +103,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 @@ -167,6 +189,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. @@ -441,6 +467,19 @@ impl RingCursor { } } +/// Local [`RingConsumer`] state needed to undo a sequence of polls. +/// +/// The checkpoint is valid until a polled chain is completed. +#[derive(Clone, Copy)] +pub struct Checkpoint { + /// Position of the next available chain. + avail_cursor: RingCursor, + /// Position of the next completion. + used_cursor: RingCursor, + /// Number of descriptors awaiting completion. + num_inflight: usize, +} + /// Producer (driver) side of a packed virtqueue. /// /// The producer submits buffer chains for the device to process and polls @@ -905,45 +944,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(()) } } @@ -1190,6 +1222,62 @@ impl RingConsumer { Ok(flags.is_avail(self.avail_cursor.wrap())) } + /// Capture the local state needed to undo subsequent polls. + pub fn poll_checkpoint(&self) -> Checkpoint { + Checkpoint { + avail_cursor: self.avail_cursor, + used_cursor: self.used_cursor, + num_inflight: self.num_inflight, + } + } + + /// Undo every chain in `ids` polled after `cp`. + /// + /// None of the chains may have been completed. On success, the next poll + /// observes the first chain again. + /// + /// # Errors + /// + /// Returns [`RingError::InvalidState`] when the IDs, cursors, or inflight + /// descriptor count do not match the checkpoint. + pub fn rollback_polls(&mut self, cp: Checkpoint, ids: &[u16]) -> Result<(), RingError> { + let desc_count = ids.iter().try_fold(0usize, |count, id| { + let chain_len = self + .id_num + .get(*id as usize) + .copied() + .filter(|len| *len != 0) + .ok_or(RingError::InvalidState)?; + + count + .checked_add(chain_len as usize) + .ok_or(RingError::InvalidState) + })?; + + let expected_inflight = cp + .num_inflight + .checked_add(desc_count) + .ok_or(RingError::InvalidState)?; + + let mut expected_cursor = cp.avail_cursor; + expected_cursor.advance_by(u16::try_from(desc_count).map_err(|_| RingError::InvalidState)?); + + if self.avail_cursor != expected_cursor + || self.used_cursor != cp.used_cursor + || self.num_inflight != expected_inflight + { + return Err(RingError::InvalidState); + } + + for id in ids { + self.id_num[*id as usize] = 0; + } + + self.avail_cursor = cp.avail_cursor; + self.num_inflight = cp.num_inflight; + Ok(()) + } + /// Submit a used descriptor and return whether to notify the driver. pub fn submit_used_with_notify( &mut self, @@ -1322,14 +1410,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(()) } } @@ -1402,10 +1503,11 @@ impl From<&Descriptor> for BufferElement { #[cfg(test)] pub(crate) mod tests { use alloc::sync::Arc; + use alloc::vec::Vec; use core::cell::UnsafeCell; use core::num::NonZeroU16; use core::ptr; - use core::sync::atomic::{AtomicU16, Ordering}; + use core::sync::atomic::{AtomicU16, AtomicUsize, Ordering}; use bytemuck::{Pod, Zeroable}; @@ -1515,6 +1617,77 @@ pub(crate) mod tests { } } + #[derive(Clone)] + struct FailingWriteMem { + inner: TestMem, + fail_at: Arc, + writes: Arc, + } + + impl FailingWriteMem { + fn new(inner: TestMem) -> Self { + Self { + inner, + fail_at: Arc::new(AtomicUsize::new(usize::MAX)), + writes: Arc::new(AtomicUsize::new(0)), + } + } + + fn fail_at(&self, write: usize) { + self.writes.store(0, Ordering::Relaxed); + self.fail_at.store(write, Ordering::Relaxed); + } + + fn allow_writes(&self) { + self.writes.store(0, Ordering::Relaxed); + self.fail_at.store(usize::MAX, Ordering::Relaxed); + } + + fn check_write(&self) -> Result<(), ()> { + let write = self.writes.fetch_add(1, Ordering::Relaxed); + if write == self.fail_at.load(Ordering::Relaxed) { + Err(()) + } else { + Ok(()) + } + } + } + + // SAFETY: FailingWriteMem delegates to TestMem and only injects errors + // before writes. + unsafe impl MemOps for FailingWriteMem { + type Error = (); + + fn read(&self, addr: u64, dst: &mut [u8]) -> Result<(), Self::Error> { + self.inner.read(addr, dst).unwrap(); + Ok(()) + } + + fn write(&self, addr: u64, src: &[u8]) -> Result<(), Self::Error> { + self.check_write()?; + self.inner.write(addr, src).unwrap(); + Ok(()) + } + + fn load_acquire(&self, addr: u64) -> Result { + Ok(self.inner.load_acquire(addr).unwrap()) + } + + fn store_release(&self, addr: u64, val: u16) -> Result<(), Self::Error> { + self.check_write()?; + self.inner.store_release(addr, val).unwrap(); + Ok(()) + } + + unsafe fn as_slice(&self, addr: u64, len: usize) -> Result<&[u8], Self::Error> { + Ok(unsafe { self.inner.as_slice(addr, len) }.unwrap()) + } + + unsafe fn as_mut_slice(&self, addr: u64, len: usize) -> Result<&mut [u8], Self::Error> { + Ok(unsafe { self.inner.as_mut_slice(addr, len) }.unwrap()) + } + } + /// Owns the descriptor table and event suppression structures pub struct OwnedRing { mem: TestMem, @@ -3227,7 +3400,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); @@ -3248,7 +3421,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); @@ -3258,6 +3431,50 @@ pub(crate) mod tests { } } + #[test] + fn test_ring_producer_failed_reset_preserves_local_state() { + let ring = make_ring(4); + let mem = FailingWriteMem::new(ring.mem()); + let mut producer = RingProducer::new(ring.layout(), mem.clone()); + + producer.submit_one(0x1000, 64, false).unwrap(); + producer.submit_one(0x2000, 128, true).unwrap(); + + let avail_cursor = producer.avail_cursor; + let used_cursor = producer.used_cursor; + let num_free = producer.num_free; + let id_free = producer.id_free.clone(); + let id_num = producer.id_num.clone(); + let event_flags_shadow = producer.event_flags_shadow; + + mem.fail_at(1); + assert!(matches!( + producer.reset(), + Err(RingError::MemError { + op: MemOp::WriteDesc, + .. + }) + )); + + assert_eq!(producer.avail_cursor, avail_cursor); + assert_eq!(producer.used_cursor, used_cursor); + assert_eq!(producer.num_free, num_free); + assert_eq!(producer.id_free, id_free); + assert_eq!(producer.id_num, id_num); + assert_eq!(producer.event_flags_shadow, event_flags_shadow); + + mem.allow_writes(); + producer.reset().unwrap(); + + let fresh = RingProducer::new(ring.layout(), mem); + assert_eq!(producer.avail_cursor, fresh.avail_cursor); + assert_eq!(producer.used_cursor, fresh.used_cursor); + assert_eq!(producer.num_free, fresh.num_free); + assert_eq!(producer.id_free, fresh.id_free); + assert_eq!(producer.id_num, fresh.id_num); + assert_eq!(producer.event_flags_shadow, fresh.event_flags_shadow); + } + #[test] fn test_ring_consumer_reset_matches_new() { let ring = make_ring(8); @@ -3273,7 +3490,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); @@ -3295,101 +3512,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); @@ -4223,193 +4349,4 @@ mod virtio_villain { } #[cfg(test)] -mod fuzz { - use quickcheck::{Arbitrary, Gen, QuickCheck}; - - use super::tests::{OwnedRing, make_consumer, make_producer}; - use super::*; - - const MAX_RING: usize = 64; - const MAX_OPS: usize = 128; - const MAX_CHAIN_LEN: usize = 8; - - #[allow(clippy::large_enum_variant)] - #[derive(Clone, Debug)] - enum Op { - /// submit one chain - Submit(BufferChain), - /// poll up to N chains - PollAvail(u8), - /// driver reclaims up to N completions - PollUsed(u8), - /// complete one previously polled chain - CompleteOne, - } - - impl Arbitrary for Op { - fn arbitrary(g: &mut Gen) -> Self { - let choice = u8::arbitrary(g) % 4; - match choice { - 0 => Op::Submit(BufferChain::arbitrary(g)), - 1 => Op::PollAvail(u8::arbitrary(g) % 8 + 1), - 2 => Op::PollUsed(u8::arbitrary(g) % 8 + 1), - 3 => Op::CompleteOne, - _ => unreachable!(), - } - } - } - - #[derive(Clone, Debug)] - struct Scenario { - table_size: usize, - ops: Vec, - } - - impl Arbitrary for Scenario { - fn arbitrary(g: &mut Gen) -> Self { - let table_size = (usize::arbitrary(g) % MAX_RING + 1).next_power_of_two(); - let num_ops = usize::arbitrary(g) % MAX_OPS + 1; - - let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); - Scenario { table_size, ops } - } - } - - impl Arbitrary for BufferElement { - fn arbitrary(g: &mut Gen) -> Self { - let addr = u64::arbitrary(g); - let len = u32::arbitrary(g); - let writable = bool::arbitrary(g); - - BufferElement { - addr, - len, - writable, - } - } - } - - impl Arbitrary for BufferChain { - fn arbitrary(g: &mut Gen) -> Self { - let chain_len = usize::arbitrary(g) % MAX_CHAIN_LEN + 1; - - let mut elems = vec![BufferElement::zeroed(); chain_len]; - let mut readables = 0; - let mut writables = 0; - - for _ in 0..chain_len { - let elem = BufferElement::arbitrary(g); - if elem.writable { - elems[chain_len - 1 - writables] = elem; - writables += 1; - } else { - elems[readables] = elem; - readables += 1; - } - } - - BufferChain { - elems: elems.into(), - split: readables, - } - } - } - - fn run_scenario(s: Scenario) -> bool { - let ring = OwnedRing::new(s.table_size); - let mut producer = make_producer(&ring); - let mut consumer = make_consumer(&ring); - - // Order logs - let mut dev_order: Vec = Vec::new(); - let mut drv_order: Vec = Vec::new(); - - // Device-tracked polled-but-not-completed IDs - let mut dev_ready: Vec<(u16, u32)> = Vec::new(); - - for op in &s.ops { - match op { - Op::Submit(chain) => { - // Submit only if space; otherwise skip - let _ = producer.submit_available(chain); - } - Op::PollAvail(n) => { - for _ in 0..*n { - if let Ok((id, chain)) = consumer.poll_available() { - dev_ready.push((id, chain.len() as u32)); - } else { - break; - } - } - } - Op::PollUsed(n) => { - for _ in 0..*n { - match producer.poll_used() { - Ok(u) => { - drv_order.push(u.id); - if producer.id_num[u.id as usize] != 0 { - return false; - } - if !producer.id_free.contains(&u.id) { - return false; - } - } - Err(RingError::WouldBlock) => break, - Err(_) => return false, - } - } - } - Op::CompleteOne => { - if let Some((id, len)) = dev_ready.pop() { - if consumer.submit_used(id, len).is_err() { - return false; - } - - dev_order.push(id); - } - } - } - - // assert invariants after each op - let outstanding: u16 = producer.id_num.iter().copied().sum(); - if outstanding as usize + producer.num_free != ring.len() { - return false; - } - - for id in producer.id_free.iter() { - if producer.id_num[*id as usize] != 0 { - return false; - } - } - } - - // Drain remaining completions and reclaims - while let Some((id, len)) = dev_ready.pop() { - if consumer.submit_used(id, len).is_err() { - return false; - } - } - - loop { - match producer.poll_used() { - Ok(u) => drv_order.push(u.id), - Err(RingError::WouldBlock) => break, - Err(_) => return false, - } - } - - true - } - - #[test] - fn prop_interleaved_with_order_verification() { - #[cfg(miri)] - let tests = 1; - #[cfg(not(miri))] - let tests = 100; - - QuickCheck::new() - .tests(tests) - .quickcheck(run_scenario as fn(Scenario) -> bool); - } -} +mod fuzz; diff --git a/src/hyperlight_common/src/virtq/ring/canonical.rs b/src/hyperlight_common/src/virtq/ring/canonical.rs new file mode 100644 index 000000000..684ecf68b --- /dev/null +++ b/src/hyperlight_common/src/virtq/ring/canonical.rs @@ -0,0 +1,611 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + */ + +//! Canonical packed virtqueue images. +//! +//! A canonical image starts at the initial wrap round. Available descriptors +//! occupy a complete prefix, unused descriptors are zero, and both event +//! structures are enabled at offset zero. This form can be restored as bytes +//! and validated before either peer resumes. +//! +//! Ring resets normalize the structures owned by each peer. Buffer range +//! policy remains with the integration through the validator callback. + +use alloc::vec::Vec; + +use bytemuck::Zeroable; +use fixedbitset::FixedBitSet; +use smallvec::SmallVec; +use thiserror::Error; + +use super::super::desc::{DescFlags, DescTable, Descriptor}; +use super::super::event::{EventFlags, EventSuppression}; +use super::super::{Layout, MemOps}; +use super::{BufferChain, BufferElement, MemOp, RingError}; + +/// Why a descriptor does not belong to a canonical packed-ring image. +#[derive(Error, Debug, Copy, Clone, PartialEq, Eq)] +pub enum DescError { + /// An unused descriptor was not completely zeroed. + #[error("unused descriptor is not zeroed")] + ExpectedZero, + /// The raw descriptor contains unsupported or reserved flag bits. + #[error("descriptor contains unknown flags")] + UnknownFlags, + /// The descriptor is not available in the initial packed-ring wrap round. + #[error("descriptor is not initially available")] + NotAvailable, + /// Indirect descriptor tables are unsupported. + #[error("indirect descriptor is unsupported")] + Indirect, + /// A chain's NEXT flag extends beyond the available descriptor prefix. + #[error("chain continues beyond the available descriptor prefix")] + ChainContinues, + /// A chain ID is outside the descriptor-table bounds. + #[error("descriptor ID is out of range")] + IdOutOfRange, + /// A tail descriptor does not carry its head descriptor's ID. + #[error("descriptor ID differs within a chain")] + IdMismatch, + /// Two available chains use the same descriptor ID. + #[error("descriptor ID is already used by another chain")] + DuplicateId, + /// A readable descriptor follows a writable descriptor. + #[error("readable descriptor follows a writable descriptor")] + ReadableAfterWritable, +} + +/// Validation failure for a canonical packed-ring image. +#[derive(Error, Debug)] +pub enum ImageError { + /// Reading the shared ring image failed. + #[error(transparent)] + Ring(#[from] RingError), + /// The caller supplied an impossible available-descriptor prefix length. + #[error("available descriptor count {available} exceeds ring capacity {capacity}")] + DescCount { + /// Number of descriptors expected to be available. + available: usize, + /// Descriptor-table capacity. + capacity: usize, + }, + /// An event-suppression structure is not the canonical enabled value. + #[error("event suppression at address 0x{addr:x} is not canonical")] + Event { + /// Address of the invalid event-suppression structure. + addr: u64, + }, + /// A descriptor violates the canonical packed-ring structure. + #[error("descriptor {index} is not canonical: {reason}")] + Desc { + /// Descriptor-table index. + index: u16, + /// Structural validation failure. + reason: DescError, + }, + /// The caller rejected a descriptor's payload range or attributes. + #[error("descriptor {index} buffer at 0x{addr:x} with length {len} was rejected")] + Buffer { + /// Descriptor-table index. + index: u16, + /// Buffer address from the descriptor. + addr: u64, + /// Buffer length from the descriptor. + len: u32, + }, +} + +impl ImageError { + fn desc_count(available: usize, capacity: usize) -> Self { + Self::DescCount { + available, + capacity, + } + } + + fn event(addr: u64) -> Self { + Self::Event { addr } + } + + fn desc(index: u16, reason: DescError) -> Self { + Self::Desc { index, reason } + } + + fn buffer(index: u16, addr: u64, len: u32) -> Self { + Self::Buffer { index, addr, len } + } +} + +/// One available descriptor chain from a validated canonical ring image. +#[derive(Debug, Clone)] +pub struct CanonChain { + id: u16, + inner: BufferChain, +} + +impl CanonChain { + fn new(id: u16, chain: BufferChain) -> Self { + Self { id, inner: chain } + } + + /// Descriptor ID shared by every buffer in the chain. + pub fn id(&self) -> u16 { + self.id + } + + /// Validated buffers in descriptor order. + pub fn buffers(&self) -> &BufferChain { + &self.inner + } + + /// Consume the image metadata and return its buffer chain. + pub fn into_buffers(self) -> BufferChain { + self.inner + } +} + +/// Validate a packed ring while neither peer can modify it. +/// +/// The first `avail_descs` descriptors must form complete available chains +/// beginning at descriptor zero. Every remaining descriptor must be zeroed, +/// and both event-suppression structures must be the canonical enabled value. +/// `validate_buf` supplies integration-specific address, length, and +/// direction bounds without embedding them in the ring implementation. +/// +/// The returned chains preserve descriptor IDs and chain boundaries for +/// cross-checking against producer and pool ownership. +/// +/// # Errors +/// +/// Returns [`ImageError`] if event state is not normalized, descriptor +/// structure is malformed, an unused descriptor is not zero, a buffer is +/// rejected by `validate_buf`, or shared memory cannot be read. +pub fn validate_canon_image( + mem: &M, + layout: Layout, + avail_descs: usize, + mut validate_buf: F, +) -> Result, ImageError> +where + M: MemOps, + F: FnMut(u16, BufferElement) -> bool, +{ + let cap = layout.desc_table_len() as usize; + if avail_descs > cap { + return Err(ImageError::desc_count(avail_descs, cap)); + } + + let canon_evt = EventSuppression::new(0, EventFlags::ENABLE); + for addr in [layout.drv_evt_addr(), layout.dev_evt_addr()] { + let evt = mem + .read_val::(addr) + .map_err(|_| RingError::mem_err(MemOp::ReadEvent, addr))?; + + if evt != canon_evt { + return Err(ImageError::event(addr)); + } + } + + // SAFETY: `Layout` validates the table base, alignment, and descriptor count. + let table = unsafe { DescTable::from_raw_parts(layout.desc_table_addr(), cap) }; + + let mut seen_ids = FixedBitSet::with_capacity(cap); + let mut chains = Vec::new(); + let mut pos = 0usize; + + while pos < avail_descs { + let head_idx = u16::try_from(pos).map_err(|_| RingError::InvalidState)?; + let (head, _) = read_canon_avail_desc(mem, &table, head_idx)?; + let id_idx = head.id as usize; + if id_idx >= cap { + return Err(ImageError::desc(head_idx, DescError::IdOutOfRange)); + } + + if seen_ids.contains(id_idx) { + return Err(ImageError::desc(head_idx, DescError::DuplicateId)); + } + + seen_ids.insert(id_idx); + + let mut elems = SmallVec::<[BufferElement; 16]>::new(); + let mut split = 0usize; + + loop { + let idx = u16::try_from(pos).map_err(|_| RingError::InvalidState)?; + let (desc, flags) = read_canon_avail_desc(mem, &table, idx)?; + if desc.id != head.id { + return Err(ImageError::desc(idx, DescError::IdMismatch)); + } + + let elem = BufferElement::from(&desc); + if !elem.writable && split != elems.len() { + return Err(ImageError::desc(idx, DescError::ReadableAfterWritable)); + } + + split += usize::from(!elem.writable); + + if !validate_buf(idx, elem) { + return Err(ImageError::buffer(idx, elem.addr, elem.len)); + } + + elems.push(elem); + pos += 1; + + if !flags.contains(DescFlags::NEXT) { + break; + } + if pos >= avail_descs { + return Err(ImageError::desc(idx, DescError::ChainContinues)); + } + } + + let canon = CanonChain::new(head.id, BufferChain { elems, split }); + chains.push(canon); + } + + let empty = Descriptor::zeroed(); + for pos in avail_descs..cap { + let idx = u16::try_from(pos).map_err(|_| RingError::InvalidState)?; + let addr = table.desc_addr(idx).ok_or(RingError::InvalidState)?; + + let desc = mem + .read_val::(addr) + .map_err(|_| RingError::mem_err(MemOp::ReadDesc, addr))?; + + if desc != empty { + return Err(ImageError::desc(idx, DescError::ExpectedZero)); + } + } + + Ok(chains) +} + +fn read_canon_avail_desc( + mem: &M, + table: &DescTable, + idx: u16, +) -> Result<(Descriptor, DescFlags), ImageError> { + let addr = table.desc_addr(idx).ok_or(RingError::InvalidState)?; + let desc = mem + .read_val::(addr) + .map_err(|_| RingError::mem_err(MemOp::ReadDesc, addr))?; + + let flags = DescFlags::from_bits(desc.flags) + .ok_or_else(|| ImageError::desc(idx, DescError::UnknownFlags))?; + + if flags.contains(DescFlags::INDIRECT) { + return Err(ImageError::desc(idx, DescError::Indirect)); + } + if !flags.is_avail(true) { + return Err(ImageError::desc(idx, DescError::NotAvailable)); + } + + Ok((desc, flags)) +} + +#[cfg(test)] +mod tests { + use super::super::BufferChainBuilder; + use super::super::tests::{OwnedRing, make_consumer, make_producer, make_ring}; + use super::*; + + fn writable_chain(base: u64, lengths: &[u32]) -> BufferChain { + BufferChainBuilder::new() + .writables(lengths.iter().scan(base, |addr, &len| { + let element = BufferElement { + addr: *addr, + len, + writable: true, + }; + *addr += len as u64; + Some(element) + })) + .build() + .unwrap() + } + + fn validate_all(ring: &OwnedRing, avail_descs: usize) -> Result, ImageError> { + validate_canon_image(&ring.mem(), ring.layout(), avail_descs, |_, _| true) + } + + #[test] + fn canon_reset_normalizes_empty_image() { + let ring = make_ring(8); + let mut producer = make_producer(&ring); + let mut consumer = make_consumer(&ring); + + producer.submit_one(0x1000, 64, true).unwrap(); + producer.enable_used_notifications_desc(3, false).unwrap(); + consumer.enable_avail_notifications_desc(5, false).unwrap(); + + producer.reset().unwrap(); + consumer.reset().unwrap(); + + assert!(validate_all(&ring, 0).unwrap().is_empty()); + assert_eq!( + ring.read_driver_event(), + EventSuppression::new(0, EventFlags::ENABLE) + ); + assert_eq!( + ring.read_device_event(), + EventSuppression::new(0, EventFlags::ENABLE) + ); + for index in 0..ring.len() as u16 { + assert_eq!(ring.read_desc(index), Descriptor::zeroed()); + } + } + + #[test] + fn canon_multi_desc_refill_is_visible_to_fresh_consumer() { + let ring = make_ring(8); + let mut producer = make_producer(&ring); + let mut consumer = make_consumer(&ring); + + producer.reset().unwrap(); + consumer.reset().unwrap(); + + let chains = [ + writable_chain(0x1000, &[64, 128, 256]), + writable_chain(0x2000, &[512, 1024]), + writable_chain(0x4000, &[64, 64, 64]), + ]; + let mut expected = Vec::new(); + for chain in &chains { + let id = producer.submit_available(chain).unwrap(); + expected.push((id, chain.len())); + } + + assert_eq!(producer.num_free(), 0); + assert_eq!(producer.avail_cursor().head(), 0); + assert!(!producer.avail_cursor().wrap()); + + let image = validate_canon_image(&ring.mem(), ring.layout(), ring.len(), |_, elem| { + elem.writable + && elem.len > 0 + && elem + .addr + .checked_add(elem.len as u64) + .is_some_and(|end| end <= 0x5000) + }) + .unwrap(); + + assert_eq!(image.len(), expected.len()); + for (validated, (id, len)) in image.iter().zip(&expected) { + assert_eq!(validated.id(), *id); + assert_eq!(validated.buffers().len(), *len); + assert!(validated.buffers().elems().iter().all(|elem| elem.writable)); + assert_eq!(producer.id_num[*id as usize] as usize, *len); + } + + let mut fresh = make_consumer(&ring); + for (expected_id, expected_len) in expected { + let (id, chain) = fresh.poll_available().unwrap(); + assert_eq!(id, expected_id); + assert_eq!(chain.len(), expected_len); + } + assert!(matches!(fresh.poll_available(), Err(RingError::WouldBlock))); + } + + #[test] + fn canon_image_rejects_invalid_ids() { + let duplicate_ring = make_ring(4); + let mut producer = make_producer(&duplicate_ring); + producer.submit_one(0x1000, 64, true).unwrap(); + producer.submit_one(0x2000, 64, true).unwrap(); + + let head_id = duplicate_ring.read_desc(0).id; + let mut duplicate = duplicate_ring.read_desc(1); + duplicate.id = head_id; + duplicate_ring.write_desc(1, duplicate); + + assert!(matches!( + validate_all(&duplicate_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::DuplicateId, + }) + )); + + let mismatch_ring = make_ring(4); + let mut producer = make_producer(&mismatch_ring); + producer + .submit_available(&writable_chain(0x3000, &[64, 64])) + .unwrap(); + + let mut tail = mismatch_ring.read_desc(1); + tail.id = tail.id.wrapping_sub(1); + mismatch_ring.write_desc(1, tail); + + assert!(matches!( + validate_all(&mismatch_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::IdMismatch, + }) + )); + + let range_ring = make_ring(4); + let mut producer = make_producer(&range_ring); + producer.submit_one(0x4000, 64, true).unwrap(); + + let mut out_of_range = range_ring.read_desc(0); + out_of_range.id = range_ring.len() as u16; + range_ring.write_desc(0, out_of_range); + + assert!(matches!( + validate_all(&range_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::IdOutOfRange, + }) + )); + } + + #[test] + fn canon_image_rejects_invalid_flags_and_wrap_state() { + let unknown_ring = make_ring(4); + let mut producer = make_producer(&unknown_ring); + producer.submit_one(0x1000, 64, true).unwrap(); + + let mut unknown = unknown_ring.read_desc(0); + unknown.flags |= 1 << 3; + unknown_ring.write_desc(0, unknown); + + assert!(matches!( + validate_all(&unknown_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::UnknownFlags, + }) + )); + + let used_ring = make_ring(4); + let mut producer = make_producer(&used_ring); + producer.submit_one(0x2000, 64, true).unwrap(); + + let mut used = used_ring.read_desc(0); + used.mark_used(true); + used_ring.write_desc(0, used); + + assert!(matches!( + validate_all(&used_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::NotAvailable, + }) + )); + + let indirect_ring = make_ring(4); + let mut producer = make_producer(&indirect_ring); + producer.submit_one(0x3000, 64, true).unwrap(); + + let mut indirect = indirect_ring.read_desc(0); + indirect.flags |= DescFlags::INDIRECT.bits(); + indirect_ring.write_desc(0, indirect); + + assert!(matches!( + validate_all(&indirect_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::Indirect, + }) + )); + } + + #[test] + fn canon_image_rejects_malformed_chain_and_unused_desc() { + let chain_ring = make_ring(4); + let mut producer = make_producer(&chain_ring); + producer + .submit_available(&writable_chain(0x1000, &[64, 64])) + .unwrap(); + + let mut tail = chain_ring.read_desc(1); + tail.flags |= DescFlags::NEXT.bits(); + chain_ring.write_desc(1, tail); + + assert!(matches!( + validate_all(&chain_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::ChainContinues, + }) + )); + + let unused_ring = make_ring(4); + let mut producer = make_producer(&unused_ring); + producer.submit_one(0x2000, 64, true).unwrap(); + unused_ring.write_desc(3, Descriptor::new(0x3000, 64, 0, DescFlags::empty())); + + assert!(matches!( + validate_all(&unused_ring, 1), + Err(ImageError::Desc { + index: 3, + reason: DescError::ExpectedZero, + }) + )); + + let direction_ring = make_ring(4); + let mut producer = make_producer(&direction_ring); + producer + .submit_available(&writable_chain(0x4000, &[64, 64])) + .unwrap(); + + let mut readable_tail = direction_ring.read_desc(1); + readable_tail.flags &= !DescFlags::WRITE.bits(); + direction_ring.write_desc(1, readable_tail); + + assert!(matches!( + validate_all(&direction_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::ReadableAfterWritable, + }) + )); + } + + #[test] + fn canon_image_rejects_noncanon_evt_and_buffer_bounds() { + let evt_ring = make_ring(4); + evt_ring + .mem() + .write_val( + evt_ring.layout().drv_evt_addr(), + EventSuppression::new(1, EventFlags::ENABLE), + ) + .unwrap(); + + match validate_all(&evt_ring, 0) { + Err(ImageError::Event { addr }) => { + let expected = evt_ring.layout().drv_evt_addr(); + assert_eq!(addr, expected); + } + other => unreachable!("unexpected result: {other:?}"), + } + + let bounds = make_ring(4); + let mut producer = make_producer(&bounds); + producer.submit_one(u64::MAX - 15, 32, true).unwrap(); + + let res = validate_canon_image(&bounds.mem(), bounds.layout(), 1, |_, element| { + element + .addr + .checked_add(element.len as u64) + .is_some_and(|end| end <= 0x8000) + }); + + match res { + Err(ImageError::Buffer { index, addr, len }) => { + assert_eq!(index, 0); + assert_eq!(addr, u64::MAX - 15); + assert_eq!(len, 32); + } + other => unreachable!("unexpected result: {other:?}"), + } + } + + #[test] + fn canon_image_rejects_avail_count_over_capacity() { + let ring = make_ring(4); + assert!(matches!( + validate_all(&ring, 5), + Err(ImageError::DescCount { + available: 5, + capacity: 4, + }) + )); + } +} diff --git a/src/hyperlight_common/src/virtq/ring/fuzz.rs b/src/hyperlight_common/src/virtq/ring/fuzz.rs new file mode 100644 index 000000000..cef6d12e0 --- /dev/null +++ b/src/hyperlight_common/src/virtq/ring/fuzz.rs @@ -0,0 +1,204 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use quickcheck::{Arbitrary, Gen, QuickCheck}; + +use super::tests::{OwnedRing, make_consumer, make_producer}; +use super::*; + +const MAX_RING: usize = 64; +const MAX_OPS: usize = 128; +const MAX_CHAIN_LEN: usize = 8; + +#[allow(clippy::large_enum_variant)] +#[derive(Clone, Debug)] +enum Op { + /// submit one chain + Submit(BufferChain), + /// poll up to N chains + PollAvail(u8), + /// driver reclaims up to N completions + PollUsed(u8), + /// complete one previously polled chain + CompleteOne, +} + +impl Arbitrary for Op { + fn arbitrary(g: &mut Gen) -> Self { + let choice = u8::arbitrary(g) % 4; + match choice { + 0 => Op::Submit(BufferChain::arbitrary(g)), + 1 => Op::PollAvail(u8::arbitrary(g) % 8 + 1), + 2 => Op::PollUsed(u8::arbitrary(g) % 8 + 1), + 3 => Op::CompleteOne, + _ => unreachable!(), + } + } +} + +#[derive(Clone, Debug)] +struct Scenario { + table_size: usize, + ops: Vec, +} + +impl Arbitrary for Scenario { + fn arbitrary(g: &mut Gen) -> Self { + let table_size = (usize::arbitrary(g) % MAX_RING + 1).next_power_of_two(); + let num_ops = usize::arbitrary(g) % MAX_OPS + 1; + + let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); + Scenario { table_size, ops } + } +} + +impl Arbitrary for BufferElement { + fn arbitrary(g: &mut Gen) -> Self { + let addr = u64::arbitrary(g); + let len = u32::arbitrary(g); + let writable = bool::arbitrary(g); + + BufferElement { + addr, + len, + writable, + } + } +} + +impl Arbitrary for BufferChain { + fn arbitrary(g: &mut Gen) -> Self { + let chain_len = usize::arbitrary(g) % MAX_CHAIN_LEN + 1; + + let mut elems = vec![BufferElement::zeroed(); chain_len]; + let mut readables = 0; + let mut writables = 0; + + for _ in 0..chain_len { + let elem = BufferElement::arbitrary(g); + if elem.writable { + elems[chain_len - 1 - writables] = elem; + writables += 1; + } else { + elems[readables] = elem; + readables += 1; + } + } + + BufferChain { + elems: elems.into(), + split: readables, + } + } +} + +fn run_scenario(s: Scenario) -> bool { + let ring = OwnedRing::new(s.table_size); + let mut producer = make_producer(&ring); + let mut consumer = make_consumer(&ring); + + // Order logs + let mut dev_order: Vec = Vec::new(); + let mut drv_order: Vec = Vec::new(); + + // Device-tracked polled-but-not-completed IDs + let mut dev_ready: Vec<(u16, u32)> = Vec::new(); + + for op in &s.ops { + match op { + Op::Submit(chain) => { + // Submit only if space; otherwise skip + let _ = producer.submit_available(chain); + } + Op::PollAvail(n) => { + for _ in 0..*n { + if let Ok((id, chain)) = consumer.poll_available() { + dev_ready.push((id, chain.len() as u32)); + } else { + break; + } + } + } + Op::PollUsed(n) => { + for _ in 0..*n { + match producer.poll_used() { + Ok(u) => { + drv_order.push(u.id); + if producer.id_num[u.id as usize] != 0 { + return false; + } + if !producer.id_free.contains(&u.id) { + return false; + } + } + Err(RingError::WouldBlock) => break, + Err(_) => return false, + } + } + } + Op::CompleteOne => { + if let Some((id, len)) = dev_ready.pop() { + if consumer.submit_used(id, len).is_err() { + return false; + } + + dev_order.push(id); + } + } + } + + // assert invariants after each op + let outstanding: u16 = producer.id_num.iter().copied().sum(); + if outstanding as usize + producer.num_free != ring.len() { + return false; + } + + for id in producer.id_free.iter() { + if producer.id_num[*id as usize] != 0 { + return false; + } + } + } + + // Drain remaining completions and reclaims + while let Some((id, len)) = dev_ready.pop() { + if consumer.submit_used(id, len).is_err() { + return false; + } + } + + loop { + match producer.poll_used() { + Ok(u) => drv_order.push(u.id), + Err(RingError::WouldBlock) => break, + Err(_) => return false, + } + } + + true +} + +#[test] +fn prop_interleaved_with_order_verification() { + #[cfg(miri)] + let tests = 1; + #[cfg(not(miri))] + let tests = 100; + + QuickCheck::new() + .tests(tests) + .quickcheck(run_scenario as fn(Scenario) -> bool); +} diff --git a/src/hyperlight_component_util/src/guest.rs b/src/hyperlight_component_util/src/guest.rs index 18e555e3d..686c2de02 100644 --- a/src/hyperlight_component_util/src/guest.rs +++ b/src/hyperlight_component_util/src/guest.rs @@ -209,12 +209,12 @@ fn emit_export_extern_decl<'a, 'b, 'c>( let marshal_result = emit_hl_marshal_result(s, ret.clone(), &ft.result); let trait_path = s.cur_trait_path(); quote! { - fn #n(fc: ::hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall) -> ::hyperlight_guest::error::Result<::alloc::vec::Vec> { + fn #n(fc: ::hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall) -> ::hyperlight_guest::error::Result<::hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue> { ::with_guest_state(|state| { #(#pds)* #(#get_instance)* let #ret = #trait_path::#n(state, #(#pus,)*); - ::core::result::Result::Ok(::hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result::<&[u8]>(&#marshal_result)) + ::core::result::Result::Ok(::hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue::VecBytes(#marshal_result)) }) } ::hyperlight_guest_bin::guest_function::register::register_function( diff --git a/src/hyperlight_component_util/src/hl.rs b/src/hyperlight_component_util/src/hl.rs index a3ffe6f6d..5cc7efc0d 100644 --- a/src/hyperlight_component_util/src/hl.rs +++ b/src/hyperlight_component_util/src/hl.rs @@ -766,7 +766,7 @@ pub fn emit_hl_marshal_param(s: &mut State, id: Ident, pt: &Value) -> TokenStrea /// are no names in it (a unit type) pub fn emit_hl_marshal_result(s: &mut State, id: Ident, rt: &etypes::Result) -> TokenStream { match rt { - None => quote! { ::alloc::vec::Vec::new() }, + None => quote! { ::alloc::vec::Vec::::new() }, Some(vt) => { let toks = emit_hl_marshal_value(s, id, vt); quote! { { #toks } } diff --git a/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs b/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs index d49e3f936..75e4129e0 100644 --- a/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs +++ b/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs @@ -35,13 +35,8 @@ pub unsafe fn alloc_phys_pages(n: u64) -> u64 { prev_base = out(reg) prev_base, ); } - // Set aside two pages at the top of the scratch region for the - // exception stack, shared state, etc - let max_avail = layout::SCRATCH_TOP_GPA - vmem::PAGE_SIZE * 2; - if prev_base - .checked_add(nbytes) - .is_none_or(|xx| xx >= max_avail as u64) - { + let limit = layout::scratch_allocator_limit_gpa(); + if super::allocation_exceeds_limit(prev_base, nbytes, limit) { unsafe { crate::exit::abort_with_code_and_message( &[ErrorCode::MallocFailed as u8], diff --git a/src/hyperlight_guest/src/arch/amd64/prim_alloc.rs b/src/hyperlight_guest/src/arch/amd64/prim_alloc.rs index e1d388c64..cbc1b75ff 100644 --- a/src/hyperlight_guest/src/arch/amd64/prim_alloc.rs +++ b/src/hyperlight_guest/src/arch/amd64/prim_alloc.rs @@ -15,6 +15,7 @@ limitations under the License. */ 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 @@ -22,7 +23,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!( @@ -31,13 +32,8 @@ pub unsafe fn alloc_phys_pages(n: u64) -> u64 { x = inout(reg) x ); } - // Set aside two pages at the top of the scratch region for the - // exception stack, shared state, etc - let max_avail = - hyperlight_common::layout::SCRATCH_TOP_GPA - hyperlight_common::vmem::PAGE_SIZE * 2; - if x.checked_add(nbytes) - .is_none_or(|xx| xx >= max_avail as u64) - { + let limit = layout::scratch_allocator_limit_gpa(); + if super::allocation_exceeds_limit(x, nbytes, limit) { unsafe { crate::exit::abort_with_code_and_message( &[ErrorCode::MallocFailed as u8], diff --git a/src/hyperlight_guest/src/error.rs b/src/hyperlight_guest/src/error.rs index 0a33bce79..681af084c 100644 --- a/src/hyperlight_guest/src/error.rs +++ b/src/hyperlight_guest/src/error.rs @@ -18,7 +18,9 @@ use alloc::format; use alloc::string::{String, ToString as _}; pub use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; +use hyperlight_common::flatbuffer_wrappers::guest_error::GuestError; use hyperlight_common::func::Error as FuncError; +use hyperlight_common::virtq::VirtqError; use {anyhow, serde_json}; pub type Result = core::result::Result; @@ -80,6 +82,24 @@ impl From for HyperlightGuestError { } } +impl From for HyperlightGuestError { + fn from(error: VirtqError) -> Self { + Self { + kind: ErrorCode::GuestError, + message: format!("virtq: {error}"), + } + } +} + +impl From for HyperlightGuestError { + fn from(error: GuestError) -> Self { + Self { + kind: error.code, + message: error.message, + } + } +} + /// Extension trait to add context to `Option` and `Result` types in guest code, /// converting them to `Result`. /// @@ -171,10 +191,10 @@ impl GuestErrorContext for core::result::Result { #[macro_export] macro_rules! bail { ($ec:expr => $($msg:tt)*) => { - return ::core::result::Result::Err($crate::error::HyperlightGuestError::new($ec, ::alloc::format!($($msg)*))); + return ::core::result::Result::Err($crate::error::HyperlightGuestError::new($ec, ::alloc::format!($($msg)*))) }; ($($msg:tt)*) => { - $crate::bail!($crate::error::ErrorCode::GuestError => $($msg)*); + $crate::bail!($crate::error::ErrorCode::GuestError => $($msg)*) }; } diff --git a/src/hyperlight_guest/src/guest_handle/handle.rs b/src/hyperlight_guest/src/guest_handle/handle.rs index 3a22ee26c..9f5bf27fe 100644 --- a/src/hyperlight_guest/src/guest_handle/handle.rs +++ b/src/hyperlight_guest/src/guest_handle/handle.rs @@ -14,17 +14,16 @@ See the License for the specific language governing permissions and limitations under the License. */ +use alloc::format; +use alloc::vec::Vec; + +use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; use hyperlight_common::mem::HyperlightPEB; +use tracing::instrument; + +use crate::error::{HyperlightGuestError, Result}; -/// A guest handle holds the `HyperlightPEB` and enables the guest to perform -/// operations like: -/// - calling host functions, -/// - accessing shared input and output buffers, -/// - writing errors, -/// - etc. -/// -/// Guests are expected to initialize this and store it. For example, you -/// could store it in a global variable. +/// Access to memory regions described by the guest's `HyperlightPEB`. #[derive(Debug, Clone, Copy, Default)] pub struct GuestHandle { peb: Option<*mut HyperlightPEB>, @@ -45,4 +44,29 @@ impl GuestHandle { pub fn peb(&self) -> Option<*mut HyperlightPEB> { self.peb } + + /// Get user memory region as bytes. + #[instrument(skip_all, level = "Trace")] + pub fn read_n_bytes_from_user_memory(&self, num: u64) -> Result> { + let peb_ptr = self.peb().unwrap(); + // SAFETY: GuestHandle is initialized with the PEB provided by the host, + // which remains valid for the guest lifetime. + let init_data = unsafe { (*peb_ptr).init_data }; + + if num > init_data.size { + return Err(HyperlightGuestError::new( + ErrorCode::GuestError, + format!( + "Requested {} bytes from user memory, but only {} bytes are available", + num, init_data.size + ), + )); + } + + // SAFETY: The PEB describes a valid user memory region and num was + // checked against its size. + let bytes = + unsafe { core::slice::from_raw_parts(init_data.ptr as *const u8, num as usize) }; + Ok(bytes.to_vec()) + } } diff --git a/src/hyperlight_guest/src/guest_handle/host_comm.rs b/src/hyperlight_guest/src/guest_handle/host_comm.rs deleted file mode 100644 index c72de8a3f..000000000 --- a/src/hyperlight_guest/src/guest_handle/host_comm.rs +++ /dev/null @@ -1,207 +0,0 @@ -/* -Copyright 2025 The Hyperlight Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -use alloc::format; -use alloc::string::ToString; -use alloc::vec::Vec; - -use flatbuffers::FlatBufferBuilder; -use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType}; -use hyperlight_common::flatbuffer_wrappers::function_types::{ - FunctionCallResult, ParameterValue, ReturnType, ReturnValue, -}; -use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; -use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData; -use hyperlight_common::flatbuffer_wrappers::guest_log_level::LogLevel; -use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity; -use hyperlight_common::outb::OutBAction; -use tracing::instrument; - -use super::handle::GuestHandle; -use crate::error::{HyperlightGuestError, Result}; -use crate::exit::out32; - -impl GuestHandle { - /// Get user memory region as bytes. - #[instrument(skip_all, level = "Trace")] - pub fn read_n_bytes_from_user_memory(&self, num: u64) -> Result> { - let peb_ptr = self.peb().unwrap(); - let user_memory_region_ptr = unsafe { (*peb_ptr).init_data.ptr as *mut u8 }; - let user_memory_region_size = unsafe { (*peb_ptr).init_data.size }; - - if num > user_memory_region_size { - Err(HyperlightGuestError::new( - ErrorCode::GuestError, - format!( - "Requested {} bytes from user memory, but only {} bytes are available", - num, user_memory_region_size - ), - )) - } else { - let user_memory_region_slice = - unsafe { core::slice::from_raw_parts(user_memory_region_ptr, num as usize) }; - let user_memory_region_bytes = user_memory_region_slice.to_vec(); - - Ok(user_memory_region_bytes) - } - } - - /// Get a return value from a host function call. - /// This usually requires a host function to be called first using - /// `call_host_function_internal`. - /// - /// When calling `call_host_function`, this function is called - /// internally to get the return value. - #[instrument(skip_all, level = "Trace")] - pub fn get_host_return_value>(&self) -> Result { - let inner = self - .try_pop_shared_input_data_into::() - .expect("Unable to deserialize a return value from host") - .into_inner(); - - match inner { - Ok(ret) => T::try_from(ret).map_err(|_| { - let expected = core::any::type_name::(); - HyperlightGuestError::new( - ErrorCode::UnsupportedParameterType, - format!("Host return value could not be converted to expected {expected}",), - ) - }), - Err(e) => Err(HyperlightGuestError { - kind: e.code, - message: e.message, - }), - } - } - - pub fn get_host_return_raw(&self) -> Result { - let inner = self - .try_pop_shared_input_data_into::() - .expect("Unable to deserialize a return value from host") - .into_inner(); - - match inner { - Ok(ret) => Ok(ret), - Err(e) => Err(HyperlightGuestError { - kind: e.code, - message: e.message, - }), - } - } - - /// Call a host function without reading its return value from shared mem. - /// This is used by both the Rust and C APIs to reduce code duplication. - /// - /// Note: The function return value must be obtained by calling - /// `get_host_return_value`. - #[instrument(skip_all, level = "Trace")] - pub fn call_host_function_without_returning_result( - &self, - function_name: &str, - parameters: Option>, - return_type: ReturnType, - ) -> Result<()> { - let estimated_capacity = - estimate_flatbuffer_capacity(function_name, parameters.as_deref().unwrap_or(&[])); - - let host_function_call = FunctionCall::new( - function_name.to_string(), - parameters, - FunctionCallType::Host, - return_type, - ); - - let mut builder = FlatBufferBuilder::with_capacity(estimated_capacity); - - let host_function_call_buffer = host_function_call.encode(&mut builder); - self.push_shared_output_data(host_function_call_buffer)?; - - unsafe { - out32(OutBAction::CallFunction as u16, 0); - } - - Ok(()) - } - - /// Call a host function with the given parameters and return type. - /// This function serializes the function call and its parameters, - /// sends it to the host, and then retrieves the return value. - /// - /// The return value is deserialized into the specified type `T`. - #[instrument(skip_all, level = "Info")] - pub fn call_host_function>( - &self, - function_name: &str, - parameters: Option>, - return_type: ReturnType, - ) -> Result { - self.call_host_function_without_returning_result(function_name, parameters, return_type)?; - self.get_host_return_value::() - } - - /// Log a message with the specified log level, source, caller, source file, and line number. - pub fn log_message( - &self, - log_level: LogLevel, - message: &str, - source: &str, - caller: &str, - source_file: &str, - line: u32, - ) { - // Closure to send log message to host - let _send_to_host = || { - let guest_log_data = GuestLogData::new( - message.to_string(), - source.to_string(), - log_level, - caller.to_string(), - source_file.to_string(), - line, - ); - - let bytes: Vec = guest_log_data - .try_into() - .expect("Failed to convert GuestLogData to bytes"); - - self.push_shared_output_data(&bytes) - .expect("Unable to push log data to shared output data"); - - unsafe { - out32(OutBAction::Log as u16, 0); - } - }; - - #[cfg(all(feature = "trace_guest", target_arch = "x86_64"))] - if hyperlight_guest_tracing::is_trace_enabled() { - // If the "trace_guest" feature is enabled and tracing is initialized, log using tracing - tracing::trace!( - event = message, - level = ?log_level, - code.filepath = source, - caller = caller, - source_file = source_file, - code.lineno = line, - ); - } else { - _send_to_host(); - } - #[cfg(not(all(feature = "trace_guest", target_arch = "x86_64")))] - { - _send_to_host(); - } - } -} diff --git a/src/hyperlight_guest/src/guest_handle/io.rs b/src/hyperlight_guest/src/guest_handle/io.rs deleted file mode 100644 index 46c1d68f6..000000000 --- a/src/hyperlight_guest/src/guest_handle/io.rs +++ /dev/null @@ -1,150 +0,0 @@ -/* -Copyright 2025 The Hyperlight Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -use alloc::format; -use alloc::string::ToString; -use core::any::type_name; -use core::slice::from_raw_parts_mut; - -use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; -use tracing::instrument; - -use super::handle::GuestHandle; -use crate::error::{HyperlightGuestError, Result}; - -impl GuestHandle { - /// Pops the top element from the shared input data buffer and returns it as a T - #[instrument(skip_all, level = "Trace")] - pub fn try_pop_shared_input_data_into(&self) -> Result - where - T: for<'a> TryFrom<&'a [u8]>, - { - let peb_ptr = self.peb().unwrap(); - let input_stack_size = unsafe { (*peb_ptr).input_stack.size as usize }; - let input_stack_ptr = unsafe { (*peb_ptr).input_stack.ptr as *mut u8 }; - - let idb = unsafe { from_raw_parts_mut(input_stack_ptr, input_stack_size) }; - - if idb.is_empty() { - return Err(HyperlightGuestError::new( - ErrorCode::GuestError, - "Got a 0-size buffer in pop_shared_input_data_into".to_string(), - )); - } - - // get relative offset to next free address - let stack_ptr_rel: u64 = - u64::from_le_bytes(idb[..8].try_into().expect("Shared input buffer too small")); - - if stack_ptr_rel as usize > input_stack_size || stack_ptr_rel < 16 { - return Err(HyperlightGuestError::new( - ErrorCode::GuestError, - format!( - "Invalid stack pointer: {} in pop_shared_input_data_into", - stack_ptr_rel - ), - )); - } - - // go back 8 bytes and read. This is the offset to the element on top of stack - let last_element_offset_rel = u64::from_le_bytes( - idb[stack_ptr_rel as usize - 8..stack_ptr_rel as usize] - .try_into() - .expect("Invalid stack pointer in pop_shared_input_data_into"), - ); - - let buffer = &idb[last_element_offset_rel as usize..]; - - // convert the buffer to T - let type_t = match T::try_from(buffer) { - Ok(t) => Ok(t), - Err(_e) => { - return Err(HyperlightGuestError::new( - ErrorCode::GuestError, - format!("Unable to convert buffer to {}", type_name::()), - )); - } - }; - - // update the stack pointer to point to the element we just popped of since that is now free - idb[..8].copy_from_slice(&last_element_offset_rel.to_le_bytes()); - - // zero out popped off buffer - idb[last_element_offset_rel as usize..stack_ptr_rel as usize].fill(0); - - type_t - } - - /// Pushes the given data onto the shared output data buffer. - pub fn push_shared_output_data(&self, data: &[u8]) -> Result<()> { - let peb_ptr = self.peb().unwrap(); - let output_stack_size = unsafe { (*peb_ptr).output_stack.size as usize }; - let output_stack_ptr = unsafe { (*peb_ptr).output_stack.ptr as *mut u8 }; - - let odb = unsafe { from_raw_parts_mut(output_stack_ptr, output_stack_size) }; - - if odb.is_empty() { - return Err(HyperlightGuestError::new( - ErrorCode::GuestError, - "Got a 0-size buffer in push_shared_output_data".to_string(), - )); - } - - // get offset to next free address on the stack - let stack_ptr_rel: u64 = - u64::from_le_bytes(odb[..8].try_into().expect("Shared output buffer too small")); - - // check if the stack pointer is within the bounds of the buffer. - // It can be equal to the size, but never greater - // It can never be less than 8. An empty buffer's stack pointer is 8 - if stack_ptr_rel as usize > output_stack_size || stack_ptr_rel < 8 { - return Err(HyperlightGuestError::new( - ErrorCode::GuestError, - format!( - "Invalid stack pointer: {} in push_shared_output_data", - stack_ptr_rel - ), - )); - } - - // check if there is enough space in the buffer - let size_required = data.len() + 8; // the data plus the pointer pointing to the data - let size_available = output_stack_size - stack_ptr_rel as usize; - if size_required > size_available { - return Err(HyperlightGuestError::new( - ErrorCode::GuestError, - format!( - "Not enough space in shared output buffer. Required: {}, Available: {}", - size_required, size_available - ), - )); - } - - // write the actual data - odb[stack_ptr_rel as usize..stack_ptr_rel as usize + data.len()].copy_from_slice(data); - - // write the offset to the newly written data, to the top of the stack - let bytes: [u8; 8] = stack_ptr_rel.to_le_bytes(); - odb[stack_ptr_rel as usize + data.len()..stack_ptr_rel as usize + data.len() + 8] - .copy_from_slice(&bytes); - - // update stack pointer to point to next free address - let new_stack_ptr_rel: u64 = (stack_ptr_rel as usize + data.len() + 8) as u64; - odb[0..8].copy_from_slice(&(new_stack_ptr_rel).to_le_bytes()); - - Ok(()) - } -} diff --git a/src/hyperlight_guest/src/layout.rs b/src/hyperlight_guest/src/layout.rs index 6d132ae7c..f7e283905 100644 --- a/src/hyperlight_guest/src/layout.rs +++ b/src/hyperlight_guest/src/layout.rs @@ -19,20 +19,42 @@ limitations under the License. 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 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 19e5ac5f2..9e8335a3c 100644 --- a/src/hyperlight_guest/src/lib.rs +++ b/src/hyperlight_guest/src/lib.rs @@ -25,10 +25,9 @@ pub mod error; pub mod exit; pub mod layout; pub mod prim_alloc; +pub mod transport; pub mod types; pub mod guest_handle { pub mod handle; - pub mod host_comm; - pub mod io; } diff --git a/src/hyperlight_guest/src/prim_alloc.rs b/src/hyperlight_guest/src/prim_alloc.rs index ed143dfb3..6ac6bc5cb 100644 --- a/src/hyperlight_guest/src/prim_alloc.rs +++ b/src/hyperlight_guest/src/prim_alloc.rs @@ -35,3 +35,20 @@ mod arch; /// latter cannot be perfectly satisfied due to the lack of per-byte /// atomic memcpy in the host. pub use arch::alloc_phys_pages; + +#[inline] +fn allocation_exceeds_limit(base: u64, len: u64, limit: u64) -> bool { + base.checked_add(len).is_none_or(|end| end > limit) +} + +#[cfg(test)] +mod tests { + use super::allocation_exceeds_limit; + + #[test] + fn allocation_may_end_at_limit() { + assert!(!allocation_exceeds_limit(0x1000, 0x2000, 0x3000)); + assert!(allocation_exceeds_limit(0x1000, 0x2001, 0x3000)); + assert!(allocation_exceeds_limit(u64::MAX, 1, u64::MAX)); + } +} diff --git a/src/hyperlight_guest/src/transport/codec.rs b/src/hyperlight_guest/src/transport/codec.rs new file mode 100644 index 000000000..6d2995e40 --- /dev/null +++ b/src/hyperlight_guest/src/transport/codec.rs @@ -0,0 +1,239 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +//! Guest-side virtqueue message decoding. + +use alloc::vec::Vec; + +use hyperlight_common::flatbuffer_wrappers::ExternalValueSource; +use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall; +use hyperlight_common::flatbuffer_wrappers::function_types::{Bytes, FunctionCallResult}; +use hyperlight_common::transport::{ + MsgHeader, MsgKind, SIZE_PREFIX_LEN, size_prefix_payload_len, size_prefixed_len, +}; +use hyperlight_common::virtq::Segments; + +use crate::bail; +use crate::error::{GuestErrorContext, Result}; + +/// Decode one H2G guest-function request payload. +/// +/// Chunked external values retain their H2G slot owners until their final +/// [`Bytes`] clone drops. +pub(super) fn decode_request(cid: u32, segments: Segments) -> Result<(u32, FunctionCall)> { + if cid == 0 { + bail!("Guest function request has correlation ID zero"); + } + + let (control, mut external_values) = decode_payload(segments)?; + let call = FunctionCall::decode(&control, &mut external_values) + .with_context(|| "failed to decode guest function request")?; + + Ok((cid, call)) +} + +/// Decode one G2H host-function response. +/// +/// Contiguous byte values are flattened into `Vec`. Chunked values retain +/// their transport-backed [`Bytes`] owners. +pub(super) fn decode_response(segments: Segments, cid: u32) -> Result { + let (header, payload) = split_header(segments)?; + if header.msg_kind() != Ok(MsgKind::Response) { + bail!("Host function response has an invalid message kind"); + } + + if header.cid != cid { + bail!("Host function response correlation ID mismatch"); + } + + let (control, mut external_values) = decode_payload(payload)?; + FunctionCallResult::decode(&control, &mut external_values) + .with_context(|| "failed to decode host function response") +} + +fn split_header(mut segments: Segments) -> Result<(MsgHeader, Segments)> { + let header = segments + .split_to(MsgHeader::SIZE) + .context("virtqueue message is missing its header")? + .into_bytes(); + + let Some(header) = MsgHeader::from_bytes(&header) else { + bail!("Virtqueue message has an invalid header"); + }; + + if usize::try_from(header.payload_len).ok() != Some(segments.len()) { + bail!("Virtqueue message payload length mismatch"); + } + + Ok((header, segments)) +} + +/// Copy FlatBuffer control data while retaining external payload owners. +fn decode_payload(mut segments: Segments) -> Result<(Vec, SegmentSource)> { + let prefix = segments + .split_to(SIZE_PREFIX_LEN) + .context("virtqueue message is missing its size prefix")? + .into_bytes(); + + let payload_len = + size_prefix_payload_len(&prefix).context("virtqueue message has an invalid prefix")?; + + let payload = segments + .split_to(payload_len) + .context("virtqueue message control data is truncated")?; + + let control_len = + size_prefixed_len(payload_len).context("virtqueue message control length overflow")?; + + let mut control = Vec::with_capacity(control_len); + control.extend_from_slice(&prefix); + + for segment in payload.iter() { + control.extend_from_slice(segment); + } + + Ok((control, SegmentSource::new(segments))) +} + +/// Supplies complete logical external values from transport segments. +struct SegmentSource { + segments: Segments, +} + +impl SegmentSource { + fn new(segments: Segments) -> Self { + Self { segments } + } + + fn take(&mut self, length: usize) -> anyhow::Result { + self.segments.split_to(length).ok_or_else(|| { + anyhow::anyhow!( + "External value requires {length} bytes, only {} remain", + self.segments.len() + ) + }) + } +} + +impl ExternalValueSource for SegmentSource { + fn take_bytes(&mut self, length: usize) -> anyhow::Result> { + let segments = self.take(length)?; + let mut value = Vec::with_capacity(length); + for segment in segments.iter() { + value.extend_from_slice(segment); + } + Ok(value) + } + + fn take_chunks(&mut self, length: usize) -> anyhow::Result> { + Ok(self.take(length)?.into_chunks()) + } + + fn finish(&mut self) -> anyhow::Result<()> { + if !self.segments.is_empty() { + anyhow::bail!( + "Virtqueue message has {} trailing external bytes", + self.segments.len() + ); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use alloc::vec; + + use flatbuffers::FlatBufferBuilder; + use hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue; + use hyperlight_common::transport::ExternalValues; + + use super::*; + + #[test] + fn response_byte_chunks_retain_transport_storage() { + let external = Bytes::from(vec![1, 2, 3, 4]); + let external_ptr = external.as_ptr(); + let result = FunctionCallResult::new(Ok(ReturnValue::ByteChunks(vec![external.clone()]))); + let mut builder = FlatBufferBuilder::new(); + let mut external_values = ExternalValues::new(); + let control = result.encode(&mut builder, &mut external_values).unwrap(); + let payload_len = control.len() + external.len(); + let header = MsgHeader::new(MsgKind::Response, 7, u32::try_from(payload_len).unwrap()); + let segments = Segments::new([ + Bytes::copy_from_slice(header.as_bytes()), + Bytes::copy_from_slice(control), + external, + ]); + + let decoded = decode_response(segments, 7).unwrap().into_inner().unwrap(); + let ReturnValue::ByteChunks(chunks) = decoded else { + panic!("expected ByteChunks response"); + }; + + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].as_ptr(), external_ptr); + assert_eq!(chunks[0].as_ref(), &[1, 2, 3, 4]); + } + + #[test] + fn segment_source_flattens_only_contiguous_values() { + let first = Bytes::from_static(b"ab"); + let second = Bytes::from_static(b"cd"); + let second_ptr = second.as_ptr(); + let mut source = SegmentSource::new(Segments::new([first, second])); + + let contiguous = source.take_bytes(3).unwrap(); + let chunks = source.take_chunks(1).unwrap(); + source.finish().unwrap(); + + assert_eq!(contiguous, b"abc"); + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].as_ref(), b"d"); + assert_eq!(chunks[0].as_ptr(), second_ptr.wrapping_add(1)); + } + + #[test] + fn request_byte_chunks_retain_transport_storage() { + use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCallType; + use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType}; + + let external = Bytes::from(vec![1, 2, 3, 4]); + let external_ptr = external.as_ptr(); + let call = FunctionCall::new( + "echo".into(), + Some(vec![ParameterValue::ByteChunks(vec![external.clone()])]), + FunctionCallType::Guest, + ReturnType::ByteChunks, + ); + let mut builder = FlatBufferBuilder::new(); + let mut external_values = ExternalValues::new(); + let control = call.encode(&mut builder, &mut external_values).unwrap(); + let segments = Segments::new([Bytes::copy_from_slice(control), external]); + + let (cid, decoded) = decode_request(9, segments).unwrap(); + let ParameterValue::ByteChunks(chunks) = + decoded.parameters.unwrap().into_iter().next().unwrap() + else { + panic!("expected ByteChunks parameter"); + }; + + assert_eq!(cid, 9); + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].as_ptr(), external_ptr); + assert_eq!(chunks[0].as_ref(), &[1, 2, 3, 4]); + } +} diff --git a/src/hyperlight_guest/src/transport/context.rs b/src/hyperlight_guest/src/transport/context.rs new file mode 100644 index 000000000..fa581a09f --- /dev/null +++ b/src/hyperlight_guest/src/transport/context.rs @@ -0,0 +1,646 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +//! Guest virtqueue context. + +use alloc::vec::Vec; +use core::result; + +use flatbuffers::FlatBufferBuilder; +use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType}; +use hyperlight_common::flatbuffer_wrappers::function_types::{ + FunctionCallResult, ParameterValue, ReturnType, ReturnValue, +}; +use hyperlight_common::flatbuffer_wrappers::guest_error::GuestError; +use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity; +use hyperlight_common::outb::OutBAction; +use hyperlight_common::transport::{EncodedMessage, ExternalValues, MsgHeader, MsgKind}; +use hyperlight_common::virtq::{ + AllocError, G2H_LOWER_SLOT_COUNT, G2H_LOWER_SLOT_SIZE, Layout, MemOps, Notifier, QueueStats, + Segments, SendChain, SlotLayout, SlotPool, Token, UsedChain, VirtqError, VirtqProducer, +}; + +use super::{GuestMemOps, codec}; +use crate::bail; +use crate::error::{GuestErrorContext, Result}; +use crate::exit::out32; + +/// G2H notifier that exits to the host to process available work. +#[derive(Clone, Copy)] +pub struct G2hNotifier; + +impl Notifier for G2hNotifier { + fn notify(&self, _stats: QueueStats) { + unsafe { + out32(OutBAction::VirtqNotify as u16, 0); + } + } +} + +/// H2G prefill does not notify before the host consumer is attached. +#[derive(Clone, Copy)] +pub struct H2gNotifier; + +impl Notifier for H2gNotifier { + 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; + +/// Work selected by one H2G dispatch entry. +pub enum DispatchAction { + /// Invoke one guest function and return its correlation ID. + Call(u32, FunctionCall), + /// Prepare canonical transport state for snapshot capture. + SnapshotCheckpoint, +} + +/// Configuration for one queue passed to [`GuestContext::new`]. +#[derive(Debug)] +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, +} + +/// Writable capacity reserved on a G2H request chain. +#[derive(Clone, Copy)] +enum ReplyCapacity { + /// The chain carries no reply. + None, + /// Reserve at least this many reply bytes. + Bounded(usize), + /// Reserve every available preferred allocation. + Available, +} + +impl ReplyCapacity { + /// Select reply capacity for one host function return type. + fn for_return_type(return_type: ReturnType) -> Self { + match return_type { + ReturnType::String | ReturnType::VecBytes | ReturnType::ByteChunks => Self::Available, + _ => Self::Bounded(G2H_LOWER_SLOT_SIZE), + } + } +} + +/// Virtqueue runtime state for guest-host communication. +pub struct GuestContext { + /// Access to the shared transport arena. + mem: GuestMemOps, + /// Guest-to-host driver. + g2h_producer: G2hProducer, + /// G2H pool state used to count retained buffers. + g2h_pool: SlotPool, + /// Host-to-guest driver. + h2g_producer: H2gProducer, + /// H2G pool state used to count retained buffers. + h2g_pool: SlotPool, + /// Size of each prefilled H2G buffer. + h2g_slot_size: usize, + /// Snapshot checkpoint mailbox GVA. + mbx_gva: u64, + /// Correlation ID assigned to the next host-function request. + next_cid: u32, + /// Used by the C API. + last_host_result: Option>, + /// Error set by a C guest function. + last_guest_error: Option, +} + +impl GuestContext { + /// Create a new context with G2H and H2G queues. + pub fn new(g2h: QueueConfig, h2g: QueueConfig, mbx_gva: u64) -> Result { + let g2h_pool = g2h_pool(g2h.pool_gva, g2h.pool_pages, g2h.buffer_size) + .with_context(|| "failed to create G2H pool")?; + let mem = GuestMemOps::for_scratch(); + let g2h_producer = VirtqProducer::new(g2h.layout, mem, G2hNotifier, g2h_pool.clone()); + + 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, H2gNotifier, h2g_pool.clone()); + + let mut ctx = Self { + mem, + g2h_producer, + g2h_pool, + h2g_producer, + h2g_pool, + h2g_slot_size: h2g.buffer_size, + mbx_gva, + next_cid: 1, + last_host_result: None, + last_guest_error: None, + }; + + ctx.prefill_h2g()?; + Ok(ctx) + } + + /// Record an error raised through the C guest API. + pub fn set_guest_error(&mut self, error: GuestError) { + self.last_guest_error = Some(error); + } + + /// Take an error raised through the C guest API. + pub fn take_guest_error(&mut self) -> Option { + self.last_guest_error.take() + } + + /// Call a host function via the G2H virtqueue. + /// + /// Slot-aligned external values use a separate readable region. The same + /// chain carries bounded writable buffers for the response. + /// + /// # Errors + /// + /// Returns an error when encoding, queue submission, host dispatch, + /// response validation, or return-value conversion fails. + pub fn call_host_function>( + &mut self, + function_name: &str, + parameters: Option>, + return_type: ReturnType, + ) -> Result { + // Encode control data separately from borrowed external byte values. + let params = parameters.as_deref().unwrap_or_default(); + let estimated_capacity = estimate_flatbuffer_capacity(function_name, params); + + let fc = FunctionCall::new( + function_name.into(), + parameters, + FunctionCallType::Host, + return_type, + ); + + let mut builder = FlatBufferBuilder::with_capacity(estimated_capacity); + let mut externals = ExternalValues::new(); + + let control = fc + .encode(&mut builder, &mut externals) + .with_context(|| "failed to encode host function call")?; + + // Frame the request and include external values in its total length. + let cid = self.allocate_cid(); + let msg = EncodedMessage::new(MsgKind::Request, cid, control, externals) + .context("G2H message length overflow")?; + + let reply_cap = ReplyCapacity::for_return_type(return_type); + + // Submit once more after forcing the host to drain on backpressure. + let token = match self.try_send(&msg, reply_cap) { + Ok(token) => token, + Err(error) if error.is_transient() => { + self.g2h_producer.notify_backpressure(); + + if let Err(error) = self.g2h_producer.reclaim() { + bail!("G2H reclaim: {error}"); + } + + match self.try_send(&msg, reply_cap) { + Ok(token) => token, + Err(error) => bail!("G2H call retry: {error}"), + } + } + Err(error) => { + bail!("G2H call: {error}"); + } + }; + + // Poll completions, skipping earlier one-way acknowledgements until + // the request reply is available. + let reply = loop { + let Some(reply) = self.g2h_producer.poll()? else { + bail!("G2H: no reply received"); + }; + if reply.token() == token { + break reply; + } + if matches!(&reply, UsedChain::Data(..)) { + bail!("G2H: unexpected reply token {:?}", reply.token()); + } + }; + + let segments = match reply { + UsedChain::Data(_, segments) => segments, + UsedChain::Ack(_) => bail!("G2H: response was ack-only"), + }; + + // Decode external ByteChunks without flattening their transport-backed + // segments. + let fcr = codec::decode_response(segments, cid)?; + let ret = fcr.into_inner()?; + + let Ok(ret) = T::try_from(ret) else { + bail!("G2H: host return value type mismatch"); + }; + + Ok(ret) + } + + /// Receive one host-to-guest dispatch action. + /// + /// External `ByteChunks` retain their owner-backed H2G slots. Contiguous + /// `VecBytes` values copy directly into their final `Vec`. + pub fn recv_h2g_dispatch(&mut self) -> Result { + self.g2h_producer + .reclaim() + .with_context(|| "G2H completion reclaim failed")?; + + let Some(used) = self.h2g_producer.poll()? else { + bail!("H2G: expected a guest function call buffer"); + }; + + let mut first = match used { + UsedChain::Data(_, segments) => segments, + UsedChain::Ack(_) => bail!("H2G: guest function call buffer was ack-only"), + }; + + let header = first + .split_to(MsgHeader::SIZE) + .context("H2G buffer is missing its message header")? + .into_bytes(); + + let Some(header) = MsgHeader::from_bytes(&header) else { + bail!("H2G buffer has an invalid message header"); + }; + + match header.msg_kind() { + Ok(MsgKind::SnapshotCheckpoint) => { + return Ok(DispatchAction::SnapshotCheckpoint); + } + Ok(MsgKind::Request) if header.cid != 0 => {} + _ => bail!("H2G buffer has invalid request framing"), + } + + let payload_len = + usize::try_from(header.payload_len).context("H2G payload length overflow")?; + + if first.len() > payload_len { + bail!("H2G first buffer exceeds the declared payload length"); + } + + let mut received = first.len(); + let mut payload = first.into_chunks(); + + while received < payload_len { + let Some(used) = self.h2g_producer.poll()? else { + bail!("H2G: expected a continuation buffer"); + }; + + let segments = match used { + UsedChain::Data(_, segments) => segments, + UsedChain::Ack(_) => bail!("H2G continuation buffer was ack-only"), + }; + + if segments.is_empty() { + bail!("H2G continuation buffer is empty"); + } + + received = received + .checked_add(segments.len()) + .context("H2G payload length overflow")?; + + if received > payload_len { + bail!("H2G buffers exceed the declared payload length"); + } + payload.extend(segments.into_chunks()); + } + + let (cid, call) = codec::decode_request(header.cid, Segments::new(payload))?; + Ok(DispatchAction::Call(cid, call)) + } + + /// Return a guest-function result and replenish H2G receive buffers. + pub fn send_h2g_result(&mut self, cid: u32, result: FunctionCallResult) -> Result<()> { + self.g2h_producer + .reclaim() + .with_context(|| "G2H response reclaim failed")?; + + { + let mut builder = FlatBufferBuilder::new(); + let mut externals = ExternalValues::new(); + + let control = result + .encode(&mut builder, &mut externals) + .with_context(|| "failed to encode guest function result")?; + + let msg = EncodedMessage::new(MsgKind::Response, cid, control, externals) + .context("G2H response length overflow")?; + + self.try_send_deferred(&msg, ReplyCapacity::None) + .with_context(|| "G2H response submission failed")?; + } + + drop(result); + self.prefill_h2g() + } + + /// Canonicalize both queues while the host consumers are stopped. + pub fn prepare_snapshot(&mut self) -> Result<()> { + self.g2h_producer.reclaim()?; + self.g2h_producer.reset()?; + self.h2g_producer.reset()?; + + // [`SlotPool`] clones share one allocation bitmap with their producer. + // Producer reset releases every allocation still tracked by queue + // bookkeeping. At this checkpoint boundary, any allocation left in the + // bitmap is therefore held by an owner-backed `Bytes` returned to guest + // code. Subtracting the free slot count from the total slot count gives the + // exact number of retained slots across both size tiers. Multiple `Bytes` + // clones or slices backed by one owner still count as one slot. + // + // This runs before H2G prefill because posted receive buffers are + // transport owned allocations and must not be counted. The result only + // answers whether retained buffers exist. It does not identify their + // addresses, capacities, or initialized lengths. + let g2h = self + .g2h_pool + .count() + .checked_sub(self.g2h_pool.num_free()) + .ok_or(VirtqError::InvalidState)?; + + let h2g = self + .h2g_pool + .count() + .checked_sub(self.h2g_pool.num_free()) + .ok_or(VirtqError::InvalidState)?; + + let guest_owned = g2h.checked_add(h2g).ok_or(VirtqError::InvalidState)?; + + // TODO: Publish a retained-buffer manifest with pool-relative offsets and + // initialized lengths so the host can snapshot sanitized payload ranges. + // The count-only mailbox currently rejects every retained-buffer snapshot. + self.mem + .write(self.mbx_gva, &guest_owned.to_le_bytes()) + .map_err(|_| VirtqError::MemoryWriteError)?; + + self.prefill_h2g()?; + Ok(()) + } + + /// Send a log message via the G2H queue. + /// + /// Current notification policy exits to the host for every log. + /// + /// # Errors + /// + /// Returns an error when the message cannot be framed or submitted. + pub fn emit_log(&mut self, log_data: &[u8]) -> Result<()> { + let message = EncodedMessage::new(MsgKind::Log, 0, log_data, ExternalValues::new()) + .context("G2H message length overflow")?; + self.send_g2h_oneshot(&message) + } + + /// Stash a host function result for later retrieval. + /// + /// Used by the C API's two-step calling convention where + /// `hl_call_host_function` and `hl_get_host_return_value_as_*` + /// are separate calls. + pub fn stash_host_result(&mut self, result: Result) { + self.last_host_result = Some(result); + } + + /// Take the stashed host return value. + /// + /// Panics if no value was stashed or if the type conversion fails. + /// If the stashed result was an error, panics with the error message. + pub fn take_host_return>(&mut self) -> T { + let value = self + .last_host_result + .take() + .expect("No host return value available") + .expect("Host function returned an error"); + + match T::try_from(value) { + Ok(value) => value, + Err(_) => panic!("Host return value type mismatch"), + } + } + + /// Publish one writable H2G chain for each currently free slot. + /// + /// Retained external values reduce the number of available receive buffers + /// until their final owner drops. + 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_without_notify(); + return Ok(()); + } + Err(error) => bail!("H2G prefill build: {error}"), + }; + + match batch.submit(chain) { + Ok(_) => {} + Err(error) if error.is_transient() => { + batch.finish_without_notify(); + return Ok(()); + } + Err(error) => bail!("H2G prefill submit: {error}"), + } + } + } + + /// Submit a one-way G2H message without polling its acknowledgement. + /// + /// Completed acknowledgements remain available for normal polling or + /// reclamation when later submissions encounter backpressure. + fn send_g2h_oneshot(&mut self, message: &EncodedMessage<'_>) -> Result<()> { + match self.try_send(message, ReplyCapacity::None) { + Ok(_) => Ok(()), + Err(error) if error.is_transient() => { + // VM exit so host drains and completes G2H entries. + self.g2h_producer.notify_backpressure(); + + if let Err(error) = self.g2h_producer.reclaim() { + bail!("G2H one-way reclaim: {error}"); + } + + match self.try_send(message, ReplyCapacity::None) { + Ok(_) => Ok(()), + Err(error) => bail!("G2H one-way retry: {error}"), + } + } + Err(error) => bail!("G2H one-way message: {error}"), + } + } + + /// Build and submit one G2H descriptor chain. + /// + /// `reply_cap` defines optional host-function reply space. + fn try_send( + &mut self, + message: &EncodedMessage<'_>, + reply_cap: ReplyCapacity, + ) -> result::Result { + let chain = self.build_g2h_chain(message, reply_cap)?; + self.g2h_producer.submit(chain) + } + + /// Submit one G2H message for polling after the existing final halt. + fn try_send_deferred( + &mut self, + message: &EncodedMessage<'_>, + reply_capacity: ReplyCapacity, + ) -> result::Result { + let chain = self.build_g2h_chain(message, reply_capacity)?; + let mut batch = self.g2h_producer.batch(); + + let token = batch.submit(chain)?; + batch.finish_without_notify(); + + Ok(token) + } + + /// Build and initialize one G2H message chain. + fn build_g2h_chain( + &self, + message: &EncodedMessage<'_>, + reply_cap: ReplyCapacity, + ) -> result::Result, VirtqError> { + let segment_len = self.g2h_producer.preferred_segment_len(); + let num_free = self.g2h_pool.num_free(); + + let lengths = || message_region_lengths(message, segment_len); + + let reply_cap = match reply_cap { + ReplyCapacity::None => None, + ReplyCapacity::Bounded(cap) => Some(cap), + ReplyCapacity::Available => Some(self.g2h_pool.max_alloc(lengths(), num_free)?), + }; + + let mut builder = self.g2h_producer.chain(); + + for len in lengths() { + builder = builder.readable(len); + } + + if let Some(cap) = reply_cap { + builder = builder.writable(cap); + } + + let mut chain = builder.build()?; + for chunk in message.chunks() { + chain.write_all(chunk)?; + } + + Ok(chain) + } + + /// Allocate a new correlation ID for a host function request. + fn allocate_cid(&mut self) -> u32 { + let cid = self.next_cid; + self.next_cid = self.next_cid.wrapping_add(1); + + if self.next_cid == 0 { + self.next_cid = 1; + } + cid + } +} + +/// Group message bytes into logical readable region lengths. +fn message_region_lengths( + message: &EncodedMessage<'_>, + segment_len: usize, +) -> impl Iterator { + let external_len = message.external_len(); + let split_external = external_len != 0 && external_len.is_multiple_of(segment_len); + + let first_len = if split_external { + message.prefix_len() + } else { + message.total_len() + }; + + core::iter::once(first_len).chain(split_external.then_some(external_len)) +} + +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 { + let count = pool_len(pages)? / buffer_size; + SlotPool::new(SlotLayout::new(base, buffer_size, count)) +} + +/// Build the tiered G2H pool. +/// +/// One page of 256-byte slots serves small control and log messages without +/// consuming configured-size slots. Complete slots in the remaining pages form +/// the upper tier. +fn g2h_pool(base: u64, pages: usize, upper_size: usize) -> result::Result { + let pool_len = pool_len(pages)?; + let lower_len = G2H_LOWER_SLOT_COUNT + .checked_mul(G2H_LOWER_SLOT_SIZE) + .ok_or(AllocError::Overflow)?; + + let upper_len = pool_len + .checked_sub(lower_len) + .ok_or(AllocError::EmptyRegion)?; + + let upper_count = upper_len / upper_size; + + let lower = SlotLayout::new(base, G2H_LOWER_SLOT_SIZE, G2H_LOWER_SLOT_COUNT); + let upper = SlotLayout::new(lower.end_addr()?, upper_size, upper_count); + SlotPool::new_tiered(lower, upper) +} + +#[cfg(test)] +mod tests { + use hyperlight_common::flatbuffer_wrappers::ExternalValueSink; + + use super::*; + + fn encoded_message(external: &[u8]) -> EncodedMessage<'_> { + let mut values = ExternalValues::new(); + values.push_bytes(external).unwrap(); + EncodedMessage::new(MsgKind::Request, 1, b"control", values).unwrap() + } + + #[test] + fn message_regions_split_only_aligned_external_values() { + let aligned = [0; 4096]; + let message = encoded_message(&aligned); + let regions = message_region_lengths(&message, 4096).collect::>(); + assert_eq!(regions, [message.prefix_len(), aligned.len()]); + + let unaligned = [0; 4095]; + let message = encoded_message(&unaligned); + let regions = message_region_lengths(&message, 4096).collect::>(); + assert_eq!(regions, [message.total_len()]); + } +} diff --git a/src/hyperlight_guest/src/transport/mem.rs b/src/hyperlight_guest/src/transport/mem.rs new file mode 100644 index 000000000..856cd355c --- /dev/null +++ b/src/hyperlight_guest/src/transport/mem.rs @@ -0,0 +1,148 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +//! 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..51378ff2e --- /dev/null +++ b/src/hyperlight_guest/src/transport/mod.rs @@ -0,0 +1,76 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +//! Guest transport context and memory access. +//! +//! Global context is installed once via [`set_global_context`] and accessed via [`with_context`]. + +mod codec; +pub mod context; +pub mod mem; + +use core::cell::RefCell; +use core::sync::atomic::{AtomicU8, Ordering}; + +pub use context::{DispatchAction, 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_ctx(f: impl FnOnce(&mut GuestContext) -> R) -> R { + assert!(is_initialized(), "transport context not initialized"); + let mut ctx = GLOBAL_CONTEXT.0.borrow_mut(); + f(ctx.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/guest_function/call.rs b/src/hyperlight_guest_bin/src/guest_function/call.rs index 82874c659..6e58e8d46 100644 --- a/src/hyperlight_guest_bin/src/guest_function/call.rs +++ b/src/hyperlight_guest_bin/src/guest_function/call.rs @@ -17,15 +17,17 @@ limitations under the License. use alloc::format; use alloc::vec::Vec; -use flatbuffers::FlatBufferBuilder; use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType}; -use hyperlight_common::flatbuffer_wrappers::function_types::{FunctionCallResult, ParameterType}; +use hyperlight_common::flatbuffer_wrappers::function_types::{ + FunctionCallResult, ParameterType, ReturnValue, +}; use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError}; -use hyperlight_guest::bail; use hyperlight_guest::error::{HyperlightGuestError, Result}; +use hyperlight_guest::transport::DispatchAction; +use hyperlight_guest::{bail, transport}; use tracing::instrument; -use crate::{GUEST_HANDLE, REGISTERED_GUEST_FUNCTIONS}; +use crate::REGISTERED_GUEST_FUNCTIONS; core::arch::global_asm!( ".weak guest_dispatch_function", @@ -34,13 +36,13 @@ core::arch::global_asm!( ); #[tracing::instrument(skip_all, parent = tracing::Span::current(), level= "Trace")] -fn guest_dispatch_function_default(function_call: FunctionCall) -> Result> { +fn guest_dispatch_function_default(function_call: FunctionCall) -> Result { let name = &function_call.function_name; bail!(ErrorCode::GuestFunctionNotFound => "No handler found for function call: {name:#?}"); } #[instrument(skip_all, level = "Info")] -pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result> { +pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result { // Validate this is a Guest Function Call if function_call.function_call_type() != FunctionCallType::Guest { return Err(HyperlightGuestError::new( @@ -73,12 +75,8 @@ pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result } else { // The given function is not registered. The guest should implement a function called // guest_dispatch_function to handle this. - - // TODO: ideally we would define a default implementation of this with weak linkage so the guest is not required - // to implement the function but its seems that weak linkage is an unstable feature so for now its probably better - // to not do that. unsafe extern "Rust" { - fn guest_dispatch_function(function_call: FunctionCall) -> Result>; + fn guest_dispatch_function(function_call: FunctionCall) -> Result; } unsafe { guest_dispatch_function(function_call) } @@ -98,30 +96,16 @@ pub(crate) fn internal_dispatch_function() { tracing::span!(tracing::Level::INFO, "internal_dispatch_function").entered() }; - let handle = unsafe { GUEST_HANDLE }; - - let function_call = handle - .try_pop_shared_input_data_into::() - .expect("Function call deserialization failed"); + let dispatch = transport::with_ctx(|ctx| ctx.recv_h2g_dispatch()) + .expect("H2G dispatch deserialization failed"); - let res = call_guest_function(function_call); - - match res { - Ok(bytes) => { - handle - .push_shared_output_data(bytes.as_slice()) - .expect("Failed to serialize function call result"); - } - Err(err) => { - let guest_error = Err(GuestError::new(err.kind, err.message)); - let fcr = FunctionCallResult::new(guest_error); - let mut builder = FlatBufferBuilder::new(); - let data = fcr.encode(&mut builder); - handle - .push_shared_output_data(data) - .expect("Failed to serialize function call result"); + let result = match dispatch { + DispatchAction::Call(cid, fc) => { + let res = call_guest_function(fc).map_err(|err| GuestError::new(err.kind, err.message)); + Some((cid, FunctionCallResult::new(res))) } - } + DispatchAction::SnapshotCheckpoint => None, + }; // All this tracing logic shall be done right before the call to `hlt` which is done after this // function returns @@ -139,4 +123,11 @@ pub(crate) fn internal_dispatch_function() { // the host, if necessary. hyperlight_guest_tracing::flush(); } + + match result { + Some((cid, result)) => transport::with_ctx(|ctx| ctx.send_h2g_result(cid, result)) + .expect("Failed to send function call result"), + None => transport::with_ctx(|ctx| ctx.prepare_snapshot()) + .expect("Failed to prepare snapshot transport"), + } } diff --git a/src/hyperlight_guest_bin/src/guest_function/definition.rs b/src/hyperlight_guest_bin/src/guest_function/definition.rs index c96f2e9fe..27449ef79 100644 --- a/src/hyperlight_guest_bin/src/guest_function/definition.rs +++ b/src/hyperlight_guest_bin/src/guest_function/definition.rs @@ -21,7 +21,6 @@ use alloc::vec::Vec; use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall; use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterType, ReturnType}; use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; -use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result; use hyperlight_common::for_each_tuple; use hyperlight_common::func::{ Function, ParameterTuple, ResultType, ReturnValue, SupportedReturnType, @@ -29,7 +28,7 @@ use hyperlight_common::func::{ use hyperlight_guest::error::{HyperlightGuestError, Result}; /// The function pointer type for Rust guest functions. -pub type GuestFunc = fn(FunctionCall) -> Result>; +pub type GuestFunc = fn(FunctionCall) -> Result; /// The definition of a function exposed from the guest to the host. /// @@ -47,7 +46,7 @@ pub struct GuestFunctionDefinition { pub function_pointer: F, } -/// Trait for functions that can be converted to a `fn(FunctionCall) -> Result>` +/// Trait for functions that can be converted to a [`GuestFunc`]. #[doc(hidden)] pub trait IntoGuestFunction where @@ -59,8 +58,8 @@ where #[doc(hidden)] const ASSERT_ZERO_SIZED: (); - /// Convert the function into a `fn(FunctionCall) -> Result>` - fn into_guest_function(self) -> fn(FunctionCall) -> Result>; + /// Convert the function into a [`GuestFunc`]. + fn into_guest_function(self) -> GuestFunc; } /// Trait for functions that can be converted to a `GuestFunctionDefinition` @@ -78,21 +77,6 @@ where ) -> GuestFunctionDefinition; } -fn into_flatbuffer_result(value: ReturnValue) -> Vec { - match value { - ReturnValue::Void(()) => get_flatbuffer_result(()), - ReturnValue::Int(i) => get_flatbuffer_result(i), - ReturnValue::UInt(u) => get_flatbuffer_result(u), - ReturnValue::Long(l) => get_flatbuffer_result(l), - ReturnValue::ULong(ul) => get_flatbuffer_result(ul), - ReturnValue::Float(f) => get_flatbuffer_result(f), - ReturnValue::Double(d) => get_flatbuffer_result(d), - ReturnValue::Bool(b) => get_flatbuffer_result(b), - ReturnValue::String(s) => get_flatbuffer_result(s.as_str()), - ReturnValue::VecBytes(v) => get_flatbuffer_result(v.as_slice()), - } -} - macro_rules! impl_host_function { ([$N:expr] ($($p:ident: $P:ident),*)) => { impl IntoGuestFunction for F @@ -133,7 +117,7 @@ macro_rules! impl_host_function { assert!(core::mem::size_of::() == 0) }; - fn into_guest_function(self) -> fn(FunctionCall) -> Result> { + fn into_guest_function(self) -> GuestFunc { |fc: FunctionCall| { // SAFETY: This is safe because: // 1. F is zero-sized (enforced by the ASSERT_ZERO_SIZED const). @@ -143,7 +127,7 @@ macro_rules! impl_host_function { let params = fc.parameters.unwrap_or_default(); let params = <($($P,)*) as ParameterTuple>::from_value(params)?; let result = Function::::call(&this, params)?; - Ok(into_flatbuffer_result(result.into_value())) + Ok(result.into_value()) } } } diff --git a/src/hyperlight_guest_bin/src/guest_logger.rs b/src/hyperlight_guest_bin/src/guest_logger.rs index 7ff3d97fc..877d42721 100644 --- a/src/hyperlight_guest_bin/src/guest_logger.rs +++ b/src/hyperlight_guest_bin/src/guest_logger.rs @@ -15,12 +15,14 @@ limitations under the License. */ use alloc::format; +use alloc::string::ToString; +use alloc::vec::Vec; +use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData; use hyperlight_common::flatbuffer_wrappers::guest_log_level::LogLevel; +use hyperlight_guest::transport; use log::{LevelFilter, Metadata, Record}; -use crate::GUEST_HANDLE; - // this is private on purpose so that `log` can only be called though the `log!` macros. struct GuestLogger {} @@ -38,11 +40,9 @@ impl log::Log for GuestLogger { fn enabled(&self, _: &Metadata) -> bool { true } - fn log(&self, record: &Record) { - let handle = unsafe { GUEST_HANDLE }; if self.enabled(record.metadata()) { - handle.log_message( + log_message( record.level().into(), format!("{}", record.args()).as_str(), record.module_path().unwrap_or("Unknown"), @@ -64,6 +64,40 @@ pub fn log_message( file: &str, line: u32, ) { - let handle = unsafe { GUEST_HANDLE }; - handle.log_message(level, message, module_path, target, file, line); + let _send_to_host = || { + let log = GuestLogData::new( + message.to_string(), + module_path.to_string(), + level, + target.to_string(), + file.to_string(), + line, + ); + let bytes: Vec = log + .try_into() + .expect("Failed to convert GuestLogData to bytes"); + + transport::with_ctx(|ctx| { + ctx.emit_log(&bytes) + .expect("Unable to send log data via virtq"); + }); + }; + + #[cfg(all(feature = "trace_guest", target_arch = "x86_64"))] + if hyperlight_guest_tracing::is_trace_enabled() { + tracing::trace!( + event = message, + level = ?level, + code.filepath = module_path, + caller = target, + source_file = file, + code.lineno = line, + ); + } else { + _send_to_host(); + } + #[cfg(not(all(feature = "trace_guest", target_arch = "x86_64")))] + { + _send_to_host(); + } } diff --git a/src/hyperlight_guest_bin/src/host_comm.rs b/src/hyperlight_guest_bin/src/host_comm.rs index 301462313..51126daa7 100644 --- a/src/hyperlight_guest_bin/src/host_comm.rs +++ b/src/hyperlight_guest_bin/src/host_comm.rs @@ -22,9 +22,9 @@ use hyperlight_common::flatbuffer_wrappers::function_types::{ ParameterValue, ReturnType, ReturnValue, }; use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; -use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result; use hyperlight_common::func::{ParameterTuple, SupportedReturnType}; use hyperlight_guest::error::{HyperlightGuestError, Result}; +use hyperlight_guest::transport; use crate::GUEST_HANDLE; @@ -36,8 +36,7 @@ pub fn call_host_function( where T: TryFrom, { - let handle = unsafe { GUEST_HANDLE }; - handle.call_host_function::(function_name, parameters, return_type) + transport::with_ctx(|ctx| ctx.call_host_function(function_name, parameters, return_type)) } pub fn call_host(function_name: impl AsRef, args: impl ParameterTuple) -> Result @@ -47,44 +46,21 @@ where call_host_function::(function_name.as_ref(), Some(args.into_value()), T::TYPE) } -pub fn call_host_function_without_returning_result( - function_name: &str, - parameters: Option>, - return_type: ReturnType, -) -> Result<()> { - let handle = unsafe { GUEST_HANDLE }; - handle.call_host_function_without_returning_result(function_name, parameters, return_type) -} - -pub fn get_host_return_value_raw() -> Result { - let handle = unsafe { GUEST_HANDLE }; - handle.get_host_return_raw() -} - -pub fn get_host_return_value>() -> Result { - let handle = unsafe { GUEST_HANDLE }; - handle.get_host_return_value::() -} - pub fn read_n_bytes_from_user_memory(num: u64) -> Result> { let handle = unsafe { GUEST_HANDLE }; handle.read_n_bytes_from_user_memory(num) } /// Print a message using the host's print function. -/// -/// This function requires memory to be setup to be used. In particular, the -/// existence of the input and output memory regions. -pub fn print_output_with_host_print(function_call: FunctionCall) -> Result> { - let handle = unsafe { GUEST_HANDLE }; +pub fn print_output_with_host_print(function_call: FunctionCall) -> Result { if let ParameterValue::String(message) = function_call.parameters.unwrap().remove(0) { - let res = handle.call_host_function::( + let res = call_host_function::( "HostPrint", Some(Vec::from(&[ParameterValue::String(message)])), ReturnType::Int, )?; - Ok(get_flatbuffer_result(res)) + Ok(ReturnValue::Int(res)) } else { Err(HyperlightGuestError::new( ErrorCode::GuestError, diff --git a/src/hyperlight_guest_bin/src/lib.rs b/src/hyperlight_guest_bin/src/lib.rs index 5df92f647..ac1a15df5 100644 --- a/src/hyperlight_guest_bin/src/lib.rs +++ b/src/hyperlight_guest_bin/src/lib.rs @@ -52,6 +52,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 @@ -262,6 +263,9 @@ pub(crate) extern "C" fn generic_init( OS_PAGE_SIZE = ops as u32; } + // Prepare transport before logging or guest initialization code can use it. + transport::initialize(); + // set up the logger let guest_log_level_filter = GuestLogFilter::try_from(max_log_level).expect("Invalid log level"); @@ -315,6 +319,7 @@ pub mod __private { pub use alloc::vec::Vec; pub use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall; + pub use hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue; pub use hyperlight_common::func::ResultType; pub use hyperlight_guest::error::HyperlightGuestError; pub use linkme; diff --git a/src/hyperlight_guest_bin/src/transport.rs b/src/hyperlight_guest_bin/src/transport.rs new file mode 100644 index 000000000..6df1f3c17 --- /dev/null +++ b/src/hyperlight_guest_bin/src/transport.rs @@ -0,0 +1,105 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +//! 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 mbx_gva = scratch_gva(arena.mbx_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, + }, + mbx_gva, + ) + .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_guest_capi/README.md b/src/hyperlight_guest_capi/README.md index 418b4ac61..d8c2c74bc 100644 --- a/src/hyperlight_guest_capi/README.md +++ b/src/hyperlight_guest_capi/README.md @@ -2,13 +2,27 @@ This is a c-api wrapper over the hyperlight-guest/hyperlight-guest-bin crate. Th For examples on how to use it, see the c [simpleguest](../tests/c_guests/c_simpleguest/). -# Important +## Byte chunks -All guest functions must return a `hl_Vec*` obtained by calling one of the `hl_flatbuffer_result_from_*` functions. These functions will return a flatbuffer encoded byte-buffer of given value, for example `hl_flatbuffer_result_from_int(int)` will return the flatbuffer representation of the given int. +`ByteChunks` parameters use an array of borrowed pointer and length spans: -## NOTE +```c +hl_ByteChunks value = call->parameters[0].value.ByteChunks; +for (uintptr_t i = 0; i < value.count; i++) { + consume(value.chunks[i].data, value.chunks[i].len); +} +``` + +The parameter view is valid until the guest function returns. A value from +`hl_get_host_return_value_as_ByteChunks` remains valid until +`hl_free_byte_chunks`. `hl_result_from_ByteChunks` copies the supplied spans. -**You may not construct and return your own `hl_Vec*`**, as the hyperlight api assumes that all returned `hl_Vec*` are constructed through calls to a `hl_flatbuffer_result_from_*` function. +# Important + +Guest function wrappers return an `hl_ReturnValue*` created by an +`hl_result_from_*` function. -Additionally, note that type `hl_Vec*` is used in two different contexts. First, `hl_Vec*` is used input-parameter-type for guest functions that take a buffer of bytes. This buffer of bytes can contain **arbitrary** bytes. Second, all guest functions return a `hl_Vec*` (it might be hidden away by c macros). These `hl_Vec*` are flatbuffer-encoded data, and are not arbitrary. +## NOTE +The `hl_result_from_*` constructors establish matching tags, union payloads, +and ownership. diff --git a/src/hyperlight_guest_capi/cbindgen.toml b/src/hyperlight_guest_capi/cbindgen.toml index 89a4b93de..67866a524 100644 --- a/src/hyperlight_guest_capi/cbindgen.toml +++ b/src/hyperlight_guest_capi/cbindgen.toml @@ -20,8 +20,11 @@ prefix_with_name = true prefix = "hl_" [export.rename] +"FfiByteChunk" = "ByteChunk" +"FfiByteChunks" = "ByteChunks" "FfiFunctionCall" = "FunctionCall" "FfiParameter" = "Parameter" "FfiParameterValue" = "ParameterValue" +"FfiReturnValue" = "ReturnValue" +"FfiReturnValueUnion" = "ReturnValueUnion" "FfiVec" = "Vec" - diff --git a/src/hyperlight_guest_capi/include/macro.h b/src/hyperlight_guest_capi/include/macro.h index 1c6dc1ff7..9e49ab549 100644 --- a/src/hyperlight_guest_capi/include/macro.h +++ b/src/hyperlight_guest_capi/include/macro.h @@ -8,9 +8,9 @@ // // Parameters: 1. A function name // 2. The return type of the function. This must be one of the variant names in hl_ReturnType -// Note: This macro does not work for functions that return VecBytes. Instead, +// Note: This macro does not work for functions that return VecBytes or ByteChunks. Instead, // use `hl_register_function_definition` directly. You'll also need to return -// a flatbuffer-encoded hl_Vec* using the various hl_flatbuffer_result_from_* functions. +// an hl_ReturnValue* using the hl_result_from_* functions. // See c_simpleguest/main.c for an example. // 3. The number of parameters the function takes // 4+ The types of the parameters the function takes. The must be one of the variant names @@ -18,9 +18,9 @@ #define HYPERLIGHT_WRAP_FUNCTION(function, return_type, paramsc, ... ) HYPERLIGHT_WRAP_FUNCTION_##paramsc(function, return_type, __VA_ARGS__) #define HYPERLIGHT_WRAP_FUNCTION_0(function, return_type, ...) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function() \ + return hl_result_from_##return_type( function() \ ); \ } \ uintptr_t _##function##_parameter_count = 0; \ @@ -29,9 +29,9 @@ hl_ParameterType _##function##_parameter_types[] = { 0 }; \ #define HYPERLIGHT_WRAP_FUNCTION_1(function, return_type, arg1) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function( \ + return hl_result_from_##return_type( function( \ function_call->parameters[0].value.arg1 \ )); \ } \ @@ -40,9 +40,9 @@ hl_ReturnType _##function##_return_type = hl_ReturnType_##return_type; \ hl_ParameterType _##function##_parameter_types[] = { hl_ParameterType_##arg1 }; \ #define HYPERLIGHT_WRAP_FUNCTION_2(function, return_type, arg1, arg2) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function( \ + return hl_result_from_##return_type( function( \ function_call->parameters[0].value.arg1, \ function_call->parameters[1].value.arg2 \ )); \ @@ -55,9 +55,9 @@ hl_ParameterType _##function##_parameter_types[] = { hl_ParameterType_##arg1, \ #define HYPERLIGHT_WRAP_FUNCTION_3(function, return_type, arg1, arg2, arg3) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function( \ + return hl_result_from_##return_type( function( \ function_call->parameters[0].value.arg1, \ function_call->parameters[1].value.arg2, \ function_call->parameters[2].value.arg3 \ @@ -72,9 +72,9 @@ hl_ParameterType _##function##_parameter_types[] = { hl_ParameterType_##arg1, \ #define HYPERLIGHT_WRAP_FUNCTION_4(function, return_type, arg1, arg2, arg3, arg4) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function( \ + return hl_result_from_##return_type( function( \ function_call->parameters[0].value.arg1, \ function_call->parameters[1].value.arg2, \ function_call->parameters[2].value.arg3, \ @@ -90,9 +90,9 @@ hl_ParameterType _##function##_parameter_types[] = { hl_ParameterType_##arg1, \ }; \ #define HYPERLIGHT_WRAP_FUNCTION_5(function, return_type, arg1, arg2, arg3, arg4, arg5) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function( \ + return hl_result_from_##return_type( function( \ function_call->parameters[0].value.arg1, \ function_call->parameters[1].value.arg2, \ function_call->parameters[2].value.arg3, \ @@ -110,9 +110,9 @@ hl_ParameterType _##function##_parameter_types[] = { hl_ParameterType_##arg1, \ }; \ #define HYPERLIGHT_WRAP_FUNCTION_6(function, return_type, arg1, arg2, arg3, arg4, arg5, arg6) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function( \ + return hl_result_from_##return_type( function( \ function_call->parameters[0].value.arg1, \ function_call->parameters[1].value.arg2, \ function_call->parameters[2].value.arg3, \ @@ -132,9 +132,9 @@ hl_ParameterType _##function##_parameter_types[] = { hl_ParameterType_##arg1, \ }; \ #define HYPERLIGHT_WRAP_FUNCTION_7(function, return_type, arg1, arg2, arg3, arg4, arg5, arg6, arg7) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function( \ + return hl_result_from_##return_type( function( \ function_call->parameters[0].value.arg1, \ function_call->parameters[1].value.arg2, \ function_call->parameters[2].value.arg3, \ @@ -156,9 +156,9 @@ hl_ParameterType _##function##_parameter_types[] = { hl_ParameterType_##arg1, \ }; \ #define HYPERLIGHT_WRAP_FUNCTION_8(function, return_type, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function( \ + return hl_result_from_##return_type( function( \ function_call->parameters[0].value.arg1, \ function_call->parameters[1].value.arg2, \ function_call->parameters[2].value.arg3, \ @@ -182,9 +182,9 @@ hl_ParameterType _##function##_parameter_types[] = { hl_ParameterType_##arg1, \ }; \ #define HYPERLIGHT_WRAP_FUNCTION_9(function, return_type, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function( \ + return hl_result_from_##return_type( function( \ function_call->parameters[0].value.arg1, \ function_call->parameters[1].value.arg2, \ function_call->parameters[2].value.arg3, \ @@ -210,9 +210,9 @@ hl_ParameterType _##function##_parameter_types[] = { hl_ParameterType_##arg1, \ }; \ #define HYPERLIGHT_WRAP_FUNCTION_10(function, return_type, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function( \ + return hl_result_from_##return_type( function( \ function_call->parameters[0].value.arg1, \ function_call->parameters[1].value.arg2, \ function_call->parameters[2].value.arg3, \ @@ -240,9 +240,9 @@ hl_ParameterType _##function##_parameter_types[] = { hl_ParameterType_##arg1, \ }; \ #define HYPERLIGHT_WRAP_FUNCTION_11(function, return_type, arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11) \ -hl_Vec *_call_##function(const hl_FunctionCall *function_call) \ +hl_ReturnValue *_call_##function(const hl_FunctionCall *function_call) \ { \ - return hl_flatbuffer_result_from_##return_type( function( \ + return hl_result_from_##return_type( function( \ function_call->parameters[0].value.arg1, \ function_call->parameters[1].value.arg2, \ function_call->parameters[2].value.arg3, \ diff --git a/src/hyperlight_guest_capi/src/dispatch.rs b/src/hyperlight_guest_capi/src/dispatch.rs index e0a8bc34c..1ef30632c 100644 --- a/src/hyperlight_guest_capi/src/dispatch.rs +++ b/src/hyperlight_guest_capi/src/dispatch.rs @@ -20,27 +20,32 @@ use alloc::vec::Vec; use core::ffi::{CStr, c_char}; use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall; -use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterType, ReturnType}; +use hyperlight_common::flatbuffer_wrappers::function_types::{ + ParameterType, ReturnType, ReturnValue, +}; use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; use hyperlight_guest::error::{HyperlightGuestError, Result}; +use hyperlight_guest::transport; use hyperlight_guest_bin::guest_function::definition::GuestFunctionDefinition; use hyperlight_guest_bin::guest_function::register::GuestFunctionRegister; -use hyperlight_guest_bin::host_comm::call_host_function_without_returning_result; +use hyperlight_guest_bin::host_comm::call_host_function; -use crate::types::{FfiFunctionCall, FfiVec}; +use crate::types::{FfiFunctionCall, FfiReturnValue, OwnedFfiFunctionCall}; static mut REGISTERED_C_GUEST_FUNCTIONS: GuestFunctionRegister = GuestFunctionRegister::new(); -type CGuestFunc = extern "C" fn(&FfiFunctionCall) -> Box; +type CGuestFunc = extern "C" fn(&FfiFunctionCall) -> *mut FfiReturnValue; unsafe extern "C" { - // NOTE *mut FfiVec must be a Box. This will be the case as long as the guest - // returns a FfiVec that they created using the c-api hl_flatbuffer_result_from_* functions. - fn c_guest_dispatch_function(function_call: &FfiFunctionCall) -> *mut FfiVec; + // The guest must return a value created by an hl_result_from_* function. + fn c_guest_dispatch_function(function_call: &FfiFunctionCall) -> *mut FfiReturnValue; } #[unsafe(no_mangle)] -pub fn guest_dispatch_function(function_call: FunctionCall) -> Result> { +pub fn guest_dispatch_function(function_call: FunctionCall) -> Result { + // Discard an error left by guest code outside the current dispatch. + let _ = transport::with_ctx(|ctx| ctx.take_guest_error()); + // Use &raw const to get an immutable reference to the static HashMap // this is to avoid the clippy warning "shared reference to mutable static" if let Some(registered_func) = @@ -54,10 +59,29 @@ pub fn guest_dispatch_function(function_call: FunctionCall) -> Result> { .collect(); registered_func.verify_parameters(&function_call_parameter_types)?; - let ffi_func_call = FfiFunctionCall::from_function_call(function_call)?; - let function_result = (registered_func.function_pointer)(&ffi_func_call); + let function_name = function_call.function_name.clone(); + let ffi_func_call = OwnedFfiFunctionCall::from_function_call(function_call)?; + let function_result = (registered_func.function_pointer)(ffi_func_call.as_ffi()); + if function_result.is_null() { + if let Some(error) = transport::with_ctx(|ctx| ctx.take_guest_error()) { + return Err(HyperlightGuestError::new(error.code, error.message)); + } + return Err(HyperlightGuestError::new( + ErrorCode::GuestError, + alloc::format!("C guest function {function_name:?} returned null"), + )); + } + + // SAFETY: the pointer is non-null and C functions return ownership. + let function_result = unsafe { Box::from_raw(function_result) }; + // SAFETY: registered C functions return values created by hl_result_from_*. + let function_result = unsafe { (*function_result).into_return_value() }; - unsafe { Ok(FfiVec::into_vec(*function_result)) } + if let Some(error) = transport::with_ctx(|ctx| ctx.take_guest_error()) { + return Err(HyperlightGuestError::new(error.code, error.message)); + } + + Ok(function_result) } else { // The given function is not registered. The guest should implement a function called c_guest_dispatch_function to handle this. @@ -65,16 +89,26 @@ pub fn guest_dispatch_function(function_call: FunctionCall) -> Result> { // to implement the function but its seems that weak linkage is an unstable feature so for now its probably better // to not do that. let function_name = function_call.function_name.clone(); - let ffi_func_call = FfiFunctionCall::from_function_call(function_call)?; - let function_result = unsafe { c_guest_dispatch_function(&ffi_func_call) }; + let ffi_func_call = OwnedFfiFunctionCall::from_function_call(function_call)?; + let function_result = unsafe { c_guest_dispatch_function(ffi_func_call.as_ffi()) }; if function_result.is_null() { + if let Some(error) = transport::with_ctx(|ctx| ctx.take_guest_error()) { + return Err(HyperlightGuestError::new(error.code, error.message)); + } Err(HyperlightGuestError::new( ErrorCode::GuestFunctionNotFound, function_name, )) } else { let result = unsafe { Box::from_raw(function_result) }; - Ok(unsafe { FfiVec::into_vec(*result) }) + // SAFETY: non-null fallback results are created by hl_result_from_*. + let result = unsafe { (*result).into_return_value() }; + + if let Some(error) = transport::with_ctx(|ctx| ctx.take_guest_error()) { + return Err(HyperlightGuestError::new(error.code, error.message)); + } + + Ok(result) } } } @@ -98,15 +132,19 @@ pub extern "C" fn hl_register_function_definition( unsafe { (&mut *(&raw mut REGISTERED_C_GUEST_FUNCTIONS)).register(func_def) }; } -/// The caller is responsible for freeing the memory associated with given `FfiFunctionCall`. +/// Call a host function. The return value can be retrieved with +/// `hl_get_host_return_value_as_*` immediately after. #[unsafe(no_mangle)] pub extern "C" fn hl_call_host_function(function_call: &FfiFunctionCall) { let parameters = unsafe { function_call.copy_parameters() }; let func_name = unsafe { function_call.copy_function_name() }; let return_type = unsafe { function_call.copy_return_type() }; - // Use the non-generic internal implementation - // The C API will then call specific getter functions to fetch the properly typed return value - let _ = call_host_function_without_returning_result(&func_name, Some(parameters), return_type) - .expect("Failed to call host function"); + let result = call_host_function::(&func_name, Some(parameters), return_type); + transport::with_ctx(|ctx| ctx.stash_host_result(result)); +} + +/// Retrieve the return value stashed by the last `hl_call_host_function`. +pub(crate) fn take_last_host_return>() -> T { + transport::with_ctx(|ctx| ctx.take_host_return::()) } diff --git a/src/hyperlight_guest_capi/src/error.rs b/src/hyperlight_guest_capi/src/error.rs index 03217600e..98482731e 100644 --- a/src/hyperlight_guest_capi/src/error.rs +++ b/src/hyperlight_guest_capi/src/error.rs @@ -16,31 +16,21 @@ limitations under the License. use core::ffi::{CStr, c_char}; -use flatbuffers::FlatBufferBuilder; -use hyperlight_common::flatbuffer_wrappers::function_types::FunctionCallResult; use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError}; -use hyperlight_guest_bin::GUEST_HANDLE; +use hyperlight_guest::transport; use crate::alloc::borrow::ToOwned; #[unsafe(no_mangle)] pub extern "C" fn hl_set_error(err: ErrorCode, message: *const c_char) { let cstr = unsafe { CStr::from_ptr(message) }; - let guest_error = Err(GuestError::new( + let guest_error = GuestError::new( err.into(), cstr.to_str() .expect("Failed to convert CStr to &str") .to_owned(), - )); - let fcr = FunctionCallResult::new(guest_error); - let mut builder = FlatBufferBuilder::new(); - let data = fcr.encode(&mut builder); - unsafe { - #[allow(static_mut_refs)] // we are single threaded - GUEST_HANDLE - .push_shared_output_data(data) - .expect("Failed to set error") - } + ); + transport::with_ctx(|ctx| ctx.set_guest_error(guest_error)); } #[unsafe(no_mangle)] diff --git a/src/hyperlight_guest_capi/src/flatbuffer.rs b/src/hyperlight_guest_capi/src/flatbuffer.rs deleted file mode 100644 index ff12400d6..000000000 --- a/src/hyperlight_guest_capi/src/flatbuffer.rs +++ /dev/null @@ -1,158 +0,0 @@ -/* -Copyright 2025 The Hyperlight Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -use alloc::boxed::Box; -use alloc::ffi::CString; -use alloc::string::String; -use alloc::vec::Vec; -use core::ffi::{CStr, c_char}; - -use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result; -use hyperlight_guest_bin::host_comm::get_host_return_value; - -use crate::types::FfiVec; - -// The reason for the capitalized type in the function names below -// is to match the names of the variants in hl_ReturnType, -// which is used in the C macros in macro.h - -#[unsafe(no_mangle)] -pub extern "C" fn hl_flatbuffer_result_from_Int(value: i32) -> Box { - let vec = get_flatbuffer_result(value); - - Box::new(unsafe { FfiVec::from_vec(vec) }) -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_flatbuffer_result_from_UInt(value: u32) -> Box { - let vec = get_flatbuffer_result(value); - - Box::new(unsafe { FfiVec::from_vec(vec) }) -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_flatbuffer_result_from_Long(value: i64) -> Box { - let vec = get_flatbuffer_result(value); - - Box::new(unsafe { FfiVec::from_vec(vec) }) -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_flatbuffer_result_from_ULong(value: u64) -> Box { - let vec = get_flatbuffer_result(value); - - Box::new(unsafe { FfiVec::from_vec(vec) }) -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_flatbuffer_result_from_Float(value: f32) -> Box { - let vec = get_flatbuffer_result(value); - - Box::new(unsafe { FfiVec::from_vec(vec) }) -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_flatbuffer_result_from_Double(value: f64) -> Box { - let vec = get_flatbuffer_result(value); - - Box::new(unsafe { FfiVec::from_vec(vec) }) -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_flatbuffer_result_from_Void() -> Box { - let vec = get_flatbuffer_result(()); - - Box::new(unsafe { FfiVec::from_vec(vec) }) -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_flatbuffer_result_from_String(value: *const c_char) -> Box { - let str = unsafe { CStr::from_ptr(value) }; - let vec = get_flatbuffer_result(str.to_string_lossy().as_ref()); - - Box::new(unsafe { FfiVec::from_vec(vec) }) -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_flatbuffer_result_from_Bytes(data: *const u8, len: usize) -> Box { - let slice = unsafe { core::slice::from_raw_parts(data, len) }; - - let vec = get_flatbuffer_result(slice); - - Box::new(unsafe { FfiVec::from_vec(vec) }) -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_flatbuffer_result_from_Bool(value: bool) -> Box { - let vec = get_flatbuffer_result(value); - - Box::new(unsafe { FfiVec::from_vec(vec) }) -} - -//--- Functions for getting values returned by host functions calls - -#[unsafe(no_mangle)] -pub extern "C" fn hl_get_host_return_value_as_Int() -> i32 { - get_host_return_value().expect("Unable to get host return value as int") -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_get_host_return_value_as_UInt() -> u32 { - get_host_return_value().expect("Unable to get host return value as uint") -} - -// the same for long, ulong -#[unsafe(no_mangle)] -pub extern "C" fn hl_get_host_return_value_as_Long() -> i64 { - get_host_return_value().expect("Unable to get host return value as long") -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_get_host_return_value_as_ULong() -> u64 { - get_host_return_value().expect("Unable to get host return value as ulong") -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_get_host_return_value_as_Bool() -> bool { - get_host_return_value().expect("Unable to get host return value as bool") -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_get_host_return_value_as_Float() -> f32 { - get_host_return_value().expect("Unable to get host return value as f32") -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_get_host_return_value_as_Double() -> f64 { - get_host_return_value().expect("Unable to get host return value as f64") -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_get_host_return_value_as_String() -> *const c_char { - let string_value: String = - get_host_return_value().expect("Unable to get host return value as string"); - - let c_string = CString::new(string_value).expect("Failed to create CString"); - c_string.into_raw() -} - -#[unsafe(no_mangle)] -pub extern "C" fn hl_get_host_return_value_as_VecBytes() -> Box { - let vec_value: Vec = - get_host_return_value().expect("Unable to get host return value as vec bytes"); - - Box::new(unsafe { FfiVec::from_vec(vec_value) }) -} diff --git a/src/hyperlight_guest_capi/src/lib.rs b/src/hyperlight_guest_capi/src/lib.rs index fcf9aa2c1..c9826a0e6 100644 --- a/src/hyperlight_guest_capi/src/lib.rs +++ b/src/hyperlight_guest_capi/src/lib.rs @@ -21,6 +21,6 @@ extern crate alloc; pub mod dispatch; pub mod error; -pub mod flatbuffer; pub mod logging; +pub mod return_value; pub mod types; diff --git a/src/hyperlight_guest_capi/src/return_value.rs b/src/hyperlight_guest_capi/src/return_value.rs new file mode 100644 index 000000000..cdf12314a --- /dev/null +++ b/src/hyperlight_guest_capi/src/return_value.rs @@ -0,0 +1,173 @@ +/* +Copyright 2025 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use alloc::boxed::Box; +use alloc::ffi::CString; +use alloc::string::String; +use alloc::vec::Vec; +use core::ffi::{CStr, c_char}; + +use hyperlight_common::flatbuffer_wrappers::function_types::Bytes; + +use crate::dispatch::take_last_host_return; +use crate::types::{FfiByteChunks, FfiReturnValue, FfiVec, OwnedFfiByteChunks}; + +// The reason for the capitalized type in the function names below +// is to match the names of the variants in hl_ReturnType, +// which is used in the C macros in macro.h + +#[unsafe(no_mangle)] +pub extern "C" fn hl_result_from_Int(value: i32) -> Box { + Box::new(FfiReturnValue::int(value)) +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_result_from_UInt(value: u32) -> Box { + Box::new(FfiReturnValue::uint(value)) +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_result_from_Long(value: i64) -> Box { + Box::new(FfiReturnValue::long(value)) +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_result_from_ULong(value: u64) -> Box { + Box::new(FfiReturnValue::ulong(value)) +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_result_from_Float(value: f32) -> Box { + Box::new(FfiReturnValue::float(value)) +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_result_from_Double(value: f64) -> Box { + Box::new(FfiReturnValue::double(value)) +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_result_from_Void() -> Box { + Box::new(FfiReturnValue::void()) +} + +#[unsafe(no_mangle)] +/// # Safety +/// +/// `value` must point to a live NUL-terminated string. +pub unsafe extern "C" fn hl_result_from_String(value: *const c_char) -> Box { + // SAFETY: callers provide a live NUL-terminated string. + let value = unsafe { CStr::from_ptr(value) }; + Box::new(FfiReturnValue::string(value)) +} + +#[unsafe(no_mangle)] +/// # Safety +/// +/// `data` must reference `len` readable bytes when `len` is nonzero. +pub unsafe extern "C" fn hl_result_from_Bytes(data: *const u8, len: usize) -> Box { + let value = if len == 0 { + Vec::new() + } else { + // SAFETY: callers provide `len` readable bytes. + unsafe { core::slice::from_raw_parts(data, len) }.to_vec() + }; + Box::new(FfiReturnValue::vec_bytes(value)) +} + +#[unsafe(no_mangle)] +/// # Safety +/// +/// Every pointer in `value` must reference its declared number of bytes. +pub unsafe extern "C" fn hl_result_from_ByteChunks(value: FfiByteChunks) -> Box { + // SAFETY: required by the caller. + Box::new(unsafe { FfiReturnValue::byte_chunks(value) }) +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_result_from_Bool(value: bool) -> Box { + Box::new(FfiReturnValue::boolean(value)) +} + +//--- Functions for getting values returned by host functions calls + +#[unsafe(no_mangle)] +pub extern "C" fn hl_get_host_return_value_as_Int() -> i32 { + take_last_host_return() +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_get_host_return_value_as_UInt() -> u32 { + take_last_host_return() +} + +// the same for long, ulong +#[unsafe(no_mangle)] +pub extern "C" fn hl_get_host_return_value_as_Long() -> i64 { + take_last_host_return() +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_get_host_return_value_as_ULong() -> u64 { + take_last_host_return() +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_get_host_return_value_as_Bool() -> bool { + take_last_host_return() +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_get_host_return_value_as_Float() -> f32 { + take_last_host_return() +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_get_host_return_value_as_Double() -> f64 { + take_last_host_return() +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_get_host_return_value_as_String() -> *const c_char { + let string_value: String = take_last_host_return(); + + let c_string = CString::new(string_value).expect("Failed to create CString"); + c_string.into_raw() +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_get_host_return_value_as_VecBytes() -> Box { + let vec_value: Vec = take_last_host_return(); + + Box::new(unsafe { FfiVec::from_vec(vec_value) }) +} + +#[unsafe(no_mangle)] +pub extern "C" fn hl_get_host_return_value_as_ByteChunks() -> *mut FfiByteChunks { + let chunks: Vec = take_last_host_return(); + + OwnedFfiByteChunks::into_raw(chunks) +} + +#[unsafe(no_mangle)] +/// # Safety +/// +/// `value` must be null or a pointer returned by +/// [`hl_get_host_return_value_as_ByteChunks`] that has not already been freed. +pub unsafe extern "C" fn hl_free_byte_chunks(value: *mut FfiByteChunks) { + // SAFETY: required by the caller. + unsafe { OwnedFfiByteChunks::free(value) }; +} diff --git a/src/hyperlight_guest_capi/src/types.rs b/src/hyperlight_guest_capi/src/types.rs index 148a113f6..250cb4068 100644 --- a/src/hyperlight_guest_capi/src/types.rs +++ b/src/hyperlight_guest_capi/src/types.rs @@ -14,11 +14,17 @@ See the License for the specific language governing permissions and limitations under the License. */ +mod byte_chunks; +pub use byte_chunks::*; + mod function_call; pub use function_call::*; mod parameter; pub use parameter::*; +mod return_value; +pub use return_value::*; + mod vec; pub use vec::*; diff --git a/src/hyperlight_guest_capi/src/types/byte_chunks.rs b/src/hyperlight_guest_capi/src/types/byte_chunks.rs new file mode 100644 index 000000000..5ab1e0f4b --- /dev/null +++ b/src/hyperlight_guest_capi/src/types/byte_chunks.rs @@ -0,0 +1,286 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use alloc::boxed::Box; +use alloc::vec::Vec; +use core::{ptr, slice}; + +use hyperlight_common::flatbuffer_wrappers::function_types::Bytes; + +/// One borrowed byte chunk exposed through the C API. +#[repr(C)] +#[derive(Copy, Clone)] +pub struct FfiByteChunk { + data: *const u8, + len: usize, +} + +impl FfiByteChunk { + #[cfg(test)] + pub(crate) fn data(&self) -> *const u8 { + self.data + } + + fn borrowed(value: &Bytes) -> Self { + Self { + data: value.as_ptr(), + len: value.len(), + } + } + + fn from_owned_vec(value: Vec) -> Self { + let value = value.into_boxed_slice(); + let len = value.len(); + let data = Box::into_raw(value) as *mut u8; + Self { data, len } + } + + /// # Safety + /// + /// `data` must reference `len` readable bytes when `len` is nonzero. + unsafe fn as_slice(&self) -> &[u8] { + if self.len == 0 { + &[] + } else { + // SAFETY: required by the caller. + unsafe { slice::from_raw_parts(self.data, self.len) } + } + } + + /// # Safety + /// + /// This chunk must have been created by [`Self::from_owned_vec`] and must + /// not have been consumed before. + unsafe fn into_owned_bytes(self) -> Bytes { + let value = ptr::slice_from_raw_parts_mut(self.data.cast_mut(), self.len); + // SAFETY: required by the caller. + let value = unsafe { Box::from_raw(value) }; + Bytes::from(value.into_vec()) + } + + /// # Safety + /// + /// This chunk must have been created by [`Self::from_owned_vec`] and must + /// not have been consumed before. + unsafe fn drop_owned(self) { + let value = ptr::slice_from_raw_parts_mut(self.data.cast_mut(), self.len); + // SAFETY: required by the caller. + drop(unsafe { Box::from_raw(value) }); + } +} + +/// A borrowed array of byte chunks exposed through the C API. +#[repr(C)] +#[derive(Copy, Clone)] +pub struct FfiByteChunks { + chunks: *const FfiByteChunk, + count: usize, +} + +impl FfiByteChunks { + /// # Safety + /// + /// `chunks` must reference `count` live descriptors when `count` is + /// nonzero. Each descriptor must reference its declared number of bytes. + pub(crate) unsafe fn copy_to_bytes(self) -> Vec { + // SAFETY: required by the caller. + unsafe { self.as_slice() } + .iter() + .map(|chunk| { + // SAFETY: required by the caller. + Bytes::copy_from_slice(unsafe { chunk.as_slice() }) + }) + .collect() + } + + /// # Safety + /// + /// `chunks` must reference `count` live descriptors when `count` is + /// nonzero. Each descriptor must reference its declared number of bytes. + pub(crate) unsafe fn copy_owned(self) -> Self { + // Copy every input before leaking any allocation into the owned view. + // SAFETY: required by the caller. + let chunks = unsafe { self.as_slice() } + .iter() + .map(|chunk| { + // SAFETY: required by the caller. + unsafe { chunk.as_slice() }.to_vec() + }) + .collect(); + Self::from_owned_chunks(chunks) + } + + /// # Safety + /// + /// This value must have been created by [`Self::copy_owned`] and must not + /// have been consumed before. + pub(crate) unsafe fn into_owned_bytes(self) -> Vec { + // SAFETY: required by the caller. + let chunks = unsafe { self.into_owned_descriptors() }; + chunks + .into_vec() + .into_iter() + .map(|chunk| { + // SAFETY: every descriptor owns an allocation created by + // `from_owned_chunks`. + unsafe { chunk.into_owned_bytes() } + }) + .collect() + } + + /// # Safety + /// + /// This value must have been created by [`Self::copy_owned`] and must not + /// have been consumed before. + pub(crate) unsafe fn drop_owned(self) { + // SAFETY: required by the caller. + let chunks = unsafe { self.into_owned_descriptors() }; + for chunk in chunks.iter().copied() { + // SAFETY: every descriptor owns an allocation created by + // `from_owned_chunks`. + unsafe { chunk.drop_owned() }; + } + } + + /// # Safety + /// + /// `chunks` must reference `count` live descriptors when `count` is + /// nonzero. + pub(crate) unsafe fn as_slice(&self) -> &[FfiByteChunk] { + if self.count == 0 { + &[] + } else { + // SAFETY: required by the caller. + unsafe { slice::from_raw_parts(self.chunks, self.count) } + } + } + + fn from_owned_chunks(chunks: Vec>) -> Self { + let chunks: Vec<_> = chunks + .into_iter() + .map(FfiByteChunk::from_owned_vec) + .collect(); + + let chunks = chunks.into_boxed_slice(); + let count = chunks.len(); + let chunks = Box::into_raw(chunks) as *mut FfiByteChunk; + + Self { chunks, count } + } + + /// # Safety + /// + /// This value must have been created by [`Self::from_owned_chunks`] and + /// must not have been consumed before. + unsafe fn into_owned_descriptors(self) -> Box<[FfiByteChunk]> { + let chunks = ptr::slice_from_raw_parts_mut(self.chunks.cast_mut(), self.count); + // SAFETY: required by the caller. + unsafe { Box::from_raw(chunks) } + } +} + +/// Owns the Rust chunks and descriptors behind one borrowed C view. +pub(crate) struct FfiByteChunksOwner { + _chunks: Vec, + descriptors: Box<[FfiByteChunk]>, +} + +impl FfiByteChunksOwner { + pub(crate) fn new(chunks: Vec) -> Self { + let descriptors = chunks + .iter() + .map(FfiByteChunk::borrowed) + .collect::>() + .into_boxed_slice(); + Self { + _chunks: chunks, + descriptors, + } + } + + pub(crate) fn view(&self) -> FfiByteChunks { + FfiByteChunks { + chunks: self.descriptors.as_ptr(), + count: self.descriptors.len(), + } + } +} + +/// Keeps a host return alive behind the public view pointer. +#[repr(C)] +pub(crate) struct OwnedFfiByteChunks { + view: FfiByteChunks, + _owner: FfiByteChunksOwner, +} + +impl OwnedFfiByteChunks { + pub(crate) fn into_raw(chunks: Vec) -> *mut FfiByteChunks { + let owner = FfiByteChunksOwner::new(chunks); + let value = Box::new(Self { + view: owner.view(), + _owner: owner, + }); + Box::into_raw(value).cast() + } + + /// # Safety + /// + /// `value` must be null or a pointer returned by [`Self::into_raw`] that + /// has not already been freed. + pub(crate) unsafe fn free(value: *mut FfiByteChunks) { + if !value.is_null() { + // SAFETY: `view` is the first field of this `repr(C)` allocation. + drop(unsafe { Box::from_raw(value.cast::()) }); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn borrowed_view_preserves_chunks_without_copying() { + let chunks = vec![Bytes::from_static(b"first"), Bytes::from_static(b"second")]; + let addresses = chunks.iter().map(Bytes::as_ptr).collect::>(); + let owner = FfiByteChunksOwner::new(chunks); + let view = owner.view(); + + // SAFETY: `owner` keeps the descriptor array and chunks alive. + let descriptors = unsafe { view.as_slice() }; + assert_eq!(descriptors.len(), 2); + assert_eq!(descriptors[0].data, addresses[0]); + assert_eq!(descriptors[1].data, addresses[1]); + } + + #[test] + fn owned_view_preserves_chunk_contents() { + let source = FfiByteChunksOwner::new(vec![ + Bytes::from_static(b"first"), + Bytes::from_static(b"second"), + ]); + + // SAFETY: `source` keeps the view live while it is copied. + let owned = unsafe { source.view().copy_owned() }; + // SAFETY: `owned` has not been consumed since `copy_owned`. + let chunks = unsafe { owned.into_owned_bytes() }; + + assert_eq!( + chunks, + vec![Bytes::from_static(b"first"), Bytes::from_static(b"second")] + ); + } +} diff --git a/src/hyperlight_guest_capi/src/types/function_call.rs b/src/hyperlight_guest_capi/src/types/function_call.rs index 1b2d4e85e..cb6078ede 100644 --- a/src/hyperlight_guest_capi/src/types/function_call.rs +++ b/src/hyperlight_guest_capi/src/types/function_call.rs @@ -25,7 +25,7 @@ use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall; use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType}; use hyperlight_guest::error::Result; -use crate::types::FfiParameter; +use crate::types::{FfiParameter, OwnedFfiParameter}; /// An FFI version of `FunctionCall` #[repr(C)] @@ -36,49 +36,75 @@ pub struct FfiFunctionCall { return_type: ReturnType, } -impl FfiFunctionCall { - /// Create a new `FfiFunctionCall` by consuming a FunctionCall. - pub fn from_function_call(value: FunctionCall) -> Result { - let leaked_function_name = CString::new(value.function_name.as_str()) - .expect("Failed to convert function name to CString") - .into_raw(); +pub(crate) struct OwnedFfiFunctionCall { + ffi: FfiFunctionCall, + _function_name: CString, + _parameters: Box<[FfiParameter]>, + _parameter_owners: Vec, +} - let (parameters, parameter_len) = match value.parameters { - Some(p) => { - let parameters: Vec = p - .into_iter() - .map(|param| FfiParameter::from_parameter_value(param).unwrap()) - .collect(); - let boxed = parameters.into_boxed_slice(); - let parameters_len = boxed.len(); - let leaked_param_vec = Box::into_raw(boxed); - (leaked_param_vec as *const FfiParameter, parameters_len) - } - None => (core::ptr::null(), 0), +impl OwnedFfiFunctionCall { + pub(crate) fn from_function_call(value: FunctionCall) -> Result { + let function_name = CString::new(value.function_name.as_str()) + .expect("Failed to convert function name to CString"); + let parameter_owners = value + .parameters + .unwrap_or_default() + .into_iter() + .map(OwnedFfiParameter::from_parameter_value) + .collect::>>()?; + let parameters = parameter_owners + .iter() + .map(OwnedFfiParameter::ffi) + .collect::>() + .into_boxed_slice(); + let parameters_len = parameters.len(); + let parameters_ptr = if parameters.is_empty() { + core::ptr::null() + } else { + parameters.as_ptr() + }; + let ffi = FfiFunctionCall { + function_name: function_name.as_ptr(), + parameters: parameters_ptr, + parameters_len, + return_type: value.expected_return_type, }; Ok(Self { - function_name: leaked_function_name, - parameters, - parameters_len: parameter_len, - return_type: value.expected_return_type, + ffi, + _function_name: function_name, + _parameters: parameters, + _parameter_owners: parameter_owners, }) } + pub(crate) fn as_ffi(&self) -> &FfiFunctionCall { + &self.ffi + } +} + +impl FfiFunctionCall { /// Copies the parameters of `self` into a new `Vec`. /// # Safety - /// `self` must be an unmodified version of what `from_function_call` returned. + /// Every pointer in `self` must reference a live value of the declared + /// length. pub unsafe fn copy_parameters(&self) -> Vec { - let slice = unsafe { slice::from_raw_parts(self.parameters, self.parameters_len) }; + let slice = if self.parameters_len == 0 { + &[] + } else { + // SAFETY: required by the caller. + unsafe { slice::from_raw_parts(self.parameters, self.parameters_len) } + }; slice .iter() .map(|param| unsafe { param.copy_to_parameter_value() }) .collect() } - /// Copies the function name of `self into a new `String`. + /// Copies the function name of `self` into a new `String`. /// # Safety - /// `self` must be an unmodified version of what `from_function_call` returned. + /// `function_name` must point to a live NUL-terminated string. pub unsafe fn copy_function_name(&self) -> String { unsafe { CStr::from_ptr(self.function_name) @@ -89,25 +115,8 @@ impl FfiFunctionCall { /// Copies the return type of `self` into a new `ReturnType`. /// # Safety - /// `self` must be an unmodified version of what `from_function_call` returned. + /// `return_type` must contain a valid [`ReturnType`] discriminant. pub unsafe fn copy_return_type(&self) -> ReturnType { self.return_type } } - -impl Drop for FfiFunctionCall { - fn drop(&mut self) { - unsafe { - if !self.function_name.is_null() { - drop(CString::from_raw(self.function_name as *mut c_char)); - } - if !self.parameters.is_null() { - let slice = Box::from_raw(slice::from_raw_parts_mut( - self.parameters as *mut FfiParameter, - self.parameters_len, - )); - drop(slice); - } - } - } -} diff --git a/src/hyperlight_guest_capi/src/types/parameter.rs b/src/hyperlight_guest_capi/src/types/parameter.rs index 048169754..3ecaf0fc4 100644 --- a/src/hyperlight_guest_capi/src/types/parameter.rs +++ b/src/hyperlight_guest_capi/src/types/parameter.rs @@ -15,12 +15,13 @@ limitations under the License. */ use alloc::ffi::CString; +use alloc::vec::Vec; use core::ffi::{CStr, c_char}; use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterType, ParameterValue}; use hyperlight_guest::error::Result; -use crate::types::FfiVec; +use crate::types::{FfiByteChunks, FfiByteChunksOwner, FfiVec}; /// A union of the value stored in a ParameterValue, used for FFI. /// On it's own, this union has no way to know which value type is stored @@ -38,46 +39,113 @@ pub union FfiParameterValue { pub Bool: bool, pub String: *mut c_char, pub VecBytes: FfiVec, + pub ByteChunks: FfiByteChunks, } -/// An owned FFI version Of `ParameterValue` +/// An FFI view of a [`ParameterValue`]. #[repr(C)] +#[derive(Clone)] #[allow(non_camel_case_types)] pub struct FfiParameter { tag: ParameterType, value: FfiParameterValue, } -impl FfiParameter { - /// Returns a new `FfiParameter` by consuming a `ParameterValue` - pub fn from_parameter_value(value: ParameterValue) -> Result { - let (tag, union) = match value { - ParameterValue::Int(v) => (ParameterType::Int, FfiParameterValue { Int: v }), - ParameterValue::UInt(v) => (ParameterType::UInt, FfiParameterValue { UInt: v }), - ParameterValue::Long(v) => (ParameterType::Long, FfiParameterValue { Long: v }), - ParameterValue::ULong(v) => (ParameterType::ULong, FfiParameterValue { ULong: v }), - ParameterValue::Float(v) => (ParameterType::Float, FfiParameterValue { Float: v }), - ParameterValue::Double(v) => (ParameterType::Double, FfiParameterValue { Double: v }), - ParameterValue::Bool(v) => (ParameterType::Bool, FfiParameterValue { Bool: v }), +enum FfiParameterOwner { + None, + String { _value: CString }, + VecBytes { _value: Vec }, + ByteChunks { _value: FfiByteChunksOwner }, +} + +pub(crate) struct OwnedFfiParameter { + ffi: FfiParameter, + _owner: FfiParameterOwner, +} + +impl OwnedFfiParameter { + pub(crate) fn from_parameter_value(value: ParameterValue) -> Result { + let (tag, union, owner) = match value { + ParameterValue::Int(v) => ( + ParameterType::Int, + FfiParameterValue { Int: v }, + FfiParameterOwner::None, + ), + ParameterValue::UInt(v) => ( + ParameterType::UInt, + FfiParameterValue { UInt: v }, + FfiParameterOwner::None, + ), + ParameterValue::Long(v) => ( + ParameterType::Long, + FfiParameterValue { Long: v }, + FfiParameterOwner::None, + ), + ParameterValue::ULong(v) => ( + ParameterType::ULong, + FfiParameterValue { ULong: v }, + FfiParameterOwner::None, + ), + ParameterValue::Float(v) => ( + ParameterType::Float, + FfiParameterValue { Float: v }, + FfiParameterOwner::None, + ), + ParameterValue::Double(v) => ( + ParameterType::Double, + FfiParameterValue { Double: v }, + FfiParameterOwner::None, + ), + ParameterValue::Bool(v) => ( + ParameterType::Bool, + FfiParameterValue { Bool: v }, + FfiParameterOwner::None, + ), ParameterValue::String(v) => { - let c_str = CString::new(v.as_str()).expect("Unable to make CString from String"); - let leaked = c_str.into_raw(); - (ParameterType::String, FfiParameterValue { String: leaked }) + let value = CString::new(v.as_str()).expect("Unable to make CString from String"); + let ptr = value.as_ptr().cast_mut(); + ( + ParameterType::String, + FfiParameterValue { String: ptr }, + FfiParameterOwner::String { _value: value }, + ) } ParameterValue::VecBytes(v) => { - let leaked = unsafe { FfiVec::from_vec(v) }; + let mut value = v; + let view = FfiVec::from_mut_slice(&mut value); ( ParameterType::VecBytes, - FfiParameterValue { VecBytes: leaked }, + FfiParameterValue { VecBytes: view }, + FfiParameterOwner::VecBytes { _value: value }, + ) + } + ParameterValue::ByteChunks(v) => { + let owner = FfiByteChunksOwner::new(v); + ( + ParameterType::ByteChunks, + FfiParameterValue { + ByteChunks: owner.view(), + }, + FfiParameterOwner::ByteChunks { _value: owner }, ) } }; - Ok(FfiParameter { tag, value: union }) + Ok(Self { + ffi: FfiParameter { tag, value: union }, + _owner: owner, + }) } + pub(crate) fn ffi(&self) -> FfiParameter { + self.ffi.clone() + } +} + +impl FfiParameter { /// Copies self into a new `ParameterValue`. /// # Safety - /// `self` must be an unmodified version of what `from_parameter_value` returned. + /// Every pointer selected by `tag` must reference a live value of the + /// declared length. pub unsafe fn copy_to_parameter_value(&self) -> ParameterValue { match self.tag { ParameterType::Int => ParameterValue::Int(unsafe { self.value.Int }), @@ -95,20 +163,42 @@ impl FfiParameter { ParameterType::VecBytes => { ParameterValue::VecBytes(unsafe { self.value.VecBytes.copy_to_vec() }) } + ParameterType::ByteChunks => { + // SAFETY: required by the caller. + ParameterValue::ByteChunks(unsafe { self.value.ByteChunks.copy_to_bytes() }) + } } } } -impl Drop for FfiParameter { - fn drop(&mut self) { - match self.tag { - ParameterType::String => unsafe { - drop(CString::from_raw(self.value.String)); - }, - ParameterType::VecBytes => unsafe { - drop(self.value.VecBytes.into_vec()); - }, - _ => {} - } +#[cfg(test)] +mod tests { + use hyperlight_common::flatbuffer_wrappers::function_types::Bytes; + + use super::*; + + #[test] + fn byte_chunks_parameter_is_borrowed_without_copying() { + let chunks = vec![Bytes::from_static(b"first"), Bytes::from_static(b"second")]; + let addresses = chunks.iter().map(Bytes::as_ptr).collect::>(); + let parameter = + OwnedFfiParameter::from_parameter_value(ParameterValue::ByteChunks(chunks)).unwrap(); + let ffi = parameter.ffi(); + + // SAFETY: `parameter` keeps its descriptor array and chunks alive. + let descriptors = unsafe { ffi.value.ByteChunks.as_slice() }; + assert_eq!(descriptors.len(), 2); + assert_eq!(descriptors[0].data(), addresses[0]); + assert_eq!(descriptors[1].data(), addresses[1]); + + // SAFETY: `parameter` keeps every pointer in `ffi` alive. + let copied = unsafe { ffi.copy_to_parameter_value() }; + assert_eq!( + copied, + ParameterValue::ByteChunks(vec![ + Bytes::from_static(b"first"), + Bytes::from_static(b"second") + ]) + ); } } diff --git a/src/hyperlight_guest_capi/src/types/return_value.rs b/src/hyperlight_guest_capi/src/types/return_value.rs new file mode 100644 index 000000000..3f2899b4a --- /dev/null +++ b/src/hyperlight_guest_capi/src/types/return_value.rs @@ -0,0 +1,185 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use alloc::borrow::ToOwned; +use alloc::ffi::CString; +use alloc::vec::Vec; +use core::ffi::{CStr, c_char}; +use core::mem::ManuallyDrop; + +use hyperlight_common::flatbuffer_wrappers::function_types::{ReturnType, ReturnValue}; + +use super::{FfiByteChunks, FfiVec}; + +/// The value held by an [`FfiReturnValue`]. +#[repr(C)] +#[derive(Copy, Clone)] +#[allow(non_camel_case_types, non_snake_case)] +pub union FfiReturnValueUnion { + pub Int: i32, + pub UInt: u32, + pub Long: i64, + pub ULong: u64, + pub Float: f32, + pub Double: f64, + pub Bool: bool, + pub String: *mut c_char, + pub VecBytes: FfiVec, + pub ByteChunks: FfiByteChunks, +} + +/// An owned FFI return value. +#[repr(C)] +#[allow(non_camel_case_types)] +pub struct FfiReturnValue { + tag: ReturnType, + value: FfiReturnValueUnion, +} + +impl FfiReturnValue { + pub fn int(value: i32) -> Self { + Self { + tag: ReturnType::Int, + value: FfiReturnValueUnion { Int: value }, + } + } + + pub fn uint(value: u32) -> Self { + Self { + tag: ReturnType::UInt, + value: FfiReturnValueUnion { UInt: value }, + } + } + + pub fn long(value: i64) -> Self { + Self { + tag: ReturnType::Long, + value: FfiReturnValueUnion { Long: value }, + } + } + + pub fn ulong(value: u64) -> Self { + Self { + tag: ReturnType::ULong, + value: FfiReturnValueUnion { ULong: value }, + } + } + + pub fn float(value: f32) -> Self { + Self { + tag: ReturnType::Float, + value: FfiReturnValueUnion { Float: value }, + } + } + + pub fn double(value: f64) -> Self { + Self { + tag: ReturnType::Double, + value: FfiReturnValueUnion { Double: value }, + } + } + + pub fn boolean(value: bool) -> Self { + Self { + tag: ReturnType::Bool, + value: FfiReturnValueUnion { Bool: value }, + } + } + + pub fn void() -> Self { + Self { + tag: ReturnType::Void, + value: FfiReturnValueUnion { Int: 0 }, + } + } + + pub fn string(value: &CStr) -> Self { + Self { + tag: ReturnType::String, + value: FfiReturnValueUnion { + String: value.to_owned().into_raw(), + }, + } + } + + pub fn vec_bytes(value: Vec) -> Self { + Self { + tag: ReturnType::VecBytes, + // SAFETY: `FfiReturnValue` reclaims the allocation when consumed or dropped. + value: FfiReturnValueUnion { + VecBytes: unsafe { FfiVec::from_vec(value) }, + }, + } + } + + /// # Safety + /// + /// Every pointer in `value` must reference its declared number of bytes. + pub unsafe fn byte_chunks(value: FfiByteChunks) -> Self { + Self { + tag: ReturnType::ByteChunks, + value: FfiReturnValueUnion { + // SAFETY: required by the caller. + ByteChunks: unsafe { value.copy_owned() }, + }, + } + } + + /// Consume this value and transfer its payload into a Rust return value. + /// + /// # Safety + /// + /// The tag and union value must be unchanged from a value created by this + /// type's constructors. + pub unsafe fn into_return_value(self) -> ReturnValue { + let value = ManuallyDrop::new(self); + // SAFETY: the contract requires the tag to identify the initialized union field. + unsafe { + match value.tag { + ReturnType::Int => ReturnValue::Int(value.value.Int), + ReturnType::UInt => ReturnValue::UInt(value.value.UInt), + ReturnType::Long => ReturnValue::Long(value.value.Long), + ReturnType::ULong => ReturnValue::ULong(value.value.ULong), + ReturnType::Float => ReturnValue::Float(value.value.Float), + ReturnType::Double => ReturnValue::Double(value.value.Double), + ReturnType::Bool => ReturnValue::Bool(value.value.Bool), + ReturnType::Void => ReturnValue::Void(()), + ReturnType::String => { + let value = CString::from_raw(value.value.String); + ReturnValue::String(value.to_string_lossy().into_owned()) + } + ReturnType::VecBytes => ReturnValue::VecBytes(value.value.VecBytes.into_vec()), + ReturnType::ByteChunks => { + ReturnValue::ByteChunks(value.value.ByteChunks.into_owned_bytes()) + } + } + } + } +} + +impl Drop for FfiReturnValue { + fn drop(&mut self) { + // SAFETY: constructors initialize the owned field selected by the tag. + unsafe { + match self.tag { + ReturnType::String => drop(CString::from_raw(self.value.String)), + ReturnType::VecBytes => drop(self.value.VecBytes.into_vec()), + ReturnType::ByteChunks => self.value.ByteChunks.drop_owned(), + _ => {} + } + } + } +} diff --git a/src/hyperlight_guest_capi/src/types/vec.rs b/src/hyperlight_guest_capi/src/types/vec.rs index aff68cef4..94b76f81e 100644 --- a/src/hyperlight_guest_capi/src/types/vec.rs +++ b/src/hyperlight_guest_capi/src/types/vec.rs @@ -29,6 +29,13 @@ pub struct FfiVec { } impl FfiVec { + pub(crate) fn from_mut_slice(value: &mut [u8]) -> Self { + Self { + data: value.as_mut_ptr(), + len: value.len(), + } + } + /// Creates a new `FfiVec` from the given Vec without copying memory. /// # Safety /// The caller must later reclaim memory by calling `into_vec`, otherwise memory will be leaked. @@ -55,21 +62,15 @@ impl FfiVec { res } - /// Copies the contents of `self` to a new independent Vec. + /// Copies the contents of `self` to a new independent `Vec`. /// # Safety - /// Self must have been obtained using `from_vec`, and must be in its original state (i.e. not modified). + /// `data` must reference `len` readable bytes when `len` is nonzero. pub unsafe fn copy_to_vec(&self) -> Vec { - // deconstruct - let slice = unsafe { slice::from_raw_parts_mut(self.data, self.len) }; - let boxed: Box<[u8]> = unsafe { Box::from_raw(slice) }; - let original = boxed.into_vec(); - // clone - let clone = original.clone(); - // reverse deconstruct - let boxed = original.into_boxed_slice(); - let leaked = Box::into_raw(boxed); - assert_eq!(self.data, leaked as *mut u8); - assert_eq!(self.len, leaked.len()); - clone + if self.len == 0 { + Vec::new() + } else { + // SAFETY: required by the caller. + unsafe { slice::from_raw_parts(self.data, self.len) }.to_vec() + } } } diff --git a/src/hyperlight_guest_macro/src/lib.rs b/src/hyperlight_guest_macro/src/lib.rs index 6f25119c2..859b4fa2c 100644 --- a/src/hyperlight_guest_macro/src/lib.rs +++ b/src/hyperlight_guest_macro/src/lib.rs @@ -204,12 +204,12 @@ pub fn main(_attr: TokenStream, item: TokenStream) -> TokenStream { /// use hyperlight_guest::error::Result; /// use hyperlight_guest::bail; /// use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall; -/// use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result; +/// use hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue; /// #[dispatch] -/// fn dispatch(fc: FunctionCall) -> Result> { +/// fn dispatch(fc: FunctionCall) -> Result { /// let name = &fc.function_name; /// if name == "greet" { -/// return Ok(get_flatbuffer_result("Hello, world!")); +/// return Ok(ReturnValue::String("Hello, world!".into())); /// } /// bail!("Unknown function: {name}"); /// } @@ -241,9 +241,9 @@ pub fn dispatch(_attr: TokenStream, item: TokenStream) -> TokenStream { const _: () = { mod wrapper { - use #crate_name::__private::{FunctionCall, HyperlightGuestError, Vec}; + use #crate_name::__private::{FunctionCall, HyperlightGuestError, ReturnValue}; #[unsafe(no_mangle)] - pub fn guest_dispatch_function(function_call: FunctionCall) -> ::core::result::Result, HyperlightGuestError> { + pub fn guest_dispatch_function(function_call: FunctionCall) -> ::core::result::Result { super::#ident(function_call) } } diff --git a/src/hyperlight_guest_tracing/src/lib.rs b/src/hyperlight_guest_tracing/src/lib.rs index 6fae94925..bc27912cf 100644 --- a/src/hyperlight_guest_tracing/src/lib.rs +++ b/src/hyperlight_guest_tracing/src/lib.rs @@ -185,26 +185,15 @@ mod trace { } } - /// Returns information about the current trace state needed by the host to read the spans. + /// Returns information about the current trace state needed by the host. + /// + /// Returns `None` if tracing code already holds the state lock. Exception and + /// abort paths can then proceed without the pending trace data. pub fn serialized_data() -> Option<(u64, u64)> { if let Some(w) = GUEST_STATE.get() && let Some(state_mutex) = w.upgrade() { - // We want to protect against re-entrancy issues produced by tracing code that locks - // the state and then causes an exception that tries to lock the state again. - // - // For example: - // - 1. A span is created, locking the state - // - 2. An exception occurs while the span is being created (e.g. not enough memory, etc.) - // - 3. The exception handler uses the tracing API to send the trace data to the host - // or just create spans/events for logging purposes. - // - 4. The tracing API tries to lock the state again, causing a deadlock. - // To avoid this, we use try_lock and if we cannot acquire the lock, we panic to signal - // the issue. - let state = state_mutex - .try_lock() - .expect("Unable to lock GuestState in `serialized_data`"); - + let state = state_mutex.try_lock()?; state.serialized_data() } else { None diff --git a/src/hyperlight_host/benches/benchmarks.rs b/src/hyperlight_host/benches/benchmarks.rs index 1982244bb..80dbe37b6 100644 --- a/src/hyperlight_host/benches/benchmarks.rs +++ b/src/hyperlight_host/benches/benchmarks.rs @@ -18,11 +18,15 @@ use std::sync::{Arc, Barrier, Mutex}; use std::thread; use std::time::{Duration, Instant}; +use anyhow::{Result, bail}; use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; use flatbuffers::FlatBufferBuilder; +use hyperlight_common::flatbuffer_wrappers::ExternalValueSource; use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType}; -use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType}; +use hyperlight_common::flatbuffer_wrappers::function_types::{Bytes, ParameterValue, ReturnType}; use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity; +use hyperlight_common::transport::ExternalValues; +use hyperlight_common::vmem::PAGE_SIZE; use hyperlight_host::GuestBinary; use hyperlight_host::mem::shared_mem::ExclusiveSharedMemory; use hyperlight_host::sandbox::{MultiUseSandbox, SandboxConfiguration, UninitializedSandbox}; @@ -407,13 +411,15 @@ fn guest_call_benchmark_large_param(c: &mut Criterion) { group.bench_function("guest_call_with_large_parameters", |b| { const SIZE: usize = 50 * 1024 * 1024; // 50 MB + const MIB: usize = 1024 * 1024; let large_vec = vec![0u8; SIZE]; let large_string = String::from_utf8(large_vec.clone()).unwrap(); let mut config = SandboxConfiguration::default(); - config.set_input_data_size(2 * SIZE + (1024 * 1024)); // 2 * SIZE + 1 MB, to allow 1MB for the rest of the serialized function call + config.set_h2g_buffer_size(4 * MIB); + config.set_h2g_pool_pages((2 * SIZE + 8 * MIB).div_ceil(PAGE_SIZE)); config.set_heap_size(SIZE as u64 * 15); - config.set_scratch_size(6 * SIZE + 4 * (1024 * 1024)); // Big enough for the IO data regions and enough of the heap to be used + config.set_scratch_size(9 * SIZE); let sandbox = UninitializedSandbox::new( GuestBinary::FilePath(simple_guest_as_pathbuf()), @@ -436,51 +442,129 @@ fn guest_call_benchmark_large_param(c: &mut Criterion) { } // ============================================================================ -// Benchmark Category: Serialization +// Benchmark Category: Function Call Codec // ============================================================================ -fn function_call_serialization_benchmark(c: &mut Criterion) { - let mut group = c.benchmark_group("function_call_serialization"); +enum BenchExternalValue<'a> { + Bytes(&'a [u8]), + Chunks(&'a [Bytes]), +} + +struct BenchExternalSource<'a> { + value: Option>, +} + +impl<'a> BenchExternalSource<'a> { + fn new(value: BenchExternalValue<'a>) -> Self { + Self { value: Some(value) } + } +} + +impl ExternalValueSource for BenchExternalSource<'_> { + fn take_bytes(&mut self, length: usize) -> Result> { + let Some(BenchExternalValue::Bytes(value)) = self.value.take() else { + bail!("expected external bytes"); + }; + if value.len() != length { + bail!( + "external byte length mismatch: expected {length}, got {}", + value.len() + ); + } + Ok(value.to_vec()) + } + + fn take_chunks(&mut self, length: usize) -> Result> { + let Some(BenchExternalValue::Chunks(value)) = self.value.take() else { + bail!("expected external byte chunks"); + }; + let actual = value.iter().map(Bytes::len).sum::(); + if actual != length { + bail!("external chunk length mismatch: expected {length}, got {actual}"); + } + Ok(value.to_vec()) + } + + fn finish(&mut self) -> Result<()> { + if self.value.is_some() { + bail!("external value was not consumed"); + } + Ok(()) + } +} - let function_call = FunctionCall::new( +fn codec_benchmark_call(parameter: ParameterValue) -> FunctionCall { + FunctionCall::new( "TestFunction".to_string(), Some(vec![ - ParameterValue::VecBytes(vec![1; 10 * 1024 * 1024]), - ParameterValue::String(String::from_utf8(vec![2; 10 * 1024 * 1024]).unwrap()), + parameter, + ParameterValue::String("argument".to_string()), ParameterValue::Int(42), - ParameterValue::UInt(100), - ParameterValue::Long(1000), - ParameterValue::ULong(2000), - ParameterValue::Float(521521.53), - ParameterValue::Double(432.53), ParameterValue::Bool(true), - ParameterValue::VecBytes(vec![1; 10 * 1024 * 1024]), - ParameterValue::String(String::from_utf8(vec![2; 10 * 1024 * 1024]).unwrap()), ]), FunctionCallType::Guest, ReturnType::Int, - ); + ) +} + +fn function_call_codec_benchmark(c: &mut Criterion) { + const PAYLOAD_SIZE: usize = 10 * 1024 * 1024; + const CHUNK_SIZE: usize = 256 * 1024; + + let vec_bytes = vec![1; PAYLOAD_SIZE]; + let byte_chunks = (0..PAYLOAD_SIZE / CHUNK_SIZE) + .map(|_| Bytes::from(vec![1; CHUNK_SIZE])) + .collect::>(); + + let vec_call = codec_benchmark_call(ParameterValue::VecBytes(vec_bytes.clone())); + let chunk_call = codec_benchmark_call(ParameterValue::ByteChunks(byte_chunks.clone())); + let mut group = c.benchmark_group("function_call_codec"); + + for (name, function_call) in [("vec_bytes", &vec_call), ("byte_chunks", &chunk_call)] { + group.bench_function(BenchmarkId::new("encode_control", name), |b| { + b.iter(|| { + let estimated_capacity = estimate_flatbuffer_capacity( + &function_call.function_name, + function_call.parameters.as_deref().unwrap_or_default(), + ); + let mut builder = FlatBufferBuilder::with_capacity(estimated_capacity); + let mut exts = ExternalValues::new(); - group.bench_function("serialize_function_call", |b| { + let control = function_call.encode(&mut builder, &mut exts).unwrap(); + std::hint::black_box((control, exts.total_len())); + }); + }); + } + + let mut builder = FlatBufferBuilder::new(); + let mut external_values = ExternalValues::new(); + + let vec_control = vec_call + .encode(&mut builder, &mut external_values) + .unwrap() + .to_vec(); + + group.bench_function("decode_vec_bytes_copy", |b| { b.iter(|| { - // We specifically want to include the time to estimate the capacity in this benchmark - let estimated_capacity = estimate_flatbuffer_capacity( - function_call.function_name.as_str(), - function_call.parameters.as_deref().unwrap_or(&[]), - ); - let mut builder = FlatBufferBuilder::with_capacity(estimated_capacity); - let serialized: &[u8] = function_call.encode(&mut builder); - std::hint::black_box(serialized); + let mut src = BenchExternalSource::new(BenchExternalValue::Bytes(&vec_bytes)); + let function_call = FunctionCall::decode(&vec_control, &mut src).unwrap(); + std::hint::black_box(function_call); }); }); - group.bench_function("deserialize_function_call", |b| { - let mut builder = FlatBufferBuilder::new(); - let bytes = function_call.clone().encode(&mut builder); + let mut builder = FlatBufferBuilder::new(); + let mut external_values = ExternalValues::new(); + + let chunk_control = chunk_call + .encode(&mut builder, &mut external_values) + .unwrap() + .to_vec(); + group.bench_function("decode_byte_chunks_owner_backed", |b| { b.iter(|| { - let deserialized: FunctionCall = bytes.try_into().unwrap(); - std::hint::black_box(deserialized); + let mut src = BenchExternalSource::new(BenchExternalValue::Chunks(&byte_chunks)); + let function_call = FunctionCall::decode(&chunk_control, &mut src).unwrap(); + std::hint::black_box(function_call); }); }); @@ -496,7 +580,7 @@ fn sample_workloads_benchmark(c: &mut Criterion) { fn bench_24k_in_8k_out(b: &mut criterion::Bencher, guest_path: std::path::PathBuf) { let mut cfg = SandboxConfiguration::default(); - cfg.set_input_data_size(25 * 1024); + cfg.set_h2g_pool_pages(8); let mut sandbox = UninitializedSandbox::new(GuestBinary::FilePath(guest_path), Some(cfg)) .unwrap() @@ -725,7 +809,7 @@ criterion_group! { guest_calls_benchmark, snapshots_benchmark, guest_call_benchmark_large_param, - function_call_serialization_benchmark, + function_call_codec_benchmark, sample_workloads_benchmark, shared_memory_benchmark, snapshot_file_benchmark diff --git a/src/hyperlight_host/src/error.rs b/src/hyperlight_host/src/error.rs index c6738374d..3b52a0b66 100644 --- a/src/hyperlight_host/src/error.rs +++ b/src/hyperlight_host/src/error.rs @@ -28,6 +28,7 @@ use crossbeam_channel::{RecvError, SendError}; use flatbuffers::InvalidFlatbuffer; use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnValue}; use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; +use hyperlight_common::virtq::VirtqError; use thiserror::Error; use crate::hypervisor::hyperlight_vm::HyperlightVmError; @@ -68,6 +69,10 @@ pub enum HyperlightError { #[error("{0}")] Error(String), + /// The virtqueue transport cannot be used safely. + #[error("Virtqueue transport error: {0}")] + TransportError(String), + /// Execution violation #[error("Non-executable address {0:#x} tried to be executed")] ExecutionAccessViolation(u64), @@ -307,6 +312,12 @@ impl From for HyperlightError { } } +impl From for HyperlightError { + fn from(error: VirtqError) -> Self { + Self::TransportError(error.to_string()) + } +} + impl From<&str> for HyperlightError { fn from(s: &str) -> Self { HyperlightError::Error(s.to_string()) @@ -348,6 +359,7 @@ impl HyperlightError { | HyperlightError::ExecutionAccessViolation(_) | HyperlightError::MemoryAccessViolation(_, _, _) | HyperlightError::MemoryRegionSizeMismatch(_, _, _) + | HyperlightError::TransportError(_) // HyperlightVmError::Restore is already handled manually in restore(), but we mark it // as poisoning here too for defense in depth. | HyperlightError::HyperlightVmError(HyperlightVmError::Restore(_)) => true, @@ -448,6 +460,14 @@ mod tests { }; use crate::sandbox::outb::HandleOutbError; + #[test] + fn virtq_error_converts_to_poisoning_transport_error() { + let error = HyperlightError::from(VirtqError::InvalidState); + + assert!(matches!(error, HyperlightError::TransportError(_))); + assert!(error.is_poison_error()); + } + /// Test that ExecutionCancelledByHost promotes to HyperlightError::ExecutionCanceledByHost #[test] fn test_promote_execution_cancelled_by_host() { diff --git a/src/hyperlight_host/src/func/mod.rs b/src/hyperlight_host/src/func/mod.rs index 426897e9f..b5aab1b76 100644 --- a/src/hyperlight_host/src/func/mod.rs +++ b/src/hyperlight_host/src/func/mod.rs @@ -29,6 +29,8 @@ pub(crate) mod host_functions; /// Re-export for `HostFunction` trait pub use host_functions::{HostFunction, Registerable}; +/// Re-export for chunk-preserving byte values +pub use hyperlight_common::flatbuffer_wrappers::function_types::Bytes; /// Re-export for `ParameterType` enum pub use hyperlight_common::flatbuffer_wrappers::function_types::ParameterType; /// Re-export for `ParameterValue` enum diff --git a/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs b/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs index 372483e3c..eb870ece2 100644 --- a/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs +++ b/src/hyperlight_host/src/hypervisor/hyperlight_vm/x86_64.rs @@ -1482,14 +1482,9 @@ mod tests { // Test VM Setup // ========================================================================== - /// Creates a test VM with the given code. This is the shared setup logic used by - /// both `hyperlight_vm()` and `create_test_vm_context()`. - fn create_test_vm_context(code: &[u8]) -> TestVmContext { - let config: SandboxConfiguration = Default::default(); - #[cfg(any(crashdump, gdb))] - let rt_cfg: SandboxRuntimeConfig = Default::default(); - - let mut layout = SandboxMemoryLayout::new(config, code.len(), 4096, None).unwrap(); + fn create_test_layout(code_size: usize) -> (SandboxMemoryLayout, Box<[u8]>) { + let config = SandboxConfiguration::default(); + let mut layout = SandboxMemoryLayout::new(config, code_size, 4096, None).unwrap(); let pt_base_gpa = layout.get_pt_base_gpa(); let pt_buf = GuestPageTableBuffer::new(pt_base_gpa as usize); @@ -1515,7 +1510,6 @@ mod tests { unsafe { vmem::map(&pt_buf, mapping) }; } - // Map the scratch region at the top of the address space let scratch_size = config.get_scratch_size(); let scratch_gpa = hyperlight_common::layout::scratch_base_gpa(scratch_size); let scratch_gva = hyperlight_common::layout::scratch_base_gva(scratch_size); @@ -1526,13 +1520,24 @@ mod tests { kind: MappingKind::Basic(BasicMapping { readable: true, writable: true, - executable: true, // Match regular codepath (map_specials) + executable: true, }), }; unsafe { vmem::map(&pt_buf, scratch_mapping) }; let pt_bytes = pt_buf.into_bytes(); layout.set_pt_size(pt_bytes.len()).unwrap(); + (layout, pt_bytes) + } + + /// Creates a test VM with the given code. This is the shared setup logic used by + /// both `hyperlight_vm()` and `create_test_vm_context()`. + fn create_test_vm_context(code: &[u8]) -> TestVmContext { + let config: SandboxConfiguration = Default::default(); + #[cfg(any(crashdump, gdb))] + let rt_cfg: SandboxRuntimeConfig = Default::default(); + + let (layout, pt_bytes) = create_test_layout(code.len()); let mem_size = layout.get_memory_size().unwrap(); let mut snapshot_contents = vec![0u8; mem_size]; @@ -2170,7 +2175,7 @@ mod tests { /// Extended test context for FXSAVE tests that need to read memory at a specific offset. struct FxsaveTestContext { ctx: TestVmContext, - /// Offset in shared memory where FXSAVE data is stored (output_data region) + /// Offset in scratch memory where FXSAVE data is stored. fxsave_offset: usize, } @@ -2205,18 +2210,19 @@ mod tests { } } - /// Creates VM with guest code that: dirtys FPU (if flag==0), does FXSAVE to buffer, sets flag=1. - /// Uses output_data region for FXSAVE buffer (like regular guest output), scratch for stack. + /// Creates VM with guest code that dirties FPU once and writes FXSAVE state. fn hyperlight_vm_with_mem_mgr_fxsave() -> FxsaveTestContext { use iced_x86::code_asm::*; - // Compute fixed addresses for FXSAVE buffer and flag. - // These are in the output_data region which starts at a known offset. - // We use a default SandboxConfiguration to get the same layout as create_test_vm_context. - let config: SandboxConfiguration = Default::default(); - let layout = SandboxMemoryLayout::new(config, 512, 4096, None).unwrap(); - let fxsave_offset = layout.get_output_data_buffer_scratch_host_offset(); - let fxsave_gva = layout.get_output_data_buffer_gva(); + const CODE_SIZE_BOUND: usize = 512; + + let (layout, _) = create_test_layout(CODE_SIZE_BOUND); + let scratch_size = layout.get_scratch_size(); + 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 fxsave_gpa = layout.get_first_free_scratch_gpa(); + let fxsave_offset = usize::try_from(fxsave_gpa - scratch_base_gpa).unwrap(); + let fxsave_gva = scratch_base_gva + fxsave_offset as u64; let flag_gva = fxsave_gva + 512; let mut a = CodeAssembler::new(64).unwrap(); @@ -2273,9 +2279,11 @@ mod tests { a.hlt().unwrap(); let code = a.assemble(0).unwrap(); + assert!(code.len() <= CODE_SIZE_BOUND); // Reuse common test setup - initialise() will run the code let ctx = create_test_vm_context(&code); + assert_eq!(ctx.hshm.layout.get_first_free_scratch_gpa(), fxsave_gpa); FxsaveTestContext { ctx, fxsave_offset } } diff --git a/src/hyperlight_host/src/mem/layout.rs b/src/hyperlight_host/src/mem/layout.rs index 6422b9b11..8561fd017 100644 --- a/src/hyperlight_host/src/mem/layout.rs +++ b/src/hyperlight_host/src/mem/layout.rs @@ -47,22 +47,23 @@ limitations under the License. //! //! 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: +//! memory, although the transport arena and copied page tables have +//! fixed positions: //! //! +-------------------------------------------+ (top of physical memory) //! | Exception Stack, Metadata | //! +-------------------------------------------+ (1 page below) //! | Scratch Memory | //! +-------------------------------------------+ -//! | Output Data | +//! | Guest Page Tables | //! +-------------------------------------------+ -//! | Input Data | +//! | Transport Arena | //! +-------------------------------------------+ (scratch size) use std::fmt::Debug; use std::mem::size_of; +use hyperlight_common::layout::{QueueDims, TransportArena}; use hyperlight_common::mem::{HyperlightPEB, PAGE_SIZE_USIZE}; use tracing::{Span, instrument}; @@ -248,10 +249,6 @@ impl ResolvedGpa { #[derive(Copy, Clone)] pub(crate) struct SandboxMemoryLayout { - /// Input data buffer size (from SandboxConfiguration). - input_data_size: usize, - /// Output data buffer size (from SandboxConfiguration). - output_data_size: usize, /// The heap size of this sandbox. heap_size: usize, /// The size of the guest code section. @@ -262,6 +259,18 @@ pub(crate) struct SandboxMemoryLayout { init_data_permissions: Option, /// The size of the scratch region in physical memory. scratch_size: usize, + /// Number of descriptors in the G2H virtqueue. + g2h_queue_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, /// 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 @@ -287,15 +296,13 @@ impl Debug for SandboxMemoryLayout { "Init Data Size", &format_args!("{:#x}", self.init_data_size), ) - .field( - "Input Data Size", - &format_args!("{:#x}", self.input_data_size), - ) - .field( - "Output Data Size", - &format_args!("{:#x}", self.output_data_size), - ) .field("Scratch Size", &format_args!("{:#x}", self.scratch_size)) + .field("G2H Queue Size", &self.g2h_queue_size) + .field("H2G Queue Size", &self.h2g_queue_size) + .field("G2H Buffer Size", &self.g2h_buffer_size) + .field("H2G Buffer Size", &self.h2g_buffer_size) + .field("G2H Pool Pages", &self.g2h_pool_pages) + .field("H2G Pool Pages", &self.h2g_pool_pages) .field("Snapshot Size", &format_args!("{:#x}", self.snapshot_size)) .field("PT Size", &format_args!("{:#x}", self.pt_size.unwrap_or(0))) .field( @@ -317,14 +324,11 @@ impl Debug for SandboxMemoryLayout { } impl SandboxMemoryLayout { - /// Whether `other` has the same layout configuration as `self`, - /// i.e. the fields that come from the guest binary and the - /// `SandboxConfiguration`. `snapshot_size` and `pt_size` are - /// excluded because they are outputs of building a snapshot blob - /// (the compacted data size and the size of the rebuilt - /// page-table tail), not configuration inputs, so they differ - /// between the sandbox's live layout and any snapshot taken - /// from it. + /// Whether `other` has the same active memory layout as `self`. + /// + /// Transport configuration participates because it determines fixed scratch + /// addresses. `snapshot_size` and `pt_size` are outputs of building a + /// snapshot blob, so they may differ between live and captured layouts. /// /// TODO: separate/remove snapshot_size and pt_size from this struct. pub(crate) fn is_compatible_with(&self, other: &Self) -> bool { @@ -332,23 +336,31 @@ impl SandboxMemoryLayout { // `SandboxMemoryLayout` fails to compile here, forcing the // author to decide whether it participates in compatibility. let Self { - input_data_size, - output_data_size, heap_size, code_size, init_data_size, init_data_permissions, scratch_size, + g2h_queue_size, + h2g_queue_size, + g2h_buffer_size, + h2g_buffer_size, + g2h_pool_pages, + h2g_pool_pages, snapshot_size: _, pt_size: _, } = self; - *input_data_size == other.input_data_size - && *output_data_size == other.output_data_size - && *heap_size == other.heap_size + *heap_size == other.heap_size && *code_size == other.code_size && *init_data_size == other.init_data_size && *init_data_permissions == other.init_data_permissions && *scratch_size == other.scratch_size + && *g2h_queue_size == other.g2h_queue_size + && *h2g_queue_size == other.h2g_queue_size + && *g2h_buffer_size == other.g2h_buffer_size + && *h2g_buffer_size == other.h2g_buffer_size + && *g2h_pool_pages == other.g2h_pool_pages + && *h2g_pool_pages == other.h2g_pool_pages } /// The maximum amount of memory a single sandbox will be allowed. @@ -361,9 +373,6 @@ impl SandboxMemoryLayout { /// The base address of the sandbox's memory. pub(crate) const BASE_ADDRESS: usize = 0x1000; - // the offset into a sandbox's input/output buffer where the stack starts - pub(crate) const STACK_POINTER_SIZE_BYTES: u64 = 8; - /// Create a new `SandboxMemoryLayout` with the given /// `SandboxConfiguration`, code size and stack/heap size. #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] @@ -378,37 +387,46 @@ impl SandboxMemoryLayout { if scratch_size > Self::MAX_MEMORY_SIZE { return Err(MemoryRequestTooBig(scratch_size, Self::MAX_MEMORY_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); + if !scratch_size.is_multiple_of(PAGE_SIZE_USIZE) { + return Err(new_error!( + "scratch size {scratch_size} must be a multiple of {PAGE_SIZE_USIZE}" + )); + } + 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 min_scratch_size = hyperlight_common::layout::min_scratch_size( + g2h_queue_size, + h2g_queue_size, + g2h_pool_pages, + h2g_pool_pages, + ); if scratch_size < min_scratch_size { return Err(MemoryRequestTooSmall(scratch_size, min_scratch_size)); } let mut ret = Self { - input_data_size, - output_data_size, heap_size, code_size, init_data_size, init_data_permissions, pt_size: None, scratch_size, + g2h_queue_size, + h2g_queue_size, + g2h_buffer_size, + h2g_buffer_size, + g2h_pool_pages, + h2g_pool_pages, snapshot_size: 0, }; ret.set_snapshot_size(ret.get_memory_size()?); Ok(ret) } - pub(crate) fn input_data_size(&self) -> usize { - self.input_data_size - } - - pub(crate) fn output_data_size(&self) -> usize { - self.output_data_size - } - pub(crate) fn heap_size(&self) -> usize { self.heap_size } @@ -429,6 +447,48 @@ impl SandboxMemoryLayout { self.scratch_size } + #[allow(dead_code)] + pub(crate) fn get_g2h_queue_size(&self) -> usize { + self.g2h_queue_size + } + + #[allow(dead_code)] + pub(crate) fn get_h2g_queue_size(&self) -> usize { + self.h2g_queue_size + } + + #[allow(dead_code)] + pub(crate) fn get_g2h_buffer_size(&self) -> usize { + self.g2h_buffer_size + } + + #[allow(dead_code)] + pub(crate) fn get_h2g_buffer_size(&self) -> usize { + self.h2g_buffer_size + } + + #[allow(dead_code)] + pub(crate) fn get_g2h_pool_pages(&self) -> usize { + self.g2h_pool_pages + } + + #[allow(dead_code)] + pub(crate) fn get_h2g_pool_pages(&self) -> usize { + self.h2g_pool_pages + } + + #[allow(clippy::expect_used)] // `new` validates these dimensions. + pub(crate) fn get_g2h_queue_dims(&self) -> QueueDims { + QueueDims::new(self.g2h_queue_size, self.g2h_pool_pages) + .expect("validated G2H queue dimensions") + } + + #[allow(clippy::expect_used)] // `new` validates these dimensions. + pub(crate) fn get_h2g_queue_dims(&self) -> QueueDims { + QueueDims::new(self.h2g_queue_size, self.h2g_pool_pages) + .expect("validated H2G queue dimensions") + } + /// Guest-visible prefix size of the snapshot blob. pub(crate) fn snapshot_size(&self) -> usize { self.snapshot_size @@ -452,10 +512,12 @@ impl SandboxMemoryLayout { /// independent field and must be set separately. pub(crate) fn set_pt_size(&mut self, size: usize) -> Result<()> { let min_fixed_scratch = hyperlight_common::layout::min_scratch_size( - self.input_data_size, - self.output_data_size, + self.g2h_queue_size, + self.h2g_queue_size, + self.g2h_pool_pages, + self.h2g_pool_pages, ); - let min_scratch = min_fixed_scratch + size; + let min_scratch = min_fixed_scratch.saturating_add(size); if self.scratch_size < min_scratch { return Err(MemoryRequestTooSmall(self.scratch_size, min_scratch)); } @@ -574,14 +636,6 @@ impl SandboxMemoryLayout { let guest_base = Self::BASE_ADDRESS as u64; let peb = HyperlightPEB { - input_stack: GuestMemoryRegion { - size: self.input_data_size as u64, - ptr: self.get_input_data_buffer_gva(), - }, - output_stack: GuestMemoryRegion { - size: self.output_data_size as u64, - ptr: self.get_output_data_buffer_gva(), - }, init_data: GuestMemoryRegion { size: (self.get_unaligned_memory_size() - self.init_data_offset()) as u64, ptr: guest_base + self.init_data_offset() as u64, @@ -606,11 +660,6 @@ impl SandboxMemoryLayout { })?; dst.copy_from_slice(bytes); - // The input and output data regions do not have their layout - // initialised here, because they are in the scratch - // region---they are instead set in - // [`SandboxMemoryManager::update_scratch_bookkeeping`]. - Ok(()) } @@ -682,31 +731,10 @@ impl SandboxMemoryLayout { Self::BASE_ADDRESS + self.guest_code_offset() } - /// Guest virtual address of the start of output data. - pub(crate) fn get_output_data_buffer_gva(&self) -> u64 { - hyperlight_common::layout::scratch_base_gva(self.scratch_size) + self.input_data_size as u64 - } - - /// Offset into the host scratch buffer of the start of the output data. - pub(crate) fn get_output_data_buffer_scratch_host_offset(&self) -> usize { - self.input_data_size - } - - /// Guest virtual address of the start of input data. - fn get_input_data_buffer_gva(&self) -> u64 { - hyperlight_common::layout::scratch_base_gva(self.scratch_size) - } - - /// Offset into the host scratch buffer of the start of the input data. - pub(crate) fn get_input_data_buffer_scratch_host_offset(&self) -> usize { - 0 - } - /// 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(hyperlight_common::vmem::PAGE_SIZE) + self.get_transport_arena().size() } /// Base GPA to which the page tables are eagerly copied on restore. @@ -715,12 +743,24 @@ impl SandboxMemoryLayout { + self.get_pt_base_scratch_offset() as u64 } - /// First GPA of the scratch region the host has not used for - /// something else. + /// First GPA available to the guest scratch allocator. pub(crate) fn get_first_free_scratch_gpa(&self) -> u64 { self.get_pt_base_gpa() + self.pt_size.unwrap_or(0) as u64 } + /// Exact transport placement in the fixed scratch prefix. + #[allow(clippy::expect_used)] // The base and dimensions are validated by `new`. + pub(crate) fn get_transport_arena(&self) -> TransportArena { + let base_gpa = hyperlight_common::layout::scratch_base_gpa(self.scratch_size); + + TransportArena::new( + base_gpa, + self.get_g2h_queue_dims(), + self.get_h2g_queue_dims(), + ) + .expect("validated virtqueue arena dimensions") + } + /// Total size of guest memory in `self`'s memory layout. fn get_unaligned_memory_size(&self) -> usize { self.init_data_offset() + self.init_data_size @@ -778,12 +818,52 @@ mod tests { ); } + #[test] + fn transport_memory_is_part_of_minimum_scratch_size() { + let mut cfg = SandboxConfiguration::default(); + let minimum = hyperlight_common::layout::min_scratch_size( + cfg.get_g2h_queue_size(), + cfg.get_h2g_queue_size(), + cfg.get_g2h_pool_pages(), + cfg.get_h2g_pool_pages(), + ); + cfg.set_scratch_size(minimum); + let mut layout = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap(); + + assert!(matches!( + layout.set_pt_size(PAGE_SIZE_USIZE), + Err(MemoryRequestTooSmall(..)) + )); + } + + #[test] + fn transport_minimum_rejects_capacity_overflow() { + let mut cfg = SandboxConfiguration::default(); + cfg.set_g2h_pool_pages(usize::MAX); + let layout = SandboxMemoryLayout::new(cfg, 4096, 0, None); + assert!(matches!(layout, Err(MemoryRequestTooSmall(_, usize::MAX)))); + } + + #[test] + fn 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_USIZE}", + SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 1 + ) + ); + } + #[test] fn test_max_memory_sandbox() { let mut cfg = SandboxConfiguration::default(); // scratch_size exceeds 16 GiB limit cfg.set_scratch_size(17 * 1024 * 1024 * 1024); - cfg.set_input_data_size(16 * 1024 * 1024 * 1024); let layout = SandboxMemoryLayout::new(cfg, 4096, 4096, None); assert!(matches!(layout.unwrap_err(), MemoryRequestTooBig(..))); } @@ -818,12 +898,16 @@ mod tests { // Each mutation must independently break compatibility. let mutators: &[fn(&mut SandboxMemoryLayout)] = &[ - |l| l.input_data_size += PAGE_SIZE_USIZE, - |l| l.output_data_size += PAGE_SIZE_USIZE, |l| l.heap_size += PAGE_SIZE_USIZE, |l| l.code_size += PAGE_SIZE_USIZE, |l| l.init_data_size += PAGE_SIZE_USIZE, |l| l.scratch_size += PAGE_SIZE_USIZE, + |l| l.g2h_queue_size *= 2, + |l| l.h2g_queue_size *= 2, + |l| l.g2h_buffer_size *= 2, + |l| l.h2g_buffer_size *= 2, + |l| l.g2h_pool_pages += 1, + |l| l.h2g_pool_pages += 1, |l| { l.init_data_permissions = Some(MemoryRegionFlags::READ); }, @@ -898,10 +982,8 @@ mod tests { ); let mut cfg = SandboxConfiguration::default(); - 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(0x30000); let layout = SandboxMemoryLayout::new(cfg, 0x1000, 0, None).unwrap(); pin_eq!(layout.guest_code_offset(), 0); @@ -911,31 +993,26 @@ 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(), 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(), 0x2000); - pin_eq!(layout.get_pt_base_scratch_offset(), 0x4000); + pin_eq!(layout.get_pt_base_scratch_offset(), 0x15000); - // The output buffer sits one input buffer past the input - // buffer in the guest's scratch view. - pin_eq!( - layout.get_output_data_buffer_gva() - layout.get_input_data_buffer_gva(), - 0x2000 - ); + 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, 0); + pin_eq!(arena.h2g_ring_addr() - scratch_base_gpa, 0x410); + pin_eq!(arena.mbx_addr() - scratch_base_gpa, 0x618); + pin_eq!(arena.g2h_pool_addr() - scratch_base_gpa, 0x1000); + pin_eq!(arena.h2g_pool_addr() - scratch_base_gpa, 0xd000); + pin_eq!(arena.end_addr() - scratch_base_gpa, 0x15000); - // The input buffer sits at the scratch base. The page tables - // sit `get_pt_base_scratch_offset` above it. With the - // `SCRATCH_TOP` pins above, these fix the absolute addresses. - pin_eq!( - layout.get_input_data_buffer_gva() - - hyperlight_common::layout::scratch_base_gva(0x10000), - 0 - ); + // The transport arena sits at the scratch base. The page tables + // follow it. With the `SCRATCH_TOP` pins above, these fix the + // absolute addresses. 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(0x30000), + 0x15000 ); // pt_size is zero here, so the first free scratch GPA equals // the page table base. @@ -944,13 +1021,11 @@ mod tests { layout.get_pt_base_gpa() ); - // A second config with different sizes shifts the offsets off - // the first config's page boundaries. + // A second snapshot layout keeps the transport prefix fixed + // relative to its scratch base. let mut cfg = SandboxConfiguration::default(); - 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(0x40000); let layout = SandboxMemoryLayout::new(cfg, 0x3000, 0, None).unwrap(); pin_eq!(layout.guest_code_offset(), 0); @@ -960,26 +1035,23 @@ mod tests { pin_eq!(layout.init_data_offset(), 0x9000); pin_eq!(layout.get_memory_size().unwrap(), 0x9000); - pin_eq!(layout.get_scratch_size(), 0x20000); + pin_eq!(layout.get_scratch_size(), 0x40000); 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(), 0x15000); - pin_eq!( - layout.get_output_data_buffer_gva() - layout.get_input_data_buffer_gva(), - 0x4000 - ); + let arena = layout.get_transport_arena(); + let scratch_base_gpa = hyperlight_common::layout::scratch_base_gpa(0x40000); + pin_eq!(arena.g2h_ring_addr() - scratch_base_gpa, 0); + pin_eq!(arena.h2g_ring_addr() - scratch_base_gpa, 0x410); + pin_eq!(arena.mbx_addr() - scratch_base_gpa, 0x618); + pin_eq!(arena.g2h_pool_addr() - scratch_base_gpa, 0x1000); + pin_eq!(arena.h2g_pool_addr() - scratch_base_gpa, 0xd000); + pin_eq!(arena.end_addr() - scratch_base_gpa, 0x15000); pin_eq!( - layout.get_input_data_buffer_gva() - - hyperlight_common::layout::scratch_base_gva(0x20000), - 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(0x40000), + 0x15000 ); 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 c93f1cac1..19a5bc957 100644 --- a/src/hyperlight_host/src/mem/mgr.rs +++ b/src/hyperlight_host/src/mem/mgr.rs @@ -15,12 +15,12 @@ limitations under the License. */ use flatbuffers::FlatBufferBuilder; -use hyperlight_common::flatbuffer_wrappers::function_call::{ - FunctionCall, validate_guest_function_call_buffer, -}; +use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall; use hyperlight_common::flatbuffer_wrappers::function_types::FunctionCallResult; -use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData; use hyperlight_common::flatbuffer_wrappers::host_function_details::HostFunctionDetails; +use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity; +use hyperlight_common::transport::{Buf, EncodedMessage, ExternalValues, MsgKind}; +use hyperlight_common::virtq::ReplyChain; use hyperlight_common::vmem::{self, PAGE_TABLE_SIZE}; #[cfg(crashdump)] use hyperlight_common::vmem::{BasicMapping, MappingKind}; @@ -30,12 +30,13 @@ 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)] use crate::mem::memory_region::{CrashDumpRegion, MemoryRegionFlags, MemoryRegionType}; use crate::sandbox::snapshot::{NextAction, Snapshot}; -use crate::{Result, new_error}; +use crate::{HyperlightError, Result, new_error}; #[cfg(crashdump)] fn mapping_kind_to_flags(kind: &MappingKind) -> (MemoryRegionFlags, MemoryRegionType) { @@ -130,9 +131,9 @@ impl ReadonlySharedMemory { } } pub(crate) use unused_hack::SnapshotSharedMemory; + /// A struct that is responsible for laying out and managing the memory /// for a given `Sandbox`. -#[derive(Clone)] pub(crate) struct SandboxMemoryManager { /// Shared memory for the Sandbox pub(crate) shared_mem: SnapshotSharedMemory, @@ -155,6 +156,29 @@ 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, + /// Correlation ID assigned to the next guest-function call. + next_guest_cid: u32, +} + +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, + next_guest_cid: self.next_guest_cid, + } + } } /// Buffer for building guest page tables during snapshot creation. @@ -290,6 +314,9 @@ where original_entrypoint: 0, abort_buffer: Vec::new(), snapshot_count: 0, + g2h_consumer: None, + h2g_consumer: None, + next_guest_cid: 1, } } @@ -297,37 +324,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 { @@ -372,6 +368,9 @@ impl SandboxMemoryManager { original_entrypoint: self.original_entrypoint, abort_buffer: self.abort_buffer, snapshot_count: self.snapshot_count, + g2h_consumer: None, + h2g_consumer: None, + next_guest_cid: self.next_guest_cid, }; let guest_mgr = SandboxMemoryManager { shared_mem: gshm, @@ -381,94 +380,317 @@ 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, + next_guest_cid: self.next_guest_cid, }; host_mgr.update_scratch_bookkeeping()?; + + if matches!(host_mgr.next_action, NextAction::Initialise(_)) { + host_mgr.create_virtq_consumers()?; + } + Ok((host_mgr, guest_mgr)) } } impl SandboxMemoryManager { - /// 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 { - self.scratch_mem.try_pop_buffer_into::( - self.layout.get_output_data_buffer_scratch_host_offset(), - self.layout.output_data_size(), - ) - } - - /// Writes a host function call result to memory - #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] - pub(crate) fn write_response_from_host_function_call( + /// Create a snapshot with the given mapped regions. + #[allow(clippy::too_many_arguments)] + pub(crate) fn snapshot( &mut self, - res: &FunctionCallResult, - ) -> Result<()> { - let mut builder = FlatBufferBuilder::new(); - let data = res.encode(&mut builder); - - self.scratch_mem.push_buffer( - self.layout.get_input_data_buffer_scratch_host_offset(), - self.layout.input_data_size(), - data, + 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, ) } - /// Writes a guest function call to memory - #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] - pub(crate) fn write_guest_function_call(&mut self, buffer: &[u8]) -> Result<()> { - validate_guest_function_call_buffer(buffer).map_err(|e| { - new_error!( - "Guest function call buffer validation failed: {}", - e.to_string() - ) - })?; - - self.scratch_mem.push_buffer( - self.layout.get_input_data_buffer_scratch_host_offset(), - self.layout.input_data_size(), - buffer, - )?; + /// Create host consumers before the guest initializes the transport. + /// + /// The consumers begin at cursor zero and observe descriptors published + /// during the first guest entry. + fn create_virtq_consumers(&mut self) -> Result<()> { + if self.g2h_consumer.is_some() || self.h2g_consumer.is_some() { + return Err(new_error!("virtqueue consumers already exist")); + } + + let (g2h, h2g) = virtq::create_consumers(&self.layout, &self.scratch_mem)?; + self.g2h_consumer = Some(g2h); + self.h2g_consumer = Some(h2g); Ok(()) } - /// Reads a function call result from memory. - /// A function call result can be either an error or a successful return value. - #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] - pub(crate) fn get_guest_function_call_result(&mut self) -> Result { - self.scratch_mem.try_pop_buffer_into::( - self.layout.get_output_data_buffer_scratch_host_offset(), - self.layout.output_data_size(), - ) + /// Restore a captured canonical transport image against this scratch mapping. + pub(crate) fn restore_virtq(&mut self, snapshot: &virtq::VirtqSnapshot) -> 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::restore(&self.layout, &self.scratch_mem, snapshot)?; + self.g2h_consumer = Some(g2h); + self.h2g_consumer = Some(h2g); + Ok(()) } - /// Read guest log data from the `SharedMemory` contained within `self` + /// Write a guest function call into the H2G virtqueue. #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] - pub(crate) fn read_guest_log_data(&mut self) -> Result { - self.scratch_mem.try_pop_buffer_into::( - self.layout.get_output_data_buffer_scratch_host_offset(), - self.layout.output_data_size(), - ) + pub(crate) fn write_guest_function_call(&mut self, call: &FunctionCall) -> Result { + let cid = self.next_guest_cid; + let params = call.parameters.as_deref().unwrap_or_default(); + let cap = estimate_flatbuffer_capacity(&call.function_name, params); + + let mut builder = FlatBufferBuilder::with_capacity(cap); + let mut externals = ExternalValues::new(); + + let control = call.encode(&mut builder, &mut externals)?; + + let Some(msg) = EncodedMessage::new(MsgKind::Request, cid, control, externals) else { + return Err(new_error!("H2G request exceeds the wire payload limit")); + }; + + self.write_h2g_message(&msg)?; + + self.next_guest_cid = cid.wrapping_add(1); + if self.next_guest_cid == 0 { + self.next_guest_cid = 1; + } + + Ok(cid) } - pub(crate) fn clear_io_buffers(&mut self) { - // Clear the output data buffer - loop { - let Ok(_) = self.scratch_mem.try_pop_buffer_into::>( - self.layout.get_output_data_buffer_scratch_host_offset(), - self.layout.output_data_size(), - ) else { - break; + fn write_h2g_message(&mut self, message: &EncodedMessage<'_>) -> Result<()> { + let Some(consumer) = self.h2g_consumer.as_mut() else { + return Err(new_error!("H2G consumer is not attached")); + }; + + let buffer_size = self.layout.get_h2g_buffer_size(); + let buffer_count = message.total_len().div_ceil(buffer_size); + + // External bytes may become owner-backed ByteChunks retained across + // calls. If they consume every posted H2G buffer, no buffer remains + // for a control call that releases them. External payloads therefore + // require one extra chain. poll_exact_with_spare checks the chain and + // leaves it available for the next call. + let spare_buffers = match message.header().msg_kind() { + Ok(MsgKind::SnapshotCheckpoint) => 0, + Ok(_) => usize::from(message.external_len() != 0), + Err(_) => unreachable!("validated upstream"), + }; + + // H2G receive buffers are writable-only, so any readable payload is malformed. + let maybe_buffers = consumer + .poll_exact_with_spare(buffer_count, spare_buffers, 0) + .map_err(|err| HyperlightError::TransportError(format!("H2G poll failed: {err}")))?; + + let Some(buffers) = maybe_buffers else { + return Err(new_error!( + "H2G capacity cannot provide {buffer_count} buffers with {spare_buffers} spare" + )); + }; + + // The message is a contiguous sequence of bytes, but the buffers are a chain of possibly + // non contiguous slices. Write the message into the buffers in order, advancing the message + // cursor as we go. + let mut message = message.as_buf(); + + for (recv, reply) in buffers { + let ReplyChain::Writable(mut buffer) = reply else { + return Err(HyperlightError::TransportError( + "H2G receive buffer is not writable".into(), + )); }; + + if buffer.desc_count() != 1 || buffer.capacity() != buffer_size { + return Err(HyperlightError::TransportError( + "H2G receive buffer has an invalid shape".into(), + )); + } + + while message.has_remaining() && buffer.remaining() != 0 { + let written = buffer.write(message.chunk()).map_err(|err| { + HyperlightError::TransportError(format!("H2G write failed: {err}")) + })?; + + message.advance(written); + } + + consumer.complete(recv, buffer).map_err(|err| { + HyperlightError::TransportError(format!("H2G completion failed: {err}")) + })?; } - // Clear the input data buffer + + debug_assert!(!message.has_remaining()); + Ok(()) + } + + /// Read a guest function result from the G2H virtqueue. + #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] + pub(crate) fn read_h2g_result_from_g2h(&mut self, cid: u32) -> Result { + let max_recv_len = self.layout.get_g2h_queue_dims().pool_len(); + + let Some(consumer) = self.g2h_consumer.as_mut() else { + return Err(HyperlightError::TransportError( + "G2H consumer is not attached".into(), + )); + }; + loop { - let Ok(_) = self.scratch_mem.try_pop_buffer_into::>( - self.layout.get_input_data_buffer_scratch_host_offset(), - self.layout.input_data_size(), - ) else { - break; + let maybe_next = consumer.poll(max_recv_len).map_err(|err| { + HyperlightError::TransportError(format!("G2H poll failed: {err}")) + })?; + + let Some((mut recv, reply)) = maybe_next else { + return Err(HyperlightError::TransportError( + "G2H has no guest function result after halt".into(), + )); }; + + let header = virtq::read_message_header(&mut recv).map_err(|err| { + HyperlightError::TransportError(format!("Failed to read G2H result header: {err}")) + })?; + + if !matches!(&reply, ReplyChain::Ack(_)) { + return Err(HyperlightError::TransportError( + "G2H result entry has writable buffers".into(), + )); + } + + match header.msg_kind() { + Ok(MsgKind::Log) => { + if header.cid != 0 { + return Err(HyperlightError::TransportError( + "G2H log has a correlation ID".into(), + )); + } + + let log = virtq::read_guest_log_data(&mut recv).map_err(|err| { + HyperlightError::TransportError(format!("Failed to read G2H log: {err}")) + })?; + + consumer.complete(recv, reply).map_err(|err| { + HyperlightError::TransportError(format!( + "Failed to complete G2H log: {err}" + )) + })?; + + crate::sandbox::outb::emit_guest_log(&log); + } + Ok(MsgKind::Response) => { + if header.cid != cid { + return Err(HyperlightError::TransportError( + "G2H guest function result correlation ID mismatch".into(), + )); + } + + let result = virtq::read_guest_function_call_result(&mut recv); + consumer.complete(recv, reply).map_err(|err| { + HyperlightError::TransportError(format!( + "Failed to complete G2H guest function result: {err}" + )) + })?; + + return result.map_err(|err| { + HyperlightError::TransportError(format!( + "Failed to decode G2H guest function result: {err}" + )) + }); + } + Ok(kind) => { + return Err(HyperlightError::TransportError(format!( + "Expected G2H guest function result, got {kind:?}" + ))); + } + Err(kind) => { + return Err(HyperlightError::TransportError(format!( + "Unknown G2H message kind {kind:#x}" + ))); + } + } + } + } + + /// Publish an internal request for guest-side snapshot canonicalization. + /// + /// The pending marker distinguishes a completed checkpoint with no retained + /// buffers from a guest that halted without publishing mailbox status. + pub(crate) fn begin_snapshot_checkpoint(&mut self) -> Result<()> { + let offset = self.snapshot_mbx_offset()?; + self.scratch_mem.write(offset, u64::MAX.to_le_bytes())?; + + let message = EncodedMessage::new_snapshot_cp(); + self.write_h2g_message(&message) + } + + /// Reset host consumers and read the guest-side snapshot status. + /// + /// Consumer reset completes the canonical queue before the status is interpreted. + /// A retained-buffer rejection therefore leaves both queues usable. The current + /// status is only a retained slot count. + /// + /// TODO: This will change to allow the guest to publish a more detailed snapshot + /// status about what buffer ranges were retained so we can include them in the + /// snapshot. For now we simply error if the guest has retained any buffers. + pub(crate) fn finish_snapshot_checkpoint(&mut self) -> Result { + let Some(g2h) = self.g2h_consumer.as_mut() else { + return Err(new_error!("G2H consumer is not attached")); + }; + + let Some(h2g) = self.h2g_consumer.as_mut() else { + return Err(new_error!("H2G consumer is not attached")); + }; + + g2h.reset()?; + h2g.reset()?; + + let offset = self.snapshot_mbx_offset()?; + let guest_owned = u64::from_le_bytes(self.scratch_mem.read(offset)?); + + if guest_owned == u64::MAX { + return Err(HyperlightError::TransportError( + "Guest did not publish snapshot checkpoint status".to_string(), + )); } + + Ok(guest_owned) + } + + /// Get the offset of the snapshot mailbox in scratch memory. + fn snapshot_mbx_offset(&self) -> Result { + let arena = self.layout.get_transport_arena(); + Ok(usize::try_from( + arena + .mbx_addr() + .checked_sub(arena.base_addr()) + .ok_or_else(|| new_error!("Snapshot mailbox precedes transport arena"))?, + )?) } /// This function restores a memory snapshot from a given snapshot. @@ -479,6 +701,18 @@ impl SandboxMemoryManager { Option>, Option, )> { + let virtq = snapshot.virtq(); + if let Some(virtq) = virtq { + virtq.preflight(snapshot.layout())?; + } else if matches!(snapshot.next_action(), NextAction::Call(_)) { + return Err(new_error!( + "running snapshot has no canonical transport state" + )); + } + + 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, @@ -521,6 +755,11 @@ impl SandboxMemoryManager { self.original_entrypoint = snapshot.original_entrypoint(); self.update_scratch_bookkeeping()?; + if let Some(virtq) = virtq { + self.restore_virtq(virtq)?; + } else if matches!(snapshot.next_action(), NextAction::Initialise(_)) { + self.create_virtq_consumers()?; + } Ok((gsnapshot, gscratch)) } @@ -555,16 +794,36 @@ impl SandboxMemoryManager { SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET, self.snapshot_count, )?; - - // Initialise the guest input and output data buffers in - // scratch memory. TODO: remove the need for this. - self.scratch_mem.write::( - self.layout.get_input_data_buffer_scratch_host_offset(), - SandboxMemoryLayout::STACK_POINTER_SIZE_BYTES, + // 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.scratch_mem.write::( - self.layout.get_output_data_buffer_scratch_host_offset(), - SandboxMemoryLayout::STACK_POINTER_SIZE_BYTES, + 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(), )?; // Copy page tables from `shared_mem` into scratch. PT bytes @@ -764,17 +1023,211 @@ impl SandboxMemoryManager { } #[cfg(test)] -#[cfg(target_arch = "x86_64")] mod tests { + use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCallType; + use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterValue, ReturnType}; + use hyperlight_common::transport::{ + MsgHeader, SIZE_PREFIX_LEN, size_prefix_payload_len, size_prefixed_len, + }; + use hyperlight_common::virtq::DescFlags; + #[cfg(target_arch = "x86_64")] use hyperlight_testing::sandbox_sizes::{LARGE_HEAP_SIZE, MEDIUM_HEAP_SIZE, SMALL_HEAP_SIZE}; + #[cfg(target_arch = "x86_64")] use hyperlight_testing::simple_guest_as_pathbuf; + use super::*; + #[cfg(target_arch = "x86_64")] use crate::GuestBinary; + use crate::mem::virtq::tests::{H2G_BUFFER_SIZE, TestVirtq, memory_layout}; + #[cfg(target_arch = "x86_64")] use crate::sandbox::SandboxConfiguration; - use crate::sandbox::snapshot::Snapshot; - /// Build a Snapshot for the given configuration and verify the + fn manager(queue: &TestVirtq) -> SandboxMemoryManager { + #[cfg(not(unshared_snapshot_mem))] + let shared_mem = + ReadonlySharedMemory::from_bytes(&vec![0; vmem::PAGE_SIZE], vmem::PAGE_SIZE).unwrap(); + + #[cfg(unshared_snapshot_mem)] + let shared_mem = ExclusiveSharedMemory::new(vmem::PAGE_SIZE) + .unwrap() + .build() + .0; + + let mut mgr = SandboxMemoryManager::new( + memory_layout(), + shared_mem, + queue.scratch.clone(), + NextAction::None, + ); + + mgr.h2g_consumer = Some(queue.h2g_consumer()); + mgr + } + + fn h2g_call(bytes: usize) -> FunctionCall { + let params = (bytes != 0).then(|| vec![ParameterValue::VecBytes(vec![0xa5; bytes])]); + FunctionCall::new( + "call".to_string(), + params, + FunctionCallType::Guest, + ReturnType::Void, + ) + } + + #[test] + fn rejects_invalid_h2g_descriptors() { + for (len, expected) in [ + (H2G_BUFFER_SIZE as u32, "Payload data too large"), + (0, "not writable"), + ] { + let queue = TestVirtq::new(); + let mut mgr = manager(&queue); + let mut desc = queue.h2g_desc(0); + + desc.flags &= !DescFlags::WRITE.bits(); + desc.len = len; + queue.set_h2g_desc(0, desc); + + let error = mgr.write_guest_function_call(&h2g_call(0)).unwrap_err(); + + assert!(error.to_string().contains(expected), "{error:#}"); + assert!(error.is_poison_error()); + assert!(matches!(error, HyperlightError::TransportError(_))); + } + } + + #[test] + fn partial_h2g_write_is_fatal() { + let queue = TestVirtq::new(); + let mut mgr = manager(&queue); + let mut desc = queue.h2g_desc(1); + + desc.addr = queue.h2g_pool.end; + queue.set_h2g_desc(1, desc); + + let error = mgr + .write_guest_function_call(&h2g_call(H2G_BUFFER_SIZE + 1024)) + .unwrap_err(); + + assert!(error.to_string().contains("Memory write"), "{error:#}"); + assert!(error.is_poison_error()); + assert!(matches!(error, HyperlightError::TransportError(_))); + assert_eq!(mgr.h2g_consumer.as_ref().unwrap().used_cursor().head(), 1); + } + + #[test] + fn insufficient_h2g_capacity_rolls_back() { + let queue = TestVirtq::new(); + let mut mgr = manager(&queue); + let cursor = mgr.h2g_consumer.as_ref().unwrap().avail_cursor(); + + let error = mgr + .write_guest_function_call(&h2g_call(H2G_BUFFER_SIZE * 4)) + .unwrap_err(); + + assert!(error.to_string().contains("H2G capacity"), "{error:#}"); + assert!(!error.is_poison_error()); + assert_eq!(mgr.h2g_consumer.as_ref().unwrap().avail_cursor(), cursor); + assert_eq!(mgr.write_guest_function_call(&h2g_call(0)).unwrap(), 1); + } + + #[test] + fn missing_g2h_result_is_fatal() { + let queue = TestVirtq::new(); + let mut mgr = manager(&queue); + mgr.g2h_consumer = Some(queue.g2h_consumer()); + + let Err(error) = mgr.read_h2g_result_from_g2h(1) else { + panic!("expected missing G2H result"); + }; + + assert!( + error + .to_string() + .contains("G2H has no guest function result") + ); + assert!(error.is_poison_error()); + assert!(matches!(error, HyperlightError::TransportError(_))); + } + + #[test] + fn writes_dense_h2g_request_and_reserves_control_buffer() { + let queue = TestVirtq::new(); + let mut mgr = manager(&queue); + let external_len = H2G_BUFFER_SIZE * 2; + let buffers: Vec<_> = (0..4).map(|index| queue.h2g_desc(index).addr).collect(); + + let cid = mgr + .write_guest_function_call(&h2g_call(external_len)) + .unwrap(); + + let used = mgr.h2g_consumer.as_ref().unwrap().avail_cursor().head(); + let wire: Vec = (0..used) + .flat_map(|index| queue.h2g_buffer(index, buffers[index as usize])) + .collect(); + + let header = MsgHeader::from_bytes(&wire[..MsgHeader::SIZE]).unwrap(); + assert_eq!(header.msg_kind(), Ok(MsgKind::Request)); + assert_eq!(header.cid, cid); + assert_eq!(header.payload_len as usize, wire.len() - MsgHeader::SIZE); + + let control = + size_prefix_payload_len(&wire[MsgHeader::SIZE..MsgHeader::SIZE + SIZE_PREFIX_LEN]) + .unwrap(); + + let control_len = size_prefixed_len(control).unwrap(); + let external = &wire[MsgHeader::SIZE + control_len..]; + assert_eq!(external, vec![0xa5; external_len]); + + let cursor = mgr.h2g_consumer.as_ref().unwrap().avail_cursor(); + let error = mgr.write_guest_function_call(&h2g_call(1)).unwrap_err(); + + assert!(error.to_string().contains("H2G capacity"), "{error:#}"); + assert_eq!(mgr.h2g_consumer.as_ref().unwrap().avail_cursor(), cursor); + assert_eq!(mgr.write_guest_function_call(&h2g_call(0)).unwrap(), 2); + } + + #[test] + fn writes_header_only_snapshot_checkpoint() { + let queue = TestVirtq::new(); + let mut mgr = manager(&queue); + let buffer = queue.h2g_desc(0).addr; + + mgr.begin_snapshot_checkpoint().unwrap(); + + let wire = queue.h2g_buffer(0, buffer); + let header = MsgHeader::from_bytes(&wire).unwrap(); + + assert_eq!(wire.len(), MsgHeader::SIZE); + assert_eq!(header.msg_kind(), Ok(MsgKind::SnapshotCheckpoint)); + assert_eq!(header.cid, 0); + assert_eq!(header.payload_len, 0); + assert_eq!(mgr.next_guest_cid, 1); + + let mbx = mgr.snapshot_mbx_offset().unwrap(); + + assert_eq!( + mgr.scratch_mem.read::<[u8; 8]>(mbx).unwrap(), + u64::MAX.to_le_bytes() + ); + } + + #[test] + fn guest_cid_wraps_without_zero() { + let queue = TestVirtq::new(); + let mut mgr = manager(&queue); + mgr.next_guest_cid = u32::MAX; + + assert_eq!( + mgr.write_guest_function_call(&h2g_call(0)).unwrap(), + u32::MAX + ); + assert_eq!(mgr.write_guest_function_call(&h2g_call(0)).unwrap(), 1); + } + + /// Build a snapshot for the given configuration and verify the /// NULL page is not mapped in its page tables. + #[cfg(target_arch = "x86_64")] fn verify_page_tables(name: &str, config: SandboxConfiguration) { let path = simple_guest_as_pathbuf(); let snapshot = Snapshot::from_env(GuestBinary::FilePath(path), config) @@ -791,6 +1244,7 @@ mod tests { } #[test] + #[cfg(target_arch = "x86_64")] fn test_page_tables_for_various_configurations() { let test_cases: [(&str, SandboxConfiguration); 4] = [ ("default", { SandboxConfiguration::default() }), @@ -816,4 +1270,18 @@ mod tests { verify_page_tables(name, config); } } + + #[test] + #[cfg(target_arch = "x86_64")] + fn build_creates_virtq_consumers_before_initialization() { + let path = simple_guest_as_pathbuf(); + let bin = GuestBinary::FilePath(path); + let snapshot = Snapshot::from_env(bin, SandboxConfiguration::default()).unwrap(); + + let mgr = SandboxMemoryManager::from_snapshot(&snapshot).unwrap(); + let (mgr, _) = mgr.build().unwrap(); + + assert!(mgr.g2h_consumer.is_some()); + assert!(mgr.h2g_consumer.is_some()); + } } diff --git a/src/hyperlight_host/src/mem/mod.rs b/src/hyperlight_host/src/mem/mod.rs index 64f5db2fe..977df2e4b 100644 --- a/src/hyperlight_host/src/mem/mod.rs +++ b/src/hyperlight_host/src/mem/mod.rs @@ -38,3 +38,5 @@ pub mod shared_mem; /// Utilities for writing shared memory tests #[cfg(all(test, not(miri)))] // uses proptest which isn't miri-compatible pub(crate) mod shared_mem_tests; +/// Host virtqueue attachment and validation. +pub(crate) mod virtq; diff --git a/src/hyperlight_host/src/mem/shared_mem.rs b/src/hyperlight_host/src/mem/shared_mem.rs index 4b706cc41..e355fc12b 100644 --- a/src/hyperlight_host/src/mem/shared_mem.rs +++ b/src/hyperlight_host/src/mem/shared_mem.rs @@ -14,12 +14,12 @@ See the License for the specific language governing permissions and limitations under the License. */ -use std::any::type_name; use std::ffi::c_void; use std::io::Error; use std::mem::{align_of, size_of}; #[cfg(target_os = "linux")] use std::ptr::null_mut; +use std::sync::atomic::Ordering; use std::sync::{Arc, RwLock}; use bytemuck::Pod; @@ -61,6 +61,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) => { @@ -91,6 +131,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 @@ -1148,6 +1199,70 @@ 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(new_error!("Invalid atomic load ordering: {:?}", 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(new_error!( + "Atomic access at offset {} is not aligned to {} bytes", + offset, + align_of::() + )); + } + + let _guard = self + .lock + .try_read() + .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?; + + // 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(new_error!("Invalid atomic store ordering: {:?}", 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(new_error!( + "Atomic access at offset {} is not aligned to {} bytes", + offset, + align_of::() + )); + } + + let _guard = self + .lock + .try_read() + .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?; + + // 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<()> { @@ -1297,145 +1412,6 @@ impl HostSharedMemory { drop(guard); Ok(()) } - - /// Pushes the given data onto shared memory to the buffer at the given offset. - /// NOTE! buffer_start_offset must point to the beginning of the buffer - #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] - pub fn push_buffer( - &mut self, - buffer_start_offset: usize, - buffer_size: usize, - data: &[u8], - ) -> Result<()> { - let stack_pointer_rel = self.read::(buffer_start_offset)? as usize; - let buffer_size_u64: u64 = buffer_size.try_into()?; - - if stack_pointer_rel > buffer_size || stack_pointer_rel < 8 { - return Err(new_error!( - "Unable to push data to buffer: Stack pointer is out of bounds. Stack pointer: {}, Buffer size: {}", - stack_pointer_rel, - buffer_size_u64 - )); - } - - let size_required = data.len() + 8; - let size_available = buffer_size - stack_pointer_rel; - - if size_required > size_available { - return Err(new_error!( - "Not enough space in buffer to push data. Required: {}, Available: {}", - size_required, - size_available - )); - } - - // get absolute - let stack_pointer_abs = stack_pointer_rel + buffer_start_offset; - - // write the actual data to the top of stack - self.copy_from_slice(data, stack_pointer_abs)?; - - // write the offset to the newly written data, to the top of stack. - // this is used when popping the stack, to know how far back to jump - self.write::(stack_pointer_abs + data.len(), stack_pointer_rel as u64)?; - - // update stack pointer to point to the next free address - self.write::( - buffer_start_offset, - (stack_pointer_rel + data.len() + 8) as u64, - )?; - Ok(()) - } - - /// Pops the given given buffer into a `T` and returns it. - /// NOTE! the data must be a size-prefixed flatbuffer, and - /// buffer_start_offset must point to the beginning of the buffer - pub fn try_pop_buffer_into( - &mut self, - buffer_start_offset: usize, - buffer_size: usize, - ) -> Result - where - T: for<'b> TryFrom<&'b [u8]>, - { - // get the stackpointer - let stack_pointer_rel = self.read::(buffer_start_offset)? as usize; - - if stack_pointer_rel > buffer_size || stack_pointer_rel < 16 { - return Err(new_error!( - "Unable to pop data from buffer: Stack pointer is out of bounds. Stack pointer: {}, Buffer size: {}", - stack_pointer_rel, - buffer_size - )); - } - - // make it absolute - let last_element_offset_abs = stack_pointer_rel + buffer_start_offset; - - // go back 8 bytes to get offset to element on top of stack - let last_element_offset_rel: usize = - self.read::(last_element_offset_abs - 8)? as usize; - - // Validate element offset (guest-writable): must be in [8, stack_pointer_rel - 16] - // to leave room for the 8-byte back-pointer plus at least 8 bytes of element data - // (the minimum for a size-prefixed flatbuffer: 4-byte prefix + 4-byte root offset). - if last_element_offset_rel > stack_pointer_rel.saturating_sub(16) - || last_element_offset_rel < 8 - { - return Err(new_error!( - "Corrupt buffer back-pointer: element offset {} is outside valid range [8, {}].", - last_element_offset_rel, - stack_pointer_rel.saturating_sub(16), - )); - } - - // make it absolute - let last_element_offset_abs = last_element_offset_rel + buffer_start_offset; - - // Max bytes the element can span (excluding the 8-byte back-pointer). - let max_element_size = stack_pointer_rel - last_element_offset_rel - 8; - - // Get the size of the flatbuffer buffer from memory - let fb_buffer_size = { - let raw_prefix = self.read::(last_element_offset_abs)?; - // flatbuffer byte arrays are prefixed by 4 bytes indicating - // the remaining size; add 4 for the prefix itself. - let total = raw_prefix.checked_add(4).ok_or_else(|| { - new_error!( - "Corrupt buffer size prefix: value {} overflows when adding 4-byte header.", - raw_prefix - ) - })?; - usize::try_from(total) - }?; - - if fb_buffer_size > max_element_size { - return Err(new_error!( - "Corrupt buffer size prefix: flatbuffer claims {} bytes but the element slot is only {} bytes.", - fb_buffer_size, - max_element_size - )); - } - - let mut result_buffer = vec![0; fb_buffer_size]; - - self.copy_to_slice(&mut result_buffer, last_element_offset_abs)?; - let to_return = T::try_from(result_buffer.as_slice()).map_err(|_e| { - new_error!( - "pop_buffer_into: failed to convert buffer to {}", - type_name::() - ) - })?; - - // update the stack pointer to point to the element we just popped off since that is now free - self.write::(buffer_start_offset, last_element_offset_rel as u64)?; - - // zero out the memory we just popped off - let num_bytes_to_zero = stack_pointer_rel - last_element_offset_rel; - self.fill(0, last_element_offset_abs, num_bytes_to_zero)?; - - Ok(to_return) - } } impl SharedMemory for HostSharedMemory { @@ -1791,6 +1767,8 @@ impl PartialEq for ReadonlySharedMemory { #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicI32, AtomicU16, Ordering}; + use hyperlight_common::mem::PAGE_SIZE_USIZE; #[cfg(not(miri))] use proptest::prelude::*; @@ -1861,6 +1839,53 @@ mod tests { assert!(hshm.fill(0, 1, usize::MAX).is_err()); } + #[test] + fn atomic_access() { + let eshm = ExclusiveSharedMemory::new(PAGE_SIZE_USIZE).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::(PAGE_SIZE_USIZE - 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_USIZE).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 = 4096; @@ -2272,192 +2297,6 @@ mod tests { } } - /// Bounds checking for `try_pop_buffer_into` against corrupt guest data. - mod try_pop_buffer_bounds { - use super::*; - - #[derive(Debug, PartialEq)] - struct RawBytes(Vec); - - impl TryFrom<&[u8]> for RawBytes { - type Error = String; - fn try_from(value: &[u8]) -> std::result::Result { - Ok(RawBytes(value.to_vec())) - } - } - - /// Create a buffer with stack pointer initialized to 8 (empty). - fn make_buffer(mem_size: usize) -> super::super::HostSharedMemory { - let eshm = ExclusiveSharedMemory::new(mem_size).unwrap(); - let (hshm, _) = eshm.build(); - hshm.write::(0, 8u64).unwrap(); - hshm - } - - #[test] - fn normal_push_pop_roundtrip() { - let mem_size = 4096; - let mut hshm = make_buffer(mem_size); - - // Size-prefixed flatbuffer-like payload: [size: u32 LE][payload] - let payload = b"hello"; - let mut data = Vec::new(); - data.extend_from_slice(&(payload.len() as u32).to_le_bytes()); - data.extend_from_slice(payload); - - hshm.push_buffer(0, mem_size, &data).unwrap(); - let result: RawBytes = hshm.try_pop_buffer_into(0, mem_size).unwrap(); - assert_eq!(result.0, data); - } - - #[test] - fn malicious_flatbuffer_size_prefix() { - let mem_size = 4096; - let mut hshm = make_buffer(mem_size); - - let payload = b"small"; - let mut data = Vec::new(); - data.extend_from_slice(&(payload.len() as u32).to_le_bytes()); - data.extend_from_slice(payload); - hshm.push_buffer(0, mem_size, &data).unwrap(); - - // Corrupt size prefix at element start (offset 8) to near u32::MAX. - hshm.write::(8, 0xFFFF_FFFBu32).unwrap(); // +4 = 0xFFFF_FFFF - - let result: Result = hshm.try_pop_buffer_into(0, mem_size); - let err_msg = format!("{}", result.unwrap_err()); - assert!( - err_msg.contains("Corrupt buffer size prefix: flatbuffer claims 4294967295 bytes but the element slot is only 9 bytes"), - "Unexpected error message: {}", - err_msg - ); - } - - #[test] - fn malicious_element_offset_too_small() { - let mem_size = 4096; - let mut hshm = make_buffer(mem_size); - - let payload = b"test"; - let mut data = Vec::new(); - data.extend_from_slice(&(payload.len() as u32).to_le_bytes()); - data.extend_from_slice(payload); - hshm.push_buffer(0, mem_size, &data).unwrap(); - - // Corrupt back-pointer (offset 16) to 0 (before valid range). - hshm.write::(16, 0u64).unwrap(); - - let result: Result = hshm.try_pop_buffer_into(0, mem_size); - let err_msg = format!("{}", result.unwrap_err()); - assert!( - err_msg.contains( - "Corrupt buffer back-pointer: element offset 0 is outside valid range [8, 8]" - ), - "Unexpected error message: {}", - err_msg - ); - } - - #[test] - fn malicious_element_offset_past_stack_pointer() { - let mem_size = 4096; - let mut hshm = make_buffer(mem_size); - - let payload = b"test"; - let mut data = Vec::new(); - data.extend_from_slice(&(payload.len() as u32).to_le_bytes()); - data.extend_from_slice(payload); - hshm.push_buffer(0, mem_size, &data).unwrap(); - - // Corrupt back-pointer (offset 16) to 9999 (past stack pointer 24). - hshm.write::(16, 9999u64).unwrap(); - - let result: Result = hshm.try_pop_buffer_into(0, mem_size); - let err_msg = format!("{}", result.unwrap_err()); - assert!( - err_msg.contains( - "Corrupt buffer back-pointer: element offset 9999 is outside valid range [8, 8]" - ), - "Unexpected error message: {}", - err_msg - ); - } - - #[test] - fn malicious_flatbuffer_size_off_by_one() { - let mem_size = 4096; - let mut hshm = make_buffer(mem_size); - - let payload = b"abcd"; - let mut data = Vec::new(); - data.extend_from_slice(&(payload.len() as u32).to_le_bytes()); - data.extend_from_slice(payload); - hshm.push_buffer(0, mem_size, &data).unwrap(); - - // Corrupt size prefix: claim 5 bytes (total 9), exceeding the 8-byte slot. - hshm.write::(8, 5u32).unwrap(); // fb_buffer_size = 5 + 4 = 9 - - let result: Result = hshm.try_pop_buffer_into(0, mem_size); - let err_msg = format!("{}", result.unwrap_err()); - assert!( - err_msg.contains("Corrupt buffer size prefix: flatbuffer claims 9 bytes but the element slot is only 8 bytes"), - "Unexpected error message: {}", - err_msg - ); - } - - /// Back-pointer just below stack_pointer causes underflow in - /// `stack_pointer_rel - last_element_offset_rel - 8`. - #[test] - fn back_pointer_near_stack_pointer_underflow() { - let mem_size = 4096; - let mut hshm = make_buffer(mem_size); - - let payload = b"test"; - let mut data = Vec::new(); - data.extend_from_slice(&(payload.len() as u32).to_le_bytes()); - data.extend_from_slice(payload); - hshm.push_buffer(0, mem_size, &data).unwrap(); - - // stack_pointer_rel = 24. Set back-pointer to 23 (> 24 - 16 = 8, so rejected). - hshm.write::(16, 23u64).unwrap(); - - let result: Result = hshm.try_pop_buffer_into(0, mem_size); - let err_msg = format!("{}", result.unwrap_err()); - assert!( - err_msg.contains( - "Corrupt buffer back-pointer: element offset 23 is outside valid range [8, 8]" - ), - "Unexpected error message: {}", - err_msg - ); - } - - /// Size prefix of 0xFFFF_FFFD causes u32 overflow: 0xFFFF_FFFD + 4 wraps. - #[test] - fn size_prefix_u32_overflow() { - let mem_size = 4096; - let mut hshm = make_buffer(mem_size); - - let payload = b"test"; - let mut data = Vec::new(); - data.extend_from_slice(&(payload.len() as u32).to_le_bytes()); - data.extend_from_slice(payload); - hshm.push_buffer(0, mem_size, &data).unwrap(); - - // Write 0xFFFF_FFFD as size prefix: checked_add(4) returns None. - hshm.write::(8, 0xFFFF_FFFDu32).unwrap(); - - let result: Result = hshm.try_pop_buffer_into(0, mem_size); - let err_msg = format!("{}", result.unwrap_err()); - assert!( - err_msg.contains("Corrupt buffer size prefix: value 4294967293 overflows when adding 4-byte header"), - "Unexpected error message: {}", - err_msg - ); - } - } - #[cfg(target_os = "linux")] mod guard_page_crash_test { use crate::mem::shared_mem::{ExclusiveSharedMemory, SharedMemory}; diff --git a/src/hyperlight_host/src/mem/virtq/codec.rs b/src/hyperlight_host/src/mem/virtq/codec.rs new file mode 100644 index 000000000..8e49db6cc --- /dev/null +++ b/src/hyperlight_host/src/mem/virtq/codec.rs @@ -0,0 +1,197 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +//! Host RPC encoding and decoding over virtqueue chains. + +use anyhow::{Context, bail}; +use flatbuffers::FlatBufferBuilder; +use hyperlight_common::flatbuffer_wrappers::ExternalValueSource; +use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall; +use hyperlight_common::flatbuffer_wrappers::function_types::{Bytes, FunctionCallResult}; +use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData; +use hyperlight_common::transport::{ + EncodedMessage, ExternalValues, MsgHeader, MsgKind, SIZE_PREFIX_LEN, size_prefix_payload_len, + size_prefixed_len, +}; +use hyperlight_common::virtq::{RecvChain, WritableChain}; + +use super::mem::HostMemOps; + +/// Copies external values from guest-writable scratch into host-owned storage. +/// +/// Chunked values become one owned chunk because host calls cannot retain +/// references into untrusted guest memory. +struct ChainExternalValues<'a> { + request: &'a mut RecvChain, +} + +impl<'a> ChainExternalValues<'a> { + fn new(request: &'a mut RecvChain) -> Self { + Self { request } + } +} + +impl ExternalValueSource for ChainExternalValues<'_> { + fn take_bytes(&mut self, length: usize) -> anyhow::Result> { + let remain = self.request.remaining(); + if length > remain { + bail!("External VecBytes requires {length} bytes, only {remain} remain"); + } + + let mut value = zeroed_vec(length)?; + self.request.read_exact(&mut value)?; + + Ok(value) + } + + fn take_chunks(&mut self, length: usize) -> anyhow::Result> { + if length == 0 { + return Ok(Vec::new()); + } + + let rem = self.request.remaining(); + if length > rem { + bail!("External ByteChunks requires {length} bytes, only {rem} remain"); + } + + let mut value = zeroed_vec(length)?; + self.request.read_exact(&mut value)?; + + Ok(vec![Bytes::from(value)]) + } + + fn finish(&mut self) -> anyhow::Result<()> { + if self.request.remaining() != 0 { + bail!( + "G2H message has {} trailing external bytes", + self.request.remaining() + ); + } + Ok(()) + } +} + +/// Decode one complete host function call from a G2H request. +/// +/// Control data and external values are copied out of guest-writable scratch. +/// Unconsumed trailing bytes are rejected. +pub(crate) fn get_host_function_call( + chain: &mut RecvChain, +) -> anyhow::Result { + let control = read_control(chain)?; + let mut exts = ChainExternalValues::new(chain); + FunctionCall::decode(&control, &mut exts) +} + +/// Read and validate one complete G2H message header. +pub(crate) fn read_message_header( + request: &mut RecvChain, +) -> anyhow::Result { + let mut bytes = [0u8; MsgHeader::SIZE]; + request.read_exact(&mut bytes)?; + + let header = MsgHeader::from_bytes(&bytes).context("G2H message has an invalid header")?; + if header.payload_len as usize != request.remaining() { + bail!("G2H message payload length mismatch"); + } + + Ok(header) +} + +/// Decode a guest-function result body after its G2H header. +pub(crate) fn read_guest_function_call_result( + request: &mut RecvChain, +) -> anyhow::Result { + let control = read_control(request)?; + let mut exts = ChainExternalValues::new(request); + FunctionCallResult::decode(&control, &mut exts) +} + +/// Encode and write a response when its complete wire message fits. +/// +/// `false` leaves the writable chain unchanged. +pub(crate) fn try_write_response( + reply: &mut WritableChain, + cid: u32, + result: &FunctionCallResult, +) -> anyhow::Result { + let mut builder = FlatBufferBuilder::new(); + let mut externals = ExternalValues::new(); + + let control = result.encode(&mut builder, &mut externals)?; + + let Some(msg) = EncodedMessage::new(MsgKind::Response, cid, control, externals) else { + bail!("Host function response length overflow"); + }; + + if msg.total_len() > reply.capacity() { + return Ok(false); + } + + for chunk in msg.chunks() { + reply.write_all(chunk)?; + } + + Ok(true) +} + +/// Decode guest log data and reject trailing external bytes. +pub(crate) fn read_guest_log_data( + chain: &mut RecvChain, +) -> anyhow::Result { + let control = read_control(chain)?; + let remain = chain.remaining(); + + if remain != 0 { + bail!("G2H log has {remain} trailing external bytes"); + } + + GuestLogData::try_from(control.as_slice()) +} + +/// Copy size-prefixed control data and leave external values unread. +fn read_control(request: &mut RecvChain) -> anyhow::Result> { + let mut prefix = [0u8; SIZE_PREFIX_LEN]; + request.read_exact(&mut prefix)?; + + let payload_len = size_prefix_payload_len(&prefix).context("G2H size prefix is invalid")?; + if payload_len > request.remaining() { + bail!( + "G2H control data declares {payload_len} bytes, only {} remain", + request.remaining() + ); + } + + let control_len = size_prefixed_len(payload_len).context("G2H control length overflow")?; + // Do not trust control_len to be small enough to allocate. + let mut control = zeroed_vec(control_len)?; + + control[..SIZE_PREFIX_LEN].copy_from_slice(&prefix); + request.read_exact(&mut control[SIZE_PREFIX_LEN..])?; + + Ok(control) +} + +/// Allocate zeroed host-owned storage without panicking on reserve failure. +pub fn zeroed_vec(length: usize) -> anyhow::Result> { + let mut value = Vec::new(); + value + .try_reserve_exact(length) + .with_context(|| format!("Failed to allocate {length} bytes"))?; + + value.resize(length, 0); + Ok(value) +} 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..f0795e23f --- /dev/null +++ b/src/hyperlight_host/src/mem/virtq/mem.rs @@ -0,0 +1,271 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +//! 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 crate::mem::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())?; + self.scratch_mem.copy_to_slice(dst, offset) + } + + fn write(&self, addr: u64, src: &[u8]) -> Result<()> { + let offset = self.to_offset(addr, src.len())?; + self.scratch_mem.copy_from_slice(src, offset) + } + + fn load_acquire(&self, addr: u64) -> Result { + let offset = self.to_offset(addr, size_of::())?; + 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::())?; + 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/mem/virtq/mod.rs b/src/hyperlight_host/src/mem/virtq/mod.rs new file mode 100644 index 000000000..8af67e5be --- /dev/null +++ b/src/hyperlight_host/src/mem/virtq/mod.rs @@ -0,0 +1,490 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +//! Host virtqueue construction and canonical snapshot validation. +//! +//! Runtime consumers bind bounded ring and pool views to the host-owned fixed +//! transport arena. They start at cursor zero before the first guest entry and +//! observe descriptors published by the guest later. +//! +//! H2G requests are written into guest-prefilled chains. G2H codec helpers copy +//! untrusted guest requests and results into host-owned values before use. +//! Shared wire framing lives in `hyperlight_common::transport`. +//! +//! Snapshot capture and restore validate canonical ring images against the +//! configured arena before exposing consumers. + +mod codec; +mod mem; +#[cfg(test)] +pub(crate) mod tests; + +use core::ops::Range; + +pub(crate) use codec::{ + get_host_function_call, read_guest_function_call_result, read_guest_log_data, + read_message_header, try_write_response, +}; +use hyperlight_common::layout::{QueueDims, TransportArena}; +use hyperlight_common::virtq::canonical::validate_canon_image; +use hyperlight_common::virtq::{ + Layout as VirtqLayout, MemOps, Notifier, QueueStats, VirtqConsumer, +}; +use mem::{HostMemOps, ImageMem}; + +use super::layout::{BaseGpaRegion, SandboxMemoryLayout}; +use super::shared_mem::{HostSharedMemory, SharedMemory}; +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 because the host completes work during the current VM exit. +#[derive(Clone, Copy)] +pub(crate) struct HostNotifier; + +impl Notifier for HostNotifier { + fn notify(&self, _stats: QueueStats) {} +} + +/// Create both host consumers before the first guest entry. +/// +/// Ring contents are not inspected because the guest has not initialized them +/// yet. Consumer cursors start at zero and observe descriptors published later. +pub(crate) fn create_consumers( + layout: &SandboxMemoryLayout, + scratch_mem: &HostSharedMemory, +) -> Result<(G2hConsumer, H2gConsumer)> { + let validator = Validator::new(layout)?; + let regions = validator.resolve_gva_regions()?; + let g2h_layout = validator.config.g2h.layout(®ions.g2h_ring)?; + let h2g_layout = validator.config.h2g.layout(®ions.h2g_ring)?; + + build_consumers(scratch_mem, regions, g2h_layout, h2g_layout) +} + +/// Validate a materialized canonical image before attaching consumers. +fn attach_canonical( + layout: &SandboxMemoryLayout, + scratch_mem: &HostSharedMemory, +) -> Result<(G2hConsumer, H2gConsumer)> { + let validator = Validator::new(layout)?; + let arena_gpa = read_published_arena_gpa(scratch_mem)?; + let regions = validator.validate_published_arena(arena_gpa)?; + + let g2h_ring_mem = HostMemOps::new(scratch_mem, regions.g2h_ring.clone())?; + let g2h_layout = validator.validate_g2h(&g2h_ring_mem, regions.g2h_ring.clone())?; + + let h2g_ring_mem = HostMemOps::new(scratch_mem, regions.h2g_ring.clone())?; + // Why Range is not Copy? + let h2g_ring = regions.h2g_ring.clone(); + let h2g_pool = regions.h2g_pool.clone(); + let h2g_layout = validator.validate_h2g(&h2g_ring_mem, h2g_ring, h2g_pool)?; + + build_consumers(scratch_mem, regions, g2h_layout, h2g_layout) +} + +/// Bind consumers to separately bounded ring and pool mappings. +fn build_consumers( + scratch_mem: &HostSharedMemory, + regions: GvaRegions, + g2h_layout: VirtqLayout, + h2g_layout: VirtqLayout, +) -> Result<(G2hConsumer, H2gConsumer)> { + let g2h_ring_mem = HostMemOps::new(scratch_mem, regions.g2h_ring)?; + let g2h_pool_mem = HostMemOps::new(scratch_mem, regions.g2h_pool)?; + let h2g_ring_mem = HostMemOps::new(scratch_mem, regions.h2g_ring)?; + let h2g_pool_mem = HostMemOps::new(scratch_mem, 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 ring state omitted from the main memory snapshot. +/// +/// Pool contents are transient and are not included. +pub(crate) fn snapshot( + layout: &SandboxMemoryLayout, + scratch_mem: &HostSharedMemory, +) -> Result { + let validator = Validator::new(layout)?; + + let arena_gpa = read_published_arena_gpa(scratch_mem)?; + let regions = validator.validate_published_arena(arena_gpa)?; + + let g2h_mem = HostMemOps::new(scratch_mem, regions.g2h_ring.clone())?; + validator.validate_g2h(&g2h_mem, regions.g2h_ring.clone())?; + + let h2g_mem = HostMemOps::new(scratch_mem, regions.h2g_ring.clone())?; + validator.validate_h2g(&h2g_mem, regions.h2g_ring.clone(), regions.h2g_pool.clone())?; + + // The vCPU is stopped, so the ring images and snapshotted guest producer + // bookkeeping describe the same instant. + Ok(VirtqSnapshot { + scratch_size: layout.get_scratch_size(), + g2h_ring: read_ring(scratch_mem, regions.g2h_ring)?, + h2g_ring: read_ring(scratch_mem, regions.h2g_ring)?, + }) +} + +/// Validate and restore one canonical transport image. +/// +/// Validation completes before restored scratch is mutated. Fresh consumers +/// start from the canonical cursor state encoded in the rings. +pub(crate) fn restore( + layout: &SandboxMemoryLayout, + scratch_mem: &HostSharedMemory, + snapshot: &VirtqSnapshot, +) -> Result<(G2hConsumer, H2gConsumer)> { + let regions = Validator::new(layout)?.validate_snapshot(snapshot)?; + + write_published_arena_gpa(scratch_mem, layout.get_transport_arena().base_addr())?; + write_ring(scratch_mem, regions.g2h_ring, &snapshot.g2h_ring)?; + write_ring(scratch_mem, regions.h2g_ring, &snapshot.h2g_ring)?; + attach_canonical(layout, scratch_mem) +} + +/// Bounded GVA regions derived from validated transport GPAs. +struct GvaRegions { + /// Guest-to-host packed ring image. + g2h_ring: Range, + /// Host-to-guest packed ring image. + h2g_ring: Range, + /// Guest-to-host descriptor buffer pool. + g2h_pool: Range, + /// Host-to-guest descriptor buffer pool. + h2g_pool: Range, +} + +#[derive(Clone, Copy)] +struct QueueConfig { + /// Address-independent queue dimensions. + dims: QueueDims, + /// Size of each buffer in the pool in bytes. + buffer_size: usize, +} + +impl QueueConfig { + fn new(dims: QueueDims, buffer_size: usize) -> Result { + if buffer_size == 0 { + return Err(new_error!("buffer size is zero")); + } + + Ok(Self { dims, buffer_size }) + } + + fn layout(&self, ring: &Range) -> Result { + // SAFETY: `ring` is derived from the validated fixed transport arena. + unsafe { VirtqLayout::from_base(ring.start, self.dims.size()) } + .map_err(|error| new_error!("invalid ring layout: {error}")) + } +} + +/// Host-owned transport dimensions. +#[derive(Clone, Copy)] +struct Config { + /// Host-requested G2H configuration. + g2h: QueueConfig, + /// Host-requested H2G configuration. + h2g: QueueConfig, + /// Fixed host-assigned transport arena. + arena: TransportArena, + /// Number of canonical single-buffer H2G receive chains. + h2g_prefill_descs: usize, +} + +impl Config { + /// Compute the host transport configuration from the memory layout. + fn from_layout(layout: &SandboxMemoryLayout) -> Result { + let g2h = QueueConfig::new(layout.get_g2h_queue_dims(), layout.get_g2h_buffer_size())?; + let h2g = QueueConfig::new(layout.get_h2g_queue_dims(), layout.get_h2g_buffer_size())?; + + let h2g_prefill_descs = + usize::from(h2g.dims.size().get()).min(h2g.dims.pool_len() / h2g.buffer_size); + + let arena = layout.get_transport_arena(); + + Ok(Self { + g2h, + h2g, + arena, + h2g_prefill_descs, + }) + } +} + +/// Canonical in-memory transport state excluded from ordinary snapshot pages. +#[derive(Debug, PartialEq, Eq)] +pub(crate) struct VirtqSnapshot { + /// Scratch size used to derive transport GVAs. + scratch_size: usize, + /// Canonical guest-to-host ring image. + g2h_ring: Vec, + /// Canonical host-to-guest ring image. + h2g_ring: Vec, +} + +impl VirtqSnapshot { + pub(crate) fn new(scratch_size: usize, g2h_ring: Vec, h2g_ring: Vec) -> Self { + Self { + scratch_size, + g2h_ring, + h2g_ring, + } + } + + pub(crate) fn scratch_size(&self) -> usize { + self.scratch_size + } + + pub(crate) fn g2h_ring(&self) -> &[u8] { + &self.g2h_ring + } + + pub(crate) fn h2g_ring(&self) -> &[u8] { + &self.h2g_ring + } + + /// Validate every captured field before mutating restored scratch. + pub(crate) fn preflight(&self, layout: &SandboxMemoryLayout) -> Result<()> { + Validator::new(layout)?.validate_snapshot(self).map(|_| ()) + } +} + +/// Validates live and captured transport images against one host layout. +struct Validator<'a> { + config: Config, + layout: &'a SandboxMemoryLayout, +} + +impl<'a> Validator<'a> { + fn new(layout: &'a SandboxMemoryLayout) -> Result { + Ok(Self { + config: Config::from_layout(layout)?, + layout, + }) + } + + /// Validate a canonical G2H image and return its layout. + fn validate_g2h(&self, mem: &M, ring: Range) -> Result { + let layout = self.config.g2h.layout(&ring)?; + + validate_canon_image(mem, layout, 0, |_, _| false) + .map_err(|error| new_error!("invalid canonical G2H image: {error}"))?; + + Ok(layout) + } + + /// Validate a canonical H2G image and return its layout. + /// + /// Each available chain contains one 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 layout = self.config.h2g.layout(&ring)?; + let bufsz = self.config.h2g.buffer_size; + let prefill = self.config.h2g_prefill_descs; + + 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}"))?; + + if image.len() != prefill || image.iter().any(|chain| chain.buffers().len() != 1) { + return Err(new_error!("invalid initial H2G receive buffers")); + } + + Ok(layout) + } + + /// Validate the published arena and return its GVA regions. + fn validate_published_arena(&self, arena_gpa: u64) -> Result { + if arena_gpa != self.config.arena.base_addr() { + return Err(new_error!("published transport arena is invalid")); + } + + self.resolve_gva_regions() + } + + fn validate_snapshot(&self, snapshot: &VirtqSnapshot) -> Result { + if snapshot.scratch_size != self.layout.get_scratch_size() { + return Err(new_error!( + "virtqueue snapshot scratch size {} does not match layout size {}", + snapshot.scratch_size, + self.layout.get_scratch_size() + )); + } + + let regions = self.resolve_gva_regions()?; + validate_ring_len("G2H", &snapshot.g2h_ring, self.config.g2h.dims.ring_len())?; + validate_ring_len("H2G", &snapshot.h2g_ring, self.config.h2g.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) + } + + /// Translate validated transport GPAs into the GVA ranges used by descriptors. + fn resolve_gva_regions(&self) -> Result { + let to_gva = |gpa| { + let resolved = self + .layout + .resolve_gpa(gpa, &[]) + .ok_or_else(|| new_error!("GPA {gpa:#x} is outside scratch"))?; + + if !matches!(resolved.base, BaseGpaRegion::Scratch(())) { + return Err(new_error!("GPA {gpa:#x} is outside scratch")); + } + + hyperlight_common::layout::scratch_base_gva(self.layout.get_scratch_size()) + .checked_add(u64::try_from(resolved.offset)?) + .ok_or_else(|| new_error!("GPA {gpa:#x} to GVA translation overflow")) + }; + + let ( + g2h_ring_addr, + h2g_ring_addr, + g2h_pool_addr, + h2g_pool_addr, + g2h_ring_len, + h2g_ring_len, + g2h_pool_len, + h2g_pool_len, + ) = ( + self.config.arena.g2h_ring_addr(), + self.config.arena.h2g_ring_addr(), + self.config.arena.g2h_pool_addr(), + self.config.arena.h2g_pool_addr(), + self.config.g2h.dims.ring_len(), + self.config.h2g.dims.ring_len(), + self.config.g2h.dims.pool_len(), + self.config.h2g.dims.pool_len(), + ); + + Ok(GvaRegions { + g2h_ring: checked_region(to_gva(g2h_ring_addr)?, g2h_ring_len, "G2H ring")?, + h2g_ring: checked_region(to_gva(h2g_ring_addr)?, h2g_ring_len, "H2G ring")?, + g2h_pool: checked_region(to_gva(g2h_pool_addr)?, g2h_pool_len, "G2H pool")?, + h2g_pool: checked_region(to_gva(h2g_pool_addr)?, h2g_pool_len, "H2G pool")?, + }) + } +} + +/// Read the transport arena GPA from scratch-top metadata. +fn read_published_arena_gpa(scratch_mem: &HostSharedMemory) -> Result { + let offset = hyperlight_common::layout::SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET as usize; + scratch_mem.read::(scratch_mem.mem_size() - offset) +} + +/// Publish the fixed transport arena GPA in scratch-top metadata. +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; + scratch_mem.write::(scratch_mem.mem_size() - offset, arena_gpa) +} + +/// Copy one ring image from its bounded scratch mapping. +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) +} + +/// Copy one validated ring image into its bounded scratch mapping. +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) +} + +/// Require a captured ring image to match its configured region exactly. +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(()) +} + +/// Build a GVA range while checking address arithmetic. +fn checked_region(start: u64, len: usize, tag: &str) -> Result> { + let end = start + .checked_add(u64::try_from(len)?) + .ok_or_else(|| new_error!("{tag} GVA range overflow"))?; + + Ok(start..end) +} diff --git a/src/hyperlight_host/src/mem/virtq/tests.rs b/src/hyperlight_host/src/mem/virtq/tests.rs new file mode 100644 index 000000000..6cdf52b9b --- /dev/null +++ b/src/hyperlight_host/src/mem/virtq/tests.rs @@ -0,0 +1,326 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use core::ops::Range; + +use hyperlight_common::layout::SCRATCH_TOP_ALLOCATOR_OFFSET; +use hyperlight_common::virtq::{ + DescFlags, Descriptor, MemOps, SlotLayout, SlotPool, VirtqProducer, +}; +use hyperlight_common::vmem; + +use super::mem::HostMemOps; +use super::*; +use crate::mem::shared_mem::{ExclusiveSharedMemory, HostSharedMemory}; +use crate::sandbox::SandboxConfiguration; + +pub(crate) const SCRATCH_SIZE: usize = 0x20_000; +pub(crate) const H2G_BUFFER_SIZE: usize = 3000; + +pub(crate) struct TestVirtq { + pub(crate) scratch: HostSharedMemory, + pub(crate) g2h_mem: HostMemOps, + pub(crate) h2g_mem: HostMemOps, + pub(crate) g2h_ring: Range, + pub(crate) h2g_ring: Range, + pub(crate) g2h_pool: Range, + pub(crate) h2g_pool: Range, + pub(crate) g2h_layout: VirtqLayout, + pub(crate) h2g_layout: VirtqLayout, +} + +impl TestVirtq { + pub(crate) fn new() -> Self { + let scratch = host_scratch(); + let layout = memory_layout(); + + let validator = Validator::new(&layout).unwrap(); + let config = validator.config; + let regions = validator.resolve_gva_regions().unwrap(); + + let g2h_layout = config.g2h.layout(®ions.g2h_ring).unwrap(); + let h2g_layout = config.h2g.layout(®ions.h2g_ring).unwrap(); + let arena = regions.g2h_ring.start..regions.h2g_pool.end; + + let mem = HostMemOps::new(&scratch, arena).unwrap(); + + let h2g_pool = SlotPool::new(SlotLayout::new( + regions.h2g_pool.start, + config.h2g.buffer_size, + config.h2g_prefill_descs, + )) + .unwrap(); + + let mut producer = VirtqProducer::new(h2g_layout, mem, HostNotifier, h2g_pool.clone()); + let mut batch = producer.batch(); + + for _ in 0..config.h2g_prefill_descs { + let chain = batch + .chain() + .writable(config.h2g.buffer_size) + .build() + .unwrap(); + batch.submit(chain).unwrap(); + } + + batch.finish_without_notify(); + + write_published_arena_gpa(&scratch, config.arena.base_addr()).unwrap(); + + Self { + g2h_mem: HostMemOps::new(&scratch, regions.g2h_ring.clone()).unwrap(), + h2g_mem: HostMemOps::new(&scratch, regions.h2g_ring.clone()).unwrap(), + scratch, + g2h_ring: regions.g2h_ring, + h2g_ring: regions.h2g_ring, + g2h_pool: regions.g2h_pool, + h2g_pool: regions.h2g_pool, + g2h_layout, + h2g_layout, + } + } + + pub(crate) fn h2g_consumer(&self) -> H2gConsumer { + attach_canonical(&memory_layout(), &self.scratch).unwrap().1 + } + + pub(crate) fn g2h_consumer(&self) -> G2hConsumer { + attach_canonical(&memory_layout(), &self.scratch).unwrap().0 + } + + fn validate(&self) -> Result<()> { + let layout = memory_layout(); + let validator = Validator::new(&layout)?; + validator.validate_g2h(&self.g2h_mem, self.g2h_ring.clone())?; + validator.validate_h2g(&self.h2g_mem, self.h2g_ring.clone(), self.h2g_pool.clone())?; + Ok(()) + } + + fn g2h_desc(&self, index: u16) -> Descriptor { + read_desc(&self.g2h_mem, self.g2h_layout, index) + } + + pub(crate) fn h2g_desc(&self, index: u16) -> Descriptor { + read_desc(&self.h2g_mem, self.h2g_layout, index) + } + + fn set_g2h_desc(&self, index: u16, desc: Descriptor) { + write_desc(&self.g2h_mem, self.g2h_layout, index, desc); + } + + pub(crate) fn set_h2g_desc(&self, index: u16, desc: Descriptor) { + write_desc(&self.h2g_mem, self.h2g_layout, index, desc); + } + + pub(crate) fn h2g_buffer(&self, index: u16, addr: u64) -> Vec { + let desc = self.h2g_desc(index); + let mut bytes = vec![0; desc.len as usize]; + let pool = HostMemOps::new(&self.scratch, self.h2g_pool.clone()).unwrap(); + pool.read(addr, &mut bytes).unwrap(); + bytes + } +} + +pub(crate) fn memory_layout() -> SandboxMemoryLayout { + let mut config = SandboxConfiguration::default(); + config.set_scratch_size(SCRATCH_SIZE); + config.set_g2h_queue_size(16); + config.set_h2g_queue_size(8); + config.set_h2g_buffer_size(H2G_BUFFER_SIZE); + config.set_g2h_pool_pages(3); + config.set_h2g_pool_pages(3); + + SandboxMemoryLayout::new(config, 4096, 0, None).unwrap() +} + +fn host_scratch() -> HostSharedMemory { + ExclusiveSharedMemory::new(SCRATCH_SIZE).unwrap().build().0 +} + +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::new(&layout).unwrap(); + let config = validator.config; + let regions = validator.resolve_gva_regions().unwrap(); + + assert_eq!( + regions.g2h_ring.end - regions.g2h_ring.start, + config.g2h.dims.ring_len() as u64 + ); + assert_eq!( + regions.h2g_ring.end - regions.h2g_ring.start, + config.h2g.dims.ring_len() as u64 + ); + assert_eq!( + regions.g2h_pool.end - regions.g2h_pool.start, + config.g2h.dims.pool_len() as u64 + ); + assert_eq!( + regions.h2g_pool.end - regions.h2g_pool.start, + config.h2g.dims.pool_len() as u64 + ); +} + +#[test] +fn rejects_incorrect_published_arena() { + let layout = memory_layout(); + let validator = Validator::new(&layout).unwrap(); + let arena_gpa = validator.config.arena.base_addr() + 1; + assert!(validator.validate_published_arena(arena_gpa).is_err()); +} + +#[test] +fn validates_initial_virtq_images() { + TestVirtq::new().validate().unwrap(); +} + +#[test] +fn snapshots_and_restores_canonical_image() { + let queue = TestVirtq::new(); + let layout = memory_layout(); + let stale_pool = [0xa5; 16]; + let pool_mem = HostMemOps::new(&queue.scratch, queue.h2g_pool.clone()).unwrap(); + + pool_mem.write(queue.h2g_pool.start, &stale_pool).unwrap(); + + let captured = snapshot(&layout, &queue.scratch).unwrap(); + let restored = host_scratch(); + let allocator = layout.get_first_free_scratch_gpa(); + let allocator_offset = restored.mem_size() - 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, queue.h2g_pool.clone()).unwrap(); + let mut pool_bytes = [0; 16]; + + restored_pool + .read(queue.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 queue = TestVirtq::new(); + let layout = memory_layout(); + let mut captured = snapshot(&layout, &queue.scratch).unwrap(); + + captured.h2g_ring.fill(0); + let restored = host_scratch(); + + assert!(restore(&layout, &restored, &captured).is_err()); + assert_eq!(read_published_arena_gpa(&restored).unwrap(), 0); +} + +#[test] +fn restores_with_grown_page_tables() { + let queue = TestVirtq::new(); + let layout = memory_layout(); + let captured = snapshot(&layout, &queue.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, &captured).unwrap(); + assert_eq!( + read_published_arena_gpa(&restored).unwrap(), + grown_layout.get_transport_arena().base_addr() + ); +} + +#[test] +fn rejects_nonzero_g2h_descriptors() { + let queue = TestVirtq::new(); + let mut desc = queue.g2h_desc(0); + desc.addr = queue.g2h_pool.start; + queue.set_g2h_desc(0, desc); + assert!(queue.validate().is_err()); +} + +#[test] +fn rejects_invalid_h2g_descriptors() { + #[derive(Debug)] + enum Corruption { + OutsidePool, + Readable, + WrongSize, + Misaligned, + Overlapping, + } + + for corruption in [ + Corruption::OutsidePool, + Corruption::Readable, + Corruption::WrongSize, + Corruption::Misaligned, + Corruption::Overlapping, + ] { + let queue = TestVirtq::new(); + let mut desc = queue.h2g_desc(0); + + let index = match corruption { + Corruption::OutsidePool => { + desc.addr = queue.g2h_pool.start; + 0 + } + Corruption::Readable => { + desc.flags &= !DescFlags::WRITE.bits(); + 0 + } + Corruption::WrongSize => { + desc.len -= 1; + 0 + } + Corruption::Misaligned => { + desc.addr += 1; + 0 + } + Corruption::Overlapping => { + let first_addr = desc.addr; + desc = queue.h2g_desc(1); + desc.addr = first_addr; + 1 + } + }; + + queue.set_h2g_desc(index, desc); + assert!(queue.validate().is_err(), "{corruption:?}"); + } +} diff --git a/src/hyperlight_host/src/sandbox/config.rs b/src/hyperlight_host/src/sandbox/config.rs index 442da8415..cde07070d 100644 --- a/src/hyperlight_host/src/sandbox/config.rs +++ b/src/hyperlight_host/src/sandbox/config.rs @@ -17,6 +17,8 @@ limitations under the License. use std::cmp::max; use std::time::Duration; +use hyperlight_common::virtq::G2H_LOWER_SLOT_SIZE; +use hyperlight_common::vmem::PAGE_SIZE; #[cfg(target_os = "linux")] use libc::c_int; use tracing::{Span, instrument}; @@ -56,12 +58,6 @@ pub struct SandboxConfiguration { /// Guest gdb debug port #[cfg(gdb)] guest_debug_info: Option, - /// The size of the memory buffer that is made available for input to the - /// Guest Binary - input_data_size: usize, - /// The size of the memory buffer that is made available for input to the - /// Guest Binary - output_data_size: usize, /// The heap size to use in the guest sandbox. If set to 0, the heap /// size will be determined from the PE file header /// @@ -86,6 +82,18 @@ pub struct SandboxConfiguration { interrupt_vcpu_sigrtmin_offset: u8, /// How much writable memory to offer the guest scratch_size: usize, + /// Number of descriptors in the G2H virtqueue. + g2h_queue_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], @@ -95,34 +103,44 @@ pub struct SandboxConfiguration { } impl SandboxConfiguration { - /// The default size of input data - pub const DEFAULT_INPUT_SIZE: usize = 0x4000; - /// The minimum size of input data - pub const MIN_INPUT_SIZE: usize = 0x2000; - /// The default size of output data - pub const DEFAULT_OUTPUT_SIZE: usize = 0x4000; - /// The minimum size of output data - pub const MIN_OUTPUT_SIZE: usize = 0x2000; /// The default interrupt retry delay pub const DEFAULT_INTERRUPT_RETRY_DELAY: Duration = Duration::from_micros(500); /// The default signal offset from `SIGRTMIN` used to determine the signal number for interrupting 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; + /// The default scratch size keeps enough dynamic space to back the default + /// heap after reserving the transport arena and page tables. + pub const DEFAULT_SCRATCH_SIZE: usize = 0x56000; + /// 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 = 12; + /// The default total number of H2G pool pages. + pub const DEFAULT_H2G_POOL_PAGES: usize = 8; + /// 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. #[cfg(target_arch = "x86_64")] pub const MAX_GUEST_MSRS: usize = 16; - #[allow(clippy::too_many_arguments)] /// Create a new configuration for a sandbox with the given sizes. #[instrument(skip_all, parent = Span::current(), level= "Trace")] fn new( - input_data_size: usize, - output_data_size: usize, heap_size_override: Option, scratch_size: usize, interrupt_retry_delay: Duration, @@ -131,10 +149,14 @@ impl SandboxConfiguration { #[cfg(crashdump)] guest_core_dump: bool, ) -> Self { Self { - input_data_size: max(input_data_size, Self::MIN_INPUT_SIZE), - output_data_size: max(output_data_size, Self::MIN_OUTPUT_SIZE), heap_size_override: heap_size_override.unwrap_or(0), scratch_size, + 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)] @@ -148,20 +170,6 @@ impl SandboxConfiguration { } } - /// Set the size of the memory buffer that is made available for input to the guest - /// the minimum value is MIN_INPUT_SIZE - #[instrument(skip_all, parent = Span::current(), level= "Trace")] - pub fn set_input_data_size(&mut self, input_data_size: usize) { - self.input_data_size = max(input_data_size, Self::MIN_INPUT_SIZE); - } - - /// Set the size of the memory buffer that is made available for output from the guest - /// the minimum value is MIN_OUTPUT_SIZE - #[instrument(skip_all, parent = Span::current(), level= "Trace")] - pub fn set_output_data_size(&mut self, output_data_size: usize) { - self.output_data_size = max(output_data_size, Self::MIN_OUTPUT_SIZE); - } - /// Set the heap size to use in the guest sandbox. If set to 0, the heap size will be determined from the PE file header #[instrument(skip_all, parent = Span::current(), level= "Trace")] pub fn set_heap_size(&mut self, heap_size: u64) { @@ -279,24 +287,106 @@ impl SandboxConfiguration { } #[instrument(skip_all, parent = Span::current(), level= "Trace")] - pub(crate) fn get_input_data_size(&self) -> usize { - self.input_data_size + pub(crate) fn get_scratch_size(&self) -> usize { + self.scratch_size } + /// Set the size of the scratch regiong #[instrument(skip_all, parent = Span::current(), level= "Trace")] - pub(crate) fn get_output_data_size(&self) -> usize { - self.output_data_size + pub fn set_scratch_size(&mut self, scratch_size: usize) { + self.scratch_size = scratch_size; } + /// Get the G2H virtqueue descriptor count. #[instrument(skip_all, parent = Span::current(), level= "Trace")] - pub(crate) fn get_scratch_size(&self) -> usize { - self.scratch_size + pub fn get_g2h_queue_size(&self) -> usize { + self.g2h_queue_size } - /// Set the size of the scratch regiong + /// 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_scratch_size(&mut self, scratch_size: usize) { - self.scratch_size = scratch_size; + 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)] @@ -323,14 +413,25 @@ 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 { #[instrument(skip_all, parent = Span::current(), level= "Trace")] fn default() -> Self { Self::new( - Self::DEFAULT_INPUT_SIZE, - Self::DEFAULT_OUTPUT_SIZE, None, Self::DEFAULT_SCRATCH_SIZE, Self::DEFAULT_INTERRUPT_RETRY_DELAY, @@ -345,6 +446,8 @@ impl Default for SandboxConfiguration { #[cfg(test)] mod tests { + use hyperlight_common::vmem::PAGE_SIZE; + #[cfg(target_arch = "x86_64")] use super::GuestMsrError; use super::SandboxConfiguration; @@ -408,12 +511,8 @@ mod tests { #[test] fn overrides() { const HEAP_SIZE_OVERRIDE: u64 = 0x50000; - const INPUT_DATA_SIZE_OVERRIDE: usize = 0x4000; - const OUTPUT_DATA_SIZE_OVERRIDE: usize = 0x4001; const SCRATCH_SIZE_OVERRIDE: usize = 0x60000; let mut cfg = SandboxConfiguration::new( - INPUT_DATA_SIZE_OVERRIDE, - OUTPUT_DATA_SIZE_OVERRIDE, Some(HEAP_SIZE_OVERRIDE), SCRATCH_SIZE_OVERRIDE, SandboxConfiguration::DEFAULT_INTERRUPT_RETRY_DELAY, @@ -433,33 +532,95 @@ mod tests { cfg.scratch_size = 0x40000; assert_eq!(2048, cfg.heap_size_override); 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] - fn min_sizes() { - let mut cfg = SandboxConfiguration::new( - SandboxConfiguration::MIN_INPUT_SIZE - 1, - SandboxConfiguration::MIN_OUTPUT_SIZE - 1, - None, - SandboxConfiguration::DEFAULT_SCRATCH_SIZE, - SandboxConfiguration::DEFAULT_INTERRUPT_RETRY_DELAY, - SandboxConfiguration::INTERRUPT_VCPU_SIGRTMIN_OFFSET, - #[cfg(gdb)] - None, - #[cfg(crashdump)] - true, - ); - assert_eq!(SandboxConfiguration::MIN_INPUT_SIZE, cfg.input_data_size); - assert_eq!(SandboxConfiguration::MIN_OUTPUT_SIZE, cfg.output_data_size); - assert_eq!(0, cfg.heap_size_override); + 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()); + } + } - cfg.set_input_data_size(SandboxConfiguration::MIN_INPUT_SIZE - 1); - cfg.set_output_data_size(SandboxConfiguration::MIN_OUTPUT_SIZE - 1); + #[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()); - assert_eq!(SandboxConfiguration::MIN_INPUT_SIZE, cfg.input_data_size); - assert_eq!(SandboxConfiguration::MIN_OUTPUT_SIZE, cfg.output_data_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 { @@ -470,21 +631,6 @@ mod tests { use crate::sandbox::config::DebugInfo; proptest! { - #[test] - fn input_data_size(size in SandboxConfiguration::MIN_INPUT_SIZE..=SandboxConfiguration::MIN_INPUT_SIZE * 10) { - let mut cfg = SandboxConfiguration::default(); - cfg.set_input_data_size(size); - prop_assert_eq!(size, cfg.get_input_data_size()); - } - - #[test] - fn output_data_size(size in SandboxConfiguration::MIN_OUTPUT_SIZE..=SandboxConfiguration::MIN_OUTPUT_SIZE * 10) { - let mut cfg = SandboxConfiguration::default(); - cfg.set_output_data_size(size); - prop_assert_eq!(size, cfg.get_output_data_size()); - } - - #[test] fn heap_size_override(size in 0x1000..=0x10000u64) { let mut cfg = SandboxConfiguration::default(); diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index 623c24667..e2961d3d7 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -19,12 +19,10 @@ use std::path::Path; use std::path::PathBuf; use std::sync::{Arc, Mutex}; -use flatbuffers::FlatBufferBuilder; use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType}; use hyperlight_common::flatbuffer_wrappers::function_types::{ ParameterValue, ReturnType, ReturnValue, }; -use hyperlight_common::flatbuffer_wrappers::util::estimate_flatbuffer_capacity; use tracing::{Span, instrument}; use super::Callable; @@ -91,6 +89,8 @@ pub struct MultiUseSandbox { /// If the current state of the sandbox has been captured in a snapshot, /// that snapshot is stored here. pub(crate) snapshot: Option>, + /// Whether queue traffic occurred since the last canonical boundary. + transport_dirty: bool, /// Optional callback to discover page table roots from guest memory. /// Given (snapshot_mem, scratch_mem, cr3), returns a list of root GPAs. /// If not set, only CR3 is used as the single root. @@ -130,6 +130,7 @@ impl MultiUseSandbox { #[cfg(gdb)] dbg_mem_access_fn, snapshot: None, + transport_dirty: false, pt_root_finder: None, } } @@ -158,10 +159,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 (`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 /// 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. @@ -234,14 +234,13 @@ impl MultiUseSandbox { if caller_supplied_config { warn_on_layout_override(&config, snapshot.layout()); } - config.set_input_data_size(snapshot.layout().input_data_size()); - 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()); let load_info = snapshot.load_info(); let mgr = crate::mem::mgr::SandboxMemoryManager::from_snapshot(&snapshot)?; let (mut hshm, gshm) = mgr.build()?; + let restore_virtq = matches!(snapshot.next_action(), super::snapshot::NextAction::Call(_)); let page_size = u32::try_from(page_size::get())? as usize; @@ -322,6 +321,13 @@ impl MultiUseSandbox { #[cfg(gdb)] let dbg_mem_wrapper = Arc::new(Mutex::new(hshm.clone())); + if restore_virtq { + let virtq = snapshot.virtq().ok_or_else(|| { + crate::new_error!("running snapshot has no canonical transport state") + })?; + hshm.restore_virtq(virtq)?; + } + let sbox = MultiUseSandbox::from_uninit( host_funcs, hshm, @@ -378,6 +384,11 @@ impl MultiUseSandbox { if let Some(snapshot) = &self.snapshot { return Ok(snapshot.clone()); } + + if self.transport_dirty { + self.checkpoint_transport_for_snapshot()?; + } + let mapped_regions_iter = self.vm.get_mapped_regions(); let mapped_regions_vec: Vec = mapped_regions_iter.cloned().collect(); // Get CR3 from the vCPU @@ -428,6 +439,44 @@ impl MultiUseSandbox { Ok(snapshot) } + fn checkpoint_transport_for_snapshot(&mut self) -> Result<()> { + if let Err(error) = self.mem_mgr.begin_snapshot_checkpoint() { + self.poisoned |= error.is_poison_error(); + return Err(error); + } + + if let Err(error) = self.vm.dispatch_call_from_host( + &mut self.mem_mgr, + &self.host_funcs, + #[cfg(gdb)] + self.dbg_mem_access_fn.clone(), + ) { + let (error, should_poison) = error.promote(); + self.poisoned |= should_poison; + return Err(error); + } + + let guest_owned = match self.mem_mgr.finish_snapshot_checkpoint() { + Ok(guest_owned) => guest_owned, + Err(error) => { + self.poisoned |= error.is_poison_error(); + return Err(error); + } + }; + + if guest_owned != 0 { + // TODO: Parse retained pool-relative ranges and initialized lengths + // from the mailbox, sanitize them, and include them in the snapshot. + // The count-only protocol cannot preserve payloads safely. + return Err(HyperlightError::Error(format!( + "Cannot snapshot while {guest_owned} transport buffers are retained" + ))); + } + + self.transport_dirty = false; + Ok(()) + } + /// Restores the sandbox's memory to a previously captured snapshot state. /// /// The snapshot's memory layout must be structurally compatible @@ -600,6 +649,7 @@ impl MultiUseSandbox { // The restored snapshot is now our most current snapshot self.snapshot = Some(snapshot.clone()); + self.transport_dirty = false; // Clear poison state when successfully restoring from snapshot. // @@ -907,7 +957,7 @@ impl MultiUseSandbox { self.vm.clear_cancel(); let res = (|| { - let estimated_capacity = estimate_flatbuffer_capacity(function_name, &args); + self.transport_dirty = true; let fc = FunctionCall::new( function_name.to_string(), @@ -916,10 +966,7 @@ impl MultiUseSandbox { return_type, ); - let mut builder = FlatBufferBuilder::with_capacity(estimated_capacity); - let buffer = fc.encode(&mut builder); - - self.mem_mgr.write_guest_function_call(buffer)?; + let cid = self.mem_mgr.write_guest_function_call(&fc)?; let dispatch_res = self.vm.dispatch_call_from_host( &mut self.mem_mgr, @@ -936,7 +983,7 @@ impl MultiUseSandbox { return Err(error); } - let guest_result = self.mem_mgr.get_guest_function_call_result()?.into_inner(); + let guest_result = self.mem_mgr.read_h2g_result_from_g2h(cid)?.into_inner(); match guest_result { Ok(val) => Ok(val), @@ -958,15 +1005,7 @@ impl MultiUseSandbox { // Clear partial abort bytes so they don't leak across calls. self.mem_mgr.abort_buffer.clear(); - // In the happy path we do not need to clear io-buffers from the host because: - // - the serialized guest function call is zeroed out by the guest during deserialization, see call to `try_pop_shared_input_data_into::()` - // - the serialized guest function result is zeroed out by us (the host) during deserialization, see `get_guest_function_call_result` - // - any serialized host function call are zeroed out by us (the host) during deserialization, see `get_host_function_call` - // - any serialized host function result is zeroed out by the guest during deserialization, see `get_host_return_value` if let Err(e) = &res { - self.mem_mgr.clear_io_buffers(); - - // Determine if we should poison the sandbox. self.poisoned |= e.is_poison_error(); } @@ -1134,16 +1173,6 @@ fn warn_on_layout_override( snapshot: &crate::mem::layout::SandboxMemoryLayout, ) { let mismatches: &[(&str, u64, u64)] = &[ - ( - "input_data_size", - caller.get_input_data_size() as u64, - snapshot.input_data_size() as u64, - ), - ( - "output_data_size", - caller.get_output_data_size() as u64, - snapshot.output_data_size() as u64, - ), ( "heap_size", caller.get_heap_size(), @@ -1173,6 +1202,7 @@ mod tests { use std::thread; use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; + use hyperlight_common::func::Bytes; use hyperlight_testing::sandbox_sizes::{LARGE_HEAP_SIZE, MEDIUM_HEAP_SIZE, SMALL_HEAP_SIZE}; use hyperlight_testing::simple_guest_as_pathbuf; @@ -1181,6 +1211,11 @@ mod tests { use crate::sandbox::SandboxConfiguration; use crate::{GuestBinary, HyperlightError, MultiUseSandbox, Result, UninitializedSandbox}; + 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: MultiUseSandbox = { @@ -1269,7 +1304,7 @@ mod tests { let _ = sbox.snapshot().unwrap(); } - /// Make sure input/output buffers are properly reset after guest call (with host call) + /// Make sure transport buffers are reclaimed after host call failures. #[test] fn host_func_error() { let path = simple_guest_as_pathbuf(); @@ -1281,7 +1316,7 @@ mod tests { .unwrap(); let mut sandbox = sandbox.evolve().unwrap(); - // will exhaust io if leaky + // Repeated calls exhaust the transport if buffers leak. for _ in 0..1000 { let result = sandbox .call::( @@ -1306,12 +1341,10 @@ mod tests { .unwrap(); } - /// Make sure input/output buffers are properly reset after guest call (with host call) + /// Make sure transport buffers are reclaimed after guest calls. #[test] - fn io_buffer_reset() { - let mut cfg = SandboxConfiguration::default(); - cfg.set_input_data_size(4096); - cfg.set_output_data_size(4096); + fn transport_buffers_are_reclaimed() { + let cfg = SandboxConfiguration::default(); let path = simple_guest_as_pathbuf(); let mut sandbox = UninitializedSandbox::new(GuestBinary::FilePath(path), Some(cfg)).unwrap(); @@ -1355,19 +1388,21 @@ 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 a 24 KiB stack and 40 KiB heap. + // This catches guest stack reset and heap leaks. #[test] fn test_with_small_stack_and_heap() { let mut cfg = SandboxConfiguration::default(); - cfg.set_heap_size(32 * 1024); + cfg.set_heap_size(40 * 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 let min_scratch = hyperlight_common::layout::min_scratch_size( - cfg.get_input_data_size(), - cfg.get_output_data_size(), + cfg.get_g2h_queue_size(), + cfg.get_h2g_queue_size(), + cfg.get_g2h_pool_pages(), + cfg.get_h2g_pool_pages(), ); cfg.set_scratch_size(min_scratch + 0x10000 + 0x10000); @@ -1422,6 +1457,87 @@ mod tests { assert_eq!(res, 0); } + #[test] + fn snapshots_checkpoint_only_dirty_transport() { + let path = simple_guest_as_pathbuf(); + let sandbox = UninitializedSandbox::new(GuestBinary::FilePath(path), None).unwrap(); + let mut sandbox = sandbox.evolve().unwrap(); + + assert!(!sandbox.transport_dirty); + sandbox.call::("AddToStatic", 5i32).unwrap(); + assert!(sandbox.transport_dirty); + + let first = sandbox.snapshot().unwrap(); + assert!(!sandbox.transport_dirty); + let cached = sandbox.snapshot().unwrap(); + assert!(Arc::ptr_eq(&first, &cached)); + assert!(!sandbox.transport_dirty); + + sandbox.call::("AddToStatic", 5i32).unwrap(); + assert!(sandbox.transport_dirty); + sandbox.snapshot().unwrap(); + assert!(!sandbox.transport_dirty); + } + + #[test] + fn snapshots_reject_retained_transport_buffers_without_poisoning() { + let path = simple_guest_as_pathbuf(); + let mut sandbox = UninitializedSandbox::new(GuestBinary::FilePath(path), None).unwrap(); + sandbox + .register("HostEchoByteChunks", |value: Vec| value) + .unwrap(); + + let mut sandbox = sandbox.evolve().unwrap(); + let retained = vec![Bytes::from(vec![0xa5; 6 * 1024])]; + + let retained_len: i32 = sandbox + .call("RetainGuestByteChunks", retained.clone()) + .unwrap(); + + assert_eq!(retained_len, 6 * 1024); + + let Err(error) = sandbox.snapshot() else { + panic!("snapshot with retained H2G buffers succeeded"); + }; + + match error { + HyperlightError::Error(message) => { + assert!(message.contains("transport buffers are retained")) + } + err => unreachable!("unexpected snapshot error: {err:#}"), + } + assert!(!sandbox.poisoned()); + assert!(sandbox.transport_dirty); + + let released_len: i32 = sandbox.call("ReleaseGuestByteChunks", ()).unwrap(); + assert_eq!(released_len, retained_len); + + sandbox.snapshot().unwrap(); + assert!(!sandbox.transport_dirty); + + let retained_len: i32 = sandbox.call("RetainHostByteChunks", retained).unwrap(); + assert_eq!(retained_len, 6 * 1024); + + let Err(error) = sandbox.snapshot() else { + panic!("snapshot with retained G2H buffers succeeded"); + }; + + match error { + HyperlightError::Error(message) => { + assert!(message.contains("transport buffers are retained")) + } + err => unreachable!("unexpected snapshot error: {err:#}"), + } + assert!(!sandbox.poisoned()); + assert!(sandbox.transport_dirty); + + let released_len: i32 = sandbox.call("ReleaseHostByteChunks", ()).unwrap(); + assert_eq!(released_len, retained_len); + + sandbox.snapshot().unwrap(); + assert!(!sandbox.transport_dirty); + } + #[test] fn test_trigger_exception_on_guest() { let usbox = @@ -1735,6 +1851,7 @@ mod tests { let snapshot = sandbox.snapshot().unwrap(); sandbox2.restore(snapshot).unwrap(); + assert_virtq_attached(&sandbox2); assert_eq!(sandbox2.call::("GetStatic", ()).unwrap(), 42); } @@ -4000,8 +4117,10 @@ 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 = MultiUseSandbox::from_snapshot(snapshot, HostFunctions::default(), None).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"); @@ -4013,6 +4132,7 @@ mod tests { let snap = Snapshot::from_env(GuestBinary::FilePath(path), SandboxConfiguration::default()) .unwrap(); + assert!(snap.virtq().is_none()); let mut sbox = MultiUseSandbox::from_snapshot(Arc::new(snap), HostFunctions::default(), None) .unwrap(); @@ -4034,6 +4154,8 @@ mod tests { let mut b = MultiUseSandbox::from_snapshot(snapshot.clone(), HostFunctions::default(), None) .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); @@ -4043,6 +4165,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/outb.rs b/src/hyperlight_host/src/sandbox/outb.rs index 4b00d52b4..0116fd7b3 100644 --- a/src/hyperlight_host/src/sandbox/outb.rs +++ b/src/hyperlight_host/src/sandbox/outb.rs @@ -16,11 +16,14 @@ limitations under the License. use std::sync::{Arc, Mutex}; -use hyperlight_common::flatbuffer_wrappers::function_types::{FunctionCallResult, ParameterValue}; +use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCallType; +use hyperlight_common::flatbuffer_wrappers::function_types::FunctionCallResult; use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError}; use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData; use hyperlight_common::flatbuffer_wrappers::guest_log_level::LogLevel; use hyperlight_common::outb::{Exception, OutBAction}; +use hyperlight_common::transport::MsgKind; +use hyperlight_common::virtq::ReplyChain; use tracing::{Span, instrument}; use super::host_funcs::FunctionRegistry; @@ -28,6 +31,7 @@ use super::host_funcs::FunctionRegistry; use crate::hypervisor::regs::CommonRegisters; use crate::mem::mgr::SandboxMemoryManager; use crate::mem::shared_mem::HostSharedMemory; +use crate::mem::virtq; #[cfg(feature = "mem_profile")] use crate::sandbox::trace::MemTraceInfo; @@ -43,8 +47,6 @@ pub enum HandleOutbError { }, #[error("Invalid outb port: {0}")] InvalidPort(String), - #[error("Failed to read guest log data: {0}")] - ReadLogData(String), #[error("Failed to read host function call: {0}")] ReadHostFunctionCall(String), #[error("Failed to acquire lock at {0}:{1} - {2}")] @@ -58,14 +60,7 @@ pub enum HandleOutbError { MemProfile(String), } -#[instrument(err(Debug), skip_all, parent = Span::current(), level="Trace")] -pub(super) fn outb_log( - mgr: &mut SandboxMemoryManager, -) -> Result<(), HandleOutbError> { - let log_data: GuestLogData = mgr - .read_guest_log_data() - .map_err(|e| HandleOutbError::ReadLogData(e.to_string()))?; - +pub(crate) fn emit_guest_log(log_data: &GuestLogData) { // Emit guest log data as a tracing event with structured fields. // // We match on the level at runtime because tracing macros determine their @@ -134,8 +129,6 @@ pub(super) fn outb_log( ); } } - - Ok(()) } const ABORT_TERMINATOR: u8 = 0xFF; @@ -204,27 +197,7 @@ pub(crate) fn handle_outb( .try_into() .map_err(|e: anyhow::Error| HandleOutbError::InvalidPort(e.to_string()))? { - OutBAction::Log => outb_log(mem_mgr), - OutBAction::CallFunction => { - let call = mem_mgr - .get_host_function_call() - .map_err(|e| HandleOutbError::ReadHostFunctionCall(e.to_string()))?; - let name = call.function_name.clone(); - let args: Vec = call.parameters.unwrap_or(vec![]); - let res = host_funcs - .try_lock() - .map_err(|e| HandleOutbError::LockFailed(file!(), line!(), e.to_string()))? - .call_host_function(&name, args) - .map_err(|e| GuestError::new(ErrorCode::HostFunctionError, e.to_string())); - - let func_result = FunctionCallResult::new(res); - - mem_mgr - .write_response_from_host_function_call(&func_result) - .map_err(|e| HandleOutbError::WriteHostFunctionResponse(e.to_string()))?; - - Ok(()) - } + OutBAction::VirtqNotify => outb_virtq_call(mem_mgr, host_funcs), OutBAction::Abort => outb_abort(mem_mgr, data), OutBAction::DebugPrint => { let ch: char = match char::from_u32(data) { @@ -245,18 +218,141 @@ pub(crate) fn handle_outb( OutBAction::TraceMemoryFree => trace_info.handle_trace_mem_free(regs, mem_mgr), } } + +/// Drain G2H messages published before this notification. +fn outb_virtq_call( + mem_mgr: &mut SandboxMemoryManager, + host_funcs: &Arc>, +) -> Result<(), HandleOutbError> { + let max_recv_len = mem_mgr.layout.get_g2h_queue_dims().pool_len(); + + let Some(consumer) = mem_mgr.g2h_consumer.as_mut() else { + return Err(HandleOutbError::ReadHostFunctionCall( + "G2H consumer is not attached".into(), + )); + }; + + // Drain entries, processing logs, until we find one call. + let (mut request, reply, header) = loop { + let maybe_next = consumer.poll(max_recv_len).map_err(|error| { + HandleOutbError::ReadHostFunctionCall(format!("G2H poll failed: {error}")) + })?; + + let Some((mut request, reply)) = maybe_next else { + // No entry can be a backpressure or prefill notification. + return Ok(()); + }; + + let header = virtq::read_message_header(&mut request) + .map_err(|error| HandleOutbError::ReadHostFunctionCall(error.to_string()))?; + + match header.msg_kind() { + Ok(MsgKind::Request) => break (request, reply, header), + Ok(MsgKind::Log) => { + if header.cid != 0 { + return Err(HandleOutbError::ReadHostFunctionCall( + "G2H log has a nonzero correlation ID".into(), + )); + } + + if !matches!(reply, ReplyChain::Ack(_)) { + return Err(HandleOutbError::ReadHostFunctionCall( + "G2H log has writable response buffers".into(), + )); + } + + let log = virtq::read_guest_log_data(&mut request) + .map_err(|error| HandleOutbError::ReadHostFunctionCall(error.to_string()))?; + + emit_guest_log(&log); + + consumer.complete(request, reply).map_err(|error| { + HandleOutbError::ReadHostFunctionCall(format!( + "G2H log completion failed: {error}" + )) + })?; + } + Ok(kind) => { + return Err(HandleOutbError::ReadHostFunctionCall(format!( + "Expected G2H request, got {kind:?}" + ))); + } + Err(kind) => { + return Err(HandleOutbError::ReadHostFunctionCall(format!( + "Unknown G2H message kind {kind:#x}" + ))); + } + } + }; + + if header.cid == 0 { + return Err(HandleOutbError::ReadHostFunctionCall( + "G2H request has correlation ID zero".into(), + )); + } + + let mut resp = reply.into_writable().map_err(|_| { + HandleOutbError::WriteHostFunctionResponse( + "G2H request has no writable response buffers".into(), + ) + })?; + + let call = virtq::get_host_function_call(&mut request) + .map_err(|error| HandleOutbError::ReadHostFunctionCall(error.to_string()))?; + + if call.function_call_type() != FunctionCallType::Host { + return Err(HandleOutbError::ReadHostFunctionCall( + "G2H request does not target a host function".into(), + )); + } + + let name = call.function_name; + let args = call.parameters.unwrap_or_default(); + + let result = host_funcs + .try_lock() + .map_err(|err| HandleOutbError::LockFailed(file!(), line!(), err.to_string()))? + .call_host_function(&name, args) + .map_err(|err| GuestError::new(ErrorCode::HostFunctionError, err.to_string())); + + let result = FunctionCallResult::new(result); + let resp_capacity = resp.capacity(); + + // Capacity is checked before writing, so an oversized result leaves the + // chain untouched and can be replaced with a bounded transport error. + if !virtq::try_write_response(&mut resp, header.cid, &result) + .map_err(|err| HandleOutbError::WriteHostFunctionResponse(err.to_string()))? + { + let fallback = FunctionCallResult::new(Err(GuestError::new( + ErrorCode::HostFunctionError, + "Host response exceeds virtqueue capacity".into(), + ))); + + // The guest must receive a response for this correlation id. Failure + // to fit even this small error makes the transport unusable. + if !virtq::try_write_response(&mut resp, header.cid, &fallback) + .map_err(|err| HandleOutbError::WriteHostFunctionResponse(err.to_string()))? + { + return Err(HandleOutbError::WriteHostFunctionResponse(format!( + "Writable response capacity {resp_capacity} cannot hold a transport error" + ))); + } + } + + consumer + .complete(request, resp) + .map_err(|err| HandleOutbError::WriteHostFunctionResponse(err.to_string()))?; + + Ok(()) +} + #[cfg(test)] mod tests { use hyperlight_common::flatbuffer_wrappers::guest_log_level::LogLevel; use hyperlight_testing::logger::{LOGGER, Logger}; - use hyperlight_testing::simple_guest_as_pathbuf; use tracing_core::callsite::rebuild_interest_cache; - use super::outb_log; - use crate::GuestBinary; - use crate::mem::mgr::SandboxMemoryManager; - use crate::sandbox::SandboxConfiguration; - use crate::sandbox::outb::GuestLogData; + use super::{GuestLogData, emit_guest_log}; use crate::testing::log_values::test_value_as_str; fn new_guest_log_data(level: LogLevel) -> GuestLogData { @@ -271,140 +367,70 @@ mod tests { } // Verifies that guest log events are forwarded to a `log` logger when no - // tracing subscriber is set. This exercises the `tracing` crate's built-in - // `log` compatibility feature, proving that consumers who only set up a - // `log` logger (not a tracing subscriber) still receive guest output. + // tracing subscriber is set. #[test] #[ignore] - fn test_log_outb_log() { + fn test_log_emit_guest_log() { Logger::initialize_test_logger(); LOGGER.set_max_level(log::LevelFilter::Off); - let sandbox_cfg = SandboxConfiguration::default(); - - let new_mgr = || { - let bin = GuestBinary::FilePath(simple_guest_as_pathbuf()); - let snapshot = crate::sandbox::snapshot::Snapshot::from_env(bin, sandbox_cfg).unwrap(); - let mgr = SandboxMemoryManager::from_snapshot(&snapshot).unwrap(); - let (hmgr, _) = mgr.build().unwrap(); - hmgr - }; - { - // We set a logger but there is no guest log data - // in memory, so expect a log operation to fail - let mut mgr = new_mgr(); - assert!(outb_log(&mut mgr).is_err()); - } - { - // Write a log message so outb_log will succeed. - // Since the logger level is set off, expect logs to be no-ops - let mut mgr = new_mgr(); - let log_msg = new_guest_log_data(LogLevel::Information); - - let guest_log_data_buffer: Vec = log_msg.try_into().unwrap(); - let offset = mgr.layout.get_output_data_buffer_scratch_host_offset(); - mgr.scratch_mem - .push_buffer( - offset, - sandbox_cfg.get_output_data_size(), - &guest_log_data_buffer, - ) - .unwrap(); - - let res = outb_log(&mut mgr); - assert!(res.is_ok()); - assert_eq!(0, LOGGER.num_log_calls()); - LOGGER.clear_log_calls(); - } - { - // now, test logging - LOGGER.set_max_level(log::LevelFilter::Trace); - let mut mgr = new_mgr(); + emit_guest_log(&new_guest_log_data(LogLevel::Information)); + assert_eq!(0, LOGGER.num_log_calls()); + LOGGER.clear_log_calls(); + + LOGGER.set_max_level(log::LevelFilter::Trace); + let levels = vec![ + LogLevel::Trace, + LogLevel::Debug, + LogLevel::Information, + LogLevel::Warning, + LogLevel::Error, + LogLevel::Critical, + LogLevel::None, + ]; + + for level in levels { LOGGER.clear_log_calls(); + emit_guest_log(&new_guest_log_data(level)); + + LOGGER.test_log_records(|log_calls| { + let expected_level: tracing::Level = match level { + LogLevel::Trace => tracing::Level::TRACE, + LogLevel::Debug => tracing::Level::DEBUG, + LogLevel::Information => tracing::Level::INFO, + LogLevel::Warning => tracing::Level::WARN, + LogLevel::Error | LogLevel::Critical => tracing::Level::ERROR, + LogLevel::None => tracing::Level::TRACE, + }; - // set up the logger and set the log level to the maximum - // possible (Trace) to ensure we're able to test all - // the possible branches of the match in outb_log - - let levels = vec![ - LogLevel::Trace, - LogLevel::Debug, - LogLevel::Information, - LogLevel::Warning, - LogLevel::Error, - LogLevel::Critical, - LogLevel::None, - ]; - for level in levels { - let layout = mgr.layout; - let log_data = new_guest_log_data(level); - - let guest_log_data_buffer: Vec = log_data.clone().try_into().unwrap(); - mgr.scratch_mem - .push_buffer( - layout.get_output_data_buffer_scratch_host_offset(), - sandbox_cfg.get_output_data_size(), - guest_log_data_buffer.as_slice(), - ) - .unwrap(); - - outb_log(&mut mgr).unwrap(); - - LOGGER.test_log_records(|log_calls| { - let expected_level: tracing::Level = match level { - LogLevel::Trace => tracing::Level::TRACE, - LogLevel::Debug => tracing::Level::DEBUG, - LogLevel::Information => tracing::Level::INFO, - LogLevel::Warning => tracing::Level::WARN, - LogLevel::Error => tracing::Level::ERROR, - LogLevel::Critical => tracing::Level::ERROR, - LogLevel::None => tracing::Level::TRACE, - }; - - assert!( - log_calls - .iter() - .filter(|log_call| { - log_call.level.as_str() == expected_level.as_str() - && log_call.args.contains("test log") - }) - .count() - == 1, - "log call did not occur for level {:?}", - level.clone() - ); - }); - } + assert_eq!( + log_calls + .iter() + .filter(|log_call| { + log_call.level.as_str() == expected_level.as_str() + && log_call.args.contains("test log") + }) + .count(), + 1, + "log call did not occur for level {level:?}" + ); + }); } } - // Tests that outb_log emits traces when a trace subscriber is set + // Tests that guest logs emit traces when a trace subscriber is set // this test is ignored because it is incompatible with other tests , specifically those which require a logger for tracing // marking this test as ignored means that running `cargo test` will not run this test but will allow a developer who runs that command // from their workstation to be successful without needed to know about test interdependencies // this test will be run explicitly as a part of the CI pipeline #[ignore] #[test] - fn test_trace_outb_log() { + fn test_trace_emit_guest_log() { Logger::initialize_log_tracer(); rebuild_interest_cache(); let subscriber = hyperlight_testing::tracing_subscriber::TracingSubscriber::new(tracing::Level::TRACE); - let sandbox_cfg = SandboxConfiguration::default(); tracing::subscriber::with_default(subscriber.clone(), || { - let new_mgr = || { - let bin = GuestBinary::FilePath(simple_guest_as_pathbuf()); - let snapshot = - crate::sandbox::snapshot::Snapshot::from_env(bin, sandbox_cfg).unwrap(); - let mgr = SandboxMemoryManager::from_snapshot(&snapshot).unwrap(); - let (hmgr, _) = mgr.build().unwrap(); - hmgr - }; - - // as a span does not exist one will be automatically created - // after that there will be an event for each log message - // we are interested only in the events for the log messages that we created - let levels = vec![ LogLevel::Trace, LogLevel::Debug, @@ -415,23 +441,11 @@ mod tests { LogLevel::None, ]; for level in levels { - let mut mgr = new_mgr(); - let layout = mgr.layout; let log_data: GuestLogData = new_guest_log_data(level); subscriber.clear(); + emit_guest_log(&log_data); - let guest_log_data_buffer: Vec = log_data.try_into().unwrap(); - mgr.scratch_mem - .push_buffer( - layout.get_output_data_buffer_scratch_host_offset(), - sandbox_cfg.get_output_data_size(), - guest_log_data_buffer.as_slice(), - ) - .unwrap(); - subscriber.clear(); - outb_log(&mut mgr).unwrap(); - - subscriber.test_trace_records(|spans, events| { + subscriber.test_trace_records(|_, events| { let expected_level = match level { LogLevel::Trace => "TRACE", LogLevel::Debug => "DEBUG", @@ -442,38 +456,6 @@ mod tests { LogLevel::None => "TRACE", }; - // We cannot get the parent span using the `current_span()` method as by the time we get to this point that span has been exited so there is no current span - // We need to make sure that the span that we created is in the spans map instead - // We are only interested in the first one that was created when calling outb_log. - - assert!(!spans.is_empty(), "expected at least one span, found none"); - - let span_value = spans - .get(&1) - .unwrap() - .as_object() - .unwrap() - .get("span") - .unwrap() - .get("attributes") - .unwrap() - .as_object() - .unwrap() - .get("metadata") - .unwrap() - .as_object() - .unwrap(); - - //test_value_as_str(span_value, "level", "INFO"); - test_value_as_str(span_value, "module_path", "hyperlight_host::sandbox::outb"); - let expected_file = if cfg!(windows) { - "src\\hyperlight_host\\src\\sandbox\\outb.rs" - } else { - "src/hyperlight_host/src/sandbox/outb.rs" - }; - test_value_as_str(span_value, "file", expected_file); - test_value_as_str(span_value, "target", "hyperlight_host::sandbox::outb"); - let mut count_matching_events = 0; for json_value in events { diff --git a/src/hyperlight_host/src/sandbox/snapshot/file/config.rs b/src/hyperlight_host/src/sandbox/snapshot/file/config.rs index 01c2f501b..46ab928eb 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/config.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file/config.rs @@ -157,7 +157,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,14 +214,18 @@ pub(super) struct OciSnapshotConfig { #[derive(Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub(super) struct MemoryLayout { - pub(super) input_data_size: usize, - pub(super) output_data_size: usize, pub(super) heap_size: usize, pub(super) code_size: usize, pub(super) init_data_size: usize, /// 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, } @@ -254,6 +258,7 @@ enum ParameterTypeRepr { String, Bool, VecBytes, + ByteChunks, } /// JSON-friendly mirror of @@ -271,6 +276,7 @@ enum ReturnTypeRepr { Bool, Void, VecBytes, + ByteChunks, } impl From<&ParameterType> for ParameterTypeRepr { @@ -285,6 +291,7 @@ impl From<&ParameterType> for ParameterTypeRepr { ParameterType::String => Self::String, ParameterType::Bool => Self::Bool, ParameterType::VecBytes => Self::VecBytes, + ParameterType::ByteChunks => Self::ByteChunks, } } } @@ -301,6 +308,7 @@ impl From for ParameterType { ParameterTypeRepr::String => Self::String, ParameterTypeRepr::Bool => Self::Bool, ParameterTypeRepr::VecBytes => Self::VecBytes, + ParameterTypeRepr::ByteChunks => Self::ByteChunks, } } } @@ -318,6 +326,7 @@ impl From<&ReturnType> for ReturnTypeRepr { ReturnType::Bool => Self::Bool, ReturnType::Void => Self::Void, ReturnType::VecBytes => Self::VecBytes, + ReturnType::ByteChunks => Self::ByteChunks, } } } @@ -335,6 +344,7 @@ impl From for ReturnType { ReturnTypeRepr::Bool => Self::Bool, ReturnTypeRepr::Void => Self::Void, ReturnTypeRepr::VecBytes => Self::VecBytes, + ReturnTypeRepr::ByteChunks => Self::ByteChunks, } } } @@ -460,12 +470,14 @@ impl OciSnapshotConfig { // checked against `snapshot_size` in `load_inner`. let max_region = SandboxMemoryLayout::MAX_MEMORY_SIZE; for (name, value) in [ - ("input_data_size", self.layout.input_data_size), - ("output_data_size", self.layout.output_data_size), ("heap_size", self.layout.heap_size), ("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!( @@ -477,6 +489,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; @@ -693,6 +754,7 @@ mod tests { ParameterType::String, ParameterType::Bool, ParameterType::VecBytes, + ParameterType::ByteChunks, ]; for p in variants { let back: ParameterType = ParameterTypeRepr::from(&p).into(); @@ -715,6 +777,7 @@ mod tests { ReturnType::Bool, ReturnType::Void, ReturnType::VecBytes, + ReturnType::ByteChunks, ]; for r in variants { let back: ReturnType = ReturnTypeRepr::from(&r).into(); @@ -767,13 +830,17 @@ mod tests { #[cfg(target_arch = "x86_64")] msrs: Vec::new(), layout: MemoryLayout { - input_data_size: 0, - output_data_size: 0, heap_size: 0, code_size: 0, 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, }, @@ -840,7 +907,7 @@ mod schema_pin { const PINNED_CALL: &str = r#"{ "hyperlight_version": "x.y.z", "arch": "x86_64", - "abi_version": 1, + "abi_version": 2, "hypervisor": "mshv", "cpu_vendor": "intel", "stack_top_gva": 3735928559, @@ -999,13 +1066,17 @@ mod schema_pin { } ], "layout": { - "input_data_size": 1, - "output_data_size": 2, "heap_size": 3, "code_size": 4, "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 }, @@ -1026,7 +1097,7 @@ mod schema_pin { const PINNED_CALL: &str = r#"{ "hyperlight_version": "x.y.z", "arch": "aarch64", - "abi_version": 1, + "abi_version": 2, "hypervisor": "mshv", "cpu_vendor": "intel", "stack_top_gva": 3735928559, @@ -1041,13 +1112,17 @@ mod schema_pin { "sp_el1": 6 }, "layout": { - "input_data_size": 1, - "output_data_size": 2, "heap_size": 3, "code_size": 4, "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 }, @@ -1086,7 +1161,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 661bd4a04..eab4e44b9 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs @@ -19,15 +19,20 @@ limitations under the License. // 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; +pub(in crate::sandbox::snapshot) const MT_TRANSPORT_V1: &str = + "application/vnd.hyperlight.snapshot.transport.v1"; +pub(in crate::sandbox::snapshot) const MT_TRANSPORT_CURRENT: &str = MT_TRANSPORT_V1; /// 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 = 1; +pub(in crate::sandbox::snapshot) const SNAPSHOT_ABI_VERSION: u32 = 2; /// 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 2a0f00d2f..a27eb57d0 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs @@ -39,13 +39,15 @@ 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, + MT_TRANSPORT_CURRENT, MT_TRANSPORT_V1, SNAPSHOT_ABI_VERSION, }; use self::reference::{OciDigest, OciReference, OciTag}; use super::{NextAction, Snapshot}; use crate::mem::layout::SandboxMemoryLayout; use crate::mem::memory_region::MemoryRegionFlags; use crate::mem::shared_mem::{ReadonlySharedMemory, SharedMemory}; +use crate::mem::virtq::VirtqSnapshot; pub(super) const OCI_LAYOUT_VERSION: &str = "1.0.0"; @@ -61,6 +63,10 @@ pub fn host_cpu_vendor_golden_tag() -> Option<&'static str> { /// `oci-layout`, `index.json`, the OCI image manifest, and the /// Hyperlight config blob. Bounds the allocation done before parsing. const MAX_JSON_BLOB_SIZE: u64 = 1024 * 1024; +const MAX_TRANSPORT_BLOB_SIZE: u64 = 2 * 1024 * 1024; +const TRANSPORT_MAGIC: [u8; 8] = *b"HLVQSNAP"; +const TRANSPORT_VERSION: u32 = 1; +const TRANSPORT_HEADER_LEN: usize = 40; /// Reject a JSON artifact larger than the cap the loader reads with /// [`read_bounded`]. The writer holds to the same cap so every layout @@ -77,6 +83,109 @@ fn check_json_blob_size(what: &str, len: usize) -> crate::Result<()> { Ok(()) } +fn encode_transport(snapshot: &VirtqSnapshot) -> crate::Result> { + let g2h_len = snapshot.g2h_ring().len(); + let h2g_len = snapshot.h2g_ring().len(); + + let total_len = TRANSPORT_HEADER_LEN + .checked_add(g2h_len) + .and_then(|len| len.checked_add(h2g_len)) + .ok_or_else(|| crate::new_error!("snapshot transport length overflow"))?; + + if total_len as u64 > MAX_TRANSPORT_BLOB_SIZE { + return Err(crate::new_error!( + "transport blob of {total_len} bytes exceeds the {MAX_TRANSPORT_BLOB_SIZE} byte maximum" + )); + } + + let mut bytes = Vec::new(); + bytes + .try_reserve_exact(total_len) + .map_err(|error| crate::new_error!("failed to allocate transport blob: {error}"))?; + + bytes.extend_from_slice(&TRANSPORT_MAGIC); + bytes.extend_from_slice(&TRANSPORT_VERSION.to_le_bytes()); + bytes.extend_from_slice(&0u32.to_le_bytes()); + bytes.extend_from_slice(&u64::try_from(snapshot.scratch_size())?.to_le_bytes()); + bytes.extend_from_slice(&u64::try_from(g2h_len)?.to_le_bytes()); + bytes.extend_from_slice(&u64::try_from(h2g_len)?.to_le_bytes()); + bytes.extend_from_slice(snapshot.g2h_ring()); + bytes.extend_from_slice(snapshot.h2g_ring()); + Ok(bytes) +} + +fn read_transport_field(bytes: &mut &[u8], field: &str) -> crate::Result<[u8; N]> { + let (value, remaining) = bytes + .split_at_checked(N) + .ok_or_else(|| crate::new_error!("snapshot transport {field} is truncated"))?; + + let mut array = [0; N]; + array.copy_from_slice(value); + + *bytes = remaining; + Ok(array) +} + +fn decode_transport(bytes: &[u8]) -> crate::Result { + let total_len = bytes.len(); + let mut bytes = bytes; + + if read_transport_field(&mut bytes, "magic")? != TRANSPORT_MAGIC { + return Err(crate::new_error!("snapshot transport magic is invalid")); + } + + let version = u32::from_le_bytes(read_transport_field(&mut bytes, "version")?); + if version != TRANSPORT_VERSION { + return Err(crate::new_error!( + "snapshot transport version mismatch: file has version {version}, this build expects {TRANSPORT_VERSION}" + )); + } + + let reserved = u32::from_le_bytes(read_transport_field(&mut bytes, "reserved field")?); + if reserved != 0 { + return Err(crate::new_error!( + "snapshot transport reserved field is nonzero" + )); + } + + let scratch_size = usize::try_from(u64::from_le_bytes(read_transport_field( + &mut bytes, + "scratch size", + )?))?; + + let g2h_len = usize::try_from(u64::from_le_bytes(read_transport_field( + &mut bytes, + "G2H ring length", + )?))?; + + let h2g_len = usize::try_from(u64::from_le_bytes(read_transport_field( + &mut bytes, + "H2G ring length", + )?))?; + + let expected_len = TRANSPORT_HEADER_LEN + .checked_add(g2h_len) + .and_then(|len| len.checked_add(h2g_len)) + .ok_or_else(|| crate::new_error!("snapshot transport length overflow"))?; + + if total_len != expected_len { + return Err(crate::new_error!( + "snapshot transport length {} does not match header length {expected_len}", + total_len + )); + } + + let (g2h_ring, h2g_ring) = bytes + .split_at_checked(g2h_len) + .ok_or_else(|| crate::new_error!("snapshot transport G2H ring is truncated"))?; + + Ok(VirtqSnapshot::new( + scratch_size, + g2h_ring.to_vec(), + h2g_ring.to_vec(), + )) +} + /// Select one manifest descriptor from `index` by `reference`. /// /// A tag matches the `org.opencontainers.image.ref.name` annotation @@ -285,6 +394,27 @@ fn open_snapshot_blob( Ok(snap_file) } +fn load_transport_blob( + blobs_dir: &Path, + transport_desc: &Descriptor, + verify_blobs: bool, +) -> crate::Result> { + let transport_hex = parse_oci_digest(transport_desc.digest())?; + let transport_path = blobs_dir.join(&transport_hex); + let bytes = read_bounded(&transport_path, MAX_TRANSPORT_BLOB_SIZE)?; + if bytes.len() as u64 != transport_desc.size() { + return Err(crate::new_error!( + "transport blob size mismatch: descriptor says {}, file is {}", + transport_desc.size(), + bytes.len() + )); + } + if verify_blobs { + verify_blob_bytes("transport", &bytes, &transport_hex)?; + } + Ok(bytes) +} + impl Snapshot { /// Save this snapshot into an OCI Image Layout directory on disk. /// The saved snapshot can be loaded later with @@ -518,6 +648,14 @@ impl Snapshot { let snapshot_digest = Digest256::from_bytes(memory_bytes); put_blob_if_absent(&blobs_dir, &snapshot_digest, memory_bytes)?; + // Transport blob: the canonical ring image omitted from memory. + let transport = self.virtq.as_ref().ok_or_else(|| { + crate::new_error!("initialized snapshot has no canonical transport state") + })?; + let transport_bytes = encode_transport(transport)?; + let transport_digest = Digest256::from_bytes(&transport_bytes); + put_blob(&blobs_dir, &transport_digest, &transport_bytes)?; + // Config blob. let cfg_digest = Digest256::from_bytes(cfg_bytes); put_blob(&blobs_dir, &cfg_digest, cfg_bytes)?; @@ -535,6 +673,12 @@ impl Snapshot { .size(memory_size as u64) .build() .map_err(|e| crate::new_error!("failed to build snapshot descriptor: {}", e))?; + let transport_descriptor = DescriptorBuilder::default() + .media_type(MediaType::Other(MT_TRANSPORT_CURRENT.to_string())) + .digest(oci_digest(&transport_digest)?) + .size(transport_bytes.len() as u64) + .build() + .map_err(|e| crate::new_error!("failed to build transport descriptor: {}", e))?; // `artifactType` is set equal to `config.mediaType` per OCI // image-spec "Guidelines for Artifact Usage". Registries // surface this on the distribution-spec referrers API. Tools @@ -544,7 +688,7 @@ impl Snapshot { .media_type(MediaType::ImageManifest) .artifact_type(MediaType::Other(MT_CONFIG_CURRENT.to_string())) .config(config_descriptor) - .layers(vec![snapshot_descriptor]) + .layers(vec![snapshot_descriptor, transport_descriptor]) .build() .map_err(|e| crate::new_error!("failed to build OCI manifest: {}", e))?; let manifest_bytes = serde_json::to_vec_pretty(&manifest) @@ -593,6 +737,10 @@ impl Snapshot { )); } }; + let transport = self.virtq.as_ref().ok_or_else(|| { + crate::new_error!("initialized snapshot has no canonical transport state") + })?; + transport.preflight(&self.layout)?; let host_functions = match &self.host_functions.host_functions { Some(v) => v.iter().map(HostFunction::from).collect(), @@ -618,13 +766,17 @@ impl Snapshot { .ok_or_else(|| crate::new_error!("snapshot has no MSR state"))? .clone(), layout: MemoryLayout { - input_data_size: l.input_data_size(), - output_data_size: l.output_data_size(), heap_size: l.heap_size(), code_size: l.code_size(), 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(), }, @@ -670,7 +822,7 @@ impl Snapshot { /// /// # Verification /// - /// This method does not check the manifest, config, or snapshot + /// This method does not check the manifest, config, memory, or transport /// blobs against their recorded sha256 digests. Load only from a /// layout you trust. /// @@ -706,7 +858,7 @@ impl Snapshot { /// Loads a snapshot like [`Snapshot::load`]. See its rustdoc for /// `path`, `reference`, portability, and the file-mutation /// hazard. This method additionally checks the manifest, config, - /// and snapshot blobs against their recorded sha256 digests + /// memory, and transport blobs against their recorded sha256 digests /// before use, at the expense of some performance. /// /// # Trust @@ -747,16 +899,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 )); } } @@ -783,9 +940,9 @@ impl Snapshot { } } let layers = manifest.layers(); - if layers.len() != 1 { + if layers.len() != 2 { return Err(crate::new_error!( - "expected exactly one OCI layer (the snapshot), found {}", + "expected exactly two OCI layers (memory and transport), found {}", layers.len() )); } @@ -801,6 +958,18 @@ impl Snapshot { )); } } + let transport_desc = &layers[1]; + let transport_media = transport_desc.media_type().to_string(); + match transport_media.as_str() { + MT_TRANSPORT_V1 => {} + other => { + return Err(crate::new_error!( + "unexpected transport layer media type {:?} (supported: {:?})", + other, + MT_TRANSPORT_V1 + )); + } + } // 4. config blob let cfg = load_config(&blobs_dir, cfg_desc, verify_blobs)?; @@ -809,13 +978,19 @@ impl Snapshot { // handle so an attacker cannot swap the file between // verification and mapping. let snap_file = open_snapshot_blob(&blobs_dir, snap_desc, cfg.memory_size, verify_blobs)?; + let transport_bytes = load_transport_blob(&blobs_dir, transport_desc, verify_blobs)?; + let virtq = decode_transport(&transport_bytes)?; // 6. Reconstruct layout. let mut sbox_cfg = crate::sandbox::SandboxConfiguration::default(); - sbox_cfg.set_input_data_size(cfg.layout.input_data_size); - 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(|| { @@ -881,6 +1056,7 @@ impl Snapshot { // 8. Build the next action + sregs back from the config. let next_action = NextAction::Call(cfg.entrypoint_addr); + virtq.preflight(&layout)?; // 9. Reconstitute host_functions metadata. let snapshot_generation = cfg.snapshot_generation; @@ -909,6 +1085,78 @@ impl Snapshot { original_entrypoint: cfg.original_entrypoint_addr, snapshot_generation, host_functions, + virtq: Some(virtq), }) } } + +#[cfg(test)] +mod transport_tests { + use super::*; + + #[test] + fn transport_blob_round_trips() { + let snapshot = VirtqSnapshot::new(0x20_000, vec![1, 2, 3], vec![4, 5]); + let bytes = encode_transport(&snapshot).unwrap(); + let decoded = decode_transport(&bytes).unwrap(); + + assert_eq!(decoded.scratch_size(), snapshot.scratch_size()); + assert_eq!(decoded.g2h_ring(), snapshot.g2h_ring()); + assert_eq!(decoded.h2g_ring(), snapshot.h2g_ring()); + } + + #[test] + fn transport_blob_rejects_header_corruption() { + let snapshot = VirtqSnapshot::new(0x20_000, vec![1], vec![2]); + let mut bytes = encode_transport(&snapshot).unwrap(); + + bytes[8..12].copy_from_slice(&TRANSPORT_VERSION.wrapping_add(1).to_le_bytes()); + assert!( + decode_transport(&bytes) + .unwrap_err() + .to_string() + .contains("version mismatch") + ); + + bytes[8..12].copy_from_slice(&TRANSPORT_VERSION.to_le_bytes()); + bytes[12] = 1; + assert!( + decode_transport(&bytes) + .unwrap_err() + .to_string() + .contains("reserved") + ); + } + + #[test] + fn transport_blob_rejects_truncated_header_fields() { + let snapshot = VirtqSnapshot::new(0x20_000, vec![1], vec![2]); + let bytes = encode_transport(&snapshot).unwrap(); + + for (len, field) in [ + (0, "magic"), + (8, "version"), + (12, "reserved field"), + (16, "scratch size"), + (24, "G2H ring length"), + (32, "H2G ring length"), + ] { + let error = decode_transport(&bytes[..len]).unwrap_err(); + assert!(error.to_string().contains(field), "{error:?}"); + } + } + + #[test] + fn transport_blob_rejects_length_mismatch() { + let snapshot = VirtqSnapshot::new(0x20_000, vec![1], vec![2]); + let mut bytes = encode_transport(&snapshot).unwrap(); + bytes.push(3); + + assert!( + decode_transport(&bytes) + .unwrap_err() + .to_string() + .contains("does not match") + ); + } +} diff --git a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs b/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs index 4e0604c86..cc84416a3 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs @@ -91,6 +91,25 @@ fn find_snapshot_blob(oci_dir: &std::path::Path) -> std::path::PathBuf { oci_dir.join("blobs").join("sha256").join(snap_digest) } +/// Locate the transport (layer 1) blob inside `oci_dir`. +fn find_transport_blob(oci_dir: &std::path::Path) -> std::path::PathBuf { + let index: Value = + serde_json::from_slice(&std::fs::read(oci_dir.join("index.json")).unwrap()).unwrap(); + let manifest_digest = index["manifests"][0]["digest"] + .as_str() + .unwrap() + .strip_prefix("sha256:") + .unwrap(); + let manifest_path = oci_dir.join("blobs").join("sha256").join(manifest_digest); + let manifest: Value = serde_json::from_slice(&std::fs::read(&manifest_path).unwrap()).unwrap(); + let transport_digest = manifest["layers"][1]["digest"] + .as_str() + .unwrap() + .strip_prefix("sha256:") + .unwrap(); + oci_dir.join("blobs").join("sha256").join(transport_digest) +} + // In-memory `from_snapshot` round-trips. #[test] @@ -1399,8 +1418,7 @@ fn save_same_tag_same_content_is_idempotent() { ); } -/// Two tags written from one in-memory snapshot share all three blobs -/// (manifest, config, snapshot). +/// Two tags written from one in-memory snapshot share all four blobs. #[test] fn save_shares_blobs_across_tags_with_identical_content() { let snap = create_snapshot(); @@ -1414,7 +1432,7 @@ fn save_shares_blobs_across_tags_with_identical_content() { .unwrap() .filter_map(|e| e.ok().map(|e| e.file_name())) .collect(); - assert_eq!(blobs.len(), 3, "expected 3 deduped blobs, got {:?}", blobs); + assert_eq!(blobs.len(), 4, "expected 4 deduped blobs, got {:?}", blobs); } /// Replacing one tag in a three-tag layout keeps the other two @@ -1574,6 +1592,28 @@ fn checked_load_rejects_snapshot_blob_byte_mutation() { ); } +#[test] +fn checked_load_rejects_transport_blob_byte_mutation() { + let snapshot = create_snapshot(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("snap"); + snapshot + .save(&path, &OciTag::new("latest").unwrap()) + .unwrap(); + + let transport_path = find_transport_blob(&path); + let mut bytes = std::fs::read(&transport_path).unwrap(); + let mid = bytes.len() / 2; + bytes[mid] ^= 0xFF; + std::fs::write(&transport_path, bytes).unwrap(); + + let err = unwrap_err_snapshot(Snapshot::checked_load( + &path, + OciTag::new("latest").unwrap(), + )); + assert_err_contains(err, "digest"); +} + /// Config-blob byte mutation must be caught by digest verification /// before any structural validator runs. #[test] @@ -1824,6 +1864,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 2"); +} + #[test] fn empty_layers_rejected() { let (_dir, path) = save_for_mutation(); @@ -1864,6 +1918,19 @@ fn unknown_snapshot_layer_media_type_rejected() { assert_err_contains(err, "snapshot layer media type"); } +#[test] +fn unknown_transport_layer_media_type_rejected() { + let (_dir, path) = save_for_mutation(); + rewrite_manifest(&path, |m| { + m["layers"][1]["mediaType"] = Value::from("application/vnd.example.unknown.v1"); + }); + let err = unwrap_err_snapshot(Snapshot::checked_load( + &path, + OciTag::new("latest").unwrap(), + )); + assert_err_contains(err, "transport layer media type"); +} + /// Annotations injected by third-party tools (cosign, ORAS, build /// pipelines) must not break load. The OCI envelope around /// `OciSnapshotConfig` is parsed via `oci-spec`'s lenient types. @@ -2286,19 +2353,23 @@ 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!(manifest["layers"].as_array().unwrap().len(), 2); assert_eq!( manifest["layers"][0]["mediaType"].as_str().unwrap(), "application/vnd.hyperlight.snapshot.memory.v1" ); + assert_eq!( + manifest["layers"][1]["mediaType"].as_str().unwrap(), + "application/vnd.hyperlight.snapshot.transport.v1" + ); // `artifactType` mirrors `config.mediaType` so registries that surface // the distribution-spec referrers API report a useful type, and tooling // 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" ); } @@ -2760,7 +2831,7 @@ fn round_trip_preserves_stack_top_gva() { fn round_trip_preserves_non_default_scratch_size() { use crate::sandbox::SandboxConfiguration; let mut cfg = SandboxConfiguration::default(); - let custom_scratch: usize = 256 * 1024; + let custom_scratch = SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 64 * 1024; cfg.set_scratch_size(custom_scratch); let mut sbox = UninitializedSandbox::new(GuestBinary::FilePath(simple_guest_as_pathbuf()), Some(cfg)) @@ -2778,6 +2849,45 @@ fn round_trip_preserves_non_default_scratch_size() { assert_eq!(loaded.layout().get_scratch_size(), custom_scratch); } +#[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 bin = GuestBinary::FilePath(simple_guest_as_pathbuf()); + let mut sbox = UninitializedSandbox::new(bin, 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(); @@ -3151,14 +3261,10 @@ fn from_snapshot_silently_ignores_layout_overrides() { let mut sbox = create_test_sandbox(); let snapshot = sbox.snapshot().unwrap(); - let original_input = snapshot.layout().input_data_size(); - let original_output = snapshot.layout().output_data_size(); let original_heap = snapshot.layout().heap_size(); let original_scratch = snapshot.layout().get_scratch_size(); let mut config = SandboxConfiguration::default(); - config.set_input_data_size(original_input * 2); - config.set_output_data_size(original_output * 2); config.set_heap_size((original_heap as u64) * 2); config.set_scratch_size(original_scratch * 2); @@ -3169,8 +3275,6 @@ fn from_snapshot_silently_ignores_layout_overrides() { sbox2.call::("GetStatic", ()).unwrap(); let new_snap = sbox2.snapshot().unwrap(); - assert_eq!(new_snap.layout().input_data_size(), original_input); - assert_eq!(new_snap.layout().output_data_size(), original_output); assert_eq!(new_snap.layout().heap_size(), original_heap); assert_eq!(new_snap.layout().get_scratch_size(), original_scratch); } diff --git a/src/hyperlight_host/src/sandbox/snapshot/mod.rs b/src/hyperlight_host/src/sandbox/snapshot/mod.rs index e1578e2e7..b2caff00a 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/mod.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/mod.rs @@ -39,6 +39,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}; @@ -123,6 +124,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 { @@ -406,6 +413,7 @@ impl Snapshot { host_functions: HostFunctionDetails { host_functions: None, }, + virtq: None, }) } @@ -432,6 +440,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()); @@ -576,6 +585,10 @@ impl Snapshot { debug_assert!(guest_visible_size.is_multiple_of(PAGE_SIZE)); 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)?, @@ -588,6 +601,7 @@ impl Snapshot { original_entrypoint, snapshot_generation, host_functions, + virtq, }) } @@ -638,6 +652,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. @@ -805,6 +823,7 @@ mod tests { 0, 1, HostFunctionDetails::default(), + None, ) .unwrap(); @@ -825,6 +844,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 ee04373a8..228298b67 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs @@ -25,12 +25,14 @@ limitations under the License. //! When an assertion fires, see `docs/snapshot-versioning.md`. use super::file::{ - MT_CONFIG_CURRENT, MT_SNAPSHOT_CURRENT, OCI_LAYOUT_VERSION, SNAPSHOT_ABI_VERSION, + MT_CONFIG_CURRENT, MT_SNAPSHOT_CURRENT, MT_TRANSPORT_CURRENT, OCI_LAYOUT_VERSION, + SNAPSHOT_ABI_VERSION, }; -const EXPECTED_ABI_VERSION: u32 = 1; -const EXPECTED_MT_CONFIG: &str = "application/vnd.hyperlight.snapshot.config.v1+json"; +const EXPECTED_ABI_VERSION: u32 = 2; +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_MT_TRANSPORT: &str = "application/vnd.hyperlight.snapshot.transport.v1"; const EXPECTED_OCI_LAYOUT_VERSION: &str = "1.0.0"; /// `assert!` with the shared tripwire failure message. The message must @@ -50,6 +52,7 @@ const _: () = { abi_assert!(SNAPSHOT_ABI_VERSION == EXPECTED_ABI_VERSION); abi_assert!(str_eq(MT_CONFIG_CURRENT, EXPECTED_MT_CONFIG)); abi_assert!(str_eq(MT_SNAPSHOT_CURRENT, EXPECTED_MT_SNAPSHOT)); + abi_assert!(str_eq(MT_TRANSPORT_CURRENT, EXPECTED_MT_TRANSPORT)); abi_assert!(str_eq(OCI_LAYOUT_VERSION, EXPECTED_OCI_LAYOUT_VERSION)); }; @@ -66,8 +69,6 @@ const _: () = { const _: () = { use hyperlight_common::outb::OutBAction; - abi_assert!(OutBAction::Log as u16 == 99); - abi_assert!(OutBAction::CallFunction as u16 == 101); abi_assert!(OutBAction::Abort as u16 == 102); abi_assert!(OutBAction::DebugPrint as u16 == 103); #[cfg(feature = "trace_guest")] @@ -76,6 +77,7 @@ const _: () = { abi_assert!(OutBAction::TraceMemoryAlloc as u16 == 105); #[cfg(feature = "mem_profile")] abi_assert!(OutBAction::TraceMemoryFree as u16 == 106); + abi_assert!(OutBAction::VirtqNotify as u16 == 109); }; const _: () = { diff --git a/src/hyperlight_host/src/sandbox/uninitialized.rs b/src/hyperlight_host/src/sandbox/uninitialized.rs index 59f89cbac..a4768951d 100644 --- a/src/hyperlight_host/src/sandbox/uninitialized.rs +++ b/src/hyperlight_host/src/sandbox/uninitialized.rs @@ -471,8 +471,6 @@ mod tests { // Non default memory configuration let cfg = { let mut cfg = SandboxConfiguration::default(); - cfg.set_input_data_size(0x1000); - cfg.set_output_data_size(0x1000); cfg.set_heap_size(0x1000); Some(cfg) }; @@ -1196,7 +1194,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); @@ -1216,37 +1214,11 @@ mod tests { let _evolved: MultiUseSandbox = sandbox.evolve().expect("Failed to evolve sandbox"); } - // Test 4: Create snapshot with custom input/output buffer sizes - { - let mut cfg = SandboxConfiguration::default(); - cfg.set_input_data_size(64 * 1024); // 64KB input - cfg.set_output_data_size(64 * 1024); // 64KB output - - let env = GuestEnvironment::new(GuestBinary::FilePath(binary_path.clone()), None); - - let snapshot = Arc::new( - Snapshot::from_env(env, cfg) - .expect("Failed to create snapshot with custom buffer sizes"), - ); - - let sandbox = UninitializedSandbox::from_snapshot( - snapshot, - None, - #[cfg(crashdump)] - Some(binary_path.clone()), - ) - .expect("Failed to create sandbox from snapshot with custom buffers"); - - let _evolved: MultiUseSandbox = sandbox.evolve().expect("Failed to evolve sandbox"); - } - - // Test 5: Create snapshot with all custom settings + // Test 4: Create snapshot with custom heap and scratch sizes { 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_input_data_size(128 * 1024); // 128KB input - cfg.set_output_data_size(128 * 1024); // 128KB output + cfg.set_scratch_size(512 * 1024); // 512KB scratch let env = GuestEnvironment::new(GuestBinary::FilePath(binary_path.clone()), None); @@ -1283,7 +1255,7 @@ mod tests { let _evolved3: MultiUseSandbox = sandbox3.evolve().expect("Failed to evolve sandbox3"); } - // Test 6: Create snapshot from binary buffer instead of file path + // Test 5: Create snapshot from binary buffer instead of file path { let binary_bytes = fs::read(&binary_path).expect("Failed to read binary file"); @@ -1303,7 +1275,7 @@ mod tests { let _evolved: MultiUseSandbox = sandbox.evolve().expect("Failed to evolve sandbox"); } - // Test 7: Register host functions on sandboxes created from snapshot + // Test 6: Register host functions on sandboxes created from snapshot { let env = GuestEnvironment::new(GuestBinary::FilePath(binary_path.clone()), None); @@ -1342,7 +1314,7 @@ mod tests { assert_eq!(result, ReturnValue::Int(30)); } - // Test 8: Create snapshot with init data (guest blob) + // Test 7: Create snapshot with init data (guest blob) { let init_data = [0xCA, 0xFE, 0xBA, 0xBE]; let guest_env = @@ -1364,7 +1336,7 @@ mod tests { let _evolved: MultiUseSandbox = sandbox.evolve().expect("Failed to evolve sandbox"); } - // Test 9: Create snapshot from existing sandbox + // Test 8: Create snapshot from existing sandbox { let env = GuestEnvironment::new(GuestBinary::FilePath(binary_path.clone()), None); let orig_snapshot = Arc::new( diff --git a/src/hyperlight_host/tests/common/mod.rs b/src/hyperlight_host/tests/common/mod.rs index 003ccdc67..55035e5a0 100644 --- a/src/hyperlight_host/tests/common/mod.rs +++ b/src/hyperlight_host/tests/common/mod.rs @@ -82,6 +82,16 @@ where f(sandbox); } +/// Runs a test with a Rust guest UninitializedSandbox using custom configuration. +pub fn with_rust_uninit_sandbox_cfg(cfg: SandboxConfiguration, f: F) +where + F: FnOnce(UninitializedSandbox), +{ + let sandbox = + UninitializedSandbox::new(GuestBinary::FilePath(rust_guest_path()), Some(cfg)).unwrap(); + f(sandbox); +} + // ============================================================================= // C guest helpers // ============================================================================= diff --git a/src/hyperlight_host/tests/integration_test.rs b/src/hyperlight_host/tests/integration_test.rs index b449ea68d..f0a84f3fd 100644 --- a/src/hyperlight_host/tests/integration_test.rs +++ b/src/hyperlight_host/tests/integration_test.rs @@ -19,6 +19,7 @@ use std::thread; use std::time::Duration; use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; +use hyperlight_common::func::Bytes; use hyperlight_common::log_level::GuestLogFilter; use hyperlight_host::sandbox::SandboxConfiguration; use hyperlight_host::{HyperlightError, MultiUseSandbox}; @@ -30,6 +31,7 @@ pub mod common; // pub to disable dead_code warning use crate::common::{ new_rust_sandbox, new_rust_uninit_sandbox, with_all_sandboxes, with_c_sandbox, with_c_uninit_sandbox, with_rust_sandbox, with_rust_sandbox_cfg, with_rust_uninit_sandbox, + with_rust_uninit_sandbox_cfg, }; // A host function cannot be interrupted, but we can at least make sure after requesting to interrupt a host call, @@ -535,7 +537,7 @@ fn guest_malloc_abort() { }); // allocate a vector (on heap) that is bigger than the heap - let heap_size = 0x6000; + let heap_size = 40 * 1024; let size_to_allocate = 0x10000; assert!( size_to_allocate > heap_size, @@ -578,45 +580,9 @@ fn guest_outb_with_invalid_port_poisons_sandbox() { }); } -#[test] -fn corrupt_output_size_prefix_rejected() { - with_rust_sandbox(|mut sbox| { - let res = sbox.call::("CorruptOutputSizePrefix", ()); - assert!( - res.is_err(), - "Expected error when guest corrupts size prefix, got: {:?}", - res, - ); - let err_msg = format!("{:?}", res.unwrap_err()); - assert!( - err_msg.contains("Corrupt buffer size prefix: flatbuffer claims 4294967295 bytes but the element slot is only 8 bytes"), - "Unexpected error message: {err_msg}" - ); - }); -} - -#[test] -fn corrupt_output_back_pointer_rejected() { - with_rust_sandbox(|mut sbox| { - let res = sbox.call::("CorruptOutputBackPointer", ()); - assert!( - res.is_err(), - "Expected error when guest corrupts back-pointer, got: {:?}", - res, - ); - let err_msg = format!("{:?}", res.unwrap_err()); - assert!( - err_msg.contains( - "Corrupt buffer back-pointer: element offset 57005 is outside valid range [8, 8]" - ), - "Unexpected error message: {err_msg}" - ); - }); -} - #[test] fn guest_panic_no_alloc() { - let heap_size = 0x6000; + let heap_size = 40 * 1024; let mut cfg = SandboxConfiguration::default(); cfg.set_heap_size(heap_size); @@ -744,19 +710,9 @@ fn recursive_stack_allocate_overflow() { #[test] #[ignore] fn log_message() { - // The magic numbers below represent the number of fixed log messages that are emitted as - // follows: - // - logs from trace level tracing spans created as logs because of the tracing `log` feature - // - 4 from evolve call (generic_init + hyperlight_main) - // - 8 from guest call - // and are multiplied because we make 6 calls to `log_test_messages` - // NOTE: These numbers need to be updated if log messages or spans are added/removed - let num_fixed_trace_log = 12 * 6; - - // Calculate fixed info logs - // - 4 logs per iteration from infrastructure at Info level (internal_dispatch_function) - // (dispatch x 1 + call_guest x 1) * 2 logs (Enter/Exit) = 4 logs - // - 6 iterations + // Each of the six sandboxes emits eight fixed records at trace level. + // Dispatch and call spans emit four fixed records at info level. + let num_fixed_trace_log = 8 * 6; let num_fixed_info_log = 4 * 6; let tests = vec![ @@ -833,6 +789,31 @@ fn log_test_messages(levelfilter: Option) { } } +#[test] +#[ignore] +fn virtq_repeated_log_delivery_small_ring() { + SimpleLogger::initialize_test_logger(); + LOGGER.clear_log_calls(); + + let mut cfg = SandboxConfiguration::default(); + cfg.set_g2h_queue_size(4); + cfg.set_g2h_pool_pages(2); + + with_rust_uninit_sandbox_cfg(cfg, |mut sandbox| { + sandbox.set_max_guest_log_level(LevelFilter::INFO); + let mut sandbox = sandbox.evolve().unwrap(); + + sandbox.call::<()>("LogMessageN", 20_i32).unwrap(); + + let count = (0..LOGGER.num_log_calls()) + .filter_map(|index| LOGGER.get_log_call(index)) + .filter(|call| call.target == "hyperlight_guest" && call.args.contains("log entry")) + .count(); + assert_eq!(count, 20); + LOGGER.clear_log_calls(); + }); +} + /// Tests whether host is able to return Bool as return type /// or not #[test] @@ -912,6 +893,40 @@ fn test_if_guest_is_able_to_get_string_return_values_from_host() { }); } +#[test] +fn c_guest_accesses_byte_chunks() { + with_c_uninit_sandbox(|mut sandbox| { + sandbox + .register("HostEchoByteChunks", |value: Vec| value) + .unwrap(); + + let mut sandbox = sandbox.evolve().unwrap(); + let expected = (0..10 * 1024) + .map(|index| (index % 251) as u8) + .collect::>(); + + let input = vec![ + Bytes::copy_from_slice(&expected[..2047]), + Bytes::copy_from_slice(&expected[2047..4097]), + Bytes::copy_from_slice(&expected[4097..]), + ]; + + for _ in 0..2 { + let output: Vec = sandbox + .call("RoundTripHostByteChunks", input.clone()) + .unwrap(); + + assert_eq!( + output + .iter() + .flat_map(|chunk| chunk.iter().copied()) + .collect::>(), + expected + ); + } + }); +} + /// Test that validates interrupt behavior with random kill timing under concurrent load /// Uses a pool of 100 sandboxes, 100 threads, and 500 iterations per thread. /// Randomly decides to kill some calls at random times during execution. diff --git a/src/hyperlight_host/tests/sandbox_host_tests.rs b/src/hyperlight_host/tests/sandbox_host_tests.rs index b1a1a9918..da2ccf699 100644 --- a/src/hyperlight_host/tests/sandbox_host_tests.rs +++ b/src/hyperlight_host/tests/sandbox_host_tests.rs @@ -17,6 +17,7 @@ use core::f64; use std::sync::mpsc::channel; use std::sync::{Arc, Mutex}; +use hyperlight_common::func::Bytes; use hyperlight_host::sandbox::SandboxConfiguration; use hyperlight_host::{ GuestBinary, HyperlightError, MultiUseSandbox, Result, UninitializedSandbox, new_error, @@ -26,7 +27,7 @@ use hyperlight_testing::simple_guest_as_pathbuf; pub mod common; // pub to disable dead_code warning use crate::common::{ with_all_sandboxes, with_all_sandboxes_cfg, with_all_sandboxes_with_writer, - with_all_uninit_sandboxes, + with_all_uninit_sandboxes, with_rust_uninit_sandbox, with_rust_uninit_sandbox_cfg, }; #[test] @@ -212,9 +213,7 @@ fn incorrect_parameter_num() { #[test] fn small_scratch_sandbox() { let mut cfg = SandboxConfiguration::default(); - cfg.set_scratch_size(0x48000); - cfg.set_input_data_size(0x24000); - cfg.set_output_data_size(0x24000); + cfg.set_scratch_size(0x1000); let a = UninitializedSandbox::new(GuestBinary::FilePath(simple_guest_as_pathbuf()), Some(cfg)); assert!(matches!( @@ -224,7 +223,7 @@ fn small_scratch_sandbox() { } #[test] -fn iostack_is_working() { +fn custom_guest_dispatch_is_working() { with_all_sandboxes(|mut sandbox| { let res: i32 = sandbox .call::("ThisIsNotARealFunctionButTheNameIsImportant", ()) @@ -325,6 +324,167 @@ fn callback_test() { callback_test_helper(); } +#[test] +fn host_external_bytes_round_trip() { + with_rust_uninit_sandbox(|mut sandbox| { + sandbox + .register("HostEchoVecBytes", |value: Vec| value) + .unwrap(); + sandbox + .register("HostEchoByteChunks", |value: Vec| value) + .unwrap(); + sandbox.register("HostNoOp", || {}).unwrap(); + let mut sandbox = sandbox.evolve().unwrap(); + let expected: Vec = (0..6 * 1024).map(|index| (index % 251) as u8).collect(); + + let contiguous: Vec = sandbox + .call("RoundTripHostVecBytes", expected.clone()) + .unwrap(); + assert_eq!(contiguous, expected); + + let input = vec![ + Bytes::copy_from_slice(&expected[..2047]), + Bytes::copy_from_slice(&expected[2047..4097]), + Bytes::copy_from_slice(&expected[4097..]), + ]; + for _ in 0..2 { + let chunks: Vec = sandbox + .call("RoundTripHostByteChunks", input.clone()) + .unwrap(); + let flattened: Vec = chunks + .iter() + .flat_map(|chunk| chunk.iter().copied()) + .collect(); + assert_eq!(flattened, expected); + } + }); +} + +#[test] +fn guest_external_bytes_round_trip_and_retention() { + with_rust_uninit_sandbox(|sandbox| { + let mut sandbox = sandbox.evolve().unwrap(); + let expected: Vec = (0..9 * 1024).map(|index| (index % 251) as u8).collect(); + + let contiguous: Vec = sandbox.call("EchoGuestVecBytes", expected.clone()).unwrap(); + assert_eq!(contiguous, expected); + + let input = vec![ + Bytes::copy_from_slice(&expected[..2047]), + Bytes::copy_from_slice(&expected[2047..4097]), + Bytes::copy_from_slice(&expected[4097..]), + ]; + let chunks: Vec = sandbox.call("EchoGuestByteChunks", input.clone()).unwrap(); + assert_eq!( + chunks + .iter() + .flat_map(|chunk| chunk.iter().copied()) + .collect::>(), + expected + ); + + let retained: Vec = (0..12_000).map(|index| (index % 251) as u8).collect(); + let retained_len: i32 = sandbox + .call( + "RetainGuestByteChunks", + vec![Bytes::copy_from_slice(&retained)], + ) + .unwrap(); + assert_eq!(retained_len as usize, retained.len()); + + let released_len: i32 = sandbox.call("ReleaseGuestByteChunks", ()).unwrap(); + assert_eq!(released_len as usize, retained.len()); + + let retried: Vec = sandbox.call("EchoGuestVecBytes", expected.clone()).unwrap(); + assert_eq!(retried, expected); + }); +} + +#[test] +fn h2g_capacity_failure_does_not_poison_sandbox() { + let mut cfg = SandboxConfiguration::default(); + cfg.set_h2g_pool_pages(4); + + with_rust_uninit_sandbox_cfg(cfg, |sandbox| { + let mut sandbox = sandbox.evolve().unwrap(); + let retained = vec![0u8; 12_000]; + + sandbox + .call::( + "RetainGuestByteChunks", + vec![Bytes::copy_from_slice(&retained)], + ) + .unwrap(); + + let error = sandbox + .call::>("EchoGuestVecBytes", vec![0u8; 9 * 1024]) + .unwrap_err(); + + assert!(error.to_string().contains("H2G capacity")); + + let released: i32 = sandbox.call("ReleaseGuestByteChunks", ()).unwrap(); + assert_eq!(released as usize, retained.len()); + }); +} + +#[test] +fn oversized_host_response_returns_transport_error() { + with_rust_uninit_sandbox(|mut sandbox| { + sandbox + .register("HostOversizedVecBytes", || vec![0u8; 64 * 1024]) + .unwrap(); + sandbox.register("HostNoOp", || {}).unwrap(); + let mut sandbox = sandbox.evolve().unwrap(); + + let error = sandbox + .call::>("GetOversizedHostVecBytes", ()) + .unwrap_err(); + assert!(matches!( + error, + HyperlightError::GuestError(_, message) + if message == "Host response exceeds virtqueue capacity" + )); + sandbox.call::<()>("RoundTripHostNoOp", ()).unwrap(); + }); +} + +#[test] +fn log_then_host_call_with_small_rings() { + let mut cfg = SandboxConfiguration::default(); + cfg.set_g2h_queue_size(4); + cfg.set_h2g_queue_size(4); + cfg.set_g2h_pool_pages(2); + + with_rust_uninit_sandbox_cfg(cfg, |mut sandbox| { + sandbox.set_max_guest_log_level(tracing_core::LevelFilter::INFO); + sandbox.register("HostNoOp", || {}).unwrap(); + let mut sandbox = sandbox.evolve().unwrap(); + + for _ in 0..20 { + sandbox.call::<()>("LogThenHostNoOp", ()).unwrap(); + } + }); +} + +#[test] +fn oversized_fixed_host_error_returns_transport_error() { + with_rust_uninit_sandbox(|mut sandbox| { + sandbox + .register("HostNoOp", || -> Result<()> { + Err(new_error!("host error {}", "x".repeat(1024))) + }) + .unwrap(); + let mut sandbox = sandbox.evolve().unwrap(); + + let error = sandbox.call::<()>("RoundTripHostNoOp", ()).unwrap_err(); + assert!(matches!( + error, + HyperlightError::GuestError(_, message) + if message == "Host response exceeds virtqueue capacity" + )); + }); +} + #[test] fn callback_test_parallel() { let handles: Vec<_> = (0..100) diff --git a/src/hyperlight_host/tests/snapshot_goldens/fixtures.rs b/src/hyperlight_host/tests/snapshot_goldens/fixtures.rs index 0b35fb325..7bb78b016 100644 --- a/src/hyperlight_host/tests/snapshot_goldens/fixtures.rs +++ b/src/hyperlight_host/tests/snapshot_goldens/fixtures.rs @@ -41,10 +41,10 @@ pub(crate) const CALL_COUNTER_BUMP: i32 = 42; /// least one region between generate-time and load-time. fn golden_config() -> SandboxConfiguration { let mut cfg = SandboxConfiguration::default(); - cfg.set_input_data_size(64 * 1024); - cfg.set_output_data_size(64 * 1024); cfg.set_heap_size(256 * 1024); cfg.set_scratch_size(512 * 1024); + cfg.set_g2h_pool_pages(16); + cfg.set_h2g_pool_pages(16); cfg } diff --git a/src/hyperlight_host/tests/snapshot_goldens/goldens_version.rs b/src/hyperlight_host/tests/snapshot_goldens/goldens_version.rs index 32572f417..c0b6c0c24 100644 --- a/src/hyperlight_host/tests/snapshot_goldens/goldens_version.rs +++ b/src/hyperlight_host/tests/snapshot_goldens/goldens_version.rs @@ -21,7 +21,7 @@ limitations under the License. //! publish. See `docs/snapshot-versioning.md`. /// Goldens version, a `vMAJOR.MINOR` string. -pub(crate) const GOLDENS_VERSION: &str = "v1.0"; +pub(crate) const GOLDENS_VERSION: &str = "v2.0"; /// Old majors kept loadable through a compatibility path, verified /// alongside `GOLDENS_VERSION`. A backwards-compatible break (Option 2) diff --git a/src/schema/function_types.fbs b/src/schema/function_types.fbs index d5c209ff6..091ff5918 100644 --- a/src/schema/function_types.fbs +++ b/src/schema/function_types.fbs @@ -57,6 +57,21 @@ table hlvecbytes { value:[ubyte]; } +// hlbytechunks is the embedded compatibility representation of a chunked byte +// value. Embedded transport does not preserve chunk boundaries. + +table hlbytechunks { + value:[ubyte]; +} + +// hlexternalbytes declares a logical byte value stored outside the FlatBuffer. +// chunked distinguishes ByteChunks from the default VecBytes logical type. + +table hlexternalbytes { + length:ulong; + chunked:bool; +} + // hlsizeprefixedbuffer is a vector of bytes prefixed with a 32 bit integer table hlsizeprefixedbuffer { @@ -64,6 +79,14 @@ table hlsizeprefixedbuffer { value:[ubyte]; } +// hlsizeprefixedbytechunks is the embedded compatibility representation of a +// chunked return byte value. Embedded transport does not preserve boundaries. + +table hlsizeprefixedbytechunks { + size:int; + value:[ubyte]; +} + // hlvoid is a void (used for functions that return nothing) table hlvoid { @@ -81,6 +104,8 @@ union ParameterValue { hlstring, hlbool, hlvecbytes, + hlexternalbytes, + hlbytechunks, } // This represents a parameter type in a function definition @@ -95,6 +120,7 @@ enum ParameterType : ubyte { hlstring, hlbool, hlvecbytes, + hlbytechunks, } enum ReturnType : ubyte { @@ -108,6 +134,7 @@ enum ReturnType : ubyte { hlbool, hlvoid, hlsizeprefixedbuffer, + hlbytechunks, } union ReturnValue { @@ -121,4 +148,6 @@ union ReturnValue { hlbool, hlvoid, hlsizeprefixedbuffer, + hlexternalbytes, + hlsizeprefixedbytechunks, } diff --git a/src/tests/c_guests/c_simpleguest/main.c b/src/tests/c_guests/c_simpleguest/main.c index 91da1973c..eb0d3f9d9 100644 --- a/src/tests/c_guests/c_simpleguest/main.c +++ b/src/tests/c_guests/c_simpleguest/main.c @@ -25,13 +25,13 @@ float echo_float(float f) { return f; } double echo_double(double d) { return d; } -hl_Vec *set_byte_array_to_zero(const hl_FunctionCall* params) { +hl_ReturnValue *set_byte_array_to_zero(const hl_FunctionCall* params) { hl_Vec input = params->parameters[0].value.VecBytes; uint8_t *x = malloc(input.len); for (uintptr_t i = 0; i < input.len; i++) { x[i] = 0; } - return hl_flatbuffer_result_from_Bytes(x, input.len); + return hl_result_from_Bytes(x, input.len); } int print_output(const char *message) { @@ -213,9 +213,9 @@ int set_static(void) { return length; } -hl_Vec *get_size_prefixed_buffer(const hl_FunctionCall* params) { +hl_ReturnValue *get_size_prefixed_buffer(const hl_FunctionCall* params) { hl_Vec input = params->parameters[0].value.VecBytes; - return hl_flatbuffer_result_from_Bytes(input.data, input.len); + return hl_result_from_Bytes(input.data, input.len); } int guest_abort_with_code(int32_t code) { @@ -239,10 +239,10 @@ int log_message(const char *message, int64_t level) { return -1; } -hl_Vec *twenty_four_k_in_eight_k_out(const hl_FunctionCall* params) { +hl_ReturnValue *twenty_four_k_in_eight_k_out(const hl_FunctionCall* params) { hl_Vec input = params->parameters[0].value.VecBytes; assert(input.len == 24 * 1024); - return hl_flatbuffer_result_from_Bytes(input.data, 8 * 1024); + return hl_result_from_Bytes(input.data, 8 * 1024); } int guest_function(const char *from_host) { @@ -332,6 +332,36 @@ const char* guest_fn_checks_if_host_returns_string_value() { return hl_get_host_return_value_as_String(); } +hl_ReturnValue *round_trip_host_byte_chunks(const hl_FunctionCall *params) { + hl_ByteChunks input = params->parameters[0].value.ByteChunks; + assert(input.count > 1); + + for (uintptr_t i = 0; i < input.count; i++) { + assert(input.chunks[i].data != NULL || input.chunks[i].len == 0); + } + + hl_Parameter host_param = { + .tag = hl_ParameterType_ByteChunks, + .value = {.ByteChunks = input}, + }; + + const hl_FunctionCall host_call = { + .function_name = "HostEchoByteChunks", + .parameters = &host_param, + .parameters_len = 1, + .return_type = hl_ReturnType_ByteChunks, + }; + hl_call_host_function(&host_call); + + hl_ByteChunks *output = hl_get_host_return_value_as_ByteChunks(); + assert(output != NULL); + assert(output->count > 1); + + hl_ReturnValue *result = hl_result_from_ByteChunks(*output); + hl_free_byte_chunks(output); + return result; +} + HYPERLIGHT_WRAP_FUNCTION(guest_fn_checks_if_host_returns_float_value, Float, 2, Float, Float) HYPERLIGHT_WRAP_FUNCTION(guest_fn_checks_if_host_returns_double_value, Double, 2, Double, Double) HYPERLIGHT_WRAP_FUNCTION(guest_fn_checks_if_host_returns_string_value, String, 0) @@ -409,16 +439,17 @@ void hyperlight_main(void) // HYPERLIGHT_REGISTER_FUNCTION macro does not work for functions that return VecBytes, // so we use hl_register_function_definition directly hl_register_function_definition("24K_in_8K_out", twenty_four_k_in_eight_k_out, 1, (hl_ParameterType[]){hl_ParameterType_VecBytes}, hl_ReturnType_VecBytes); + hl_register_function_definition("RoundTripHostByteChunks", round_trip_host_byte_chunks, 1, (hl_ParameterType[]){hl_ParameterType_ByteChunks}, hl_ReturnType_ByteChunks); } // This dispatch function is only used when the host dispatches a guest function // call but there is no registered guest function with the given name. -hl_Vec *c_guest_dispatch_function(const hl_FunctionCall *function_call) { +hl_ReturnValue *c_guest_dispatch_function(const hl_FunctionCall *function_call) { const char *func_name = function_call->function_name; if (strcmp(func_name, "ThisIsNotARealFunctionButTheNameIsImportant") == 0) { // TODO DO A LOG HERE - // This is special case for test `iostack_is_working - return hl_flatbuffer_result_from_Int(99); + // This is a special case for test `custom_guest_dispatch_is_working`. + return hl_result_from_Int(99); } return NULL; diff --git a/src/tests/rust_guests/simpleguest/src/main.rs b/src/tests/rust_guests/simpleguest/src/main.rs index fc0ce416b..bb191a5bc 100644 --- a/src/tests/rust_guests/simpleguest/src/main.rs +++ b/src/tests/rust_guests/simpleguest/src/main.rs @@ -38,11 +38,10 @@ use core::sync::atomic::{AtomicU64, Ordering}; use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType}; use hyperlight_common::flatbuffer_wrappers::function_types::{ - ParameterType, ParameterValue, ReturnType, ReturnValue, + Bytes, ParameterType, ParameterValue, ReturnType, ReturnValue, }; use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; use hyperlight_common::flatbuffer_wrappers::guest_log_level::LogLevel; -use hyperlight_common::flatbuffer_wrappers::util::get_flatbuffer_result; use hyperlight_common::log_level::GuestLogFilter; use hyperlight_common::vmem::{BasicMapping, MappingKind}; use hyperlight_guest::error::{HyperlightGuestError, Result}; @@ -52,11 +51,10 @@ use hyperlight_guest_bin::exception::arch::{Context, ExceptionInfo}; use hyperlight_guest_bin::guest_function::definition::{GuestFunc, GuestFunctionDefinition}; use hyperlight_guest_bin::guest_function::register::register_function; use hyperlight_guest_bin::host_comm::{ - call_host_function, call_host_function_without_returning_result, get_host_return_value_raw, - print_output_with_host_print, read_n_bytes_from_user_memory, + call_host_function, print_output_with_host_print, read_n_bytes_from_user_memory, }; use hyperlight_guest_bin::memory::malloc; -use hyperlight_guest_bin::{GUEST_HANDLE, guest_function, guest_logger, host_function}; +use hyperlight_guest_bin::{guest_function, guest_logger, host_function}; // `log` is intentionally kept here: the LogMessage guest function exercises the // guest-side `log` crate path to verify that guests using `log` are still supported. use log::LevelFilter; @@ -403,6 +401,58 @@ fn get_size_prefixed_buffer(data: Vec) -> Vec { data } +#[guest_function("EchoGuestVecBytes")] +fn echo_guest_vec_bytes(data: Vec) -> Vec { + data +} + +#[guest_function("EchoGuestByteChunks")] +fn echo_guest_byte_chunks(data: Vec) -> Vec { + data +} + +static mut RETAINED_GUEST_CHUNKS: Option> = None; +static mut RETAINED_HOST_CHUNKS: Option> = None; + +#[guest_function("RetainGuestByteChunks")] +fn retain_guest_byte_chunks(data: Vec) -> i32 { + let len = data.iter().map(Bytes::len).sum::(); + // SAFETY: the guest is single threaded, so the static has no concurrent access. + unsafe { RETAINED_GUEST_CHUNKS = Some(data) }; + len as i32 +} + +#[guest_function("ReleaseGuestByteChunks")] +fn release_guest_byte_chunks() -> i32 { + // SAFETY: the guest is single threaded, so the static has no concurrent access. + #[allow(static_mut_refs)] + unsafe { + RETAINED_GUEST_CHUNKS.take().map_or(0, |chunks| { + chunks.iter().map(Bytes::len).sum::() as i32 + }) + } +} + +#[guest_function("RetainHostByteChunks")] +fn retain_host_byte_chunks(data: Vec) -> Result { + let chunks = host_echo_byte_chunks(data)?; + let len = chunks.iter().map(Bytes::len).sum::(); + // SAFETY: the guest is single threaded, so the static has no concurrent access. + unsafe { RETAINED_HOST_CHUNKS = Some(chunks) }; + Ok(len as i32) +} + +#[guest_function("ReleaseHostByteChunks")] +fn release_host_byte_chunks() -> i32 { + // SAFETY: the guest is single threaded, so the static has no concurrent access. + #[allow(static_mut_refs)] + unsafe { + RETAINED_HOST_CHUNKS.take().map_or(0, |chunks| { + chunks.iter().map(Bytes::len).sum::() as i32 + }) + } +} + #[guest_function("EchoI32")] fn echo_i32(v: i32) -> i32 { v @@ -458,6 +508,12 @@ fn host_echo_string(v: String) -> Result; #[host_function("HostEchoVecBytes")] fn host_echo_vec_bytes(v: Vec) -> Result>; +#[host_function("HostEchoByteChunks")] +fn host_echo_byte_chunks(v: Vec) -> Result>; + +#[host_function("HostOversizedVecBytes")] +fn host_oversized_vec_bytes() -> Result>; + #[host_function("HostNoOp")] fn host_noop() -> Result<()>; @@ -506,11 +562,29 @@ fn round_trip_host_vec_bytes(v: Vec) -> Result> { host_echo_vec_bytes(v) } +#[guest_function("RoundTripHostByteChunks")] +fn round_trip_host_byte_chunks(v: Vec) -> Result> { + let chunks = host_echo_byte_chunks(v)?; + host_noop()?; + Ok(chunks) +} + +#[guest_function("GetOversizedHostVecBytes")] +fn get_oversized_host_vec_bytes() -> Result> { + host_oversized_vec_bytes() +} + #[guest_function("RoundTripHostNoOp")] fn round_trip_host_noop() -> Result<()> { host_noop() } +#[guest_function("LogThenHostNoOp")] +fn log_then_host_noop() -> Result<()> { + log::info!("log before host call"); + host_noop() +} + static mut HEAP_PATTERN: Option> = None; #[guest_function("AllocAndWritePattern")] @@ -629,6 +703,13 @@ fn log_message(message: String, level: i32) { } } +#[guest_function("LogMessageN")] +fn log_message_n(count: i32) { + for i in 0..count { + log::info!("log entry {}", i); + } +} + #[guest_function("TriggerException")] fn trigger_exception() { // trigger an undefined instruction exception @@ -1471,41 +1552,8 @@ fn fuzz_guest_trace(max_depth: u32, msg: String) -> u32 { fuzz_traced_function(0, max_depth, &msg) } -#[guest_function("CorruptOutputSizePrefix")] -fn corrupt_output_size_prefix() -> i32 { - unsafe { - let peb_ptr = core::ptr::addr_of!(GUEST_HANDLE).read().peb().unwrap(); - let output_stack_ptr = (*peb_ptr).output_stack.ptr as *mut u8; - - // Write a fake stack entry with a ~4 GB size prefix (0xFFFF_FFFB + 4). - let buf = core::slice::from_raw_parts_mut(output_stack_ptr, 24); - buf[0..8].copy_from_slice(&24_u64.to_le_bytes()); - buf[8..12].copy_from_slice(&0xFFFF_FFFBu32.to_le_bytes()); - buf[12..16].copy_from_slice(&[0u8; 4]); - buf[16..24].copy_from_slice(&8_u64.to_le_bytes()); - outb_with_port(hyperlight_common::outb::VmAction::Halt as u32, 0u32); - unreachable!(); - } -} - -#[guest_function("CorruptOutputBackPointer")] -fn corrupt_output_back_pointer() -> i32 { - unsafe { - let peb_ptr = core::ptr::addr_of!(GUEST_HANDLE).read().peb().unwrap(); - let output_stack_ptr = (*peb_ptr).output_stack.ptr as *mut u8; - - // Write a fake stack entry with back-pointer 0xDEAD (past stack pointer 24). - let buf = core::slice::from_raw_parts_mut(output_stack_ptr, 24); - buf[0..8].copy_from_slice(&24_u64.to_le_bytes()); - buf[8..16].copy_from_slice(&[0u8; 8]); - buf[16..24].copy_from_slice(&0xDEAD_u64.to_le_bytes()); - outb_with_port(hyperlight_common::outb::VmAction::Halt as u32, 0u32); - unreachable!(); - } -} - // Interprets the given guest function call as a host function call and dispatches it to the host. -fn fuzz_host_function(func: FunctionCall) -> Result> { +fn fuzz_host_function(func: FunctionCall) -> Result { let mut params = func.parameters.unwrap(); // first parameter must be string (the name of the host function to call) let host_func_name = match params.remove(0) { @@ -1520,43 +1568,12 @@ fn fuzz_host_function(func: FunctionCall) -> Result> { } }; - // Because we do not know at compile time the actual return type of the host function to be called - // we cannot use the `call_host_function` generic function. - // We need to use the `call_host_function_without_returning_result` function that does not retrieve the return - // value - call_host_function_without_returning_result( - &host_func_name, - Some(params), - func.expected_return_type, - ) - .expect("failed to call host function"); - - let host_return = get_host_return_value_raw(); - match host_return { - Ok(return_value) => match return_value { - ReturnValue::Int(i) => Ok(get_flatbuffer_result(i)), - ReturnValue::UInt(i) => Ok(get_flatbuffer_result(i)), - ReturnValue::Long(i) => Ok(get_flatbuffer_result(i)), - ReturnValue::ULong(i) => Ok(get_flatbuffer_result(i)), - ReturnValue::Float(i) => Ok(get_flatbuffer_result(i)), - ReturnValue::Double(i) => Ok(get_flatbuffer_result(i)), - ReturnValue::String(str) => Ok(get_flatbuffer_result(str.as_str())), - ReturnValue::Bool(bool) => Ok(get_flatbuffer_result(bool)), - ReturnValue::Void(()) => Ok(get_flatbuffer_result(())), - ReturnValue::VecBytes(byte) => Ok(get_flatbuffer_result(byte.as_slice())), - }, - Err(e) => Err(e), - } + call_host_function::(&host_func_name, Some(params), func.expected_return_type) } #[hyperlight_guest_bin::dispatch] #[instrument(skip_all, parent = Span::current(), level= "Trace")] -fn dispatch(function_call: FunctionCall) -> Result> { - // This test checks the stack behavior of the input/output buffer - // by calling the host before serializing the function call. - // If the stack is not working correctly, the input or output buffer will be - // overwritten before the function call is serialized, and we will not be able - // to verify that the function call name is "ThisIsNotARealFunctionButTheNameIsImportant" +fn dispatch(function_call: FunctionCall) -> Result { if function_call.function_name == "FuzzHostFunc" { return fuzz_host_function(function_call); } @@ -1592,5 +1609,5 @@ fn dispatch(function_call: FunctionCall) -> Result> { )); } - Ok(get_flatbuffer_result(99)) + Ok(ReturnValue::Int(99)) } diff --git a/src/tests/rust_guests/witguest/src/main.rs b/src/tests/rust_guests/witguest/src/main.rs index 307d9cca6..d9ec18e18 100644 --- a/src/tests/rust_guests/witguest/src/main.rs +++ b/src/tests/rust_guests/witguest/src/main.rs @@ -234,10 +234,11 @@ pub extern "C" fn hyperlight_main() { use ::alloc::vec::Vec; use ::hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall; +use ::hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue; use ::hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; use ::hyperlight_guest::error::{HyperlightGuestError, Result}; #[no_mangle] -pub fn guest_dispatch_function(function_call: FunctionCall) -> Result> { +pub fn guest_dispatch_function(function_call: FunctionCall) -> Result { Err(HyperlightGuestError::new( ErrorCode::GuestFunctionNotFound, function_call.function_name.clone(),