fix(mm): add VmIo retry logic for lazy allocation

api/src/io.rs

@@ -1,8 +1,11 @@
 use core::mem::{self, MaybeUninit};
 
 use axerrno::{AxError, AxResult};
+use axhal::paging::MappingFlags;
 use axio::{Buf, BufMut, Read, Write};
 use bytemuck::AnyBitPattern;
+use memory_addr::{MemoryAddr, VirtAddr};
+use starry_core::task::AsThread;
 use starry_vm::{VmPtr, vm_read_slice, vm_write_slice};
 
 #[repr(C)]
@@ -145,10 +148,28 @@ impl Write for IoVectorBufIo {
             if len == 0 {
                 break;
             }
-            vm_write_slice(
-                iov.iov_base.wrapping_add(self.offset),
-                &buf[count..count + len],
-            )?;
+
+            // Pre-populate pages to support lazy allocation
+            let dst_ptr = iov.iov_base.wrapping_add(self.offset);
+            let start_addr = VirtAddr::from_ptr_of(dst_ptr);
+            {
+                let curr = axtask::current();
+                let mut aspace = curr.as_thread().proc_data.aspace.lock();
+                let page_start = start_addr.align_down_4k();
+                let page_end = (start_addr + len).align_up_4k();
+                if aspace
+                    .populate_area(
+                        page_start,
+                        page_end - page_start,
+                        MappingFlags::READ | MappingFlags::WRITE | MappingFlags::USER,
+                    )
+                    .is_err()
+                {
+                    return Err(AxError::BadAddress);
+                }
+            }
+
+            vm_write_slice(dst_ptr, &buf[count..count + len])?;
             self.offset += len;
             self.inner.len -= len;
             count += len;

api/src/mm.rs

@@ -330,6 +330,22 @@ impl Write for VmBytesMut {
     /// Writes bytes from the provided buffer into the VM's memory.
     fn write(&mut self, buf: &[u8]) -> axio::Result<usize> {
         let len = self.len.min(buf.len());
+        if len == 0 {
+            return Ok(0);
+        }
+
+        let start_addr = VirtAddr::from_ptr_of(self.ptr);
+        // Ensure pages are allocated (for lazy-allocated stack/heap regions)
+        if check_region(
+            start_addr,
+            Layout::from_size_align(len, 1).unwrap(),
+            MappingFlags::READ | MappingFlags::WRITE | MappingFlags::USER,
+        )
+        .is_err()
+        {
+            return Err(axio::Error::InvalidData);
+        }
+
         vm_write_slice(self.ptr, &buf[..len])?;
         self.ptr = self.ptr.wrapping_add(len);
         self.len -= len;

core/src/mm.rs

@@ -3,7 +3,7 @@
 use alloc::{borrow::ToOwned, string::String, vec, vec::Vec};
 use core::{
     ffi::CStr,
-    hint::unlikely,
+    hint::{likely, unlikely},
     iter,
     mem::MaybeUninit,
     sync::atomic::{AtomicBool, Ordering},
@@ -19,6 +19,7 @@ use axhal::{
 };
 use axmm::{AddrSpace, backend::Backend};
 use axsync::Mutex;
+use axtask::current;
 use extern_trait::extern_trait;
 use kernel_elf_parser::{AuxEntry, ELFHeaders, ELFHeadersBuilder, ELFParser, app_stack_region};
 use kernel_guard::IrqSave;
@@ -27,7 +28,10 @@ use ouroboros::self_referencing;
 use starry_vm::{VmError, VmIo, VmResult};
 use uluru::LRUCache;
 
-use crate::config::{USER_SPACE_BASE, USER_SPACE_SIZE};
+use crate::{
+    config::{USER_SPACE_BASE, USER_SPACE_SIZE},
+    task::AsThread,
+};
 
 /// Creates a new empty user address space.
 pub fn new_user_aspace_empty() -> AxResult<AddrSpace> {
@@ -122,11 +126,14 @@ fn map_elf<'a>(
             ph.offset,
             Some(ph.offset + ph.file_size),
         );
+
+        let seg_flags = mapping_flags(ph.flags);
+
         uspace.map(
             seg_start.align_down_4k(),
             seg_align_size,
-            mapping_flags(ph.flags),
-            false,
+            seg_flags,
+            true, // Populate ELF segments immediately
             backend,
         )?;
 
@@ -201,48 +208,51 @@ impl ElfLoader {
         uspace.clear();
         map_trampoline(uspace)?;
 
-        let entry = self.0.front().unwrap();
-        let ldso = if let Some(header) = entry
-            .borrow_elf()
-            .ph
-            .iter()
-            .find(|ph| ph.get_type() == Ok(xmas_elf::program::Type::Interp))
-        {
-            let cache = entry.borrow_cache();
-            let mut data = vec![0; header.file_size as usize];
-            let read = cache.read_at(&mut data.as_mut_slice(), header.offset)?;
-            assert_eq!(data.len(), read);
-
-            let ldso = CStr::from_bytes_with_nul(&data)
-                .ok()
-                .and_then(|cstr| cstr.to_str().ok())
-                .ok_or(AxError::InvalidInput)?;
-            debug!("Loading dynamic linker: {ldso}");
-            Some(ldso.to_owned())
-        } else {
-            None
+        let mut ldso_path = None;
+        let entry_ptr = {
+            let entry = self.0.front().unwrap();
+            if let Some(header) = entry
+                .borrow_elf()
+                .ph
+                .iter()
+                .find(|ph| ph.get_type() == Ok(xmas_elf::program::Type::Interp))
+            {
+                let cache = entry.borrow_cache();
+                let mut data = vec![0; header.file_size as usize];
+                let read = cache.read_at(&mut data.as_mut_slice(), header.offset)?;
+                assert_eq!(data.len(), read);
+
+                let ldso = CStr::from_bytes_with_nul(&data)
+                    .ok()
+                    .and_then(|cstr| cstr.to_str().ok())
+                    .ok_or(AxError::InvalidInput)?;
+                debug!("Loading dynamic linker: {ldso}");
+                ldso_path = Some(ldso.to_owned());
+            }
+            entry as *const ElfCacheEntry
         };
 
-        let (elf, ldso) = if let Some(ldso) = ldso {
+        let mut ldso_entry_ptr = None;
+        if let Some(ldso) = ldso_path.as_ref() {
             let loc = FS_CONTEXT.lock().resolve(ldso)?;
             if !self.0.touch(|e| e.borrow_cache().location().ptr_eq(&loc)) {
                 let e = ElfCacheEntry::load(loc)?.map_err(|_| AxError::InvalidInput)?;
                 self.0.insert(e);
             }
+            ldso_entry_ptr = Some(self.0.front().unwrap() as *const ElfCacheEntry);
+        }
 
-            let mut iter = self.0.iter();
-            let ldso = iter.next().unwrap();
-            let elf = iter.next().unwrap();
-            (elf, Some(ldso))
+        let elf = unsafe { map_elf(uspace, crate::config::USER_SPACE_BASE, &*entry_ptr)? };
+        let ldso = if let (Some(ldso_entry_ptr), Some(_ldso_path)) =
+            (ldso_entry_ptr, ldso_path.as_ref())
+        {
+            Some(map_elf(uspace, crate::config::USER_INTERP_BASE, unsafe {
+                &*ldso_entry_ptr
+            })?)
         } else {
-            (entry, None)
+            None
         };
 
-        let elf = map_elf(uspace, crate::config::USER_SPACE_BASE, elf)?;
-        let ldso = ldso
-            .map(|elf| map_elf(uspace, crate::config::USER_INTERP_BASE, elf))
-            .transpose()?;
-
         let entry = VirtAddr::from_usize(
             ldso.as_ref()
                 .map_or_else(|| elf.entry(), |ldso| ldso.entry()),
@@ -388,25 +398,91 @@ unsafe impl VmIo for Vm {
 
     fn read(&mut self, start: usize, buf: &mut [MaybeUninit<u8>]) -> VmResult {
         check_access(start, buf.len())?;
+
+        // First attempt to read
         let failed_at = access_user_memory(|| unsafe {
             user_copy(buf.as_mut_ptr() as *mut _, start as _, buf.len())
         });
-        if unlikely(failed_at != 0) {
-            Err(VmError::AccessDenied)
-        } else {
-            Ok(())
+
+        if likely(failed_at == 0) {
+            return Ok(());
         }
+
+        // Read failed, try to populate pages and retry
+        let fail_offset = buf.len() - failed_at;
+        let fail_addr = VirtAddr::from(start + fail_offset);
+
+        if let Some(thr) = current().try_as_thread() {
+            let page_start = fail_addr.align_down_4k();
+            let page_end = VirtAddr::from(start + buf.len()).align_up_4k();
+            let populate_size = page_end - page_start;
+
+            let populate_result = {
+                let mut aspace = thr.proc_data.aspace.lock();
+                aspace.populate_area(
+                    page_start,
+                    populate_size,
+                    MappingFlags::READ | MappingFlags::USER,
+                )
+            };
+
+            if populate_result.is_ok() {
+                // Retry read
+                let failed_at2 = access_user_memory(|| unsafe {
+                    user_copy(buf.as_mut_ptr() as *mut _, start as _, buf.len())
+                });
+
+                if failed_at2 == 0 {
+                    return Ok(());
+                }
+            }
+        }
+
+        Err(VmError::AccessDenied)
     }
 
     fn write(&mut self, start: usize, buf: &[u8]) -> VmResult {
         check_access(start, buf.len())?;
+
+        // First attempt to write
         let failed_at = access_user_memory(|| unsafe {
             user_copy(start as _, buf.as_ptr() as *const _, buf.len())
         });
-        if unlikely(failed_at != 0) {
-            Err(VmError::AccessDenied)
-        } else {
-            Ok(())
+
+        if likely(failed_at == 0) {
+            return Ok(());
         }
+
+        // Write failed, try to populate pages and retry
+        let fail_offset = buf.len() - failed_at;
+        let fail_addr = VirtAddr::from(start + fail_offset);
+
+        if let Some(thr) = current().try_as_thread() {
+            let page_start = fail_addr.align_down_4k();
+            let page_end = VirtAddr::from(start + buf.len()).align_up_4k();
+            let populate_size = page_end - page_start;
+
+            let populate_result = {
+                let mut aspace = thr.proc_data.aspace.lock();
+                aspace.populate_area(
+                    page_start,
+                    populate_size,
+                    MappingFlags::READ | MappingFlags::WRITE | MappingFlags::USER,
+                )
+            };
+
+            if populate_result.is_ok() {
+                // Retry write
+                let failed_at2 = access_user_memory(|| unsafe {
+                    user_copy(start as _, buf.as_ptr() as *const _, buf.len())
+                });
+
+                if failed_at2 == 0 {
+                    return Ok(());
+                }
+            }
+        }
+
+        Err(VmError::AccessDenied)
     }
 }