@@ -21,15 +21,20 @@ use rand::{RngCore, rng};
2121use tracing:: { Span , instrument} ;
2222
2323use super :: memory_region:: MemoryRegionType :: {
24- Code , GuardPage , Heap , HostFunctionDefinitions , InputData , OutputData , PageTables , Peb , Stack ,
24+ Code , GuardPage , Heap , HostFunctionDefinitions , InitData , InputData , OutputData , PageTables ,
25+ Peb , Stack ,
26+ } ;
27+ use super :: memory_region:: {
28+ DEFAULT_GUEST_BLOB_MEM_FLAGS , MemoryRegion , MemoryRegionFlags , MemoryRegionVecBuilder ,
2529} ;
26- use super :: memory_region:: { MemoryRegion , MemoryRegionFlags , MemoryRegionVecBuilder } ;
2730use super :: mgr:: AMOUNT_OF_MEMORY_PER_PT ;
2831use super :: shared_mem:: { ExclusiveSharedMemory , GuestSharedMemory , SharedMemory } ;
2932use crate :: error:: HyperlightError :: { GuestOffsetIsInvalid , MemoryRequestTooBig } ;
3033use crate :: sandbox:: SandboxConfiguration ;
3134use crate :: { Result , new_error} ;
3235
36+ // +-------------------------------------------+
37+ // | Init Data | (GuestBlob size)
3338// +-------------------------------------------+
3439// | Guest (User) Stack |
3540// +-------------------------------------------+
@@ -56,6 +61,8 @@ use crate::{Result, new_error};
5661// | PML4 |
5762// +-------------------------------------------+ 0x0_000
5863
64+ /// - `InitData` - some extra data that can be loaded onto the sandbox during
65+ /// initialization.
5966///
6067/// - `HostDefinitions` - the length of this is the `HostFunctionDefinitionSize`
6168/// field from `SandboxConfiguration`
@@ -82,6 +89,7 @@ pub(crate) struct SandboxMemoryLayout {
8289 pub ( super ) stack_size : usize ,
8390 /// The heap size of this sandbox.
8491 pub ( super ) heap_size : usize ,
92+ init_data_size : usize ,
8593
8694 /// The following fields are offsets to the actual PEB struct fields.
8795 /// They are used when writing the PEB struct itself
@@ -103,6 +111,7 @@ pub(crate) struct SandboxMemoryLayout {
103111 guest_heap_buffer_offset : usize ,
104112 guard_page_offset : usize ,
105113 guest_user_stack_buffer_offset : usize , // the lowest address of the user stack
114+ init_data_offset : usize ,
106115
107116 // other
108117 pub ( crate ) peb_address : usize ,
@@ -111,6 +120,7 @@ pub(crate) struct SandboxMemoryLayout {
111120 total_page_table_size : usize ,
112121 // The offset in the sandbox memory where the code starts
113122 guest_code_offset : usize ,
123+ pub ( crate ) init_data_permissions : Option < MemoryRegionFlags > ,
114124}
115125
116126impl Debug for SandboxMemoryLayout {
@@ -122,6 +132,10 @@ impl Debug for SandboxMemoryLayout {
122132 )
123133 . field ( "Stack Size" , & format_args ! ( "{:#x}" , self . stack_size) )
124134 . field ( "Heap Size" , & format_args ! ( "{:#x}" , self . heap_size) )
135+ . field (
136+ "Init Data Size" ,
137+ & format_args ! ( "{:#x}" , self . init_data_size) ,
138+ )
125139 . field ( "PEB Address" , & format_args ! ( "{:#x}" , self . peb_address) )
126140 . field ( "PEB Offset" , & format_args ! ( "{:#x}" , self . peb_offset) )
127141 . field ( "Code Size" , & format_args ! ( "{:#x}" , self . code_size) )
@@ -181,6 +195,10 @@ impl Debug for SandboxMemoryLayout {
181195 "Guest User Stack Buffer Offset" ,
182196 & format_args ! ( "{:#x}" , self . guest_user_stack_buffer_offset) ,
183197 )
198+ . field (
199+ "Init Data Offset" ,
200+ & format_args ! ( "{:#x}" , self . init_data_offset) ,
201+ )
184202 . field (
185203 "Page Table Size" ,
186204 & format_args ! ( "{:#x}" , self . total_page_table_size) ,
@@ -231,6 +249,8 @@ impl SandboxMemoryLayout {
231249 code_size : usize ,
232250 stack_size : usize ,
233251 heap_size : usize ,
252+ init_data_size : usize ,
253+ init_data_permissions : Option < MemoryRegionFlags > ,
234254 ) -> Result < Self > {
235255 let total_page_table_size =
236256 Self :: get_total_page_table_size ( cfg, code_size, stack_size, heap_size) ;
@@ -275,6 +295,7 @@ impl SandboxMemoryLayout {
275295 let guest_user_stack_buffer_offset = guard_page_offset + PAGE_SIZE_USIZE ;
276296 // round up stack size to page size. This is needed for MemoryRegion
277297 let stack_size_rounded = round_up_to ( stack_size, PAGE_SIZE_USIZE ) ;
298+ let init_data_offset = guest_user_stack_buffer_offset + stack_size_rounded;
278299
279300 Ok ( Self {
280301 peb_offset,
@@ -299,6 +320,9 @@ impl SandboxMemoryLayout {
299320 guard_page_offset,
300321 total_page_table_size,
301322 guest_code_offset,
323+ init_data_offset,
324+ init_data_size,
325+ init_data_permissions,
302326 } )
303327 }
304328
@@ -444,7 +468,7 @@ impl SandboxMemoryLayout {
444468 /// layout.
445469 #[ instrument( skip_all, parent = Span :: current( ) , level= "Trace" ) ]
446470 fn get_unaligned_memory_size ( & self ) -> usize {
447- self . get_top_of_user_stack_offset ( ) + self . stack_size
471+ self . init_data_offset + self . init_data_size
448472 }
449473
450474 /// get the code offset
@@ -682,12 +706,31 @@ impl SandboxMemoryLayout {
682706 }
683707
684708 // stack
685- let final_offset = builder. push_page_aligned (
709+ let init_data_offset = builder. push_page_aligned (
686710 self . get_guest_stack_size ( ) ,
687711 MemoryRegionFlags :: READ | MemoryRegionFlags :: WRITE ,
688712 Stack ,
689713 ) ;
690714
715+ let expected_init_data_offset = TryInto :: < usize > :: try_into ( self . init_data_offset ) ?;
716+
717+ if init_data_offset != expected_init_data_offset {
718+ return Err ( new_error ! (
719+ "Init Data offset does not match expected Init Data offset expected: {}, actual: {}" ,
720+ expected_init_data_offset,
721+ init_data_offset
722+ ) ) ;
723+ }
724+
725+ let final_offset = if self . init_data_size > 0 {
726+ let mem_flags = self
727+ . init_data_permissions
728+ . unwrap_or ( DEFAULT_GUEST_BLOB_MEM_FLAGS ) ;
729+ builder. push_page_aligned ( self . init_data_size , mem_flags, InitData )
730+ } else {
731+ init_data_offset
732+ } ;
733+
691734 let expected_final_offset = TryInto :: < usize > :: try_into ( self . get_memory_size ( ) ?) ?;
692735
693736 if final_offset != expected_final_offset {
@@ -701,6 +744,16 @@ impl SandboxMemoryLayout {
701744 Ok ( builder. build ( ) )
702745 }
703746
747+ #[ instrument( err( Debug ) , skip_all, parent = Span :: current( ) , level= "Trace" ) ]
748+ pub ( crate ) fn write_init_data (
749+ & self ,
750+ shared_mem : & mut ExclusiveSharedMemory ,
751+ bytes : & [ u8 ] ,
752+ ) -> Result < ( ) > {
753+ shared_mem. copy_from_slice ( bytes, self . init_data_offset ) ?;
754+ Ok ( ( ) )
755+ }
756+
704757 /// Write the finished memory layout to `shared_mem` and return
705758 /// `Ok` if successful.
706759 ///
@@ -869,7 +922,8 @@ mod tests {
869922 #[ test]
870923 fn test_get_memory_size ( ) {
871924 let sbox_cfg = SandboxConfiguration :: default ( ) ;
872- let sbox_mem_layout = SandboxMemoryLayout :: new ( sbox_cfg, 4096 , 2048 , 4096 ) . unwrap ( ) ;
925+ let sbox_mem_layout =
926+ SandboxMemoryLayout :: new ( sbox_cfg, 4096 , 2048 , 4096 , 0 , None ) . unwrap ( ) ;
873927 assert_eq ! (
874928 sbox_mem_layout. get_memory_size( ) . unwrap( ) ,
875929 get_expected_memory_size( & sbox_mem_layout)
0 commit comments