GitFOSS

#include "module.h"
#include <stdint.h>
#include "timer.h"

extern void vga_write(const char*);
extern void idt_init(void);
extern void irq_enable(void);
extern void *kalloc(uint64_t);
extern int kbd_read_scancode(void);
extern void register_irq(int irq, void (*handler)(void));
extern void ps2_init(struct kernel_api *api);

/* Declare module section symbols */
extern char __modules_start;
extern char __modules_end;

/* Long mode setup functions */
extern void enable_long_mode(void);

/* Page table for identity mapping in long mode */
static uint64_t pml4[512] __attribute__((aligned(4096)));
static uint64_t pdpt[512] __attribute__((aligned(4096)));
static uint64_t pd[512] __attribute__((aligned(4096)));

void setup_paging(void) {
    /* Clear page tables */
    for (int i = 0; i < 512; i++) {
        pml4[i] = 0;
        pdpt[i] = 0;
        pd[i] = 0;
    }

    /* Set up PML4 -> PDPT -> PD (identity mapping for first 2MB) */
    pml4[0] = (uint64_t)pdpt | 0x03;    /* Present, Writable */
    pdpt[0] = (uint64_t)pd | 0x03;      /* Present, Writable */
    
    /* Map first 2MB as 2MB pages */
    for (int i = 0; i < 4; i++) {
        pd[i] = (i * 0x200000) | 0x83;  /* Present, Writable, 2MB page */
    }

    /* Load page table base */
    asm volatile ("mov %0, %%cr3" : : "r"(pml4));
}

void enable_long_mode(void) {
    /* Enable PAE in CR4 */
    uint64_t cr4;
    asm volatile ("mov %%cr4, %0" : "=r"(cr4));
    cr4 |= (1 << 5);  /* PAE bit */
    asm volatile ("mov %0, %%cr4" : : "r"(cr4));

    /* Set up paging */
    setup_paging();

    /* Enable long mode in EFER MSR */
    uint32_t efer_lo, efer_hi = 0;
    asm volatile ("rdmsr" : "=a"(efer_lo), "=d"(efer_hi) : "c"(0xC0000080));
    efer_lo |= (1 << 0);  /* LME bit */
    asm volatile ("wrmsr" : : "a"(efer_lo), "d"(efer_hi), "c"(0xC0000080));

    /* Enable paging and protected mode */
    uint64_t cr0;
    asm volatile ("mov %%cr0, %0" : "=r"(cr0));
    cr0 |= (1 << 0);    /* PE (protected mode) */
    cr0 |= (1 << 31);  /* PG (paging) */
    asm volatile ("mov %0, %%cr0" : : "r"(cr0));

    /* Far jump to enable 64-bit mode - use wrapper function */
    extern void jump_to_long_mode(void);
    jump_to_long_mode();
}

void kmain(void) {    
    /* Enable long mode before anything else */
    enable_long_mode();
    
    idt_init();
    irq_enable();
    
    vga_write("CrystalOS kernel\n");

    struct kernel_api api;
    api.vga_write = vga_write;
    api.alloc = kalloc;
    api.register_irq = register_irq;

    /* Initialize PS/2 keyboard IRQ handling */
    ps2_init(&api);

    /* Initialize timer and register its IRQ handler */
    timer_init();
    register_irq(0, timer_irq_handler); /* IRQ 0 = timer */

    char *p = &__modules_start;
    while (p < &__modules_end) {
        void (*initfn)(void*) = *(void(**)(void*))p;
        if (initfn) initfn(&api);
        p += sizeof(void*);
    }

    /* Idle loop */
    while (1) {
        asm volatile ("hlt");
    }
}