paging.h

Go to the documentation of this file.

#pragma once
 
#include <kernel/cpu/regs.h>
#include <kernel/mem/paging_types.h>
 
#include <assert.h>
#include <stdbool.h>
#include <stdlib.h>
#include <string.h>
#include <sys/proc.h>
 
#ifdef __BOOT__
#include <efi.h>
#include <efilib.h>
#endif
 
static inline void tlb_invalidate(void* virtAddr, uint64_t pageCount)
{
    if (pageCount == 0)
    {
        return;
    }
 
    if (pageCount > 16)
    {
        cr3_write(cr3_read());
    }
    else
    {
        for (uint64_t i = 0; i < pageCount; i++)
        {
            asm volatile("invlpg (%0)" ::"r"(virtAddr + i * PAGE_SIZE) : "memory");
        }
    }
}
 
static inline uintptr_t pml_accessible_addr(pml_entry_t entry)
{
#ifdef __BOOT__
    return entry.addr << PML_ADDR_OFFSET_BITS;
#else
    return PML_LOWER_TO_HIGHER(entry.addr << PML_ADDR_OFFSET_BITS);
#endif
}
 
static inline bool pml_is_empty(pml_t* pml)
{
    for (pml_index_t i = 0; i < PML_INDEX_AMOUNT; i++)
    {
        if (pml->entries[i].raw != 0)
        {
            return false;
        }
    }
    return true;
}
 
static inline uint64_t pml_new(page_table_t* table, pml_t** outPml)
{
    pml_t* pml;
    if (table->allocPages((void**)&pml, 1) == ERR)
    {
        return ERR;
    }
#ifdef __BOOT__
    SetMem(pml, PAGE_SIZE, 0);
#else
    memset(pml, 0, PAGE_SIZE);
#endif
    *outPml = pml;
    return 0;
}
 
static inline void pml_free(page_table_t* table, pml_t* pml, pml_level_t level)
{
    if (level < 0)
    {
        return;
    }
 
    for (pml_index_t i = 0; i < PML_INDEX_AMOUNT; i++)
    {
        pml_entry_t* entry = &pml->entries[i];
        if (!entry->present)
        {
            continue;
        }
 
        if (level > PML1)
        {
            pml_free(table, (pml_t*)pml_accessible_addr(*entry), level - 1);
        }
        else if (entry->owned)
        {
            void* addr = (void*)pml_accessible_addr(*entry);
            table->freePages(&addr, 1);
        }
    }
 
    table->freePages((void**)&pml, 1);
}
 
static inline uint64_t page_table_init(page_table_t* table, pml_alloc_pages_t allocPages, pml_free_pages_t freePages)
{
    table->allocPages = allocPages;
    table->freePages = freePages;
    if (pml_new(table, &table->pml4) == ERR)
    {
        return ERR;
    }
    return 0;
}
 
static inline void page_table_deinit(page_table_t* table)
{
    pml_free(table, table->pml4, PML4);
}
 
static inline void page_table_load(page_table_t* table)
{
    uint64_t cr3 = PML_ENSURE_LOWER_HALF(table->pml4);
    if (cr3 != cr3_read())
    {
        cr3_write(cr3);
    }
}
 
static inline uint64_t page_table_get_pml(page_table_t* table, pml_t* currentPml, pml_index_t index, pml_flags_t flags,
    pml_t** outPml)
{
    pml_entry_t* entry = &currentPml->entries[index];
    if (entry->present)
    {
        *outPml = (pml_t*)pml_accessible_addr(*entry);
        return 0;
    }
    else if (flags & PML_PRESENT)
    {
        pml_t* nextPml;
        if (pml_new(table, &nextPml) == ERR)
        {
            return ERR;
        }
        currentPml->entries[index].raw = (flags & PML_FLAGS_MASK) | (PML_ENSURE_LOWER_HALF(nextPml) & PML_ADDR_MASK);
        *outPml = nextPml;
        return 0;
    }
 
    return ERR;
}
 
typedef struct
{
    pml_t* pml3;
    pml_t* pml2;
    pml_t* pml1;
    bool pml3Valid;
    bool pml2Valid;
    bool pml1Valid;
    pml_index_t oldIdx3;
    pml_index_t oldIdx2;
    pml_index_t oldIdx1;
    pml_entry_t* entry;
} page_table_traverse_t;
 
#define PAGE_TABLE_TRAVERSE_CREATE \
    { \
        .pml3Valid = false, \
        .pml2Valid = false, \
        .pml1Valid = false, \
    }
 
