Untitled

/*
    Copyright (C) 2024-25 Amar Djulovic

    This file is part of floppaOS.

    floppaOS is free software: you can redistribute it and/or modify it under
    the terms of the GNU General Public License as published by the Free Software Foundation,
    either version 3 of the License, or (at your option) any later version.

    floppaOS is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
    See the GNU General Public License for more details.
*/

#include "vmm.h"
#include "pmm.h"
#include "paging.h"
#include "utils.h"
#include "../lib/logging.h"
#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#define KERNEL_VADDR_BASE  0xC0000000
#define USER_SPACE_END     0xBFFFFFFF
#define GNU_INLINE      __attribute__((always_inline))
#define PAGE_ALIGN(addr)   (((addr) + 0xFFF) & ~0xFFF)
#define RECURSIVE_PD      ((PDE*)0xFFFFF000)
#define RECURSIVE_PT(i)   ((PTE*)(0xFFC00000 + ((i) << 12)))
typedef struct vm_region {
    uintptr_t start;
    uintptr_t end;
    PageAttributes attrs;
    struct vm_region* next;
} vm_region_t;

static vm_region_t* vm_regions = NULL;
uintptr_t kernel_heap_start;
uintptr_t kernel_heap_end;

extern PDE pd[PAGE_DIRECTORY_SIZE];
extern PDE kernel_pd[PAGE_DIRECTORY_SIZE];
extern uint8_t _kernel_start;
extern uint8_t _kernel_end;

static inline GNU_INLINE uint32_t _pd_index(uintptr_t addr) {
    return (addr >> 22) & 0x3FF;
}

static inline GNU_INLINE uint32_t _pt_index(uintptr_t addr) {
    return (addr >> 12) & 0x3FF;
}

static inline GNU_INLINE uintptr_t _align_down(uintptr_t addr) {
    return addr & ~(PAGE_SIZE - 1);
}

static inline GNU_INLINE uintptr_t _align_up(uintptr_t addr) {
    return PAGE_ALIGN(addr);
}

static void _zero_pt(PTE* pt) {
    flop_memset(pt, 0, PAGE_TABLE_SIZE * sizeof(PTE));
}


static PTE* _alloc_pt(void) {
    void* pt = pmm_alloc_page();
    if (!pt) return NULL;
    return (PTE*)pt;
}

static int _pd_lvl(PTE* tables[2], uint32_t indices[2], PageAttributes attrs) {
    if (!tables[0][indices[0]].present) {
        PTE* new_pt = (PTE*)pmm_alloc_page();
        if (!new_pt) return -1;
        _zero_pt((PTE*)((uintptr_t)new_pt));

        PageAttributes new_attrs = {
            .present = 1,
            .rw = 1,
            .user = attrs.user,
            .frame_addr = ((uintptr_t)new_pt) >> 12
        };
        SET_PF((PTE*)&tables[0][indices[0]], new_attrs);
    }
    tables[1] = RECURSIVE_PT(indices[0]);
    return 0;
}

static void _pt_lvl(PTE* tables[2], uint32_t indices[2], uintptr_t paddr, PageAttributes attrs) {
    if (tables[1][indices[1]].present) {
        pmm_free_page((void*)(tables[1][indices[1]].frame_addr << 12));
    }
    attrs.frame_addr = paddr >> 12;
    SET_PF(&tables[1][indices[1]], attrs);
}

int map_page(uintptr_t vaddr, uintptr_t paddr, PageAttributes attrs) {
    vaddr = _align_down(vaddr);
    paddr = _align_down(paddr);
    if (vaddr == 0) return -1;

    uint32_t indices[2] = {_pd_index(vaddr), _pt_index(vaddr)};
    PTE* tables[2] = {RECURSIVE_PD, NULL};
    for (int i = 0; i < 2; i++) { // walk page table levels
        switch (i) {
            case 0:
                _pd_lvl(tables, indices, attrs);
            case 1:
                _pt_lvl(tables, indices, paddr, attrs);
        }
    }

    asm volatile ("invlpg (%0)" : : "a" (vaddr));
    return 0;
}

