diff --git a/api/src/io.rs b/api/src/io.rs index c7bd4eaf..10af38cc 100644 --- a/api/src/io.rs +++ b/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; diff --git a/api/src/mm.rs b/api/src/mm.rs index d6a36cad..2259c93e 100644 --- a/api/src/mm.rs +++ b/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 { 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; diff --git a/core/src/mm.rs b/core/src/mm.rs index e27e893e..a38844ab 100644 --- a/core/src/mm.rs +++ b/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 { @@ -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]) -> 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) } }