static inline uint64_t page_table_traverse(page_table_t* table, page_table_traverse_t* traverse, uintptr_t virtAddr,
    pml_flags_t flags)
{
    pml_index_t newIdx3 = PML_ADDR_TO_INDEX(virtAddr, PML4);
    if (!traverse->pml3Valid || traverse->oldIdx3 != newIdx3)
    {
        if (page_table_get_pml(table, table->pml4, newIdx3, (flags | PML_WRITE | PML_USER) & ~PML_GLOBAL,
                &traverse->pml3) == ERR)
        {
            return ERR;
        }
        traverse->oldIdx3 = newIdx3;
        traverse->pml2Valid = false; // Invalidate cache for lower levels
    }
 
    pml_index_t newIdx2 = PML_ADDR_TO_INDEX(virtAddr, PML3);
    if (!traverse->pml2Valid || traverse->oldIdx2 != newIdx2)
    {
        if (page_table_get_pml(table, traverse->pml3, newIdx2, flags | PML_WRITE | PML_USER, &traverse->pml2) == ERR)
        {
            return ERR;
        }
        traverse->oldIdx2 = newIdx2;
        traverse->pml1Valid = false; // Invalidate cache for lower levels
    }
 
    pml_index_t newIdx1 = PML_ADDR_TO_INDEX(virtAddr, PML2);
    if (!traverse->pml1Valid || traverse->oldIdx1 != newIdx1)
    {
        if (page_table_get_pml(table, traverse->pml2, newIdx1, flags | PML_WRITE | PML_USER, &traverse->pml1) == ERR)
        {
            return ERR;
        }
        traverse->oldIdx1 = newIdx1;
    }
 
    traverse->entry = &traverse->pml1->entries[PML_ADDR_TO_INDEX(virtAddr, PML1)];
    return 0;
}
 
static inline uint64_t page_table_get_phys_addr(page_table_t* table, const void* virtAddr, void** outPhysAddr)
{
    uint64_t offset = ((uint64_t)virtAddr) % PAGE_SIZE;
    virtAddr = (void*)ROUND_DOWN(virtAddr, PAGE_SIZE);
 
    page_table_traverse_t traverse = PAGE_TABLE_TRAVERSE_CREATE;
 
    if (page_table_traverse(table, &traverse, (uintptr_t)virtAddr, PML_NONE) == ERR)
    {
        return ERR;
    }
 
    if (!traverse.entry->present)
    {
        return ERR;
    }
 
    *outPhysAddr = (void*)((traverse.entry->addr << PML_ADDR_OFFSET_BITS) + offset);
    return 0;
}
 
static inline bool page_table_is_mapped(page_table_t* table, const void* virtAddr, uint64_t pageAmount)
{
    page_table_traverse_t traverse = PAGE_TABLE_TRAVERSE_CREATE;
    for (uint64_t i = 0; i < pageAmount; i++)
    {
        if (page_table_traverse(table, &traverse, (uintptr_t)virtAddr + i * PAGE_SIZE, PML_NONE) == ERR)
        {
            return false;
        }
 
        if (!traverse.entry->present)
        {
            return false;
        }
    }
 
    return true;
}
 
static inline bool page_table_is_unmapped(page_table_t* table, const void* virtAddr, uint64_t pageAmount)
{
    page_table_traverse_t traverse = PAGE_TABLE_TRAVERSE_CREATE;
 
    for (uint64_t i = 0; i < pageAmount; i++)
    {
        if (page_table_traverse(table, &traverse, (uintptr_t)virtAddr + i * PAGE_SIZE, PML_NONE) == ERR)
        {
            continue;
        }
 
        if (traverse.entry->present)
        {
            return false;
        }
    }
 
    return true;
}
 
static inline uint64_t page_table_map(page_table_t* table, void* virtAddr, void* physAddr, uint64_t pageAmount,
    pml_flags_t flags, pml_callback_id_t callbackId)
{
    if (!(flags & PML_PRESENT))
    {
        return ERR;
    }
 
    page_table_traverse_t traverse = PAGE_TABLE_TRAVERSE_CREATE;
 
    for (uint64_t i = 0; i < pageAmount; i++)
    {
        if (page_table_traverse(table, &traverse, (uintptr_t)virtAddr, flags) == ERR)
        {
            return ERR;
        }
 
        if (traverse.entry->present)
        {
            return ERR;
        }
 
        traverse.entry->raw = flags;
        traverse.entry->addr = ((uintptr_t)PML_ENSURE_LOWER_HALF(physAddr)) >> PML_ADDR_OFFSET_BITS;
        traverse.entry->lowCallbackId = callbackId & 1;
        traverse.entry->highCallbackId = callbackId >> 1;
 
        physAddr = (void*)((uintptr_t)physAddr + PAGE_SIZE);
        virtAddr = (void*)((uintptr_t)virtAddr + PAGE_SIZE);
    }
 
    return 0;
}
 