int unmap_page(uintptr_t vaddr) {
    vaddr = _align_down(vaddr);
    uint32_t pd_idx = _pd_index(vaddr);
    uint32_t pt_idx = _pt_index(vaddr);

    if (!RECURSIVE_PD[pd_idx].present) return -1;
    PTE* pt = RECURSIVE_PT(pd_idx);

    if (!pt[pt_idx].present) return -1;
    pmm_free_page((void*)(pt[pt_idx].frame_addr << 12));
    PageAttributes zero = {0};
    SET_PF(&pt[pt_idx], zero);
    bool pt_empty = true;
    for (int i = 0; i < PAGE_TABLE_SIZE; i++) {
        if (pt[i].present) {
            pt_empty = false;
            break;
        }
    }
    if (pt_empty) {
        pmm_free_page((void*)(RECURSIVE_PD[pd_idx].table_addr << 12));
        PageAttributes zero_pd = {0};
        SET_PF((PTE*)&RECURSIVE_PD[pd_idx], zero_pd);
    }

    asm volatile ("invlpg (%0)" : : "a" (vaddr));
    return 0;
}


int unmap_range(uintptr_t vaddr, size_t size) {
    size_t pages = _align_up(size) / PAGE_SIZE;

    vm_region_t *cur = vm_regions, *prev = NULL;

    while (cur) {
        // check if the range is within the current region
        if (cur->start == vaddr && cur->end == vaddr + size) {

            if (prev) {
                prev->next = cur->next;

            } else {
                vm_regions = cur->next;
            }
            break;
        }

        // move to next region
        prev = cur;
        cur = cur->next;
    }

    for (size_t i = 0; i < pages; i++) {
        uintptr_t va = vaddr + i * PAGE_SIZE;
        if (unmap_page(va) < 0) {
            // todo: handle error
            return -1;
        }
    }
    return 0;
}

int map_range(uintptr_t vaddr, uintptr_t paddr, size_t size, PageAttributes attrs) {
    size_t pages = _align_up(size) / PAGE_SIZE;

    vm_region_t* r = _make_region(vaddr, vaddr + size, attrs);
    if (!r) {
        return -1;
    }
    r->next = vm_regions;
    vm_regions = r;

    for (size_t i = 0; i < pages; i++) {
        uintptr_t va = vaddr + i * PAGE_SIZE;
        uintptr_t pa = paddr + i * PAGE_SIZE;
        if (map_page(va, pa, attrs) < 0) {
            // rollback if failed 
            // prevent mem leaks
            // todo: actual error handling
            for (size_t j = 0; j < i; j++) {
                uintptr_t undo_va = vaddr + j * PAGE_SIZE;
                unmap_page(undo_va);
            }
            vm_regions = r->next;
            return -1;
        }
    }
    return 0;
}

static vm_region_t* _make_region(uintptr_t start, uintptr_t end, PageAttributes attrs) {
    vm_region_t* r = pmm_alloc_page();
    if (!r) return NULL;
    r->start = start;
    r->end = end;
    r->attrs = attrs;
    r->next = NULL;
    return r;
}

static void _clear_pd(void) {
    for (int i = 0; i < PAGE_DIRECTORY_SIZE; i++) {
        PageAttributes zero = {
            0
        };
        SET_PF((PTE*)&pd[i], zero);
    }
}

static void _id_map(void) {
    for (uintptr_t addr = 0; addr < 0x100000; addr += PAGE_SIZE) {
        PageAttributes id_map_pg_attrs= { 
            .present = 1, 
            .rw = 1, 
            .user = 0, 
            .frame_addr = addr >> 12 };
        map_page(addr, addr, id_map_pg_attrs);
    }
    // todo: check if low memory is mapped correctly
}

static void _map_kernel(uintptr_t pstart, size_t size) {
    for (uintptr_t i = 0; i < size; i += PAGE_SIZE) {
        uintptr_t vaddr = KERNEL_VADDR_BASE + i;
        uintptr_t paddr = pstart + i;
        PageAttributes k_pg_attrs = { 
            .present = 1, 
            .rw = 1, 
            .user = 0, 
            .frame_addr = paddr >> 12 };
        map_page(vaddr, paddr, k_pg_attrs);
    }
}

static void _init_heap(uintptr_t kernel_size) {
    // place heap right after the kernel size in the kernel virtual address space
    
    kernel_heap_start = KERNEL_VADDR_BASE + kernel_size;
    kernel_heap_end = kernel_heap_start;
}

