From 34467341defe3285c34012d2dd27411dae8c7461 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Mon, 20 Jul 2026 19:06:08 +0200 Subject: [PATCH 01/34] feat(virtq): add external byte transport foundations Add distinct VecBytes and ByteChunks function value types while preserving embedded stack encoding. Introduce the external-byte marker for future virtq codecs and simplify the virtq message header. Signed-off-by: Tomasz Andrzejak --- .../src/flatbuffer_wrappers/function_call.rs | 51 +++++- .../src/flatbuffer_wrappers/function_types.rs | 171 +++++++++++++++++- .../src/flatbuffer_wrappers/util.rs | 80 +++++++- .../generated/hlbytechunks_generated.rs | 124 +++++++++++++ .../generated/hlexternalbytes_generated.rs | 140 ++++++++++++++ .../hlsizeprefixedbytechunks_generated.rs | 150 +++++++++++++++ .../generated/parameter_generated.rs | 58 ++++++ .../generated/parameter_type_generated.rs | 10 +- .../generated/parameter_value_generated.rs | 14 +- .../generated/return_type_generated.rs | 10 +- .../generated/return_value_box_generated.rs | 127 ++++++------- .../generated/return_value_generated.rs | 14 +- src/hyperlight_common/src/flatbuffers/mod.rs | 6 + src/hyperlight_common/src/func/mod.rs | 2 + src/hyperlight_common/src/func/param_type.rs | 22 ++- src/hyperlight_common/src/func/ret_type.rs | 25 +++ src/hyperlight_common/src/virtq/msg.rs | 96 +++++----- src/hyperlight_component_util/src/hl.rs | 2 +- .../src/guest_function/definition.rs | 1 + src/hyperlight_guest_capi/include/macro.h | 2 +- src/hyperlight_guest_capi/src/flatbuffer.rs | 21 ++- .../src/types/parameter.rs | 15 ++ src/hyperlight_host/src/func/mod.rs | 2 + .../src/sandbox/snapshot/file/config.rs | 8 + src/schema/function_types.fbs | 29 +++ src/tests/rust_guests/simpleguest/src/main.rs | 1 + 26 files changed, 1054 insertions(+), 127 deletions(-) create mode 100644 src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlbytechunks_generated.rs create mode 100644 src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlexternalbytes_generated.rs create mode 100644 src/hyperlight_common/src/flatbuffers/hyperlight/generated/hlsizeprefixedbytechunks_generated.rs diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs b/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs index 056ced8e0..f31e96b8d 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs @@ -23,12 +23,13 @@ use flatbuffers::{FlatBufferBuilder, WIPOffset, size_prefixed_root}; use tracing::{Span, instrument}; use super::function_types::{ParameterValue, ReturnType}; +use super::util::byte_chunks_to_bytes; 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, hlbytechunks, hlbytechunksArgs, + hldouble, hldoubleArgs, hlfloat, hlfloatArgs, hlint, hlintArgs, hllong, hllongArgs, hlstring, + hlstringArgs, hluint, hluintArgs, hlulong, hlulongArgs, hlvecbytes, hlvecbytesArgs, }; /// The type of function call. @@ -193,6 +194,23 @@ impl FunctionCall { }, ) } + ParameterValue::ByteChunks(v) => { + let value = byte_chunks_to_bytes(v); + let vec_bytes = builder.create_vector(value.as_ref()); + let hlbytechunks = hlbytechunks::create( + builder, + &hlbytechunksArgs { + value: Some(vec_bytes), + }, + ); + Parameter::create( + builder, + &ParameterArgs { + value_type: FbParameterValue::hlbytechunks, + value: Some(hlbytechunks.as_union_value()), + }, + ) + } }) .collect(); Some(builder.create_vector(¶meter_offsets)) @@ -277,7 +295,7 @@ mod tests { use alloc::vec; use super::*; - use crate::flatbuffer_wrappers::function_types::ReturnType; + use crate::flatbuffer_wrappers::function_types::{Bytes, ReturnType}; #[test] fn read_from_flatbuffer() -> Result<()> { @@ -327,4 +345,29 @@ mod tests { Ok(()) } + + #[test] + fn embedded_byte_parameters_round_trip_as_distinct_logical_types() { + let mut builder = FlatBufferBuilder::new(); + let parameters = vec![ + ParameterValue::VecBytes(vec![1, 2, 3]), + ParameterValue::ByteChunks(vec![Bytes::from_static(&[4, 5]), Bytes::from_static(&[6])]), + ]; + let encoded = FunctionCall::new( + "bytes".to_string(), + Some(parameters), + FunctionCallType::Host, + ReturnType::VecBytes, + ) + .encode(&mut builder); + + let decoded = FunctionCall::try_from(encoded).unwrap(); + assert_eq!( + decoded.parameters, + Some(vec![ + ParameterValue::VecBytes(vec![1, 2, 3]), + ParameterValue::ByteChunks(vec![Bytes::from_static(&[4, 5, 6])]), + ]) + ); + } } diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs b/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs index 42c7ff823..3db72780a 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs @@ -18,19 +18,23 @@ 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::guest_error::GuestError; +#[cfg(feature = "fuzzing")] +use super::util::arbitrary_byte_chunks; +use super::util::{byte_chunks_from_bytes, byte_chunks_to_bytes}; 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, + hlsizeprefixedbufferArgs, hlsizeprefixedbytechunks, hlsizeprefixedbytechunksArgs, hlstring, + hlstringArgs, hluint, hluintArgs, hlulong, hlulongArgs, hlvoid, hlvoidArgs, }; pub struct FunctionCallResult(core::result::Result); @@ -95,6 +99,21 @@ impl FunctionCallResult { FbReturnValue::hlsizeprefixedbuffer, ) } + ReturnValue::ByteChunks(v) => { + let value = byte_chunks_to_bytes(v); + let val = builder.create_vector(value.as_ref()); + let off = hlsizeprefixedbytechunks::create( + builder, + &hlsizeprefixedbytechunksArgs { + value: Some(val), + size: value.len() as i32, + }, + ); + ( + Some(off.as_union_value()), + FbReturnValue::hlsizeprefixedbytechunks, + ) + } ReturnValue::Void(()) => { let off = hlvoid::create(builder, &hlvoidArgs {}); (Some(off.as_union_value()), FbReturnValue::hlvoid) @@ -204,6 +223,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 +254,8 @@ pub enum ParameterType { Bool, /// `Vec` VecBytes, + /// One complete chunk-preserving byte value. + ByteChunks, } /// Supported return types with values from function calling. @@ -253,6 +281,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 +314,8 @@ pub enum ReturnType { Void, /// `Vec` VecBytes, + /// One complete chunk-preserving byte value. + ByteChunks, } impl From<&ParameterValue> for ParameterType { @@ -296,6 +331,7 @@ impl From<&ParameterValue> for ParameterType { ParameterValue::String(_) => ParameterType::String, ParameterValue::Bool(_) => ParameterType::Bool, ParameterValue::VecBytes(_) => ParameterType::VecBytes, + ParameterValue::ByteChunks(_) => ParameterType::ByteChunks, } } } @@ -334,6 +370,14 @@ impl TryFrom> for ParameterValue { FbParameterValue::hlvecbytes => param.value_as_hlvecbytes().map(|hlvecbytes| { ParameterValue::VecBytes(hlvecbytes.value().unwrap_or_default().bytes().to_vec()) }), + FbParameterValue::hlbytechunks => param.value_as_hlbytechunks().map(|hlbytechunks| { + ParameterValue::ByteChunks(byte_chunks_from_bytes(Bytes::copy_from_slice( + hlbytechunks.value().unwrap_or_default().bytes(), + ))) + }), + FbParameterValue::hlexternalbytes => { + bail!("External byte parameter requires an external value source") + } other => { bail!("Unexpected flatbuffer parameter value type: {:?}", other); } @@ -355,6 +399,7 @@ impl From for FbParameterType { ParameterType::String => FbParameterType::hlstring, ParameterType::Bool => FbParameterType::hlbool, ParameterType::VecBytes => FbParameterType::hlvecbytes, + ParameterType::ByteChunks => FbParameterType::hlbytechunks, } } } @@ -373,6 +418,7 @@ impl From for FbReturnType { ReturnType::Bool => FbReturnType::hlbool, ReturnType::Void => FbReturnType::hlvoid, ReturnType::VecBytes => FbReturnType::hlsizeprefixedbuffer, + ReturnType::ByteChunks => FbReturnType::hlbytechunks, } } } @@ -391,6 +437,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 +460,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 +585,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 +713,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"))] @@ -731,6 +801,17 @@ impl TryFrom> for ReturnValue { }; Ok(ReturnValue::VecBytes(hlvecbytes.unwrap_or(Vec::new()))) } + FbReturnValue::hlsizeprefixedbytechunks => { + let value = return_value_box + .value_as_hlsizeprefixedbytechunks() + .and_then(|value| value.value()) + .map(|value| byte_chunks_from_bytes(Bytes::copy_from_slice(value.bytes()))) + .unwrap_or_default(); + Ok(ReturnValue::ByteChunks(value)) + } + FbReturnValue::hlexternalbytes => { + bail!("External byte return requires an external value source") + } other => { bail!("Unexpected flatbuffer return value type: {:?}", other) } @@ -927,6 +1008,35 @@ impl TryFrom<&ReturnValue> for Vec { builder.finish_size_prefixed(fcr, None); builder.finished_data().to_vec() } + ReturnValue::ByteChunks(v) => { + let off = { + let value = byte_chunks_to_bytes(v); + let val = builder.create_vector(value.as_ref()); + hlsizeprefixedbytechunks::create( + &mut builder, + &hlsizeprefixedbytechunksArgs { + value: Some(val), + size: value.len() as i32, + }, + ) + }; + let rv_box = ReturnValueBox::create( + &mut builder, + &ReturnValueBoxArgs { + value: Some(off.as_union_value()), + value_type: FbReturnValue::hlsizeprefixedbytechunks, + }, + ); + 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( @@ -954,10 +1064,14 @@ impl TryFrom<&ReturnValue> for Vec { #[cfg(test)] mod tests { + use alloc::vec; + use flatbuffers::FlatBufferBuilder; use super::super::guest_error::ErrorCode; + use super::super::util::{byte_chunks_to_vec, get_flatbuffer_result}; use super::*; + use crate::flatbuffers::hyperlight::generated::{hlexternalbytes, hlexternalbytesArgs}; #[test] fn encode_success_result() { @@ -983,4 +1097,57 @@ mod tests { assert_eq!(error.code, test_error.code); assert_eq!(error.message, test_error.message); } + + #[test] + fn embedded_byte_chunks_return_round_trips() { + let mut builder = FlatBufferBuilder::new(); + let expected = vec![Bytes::from_static(b"hello"), Bytes::from_static(b" world")]; + let encoded = + FunctionCallResult::new(Ok(ReturnValue::ByteChunks(expected))).encode(&mut builder); + + let decoded = FunctionCallResult::try_from(encoded) + .unwrap() + .into_inner() + .unwrap(); + let ReturnValue::ByteChunks(decoded) = decoded else { + panic!("expected byte chunks return value"); + }; + assert_eq!(byte_chunks_to_vec(&decoded), b"hello world"); + } + + #[test] + fn direct_byte_chunks_return_encoding_preserves_logical_type() { + let encoded = get_flatbuffer_result(vec![ + Bytes::from_static(b"hello"), + Bytes::from_static(b" world"), + ]); + + let decoded = FunctionCallResult::try_from(encoded.as_slice()) + .unwrap() + .into_inner() + .unwrap(); + assert!(matches!(decoded, ReturnValue::ByteChunks(_))); + } + + #[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)); + } } diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/util.rs b/src/hyperlight_common/src/flatbuffer_wrappers/util.rs index 2ee32c9a3..d942bceb1 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/util.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/util.rs @@ -14,8 +14,10 @@ See the License for the specific language governing permissions and limitations under the License. */ +use alloc::vec; use alloc::vec::Vec; +use bytes::Bytes; use flatbuffers::FlatBufferBuilder; use crate::flatbuffer_wrappers::function_types::ParameterValue; @@ -26,7 +28,9 @@ use crate::flatbuffers::hyperlight::generated::{ 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, + hlsizeprefixedbufferArgs as FbhlsizeprefixedbufferArgs, + hlsizeprefixedbytechunks as Fbhlsizeprefixedbytechunks, + hlsizeprefixedbytechunksArgs as FbhlsizeprefixedbytechunksArgs, hlstring as Fbhlstring, hlstringArgs as FbhlstringArgs, hluint as Fbhluint, hluintArgs as FbhluintArgs, hlulong as Fbhlulong, hlulongArgs as FbhlulongArgs, hlvoid as Fbhlvoid, hlvoidArgs as FbhlvoidArgs, @@ -113,6 +117,31 @@ impl FlatbufferSerializable for &[u8] { } } +impl FlatbufferSerializable for Vec { + fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { + let value = byte_chunks_to_bytes(self); + let vec_off = builder.create_vector(value.as_ref()); + let buf_off = Fbhlsizeprefixedbytechunks::create( + builder, + &FbhlsizeprefixedbytechunksArgs { + size: value.len() as i32, + value: Some(vec_off), + }, + ); + let rv_box = ReturnValueBox::create( + builder, + &ReturnValueBoxArgs { + value_type: FbReturnValue::hlsizeprefixedbytechunks, + 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 }); @@ -260,6 +289,7 @@ pub fn estimate_flatbuffer_capacity(function_name: &str, args: &[ParameterValue] estimated_capacity += match arg { ParameterValue::String(s) => s.len() + 20, ParameterValue::VecBytes(v) => v.len() + 20, + ParameterValue::ByteChunks(v) => byte_chunks_len(v) + 20, ParameterValue::Int(_) | ParameterValue::UInt(_) => 16, ParameterValue::Long(_) | ParameterValue::ULong(_) => 20, ParameterValue::Float(_) => 16, @@ -272,6 +302,54 @@ 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() +} + +/// 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; 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/virtq/msg.rs b/src/hyperlight_common/src/virtq/msg.rs index 090c2eb5b..d0e937bfc 100644 --- a/src/hyperlight_common/src/virtq/msg.rs +++ b/src/hyperlight_common/src/virtq/msg.rs @@ -16,11 +16,10 @@ 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; +//! Every message chain on both the G2H and H2G queues starts with this fixed +//! 8-byte header, enabling message type discrimination and request/response +//! correlation. Payload lengths come from the size-prefixed FlatBuffer and its +//! external-byte declarations. /// Message types for the virtqueue wire protocol. #[repr(u8)] @@ -56,49 +55,27 @@ impl TryFrom for MsgKind { } } -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)] +/// Wire header for all virtqueue messages. +#[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable)] #[repr(C)] pub struct VirtqMsgHeader { /// Discriminates the message type. pub kind: u8, - /// Per-message flags (see [`MsgFlags`]). - pub flags: u8, + /// keep the header 8 bytes long and aligned to 4 bytes. + reserved: [u8; 3], /// 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, + pub cid: 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 { + /// Create a message header. + pub const fn new(kind: MsgKind, cid: u32) -> Self { Self { kind: kind as u8, - flags: flags.bits(), - req_id, - payload_len, + reserved: [0; 3], + cid, } } @@ -107,14 +84,47 @@ impl VirtqMsgHeader { MsgKind::try_from(self.kind) } - /// Interpret the raw flags field as [`MsgFlags`]. - pub fn msg_flags(&self) -> MsgFlags { - MsgFlags::from_bits_truncate(self.flags) + /// Return the wire representation. + pub fn as_bytes(&self) -> &[u8] { + bytemuck::bytes_of(self) } - /// 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 + /// 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) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn header_contains_only_kind_and_cid() { + let header = VirtqMsgHeader::new(MsgKind::Response, 0x1234_5678); + + assert_eq!(VirtqMsgHeader::SIZE, 8); + assert_eq!(header.msg_kind(), Ok(MsgKind::Response)); + assert_eq!(header.cid, 0x1234_5678); + assert_eq!(header.reserved, [0; 3]); + } + + #[test] + fn rejects_invalid_wire_headers() { + let header = VirtqMsgHeader::new(MsgKind::Request, 1); + let mut bytes = [0; VirtqMsgHeader::SIZE]; + bytes.copy_from_slice(header.as_bytes()); + + bytes[1] = 1; + assert_eq!(VirtqMsgHeader::from_bytes(&bytes), None); + + bytes[1] = 0; + bytes[0] = u8::MAX; + assert_eq!(VirtqMsgHeader::from_bytes(&bytes), None); + assert_eq!(VirtqMsgHeader::from_bytes(&bytes[..7]), None); } } 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_bin/src/guest_function/definition.rs b/src/hyperlight_guest_bin/src/guest_function/definition.rs index c96f2e9fe..46347f016 100644 --- a/src/hyperlight_guest_bin/src/guest_function/definition.rs +++ b/src/hyperlight_guest_bin/src/guest_function/definition.rs @@ -90,6 +90,7 @@ fn into_flatbuffer_result(value: ReturnValue) -> Vec { 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()), + ReturnValue::ByteChunks(v) => get_flatbuffer_result(v), } } diff --git a/src/hyperlight_guest_capi/include/macro.h b/src/hyperlight_guest_capi/include/macro.h index 1c6dc1ff7..99ca16790 100644 --- a/src/hyperlight_guest_capi/include/macro.h +++ b/src/hyperlight_guest_capi/include/macro.h @@ -8,7 +8,7 @@ // // 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. // See c_simpleguest/main.c for an example. diff --git a/src/hyperlight_guest_capi/src/flatbuffer.rs b/src/hyperlight_guest_capi/src/flatbuffer.rs index ff12400d6..77fb7be67 100644 --- a/src/hyperlight_guest_capi/src/flatbuffer.rs +++ b/src/hyperlight_guest_capi/src/flatbuffer.rs @@ -20,7 +20,10 @@ 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_common::flatbuffer_wrappers::function_types::Bytes; +use hyperlight_common::flatbuffer_wrappers::util::{ + byte_chunks_from_vec, byte_chunks_to_vec, get_flatbuffer_result, +}; use hyperlight_guest_bin::host_comm::get_host_return_value; use crate::types::FfiVec; @@ -95,6 +98,14 @@ pub extern "C" fn hl_flatbuffer_result_from_Bytes(data: *const u8, len: usize) - Box::new(unsafe { FfiVec::from_vec(vec) }) } +#[unsafe(no_mangle)] +pub extern "C" fn hl_flatbuffer_result_from_ByteChunks(data: *const u8, len: usize) -> Box { + let slice = unsafe { core::slice::from_raw_parts(data, len) }; + let vec = get_flatbuffer_result(byte_chunks_from_vec(slice.to_vec())); + + 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); @@ -156,3 +167,11 @@ pub extern "C" fn hl_get_host_return_value_as_VecBytes() -> Box { Box::new(unsafe { FfiVec::from_vec(vec_value) }) } + +#[unsafe(no_mangle)] +pub extern "C" fn hl_get_host_return_value_as_ByteChunks() -> Box { + let chunks: Vec = + get_host_return_value().expect("Unable to get host return value as byte chunks"); + + Box::new(unsafe { FfiVec::from_vec(byte_chunks_to_vec(&chunks)) }) +} diff --git a/src/hyperlight_guest_capi/src/types/parameter.rs b/src/hyperlight_guest_capi/src/types/parameter.rs index 048169754..ead032658 100644 --- a/src/hyperlight_guest_capi/src/types/parameter.rs +++ b/src/hyperlight_guest_capi/src/types/parameter.rs @@ -18,6 +18,7 @@ use alloc::ffi::CString; use core::ffi::{CStr, c_char}; use hyperlight_common::flatbuffer_wrappers::function_types::{ParameterType, ParameterValue}; +use hyperlight_common::flatbuffer_wrappers::util::{byte_chunks_from_vec, byte_chunks_to_vec}; use hyperlight_guest::error::Result; use crate::types::FfiVec; @@ -38,6 +39,7 @@ pub union FfiParameterValue { pub Bool: bool, pub String: *mut c_char, pub VecBytes: FfiVec, + pub ByteChunks: FfiVec, } /// An owned FFI version Of `ParameterValue` @@ -71,6 +73,13 @@ impl FfiParameter { FfiParameterValue { VecBytes: leaked }, ) } + ParameterValue::ByteChunks(v) => { + let leaked = unsafe { FfiVec::from_vec(byte_chunks_to_vec(&v)) }; + ( + ParameterType::ByteChunks, + FfiParameterValue { ByteChunks: leaked }, + ) + } }; Ok(FfiParameter { tag, value: union }) } @@ -95,6 +104,9 @@ impl FfiParameter { ParameterType::VecBytes => { ParameterValue::VecBytes(unsafe { self.value.VecBytes.copy_to_vec() }) } + ParameterType::ByteChunks => ParameterValue::ByteChunks(byte_chunks_from_vec(unsafe { + self.value.ByteChunks.copy_to_vec() + })), } } } @@ -108,6 +120,9 @@ impl Drop for FfiParameter { ParameterType::VecBytes => unsafe { drop(self.value.VecBytes.into_vec()); }, + ParameterType::ByteChunks => unsafe { + drop(self.value.ByteChunks.into_vec()); + }, _ => {} } } 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/sandbox/snapshot/file/config.rs b/src/hyperlight_host/src/sandbox/snapshot/file/config.rs index 01c2f501b..efae026d2 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/config.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file/config.rs @@ -254,6 +254,7 @@ enum ParameterTypeRepr { String, Bool, VecBytes, + ByteChunks, } /// JSON-friendly mirror of @@ -271,6 +272,7 @@ enum ReturnTypeRepr { Bool, Void, VecBytes, + ByteChunks, } impl From<&ParameterType> for ParameterTypeRepr { @@ -285,6 +287,7 @@ impl From<&ParameterType> for ParameterTypeRepr { ParameterType::String => Self::String, ParameterType::Bool => Self::Bool, ParameterType::VecBytes => Self::VecBytes, + ParameterType::ByteChunks => Self::ByteChunks, } } } @@ -301,6 +304,7 @@ impl From for ParameterType { ParameterTypeRepr::String => Self::String, ParameterTypeRepr::Bool => Self::Bool, ParameterTypeRepr::VecBytes => Self::VecBytes, + ParameterTypeRepr::ByteChunks => Self::ByteChunks, } } } @@ -318,6 +322,7 @@ impl From<&ReturnType> for ReturnTypeRepr { ReturnType::Bool => Self::Bool, ReturnType::Void => Self::Void, ReturnType::VecBytes => Self::VecBytes, + ReturnType::ByteChunks => Self::ByteChunks, } } } @@ -335,6 +340,7 @@ impl From for ReturnType { ReturnTypeRepr::Bool => Self::Bool, ReturnTypeRepr::Void => Self::Void, ReturnTypeRepr::VecBytes => Self::VecBytes, + ReturnTypeRepr::ByteChunks => Self::ByteChunks, } } } @@ -693,6 +699,7 @@ mod tests { ParameterType::String, ParameterType::Bool, ParameterType::VecBytes, + ParameterType::ByteChunks, ]; for p in variants { let back: ParameterType = ParameterTypeRepr::from(&p).into(); @@ -715,6 +722,7 @@ mod tests { ReturnType::Bool, ReturnType::Void, ReturnType::VecBytes, + ReturnType::ByteChunks, ]; for r in variants { let back: ReturnType = ReturnTypeRepr::from(&r).into(); 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/rust_guests/simpleguest/src/main.rs b/src/tests/rust_guests/simpleguest/src/main.rs index fc0ce416b..a40f6ca0b 100644 --- a/src/tests/rust_guests/simpleguest/src/main.rs +++ b/src/tests/rust_guests/simpleguest/src/main.rs @@ -1544,6 +1544,7 @@ fn fuzz_host_function(func: FunctionCall) -> Result> { ReturnValue::Bool(bool) => Ok(get_flatbuffer_result(bool)), ReturnValue::Void(()) => Ok(get_flatbuffer_result(())), ReturnValue::VecBytes(byte) => Ok(get_flatbuffer_result(byte.as_slice())), + ReturnValue::ByteChunks(chunks) => Ok(get_flatbuffer_result(chunks)), }, Err(e) => Err(e), } From 120e67b85d9458cfee3051d6d505d69ec17f0904 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Tue, 21 Jul 2026 11:19:50 +0200 Subject: [PATCH 02/34] feat(virtq): add external function value codecs Add external-aware function call and result encoding and decoding. This changes are transitional only and next commits will unify embedded and external (de)coding paths. Signed-off-by: Tomasz Andrzejak --- .../src/flatbuffer_wrappers/codec.rs | 45 ++ .../src/flatbuffer_wrappers/function_call.rs | 417 +++++++++++++++++- .../src/flatbuffer_wrappers/function_types.rs | 346 ++++++++++++++- .../src/flatbuffer_wrappers/mod.rs | 3 + .../src/flatbuffer_wrappers/util.rs | 7 + 5 files changed, 809 insertions(+), 9 deletions(-) create mode 100644 src/hyperlight_common/src/flatbuffer_wrappers/codec.rs 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..41052b9ac --- /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 markers are 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 f31e96b8d..7270f191c 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs @@ -22,14 +22,16 @@ use flatbuffers::{FlatBufferBuilder, WIPOffset, size_prefixed_root}; #[cfg(feature = "tracing")] use tracing::{Span, instrument}; -use super::function_types::{ParameterValue, ReturnType}; -use super::util::byte_chunks_to_bytes; +use super::codec::{ExternalValueSink, ExternalValueSource}; +use super::function_types::{ParameterValue, ReturnType, decode_external_parameter_value}; +use super::util::{byte_chunks_to_bytes, 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, hlbytechunks, hlbytechunksArgs, - hldouble, hldoubleArgs, hlfloat, hlfloatArgs, hlint, hlintArgs, hllong, hllongArgs, hlstring, - hlstringArgs, hluint, hluintArgs, hlulong, hlulongArgs, hlvecbytes, hlvecbytesArgs, + hldouble, hldoubleArgs, hlexternalbytes, hlexternalbytesArgs, hlfloat, hlfloatArgs, hlint, + hlintArgs, hllong, hllongArgs, hlstring, hlstringArgs, hluint, hluintArgs, hlulong, + hlulongArgs, hlvecbytes, hlvecbytesArgs, }; /// The type of function call. @@ -230,6 +232,222 @@ impl FunctionCall { builder.finish_size_prefixed(function_call, None); builder.finished_data() } + + /// Encodes byte parameters as external markers and sends their payloads to + /// `external_values` in parameter order. + pub fn encode_external<'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 { + FunctionCallType::Guest => FbFunctionCallType::guest, + FunctionCallType::Host => FbFunctionCallType::host, + }; + + let expected_return_type = self.expected_return_type.into(); + + let parameters = match &self.parameters { + Some(parameters) if !parameters.is_empty() => { + let parameter_offsets: Vec> = parameters + .iter() + .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::>>()?; + Some(builder.create_vector(¶meter_offsets)) + } + _ => None, + }; + + let function_call = FbFunctionCall::create( + builder, + &FbFunctionCallArgs { + function_name: Some(function_name), + parameters, + function_call_type, + expected_return_type, + }, + ); + builder.finish_size_prefixed(function_call, None); + Ok(builder.finished_data()) + } + + /// Decodes a function call using `external_values` for external byte + /// markers. + pub fn decode_external(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(); + let function_call_type = match function_call_fb.function_call_type() { + FbFunctionCallType::guest => FunctionCallType::Guest, + FbFunctionCallType::host => FunctionCallType::Host, + other => { + bail!("Invalid function call type: {:?}", other); + } + }; + let expected_return_type = function_call_fb.expected_return_type().try_into()?; + + let parameters = function_call_fb + .parameters() + .map(|parameters| { + parameters + .iter() + .map(|parameter| decode_external_parameter_value(parameter, external_values)) + .collect::>>() + }) + .transpose()?; + + external_values.finish()?; + Ok(Self { + function_name: function_name.to_string(), + parameters, + function_call_type, + expected_return_type, + }) + } } #[cfg_attr(feature = "tracing", instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace"))] @@ -292,11 +510,75 @@ impl TryFrom<&[u8]> for FunctionCall { #[cfg(test)] mod tests { + use alloc::collections::VecDeque; use alloc::vec; use super::*; 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(); @@ -370,4 +652,131 @@ mod tests { ]) ); } + + #[test] + fn external_byte_parameters_round_trip_in_parameter_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_external(&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 marker = parameter.value_as_hlexternalbytes().unwrap(); + assert_eq!(marker.length(), length); + assert_eq!(marker.chunked(), chunked); + } + + assert!(FunctionCall::try_from(encoded).is_err()); + let decoded = FunctionCall::decode_external(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 external_encoding_matches_embedded_encoding_without_byte_parameters() { + let call = FunctionCall::new( + "scalars".to_string(), + Some(vec![ + ParameterValue::Int(42), + ParameterValue::String("value".to_string()), + ]), + FunctionCallType::Guest, + ReturnType::Bool, + ); + let mut embedded_builder = FlatBufferBuilder::new(); + let embedded = call.encode(&mut embedded_builder).to_vec(); + + let mut external_builder = FlatBufferBuilder::new(); + let mut external_values = TestExternalValues::default(); + let external = call + .encode_external(&mut external_builder, &mut external_values) + .unwrap(); + + assert_eq!(external, embedded); + assert!(external_values.values.is_empty()); + } + + #[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_external(&mut builder, &mut encoded_values) + .unwrap(); + + let mut missing = TestExternalValues::default(); + assert!(FunctionCall::decode_external(encoded, &mut missing).is_err()); + + let mut wrong_type = + TestExternalValues::from_values([TestExternalValue::ByteChunks(vec![ + Bytes::from_static(b"123"), + ])]); + assert!(FunctionCall::decode_external(encoded, &mut wrong_type).is_err()); + + let mut wrong_length = + TestExternalValues::from_values([TestExternalValue::VecBytes(vec![1, 2])]); + assert!(FunctionCall::decode_external(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_external(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 3db72780a..a0117f437 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs @@ -23,18 +23,20 @@ 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::{byte_chunks_from_bytes, byte_chunks_to_bytes}; +use super::util::{byte_chunks_from_bytes, byte_chunks_to_bytes, 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, hlsizeprefixedbytechunks, hlsizeprefixedbytechunksArgs, hlstring, - hlstringArgs, hluint, hluintArgs, hlulong, hlulongArgs, hlvoid, hlvoidArgs, + ReturnValueBox, ReturnValueBoxArgs, hlbool, hlboolArgs, hldouble, hldoubleArgs, + hlexternalbytes, hlexternalbytesArgs, hlfloat, hlfloatArgs, hlint, hlintArgs, hllong, + hllongArgs, hlsizeprefixedbuffer, hlsizeprefixedbufferArgs, hlsizeprefixedbytechunks, + hlsizeprefixedbytechunksArgs, hlstring, hlstringArgs, hluint, hluintArgs, hlulong, hlulongArgs, + hlvoid, hlvoidArgs, }; pub struct FunctionCallResult(core::result::Result); @@ -154,6 +156,61 @@ impl FunctionCallResult { } } } + + /// Encodes byte returns as external markers and sends their payload to + /// `external_values`. + /// + /// Non-byte returns and guest errors retain their existing embedded + /// encoding. + pub fn encode_external<'a, 'b, S>( + &'a self, + builder: &'b mut flatbuffers::FlatBufferBuilder, + external_values: &mut S, + ) -> Result<&'b [u8]> + where + S: ExternalValueSink<'a> + ?Sized, + { + let Ok(return_value) = &self.0 else { + return Ok(self.encode(builder)); + }; + + let (length, chunked) = match return_value { + ReturnValue::VecBytes(value) => { + let length = u64::try_from(value.len()) + .map_err(|_| anyhow!("External VecBytes length does not fit in u64"))?; + external_values.push_bytes(value)?; + (length, false) + } + ReturnValue::ByteChunks(value) => { + let length = try_byte_chunks_len(value) + .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(value)?; + (length, true) + } + _ => return Ok(self.encode(builder)), + }; + + let value = hlexternalbytes::create(builder, &hlexternalbytesArgs { length, chunked }); + let return_value = ReturnValueBox::create( + builder, + &ReturnValueBoxArgs { + value: Some(value.as_union_value()), + value_type: FbReturnValue::hlexternalbytes, + }, + ); + let result = FbFunctionCallResult::create( + builder, + &FbFunctionCallResultArgs { + result: Some(return_value.as_union_value()), + result_type: FunctionCallResultType::ReturnValueBox, + }, + ); + builder.finish_size_prefixed(result, None); + Ok(builder.finished_data()) + } + pub fn new(value: core::result::Result) -> Self { FunctionCallResult(value) } @@ -161,6 +218,44 @@ impl FunctionCallResult { pub fn into_inner(self) -> core::result::Result { self.0 } + + /// Decodes a function-call result using `external_values` for external byte + /// markers. + pub fn decode_external(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))?; + + 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") + })?; + Ok(decode_external_return_value(boxed, external_values)?) + } + FunctionCallResultType::GuestError => { + let guest_error_table = function_call_result_fb + .result_as_guest_error() + .ok_or_else(|| anyhow!("Failed to get GuestError from function call result"))?; + let code = guest_error_table.code(); + let message = guest_error_table + .message() + .map(|s| s.to_string()) + .unwrap_or_default(); + Err(GuestError::new(code.into(), message)) + } + other => { + bail!("Unexpected function call result type: {:?}", other) + } + }; + + external_values.finish()?; + Ok(FunctionCallResult(result)) + } } impl TryFrom<&[u8]> for FunctionCallResult { @@ -318,6 +413,98 @@ pub enum ReturnType { ByteChunks, } +pub(crate) fn decode_external_parameter_value( + parameter: Parameter<'_>, + external_values: &mut S, +) -> Result +where + S: ExternalValueSource + ?Sized, +{ + if parameter.value_type() != FbParameterValue::hlexternalbytes { + return parameter.try_into(); + } + + let marker = parameter + .value_as_hlexternalbytes() + .ok_or_else(|| anyhow!("Failed to get external byte parameter marker"))?; + let length = usize::try_from(marker.length()).map_err(|_| { + anyhow!( + "External byte parameter length {} does not fit in usize", + marker.length() + ) + })?; + + if marker.chunked() { + let value = external_values.take_chunks(length)?; + let actual_length = try_byte_chunks_len(&value) + .ok_or_else(|| anyhow!("External ByteChunks parameter length overflow"))?; + if actual_length != length { + bail!( + "External ByteChunks parameter length mismatch: declared {}, received {}", + length, + actual_length + ); + } + Ok(ParameterValue::ByteChunks(value)) + } else { + let value = external_values.take_bytes(length)?; + if value.len() != length { + bail!( + "External VecBytes parameter length mismatch: declared {}, received {}", + length, + value.len() + ); + } + Ok(ParameterValue::VecBytes(value)) + } +} + +fn decode_external_return_value( + return_value: ReturnValueBox<'_>, + external_values: &mut S, +) -> Result +where + S: ExternalValueSource + ?Sized, +{ + if return_value.value_type() != FbReturnValue::hlexternalbytes { + return return_value.try_into(); + } + + let marker = return_value + .value_as_hlexternalbytes() + .ok_or_else(|| anyhow!("Failed to get external byte return marker"))?; + let length = usize::try_from(marker.length()).map_err(|_| { + anyhow!( + "External byte return length {} does not fit in usize", + marker.length() + ) + })?; + + if marker.chunked() { + let value = external_values.take_chunks(length)?; + let actual_length = try_byte_chunks_len(&value) + .ok_or_else(|| anyhow!("External ByteChunks return length overflow"))?; + if actual_length != length { + bail!( + "External ByteChunks return length mismatch: declared {}, received {}", + length, + actual_length + ); + } + Ok(ReturnValue::ByteChunks(value)) + } else { + let value = external_values.take_bytes(length)?; + if value.len() != length { + bail!( + "External VecBytes return length mismatch: declared {}, received {}", + length, + value.len() + ); + } + Ok(ReturnValue::VecBytes(value)) + } +} + impl From<&ParameterValue> for ParameterType { #[cfg_attr(feature = "tracing", instrument(skip_all, parent = Span::current(), level= "Trace"))] fn from(value: &ParameterValue) -> Self { @@ -1064,6 +1251,7 @@ impl TryFrom<&ReturnValue> for Vec { #[cfg(test)] mod tests { + use alloc::collections::VecDeque; use alloc::vec; use flatbuffers::FlatBufferBuilder; @@ -1073,6 +1261,69 @@ mod tests { 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(); @@ -1150,4 +1401,89 @@ mod tests { assert!(!round_trip(false)); assert!(round_trip(true)); } + + #[test] + fn external_byte_returns_round_trip_without_embedding_payloads() { + 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_external(&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 marker = 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!(marker.length(), length as u64); + assert_eq!(marker.chunked(), chunked); + + assert!(FunctionCallResult::try_from(encoded).is_err()); + let decoded = FunctionCallResult::decode_external(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_external(&mut builder, &mut encoded_values) + .unwrap(); + + let mut missing = TestExternalValues::default(); + assert!(FunctionCallResult::decode_external(encoded, &mut missing).is_err()); + + let mut wrong_type = + TestExternalValues::from_values([TestExternalValue::VecBytes(vec![1, 2, 3])]); + assert!(FunctionCallResult::decode_external(encoded, &mut wrong_type).is_err()); + + let mut wrong_length = + TestExternalValues::from_values([TestExternalValue::ByteChunks(vec![ + Bytes::from_static(b"12"), + ])]); + assert!(FunctionCallResult::decode_external(encoded, &mut wrong_length).is_err()); + } + + #[test] + fn external_result_decoder_rejects_unused_values() { + let result = FunctionCallResult::new(Ok(ReturnValue::Int(42))); + let mut embedded_builder = FlatBufferBuilder::new(); + let embedded = result.encode(&mut embedded_builder).to_vec(); + + let mut external_builder = FlatBufferBuilder::new(); + let mut external_values = TestExternalValues::default(); + let external = result + .encode_external(&mut external_builder, &mut external_values) + .unwrap(); + assert_eq!(external, embedded); + assert!(external_values.values.is_empty()); + + external_values + .values + .push_back(TestExternalValue::VecBytes(Vec::new())); + assert!(FunctionCallResult::decode_external(external, &mut external_values).is_err()); + } } 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 d942bceb1..57213d1ba 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/util.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/util.rs @@ -321,6 +321,13 @@ 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. From e0a51c02e1608acd57c8d2cae7d6a1eec3c748fa Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Tue, 21 Jul 2026 17:43:28 +0200 Subject: [PATCH 03/34] feat(virtq): add stateful chain byte streams Read received payloads directly into caller owned storage and require paired completion of readable/writable chains. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/benches/common/mod.rs | 8 +- src/hyperlight_common/src/virtq/buffer.rs | 54 ++ .../src/virtq/concurrency.rs | 16 +- src/hyperlight_common/src/virtq/consumer.rs | 715 ++++++++++++++---- src/hyperlight_common/src/virtq/mod.rs | 105 +-- src/hyperlight_common/src/virtq/producer.rs | 198 +++-- 6 files changed, 785 insertions(+), 311 deletions(-) diff --git a/src/hyperlight_common/benches/common/mod.rs b/src/hyperlight_common/benches/common/mod.rs index bd0930fcd..7461de6fa 100644 --- a/src/hyperlight_common/benches/common/mod.rs +++ b/src/hyperlight_common/benches/common/mod.rs @@ -232,8 +232,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 +258,13 @@ where let token = pair.producer.submit(chain).unwrap(); let (recv, reply) = pair.consumer.poll(request.len()).unwrap().unwrap(); - black_box(recv.segments().segment_count()); + black_box(recv.len()); let ReplyChain::Writable(mut writable) = reply else { panic!("expected writable reply"); }; writable.write_all(response).unwrap(); - pair.consumer.complete(writable).unwrap(); + pair.consumer.complete(recv, writable).unwrap(); let used = pair.producer.poll().unwrap().unwrap(); debug_assert_eq!(used.token(), token); diff --git a/src/hyperlight_common/src/virtq/buffer.rs b/src/hyperlight_common/src/virtq/buffer.rs index 0d1ca4d6a..e93cebda1 100644 --- a/src/hyperlight_common/src/virtq/buffer.rs +++ b/src/hyperlight_common/src/virtq/buffer.rs @@ -187,6 +187,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()) @@ -472,6 +500,32 @@ mod tests { assert_eq!(cursor.chunk(), b"world"); } + #[test] + fn segments_split_to_shares_boundary_segment() { + let boundary = Bytes::from(vec![b'd', b'e', b'f']); + let boundary_ptr = boundary.as_ptr(); + let mut segments = Segments::new([ + Bytes::from_static(b"abc"), + boundary, + Bytes::from_static(b"ghi"), + ]); + + let prefix = segments.split_to(5).unwrap(); + + assert_eq!(prefix.segment_count(), 2); + assert_eq!(prefix.to_bytes().as_ref(), b"abcde"); + assert_eq!(prefix.as_slice()[1].as_ptr(), boundary_ptr); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.to_bytes().as_ref(), b"fghi"); + assert_eq!( + segments.as_slice()[0].as_ptr(), + boundary_ptr.wrapping_add(2) + ); + + assert!(segments.split_to(5).is_none()); + assert_eq!(segments.to_bytes().as_ref(), b"fghi"); + } + #[test] fn segments_into_bytes_reuses_single_segment() { let segment = Bytes::from(vec![1, 2, 3, 4]); diff --git a/src/hyperlight_common/src/virtq/concurrency.rs b/src/hyperlight_common/src/virtq/concurrency.rs index ad58b05cb..33a369d5d 100644 --- a/src/hyperlight_common/src/virtq/concurrency.rs +++ b/src/hyperlight_common/src/virtq/concurrency.rs @@ -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(); @@ -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(); @@ -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(); diff --git a/src/hyperlight_common/src/virtq/consumer.rs b/src/hyperlight_common/src/virtq/consumer.rs index d4fe0ff45..22ce9bc13 100644 --- a/src/hyperlight_common/src/virtq/consumer.rs +++ b/src/hyperlight_common/src/virtq/consumer.rs @@ -15,6 +15,7 @@ limitations under the License. */ use alloc::vec; +use core::fmt; use bytes::Bytes; use fixedbitset::FixedBitSet; @@ -22,46 +23,162 @@ use smallvec::SmallVec; use super::*; -type WritableElems = SmallVec<[BufferElement; 2]>; - -/// Payload received from the producer, safely copied out of shared memory. +/// Stateful reader over device-readable descriptors received from the producer. /// -/// Created by [`VirtqConsumer::poll`]. Device-readable segments are eagerly -/// copied during poll using [`MemOps::read`] (volatile on the host side), so -/// accessing data requires no unsafe code and no references into shared -/// memory. Segment boundaries are preserved in [`Segments`]. -#[derive(Debug, Clone)] -pub struct RecvChain { - token: Token, - segments: Segments, +/// Reads copy directly from shared memory into caller-provided final storage. +/// The chain must be returned together with its paired [`ReplyChain`] through +/// [`VirtqConsumer::complete`] before the descriptors can be reused. +#[must_use = "dropping without completing leaks the descriptor"] +pub struct RecvChain { + state: ChainState, } -impl RecvChain { +impl fmt::Debug for RecvChain { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RecvChain") + .field("token", &self.state.token) + .field("elems", &self.state.elems) + .field("len", &self.state.total) + .field("consumed", &self.state.position) + .field("desc_index", &self.state.desc_idx) + .field("desc_offset", &self.state.desc_off) + .finish() + } +} + +impl RecvChain { + fn new(mem: M, token: Token, elems: ChainElems, len: usize) -> Self { + Self { + state: ChainState::new(mem, token, elems, len), + } + } + /// The token identifying this chain. + #[inline] pub fn token(&self) -> Token { - self.token + self.state.token() + } + + /// Total readable payload length. + #[inline] + pub fn len(&self) -> usize { + self.state.total() } - /// The chain payload as ordered byte segments. - pub fn segments(&self) -> &Segments { - &self.segments + /// Whether this chain has no readable payload. + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 } - /// Consume the chain, taking ownership of the segments. - pub fn into_segments(self) -> Segments { - self.segments + /// Number of bytes consumed by the stateful reader. + #[inline] + pub fn consumed(&self) -> usize { + self.state.position() + } + + /// Number of bytes still available to the stateful reader. + #[inline] + pub fn remaining(&self) -> usize { + self.state.remaining() } - /// Return the chain payload as contiguous bytes. + /// Read bytes sequentially across descriptor boundaries. /// - /// Returns empty [`Bytes`] when the chain has no readable buffers. - pub fn to_bytes(&self) -> Bytes { - self.segments.to_bytes() + /// Returns the number of bytes copied, which may be smaller than `buf.len()` + /// at the end of the chain. If a later memory read fails, the cursor remains + /// advanced past any earlier chunks copied by the same call. + pub fn read(&mut self, buf: &mut [u8]) -> Result { + let len = buf.len().min(self.remaining()); + let mut dst = &mut buf[..len]; + let mut read = 0; + + while !dst.is_empty() { + let Some(elem) = self.state.current_elem() else { + break; + }; + + let desc_len = elem.len as usize; + let desc_offset = self.state.desc_offset(); + let len = (desc_len - desc_offset).min(dst.len()); + let (current, rest) = dst.split_at_mut(len); + + let addr = elem + .addr + .checked_add(desc_offset as u64) + .ok_or(VirtqError::MemoryReadError)?; + + self.state + .mem + .read(addr, current) + .map_err(|_| VirtqError::MemoryReadError)?; + + self.state.advance(len); + read += len; + dst = rest; + } + + Ok(read) } - /// Consume the chain and return the payload as contiguous bytes. - pub fn into_bytes(self) -> Bytes { - self.segments.into_bytes() + /// Read exactly `buf.len()` bytes or return an error. + #[inline] + pub fn read_exact(&mut self, buf: &mut [u8]) -> Result<&mut Self, VirtqError> { + if buf.len() > self.remaining() { + return Err(VirtqError::ReceiveTooShort { + requested: buf.len(), + remaining: self.remaining(), + }); + } + + let read = self.read(buf)?; + debug_assert_eq!(read, buf.len()); + Ok(self) + } + + /// Copy the complete payload into descriptor-preserving owned segments. + /// + /// This does not change the stateful read position. Each call takes a new + /// snapshot of shared memory; callers should validate and use the returned + /// owned value rather than reading the same untrusted payload again. + pub fn to_segments(&self) -> Result { + let mut segments = SmallVec::<[Bytes; 4]>::new(); + + for elem in &self.state.elems { + let mut buf = vec![0u8; elem.len as usize]; + self.state + .mem + .read(elem.addr, &mut buf) + .map_err(|_| VirtqError::MemoryReadError)?; + segments.push(Bytes::from(buf)); + } + + Ok(Segments::from_smallvec(segments)) + } + + /// Copy the complete payload directly into one contiguous allocation. + /// + /// This does not change the stateful read position. Each call takes a new + /// snapshot of shared memory; callers should validate and use the returned + /// owned value rather than reading the same untrusted payload again. + pub fn to_bytes(&self) -> Result { + if self.is_empty() { + return Ok(Bytes::new()); + } + + let mut buf = vec![0u8; self.len()]; + let mut offset = 0; + + for elem in &self.state.elems { + let end = offset + elem.len as usize; + self.state + .mem + .read(elem.addr, &mut buf[offset..end]) + .map_err(|_| VirtqError::MemoryReadError)?; + offset = end; + } + + Ok(Bytes::from(buf)) } } @@ -75,13 +192,15 @@ pub enum ReplyChain { /// Use the `write*` methods on [`WritableChain`] to fill the /// response buffer. Writable(WritableChain), - /// Ack-only reply (for chains with only readable buffers). No response buffer. - /// Just pass back to [`VirtqConsumer::complete`] to acknowledge. + /// Ack-only reply (for chains with only readable buffers). No response + /// buffer. Pass it back with the paired [`RecvChain`] through + /// [`VirtqConsumer::complete`] to acknowledge. Ack(AckChain), } impl ReplyChain { /// The token identifying this reply. + #[inline] pub fn token(&self) -> Token { match self { ReplyChain::Writable(wc) => wc.token(), @@ -90,9 +209,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 +236,44 @@ impl ReplyChain { /// ```ignore /// if let ReplyChain::Writable(mut wc) = reply { /// wc.write_all(b"response data")?; -/// consumer.complete(wc)?; +/// consumer.complete(recv, wc)?; /// } /// ``` #[must_use = "dropping without completing leaks the descriptor"] pub struct WritableChain { - mem: M, - token: Token, - elems: WritableElems, - capacity: usize, - written: usize, + state: ChainState, } impl WritableChain { - fn new(mem: M, token: Token, elems: WritableElems) -> Self { + fn new(mem: M, token: Token, elems: ChainElems) -> Self { let capacity = elems.iter().map(|elem| elem.len as usize).sum(); Self { - mem, - token, - elems, - capacity, - written: 0, + state: ChainState::new(mem, token, elems, capacity), } } /// The token identifying this writable reply. + #[inline] pub fn token(&self) -> Token { - self.token + self.state.token() } /// Total reply capacity in bytes. + #[inline] pub fn capacity(&self) -> usize { - self.capacity + self.state.total() } /// Number of bytes written so far. + #[inline] pub fn written(&self) -> usize { - self.written + self.state.position() } /// Remaining reply capacity. + #[inline] pub fn remaining(&self) -> usize { - self.capacity() - self.written() + self.state.remaining() } /// Write bytes into writable buffers, returning how many were written. @@ -165,15 +281,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 +323,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 +339,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 +360,7 @@ impl AckChain { Self { token } } + #[inline] pub fn token(&self) -> Token { self.token } @@ -234,29 +377,30 @@ impl AckChain { /// let mut consumer = VirtqConsumer::new(layout, mem, notifier); /// /// // Poll and process -/// while let Some((chain, reply)) = consumer.poll(MAX_RECV_LEN)? { -/// let data = chain.to_bytes(); +/// while let Some((recv, reply)) = consumer.poll(MAX_RECV_LEN)? { +/// let data = recv.to_bytes()?; /// match reply { /// ReplyChain::Writable(mut wc) => { /// let response = handle_request(data); /// wc.write_all(&response)?; -/// consumer.complete(wc)?; +/// consumer.complete(recv, wc)?; /// } /// ReplyChain::Ack(ack) => { -/// consumer.complete(ack)?; +/// consumer.complete(recv, ack)?; /// } /// } /// } /// /// // Or defer completions /// let mut pending = Vec::new(); -/// while let Some((chain, reply)) = consumer.poll(MAX_RECV_LEN)? { -/// pending.push((process(chain), reply)); +/// while let Some((recv, reply)) = consumer.poll(MAX_RECV_LEN)? { +/// let result = process(&recv); +/// pending.push((result, recv, reply)); /// } /// -/// for (result, reply) in pending { +/// for (result, recv, reply) in pending { /// // ... complete later ... -/// consumer.complete(reply)?; +/// consumer.complete(recv, reply)?; /// } /// ``` pub struct VirtqConsumer { @@ -288,29 +432,29 @@ impl VirtqConsumer { /// Poll for a single incoming chain from the driver. /// - /// Returns a [`RecvChain`] (copied data) and a [`ReplyChain`] (writable reply - /// capacity or ack token). Both are independent owned values with no borrow - /// on the consumer. + /// Returns a stateful [`RecvChain`] reader and a [`ReplyChain`] writable + /// reply or ack capability. Both are independent owned values with no + /// borrow on the consumer, but they must be returned together through + /// [`complete`](Self::complete). /// - /// On [`VirtqError::BadChain`], [`VirtqError::PayloadTooLarge`], and - /// [`VirtqError::MemoryReadError`] the descriptor is returned to the driver - /// (completed with zero length) before the error is propagated, so a - /// rejected chain does not leak. + /// On [`VirtqError::BadChain`] and [`VirtqError::PayloadTooLarge`] the + /// descriptor is returned to the driver (completed with zero length) before + /// the error is propagated, so a rejected chain does not leak. /// /// # Arguments /// - /// * `max_recv_len` - Maximum receive payload size to copy. Payloads larger + /// * `max_recv_len` - Maximum readable payload size. Payloads larger /// than this return [`VirtqError::PayloadTooLarge`]. /// /// # Errors /// /// - [`VirtqError::BadChain`] - Descriptor chain format not recognized /// - [`VirtqError::InvalidState`] - Descriptor ID collision (driver bug) - /// - [`VirtqError::MemoryReadError`] - Failed to read chain payload from shared memory + #[allow(clippy::type_complexity)] pub fn poll( &mut self, max_recv_len: usize, - ) -> Result)>, VirtqError> { + ) -> Result, ReplyChain)>, VirtqError> { let (id, chain) = match self.inner.poll_available() { Ok(x) => x, Err(RingError::WouldBlock) => return Ok(None), @@ -354,20 +498,19 @@ impl VirtqConsumer { )); } - // Copy chain payload from shared memory - let data = match self.read_elements(readables) { - Ok(d) => d, - Err(e) => return Err(self.abort_chain(id, e)), - }; - - let chain = RecvChain { + let chain = RecvChain::new( + self.inner.mem().clone(), token, - segments: data, - }; + readables.iter().copied().collect(), + recv_len, + ); let reply = if !writables.is_empty() { - let mem = self.inner.mem().clone(); - let writable = WritableChain::new(mem, token, writables.iter().copied().collect()); + let writable = WritableChain::new( + self.inner.mem().clone(), + token, + writables.iter().copied().collect(), + ); ReplyChain::Writable(writable) } else { let ack = AckChain::new(token); @@ -377,14 +520,29 @@ impl VirtqConsumer { Ok(Some((chain, reply))) } - /// Submit a reply/ack for a received chain back to the ring. + /// Submit both halves of a received chain back to the ring. + /// + /// Consuming the [`RecvChain`] prevents further reads once its descriptors + /// can be reused by the producer. `reply` accepts both [`WritableChain`] + /// (with written byte count) and [`AckChain`] (zero-length) through + /// [`ReplyChain`]. The two halves must have matching tokens. /// - /// Accepts both [`WritableChain`] (with written byte count) and - /// [`AckChain`] (zero-length) via the [`ReplyChain`] enum. - /// Clears the inflight slot and notifies the producer if event - /// suppression allows. - pub fn complete(&mut self, reply: impl Into>) -> Result<(), VirtqError> { + /// A mismatched pair returns [`VirtqError::InvalidState`] without returning + /// either descriptor. This fails closed: the descriptors remain in flight + /// because completing either could invalidate another still-live + /// [`RecvChain`]. + pub fn complete( + &mut self, + recv: impl Into>, + reply: impl Into>, + ) -> Result<(), VirtqError> { + let recv = recv.into(); let reply = reply.into(); + + if recv.token() != reply.token() { + return Err(VirtqError::InvalidState); + } + let id = reply.token().id; let written = u32::try_from(reply.written()).map_err(|_| VirtqError::ReplyTooLarge)?; @@ -475,63 +633,111 @@ impl VirtqConsumer { Ok(()) } - /// Read readable buffer elements from shared memory into `Bytes`. - fn read_elements(&self, elems: &[BufferElement]) -> Result { - let mut segments = SmallVec::<[Bytes; 4]>::new(); - - for elem in elems { - let mut buf = vec![0u8; elem.len as usize]; - self.inner - .mem() - .read(elem.addr, &mut buf) - .map_err(|_| VirtqError::MemoryReadError)?; - segments.push(Bytes::from(buf)); + /// Reset ring and inflight state to initial values. + /// + /// Fails while a polled chain has not yet been completed, preventing a + /// live [`RecvChain`] from reading descriptors after reset and reuse. + /// + /// # Errors + /// + /// - [`VirtqError::InvalidState`] - one or more chains are still in flight + pub fn reset(&mut self) -> Result<(), VirtqError> { + if self.inflight.ones().next().is_some() { + return Err(VirtqError::InvalidState); } - Ok(Segments::from_smallvec(segments)) - } - - /// Reset ring and inflight state to initial values. - pub fn reset(&mut self) { self.inner.reset(); self.inflight.clear(); + Ok(()) } } -fn write_elements( - mem: &M, - elems: &[BufferElement], - offset: usize, - buf: &[u8], -) -> Result { - let capacity: usize = elems.iter().map(|elem| elem.len as usize).sum(); - let mut src = &buf[..buf.len().min(capacity.saturating_sub(offset))]; - let mut written = 0; - let mut skip = offset; +type ChainElems = SmallVec<[BufferElement; 4]>; - for elem in elems { - if src.is_empty() { - break; - } +struct ChainState { + mem: M, + token: Token, + elems: ChainElems, + total: usize, + position: usize, + desc_idx: usize, + desc_off: usize, +} - let elem_len = elem.len as usize; - if skip >= elem_len { - skip -= elem_len; - continue; - } +impl ChainState { + fn new(mem: M, token: Token, elems: ChainElems, total: usize) -> Self { + let mut state = Self { + mem, + token, + elems, + total, + position: 0, + desc_idx: 0, + desc_off: 0, + }; + state.rewind(); + state + } - let elem_offset = skip; - skip = 0; - let n = (elem_len - elem_offset).min(src.len()); - let addr = elem.addr + elem_offset as u64; + #[inline] + fn token(&self) -> Token { + self.token + } - mem.write(addr, &src[..n])?; + #[inline] + fn total(&self) -> usize { + self.total + } + + #[inline] + fn position(&self) -> usize { + self.position + } + + #[inline] + fn remaining(&self) -> usize { + self.total - self.position + } + + #[inline(always)] + fn desc_len(&self) -> usize { + self.elems + .get(self.desc_idx) + .map(|elem| elem.len as usize) + .unwrap_or(0) + } - written += n; - src = &src[n..]; + #[inline(always)] + fn desc_offset(&self) -> usize { + self.desc_off } - Ok(written) + #[inline(always)] + fn current_elem(&self) -> Option { + self.elems.get(self.desc_idx).copied() + } + + #[inline(always)] + fn advance(&mut self, len: usize) { + debug_assert!(len <= self.desc_len() - self.desc_off); + self.desc_off += len; + self.position += len; + + while self.current_elem().is_some() && self.desc_off == self.desc_len() { + self.desc_idx += 1; + self.desc_off = 0; + } + } + + fn rewind(&mut self) { + self.position = 0; + self.desc_off = 0; + self.desc_idx = self + .elems + .iter() + .position(|elem| elem.len != 0) + .unwrap_or(self.elems.len()); + } } impl From> for ReplyChain { @@ -549,15 +755,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 +817,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 +836,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 +852,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 +861,78 @@ mod tests { wc.write_all(b"response").unwrap(); assert_eq!(wc.written(), 8); assert_eq!(wc.remaining(), 56); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); } else { panic!("expected Writable reply for recv+reply chain"); } } + #[test] + fn test_recv_reads_across_descriptor_boundaries() { + let ring = make_ring(16); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + let mut se = producer.chain().readable(4).readable(4).build().unwrap(); + se.write_all(b"abcdefgh").unwrap(); + producer.submit(se).unwrap(); + + let (mut recv, reply) = poll_data(&mut consumer); + assert_eq!(recv.len(), 8); + assert_eq!(recv.remaining(), 8); + let mut first = [0u8; 2]; + recv.read_exact(&mut first).unwrap(); + assert_eq!(&first, b"ab"); + assert_eq!(recv.consumed(), 2); + assert_eq!(recv.remaining(), 6); + + let mut second = [0u8; 3]; + recv.read_exact(&mut second).unwrap(); + assert_eq!(&second, b"cde"); + + let mut too_long = [0u8; 4]; + assert!(matches!( + recv.read_exact(&mut too_long), + Err(VirtqError::ReceiveTooShort { + requested: 4, + remaining: 3 + }) + )); + + let mut final_buf = [0u8; 4]; + assert_eq!(recv.read(&mut final_buf).unwrap(), 3); + assert_eq!(&final_buf[..3], b"fgh"); + assert_eq!(recv.read(&mut final_buf).unwrap(), 0); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"abcdefgh"); + + consumer.complete(recv, reply).unwrap(); + } + + #[test] + fn test_poll_defers_payload_reads() { + let ring = make_ring(16); + let mem = ring.mem(); + let mut ring_producer = make_producer(&ring); + let payload_addr = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + mem.write(payload_addr, b"data").unwrap(); + + let chain = BufferChainBuilder::new() + .readable(payload_addr, 4) + .build() + .unwrap(); + ring_producer.submit_available(&chain).unwrap(); + + let guarded_mem = FailingPayloadReadMem { + inner: mem, + payload_addr, + payload_len: 4, + }; + let mut consumer = VirtqConsumer::new(ring.layout(), guarded_mem, TestNotifier::new()); + + let (recv, reply) = consumer.poll(4).unwrap().unwrap(); + assert!(matches!(recv.to_bytes(), Err(VirtqError::MemoryReadError))); + consumer.complete(recv, reply).unwrap(); + } + #[test] fn test_writable_partial_write() { let ring = make_ring(16); @@ -625,13 +941,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"); } @@ -692,8 +1008,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 +1060,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 +1080,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 +1090,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 +1112,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 +1120,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 +1134,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 +1181,43 @@ mod tests { let se2 = producer.chain().writable(16).build().unwrap(); producer.submit(se2).unwrap(); - let (_e1, c1) = poll_data(&mut consumer); - let (_e2, c2) = poll_data(&mut consumer); + let (e1, c1) = poll_data(&mut consumer); + let (e2, c2) = poll_data(&mut consumer); // Complete in reverse order - consumer.complete(c2).unwrap(); - consumer.complete(c1).unwrap(); + consumer.complete(e2, c2).unwrap(); + consumer.complete(e1, c1).unwrap(); } #[test] - fn test_recv_into_bytes() { + fn test_mismatched_completion_fails_closed() { + let ring = make_ring(16); + let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); + + let mut first = producer.chain().readable(1).build().unwrap(); + first.write_all(b"a").unwrap(); + producer.submit(first).unwrap(); + + let mut second = producer.chain().readable(1).build().unwrap(); + second.write_all(b"b").unwrap(); + producer.submit(second).unwrap(); + + let (recv1, reply1) = poll_data(&mut consumer); + let (recv2, reply2) = poll_data(&mut consumer); + + assert!(matches!( + consumer.complete(recv1, reply2), + Err(VirtqError::InvalidState) + )); + assert_eq!(consumer.inflight.count_ones(..), 2); + assert!(producer.poll().unwrap().is_none()); + assert!(matches!(consumer.reset(), Err(VirtqError::InvalidState))); + + drop((recv2, reply1)); + } + + #[test] + fn test_recv_to_bytes_preserves_reader_position() { let ring = make_ring(16); let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); @@ -845,10 +1225,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 +1244,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,13 +1268,13 @@ mod tests { let se2 = producer.chain().writable(16).build().unwrap(); producer.submit(se2).unwrap(); - let (_e1, c1) = poll_data(&mut consumer); - let (_e2, c2) = poll_data(&mut consumer); + let (e1, c1) = poll_data(&mut consumer); + let (e2, c2) = poll_data(&mut consumer); // Complete both before reset - consumer.complete(c1).unwrap(); - consumer.complete(c2).unwrap(); + consumer.complete(e1, c1).unwrap(); + consumer.complete(e2, c2).unwrap(); - consumer.reset(); + consumer.reset().unwrap(); assert_eq!(consumer.inflight.count_ones(..), 0); assert_eq!(consumer.inner.num_inflight(), 0); diff --git a/src/hyperlight_common/src/virtq/mod.rs b/src/hyperlight_common/src/virtq/mod.rs index e676e7cc7..895bc6dc6 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)?; //! } //! ``` //! @@ -204,6 +205,8 @@ pub enum VirtqError { BadChain, #[error("Payload data too large: received {recv} bytes, limit {limit} bytes")] PayloadTooLarge { recv: usize, limit: usize }, + #[error("Receive data too short: requested {requested} bytes, only {remaining} bytes remain")] + ReceiveTooShort { requested: usize, remaining: usize }, #[error("Reply data too large for allocated buffer")] ReplyTooLarge, #[error("Internal state error")] @@ -601,7 +604,7 @@ mod tests { fn poll_received( consumer: &mut VirtqConsumer, - ) -> (RecvChain, ReplyChain) { + ) -> (RecvChain, ReplyChain) { consumer.poll(1024).unwrap().unwrap() } @@ -630,8 +633,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 +666,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 +752,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 +787,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 +810,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 +832,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 +851,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 +875,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(); @@ -937,8 +940,8 @@ mod tests { // Consumer acks all entries while let Some(result) = consumer.poll(1024).unwrap() { - let (_, reply) = result; - consumer.complete(reply).unwrap(); + let (recv, reply) = result; + consumer.complete(recv, reply).unwrap(); } // Reclaim should free ring slots without losing data @@ -958,12 +961,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 +989,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 +1024,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 +1057,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 +1073,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 +1103,8 @@ mod tests { // Consumer acks all while let Some(result) = consumer.poll(1024).unwrap() { - let (_, reply) = result; - consumer.complete(reply).unwrap(); + let (recv, reply) = result; + consumer.complete(recv, reply).unwrap(); } // Reclaim frees ring slots; empty completions are discarded diff --git a/src/hyperlight_common/src/virtq/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index 1d3750e6c..ad6cb7e4c 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -752,10 +752,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 +777,26 @@ impl SendChain { Inflight { token, chain } } - /// Number of producer-written readable segments in this chain. - pub fn segment_count(&self) -> usize { + /// Number of producer-written readable descriptors in this chain. + #[inline] + pub fn desc_count(&self) -> usize { self.chain().readables().len() } /// Total producer-written readable capacity in bytes. + #[inline] pub fn capacity(&self) -> usize { self.rd_capacity } /// Number of producer-written readable bytes written so far. + #[inline] pub fn written(&self) -> usize { self.rd_written } /// Remaining producer-written readable capacity. + #[inline] pub fn remaining(&self) -> usize { self.capacity() - self.written() } @@ -800,14 +806,15 @@ impl SendChain { /// Appends at the current aggregate write position and scatters across /// readable segments in chain order. Uses [`MemOps::write`] (volatile on /// host side). If `buf` is larger than the remaining capacity, writes as - /// many bytes as will fit. + /// many bytes as will fit. If a later memory write fails, the cursor and + /// written length retain any earlier chunks written by the same call. /// /// # Errors /// /// - [`VirtqError::NoPayloadSegment`] - no readable buffer allocated /// - [`VirtqError::MemoryWriteError`] - underlying write failed pub fn write(&mut self, buf: &[u8]) -> Result { - if self.segment_count() == 0 { + if self.desc_count() == 0 { return Err(VirtqError::NoPayloadSegment); } @@ -816,40 +823,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 +865,9 @@ impl SendChain { /// - [`VirtqError::PayloadTooLarge`] - buf exceeds remaining capacity /// - [`VirtqError::NoPayloadSegment`] - no readable buffer allocated /// - [`VirtqError::MemoryWriteError`] - underlying write failed + #[inline] pub fn write_all(&mut self, buf: &[u8]) -> Result<&mut Self, VirtqError> { - if self.segment_count() == 0 { + if self.desc_count() == 0 { return Err(VirtqError::NoPayloadSegment); } @@ -1008,7 +1010,7 @@ mod tests { fn poll_received( consumer: &mut VirtqConsumer, - ) -> (RecvChain, ReplyChain) { + ) -> (RecvChain, ReplyChain) { consumer.poll(1024).unwrap().unwrap() } @@ -1062,7 +1064,7 @@ mod tests { let (producer, _consumer, _notifier) = make_test_producer(&ring); let se = producer.chain().readable(16).writable(32).build().unwrap(); - assert_eq!(se.segment_count(), 1); + assert_eq!(se.desc_count(), 1); assert_eq!(se.capacity(), 16); } @@ -1085,11 +1087,38 @@ mod tests { let token = producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); assert_eq!(recv.token(), token); - assert_eq!(recv.to_bytes().as_ref(), b"hello world"); - assert_eq!(recv.segments().segment_count(), 2); - assert_eq!(recv.segments().as_slice()[0].as_ref(), b"hello"); - assert_eq!(recv.segments().as_slice()[1].as_ref(), b" world"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello world"); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.as_slice()[0].as_ref(), b"hello"); + assert_eq!(segments.as_slice()[1].as_ref(), b" world"); + consumer.complete(recv, reply).unwrap(); + } + + #[test] + fn test_chain_multi_readable_appends_across_calls() { + let ring = make_ring(16); + let layout = ring.layout(); + let mem = ring.mem(); + let pool_base = mem.base_addr() + Layout::query_size(ring.len()) as u64 + 0x100; + let pool = TestPool::new_with_max_alloc_len(pool_base, 0x8000, 4); + let notifier = TestNotifier::new(); + let mut producer = VirtqProducer::new(layout, mem.clone(), notifier.clone(), pool); + let mut consumer = VirtqConsumer::new(layout, mem, notifier); + + let mut send = producer.chain().readable(8).build().unwrap(); + send.write_all(b"abc").unwrap(); + send.write_all(b"def").unwrap(); + assert_eq!(send.written(), 6); + assert_eq!(send.remaining(), 2); + + producer.submit(send).unwrap(); + let (recv, reply) = poll_received(&mut consumer); + let segments = recv.to_segments().unwrap(); + assert_eq!(segments.segment_count(), 2); + assert_eq!(segments.as_slice()[0].as_ref(), b"abcd"); + assert_eq!(segments.as_slice()[1].as_ref(), b"ef"); + consumer.complete(recv, reply).unwrap(); } #[test] @@ -1105,7 +1134,7 @@ mod tests { let mut se = producer.chain().readable(10).writable(32).build().unwrap(); - assert_eq!(se.segment_count(), 3); + assert_eq!(se.desc_count(), 3); assert_eq!(se.capacity(), 10); se.write_all(b"abcdefghij").unwrap(); @@ -1114,12 +1143,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] @@ -1147,14 +1177,14 @@ mod tests { let se = producer.chain().writable(10).build().unwrap(); let token = producer.submit(se).unwrap(); - let (_recv, reply) = poll_received(&mut consumer); + let (recv, reply) = poll_received(&mut consumer); let ReplyChain::Writable(mut wc) = reply else { panic!("expected writable reply"); }; assert_eq!(wc.capacity(), 10); wc.write_all(b"abcdefghij").unwrap(); - consumer.complete(wc).unwrap(); + consumer.complete(recv, wc).unwrap(); let used = producer.poll().unwrap().unwrap(); assert_eq!(used.token(), token); @@ -1186,8 +1216,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 +1239,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 +1256,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 +1270,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 +1296,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 +1320,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 +1374,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 +1394,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 +1410,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 +1429,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 +1449,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] @@ -1444,7 +1476,7 @@ mod tests { let producer = VirtqProducer::new(layout, mem, notifier, pool); let mut se = producer.chain().readable(8).build().unwrap(); - assert_eq!(se.segment_count(), 2); + assert_eq!(se.desc_count(), 2); assert!(matches!( se.with_seg(2, |_| Ok::(0)), Err(VirtqError::NoPayloadSegment) @@ -1588,12 +1620,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] @@ -1642,11 +1674,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 +1700,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 +1718,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 +1747,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"); } @@ -1776,8 +1808,8 @@ mod tests { producer.submit(se).unwrap(); let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().as_ref(), b"hello"); - consumer.complete(reply).unwrap(); + assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello"); + consumer.complete(recv, reply).unwrap(); let _ = producer.poll().unwrap().unwrap(); // Now reset From e1342f2d7c5db8afe9df5e93a7e07182a3285c16 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Wed, 22 Jul 2026 14:52:30 +0200 Subject: [PATCH 04/34] refactor(virtq): remove reset api and harden allocation rollback Make pool restoration transactional. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/virtq/buffer.rs | 47 +----- src/hyperlight_common/src/virtq/pool.rs | 79 ++------- src/hyperlight_common/src/virtq/producer.rs | 169 ++++++++------------ 3 files changed, 85 insertions(+), 210 deletions(-) diff --git a/src/hyperlight_common/src/virtq/buffer.rs b/src/hyperlight_common/src/virtq/buffer.rs index e93cebda1..3b8a49c06 100644 --- a/src/hyperlight_common/src/virtq/buffer.rs +++ b/src/hyperlight_common/src/virtq/buffer.rs @@ -66,9 +66,6 @@ pub trait BufferProvider { /// Free a previously allocated segment by start address. fn dealloc(&self, addr: u64) -> Result<(), AllocError>; - /// Reset the pool to initial state. - fn reset(&self) {} - /// Allocate scatter/gather segments for a logical payload of `total_len` bytes. fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { if total_len == 0 { @@ -114,9 +111,6 @@ impl BufferProvider for Rc { fn dealloc(&self, addr: u64) -> Result<(), AllocError> { (**self).dealloc(addr) } - fn reset(&self) { - (**self).reset() - } fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { (**self).alloc_sg(total_len) } @@ -132,9 +126,6 @@ impl BufferProvider for Arc { fn dealloc(&self, addr: u64) -> Result<(), AllocError> { (**self).dealloc(addr) } - fn reset(&self) { - (**self).reset() - } fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { (**self).alloc_sg(total_len) } @@ -350,16 +341,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, } @@ -367,9 +350,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. @@ -379,35 +360,15 @@ impl OwnedAlloc

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

{ fn drop(&mut self) { - if let Some(Inner { pool, alloc }) = self.inner.take() { - let result = pool.dealloc(alloc.addr); - debug_assert!(result.is_ok(), "OwnedAlloc drop dealloc failed: {result:?}"); - } + let result = self.pool.dealloc(self.alloc.addr); + debug_assert!(result.is_ok(), "OwnedAlloc drop dealloc failed: {result:?}"); } } diff --git a/src/hyperlight_common/src/virtq/pool.rs b/src/hyperlight_common/src/virtq/pool.rs index 3f0acc963..f54af6bb5 100644 --- a/src/hyperlight_common/src/virtq/pool.rs +++ b/src/hyperlight_common/src/virtq/pool.rs @@ -246,12 +246,6 @@ impl Slab { fn contains(&self, addr: u64) -> bool { self.range().contains(&addr) } - - fn reset(&mut self) { - self.used_slots.clear(); - self.run_starts.clear(); - self.last_free_run = None; - } } #[cfg(test)] @@ -404,12 +398,6 @@ impl BufferProvider for BufferPool { fn dealloc(&self, addr: u64) -> Result<(), AllocError> { self.inner.borrow_mut().dealloc_addr(addr) } - - fn reset(&self) { - let mut inner = self.inner.borrow_mut(); - inner.lower.reset(); - inner.upper.reset(); - } } impl BufferPool { @@ -557,26 +545,21 @@ impl RecycleList { /// Rebuild state so that exactly the addresses in `allocated` are marked /// live and every other slot is free. /// - /// On error the pool is left in an indeterminate state and should be - /// [`reset`](Self::reset) before reuse. + /// On error the pool is left unchanged. fn restore_allocated(&mut self, allocated: &[u64]) -> Result<(), AllocError> { - self.allocated.clear(); + let mut restored = FixedBitSet::with_capacity(self.count); for &addr in allocated { let slot = self.slot_of(addr)?; - if self.allocated.contains(slot) { + if restored.contains(slot) { return Err(AllocError::InvalidFree(addr, self.slot_size)); } - self.allocated.insert(slot); + restored.insert(slot); } + self.allocated = restored; self.rebuild_free(); Ok(()) } - fn reset(&mut self) { - self.allocated.clear(); - self.rebuild_free(); - } - /// Repopulate the free list with every slot whose allocated bit is clear. fn rebuild_free(&mut self) { self.free.clear(); @@ -632,7 +615,8 @@ impl RecyclePool { } /// Rebuild pool state so that every address in `allocated` is removed from - /// the free list, matching externally known inflight state. + /// the free list, matching externally known inflight state. Validation is + /// transactional: an error leaves the existing pool state unchanged. pub fn restore_allocated(&self, allocated: &[u64]) -> Result<(), AllocError> { self.inner.borrow_mut().restore_allocated(allocated) } @@ -687,10 +671,6 @@ impl BufferProvider for RecyclePool { fn dealloc(&self, addr: u64) -> Result<(), AllocError> { self.inner.borrow_mut().dealloc_addr(addr) } - - fn reset(&self) { - self.inner.borrow_mut().reset() - } } #[cfg(test)] @@ -872,40 +852,6 @@ mod tests { assert!(pool.inner.borrow().upper.contains(alloc2.addr)); } - #[test] - fn test_buffer_pool_reset_returns_to_initial_state() { - let pool = make_pool::<256, 4096>(0x20000); - - // Allocate from both tiers - let a1 = pool.inner.borrow_mut().alloc(128).unwrap(); - let a2 = pool.inner.borrow_mut().alloc(4096).unwrap(); - assert!(a1.len > 0); - assert!(a2.len > 0); - - pool.reset(); - - let inner = pool.inner.borrow(); - assert_eq!(inner.lower.free_bytes(), inner.lower.capacity()); - assert_eq!(inner.upper.free_bytes(), inner.upper.capacity()); - } - - #[test] - fn test_buffer_pool_reset_allows_reallocation() { - let pool = make_pool::<256, 4096>(0x20000); - - // Fill up some allocations - let mut allocs = Vec::new(); - for _ in 0..5 { - allocs.push(pool.inner.borrow_mut().alloc(256).unwrap()); - } - - pool.reset(); - - // Should be able to allocate as if fresh - let a = pool.inner.borrow_mut().alloc(256).unwrap(); - assert!(a.len > 0); - } - #[test] fn test_pool_dealloc_addr_routes_to_correct_tier() { let pool = make_pool::<256, 4096>(0x20000); @@ -983,8 +929,13 @@ mod tests { #[test] fn test_recycle_pool_restore_allocated_invalid_addr_returns_error() { let pool = make_recycle_pool(4, 4096); - let result = pool.restore_allocated(&[0xDEAD]); + pool.restore_allocated(&[0x80000]).unwrap(); + + let result = pool.restore_allocated(&[0x81000, 0xDEAD]); assert!(result.is_err()); + assert_eq!(pool.num_free(), 3); + assert_eq!(pool.allocation_len(0x80000).unwrap(), 4096); + assert!(pool.allocation_len(0x81000).is_err()); } #[test] @@ -1018,15 +969,13 @@ mod tests { } #[test] - fn test_recycle_pool_restore_allocated_resets_first() { + fn test_recycle_pool_restore_allocated_replaces_state() { let pool = make_recycle_pool(4, 4096); - // Allocate some slots let _ = pool.alloc(4096).unwrap(); let _ = pool.alloc(4096).unwrap(); assert_eq!(pool.num_free(), 2); - // restore_allocated resets then removes - so 4 - 1 = 3 pool.restore_allocated(&[0x80000]).unwrap(); assert_eq!(pool.num_free(), 3); } diff --git a/src/hyperlight_common/src/virtq/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index ad6cb7e4c..635be0816 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -162,7 +162,11 @@ where } } - fn dealloc_elems( + /// Retire allocations from a completed descriptor chain. + /// + /// Every element is attempted so one deallocation failure does not strand + /// later allocations; the first failure is returned after cleanup. + fn retire_elems( &self, elems: impl IntoIterator, ) -> Result<(), VirtqError> { @@ -302,39 +306,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 @@ -406,7 +377,7 @@ where let written = used.len as usize; let Inflight { token, chain } = inf; - self.dealloc_elems(chain.readables().iter().copied())?; + self.retire_elems(chain.readables().iter().copied())?; let used = if chain.writables().is_empty() { UsedChain::Ack(token) @@ -439,14 +410,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); } } @@ -469,7 +440,7 @@ where sgs.push(Bytes::from_owner(owner)); } - self.dealloc_elems(free)?; + self.retire_elems(free)?; Ok(Segments::from_smallvec(sgs)) } @@ -608,7 +579,8 @@ impl ChainBuilder { return Err(VirtqError::InvalidState); } - let mut rollback = Rollback::new(&self.pool); + let rd_capacity = self.rd_caps.iter().sum(); + let mut allocs = AllocTxn::new(&self.pool); let mut rd_caps = SmallVec::<[usize; 4]>::new(); let mut rd_elems = SmallVec::<[BufferElement; 4]>::new(); let mut wr_elems = SmallVec::<[BufferElement; 4]>::new(); @@ -617,7 +589,7 @@ impl ChainBuilder { // The buffer element lengths are initialized to zero and updated as the // `SendChain` writes. for &cap in &self.rd_caps { - let sgs = self.pool.alloc_sg(cap)?; + let sgs = allocs.alloc_sg(cap)?; let mut remaining = cap; for alloc in sgs { @@ -631,7 +603,6 @@ impl ChainBuilder { writable: false, }); remaining -= seg_cap; - rollback.allocs.push(alloc); } if remaining != 0 { @@ -643,7 +614,7 @@ impl ChainBuilder { // Writable buffer elements are initialized with their full capacity for the device to // write into. for &cap in &self.wr_caps { - let sgs = self.pool.alloc_sg(cap)?; + let sgs = allocs.alloc_sg(cap)?; for alloc in sgs { let len = checked_descriptor_len(alloc.len)?; wr_elems.push(BufferElement { @@ -651,7 +622,6 @@ impl ChainBuilder { len, writable: true, }); - rollback.allocs.push(alloc); } } @@ -660,43 +630,61 @@ impl ChainBuilder { .writables(wr_elems) .build()?; - rollback.release(); + allocs.commit(); Ok(SendChain { mem: self.mem, pool: self.pool, chain: Some(chain), rd_caps, - rd_capacity: self.rd_caps.iter().sum(), + rd_capacity, rd_written: 0, write_mode: WriteMode::Unset, }) } } -struct Rollback<'a, P: BufferProvider> { +/// Build-scoped allocation transaction. +/// +/// `SendChain` and `Inflight` intentionally retain lightweight descriptor +/// metadata instead of one allocation guard and cloned pool handle per +/// descriptor. While a valid `BufferChain` is being built, this transaction +/// provides aggregate RAII: it records every allocated address before the +/// caller can perform fallible validation and returns them all on drop. +/// [`commit`](Self::commit) disarms rollback once `SendChain` can take +/// responsibility for reclaiming the completed chain. +struct AllocTxn<'a, P: BufferProvider> { pool: &'a P, - allocs: SmallVec<[Allocation; 8]>, + addrs: SmallVec<[u64; 8]>, } -impl<'a, P: BufferProvider> Rollback<'a, P> { +impl<'a, P: BufferProvider> AllocTxn<'a, P> { fn new(pool: &'a P) -> Self { Self { pool, - allocs: SmallVec::new(), + addrs: SmallVec::new(), } } - fn release(mut self) { - self.allocs.clear(); + fn alloc_sg(&mut self, total_len: usize) -> Result, AllocError> { + let allocs = self.pool.alloc_sg(total_len)?; + self.addrs.extend(allocs.iter().map(|alloc| alloc.addr)); + Ok(allocs) + } + + fn commit(mut self) { + self.addrs.clear(); } } -impl Drop for Rollback<'_, P> { +impl Drop for AllocTxn<'_, P> { fn drop(&mut self) { - for alloc in self.allocs.drain(..) { - let result = self.pool.dealloc(alloc.addr); - debug_assert!(result.is_ok(), "rollback dealloc failed: {result:?}"); + for addr in self.addrs.drain(..) { + let result = self.pool.dealloc(addr); + debug_assert!( + result.is_ok(), + "allocation rollback dealloc failed: {result:?}" + ); } } } @@ -1557,6 +1545,30 @@ mod tests { assert!(tok.id < 16); } + #[cfg(target_pointer_width = "64")] + #[test] + fn test_chain_build_rolls_back_unrepresentable_allocations() { + let ring = make_ring(16); + let slot_size = u32::MAX as usize + 1; + let pool = RecyclePool::new(0, slot_size, slot_size).unwrap(); + let mem = ring.mem(); + let producer = VirtqProducer::new(ring.layout(), mem, TestNotifier::new(), pool.clone()); + + assert!(matches!( + producer.chain().readable(1).build(), + Err(VirtqError::PayloadTooLarge { recv, limit }) + if recv == slot_size && limit == u32::MAX as usize + )); + assert_eq!(pool.num_free(), 1); + + assert!(matches!( + producer.chain().writable(1).build(), + Err(VirtqError::PayloadTooLarge { recv, limit }) + if recv == slot_size && limit == u32::MAX as usize + )); + assert_eq!(pool.num_free(), 1); + } + #[test] fn test_submit_notifies() { let ring = make_ring(16); @@ -1797,51 +1809,4 @@ mod tests { assert_eq!(producer.inner.num_inflight(), 1); } - #[test] - fn test_virtq_producer_reset() { - let ring = make_ring(16); - let (mut producer, mut consumer, _notifier) = make_test_producer(&ring); - - // Submit and complete a round trip - let mut se = producer.chain().readable(32).writable(64).build().unwrap(); - se.write_all(b"hello").unwrap(); - producer.submit(se).unwrap(); - - let (recv, reply) = poll_received(&mut consumer); - assert_eq!(recv.to_bytes().unwrap().as_ref(), b"hello"); - consumer.complete(recv, reply).unwrap(); - let _ = producer.poll().unwrap().unwrap(); - - // Now reset - // SAFETY: the used chain was dropped before reset and no peer can - // access the reset test ring concurrently. - unsafe { - producer.reset(); - } - - // All inflight slots should be cleared - assert_eq!(producer.inner.num_inflight(), 0); - // Ring state should be back to initial - assert_eq!(producer.inner.num_free(), producer.inner.len()); - } - - #[test] - fn test_virtq_producer_reset_clears_inflight() { - let ring = make_ring(16); - let (mut producer, _consumer, _notifier) = make_test_producer(&ring); - - // Submit without completing - let se = producer.chain().writable(64).build().unwrap(); - producer.submit(se).unwrap(); - - assert_eq!(producer.inner.num_inflight(), 1); - - // SAFETY: no peer can access the reset test ring concurrently. - unsafe { - producer.reset(); - } - - assert_eq!(producer.inner.num_inflight(), 0); - assert_eq!(producer.inner.num_free(), producer.inner.len()); - } } From 4b23b8e15241d13e4db4d567784d6ed65752799c Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Wed, 22 Jul 2026 17:53:39 +0200 Subject: [PATCH 05/34] feat(virtq): add tiered fixed slot allocation Split and rename the pool implementations, add explicit lower/upper SlotPool regions. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/benches/buffer_pool.rs | 38 +- src/hyperlight_common/benches/common/mod.rs | 23 +- src/hyperlight_common/benches/virtq_api.rs | 32 +- src/hyperlight_common/src/virtq/buffer.rs | 111 +- .../src/virtq/concurrency.rs | 10 +- src/hyperlight_common/src/virtq/pool.rs | 1345 ++--------------- src/hyperlight_common/src/virtq/pool/fuzz.rs | 382 +++++ src/hyperlight_common/src/virtq/pool/run.rs | 379 +++++ src/hyperlight_common/src/virtq/pool/slot.rs | 391 +++++ src/hyperlight_common/src/virtq/pool/tests.rs | 515 +++++++ src/hyperlight_common/src/virtq/producer.rs | 3 +- 11 files changed, 1848 insertions(+), 1381 deletions(-) create mode 100644 src/hyperlight_common/src/virtq/pool/fuzz.rs create mode 100644 src/hyperlight_common/src/virtq/pool/run.rs create mode 100644 src/hyperlight_common/src/virtq/pool/slot.rs create mode 100644 src/hyperlight_common/src/virtq/pool/tests.rs diff --git a/src/hyperlight_common/benches/buffer_pool.rs b/src/hyperlight_common/benches/buffer_pool.rs index 31c0bb566..5b9ef47a6 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,7 +149,7 @@ fn bench_segmented_payload(c: &mut Criterion) { BenchmarkId::from_parameter(payload_size), &payload_size, |b, &payload_size| { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); b.iter(|| { let sgs = pool.alloc_sg(black_box(payload_size)).unwrap(); for sg in sgs { @@ -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,7 +185,8 @@ fn bench_recycle_pool(c: &mut Criterion) { }); group.bench_function("alloc_dealloc_1500", |b| { - let pool = RecyclePool::new(0x80000, 4 * 1024 * 1024, 4096).unwrap(); + let layout = SlotLayout::new(0x80000, 4096, 1024); + let pool = SlotPool::new(layout).unwrap(); b.iter(|| { let alloc = pool.alloc(black_box(1500)).unwrap(); pool.dealloc(alloc.addr).unwrap(); @@ -191,7 +194,8 @@ fn bench_recycle_pool(c: &mut Criterion) { }); group.bench_function("alloc_sg_64k", |b| { - let pool = RecyclePool::new(0x80000, 4 * 1024 * 1024, 4096).unwrap(); + let layout = SlotLayout::new(0x80000, 4096, 1024); + let pool = SlotPool::new(layout).unwrap(); b.iter(|| { let sgs = pool.alloc_sg(black_box(64 * 1024)).unwrap(); for sg in sgs { @@ -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 7461de6fa..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()) 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/virtq/buffer.rs b/src/hyperlight_common/src/virtq/buffer.rs index 3b8a49c06..fd5570ba1 100644 --- a/src/hyperlight_common/src/virtq/buffer.rs +++ b/src/hyperlight_common/src/virtq/buffer.rs @@ -14,122 +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>; - - /// Allocate scatter/gather segments for a logical payload of `total_len` bytes. - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - if total_len == 0 { - return Err(AllocError::InvalidArg); - } - - let seg_cap = self.max_alloc_len(); - if seg_cap == 0 { - return Err(AllocError::InvalidArg); - } - - let mut rem = total_len; - let mut sgs = SmallVec::<[Allocation; 4]>::new(); - - while rem > 0 { - let len = rem.min(seg_cap); - match self.alloc(len) { - Ok(alloc) => { - sgs.push(alloc); - rem -= len; - } - Err(err) => { - for sg in sgs { - let _res = self.dealloc(sg.addr); - debug_assert!(_res.is_ok(), "dealloc failed: {_res:?}"); - } - return Err(err); - } - } - } - - Ok(sgs) - } -} - -impl BufferProvider for Rc { - fn max_alloc_len(&self) -> usize { - (**self).max_alloc_len() - } - fn alloc(&self, len: usize) -> Result { - (**self).alloc(len) - } - fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - (**self).dealloc(addr) - } - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - (**self).alloc_sg(total_len) - } -} - -impl BufferProvider for Arc { - fn max_alloc_len(&self) -> usize { - (**self).max_alloc_len() - } - fn alloc(&self, len: usize) -> Result { - (**self).alloc(len) - } - fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - (**self).dealloc(addr) - } - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - (**self).alloc_sg(total_len) - } -} +use super::pool::{AllocError, Allocation, BufferProvider}; /// Ordered byte segments that make up one virtqueue payload. /// diff --git a/src/hyperlight_common/src/virtq/concurrency.rs b/src/hyperlight_common/src/virtq/concurrency.rs index 33a369d5d..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); @@ -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); @@ -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); @@ -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/pool.rs b/src/hyperlight_common/src/virtq/pool.rs index f54af6bb5..76f6416e1 100644 --- a/src/hyperlight_common/src/virtq/pool.rs +++ b/src/hyperlight_common/src/virtq/pool.rs @@ -13,1288 +13,189 @@ 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; +mod run; +mod slot; - if alloc.addr < run_end && run.addr < new_end { - self.last_free_run = None; - } - } +pub use run::RunPool; +#[cfg(all(test, loom))] +pub use run::RunPoolSync; +pub use slot::{SlotLayout, SlotPool}; + +/// Buffer allocation failure. +#[derive(Debug, Error, Copy, Clone)] +pub enum AllocError { + /// A region does not meet its required alignment. + #[error("Invalid region addr {0}")] + InvalidAlign(u64), + /// An address does not identify a live allocation. + #[error("Invalid free addr {0} and size {1}")] + InvalidFree(u64, usize), + /// An argument is zero or otherwise invalid. + #[error("Invalid argument")] + InvalidArg, + /// A region cannot hold any allocation. + #[error("Empty region")] + EmptyRegion, + /// No currently free allocation can satisfy the request. + #[error("No space available")] + NoSpace, + /// The request exceeds the pool's allocation capacity. + #[error("Requested size exceeds pool capacity")] + OutOfMemory, + /// Address or size arithmetic overflowed. + #[error("Overflow")] + Overflow, +} + +/// One allocation returned by a [`BufferProvider`]. +#[derive(Debug, Clone, Copy)] +pub struct Allocation { + /// Starting address of the allocation. + pub addr: u64, + /// Capacity in bytes, rounded up according to the provider's policy. + pub len: usize, +} + +/// Allocates and reclaims virtqueue payload buffers. +pub trait BufferProvider { + /// Preferred maximum size of one allocation segment. + fn max_alloc_len(&self) -> usize { + usize::MAX } - fn find_slots(&mut self, slots_num: usize) -> Option { - debug_assert!(slots_num > 0); + /// Allocate one buffer that can hold at least `len` bytes. + fn alloc(&self, len: usize) -> Result; - if let Some(alloc) = self.last_free_run - && alloc.len >= slots_num * N - { - let pos = self.slot_of(alloc.addr); - let _ = self.last_free_run.take(); - return Some(pos); - } - - let total = self.used_slots.len(); - self.used_slots.zeroes().find(|&next_free| { - let end = next_free + slots_num; - end <= total && self.used_slots.count_zeroes(next_free..end) == slots_num - }) - } + /// Free a previously allocated segment by start address. + fn dealloc(&self, addr: u64) -> Result<(), AllocError>; - fn alloc(&mut self, len: usize) -> Result { - if len == 0 { + /// Allocate scatter/gather segments for a logical payload of `total_len` bytes. + fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { + if total_len == 0 { return Err(AllocError::InvalidArg); } - let total = self.used_slots.len(); - let need_slots = len.div_ceil(N); - if need_slots > total { - return Err(AllocError::OutOfMemory); + let seg_cap = self.max_alloc_len(); + if seg_cap == 0 { + return Err(AllocError::InvalidArg); } - let idx = self.find_slots(need_slots).ok_or(AllocError::NoSpace)?; - self.used_slots.insert_range(idx..idx + need_slots); - self.run_starts.insert(idx); - let addr = self.addr_of(idx).ok_or(AllocError::Overflow)?; - - let alloc = Allocation { - addr, - len: need_slots * N, - }; - - self.maybe_invalidate_last_run(alloc); - Ok(alloc) - } - - fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { - let start = self.checked_slot_of(addr, 0)?; - let run_slots = self - .live_run_slots_at(start) - .ok_or(AllocError::InvalidFree(addr, 0))?; - self.dealloc_run(start, run_slots, addr) - } - - fn dealloc_run(&mut self, start: usize, run_slots: usize, addr: u64) -> Result<(), AllocError> { - let len = run_slots * N; - self.used_slots.remove_range(start..start + run_slots); - self.run_starts.set(start, false); - self.last_free_run = Some(Allocation { addr, len }); - Ok(()) - } - - fn allocation_len(&self, addr: u64) -> Result { - let start = self.checked_slot_of(addr, 0)?; - let run_slots = self - .live_run_slots_at(start) - .ok_or(AllocError::InvalidFree(addr, 0))?; - Ok(run_slots * N) - } - - fn capacity(&self) -> usize { - self.used_slots.len() * N - } - - fn range(&self) -> core::ops::Range { - self.base_addr..self.base_addr + self.capacity() as u64 - } - - fn contains(&self, addr: u64) -> bool { - self.range().contains(&addr) - } -} - -#[cfg(test)] -impl Slab { - fn free_bytes(&self) -> usize { - (self.used_slots.len() - self.used_slots.count_ones(..)) * N - } -} - -#[inline] -fn align_up(val: usize, align: usize) -> Result { - if align == 0 { - return Err(AllocError::InvalidArg); - } - - val.checked_next_multiple_of(align) - .ok_or(AllocError::Overflow) -} - -#[derive(Debug)] -struct Inner { - lower: Slab, - upper: Slab, -} - -// SAFETY: only sound for single-threaded (guest-side) access; see the -// type-level invariant on `SendWrap`. -unsafe impl Send for SendWrap>>> {} - -/// Two tier buffer pool with small and large slabs. -#[derive(Debug, Clone)] -pub struct BufferPool { - inner: SendWrap>>>, -} - -impl BufferPool { - /// Create a new buffer pool over a fixed region. - pub fn new(base_addr: u64, region_len: usize) -> Result { - let inner = Inner::::new(base_addr, region_len)?; - Ok(Self { - inner: SendWrap(Rc::new(RefCell::new(inner))), - }) - } -} - -impl BufferPool { - /// Upper slab slot size in bytes. - pub const fn upper_slot_size() -> usize { - 4096 - } - - /// Lower slab slot size in bytes. - pub const fn lower_slot_size() -> usize { - 256 - } -} - -#[cfg(all(test, loom))] -#[derive(Debug, Clone)] -pub struct BufferPoolSync { - inner: std::sync::Arc>>, -} - -#[cfg(all(test, loom))] -impl BufferPoolSync { - /// Create a new buffer pool over a fixed region. - pub fn new(base_addr: u64, region_len: usize) -> Result { - let inner = Inner::::new(base_addr, region_len)?; - Ok(Self { - inner: std::sync::Arc::new(std::sync::Mutex::new(inner)), - }) - } -} - -impl Inner { - /// Create a new buffer pool over a fixed region. - pub fn new(base_addr: u64, region_len: usize) -> Result { - const LOWER_FRACTION: usize = 8; - - let base = usize::try_from(base_addr).map_err(|_| AllocError::Overflow)?; - let region_end = base.checked_add(region_len).ok_or(AllocError::Overflow)?; - - let lower_base = align_up(base, L)?; - let usable = region_end - .checked_sub(lower_base) - .ok_or(AllocError::EmptyRegion)?; - - let lower_region = usable / LOWER_FRACTION; - let lower = Slab::::new(lower_base as u64, lower_region)?; - - let upper_base = lower_base - .checked_add(lower.capacity()) - .ok_or(AllocError::Overflow)?; + let mut rem = total_len; + let mut sgs = SmallVec::<[Allocation; 4]>::new(); - let upper_base = align_up(upper_base, U)?; - let upper_region = region_end - .checked_sub(upper_base) - .ok_or(AllocError::EmptyRegion)?; - - let upper = Slab::::new(upper_base as u64, upper_region)?; - Ok(Self { lower, upper }) - } - - /// Allocate at least `len` bytes. - pub fn alloc(&mut self, len: usize) -> Result { - if len <= L { - match self.lower.alloc(len) { - Ok(alloc) => return Ok(alloc), - Err(AllocError::NoSpace) => {} - Err(e) => return Err(e), + while rem > 0 { + let len = rem.min(seg_cap); + match self.alloc(len) { + Ok(alloc) => { + sgs.push(alloc); + rem -= len; + } + Err(err) => { + for sg in sgs { + let result = self.dealloc(sg.addr); + debug_assert!(result.is_ok(), "dealloc failed: {result:?}"); + } + return Err(err); + } } } - // fallback to upper slab - self.upper.alloc(len) - } - - /// Free a previously allocated block by its start address. - pub fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { - if self.lower.contains(addr) { - self.lower.dealloc_addr(addr) - } else { - self.upper.dealloc_addr(addr) - } - } - - /// Capacity of a live allocation by its start address. - pub fn allocation_len(&self, addr: u64) -> Result { - if self.lower.contains(addr) { - self.lower.allocation_len(addr) - } else { - self.upper.allocation_len(addr) - } + Ok(sgs) } } -impl BufferProvider for BufferPool { +impl BufferProvider for Rc { fn max_alloc_len(&self) -> usize { - U + (**self).max_alloc_len() } fn alloc(&self, len: usize) -> Result { - self.inner.borrow_mut().alloc(len) - } - - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - Ok(smallvec::smallvec![self.alloc(total_len)?]) + (**self).alloc(len) } fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) + (**self).dealloc(addr) } -} -impl BufferPool { - /// Free a previously allocated block by its start address. - pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) - } - - /// Capacity of a live allocation by its start address. - pub fn allocation_len(&self, addr: u64) -> Result { - self.inner.borrow().allocation_len(addr) + fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { + (**self).alloc_sg(total_len) } } -#[cfg(all(test, loom))] -impl BufferProvider for BufferPoolSync { +impl BufferProvider for Arc { fn max_alloc_len(&self) -> usize { - U + (**self).max_alloc_len() } fn alloc(&self, len: usize) -> Result { - self.inner.lock().expect("poisoned mutex").alloc(len) - } - - fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { - Ok(smallvec::smallvec![self.alloc(total_len)?]) + (**self).alloc(len) } fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - self.inner - .lock() - .expect("poisoned mutex") - .dealloc_addr(addr) - } -} - -/// Single-tier fixed-slot free list. -/// -/// Tracks a fixed set of equal-sized buffer slots. Allocation pops a free slot -/// and deallocation returns it, both O(1). A [`FixedBitSet`] records which slots -/// are currently allocated, so double frees and frees of unknown addresses are -/// rejected without scanning the free list. -struct RecycleList { - base_addr: u64, - slot_size: usize, - count: usize, - /// Free slot addresses, popped/pushed LIFO. - free: SmallVec<[u64; 64]>, - /// One bit per slot index; set means the slot is currently handed out. - allocated: FixedBitSet, -} - -// SAFETY: only sound for single-threaded (guest-side) access; see the -// type-level invariant on `SendWrap`. -unsafe impl Send for SendWrap>> {} - -impl RecycleList { - fn new(base_addr: u64, region_len: usize, slot_size: usize) -> Result { - if slot_size == 0 { - return Err(AllocError::InvalidArg); - } - - let count = region_len / slot_size; - if count == 0 { - return Err(AllocError::EmptyRegion); - } - - let mut free = SmallVec::with_capacity(count); - for i in 0..count { - free.push(base_addr + (i * slot_size) as u64); - } - - Ok(Self { - base_addr, - slot_size, - count, - free, - allocated: FixedBitSet::with_capacity(count), - }) - } - - fn end(&self) -> u64 { - self.base_addr + (self.count * self.slot_size) as u64 - } - - fn contains(&self, addr: u64) -> bool { - (self.base_addr..self.end()).contains(&addr) - } - - /// Validate that `addr` names a slot start within the region. - fn slot_of(&self, addr: u64) -> Result { - if !self.contains(addr) { - return Err(AllocError::InvalidFree(addr, 0)); - } - - let off = addr - self.base_addr; - if !off.is_multiple_of(self.slot_size as u64) { - return Err(AllocError::InvalidFree(addr, 0)); - } - - Ok((off / self.slot_size as u64) as usize) - } - - /// Validate that `addr` is a live (currently allocated) slot start. - fn live_slot_of(&self, addr: u64) -> Result { - let slot = self.slot_of(addr)?; - if !self.allocated.contains(slot) { - return Err(AllocError::InvalidFree(addr, 0)); - } - Ok(slot) - } - - fn alloc(&mut self, len: usize) -> Result { - if len == 0 { - return Err(AllocError::InvalidArg); - } - if len > self.slot_size { - return Err(AllocError::OutOfMemory); - } - - let addr = self.free.pop().ok_or(AllocError::NoSpace)?; - // Safety of the index: `addr` came from `free`, which only ever holds - // valid slot starts. - self.allocated - .insert(((addr - self.base_addr) / self.slot_size as u64) as usize); - - Ok(Allocation { - addr, - len: self.slot_size, - }) - } - - fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { - let slot = self.live_slot_of(addr)?; - self.allocated.set(slot, false); - self.free.push(addr); - Ok(()) - } - - fn allocation_len(&self, addr: u64) -> Result { - self.live_slot_of(addr)?; - Ok(self.slot_size) - } - - /// Rebuild state so that exactly the addresses in `allocated` are marked - /// live and every other slot is free. - /// - /// On error the pool is left unchanged. - fn restore_allocated(&mut self, allocated: &[u64]) -> Result<(), AllocError> { - let mut restored = FixedBitSet::with_capacity(self.count); - for &addr in allocated { - let slot = self.slot_of(addr)?; - if restored.contains(slot) { - return Err(AllocError::InvalidFree(addr, self.slot_size)); - } - restored.insert(slot); - } - self.allocated = restored; - self.rebuild_free(); - Ok(()) - } - - /// Repopulate the free list with every slot whose allocated bit is clear. - fn rebuild_free(&mut self) { - self.free.clear(); - for i in 0..self.count { - if !self.allocated.contains(i) { - self.free.push(self.base_addr + (i * self.slot_size) as u64); - } - } - } - - fn slot_addr(&self, index: usize) -> Option { - (index < self.count).then(|| self.base_addr + (index * self.slot_size) as u64) + (**self).dealloc(addr) } - fn num_free(&self) -> usize { - self.free.len() + fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { + (**self).alloc_sg(total_len) } } -/// A recycling buffer provider with fixed-size slots. +/// Wrapper asserting `Send` for an inner value that is only ever accessed from +/// a single thread. /// -/// Holds a fixed set of equal-sized buffer addresses in a free list. Alloc and -/// dealloc are O(1). It is intended for bounded scatter/gather descriptor -/// segments that are pre-allocated and recycled after use: -/// [`alloc_sg`](BufferProvider::alloc_sg) splits a logical payload into -/// `ceil(total_len / slot_size)` fixed-size segments. -#[derive(Clone)] -pub struct RecyclePool { - inner: SendWrap>>, -} - -impl RecyclePool { - /// Create a recycling pool of `slot_size`-byte slots over a fixed region. - /// - /// The base address is aligned up to `slot_size`; the slot count is based - /// on the remaining usable region after alignment. - pub fn new(base_addr: u64, region_len: usize, slot_size: usize) -> Result { - if slot_size == 0 { - return Err(AllocError::InvalidArg); - } - - let base = usize::try_from(base_addr).map_err(|_| AllocError::Overflow)?; - let region_end = base.checked_add(region_len).ok_or(AllocError::Overflow)?; - let aligned = align_up(base, slot_size)?; - let usable = region_end - .checked_sub(aligned) - .ok_or(AllocError::EmptyRegion)?; - let list = RecycleList::new(aligned as u64, usable, slot_size)?; - - Ok(Self { - inner: SendWrap(Rc::new(RefCell::new(list))), - }) - } - - /// Rebuild pool state so that every address in `allocated` is removed from - /// the free list, matching externally known inflight state. Validation is - /// transactional: an error leaves the existing pool state unchanged. - pub fn restore_allocated(&self, allocated: &[u64]) -> Result<(), AllocError> { - self.inner.borrow_mut().restore_allocated(allocated) - } - - /// Compute the address of slot `index`. - /// - /// Returns `None` if `index >= count`. - pub fn slot_addr(&self, index: usize) -> Option { - self.inner.borrow().slot_addr(index) - } - - /// Number of free slots. - pub fn num_free(&self) -> usize { - self.inner.borrow().num_free() - } - - /// Free a previously allocated slot by address. - pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) - } - - /// Capacity of a live allocation by its start address. - pub fn allocation_len(&self, addr: u64) -> Result { - self.inner.borrow().allocation_len(addr) - } - - /// Base address of the pool region. - pub fn base_addr(&self) -> u64 { - self.inner.borrow().base_addr - } - - /// Slot size in bytes. - pub fn slot_size(&self) -> usize { - self.inner.borrow().slot_size - } +/// [`RunPool`] and [`SlotPool`] hold their state in an `Rc>`, which +/// is neither `Send` nor `Sync`. Their allocations are exposed as zero-copy +/// reply payloads through [`Bytes::from_owner`](bytes::Bytes::from_owner), +/// whose owner bound is `Send + 'static`; this wrapper exists solely so the +/// pools can satisfy that bound. +/// +/// # Safety +/// +/// The `Send` assertion is only sound while the wrapped value and every +/// `Bytes` handed out from it stay on a single thread. Hyperlight guests are +/// single-threaded, so this holds for guest-side use. It is unsound to move a +/// pool or reply `Bytes` to another thread. +#[derive(Debug)] +struct SendWrap(T); - /// Number of slots in the pool. - pub fn count(&self) -> usize { - self.inner.borrow().count +impl Clone for SendWrap { + fn clone(&self) -> Self { + Self(self.0.clone()) } } -impl BufferProvider for RecyclePool { - fn max_alloc_len(&self) -> usize { - self.inner.borrow().slot_size - } - - fn alloc(&self, len: usize) -> Result { - self.inner.borrow_mut().alloc(len) - } +impl Deref for SendWrap { + type Target = T; - fn dealloc(&self, addr: u64) -> Result<(), AllocError> { - self.inner.borrow_mut().dealloc_addr(addr) + fn deref(&self) -> &T { + &self.0 } } -#[cfg(test)] -mod tests { - use super::*; - - fn make_pool(size: usize) -> BufferPool { - let base = align_up(0x10000, L.max(U)).unwrap() as u64; - BufferPool::::new(base, size).unwrap() - } - - fn make_recycle_pool(slot_count: usize, slot_size: usize) -> RecyclePool { - let base = 0x80000u64; - RecyclePool::new(base, slot_count * slot_size, slot_size).unwrap() - } - - #[test] - fn test_pool_new_success() { - let pool = BufferPool::<256, 4096>::new(0x10000, 1024 * 1024).unwrap(); - assert!(pool.inner.borrow().lower.capacity() > 0); - assert!(pool.inner.borrow().upper.capacity() > 0); - } - - #[test] - fn test_pool_alloc_small_to_lower() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(128).unwrap(); - - // Should come from lower slab - assert!(pool.inner.borrow().lower.contains(alloc.addr)); - assert_eq!(alloc.len, 256); - } - - #[test] - fn test_pool_alloc_large_to_upper() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(1500).unwrap(); - - // Should come from upper slab - assert!(pool.inner.borrow().upper.contains(alloc.addr)); - assert_eq!(alloc.len, 4096); - } - - #[test] - fn test_pool_alloc_fallback_to_upper() { - let pool = make_pool::<256, 4096>(1024 * 1024); - - // Fill lower slab completely - let mut allocations = Vec::new(); - while pool.inner.borrow().lower.free_bytes() > 0 { - allocations.push(pool.inner.borrow_mut().lower.alloc(256).unwrap()); - } - - // Small allocation should fallback to upper slab - let alloc = pool.alloc(128).unwrap(); - assert!(pool.inner.borrow().upper.contains(alloc.addr)); - } - - #[test] - fn test_pool_free_from_lower() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(128).unwrap(); - - let free_before = pool.inner.borrow().lower.free_bytes(); - pool.dealloc(alloc.addr).unwrap(); - assert_eq!( - pool.inner.borrow().lower.free_bytes(), - free_before + alloc.len - ); - } - - #[test] - fn test_pool_free_from_upper() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let alloc = pool.alloc(1500).unwrap(); - - let free_before = pool.inner.borrow().upper.free_bytes(); - pool.dealloc(alloc.addr).unwrap(); - assert_eq!( - pool.inner.borrow().upper.free_bytes(), - free_before + alloc.len - ); - } - - #[test] - fn test_pool_stress_many_allocations() { - let pool = make_pool::<256, 4096>(4 * 1024 * 1024); - let mut allocations = Vec::new(); - - // Allocate many buffers - for i in 0..100 { - let size = if i % 2 == 0 { 128 } else { 1500 }; - allocations.push(pool.alloc(size).unwrap()); - } - - // Free half of them - for i in (0..100).step_by(2) { - pool.dealloc(allocations[i].addr).unwrap(); - } - - // Should be able to allocate again - for i in 0..50 { - let size = if i % 2 == 0 { 128 } else { 1500 }; - let _alloc = pool.alloc(size).unwrap(); - } - } - - #[test] - fn test_pool_mixed_workload() { - let pool = make_pool::<256, 4096>(2 * 1024 * 1024); - - // Simulate virtio-net workload - let desc_buf = pool.alloc(64).unwrap(); // Control message - let rx_buf1 = pool.alloc(1500).unwrap(); // MTU packet - let rx_buf2 = pool.alloc(1500).unwrap(); // MTU packet - let tx_buf = pool.alloc(4096).unwrap(); // Large buffer - - // Free and reallocate - pool.dealloc(rx_buf1.addr).unwrap(); - let rx_buf3 = pool.alloc(1500).unwrap(); - - // Should reuse freed buffer (LIFO) - assert_eq!(rx_buf3.addr, rx_buf1.addr); - - pool.dealloc(desc_buf.addr).unwrap(); - pool.dealloc(rx_buf2.addr).unwrap(); - pool.dealloc(rx_buf3.addr).unwrap(); - pool.dealloc(tx_buf.addr).unwrap(); - } - - #[test] - fn test_pool_zero_allocation_error() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let result = pool.alloc(0); - assert!(matches!(result, Err(AllocError::InvalidArg))); - } - - #[test] - fn test_pool_too_large_allocation() { - let pool = make_pool::<256, 4096>(1024 * 1024); - let result = pool.alloc(2 * 1024 * 1024); // Larger than pool - assert!(matches!(result, Err(AllocError::OutOfMemory))); - } - - #[test] - fn test_align_up_helper() { - assert_eq!(align_up(0, 256).unwrap(), 0); - assert_eq!(align_up(1, 256).unwrap(), 256); - assert_eq!(align_up(256, 256).unwrap(), 256); - assert_eq!(align_up(257, 256).unwrap(), 512); - assert_eq!(align_up(511, 256).unwrap(), 512); - assert_eq!(align_up(512, 256).unwrap(), 512); - assert!(matches!(align_up(1, 0), Err(AllocError::InvalidArg))); - assert!(matches!( - align_up(usize::MAX, 256), - Err(AllocError::Overflow) - )); - } - - #[test] - fn test_recycle_pool_alignment_subtracts_padding() { - let pool = RecyclePool::new(0x80001, 8192, 4096).unwrap(); - - assert_eq!(pool.base_addr(), 0x81000); - assert_eq!(pool.count(), 1); - } - - // Edge case: allocation exactly at boundary - #[test] - fn test_pool_boundary_allocation() { - let pool = make_pool::<256, 4096>(1024 * 1024); - - // Allocate exactly at boundary - let alloc = pool.alloc(256).unwrap(); - assert!(pool.inner.borrow().lower.contains(alloc.addr)); - - // Allocate just over boundary - let alloc2 = pool.alloc(257).unwrap(); - assert!(pool.inner.borrow().upper.contains(alloc2.addr)); - } - - #[test] - fn test_pool_dealloc_addr_routes_to_correct_tier() { - let pool = make_pool::<256, 4096>(0x20000); - let lower = pool.alloc(128).unwrap(); - let upper = pool.alloc(1024).unwrap(); - - assert_eq!(pool.allocation_len(lower.addr).unwrap(), 256); - assert_eq!(pool.allocation_len(upper.addr).unwrap(), 4096); - - pool.dealloc_addr(lower.addr).unwrap(); - pool.dealloc_addr(upper.addr).unwrap(); - } - - #[test] - fn test_buffer_pool_alloc_sg_uses_one_contiguous_run() { - let pool = make_pool::<256, 4096>(0x20000); - let sgs = pool.alloc_sg(4096 * 2 + 1).unwrap(); - - assert_eq!(sgs.len(), 1); - assert_eq!(sgs[0].len, 4096 * 3); - - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); - } - } - - #[test] - fn test_buffer_pool_alloc_sg_large_run() { - let pool = make_pool::<256, 4096>(0x20000); - let sgs = pool.alloc_sg(8192).unwrap(); - - assert_eq!(sgs.len(), 1); - assert_eq!(sgs[0].len, 8192); - - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); - } - } - - #[test] - fn test_recycle_pool_alloc_sg_splits() { - let pool = make_recycle_pool(8, 4096); - let sgs = pool.alloc_sg(4096 * 2 + 1).unwrap(); - - assert_eq!(sgs.len(), 3); - assert_eq!(sgs[0].len, 4096); - assert_eq!(sgs[1].len, 4096); - assert_eq!(sgs[2].len, 4096); - - for sg in sgs { - pool.dealloc(sg.addr).unwrap(); - } - } - - #[test] - fn test_recycle_pool_restore_allocated_removes_from_free_list() { - let pool = make_recycle_pool(4, 4096); - assert_eq!(pool.num_free(), 4); - - let addrs = [0x80000, 0x81000]; // slots 0 and 1 - pool.restore_allocated(&addrs).unwrap(); - assert_eq!(pool.num_free(), 2); - - // Allocating should only return the two remaining slots - let a1 = pool.alloc(4096).unwrap(); - let a2 = pool.alloc(4096).unwrap(); - assert!(pool.alloc(4096).is_err()); - - // The allocated addresses should be the non-restored ones - let mut got = [a1.addr, a2.addr]; - got.sort(); - assert_eq!(got, [0x82000, 0x83000]); - } - - #[test] - fn test_recycle_pool_restore_allocated_invalid_addr_returns_error() { - let pool = make_recycle_pool(4, 4096); - pool.restore_allocated(&[0x80000]).unwrap(); - - let result = pool.restore_allocated(&[0x81000, 0xDEAD]); - assert!(result.is_err()); - assert_eq!(pool.num_free(), 3); - assert_eq!(pool.allocation_len(0x80000).unwrap(), 4096); - assert!(pool.allocation_len(0x81000).is_err()); - } - - #[test] - fn test_recycle_pool_restore_allocated_then_dealloc_roundtrip() { - let pool = make_recycle_pool(4, 4096); - let addr = 0x81000u64; - - pool.restore_allocated(&[addr]).unwrap(); - assert_eq!(pool.num_free(), 3); - - // Dealloc the restored address - pool.dealloc(addr).unwrap(); - assert_eq!(pool.num_free(), 4); - } - - #[test] - fn test_recycle_pool_restore_allocated_all_slots() { - let pool = make_recycle_pool(4, 4096); - let addrs: Vec = (0..4).map(|i| 0x80000 + i * 4096).collect(); - - pool.restore_allocated(&addrs).unwrap(); - assert_eq!(pool.num_free(), 0); - assert!(pool.alloc(4096).is_err()); - } - - #[test] - fn test_recycle_pool_restore_allocated_empty_list_is_noop() { - let pool = make_recycle_pool(4, 4096); - pool.restore_allocated(&[]).unwrap(); - assert_eq!(pool.num_free(), 4); - } - - #[test] - fn test_recycle_pool_restore_allocated_replaces_state() { - let pool = make_recycle_pool(4, 4096); - - let _ = pool.alloc(4096).unwrap(); - let _ = pool.alloc(4096).unwrap(); - assert_eq!(pool.num_free(), 2); - - pool.restore_allocated(&[0x80000]).unwrap(); - assert_eq!(pool.num_free(), 3); - } - - #[test] - fn test_recycle_pool_dealloc_out_of_range() { - let pool = make_recycle_pool(4, 4096); - let _ = pool.alloc(4096).unwrap(); - - assert!(matches!( - pool.dealloc(0xDEAD), - Err(AllocError::InvalidFree(0xDEAD, 0)) - )); - } - - #[test] - fn test_recycle_pool_dealloc_misaligned() { - let pool = make_recycle_pool(4, 4096); - let _ = pool.alloc(4096).unwrap(); - - assert!(matches!( - pool.dealloc(0x80001), - Err(AllocError::InvalidFree(0x80001, 0)) - )); - } - - #[test] - fn test_recycle_pool_dealloc_double_free() { - let pool = make_recycle_pool(4, 4096); - let a = pool.alloc(4096).unwrap(); - pool.dealloc(a.addr).unwrap(); - - // Second dealloc should fail - address is already in the free list - assert!(matches!( - pool.dealloc(a.addr), - Err(AllocError::InvalidFree(_, _)) - )); - } - - #[test] - fn test_recycle_pool_alloc_sg_rolls_back_on_failure() { - let pool = make_recycle_pool(2, 4096); - - assert!(matches!(pool.alloc_sg(4096 * 3), Err(AllocError::NoSpace))); - assert_eq!(pool.num_free(), 2); - - let alloc = pool.alloc(4096).unwrap(); - assert_eq!(pool.num_free(), 1); - pool.dealloc(alloc.addr).unwrap(); - } - - #[test] - fn test_recycle_pool_dealloc_addr_and_allocation_len() { - let pool = make_recycle_pool(4, 4096); - let alloc = pool.alloc(4096).unwrap(); - - assert_eq!(pool.allocation_len(alloc.addr).unwrap(), 4096); - pool.dealloc_addr(alloc.addr).unwrap(); - assert!(matches!( - pool.allocation_len(alloc.addr), - Err(AllocError::InvalidFree(_, 0)) - )); - } - - #[test] - fn test_recycle_pool_random_order_dealloc() { - let pool = make_recycle_pool(8, 4096); - - let mut allocs: Vec = (0..8).map(|_| pool.alloc(4096).unwrap()).collect(); - assert_eq!(pool.num_free(), 0); - - // Dealloc in reverse order - allocs.reverse(); - for a in &allocs { - pool.dealloc(a.addr).unwrap(); - } - assert_eq!(pool.num_free(), 8); - - // All slots should be re-allocatable - let reallocs: Vec = (0..8).map(|_| pool.alloc(4096).unwrap()).collect(); - assert_eq!(pool.num_free(), 0); - - // Verify all addresses are distinct - let mut addrs: Vec = reallocs.iter().map(|a| a.addr).collect(); - addrs.sort(); - addrs.dedup(); - assert_eq!(addrs.len(), 8); - } - - #[test] - fn test_recycle_pool_interleaved_alloc_dealloc_order() { - let pool = make_recycle_pool(4, 4096); - - let a0 = pool.alloc(4096).unwrap(); - let a1 = pool.alloc(4096).unwrap(); - let a2 = pool.alloc(4096).unwrap(); - let a3 = pool.alloc(4096).unwrap(); - assert_eq!(pool.num_free(), 0); - - // Free middle slots first (out of allocation order) - pool.dealloc(a2.addr).unwrap(); - pool.dealloc(a0.addr).unwrap(); - assert_eq!(pool.num_free(), 2); - - // Re-alloc gets the out-of-order slots back (LIFO) - let b0 = pool.alloc(4096).unwrap(); - assert_eq!(b0.addr, a0.addr); - let b1 = pool.alloc(4096).unwrap(); - assert_eq!(b1.addr, a2.addr); - - // Free everything in yet another order - pool.dealloc(a1.addr).unwrap(); - pool.dealloc(b0.addr).unwrap(); - pool.dealloc(b1.addr).unwrap(); - pool.dealloc(a3.addr).unwrap(); - assert_eq!(pool.num_free(), 4); - - // All 4 original addresses should be available - let mut final_addrs: Vec = (0..4).map(|_| pool.alloc(4096).unwrap().addr).collect(); - final_addrs.sort(); - let expected: Vec = (0..4).map(|i| 0x80000 + i * 4096).collect(); - assert_eq!(final_addrs, expected); +#[inline] +fn align_up(val: usize, align: usize) -> Result { + if align == 0 { + return Err(AllocError::InvalidArg); } - #[test] - fn test_recycle_pool_dealloc_order_independent_of_alloc_order() { - let pool = make_recycle_pool(6, 256); - - // Allocate all - let allocs: Vec = (0..6).map(|_| pool.alloc(256).unwrap()).collect(); - - // Dealloc in scattered order: 4, 1, 5, 0, 3, 2 - let order = [4, 1, 5, 0, 3, 2]; - for &i in &order { - pool.dealloc(allocs[i].addr).unwrap(); - } - assert_eq!(pool.num_free(), 6); - - // Re-allocate all and verify we get back the full set - let mut realloc_addrs: Vec = (0..6).map(|_| pool.alloc(256).unwrap().addr).collect(); - realloc_addrs.sort(); - - let mut orig_addrs: Vec = allocs.iter().map(|a| a.addr).collect(); - orig_addrs.sort(); - - assert_eq!(realloc_addrs, orig_addrs); - } + val.checked_next_multiple_of(align) + .ok_or(AllocError::Overflow) } #[cfg(test)] -mod fuzz { - use quickcheck::{Arbitrary, Gen, QuickCheck}; +mod tests; - use super::*; - - const MAX_OPS: usize = 10; - const MAX_ALLOC_SIZE: usize = 8192; - - #[derive(Clone, Debug)] - enum Op { - Alloc(usize), - AllocSg(usize), - Dealloc(usize), - } - - impl Arbitrary for Op { - fn arbitrary(g: &mut Gen) -> Self { - match u8::arbitrary(g) % 3 { - 0 => Op::Alloc(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), - 1 => Op::AllocSg(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), - 2 => Op::Dealloc(usize::arbitrary(g)), - _ => unreachable!(), - } - } - } - - #[derive(Clone, Debug)] - struct Scenario { - pool_size: usize, - ops: Vec, - } - - impl Arbitrary for Scenario { - fn arbitrary(g: &mut Gen) -> Self { - let pool_size = (usize::arbitrary(g) % (4 * 1024 * 1024)) + (1024 * 1024); - let num_ops = usize::arbitrary(g) % MAX_OPS + 1; - let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); - - Scenario { pool_size, ops } - } - } - - fn run_scenario(s: Scenario) -> bool { - let base = align_up(0x10000, 4096).unwrap() as u64; - let pool = match BufferPool::<256, 4096>::new(base, s.pool_size) { - Ok(p) => p, - Err(_) => return true, - }; - - let mut allocations: Vec = Vec::new(); - - for op in &s.ops { - match op { - Op::Alloc(size) => match pool.alloc(*size) { - Ok(alloc) => { - assert!(alloc.len >= *size); - allocations.push(alloc); - } - Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} - Err(_) => { - return false; - } - }, - Op::AllocSg(size) => match pool.alloc_sg(*size) { - Ok(sgs) => { - let total: usize = sgs.iter().map(|sg| sg.len).sum(); - assert!(total >= *size); - allocations.extend(sgs); - } - Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} - Err(_) => { - return false; - } - }, - Op::Dealloc(idx) => { - if allocations.is_empty() { - continue; - } - - let idx = idx % allocations.len(); - let alloc = allocations.swap_remove(idx); - - match pool.dealloc(alloc.addr) { - Ok(_) => {} - Err(_) => return false, - } - } - } - - if check_pool_invariants(&pool, &allocations).is_err() { - return false; - } - } - - // Cleanup - for alloc in &allocations { - if pool.dealloc(alloc.addr).is_err() { - return false; - } - } - - check_pool_invariants(&pool, &allocations).is_ok() - } - - fn check_slab_invariants(slab: &Slab) -> Result<(), &'static str> { - let used = slab.used_slots.count_ones(..); - let free = slab.used_slots.count_zeroes(..); - if used + free != slab.used_slots.len() { - return Err("used + free != total slots"); - } - - let expected_free = free * N; - if slab.free_bytes() != expected_free { - return Err("free_bytes doesn't match bitmap"); - } - - if let Some(alloc) = slab.last_free_run { - if alloc.len == 0 || alloc.len % N != 0 { - return Err("last_free_run has invalid length"); - } - if !slab.contains(alloc.addr) { - return Err("last_free_run addr outside range"); - } - } - - Ok(()) - } - - fn check_pool_invariants( - pool: &BufferPool, - allocations: &[Allocation], - ) -> Result<(), &'static str> { - check_slab_invariants(&pool.inner.borrow().lower)?; - check_slab_invariants(&pool.inner.borrow().upper)?; - - if pool.inner.borrow().lower.range().end > pool.inner.borrow().upper.range().start { - return Err("lower and upper ranges overlap"); - } - - let mut seen = std::collections::HashSet::new(); - - for alloc in allocations { - if !pool.inner.borrow().lower.contains(alloc.addr) - && !pool.inner.borrow().upper.contains(alloc.addr) - { - return Err("allocation address outside pool ranges"); - } - - if alloc.len % L != 0 && alloc.len % U != 0 { - return Err("allocation length not aligned to any tier"); - } - - if !seen.insert(alloc.addr) { - return Err("duplicate allocation address in tracking"); - } - } - - Ok(()) - } - - #[test] - fn prop_allocator_invariants() { - #[cfg(miri)] - let tests = 10; - #[cfg(not(miri))] - let tests = 1000; - - QuickCheck::new() - .tests(tests) - .quickcheck(run_scenario as fn(Scenario) -> bool); - } -} +#[cfg(test)] +mod fuzz; diff --git a/src/hyperlight_common/src/virtq/pool/fuzz.rs b/src/hyperlight_common/src/virtq/pool/fuzz.rs new file mode 100644 index 000000000..711808841 --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/fuzz.rs @@ -0,0 +1,382 @@ +/* +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), + AllocSg(usize), + Dealloc(usize), +} + +impl Arbitrary for Op { + fn arbitrary(g: &mut Gen) -> Self { + match u8::arbitrary(g) % 3 { + 0 => Op::Alloc(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), + 1 => Op::AllocSg(usize::arbitrary(g) % MAX_ALLOC_SIZE + 1), + 2 => Op::Dealloc(usize::arbitrary(g)), + _ => unreachable!(), + } + } +} + +#[derive(Clone, Debug)] +struct RunScenario { + pool_size: usize, + ops: Vec, +} + +impl Arbitrary for RunScenario { + fn arbitrary(g: &mut Gen) -> Self { + let pool_size = (usize::arbitrary(g) % (4 * 1024 * 1024)) + (1024 * 1024); + let num_ops = usize::arbitrary(g) % MAX_OPS + 1; + let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); + + RunScenario { pool_size, ops } + } +} + +fn run_provider_ops(pool: &P, ops: &[Op], check: F) -> bool +where + P: BufferProvider, + F: Fn(&P, &[Allocation]) -> Result<(), &'static str>, +{ + let mut allocations: Vec = Vec::new(); + + for op in ops { + match op { + Op::Alloc(size) => match pool.alloc(*size) { + Ok(alloc) => { + if alloc.len < *size + || allocations + .iter() + .any(|existing| existing.addr == alloc.addr) + { + return false; + } + allocations.push(alloc); + } + Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} + Err(_) => return false, + }, + Op::AllocSg(size) => match pool.alloc_sg(*size) { + Ok(sgs) => { + let mut total = 0usize; + for sg in sgs { + let Some(next_total) = total.checked_add(sg.len) else { + return false; + }; + if allocations.iter().any(|existing| existing.addr == sg.addr) { + return false; + } + total = next_total; + allocations.push(sg); + } + if total < *size { + return false; + } + } + Err(AllocError::NoSpace | AllocError::OutOfMemory) => {} + Err(_) => return false, + }, + Op::Dealloc(index) => { + if !allocations.is_empty() { + let index = index % allocations.len(); + if pool.dealloc(allocations[index].addr).is_err() { + return false; + } + allocations.swap_remove(index); + } + } + } + + if check(pool, &allocations).is_err() { + return false; + } + } + + while let Some(alloc) = allocations.pop() { + if pool.dealloc(alloc.addr).is_err() { + return false; + } + } + + check(pool, &allocations).is_ok() +} + +fn run_pool_scenario(scenario: RunScenario) -> bool { + let base = align_up(0x10000, 4096).unwrap() as u64; + let pool = match RunPool::<256, 4096>::new(base, scenario.pool_size) { + Ok(pool) => pool, + Err(_) => return true, + }; + + run_provider_ops(&pool, &scenario.ops, check_run_pool_invariants) +} + +fn check_run_tier_invariants( + tier: &RunTier, + allocations: &[Allocation], +) -> Result<(), &'static str> { + let mut expected_used = HashSet::new(); + let mut expected_starts = HashSet::new(); + + for alloc in allocations.iter().filter(|alloc| tier.contains(alloc.addr)) { + let offset = usize::try_from(alloc.addr - tier.base_addr) + .map_err(|_| "allocation offset overflows usize")?; + if alloc.len == 0 || !offset.is_multiple_of(N) || !alloc.len.is_multiple_of(N) { + return Err("allocation is not tier-aligned"); + } + + let start = offset / N; + let slots = alloc.len / N; + let end = start + .checked_add(slots) + .ok_or("allocation slot range overflow")?; + if end > tier.used_slots.len() || !expected_starts.insert(start) { + return Err("allocation run is invalid"); + } + for slot in start..end { + if !expected_used.insert(slot) { + return Err("allocation runs overlap"); + } + } + } + + for slot in 0..tier.used_slots.len() { + if tier.used_slots.contains(slot) != expected_used.contains(&slot) { + return Err("used bitmap does not match live allocations"); + } + if tier.run_starts.contains(slot) != expected_starts.contains(&slot) { + return Err("run-start bitmap does not match live allocations"); + } + } + + if tier.free_bytes() != (tier.used_slots.len() - expected_used.len()) * N { + return Err("free_bytes does not match live allocations"); + } + + if let Some(free_run) = tier.last_free_run { + if !tier.contains(free_run.addr) { + return Err("cached free-run address outside tier"); + } + let offset = usize::try_from(free_run.addr - tier.base_addr) + .map_err(|_| "cached free-run offset overflows usize")?; + if free_run.len == 0 || !offset.is_multiple_of(N) || !free_run.len.is_multiple_of(N) { + return Err("cached free run is not tier-aligned"); + } + + let start = offset / N; + let end = start + .checked_add(free_run.len / N) + .ok_or("cached free-run range overflow")?; + if end > tier.used_slots.len() || (start..end).any(|slot| tier.used_slots.contains(slot)) { + return Err("cached free run overlaps live allocations"); + } + } + + Ok(()) +} + +fn check_run_pool_invariants( + pool: &RunPool, + allocations: &[Allocation], +) -> Result<(), &'static str> { + let inner = pool.inner.borrow(); + if inner.lower.range().end > inner.upper.range().start { + return Err("lower and upper ranges overlap"); + } + + let mut seen = HashSet::new(); + for alloc in allocations { + let in_lower = inner.lower.contains(alloc.addr); + let in_upper = inner.upper.contains(alloc.addr); + if in_lower == in_upper { + return Err("allocation does not belong to exactly one tier"); + } + if !seen.insert(alloc.addr) { + return Err("duplicate allocation address in tracking"); + } + } + + check_run_tier_invariants(&inner.lower, allocations)?; + check_run_tier_invariants(&inner.upper, allocations) +} + +#[derive(Clone, Debug)] +struct SlotScenario { + tiered: bool, + lower_count: usize, + upper_count: usize, + ops: Vec, +} + +impl Arbitrary for SlotScenario { + fn arbitrary(g: &mut Gen) -> Self { + let tiered = bool::arbitrary(g); + let lower_count = usize::arbitrary(g) % MAX_TIER_SLOTS + 1; + let upper_count = usize::arbitrary(g) % MAX_TIER_SLOTS + 1; + let num_ops = usize::arbitrary(g) % MAX_OPS + 1; + let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); + + Self { + tiered, + lower_count, + upper_count, + ops, + } + } +} + +fn make_slot_pool(scenario: &SlotScenario) -> SlotPool { + if scenario.tiered { + let lower = SlotLayout::new(LOWER_BASE, LOWER_SLOT_SIZE, scenario.lower_count); + let upper = SlotLayout::new(UPPER_BASE, UPPER_SLOT_SIZE, scenario.upper_count); + SlotPool::new_tiered(lower, upper).unwrap() + } else { + let layout = SlotLayout::new(UPPER_BASE, UPPER_SLOT_SIZE, scenario.upper_count); + SlotPool::new(layout).unwrap() + } +} + +fn run_slot_pool_scenario(scenario: SlotScenario) -> bool { + let pool = make_slot_pool(&scenario); + run_provider_ops(&pool, &scenario.ops, check_slot_pool_invariants) +} + +fn layout_contains(layout: SlotLayout, addr: u64) -> bool { + let Ok(end) = layout.end_addr() else { + return false; + }; + (layout.base_addr..end).contains(&addr) +} + +fn slot_capacity(pool: &SlotPool, addr: u64) -> Option { + let (lower, upper) = pool.layouts(); + if let Some(lower) = lower + && layout_contains(lower, addr) + { + return Some(lower.slot_size); + } + layout_contains(upper, addr).then_some(upper.slot_size) +} + +fn check_slot_pool_invariants( + pool: &SlotPool, + allocations: &[Allocation], +) -> Result<(), &'static str> { + let mut expected_live = BTreeMap::new(); + for alloc in allocations { + if expected_live.insert(alloc.addr, alloc.len).is_some() { + return Err("duplicate allocation address in tracking"); + } + } + + let live = pool.live_addrs(); + let expected_addrs: Vec = expected_live.keys().copied().collect(); + if live != expected_addrs || live.windows(2).any(|pair| pair[0] >= pair[1]) { + return Err("live addresses are not unique and deterministic"); + } + if pool.num_free() + live.len() != pool.count() { + return Err("free + live != total slots"); + } + + let (lower, upper) = pool.layouts(); + let expected_base = lower.map_or(upper.base_addr, |layout| layout.base_addr); + if pool.base_addr() != expected_base || pool.slot_size() != upper.slot_size { + return Err("reported pool layout is inconsistent"); + } + + let mut expected_count = upper.slot_count; + if let Some(lower) = lower { + if lower.slot_size >= upper.slot_size + || lower.end_addr().map_err(|_| "lower layout overflow")? > upper.base_addr + { + return Err("tier layout is invalid"); + } + expected_count += lower.slot_count; + } + if pool.count() != expected_count || pool.slot_addr(pool.count()).is_some() { + return Err("reported slot count is inconsistent"); + } + + let mut seen = HashSet::new(); + for index in 0..pool.count() { + let Some(addr) = pool.slot_addr(index) else { + return Err("missing slot address"); + }; + if !seen.insert(addr) { + return Err("duplicate slot address"); + } + let Some(capacity) = slot_capacity(pool, addr) else { + return Err("slot address outside layout"); + }; + + match expected_live.get(&addr) { + Some(expected_capacity) => { + if *expected_capacity != capacity + || pool.allocation_len(addr).ok() != Some(capacity) + { + return Err("live slot capacity is inconsistent"); + } + } + None if pool.allocation_len(addr).is_ok() => { + return Err("free slot reported as live"); + } + None => {} + } + } + + Ok(()) +} + +#[test] +fn prop_run_pool_invariants() { + #[cfg(miri)] + let tests = 10; + #[cfg(not(miri))] + let tests = 1000; + + QuickCheck::new() + .tests(tests) + .quickcheck(run_pool_scenario as fn(RunScenario) -> bool); +} + +#[test] +fn prop_slot_pool_invariants() { + #[cfg(miri)] + let tests = 10; + #[cfg(not(miri))] + let tests = 1000; + + QuickCheck::new() + .tests(tests) + .quickcheck(run_slot_pool_scenario as fn(SlotScenario) -> bool); +} diff --git a/src/hyperlight_common/src/virtq/pool/run.rs b/src/hyperlight_common/src/virtq/pool/run.rs new file mode 100644 index 000000000..cbd75f3e9 --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/run.rs @@ -0,0 +1,379 @@ +/* +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 smallvec::SmallVec; + +use super::{AllocError, Allocation, BufferProvider, SendWrap, align_up}; + +#[derive(Debug, Clone)] +pub(super) struct Tier { + pub(super) base_addr: u64, + pub(super) used_slots: FixedBitSet, + pub(super) run_starts: FixedBitSet, + pub(super) last_free_run: Option, +} + +impl Tier { + fn new(base_addr: u64, region_len: usize) -> Result { + let usable = region_len - (region_len % N); + let num_slots = usable / N; + let used_slots = FixedBitSet::with_capacity(num_slots); + let run_starts = FixedBitSet::with_capacity(num_slots); + + if !base_addr.is_multiple_of(N as u64) { + return Err(AllocError::InvalidAlign(base_addr)); + } + if num_slots == 0 { + return Err(AllocError::EmptyRegion); + } + + Ok(Self { + base_addr, + used_slots, + run_starts, + last_free_run: None, + }) + } + + fn addr_of(&self, slot_idx: usize) -> Option { + self.base_addr + .checked_add((slot_idx as u64).checked_mul(N as u64)?) + } + + fn slot_of(&self, addr: u64) -> usize { + let off = (addr - self.base_addr) as usize; + off / N + } + + fn checked_slot_of(&self, addr: u64, len: usize) -> Result { + if addr < self.base_addr { + return Err(AllocError::InvalidFree(addr, len)); + } + + let off = (addr - self.base_addr) as usize; + if !off.is_multiple_of(N) { + return Err(AllocError::InvalidFree(addr, len)); + } + + let slot = off / N; + if slot >= self.used_slots.len() { + return Err(AllocError::InvalidFree(addr, len)); + } + + Ok(slot) + } + + fn live_run_slots_at(&self, start: usize) -> Option { + if start >= self.used_slots.len() + || !self.used_slots.contains(start) + || !self.run_starts.contains(start) + { + return None; + } + + let mut end = start + 1; + while end < self.used_slots.len() + && self.used_slots.contains(end) + && !self.run_starts.contains(end) + { + end += 1; + } + + Some(end - start) + } + + fn maybe_invalidate_last_run(&mut self, alloc: Allocation) { + if let Some(run) = &self.last_free_run { + let new_end = alloc.addr + alloc.len as u64; + let run_end = run.addr + run.len as u64; + + if alloc.addr < run_end && run.addr < new_end { + self.last_free_run = None; + } + } + } + + fn find_slots(&mut self, slots_num: usize) -> Option { + debug_assert!(slots_num > 0); + + if let Some(alloc) = self.last_free_run + && alloc.len >= slots_num * N + { + let pos = self.slot_of(alloc.addr); + let _ = self.last_free_run.take(); + return Some(pos); + } + + let total = self.used_slots.len(); + self.used_slots.zeroes().find(|&next_free| { + let end = next_free + slots_num; + end <= total && self.used_slots.count_zeroes(next_free..end) == slots_num + }) + } + + pub(super) fn alloc(&mut self, len: usize) -> Result { + if len == 0 { + return Err(AllocError::InvalidArg); + } + + let total = self.used_slots.len(); + let need_slots = len.div_ceil(N); + if need_slots > total { + return Err(AllocError::OutOfMemory); + } + + let idx = self.find_slots(need_slots).ok_or(AllocError::NoSpace)?; + self.used_slots.insert_range(idx..idx + need_slots); + self.run_starts.insert(idx); + let addr = self.addr_of(idx).ok_or(AllocError::Overflow)?; + + let alloc = Allocation { + addr, + len: need_slots * N, + }; + + self.maybe_invalidate_last_run(alloc); + Ok(alloc) + } + + fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { + let start = self.checked_slot_of(addr, 0)?; + let run_slots = self + .live_run_slots_at(start) + .ok_or(AllocError::InvalidFree(addr, 0))?; + self.dealloc_run(start, run_slots, addr) + } + + fn dealloc_run(&mut self, start: usize, run_slots: usize, addr: u64) -> Result<(), AllocError> { + let len = run_slots * N; + self.used_slots.remove_range(start..start + run_slots); + self.run_starts.set(start, false); + self.last_free_run = Some(Allocation { addr, len }); + Ok(()) + } + + fn allocation_len(&self, addr: u64) -> Result { + let start = self.checked_slot_of(addr, 0)?; + let run_slots = self + .live_run_slots_at(start) + .ok_or(AllocError::InvalidFree(addr, 0))?; + Ok(run_slots * N) + } + + pub(super) fn capacity(&self) -> usize { + self.used_slots.len() * N + } + + pub(super) fn range(&self) -> core::ops::Range { + self.base_addr..self.base_addr + self.capacity() as u64 + } + + pub(super) fn contains(&self, addr: u64) -> bool { + self.range().contains(&addr) + } +} + +#[cfg(test)] +impl Tier { + pub(super) fn free_bytes(&self) -> usize { + (self.used_slots.len() - self.used_slots.count_ones(..)) * N + } +} + +#[derive(Debug)] +pub(super) struct Inner { + pub(super) lower: Tier, + pub(super) upper: Tier, +} + +// SAFETY: only sound for single-threaded (guest-side) access; see the +// type-level invariant on `SendWrap`. +unsafe impl Send for SendWrap>>> {} + +/// Two-tier pool for variable-sized contiguous runs. +#[derive(Debug, Clone)] +pub struct RunPool { + pub(super) inner: SendWrap>>>, +} + +impl RunPool { + /// Create a new run pool over a fixed region. + pub fn new(base_addr: u64, region_len: usize) -> Result { + let inner = Inner::::new(base_addr, region_len)?; + Ok(Self { + inner: SendWrap(Rc::new(RefCell::new(inner))), + }) + } +} + +impl RunPool { + /// Upper tier slot size in bytes. + pub const fn upper_slot_size() -> usize { + 4096 + } + + /// Lower tier slot size in bytes. + pub const fn lower_slot_size() -> usize { + 256 + } +} + +#[cfg(all(test, loom))] +#[derive(Debug, Clone)] +pub struct RunPoolSync { + inner: std::sync::Arc>>, +} + +#[cfg(all(test, loom))] +impl RunPoolSync { + /// Create a new synchronized run pool over a fixed region. + pub fn new(base_addr: u64, region_len: usize) -> Result { + let inner = Inner::::new(base_addr, region_len)?; + Ok(Self { + inner: std::sync::Arc::new(std::sync::Mutex::new(inner)), + }) + } +} + +impl Inner { + /// Create new run-pool state over a fixed region. + pub fn new(base_addr: u64, region_len: usize) -> Result { + const LOWER_FRACTION: usize = 8; + + let base = usize::try_from(base_addr).map_err(|_| AllocError::Overflow)?; + let region_end = base.checked_add(region_len).ok_or(AllocError::Overflow)?; + + let lower_base = align_up(base, L)?; + let usable = region_end + .checked_sub(lower_base) + .ok_or(AllocError::EmptyRegion)?; + + let lower_region = usable / LOWER_FRACTION; + let lower = Tier::::new(lower_base as u64, lower_region)?; + + let upper_base = lower_base + .checked_add(lower.capacity()) + .ok_or(AllocError::Overflow)?; + + let upper_base = align_up(upper_base, U)?; + let upper_region = region_end + .checked_sub(upper_base) + .ok_or(AllocError::EmptyRegion)?; + + let upper = Tier::::new(upper_base as u64, upper_region)?; + Ok(Self { lower, upper }) + } + + /// Allocate at least `len` bytes. + pub fn alloc(&mut self, len: usize) -> Result { + if len <= L { + match self.lower.alloc(len) { + Ok(alloc) => return Ok(alloc), + Err(AllocError::NoSpace) => {} + Err(e) => return Err(e), + } + } + + // Fall back to the upper tier. + self.upper.alloc(len) + } + + /// Free a previously allocated block by its start address. + pub fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { + if self.lower.contains(addr) { + self.lower.dealloc_addr(addr) + } else { + self.upper.dealloc_addr(addr) + } + } + + /// Capacity of a live allocation by its start address. + pub fn allocation_len(&self, addr: u64) -> Result { + if self.lower.contains(addr) { + self.lower.allocation_len(addr) + } else { + self.upper.allocation_len(addr) + } + } +} + +impl BufferProvider for RunPool { + fn max_alloc_len(&self) -> usize { + U + } + + fn alloc(&self, len: usize) -> Result { + self.inner.borrow_mut().alloc(len) + } + + fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { + Ok(smallvec::smallvec![self.alloc(total_len)?]) + } + + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } +} + +impl RunPool { + /// Free a previously allocated block by its start address. + pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } + + /// Capacity of a live allocation by its start address. + pub fn allocation_len(&self, addr: u64) -> Result { + self.inner.borrow().allocation_len(addr) + } +} + +#[cfg(all(test, loom))] +impl BufferProvider for RunPoolSync { + fn max_alloc_len(&self) -> usize { + U + } + + fn alloc(&self, len: usize) -> Result { + self.inner.lock().expect("poisoned mutex").alloc(len) + } + + fn alloc_sg(&self, total_len: usize) -> Result, AllocError> { + Ok(smallvec::smallvec![self.alloc(total_len)?]) + } + + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { + self.inner + .lock() + .expect("poisoned mutex") + .dealloc_addr(addr) + } +} diff --git a/src/hyperlight_common/src/virtq/pool/slot.rs b/src/hyperlight_common/src/virtq/pool/slot.rs new file mode 100644 index 000000000..7689e8e01 --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/slot.rs @@ -0,0 +1,391 @@ +/* +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, BufferProvider, SendWrap}; + +/// Exact memory layout for one [`SlotPool`] tier. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SlotLayout { + /// Start of the first slot. + pub base_addr: u64, + /// Capacity of each slot. + pub slot_size: usize, + /// Number of slots. + pub slot_count: usize, +} + +impl SlotLayout { + /// Describe exact fixed-slot placement. + pub const fn new(base_addr: u64, slot_size: usize, slot_count: usize) -> Self { + Self { + base_addr, + slot_size, + slot_count, + } + } + + /// Total bytes occupied by the slots. + pub fn byte_len(self) -> Result { + self.slot_size + .checked_mul(self.slot_count) + .ok_or(AllocError::Overflow) + } + + /// Exclusive end address. + pub fn end_addr(self) -> Result { + self.base_addr + .checked_add(u64::try_from(self.byte_len()?).map_err(|_| AllocError::Overflow)?) + .ok_or(AllocError::Overflow) + } +} + +/// Single-tier fixed-slot free list. +/// +/// Tracks a fixed set of equal-sized buffer slots. Allocation pops a free slot +/// and deallocation returns it, both O(1). A [`FixedBitSet`] records which slots +/// are currently allocated, so double frees and frees of unknown addresses are +/// rejected without scanning the free list. +struct Tier { + /// Start of this tier's backing memory. + base_addr: u64, + /// Capacity of this slot. + slot_size: usize, + /// Number of slots in this tier. + count: usize, + /// Free slot addresses, popped/pushed LIFO. + free: SmallVec<[u64; 64]>, + /// One bit per slot index; set means the slot is currently handed out. + allocated: FixedBitSet, +} + +// SAFETY: only sound for single-threaded (guest-side) access; see the +// type-level invariant on `SendWrap`. +unsafe impl Send for SendWrap>> {} + +impl Tier { + fn from_layout(layout: SlotLayout) -> Result { + if layout.slot_size == 0 { + return Err(AllocError::InvalidArg); + } + + if layout.slot_count == 0 { + return Err(AllocError::EmptyRegion); + } + + layout.end_addr()?; + + let mut free = SmallVec::with_capacity(layout.slot_count); + for i in 0..layout.slot_count { + free.push(layout.base_addr + (i * layout.slot_size) as u64); + } + + Ok(Self { + base_addr: layout.base_addr, + slot_size: layout.slot_size, + count: layout.slot_count, + free, + allocated: FixedBitSet::with_capacity(layout.slot_count), + }) + } + + fn end(&self) -> u64 { + self.base_addr + (self.count * self.slot_size) as u64 + } + + fn contains(&self, addr: u64) -> bool { + (self.base_addr..self.end()).contains(&addr) + } + + /// Validate that `addr` names a slot start within the region. + fn slot_of(&self, addr: u64) -> Result { + if !self.contains(addr) { + return Err(AllocError::InvalidFree(addr, 0)); + } + + let off = addr - self.base_addr; + if !off.is_multiple_of(self.slot_size as u64) { + return Err(AllocError::InvalidFree(addr, 0)); + } + + Ok((off / self.slot_size as u64) as usize) + } + + /// Validate that `addr` is a live (currently allocated) slot start. + fn live_slot_of(&self, addr: u64) -> Result { + let slot = self.slot_of(addr)?; + if !self.allocated.contains(slot) { + return Err(AllocError::InvalidFree(addr, 0)); + } + Ok(slot) + } + + fn alloc(&mut self, len: usize) -> Result { + if len == 0 { + return Err(AllocError::InvalidArg); + } + if len > self.slot_size { + return Err(AllocError::OutOfMemory); + } + + let addr = self.free.pop().ok_or(AllocError::NoSpace)?; + // Safety of the index: `addr` came from `free`, which only ever holds + // valid slot starts. + self.allocated + .insert(((addr - self.base_addr) / self.slot_size as u64) as usize); + + Ok(Allocation { + addr, + len: self.slot_size, + }) + } + + fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { + let slot = self.live_slot_of(addr)?; + self.allocated.set(slot, false); + self.free.push(addr); + Ok(()) + } + + fn allocation_len(&self, addr: u64) -> Result { + self.live_slot_of(addr)?; + Ok(self.slot_size) + } + + fn slot_addr(&self, index: usize) -> Option { + (index < self.count).then(|| self.base_addr + (index * self.slot_size) as u64) + } + + fn num_free(&self) -> usize { + self.free.len() + } + + fn append_live_addrs(&self, addrs: &mut Vec) { + addrs.extend( + self.allocated + .ones() + .map(|slot| self.base_addr + (slot * self.slot_size) as u64), + ); + } + + fn layout(&self) -> SlotLayout { + SlotLayout::new(self.base_addr, self.slot_size, self.count) + } +} + +struct Inner { + lower: Option, + upper: Tier, +} + +impl Inner { + fn new(lower: Option, upper: SlotLayout) -> Result { + let lower = lower.map(Tier::from_layout).transpose()?; + let upper = Tier::from_layout(upper)?; + + if let Some(lower) = &lower + && (lower.slot_size >= upper.slot_size || lower.end() > upper.base_addr) + { + return Err(AllocError::InvalidArg); + } + + Ok(Self { lower, upper }) + } + + fn max_alloc_len(&self) -> usize { + self.upper.slot_size + } + + fn alloc(&mut self, len: usize) -> Result { + if let Some(lower) = &mut self.lower + && len <= lower.slot_size + { + match lower.alloc(len) { + Ok(alloc) => return Ok(alloc), + Err(AllocError::NoSpace) => {} + Err(err) => return Err(err), + } + } + + self.upper.alloc(len) + } + + fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { + if let Some(lower) = &mut self.lower + && lower.contains(addr) + { + return lower.dealloc_addr(addr); + } + self.upper.dealloc_addr(addr) + } + + fn allocation_len(&self, addr: u64) -> Result { + if let Some(lower) = &self.lower + && lower.contains(addr) + { + return lower.allocation_len(addr); + } + self.upper.allocation_len(addr) + } + + fn slot_addr(&self, index: usize) -> Option { + if let Some(lower) = &self.lower { + if index < lower.count { + return lower.slot_addr(index); + } + return self.upper.slot_addr(index - lower.count); + } + self.upper.slot_addr(index) + } + + fn live_addrs(&self) -> Vec { + let mut addrs = Vec::with_capacity(self.count() - self.num_free()); + if let Some(lower) = &self.lower { + lower.append_live_addrs(&mut addrs); + } + self.upper.append_live_addrs(&mut addrs); + addrs + } + + fn base_addr(&self) -> u64 { + self.lower + .as_ref() + .map_or(self.upper.base_addr, |lower| lower.base_addr) + } + + fn count(&self) -> usize { + self.lower.as_ref().map_or(0, |lower| lower.count) + self.upper.count + } + + fn num_free(&self) -> usize { + self.lower.as_ref().map_or(0, Tier::num_free) + self.upper.num_free() + } + + fn layouts(&self) -> (Option, SlotLayout) { + (self.lower.as_ref().map(Tier::layout), self.upper.layout()) + } +} + +/// A buffer pool with one or two fixed-slot tiers. +/// +/// Allocation and deallocation are O(1) per slot. Eligible allocations first +/// try the optional lower tier and fall back to the required upper tier when +/// the lower tier is full. [`alloc_sg`](BufferProvider::alloc_sg) splits logical +/// payloads into bounded descriptor segments. +#[derive(Clone)] +pub struct SlotPool { + inner: SendWrap>>, +} + +impl SlotPool { + /// Create a single-tier recycling pool from exact slot placement. + pub fn new(layout: SlotLayout) -> Result { + Self::from_layouts(None, layout) + } + + /// Create a two-tier recycling pool from exact lower and upper layouts. + /// + /// The lower layout must precede the upper layout without overlap, and its + /// slot size must be strictly smaller. + pub fn new_tiered(lower: SlotLayout, upper: SlotLayout) -> Result { + Self::from_layouts(Some(lower), upper) + } + + fn from_layouts(lower: Option, upper: SlotLayout) -> Result { + let inner = Inner::new(lower, upper)?; + Ok(Self { + inner: SendWrap(Rc::new(RefCell::new(inner))), + }) + } + + /// Return every live slot address in deterministic tier and index order. + pub fn live_addrs(&self) -> Vec { + self.inner.borrow().live_addrs() + } + + /// Return the lower and upper tier layouts. + pub fn layouts(&self) -> (Option, SlotLayout) { + self.inner.borrow().layouts() + } + + /// Compute the address of slot `index`, with lower-tier slots first. + /// + /// Returns `None` if `index >= count`. + pub fn slot_addr(&self, index: usize) -> Option { + self.inner.borrow().slot_addr(index) + } + + /// Total number of free slots across all tiers. + pub fn num_free(&self) -> usize { + self.inner.borrow().num_free() + } + + /// Free a previously allocated slot by address. + pub fn dealloc_addr(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } + + /// Capacity of a live allocation by its start address. + pub fn allocation_len(&self, addr: u64) -> Result { + self.inner.borrow().allocation_len(addr) + } + + /// Base address of the first managed tier. + pub fn base_addr(&self) -> u64 { + self.inner.borrow().base_addr() + } + + /// Maximum slot size in bytes. + pub fn slot_size(&self) -> usize { + self.inner.borrow().max_alloc_len() + } + + /// Total number of slots across all tiers. + pub fn count(&self) -> usize { + self.inner.borrow().count() + } +} + +impl BufferProvider for SlotPool { + fn max_alloc_len(&self) -> usize { + self.inner.borrow().max_alloc_len() + } + + fn alloc(&self, len: usize) -> Result { + self.inner.borrow_mut().alloc(len) + } + + fn dealloc(&self, addr: u64) -> Result<(), AllocError> { + self.inner.borrow_mut().dealloc_addr(addr) + } +} diff --git a/src/hyperlight_common/src/virtq/pool/tests.rs b/src/hyperlight_common/src/virtq/pool/tests.rs new file mode 100644 index 000000000..87c917e7c --- /dev/null +++ b/src/hyperlight_common/src/virtq/pool/tests.rs @@ -0,0 +1,515 @@ +/* +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() +} + +#[test] +fn test_run_pool_new_success() { + let pool = RunPool::<256, 4096>::new(0x10000, 1024 * 1024).unwrap(); + assert!(pool.inner.borrow().lower.capacity() > 0); + assert!(pool.inner.borrow().upper.capacity() > 0); +} + +#[test] +fn test_run_pool_alloc_small_to_lower() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(128).unwrap(); + + // Should come from the lower tier. + assert!(pool.inner.borrow().lower.contains(alloc.addr)); + assert_eq!(alloc.len, 256); +} + +#[test] +fn test_run_pool_alloc_large_to_upper() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(1500).unwrap(); + + // Should come from the upper tier. + assert!(pool.inner.borrow().upper.contains(alloc.addr)); + assert_eq!(alloc.len, 4096); +} + +#[test] +fn test_run_pool_alloc_fallback_to_upper() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + + // Fill the lower tier completely. + let mut allocations = Vec::new(); + while pool.inner.borrow().lower.free_bytes() > 0 { + allocations.push(pool.inner.borrow_mut().lower.alloc(256).unwrap()); + } + + // Small allocation should fall back to the upper tier. + let alloc = pool.alloc(128).unwrap(); + assert!(pool.inner.borrow().upper.contains(alloc.addr)); +} + +#[test] +fn test_run_pool_free_from_lower() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(128).unwrap(); + + let free_before = pool.inner.borrow().lower.free_bytes(); + pool.dealloc(alloc.addr).unwrap(); + assert_eq!( + pool.inner.borrow().lower.free_bytes(), + free_before + alloc.len + ); +} + +#[test] +fn test_run_pool_free_from_upper() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let alloc = pool.alloc(1500).unwrap(); + + let free_before = pool.inner.borrow().upper.free_bytes(); + pool.dealloc(alloc.addr).unwrap(); + assert_eq!( + pool.inner.borrow().upper.free_bytes(), + free_before + alloc.len + ); +} + +#[test] +fn test_run_pool_stress_many_allocations() { + let pool = make_run_pool::<256, 4096>(4 * 1024 * 1024); + let mut allocations = Vec::new(); + + // Allocate many buffers + for i in 0..100 { + let size = if i % 2 == 0 { 128 } else { 1500 }; + allocations.push(pool.alloc(size).unwrap()); + } + + // Free half of them + for i in (0..100).step_by(2) { + pool.dealloc(allocations[i].addr).unwrap(); + } + + // Should be able to allocate again + for i in 0..50 { + let size = if i % 2 == 0 { 128 } else { 1500 }; + let _alloc = pool.alloc(size).unwrap(); + } +} + +#[test] +fn test_run_pool_mixed_workload() { + let pool = make_run_pool::<256, 4096>(2 * 1024 * 1024); + + // Simulate virtio-net workload + let desc_buf = pool.alloc(64).unwrap(); // Control message + let rx_buf1 = pool.alloc(1500).unwrap(); // MTU packet + let rx_buf2 = pool.alloc(1500).unwrap(); // MTU packet + let tx_buf = pool.alloc(4096).unwrap(); // Large buffer + + // Free and reallocate + pool.dealloc(rx_buf1.addr).unwrap(); + let rx_buf3 = pool.alloc(1500).unwrap(); + + // Should reuse freed buffer (LIFO) + assert_eq!(rx_buf3.addr, rx_buf1.addr); + + pool.dealloc(desc_buf.addr).unwrap(); + pool.dealloc(rx_buf2.addr).unwrap(); + pool.dealloc(rx_buf3.addr).unwrap(); + pool.dealloc(tx_buf.addr).unwrap(); +} + +#[test] +fn test_run_pool_zero_allocation_error() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let result = pool.alloc(0); + assert!(matches!(result, Err(AllocError::InvalidArg))); +} + +#[test] +fn test_run_pool_too_large_allocation() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + let result = pool.alloc(2 * 1024 * 1024); // Larger than pool + assert!(matches!(result, Err(AllocError::OutOfMemory))); +} + +#[test] +fn test_align_up_helper() { + assert_eq!(align_up(0, 256).unwrap(), 0); + assert_eq!(align_up(1, 256).unwrap(), 256); + assert_eq!(align_up(256, 256).unwrap(), 256); + assert_eq!(align_up(257, 256).unwrap(), 512); + assert_eq!(align_up(511, 256).unwrap(), 512); + assert_eq!(align_up(512, 256).unwrap(), 512); + assert!(matches!(align_up(1, 0), Err(AllocError::InvalidArg))); + assert!(matches!( + align_up(usize::MAX, 256), + Err(AllocError::Overflow) + )); +} + +#[test] +fn test_slot_pool_preserves_exact_base() { + let layout = SlotLayout::new(0x80001, 4096, 2); + let pool = SlotPool::new(layout).unwrap(); + + assert_eq!(pool.base_addr(), 0x80001); + assert_eq!(pool.count(), 2); + assert_eq!(pool.slot_addr(0), Some(0x80001)); + assert_eq!(pool.slot_addr(1), Some(0x81001)); +} + +#[test] +fn test_tiered_slot_pool_reports_layouts() { + let lower = SlotLayout::new(0x80001, 0x100, 2); + let upper = SlotLayout::new(0x90001, 0x1000, 2); + let pool = SlotPool::new_tiered(lower, upper).unwrap(); + + let (lower, upper) = pool.layouts(); + assert_eq!(lower, Some(SlotLayout::new(0x80001, 0x100, 2))); + assert_eq!(upper, SlotLayout::new(0x90001, 0x1000, 2)); + assert_eq!(pool.base_addr(), 0x80001); + assert_eq!(pool.slot_size(), 0x1000); + assert_eq!(pool.count(), 4); + assert_eq!(pool.slot_addr(0), Some(0x80001)); + assert_eq!(pool.slot_addr(1), Some(0x80101)); + assert_eq!(pool.slot_addr(2), Some(0x90001)); + assert_eq!(pool.slot_addr(3), Some(0x91001)); + assert_eq!(pool.slot_addr(4), None); +} + +#[test] +fn test_tiered_slot_pool_rejects_invalid_layout() { + let lower = SlotLayout::new(0x80000, 0x100, 32); + let overlapping_upper = SlotLayout::new(0x81000, 0x1000, 2); + let overlapping = SlotPool::new_tiered(lower, overlapping_upper); + assert!(matches!(overlapping, Err(AllocError::InvalidArg))); + + let lower = SlotLayout::new(0x80000, 0x1000, 2); + let smaller_upper = SlotLayout::new(0x90000, 0x100, 32); + let reversed_sizes = SlotPool::new_tiered(lower, smaller_upper); + assert!(matches!(reversed_sizes, Err(AllocError::InvalidArg))); +} + +#[test] +fn test_tiered_slot_pool_routes_by_size() { + let pool = make_tiered_slot_pool(2, 2); + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(257).unwrap(); + + assert_eq!(lower.len, 256); + assert!((0x80000..0x80200).contains(&lower.addr)); + assert_eq!(upper.len, 4096); + assert!((0x90000..0x92000).contains(&upper.addr)); + assert_eq!(pool.allocation_len(lower.addr).unwrap(), 256); + assert_eq!(pool.allocation_len(upper.addr).unwrap(), 4096); +} + +#[test] +fn test_tiered_slot_pool_lower_falls_back_when_full() { + let pool = make_tiered_slot_pool(1, 2); + + let lower = pool.alloc(128).unwrap(); + let fallback = pool.alloc(128).unwrap(); + + assert_eq!(lower.len, 256); + assert_eq!(fallback.len, 4096); + assert!((0x90000..0x92000).contains(&fallback.addr)); +} + +#[test] +fn test_tiered_slot_pool_does_not_mask_lower_errors() { + let pool = make_tiered_slot_pool(1, 1); + + assert!(matches!(pool.alloc(0), Err(AllocError::InvalidArg))); + assert!(matches!(pool.alloc(4097), Err(AllocError::OutOfMemory))); + assert_eq!(pool.num_free(), 2); +} + +#[test] +fn test_tiered_slot_pool_alloc_sg_uses_both_tiers() { + let pool = make_tiered_slot_pool(1, 2); + let sgs = pool.alloc_sg(4096 + 128).unwrap(); + + assert_eq!(sgs.len(), 2); + assert_eq!(sgs[0].len, 4096); + assert_eq!(sgs[1].len, 256); + assert!((0x90000..0x92000).contains(&sgs[0].addr)); + assert!((0x80000..0x80100).contains(&sgs[1].addr)); + + for sg in sgs { + pool.dealloc(sg.addr).unwrap(); + } + assert_eq!(pool.num_free(), 3); +} + +#[test] +fn test_tiered_slot_pool_dealloc_routes_by_region() { + let pool = make_tiered_slot_pool(1, 1); + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(1024).unwrap(); + + pool.dealloc(lower.addr).unwrap(); + pool.dealloc(upper.addr).unwrap(); + assert_eq!(pool.num_free(), 2); + assert!(matches!( + pool.dealloc(lower.addr), + Err(AllocError::InvalidFree(_, _)) + )); + assert!(matches!( + pool.dealloc(0x88000), + Err(AllocError::InvalidFree(_, _)) + )); +} + +// Edge case: allocation exactly at boundary +#[test] +fn test_run_pool_boundary_allocation() { + let pool = make_run_pool::<256, 4096>(1024 * 1024); + + // Allocate exactly at boundary + let alloc = pool.alloc(256).unwrap(); + assert!(pool.inner.borrow().lower.contains(alloc.addr)); + + // Allocate just over boundary + let alloc2 = pool.alloc(257).unwrap(); + assert!(pool.inner.borrow().upper.contains(alloc2.addr)); +} + +#[test] +fn test_run_pool_dealloc_addr_routes_to_correct_tier() { + let pool = make_run_pool::<256, 4096>(0x20000); + let lower = pool.alloc(128).unwrap(); + let upper = pool.alloc(1024).unwrap(); + + assert_eq!(pool.allocation_len(lower.addr).unwrap(), 256); + assert_eq!(pool.allocation_len(upper.addr).unwrap(), 4096); + + pool.dealloc_addr(lower.addr).unwrap(); + pool.dealloc_addr(upper.addr).unwrap(); +} + +#[test] +fn test_run_pool_alloc_sg_uses_one_contiguous_run() { + let pool = make_run_pool::<256, 4096>(0x20000); + let sgs = pool.alloc_sg(4096 * 2 + 1).unwrap(); + + assert_eq!(sgs.len(), 1); + assert_eq!(sgs[0].len, 4096 * 3); + + for sg in sgs { + pool.dealloc(sg.addr).unwrap(); + } +} + +#[test] +fn test_run_pool_alloc_sg_large_run() { + let pool = make_run_pool::<256, 4096>(0x20000); + let sgs = pool.alloc_sg(8192).unwrap(); + + assert_eq!(sgs.len(), 1); + assert_eq!(sgs[0].len, 8192); + + for sg in sgs { + pool.dealloc(sg.addr).unwrap(); + } +} + +#[test] +fn test_slot_pool_alloc_sg_splits() { + let pool = make_slot_pool(8, 4096); + let sgs = pool.alloc_sg(4096 * 2 + 1).unwrap(); + + assert_eq!(sgs.len(), 3); + assert_eq!(sgs[0].len, 4096); + assert_eq!(sgs[1].len, 4096); + assert_eq!(sgs[2].len, 4096); + + for sg in sgs { + pool.dealloc(sg.addr).unwrap(); + } +} + +#[test] +fn test_tiered_slot_pool_live_addrs_are_deterministic() { + let pool = make_tiered_slot_pool(2, 2); + let lower_high = pool.alloc(128).unwrap(); + let upper_high = pool.alloc(1024).unwrap(); + let lower_low = pool.alloc(128).unwrap(); + + assert_eq!( + pool.live_addrs(), + vec![lower_low.addr, lower_high.addr, upper_high.addr] + ); +} + +#[test] +fn test_slot_pool_dealloc_out_of_range() { + let pool = make_slot_pool(4, 4096); + let _ = pool.alloc(4096).unwrap(); + + assert!(matches!( + pool.dealloc(0xDEAD), + Err(AllocError::InvalidFree(0xDEAD, 0)) + )); +} + +#[test] +fn test_slot_pool_dealloc_misaligned() { + let pool = make_slot_pool(4, 4096); + let _ = pool.alloc(4096).unwrap(); + + assert!(matches!( + pool.dealloc(0x80001), + Err(AllocError::InvalidFree(0x80001, 0)) + )); +} + +#[test] +fn test_slot_pool_dealloc_double_free() { + let pool = make_slot_pool(4, 4096); + let a = pool.alloc(4096).unwrap(); + pool.dealloc(a.addr).unwrap(); + + // Second dealloc should fail - address is already in the free list + assert!(matches!( + pool.dealloc(a.addr), + Err(AllocError::InvalidFree(_, _)) + )); +} + +#[test] +fn test_slot_pool_alloc_sg_rolls_back_on_failure() { + let pool = make_slot_pool(2, 4096); + + assert!(matches!(pool.alloc_sg(4096 * 3), Err(AllocError::NoSpace))); + assert_eq!(pool.num_free(), 2); + + let alloc = pool.alloc(4096).unwrap(); + assert_eq!(pool.num_free(), 1); + pool.dealloc(alloc.addr).unwrap(); +} + +#[test] +fn test_slot_pool_dealloc_addr_and_allocation_len() { + let pool = make_slot_pool(4, 4096); + let alloc = pool.alloc(4096).unwrap(); + + assert_eq!(pool.allocation_len(alloc.addr).unwrap(), 4096); + pool.dealloc_addr(alloc.addr).unwrap(); + assert!(matches!( + pool.allocation_len(alloc.addr), + Err(AllocError::InvalidFree(_, 0)) + )); +} + +#[test] +fn test_slot_pool_random_order_dealloc() { + let pool = make_slot_pool(8, 4096); + + let mut allocs: Vec = (0..8).map(|_| pool.alloc(4096).unwrap()).collect(); + assert_eq!(pool.num_free(), 0); + + // Dealloc in reverse order + allocs.reverse(); + for a in &allocs { + pool.dealloc(a.addr).unwrap(); + } + assert_eq!(pool.num_free(), 8); + + // All slots should be re-allocatable + let reallocs: Vec = (0..8).map(|_| pool.alloc(4096).unwrap()).collect(); + assert_eq!(pool.num_free(), 0); + + // Verify all addresses are distinct + let mut addrs: Vec = reallocs.iter().map(|a| a.addr).collect(); + addrs.sort(); + addrs.dedup(); + assert_eq!(addrs.len(), 8); +} + +#[test] +fn test_slot_pool_interleaved_alloc_dealloc_order() { + let pool = make_slot_pool(4, 4096); + + let a0 = pool.alloc(4096).unwrap(); + let a1 = pool.alloc(4096).unwrap(); + let a2 = pool.alloc(4096).unwrap(); + let a3 = pool.alloc(4096).unwrap(); + assert_eq!(pool.num_free(), 0); + + // Free middle slots first (out of allocation order) + pool.dealloc(a2.addr).unwrap(); + pool.dealloc(a0.addr).unwrap(); + assert_eq!(pool.num_free(), 2); + + // Re-alloc gets the out-of-order slots back (LIFO) + let b0 = pool.alloc(4096).unwrap(); + assert_eq!(b0.addr, a0.addr); + let b1 = pool.alloc(4096).unwrap(); + assert_eq!(b1.addr, a2.addr); + + // Free everything in yet another order + pool.dealloc(a1.addr).unwrap(); + pool.dealloc(b0.addr).unwrap(); + pool.dealloc(b1.addr).unwrap(); + pool.dealloc(a3.addr).unwrap(); + assert_eq!(pool.num_free(), 4); + + // All 4 original addresses should be available + let mut final_addrs: Vec = (0..4).map(|_| pool.alloc(4096).unwrap().addr).collect(); + final_addrs.sort(); + let expected: Vec = (0..4).map(|i| 0x80000 + i * 4096).collect(); + assert_eq!(final_addrs, expected); +} + +#[test] +fn test_slot_pool_dealloc_order_independent_of_alloc_order() { + let pool = make_slot_pool(6, 256); + + // Allocate all + let allocs: Vec = (0..6).map(|_| pool.alloc(256).unwrap()).collect(); + + // Dealloc in scattered order: 4, 1, 5, 0, 3, 2 + let order = [4, 1, 5, 0, 3, 2]; + for &i in &order { + pool.dealloc(allocs[i].addr).unwrap(); + } + assert_eq!(pool.num_free(), 6); + + // Re-allocate all and verify we get back the full set + let mut realloc_addrs: Vec = (0..6).map(|_| pool.alloc(256).unwrap().addr).collect(); + realloc_addrs.sort(); + + let mut orig_addrs: Vec = allocs.iter().map(|a| a.addr).collect(); + orig_addrs.sort(); + + assert_eq!(realloc_addrs, orig_addrs); +} diff --git a/src/hyperlight_common/src/virtq/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index 635be0816..9677dbe84 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -1550,7 +1550,8 @@ mod tests { fn test_chain_build_rolls_back_unrepresentable_allocations() { let ring = make_ring(16); let slot_size = u32::MAX as usize + 1; - let pool = RecyclePool::new(0, slot_size, slot_size).unwrap(); + let layout = SlotLayout::new(0, slot_size, 1); + let pool = SlotPool::new(layout).unwrap(); let mem = ring.mem(); let producer = VirtqProducer::new(ring.layout(), mem, TestNotifier::new(), pool.clone()); From ab25a50a5a3d2ab5be7d0ea6811edbba9c0006b1 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Thu, 23 Jul 2026 14:55:22 +0200 Subject: [PATCH 06/34] feat(virtq): add canonical packed ring images Define deterministic producer and consumer reset behavior. Validate canonical events, descriptor chains, IDs, buffer policy, and unused descriptors. Signed-off-by: Tomasz Andrzejak --- fuzz/README.md | 2 +- fuzz/fuzz_targets/virtq_packed_ring.rs | 81 ++- src/hyperlight_common/src/virtq/consumer.rs | 3 +- src/hyperlight_common/src/virtq/desc.rs | 15 + src/hyperlight_common/src/virtq/event.rs | 9 + src/hyperlight_common/src/virtq/producer.rs | 1 - src/hyperlight_common/src/virtq/ring.rs | 488 +++++--------- .../src/virtq/ring/canonical.rs | 611 ++++++++++++++++++ src/hyperlight_common/src/virtq/ring/fuzz.rs | 204 ++++++ 9 files changed, 1076 insertions(+), 338 deletions(-) create mode 100644 src/hyperlight_common/src/virtq/ring/canonical.rs create mode 100644 src/hyperlight_common/src/virtq/ring/fuzz.rs diff --git a/fuzz/README.md b/fuzz/README.md index b08611786..2144ad6c1 100644 --- a/fuzz/README.md +++ b/fuzz/README.md @@ -10,7 +10,7 @@ which evaluates to the following command `cargo +nightly fuzz run fuzz_host_prin As per Microsoft's Offensive Research & Security Engineering (MORSE) team, all host exposed functions that receive or interact with guest data must be continuously fuzzed for, at least, 500 million fuzz test cases without any crashes. Because `cargo-fuzz` doesn't support setting a maximum number of iterations; instead, we use the `--max_total_time` flag to set a maximum time to run the fuzzer. We have a GitHub action (acting like a CRON job) that runs the fuzzers for 24 hours every week. -Currently, we fuzz the parameters and return type to a hardcoded `PrintOutput` guest function, the `HostPrint` host function, and the packed virtqueue ring parser. We plan to add more fuzzers in the future. +Currently, we fuzz the parameters and return type to a hardcoded `PrintOutput` guest function, the `HostPrint` host function, the packed virtqueue ring parser, and canonical ring image validation. We plan to add more fuzzers in the future. ## On Failure diff --git a/fuzz/fuzz_targets/virtq_packed_ring.rs b/fuzz/fuzz_targets/virtq_packed_ring.rs index 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/src/virtq/consumer.rs b/src/hyperlight_common/src/virtq/consumer.rs index 22ce9bc13..f869f2229 100644 --- a/src/hyperlight_common/src/virtq/consumer.rs +++ b/src/hyperlight_common/src/virtq/consumer.rs @@ -641,12 +641,13 @@ impl VirtqConsumer { /// # Errors /// /// - [`VirtqError::InvalidState`] - one or more chains are still in flight + /// - [`VirtqError::RingError`] - device-event normalization failed pub fn reset(&mut self) -> Result<(), VirtqError> { if self.inflight.ones().next().is_some() { return Err(VirtqError::InvalidState); } - self.inner.reset(); + self.inner.reset()?; self.inflight.clear(); Ok(()) } diff --git a/src/hyperlight_common/src/virtq/desc.rs b/src/hyperlight_common/src/virtq/desc.rs index 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/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index 9677dbe84..3f46786c3 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -1809,5 +1809,4 @@ mod tests { )); assert_eq!(producer.inner.num_inflight(), 1); } - } diff --git a/src/hyperlight_common/src/virtq/ring.rs b/src/hyperlight_common/src/virtq/ring.rs index 060c70765..1cb2f2635 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}; @@ -167,6 +169,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. @@ -905,45 +911,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(()) } } @@ -1322,14 +1321,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 +1414,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 +1528,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 +3311,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 +3332,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 +3342,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 +3401,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 +3423,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 +4260,4 @@ mod virtio_villain { } #[cfg(test)] -mod fuzz { - use quickcheck::{Arbitrary, Gen, QuickCheck}; - - use super::tests::{OwnedRing, make_consumer, make_producer}; - use super::*; - - const MAX_RING: usize = 64; - const MAX_OPS: usize = 128; - const MAX_CHAIN_LEN: usize = 8; - - #[allow(clippy::large_enum_variant)] - #[derive(Clone, Debug)] - enum Op { - /// submit one chain - Submit(BufferChain), - /// poll up to N chains - PollAvail(u8), - /// driver reclaims up to N completions - PollUsed(u8), - /// complete one previously polled chain - CompleteOne, - } - - impl Arbitrary for Op { - fn arbitrary(g: &mut Gen) -> Self { - let choice = u8::arbitrary(g) % 4; - match choice { - 0 => Op::Submit(BufferChain::arbitrary(g)), - 1 => Op::PollAvail(u8::arbitrary(g) % 8 + 1), - 2 => Op::PollUsed(u8::arbitrary(g) % 8 + 1), - 3 => Op::CompleteOne, - _ => unreachable!(), - } - } - } - - #[derive(Clone, Debug)] - struct Scenario { - table_size: usize, - ops: Vec, - } - - impl Arbitrary for Scenario { - fn arbitrary(g: &mut Gen) -> Self { - let table_size = (usize::arbitrary(g) % MAX_RING + 1).next_power_of_two(); - let num_ops = usize::arbitrary(g) % MAX_OPS + 1; - - let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); - Scenario { table_size, ops } - } - } - - impl Arbitrary for BufferElement { - fn arbitrary(g: &mut Gen) -> Self { - let addr = u64::arbitrary(g); - let len = u32::arbitrary(g); - let writable = bool::arbitrary(g); - - BufferElement { - addr, - len, - writable, - } - } - } - - impl Arbitrary for BufferChain { - fn arbitrary(g: &mut Gen) -> Self { - let chain_len = usize::arbitrary(g) % MAX_CHAIN_LEN + 1; - - let mut elems = vec![BufferElement::zeroed(); chain_len]; - let mut readables = 0; - let mut writables = 0; - - for _ in 0..chain_len { - let elem = BufferElement::arbitrary(g); - if elem.writable { - elems[chain_len - 1 - writables] = elem; - writables += 1; - } else { - elems[readables] = elem; - readables += 1; - } - } - - BufferChain { - elems: elems.into(), - split: readables, - } - } - } - - fn run_scenario(s: Scenario) -> bool { - let ring = OwnedRing::new(s.table_size); - let mut producer = make_producer(&ring); - let mut consumer = make_consumer(&ring); - - // Order logs - let mut dev_order: Vec = Vec::new(); - let mut drv_order: Vec = Vec::new(); - - // Device-tracked polled-but-not-completed IDs - let mut dev_ready: Vec<(u16, u32)> = Vec::new(); - - for op in &s.ops { - match op { - Op::Submit(chain) => { - // Submit only if space; otherwise skip - let _ = producer.submit_available(chain); - } - Op::PollAvail(n) => { - for _ in 0..*n { - if let Ok((id, chain)) = consumer.poll_available() { - dev_ready.push((id, chain.len() as u32)); - } else { - break; - } - } - } - Op::PollUsed(n) => { - for _ in 0..*n { - match producer.poll_used() { - Ok(u) => { - drv_order.push(u.id); - if producer.id_num[u.id as usize] != 0 { - return false; - } - if !producer.id_free.contains(&u.id) { - return false; - } - } - Err(RingError::WouldBlock) => break, - Err(_) => return false, - } - } - } - Op::CompleteOne => { - if let Some((id, len)) = dev_ready.pop() { - if consumer.submit_used(id, len).is_err() { - return false; - } - - dev_order.push(id); - } - } - } - - // assert invariants after each op - let outstanding: u16 = producer.id_num.iter().copied().sum(); - if outstanding as usize + producer.num_free != ring.len() { - return false; - } - - for id in producer.id_free.iter() { - if producer.id_num[*id as usize] != 0 { - return false; - } - } - } - - // Drain remaining completions and reclaims - while let Some((id, len)) = dev_ready.pop() { - if consumer.submit_used(id, len).is_err() { - return false; - } - } - - loop { - match producer.poll_used() { - Ok(u) => drv_order.push(u.id), - Err(RingError::WouldBlock) => break, - Err(_) => return false, - } - } - - true - } - - #[test] - fn prop_interleaved_with_order_verification() { - #[cfg(miri)] - let tests = 1; - #[cfg(not(miri))] - let tests = 100; - - QuickCheck::new() - .tests(tests) - .quickcheck(run_scenario as fn(Scenario) -> bool); - } -} +mod fuzz; diff --git a/src/hyperlight_common/src/virtq/ring/canonical.rs b/src/hyperlight_common/src/virtq/ring/canonical.rs new file mode 100644 index 000000000..684ecf68b --- /dev/null +++ b/src/hyperlight_common/src/virtq/ring/canonical.rs @@ -0,0 +1,611 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + */ + +//! Canonical packed virtqueue images. +//! +//! A canonical image starts at the initial wrap round. Available descriptors +//! occupy a complete prefix, unused descriptors are zero, and both event +//! structures are enabled at offset zero. This form can be restored as bytes +//! and validated before either peer resumes. +//! +//! Ring resets normalize the structures owned by each peer. Buffer range +//! policy remains with the integration through the validator callback. + +use alloc::vec::Vec; + +use bytemuck::Zeroable; +use fixedbitset::FixedBitSet; +use smallvec::SmallVec; +use thiserror::Error; + +use super::super::desc::{DescFlags, DescTable, Descriptor}; +use super::super::event::{EventFlags, EventSuppression}; +use super::super::{Layout, MemOps}; +use super::{BufferChain, BufferElement, MemOp, RingError}; + +/// Why a descriptor does not belong to a canonical packed-ring image. +#[derive(Error, Debug, Copy, Clone, PartialEq, Eq)] +pub enum DescError { + /// An unused descriptor was not completely zeroed. + #[error("unused descriptor is not zeroed")] + ExpectedZero, + /// The raw descriptor contains unsupported or reserved flag bits. + #[error("descriptor contains unknown flags")] + UnknownFlags, + /// The descriptor is not available in the initial packed-ring wrap round. + #[error("descriptor is not initially available")] + NotAvailable, + /// Indirect descriptor tables are unsupported. + #[error("indirect descriptor is unsupported")] + Indirect, + /// A chain's NEXT flag extends beyond the available descriptor prefix. + #[error("chain continues beyond the available descriptor prefix")] + ChainContinues, + /// A chain ID is outside the descriptor-table bounds. + #[error("descriptor ID is out of range")] + IdOutOfRange, + /// A tail descriptor does not carry its head descriptor's ID. + #[error("descriptor ID differs within a chain")] + IdMismatch, + /// Two available chains use the same descriptor ID. + #[error("descriptor ID is already used by another chain")] + DuplicateId, + /// A readable descriptor follows a writable descriptor. + #[error("readable descriptor follows a writable descriptor")] + ReadableAfterWritable, +} + +/// Validation failure for a canonical packed-ring image. +#[derive(Error, Debug)] +pub enum ImageError { + /// Reading the shared ring image failed. + #[error(transparent)] + Ring(#[from] RingError), + /// The caller supplied an impossible available-descriptor prefix length. + #[error("available descriptor count {available} exceeds ring capacity {capacity}")] + DescCount { + /// Number of descriptors expected to be available. + available: usize, + /// Descriptor-table capacity. + capacity: usize, + }, + /// An event-suppression structure is not the canonical enabled value. + #[error("event suppression at address 0x{addr:x} is not canonical")] + Event { + /// Address of the invalid event-suppression structure. + addr: u64, + }, + /// A descriptor violates the canonical packed-ring structure. + #[error("descriptor {index} is not canonical: {reason}")] + Desc { + /// Descriptor-table index. + index: u16, + /// Structural validation failure. + reason: DescError, + }, + /// The caller rejected a descriptor's payload range or attributes. + #[error("descriptor {index} buffer at 0x{addr:x} with length {len} was rejected")] + Buffer { + /// Descriptor-table index. + index: u16, + /// Buffer address from the descriptor. + addr: u64, + /// Buffer length from the descriptor. + len: u32, + }, +} + +impl ImageError { + fn desc_count(available: usize, capacity: usize) -> Self { + Self::DescCount { + available, + capacity, + } + } + + fn event(addr: u64) -> Self { + Self::Event { addr } + } + + fn desc(index: u16, reason: DescError) -> Self { + Self::Desc { index, reason } + } + + fn buffer(index: u16, addr: u64, len: u32) -> Self { + Self::Buffer { index, addr, len } + } +} + +/// One available descriptor chain from a validated canonical ring image. +#[derive(Debug, Clone)] +pub struct CanonChain { + id: u16, + inner: BufferChain, +} + +impl CanonChain { + fn new(id: u16, chain: BufferChain) -> Self { + Self { id, inner: chain } + } + + /// Descriptor ID shared by every buffer in the chain. + pub fn id(&self) -> u16 { + self.id + } + + /// Validated buffers in descriptor order. + pub fn buffers(&self) -> &BufferChain { + &self.inner + } + + /// Consume the image metadata and return its buffer chain. + pub fn into_buffers(self) -> BufferChain { + self.inner + } +} + +/// Validate a packed ring while neither peer can modify it. +/// +/// The first `avail_descs` descriptors must form complete available chains +/// beginning at descriptor zero. Every remaining descriptor must be zeroed, +/// and both event-suppression structures must be the canonical enabled value. +/// `validate_buf` supplies integration-specific address, length, and +/// direction bounds without embedding them in the ring implementation. +/// +/// The returned chains preserve descriptor IDs and chain boundaries for +/// cross-checking against producer and pool ownership. +/// +/// # Errors +/// +/// Returns [`ImageError`] if event state is not normalized, descriptor +/// structure is malformed, an unused descriptor is not zero, a buffer is +/// rejected by `validate_buf`, or shared memory cannot be read. +pub fn validate_canon_image( + mem: &M, + layout: Layout, + avail_descs: usize, + mut validate_buf: F, +) -> Result, ImageError> +where + M: MemOps, + F: FnMut(u16, BufferElement) -> bool, +{ + let cap = layout.desc_table_len() as usize; + if avail_descs > cap { + return Err(ImageError::desc_count(avail_descs, cap)); + } + + let canon_evt = EventSuppression::new(0, EventFlags::ENABLE); + for addr in [layout.drv_evt_addr(), layout.dev_evt_addr()] { + let evt = mem + .read_val::(addr) + .map_err(|_| RingError::mem_err(MemOp::ReadEvent, addr))?; + + if evt != canon_evt { + return Err(ImageError::event(addr)); + } + } + + // SAFETY: `Layout` validates the table base, alignment, and descriptor count. + let table = unsafe { DescTable::from_raw_parts(layout.desc_table_addr(), cap) }; + + let mut seen_ids = FixedBitSet::with_capacity(cap); + let mut chains = Vec::new(); + let mut pos = 0usize; + + while pos < avail_descs { + let head_idx = u16::try_from(pos).map_err(|_| RingError::InvalidState)?; + let (head, _) = read_canon_avail_desc(mem, &table, head_idx)?; + let id_idx = head.id as usize; + if id_idx >= cap { + return Err(ImageError::desc(head_idx, DescError::IdOutOfRange)); + } + + if seen_ids.contains(id_idx) { + return Err(ImageError::desc(head_idx, DescError::DuplicateId)); + } + + seen_ids.insert(id_idx); + + let mut elems = SmallVec::<[BufferElement; 16]>::new(); + let mut split = 0usize; + + loop { + let idx = u16::try_from(pos).map_err(|_| RingError::InvalidState)?; + let (desc, flags) = read_canon_avail_desc(mem, &table, idx)?; + if desc.id != head.id { + return Err(ImageError::desc(idx, DescError::IdMismatch)); + } + + let elem = BufferElement::from(&desc); + if !elem.writable && split != elems.len() { + return Err(ImageError::desc(idx, DescError::ReadableAfterWritable)); + } + + split += usize::from(!elem.writable); + + if !validate_buf(idx, elem) { + return Err(ImageError::buffer(idx, elem.addr, elem.len)); + } + + elems.push(elem); + pos += 1; + + if !flags.contains(DescFlags::NEXT) { + break; + } + if pos >= avail_descs { + return Err(ImageError::desc(idx, DescError::ChainContinues)); + } + } + + let canon = CanonChain::new(head.id, BufferChain { elems, split }); + chains.push(canon); + } + + let empty = Descriptor::zeroed(); + for pos in avail_descs..cap { + let idx = u16::try_from(pos).map_err(|_| RingError::InvalidState)?; + let addr = table.desc_addr(idx).ok_or(RingError::InvalidState)?; + + let desc = mem + .read_val::(addr) + .map_err(|_| RingError::mem_err(MemOp::ReadDesc, addr))?; + + if desc != empty { + return Err(ImageError::desc(idx, DescError::ExpectedZero)); + } + } + + Ok(chains) +} + +fn read_canon_avail_desc( + mem: &M, + table: &DescTable, + idx: u16, +) -> Result<(Descriptor, DescFlags), ImageError> { + let addr = table.desc_addr(idx).ok_or(RingError::InvalidState)?; + let desc = mem + .read_val::(addr) + .map_err(|_| RingError::mem_err(MemOp::ReadDesc, addr))?; + + let flags = DescFlags::from_bits(desc.flags) + .ok_or_else(|| ImageError::desc(idx, DescError::UnknownFlags))?; + + if flags.contains(DescFlags::INDIRECT) { + return Err(ImageError::desc(idx, DescError::Indirect)); + } + if !flags.is_avail(true) { + return Err(ImageError::desc(idx, DescError::NotAvailable)); + } + + Ok((desc, flags)) +} + +#[cfg(test)] +mod tests { + use super::super::BufferChainBuilder; + use super::super::tests::{OwnedRing, make_consumer, make_producer, make_ring}; + use super::*; + + fn writable_chain(base: u64, lengths: &[u32]) -> BufferChain { + BufferChainBuilder::new() + .writables(lengths.iter().scan(base, |addr, &len| { + let element = BufferElement { + addr: *addr, + len, + writable: true, + }; + *addr += len as u64; + Some(element) + })) + .build() + .unwrap() + } + + fn validate_all(ring: &OwnedRing, avail_descs: usize) -> Result, ImageError> { + validate_canon_image(&ring.mem(), ring.layout(), avail_descs, |_, _| true) + } + + #[test] + fn canon_reset_normalizes_empty_image() { + let ring = make_ring(8); + let mut producer = make_producer(&ring); + let mut consumer = make_consumer(&ring); + + producer.submit_one(0x1000, 64, true).unwrap(); + producer.enable_used_notifications_desc(3, false).unwrap(); + consumer.enable_avail_notifications_desc(5, false).unwrap(); + + producer.reset().unwrap(); + consumer.reset().unwrap(); + + assert!(validate_all(&ring, 0).unwrap().is_empty()); + assert_eq!( + ring.read_driver_event(), + EventSuppression::new(0, EventFlags::ENABLE) + ); + assert_eq!( + ring.read_device_event(), + EventSuppression::new(0, EventFlags::ENABLE) + ); + for index in 0..ring.len() as u16 { + assert_eq!(ring.read_desc(index), Descriptor::zeroed()); + } + } + + #[test] + fn canon_multi_desc_refill_is_visible_to_fresh_consumer() { + let ring = make_ring(8); + let mut producer = make_producer(&ring); + let mut consumer = make_consumer(&ring); + + producer.reset().unwrap(); + consumer.reset().unwrap(); + + let chains = [ + writable_chain(0x1000, &[64, 128, 256]), + writable_chain(0x2000, &[512, 1024]), + writable_chain(0x4000, &[64, 64, 64]), + ]; + let mut expected = Vec::new(); + for chain in &chains { + let id = producer.submit_available(chain).unwrap(); + expected.push((id, chain.len())); + } + + assert_eq!(producer.num_free(), 0); + assert_eq!(producer.avail_cursor().head(), 0); + assert!(!producer.avail_cursor().wrap()); + + let image = validate_canon_image(&ring.mem(), ring.layout(), ring.len(), |_, elem| { + elem.writable + && elem.len > 0 + && elem + .addr + .checked_add(elem.len as u64) + .is_some_and(|end| end <= 0x5000) + }) + .unwrap(); + + assert_eq!(image.len(), expected.len()); + for (validated, (id, len)) in image.iter().zip(&expected) { + assert_eq!(validated.id(), *id); + assert_eq!(validated.buffers().len(), *len); + assert!(validated.buffers().elems().iter().all(|elem| elem.writable)); + assert_eq!(producer.id_num[*id as usize] as usize, *len); + } + + let mut fresh = make_consumer(&ring); + for (expected_id, expected_len) in expected { + let (id, chain) = fresh.poll_available().unwrap(); + assert_eq!(id, expected_id); + assert_eq!(chain.len(), expected_len); + } + assert!(matches!(fresh.poll_available(), Err(RingError::WouldBlock))); + } + + #[test] + fn canon_image_rejects_invalid_ids() { + let duplicate_ring = make_ring(4); + let mut producer = make_producer(&duplicate_ring); + producer.submit_one(0x1000, 64, true).unwrap(); + producer.submit_one(0x2000, 64, true).unwrap(); + + let head_id = duplicate_ring.read_desc(0).id; + let mut duplicate = duplicate_ring.read_desc(1); + duplicate.id = head_id; + duplicate_ring.write_desc(1, duplicate); + + assert!(matches!( + validate_all(&duplicate_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::DuplicateId, + }) + )); + + let mismatch_ring = make_ring(4); + let mut producer = make_producer(&mismatch_ring); + producer + .submit_available(&writable_chain(0x3000, &[64, 64])) + .unwrap(); + + let mut tail = mismatch_ring.read_desc(1); + tail.id = tail.id.wrapping_sub(1); + mismatch_ring.write_desc(1, tail); + + assert!(matches!( + validate_all(&mismatch_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::IdMismatch, + }) + )); + + let range_ring = make_ring(4); + let mut producer = make_producer(&range_ring); + producer.submit_one(0x4000, 64, true).unwrap(); + + let mut out_of_range = range_ring.read_desc(0); + out_of_range.id = range_ring.len() as u16; + range_ring.write_desc(0, out_of_range); + + assert!(matches!( + validate_all(&range_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::IdOutOfRange, + }) + )); + } + + #[test] + fn canon_image_rejects_invalid_flags_and_wrap_state() { + let unknown_ring = make_ring(4); + let mut producer = make_producer(&unknown_ring); + producer.submit_one(0x1000, 64, true).unwrap(); + + let mut unknown = unknown_ring.read_desc(0); + unknown.flags |= 1 << 3; + unknown_ring.write_desc(0, unknown); + + assert!(matches!( + validate_all(&unknown_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::UnknownFlags, + }) + )); + + let used_ring = make_ring(4); + let mut producer = make_producer(&used_ring); + producer.submit_one(0x2000, 64, true).unwrap(); + + let mut used = used_ring.read_desc(0); + used.mark_used(true); + used_ring.write_desc(0, used); + + assert!(matches!( + validate_all(&used_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::NotAvailable, + }) + )); + + let indirect_ring = make_ring(4); + let mut producer = make_producer(&indirect_ring); + producer.submit_one(0x3000, 64, true).unwrap(); + + let mut indirect = indirect_ring.read_desc(0); + indirect.flags |= DescFlags::INDIRECT.bits(); + indirect_ring.write_desc(0, indirect); + + assert!(matches!( + validate_all(&indirect_ring, 1), + Err(ImageError::Desc { + index: 0, + reason: DescError::Indirect, + }) + )); + } + + #[test] + fn canon_image_rejects_malformed_chain_and_unused_desc() { + let chain_ring = make_ring(4); + let mut producer = make_producer(&chain_ring); + producer + .submit_available(&writable_chain(0x1000, &[64, 64])) + .unwrap(); + + let mut tail = chain_ring.read_desc(1); + tail.flags |= DescFlags::NEXT.bits(); + chain_ring.write_desc(1, tail); + + assert!(matches!( + validate_all(&chain_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::ChainContinues, + }) + )); + + let unused_ring = make_ring(4); + let mut producer = make_producer(&unused_ring); + producer.submit_one(0x2000, 64, true).unwrap(); + unused_ring.write_desc(3, Descriptor::new(0x3000, 64, 0, DescFlags::empty())); + + assert!(matches!( + validate_all(&unused_ring, 1), + Err(ImageError::Desc { + index: 3, + reason: DescError::ExpectedZero, + }) + )); + + let direction_ring = make_ring(4); + let mut producer = make_producer(&direction_ring); + producer + .submit_available(&writable_chain(0x4000, &[64, 64])) + .unwrap(); + + let mut readable_tail = direction_ring.read_desc(1); + readable_tail.flags &= !DescFlags::WRITE.bits(); + direction_ring.write_desc(1, readable_tail); + + assert!(matches!( + validate_all(&direction_ring, 2), + Err(ImageError::Desc { + index: 1, + reason: DescError::ReadableAfterWritable, + }) + )); + } + + #[test] + fn canon_image_rejects_noncanon_evt_and_buffer_bounds() { + let evt_ring = make_ring(4); + evt_ring + .mem() + .write_val( + evt_ring.layout().drv_evt_addr(), + EventSuppression::new(1, EventFlags::ENABLE), + ) + .unwrap(); + + match validate_all(&evt_ring, 0) { + Err(ImageError::Event { addr }) => { + let expected = evt_ring.layout().drv_evt_addr(); + assert_eq!(addr, expected); + } + other => unreachable!("unexpected result: {other:?}"), + } + + let bounds = make_ring(4); + let mut producer = make_producer(&bounds); + producer.submit_one(u64::MAX - 15, 32, true).unwrap(); + + let res = validate_canon_image(&bounds.mem(), bounds.layout(), 1, |_, element| { + element + .addr + .checked_add(element.len as u64) + .is_some_and(|end| end <= 0x8000) + }); + + match res { + Err(ImageError::Buffer { index, addr, len }) => { + assert_eq!(index, 0); + assert_eq!(addr, u64::MAX - 15); + assert_eq!(len, 32); + } + other => unreachable!("unexpected result: {other:?}"), + } + } + + #[test] + fn canon_image_rejects_avail_count_over_capacity() { + let ring = make_ring(4); + assert!(matches!( + validate_all(&ring, 5), + Err(ImageError::DescCount { + available: 5, + capacity: 4, + }) + )); + } +} diff --git a/src/hyperlight_common/src/virtq/ring/fuzz.rs b/src/hyperlight_common/src/virtq/ring/fuzz.rs new file mode 100644 index 000000000..cef6d12e0 --- /dev/null +++ b/src/hyperlight_common/src/virtq/ring/fuzz.rs @@ -0,0 +1,204 @@ +/* +Copyright 2026 The Hyperlight Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +use quickcheck::{Arbitrary, Gen, QuickCheck}; + +use super::tests::{OwnedRing, make_consumer, make_producer}; +use super::*; + +const MAX_RING: usize = 64; +const MAX_OPS: usize = 128; +const MAX_CHAIN_LEN: usize = 8; + +#[allow(clippy::large_enum_variant)] +#[derive(Clone, Debug)] +enum Op { + /// submit one chain + Submit(BufferChain), + /// poll up to N chains + PollAvail(u8), + /// driver reclaims up to N completions + PollUsed(u8), + /// complete one previously polled chain + CompleteOne, +} + +impl Arbitrary for Op { + fn arbitrary(g: &mut Gen) -> Self { + let choice = u8::arbitrary(g) % 4; + match choice { + 0 => Op::Submit(BufferChain::arbitrary(g)), + 1 => Op::PollAvail(u8::arbitrary(g) % 8 + 1), + 2 => Op::PollUsed(u8::arbitrary(g) % 8 + 1), + 3 => Op::CompleteOne, + _ => unreachable!(), + } + } +} + +#[derive(Clone, Debug)] +struct Scenario { + table_size: usize, + ops: Vec, +} + +impl Arbitrary for Scenario { + fn arbitrary(g: &mut Gen) -> Self { + let table_size = (usize::arbitrary(g) % MAX_RING + 1).next_power_of_two(); + let num_ops = usize::arbitrary(g) % MAX_OPS + 1; + + let ops = (0..num_ops).map(|_| Op::arbitrary(g)).collect(); + Scenario { table_size, ops } + } +} + +impl Arbitrary for BufferElement { + fn arbitrary(g: &mut Gen) -> Self { + let addr = u64::arbitrary(g); + let len = u32::arbitrary(g); + let writable = bool::arbitrary(g); + + BufferElement { + addr, + len, + writable, + } + } +} + +impl Arbitrary for BufferChain { + fn arbitrary(g: &mut Gen) -> Self { + let chain_len = usize::arbitrary(g) % MAX_CHAIN_LEN + 1; + + let mut elems = vec![BufferElement::zeroed(); chain_len]; + let mut readables = 0; + let mut writables = 0; + + for _ in 0..chain_len { + let elem = BufferElement::arbitrary(g); + if elem.writable { + elems[chain_len - 1 - writables] = elem; + writables += 1; + } else { + elems[readables] = elem; + readables += 1; + } + } + + BufferChain { + elems: elems.into(), + split: readables, + } + } +} + +fn run_scenario(s: Scenario) -> bool { + let ring = OwnedRing::new(s.table_size); + let mut producer = make_producer(&ring); + let mut consumer = make_consumer(&ring); + + // Order logs + let mut dev_order: Vec = Vec::new(); + let mut drv_order: Vec = Vec::new(); + + // Device-tracked polled-but-not-completed IDs + let mut dev_ready: Vec<(u16, u32)> = Vec::new(); + + for op in &s.ops { + match op { + Op::Submit(chain) => { + // Submit only if space; otherwise skip + let _ = producer.submit_available(chain); + } + Op::PollAvail(n) => { + for _ in 0..*n { + if let Ok((id, chain)) = consumer.poll_available() { + dev_ready.push((id, chain.len() as u32)); + } else { + break; + } + } + } + Op::PollUsed(n) => { + for _ in 0..*n { + match producer.poll_used() { + Ok(u) => { + drv_order.push(u.id); + if producer.id_num[u.id as usize] != 0 { + return false; + } + if !producer.id_free.contains(&u.id) { + return false; + } + } + Err(RingError::WouldBlock) => break, + Err(_) => return false, + } + } + } + Op::CompleteOne => { + if let Some((id, len)) = dev_ready.pop() { + if consumer.submit_used(id, len).is_err() { + return false; + } + + dev_order.push(id); + } + } + } + + // assert invariants after each op + let outstanding: u16 = producer.id_num.iter().copied().sum(); + if outstanding as usize + producer.num_free != ring.len() { + return false; + } + + for id in producer.id_free.iter() { + if producer.id_num[*id as usize] != 0 { + return false; + } + } + } + + // Drain remaining completions and reclaims + while let Some((id, len)) = dev_ready.pop() { + if consumer.submit_used(id, len).is_err() { + return false; + } + } + + loop { + match producer.poll_used() { + Ok(u) => drv_order.push(u.id), + Err(RingError::WouldBlock) => break, + Err(_) => return false, + } + } + + true +} + +#[test] +fn prop_interleaved_with_order_verification() { + #[cfg(miri)] + let tests = 1; + #[cfg(not(miri))] + let tests = 100; + + QuickCheck::new() + .tests(tests) + .quickcheck(run_scenario as fn(Scenario) -> bool); +} From f2274f1caeef39c5e5863daa5e4fd06e81ba8473 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Fri, 24 Jul 2026 16:56:12 +0200 Subject: [PATCH 07/34] refactor(layout): model scratch-top metadata Represent scratch bookkeeping with one repr(C) layout. Derive offsets and assert the host/guest ABI at compile time. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/layout.rs | 53 +++++++++++++++++++++++++---- 1 file changed, 47 insertions(+), 6 deletions(-) diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index bf25a2e0c..d7f3cc971 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -14,6 +14,8 @@ See the License for the specific language governing permissions and limitations under the License. */ +use core::mem::{offset_of, size_of}; + #[cfg_attr(target_arch = "x86_64", path = "arch/amd64/layout.rs")] #[cfg_attr(target_arch = "aarch64", path = "arch/aarch64/layout.rs")] mod arch; @@ -22,12 +24,51 @@ pub use arch::{ SCRATCH_TOP_GPA, SCRATCH_TOP_GVA, SNAPSHOT_PT_GVA_MAX, SNAPSHOT_PT_GVA_MIN, io_page, }; -// offsets down from the top of scratch memory for various things -pub const SCRATCH_TOP_SIZE_OFFSET: u64 = 0x08; -pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = 0x10; -pub const SCRATCH_TOP_SNAPSHOT_PT_GPA_BASE_OFFSET: u64 = 0x18; -pub const SCRATCH_TOP_SNAPSHOT_GENERATION_OFFSET: u64 = 0x20; -pub const SCRATCH_TOP_EXN_STACK_OFFSET: u64 = 0x30; +const EXN_STACK_ALIGNMENT: usize = 16; + +// Fields are listed in ascending-address order. Public offsets are measured +// down from the top of scratch memory. +#[repr(C)] +struct ScratchTopMetadata { + /// Keep the exception stack pointer aligned 16 bytes aligned. + _alignment_padding: [u8; 8], + /// Reserved for future scratch metadata. + _reserved: 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 +} + +const SCRATCH_TOP_RESERVED_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, _reserved)); +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_RESERVED_OFFSET == 0x28); + assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x30); +}; pub fn scratch_base_gpa(size: usize) -> u64 { (SCRATCH_TOP_GPA - size + 1) as u64 From f2c3f46f9ae815c31cf3bbb7595fffa7e1fd3a0b Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Fri, 24 Jul 2026 17:08:54 +0200 Subject: [PATCH 08/34] fix(guest): correct scratch allocator boundary Use the first GPA of the reserved pages as an exclusive limit. Accept allocations ending at the limit and reject address overflow. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/layout.rs | 7 +++++++ .../src/arch/aarch64/prim_alloc.rs | 9 ++------- .../src/arch/amd64/prim_alloc.rs | 12 ++++-------- src/hyperlight_guest/src/prim_alloc.rs | 17 +++++++++++++++++ 4 files changed, 30 insertions(+), 15 deletions(-) diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index d7f3cc971..50f4ca4eb 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -25,6 +25,8 @@ pub use arch::{ }; const EXN_STACK_ALIGNMENT: usize = 16; +/// Pages reserved for the exception stack and scratch-top metadata. +pub const SCRATCH_TOP_RESERVED_PAGES: usize = 2; // Fields are listed in ascending-address order. Public offsets are measured // down from the top of scratch memory. @@ -70,6 +72,11 @@ const _: () = { assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x30); }; +/// Exclusive upper GPA boundary for dynamic scratch allocations. +pub const fn scratch_allocator_limit_gpa() -> u64 { + (SCRATCH_TOP_GPA + 1 - SCRATCH_TOP_RESERVED_PAGES * crate::vmem::PAGE_SIZE) as u64 +} + pub fn scratch_base_gpa(size: usize) -> u64 { (SCRATCH_TOP_GPA - size + 1) as u64 } diff --git a/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs b/src/hyperlight_guest/src/arch/aarch64/prim_alloc.rs index 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/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)); + } +} From a12ee21c048147ffd628df94ea4f7297d29dc53a Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Tue, 28 Jul 2026 00:20:39 +0200 Subject: [PATCH 09/34] feat(virtq): define virtq transport metadata Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/layout.rs | 49 ++++++++++++++++++++++++----- src/hyperlight_guest/src/layout.rs | 41 +++++++++++++++++++----- 2 files changed, 74 insertions(+), 16 deletions(-) diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index 50f4ca4eb..ec2852520 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -32,10 +32,22 @@ pub const SCRATCH_TOP_RESERVED_PAGES: usize = 2; // down from the top of scratch memory. #[repr(C)] struct ScratchTopMetadata { - /// Keep the exception stack pointer aligned 16 bytes aligned. - _alignment_padding: [u8; 8], - /// Reserved for future scratch metadata. - _reserved: u64, + /// Number of pages reserved for the H2G pool. + h2g_pool_pages: u64, + /// Guest-published GPA of the H2G pool. + h2g_pool_gpa: u64, + /// Host-published GPA of the H2G ring. + h2g_ring_gpa: u64, + /// Host-published H2G descriptor count. + h2g_queue_depth: u64, + /// Number of pages reserved for the G2H pool. + g2h_pool_pages: u64, + /// Guest-published GPA of the G2H pool. + g2h_pool_gpa: u64, + /// Host-published GPA of the G2H ring. + g2h_ring_gpa: u64, + /// Host-published G2H descriptor count. + g2h_queue_depth: u64, /// Generation of the snapshot backing the sandbox. snapshot_generation: u64, /// GPA of the snapshot page-table copy in scratch memory. @@ -50,8 +62,22 @@ const fn scratch_top_offset(field_offset: usize) -> u64 { (size_of::() - field_offset) as u64 } -const SCRATCH_TOP_RESERVED_OFFSET: u64 = - scratch_top_offset(offset_of!(ScratchTopMetadata, _reserved)); +pub const SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_queue_depth)); +pub const SCRATCH_TOP_G2H_RING_GPA_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_ring_gpa)); +pub const SCRATCH_TOP_G2H_POOL_GPA_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_gpa)); +pub const SCRATCH_TOP_G2H_POOL_PAGES_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_pages)); +pub const SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_queue_depth)); +pub const SCRATCH_TOP_H2G_RING_GPA_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_ring_gpa)); +pub const SCRATCH_TOP_H2G_POOL_GPA_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_gpa)); +pub const SCRATCH_TOP_H2G_POOL_PAGES_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_pages)); pub const SCRATCH_TOP_SIZE_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, scratch_size)); pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = @@ -68,8 +94,15 @@ const _: () = { 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_RESERVED_OFFSET == 0x28); - assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x30); + assert!(SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET == 0x28); + assert!(SCRATCH_TOP_G2H_RING_GPA_OFFSET == 0x30); + assert!(SCRATCH_TOP_G2H_POOL_GPA_OFFSET == 0x38); + assert!(SCRATCH_TOP_G2H_POOL_PAGES_OFFSET == 0x40); + assert!(SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET == 0x48); + assert!(SCRATCH_TOP_H2G_RING_GPA_OFFSET == 0x50); + assert!(SCRATCH_TOP_H2G_POOL_GPA_OFFSET == 0x58); + assert!(SCRATCH_TOP_H2G_POOL_PAGES_OFFSET == 0x60); + assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x60); }; /// Exclusive upper GPA boundary for dynamic scratch allocations. diff --git a/src/hyperlight_guest/src/layout.rs b/src/hyperlight_guest/src/layout.rs index 6d132ae7c..c80edf5f5 100644 --- a/src/hyperlight_guest/src/layout.rs +++ b/src/hyperlight_guest/src/layout.rs @@ -19,20 +19,45 @@ 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_depth_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET) +} +pub fn g2h_ring_gpa_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_RING_GPA_OFFSET) +} +pub fn h2g_ring_gpa_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_RING_GPA_OFFSET) +} +pub fn g2h_pool_gpa_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_POOL_GPA_OFFSET) +} +pub fn h2g_pool_gpa_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_POOL_GPA_OFFSET) +} +pub fn g2h_pool_pages_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_POOL_PAGES_OFFSET) +} +pub fn h2g_queue_depth_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET) +} +pub fn h2g_pool_pages_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_POOL_PAGES_OFFSET) } pub use arch::{scratch_base_gpa, scratch_base_gva}; From 3e81faa8294ab8f512928fff05946c7a0c674a1d Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Tue, 28 Jul 2026 19:39:07 +0200 Subject: [PATCH 10/34] feat(virtq): implement host side memory access Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/virtq/consumer.rs | 30 ++- src/hyperlight_host/src/mem/mod.rs | 2 + src/hyperlight_host/src/mem/shared_mem.rs | 165 ++++++++++++++++ src/hyperlight_host/src/mem/virtq_mem.rs | 207 ++++++++++++++++++++ 4 files changed, 395 insertions(+), 9 deletions(-) create mode 100644 src/hyperlight_host/src/mem/virtq_mem.rs diff --git a/src/hyperlight_common/src/virtq/consumer.rs b/src/hyperlight_common/src/virtq/consumer.rs index f869f2229..9135b7096 100644 --- a/src/hyperlight_common/src/virtq/consumer.rs +++ b/src/hyperlight_common/src/virtq/consumer.rs @@ -405,6 +405,7 @@ impl AckChain { /// ``` pub struct VirtqConsumer { inner: RingConsumer, + mem: M, notifier: N, inflight: FixedBitSet, next_token: u32, @@ -419,11 +420,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, @@ -499,18 +506,16 @@ impl VirtqConsumer { } let chain = RecvChain::new( - self.inner.mem().clone(), + self.mem.clone(), token, readables.iter().copied().collect(), recv_len, ); let reply = if !writables.is_empty() { - let writable = WritableChain::new( - self.inner.mem().clone(), - token, - writables.iter().copied().collect(), - ); + let mem = self.mem.clone(); + let elems = writables.iter().copied().collect(); + let writable = WritableChain::new(mem, token, elems); ReplyChain::Writable(writable) } else { let ack = AckChain::new(token); @@ -922,12 +927,19 @@ mod tests { .unwrap(); ring_producer.submit_available(&chain).unwrap(); - let guarded_mem = FailingPayloadReadMem { + let ring_mem = FailingPayloadReadMem { + inner: mem.clone(), + payload_addr, + payload_len: 0, + }; + let mem = FailingPayloadReadMem { inner: mem, payload_addr, payload_len: 4, }; - let mut consumer = VirtqConsumer::new(ring.layout(), guarded_mem, TestNotifier::new()); + + let mut consumer = + VirtqConsumer::new_split(ring.layout(), ring_mem, mem, TestNotifier::new()); let (recv, reply) = consumer.poll(4).unwrap().unwrap(); assert!(matches!(recv.to_bytes(), Err(VirtqError::MemoryReadError))); diff --git a/src/hyperlight_host/src/mem/mod.rs b/src/hyperlight_host/src/mem/mod.rs index 64f5db2fe..8130e17b6 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; +#[allow(dead_code)] +pub(crate) mod virtq_mem; diff --git a/src/hyperlight_host/src/mem/shared_mem.rs b/src/hyperlight_host/src/mem/shared_mem.rs index 4b706cc41..cd4ed757b 100644 --- a/src/hyperlight_host/src/mem/shared_mem.rs +++ b/src/hyperlight_host/src/mem/shared_mem.rs @@ -20,6 +20,7 @@ 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 +62,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 +132,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 +1200,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<()> { @@ -1791,6 +1907,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 +1979,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; 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..d7ecfdaf2 --- /dev/null +++ b/src/hyperlight_host/src/mem/virtq_mem.rs @@ -0,0 +1,207 @@ +/* +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`] access to a bounded scratch region. +//! +//! Every operation uses [`HostSharedMemory`]'s checked API and acquires its +//! lifecycle read lock. This preserves exclusive-memory coordination but makes +//! descriptor traversal pay for one lock acquisition per field access. + +use core::mem::size_of; +use core::ops::Range; +use core::sync::atomic::{AtomicU16, Ordering}; + +use hyperlight_common::layout::scratch_base_gva; +use hyperlight_common::virtq::MemOps; + +use super::shared_mem::{HostSharedMemory, SharedMemory}; +use crate::{HyperlightError, Result, new_error}; + +/// Host virtqueue memory access confined to one scratch GVA range. +/// +/// Accepted guest virtual addresses are translated relative to +/// `scratch_base_gva` and delegated to `scratch_mem`. Separate instances +/// confine ring metadata and payload pools independently. Clones share the +/// backing mapping and lifecycle lock while retaining the same range. +#[derive(Clone)] +pub(crate) struct HostMemOps { + /// Shared scratch mapping used for checked memory operations. + scratch_mem: HostSharedMemory, + /// Guest virtual address corresponding to offset zero in `scratch_mem`. + scratch_base_gva: u64, + /// End-exclusive guest virtual address range accepted by this accessor. + region: Range, +} + +impl HostMemOps { + /// Create a memory accessor for `region`. + pub(crate) fn new(scratch: &HostSharedMemory, region: Range) -> Result { + let scratch_size = scratch.mem_size(); + let scratch_base_gva = scratch_base_gva(scratch_size); + + let scratch_end = u64::try_from(scratch_size) + .ok() + .and_then(|size| scratch_base_gva.checked_add(size)); + + if scratch_end.is_none_or(|end| region.end > end) + || region.start >= region.end + || region.start < scratch_base_gva + { + return Err(new_error!( + "region [{:#x}, {:#x}) is outside scratch at {:#x} with size {}", + region.start, + region.end, + scratch_base_gva, + scratch_size + )); + } + + Ok(Self { + scratch_mem: scratch.clone(), + scratch_base_gva, + region, + }) + } + + fn to_offset(&self, addr: u64, len: usize) -> Result { + let out_of_bounds = || { + new_error!( + "address {:#x} with length {} is outside region [{:#x}, {:#x})", + addr, + len, + self.region.start, + self.region.end + ) + }; + + let access_end = u64::try_from(len) + .ok() + .and_then(|len| addr.checked_add(len)); + + if addr < self.region.start || access_end.is_none_or(|end| end > self.region.end) { + return Err(out_of_bounds()); + } + + addr.checked_sub(self.scratch_base_gva) + .and_then(|offset| usize::try_from(offset).ok()) + .ok_or_else(out_of_bounds) + } +} + +// TODO: Hold one HostSharedMemory read guard across a virtq transaction. +// Descriptor metadata requires several reads and writes, so locking every +// operation scales with chain length and dominates the cached metadata path. + +// SAFETY: HostMemOps rejects accesses outside its assigned region. The backing +// HostSharedMemory keeps the mapping alive, bounds-checks each operation, and +// coordinates every byte and atomic access with exclusive memory operations. +unsafe impl MemOps for HostMemOps { + type Error = HyperlightError; + + fn read(&self, addr: u64, dst: &mut [u8]) -> Result<()> { + let offset = self.to_offset(addr, dst.len())?; + 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")) + } +} + +#[cfg(test)] +mod tests { + use hyperlight_common::virtq::MemOps; + + use super::*; + use crate::mem::shared_mem::ExclusiveSharedMemory; + + const SCRATCH_SIZE: usize = 0x4000; + + fn scratch_base() -> u64 { + scratch_base_gva(SCRATCH_SIZE) + } + + fn region() -> Range { + let scratch_base = scratch_base(); + scratch_base + 0x1000..scratch_base + 0x2000 + } + + fn host_mem_ops() -> HostMemOps { + let scratch = ExclusiveSharedMemory::new(SCRATCH_SIZE).unwrap(); + let (scratch, _) = scratch.build(); + HostMemOps::new(&scratch, region()).unwrap() + } + + #[test] + fn accesses_only_assigned_region() { + let mem = host_mem_ops(); + let region = region(); + + mem.write(region.start, &[1, 2, 3, 4]).unwrap(); + let mut bytes = [0; 4]; + mem.read(region.start, &mut bytes).unwrap(); + assert_eq!(bytes, [1, 2, 3, 4]); + + assert!(mem.read(region.start - 1, &mut [0]).is_err()); + assert!(mem.write(region.end - 1, &[1, 2]).is_err()); + assert!(mem.read(region.end, &mut [0]).is_err()); + assert!(mem.read(u64::MAX, &mut [0]).is_err()); + } + + #[test] + fn atomics_use_shared_memory_checks() { + let mem = host_mem_ops(); + let region = region(); + + mem.store_release(region.start, 0x1234).unwrap(); + assert_eq!(mem.load_acquire(region.start).unwrap(), 0x1234); + assert!(mem.load_acquire(region.start + 1).is_err()); + assert!(mem.load_acquire(region.end - 1).is_err()); + } + + #[test] + fn rejects_regions_outside_scratch() { + let scratch = ExclusiveSharedMemory::new(SCRATCH_SIZE).unwrap(); + let (scratch, _) = scratch.build(); + let scratch_base = scratch_base(); + let scratch_end = scratch_base + SCRATCH_SIZE as u64; + + assert!(HostMemOps::new(&scratch, scratch_base - 1..scratch_base).is_err()); + assert!(HostMemOps::new(&scratch, scratch_end - 1..scratch_end + 1).is_err()); + } +} From c0e53a5fee6c10721e0a24dd911579adc0d00cb1 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Wed, 29 Jul 2026 10:51:59 +0200 Subject: [PATCH 11/34] feat(virtq): configure transport geometry Define directional queue depths, buffer sizes, and pool page counts. Account for guest allocated rings and pools in minimum scratch calculations. Publish the transport contract through scratch-top metadata. Signed-off-by: Tomasz Andrzejak --- CHANGELOG.md | 2 + .../src/arch/aarch64/layout.rs | 8 +- .../src/arch/amd64/layout.rs | 10 +- src/hyperlight_common/src/layout.rs | 76 +++++- src/hyperlight_common/src/virtq/mod.rs | 7 + src/hyperlight_common/src/virtq/pool/slot.rs | 32 ++- src/hyperlight_common/src/virtq/pool/tests.rs | 18 ++ src/hyperlight_guest/src/layout.rs | 6 + src/hyperlight_host/src/mem/layout.rs | 159 ++++++++++-- src/hyperlight_host/src/sandbox/config.rs | 239 +++++++++++++++++- .../src/sandbox/initialized_multi_use.rs | 4 + 11 files changed, 519 insertions(+), 42 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0546852eb..e3c64b3c0 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. diff --git a/src/hyperlight_common/src/arch/aarch64/layout.rs b/src/hyperlight_common/src/arch/aarch64/layout.rs index cb32cfe8e..3f8dc8a04 100644 --- a/src/hyperlight_common/src/arch/aarch64/layout.rs +++ b/src/hyperlight_common/src/arch/aarch64/layout.rs @@ -28,7 +28,9 @@ pub const fn io_page() -> Option<(crate::vmem::PhysAddr, crate::vmem::VirtAddr)> Some((IO_PAGE_GPA, IO_PAGE_GVA)) } -pub fn min_scratch_size(input_data_size: usize, output_data_size: usize) -> usize { - (input_data_size + output_data_size).next_multiple_of(crate::vmem::PAGE_SIZE) - + 12 * crate::vmem::PAGE_SIZE +pub(super) fn min_scratch_size(input_data_size: usize, output_data_size: usize) -> Option { + input_data_size + .checked_add(output_data_size)? + .checked_next_multiple_of(crate::vmem::PAGE_SIZE)? + .checked_add(12 * crate::vmem::PAGE_SIZE) } diff --git a/src/hyperlight_common/src/arch/amd64/layout.rs b/src/hyperlight_common/src/arch/amd64/layout.rs index 4237caf51..26442a0de 100644 --- a/src/hyperlight_common/src/arch/amd64/layout.rs +++ b/src/hyperlight_common/src/arch/amd64/layout.rs @@ -41,8 +41,10 @@ pub fn io_page() -> Option<(u64, u64)> { /// - A page for the smallest possible non-exception stack /// - (up to) 3 pages for mapping that /// - Two pages for the exception stack and metadata -/// - A page-aligned amount of memory for I/O buffers (for now) -pub fn min_scratch_size(input_data_size: usize, output_data_size: usize) -> usize { - (input_data_size + output_data_size).next_multiple_of(crate::vmem::PAGE_SIZE) - + 12 * crate::vmem::PAGE_SIZE +/// - A page-aligned amount of memory for I/O buffers +pub(super) fn min_scratch_size(input_data_size: usize, output_data_size: usize) -> Option { + input_data_size + .checked_add(output_data_size)? + .checked_next_multiple_of(crate::vmem::PAGE_SIZE)? + .checked_add(12 * crate::vmem::PAGE_SIZE) } diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index ec2852520..a0b4ad98a 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -24,6 +24,8 @@ pub use arch::{ SCRATCH_TOP_GPA, SCRATCH_TOP_GVA, SNAPSHOT_PT_GVA_MAX, SNAPSHOT_PT_GVA_MIN, io_page, }; +use crate::virtq; + const EXN_STACK_ALIGNMENT: usize = 16; /// Pages reserved for the exception stack and scratch-top metadata. pub const SCRATCH_TOP_RESERVED_PAGES: usize = 2; @@ -32,19 +34,23 @@ pub const SCRATCH_TOP_RESERVED_PAGES: usize = 2; // down from the top of scratch memory. #[repr(C)] struct ScratchTopMetadata { + /// Host-published capacity of each H2G buffer. + h2g_buffer_size: u64, /// Number of pages reserved for the H2G pool. h2g_pool_pages: u64, /// Guest-published GPA of the H2G pool. h2g_pool_gpa: u64, - /// Host-published GPA of the H2G ring. + /// Guest-published GPA of the H2G ring. h2g_ring_gpa: u64, /// Host-published H2G descriptor count. h2g_queue_depth: u64, + /// Host-published capacity of each G2H upper-tier buffer. + g2h_buffer_size: u64, /// Number of pages reserved for the G2H pool. g2h_pool_pages: u64, /// Guest-published GPA of the G2H pool. g2h_pool_gpa: u64, - /// Host-published GPA of the G2H ring. + /// Guest-published GPA of the G2H ring. g2h_ring_gpa: u64, /// Host-published G2H descriptor count. g2h_queue_depth: u64, @@ -70,6 +76,8 @@ pub const SCRATCH_TOP_G2H_POOL_GPA_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_gpa)); pub const SCRATCH_TOP_G2H_POOL_PAGES_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_pages)); +pub const SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_buffer_size)); pub const SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_queue_depth)); pub const SCRATCH_TOP_H2G_RING_GPA_OFFSET: u64 = @@ -78,6 +86,8 @@ pub const SCRATCH_TOP_H2G_POOL_GPA_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_gpa)); pub const SCRATCH_TOP_H2G_POOL_PAGES_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_pages)); +pub const SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_buffer_size)); pub const SCRATCH_TOP_SIZE_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, scratch_size)); pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = @@ -98,11 +108,13 @@ const _: () = { assert!(SCRATCH_TOP_G2H_RING_GPA_OFFSET == 0x30); assert!(SCRATCH_TOP_G2H_POOL_GPA_OFFSET == 0x38); assert!(SCRATCH_TOP_G2H_POOL_PAGES_OFFSET == 0x40); - assert!(SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET == 0x48); - assert!(SCRATCH_TOP_H2G_RING_GPA_OFFSET == 0x50); - assert!(SCRATCH_TOP_H2G_POOL_GPA_OFFSET == 0x58); - assert!(SCRATCH_TOP_H2G_POOL_PAGES_OFFSET == 0x60); - assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x60); + assert!(SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET == 0x48); + assert!(SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET == 0x50); + assert!(SCRATCH_TOP_H2G_RING_GPA_OFFSET == 0x58); + assert!(SCRATCH_TOP_H2G_POOL_GPA_OFFSET == 0x60); + assert!(SCRATCH_TOP_H2G_POOL_PAGES_OFFSET == 0x68); + assert!(SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET == 0x70); + assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x70); }; /// Exclusive upper GPA boundary for dynamic scratch allocations. @@ -118,4 +130,52 @@ pub fn scratch_base_gva(size: usize) -> u64 { } /// Compute the minimum scratch region size needed for a sandbox. -pub use arch::min_scratch_size; +/// +/// The transport allowance contains one page-backed ring arena and both +/// page-backed buffer pools. The result saturates at [`usize::MAX`]. +pub fn min_scratch_size( + input_data_size: usize, + output_data_size: usize, + g2h_queue_depth: usize, + h2g_queue_depth: usize, + g2h_pool_pages: usize, + h2g_pool_pages: usize, +) -> usize { + let size = arch::min_scratch_size(input_data_size, output_data_size).and_then(|fixed| { + let h2g_ring_offset = virtq::Layout::query_size(g2h_queue_depth) + .checked_next_multiple_of(virtq::Descriptor::ALIGN)?; + + let ring_pages = h2g_ring_offset + .checked_add(virtq::Layout::query_size(h2g_queue_depth))? + .checked_next_multiple_of(crate::vmem::PAGE_SIZE)?; + + let pool_size = g2h_pool_pages + .checked_add(h2g_pool_pages)? + .checked_mul(crate::vmem::PAGE_SIZE)?; + + fixed.checked_add(ring_pages)?.checked_add(pool_size) + }); + + size.unwrap_or(usize::MAX) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn minimum_scratch_includes_ring_arena_and_pools() { + let fixed = arch::min_scratch_size(0, 0).unwrap(); + let transport_pages = 1 + 8 + 4; + + assert_eq!( + fixed + transport_pages * crate::vmem::PAGE_SIZE, + min_scratch_size(0, 0, 64, 32, 8, 4) + ); + } + + #[test] + fn minimum_scratch_saturates_on_overflow() { + assert_eq!(usize::MAX, min_scratch_size(0, 0, 64, 32, usize::MAX, 4)); + } +} diff --git a/src/hyperlight_common/src/virtq/mod.rs b/src/hyperlight_common/src/virtq/mod.rs index 895bc6dc6..63e73f21f 100644 --- a/src/hyperlight_common/src/virtq/mod.rs +++ b/src/hyperlight_common/src/virtq/mod.rs @@ -185,6 +185,13 @@ pub use producer::*; pub use ring::*; use thiserror::Error; +/// Capacity of each fixed G2H lower-tier slot. +pub const G2H_LOWER_SLOT_SIZE: usize = 256; +/// Number of G2H lower-tier slots occupying the first pool page. +pub const G2H_LOWER_SLOT_COUNT: usize = crate::vmem::PAGE_SIZE / G2H_LOWER_SLOT_SIZE; + +const _: () = assert!(G2H_LOWER_SLOT_COUNT * G2H_LOWER_SLOT_SIZE == crate::vmem::PAGE_SIZE); + /// A trait for notifying the consumer about virtqueue events. pub trait Notifier { fn notify(&self, stats: QueueStats); diff --git a/src/hyperlight_common/src/virtq/pool/slot.rs b/src/hyperlight_common/src/virtq/pool/slot.rs index 7689e8e01..8bfc9c065 100644 --- a/src/hyperlight_common/src/virtq/pool/slot.rs +++ b/src/hyperlight_common/src/virtq/pool/slot.rs @@ -213,13 +213,34 @@ impl Inner { let lower = lower.map(Tier::from_layout).transpose()?; let upper = Tier::from_layout(upper)?; - if let Some(lower) = &lower - && (lower.slot_size >= upper.slot_size || lower.end() > upper.base_addr) - { + let Some(lower) = lower else { + return Ok(Self { lower: None, upper }); + }; + + if lower.slot_size > upper.slot_size || lower.end() > upper.base_addr { return Err(AllocError::InvalidArg); } - Ok(Self { lower, upper }) + if lower.slot_size == upper.slot_size { + if lower.end() != upper.base_addr { + return Err(AllocError::InvalidArg); + } + + let count = lower + .count + .checked_add(upper.count) + .ok_or(AllocError::Overflow)?; + let layout = SlotLayout::new(lower.base_addr, lower.slot_size, count); + return Ok(Self { + lower: None, + upper: Tier::from_layout(layout)?, + }); + } + + Ok(Self { + lower: Some(lower), + upper, + }) } fn max_alloc_len(&self) -> usize { @@ -316,7 +337,8 @@ impl SlotPool { /// Create a two-tier recycling pool from exact lower and upper layouts. /// /// The lower layout must precede the upper layout without overlap, and its - /// slot size must be strictly smaller. + /// slot size must not exceed the upper slot size. Adjacent equal-sized + /// layouts form one tier. pub fn new_tiered(lower: SlotLayout, upper: SlotLayout) -> Result { Self::from_layouts(Some(lower), upper) } diff --git a/src/hyperlight_common/src/virtq/pool/tests.rs b/src/hyperlight_common/src/virtq/pool/tests.rs index 87c917e7c..ed4216c2d 100644 --- a/src/hyperlight_common/src/virtq/pool/tests.rs +++ b/src/hyperlight_common/src/virtq/pool/tests.rs @@ -205,6 +205,19 @@ fn test_tiered_slot_pool_reports_layouts() { assert_eq!(pool.slot_addr(4), None); } +#[test] +fn test_tiered_slot_pool_combines_contiguous_equal_sized_layouts() { + let lower = SlotLayout::new(0x80000, 0x100, 2); + let upper = SlotLayout::new(0x80200, 0x100, 3); + let pool = SlotPool::new_tiered(lower, upper).unwrap(); + + assert_eq!(pool.layouts(), (None, SlotLayout::new(0x80000, 0x100, 5))); + assert_eq!(pool.base_addr(), 0x80000); + assert_eq!(pool.slot_size(), 0x100); + assert_eq!(pool.count(), 5); + assert_eq!(pool.slot_addr(4), Some(0x80400)); +} + #[test] fn test_tiered_slot_pool_rejects_invalid_layout() { let lower = SlotLayout::new(0x80000, 0x100, 32); @@ -212,6 +225,11 @@ fn test_tiered_slot_pool_rejects_invalid_layout() { let overlapping = SlotPool::new_tiered(lower, overlapping_upper); assert!(matches!(overlapping, Err(AllocError::InvalidArg))); + let lower = SlotLayout::new(0x80000, 0x100, 2); + let separated_upper = SlotLayout::new(0x80300, 0x100, 2); + let separated = SlotPool::new_tiered(lower, separated_upper); + assert!(matches!(separated, Err(AllocError::InvalidArg))); + let lower = SlotLayout::new(0x80000, 0x1000, 2); let smaller_upper = SlotLayout::new(0x90000, 0x100, 32); let reversed_sizes = SlotPool::new_tiered(lower, smaller_upper); diff --git a/src/hyperlight_guest/src/layout.rs b/src/hyperlight_guest/src/layout.rs index c80edf5f5..860ab7c16 100644 --- a/src/hyperlight_guest/src/layout.rs +++ b/src/hyperlight_guest/src/layout.rs @@ -54,10 +54,16 @@ pub fn h2g_pool_gpa_gva() -> *mut u64 { pub fn g2h_pool_pages_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_POOL_PAGES_OFFSET) } +pub fn g2h_buffer_size_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET) +} pub fn h2g_queue_depth_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET) } pub fn h2g_pool_pages_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_POOL_PAGES_OFFSET) } +pub fn h2g_buffer_size_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET) +} pub use arch::{scratch_base_gpa, scratch_base_gva}; diff --git a/src/hyperlight_host/src/mem/layout.rs b/src/hyperlight_host/src/mem/layout.rs index 6422b9b11..28817f2a2 100644 --- a/src/hyperlight_host/src/mem/layout.rs +++ b/src/hyperlight_host/src/mem/layout.rs @@ -262,6 +262,18 @@ pub(crate) struct SandboxMemoryLayout { init_data_permissions: Option, /// The size of the scratch region in physical memory. scratch_size: usize, + /// Number of descriptors in the G2H virtqueue. + g2h_queue_depth: usize, + /// Number of descriptors in the H2G virtqueue. + h2g_queue_depth: usize, + /// Capacity of each G2H upper-tier buffer. + g2h_buffer_size: usize, + /// Capacity of each H2G buffer. + h2g_buffer_size: usize, + /// Number of pages in the G2H buffer pool. + g2h_pool_pages: usize, + /// Number of pages in the H2G buffer pool. + h2g_pool_pages: usize, /// Size of the primary guest memory region at `BASE_ADDRESS` /// (code, PEB, heap, init data). For a snapshot-backed layout /// this is also the guest-visible prefix of the host snapshot @@ -296,6 +308,12 @@ impl Debug for SandboxMemoryLayout { &format_args!("{:#x}", self.output_data_size), ) .field("Scratch Size", &format_args!("{:#x}", self.scratch_size)) + .field("G2H Queue Depth", &self.g2h_queue_depth) + .field("H2G Queue Depth", &self.h2g_queue_depth) + .field("G2H Buffer Size", &self.g2h_buffer_size) + .field("H2G Buffer Size", &self.h2g_buffer_size) + .field("G2H Pool Pages", &self.g2h_pool_pages) + .field("H2G Pool Pages", &self.h2g_pool_pages) .field("Snapshot Size", &format_args!("{:#x}", self.snapshot_size)) .field("PT Size", &format_args!("{:#x}", self.pt_size.unwrap_or(0))) .field( @@ -317,14 +335,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 does not participate while snapshots use the + /// stack transport. `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 { @@ -339,6 +354,12 @@ impl SandboxMemoryLayout { init_data_size, init_data_permissions, scratch_size, + g2h_queue_depth: _, + h2g_queue_depth: _, + g2h_buffer_size: _, + h2g_buffer_size: _, + g2h_pool_pages: _, + h2g_pool_pages: _, snapshot_size: _, pt_size: _, } = self; @@ -380,8 +401,20 @@ impl SandboxMemoryLayout { } let input_data_size = cfg.get_input_data_size(); let output_data_size = cfg.get_output_data_size(); - let min_scratch_size = - hyperlight_common::layout::min_scratch_size(input_data_size, output_data_size); + let g2h_queue_depth = cfg.get_g2h_queue_depth(); + let h2g_queue_depth = cfg.get_h2g_queue_depth(); + let g2h_buffer_size = cfg.get_g2h_buffer_size(); + let h2g_buffer_size = cfg.get_h2g_buffer_size(); + let g2h_pool_pages = cfg.get_g2h_pool_pages(); + let h2g_pool_pages = cfg.get_h2g_pool_pages(); + let min_scratch_size = hyperlight_common::layout::min_scratch_size( + input_data_size, + output_data_size, + g2h_queue_depth, + h2g_queue_depth, + g2h_pool_pages, + h2g_pool_pages, + ); if scratch_size < min_scratch_size { return Err(MemoryRequestTooSmall(scratch_size, min_scratch_size)); } @@ -395,6 +428,12 @@ impl SandboxMemoryLayout { init_data_permissions, pt_size: None, scratch_size, + g2h_queue_depth, + h2g_queue_depth, + g2h_buffer_size, + h2g_buffer_size, + g2h_pool_pages, + h2g_pool_pages, snapshot_size: 0, }; ret.set_snapshot_size(ret.get_memory_size()?); @@ -429,6 +468,36 @@ impl SandboxMemoryLayout { self.scratch_size } + #[allow(dead_code)] + pub(crate) fn get_g2h_queue_depth(&self) -> usize { + self.g2h_queue_depth + } + + #[allow(dead_code)] + pub(crate) fn get_h2g_queue_depth(&self) -> usize { + self.h2g_queue_depth + } + + #[allow(dead_code)] + pub(crate) fn get_g2h_buffer_size(&self) -> usize { + self.g2h_buffer_size + } + + #[allow(dead_code)] + pub(crate) fn get_h2g_buffer_size(&self) -> usize { + self.h2g_buffer_size + } + + #[allow(dead_code)] + pub(crate) fn get_g2h_pool_pages(&self) -> usize { + self.g2h_pool_pages + } + + #[allow(dead_code)] + pub(crate) fn get_h2g_pool_pages(&self) -> usize { + self.h2g_pool_pages + } + /// Guest-visible prefix size of the snapshot blob. pub(crate) fn snapshot_size(&self) -> usize { self.snapshot_size @@ -454,8 +523,12 @@ impl SandboxMemoryLayout { let min_fixed_scratch = hyperlight_common::layout::min_scratch_size( self.input_data_size, self.output_data_size, + self.g2h_queue_depth, + self.h2g_queue_depth, + self.g2h_pool_pages, + self.h2g_pool_pages, ); - let min_scratch = min_fixed_scratch + size; + let min_scratch = min_fixed_scratch.saturating_add(size); if self.scratch_size < min_scratch { return Err(MemoryRequestTooSmall(self.scratch_size, min_scratch)); } @@ -715,8 +788,7 @@ impl SandboxMemoryLayout { + self.get_pt_base_scratch_offset() as u64 } - /// First GPA of the scratch region the host has not used for - /// something else. + /// First GPA available to the guest scratch allocator. pub(crate) fn get_first_free_scratch_gpa(&self) -> u64 { self.get_pt_base_gpa() + self.pt_size.unwrap_or(0) as u64 } @@ -778,6 +850,34 @@ mod tests { ); } + #[test] + fn transport_memory_is_part_of_minimum_scratch_size() { + let mut cfg = SandboxConfiguration::default(); + let minimum = hyperlight_common::layout::min_scratch_size( + cfg.get_input_data_size(), + cfg.get_output_data_size(), + cfg.get_g2h_queue_depth(), + cfg.get_h2g_queue_depth(), + cfg.get_g2h_pool_pages(), + cfg.get_h2g_pool_pages(), + ); + cfg.set_scratch_size(minimum); + let mut layout = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap(); + + assert!(matches!( + layout.set_pt_size(PAGE_SIZE_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 test_max_memory_sandbox() { let mut cfg = SandboxConfiguration::default(); @@ -840,6 +940,23 @@ mod tests { } } + #[test] + fn is_compatible_with_ignores_inactive_transport_configuration() { + let base = + SandboxMemoryLayout::new(SandboxConfiguration::default(), 4096, 0, None).unwrap(); + let mut cfg = SandboxConfiguration::default(); + cfg.set_g2h_queue_depth(128); + cfg.set_h2g_queue_depth(16); + cfg.set_g2h_buffer_size(16 * 1024); + cfg.set_h2g_buffer_size(8 * 1024); + cfg.set_g2h_pool_pages(16); + cfg.set_h2g_pool_pages(8); + let other = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap(); + + assert!(base.is_compatible_with(&other)); + assert!(other.is_compatible_with(&base)); + } + /// Pinned region offsets. These methods place every region that a /// restored snapshot is interpreted against, so a change shifts /// where the loader reads captured bytes and breaks existing @@ -901,7 +1018,7 @@ mod tests { cfg.set_input_data_size(0x2000); cfg.set_output_data_size(0x2000); cfg.set_heap_size(0x2000); - cfg.set_scratch_size(0x10000); + cfg.set_scratch_size(0x20000); let layout = SandboxMemoryLayout::new(cfg, 0x1000, 0, None).unwrap(); pin_eq!(layout.guest_code_offset(), 0); @@ -911,7 +1028,7 @@ mod tests { pin_eq!(layout.init_data_offset(), 0x4000); pin_eq!(layout.get_memory_size().unwrap(), 0x4000); - pin_eq!(layout.get_scratch_size(), 0x10000); + pin_eq!(layout.get_scratch_size(), 0x20000); pin_eq!(layout.get_pt_size(), 0); pin_eq!(layout.get_input_data_buffer_scratch_host_offset(), 0); @@ -930,11 +1047,11 @@ mod tests { // `SCRATCH_TOP` pins above, these fix the absolute addresses. pin_eq!( layout.get_input_data_buffer_gva() - - hyperlight_common::layout::scratch_base_gva(0x10000), + - hyperlight_common::layout::scratch_base_gva(0x20000), 0 ); pin_eq!( - layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x10000), + layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x20000), 0x4000 ); // pt_size is zero here, so the first free scratch GPA equals @@ -950,7 +1067,7 @@ mod tests { cfg.set_input_data_size(0x4000); cfg.set_output_data_size(0x2000); cfg.set_heap_size(0x5000); - cfg.set_scratch_size(0x20000); + cfg.set_scratch_size(0x30000); let layout = SandboxMemoryLayout::new(cfg, 0x3000, 0, None).unwrap(); pin_eq!(layout.guest_code_offset(), 0); @@ -960,7 +1077,7 @@ 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(), 0x30000); pin_eq!(layout.get_pt_size(), 0); pin_eq!(layout.get_input_data_buffer_scratch_host_offset(), 0); @@ -974,11 +1091,11 @@ mod tests { pin_eq!( layout.get_input_data_buffer_gva() - - hyperlight_common::layout::scratch_base_gva(0x20000), + - hyperlight_common::layout::scratch_base_gva(0x30000), 0 ); pin_eq!( - layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x20000), + layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x30000), 0x6000 ); pin_eq!( diff --git a/src/hyperlight_host/src/sandbox/config.rs b/src/hyperlight_host/src/sandbox/config.rs index 442da8415..dea59cb6e 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}; @@ -86,6 +88,18 @@ pub struct SandboxConfiguration { interrupt_vcpu_sigrtmin_offset: u8, /// How much writable memory to offer the guest scratch_size: usize, + /// Number of descriptors in the G2H virtqueue. + g2h_queue_depth: usize, + /// Number of descriptors in the H2G virtqueue. + h2g_queue_depth: usize, + /// Capacity of each G2H upper-tier buffer. + g2h_buffer_size: usize, + /// Capacity of each H2G buffer. + h2g_buffer_size: usize, + /// Number of pages in the G2H buffer pool. + g2h_pool_pages: usize, + /// Number of pages in the H2G buffer pool. + h2g_pool_pages: usize, /// Declared guest MSRs, stored inline to keep this type `Copy`. #[cfg(target_arch = "x86_64")] guest_msrs: [u32; Self::MAX_GUEST_MSRS], @@ -110,7 +124,27 @@ impl SandboxConfiguration { /// The default heap size of a hyperlight sandbox pub const DEFAULT_HEAP_SIZE: u64 = 131072; /// The default size of the scratch region - pub const DEFAULT_SCRATCH_SIZE: usize = 0x48000; + pub const DEFAULT_SCRATCH_SIZE: usize = 0x55000; + /// The default G2H virtqueue descriptor count. + pub const DEFAULT_G2H_QUEUE_DEPTH: usize = 64; + /// The default H2G virtqueue descriptor count. + pub const DEFAULT_H2G_QUEUE_DEPTH: usize = 32; + /// The default G2H upper-tier buffer size. + pub const DEFAULT_G2H_BUFFER_SIZE: usize = PAGE_SIZE; + /// The default H2G buffer size. + pub const DEFAULT_H2G_BUFFER_SIZE: usize = PAGE_SIZE; + /// The default total number of G2H pool pages. + pub const DEFAULT_G2H_POOL_PAGES: usize = 8; + /// The default total number of H2G pool pages. + pub const DEFAULT_H2G_POOL_PAGES: usize = 4; + /// The minimum G2H virtqueue descriptor count. + const MIN_QUEUE_DEPTH: usize = 2; + /// The maximum G2H virtqueue descriptor count. + const MAX_QUEUE_DEPTH: usize = 32_768; + /// The minimum configured transport buffer size. + const MIN_BUFFER_SIZE: usize = G2H_LOWER_SLOT_SIZE; + /// The maximum configured transport buffer size. + const MAX_BUFFER_SIZE: usize = u32::MAX as usize; /// Maximum number of distinct guest MSRs that can be declared. /// KVM supports at most 16 MSR filter ranges. Each index may require its /// own range, so 16 is the portable limit across backends. @@ -135,6 +169,12 @@ impl SandboxConfiguration { output_data_size: max(output_data_size, Self::MIN_OUTPUT_SIZE), heap_size_override: heap_size_override.unwrap_or(0), scratch_size, + g2h_queue_depth: Self::DEFAULT_G2H_QUEUE_DEPTH, + h2g_queue_depth: Self::DEFAULT_H2G_QUEUE_DEPTH, + g2h_buffer_size: Self::DEFAULT_G2H_BUFFER_SIZE, + h2g_buffer_size: Self::DEFAULT_H2G_BUFFER_SIZE, + g2h_pool_pages: Self::DEFAULT_G2H_POOL_PAGES, + h2g_pool_pages: Self::DEFAULT_H2G_POOL_PAGES, interrupt_retry_delay, interrupt_vcpu_sigrtmin_offset, #[cfg(gdb)] @@ -299,6 +339,98 @@ impl SandboxConfiguration { self.scratch_size = scratch_size; } + /// Get the G2H virtqueue descriptor count. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_g2h_queue_depth(&self) -> usize { + self.g2h_queue_depth + } + + /// Set the G2H virtqueue descriptor count. + /// + /// Values are rounded up to a power of two in `2..=32768`. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_g2h_queue_depth(&mut self, depth: usize) { + self.g2h_queue_depth = Self::normalize_queue_depth(depth); + } + + /// Get the H2G virtqueue descriptor count. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_h2g_queue_depth(&self) -> usize { + self.h2g_queue_depth + } + + /// Set the H2G virtqueue descriptor count. + /// + /// Values are rounded up to a power of two in `2..=32768`. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_h2g_queue_depth(&mut self, depth: usize) { + self.h2g_queue_depth = Self::normalize_queue_depth(depth); + } + + /// Get the capacity of each G2H upper-tier buffer. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_g2h_buffer_size(&self) -> usize { + self.g2h_buffer_size + } + + /// Set the capacity of each G2H upper-tier buffer. + /// + /// Values are clamped to `256..=u32::MAX`. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_g2h_buffer_size(&mut self, size: usize) { + self.g2h_buffer_size = size.clamp(Self::MIN_BUFFER_SIZE, Self::MAX_BUFFER_SIZE); + self.g2h_pool_pages = max( + self.g2h_pool_pages, + Self::min_g2h_pool_pages(self.g2h_buffer_size), + ); + } + + /// Get the capacity of each H2G buffer. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_h2g_buffer_size(&self) -> usize { + self.h2g_buffer_size + } + + /// Set the capacity of each H2G buffer. + /// + /// Values are clamped to `256..=u32::MAX`. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_h2g_buffer_size(&mut self, size: usize) { + self.h2g_buffer_size = size.clamp(Self::MIN_BUFFER_SIZE, Self::MAX_BUFFER_SIZE); + self.h2g_pool_pages = max( + self.h2g_pool_pages, + Self::min_h2g_pool_pages(self.h2g_buffer_size), + ); + } + + /// Get the total number of G2H pool pages. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_g2h_pool_pages(&self) -> usize { + self.g2h_pool_pages + } + + /// Set the total number of G2H pool pages. + /// + /// The pool contains one lower-tier page and at least one upper buffer. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_g2h_pool_pages(&mut self, pages: usize) { + self.g2h_pool_pages = max(pages, Self::min_g2h_pool_pages(self.g2h_buffer_size)); + } + + /// Get the total number of H2G pool pages. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn get_h2g_pool_pages(&self) -> usize { + self.h2g_pool_pages + } + + /// Set the total number of H2G pool pages. + /// + /// The pool contains at least one H2G buffer. + #[instrument(skip_all, parent = Span::current(), level= "Trace")] + pub fn set_h2g_pool_pages(&mut self, pages: usize) { + self.h2g_pool_pages = max(pages, Self::min_h2g_pool_pages(self.h2g_buffer_size)); + } + #[cfg(crashdump)] #[instrument(skip_all, parent = Span::current(), level= "Trace")] pub(crate) fn get_guest_core_dump(&self) -> bool { @@ -323,6 +455,20 @@ impl SandboxConfiguration { self.heap_size_override_opt() .unwrap_or(Self::DEFAULT_HEAP_SIZE) } + + fn normalize_queue_depth(depth: usize) -> usize { + depth + .clamp(Self::MIN_QUEUE_DEPTH, Self::MAX_QUEUE_DEPTH) + .next_power_of_two() + } + + fn min_g2h_pool_pages(buffer_size: usize) -> usize { + 1 + Self::min_h2g_pool_pages(buffer_size) + } + + fn min_h2g_pool_pages(buffer_size: usize) -> usize { + buffer_size.div_ceil(PAGE_SIZE) + } } impl Default for SandboxConfiguration { @@ -347,6 +493,8 @@ impl Default for SandboxConfiguration { mod tests { #[cfg(target_arch = "x86_64")] use super::GuestMsrError; + use hyperlight_common::vmem::PAGE_SIZE; + use super::SandboxConfiguration; #[test] @@ -435,6 +583,30 @@ mod tests { assert_eq!(0x40000, cfg.scratch_size); assert_eq!(INPUT_DATA_SIZE_OVERRIDE, cfg.input_data_size); assert_eq!(OUTPUT_DATA_SIZE_OVERRIDE, cfg.output_data_size); + assert_eq!( + SandboxConfiguration::DEFAULT_G2H_QUEUE_DEPTH, + cfg.get_g2h_queue_depth() + ); + assert_eq!( + SandboxConfiguration::DEFAULT_H2G_QUEUE_DEPTH, + cfg.get_h2g_queue_depth() + ); + assert_eq!( + SandboxConfiguration::DEFAULT_G2H_BUFFER_SIZE, + cfg.get_g2h_buffer_size() + ); + assert_eq!( + SandboxConfiguration::DEFAULT_H2G_BUFFER_SIZE, + cfg.get_h2g_buffer_size() + ); + assert_eq!( + SandboxConfiguration::DEFAULT_G2H_POOL_PAGES, + cfg.get_g2h_pool_pages() + ); + assert_eq!( + SandboxConfiguration::DEFAULT_H2G_POOL_PAGES, + cfg.get_h2g_pool_pages() + ); } #[test] @@ -462,6 +634,71 @@ mod tests { assert_eq!(SandboxConfiguration::MIN_OUTPUT_SIZE, cfg.output_data_size); } + #[test] + fn queue_depths_are_normalized() { + let mut cfg = SandboxConfiguration::default(); + for (depth, expected) in [ + (0, 2), + (1, 2), + (2, 2), + (3, 4), + (32_767, 32_768), + (32_768, 32_768), + (32_769, 32_768), + (usize::MAX, 32_768), + ] { + cfg.set_g2h_queue_depth(depth); + cfg.set_h2g_queue_depth(depth); + assert_eq!(expected, cfg.get_g2h_queue_depth()); + assert_eq!(expected, cfg.get_h2g_queue_depth()); + } + } + + #[test] + fn buffer_sizes_are_normalized_without_page_rounding() { + let mut cfg = SandboxConfiguration::default(); + + cfg.set_g2h_buffer_size(0); + cfg.set_h2g_buffer_size(0); + assert_eq!(256, cfg.get_g2h_buffer_size()); + assert_eq!(256, cfg.get_h2g_buffer_size()); + + cfg.set_g2h_buffer_size(3000); + cfg.set_h2g_buffer_size(3001); + assert_eq!(3000, cfg.get_g2h_buffer_size()); + assert_eq!(3001, cfg.get_h2g_buffer_size()); + + cfg.set_g2h_buffer_size(usize::MAX); + cfg.set_h2g_buffer_size(usize::MAX); + assert_eq!(u32::MAX as usize, cfg.get_g2h_buffer_size()); + assert_eq!(u32::MAX as usize, cfg.get_h2g_buffer_size()); + } + + #[test] + fn pool_page_counts_are_normalized() { + let mut cfg = SandboxConfiguration::default(); + + cfg.set_g2h_pool_pages(0); + cfg.set_h2g_pool_pages(0); + assert_eq!(2, cfg.get_g2h_pool_pages()); + assert_eq!(1, cfg.get_h2g_pool_pages()); + + cfg.set_g2h_buffer_size(PAGE_SIZE + 1); + cfg.set_h2g_buffer_size(PAGE_SIZE + 1); + assert_eq!(3, cfg.get_g2h_pool_pages()); + assert_eq!(2, cfg.get_h2g_pool_pages()); + + cfg.set_g2h_pool_pages(2); + cfg.set_h2g_pool_pages(1); + assert_eq!(3, cfg.get_g2h_pool_pages()); + assert_eq!(2, cfg.get_h2g_pool_pages()); + + cfg.set_g2h_pool_pages(4); + cfg.set_h2g_pool_pages(3); + assert_eq!(4, cfg.get_g2h_pool_pages()); + assert_eq!(3, cfg.get_h2g_pool_pages()); + } + mod proptests { use proptest::prelude::*; diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index 623c24667..54207deed 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -1368,6 +1368,10 @@ mod tests { let min_scratch = hyperlight_common::layout::min_scratch_size( cfg.get_input_data_size(), cfg.get_output_data_size(), + cfg.get_g2h_queue_depth(), + cfg.get_h2g_queue_depth(), + cfg.get_g2h_pool_pages(), + cfg.get_h2g_pool_pages(), ); cfg.set_scratch_size(min_scratch + 0x10000 + 0x10000); From e9bce77084160bce6281b01dddda776cfe750774 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Thu, 30 Jul 2026 17:58:00 +0200 Subject: [PATCH 12/34] feat(virtq): initialize runtime transport This patch adds guest owned G2H and H2G rings and pools during guest initialization, and prefill H2G receive capacity. The guest then publishes their gpas through scratch metadata. The patch is also validates allocation order, scratch ownership, and canonical ring images before installing either host consumer. Signed-off-by: Tomasz Andrzejak --- src/hyperlight_common/src/layout.rs | 267 ++++++-- src/hyperlight_common/src/virtq/mod.rs | 17 + src/hyperlight_guest/src/error.rs | 10 + src/hyperlight_guest/src/layout.rs | 13 +- src/hyperlight_guest/src/lib.rs | 1 + src/hyperlight_guest/src/transport/context.rs | 152 +++++ src/hyperlight_guest/src/transport/mem.rs | 148 +++++ src/hyperlight_guest/src/transport/mod.rs | 75 +++ src/hyperlight_guest_bin/src/lib.rs | 4 + src/hyperlight_guest_bin/src/transport.rs | 103 ++++ src/hyperlight_host/src/mem/layout.rs | 82 ++- src/hyperlight_host/src/mem/mgr.rs | 85 ++- src/hyperlight_host/src/mem/mod.rs | 2 + src/hyperlight_host/src/mem/virtq.rs | 581 ++++++++++++++++++ .../src/sandbox/initialized_multi_use.rs | 8 + .../src/sandbox/uninitialized_evolve.rs | 8 + 16 files changed, 1500 insertions(+), 56 deletions(-) create mode 100644 src/hyperlight_guest/src/transport/context.rs create mode 100644 src/hyperlight_guest/src/transport/mem.rs create mode 100644 src/hyperlight_guest/src/transport/mod.rs create mode 100644 src/hyperlight_guest_bin/src/transport.rs create mode 100644 src/hyperlight_host/src/mem/virtq.rs diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index a0b4ad98a..e0a597103 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -15,6 +15,7 @@ limitations under the License. */ use core::mem::{offset_of, size_of}; +use core::num::{NonZeroU16, NonZeroUsize}; #[cfg_attr(target_arch = "x86_64", path = "arch/amd64/layout.rs")] #[cfg_attr(target_arch = "aarch64", path = "arch/aarch64/layout.rs")] @@ -34,26 +35,22 @@ pub const SCRATCH_TOP_RESERVED_PAGES: usize = 2; // 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, - /// Guest-published GPA of the H2G pool. - h2g_pool_gpa: u64, - /// Guest-published GPA of the H2G ring. - h2g_ring_gpa: u64, /// Host-published H2G descriptor count. h2g_queue_depth: u64, /// Host-published capacity of each G2H upper-tier buffer. g2h_buffer_size: u64, /// Number of pages reserved for the G2H pool. g2h_pool_pages: u64, - /// Guest-published GPA of the G2H pool. - g2h_pool_gpa: u64, - /// Guest-published GPA of the G2H ring. - g2h_ring_gpa: u64, /// Host-published G2H descriptor count. g2h_queue_depth: u64, + /// Host-published GPA of the fixed transport arena. + transport_arena_gpa: u64, /// Generation of the snapshot backing the sandbox. snapshot_generation: u64, /// GPA of the snapshot page-table copy in scratch memory. @@ -70,24 +67,18 @@ const fn scratch_top_offset(field_offset: usize) -> u64 { pub const SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_queue_depth)); -pub const SCRATCH_TOP_G2H_RING_GPA_OFFSET: u64 = - scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_ring_gpa)); -pub const SCRATCH_TOP_G2H_POOL_GPA_OFFSET: u64 = - scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_gpa)); pub const SCRATCH_TOP_G2H_POOL_PAGES_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_pool_pages)); pub const SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, g2h_buffer_size)); pub const SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_queue_depth)); -pub const SCRATCH_TOP_H2G_RING_GPA_OFFSET: u64 = - scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_ring_gpa)); -pub const SCRATCH_TOP_H2G_POOL_GPA_OFFSET: u64 = - scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_gpa)); pub const SCRATCH_TOP_H2G_POOL_PAGES_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_pool_pages)); pub const SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, h2g_buffer_size)); +pub const SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET: u64 = + scratch_top_offset(offset_of!(ScratchTopMetadata, transport_arena_gpa)); pub const SCRATCH_TOP_SIZE_OFFSET: u64 = scratch_top_offset(offset_of!(ScratchTopMetadata, scratch_size)); pub const SCRATCH_TOP_ALLOCATOR_OFFSET: u64 = @@ -104,17 +95,14 @@ const _: () = { 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_G2H_QUEUE_DEPTH_OFFSET == 0x28); - assert!(SCRATCH_TOP_G2H_RING_GPA_OFFSET == 0x30); - assert!(SCRATCH_TOP_G2H_POOL_GPA_OFFSET == 0x38); - assert!(SCRATCH_TOP_G2H_POOL_PAGES_OFFSET == 0x40); - assert!(SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET == 0x48); - assert!(SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET == 0x50); - assert!(SCRATCH_TOP_H2G_RING_GPA_OFFSET == 0x58); - assert!(SCRATCH_TOP_H2G_POOL_GPA_OFFSET == 0x60); - assert!(SCRATCH_TOP_H2G_POOL_PAGES_OFFSET == 0x68); - assert!(SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET == 0x70); - assert!(SCRATCH_TOP_EXN_STACK_OFFSET == 0x70); + assert!(SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET == 0x28); + assert!(SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET == 0x30); + assert!(SCRATCH_TOP_G2H_POOL_PAGES_OFFSET == 0x38); + assert!(SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET == 0x40); + assert!(SCRATCH_TOP_H2G_QUEUE_DEPTH_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. @@ -131,7 +119,7 @@ pub fn scratch_base_gva(size: usize) -> u64 { /// Compute the minimum scratch region size needed for a sandbox. /// -/// The transport allowance contains one page-backed ring arena and both +/// The fixed transport prefix contains one page-backed ring arena and both /// page-backed buffer pools. The result saturates at [`usize::MAX`]. pub fn min_scratch_size( input_data_size: usize, @@ -142,27 +130,231 @@ pub fn min_scratch_size( h2g_pool_pages: usize, ) -> usize { let size = arch::min_scratch_size(input_data_size, output_data_size).and_then(|fixed| { - let h2g_ring_offset = virtq::Layout::query_size(g2h_queue_depth) + let g2h = QueueDims::new(g2h_queue_depth, g2h_pool_pages)?; + let h2g = QueueDims::new(h2g_queue_depth, h2g_pool_pages)?; + + let transport_len = TransportArena::checked_query_size(g2h, h2g)?; + fixed.checked_add(transport_len) + }); + + size.unwrap_or(usize::MAX) +} + +/// Validated address independent dimensions for one transport queue. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct QueueDims { + depth: NonZeroU16, + pool_pages: NonZeroUsize, +} + +impl QueueDims { + /// Validate one queue descriptor count and pool page count. + pub fn new(depth: usize, pool_pages: usize) -> Option { + let depth = u16::try_from(depth).ok()?; + let depth = NonZeroU16::new(depth)?; + + if !depth.get().is_power_of_two() { + return None; + } + + let pool_pages = NonZeroUsize::new(pool_pages)?; + Some(Self { depth, pool_pages }) + } + + /// Number of descriptors in the queue. + pub const fn depth(&self) -> NonZeroU16 { + self.depth + } + + /// Number of pages in the queue's buffer pool. + pub const fn pool_pages(&self) -> NonZeroUsize { + self.pool_pages + } + + /// Compute the ring length, returning `None` on arithmetic overflow. + pub fn checked_ring_len(&self) -> Option { + virtq::Layout::checked_query_size(usize::from(self.depth.get())) + } + + /// Compute the pool length, returning `None` on arithmetic overflow. + pub fn checked_pool_len(&self) -> Option { + self.pool_pages.get().checked_mul(crate::vmem::PAGE_SIZE) + } +} + +/// Addresses of both rings and pools in one fixed transport arena. +/// +/// The G2H ring begins at the arena base. The H2G ring is descriptor aligned. +/// Both pools are page aligned. +/// +/// ```text +/// +----------+------------+----------+-----+----------+----------+ +/// | G2H ring | align pad | H2G ring | pad | G2H pool | H2G pool | +/// +----------+------------+----------+-----+----------+----------+ +/// ``` +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct TransportArena { + /// Address of the G2H ring and base of the arena. + g2h_ring_addr: u64, + /// Address of the H2G ring. + h2g_ring_addr: u64, + /// Address of the G2H pool. + g2h_pool_addr: u64, + /// Address of the H2G pool. + h2g_pool_addr: u64, + /// Page-aligned length occupied by both rings. + ring_span_len: usize, + /// Total page-aligned arena length. + len: usize, +} + +impl TransportArena { + /// Derive one transport arena from its base address and queue dimensions. + pub fn new(base_addr: u64, g2h: QueueDims, h2g: QueueDims) -> Option { + if !base_addr.is_multiple_of(crate::vmem::PAGE_SIZE as u64) { + return None; + } + + let h2g_ring_offset = g2h + .checked_ring_len()? .checked_next_multiple_of(virtq::Descriptor::ALIGN)?; - let ring_pages = h2g_ring_offset - .checked_add(virtq::Layout::query_size(h2g_queue_depth))? + let g2h_pool_offset = h2g_ring_offset + .checked_add(h2g.checked_ring_len()?)? .checked_next_multiple_of(crate::vmem::PAGE_SIZE)?; - let pool_size = g2h_pool_pages - .checked_add(h2g_pool_pages)? - .checked_mul(crate::vmem::PAGE_SIZE)?; + let g2h_pool_len = g2h.checked_pool_len()?; + let h2g_pool_offset = g2h_pool_offset.checked_add(g2h_pool_len)?; - fixed.checked_add(ring_pages)?.checked_add(pool_size) - }); + let h2g_pool_len = h2g.checked_pool_len()?; + let len = h2g_pool_offset.checked_add(h2g_pool_len)?; - size.unwrap_or(usize::MAX) + let addr = |offset: usize| base_addr.checked_add(u64::try_from(offset).ok()?); + let _end_addr = addr(len)?; + + Some(Self { + g2h_ring_addr: base_addr, + h2g_ring_addr: addr(h2g_ring_offset)?, + g2h_pool_addr: addr(g2h_pool_offset)?, + h2g_pool_addr: addr(h2g_pool_offset)?, + ring_span_len: g2h_pool_offset, + len, + }) + } + + /// Compute the total arena size without assigning an address. + pub fn checked_query_size(g2h: QueueDims, h2g: QueueDims) -> Option { + Some(Self::new(0, g2h, h2g)?.len) + } + + /// Base address of the arena. + pub const fn base_addr(&self) -> u64 { + self.g2h_ring_addr + } + + /// Address of the G2H ring. + pub const fn g2h_ring_addr(&self) -> u64 { + self.g2h_ring_addr + } + + /// Address of the H2G ring. + pub const fn h2g_ring_addr(&self) -> u64 { + self.h2g_ring_addr + } + + /// Address of the G2H pool. + pub const fn g2h_pool_addr(&self) -> u64 { + self.g2h_pool_addr + } + + /// Address of the H2G pool. + pub const fn h2g_pool_addr(&self) -> u64 { + self.h2g_pool_addr + } + + /// Page-aligned length occupied by both rings. + pub const fn ring_span_len(&self) -> usize { + self.ring_span_len + } + + /// Total page-aligned arena length. + pub const fn size(&self) -> usize { + self.len + } + + /// Exclusive end address of the arena. + pub const fn end_addr(&self) -> u64 { + self.g2h_ring_addr + self.len as u64 + } + + /// Convert the arena's absolute addresses into offsets from the arena base. + pub fn to_offsets(&self) -> (usize, usize, usize, usize) { + // Already validated by `TransportArena::new`. + let to_offset = |addr| usize::try_from(addr - self.g2h_ring_addr).unwrap(); + + ( + to_offset(self.h2g_ring_addr), + to_offset(self.g2h_pool_addr), + to_offset(self.h2g_pool_addr), + self.len, + ) + } } #[cfg(test)] mod tests { use super::*; + #[test] + fn transport_arena_derives_aligned_regions() { + let base = 0x1_0000; + let g2h = QueueDims::new(64, 8).unwrap(); + let h2g = QueueDims::new(32, 4).unwrap(); + let arena = TransportArena::new(base, g2h, h2g).unwrap(); + + assert_eq!(arena.g2h_ring_addr(), base); + assert!( + arena + .h2g_ring_addr() + .is_multiple_of(virtq::Descriptor::ALIGN as u64) + ); + assert!( + arena + .g2h_pool_addr() + .is_multiple_of(crate::vmem::PAGE_SIZE as u64) + ); + assert_eq!( + arena.h2g_pool_addr(), + base + 9 * crate::vmem::PAGE_SIZE as u64 + ); + assert_eq!(arena.end_addr(), base + 13 * crate::vmem::PAGE_SIZE as u64); + assert_eq!( + arena.to_offsets(), + ( + 0x410, + crate::vmem::PAGE_SIZE, + 9 * crate::vmem::PAGE_SIZE, + 13 * crate::vmem::PAGE_SIZE, + ) + ); + assert_eq!(arena.ring_span_len(), crate::vmem::PAGE_SIZE); + assert_eq!(arena.size(), 13 * crate::vmem::PAGE_SIZE); + assert_eq!( + TransportArena::checked_query_size(g2h, h2g), + Some(arena.size()) + ); + assert_eq!(TransportArena::new(base + 1, g2h, h2g), None); + assert_eq!(QueueDims::new(3, 8), None); + assert_eq!(QueueDims::new(64, 0), None); + assert_eq!(QueueDims::new(usize::MAX, 8), None); + let oversized = QueueDims::new(64, usize::MAX).unwrap(); + assert_eq!(TransportArena::new(base, oversized, h2g), None); + assert_eq!( + TransportArena::new(u64::MAX - crate::vmem::PAGE_SIZE as u64 + 1, g2h, h2g,), + None + ); + } + #[test] fn minimum_scratch_includes_ring_arena_and_pools() { let fixed = arch::min_scratch_size(0, 0).unwrap(); @@ -177,5 +369,6 @@ mod tests { #[test] fn minimum_scratch_saturates_on_overflow() { assert_eq!(usize::MAX, min_scratch_size(0, 0, 64, 32, usize::MAX, 4)); + assert_eq!(usize::MAX, min_scratch_size(0, 0, usize::MAX, 32, 8, 4)); } } diff --git a/src/hyperlight_common/src/virtq/mod.rs b/src/hyperlight_common/src/virtq/mod.rs index 63e73f21f..9de87d5b8 100644 --- a/src/hyperlight_common/src/virtq/mod.rs +++ b/src/hyperlight_common/src/virtq/mod.rs @@ -282,6 +282,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. /// @@ -358,6 +363,18 @@ impl Layout { dev_evt_offset + event_size } + + /// Calculate the ring size, returning `None` on arithmetic overflow. + pub fn checked_query_size(num_descs: usize) -> Option { + let desc_size = num_descs.checked_mul(Descriptor::SIZE)?; + let event_size = EventSuppression::SIZE; + let align = EventSuppression::ALIGN; + + let drv_evt_offset = align_up_checked(desc_size, align)?; + let dev_evt_offset = align_up_checked(drv_evt_offset.checked_add(event_size)?, align)?; + + dev_evt_offset.checked_add(event_size) + } } /// Statistics about the current virtqueue state. diff --git a/src/hyperlight_guest/src/error.rs b/src/hyperlight_guest/src/error.rs index 0a33bce79..92cd078d0 100644 --- a/src/hyperlight_guest/src/error.rs +++ b/src/hyperlight_guest/src/error.rs @@ -19,6 +19,7 @@ use alloc::string::{String, ToString as _}; pub use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode; use hyperlight_common::func::Error as FuncError; +use hyperlight_common::virtq::VirtqError; use {anyhow, serde_json}; pub type Result = core::result::Result; @@ -80,6 +81,15 @@ impl From for HyperlightGuestError { } } +impl From for HyperlightGuestError { + fn from(error: VirtqError) -> Self { + Self { + kind: ErrorCode::GuestError, + message: format!("virtq: {error}"), + } + } +} + /// Extension trait to add context to `Option` and `Result` types in guest code, /// converting them to `Result`. /// diff --git a/src/hyperlight_guest/src/layout.rs b/src/hyperlight_guest/src/layout.rs index 860ab7c16..09c9fcf0a 100644 --- a/src/hyperlight_guest/src/layout.rs +++ b/src/hyperlight_guest/src/layout.rs @@ -39,17 +39,8 @@ pub fn snapshot_generation_gva() -> *mut u64 { pub fn g2h_queue_depth_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET) } -pub fn g2h_ring_gpa_gva() -> *mut u64 { - scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_RING_GPA_OFFSET) -} -pub fn h2g_ring_gpa_gva() -> *mut u64 { - scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_RING_GPA_OFFSET) -} -pub fn g2h_pool_gpa_gva() -> *mut u64 { - scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_POOL_GPA_OFFSET) -} -pub fn h2g_pool_gpa_gva() -> *mut u64 { - scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_H2G_POOL_GPA_OFFSET) +pub fn transport_arena_gpa_gva() -> *mut u64 { + scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET) } pub fn g2h_pool_pages_gva() -> *mut u64 { scratch_top_gva(hyperlight_common::layout::SCRATCH_TOP_G2H_POOL_PAGES_OFFSET) diff --git a/src/hyperlight_guest/src/lib.rs b/src/hyperlight_guest/src/lib.rs index 19e5ac5f2..de2c23ee3 100644 --- a/src/hyperlight_guest/src/lib.rs +++ b/src/hyperlight_guest/src/lib.rs @@ -25,6 +25,7 @@ pub mod error; pub mod exit; pub mod layout; pub mod prim_alloc; +pub mod transport; pub mod types; pub mod guest_handle { diff --git a/src/hyperlight_guest/src/transport/context.rs b/src/hyperlight_guest/src/transport/context.rs new file mode 100644 index 000000000..f47d120ca --- /dev/null +++ b/src/hyperlight_guest/src/transport/context.rs @@ -0,0 +1,152 @@ +/* +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 core::result; + +use hyperlight_common::virtq::{ + AllocError, G2H_LOWER_SLOT_COUNT, G2H_LOWER_SLOT_SIZE, Layout, Notifier, QueueStats, + SlotLayout, SlotPool, VirtqProducer, +}; + +use super::GuestMemOps; +use crate::error::{GuestErrorContext, Result}; + +/// Guest-side notifier for polled transport operation. +#[derive(Clone, Copy)] +pub struct GuestNotifier; + +impl Notifier for GuestNotifier { + fn notify(&self, _stats: QueueStats) {} +} + +/// Type alias for the guest-side G2H producer. +pub type G2hProducer = VirtqProducer; + +/// Type alias for the guest-side H2G producer. +pub type H2gProducer = VirtqProducer; + +/// Configuration for one queue passed to [`GuestContext::new`]. +pub struct QueueConfig { + /// Ring descriptor layout in shared memory. + pub layout: Layout, + /// Base GVA of the buffer pool region. + pub pool_gva: u64, + /// Number of pages in the buffer pool. + pub pool_pages: usize, + /// Size of each upper-tier buffer. + pub buffer_size: usize, +} + +/// Virtqueue runtime state for guest-host communication. +pub struct GuestContext { + /// Guest-to-host driver. + _g2h_producer: G2hProducer, + /// Host-to-guest driver. + h2g_producer: H2gProducer, + /// Size of each prefilled H2G buffer. + h2g_slot_size: usize, +} + +impl GuestContext { + /// Create a new context with G2H and H2G queues. + pub fn new(g2h: QueueConfig, h2g: QueueConfig) -> Result { + Self::with_mem(g2h, h2g, GuestMemOps::for_scratch()) + } + + /// Create a new context with memory access provided. + fn with_mem(g2h: QueueConfig, h2g: QueueConfig, mem: GuestMemOps) -> Result { + let g2h_pool = g2h_pool(g2h.pool_gva, g2h.pool_pages, g2h.buffer_size) + .with_context(|| "failed to create G2H pool")?; + let g2h_producer = VirtqProducer::new(g2h.layout, mem, GuestNotifier, g2h_pool); + + let h2g_pool = h2g_pool(h2g.pool_gva, h2g.pool_pages, h2g.buffer_size) + .with_context(|| "failed to create H2G slot pool")?; + let h2g_producer = VirtqProducer::new(h2g.layout, mem, GuestNotifier, h2g_pool); + + let mut ctx = Self { + _g2h_producer: g2h_producer, + h2g_producer, + h2g_slot_size: h2g.buffer_size, + }; + + ctx.prefill_h2g().expect("H2G initial prefill failed"); + Ok(ctx) + } + + /// Pre-fill H2G with writable buffers until its ring or pool is full. + fn prefill_h2g(&mut self) -> Result<()> { + let mut batch = self.h2g_producer.batch(); + + loop { + let chain = match batch.chain().writable(self.h2g_slot_size).build() { + Ok(chain) => chain, + Err(error) if error.is_transient() => { + batch.finish()?; + return Ok(()); + } + Err(error) => return Err(error.into()), + }; + + match batch.submit(chain) { + Ok(_) => {} + Err(error) if error.is_transient() => { + batch.finish()?; + return Ok(()); + } + Err(error) => return Err(error.into()), + } + } + } +} + +fn pool_len(pages: usize) -> result::Result { + pages + .checked_mul(hyperlight_common::vmem::PAGE_SIZE) + .ok_or(AllocError::Overflow) +} + +/// Build the uniform H2G pool. +/// +/// Every preposted receive buffer has the configured size so the host sees one +/// predictable capacity for guest calls. +fn h2g_pool(base: u64, pages: usize, buffer_size: usize) -> result::Result { + let count = pool_len(pages)? / buffer_size; + SlotPool::new(SlotLayout::new(base, buffer_size, count)) +} + +/// Build the tiered G2H pool. +/// +/// One page of 256-byte slots serves small control and log messages without +/// consuming configured-size slots. Complete slots in the remaining pages form +/// the upper tier. +fn g2h_pool(base: u64, pages: usize, upper_size: usize) -> result::Result { + let pool_len = pool_len(pages)?; + let lower_len = G2H_LOWER_SLOT_COUNT + .checked_mul(G2H_LOWER_SLOT_SIZE) + .ok_or(AllocError::Overflow)?; + + let upper_len = pool_len + .checked_sub(lower_len) + .ok_or(AllocError::EmptyRegion)?; + + let upper_count = upper_len / upper_size; + + let lower = SlotLayout::new(base, G2H_LOWER_SLOT_SIZE, G2H_LOWER_SLOT_COUNT); + let upper = SlotLayout::new(lower.end_addr()?, upper_size, upper_count); + SlotPool::new_tiered(lower, upper) +} diff --git a/src/hyperlight_guest/src/transport/mem.rs b/src/hyperlight_guest/src/transport/mem.rs new file mode 100644 index 000000000..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..6290089fd --- /dev/null +++ b/src/hyperlight_guest/src/transport/mod.rs @@ -0,0 +1,75 @@ +/* +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`]. + +pub mod context; +pub mod mem; + +use core::cell::RefCell; +use core::sync::atomic::{AtomicU8, Ordering}; + +pub use context::{GuestContext, QueueConfig}; +pub use mem::GuestMemOps; + +const UNINITIALIZED: u8 = 0; +const INITIALIZED: u8 = 1; + +static INIT_STATE: AtomicU8 = AtomicU8::new(UNINITIALIZED); +static GLOBAL_CONTEXT: SyncWrap>> = SyncWrap(RefCell::new(None)); + +struct SyncWrap(T); + +// SAFETY: Hyperlight guests have one vCPU and serialize guest entry. +unsafe impl Sync for SyncWrap {} + +/// Whether the virtqueue context is installed. +pub fn is_initialized() -> bool { + INIT_STATE.load(Ordering::Acquire) == INITIALIZED +} + +/// Run a closure with the global virtqueue context. +/// +/// # Panics +/// +/// Panics if the context is uninitialized or already borrowed. +pub fn with_context(f: impl FnOnce(&mut GuestContext) -> R) -> R { + assert!(is_initialized(), "transport context not initialized"); + let mut context = GLOBAL_CONTEXT.0.borrow_mut(); + f(context.as_mut().expect("transport context missing")) +} + +/// Install the global transport context. +/// +/// # Panics +/// +/// Panics if a context was already installed. +pub fn set_global_context(context: GuestContext) { + assert!( + INIT_STATE + .compare_exchange( + UNINITIALIZED, + INITIALIZED, + Ordering::SeqCst, + Ordering::SeqCst, + ) + .is_ok(), + "virtqueue context already initialized" + ); + *GLOBAL_CONTEXT.0.borrow_mut() = Some(context); +} diff --git a/src/hyperlight_guest_bin/src/lib.rs b/src/hyperlight_guest_bin/src/lib.rs index 5df92f647..83afe45c5 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 @@ -287,6 +288,9 @@ pub(crate) extern "C" fn generic_init( registration(); } + // Prepare transport before guest code starts. + transport::initialize(); + unsafe { hyperlight_main(); } diff --git a/src/hyperlight_guest_bin/src/transport.rs b/src/hyperlight_guest_bin/src/transport.rs new file mode 100644 index 000000000..ae51f1ac7 --- /dev/null +++ b/src/hyperlight_guest_bin/src/transport.rs @@ -0,0 +1,103 @@ +/* +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 (depth, pages, g2h_bufsz) = read_published_g2h(); + let g2h = QueueDims::new(depth, pages).expect("invalid G2H queue dimensions"); + + let (depth, pages, h2g_bufsz) = read_published_h2g(); + let h2g = QueueDims::new(depth, pages).expect("invalid H2G queue dimensions"); + + assert!(g2h_bufsz > 0 && h2g_bufsz > 0); + + let arena = TransportArena::new(transport_arena_gpa, g2h, h2g).expect("invalid virtq arena"); + let g2h_pages = g2h.pool_pages().get(); + let h2g_pages = h2g.pool_pages().get(); + + let g2h_ring_gva = scratch_gva(arena.g2h_ring_addr()); + let h2g_ring_gva = scratch_gva(arena.h2g_ring_addr()); + let g2h_pool_gva = scratch_gva(arena.g2h_pool_addr()); + let h2g_pool_gva = scratch_gva(arena.h2g_pool_addr()); + + let g2h_layout = + unsafe { Layout::from_base(g2h_ring_gva, g2h.depth()) }.expect("G2H layout is invalid"); + let h2g_layout = + unsafe { Layout::from_base(h2g_ring_gva, h2g.depth()) }.expect("H2G layout is invalid"); + + // Build the queues and prefill H2G before exposing either queue to the host. + let context = GuestContext::new( + QueueConfig { + layout: g2h_layout, + pool_gva: g2h_pool_gva, + pool_pages: g2h_pages, + buffer_size: g2h_bufsz, + }, + QueueConfig { + layout: h2g_layout, + pool_gva: h2g_pool_gva, + pool_pages: h2g_pages, + buffer_size: h2g_bufsz, + }, + ) + .expect("failed to create guest context"); + + guest_transport::set_global_context(context); +} + +fn scratch_gva(gpa: u64) -> u64 { + let ptr = phys_to_virt(gpa).expect("transport GPA is outside scratch"); + u64::try_from(ptr as usize).expect("transport GVA exceeds u64") +} + +fn read_published_g2h() -> (usize, usize, usize) { + // SAFETY: Generic initialization has mapped writable scratch metadata. + let depth_raw = unsafe { layout::g2h_queue_depth_gva().read_volatile() }; + let pages_raw = unsafe { layout::g2h_pool_pages_gva().read_volatile() }; + let bufsz_raw = unsafe { layout::g2h_buffer_size_gva().read_volatile() }; + + let depth = usize::try_from(depth_raw).expect("G2H queue depth exceeds usize"); + let pages = usize::try_from(pages_raw).expect("G2H pool page count exceeds usize"); + let bufsz = usize::try_from(bufsz_raw).expect("G2H buffer size exceeds usize"); + + (depth, pages, bufsz) +} + +fn read_published_h2g() -> (usize, usize, usize) { + // SAFETY: Generic initialization has mapped writable scratch metadata. + let depth_raw = unsafe { layout::h2g_queue_depth_gva().read_volatile() }; + let pages_raw = unsafe { layout::h2g_pool_pages_gva().read_volatile() }; + let bufsz_raw = unsafe { layout::h2g_buffer_size_gva().read_volatile() }; + + let depth = usize::try_from(depth_raw).expect("H2G queue depth exceeds usize"); + let pages = usize::try_from(pages_raw).expect("H2G pool page count exceeds usize"); + let bufsz = usize::try_from(bufsz_raw).expect("H2G buffer size exceeds usize"); + + (depth, pages, bufsz) +} diff --git a/src/hyperlight_host/src/mem/layout.rs b/src/hyperlight_host/src/mem/layout.rs index 28817f2a2..635d2d7d7 100644 --- a/src/hyperlight_host/src/mem/layout.rs +++ b/src/hyperlight_host/src/mem/layout.rs @@ -48,13 +48,17 @@ 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: +//! privileges fixed input, output, and transport regions: //! //! +-------------------------------------------+ (top of physical memory) //! | Exception Stack, Metadata | //! +-------------------------------------------+ (1 page below) //! | Scratch Memory | //! +-------------------------------------------+ +//! | Guest Page Tables | +//! +-------------------------------------------+ +//! | Transport Arena | +//! +-------------------------------------------+ //! | Output Data | //! +-------------------------------------------+ //! | Input Data | @@ -63,6 +67,7 @@ limitations under the License. use std::fmt::Debug; use std::mem::size_of; +use hyperlight_common::layout::TransportArena; use hyperlight_common::mem::{HyperlightPEB, PAGE_SIZE_USIZE}; use tracing::{Span, instrument}; @@ -399,6 +404,11 @@ impl SandboxMemoryLayout { if scratch_size > Self::MAX_MEMORY_SIZE { return Err(MemoryRequestTooBig(scratch_size, Self::MAX_MEMORY_SIZE)); } + if !scratch_size.is_multiple_of(PAGE_SIZE_USIZE) { + return Err(new_error!( + "scratch size {scratch_size} must be a multiple of {PAGE_SIZE_USIZE}" + )); + } let input_data_size = cfg.get_input_data_size(); let output_data_size = cfg.get_output_data_size(); let g2h_queue_depth = cfg.get_g2h_queue_depth(); @@ -498,6 +508,16 @@ impl SandboxMemoryLayout { self.h2g_pool_pages } + pub(crate) fn get_g2h_queue_dims(&self) -> hyperlight_common::layout::QueueDims { + hyperlight_common::layout::QueueDims::new(self.g2h_queue_depth, self.g2h_pool_pages) + .expect("validated G2H queue dimensions") + } + + pub(crate) fn get_h2g_queue_dims(&self) -> hyperlight_common::layout::QueueDims { + hyperlight_common::layout::QueueDims::new(self.h2g_queue_depth, self.h2g_pool_pages) + .expect("validated H2G queue dimensions") + } + /// Guest-visible prefix size of the snapshot blob. pub(crate) fn snapshot_size(&self) -> usize { self.snapshot_size @@ -778,8 +798,7 @@ impl SandboxMemoryLayout { /// Offset from the beginning of the scratch region to the location /// where page tables are eagerly copied on restore. pub(crate) fn get_pt_base_scratch_offset(&self) -> usize { - (self.input_data_size + self.output_data_size) - .next_multiple_of(hyperlight_common::vmem::PAGE_SIZE) + self.get_virtq_base_scratch_offset() + self.get_transport_arena().size() } /// Base GPA to which the page tables are eagerly copied on restore. @@ -793,6 +812,24 @@ impl SandboxMemoryLayout { self.get_pt_base_gpa() + self.pt_size.unwrap_or(0) as u64 } + fn get_virtq_base_scratch_offset(&self) -> usize { + (self.input_data_size + self.output_data_size) + .next_multiple_of(hyperlight_common::vmem::PAGE_SIZE) + } + + /// Exact transport placement in the fixed scratch prefix. + pub(crate) fn get_transport_arena(&self) -> TransportArena { + let base_gpa = hyperlight_common::layout::scratch_base_gpa(self.scratch_size) + + self.get_virtq_base_scratch_offset() as u64; + + TransportArena::new( + base_gpa, + self.get_g2h_queue_dims(), + self.get_h2g_queue_dims(), + ) + .expect("validated virtqueue arena dimensions") + } + /// Total size of guest memory in `self`'s memory layout. fn get_unaligned_memory_size(&self) -> usize { self.init_data_offset() + self.init_data_size @@ -878,6 +915,21 @@ mod tests { assert!(matches!(layout, Err(MemoryRequestTooSmall(_, usize::MAX)))); } + #[test] + fn rejects_unaligned_scratch_size() { + let mut cfg = SandboxConfiguration::default(); + cfg.set_scratch_size(SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 1); + + let error = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap_err(); + assert_eq!( + error.to_string(), + format!( + "scratch size {} must be a multiple of {PAGE_SIZE_USIZE}", + SandboxConfiguration::DEFAULT_SCRATCH_SIZE + 1 + ) + ); + } + #[test] fn test_max_memory_sandbox() { let mut cfg = SandboxConfiguration::default(); @@ -1033,7 +1085,15 @@ mod tests { pin_eq!(layout.get_input_data_buffer_scratch_host_offset(), 0); pin_eq!(layout.get_output_data_buffer_scratch_host_offset(), 0x2000); - pin_eq!(layout.get_pt_base_scratch_offset(), 0x4000); + pin_eq!(layout.get_pt_base_scratch_offset(), 0x11000); + + let arena = layout.get_transport_arena(); + let scratch_base_gpa = hyperlight_common::layout::scratch_base_gpa(0x20000); + pin_eq!(arena.g2h_ring_addr() - scratch_base_gpa, 0x4000); + pin_eq!(arena.h2g_ring_addr() - scratch_base_gpa, 0x4410); + pin_eq!(arena.g2h_pool_addr() - scratch_base_gpa, 0x5000); + pin_eq!(arena.h2g_pool_addr() - scratch_base_gpa, 0xd000); + pin_eq!(arena.end_addr() - scratch_base_gpa, 0x11000); // The output buffer sits one input buffer past the input // buffer in the guest's scratch view. @@ -1052,7 +1112,7 @@ mod tests { ); pin_eq!( layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x20000), - 0x4000 + 0x11000 ); // pt_size is zero here, so the first free scratch GPA equals // the page table base. @@ -1082,7 +1142,15 @@ mod tests { pin_eq!(layout.get_input_data_buffer_scratch_host_offset(), 0); pin_eq!(layout.get_output_data_buffer_scratch_host_offset(), 0x4000); - pin_eq!(layout.get_pt_base_scratch_offset(), 0x6000); + pin_eq!(layout.get_pt_base_scratch_offset(), 0x13000); + + let arena = layout.get_transport_arena(); + let scratch_base_gpa = hyperlight_common::layout::scratch_base_gpa(0x30000); + pin_eq!(arena.g2h_ring_addr() - scratch_base_gpa, 0x6000); + pin_eq!(arena.h2g_ring_addr() - scratch_base_gpa, 0x6410); + pin_eq!(arena.g2h_pool_addr() - scratch_base_gpa, 0x7000); + pin_eq!(arena.h2g_pool_addr() - scratch_base_gpa, 0xf000); + pin_eq!(arena.end_addr() - scratch_base_gpa, 0x13000); pin_eq!( layout.get_output_data_buffer_gva() - layout.get_input_data_buffer_gva(), @@ -1096,7 +1164,7 @@ mod tests { ); pin_eq!( layout.get_pt_base_gpa() - hyperlight_common::layout::scratch_base_gpa(0x30000), - 0x6000 + 0x13000 ); pin_eq!( layout.get_first_free_scratch_gpa(), diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs index c93f1cac1..17df5bb7d 100644 --- a/src/hyperlight_host/src/mem/mgr.rs +++ b/src/hyperlight_host/src/mem/mgr.rs @@ -30,6 +30,7 @@ use super::layout::SandboxMemoryLayout; use super::shared_mem::{ ExclusiveSharedMemory, GuestSharedMemory, HostSharedMemory, ReadonlySharedMemory, SharedMemory, }; +use super::virtq::{self, G2hConsumer, H2gConsumer}; use crate::hypervisor::regs::CommonSpecialRegisters; use crate::mem::memory_region::MemoryRegion; #[cfg(crashdump)] @@ -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,26 @@ pub(crate) struct SandboxMemoryManager { /// restored snapshot's own generation number so the guest-visible /// counter tracks which snapshot the sandbox is a clone of. pub(crate) snapshot_count: u64, + /// G2H consumer bound to the current scratch mapping. + pub(crate) g2h_consumer: Option, + /// H2G consumer bound to the current scratch mapping. + pub(crate) h2g_consumer: Option, +} + +impl Clone for SandboxMemoryManager { + fn clone(&self) -> Self { + Self { + shared_mem: self.shared_mem.clone(), + scratch_mem: self.scratch_mem.clone(), + layout: self.layout, + next_action: self.next_action, + original_entrypoint: self.original_entrypoint, + abort_buffer: self.abort_buffer.clone(), + snapshot_count: self.snapshot_count, + g2h_consumer: None, + h2g_consumer: None, + } + } } /// Buffer for building guest page tables during snapshot creation. @@ -290,6 +311,8 @@ where original_entrypoint: 0, abort_buffer: Vec::new(), snapshot_count: 0, + g2h_consumer: None, + h2g_consumer: None, } } @@ -372,6 +395,8 @@ impl SandboxMemoryManager { original_entrypoint: self.original_entrypoint, abort_buffer: self.abort_buffer, snapshot_count: self.snapshot_count, + g2h_consumer: None, + h2g_consumer: None, }; let guest_mgr = SandboxMemoryManager { shared_mem: gshm, @@ -381,6 +406,8 @@ impl SandboxMemoryManager { original_entrypoint: self.original_entrypoint, abort_buffer: Vec::new(), // Guest doesn't need abort buffer snapshot_count: self.snapshot_count, + g2h_consumer: None, + h2g_consumer: None, }; host_mgr.update_scratch_bookkeeping()?; Ok((host_mgr, guest_mgr)) @@ -388,6 +415,27 @@ impl SandboxMemoryManager { } impl SandboxMemoryManager { + /// Attach host consumers to a guest-produced initial transport image. + /// + /// Before guest initialization, the host publishes queue dimensions and the + /// transport arena GPA. The guest derives and initializes every fixed region + /// without consuming dynamic scratch. + /// + /// This method runs after the initialization VM exit. It checks the + /// published arena against the host layout, derives bounded GVA views, + /// and validates each directional ring before exposing either consumer. + /// Fresh sandboxes and pre-initialization restores use this path. + pub(crate) fn attach_virtq(&mut self) -> Result<()> { + if self.g2h_consumer.is_some() || self.h2g_consumer.is_some() { + return Err(new_error!("virtqueue consumers are already attached")); + } + + let (g2h, h2g) = virtq::attach(&self.layout, &self.scratch_mem)?; + self.g2h_consumer = Some(g2h); + self.h2g_consumer = Some(h2g); + Ok(()) + } + /// Reads a host function call from memory #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] pub(crate) fn get_host_function_call(&mut self) -> Result { @@ -479,6 +527,9 @@ impl SandboxMemoryManager { Option>, Option, )> { + self.g2h_consumer = None; + self.h2g_consumer = None; + let gsnapshot = if *snapshot.memory() == self.shared_mem { // If the snapshot memory is already the correct memory, // which is readonly, don't bother with restoring it, @@ -556,6 +607,38 @@ impl SandboxMemoryManager { self.snapshot_count, )?; + // Record the G2H and H2G queue depths, pool page counts, and buffer sizes. + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_G2H_QUEUE_DEPTH_OFFSET, + u64::try_from(self.layout.get_g2h_queue_depth())?, + )?; + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_G2H_POOL_PAGES_OFFSET, + u64::try_from(self.layout.get_g2h_pool_pages())?, + )?; + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_G2H_BUFFER_SIZE_OFFSET, + u64::try_from(self.layout.get_g2h_buffer_size())?, + )?; + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_H2G_QUEUE_DEPTH_OFFSET, + u64::try_from(self.layout.get_h2g_queue_depth())?, + )?; + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_H2G_POOL_PAGES_OFFSET, + u64::try_from(self.layout.get_h2g_pool_pages())?, + )?; + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_H2G_BUFFER_SIZE_OFFSET, + u64::try_from(self.layout.get_h2g_buffer_size())?, + )?; + + let transport_arena = self.layout.get_transport_arena(); + self.update_scratch_bookkeeping_item( + SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET, + transport_arena.base_addr(), + )?; + // Initialise the guest input and output data buffers in // scratch memory. TODO: remove the need for this. self.scratch_mem.write::( diff --git a/src/hyperlight_host/src/mem/mod.rs b/src/hyperlight_host/src/mem/mod.rs index 8130e17b6..fe63d83a7 100644 --- a/src/hyperlight_host/src/mem/mod.rs +++ b/src/hyperlight_host/src/mem/mod.rs @@ -38,5 +38,7 @@ pub mod shared_mem; /// Utilities for writing shared memory tests #[cfg(all(test, not(miri)))] // uses proptest which isn't miri-compatible pub(crate) mod shared_mem_tests; +/// Host virtqueue attachment and validation. +pub(crate) mod virtq; #[allow(dead_code)] pub(crate) mod virtq_mem; diff --git a/src/hyperlight_host/src/mem/virtq.rs b/src/hyperlight_host/src/mem/virtq.rs new file mode 100644 index 000000000..0607bff0d --- /dev/null +++ b/src/hyperlight_host/src/mem/virtq.rs @@ -0,0 +1,581 @@ +/* +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 attachment. +//! +//! The host publishes one transport arena address in scratch-top metadata. Guest +//! initialization builds both queues in those fixed regions. This module +//! validates the complete initial image before returning either consumer. + +use core::ops::Range; + +use hyperlight_common::layout::{QueueDims, TransportArena}; +use hyperlight_common::virtq::canonical::validate_canon_image; +use hyperlight_common::virtq::{ + Layout as VirtqLayout, MemOps, Notifier, QueueStats, VirtqConsumer, +}; + +use super::layout::{BaseGpaRegion, SandboxMemoryLayout}; +use super::shared_mem::{HostSharedMemory, SharedMemory}; +use super::virtq_mem::HostMemOps; +use crate::{Result, new_error}; + +/// Host-side G2H virtqueue consumer. +pub(crate) type G2hConsumer = VirtqConsumer; +/// Host-side H2G virtqueue consumer. +pub(crate) type H2gConsumer = VirtqConsumer; + +/// No-op notifier for polled host transport. +#[derive(Clone, Copy)] +pub(crate) struct HostNotifier; + +impl Notifier for HostNotifier { + fn notify(&self, _stats: QueueStats) {} +} + +/// Build both host consumers from a guest-produced initial transport image. +/// +/// The consumers are returned only after the host-assigned arena and both +/// directional ring images have passed validation. +pub(crate) fn attach( + layout: &SandboxMemoryLayout, + scratch_mem: &HostSharedMemory, +) -> Result<(G2hConsumer, H2gConsumer)> { + let validator = Validator::new(layout)?; + let arena_gpa = read_published_arena_gpa(scratch_mem)?; + let regions = validator.validate_published_arena(arena_gpa)?; + + let g2h_ring_mem = HostMemOps::new(scratch_mem, regions.g2h_ring.clone())?; + let g2h_pool_mem = HostMemOps::new(scratch_mem, regions.g2h_pool)?; + let g2h_layout = validator.validate_g2h(&g2h_ring_mem, regions.g2h_ring)?; + + let h2g_ring_mem = HostMemOps::new(scratch_mem, regions.h2g_ring.clone())?; + let h2g_pool_mem = HostMemOps::new(scratch_mem, regions.h2g_pool.clone())?; + let h2g_layout = validator.validate_h2g(&h2g_ring_mem, regions.h2g_ring, regions.h2g_pool)?; + + Ok(( + VirtqConsumer::new_split(g2h_layout, g2h_ring_mem, g2h_pool_mem, HostNotifier), + VirtqConsumer::new_split(h2g_layout, h2g_ring_mem, h2g_pool_mem, HostNotifier), + )) +} + +/// Bounded GVA regions derived from validated transport GPAs. +struct GvaRegions { + g2h_ring: Range, + h2g_ring: Range, + g2h_pool: Range, + h2g_pool: Range, +} + +#[derive(Clone, Copy)] +struct QueueConfig { + /// Address-independent queue dimensions. + dims: QueueDims, + /// Size of the ring image in bytes including event suppressions. + ring_len: usize, + /// Size of the buffer pool in bytes. + pool_len: usize, + /// Size of each buffer in the pool in bytes. + buffer_size: usize, +} + +impl QueueConfig { + fn new(dims: QueueDims, buffer_size: usize) -> Result { + let ring_len = dims + .checked_ring_len() + .ok_or_else(|| new_error!("ring size overflow"))?; + let pool_len = dims + .checked_pool_len() + .ok_or_else(|| new_error!("pool size overflow"))?; + + if buffer_size == 0 { + return Err(new_error!("buffer size is zero")); + } + + Ok(Self { + dims, + ring_len, + pool_len, + buffer_size, + }) + } +} + +/// Host-owned transport dimensions. +#[derive(Clone, Copy)] +struct Config { + /// Host-requested G2H configuration. + g2h: QueueConfig, + /// Host-requested H2G configuration. + h2g: QueueConfig, + /// Fixed host-assigned transport arena. + arena: TransportArena, + /// Number of one-descriptor chains posted before the H2G ring or pool fills. + h2g_prefill_chains: usize, +} + +impl Config { + /// Compute the host transport configuration from the memory layout. + fn from_layout(layout: &SandboxMemoryLayout) -> Result { + let g2h = QueueConfig::new(layout.get_g2h_queue_dims(), layout.get_g2h_buffer_size())?; + + let h2g = QueueConfig::new(layout.get_h2g_queue_dims(), layout.get_h2g_buffer_size())?; + + let h2g_prefill_chains = + usize::from(h2g.dims.depth().get()).min(h2g.pool_len / h2g.buffer_size); + let arena = layout.get_transport_arena(); + + Ok(Self { + g2h, + h2g, + arena, + h2g_prefill_chains, + }) + } +} + +/// Validates one initial transport image against one host layout. +struct Validator<'a> { + config: Config, + layout: &'a SandboxMemoryLayout, +} + +impl<'a> Validator<'a> { + fn new(layout: &'a SandboxMemoryLayout) -> Result { + Ok(Self { + config: Config::from_layout(layout)?, + layout, + }) + } + + /// Validate the initial G2H queue and return its layout. + fn validate_g2h(&self, mem: &M, ring: Range) -> Result { + // SAFETY: `ring` spans the configured image and `mem` keeps that image + // valid for the duration of validation. + let layout = unsafe { VirtqLayout::from_base(ring.start, self.config.g2h.dims.depth()) } + .map_err(|error| new_error!("invalid G2H ring layout: {error}"))?; + + validate_canon_image(mem, layout, 0, |_, _| false) + .map_err(|error| new_error!("invalid canonical G2H image: {error}"))?; + + Ok(layout) + } + + /// Validate the initial H2G queue and return its layout. + /// + /// Every available chain contains one configured size writable descriptor. + /// Descriptors must name distinct, slot-aligned ranges inside the H2G pool. + fn validate_h2g( + &self, + mem: &M, + ring: Range, + pool: Range, + ) -> Result { + // SAFETY: `ring` spans the configured image and `mem` keeps that image + // valid for the duration of validation. + let layout = unsafe { VirtqLayout::from_base(ring.start, self.config.h2g.dims.depth()) } + .map_err(|error| new_error!("invalid H2G ring layout: {error}"))?; + + let bufsz = self.config.h2g.buffer_size; + let prefill = self.config.h2g_prefill_chains; + + if prefill == 0 { + return Err(new_error!("H2G pool has no complete buffers")); + } + + // Record the accepted descriptor ranges to detect overlaps. + let mut accepted: Vec> = Vec::with_capacity(prefill); + + let image = validate_canon_image(mem, layout, prefill, |_, elem| { + let Ok(bufsz_u64) = u64::try_from(bufsz) else { + return false; + }; + + // all descriptors must be writable and match the configured buffer size + if !elem.writable || usize::try_from(elem.len).ok() != Some(bufsz) { + return false; + } + + let Some(offset) = elem.addr.checked_sub(pool.start) else { + return false; + }; + let Some(end) = elem.addr.checked_add(u64::from(elem.len)) else { + return false; + }; + + // all descriptors must be slot-aligned and remain inside the pool + if !offset.is_multiple_of(bufsz_u64) || end > pool.end { + return false; + } + + let buf = elem.addr..end; + + // all descriptors must name distinct ranges + if accepted + .iter() + .any(|other| buf.start < other.end && other.start < buf.end) + { + return false; + } + + accepted.push(buf); + true + }) + .map_err(|error| new_error!("invalid canonical H2G image: {error}"))?; + + // compare the number of accepted chains to the expected prefill count + if image.len() != prefill { + return Err(new_error!("invalid initial H2G chains")); + } + + Ok(layout) + } + + /// Validate the published arena and return its GVA regions. + fn validate_published_arena(&self, arena_gpa: u64) -> Result { + if arena_gpa != self.config.arena.base_addr() { + return Err(new_error!("published transport arena is invalid")); + } + + self.resolve_gva_regions() + } + + /// Translate validated transport GPAs into the GVA ranges used by descriptors. + fn resolve_gva_regions(&self) -> Result { + let to_gva = |gpa| { + let resolved = self + .layout + .resolve_gpa(gpa, &[]) + .ok_or_else(|| new_error!("GPA {gpa:#x} is outside scratch"))?; + + if !matches!(resolved.base, BaseGpaRegion::Scratch(())) { + return Err(new_error!("GPA {gpa:#x} is outside scratch")); + } + + hyperlight_common::layout::scratch_base_gva(self.layout.get_scratch_size()) + .checked_add(u64::try_from(resolved.offset)?) + .ok_or_else(|| new_error!("GPA {gpa:#x} to GVA translation overflow")) + }; + + let ( + g2h_ring_addr, + h2g_ring_addr, + g2h_pool_addr, + h2g_pool_addr, + g2h_ring_len, + h2g_ring_len, + g2h_pool_len, + h2g_pool_len, + ) = ( + self.config.arena.g2h_ring_addr(), + self.config.arena.h2g_ring_addr(), + self.config.arena.g2h_pool_addr(), + self.config.arena.h2g_pool_addr(), + self.config.g2h.ring_len, + self.config.h2g.ring_len, + self.config.g2h.pool_len, + self.config.h2g.pool_len, + ); + + Ok(GvaRegions { + g2h_ring: checked_region(to_gva(g2h_ring_addr)?, g2h_ring_len, "G2H ring")?, + h2g_ring: checked_region(to_gva(h2g_ring_addr)?, h2g_ring_len, "H2G ring")?, + g2h_pool: checked_region(to_gva(g2h_pool_addr)?, g2h_pool_len, "G2H pool")?, + h2g_pool: checked_region(to_gva(h2g_pool_addr)?, h2g_pool_len, "H2G pool")?, + }) + } +} + +/// Read the transport arena GPA from scratch-top metadata. +fn read_published_arena_gpa(scratch_mem: &HostSharedMemory) -> Result { + let offset = hyperlight_common::layout::SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET as usize; + scratch_mem.read::(scratch_mem.mem_size() - offset) +} + +#[cfg(test)] +fn write_published_arena_gpa(scratch_mem: &HostSharedMemory, arena_gpa: u64) -> Result<()> { + let offset = hyperlight_common::layout::SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET as usize; + scratch_mem.write::(scratch_mem.mem_size() - offset, arena_gpa) +} + +fn checked_region(start: u64, len: usize, tag: &str) -> Result> { + let end = start + .checked_add(u64::try_from(len)?) + .ok_or_else(|| new_error!("{tag} GVA range overflow"))?; + + Ok(start..end) +} + +#[cfg(test)] +mod tests { + use core::num::NonZeroU16; + + use hyperlight_common::virtq::{ + DescFlags, Descriptor, MemOps, SlotLayout, SlotPool, VirtqProducer, + }; + use hyperlight_common::vmem; + + use super::*; + use crate::mem::shared_mem::ExclusiveSharedMemory; + use crate::sandbox::SandboxConfiguration; + + const SCRATCH_SIZE: usize = 0x20_000; + const G2H_DEPTH: u16 = 16; + const H2G_DEPTH: u16 = 8; + const G2H_POOL_PAGES: usize = 3; + const H2G_POOL_PAGES: usize = 2; + const H2G_BUFFER_SIZE: usize = 3000; + + fn memory_layout() -> SandboxMemoryLayout { + let mut config = SandboxConfiguration::default(); + config.set_scratch_size(SCRATCH_SIZE); + config.set_g2h_queue_depth(G2H_DEPTH as usize); + config.set_h2g_queue_depth(H2G_DEPTH as usize); + config.set_h2g_buffer_size(H2G_BUFFER_SIZE); + config.set_g2h_pool_pages(G2H_POOL_PAGES); + config.set_h2g_pool_pages(H2G_POOL_PAGES); + SandboxMemoryLayout::new(config, 4096, 0, None).unwrap() + } + + fn attach_config() -> Config { + Config::from_layout(&memory_layout()).unwrap() + } + + fn host_scratch() -> HostSharedMemory { + let scratch = ExclusiveSharedMemory::new(SCRATCH_SIZE).unwrap(); + scratch.build().0 + } + + fn validate_published(arena_gpa: u64, config: Config) -> Result { + let layout = memory_layout(); + Validator { + config, + layout: &layout, + } + .validate_published_arena(arena_gpa) + } + + struct PreparedVirtq { + g2h_mem: HostMemOps, + h2g_mem: HostMemOps, + g2h_ring: Range, + h2g_ring: Range, + g2h_pool: Range, + h2g_pool: Range, + g2h_layout: VirtqLayout, + h2g_layout: VirtqLayout, + } + + fn prepared_virtq() -> PreparedVirtq { + let scratch = ExclusiveSharedMemory::new(SCRATCH_SIZE).unwrap(); + let (scratch, _) = scratch.build(); + + let layout = memory_layout(); + let config = Config::from_layout(&layout).unwrap(); + let scratch_base_gpa = hyperlight_common::layout::scratch_base_gpa(SCRATCH_SIZE); + let scratch_base_gva = hyperlight_common::layout::scratch_base_gva(SCRATCH_SIZE); + let to_gva = |gpa| scratch_base_gva + (gpa - scratch_base_gpa); + + let ring_base = to_gva(config.arena.g2h_ring_addr()); + let h2g_base = to_gva(config.arena.h2g_ring_addr()); + let g2h_pool_base = to_gva(config.arena.g2h_pool_addr()); + let g2h_pool_end = g2h_pool_base + (G2H_POOL_PAGES * vmem::PAGE_SIZE) as u64; + let h2g_pool_base = to_gva(config.arena.h2g_pool_addr()); + let h2g_pool_end = h2g_pool_base + (H2G_POOL_PAGES * vmem::PAGE_SIZE) as u64; + + // SAFETY: The scratch mapping covers both ring layouts. + let g2h_layout = unsafe { + VirtqLayout::from_base(ring_base, NonZeroU16::new(G2H_DEPTH).unwrap()).unwrap() + }; + // SAFETY: The scratch mapping covers both ring layouts. + let h2g_layout = unsafe { + VirtqLayout::from_base(h2g_base, NonZeroU16::new(H2G_DEPTH).unwrap()).unwrap() + }; + + let mem = HostMemOps::new(&scratch, ring_base..h2g_pool_end).unwrap(); + let h2g_prefill_chains = (H2G_POOL_PAGES * vmem::PAGE_SIZE) / H2G_BUFFER_SIZE; + + let h2g_pool = SlotPool::new(SlotLayout::new( + h2g_pool_base, + H2G_BUFFER_SIZE, + h2g_prefill_chains, + )) + .unwrap(); + + let mut h2g = VirtqProducer::new(h2g_layout, mem, HostNotifier, h2g_pool.clone()); + let mut batch = h2g.batch(); + + for _ in 0..h2g_pool.num_free() { + let chain = batch.chain().writable(H2G_BUFFER_SIZE).build().unwrap(); + batch.submit(chain).unwrap(); + } + + batch.finish().unwrap(); + write_published_arena_gpa(&scratch, config.arena.base_addr()).unwrap(); + + let g2h_ring = ring_base..ring_base + VirtqLayout::query_size(G2H_DEPTH as usize) as u64; + let h2g_ring = h2g_base..h2g_base + VirtqLayout::query_size(H2G_DEPTH as usize) as u64; + let g2h_pool = g2h_pool_base..g2h_pool_end; + let h2g_pool = h2g_pool_base..h2g_pool_end; + let g2h_mem = HostMemOps::new(&scratch, g2h_ring.clone()).unwrap(); + let h2g_mem = HostMemOps::new(&scratch, h2g_ring.clone()).unwrap(); + + PreparedVirtq { + g2h_mem, + h2g_mem, + g2h_ring, + h2g_ring, + g2h_pool, + h2g_pool, + g2h_layout, + h2g_layout, + } + } + + fn validate(prepared: &PreparedVirtq) -> Result<()> { + let layout = memory_layout(); + let validator = Validator::new(&layout)?; + + validator.validate_g2h(&prepared.g2h_mem, prepared.g2h_ring.clone())?; + validator.validate_h2g( + &prepared.h2g_mem, + prepared.h2g_ring.clone(), + prepared.h2g_pool.clone(), + )?; + Ok(()) + } + + fn read_desc(mem: &HostMemOps, layout: VirtqLayout, index: u16) -> Descriptor { + mem.read_val(layout.desc_table_addr() + u64::from(index) * Descriptor::SIZE as u64) + .unwrap() + } + + fn write_desc(mem: &HostMemOps, layout: VirtqLayout, index: u16, desc: Descriptor) { + mem.write_val( + layout.desc_table_addr() + u64::from(index) * Descriptor::SIZE as u64, + desc, + ) + .unwrap(); + } + + #[test] + fn validates_host_placed_regions() { + let config = attach_config(); + let regions = validate_published(config.arena.base_addr(), config).unwrap(); + + assert_eq!( + regions.g2h_ring.end - regions.g2h_ring.start, + config.g2h.ring_len as u64 + ); + assert_eq!( + regions.h2g_ring.end - regions.h2g_ring.start, + config.h2g.ring_len as u64 + ); + assert_eq!( + regions.g2h_pool.end - regions.g2h_pool.start, + config.g2h.pool_len as u64 + ); + assert_eq!( + regions.h2g_pool.end - regions.h2g_pool.start, + config.h2g.pool_len as u64 + ); + } + + #[test] + fn rejects_invalid_published_regions() { + let config = attach_config(); + let arena_gpa = config.arena.base_addr() + 1; + assert!(validate_published(arena_gpa, config).is_err()); + } + + #[test] + fn rejects_published_region_overflow() { + let config = attach_config(); + assert!(validate_published(u64::MAX, config).is_err()); + } + + #[test] + fn rejects_untranslatable_or_overflowing_gva_regions() { + let config = attach_config(); + let arena_gpa = config.arena.base_addr(); + let invalid = hyperlight_common::layout::scratch_base_gpa(SCRATCH_SIZE) - 1; + assert!(validate_published(invalid, config).is_err()); + + let mut config = config; + config.g2h.ring_len = usize::MAX; + assert!(validate_published(arena_gpa, config).is_err()); + } + + #[test] + fn validates_initial_virtq_images() { + validate(&prepared_virtq()).unwrap(); + } + + #[test] + fn rejects_h2g_descriptors_outside_pool() { + let prepared = prepared_virtq(); + let mut desc = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 0); + desc.addr = prepared.g2h_pool.start; + write_desc(&prepared.h2g_mem, prepared.h2g_layout, 0, desc); + assert!(validate(&prepared).is_err()); + } + + #[test] + fn rejects_nonzero_g2h_descriptors() { + let prepared = prepared_virtq(); + let mut desc = read_desc(&prepared.g2h_mem, prepared.g2h_layout, 0); + desc.addr = prepared.g2h_pool.start; + write_desc(&prepared.g2h_mem, prepared.g2h_layout, 0, desc); + assert!(validate(&prepared).is_err()); + } + + #[test] + fn rejects_readable_h2g_descriptor() { + let prepared = prepared_virtq(); + let mut desc = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 0); + desc.flags &= !DescFlags::WRITE.bits(); + write_desc(&prepared.h2g_mem, prepared.h2g_layout, 0, desc); + assert!(validate(&prepared).is_err()); + } + + #[test] + fn rejects_invalid_h2g_size() { + let prepared = prepared_virtq(); + let mut desc = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 0); + desc.len -= 1; + write_desc(&prepared.h2g_mem, prepared.h2g_layout, 0, desc); + assert!(validate(&prepared).is_err()); + } + + #[test] + fn rejects_misaligned_h2g_descriptor() { + let prepared = prepared_virtq(); + let mut desc = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 0); + desc.addr += 1; + write_desc(&prepared.h2g_mem, prepared.h2g_layout, 0, desc); + assert!(validate(&prepared).is_err()); + } + + #[test] + fn rejects_overlapping_h2g_descriptors() { + let prepared = prepared_virtq(); + let first = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 0); + let mut second = read_desc(&prepared.h2g_mem, prepared.h2g_layout, 1); + second.addr = first.addr; + write_desc(&prepared.h2g_mem, prepared.h2g_layout, 1, second); + assert!(validate(&prepared).is_err()); + } +} diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index 54207deed..98bfb8f3f 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -242,6 +242,10 @@ impl MultiUseSandbox { let mgr = crate::mem::mgr::SandboxMemoryManager::from_snapshot(&snapshot)?; let (mut hshm, gshm) = mgr.build()?; + let attach_virtq = matches!( + snapshot.next_action(), + super::snapshot::NextAction::Initialise(_) + ); let page_size = u32::try_from(page_size::get())? as usize; @@ -322,6 +326,10 @@ impl MultiUseSandbox { #[cfg(gdb)] let dbg_mem_wrapper = Arc::new(Mutex::new(hshm.clone())); + if attach_virtq { + hshm.attach_virtq()?; + } + let sbox = MultiUseSandbox::from_uninit( host_funcs, hshm, diff --git a/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs index ddf407cc3..4c19b3830 100644 --- a/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs +++ b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs @@ -38,6 +38,10 @@ use crate::{MultiUseSandbox, Result, UninitializedSandbox}; #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")] pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result { let (mut hshm, gshm) = u_sbox.mgr.build()?; + let attach_virtq = matches!( + hshm.next_action, + crate::sandbox::snapshot::NextAction::Initialise(_) + ); // Get the host page size. Narrowed to u32 because the guest ABI // passes it via a 32-bit register (rdx), but widened back to usize @@ -109,6 +113,10 @@ pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result Date: Fri, 31 Jul 2026 14:22:21 +0200 Subject: [PATCH 13/34] feat(snapshot): preserve canonical virtq state Persist validated G2H and H2G ring images in running snapshots. Restore fixed transport allocations and install fresh host consumers before sandbox execution. Signed-off-by: Tomasz Andrzejak --- CHANGELOG.md | 2 + docs/snapshot-oci-format.md | 8 +- docs/snapshot-versioning.md | 5 +- src/hyperlight_host/src/mem/layout.rs | 45 ++--- src/hyperlight_host/src/mem/mgr.rs | 90 +++++---- src/hyperlight_host/src/mem/virtq.rs | 173 +++++++++++++++++- src/hyperlight_host/src/mem/virtq_mem.rs | 72 +++++++- .../src/sandbox/initialized_multi_use.rs | 15 ++ .../src/sandbox/snapshot/file/config.rs | 85 ++++++++- .../src/sandbox/snapshot/file/media_types.rs | 6 +- .../src/sandbox/snapshot/file/mod.rs | 29 ++- .../src/sandbox/snapshot/file_tests.rs | 58 +++++- .../src/sandbox/snapshot/mod.rs | 20 ++ .../src/sandbox/snapshot/tripwires.rs | 4 +- .../tests/snapshot_goldens/goldens_version.rs | 2 +- 15 files changed, 528 insertions(+), 86 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e3c64b3c0..d04d7dc0f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,8 @@ 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()`. +* Place virtqueue rings and pools in host-owned scratch before page tables. + Snapshot ABI 2 rejects snapshots created with earlier layouts. ### Removed diff --git a/docs/snapshot-oci-format.md b/docs/snapshot-oci-format.md index 971b3c868..e77b892b9 100644 --- a/docs/snapshot-oci-format.md +++ b/docs/snapshot-oci-format.md @@ -31,11 +31,11 @@ Three blob kinds per tag: * **manifest** (`application/vnd.oci.image.manifest.v1+json`). Tiny JSON pointer record selected via `index.json`. References one config and one layer by digest. -* **config** (`application/vnd.hyperlight.snapshot.config.v1+json`). The +* **config** (`application/vnd.hyperlight.snapshot.config.v2+json`). The snapshot descriptor: arch, hypervisor, CPU vendor, ABI version, - resume address and captured registers, memory layout, registered - host functions, snapshot generation counter. Loaded eagerly and - fully parsed. + resume address and captured registers, memory and transport layout, + registered host functions, snapshot generation counter. Loaded + eagerly and fully parsed. * **layer / memory** (`application/vnd.hyperlight.snapshot.memory.v1`). The raw guest memory image, exactly `memory_size` bytes. mmap'd on restore. diff --git a/docs/snapshot-versioning.md b/docs/snapshot-versioning.md index f855c1576..5c12d2d4c 100644 --- a/docs/snapshot-versioning.md +++ b/docs/snapshot-versioning.md @@ -23,8 +23,8 @@ A snapshot carries three independently evolvable version markers: `MT_SNAPSHOT_CURRENT`. This is the on-wire format of the snapshot blob: framing, section ordering, alignment, dirty/zero-page elision, anything about how the bytes are packed inside the OCI layer. -* **Config schema**, `MT_CONFIG_V1` - (`application/vnd.hyperlight.snapshot.config.v1+json`), aliased as +* **Config schema**, `MT_CONFIG_V2` + (`application/vnd.hyperlight.snapshot.config.v2+json`), aliased as `MT_CONFIG_CURRENT`. This is the JSON shape of the config blob: field names, types, required vs optional, the descriptors the loader needs in order to reconstruct the sandbox (memory sizes, buffer @@ -382,4 +382,3 @@ major: * The loader accepts the old `abi_version` (Option 2 step 4), so the old golden loads. * Register the host functions the old golden's checks call. - diff --git a/src/hyperlight_host/src/mem/layout.rs b/src/hyperlight_host/src/mem/layout.rs index 635d2d7d7..a762c64f3 100644 --- a/src/hyperlight_host/src/mem/layout.rs +++ b/src/hyperlight_host/src/mem/layout.rs @@ -342,8 +342,8 @@ impl Debug for SandboxMemoryLayout { impl SandboxMemoryLayout { /// Whether `other` has the same active memory layout as `self`. /// - /// Transport configuration does not participate while snapshots use the - /// stack transport. `snapshot_size` and `pt_size` are outputs of building a + /// 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. @@ -359,12 +359,12 @@ impl SandboxMemoryLayout { init_data_size, init_data_permissions, scratch_size, - g2h_queue_depth: _, - h2g_queue_depth: _, - g2h_buffer_size: _, - h2g_buffer_size: _, - g2h_pool_pages: _, - h2g_pool_pages: _, + g2h_queue_depth, + h2g_queue_depth, + g2h_buffer_size, + h2g_buffer_size, + g2h_pool_pages, + h2g_pool_pages, snapshot_size: _, pt_size: _, } = self; @@ -375,6 +375,12 @@ impl SandboxMemoryLayout { && *init_data_size == other.init_data_size && *init_data_permissions == other.init_data_permissions && *scratch_size == other.scratch_size + && *g2h_queue_depth == other.g2h_queue_depth + && *h2g_queue_depth == other.h2g_queue_depth + && *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. @@ -976,6 +982,12 @@ mod tests { |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_depth *= 2, + |l| l.h2g_queue_depth *= 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); }, @@ -992,23 +1004,6 @@ mod tests { } } - #[test] - fn is_compatible_with_ignores_inactive_transport_configuration() { - let base = - SandboxMemoryLayout::new(SandboxConfiguration::default(), 4096, 0, None).unwrap(); - let mut cfg = SandboxConfiguration::default(); - cfg.set_g2h_queue_depth(128); - cfg.set_h2g_queue_depth(16); - cfg.set_g2h_buffer_size(16 * 1024); - cfg.set_h2g_buffer_size(8 * 1024); - cfg.set_g2h_pool_pages(16); - cfg.set_h2g_pool_pages(8); - let other = SandboxMemoryLayout::new(cfg, 4096, 0, None).unwrap(); - - assert!(base.is_compatible_with(&other)); - assert!(other.is_compatible_with(&base)); - } - /// Pinned region offsets. These methods place every region that a /// restored snapshot is interpreted against, so a change shifts /// where the loader reads captured bytes and breaks existing diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs index 17df5bb7d..cb08832d1 100644 --- a/src/hyperlight_host/src/mem/mgr.rs +++ b/src/hyperlight_host/src/mem/mgr.rs @@ -320,37 +320,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 { @@ -415,6 +384,44 @@ impl SandboxMemoryManager { } impl SandboxMemoryManager { + /// Create a snapshot with the given mapped regions. + #[allow(clippy::too_many_arguments)] + pub(crate) fn snapshot( + &mut self, + mapped_regions: Vec, + root_pt_gpas: &[u64], + rsp_gva: u64, + sregs: CommonSpecialRegisters, + #[cfg(target_arch = "x86_64")] msrs: Vec, + next_action: NextAction, + host_functions: HostFunctionDetails, + ) -> Result { + let virtq = match (&self.g2h_consumer, &self.h2g_consumer) { + (Some(_), Some(_)) => Some(virtq::snapshot(&self.layout, &self.scratch_mem)?), + (None, None) => None, + _ => return Err(new_error!("virtqueue consumer ownership is incomplete")), + }; + + self.snapshot_count += 1; + Snapshot::new( + &mut self.shared_mem, + &mut self.scratch_mem, + self.layout, + crate::mem::exe::LoadInfo::dummy(), + mapped_regions, + root_pt_gpas, + rsp_gva, + sregs, + #[cfg(target_arch = "x86_64")] + msrs, + next_action, + self.original_entrypoint, + self.snapshot_count, + host_functions, + virtq, + ) + } + /// Attach host consumers to a guest-produced initial transport image. /// /// Before guest initialization, the host publishes queue dimensions and the @@ -436,6 +443,22 @@ impl SandboxMemoryManager { Ok(()) } + /// Restore a captured canonical transport image against this scratch mapping. + pub(crate) fn restore_virtq(&mut self, snapshot: Option<&virtq::VirtqSnapshot>) -> Result<()> { + let Some(snapshot) = snapshot else { + return Ok(()); + }; + + if self.g2h_consumer.is_some() || self.h2g_consumer.is_some() { + return Err(new_error!("virtqueue consumers are already attached")); + } + + let (g2h, h2g) = virtq::restore(&self.layout, &self.scratch_mem, snapshot)?; + self.g2h_consumer = Some(g2h); + self.h2g_consumer = Some(h2g); + Ok(()) + } + /// Reads a host function call from memory #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] pub(crate) fn get_host_function_call(&mut self) -> Result { @@ -527,6 +550,10 @@ impl SandboxMemoryManager { Option>, Option, )> { + if let Some(virtq) = snapshot.virtq() { + virtq.preflight(snapshot.layout())?; + } + self.g2h_consumer = None; self.h2g_consumer = None; @@ -572,6 +599,7 @@ impl SandboxMemoryManager { self.original_entrypoint = snapshot.original_entrypoint(); self.update_scratch_bookkeeping()?; + self.restore_virtq(snapshot.virtq())?; Ok((gsnapshot, gscratch)) } diff --git a/src/hyperlight_host/src/mem/virtq.rs b/src/hyperlight_host/src/mem/virtq.rs index 0607bff0d..534934db0 100644 --- a/src/hyperlight_host/src/mem/virtq.rs +++ b/src/hyperlight_host/src/mem/virtq.rs @@ -30,7 +30,7 @@ use hyperlight_common::virtq::{ use super::layout::{BaseGpaRegion, SandboxMemoryLayout}; use super::shared_mem::{HostSharedMemory, SharedMemory}; -use super::virtq_mem::HostMemOps; +use super::virtq_mem::{HostMemOps, ImageMem}; use crate::{Result, new_error}; /// Host-side G2H virtqueue consumer. @@ -72,6 +72,45 @@ pub(crate) fn attach( )) } +/// Capture the canonical transport state omitted from the main memory snapshot. +pub(crate) fn snapshot( + layout: &SandboxMemoryLayout, + scratch_mem: &HostSharedMemory, +) -> Result { + let validator = Validator::new(layout)?; + + let arena_gpa = read_published_arena_gpa(scratch_mem)?; + let regions = validator.validate_published_arena(arena_gpa)?; + + let g2h_mem = HostMemOps::new(scratch_mem, regions.g2h_ring.clone())?; + validator.validate_g2h(&g2h_mem, regions.g2h_ring.clone())?; + + let h2g_mem = HostMemOps::new(scratch_mem, regions.h2g_ring.clone())?; + validator.validate_h2g(&h2g_mem, regions.h2g_ring.clone(), regions.h2g_pool.clone())?; + + // The vCPU is stopped, so the ring images and snapshotted guest producer + // bookkeeping describe the same instant. + Ok(VirtqSnapshot { + scratch_size: layout.get_scratch_size(), + g2h_ring: read_ring(scratch_mem, regions.g2h_ring)?, + h2g_ring: read_ring(scratch_mem, regions.h2g_ring)?, + }) +} + +/// Restore one captured canonical transport image and return fresh consumers. +pub(crate) fn restore( + layout: &SandboxMemoryLayout, + scratch_mem: &HostSharedMemory, + snapshot: &VirtqSnapshot, +) -> Result<(G2hConsumer, H2gConsumer)> { + let regions = Validator::new(layout)?.validate_snapshot(snapshot)?; + + write_published_arena_gpa(scratch_mem, layout.get_transport_arena().base_addr())?; + write_ring(scratch_mem, regions.g2h_ring, &snapshot.g2h_ring)?; + write_ring(scratch_mem, regions.h2g_ring, &snapshot.h2g_ring)?; + attach(layout, scratch_mem) +} + /// Bounded GVA regions derived from validated transport GPAs. struct GvaRegions { g2h_ring: Range, @@ -147,7 +186,22 @@ impl Config { } } -/// Validates one initial transport image against one host layout. +/// Canonical in-memory transport state excluded from ordinary snapshot pages. +#[derive(Debug, PartialEq, Eq)] +pub(crate) struct VirtqSnapshot { + scratch_size: usize, + g2h_ring: Vec, + h2g_ring: Vec, +} + +impl VirtqSnapshot { + /// Validate every captured field before mutating restored scratch. + pub(crate) fn preflight(&self, layout: &SandboxMemoryLayout) -> Result<()> { + Validator::new(layout)?.validate_snapshot(self).map(|_| ()) + } +} + +/// Validates live and captured transport images against one host layout. struct Validator<'a> { config: Config, layout: &'a SandboxMemoryLayout, @@ -253,6 +307,28 @@ impl<'a> Validator<'a> { self.resolve_gva_regions() } + fn validate_snapshot(&self, snapshot: &VirtqSnapshot) -> Result { + if snapshot.scratch_size != self.layout.get_scratch_size() { + return Err(new_error!( + "virtqueue snapshot scratch size {} does not match layout size {}", + snapshot.scratch_size, + self.layout.get_scratch_size() + )); + } + + let regions = self.resolve_gva_regions()?; + validate_ring_len("G2H", &snapshot.g2h_ring, self.config.g2h.ring_len)?; + validate_ring_len("H2G", &snapshot.h2g_ring, self.config.h2g.ring_len)?; + + let g2h_mem = ImageMem::new(regions.g2h_ring.start, &snapshot.g2h_ring); + self.validate_g2h(&g2h_mem, regions.g2h_ring.clone())?; + + let h2g_mem = ImageMem::new(regions.h2g_ring.start, &snapshot.h2g_ring); + self.validate_h2g(&h2g_mem, regions.h2g_ring.clone(), regions.h2g_pool.clone())?; + + Ok(regions) + } + /// Translate validated transport GPAs into the GVA ranges used by descriptors. fn resolve_gva_regions(&self) -> Result { let to_gva = |gpa| { @@ -305,12 +381,41 @@ fn read_published_arena_gpa(scratch_mem: &HostSharedMemory) -> Result { scratch_mem.read::(scratch_mem.mem_size() - offset) } -#[cfg(test)] fn write_published_arena_gpa(scratch_mem: &HostSharedMemory, arena_gpa: u64) -> Result<()> { let offset = hyperlight_common::layout::SCRATCH_TOP_TRANSPORT_ARENA_GPA_OFFSET as usize; scratch_mem.write::(scratch_mem.mem_size() - offset, arena_gpa) } +fn read_ring(scratch_mem: &HostSharedMemory, ring: Range) -> Result> { + let len = usize::try_from( + ring.end + .checked_sub(ring.start) + .ok_or_else(|| new_error!("invalid ring range"))?, + )?; + + let mem = HostMemOps::new(scratch_mem, ring.clone())?; + let mut bytes = vec![0; len]; + mem.read(ring.start, &mut bytes)?; + + Ok(bytes) +} + +fn write_ring(scratch_mem: &HostSharedMemory, ring: Range, bytes: &[u8]) -> Result<()> { + validate_ring_len("restored", bytes, usize::try_from(ring.end - ring.start)?)?; + let mem = HostMemOps::new(scratch_mem, ring.clone())?; + mem.write(ring.start, bytes) +} + +fn validate_ring_len(direction: &str, bytes: &[u8], expected: usize) -> Result<()> { + if bytes.len() != expected { + return Err(new_error!( + "{direction} snapshot ring length {} and expected length {expected}", + bytes.len() + )); + } + Ok(()) +} + fn checked_region(start: u64, len: usize, tag: &str) -> Result> { let end = start .checked_add(u64::try_from(len)?) @@ -369,6 +474,7 @@ mod tests { } struct PreparedVirtq { + scratch: HostSharedMemory, g2h_mem: HostMemOps, h2g_mem: HostMemOps, g2h_ring: Range, @@ -434,6 +540,7 @@ mod tests { let h2g_mem = HostMemOps::new(&scratch, h2g_ring.clone()).unwrap(); PreparedVirtq { + scratch, g2h_mem, h2g_mem, g2h_ring, @@ -524,6 +631,66 @@ mod tests { validate(&prepared_virtq()).unwrap(); } + #[test] + fn snapshots_and_restores_canonical_image() { + let prepared = prepared_virtq(); + let layout = memory_layout(); + let stale_pool = [0xa5; 16]; + let pool_mem = HostMemOps::new(&prepared.scratch, prepared.h2g_pool.clone()).unwrap(); + pool_mem + .write(prepared.h2g_pool.start, &stale_pool) + .unwrap(); + + let captured = snapshot(&layout, &prepared.scratch).unwrap(); + let restored = host_scratch(); + let allocator = layout.get_first_free_scratch_gpa(); + let allocator_offset = + restored.mem_size() - hyperlight_common::layout::SCRATCH_TOP_ALLOCATOR_OFFSET as usize; + restored.write::(allocator_offset, allocator).unwrap(); + + restore(&layout, &restored, &captured).unwrap(); + let restored_snapshot = snapshot(&layout, &restored).unwrap(); + let restored_pool = HostMemOps::new(&restored, prepared.h2g_pool.clone()).unwrap(); + let mut pool_bytes = [0; 16]; + restored_pool + .read(prepared.h2g_pool.start, &mut pool_bytes) + .unwrap(); + + assert_eq!(restored_snapshot, captured); + assert_eq!(restored.read::(allocator_offset).unwrap(), allocator); + assert_eq!(pool_bytes, [0; 16]); + } + + #[test] + fn rejects_corrupt_snapshot_ring_before_restore() { + let prepared = prepared_virtq(); + let layout = memory_layout(); + let mut snapshot = snapshot(&layout, &prepared.scratch).unwrap(); + snapshot.h2g_ring.fill(0); + let restored = host_scratch(); + + assert!(restore(&layout, &restored, &snapshot).is_err()); + assert_eq!(read_published_arena_gpa(&restored).unwrap(), 0); + } + + #[test] + fn restores_with_grown_page_tables() { + let prepared = prepared_virtq(); + let layout = memory_layout(); + let snapshot = snapshot(&layout, &prepared.scratch).unwrap(); + let mut grown_layout = layout; + grown_layout + .set_pt_size(layout.get_pt_size() + vmem::PAGE_SIZE) + .unwrap(); + let restored = host_scratch(); + + restore(&grown_layout, &restored, &snapshot).unwrap(); + assert_eq!( + read_published_arena_gpa(&restored).unwrap(), + grown_layout.get_transport_arena().base_addr() + ); + } + #[test] fn rejects_h2g_descriptors_outside_pool() { let prepared = prepared_virtq(); diff --git a/src/hyperlight_host/src/mem/virtq_mem.rs b/src/hyperlight_host/src/mem/virtq_mem.rs index d7ecfdaf2..dd66e7b8e 100644 --- a/src/hyperlight_host/src/mem/virtq_mem.rs +++ b/src/hyperlight_host/src/mem/virtq_mem.rs @@ -14,11 +14,11 @@ See the License for the specific language governing permissions and limitations under the License. */ -//! Host [`MemOps`] access to a bounded scratch region. +//! Host [`MemOps`] implementations for live scratch and captured ring images. //! -//! Every operation uses [`HostSharedMemory`]'s checked API and acquires its -//! lifecycle read lock. This preserves exclusive-memory coordination but makes -//! descriptor traversal pay for one lock acquisition per field access. +//! Live scratch operations use [`HostSharedMemory`]'s checked API and acquire +//! its lifecycle read lock. This preserves exclusive-memory coordination but +//! makes descriptor traversal pay for one lock acquisition per field access. use core::mem::size_of; use core::ops::Range; @@ -143,6 +143,70 @@ unsafe impl MemOps for HostMemOps { } } +/// Read-only [`MemOps`] view over a captured ring image. +/// +/// Snapshot preflight must validate captured bytes before writing them into +/// restored scratch. This view maps the image to its captured ring GVA, letting +/// the same directional validators handle snapshots and live [`HostMemOps`]. +pub(super) struct ImageMem<'a> { + base: u64, + bytes: &'a [u8], +} + +impl<'a> ImageMem<'a> { + pub(super) fn new(base: u64, bytes: &'a [u8]) -> Self { + Self { base, bytes } + } + + fn offset(&self, addr: u64, len: usize) -> Result { + let out_of_bounds = || new_error!("image memory access is out of bounds"); + // VirtqLayout uses absolute GVAs, while the captured image starts at index zero. + let offset = addr.checked_sub(self.base).ok_or_else(&out_of_bounds)?; + let offset = usize::try_from(offset).map_err(|_| out_of_bounds())?; + let end = offset.checked_add(len).ok_or_else(&out_of_bounds)?; + + (end <= self.bytes.len()) + .then_some(offset) + .ok_or_else(out_of_bounds) + } +} + +// SAFETY: ImageMem provides immutable access only within `bytes`. Write +// operations fail, and the backing slice outlives every returned shared slice. +unsafe impl MemOps for ImageMem<'_> { + type Error = HyperlightError; + + fn read(&self, addr: u64, dst: &mut [u8]) -> Result<()> { + let offset = self.offset(addr, dst.len())?; + dst.copy_from_slice(&self.bytes[offset..offset + dst.len()]); + Ok(()) + } + + fn load_acquire(&self, addr: u64) -> Result { + let mut bytes = [0; size_of::()]; + self.read(addr, &mut bytes)?; + Ok(u16::from_ne_bytes(bytes)) + } + + unsafe fn as_slice(&self, addr: u64, len: usize) -> Result<&[u8]> { + let offset = self.offset(addr, len)?; + Ok(&self.bytes[offset..offset + len]) + } + + fn write(&self, _addr: u64, _src: &[u8]) -> Result<()> { + Err(new_error!("image memory is read-only")) + } + + fn store_release(&self, _addr: u64, _val: u16) -> Result<()> { + Err(new_error!("image memory is read-only")) + } + + #[allow(clippy::mut_from_ref)] + unsafe fn as_mut_slice(&self, _addr: u64, _len: usize) -> Result<&mut [u8]> { + Err(new_error!("image memory is read-only")) + } +} + #[cfg(test)] mod tests { use hyperlight_common::virtq::MemOps; diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index 98bfb8f3f..939313c98 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -328,6 +328,8 @@ impl MultiUseSandbox { if attach_virtq { hshm.attach_virtq()?; + } else { + hshm.restore_virtq(snapshot.virtq())?; } let sbox = MultiUseSandbox::from_uninit( @@ -1189,6 +1191,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 = { @@ -1747,6 +1754,7 @@ mod tests { let snapshot = sandbox.snapshot().unwrap(); sandbox2.restore(snapshot).unwrap(); + assert_virtq_attached(&sandbox2); assert_eq!(sandbox2.call::("GetStatic", ()).unwrap(), 42); } @@ -4012,8 +4020,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"); @@ -4025,6 +4035,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(); @@ -4046,6 +4057,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); @@ -4055,6 +4068,8 @@ mod tests { a.restore(snapshot.clone()).unwrap(); b.restore(snapshot).unwrap(); + super::assert_virtq_attached(&a); + super::assert_virtq_attached(&b); assert_eq!(a.call::("GetStatic", ()).unwrap(), 3); assert_eq!(b.call::("GetStatic", ()).unwrap(), 3); } diff --git a/src/hyperlight_host/src/sandbox/snapshot/file/config.rs b/src/hyperlight_host/src/sandbox/snapshot/file/config.rs index efae026d2..35b4a392e 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 @@ -222,6 +222,12 @@ pub(super) struct MemoryLayout { /// Memory region flag bits. `None` means default permissions. pub(super) init_data_permissions: Option, pub(super) scratch_size: usize, + pub(super) g2h_queue_depth: usize, + pub(super) h2g_queue_depth: usize, + pub(super) g2h_buffer_size: usize, + pub(super) h2g_buffer_size: usize, + pub(super) g2h_pool_pages: usize, + pub(super) h2g_pool_pages: usize, pub(super) snapshot_size: usize, pub(super) pt_size: Option, } @@ -472,6 +478,10 @@ impl OciSnapshotConfig { ("code_size", self.layout.code_size), ("init_data_size", self.layout.init_data_size), ("scratch_size", self.layout.scratch_size), + ("g2h_buffer_size", self.layout.g2h_buffer_size), + ("h2g_buffer_size", self.layout.h2g_buffer_size), + ("g2h_pool_pages", self.layout.g2h_pool_pages), + ("h2g_pool_pages", self.layout.h2g_pool_pages), ] { if value > max_region { return Err(crate::new_error!( @@ -483,6 +493,55 @@ impl OciSnapshotConfig { } } + let mut transport = crate::sandbox::SandboxConfiguration::default(); + transport.set_g2h_queue_depth(self.layout.g2h_queue_depth); + transport.set_h2g_queue_depth(self.layout.h2g_queue_depth); + transport.set_g2h_buffer_size(self.layout.g2h_buffer_size); + transport.set_h2g_buffer_size(self.layout.h2g_buffer_size); + transport.set_g2h_pool_pages(self.layout.g2h_pool_pages); + transport.set_h2g_pool_pages(self.layout.h2g_pool_pages); + + for (name, saved, normalized) in [ + ( + "g2h_queue_depth", + self.layout.g2h_queue_depth, + transport.get_g2h_queue_depth(), + ), + ( + "h2g_queue_depth", + self.layout.h2g_queue_depth, + transport.get_h2g_queue_depth(), + ), + ( + "g2h_buffer_size", + self.layout.g2h_buffer_size, + transport.get_g2h_buffer_size(), + ), + ( + "h2g_buffer_size", + self.layout.h2g_buffer_size, + transport.get_h2g_buffer_size(), + ), + ( + "g2h_pool_pages", + self.layout.g2h_pool_pages, + transport.get_g2h_pool_pages(), + ), + ( + "h2g_pool_pages", + self.layout.h2g_pool_pages, + transport.get_h2g_pool_pages(), + ), + ] { + if saved != normalized { + return Err(crate::new_error!( + "snapshot layout field {} ({}) is not a valid transport value", + name, + saved + )); + } + } + // The saved dispatch entrypoint must be in the executable code // region. Code occupies the page-rounded prefix of the snapshot. let code_lo = SandboxMemoryLayout::BASE_ADDRESS as u64; @@ -782,6 +841,12 @@ mod tests { init_data_size: 0, init_data_permissions: None, scratch_size: 0, + g2h_queue_depth: 64, + h2g_queue_depth: 32, + g2h_buffer_size: PAGE_SIZE, + h2g_buffer_size: PAGE_SIZE, + g2h_pool_pages: 8, + h2g_pool_pages: 4, snapshot_size: PAGE_SIZE, pt_size: None, }, @@ -848,7 +913,7 @@ mod schema_pin { const PINNED_CALL: &str = r#"{ "hyperlight_version": "x.y.z", "arch": "x86_64", - "abi_version": 1, + "abi_version": 2, "hypervisor": "mshv", "cpu_vendor": "intel", "stack_top_gva": 3735928559, @@ -1014,6 +1079,12 @@ mod schema_pin { "init_data_size": 5, "init_data_permissions": null, "scratch_size": 8, + "g2h_queue_depth": 64, + "h2g_queue_depth": 32, + "g2h_buffer_size": 4096, + "h2g_buffer_size": 4096, + "g2h_pool_pages": 8, + "h2g_pool_pages": 4, "snapshot_size": 9, "pt_size": null }, @@ -1034,7 +1105,7 @@ mod schema_pin { const PINNED_CALL: &str = r#"{ "hyperlight_version": "x.y.z", "arch": "aarch64", - "abi_version": 1, + "abi_version": 2, "hypervisor": "mshv", "cpu_vendor": "intel", "stack_top_gva": 3735928559, @@ -1056,6 +1127,12 @@ mod schema_pin { "init_data_size": 5, "init_data_permissions": null, "scratch_size": 8, + "g2h_queue_depth": 64, + "h2g_queue_depth": 32, + "g2h_buffer_size": 4096, + "h2g_buffer_size": 4096, + "g2h_pool_pages": 8, + "h2g_pool_pages": 4, "snapshot_size": 9, "pt_size": null }, @@ -1094,7 +1171,7 @@ mod schema_pin { assert_eq!( actual_value, pinned_value, "Snapshot config JSON schema changed. If the change can break \ - existing snapshots on disk, bump `MT_CONFIG_V1` in \ + existing snapshots on disk, bump `MT_CONFIG_CURRENT` in \ `super::media_types` and follow `docs/snapshot-versioning.md`. \ Either way, paste the actual output below into the matching \ `PINNED_*`.\n\nactual:\n{actual}" diff --git a/src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs b/src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs index 661bd4a04..a6191f6a1 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file/media_types.rs @@ -19,7 +19,9 @@ 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; @@ -27,7 +29,7 @@ pub(in crate::sandbox::snapshot) const MT_SNAPSHOT_CURRENT: &str = MT_SNAPSHOT_V /// ABI version for the snapshot memory blob. Bumped when the /// host-guest contract for the snapshot bytes changes. See /// docs/snapshot-versioning.md. -pub(in crate::sandbox::snapshot) const SNAPSHOT_ABI_VERSION: u32 = 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..d400ed59c 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file/mod.rs @@ -39,7 +39,8 @@ use self::media_types::{ ANNOTATION_ARCH, ANNOTATION_CPU, ANNOTATION_HYPERVISOR, ANNOTATION_REF_NAME, }; pub(super) use self::media_types::{ - MT_CONFIG_CURRENT, MT_CONFIG_V1, MT_SNAPSHOT_CURRENT, MT_SNAPSHOT_V1, SNAPSHOT_ABI_VERSION, + MT_CONFIG_CURRENT, MT_CONFIG_V1, MT_CONFIG_V2, MT_SNAPSHOT_CURRENT, MT_SNAPSHOT_V1, + SNAPSHOT_ABI_VERSION, }; use self::reference::{OciDigest, OciReference, OciTag}; use super::{NextAction, Snapshot}; @@ -625,6 +626,12 @@ impl Snapshot { init_data_size: l.init_data_size(), init_data_permissions: l.init_data_permissions().map(|f| f.bits()), scratch_size: l.get_scratch_size(), + g2h_queue_depth: l.get_g2h_queue_depth(), + h2g_queue_depth: l.get_h2g_queue_depth(), + g2h_buffer_size: l.get_g2h_buffer_size(), + h2g_buffer_size: l.get_h2g_buffer_size(), + g2h_pool_pages: l.get_g2h_pool_pages(), + h2g_pool_pages: l.get_h2g_pool_pages(), snapshot_size: l.snapshot_size(), pt_size: l.pt_size(), }, @@ -747,16 +754,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 )); } } @@ -816,6 +828,12 @@ impl Snapshot { sbox_cfg.set_output_data_size(cfg.layout.output_data_size); sbox_cfg.set_heap_size(cfg.layout.heap_size as u64); sbox_cfg.set_scratch_size(cfg.layout.scratch_size); + sbox_cfg.set_g2h_queue_depth(cfg.layout.g2h_queue_depth); + sbox_cfg.set_h2g_queue_depth(cfg.layout.h2g_queue_depth); + sbox_cfg.set_g2h_buffer_size(cfg.layout.g2h_buffer_size); + sbox_cfg.set_h2g_buffer_size(cfg.layout.h2g_buffer_size); + sbox_cfg.set_g2h_pool_pages(cfg.layout.g2h_pool_pages); + sbox_cfg.set_h2g_pool_pages(cfg.layout.h2g_pool_pages); let init_data_perms = match cfg.layout.init_data_permissions { None => None, Some(bits) => Some(MemoryRegionFlags::from_bits(bits).ok_or_else(|| { @@ -909,6 +927,7 @@ impl Snapshot { original_entrypoint: cfg.original_entrypoint_addr, snapshot_generation, host_functions, + virtq: None, }) } } diff --git a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs b/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs index 4e0604c86..4d4c7d977 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/file_tests.rs @@ -1824,6 +1824,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(); @@ -2286,7 +2300,7 @@ fn manifest_uses_correct_config_and_layer_media_types() { serde_json::from_slice(&std::fs::read(manifest_path(&path)).unwrap()).unwrap(); assert_eq!( manifest["config"]["mediaType"].as_str().unwrap(), - "application/vnd.hyperlight.snapshot.config.v1+json" + "application/vnd.hyperlight.snapshot.config.v2+json" ); assert_eq!(manifest["layers"].as_array().unwrap().len(), 1); assert_eq!( @@ -2298,7 +2312,7 @@ fn manifest_uses_correct_config_and_layer_media_types() { // that falls back to `config.mediaType` sees the same value. assert_eq!( manifest["artifactType"].as_str().unwrap(), - "application/vnd.hyperlight.snapshot.config.v1+json" + "application/vnd.hyperlight.snapshot.config.v2+json" ); } @@ -2778,6 +2792,46 @@ 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_depth(128); + cfg.set_h2g_queue_depth(16); + cfg.set_g2h_buffer_size(8192); + cfg.set_h2g_buffer_size(2048); + cfg.set_g2h_pool_pages(16); + cfg.set_h2g_pool_pages(6); + + let mut sbox = UninitializedSandbox::new( + GuestBinary::FilePath(simple_guest_as_string().unwrap()), + Some(cfg), + ) + .unwrap() + .evolve() + .unwrap(); + let snapshot = sbox.snapshot().unwrap(); + let expected = snapshot.layout().get_transport_arena(); + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("layout"); + snapshot + .save(&path, &OciTag::new("latest").unwrap()) + .unwrap(); + let loaded = Snapshot::checked_load(&path, OciTag::new("latest").unwrap()).unwrap(); + + assert_eq!(loaded.layout().get_g2h_queue_depth(), 128); + assert_eq!(loaded.layout().get_h2g_queue_depth(), 16); + assert_eq!(loaded.layout().get_g2h_buffer_size(), 8192); + assert_eq!(loaded.layout().get_h2g_buffer_size(), 2048); + assert_eq!(loaded.layout().get_g2h_pool_pages(), 16); + assert_eq!(loaded.layout().get_h2g_pool_pages(), 6); + assert_eq!(loaded.layout().get_transport_arena(), expected); +} + #[test] fn snapshot_config_records_entrypoint_and_sregs() { let snap = create_snapshot(); diff --git a/src/hyperlight_host/src/sandbox/snapshot/mod.rs b/src/hyperlight_host/src/sandbox/snapshot/mod.rs index 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..807402efe 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs @@ -28,8 +28,8 @@ use super::file::{ MT_CONFIG_CURRENT, MT_SNAPSHOT_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_OCI_LAYOUT_VERSION: &str = "1.0.0"; diff --git a/src/hyperlight_host/tests/snapshot_goldens/goldens_version.rs b/src/hyperlight_host/tests/snapshot_goldens/goldens_version.rs index 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) From 738caaa01539be1a549b4fa67ca05341e1f75088 Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Mon, 3 Aug 2026 18:51:31 +0200 Subject: [PATCH 14/34] feat(virtq): implement guest-to-host virtqueue transport Initialize queues before first guest entry. Carry host calls, logs, external values, and bounded responses over G2H chains. Signed-off-by: Tomasz Andrzejak --- .../src/flatbuffer_wrappers/codec.rs | 73 ++++ .../src/flatbuffer_wrappers/mod.rs | 2 +- src/hyperlight_common/src/layout.rs | 31 +- src/hyperlight_common/src/outb.rs | 3 + src/hyperlight_common/src/virtq/buffer.rs | 19 + src/hyperlight_common/src/virtq/consumer.rs | 2 + src/hyperlight_common/src/virtq/mod.rs | 2 + src/hyperlight_common/src/virtq/msg.rs | 105 +++++- src/hyperlight_common/src/virtq/pool.rs | 2 +- src/hyperlight_common/src/virtq/pool/slot.rs | 97 +++++ src/hyperlight_common/src/virtq/pool/tests.rs | 65 ++++ src/hyperlight_common/src/virtq/producer.rs | 6 + src/hyperlight_guest/src/error.rs | 14 +- .../src/guest_handle/host_comm.rs | 103 +----- src/hyperlight_guest/src/transport/context.rs | 329 +++++++++++++++-- src/hyperlight_guest/src/transport/mod.rs | 1 + .../src/transport/response.rs | 181 +++++++++ src/hyperlight_guest_bin/src/host_comm.rs | 28 +- src/hyperlight_guest_bin/src/lib.rs | 6 +- src/hyperlight_guest_capi/src/dispatch.rs | 21 +- src/hyperlight_guest_capi/src/flatbuffer.rs | 25 +- src/hyperlight_host/src/mem/mgr.rs | 38 +- src/hyperlight_host/src/mem/virtq.rs | 344 ++++++++++++++---- .../src/sandbox/initialized_multi_use.rs | 9 +- src/hyperlight_host/src/sandbox/outb.rs | 129 ++++++- .../src/sandbox/snapshot/tripwires.rs | 1 + .../src/sandbox/uninitialized_evolve.rs | 8 - src/hyperlight_host/tests/common/mod.rs | 10 + src/hyperlight_host/tests/integration_test.rs | 26 ++ .../tests/sandbox_host_tests.rs | 96 ++++- src/tests/rust_guests/simpleguest/src/main.rs | 51 ++- 31 files changed, 1541 insertions(+), 286 deletions(-) create mode 100644 src/hyperlight_guest/src/transport/response.rs diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/codec.rs b/src/hyperlight_common/src/flatbuffer_wrappers/codec.rs index 41052b9ac..b19f535fe 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/codec.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/codec.rs @@ -19,6 +19,79 @@ use alloc::vec::Vec; use anyhow::Result; use bytes::Bytes; +/// One borrowed external byte value emitted alongside a FlatBuffer. +#[derive(Clone, Copy, Debug)] +pub enum ExternalValueRef<'a> { + /// A logically contiguous value. + Bytes(&'a [u8]), + /// A logically chunked value. + Chunks(&'a [Bytes]), +} + +impl ExternalValueRef<'_> { + /// Total byte length of this value. + pub fn len(&self) -> Option { + match self { + Self::Bytes(value) => Some(value.len()), + Self::Chunks(chunks) => chunks + .iter() + .try_fold(0usize, |len, chunk| len.checked_add(chunk.len())), + } + } + + /// Whether this value contains no bytes. + pub fn is_empty(&self) -> bool { + self.len() == Some(0) + } +} + +/// Borrowed external values collected while encoding a FlatBuffer. +#[derive(Debug, Default)] +pub struct ExternalValueRefs<'a> { + values: Vec>, +} + +impl<'a> ExternalValueRefs<'a> { + /// Create an empty collection. + pub fn new() -> Self { + Self::default() + } + + /// Borrow the collected values in wire order. + pub fn as_slice(&self) -> &[ExternalValueRef<'a>] { + &self.values + } + + /// Number of collected logical values. + pub fn len(&self) -> usize { + self.values.len() + } + + /// Whether no logical values were collected. + pub fn is_empty(&self) -> bool { + self.values.is_empty() + } + + /// Total byte length of all collected values. + pub fn total_len(&self) -> Option { + self.values + .iter() + .try_fold(0usize, |len, value| len.checked_add(value.len()?)) + } +} + +impl<'a> ExternalValueSink<'a> for ExternalValueRefs<'a> { + fn push_bytes(&mut self, value: &'a [u8]) -> Result<()> { + self.values.push(ExternalValueRef::Bytes(value)); + Ok(()) + } + + fn push_chunks(&mut self, value: &'a [Bytes]) -> Result<()> { + self.values.push(ExternalValueRef::Chunks(value)); + Ok(()) + } +} + /// Receives external byte values while their FlatBuffer markers are encoded. /// /// Values are delivered in their logical order without flattening chunked diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs b/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs index d013cb78c..0a7770e27 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs @@ -31,4 +31,4 @@ pub mod host_function_definition; pub mod host_function_details; pub mod util; -pub use codec::{ExternalValueSink, ExternalValueSource}; +pub use codec::{ExternalValueRef, ExternalValueRefs, ExternalValueSink, ExternalValueSource}; diff --git a/src/hyperlight_common/src/layout.rs b/src/hyperlight_common/src/layout.rs index e0a597103..0de99c580 100644 --- a/src/hyperlight_common/src/layout.rs +++ b/src/hyperlight_common/src/layout.rs @@ -158,6 +158,10 @@ impl QueueDims { } let pool_pages = NonZeroUsize::new(pool_pages)?; + pool_pages.get().checked_mul(crate::vmem::PAGE_SIZE)?; + + virtq::Layout::checked_query_size(usize::from(depth.get()))?; + Some(Self { depth, pool_pages }) } @@ -171,14 +175,14 @@ impl QueueDims { self.pool_pages } - /// Compute the ring length, returning `None` on arithmetic overflow. - pub fn checked_ring_len(&self) -> Option { - virtq::Layout::checked_query_size(usize::from(self.depth.get())) + /// Ring length in bytes. + pub const fn ring_len(&self) -> usize { + virtq::Layout::query_size(self.depth.get() as usize) } - /// Compute the pool length, returning `None` on arithmetic overflow. - pub fn checked_pool_len(&self) -> Option { - self.pool_pages.get().checked_mul(crate::vmem::PAGE_SIZE) + /// Buffer pool length in bytes. + pub const fn pool_len(&self) -> usize { + self.pool_pages.get() * crate::vmem::PAGE_SIZE } } @@ -216,18 +220,16 @@ impl TransportArena { } let h2g_ring_offset = g2h - .checked_ring_len()? + .ring_len() .checked_next_multiple_of(virtq::Descriptor::ALIGN)?; let g2h_pool_offset = h2g_ring_offset - .checked_add(h2g.checked_ring_len()?)? + .checked_add(h2g.ring_len())? .checked_next_multiple_of(crate::vmem::PAGE_SIZE)?; - let g2h_pool_len = g2h.checked_pool_len()?; - let h2g_pool_offset = g2h_pool_offset.checked_add(g2h_pool_len)?; + let h2g_pool_offset = g2h_pool_offset.checked_add(g2h.pool_len())?; - let h2g_pool_len = h2g.checked_pool_len()?; - let len = h2g_pool_offset.checked_add(h2g_pool_len)?; + let len = h2g_pool_offset.checked_add(h2g.pool_len())?; let addr = |offset: usize| base_addr.checked_add(u64::try_from(offset).ok()?); let _end_addr = addr(len)?; @@ -312,6 +314,8 @@ mod tests { 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 @@ -347,8 +351,7 @@ mod tests { assert_eq!(QueueDims::new(3, 8), None); assert_eq!(QueueDims::new(64, 0), None); assert_eq!(QueueDims::new(usize::MAX, 8), None); - let oversized = QueueDims::new(64, usize::MAX).unwrap(); - assert_eq!(TransportArena::new(base, oversized, h2g), None); + assert_eq!(QueueDims::new(64, usize::MAX), None); assert_eq!( TransportArena::new(u64::MAX - crate::vmem::PAGE_SIZE as u64 + 1, g2h, h2g,), None diff --git a/src/hyperlight_common/src/outb.rs b/src/hyperlight_common/src/outb.rs index 3bfb99848..9e4130c26 100644 --- a/src/hyperlight_common/src/outb.rs +++ b/src/hyperlight_common/src/outb.rs @@ -94,6 +94,7 @@ impl TryFrom for Exception { /// - 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, @@ -105,6 +106,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`) @@ -137,6 +139,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)), } } diff --git a/src/hyperlight_common/src/virtq/buffer.rs b/src/hyperlight_common/src/virtq/buffer.rs index fd5570ba1..1d5345980 100644 --- a/src/hyperlight_common/src/virtq/buffer.rs +++ b/src/hyperlight_common/src/virtq/buffer.rs @@ -128,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())); @@ -390,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/consumer.rs b/src/hyperlight_common/src/virtq/consumer.rs index 9135b7096..21b39da82 100644 --- a/src/hyperlight_common/src/virtq/consumer.rs +++ b/src/hyperlight_common/src/virtq/consumer.rs @@ -654,6 +654,7 @@ impl VirtqConsumer { self.inner.reset()?; self.inflight.clear(); + self.next_token = 0; Ok(()) } } @@ -1291,5 +1292,6 @@ mod tests { 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/mod.rs b/src/hyperlight_common/src/virtq/mod.rs index 9de87d5b8..01ca7dc75 100644 --- a/src/hyperlight_common/src/virtq/mod.rs +++ b/src/hyperlight_common/src/virtq/mod.rs @@ -952,6 +952,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(); @@ -971,6 +972,7 @@ mod tests { // 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"); diff --git a/src/hyperlight_common/src/virtq/msg.rs b/src/hyperlight_common/src/virtq/msg.rs index d0e937bfc..1f19988b0 100644 --- a/src/hyperlight_common/src/virtq/msg.rs +++ b/src/hyperlight_common/src/virtq/msg.rs @@ -14,13 +14,18 @@ See the License for the specific language governing permissions and limitations under the License. */ -//! Wire format header for all virtqueue messages. +//! Wire framing for virtqueue messages. //! //! Every message chain on both the G2H and H2G queues starts with this fixed //! 8-byte header, enabling message type discrimination and request/response //! correlation. Payload lengths come from the size-prefixed FlatBuffer and its //! external-byte declarations. +use crate::flatbuffer_wrappers::{ExternalValueRef, ExternalValueRefs}; + +/// 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)] @@ -99,9 +104,76 @@ impl VirtqMsgHeader { } } +/// Borrowed wire message split into transport-ready chunks. +#[derive(Debug)] +pub struct EncodedMessage<'a> { + header: VirtqMsgHeader, + control: &'a [u8], + externals: ExternalValueRefs<'a>, + wire_len: usize, +} + +impl<'a> EncodedMessage<'a> { + /// Build a message, returning `None` if its wire length overflows. + pub fn new( + kind: MsgKind, + cid: u32, + control: &'a [u8], + externals: ExternalValueRefs<'a>, + ) -> Option { + let wire_len = VirtqMsgHeader::SIZE + .checked_add(control.len())? + .checked_add(externals.total_len()?)?; + + Some(Self { + header: VirtqMsgHeader::new(kind, cid), + control, + externals, + wire_len, + }) + } + + /// Visit wire chunks in transmission order. + pub fn try_for_each_chunk( + &self, + mut visit: impl FnMut(&[u8]) -> Result<(), E>, + ) -> Result<(), E> { + visit(self.header.as_bytes())?; + visit(self.control)?; + + for val in self.externals.as_slice() { + match val { + ExternalValueRef::Bytes(value) => visit(value)?, + ExternalValueRef::Chunks(chunks) => chunks.iter().try_for_each(|c| visit(c))?, + } + } + + Ok(()) + } + + /// Total wire length of all chunks. + pub const fn wire_len(&self) -> usize { + self.wire_len + } +} + +/// 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_only_kind_and_cid() { @@ -127,4 +199,35 @@ mod tests { assert_eq!(VirtqMsgHeader::from_bytes(&bytes), None); assert_eq!(VirtqMsgHeader::from_bytes(&bytes[..7]), None); } + + #[test] + fn encoded_message_visits_wire_chunks_in_order() { + let chunks = [ + bytes::Bytes::from_static(b"ef"), + bytes::Bytes::from_static(b"gh"), + ]; + let mut external_values = ExternalValueRefs::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 mut visited = Vec::new(); + message + .try_for_each_chunk(|chunk| { + visited.push(chunk.to_vec()); + Ok::<_, ()>(()) + }) + .unwrap(); + + assert_eq!(message.wire_len(), VirtqMsgHeader::SIZE + 8); + assert_eq!(visited[1..], [b"ab", b"cd", b"ef", b"gh"]); + } + + #[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/pool.rs b/src/hyperlight_common/src/virtq/pool.rs index 76f6416e1..15df198aa 100644 --- a/src/hyperlight_common/src/virtq/pool.rs +++ b/src/hyperlight_common/src/virtq/pool.rs @@ -31,7 +31,7 @@ mod slot; pub use run::RunPool; #[cfg(all(test, loom))] pub use run::RunPoolSync; -pub use slot::{SlotLayout, SlotPool}; +pub use slot::{AllocationPlan, SlotLayout, SlotPool}; /// Buffer allocation failure. #[derive(Debug, Error, Copy, Clone)] diff --git a/src/hyperlight_common/src/virtq/pool/slot.rs b/src/hyperlight_common/src/virtq/pool/slot.rs index 8bfc9c065..8d97195d5 100644 --- a/src/hyperlight_common/src/virtq/pool/slot.rs +++ b/src/hyperlight_common/src/virtq/pool/slot.rs @@ -70,6 +70,33 @@ impl SlotLayout { } } +/// Slot usage for one prospective scatter/gather allocation. +/// +/// The plan reflects current pool availability without reserving slots. +/// Another allocation can invalidate it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AllocationPlan { + lower_slots: usize, + upper_slots: usize, +} + +impl AllocationPlan { + /// Number of lower-tier slots the allocation would consume. + pub fn lower_slots(self) -> usize { + self.lower_slots + } + + /// Number of upper-tier slots the allocation would consume. + pub fn upper_slots(self) -> usize { + self.upper_slots + } + + /// Total number of slots the allocation would consume. + pub fn num_slots(self) -> usize { + self.lower_slots + self.upper_slots + } +} + /// Single-tier fixed-slot free list. /// /// Tracks a fixed set of equal-sized buffer slots. Allocation pops a free slot @@ -261,6 +288,39 @@ impl Inner { self.upper.alloc(len) } + fn plan_alloc(&self, total_len: usize) -> Result { + if total_len == 0 { + return Err(AllocError::InvalidArg); + } + + let upper_size = self.upper.slot_size; + let mut upper_slots = total_len / upper_size; + + let tail_len = total_len % upper_size; + let mut lower_slots = 0; + + if tail_len != 0 { + if self + .lower + .as_ref() + .is_some_and(|lower| tail_len <= lower.slot_size && lower.num_free() != 0) + { + lower_slots = 1; + } else { + upper_slots = upper_slots.checked_add(1).ok_or(AllocError::Overflow)?; + } + } + + if upper_slots > self.upper.num_free() { + return Err(AllocError::NoSpace); + } + + Ok(AllocationPlan { + lower_slots, + upper_slots, + }) + } + fn dealloc_addr(&mut self, addr: u64) -> Result<(), AllocError> { if let Some(lower) = &mut self.lower && lower.contains(addr) @@ -350,6 +410,19 @@ impl SlotPool { }) } + /// Plan how the next scatter/gather allocation would use each tier. + /// + /// The result follows [`BufferProvider::alloc_sg`] segmentation and + /// lower-tier fallback without changing free-list order. + /// + /// # Errors + /// + /// Returns an error when `total_len` is zero, arithmetic overflows, or the + /// current free slots cannot satisfy the allocation. + pub fn plan_alloc(&self, total_len: usize) -> Result { + self.inner.borrow().plan_alloc(total_len) + } + /// Return every live slot address in deterministic tier and index order. pub fn live_addrs(&self) -> Vec { self.inner.borrow().live_addrs() @@ -372,6 +445,16 @@ impl SlotPool { self.inner.borrow().num_free() } + /// 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) @@ -392,6 +475,20 @@ impl SlotPool { self.inner.borrow().max_alloc_len() } + /// Slot size in bytes for the lower tier, if present. + pub fn lower_slot_size(&self) -> Option { + self.inner + .borrow() + .lower + .as_ref() + .map(|lower| lower.slot_size) + } + + /// Maximum slot size in bytes for the upper tier. + pub fn upper_slot_size(&self) -> usize { + self.inner.borrow().upper.slot_size + } + /// Total number of slots across all tiers. pub fn count(&self) -> usize { self.inner.borrow().count() diff --git a/src/hyperlight_common/src/virtq/pool/tests.rs b/src/hyperlight_common/src/virtq/pool/tests.rs index ed4216c2d..9f1ac368f 100644 --- a/src/hyperlight_common/src/virtq/pool/tests.rs +++ b/src/hyperlight_common/src/virtq/pool/tests.rs @@ -215,6 +215,8 @@ fn test_tiered_slot_pool_combines_contiguous_equal_sized_layouts() { assert_eq!(pool.base_addr(), 0x80000); assert_eq!(pool.slot_size(), 0x100); assert_eq!(pool.count(), 5); + assert_eq!(pool.num_free_lower(), 0); + assert_eq!(pool.num_free_upper(), 5); assert_eq!(pool.slot_addr(4), Some(0x80400)); } @@ -271,6 +273,24 @@ fn test_tiered_slot_pool_does_not_mask_lower_errors() { 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_alloc_sg_uses_both_tiers() { let pool = make_tiered_slot_pool(1, 2); @@ -288,6 +308,51 @@ fn test_tiered_slot_pool_alloc_sg_uses_both_tiers() { assert_eq!(pool.num_free(), 3); } +#[test] +fn test_tiered_slot_pool_plans_alloc_sg_without_allocating() { + let pool = make_tiered_slot_pool(2, 3); + let next_lower = pool.slot_addr(1).unwrap(); + + let plan = pool.plan_alloc(4096 + 128).unwrap(); + assert_eq!(plan.lower_slots(), 1); + assert_eq!(plan.upper_slots(), 1); + assert_eq!(plan.num_slots(), 2); + assert_eq!(pool.num_free_lower(), 2); + assert_eq!(pool.num_free_upper(), 3); + + let lower = pool.alloc(128).unwrap(); + assert_eq!(lower.addr, next_lower); + let lower_allocations = [lower, pool.alloc(128).unwrap()]; + let plan = pool.plan_alloc(128).unwrap(); + assert_eq!(plan.lower_slots(), 0); + assert_eq!(plan.upper_slots(), 1); + + for allocation in lower_allocations { + pool.dealloc(allocation.addr).unwrap(); + } +} + +#[test] +fn test_single_tier_slot_pool_plans_all_segments_as_upper() { + let pool = SlotPool::new(SlotLayout::new(0x80000, 256, 3)).unwrap(); + + let plan = pool.plan_alloc(257).unwrap(); + assert_eq!(plan.lower_slots(), 0); + assert_eq!(plan.upper_slots(), 2); + assert_eq!(plan.num_slots(), 2); +} + +#[test] +fn test_slot_pool_plan_rejects_invalid_or_unavailable_allocations() { + let pool = make_tiered_slot_pool(1, 1); + + assert!(matches!(pool.plan_alloc(0), Err(AllocError::InvalidArg))); + assert!(matches!( + pool.plan_alloc(4096 + 257), + Err(AllocError::NoSpace) + )); +} + #[test] fn test_tiered_slot_pool_dealloc_routes_by_region() { let pool = make_tiered_slot_pool(1, 1); diff --git a/src/hyperlight_common/src/virtq/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index 3f46786c3..369e9c447 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -277,6 +277,12 @@ 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() + } + /// Configure event suppression for used buffer notifications. /// /// This controls when the device (consumer) signals us about completed buffers: diff --git a/src/hyperlight_guest/src/error.rs b/src/hyperlight_guest/src/error.rs index 92cd078d0..681af084c 100644 --- a/src/hyperlight_guest/src/error.rs +++ b/src/hyperlight_guest/src/error.rs @@ -18,6 +18,7 @@ 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}; @@ -90,6 +91,15 @@ impl From for HyperlightGuestError { } } +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`. /// @@ -181,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/host_comm.rs b/src/hyperlight_guest/src/guest_handle/host_comm.rs index c72de8a3f..73c33ddcc 100644 --- a/src/hyperlight_guest/src/guest_handle/host_comm.rs +++ b/src/hyperlight_guest/src/guest_handle/host_comm.rs @@ -18,21 +18,17 @@ 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, + 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; +use crate::transport; impl GuestHandle { /// Get user memory region as bytes. @@ -59,83 +55,6 @@ impl GuestHandle { } } - /// 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. @@ -148,8 +67,9 @@ impl GuestHandle { parameters: Option>, return_type: ReturnType, ) -> Result { - self.call_host_function_without_returning_result(function_name, parameters, return_type)?; - self.get_host_return_value::() + transport::with_context(|context| { + context.call_host_function(function_name, parameters, return_type) + }) } /// Log a message with the specified log level, source, caller, source file, and line number. @@ -162,7 +82,7 @@ impl GuestHandle { source_file: &str, line: u32, ) { - // Closure to send log message to host + // Closure to send log message to host via G2H virtqueue let _send_to_host = || { let guest_log_data = GuestLogData::new( message.to_string(), @@ -177,12 +97,11 @@ impl GuestHandle { .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); - } + transport::with_context(|context| { + context + .emit_log(&bytes) + .expect("Unable to send log data via virtq"); + }); }; #[cfg(all(feature = "trace_guest", target_arch = "x86_64"))] diff --git a/src/hyperlight_guest/src/transport/context.rs b/src/hyperlight_guest/src/transport/context.rs index f47d120ca..443877a2b 100644 --- a/src/hyperlight_guest/src/transport/context.rs +++ b/src/hyperlight_guest/src/transport/context.rs @@ -16,29 +16,53 @@ limitations under the License. //! Guest virtqueue context. +use alloc::vec::Vec; use core::result; +use flatbuffers::FlatBufferBuilder; +use hyperlight_common::flatbuffer_wrappers::ExternalValueRefs; +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 hyperlight_common::outb::OutBAction; +use hyperlight_common::virtq::msg::{EncodedMessage, MsgKind}; use hyperlight_common::virtq::{ - AllocError, G2H_LOWER_SLOT_COUNT, G2H_LOWER_SLOT_SIZE, Layout, Notifier, QueueStats, - SlotLayout, SlotPool, VirtqProducer, + AllocError, BufferProvider, G2H_LOWER_SLOT_COUNT, G2H_LOWER_SLOT_SIZE, Layout, Notifier, + QueueStats, SlotLayout, SlotPool, Token, UsedChain, VirtqError, VirtqProducer, }; -use super::GuestMemOps; +use super::{GuestMemOps, response}; +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); + } + } +} -/// Guest-side notifier for polled transport operation. +/// H2G prefill does not notify before the host consumer is attached. #[derive(Clone, Copy)] -pub struct GuestNotifier; +pub struct H2gNotifier; -impl Notifier for GuestNotifier { +impl Notifier for H2gNotifier { fn notify(&self, _stats: QueueStats) {} } /// Type alias for the guest-side G2H producer. -pub type G2hProducer = VirtqProducer; +pub type G2hProducer = VirtqProducer; /// Type alias for the guest-side H2G producer. -pub type H2gProducer = VirtqProducer; +pub type H2gProducer = VirtqProducer; /// Configuration for one queue passed to [`GuestContext::new`]. pub struct QueueConfig { @@ -55,40 +79,178 @@ pub struct QueueConfig { /// Virtqueue runtime state for guest-host communication. pub struct GuestContext { /// Guest-to-host driver. - _g2h_producer: G2hProducer, + g2h_producer: G2hProducer, + /// G2H pool state used to size writable replies. + g2h_pool: SlotPool, /// Host-to-guest driver. h2g_producer: H2gProducer, /// Size of each prefilled H2G buffer. h2g_slot_size: usize, + /// Correlation ID assigned to the next host-function request. + next_cid: u32, + /// Used by the C API. + last_host_result: Option>, } impl GuestContext { /// Create a new context with G2H and H2G queues. pub fn new(g2h: QueueConfig, h2g: QueueConfig) -> Result { - Self::with_mem(g2h, h2g, GuestMemOps::for_scratch()) - } - - /// Create a new context with memory access provided. - fn with_mem(g2h: QueueConfig, h2g: QueueConfig, mem: GuestMemOps) -> Result { let g2h_pool = g2h_pool(g2h.pool_gva, g2h.pool_pages, g2h.buffer_size) .with_context(|| "failed to create G2H pool")?; - let g2h_producer = VirtqProducer::new(g2h.layout, mem, GuestNotifier, g2h_pool); + let 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, GuestNotifier, h2g_pool); + let h2g_producer = VirtqProducer::new(h2g.layout, mem, H2gNotifier, h2g_pool); let mut ctx = Self { - _g2h_producer: g2h_producer, + g2h_producer, + g2h_pool, h2g_producer, h2g_slot_size: h2g.buffer_size, + next_cid: 1, + last_host_result: None, }; - ctx.prefill_h2g().expect("H2G initial prefill failed"); + ctx.prefill_h2g()?; Ok(ctx) } - /// Pre-fill H2G with writable buffers until its ring or pool is full. + /// Call a host function via the G2H virtqueue. + /// + /// Control data and borrowed external values form one readable request. + /// 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 = ExternalValueRefs::new(); + + let control = fc + .encode_external(&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 message = EncodedMessage::new(MsgKind::Request, cid, control, externals) + .context("G2H message length overflow")?; + + // Reserve response capacity from the ring and pool state that remains + // after this request. + let reply_cap = self.reply_capacity(message.wire_len(), return_type)?; + + // Submit once more after forcing the host to drain on backpressure. + let token = match self.try_send(&message, Some(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(&message, Some(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 = response::decode(segments, cid)?; + let ret = fcr.into_inner()?; + + let Ok(ret) = T::try_from(ret) else { + bail!("G2H: host return value type mismatch"); + }; + + Ok(ret) + } + + /// 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, ExternalValueRefs::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"), + } + } + + /// Pre-fill the H2G queue with writable-only descriptors so the host + /// can write incoming call payloads into them. fn prefill_h2g(&mut self) -> Result<()> { let mut batch = self.h2g_producer.batch(); @@ -99,7 +261,9 @@ impl GuestContext { batch.finish()?; return Ok(()); } - Err(error) => return Err(error.into()), + Err(error) => { + bail!("H2G prefill build: {error}"); + } }; match batch.submit(chain) { @@ -108,9 +272,132 @@ impl GuestContext { batch.finish()?; return Ok(()); } - Err(error) => return Err(error.into()), + Err(error) => { + bail!("H2G prefill submit: {error}"); + } + } + } + } + + /// Size writable reply buffers after accounting for the request. + /// + /// Variable returns use every remaining upper-tier slot allowed by ring + /// descriptor headroom. Fixed returns reserve one lower-sized allocation, + /// with the pool's normal upper-tier fallback. + fn reply_capacity(&mut self, req_len: usize, return_type: ReturnType) -> Result { + if let Some(cap) = self.try_reply_capacity(req_len, return_type)? { + return Ok(cap); + } + + // A deferred batch can own all available slots or descriptors. Force + // one host drain and reclaim before reporting insufficient capacity. + if self.g2h_producer.num_inflight() != 0 { + self.g2h_producer.notify_backpressure(); + + if let Err(error) = self.g2h_producer.reclaim() { + bail!("G2H reclaim: {error}"); + } + + if let Some(capacity) = self.try_reply_capacity(req_len, return_type)? { + return Ok(capacity); + } + } + + bail!("No virtqueue capacity remains for a host function response"); + } + + /// Calculate reply capacity without allocating request buffers. + /// + /// `None` means current ring or pool occupancy leaves no response capacity. + fn try_reply_capacity(&self, req_len: usize, return_type: ReturnType) -> Result> { + let plan = match self.g2h_pool.plan_alloc(req_len) { + Ok(request) => request, + Err(AllocError::NoSpace) => return Ok(None), + Err(error) => bail!("G2H request allocation plan: {error}"), + }; + + let free_descs = self + .g2h_producer + .num_free() + .saturating_sub(plan.num_slots()); + + let (free_slots, slot_size) = match return_type { + ReturnType::String | ReturnType::VecBytes | ReturnType::ByteChunks => ( + self.g2h_pool.num_free_upper() - plan.upper_slots(), + self.g2h_pool.max_alloc_len(), + ), + _ => ( + usize::from(self.g2h_pool.num_free() > plan.num_slots()), + self.g2h_pool + .lower_slot_size() + .unwrap_or_else(|| self.g2h_pool.max_alloc_len()), + ), + }; + + let capacity = free_slots.min(free_descs) * slot_size; + Ok((capacity != 0).then_some(capacity)) + } + + /// 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, 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, 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. + /// + /// The readable region contains the header, control message, and external + /// values. `reply_cap` adds a writable region for a host function reply. + fn try_send( + &mut self, + message: &EncodedMessage<'_>, + reply_cap: Option, + ) -> result::Result { + // Allocate the readable request and optional writable reply together. + let chain = self.g2h_producer.chain().readable(message.wire_len()); + + let mut chain = match reply_cap { + Some(reply_cap) => chain.writable(reply_cap), + None => chain, + } + .build()?; + + message.try_for_each_chunk(|chunk| { + chain.write_all(chunk)?; + Ok::<(), VirtqError>(()) + })?; + + // Transfer the initialized chain to the producer. + self.g2h_producer.submit(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 } } diff --git a/src/hyperlight_guest/src/transport/mod.rs b/src/hyperlight_guest/src/transport/mod.rs index 6290089fd..4e4a318d5 100644 --- a/src/hyperlight_guest/src/transport/mod.rs +++ b/src/hyperlight_guest/src/transport/mod.rs @@ -20,6 +20,7 @@ limitations under the License. pub mod context; pub mod mem; +mod response; use core::cell::RefCell; use core::sync::atomic::{AtomicU8, Ordering}; diff --git a/src/hyperlight_guest/src/transport/response.rs b/src/hyperlight_guest/src/transport/response.rs new file mode 100644 index 000000000..a88f28542 --- /dev/null +++ b/src/hyperlight_guest/src/transport/response.rs @@ -0,0 +1,181 @@ +/* +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 G2H response decoding. + +use alloc::vec::Vec; + +use hyperlight_common::flatbuffer_wrappers::ExternalValueSource; +use hyperlight_common::flatbuffer_wrappers::function_types::{Bytes, FunctionCallResult}; +use hyperlight_common::virtq::Segments; +use hyperlight_common::virtq::msg::{ + MsgKind, SIZE_PREFIX_LEN, VirtqMsgHeader, size_prefix_payload_len, size_prefixed_len, +}; + +use crate::bail; +use crate::error::{GuestErrorContext, Result}; + +/// Decode `header | size-prefixed control | external values`. +/// +/// Contiguous byte values are flattened into `Vec`. Chunked values retain +/// their transport-backed `Bytes` owners and return pool slots when dropped. +pub(super) fn decode(mut segments: Segments, cid: u32) -> Result { + // Validate the transport envelope before interpreting the response body. + let header = segments + .split_to(VirtqMsgHeader::SIZE) + .context("host function response is missing its header")? + .into_bytes(); + + let Some(header) = VirtqMsgHeader::from_bytes(&header) else { + bail!("Host function response has an invalid header"); + }; + + 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"); + } + + // Flatten only the FlatBuffer control data. External byte values remain + // segmented for `SegmentSource`. + let prefix = segments + .split_to(SIZE_PREFIX_LEN) + .context("host function response is missing its size prefix")? + .into_bytes(); + + let payload_len = + size_prefix_payload_len(&prefix).context("host function response has an invalid prefix")?; + + let payload = segments + .split_to(payload_len) + .context("host function response control data is truncated")?; + + let control_len = + size_prefixed_len(payload_len).context("host function response 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); + } + + let mut external_values = SegmentSource::new(segments); + FunctionCallResult::decode_external(&control, &mut external_values) + .with_context(|| "failed to decode host function response") +} + +/// 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!( + "Host function response has {} trailing external bytes", + self.segments.len() + ); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use alloc::vec; + + use flatbuffers::FlatBufferBuilder; + use hyperlight_common::flatbuffer_wrappers::ExternalValueRefs; + use hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue; + + 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 = ExternalValueRefs::new(); + let control = result + .encode_external(&mut builder, &mut external_values) + .unwrap(); + let header = VirtqMsgHeader::new(MsgKind::Response, 7); + let segments = Segments::new([ + Bytes::copy_from_slice(header.as_bytes()), + Bytes::copy_from_slice(control), + external, + ]); + + let decoded = decode(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)); + } +} diff --git a/src/hyperlight_guest_bin/src/host_comm.rs b/src/hyperlight_guest_bin/src/host_comm.rs index 301462313..e7c026cbc 100644 --- a/src/hyperlight_guest_bin/src/host_comm.rs +++ b/src/hyperlight_guest_bin/src/host_comm.rs @@ -25,6 +25,7 @@ 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 +37,9 @@ pub fn call_host_function( where T: TryFrom, { - let handle = unsafe { GUEST_HANDLE }; - handle.call_host_function::(function_name, parameters, return_type) + transport::with_context(|context| { + context.call_host_function(function_name, parameters, return_type) + }) } pub fn call_host(function_name: impl AsRef, args: impl ParameterTuple) -> Result @@ -47,25 +49,6 @@ 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) @@ -76,9 +59,8 @@ pub fn read_n_bytes_from_user_memory(num: u64) -> Result> { /// 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 }; 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, diff --git a/src/hyperlight_guest_bin/src/lib.rs b/src/hyperlight_guest_bin/src/lib.rs index 83afe45c5..80c53f457 100644 --- a/src/hyperlight_guest_bin/src/lib.rs +++ b/src/hyperlight_guest_bin/src/lib.rs @@ -263,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"); @@ -288,9 +291,6 @@ pub(crate) extern "C" fn generic_init( registration(); } - // Prepare transport before guest code starts. - transport::initialize(); - unsafe { hyperlight_main(); } diff --git a/src/hyperlight_guest_capi/src/dispatch.rs b/src/hyperlight_guest_capi/src/dispatch.rs index e0a8bc34c..dc8d307b3 100644 --- a/src/hyperlight_guest_capi/src/dispatch.rs +++ b/src/hyperlight_guest_capi/src/dispatch.rs @@ -20,12 +20,15 @@ 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}; static mut REGISTERED_C_GUEST_FUNCTIONS: GuestFunctionRegister = @@ -98,15 +101,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_context(|context| context.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_context(|context| context.take_host_return::()) } diff --git a/src/hyperlight_guest_capi/src/flatbuffer.rs b/src/hyperlight_guest_capi/src/flatbuffer.rs index 77fb7be67..c648a421e 100644 --- a/src/hyperlight_guest_capi/src/flatbuffer.rs +++ b/src/hyperlight_guest_capi/src/flatbuffer.rs @@ -24,8 +24,8 @@ use hyperlight_common::flatbuffer_wrappers::function_types::Bytes; use hyperlight_common::flatbuffer_wrappers::util::{ byte_chunks_from_vec, byte_chunks_to_vec, get_flatbuffer_result, }; -use hyperlight_guest_bin::host_comm::get_host_return_value; +use crate::dispatch::take_last_host_return; use crate::types::FfiVec; // The reason for the capitalized type in the function names below @@ -117,44 +117,43 @@ pub extern "C" fn hl_flatbuffer_result_from_Bool(value: bool) -> Box { #[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") + take_last_host_return() } #[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") + take_last_host_return() } // 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") + take_last_host_return() } #[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") + take_last_host_return() } #[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") + take_last_host_return() } #[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") + take_last_host_return() } #[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") + 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 = - get_host_return_value().expect("Unable to get host return value as string"); + let string_value: String = take_last_host_return(); let c_string = CString::new(string_value).expect("Failed to create CString"); c_string.into_raw() @@ -162,16 +161,14 @@ pub extern "C" fn hl_get_host_return_value_as_String() -> *const c_char { #[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"); + 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() -> Box { - let chunks: Vec = - get_host_return_value().expect("Unable to get host return value as byte chunks"); + let chunks: Vec = take_last_host_return(); Box::new(unsafe { FfiVec::from_vec(byte_chunks_to_vec(&chunks)) }) } diff --git a/src/hyperlight_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs index cb08832d1..3219fbbce 100644 --- a/src/hyperlight_host/src/mem/mgr.rs +++ b/src/hyperlight_host/src/mem/mgr.rs @@ -379,6 +379,11 @@ impl SandboxMemoryManager { h2g_consumer: None, }; host_mgr.update_scratch_bookkeeping()?; + + if matches!(host_mgr.next_action, NextAction::Initialise(_)) { + host_mgr.create_virtq_consumers()?; + } + Ok((host_mgr, guest_mgr)) } } @@ -422,22 +427,16 @@ impl SandboxMemoryManager { ) } - /// Attach host consumers to a guest-produced initial transport image. + /// Create host consumers before the guest initializes the transport. /// - /// Before guest initialization, the host publishes queue dimensions and the - /// transport arena GPA. The guest derives and initializes every fixed region - /// without consuming dynamic scratch. - /// - /// This method runs after the initialization VM exit. It checks the - /// published arena against the host layout, derives bounded GVA views, - /// and validates each directional ring before exposing either consumer. - /// Fresh sandboxes and pre-initialization restores use this path. - pub(crate) fn attach_virtq(&mut self) -> Result<()> { + /// 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 are already attached")); + return Err(new_error!("virtqueue consumers already exist")); } - let (g2h, h2g) = virtq::attach(&self.layout, &self.scratch_mem)?; + let (g2h, h2g) = virtq::create_consumers(&self.layout, &self.scratch_mem)?; self.g2h_consumer = Some(g2h); self.h2g_consumer = Some(h2g); Ok(()) @@ -880,6 +879,7 @@ mod tests { use hyperlight_testing::sandbox_sizes::{LARGE_HEAP_SIZE, MEDIUM_HEAP_SIZE, SMALL_HEAP_SIZE}; use hyperlight_testing::simple_guest_as_pathbuf; + use super::SandboxMemoryManager; use crate::GuestBinary; use crate::sandbox::SandboxConfiguration; use crate::sandbox::snapshot::Snapshot; @@ -927,4 +927,18 @@ mod tests { verify_page_tables(name, config); } } + + #[test] + fn build_creates_virtq_consumers_before_initialization() { + let path = simple_guest_as_string().expect("failed to get simple guest path"); + let snapshot = + Snapshot::from_env(GuestBinary::FilePath(path), 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/virtq.rs b/src/hyperlight_host/src/mem/virtq.rs index 534934db0..3f913b2b6 100644 --- a/src/hyperlight_host/src/mem/virtq.rs +++ b/src/hyperlight_host/src/mem/virtq.rs @@ -14,18 +14,35 @@ See the License for the specific language governing permissions and limitations under the License. */ -//! Host virtqueue attachment. +//! Host virtqueue construction, I/O, and canonical snapshot validation. //! -//! The host publishes one transport arena address in scratch-top metadata. Guest -//! initialization builds both queues in those fixed regions. This module -//! validates the complete initial image before returning either consumer. +//! 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. +//! +//! G2H codec helpers copy untrusted request data into host-owned values before +//! dispatch. Shared wire framing lives in `hyperlight_common::virtq::msg`. +//! +//! Snapshot capture and restore validate canonical ring images against the +//! configured arena before exposing consumers. use core::ops::Range; +use anyhow::{Context, bail}; +use flatbuffers::FlatBufferBuilder; +use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCall; +use hyperlight_common::flatbuffer_wrappers::function_types::{Bytes, FunctionCallResult}; +use hyperlight_common::flatbuffer_wrappers::guest_error::{ErrorCode, GuestError}; +use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData; +use hyperlight_common::flatbuffer_wrappers::{ExternalValueRefs, ExternalValueSource}; use hyperlight_common::layout::{QueueDims, TransportArena}; use hyperlight_common::virtq::canonical::validate_canon_image; +use hyperlight_common::virtq::msg::{ + EncodedMessage, MsgKind, SIZE_PREFIX_LEN, size_prefix_payload_len, size_prefixed_len, +}; use hyperlight_common::virtq::{ - Layout as VirtqLayout, MemOps, Notifier, QueueStats, VirtqConsumer, + Layout as VirtqLayout, MemOps, Notifier, QueueStats, RecvChain, VirtqConsumer, VirtqError, + WritableChain, }; use super::layout::{BaseGpaRegion, SandboxMemoryLayout}; @@ -38,7 +55,7 @@ pub(crate) type G2hConsumer = VirtqConsumer; /// Host-side H2G virtqueue consumer. pub(crate) type H2gConsumer = VirtqConsumer; -/// No-op notifier for polled host transport. +/// No-op notifier because the host completes work during the current VM exit. #[derive(Clone, Copy)] pub(crate) struct HostNotifier; @@ -46,11 +63,24 @@ impl Notifier for HostNotifier { fn notify(&self, _stats: QueueStats) {} } -/// Build both host consumers from a guest-produced initial transport image. +/// Create both host consumers before the first guest entry. /// -/// The consumers are returned only after the host-assigned arena and both -/// directional ring images have passed validation. -pub(crate) fn attach( +/// 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, "G2H")?; + let h2g_layout = validator.config.h2g.layout(®ions.h2g_ring, "H2G")?; + + 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)> { @@ -59,12 +89,29 @@ pub(crate) fn attach( let regions = validator.validate_published_arena(arena_gpa)?; let g2h_ring_mem = HostMemOps::new(scratch_mem, regions.g2h_ring.clone())?; - let g2h_pool_mem = HostMemOps::new(scratch_mem, regions.g2h_pool)?; - let g2h_layout = validator.validate_g2h(&g2h_ring_mem, regions.g2h_ring)?; + let g2h_layout = validator.validate_g2h(&g2h_ring_mem, regions.g2h_ring.clone())?; let h2g_ring_mem = HostMemOps::new(scratch_mem, regions.h2g_ring.clone())?; - let h2g_pool_mem = HostMemOps::new(scratch_mem, regions.h2g_pool.clone())?; - let h2g_layout = validator.validate_h2g(&h2g_ring_mem, regions.h2g_ring, regions.h2g_pool)?; + let h2g_layout = validator.validate_h2g( + &h2g_ring_mem, + regions.h2g_ring.clone(), + regions.h2g_pool.clone(), + )?; + + 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), @@ -72,7 +119,9 @@ pub(crate) fn attach( )) } -/// Capture the canonical transport state omitted from the main memory snapshot. +/// 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, @@ -97,7 +146,10 @@ pub(crate) fn snapshot( }) } -/// Restore one captured canonical transport image and return fresh consumers. +/// 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, @@ -108,14 +160,68 @@ pub(crate) fn restore( write_published_arena_gpa(scratch_mem, layout.get_transport_arena().base_addr())?; write_ring(scratch_mem, regions.g2h_ring, &snapshot.g2h_ring)?; write_ring(scratch_mem, regions.h2g_ring, &snapshot.h2g_ring)?; - attach(layout, scratch_mem) + attach_canonical(layout, scratch_mem) +} + +/// 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 external_values = ChainExternalValues::new(chain); + FunctionCall::decode_external(&control, &mut external_values) +} + +/// Encode a host function result into a G2H reply. +/// +/// A result that exceeds the writable capacity is replaced with a bounded +/// transport error. An error is returned if that fallback also cannot fit. +pub(crate) fn write_response_from_host_function_call( + chain: &mut WritableChain, + cid: u32, + result: &FunctionCallResult, +) -> anyhow::Result<()> { + if try_write_response_from_host_function_call(chain, cid, result)? { + return Ok(()); + } + + let error = FunctionCallResult::new(Err(GuestError::new( + ErrorCode::HostFunctionError, + "Host response exceeds virtqueue capacity".into(), + ))); + + if !try_write_response_from_host_function_call(chain, cid, &error)? { + bail!( + "Writable response capacity {} cannot hold a transport error", + chain.capacity() + ); + } + Ok(()) +} + +/// 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)?; + if chain.remaining() != 0 { + bail!("G2H log has {} trailing external bytes", chain.remaining()); + } + GuestLogData::try_from(control.as_slice()) } /// 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, } @@ -123,33 +229,23 @@ struct GvaRegions { struct QueueConfig { /// Address-independent queue dimensions. dims: QueueDims, - /// Size of the ring image in bytes including event suppressions. - ring_len: usize, - /// Size of the buffer pool in bytes. - pool_len: usize, /// Size of each buffer in the pool in bytes. buffer_size: usize, } impl QueueConfig { fn new(dims: QueueDims, buffer_size: usize) -> Result { - let ring_len = dims - .checked_ring_len() - .ok_or_else(|| new_error!("ring size overflow"))?; - let pool_len = dims - .checked_pool_len() - .ok_or_else(|| new_error!("pool size overflow"))?; - if buffer_size == 0 { return Err(new_error!("buffer size is zero")); } - Ok(Self { - dims, - ring_len, - pool_len, - buffer_size, - }) + Ok(Self { dims, buffer_size }) + } + + fn layout(&self, ring: &Range, direction: &str) -> Result { + // SAFETY: `ring` is derived from the validated fixed transport arena. + unsafe { VirtqLayout::from_base(ring.start, self.dims.depth()) } + .map_err(|error| new_error!("invalid {direction} ring layout: {error}")) } } @@ -170,11 +266,11 @@ impl Config { /// Compute the host transport configuration from the memory layout. fn from_layout(layout: &SandboxMemoryLayout) -> Result { let g2h = QueueConfig::new(layout.get_g2h_queue_dims(), layout.get_g2h_buffer_size())?; - let h2g = QueueConfig::new(layout.get_h2g_queue_dims(), layout.get_h2g_buffer_size())?; let h2g_prefill_chains = - usize::from(h2g.dims.depth().get()).min(h2g.pool_len / h2g.buffer_size); + usize::from(h2g.dims.depth().get()).min(h2g.dims.pool_len() / h2g.buffer_size); + let arena = layout.get_transport_arena(); Ok(Self { @@ -189,8 +285,11 @@ impl Config { /// 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, } @@ -215,12 +314,9 @@ impl<'a> Validator<'a> { }) } - /// Validate the initial G2H queue and return its layout. + /// Validate a canonical G2H image and return its layout. fn validate_g2h(&self, mem: &M, ring: Range) -> Result { - // SAFETY: `ring` spans the configured image and `mem` keeps that image - // valid for the duration of validation. - let layout = unsafe { VirtqLayout::from_base(ring.start, self.config.g2h.dims.depth()) } - .map_err(|error| new_error!("invalid G2H ring layout: {error}"))?; + let layout = self.config.g2h.layout(&ring, "G2H")?; validate_canon_image(mem, layout, 0, |_, _| false) .map_err(|error| new_error!("invalid canonical G2H image: {error}"))?; @@ -228,7 +324,7 @@ impl<'a> Validator<'a> { Ok(layout) } - /// Validate the initial H2G queue and return its layout. + /// Validate a canonical H2G image and return its layout. /// /// Every available chain contains one configured size writable descriptor. /// Descriptors must name distinct, slot-aligned ranges inside the H2G pool. @@ -238,10 +334,7 @@ impl<'a> Validator<'a> { ring: Range, pool: Range, ) -> Result { - // SAFETY: `ring` spans the configured image and `mem` keeps that image - // valid for the duration of validation. - let layout = unsafe { VirtqLayout::from_base(ring.start, self.config.h2g.dims.depth()) } - .map_err(|error| new_error!("invalid H2G ring layout: {error}"))?; + let layout = self.config.h2g.layout(&ring, "H2G")?; let bufsz = self.config.h2g.buffer_size; let prefill = self.config.h2g_prefill_chains; @@ -317,8 +410,8 @@ impl<'a> Validator<'a> { } let regions = self.resolve_gva_regions()?; - validate_ring_len("G2H", &snapshot.g2h_ring, self.config.g2h.ring_len)?; - validate_ring_len("H2G", &snapshot.h2g_ring, self.config.h2g.ring_len)?; + 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())?; @@ -360,10 +453,10 @@ impl<'a> Validator<'a> { self.config.arena.h2g_ring_addr(), self.config.arena.g2h_pool_addr(), self.config.arena.h2g_pool_addr(), - self.config.g2h.ring_len, - self.config.h2g.ring_len, - self.config.g2h.pool_len, - self.config.h2g.pool_len, + 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 { @@ -375,17 +468,65 @@ impl<'a> Validator<'a> { } } +/// 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> { + validate_external_length("VecBytes", length, self.request.remaining())?; + let mut value = zeroed_vec(length, "external VecBytes")?; + + self.request.read_exact(&mut value)?; + Ok(value) + } + + fn take_chunks(&mut self, length: usize) -> anyhow::Result> { + if length == 0 { + return Ok(Vec::new()); + } + + validate_external_length("ByteChunks", length, self.request.remaining())?; + let mut value = zeroed_vec(length, "external ByteChunks")?; + 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(()) + } +} + /// 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 @@ -400,12 +541,14 @@ fn read_ring(scratch_mem: &HostSharedMemory, ring: Range) -> Result 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!( @@ -416,6 +559,7 @@ fn validate_ring_len(direction: &str, bytes: &[u8], expected: usize) -> Result<( 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)?) @@ -424,6 +568,75 @@ fn checked_region(start: u64, len: usize, tag: &str) -> Result> { Ok(start..end) } +/// Write a response only when the complete wire message fits. +/// +/// `false` means no bytes were written, allowing the caller to try a bounded +/// transport error. Encoding and chain write failures are returned as errors. +fn try_write_response_from_host_function_call( + reply: &mut WritableChain, + cid: u32, + result: &FunctionCallResult, +) -> anyhow::Result { + let mut builder = FlatBufferBuilder::new(); + let mut external_values = ExternalValueRefs::new(); + + let control = result.encode_external(&mut builder, &mut external_values)?; + let message = EncodedMessage::new(MsgKind::Response, cid, control, external_values) + .context("Host function response length overflow")?; + + if message.wire_len() > reply.capacity() { + return Ok(false); + } + + message.try_for_each_chunk(|chunk| { + reply.write_all(chunk)?; + Ok::<(), VirtqError>(()) + })?; + + Ok(true) +} + +/// 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).expect("size prefix length is fixed"); + 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, "G2H control data")?; + + 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. +fn zeroed_vec(length: usize, what: &str) -> anyhow::Result> { + let mut value = Vec::new(); + value + .try_reserve_exact(length) + .with_context(|| format!("Failed to allocate {length} bytes for {what}"))?; + + value.resize(length, 0); + Ok(value) +} + +/// Validate a declared external length before allocating its storage. +fn validate_external_length(kind: &str, length: usize, remaining: usize) -> anyhow::Result<()> { + if length > remaining { + bail!("External {kind} requires {length} bytes, only {remaining} remain"); + } + Ok(()) +} + #[cfg(test)] mod tests { use core::num::NonZeroU16; @@ -444,6 +657,18 @@ mod tests { const H2G_POOL_PAGES: usize = 2; const H2G_BUFFER_SIZE: usize = 3000; + #[test] + fn external_length_is_bounded_before_allocation() { + assert!(validate_external_length("VecBytes", usize::MAX, 16).is_err()); + assert!(validate_external_length("ByteChunks", 17, 16).is_err()); + assert!(validate_external_length("VecBytes", 16, 16).is_ok()); + } + + #[test] + fn oversized_allocation_fails_without_panicking() { + assert!(zeroed_vec(usize::MAX, "test buffer").is_err()); + } + fn memory_layout() -> SandboxMemoryLayout { let mut config = SandboxConfiguration::default(); config.set_scratch_size(SCRATCH_SIZE); @@ -585,19 +810,19 @@ mod tests { assert_eq!( regions.g2h_ring.end - regions.g2h_ring.start, - config.g2h.ring_len as u64 + config.g2h.dims.ring_len() as u64 ); assert_eq!( regions.h2g_ring.end - regions.h2g_ring.start, - config.h2g.ring_len as u64 + config.h2g.dims.ring_len() as u64 ); assert_eq!( regions.g2h_pool.end - regions.g2h_pool.start, - config.g2h.pool_len as u64 + config.g2h.dims.pool_len() as u64 ); assert_eq!( regions.h2g_pool.end - regions.h2g_pool.start, - config.h2g.pool_len as u64 + config.h2g.dims.pool_len() as u64 ); } @@ -615,15 +840,10 @@ mod tests { } #[test] - fn rejects_untranslatable_or_overflowing_gva_regions() { + fn rejects_untranslatable_gva_regions() { let config = attach_config(); - let arena_gpa = config.arena.base_addr(); let invalid = hyperlight_common::layout::scratch_base_gpa(SCRATCH_SIZE) - 1; assert!(validate_published(invalid, config).is_err()); - - let mut config = config; - config.g2h.ring_len = usize::MAX; - assert!(validate_published(arena_gpa, config).is_err()); } #[test] diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index 939313c98..f7b765740 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -242,10 +242,7 @@ impl MultiUseSandbox { let mgr = crate::mem::mgr::SandboxMemoryManager::from_snapshot(&snapshot)?; let (mut hshm, gshm) = mgr.build()?; - let attach_virtq = matches!( - snapshot.next_action(), - super::snapshot::NextAction::Initialise(_) - ); + let restore_virtq = matches!(snapshot.next_action(), super::snapshot::NextAction::Call(_)); let page_size = u32::try_from(page_size::get())? as usize; @@ -326,9 +323,7 @@ impl MultiUseSandbox { #[cfg(gdb)] let dbg_mem_wrapper = Arc::new(Mutex::new(hshm.clone())); - if attach_virtq { - hshm.attach_virtq()?; - } else { + if restore_virtq { hshm.restore_virtq(snapshot.virtq())?; } diff --git a/src/hyperlight_host/src/sandbox/outb.rs b/src/hyperlight_host/src/sandbox/outb.rs index 4b00d52b4..3cd703ac1 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_call::FunctionCallType; use hyperlight_common::flatbuffer_wrappers::function_types::{FunctionCallResult, ParameterValue}; 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::virtq::ReplyChain; +use hyperlight_common::virtq::msg::{MsgKind, VirtqMsgHeader}; 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; @@ -65,7 +69,11 @@ pub(super) fn outb_log( let log_data: GuestLogData = mgr .read_guest_log_data() .map_err(|e| HandleOutbError::ReadLogData(e.to_string()))?; + emit_guest_log(&log_data); + Ok(()) +} +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 +142,6 @@ pub(super) fn outb_log( ); } } - - Ok(()) } const ABORT_TERMINATOR: u8 = 0xFF; @@ -225,6 +231,7 @@ pub(crate) fn handle_outb( 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,6 +252,124 @@ 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 a request. + let (mut request, mut 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 mut header = [0u8; VirtqMsgHeader::SIZE]; + request.read_exact(&mut header).map_err(|error| { + HandleOutbError::ReadHostFunctionCall(format!("G2H header read failed: {error}")) + })?; + + let Some(header) = VirtqMsgHeader::from_bytes(&header) else { + return Err(HandleOutbError::ReadHostFunctionCall( + "Invalid G2H header".into(), + )); + }; + + match header.msg_kind() { + 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(MsgKind::Request) => break (request, reply, header), + 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 ReplyChain::Writable(resp) = &mut reply else { + return 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())); + + virtq::write_response_from_host_function_call( + resp, + header.cid, + &FunctionCallResult::new(result), + ) + .map_err(|err| HandleOutbError::WriteHostFunctionResponse(err.to_string()))?; + + consumer + .complete(request, reply) + .map_err(|err| HandleOutbError::WriteHostFunctionResponse(err.to_string()))?; + + Ok(()) +} + #[cfg(test)] mod tests { use hyperlight_common::flatbuffer_wrappers::guest_log_level::LogLevel; diff --git a/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs b/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs index 807402efe..c06e1fbfb 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs @@ -76,6 +76,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_evolve.rs b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs index 4c19b3830..ddf407cc3 100644 --- a/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs +++ b/src/hyperlight_host/src/sandbox/uninitialized_evolve.rs @@ -38,10 +38,6 @@ use crate::{MultiUseSandbox, Result, UninitializedSandbox}; #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")] pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result { let (mut hshm, gshm) = u_sbox.mgr.build()?; - let attach_virtq = matches!( - hshm.next_action, - crate::sandbox::snapshot::NextAction::Initialise(_) - ); // Get the host page size. Narrowed to u32 because the guest ABI // passes it via a 32-bit register (rdx), but widened back to usize @@ -113,10 +109,6 @@ pub(super) fn evolve_impl_multi_use(u_sbox: UninitializedSandbox) -> Result(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..4c02dc05b 100644 --- a/src/hyperlight_host/tests/integration_test.rs +++ b/src/hyperlight_host/tests/integration_test.rs @@ -30,6 +30,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, @@ -833,6 +834,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_depth(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] diff --git a/src/hyperlight_host/tests/sandbox_host_tests.rs b/src/hyperlight_host/tests/sandbox_host_tests.rs index b1a1a9918..d863b88c8 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] @@ -325,6 +326,99 @@ 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 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_g2h_ring() { + let mut cfg = SandboxConfiguration::default(); + cfg.set_g2h_queue_depth(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/tests/rust_guests/simpleguest/src/main.rs b/src/tests/rust_guests/simpleguest/src/main.rs index a40f6ca0b..c68118787 100644 --- a/src/tests/rust_guests/simpleguest/src/main.rs +++ b/src/tests/rust_guests/simpleguest/src/main.rs @@ -38,7 +38,7 @@ 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; @@ -52,8 +52,7 @@ 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}; @@ -458,6 +457,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 +511,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 +652,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 @@ -1520,18 +1550,9 @@ 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(); + let host_return = + call_host_function::(&host_func_name, Some(params), func.expected_return_type); + match host_return { Ok(return_value) => match return_value { ReturnValue::Int(i) => Ok(get_flatbuffer_result(i)), From 76d078f8c03e7cf5c7e4f3f90592575571640efb Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Mon, 3 Aug 2026 19:56:42 +0200 Subject: [PATCH 15/34] feat(virtq): remove legacy guest-to-host stack transport Signed-off-by: Tomasz Andrzejak --- CHANGELOG.md | 2 + Justfile | 4 +- .../flatbuffer_wrappers/guest_log_level.rs | 2 +- src/hyperlight_common/src/outb.rs | 17 +- src/hyperlight_host/src/mem/mgr.rs | 40 +-- .../src/sandbox/initialized_multi_use.rs | 2 - src/hyperlight_host/src/sandbox/outb.rs | 247 ++++-------------- .../src/sandbox/snapshot/tripwires.rs | 2 - 8 files changed, 64 insertions(+), 252 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d04d7dc0f..b532e9d4b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). * **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()`. * 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 host function calls to use the guest-to-host + virtqueue protocol. ### Removed 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/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/outb.rs b/src/hyperlight_common/src/outb.rs index 9e4130c26..d9a2386a1 100644 --- a/src/hyperlight_common/src/outb.rs +++ b/src/hyperlight_common/src/outb.rs @@ -87,8 +87,6 @@ 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 @@ -96,8 +94,6 @@ impl TryFrom for Exception { /// - 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")] @@ -129,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")] @@ -155,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_host/src/mem/mgr.rs b/src/hyperlight_host/src/mem/mgr.rs index 3219fbbce..b173dd373 100644 --- a/src/hyperlight_host/src/mem/mgr.rs +++ b/src/hyperlight_host/src/mem/mgr.rs @@ -14,12 +14,8 @@ See the License for the specific language governing permissions and 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::validate_guest_function_call_buffer; 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::vmem::{self, PAGE_TABLE_SIZE}; #[cfg(crashdump)] @@ -458,31 +454,6 @@ impl SandboxMemoryManager { Ok(()) } - /// Reads a host function call from memory - #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")] - pub(crate) fn get_host_function_call(&mut self) -> Result { - 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( - &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, - ) - } - /// 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<()> { @@ -511,15 +482,6 @@ impl SandboxMemoryManager { ) } - /// Read guest log data from the `SharedMemory` contained within `self` - #[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 clear_io_buffers(&mut self) { // Clear the output data buffer loop { diff --git a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs index f7b765740..c7bc0aa25 100644 --- a/src/hyperlight_host/src/sandbox/initialized_multi_use.rs +++ b/src/hyperlight_host/src/sandbox/initialized_multi_use.rs @@ -966,8 +966,6 @@ impl MultiUseSandbox { // 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(); diff --git a/src/hyperlight_host/src/sandbox/outb.rs b/src/hyperlight_host/src/sandbox/outb.rs index 3cd703ac1..767b07846 100644 --- a/src/hyperlight_host/src/sandbox/outb.rs +++ b/src/hyperlight_host/src/sandbox/outb.rs @@ -17,7 +17,7 @@ limitations under the License. use std::sync::{Arc, Mutex}; use hyperlight_common::flatbuffer_wrappers::function_call::FunctionCallType; -use hyperlight_common::flatbuffer_wrappers::function_types::{FunctionCallResult, ParameterValue}; +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; @@ -47,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}")] @@ -62,17 +60,6 @@ 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()))?; - emit_guest_log(&log_data); - Ok(()) -} - fn emit_guest_log(log_data: &GuestLogData) { // Emit guest log data as a tracing event with structured fields. // @@ -210,27 +197,6 @@ 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 => { @@ -374,14 +340,9 @@ fn outb_virtq_call( 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 { @@ -396,140 +357,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()); + 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(); - } - { - // now, test logging - LOGGER.set_max_level(log::LevelFilter::Trace); - let mut mgr = new_mgr(); - LOGGER.clear_log_calls(); - - // 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, - }; + 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, + }; - 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, @@ -540,23 +431,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", @@ -567,38 +446,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/tripwires.rs b/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs index c06e1fbfb..83a3f4463 100644 --- a/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs +++ b/src/hyperlight_host/src/sandbox/snapshot/tripwires.rs @@ -66,8 +66,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")] From 69a63edf0884af08575617eede3391666c24f13e Mon Sep 17 00:00:00 2001 From: Tomasz Andrzejak Date: Mon, 10 Aug 2026 16:59:10 +0200 Subject: [PATCH 16/34] feat(virtq): use virtqueues for guest comms Route guest calls, host calls, results, and logs through bidirectional packed virtqueues. Stream dense messages across exact chain batches and reserve H2G capacity for retained external values. Keep external byte returns typed through guest dispatch and support owner-backed byte chunks in Rust and C guests. Store canonical rings in a versioned OCI transport layer and validate them during snapshot load and restore. Signed-off-by: Tomasz Andrzejak --- CHANGELOG.md | 8 +- docs/snapshot-oci-format.md | 18 +- docs/snapshot-versioning.md | 8 +- .../src/flatbuffer_wrappers/codec.rs | 73 -- .../src/flatbuffer_wrappers/function_call.rs | 305 +----- .../src/flatbuffer_wrappers/function_types.rs | 535 ++-------- .../src/flatbuffer_wrappers/mod.rs | 2 +- .../src/flatbuffer_wrappers/util.rs | 265 +---- src/hyperlight_common/src/lib.rs | 4 + src/hyperlight_common/src/transport.rs | 404 ++++++++ src/hyperlight_common/src/virtq/consumer.rs | 264 ++++- src/hyperlight_common/src/virtq/mod.rs | 1 - src/hyperlight_common/src/virtq/msg.rs | 233 ----- src/hyperlight_common/src/virtq/producer.rs | 29 +- src/hyperlight_common/src/virtq/ring.rs | 69 ++ src/hyperlight_component_util/src/guest.rs | 4 +- .../src/guest_handle/host_comm.rs | 9 +- .../src/transport/{response.rs => codec.rs} | 122 ++- src/hyperlight_guest/src/transport/context.rs | 186 +++- src/hyperlight_guest/src/transport/mod.rs | 8 +- .../src/guest_function/call.rs | 48 +- .../src/guest_function/definition.rs | 29 +- src/hyperlight_guest_bin/src/host_comm.rs | 12 +- src/hyperlight_guest_bin/src/lib.rs | 1 + src/hyperlight_guest_capi/README.md | 9 +- src/hyperlight_guest_capi/cbindgen.toml | 3 +- src/hyperlight_guest_capi/include/macro.h | 50 +- src/hyperlight_guest_capi/src/dispatch.rs | 51 +- src/hyperlight_guest_capi/src/error.rs | 18 +- src/hyperlight_guest_capi/src/lib.rs | 2 +- .../src/{flatbuffer.rs => return_value.rs} | 102 +- src/hyperlight_guest_capi/src/types.rs | 3 + .../src/types/return_value.rs | 183 ++++ src/hyperlight_guest_macro/src/lib.rs | 10 +- src/hyperlight_host/benches/benchmarks.rs | 149 ++- src/hyperlight_host/src/error.rs | 5 + src/hyperlight_host/src/mem/layout.rs | 10 +- src/hyperlight_host/src/mem/mgr.rs | 412 +++++++- src/hyperlight_host/src/mem/mod.rs | 2 - src/hyperlight_host/src/mem/virtq.rs | 968 ------------------ src/hyperlight_host/src/mem/virtq/codec.rs | 197 ++++ .../src/mem/{virtq_mem.rs => virtq/mem.rs} | 2 +- src/hyperlight_host/src/mem/virtq/mod.rs | 492 +++++++++ src/hyperlight_host/src/mem/virtq/tests.rs | 313 ++++++ .../src/sandbox/initialized_multi_use.rs | 22 +- src/hyperlight_host/src/sandbox/outb.rs | 62 +- .../src/sandbox/snapshot/file/media_types.rs | 3 + .../src/sandbox/snapshot/file/mod.rs | 206 +++- .../src/sandbox/snapshot/file_tests.rs | 65 +- .../src/sandbox/snapshot/tripwires.rs | 5 +- src/hyperlight_host/tests/integration_test.rs | 36 - .../tests/sandbox_host_tests.rs | 69 +- src/tests/c_guests/c_simpleguest/main.c | 18 +- src/tests/rust_guests/simpleguest/src/main.rs | 98 +- src/tests/rust_guests/witguest/src/main.rs | 3 +- 55 files changed, 3429 insertions(+), 2776 deletions(-) create mode 100644 src/hyperlight_common/src/transport.rs delete mode 100644 src/hyperlight_common/src/virtq/msg.rs rename src/hyperlight_guest/src/transport/{response.rs => codec.rs} (53%) rename src/hyperlight_guest_capi/src/{flatbuffer.rs => return_value.rs} (55%) create mode 100644 src/hyperlight_guest_capi/src/types/return_value.rs delete mode 100644 src/hyperlight_host/src/mem/virtq.rs create mode 100644 src/hyperlight_host/src/mem/virtq/codec.rs rename src/hyperlight_host/src/mem/{virtq_mem.rs => virtq/mem.rs} (99%) create mode 100644 src/hyperlight_host/src/mem/virtq/mod.rs create mode 100644 src/hyperlight_host/src/mem/virtq/tests.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index b532e9d4b..74da65fa9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,12 +17,16 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). * **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()`. * 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 host function calls to use the guest-to-host - virtqueue protocol. +* 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. ### Removed ### 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/docs/snapshot-oci-format.md b/docs/snapshot-oci-format.md index e77b892b9..ed5ec1ebb 100644 --- a/docs/snapshot-oci-format.md +++ b/docs/snapshot-oci-format.md @@ -24,21 +24,25 @@ 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. + 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 and transport layout, - registered host functions, snapshot generation counter. Loaded + 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. Blob filenames are the sha256 of the blob bytes, so identical blobs across tags are stored once. @@ -55,8 +59,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 +102,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 5c12d2d4c..9b72dbc5e 100644 --- a/docs/snapshot-versioning.md +++ b/docs/snapshot-versioning.md @@ -7,7 +7,7 @@ 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 @@ -23,6 +23,10 @@ A snapshot carries three independently evolvable version markers: `MT_SNAPSHOT_CURRENT`. This is the on-wire format of the snapshot blob: framing, section ordering, alignment, dirty/zero-page elision, anything about how the bytes are packed inside the OCI layer. +* **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: @@ -30,7 +34,7 @@ A snapshot carries three independently evolvable version markers: 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`. diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/codec.rs b/src/hyperlight_common/src/flatbuffer_wrappers/codec.rs index b19f535fe..41052b9ac 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/codec.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/codec.rs @@ -19,79 +19,6 @@ use alloc::vec::Vec; use anyhow::Result; use bytes::Bytes; -/// One borrowed external byte value emitted alongside a FlatBuffer. -#[derive(Clone, Copy, Debug)] -pub enum ExternalValueRef<'a> { - /// A logically contiguous value. - Bytes(&'a [u8]), - /// A logically chunked value. - Chunks(&'a [Bytes]), -} - -impl ExternalValueRef<'_> { - /// Total byte length of this value. - pub fn len(&self) -> Option { - match self { - Self::Bytes(value) => Some(value.len()), - Self::Chunks(chunks) => chunks - .iter() - .try_fold(0usize, |len, chunk| len.checked_add(chunk.len())), - } - } - - /// Whether this value contains no bytes. - pub fn is_empty(&self) -> bool { - self.len() == Some(0) - } -} - -/// Borrowed external values collected while encoding a FlatBuffer. -#[derive(Debug, Default)] -pub struct ExternalValueRefs<'a> { - values: Vec>, -} - -impl<'a> ExternalValueRefs<'a> { - /// Create an empty collection. - pub fn new() -> Self { - Self::default() - } - - /// Borrow the collected values in wire order. - pub fn as_slice(&self) -> &[ExternalValueRef<'a>] { - &self.values - } - - /// Number of collected logical values. - pub fn len(&self) -> usize { - self.values.len() - } - - /// Whether no logical values were collected. - pub fn is_empty(&self) -> bool { - self.values.is_empty() - } - - /// Total byte length of all collected values. - pub fn total_len(&self) -> Option { - self.values - .iter() - .try_fold(0usize, |len, value| len.checked_add(value.len()?)) - } -} - -impl<'a> ExternalValueSink<'a> for ExternalValueRefs<'a> { - fn push_bytes(&mut self, value: &'a [u8]) -> Result<()> { - self.values.push(ExternalValueRef::Bytes(value)); - Ok(()) - } - - fn push_chunks(&mut self, value: &'a [Bytes]) -> Result<()> { - self.values.push(ExternalValueRef::Chunks(value)); - Ok(()) - } -} - /// Receives external byte values while their FlatBuffer markers are encoded. /// /// Values are delivered in their logical order without flattening chunked diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs b/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs index 7270f191c..14208a07c 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/function_call.rs @@ -17,21 +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::codec::{ExternalValueSink, ExternalValueSource}; -use super::function_types::{ParameterValue, ReturnType, decode_external_parameter_value}; -use super::util::{byte_chunks_to_bytes, try_byte_chunks_len}; +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, hlbytechunks, hlbytechunksArgs, - hldouble, hldoubleArgs, hlexternalbytes, hlexternalbytesArgs, 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. @@ -76,166 +75,8 @@ 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] { - let function_name = builder.create_string(&self.function_name); - - let function_call_type = match self.function_call_type { - FunctionCallType::Guest => FbFunctionCallType::guest, - FunctionCallType::Host => FbFunctionCallType::host, - }; - - 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 - .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()), - }, - ) - } - ParameterValue::ByteChunks(v) => { - let value = byte_chunks_to_bytes(v); - let vec_bytes = builder.create_vector(value.as_ref()); - let hlbytechunks = hlbytechunks::create( - builder, - &hlbytechunksArgs { - value: Some(vec_bytes), - }, - ); - Parameter::create( - builder, - &ParameterArgs { - value_type: FbParameterValue::hlbytechunks, - value: Some(hlbytechunks.as_union_value()), - }, - ) - } - }) - .collect(); - Some(builder.create_vector(¶meter_offsets)) - } - _ => None, - }; - - let function_call = FbFunctionCall::create( - builder, - &FbFunctionCallArgs { - function_name: Some(function_name), - parameters, - function_call_type, - expected_return_type, - }, - ); - builder.finish_size_prefixed(function_call, None); - builder.finished_data() - } - - /// Encodes byte parameters as external markers and sends their payloads to - /// `external_values` in parameter order. - pub fn encode_external<'a, 'b, S>( + /// 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, @@ -412,9 +253,8 @@ impl FunctionCall { Ok(builder.finished_data()) } - /// Decodes a function call using `external_values` for external byte - /// markers. - pub fn decode_external(value: &[u8], external_values: &mut S) -> Result + /// Decode control data and consume external byte parameters. + pub fn decode(value: &[u8], external_values: &mut S) -> Result where S: ExternalValueSource + ?Sized, { @@ -435,7 +275,7 @@ impl FunctionCall { .map(|parameters| { parameters .iter() - .map(|parameter| decode_external_parameter_value(parameter, external_values)) + .map(|parameter| decode_parameter_value(parameter, external_values)) .collect::>>() }) .transpose()?; @@ -450,64 +290,6 @@ impl FunctionCall { } } -#[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 { - 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(); - let function_call_type = match function_call_fb.function_call_type() { - FbFunctionCallType::guest => FunctionCallType::Guest, - FbFunctionCallType::host => FunctionCallType::Host, - other => { - bail!("Invalid function call type: {:?}", other); - } - }; - let expected_return_type = function_call_fb.expected_return_type().try_into()?; - - let parameters = function_call_fb - .parameters() - .map(|v| { - v.iter() - .map(|p| p.try_into()) - .collect::>>() - }) - .transpose()?; - - Ok(Self { - function_name: function_name.to_string(), - parameters, - function_call_type, - expected_return_type, - }) - } -} - #[cfg(test)] mod tests { use alloc::collections::VecDeque; @@ -582,6 +364,7 @@ mod tests { #[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![ @@ -601,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(); @@ -629,32 +412,7 @@ mod tests { } #[test] - fn embedded_byte_parameters_round_trip_as_distinct_logical_types() { - let mut builder = FlatBufferBuilder::new(); - let parameters = vec![ - ParameterValue::VecBytes(vec![1, 2, 3]), - ParameterValue::ByteChunks(vec![Bytes::from_static(&[4, 5]), Bytes::from_static(&[6])]), - ]; - let encoded = FunctionCall::new( - "bytes".to_string(), - Some(parameters), - FunctionCallType::Host, - ReturnType::VecBytes, - ) - .encode(&mut builder); - - let decoded = FunctionCall::try_from(encoded).unwrap(); - assert_eq!( - decoded.parameters, - Some(vec![ - ParameterValue::VecBytes(vec![1, 2, 3]), - ParameterValue::ByteChunks(vec![Bytes::from_static(&[4, 5, 6])]), - ]) - ); - } - - #[test] - fn external_byte_parameters_round_trip_in_parameter_order() { + fn byte_parameters_round_trip_in_external_value_order() { let mut builder = FlatBufferBuilder::new(); let expected_parameters = vec![ ParameterValue::Int(7), @@ -680,9 +438,7 @@ mod tests { ReturnType::ByteChunks, ); let mut external_values = TestExternalValues::default(); - let encoded = call - .encode_external(&mut builder, &mut external_values) - .unwrap(); + let encoded = call.encode(&mut builder, &mut external_values).unwrap(); assert!(encoded.len() < 4096); assert_eq!( @@ -713,8 +469,7 @@ mod tests { assert_eq!(marker.chunked(), chunked); } - assert!(FunctionCall::try_from(encoded).is_err()); - let decoded = FunctionCall::decode_external(encoded, &mut external_values).unwrap(); + 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); @@ -723,7 +478,7 @@ mod tests { } #[test] - fn external_encoding_matches_embedded_encoding_without_byte_parameters() { + fn scalar_call_uses_no_external_values() { let call = FunctionCall::new( "scalars".to_string(), Some(vec![ @@ -733,17 +488,13 @@ mod tests { FunctionCallType::Guest, ReturnType::Bool, ); - let mut embedded_builder = FlatBufferBuilder::new(); - let embedded = call.encode(&mut embedded_builder).to_vec(); - - let mut external_builder = FlatBufferBuilder::new(); + let mut builder = FlatBufferBuilder::new(); let mut external_values = TestExternalValues::default(); - let external = call - .encode_external(&mut external_builder, &mut external_values) - .unwrap(); + let encoded = call.encode(&mut builder, &mut external_values).unwrap(); - assert_eq!(external, embedded); assert!(external_values.values.is_empty()); + let decoded = FunctionCall::decode(encoded, &mut external_values).unwrap(); + assert_eq!(decoded.function_name, "scalars"); } #[test] @@ -756,27 +507,25 @@ mod tests { ReturnType::Void, ); let mut encoded_values = TestExternalValues::default(); - let encoded = call - .encode_external(&mut builder, &mut encoded_values) - .unwrap(); + let encoded = call.encode(&mut builder, &mut encoded_values).unwrap(); let mut missing = TestExternalValues::default(); - assert!(FunctionCall::decode_external(encoded, &mut missing).is_err()); + 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_external(encoded, &mut wrong_type).is_err()); + assert!(FunctionCall::decode(encoded, &mut wrong_type).is_err()); let mut wrong_length = TestExternalValues::from_values([TestExternalValue::VecBytes(vec![1, 2])]); - assert!(FunctionCall::decode_external(encoded, &mut wrong_length).is_err()); + 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_external(encoded, &mut extra).is_err()); + 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 a0117f437..0f7af656b 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/function_types.rs @@ -27,29 +27,29 @@ use super::codec::{ExternalValueSink, ExternalValueSource}; use super::guest_error::GuestError; #[cfg(feature = "fuzzing")] use super::util::arbitrary_byte_chunks; -use super::util::{byte_chunks_from_bytes, byte_chunks_to_bytes, try_byte_chunks_len}; +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, hlexternalbytes, hlexternalbytesArgs, hlfloat, hlfloatArgs, hlint, hlintArgs, hllong, - hllongArgs, hlsizeprefixedbuffer, hlsizeprefixedbufferArgs, hlsizeprefixedbytechunks, - hlsizeprefixedbytechunksArgs, hlstring, hlstringArgs, hluint, hluintArgs, hlulong, hlulongArgs, - hlvoid, hlvoidArgs, + 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 @@ -88,33 +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 value = byte_chunks_to_bytes(v); - let val = builder.create_vector(value.as_ref()); - let off = hlsizeprefixedbytechunks::create( + 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, - &hlsizeprefixedbytechunksArgs { - value: Some(val), - size: value.len() as i32, + &hlexternalbytesArgs { + length, + chunked: true, }, ); - ( - Some(off.as_union_value()), - FbReturnValue::hlsizeprefixedbytechunks, - ) + (Some(off.as_union_value()), FbReturnValue::hlexternalbytes) } ReturnValue::Void(()) => { let off = hlvoid::create(builder, &hlvoidArgs {}); @@ -131,7 +134,7 @@ impl FunctionCallResult { }, ); builder.finish_size_prefixed(fcr, None); - builder.finished_data() + Ok(builder.finished_data()) } Err(ge) => { // Encode GuestError @@ -152,65 +155,11 @@ impl FunctionCallResult { }, ); builder.finish_size_prefixed(fcr, None); - builder.finished_data() + Ok(builder.finished_data()) } } } - /// Encodes byte returns as external markers and sends their payload to - /// `external_values`. - /// - /// Non-byte returns and guest errors retain their existing embedded - /// encoding. - pub fn encode_external<'a, 'b, S>( - &'a self, - builder: &'b mut flatbuffers::FlatBufferBuilder, - external_values: &mut S, - ) -> Result<&'b [u8]> - where - S: ExternalValueSink<'a> + ?Sized, - { - let Ok(return_value) = &self.0 else { - return Ok(self.encode(builder)); - }; - - let (length, chunked) = match return_value { - ReturnValue::VecBytes(value) => { - let length = u64::try_from(value.len()) - .map_err(|_| anyhow!("External VecBytes length does not fit in u64"))?; - external_values.push_bytes(value)?; - (length, false) - } - ReturnValue::ByteChunks(value) => { - let length = try_byte_chunks_len(value) - .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(value)?; - (length, true) - } - _ => return Ok(self.encode(builder)), - }; - - let value = hlexternalbytes::create(builder, &hlexternalbytesArgs { length, chunked }); - let return_value = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value: Some(value.as_union_value()), - value_type: FbReturnValue::hlexternalbytes, - }, - ); - let result = FbFunctionCallResult::create( - builder, - &FbFunctionCallResultArgs { - result: Some(return_value.as_union_value()), - result_type: FunctionCallResultType::ReturnValueBox, - }, - ); - builder.finish_size_prefixed(result, None); - Ok(builder.finished_data()) - } - pub fn new(value: core::result::Result) -> Self { FunctionCallResult(value) } @@ -219,9 +168,8 @@ impl FunctionCallResult { self.0 } - /// Decodes a function-call result using `external_values` for external byte - /// markers. - pub fn decode_external(value: &[u8], external_values: &mut S) -> Result + /// Decode control data and consume an external byte return. + pub fn decode(value: &[u8], external_values: &mut S) -> Result where S: ExternalValueSource + ?Sized, { @@ -235,7 +183,7 @@ impl FunctionCallResult { .ok_or_else(|| { anyhow!("Failed to get ReturnValueBox from function call result") })?; - Ok(decode_external_return_value(boxed, external_values)?) + Ok(decode_return_value(boxed, external_values)?) } FunctionCallResultType::GuestError => { let guest_error_table = function_call_result_fb @@ -258,44 +206,6 @@ impl FunctionCallResult { } } -impl TryFrom<&[u8]> for FunctionCallResult { - type Error = Error; - - fn try_from(value: &[u8]) -> Result { - 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() { - 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))) - } - FunctionCallResultType::GuestError => { - let guest_error_table = function_call_result_fb - .result_as_guest_error() - .ok_or_else(|| anyhow!("Failed to get GuestError from function call result"))?; - let code = guest_error_table.code(); - let message = guest_error_table - .message() - .map(|s| s.to_string()) - .unwrap_or_default(); - Ok(FunctionCallResult(Err(GuestError::new( - code.into(), - message, - )))) - } - other => { - bail!("Unexpected function call result type: {:?}", other) - } - } - } -} - /// Supported parameter types with values for function calling. #[cfg_attr(feature = "fuzzing", derive(arbitrary::Arbitrary))] #[derive(Debug, Clone, PartialEq)] @@ -413,15 +323,19 @@ pub enum ReturnType { ByteChunks, } -pub(crate) fn decode_external_parameter_value( +pub(crate) fn decode_parameter_value( parameter: Parameter<'_>, external_values: &mut S, ) -> Result where S: ExternalValueSource + ?Sized, { - if parameter.value_type() != FbParameterValue::hlexternalbytes { - return parameter.try_into(); + match parameter.value_type() { + FbParameterValue::hlexternalbytes => {} + FbParameterValue::hlvecbytes | FbParameterValue::hlbytechunks => { + bail!("Embedded byte parameters are not supported") + } + _ => return parameter.try_into(), } let marker = parameter @@ -459,15 +373,19 @@ where } } -fn decode_external_return_value( +fn decode_return_value( return_value: ReturnValueBox<'_>, external_values: &mut S, ) -> Result where S: ExternalValueSource + ?Sized, { - if return_value.value_type() != FbReturnValue::hlexternalbytes { - return return_value.try_into(); + match return_value.value_type() { + FbReturnValue::hlexternalbytes => {} + FbReturnValue::hlsizeprefixedbuffer | FbReturnValue::hlsizeprefixedbytechunks => { + bail!("Embedded byte returns are not supported") + } + _ => return return_value.try_into(), } let marker = return_value @@ -554,14 +472,9 @@ 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::hlbytechunks => param.value_as_hlbytechunks().map(|hlbytechunks| { - ParameterValue::ByteChunks(byte_chunks_from_bytes(Bytes::copy_from_slice( - hlbytechunks.value().unwrap_or_default().bytes(), - ))) - }), + FbParameterValue::hlvecbytes | FbParameterValue::hlbytechunks => { + bail!("Embedded byte parameters are not supported") + } FbParameterValue::hlexternalbytes => { bail!("External byte parameter requires an external value source") } @@ -979,22 +892,8 @@ 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::hlsizeprefixedbytechunks => { - let value = return_value_box - .value_as_hlsizeprefixedbytechunks() - .and_then(|value| value.value()) - .map(|value| byte_chunks_from_bytes(Bytes::copy_from_slice(value.bytes()))) - .unwrap_or_default(); - Ok(ReturnValue::ByteChunks(value)) + FbReturnValue::hlsizeprefixedbuffer | FbReturnValue::hlsizeprefixedbytechunks => { + bail!("Embedded byte returns are not supported") } FbReturnValue::hlexternalbytes => { bail!("External byte return requires an external value source") @@ -1006,249 +905,6 @@ 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::ByteChunks(v) => { - let off = { - let value = byte_chunks_to_bytes(v); - let val = builder.create_vector(value.as_ref()); - hlsizeprefixedbytechunks::create( - &mut builder, - &hlsizeprefixedbytechunksArgs { - value: Some(val), - size: value.len() as i32, - }, - ) - }; - let rv_box = ReturnValueBox::create( - &mut builder, - &ReturnValueBoxArgs { - value: Some(off.as_union_value()), - value_type: FbReturnValue::hlsizeprefixedbytechunks, - }, - ); - 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; @@ -1257,7 +913,6 @@ mod tests { use flatbuffers::FlatBufferBuilder; use super::super::guest_error::ErrorCode; - use super::super::util::{byte_chunks_to_vec, get_flatbuffer_result}; use super::*; use crate::flatbuffers::hyperlight::generated::{hlexternalbytes, hlexternalbytesArgs}; @@ -1327,9 +982,13 @@ mod tests { #[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::try_from(test_data).unwrap(); + let function_call_result = + FunctionCallResult::decode(test_data, &mut external_values).unwrap(); let result = function_call_result.into_inner().unwrap(); assert_eq!(result, ReturnValue::Int(42)); } @@ -1341,45 +1000,18 @@ mod tests { ErrorCode::GuestFunctionNotFound, "Function not found".to_string(), ); - let test_data = FunctionCallResult::new(Err(test_error.clone())).encode(&mut builder); + 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::try_from(test_data).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 embedded_byte_chunks_return_round_trips() { - let mut builder = FlatBufferBuilder::new(); - let expected = vec![Bytes::from_static(b"hello"), Bytes::from_static(b" world")]; - let encoded = - FunctionCallResult::new(Ok(ReturnValue::ByteChunks(expected))).encode(&mut builder); - - let decoded = FunctionCallResult::try_from(encoded) - .unwrap() - .into_inner() - .unwrap(); - let ReturnValue::ByteChunks(decoded) = decoded else { - panic!("expected byte chunks return value"); - }; - assert_eq!(byte_chunks_to_vec(&decoded), b"hello world"); - } - - #[test] - fn direct_byte_chunks_return_encoding_preserves_logical_type() { - let encoded = get_flatbuffer_result(vec![ - Bytes::from_static(b"hello"), - Bytes::from_static(b" world"), - ]); - - let decoded = FunctionCallResult::try_from(encoded.as_slice()) - .unwrap() - .into_inner() - .unwrap(); - assert!(matches!(decoded, ReturnValue::ByteChunks(_))); - } - #[test] fn external_bytes_marks_chunked_values_only() { fn round_trip(chunked: bool) -> bool { @@ -1403,7 +1035,7 @@ mod tests { } #[test] - fn external_byte_returns_round_trip_without_embedding_payloads() { + fn byte_returns_round_trip_as_external_values() { for expected in [ ReturnValue::VecBytes(vec![0xa5; 4096]), ReturnValue::ByteChunks(vec![ @@ -1416,7 +1048,7 @@ mod tests { let mut builder = FlatBufferBuilder::new(); let mut external_values = TestExternalValues::default(); let encoded = FunctionCallResult::new(Ok(expected.clone())) - .encode_external(&mut builder, &mut external_values) + .encode(&mut builder, &mut external_values) .unwrap(); assert!(encoded.len() < 4096); @@ -1432,8 +1064,7 @@ mod tests { assert_eq!(marker.length(), length as u64); assert_eq!(marker.chunked(), chunked); - assert!(FunctionCallResult::try_from(encoded).is_err()); - let decoded = FunctionCallResult::decode_external(encoded, &mut external_values) + let decoded = FunctionCallResult::decode(encoded, &mut external_values) .unwrap() .into_inner() .unwrap(); @@ -1450,40 +1081,34 @@ mod tests { FunctionCallResult::new(Ok(ReturnValue::ByteChunks(vec![Bytes::from_static( b"123", )]))) - .encode_external(&mut builder, &mut encoded_values) + .encode(&mut builder, &mut encoded_values) .unwrap(); let mut missing = TestExternalValues::default(); - assert!(FunctionCallResult::decode_external(encoded, &mut missing).is_err()); + assert!(FunctionCallResult::decode(encoded, &mut missing).is_err()); let mut wrong_type = TestExternalValues::from_values([TestExternalValue::VecBytes(vec![1, 2, 3])]); - assert!(FunctionCallResult::decode_external(encoded, &mut wrong_type).is_err()); + 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_external(encoded, &mut wrong_length).is_err()); + 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 embedded_builder = FlatBufferBuilder::new(); - let embedded = result.encode(&mut embedded_builder).to_vec(); - - let mut external_builder = FlatBufferBuilder::new(); + let mut builder = FlatBufferBuilder::new(); let mut external_values = TestExternalValues::default(); - let external = result - .encode_external(&mut external_builder, &mut external_values) - .unwrap(); - assert_eq!(external, embedded); + 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_external(external, &mut external_values).is_err()); + assert!(FunctionCallResult::decode(encoded, &mut external_values).is_err()); } } diff --git a/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs b/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs index 0a7770e27..d013cb78c 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/mod.rs @@ -31,4 +31,4 @@ pub mod host_function_definition; pub mod host_function_details; pub mod util; -pub use codec::{ExternalValueRef, ExternalValueRefs, ExternalValueSink, ExternalValueSource}; +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 57213d1ba..cbc65eb62 100644 --- a/src/hyperlight_common/src/flatbuffer_wrappers/util.rs +++ b/src/hyperlight_common/src/flatbuffer_wrappers/util.rs @@ -18,251 +18,10 @@ use alloc::vec; use alloc::vec::Vec; use bytes::Bytes; -use flatbuffers::FlatBufferBuilder; 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, - hlsizeprefixedbytechunks as Fbhlsizeprefixedbytechunks, - hlsizeprefixedbytechunksArgs as FbhlsizeprefixedbytechunksArgs, 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 Vec { - fn serialize(&self, builder: &mut FlatBufferBuilder) -> FbFunctionCallResultArgs { - let value = byte_chunks_to_bytes(self); - let vec_off = builder.create_vector(value.as_ref()); - let buf_off = Fbhlsizeprefixedbytechunks::create( - builder, - &FbhlsizeprefixedbytechunksArgs { - size: value.len() as i32, - value: Some(vec_off), - }, - ); - let rv_box = ReturnValueBox::create( - builder, - &ReturnValueBoxArgs { - value_type: FbReturnValue::hlsizeprefixedbytechunks, - 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). @@ -273,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; @@ -288,8 +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::ByteChunks(v) => byte_chunks_len(v) + 20, + ParameterValue::VecBytes(_) | ParameterValue::ByteChunks(_) => 20, ParameterValue::Int(_) | ParameterValue::UInt(_) => 16, ParameterValue::Long(_) | ParameterValue::ULong(_) => 20, ParameterValue::Float(_) => 16, @@ -363,9 +121,12 @@ mod tests { 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::ExternalValueRefs; /// Helper function to check that estimation is within reasonable bounds (±25%) fn assert_estimation_accuracy( @@ -384,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 = ExternalValueRefs::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; @@ -413,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/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/transport.rs b/src/hyperlight_common/src/transport.rs new file mode 100644 index 000000000..7f1666638 --- /dev/null +++ b/src/hyperlight_common/src/transport.rs @@ -0,0 +1,404 @@ +/* +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, +} + +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), + } + } +} + +/// 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: ExternalValueRefs<'a>, + wire_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: ExternalValueRefs<'a>, + ) -> Option { + let payload_len = control.len() + externals.total_len(); + let payload_len = u32::try_from(payload_len).ok()?; + let wire_len = MsgHeader::SIZE + payload_len as usize; + + Some(Self { + header: MsgHeader::new(kind, cid, payload_len), + control, + externals, + wire_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.wire_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() + } + + /// 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.wire_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 &[]; + } + + 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 { + 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 ExternalValueRefs<'a> { + chunks: Vec<&'a [u8]>, +} + +impl<'a> ExternalValueRefs<'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 fn total_len(&self) -> usize { + self.chunks.iter().map(|chunk| chunk.len()).sum() + } +} + +impl<'a> ExternalValueSink<'a> for ExternalValueRefs<'a> { + fn push_bytes(&mut self, value: &'a [u8]) -> Result<()> { + if !value.is_empty() { + self.chunks.push(value); + } + Ok(()) + } + + fn push_chunks(&mut self, value: &'a [Bytes]) -> Result<()> { + self.chunks.extend( + value + .iter() + .map(Bytes::as_ref) + .filter(|chunk| !chunk.is_empty()), + ); + 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 = ExternalValueRefs::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.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 = ExternalValueRefs::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 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/consumer.rs b/src/hyperlight_common/src/virtq/consumer.rs index 21b39da82..48d77b401 100644 --- a/src/hyperlight_common/src/virtq/consumer.rs +++ b/src/hyperlight_common/src/virtq/consumer.rs @@ -15,6 +15,7 @@ limitations under the License. */ use alloc::vec; +use alloc::vec::Vec; use core::fmt; use bytes::Bytes; @@ -198,6 +199,12 @@ pub enum ReplyChain { 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] @@ -264,6 +271,12 @@ impl WritableChain { 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 { @@ -457,11 +470,7 @@ impl VirtqConsumer { /// /// - [`VirtqError::BadChain`] - Descriptor chain format not recognized /// - [`VirtqError::InvalidState`] - Descriptor ID collision (driver bug) - #[allow(clippy::type_complexity)] - pub fn poll( - &mut self, - max_recv_len: usize, - ) -> Result, ReplyChain)>, 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), @@ -525,6 +534,80 @@ impl VirtqConsumer { Ok(Some((chain, reply))) } + /// Poll exactly `count` chains without consuming a partial batch. + /// + /// 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 @@ -657,6 +740,28 @@ impl VirtqConsumer { self.next_token = 0; Ok(()) } + + 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(()) + } } type ChainElems = SmallVec<[BufferElement; 4]>; @@ -984,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); diff --git a/src/hyperlight_common/src/virtq/mod.rs b/src/hyperlight_common/src/virtq/mod.rs index 01ca7dc75..c47439d82 100644 --- a/src/hyperlight_common/src/virtq/mod.rs +++ b/src/hyperlight_common/src/virtq/mod.rs @@ -165,7 +165,6 @@ mod buffer; mod consumer; mod desc; mod event; -pub mod msg; mod pool; mod producer; mod ring; diff --git a/src/hyperlight_common/src/virtq/msg.rs b/src/hyperlight_common/src/virtq/msg.rs deleted file mode 100644 index 1f19988b0..000000000 --- a/src/hyperlight_common/src/virtq/msg.rs +++ /dev/null @@ -1,233 +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 framing for virtqueue messages. -//! -//! Every message chain on both the G2H and H2G queues starts with this fixed -//! 8-byte header, enabling message type discrimination and request/response -//! correlation. Payload lengths come from the size-prefixed FlatBuffer and its -//! external-byte declarations. - -use crate::flatbuffer_wrappers::{ExternalValueRef, ExternalValueRefs}; - -/// 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, -} - -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), - } - } -} - -/// Wire header for all virtqueue messages. -#[derive(Debug, Clone, Copy, PartialEq, Eq, bytemuck::Pod, bytemuck::Zeroable)] -#[repr(C)] -pub struct VirtqMsgHeader { - /// Discriminates the message type. - pub kind: u8, - /// keep the header 8 bytes long and aligned to 4 bytes. - reserved: [u8; 3], - /// Caller-assigned correlation ID. Responses echo the request's ID. - pub cid: u32, -} - -impl VirtqMsgHeader { - pub const SIZE: usize = core::mem::size_of::(); - - /// Create a message header. - pub const fn new(kind: MsgKind, cid: u32) -> Self { - Self { - kind: kind as u8, - reserved: [0; 3], - cid, - } - } - - /// 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: VirtqMsgHeader, - control: &'a [u8], - externals: ExternalValueRefs<'a>, - wire_len: usize, -} - -impl<'a> EncodedMessage<'a> { - /// Build a message, returning `None` if its wire length overflows. - pub fn new( - kind: MsgKind, - cid: u32, - control: &'a [u8], - externals: ExternalValueRefs<'a>, - ) -> Option { - let wire_len = VirtqMsgHeader::SIZE - .checked_add(control.len())? - .checked_add(externals.total_len()?)?; - - Some(Self { - header: VirtqMsgHeader::new(kind, cid), - control, - externals, - wire_len, - }) - } - - /// Visit wire chunks in transmission order. - pub fn try_for_each_chunk( - &self, - mut visit: impl FnMut(&[u8]) -> Result<(), E>, - ) -> Result<(), E> { - visit(self.header.as_bytes())?; - visit(self.control)?; - - for val in self.externals.as_slice() { - match val { - ExternalValueRef::Bytes(value) => visit(value)?, - ExternalValueRef::Chunks(chunks) => chunks.iter().try_for_each(|c| visit(c))?, - } - } - - Ok(()) - } - - /// Total wire length of all chunks. - pub const fn wire_len(&self) -> usize { - self.wire_len - } -} - -/// 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_only_kind_and_cid() { - let header = VirtqMsgHeader::new(MsgKind::Response, 0x1234_5678); - - assert_eq!(VirtqMsgHeader::SIZE, 8); - assert_eq!(header.msg_kind(), Ok(MsgKind::Response)); - assert_eq!(header.cid, 0x1234_5678); - assert_eq!(header.reserved, [0; 3]); - } - - #[test] - fn rejects_invalid_wire_headers() { - let header = VirtqMsgHeader::new(MsgKind::Request, 1); - let mut bytes = [0; VirtqMsgHeader::SIZE]; - bytes.copy_from_slice(header.as_bytes()); - - bytes[1] = 1; - assert_eq!(VirtqMsgHeader::from_bytes(&bytes), None); - - bytes[1] = 0; - bytes[0] = u8::MAX; - assert_eq!(VirtqMsgHeader::from_bytes(&bytes), None); - assert_eq!(VirtqMsgHeader::from_bytes(&bytes[..7]), None); - } - - #[test] - fn encoded_message_visits_wire_chunks_in_order() { - let chunks = [ - bytes::Bytes::from_static(b"ef"), - bytes::Bytes::from_static(b"gh"), - ]; - let mut external_values = ExternalValueRefs::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 mut visited = Vec::new(); - message - .try_for_each_chunk(|chunk| { - visited.push(chunk.to_vec()); - Ok::<_, ()>(()) - }) - .unwrap(); - - assert_eq!(message.wire_len(), VirtqMsgHeader::SIZE + 8); - assert_eq!(visited[1..], [b"ab", b"cd", b"ef", b"gh"]); - } - - #[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/producer.rs b/src/hyperlight_common/src/virtq/producer.rs index 369e9c447..e7c4cea2d 100644 --- a/src/hyperlight_common/src/virtq/producer.rs +++ b/src/hyperlight_common/src/virtq/producer.rs @@ -479,9 +479,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, @@ -527,6 +527,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. @@ -1683,6 +1689,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); diff --git a/src/hyperlight_common/src/virtq/ring.rs b/src/hyperlight_common/src/virtq/ring.rs index 1cb2f2635..cad27d47c 100644 --- a/src/hyperlight_common/src/virtq/ring.rs +++ b/src/hyperlight_common/src/virtq/ring.rs @@ -447,6 +447,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 @@ -1189,6 +1202,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, 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_guest/src/guest_handle/host_comm.rs b/src/hyperlight_guest/src/guest_handle/host_comm.rs index 73c33ddcc..4581c23f2 100644 --- a/src/hyperlight_guest/src/guest_handle/host_comm.rs +++ b/src/hyperlight_guest/src/guest_handle/host_comm.rs @@ -67,9 +67,7 @@ impl GuestHandle { parameters: Option>, return_type: ReturnType, ) -> Result { - transport::with_context(|context| { - context.call_host_function(function_name, parameters, return_type) - }) + transport::with_ctx(|ctx| ctx.call_host_function(function_name, parameters, return_type)) } /// Log a message with the specified log level, source, caller, source file, and line number. @@ -97,9 +95,8 @@ impl GuestHandle { .try_into() .expect("Failed to convert GuestLogData to bytes"); - transport::with_context(|context| { - context - .emit_log(&bytes) + transport::with_ctx(|ctx| { + ctx.emit_log(&bytes) .expect("Unable to send log data via virtq"); }); }; diff --git a/src/hyperlight_guest/src/transport/response.rs b/src/hyperlight_guest/src/transport/codec.rs similarity index 53% rename from src/hyperlight_guest/src/transport/response.rs rename to src/hyperlight_guest/src/transport/codec.rs index a88f28542..3af1edb61 100644 --- a/src/hyperlight_guest/src/transport/response.rs +++ b/src/hyperlight_guest/src/transport/codec.rs @@ -14,35 +14,43 @@ See the License for the specific language governing permissions and limitations under the License. */ -//! Guest G2H response decoding. +//! 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::virtq::Segments; -use hyperlight_common::virtq::msg::{ - MsgKind, SIZE_PREFIX_LEN, VirtqMsgHeader, size_prefix_payload_len, size_prefixed_len, +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 `header | size-prefixed control | external values`. +/// Decode one H2G guest-function request payload. /// -/// Contiguous byte values are flattened into `Vec`. Chunked values retain -/// their transport-backed `Bytes` owners and return pool slots when dropped. -pub(super) fn decode(mut segments: Segments, cid: u32) -> Result { - // Validate the transport envelope before interpreting the response body. - let header = segments - .split_to(VirtqMsgHeader::SIZE) - .context("host function response is missing its header")? - .into_bytes(); +/// 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 Some(header) = VirtqMsgHeader::from_bytes(&header) else { - bail!("Host function response has an invalid header"); - }; + 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"); } @@ -51,22 +59,44 @@ pub(super) fn decode(mut segments: Segments, cid: u32) -> Result 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("host function response is missing its size prefix")? + .context("virtqueue message is missing its size prefix")? .into_bytes(); let payload_len = - size_prefix_payload_len(&prefix).context("host function response has an invalid prefix")?; + size_prefix_payload_len(&prefix).context("virtqueue message has an invalid prefix")?; let payload = segments .split_to(payload_len) - .context("host function response control data is truncated")?; + .context("virtqueue message control data is truncated")?; let control_len = - size_prefixed_len(payload_len).context("host function response control length overflow")?; + size_prefixed_len(payload_len).context("virtqueue message control length overflow")?; let mut control = Vec::with_capacity(control_len); control.extend_from_slice(&prefix); @@ -75,9 +105,7 @@ pub(super) fn decode(mut segments: Segments, cid: u32) -> Result anyhow::Result<()> { if !self.segments.is_empty() { anyhow::bail!( - "Host function response has {} trailing external bytes", + "Virtqueue message has {} trailing external bytes", self.segments.len() ); } @@ -130,8 +158,8 @@ mod tests { use alloc::vec; use flatbuffers::FlatBufferBuilder; - use hyperlight_common::flatbuffer_wrappers::ExternalValueRefs; use hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue; + use hyperlight_common::transport::ExternalValueRefs; use super::*; @@ -142,17 +170,16 @@ mod tests { let result = FunctionCallResult::new(Ok(ReturnValue::ByteChunks(vec![external.clone()]))); let mut builder = FlatBufferBuilder::new(); let mut external_values = ExternalValueRefs::new(); - let control = result - .encode_external(&mut builder, &mut external_values) - .unwrap(); - let header = VirtqMsgHeader::new(MsgKind::Response, 7); + 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(segments, 7).unwrap().into_inner().unwrap(); + let decoded = decode_response(segments, 7).unwrap().into_inner().unwrap(); let ReturnValue::ByteChunks(chunks) = decoded else { panic!("expected ByteChunks response"); }; @@ -178,4 +205,35 @@ mod tests { 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 = ExternalValueRefs::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 index 443877a2b..d3f1f940c 100644 --- a/src/hyperlight_guest/src/transport/context.rs +++ b/src/hyperlight_guest/src/transport/context.rs @@ -20,20 +20,21 @@ use alloc::vec::Vec; use core::result; use flatbuffers::FlatBufferBuilder; -use hyperlight_common::flatbuffer_wrappers::ExternalValueRefs; use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType}; use hyperlight_common::flatbuffer_wrappers::function_types::{ - ParameterValue, ReturnType, ReturnValue, + 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::virtq::msg::{EncodedMessage, MsgKind}; +use hyperlight_common::transport::{EncodedMessage, ExternalValueRefs, MsgHeader, MsgKind}; use hyperlight_common::virtq::{ AllocError, BufferProvider, G2H_LOWER_SLOT_COUNT, G2H_LOWER_SLOT_SIZE, Layout, Notifier, - QueueStats, SlotLayout, SlotPool, Token, UsedChain, VirtqError, VirtqProducer, + QueueStats, Segments, SendChain, SlotLayout, SlotPool, Token, UsedChain, VirtqError, + VirtqProducer, }; -use super::{GuestMemOps, response}; +use super::{GuestMemOps, codec}; use crate::bail; use crate::error::{GuestErrorContext, Result}; use crate::exit::out32; @@ -90,6 +91,8 @@ pub struct GuestContext { 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 { @@ -102,7 +105,7 @@ impl GuestContext { 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); + let h2g_producer = VirtqProducer::new(h2g.layout, mem, H2gNotifier, h2g_pool.clone()); let mut ctx = Self { g2h_producer, @@ -111,12 +114,23 @@ impl GuestContext { h2g_slot_size: h2g.buffer_size, 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. /// /// Control data and borrowed external values form one readable request. @@ -147,20 +161,20 @@ impl GuestContext { let mut externals = ExternalValueRefs::new(); let control = fc - .encode_external(&mut builder, &mut externals) + .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 message = EncodedMessage::new(MsgKind::Request, cid, control, externals) + let msg = EncodedMessage::new(MsgKind::Request, cid, control, externals) .context("G2H message length overflow")?; // Reserve response capacity from the ring and pool state that remains // after this request. - let reply_cap = self.reply_capacity(message.wire_len(), return_type)?; + let reply_cap = self.reply_capacity(msg.total_len(), return_type)?; // Submit once more after forcing the host to drain on backpressure. - let token = match self.try_send(&message, Some(reply_cap)) { + let token = match self.try_send(&msg, Some(reply_cap)) { Ok(token) => token, Err(error) if error.is_transient() => { self.g2h_producer.notify_backpressure(); @@ -169,7 +183,7 @@ impl GuestContext { bail!("G2H reclaim: {error}"); } - match self.try_send(&message, Some(reply_cap)) { + match self.try_send(&msg, Some(reply_cap)) { Ok(token) => token, Err(error) => bail!("G2H call retry: {error}"), } @@ -200,7 +214,7 @@ impl GuestContext { // Decode external ByteChunks without flattening their transport-backed // segments. - let fcr = response::decode(segments, cid)?; + let fcr = codec::decode_response(segments, cid)?; let ret = fcr.into_inner()?; let Ok(ret) = T::try_from(ret) else { @@ -210,6 +224,98 @@ impl GuestContext { Ok(ret) } + /// Receive one host-to-guest function call. + /// + /// External `ByteChunks` retain their owner-backed H2G slots. Contiguous + /// `VecBytes` values copy directly into their final `Vec`. + pub fn recv_h2g_call(&mut self) -> Result<(u32, FunctionCall)> { + 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"); + }; + if header.msg_kind() != Ok(MsgKind::Request) || 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()); + } + + codec::decode_request(header.cid, Segments::new(payload)) + } + + /// 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 external_values = ExternalValueRefs::new(); + + let control = result + .encode(&mut builder, &mut external_values) + .with_context(|| "failed to encode guest function result")?; + + let msg = EncodedMessage::new(MsgKind::Response, cid, control, external_values) + .context("G2H response length overflow")?; + + self.try_send_deferred(&msg, None) + .with_context(|| "G2H response submission failed")?; + } + + drop(result); + self.prefill_h2g() + } + /// Send a log message via the G2H queue. /// /// Current notification policy exits to the host for every log. @@ -249,8 +355,10 @@ impl GuestContext { } } - /// Pre-fill the H2G queue with writable-only descriptors so the host - /// can write incoming call payloads into them. + /// 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(); @@ -258,23 +366,19 @@ impl GuestContext { let chain = match batch.chain().writable(self.h2g_slot_size).build() { Ok(chain) => chain, Err(error) if error.is_transient() => { - batch.finish()?; + batch.finish_without_notify(); return Ok(()); } - Err(error) => { - bail!("H2G prefill build: {error}"); - } + Err(error) => bail!("H2G prefill build: {error}"), }; match batch.submit(chain) { Ok(_) => {} Err(error) if error.is_transient() => { - batch.finish()?; + batch.finish_without_notify(); return Ok(()); } - Err(error) => { - bail!("H2G prefill submit: {error}"); - } + Err(error) => bail!("H2G prefill submit: {error}"), } } } @@ -371,8 +475,33 @@ impl GuestContext { message: &EncodedMessage<'_>, reply_cap: Option, ) -> 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_cap: Option, + ) -> result::Result { + let chain = self.build_g2h_chain(message, reply_cap)?; + 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: Option, + ) -> result::Result, VirtqError> { // Allocate the readable request and optional writable reply together. - let chain = self.g2h_producer.chain().readable(message.wire_len()); + let chain = self.g2h_producer.chain().readable(message.total_len()); let mut chain = match reply_cap { Some(reply_cap) => chain.writable(reply_cap), @@ -380,13 +509,11 @@ impl GuestContext { } .build()?; - message.try_for_each_chunk(|chunk| { + for chunk in message.chunks() { chain.write_all(chunk)?; - Ok::<(), VirtqError>(()) - })?; + } - // Transfer the initialized chain to the producer. - self.g2h_producer.submit(chain) + Ok(chain) } /// Allocate a new correlation ID for a host function request. @@ -409,8 +536,7 @@ fn pool_len(pages: usize) -> result::Result { /// Build the uniform H2G pool. /// -/// Every preposted receive buffer has the configured size so the host sees one -/// predictable capacity for guest calls. +/// 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)) diff --git a/src/hyperlight_guest/src/transport/mod.rs b/src/hyperlight_guest/src/transport/mod.rs index 4e4a318d5..625d96d1a 100644 --- a/src/hyperlight_guest/src/transport/mod.rs +++ b/src/hyperlight_guest/src/transport/mod.rs @@ -18,9 +18,9 @@ limitations under the License. //! //! Global context is installed once via [`set_global_context`] and accessed via [`with_context`]. +mod codec; pub mod context; pub mod mem; -mod response; use core::cell::RefCell; use core::sync::atomic::{AtomicU8, Ordering}; @@ -49,10 +49,10 @@ pub fn is_initialized() -> bool { /// # Panics /// /// Panics if the context is uninitialized or already borrowed. -pub fn with_context(f: impl FnOnce(&mut GuestContext) -> R) -> R { +pub fn with_ctx(f: impl FnOnce(&mut GuestContext) -> R) -> R { assert!(is_initialized(), "transport context not initialized"); - let mut context = GLOBAL_CONTEXT.0.borrow_mut(); - f(context.as_mut().expect("transport context missing")) + let mut ctx = GLOBAL_CONTEXT.0.borrow_mut(); + f(ctx.as_mut().expect("transport context missing")) } /// Install the global transport 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..8d983ee80 100644 --- a/src/hyperlight_guest_bin/src/guest_function/call.rs +++ b/src/hyperlight_guest_bin/src/guest_function/call.rs @@ -17,15 +17,16 @@ 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::{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 +35,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 +74,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 +95,12 @@ 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::() + let (cid, function_call) = transport::with_ctx(|ctx| ctx.recv_h2g_call()) .expect("Function call 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 = call_guest_function(function_call) + .map_err(|error| GuestError::new(error.kind, error.message)); + let result = FunctionCallResult::new(result); // All this tracing logic shall be done right before the call to `hlt` which is done after this // function returns @@ -139,4 +118,7 @@ pub(crate) fn internal_dispatch_function() { // the host, if necessary. hyperlight_guest_tracing::flush(); } + + transport::with_ctx(|ctx| ctx.send_h2g_result(cid, result)) + .expect("Failed to send function call result"); } diff --git a/src/hyperlight_guest_bin/src/guest_function/definition.rs b/src/hyperlight_guest_bin/src/guest_function/definition.rs index 46347f016..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,22 +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()), - ReturnValue::ByteChunks(v) => get_flatbuffer_result(v), - } -} - macro_rules! impl_host_function { ([$N:expr] ($($p:ident: $P:ident),*)) => { impl IntoGuestFunction for F @@ -134,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). @@ -144,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/host_comm.rs b/src/hyperlight_guest_bin/src/host_comm.rs index e7c026cbc..51126daa7 100644 --- a/src/hyperlight_guest_bin/src/host_comm.rs +++ b/src/hyperlight_guest_bin/src/host_comm.rs @@ -22,7 +22,6 @@ 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; @@ -37,9 +36,7 @@ pub fn call_host_function( where T: TryFrom, { - transport::with_context(|context| { - context.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 @@ -55,10 +52,7 @@ pub fn read_n_bytes_from_user_memory(num: u64) -> Result> { } /// 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> { +pub fn print_output_with_host_print(function_call: FunctionCall) -> Result { if let ParameterValue::String(message) = function_call.parameters.unwrap().remove(0) { let res = call_host_function::( "HostPrint", @@ -66,7 +60,7 @@ pub fn print_output_with_host_print(function_call: FunctionCall) -> Result