static inline uint64_t page_table_map_pages(page_table_t* table, void* virtAddr, void** pages, uint64_t pageAmount,
    pml_flags_t flags, pml_callback_id_t callbackId)
{
    if (!(flags & PML_PRESENT))
    {
        return ERR;
    }
 
    page_table_traverse_t traverse = PAGE_TABLE_TRAVERSE_CREATE;
 
    for (uint64_t i = 0; i < pageAmount; i++)
    {
        if (page_table_traverse(table, &traverse, (uintptr_t)virtAddr, flags) == ERR)
        {
            return ERR;
        }
 
        if (traverse.entry->present)
        {
            return ERR;
        }
 
        traverse.entry->raw = flags;
        traverse.entry->addr = ((uintptr_t)PML_ENSURE_LOWER_HALF(pages[i])) >> PML_ADDR_OFFSET_BITS;
        traverse.entry->lowCallbackId = callbackId & 1;
        traverse.entry->highCallbackId = callbackId >> 1;
 
        virtAddr = (void*)((uintptr_t)virtAddr + PAGE_SIZE);
    }
 
    return 0;
}
 
static inline void page_table_unmap(page_table_t* table, void* virtAddr, uint64_t pageAmount)
{
    page_table_traverse_t traverse = PAGE_TABLE_TRAVERSE_CREATE;
 
    for (uint64_t i = 0; i < pageAmount; i++)
    {
        if (page_table_traverse(table, &traverse, (uintptr_t)virtAddr + i * PAGE_SIZE, PML_NONE) == ERR)
        {
            continue;
        }
 
        if (!traverse.entry->present)
        {
            continue;
        }
 
        traverse.entry->present = 0;
    }
 
    tlb_invalidate(virtAddr, pageAmount);
}
 
typedef struct
{
    void* pages[PML_PAGE_BUFFER_SIZE];
    uint64_t pageCount;
} page_table_page_buffer_t;
 
static inline void page_table_page_buffer_push(page_table_t* table, page_table_page_buffer_t* buffer, void* address)
{
    buffer->pages[buffer->pageCount] = address;
    buffer->pageCount++;
 
    if (buffer->pageCount >= PML_PAGE_BUFFER_SIZE)
    {
        table->freePages(buffer->pages, buffer->pageCount);
        buffer->pageCount = 0;
    }
}
 
static inline void page_table_page_buffer_flush(page_table_t* table, page_table_page_buffer_t* buffer)
{
    if (buffer->pageCount > 0)
    {
        table->freePages(buffer->pages, buffer->pageCount);
        buffer->pageCount = 0;
    }
}
 
static inline void page_table_clear_pml1_pml2_pml3(page_table_t* table, page_table_traverse_t* prevTraverse,
    page_table_traverse_t* traverse, page_table_page_buffer_t* pageBuffer)
{
    if (prevTraverse->pml1Valid && prevTraverse->pml1 != traverse->pml1 && pml_is_empty(prevTraverse->pml1))
    {
        page_table_page_buffer_push(table, pageBuffer, prevTraverse->pml1);
        prevTraverse->pml2->entries[prevTraverse->oldIdx1].raw = 0;
        if (prevTraverse->pml2Valid && prevTraverse->pml2 != traverse->pml2 && pml_is_empty(prevTraverse->pml2))
        {
            page_table_page_buffer_push(table, pageBuffer, prevTraverse->pml2);
            prevTraverse->pml3->entries[prevTraverse->oldIdx2].raw = 0;
            if (prevTraverse->pml3Valid && prevTraverse->pml3 != traverse->pml3 && pml_is_empty(prevTraverse->pml3))
            {
                page_table_page_buffer_push(table, pageBuffer, prevTraverse->pml3);
                table->pml4->entries[prevTraverse->oldIdx3].raw = 0;
            }
        }
    }
}
 
