early alloc stack and refactor brk (#1066)

* early init stack and refactor brk

* cleanup: remove unused vars, fix stale comments, and apply formatting

* quick fix

* address review comments: use u64::MAX for fd=-1 and add mmap flag comments

Author: qianxichen233

src/cage/src/memory/memory.rs

@@ -142,16 +142,16 @@ pub fn fork_vmmap(parent_cageid: u64, child_cageid: u64) {
     // update program break for child
     drop(child_vmmap);
     let mut child_vmmap = child_cage.vmmap.write();
-    child_vmmap.set_program_break(parent_vmmap.program_break);
+    child_vmmap.set_heap_start(parent_vmmap.heap_start);
 }
 
 // set the wasm linear memory base address to vmmap
-pub fn init_vmmap(cageid: u64, base_address: usize, program_break: Option<u32>) {
+pub fn init_vmmap(cageid: u64, base_address: usize, heap_start: Option<u32>) {
     let cage = get_cage(cageid).unwrap();
     let mut vmmap = cage.vmmap.write();
     vmmap.set_base_address(base_address);
-    if program_break.is_some() {
-        vmmap.set_program_break(program_break.unwrap());
+    if heap_start.is_some() {
+        vmmap.set_heap_start(heap_start.unwrap());
     }
 }
 

src/cage/src/memory/vmmap.rs

@@ -255,7 +255,7 @@ pub struct Vmmap {
 
     pub start_address: u32, // start address of valid vmmap address range
     pub end_address: u32,   // end address of valid vmmap address range
-    pub program_break: u32, // program break (i.e. heap bottom) of the memory
+    pub heap_start: u32,    // start of heap memory
 }
 
 #[allow(dead_code)]
@@ -269,7 +269,7 @@ impl Vmmap {
             base_address: None,
             start_address: 0,
             end_address: DEFAULT_VMMAP_SIZE,
-            program_break: 0,
+            heap_start: 0,
         }
     }
 
@@ -284,7 +284,7 @@ impl Vmmap {
         self.base_address = None;
         self.start_address = 0;
         self.end_address = DEFAULT_VMMAP_SIZE;
-        self.program_break = 0;
+        self.heap_start = 0;
     }
 
     /// Rounds up a page number to the nearest multiple of pages_per_map
@@ -320,12 +320,12 @@ impl Vmmap {
         self.base_address = Some(base_address);
     }
 
-    /// Sets the program break for the memory
+    /// Sets the heap start page number for the memory
     ///
     /// Arguments:
-    /// - program_break: The program break to set
-    pub fn set_program_break(&mut self, program_break: u32) {
-        self.program_break = program_break;
+    /// - heap_start: The page number at which the heap begins
+    pub fn set_heap_start(&mut self, heap_start: u32) {
+        self.heap_start = heap_start;
     }
 
     /// Converts a user address to a system address

src/glibc/csu/libc-tls.c

@@ -169,36 +169,7 @@ __libc_setup_tls (void)
   tlsblock = (void *) (((uintptr_t) tlsblock + max_align - 1)
 		       & ~(max_align - 1));
 
-  /* Initialize the dtv.  [0] is the length, [1] the generation counter.  */
-  _dl_static_dtv[0].counter = (sizeof (_dl_static_dtv) / sizeof (_dl_static_dtv[0])) - 2;
-  // _dl_static_dtv[1].counter = 0;		would be needed if not already done
-
-  /* Initialize the TLS block.  */
-#if TLS_TCB_AT_TP
-  _dl_static_dtv[2].pointer.val = ((char *) tlsblock + tcb_offset
-			       - roundup (memsz, align ?: 1));
-  main_map->l_tls_offset = roundup (memsz, align ?: 1);
-#elif TLS_DTV_AT_TP
-  _dl_static_dtv[2].pointer.val = (char *) tlsblock + tcb_offset;
-  main_map->l_tls_offset = tcb_offset;
-#else
-# error "Either TLS_TCB_AT_TP or TLS_DTV_AT_TP must be defined"
-#endif
-  _dl_static_dtv[2].pointer.to_free = NULL;
-  /* sbrk gives us zero'd memory, so we don't need to clear the remainder.  */
-  memcpy (_dl_static_dtv[2].pointer.val, initimage, filesz);
-
-  /* Install the pointer to the dtv.  */
-
-  /* Initialize the thread pointer.  */
-#if TLS_TCB_AT_TP
-  INSTALL_DTV ((char *) tlsblock + tcb_offset, _dl_static_dtv);
-
-  // call_tls_init_tp ((char *) tlsblock + tcb_offset);
-#elif TLS_DTV_AT_TP
-  INSTALL_DTV (tlsblock, _dl_static_dtv);
-  // call_tls_init_tp (tlsblock);
-#endif
+  // lind-wasm: remove DTV related logic since lind do not use DTV
 
   /* Update the executable's link map with enough information to make
      the TLS routines happy.  */

src/lind-boot/src/lind_wasmtime/execute.rs

@@ -22,7 +22,8 @@ use wasmtime_lind_3i::{VmCtxWrapper, init_vmctx_pool, rm_vmctx, set_vmctx, set_v
 use wasmtime_lind_common::LindEnviron;
 use wasmtime_lind_dylink::DynamicLoader;
 use wasmtime_lind_multi_process::{
-    CAGE_START_ID, LindCtx, THREAD_START_ID, attach_shared_memory, get_memory_base,
+    CAGE_START_ID, LindCtx, THREAD_START_ID, attach_shared_memory, early_init_stack,
+    get_memory_base,
 };
 use wasmtime_lind_utils::symbol_table::SymbolMap;
 use wasmtime_lind_utils::{LindCageManager, LindGOT};
@@ -243,6 +244,13 @@ pub fn execute_with_lind(
     if dylink_metadata.dylink_enabled {
         let lind_got = dylink_metadata.got.as_ref().unwrap();
 
+        early_init_stack(
+            cageid,
+            GUARD_SIZE as i32,
+            (GUARD_SIZE + DEFAULT_STACKSIZE) as i32,
+        )
+        .unwrap();
+
         // For each module (including the main module),
         // register its GOT imports with the shared LindGOT instance.
         //

src/rawposix/src/fs_calls.rs

@@ -1199,19 +1199,42 @@ pub extern "C" fn brk_syscall(
     let cage = get_cage(cageid).unwrap();
 
     let mut vmmap = cage.vmmap.write();
-    let heap = vmmap.find_page(HEAP_ENTRY_INDEX).unwrap().clone();
+    let heap_opt = vmmap.find_page(vmmap.heap_start);
 
-    assert!(heap.npages == vmmap.program_break);
+    let heap = if heap_opt.is_none() {
+        // if heap page is not found, create an empty heap entry with 0 size
+        cage::VmmapEntry::new(
+            vmmap.heap_start,
+            0,
+            (PROT_READ | PROT_WRITE),
+            (PROT_READ | PROT_WRITE),
+            (MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED) as i32,
+            false,
+            0,
+            0,
+            cageid,
+            MemoryBackingType::Anonymous,
+        )
+    } else {
+        heap_opt.unwrap().clone()
+    };
+
+    assert!(heap.page_num == vmmap.heap_start);
+
+    let old_brk_page = heap.page_num + heap.npages;
 
     // passing 0 to brk will always return the current brk
     if brk == 0 {
-        return (PAGESIZE * heap.npages) as i32;
+        return (PAGESIZE * old_brk_page) as i32;
     }
-
-    let old_brk_page = heap.npages;
     // round up the break to multiple of pages
     let brk_page = (round_up_page(brk as u64) >> PAGESHIFT) as u32;
 
+    // shrink heap below heap start is not allowed
+    if brk_page < vmmap.heap_start {
+        return syscall_error(Errno::ENOMEM, "brk", "no memory");
+    }
+
     // if we are incrementing program break, we need to check if we have enough space
     if brk_page > old_brk_page {
         if vmmap.check_existing_mapping(old_brk_page, brk_page - old_brk_page, 0) {
@@ -1220,12 +1243,12 @@ pub extern "C" fn brk_syscall(
     }
 
     // remove the old entries since new entry is overlapping with it.
-    vmmap.remove_entry(0, old_brk_page);
+    vmmap.remove_entry(heap.page_num, heap.npages);
 
     // update vmmap entry
     vmmap.add_entry_with_overwrite(
-        0,
-        brk_page,
+        heap.page_num,
+        brk_page - heap.page_num,
         heap.prot,
         heap.maxprot,
         heap.flags,
@@ -1241,8 +1264,6 @@ pub extern "C" fn brk_syscall(
     let new_heap_end_usr = (brk_page * PAGESIZE) as u32;
     let new_heap_end_sys = vmmap.user_to_sys(new_heap_end_usr) as *mut u8;
 
-    vmmap.set_program_break(brk_page);
-
     drop(vmmap);
 
     // if new brk is larger than old brk

src/sysdefs/src/constants/lind_platform_const.rs

@@ -86,7 +86,7 @@ pub const THREEI_CAGEID: u64 = 999999;
 /// Default stack size assigned to each cage
 pub const DEFAULT_STACKSIZE: u32 = 8388608; // 8 MB
 /// Size of guard pages
-pub const GUARD_SIZE: u32 = 1024; // 1 KB
+pub const GUARD_SIZE: u32 = 4096; // 4 KB
 
 /// The starting index for function tables of wasm modules in Lind.
 /// function index of 1 must be reserved for SIG_IGN constant for signal handling

src/wasmtime/crates/lind-multi-process/src/lib.rs

@@ -7,7 +7,7 @@ use std::ffi::c_void;
 use std::ptr::NonNull;
 use sysdefs::constants::lind_platform_const::{UNUSED_ARG, UNUSED_ID, UNUSED_NAME};
 use sysdefs::constants::syscall_const::{EXEC_SYSCALL, EXIT_SYSCALL, FORK_SYSCALL};
-use sysdefs::constants::{Errno, MAX_SHEBANG_DEPTH};
+use sysdefs::constants::{Errno, MAX_SHEBANG_DEPTH, MMAP_SYSCALL};
 use sysdefs::logging::lind_debug_panic;
 use sysdefs::{constants::sys_const, data::sys_struct};
 use threei::{threei::make_syscall, threei_const};
@@ -2212,6 +2212,40 @@ pub fn attach_shared_memory<
     Err(anyhow!("Main Module does not contain a shared memory"))
 }
 
+pub fn early_init_stack(cageid: u64, stack_start: i32, stack_end: i32) -> Result<()> {
+    // TODO: currently we explicitly allocate first guard page (0-stack_start)
+    // due to a known issue. This should be fixed in the future and only allocate
+    // the actual stack space (stack_start-stack_end)
+
+    let ret = make_syscall(
+        cageid,                // self cageid
+        (MMAP_SYSCALL) as u64, // syscall num
+        0, // since wasmtime operates with lower level memory, it always interacts with underlying os
+        cageid, // target cageid (should be same)
+        0, // map from address 0 (includes the leading guard page per the TODO above)
+        cageid,
+        stack_end as u64, // length: guard page + stack
+        cageid,
+        // PROT_READ | PROT_WRITE: stack needs both read and write access
+        (typemap::PROT_READ | typemap::PROT_WRITE) as u64,
+        cageid,
+        // MAP_PRIVATE: changes are not shared; MAP_ANONYMOUS: not file-backed;
+        // MAP_FIXED: must map at the exact address (address 0 for guard page + stack region)
+        (typemap::MAP_PRIVATE | typemap::MAP_ANONYMOUS | typemap::MAP_FIXED) as u64,
+        cageid,
+        u64::MAX, // fd: -1 (required for MAP_ANONYMOUS mappings)
+        cageid,
+        0,
+        cageid,
+    );
+
+    if ret < 0 {
+        return Err(anyhow!("failed to allocate stack"));
+    }
+
+    Ok(())
+}
+
 // check if the module has the necessary exported Asyncify functions
 fn support_asyncify(module: &Module) -> bool {
     module.get_export(ASYNCIFY_START_UNWIND).is_some()

src/wasmtime/crates/wasmtime/src/runtime/instance.rs

@@ -340,7 +340,7 @@ impl Instance {
                 drop(memory_iter);
                 let memory_base = memory.data_ptr(&mut *store) as usize;
 
-                let required_memory_size = if dylink_enabled {
+                let (start_addr, required_memory_size) = if dylink_enabled {
                     // for dynamic builds, we manually calcuate the initial memory region
                     // which is stack size + data region size
                     let module_meminfo = module.dylink_meminfo().unwrap();
@@ -359,7 +359,7 @@ impl Instance {
                         rounded_stack_size, rounded_data_size
                     );
 
-                    rounded_stack_size + rounded_data_size
+                    (rounded_stack_size, rounded_data_size)
                 } else {
                     // retrieve the initial memory size for static module
                     let plans = module.compiled_module().module().memory_plans.clone();
@@ -377,12 +377,15 @@ impl Instance {
 
                     let minimal_size = minimal_pages << PAGESHIFT;
 
-                    minimal_size
-                        .try_into()
-                        .expect("allocated memory is larger than 4GB")
+                    (
+                        0,
+                        minimal_size
+                            .try_into()
+                            .expect("allocated memory is larger than 4GB"),
+                    )
                 };
 
-                let required_memory_page = required_memory_size >> PAGESHIFT;
+                let required_memory_page = (start_addr + required_memory_size) >> PAGESHIFT;
 
                 init_vmmap(cageid, memory_base, Some(required_memory_page));
                 // Allocated memory should include stack AND constant data region
@@ -394,7 +397,7 @@ impl Instance {
                     (MMAP_SYSCALL) as u64, // syscall num
                     0, // since wasmtime operates with lower level memory, it always interacts with underlying os
                     cageid, // target cageid (should be same)
-                    0, // the first memory region starts from 0
+                    start_addr as u64,
                     cageid,
                     required_memory_size as u64,
                     cageid,