Add page tfault handler to support lazy mapping

src/main.rs

@@ -6,20 +6,14 @@ extern crate log;
 extern crate alloc;
 extern crate axstd;
 
+mod mm;
 mod syscall;
 mod task;
 
 use alloc::sync::Arc;
 
 use axhal::arch::UspaceContext;
-use axhal::mem::virt_to_phys;
-use axhal::paging::MappingFlags;
-use axmm::AddrSpace;
 use axsync::Mutex;
-use axtask::{AxTaskRef, TaskExtRef, TaskInner};
-use memory_addr::VirtAddr;
-
-use self::task::TaskExt;
 
 const USER_STACK_SIZE: usize = 4096;
 const KERNEL_STACK_SIZE: usize = 0x40000; // 256 KiB
@@ -30,6 +24,7 @@ fn app_main(arg0: usize) {
             "2:",
             "int3",
             "mov rax, r12",
+            "push rax",
             "syscall",
             "add r12, 1",
             "jmp 2b",
@@ -42,61 +37,13 @@ fn app_main(arg0: usize) {
     }
 }
 
-fn spawn_user_task(aspace: Arc<Mutex<AddrSpace>>, uctx: UspaceContext) -> AxTaskRef {
-    let mut task = TaskInner::new(
-        || {
-            let curr = axtask::current();
-            let kstack_top = curr.kernel_stack_top().unwrap();
-            info!(
-                "Enter user space: entry={:#x}, ustack={:#x}, kstack={:#x}",
-                curr.task_ext().uctx.get_ip(),
-                curr.task_ext().uctx.get_sp(),
-                kstack_top,
-            );
-            unsafe { curr.task_ext().uctx.enter_uspace(kstack_top) };
-        },
-        "".into(),
-        KERNEL_STACK_SIZE,
-    );
-    task.ctx_mut()
-        .set_page_table_root(aspace.lock().page_table_root());
-    task.init_task_ext(TaskExt::new(uctx, aspace));
-    axtask::spawn_task(task)
-}
-
 #[no_mangle]
-fn main() -> ! {
-    let entry = VirtAddr::from(app_main as usize);
-    let entry_paddr_align = virt_to_phys(entry.align_down_4k());
-    let entry_vaddr_align = VirtAddr::from(0x1000);
-    let entry_vaddr = entry_vaddr_align + entry.align_offset_4k();
-
-    let mut uspace = axmm::new_user_aspace().unwrap();
-    let ustack_top = uspace.end();
-    let ustack_vaddr = ustack_top - USER_STACK_SIZE;
-    uspace
-        .map_linear(
-            entry_vaddr_align,
-            entry_paddr_align,
-            4096,
-            MappingFlags::READ | MappingFlags::EXECUTE | MappingFlags::USER,
-        )
-        .unwrap();
-    uspace
-        .map_alloc(
-            ustack_vaddr,
-            4096,
-            MappingFlags::READ | MappingFlags::WRITE | MappingFlags::USER,
-            false,
-        )
-        .unwrap();
-
-    info!("New user address space: {:#x?}", uspace);
-    spawn_user_task(
+fn main() {
+    let (entry_vaddr, ustack_top, uspace) = mm::load_user_app(app_main).unwrap();
+    let user_task = task::spawn_user_task(
         Arc::new(Mutex::new(uspace)),
         UspaceContext::new(entry_vaddr.into(), ustack_top, 2333),
     );
-
-    axtask::WaitQueue::new().wait();
-    unreachable!()
+    let exit_code = user_task.join();
+    info!("User task exited with code: {:?}", exit_code);
 }

src/mm.rs

@@ -0,0 +1,51 @@
+use axerrno::AxResult;
+use memory_addr::VirtAddr;
+
+use axhal::mem::virt_to_phys;
+use axhal::paging::MappingFlags;
+use axhal::trap::{register_trap_handler, PAGE_FAULT};
+use axmm::AddrSpace;
+use axtask::TaskExtRef;
+
+pub fn load_user_app(entry_fn: fn(usize)) -> AxResult<(VirtAddr, VirtAddr, AddrSpace)> {
+    let entry_fn_kvaddr = VirtAddr::from(entry_fn as usize);
+    let load_vaddr = VirtAddr::from(0x1000);
+    let load_paddr = virt_to_phys(entry_fn_kvaddr.align_down_4k());
+    let entry_vaddr = load_vaddr + entry_fn_kvaddr.align_offset_4k();
+
+    let mut uspace = axmm::new_user_aspace()?;
+    let ustack_top = uspace.end();
+    let ustack_vaddr = ustack_top - crate::USER_STACK_SIZE;
+    uspace.map_linear(
+        load_vaddr,
+        load_paddr,
+        4096,
+        MappingFlags::READ | MappingFlags::EXECUTE | MappingFlags::USER,
+    )?;
+    uspace.map_alloc(
+        ustack_vaddr,
+        crate::USER_STACK_SIZE,
+        MappingFlags::READ | MappingFlags::WRITE | MappingFlags::USER,
+        false,
+    )?;
+    info!("New user address space: {:#x?}", uspace);
+    Ok((entry_vaddr, ustack_top, uspace))
+}
+
+#[register_trap_handler(PAGE_FAULT)]
+fn handle_page_fault(vaddr: VirtAddr, access_flags: MappingFlags, is_user: bool) -> bool {
+    if is_user {
+        if !axtask::current()
+            .task_ext()
+            .aspace
+            .lock()
+            .handle_page_fault(vaddr, access_flags)
+        {
+            warn!("{}: segmentation fault, exit!", axtask::current().id_name());
+            axtask::exit(-1);
+        }
+        true
+    } else {
+        false
+    }
+}

src/syscall.rs

@@ -20,7 +20,7 @@ fn sys_clone(tf: &TrapFrame, newsp: usize) -> usize {
     let mut uctx = UspaceContext::from(tf);
     uctx.set_sp(newsp);
     uctx.set_retval(0);
-    let new_task = super::spawn_user_task(aspace, uctx);
+    let new_task = crate::task::spawn_user_task(aspace, uctx);
     new_task.id().as_u64() as usize
 }
 

src/task.rs

@@ -3,6 +3,7 @@ use alloc::sync::Arc;
 use axhal::arch::UspaceContext;
 use axmm::AddrSpace;
 use axsync::Mutex;
+use axtask::{AxTaskRef, TaskExtRef, TaskInner};
 
 /// Task extended data for the monolithic kernel.
 pub struct TaskExt {
@@ -25,3 +26,25 @@ impl TaskExt {
 }
 
 axtask::def_task_ext!(TaskExt);
+
+pub fn spawn_user_task(aspace: Arc<Mutex<AddrSpace>>, uctx: UspaceContext) -> AxTaskRef {
+    let mut task = TaskInner::new(
+        || {
+            let curr = axtask::current();
+            let kstack_top = curr.kernel_stack_top().unwrap();
+            info!(
+                "Enter user space: entry={:#x}, ustack={:#x}, kstack={:#x}",
+                curr.task_ext().uctx.get_ip(),
+                curr.task_ext().uctx.get_sp(),
+                kstack_top,
+            );
+            unsafe { curr.task_ext().uctx.enter_uspace(kstack_top) };
+        },
+        "userboot".into(),
+        crate::KERNEL_STACK_SIZE,
+    );
+    task.ctx_mut()
+        .set_page_table_root(aspace.lock().page_table_root());
+    task.init_task_ext(TaskExt::new(uctx, aspace));
+    axtask::spawn_task(task)
+}