static inline void page_table_clear(page_table_t* table, void* virtAddr, uint64_t pageAmount)
{
    page_table_page_buffer_t pageBuffer = {.pageCount = 0};
 
    page_table_traverse_t prevTraverse = PAGE_TABLE_TRAVERSE_CREATE;
    page_table_traverse_t traverse = PAGE_TABLE_TRAVERSE_CREATE;
    for (uint64_t i = 0; i < pageAmount; i++)
    {
        uintptr_t currentVirtAddr = (uintptr_t)virtAddr + i * PAGE_SIZE;
 
        page_table_clear_pml1_pml2_pml3(table, &prevTraverse, &traverse, &pageBuffer);
 
        if (page_table_traverse(table, &traverse, currentVirtAddr, PML_NONE) == ERR)
        {
            prevTraverse.pml1Valid = false;
            prevTraverse.pml2Valid = false;
            prevTraverse.pml3Valid = false;
            continue;
        }
        prevTraverse = traverse;
 
        if (traverse.entry->present)
        {
            continue;
        }
 
        if (traverse.entry->owned)
        {
            page_table_page_buffer_push(table, &pageBuffer, (void*)pml_accessible_addr(*traverse.entry));
        }
 
        traverse.entry->raw = 0;
    }
 
    page_table_clear_pml1_pml2_pml3(table, &prevTraverse, &traverse, &pageBuffer);
    page_table_page_buffer_flush(table, &pageBuffer);
}
 
static inline void page_table_collect_callbacks(page_table_t* table, void* virtAddr, uint64_t pageAmount,
    uint64_t* callbacks)
{
    page_table_traverse_t traverse = PAGE_TABLE_TRAVERSE_CREATE;
 
    for (uint64_t i = 0; i < pageAmount; i++)
    {
        if (page_table_traverse(table, &traverse, (uintptr_t)virtAddr + i * PAGE_SIZE, PML_NONE) == ERR)
        {
            continue;
        }
 
        if (!traverse.entry->present)
        {
            continue;
        }
 
        pml_callback_id_t callbackId = traverse.entry->lowCallbackId | (traverse.entry->highCallbackId << 1);
        if (callbackId != PML_CALLBACK_NONE)
        {
            callbacks[callbackId]++;
        }
    }
}
 
static inline uint64_t page_table_set_flags(page_table_t* table, void* virtAddr, uint64_t pageAmount, pml_flags_t flags)
{
    page_table_traverse_t traverse = PAGE_TABLE_TRAVERSE_CREATE;
 
    for (uint64_t i = 0; i < pageAmount; i++)
    {
        if (page_table_traverse(table, &traverse, (uintptr_t)virtAddr + i * PAGE_SIZE, PML_NONE) == ERR)
        {
            continue;
        }
 
        if (!traverse.entry->present)
        {
            return ERR;
        }
 
        if (traverse.entry->owned)
        {
            flags |= PML_OWNED;
        }
 
        // Bit magic to only update the flags while preserving the address and callback ID.
        traverse.entry->raw = (traverse.entry->raw & ~PML_FLAGS_MASK) | (flags & PML_FLAGS_MASK);
    }
 
    tlb_invalidate(virtAddr, pageAmount);
    return 0;
}
 
