diff --git a/kernel/src/arch/riscv/vspace/kmem.rs b/kernel/src/arch/riscv/vspace/kmem.rs
index 27a8856..03dddd0 100644
--- a/kernel/src/arch/riscv/vspace/kmem.rs
+++ b/kernel/src/arch/riscv/vspace/kmem.rs
@@ -115,7 +115,7 @@ pub unsafe fn setup_kernel_paging() {
                 let start = concat_idents!($section, _START).as_virt_addr().align_down(PAGE_SIZE);
                 let end   = concat_idents!($section, _END).as_virt_addr().align_up(PAGE_SIZE);
                 debug!("Mapping section {}:\t[{:?}, {:?}]", stringify!($section), start, end);
-                map_range(&mut kernel_pt, start, kernel_virt_to_phys(start), (end - start).as_usize(), $attr);
+                map_range(alloc_page, &mut kernel_pt, start, kernel_virt_to_phys(start), (end - start).as_usize(), $attr);
             )+
         };
     }
@@ -130,6 +130,7 @@ pub unsafe fn setup_kernel_paging() {
     // map 4 GiB physical memory
     // TODO: walk fdt to get all memory region? put it in driver init
     map_range(
+        alloc_page,
         &mut kernel_pt,
         mmap_phys_to_virt(PhysAddr(0)),
         PhysAddr(0),
@@ -163,7 +164,7 @@ pub fn is_kernel_pagetable_allocated() -> bool {
 unsafe fn map_kernel(from: VirtAddr, to: PhysAddr, size: usize, attr: MapAttr) {
     let mut guard = KERNEL_PAGETABLE.lock();
     let pt = guard.as_mut().unwrap();
-    map_range(pt, from, to, size, attr);
+    map_range(alloc_page, pt, from, to, size, attr);
 }
 
 pub fn copy_kernel_pagetable(root: &CapEntry) {
diff --git a/kernel/src/arch/riscv/vspace/utils.rs b/kernel/src/arch/riscv/vspace/utils.rs
index 48e744e..50f616a 100644
--- a/kernel/src/arch/riscv/vspace/utils.rs
+++ b/kernel/src/arch/riscv/vspace/utils.rs
@@ -28,7 +28,14 @@ pub fn alloc_page() -> PhysAddr {
     addr
 }
 
-pub fn map_range(pt: &mut Table<Level0>, from: VirtAddr, to: PhysAddr, size: usize, attr: MapAttr) {
+pub fn map_range(
+    mut alloc: impl FnMut() -> PhysAddr,
+    pt: &mut Table<Level0>,
+    from: VirtAddr,
+    to: PhysAddr,
+    size: usize,
+    attr: MapAttr,
+) {
     let mut virt_start = from;
     let mut phys_start = to;
     let virt_end = from + size;
@@ -56,7 +63,7 @@ pub fn map_range(pt: &mut Table<Level0>, from: VirtAddr, to: PhysAddr, size: usi
             unsafe {
                 Table::<$level>::new(mmap_phys_to_virt(
                     $pt.lookup_mut($addr).map(|e| e.paddr()).unwrap_or_else(|| {
-                        let page = alloc_page();
+                        let page = alloc();
                         let addr = $addr.align_down(<$level as TableLevel>::PREVIOUS::LEVEL_SIZE);
                         debug!("Creating new pagetable:\t{:?} -> {:?}", addr, page);
                         $pt.map(addr, page, MapAttr::PAGE_TABLE).unwrap();