fix: enhance memory management with improved page handling and address space management

arch/x86-family/mm/address_space.c

@@ -40,7 +40,7 @@
 
 
 
-static x86_page_t __mm_copy_data(x86_page_t* __s, size_t* size, bool on_demand, int level) {
+static x86_page_t __mm_copy_data(x86_page_t* __s, size_t* size, int level) {
 
     DEBUG_ASSERT(__s);
     DEBUG_ASSERT(size);
@@ -71,22 +71,12 @@ static x86_page_t __mm_copy_data(x86_page_t* __s, size_t* size, bool on_demand,
     *size += pagesize;
 
 
-#if defined(CONFIG_DEMAND_PAGING)
-    if (on_demand) {
+    uintptr_t page = __alloc_frame(pagesize, false);
 
-        return (*__s = (*__s & ~(X86_MMU_PG_RW | X86_MMU_PG_AP_TP_MASK)) | X86_MMU_PG_AP_TP_COW);
+    memcpy((void*)arch_vmm_p2v(page, ARCH_VMM_AREA_HEAP), (void*)arch_vmm_p2v(*__s & X86_MMU_ADDRESS_MASK, ARCH_VMM_AREA_HEAP), (size_t)pagesize);
 
-    } else
-#endif
-    {
-
-        uintptr_t page = __alloc_frame(pagesize, false);
-
-        memcpy((void*)arch_vmm_p2v(page, ARCH_VMM_AREA_HEAP), (void*)arch_vmm_p2v(*__s & X86_MMU_ADDRESS_MASK, ARCH_VMM_AREA_HEAP), (size_t)pagesize);
 
-
-        return page | X86_MMU_PG_AP_PFB | (*__s & ~X86_MMU_ADDRESS_MASK);
-    }
+    return page | X86_MMU_PG_AP_PFB | (*__s & ~X86_MMU_ADDRESS_MASK);
 }
 
 
@@ -107,13 +97,13 @@ static void __mm_copy_page(x86_page_t* __s, x86_page_t* __d, size_t* size, int l
         if (flags & ARCH_VMM_CLONE_USERSPACE) {
 
 
-            if ((*__s & X86_MMU_PG_AP_TP_MASK) == X86_MMU_PG_AP_TP_COW) {
+            if (((*__s & X86_MMU_PG_AP_TP_MASK) == X86_MMU_PG_AP_TP_COW) && ((*__s & X86_MMU_ADDRESS_MASK) == 0)) {
 
                 *__d = *__s;
 
             } else {
 
-                *__d = __mm_copy_data(__s, size, (flags & ARCH_VMM_CLONE_DEMAND), level);
+                *__d = __mm_copy_data(__s, size, level);
             }
         }
     }
@@ -231,10 +221,9 @@ static void __mm_free_table(uintptr_t __s, int level) {
 
             __mm_free_table(((uintptr_t)s[i]) & X86_MMU_ADDRESS_MASK, level - 1);
 
-            // FIXME: not safe to free frame here, as it may be used by system page tables
-            // // if(s[i] & X86_MMU_PT_AP_PFB) {
-            // //     __free_frame(((uintptr_t) s[i]) & X86_MMU_ADDRESS_MASK, X86_MMU_PAGESIZE);
-            // // }
+            if (s[i] & X86_MMU_PT_AP_PFB) {
+                __free_frame(((uintptr_t)s[i]) & X86_MMU_ADDRESS_MASK, X86_MMU_PAGESIZE);
+            }
         }
     }
 }
@@ -253,6 +242,7 @@ __returns_nonnull vmm_address_space_t* arch_vmm_create_address_space(vmm_address
 
 
     dest->pm = __alloc_frame(X86_MMU_PAGESIZE, true);
+    dest->flags = ARCH_VMM_ADDRESS_SPACE_OWNED;
 
 
 
@@ -305,7 +295,9 @@ void arch_vmm_free_address_space(vmm_address_space_t* space) {
     }
 
 
-    // TODO: free all mappings
+    if (!(space->flags & ARCH_VMM_ADDRESS_SPACE_OWNED)) {
+        return;
+    }
 
 #if defined(__x86_64__)
     scoped_lock(&space->lock) {
@@ -320,11 +312,12 @@ void arch_vmm_free_address_space(vmm_address_space_t* space) {
 #endif
 
 
-    // FIXME: maybe unsafe
-    //__free_frame(space->pm, X86_MMU_PAGESIZE);
+    __free_frame(space->pm, X86_MMU_PAGESIZE);
 
     space->pm              = 0UL;
     space->size            = 0UL;
     space->mmap.heap_start = 0UL;
     space->mmap.heap_end   = 0UL;
+
+    kfree(space);
 }

arch/x86-family/mm/getphysaddr.c

@@ -41,6 +41,8 @@
 __nonnull(1) uintptr_t arch_vmm_getphysaddr(vmm_address_space_t* space, uintptr_t virtaddr) {
 
 
+retry:
+
     uintptr_t pagesize = X86_MMU_PAGESIZE;
 
 
@@ -125,8 +127,11 @@ __nonnull(1) uintptr_t arch_vmm_getphysaddr(vmm_address_space_t* space, uintptr_
         DEBUG_ASSERT((*d != X86_MMU_CLEAR) && "Page unmapped");
 
 
-        if (unlikely((*d & X86_MMU_PG_AP_TP_MASK) != X86_MMU_PG_AP_TP_PAGE))
+        if (unlikely((*d & X86_MMU_PG_AP_TP_MASK) != X86_MMU_PG_AP_TP_PAGE)) {
+            spinlock_unlock(&space->lock);
             PANIC_ASSERT(pagefault_handle(current_cpu->frame, virtaddr) == 0);
+            goto retry;
+        }
 
 
         DEBUG_ASSERT(((*d & X86_MMU_PG_AP_TP_MASK) == X86_MMU_PG_AP_TP_PAGE) && "Page bad type");

arch/x86-family/mm/pagefault.c

@@ -40,6 +40,8 @@
 
 __nonnull(1) int pagefault_handle(interrupt_frame_t* frame, uintptr_t cr2) {
 
+    bool locked = false;
+    vmm_address_space_t* space = NULL;
 
 #if DEBUG_LEVEL_TRACE
 
@@ -63,6 +65,12 @@ __nonnull(1) int pagefault_handle(interrupt_frame_t* frame, uintptr_t cr2) {
     if (unlikely(!pm))
         PFE("no memory mapping", 0L);
 
+    if (current_task && current_task->address_space && current_task->address_space->pm == pm) {
+        space = current_task->address_space;
+        spinlock_lock(&space->lock);
+        locked = true;
+    }
+
 
     uintptr_t pagesize = X86_MMU_PAGESIZE;
     uintptr_t s        = cr2;
@@ -120,7 +128,7 @@ __nonnull(1) int pagefault_handle(interrupt_frame_t* frame, uintptr_t cr2) {
 #elif defined(__i386__)
 
         /* CR3-L2 */
-        { d = &((x86_page_t*)arch_vmm_p2v(space->pm, ARCH_VMM_AREA_HEAP))[(s >> 22) & 0x3FF]; }
+        { d = &((x86_page_t*)arch_vmm_p2v(pm, ARCH_VMM_AREA_HEAP))[(s >> 22) & 0x3FF]; }
 
 
         /* HUGE_4MB */
@@ -129,7 +137,7 @@ __nonnull(1) int pagefault_handle(interrupt_frame_t* frame, uintptr_t cr2) {
             /* PD-L1 */
             {
                 if (*d == X86_MMU_CLEAR)
-                    PFE("PD-L1 doesn't not exist");
+                    PFE("PD-L1 doesn't not exist", *d);
 
                 d = &((x86_page_t*)arch_vmm_p2v(*d & X86_MMU_ADDRESS_MASK, ARCH_VMM_AREA_HEAP))[(s >> 12) & 0x3FF];
             }
@@ -153,21 +161,30 @@ __nonnull(1) int pagefault_handle(interrupt_frame_t* frame, uintptr_t cr2) {
 
             //! Handle Copy on Write
 
-            uintptr_t page = __alloc_frame(pagesize, false);
+            uintptr_t old = *d;
+            uintptr_t page = __alloc_frame(pagesize, (old & X86_MMU_ADDRESS_MASK) == 0);
 
-            if ((*d & X86_MMU_ADDRESS_MASK) != 0) {
+            if ((old & X86_MMU_ADDRESS_MASK) != 0) {
 
-                memcpy((void*)arch_vmm_p2v(page, ARCH_VMM_AREA_HEAP), (void*)arch_vmm_p2v(*d & X86_MMU_ADDRESS_MASK, ARCH_VMM_AREA_HEAP), (size_t)pagesize);
+                memcpy((void*)arch_vmm_p2v(page, ARCH_VMM_AREA_HEAP), (void*)arch_vmm_p2v(old & X86_MMU_ADDRESS_MASK, ARCH_VMM_AREA_HEAP), (size_t)pagesize);
 
                 page |= X86_MMU_PG_RW;
             }
 
-            *d = page | X86_MMU_PG_P | X86_MMU_PG_AP_PFB | X86_MMU_PG_AP_TP_PAGE | ((*d & ~X86_MMU_ADDRESS_MASK) & ~(X86_MMU_PG_AP_TP_MASK));
+            *d = page | X86_MMU_PG_P | X86_MMU_PG_AP_PFB | X86_MMU_PG_AP_TP_PAGE | ((old & ~X86_MMU_ADDRESS_MASK) & ~(X86_MMU_PG_AP_TP_MASK));
         }
     }
 
+    __asm__ __volatile__("invlpg (%0)" ::"r"(cr2) : "memory");
+
+    if (locked) {
+        spinlock_unlock(&space->lock);
+    }
+
 
-    current_task->rusage.ru_majflt++;
+    if (current_task) {
+        current_task->rusage.ru_majflt++;
+    }
 
 
 #if DEBUG_LEVEL_TRACE
@@ -180,6 +197,10 @@ __nonnull(1) int pagefault_handle(interrupt_frame_t* frame, uintptr_t cr2) {
 
 pfe:
 
+    if (locked) {
+        spinlock_unlock(&space->lock);
+    }
+
     if (x86_intr_is_user_mode(frame))
         return -1;
 

arch/x86-family/mm/vm_access.c

@@ -130,6 +130,11 @@ __nonnull(1) int arch_vmm_access(vmm_address_space_t* space, uintptr_t virtaddr,
 
         /* Page Table */
         {
+            check_or_fail(*d != X86_MMU_CLEAR);
+
+            if (!(mode & S_OK) && current_cpu->frame && x86_intr_is_user_mode(current_cpu->frame)) {
+                check_or_fail(*d & X86_MMU_PG_U);
+            }
 
             if (mode & R_OK) {
                 if (!(*d & X86_MMU_PG_P)) {

arch/x86-family/mm/vm_lock.c

@@ -51,6 +51,6 @@ __nonnull(1) void arch_vmm_lock(vmm_address_space_t* space, uintptr_t virtaddr,
 
 #if defined(CONFIG_X86_ENABLE_SMAP)
     if (cpu_has(current_cpu->id, X86_FEATURE_SMAP))
-        x86_set_cr4(x86_get_cr4() & ~(X86_CR4_SMAP_MASK));
+        x86_stac();
 #endif
 }

arch/x86-family/mm/vm_map.c

@@ -262,7 +262,7 @@ __nonnull(1) uintptr_t arch_vmm_map(vmm_address_space_t* space, uintptr_t virtad
                 if (flags & ARCH_VMM_MAP_DEMAND)
                     *d = X86_MMU_PG_AP_TP_COW | (b & ~X86_MMU_PG_P);
                 else
-                    *d = __alloc_frame(pagesize, false) | X86_MMU_PG_AP_PFB | b;
+                    *d = __alloc_frame(pagesize, true) | X86_MMU_PG_AP_PFB | b;
             }
         }
 

arch/x86-family/mm/vm_protect.c

@@ -81,23 +81,23 @@ __nonnull(1) uintptr_t arch_vmm_mprotect(vmm_address_space_t* space, uintptr_t v
 
 
 
-    uint64_t b = X86_MMU_PG_P;
+    uint64_t base = X86_MMU_PG_P;
 
 
     if (flags & ARCH_VMM_MAP_DISABLED)
-        b &= ~X86_MMU_PG_P;
+        base &= ~X86_MMU_PG_P;
 
     if (flags & ARCH_VMM_MAP_RDWR)
-        b |= X86_MMU_PG_RW;
+        base |= X86_MMU_PG_RW;
 
     if (flags & ARCH_VMM_MAP_USER)
-        b |= X86_MMU_PG_U;
+        base |= X86_MMU_PG_U;
 
     if (flags & ARCH_VMM_MAP_UNCACHED)
-        b |= X86_MMU_PG_CD;
+        base |= X86_MMU_PG_CD;
 
     if (flags & ARCH_VMM_MAP_SHARED)
-        b |= X86_MMU_PG_G;
+        base |= X86_MMU_PG_G;
 
 
 
@@ -106,23 +106,23 @@ __nonnull(1) uintptr_t arch_vmm_mprotect(vmm_address_space_t* space, uintptr_t v
     //* Set No-Execute Bit
     if (flags & ARCH_VMM_MAP_NOEXEC)
         if (boot_cpu_has(X86_FEATURE_NX))
-            b |= X86_MMU_PT_NX; /* NX */
+            base |= X86_MMU_PT_NX; /* NX */
 
 #endif
 
     if (flags & ARCH_VMM_MAP_HUGETLB) {
 
-        b |= X86_MMU_PG_PS;
+        base |= X86_MMU_PG_PS;
 
         if (flags & ARCH_VMM_MAP_VIDEO_MEMORY)
             if (boot_cpu_has(X86_FEATURE_PAT))
-                b |= X86_MMU_PG_PAT; /* WC */
+                base |= X86_MMU_PG_PAT; /* WC */
 
     } else {
 
         if (flags & ARCH_VMM_MAP_VIDEO_MEMORY)
             if (boot_cpu_has(X86_FEATURE_PAT))
-                b |= X86_MMU_PT_PAT; /* WC */
+                base |= X86_MMU_PT_PAT; /* WC */
     }
 
 
@@ -203,10 +203,13 @@ __nonnull(1) uintptr_t arch_vmm_mprotect(vmm_address_space_t* space, uintptr_t v
             DEBUG_ASSERT((*d != X86_MMU_CLEAR) && "Page unmapped");
 
 
-            if (!(*d & X86_MMU_PG_P))
-                b &= ~X86_MMU_PG_P;
+            x86_page_t old = *d;
+            uint64_t page_flags = base;
 
-            *d = (*d & X86_MMU_ADDRESS_MASK) | (*d & X86_MMU_PG_AP_TP_MASK) | b;
+            if (!(flags & ARCH_VMM_MAP_DISABLED) && !(old & X86_MMU_PG_P) && ((old & X86_MMU_PG_AP_TP_MASK) == X86_MMU_PG_AP_TP_COW))
+                page_flags &= ~X86_MMU_PG_P;
+
+            *d = (old & X86_MMU_ADDRESS_MASK) | (old & (X86_MMU_PG_AP_TP_MASK | X86_MMU_PG_AP_PFB)) | page_flags;
 
 
 

arch/x86-family/mm/vm_unlock.c

@@ -50,6 +50,6 @@ __nonnull(1) void arch_vmm_unlock(vmm_address_space_t* space, uintptr_t virtaddr
 
 #if defined(CONFIG_X86_ENABLE_SMAP)
     if (cpu_has(current_cpu->id, X86_FEATURE_SMAP))
-        x86_set_cr4(x86_get_cr4() | X86_CR4_SMAP_MASK);
+        x86_clac();
 #endif
 }

arch/x86-family/mm/vm_unmap.c

@@ -154,7 +154,7 @@ __nonnull(1) uintptr_t arch_vmm_unmap(vmm_address_space_t* space, uintptr_t virt
         }
 
 
-        __asm__ __volatile__("invlpg (%0)" ::"r"(virtaddr) : "memory");
+        __asm__ __volatile__("invlpg (%0)" ::"r"(s) : "memory");
 
         space->size -= pagesize >> 12;
     }

include/aplus/memory.h

@@ -90,6 +90,8 @@
 
     #define ARCH_VMM_CLONE_NEW_SPACE (0)
 
+    #define ARCH_VMM_ADDRESS_SPACE_OWNED (1 << 0)
+
 
 
 typedef struct {
@@ -107,6 +109,7 @@ typedef struct vmm_address_space {
     uintptr_t pm;
     size_t size;
     size_t refcount;
+    uint32_t flags;
 
     struct {
 

include/arch/x86/cpu.h

@@ -650,6 +650,14 @@ static inline void x86_swapgs() {
     __asm__ __volatile__("swapgs");
 }
 
+static inline void x86_stac() {
+    __asm__ __volatile__("stac" ::: "memory");
+}
+
+static inline void x86_clac() {
+    __asm__ __volatile__("clac" ::: "memory");
+}
+
 
 /*!
  * @brief Write to Extended Control Register.

include/arch/x86/vmm.h

@@ -85,7 +85,7 @@
 
     #if defined(__x86_64__)
 typedef uint64_t x86_page_t;
-    #elif
+    #elif defined(__i386__)
 typedef uint32_t x86_page_t;
     #endif