static inline uint64_t page_table_find_unmapped_region(page_table_t* table, void* startAddr, void* endAddr,
    uint64_t pageAmount, void** outAddr)
{
    uintptr_t currentAddr = ROUND_DOWN((uintptr_t)startAddr, PAGE_SIZE);
    uintptr_t end = (uintptr_t)endAddr;
 
    if (pageAmount >= (PML3_SIZE / PAGE_SIZE))
    {
        while (currentAddr < end)
        {
            pml_index_t idx4 = PML_ADDR_TO_INDEX(currentAddr, PML4);
            pml_index_t idx3 = PML_ADDR_TO_INDEX(currentAddr, PML3);
 
            pml_entry_t* entry4 = &table->pml4->entries[idx4];
            if (!entry4->present)
            {
                *outAddr = (void*)currentAddr;
                return 0;
            }
 
            pml_t* pml3 = (pml_t*)pml_accessible_addr(*entry4);
            pml_entry_t* entry3 = &pml3->entries[idx3];
 
            if (!entry3->present)
            {
                *outAddr = (void*)currentAddr;
                return 0;
            }
 
            currentAddr = ROUND_UP(currentAddr + 1, PML3_SIZE);
        }
        return ERR;
    }
 
    if (pageAmount >= (PML2_SIZE / PAGE_SIZE))
    {
        while (currentAddr < end)
        {
            pml_index_t idx4 = PML_ADDR_TO_INDEX(currentAddr, PML4);
            pml_entry_t* entry4 = &table->pml4->entries[idx4];
 
            if (!entry4->present)
            {
                *outAddr = (void*)currentAddr;
                return 0;
            }
 
            pml_t* pml3 = (pml_t*)pml_accessible_addr(*entry4);
            pml_index_t idx3 = PML_ADDR_TO_INDEX(currentAddr, PML3);
            pml_entry_t* entry3 = &pml3->entries[idx3];
 
            if (!entry3->present)
            {
                *outAddr = (void*)currentAddr;
                return 0;
            }
 
            pml_t* pml2 = (pml_t*)pml_accessible_addr(*entry3);
            pml_index_t idx2 = PML_ADDR_TO_INDEX(currentAddr, PML2);
            pml_entry_t* entry2 = &pml2->entries[idx2];
 
            if (!entry2->present)
            {
                *outAddr = (void*)currentAddr;
                return 0;
            }
 
            currentAddr = ROUND_UP(currentAddr + 1, PML2_SIZE);
        }
        return ERR;
    }
 
    uintptr_t regionStart = 0;
    uint64_t consecutiveUnmapped = 0;
 
    while (currentAddr < end)
    {
        pml_index_t idx4 = PML_ADDR_TO_INDEX(currentAddr, PML4);
        pml_entry_t* entry4 = &table->pml4->entries[idx4];
 
        if (!entry4->present)
        {
            if (consecutiveUnmapped == 0)
            {
                regionStart = currentAddr;
            }
 
            uintptr_t skipTo = PML_INDEX_TO_ADDR(idx4 + 1, PML4);
            uint64_t skippedPages = (MIN(skipTo, end) - currentAddr) / PAGE_SIZE;
            consecutiveUnmapped += skippedPages;
 
            if (consecutiveUnmapped >= pageAmount)
            {
                *outAddr = (void*)regionStart;
                return 0;
            }
 
            currentAddr = skipTo;
            continue;
        }
 
        pml_t* pml3 = (pml_t*)pml_accessible_addr(*entry4);
        pml_index_t idx3 = PML_ADDR_TO_INDEX(currentAddr, PML3);
        pml_entry_t* entry3 = &pml3->entries[idx3];
 
        if (!entry3->present)
        {
            if (consecutiveUnmapped == 0)
            {
                regionStart = currentAddr;
            }
 
            uint64_t skippedPages = PML3_SIZE / PAGE_SIZE;
            consecutiveUnmapped += skippedPages;
 
            if (consecutiveUnmapped >= pageAmount)
            {
                *outAddr = (void*)regionStart;
                return 0;
            }
 
            currentAddr = ROUND_UP(currentAddr + 1, PML3_SIZE);
            continue;
        }
 
        pml_t* pml2 = (pml_t*)pml_accessible_addr(*entry3);
        pml_index_t idx2 = PML_ADDR_TO_INDEX(currentAddr, PML2);
        pml_entry_t* entry2 = &pml2->entries[idx2];
 
        if (!entry2->present)
        {
            if (consecutiveUnmapped == 0)
            {
                regionStart = currentAddr;
            }
 
            uint64_t skippedPages = PML2_SIZE / PAGE_SIZE;
            consecutiveUnmapped += skippedPages;
 
            if (consecutiveUnmapped >= pageAmount)
            {
                *outAddr = (void*)regionStart;
                return 0;
            }
 
            currentAddr = ROUND_UP(currentAddr + 1, PML2_SIZE);
            continue;
        }
 
        pml_t* pml1 = (pml_t*)pml_accessible_addr(*entry2);
        pml_index_t idx1 = PML_ADDR_TO_INDEX(currentAddr, PML1);
 
        for (; idx1 < PML_INDEX_AMOUNT && currentAddr < end; idx1++, currentAddr += PAGE_SIZE)
        {
            if (!pml1->entries[idx1].present)
            {
                if (consecutiveUnmapped == 0)
                {
                    regionStart = currentAddr;
                }
                consecutiveUnmapped++;
 
                if (consecutiveUnmapped >= pageAmount)
                {
                    *outAddr = (void*)regionStart;
                    return 0;
                }
            }
            else
            {
                consecutiveUnmapped = 0;
            }
        }
    }
 
    return ERR;
}
 
static inline bool page_table_is_pinned(page_table_t* table, const void* virtAddr, uint64_t pageAmount)
{
    page_table_traverse_t traverse = PAGE_TABLE_TRAVERSE_CREATE;
 
    for (uint64_t i = 0; i < pageAmount; i++)
    {
        if (page_table_traverse(table, &traverse, (uintptr_t)virtAddr + i * PAGE_SIZE, PML_NONE) == ERR)
        {
            continue;
        }
 
        if (!traverse.entry->present)
        {
            continue;
        }
 
        if (traverse.entry->pinned)
        {
            return true;
        }
    }
 
    return false;
}