int vmm_init(void) {

    // get kernel end and start from linker.ld
    uintptr_t pstart = (uintptr_t)&_kernel_start; 
    uintptr_t pend   = (uintptr_t)&_kernel_end;
    uintptr_t ksize  = PAGE_ALIGN(pend - pstart); 
    _clear_pd();
    _id_map();
    _map_kernel(pstart, ksize);
    _init_heap(ksize); 
    _flush_tlb();
    log("vmm: initialized\n", GREEN);
    return 0;
}

uintptr_t virt_to_phys(uintptr_t vaddr) {
    vaddr = _align_down(vaddr);
    uint32_t pd_idx = _pd_index(vaddr);
    uint32_t pt_idx = _pt_index(vaddr);

    if (!RECURSIVE_PD[pd_idx].present) return 0;
    PTE* pt = RECURSIVE_PT(pd_idx);
    if (!pt[pt_idx].present) return 0;

    return (pt[pt_idx].frame_addr << 12) | (vaddr & 0xFFF);
}

uintptr_t phys_to_virt(uintptr_t paddr) {
    paddr = _align_down(paddr);
    for (uint32_t pd_idx = 0; pd_idx < PAGE_DIRECTORY_SIZE; pd_idx++) {
        if (!RECURSIVE_PD[pd_idx].present) continue;
        PTE* pt = RECURSIVE_PT(pd_idx);
        for (uint32_t pt_idx = 0; pt_idx < PAGE_TABLE_SIZE; pt_idx++) {
            if (!pt[pt_idx].present) continue;
            if ((pt[pt_idx].frame_addr << 12) == paddr) {
                return (pd_idx << 22) | (pt_idx << 12) | (paddr & 0xFFF);
            }
        }
    }
    return 0;
}
int map_kernel_space(uintptr_t vaddr, size_t size) {
    if (vaddr < KERNEL_VADDR_BASE) {
        return -1;
    }

    PageAttributes attrs = {
        .present = 1,
        .rw = 1,
        .user = 0
    };

    uintptr_t phys_addr = (uintptr_t)pmm_alloc_pages(0, size / PAGE_SIZE);
    return map_range(vaddr, phys_addr, size, attrs);
}

int map_user_space(uintptr_t vaddr, size_t size) {
    if (vaddr >= USER_SPACE_END) {
        return -1;
    }

    PageAttributes attrs = {
        .present = 1,
        .rw = 1,
        .user = 1
    };

    uintptr_t phys_addr = (uintptr_t)pmm_alloc_pages(0, size / PAGE_SIZE);
    return map_range(vaddr, phys_addr, size, attrs);
}
int vmm_is_mapped(uintptr_t vaddr) {
    uint32_t pd_idx = _pd_index(vaddr);
    uint32_t pt_idx = _pt_index(vaddr);

    if (!kernel_pd[pd_idx].present) return 0;

    PTE* pt = (PTE*)(kernel_pd[pd_idx].table_addr << 12);
    return pt[pt_idx].present;
}
int vmm_protect(uintptr_t vaddr, size_t size, PageAttributes new_attrs) {
    size_t pages = _align_up(size) / PAGE_SIZE;

    for (size_t i = 0; i < pages; i++) {
        uintptr_t va = vaddr + i * PAGE_SIZE;
        uint32_t pd_idx = _pd_index(va);
        uint32_t pt_idx = _pt_index(va);

        if (!kernel_pd[pd_idx].present)
            return -1;

        PTE* pt = (PTE*)(kernel_pd[pd_idx].table_addr << 12);
        if (!pt[pt_idx].present)
            return -1;

        new_attrs.frame_addr = pt[pt_idx].frame_addr;
        SET_PF(&pt[pt_idx], new_attrs);
    }

    _flush_tlb();
    return 0;
}
void vmm_dump_mappings(void) {
    vm_region_t* cur = vm_regions;
    while (cur) {
        log_step("vmm: region", CYAN);
        log_addr("  start", cur->start);
        log_addr("  end", cur->end);
        log("  attrs: ", WHITE);
        if (cur->attrs.user) log("user ", GREEN);
        if (cur->attrs.rw)   log("rw ", GREEN);
        if (cur->attrs.present) log("present ", GREEN);
        log("\n", WHITE);
        cur = cur->next;
    }
}
uintptr_t vmm_find_free_region(uintptr_t min, size_t size) {
    uintptr_t vaddr = min;
    while (vaddr + size < USER_SPACE_END) {
        if (!vmm_is_mapped(vaddr)) return vaddr;
        vaddr += PAGE_SIZE;
    }
    return 0;
}
PDE* vmm_clone_page_directory(void) {
    PDE* new_pd = (PDE*)pmm_alloc_page();
    if (!new_pd)
        return NULL;

    for (uint32_t i = 0; i < PAGE_DIRECTORY_SIZE; i++) {
        if ((i << 22) >= KERNEL_VADDR_BASE) {
            new_pd[i] = kernel_pd[i];
            continue;
        }

        if (!kernel_pd[i].present)
            continue;

        PTE* src_pt = (PTE*)(kernel_pd[i].table_addr << 12);
        PTE* new_pt = (PTE*)pmm_alloc_page();
        if (!new_pt)
            return NULL;

        _zero_pt(new_pt);

        for (uint32_t j = 0; j < PAGE_TABLE_SIZE; j++) {
            if (!src_pt[j].present)
                continue;

            void* new_frame = pmm_alloc_page();
            if (!new_frame)
                return NULL;

            void* old_frame = (void*)(src_pt[j].frame_addr << 12);
            flop_memcpy(new_frame, old_frame, PAGE_SIZE);

            PageAttributes attrs = {
                .present = src_pt[j].present,
                .rw = src_pt[j].rw,
                .user = src_pt[j].user,
                .frame_addr = ((uintptr_t)new_frame) >> 12
            };

            SET_PF(&new_pt[j], attrs);
        }

        PageAttributes pt_attrs = {
            .present = 1,
            .rw = 1,
            .user = 1,
            .frame_addr = ((uintptr_t)new_pt) >> 12
        };

        SET_PF((PTE*)&new_pd[i], pt_attrs);
    }

    return new_pd;
}

static PageAttributes pd_entry_get_attrs(PDE* pd, uint32_t pd_idx) {
    PageAttributes attrs = {0};
    if (!pd[pd_idx].present)
        return attrs;
    attrs.present = pd[pd_idx].present;
    attrs.rw = pd[pd_idx].rw;
    attrs.user = pd[pd_idx].user;
    attrs.frame_addr = pd[pd_idx].table_addr;
    return attrs;
}

PageAttributes vmm_pd_attrs_for_vaddr(uintptr_t vaddr) {
    uint32_t pd_idx = _pd_index(vaddr);
    return pd_entry_get_attrs(kernel_pd, pd_idx);
}

// get page attributes for a pt entry 
static PageAttributes pt_entry_get_attrs(PDE* pd, uint32_t pd_idx, uint32_t pt_idx) {
    PageAttributes attrs = {0};
    if (!pd[pd_idx].present)
        return attrs;
    PTE* pt = (PTE*)(pd[pd_idx].table_addr << 12);
    if (!pt[pt_idx].present)
        return attrs;
    attrs.present = pt[pt_idx].present;
    attrs.rw = pt[pt_idx].rw;
    attrs.user = pt[pt_idx].user;
    attrs.frame_addr = pt[pt_idx].frame_addr;
    return attrs;
}

PageAttributes vmm_pt_attrs_for_vaddr(uintptr_t vaddr) {
    uint32_t pd_idx = _pd_index(vaddr);
    uint32_t pt_idx = _pt_index(vaddr);
    return pt_entry_get_attrs(kernel_pd, pd_idx, pt_idx);
}

// Find region containing vaddr
static vm_region_t* find_region(uintptr_t vaddr) {
    vm_region_t* cur = vm_regions;
    while (cur) {
        if (vaddr >= cur->start && vaddr < cur->end)
            return cur;
        cur = cur->next;
    }
    return NULL;
}

// Get region attributes for vaddr
PageAttributes vaddr_fetch_attrs(uintptr_t vaddr) {
    vm_region_t* r = find_region(vaddr);
    if (!r) {
        PageAttributes zero = {0};
        return zero;
    }
    return r->attrs;
}