/*++

Copyright (c) 2024, Quinn Stephens.
Provided under the BSD 3-Clause license.

Module Name:

   arch.c

Abstract:

   Provides x86 architecture-specific routines.

--*/

#include <ntrtl.h>
#include "bootlib.h"

// TODO: Move these to a header file?
#define CR0_PG (1 << 31)
#define CR4_OSFXSR (1 << 9)
#define CR4_LA57 (1 << 12)
#define IA32_EFER_MSR 0xC0000080
#define IA32_EFER_LME (1 << 10)

extern ULONG BlPlatformFlags;
extern PBOOT_FIRMWARE_DATA EfiFirmwareParameters;
EXECUTION_CONTEXT ApplicationExecutionContext;
EXECUTION_CONTEXT FirmwareExecutionContext;
PEXECUTION_CONTEXT CurrentExecutionContext;

VOID
Archpx64EnableInterruptsAsm (
    );

/*++

Routine Description:

    Enables interrupts.

    Implemented in ctx.S.

Arguments:

    None.

Return Value:

    None.

--*/

VOID
Archpx64DisableInterruptsAsm (
    );

/*++

Routine Description:

    Disables interrupts.

    Implemented in ctx.S.

Arguments:

    None.

Return Value:

    None.

--*/

VOID
BlpArchGetDescriptorTableContext (
    PDESCRIPTOR_TABLE_CONTEXT Context
    );

/*++

Routine Description:

    Loads current descriptor values into a
    descriptor table context structure.

    Implemented in ctx.S.

Arguments:

    Context - Pointer to the context structure.

Return Value:

    None.

--*/

VOID
ArchSetDescriptorTableContext (
    PDESCRIPTOR_TABLE_CONTEXT Context
    );

/*++

Routine Description:

    Loads current descriptor values from a
    descriptor table context structure.

    Implemented in ctx.S.

Arguments:

    Context - Pointer to the context structure.

Return Value:

    None.

--*/

BOOLEAN
BlArchIsFiveLevelPagingActive (
    )

{
    ULONG_PTR Cr0, Cr4;
    ULONG Efer;

    //
    // Paging must be enabled.
    //
    asm volatile("mov %%cr0, %0" :"=r"(Cr0));
    if (!(Cr0 & CR0_PG)) {
        return FALSE;
    }

    //
    // Long mode must be enabled.
    //
    asm volatile("rdmsr" :"=a"(Efer) :"c"(IA32_EFER_MSR));
    if (!(Efer & IA32_EFER_LME)) {
        return FALSE;
    }

    //
    // 57-bit linear addresses must be enabled.
    // If LA57 is enabled, 5-level paging is enabled.
    //
    asm volatile("mov %%cr4, %0" :"=r"(Cr4));
    if (!(Cr4 & CR4_LA57)) {
        return FALSE;
    }

    return TRUE;
}

NTSTATUS
ArchInitializeContext (
    IN OUT PEXECUTION_CONTEXT Context
    )

/*++

Routine Description:

    Initializes an execution context.

Arguments:

    Context - Pointer to the context structure.

Return Value:

    STATUS_SUCCESS if successful.

--*/

{
    ULONG_PTR Cr4;

    if (Context->Type == ExecutionContextFirmware) {
        Context->Flags &= ~(EXECUTION_CONTEXT_PAGING_ENABLED | 0x01);
        Context->Flags |= EXECUTION_CONTEXT_INTERRUPTS_ENABLED;

        //
        // Use context data from firmware.
        //
        Context->Cr3 = EfiFirmwareParameters->Cr3;
        RtlCopyMemory(&Context->DescriptorTableContext, &EfiFirmwareParameters->DescriptorTableContext, sizeof(DESCRIPTOR_TABLE_CONTEXT));

        return STATUS_SUCCESS;
    }

    Context->Flags &= ~EXECUTION_CONTEXT_INTERRUPTS_ENABLED;
    Context->Flags |= 0x01;

    //
    // Use current context.
    //
    asm volatile("mov %%cr3, %0" :"=r"(Context->Cr3));
    BlpArchGetDescriptorTableContext(&Context->DescriptorTableContext);

    //
    // Check if 5-level paging is active.
    //
    if (!BlArchIsFiveLevelPagingActive()) {
        Context->Flags &= ~EXECUTION_CONTEXT_PAGING_ENABLED;

        //
        // Enable SSE and FXSAVE/FXRSTOR.
        //
        asm volatile("mov %%cr4, %0" :"=r"(Cr4));
        Cr4 |= CR4_OSFXSR;
        asm volatile("mov %0, %%cr4" ::"r"(Cr4));

        return STATUS_SUCCESS;
    }

    //
    // TODO: Not sure what is supposed to happen here.
    //
    DebugPrint(L"ArchInitializeContext(): 5-level paging support not implemented\r\n");
    return STATUS_NOT_IMPLEMENTED;
}

VOID
ArchSetPagingContext(
    IN PEXECUTION_CONTEXT NewContext,
    IN PEXECUTION_CONTEXT CurrentContext
    )

/*++

Routine Description:

    Loads the current paging context.

Arguments:

    NewContext - The context to switch to.

    CurrentContext - The currently loaded context.

Return Value:

    None.

--*/

{
    BOOLEAN PagingEnabled;

    //
    // Check if paging is enabled.
    //
    if (CurrentContext != NULL) {
        PagingEnabled = CurrentContext->Flags & EXECUTION_CONTEXT_PAGING_ENABLED ? TRUE:FALSE;
    } else {
        PagingEnabled = BlArchIsFiveLevelPagingActive() ? TRUE:FALSE;
    }

    //
    // If paging is not being enabled/disabled,
    // just switch CR3.
    //
    if (PagingEnabled == (NewContext->Flags & EXECUTION_CONTEXT_PAGING_ENABLED ? TRUE:FALSE)) {
        asm volatile("mov %0, %%cr3" ::"r"(NewContext->Cr3));
        return;
    }

    //
    // TODO: Finish this routine.
    //
    DebugPrint(L"ArchSetPagingContext(): Paging mode enable/disable not implemented\r\n");
}

VOID
ArchSwitchContext (
    IN PEXECUTION_CONTEXT NewContext,
    IN PEXECUTION_CONTEXT CurrentContext
    )

/*++

Routine Description:

    Switches to a specified execution context.

Arguments:

    NewContext - The context to switch to.

    CurrentContext - The currently loaded context.

Return Value:

    None.

--*/

{
    if (CurrentContext != NULL && CurrentContext->Flags & EXECUTION_CONTEXT_INTERRUPTS_ENABLED) {
        Archpx64DisableInterruptsAsm();
    }

    //
    // Set descriptor table and paging contexts,
    // in the correct order.
    //
    if (NewContext->Type == ExecutionContextFirmware) {
        ArchSetPagingContext(NewContext, CurrentContext);
        ArchSetDescriptorTableContext(&NewContext->DescriptorTableContext);
    } else {
        ArchSetDescriptorTableContext(&NewContext->DescriptorTableContext);
        ArchSetPagingContext(NewContext, CurrentContext);
    }

    if (NewContext->Flags & EXECUTION_CONTEXT_INTERRUPTS_ENABLED) {
        Archpx64EnableInterruptsAsm();
    }
}

VOID
BlpArchSwitchContext (
    IN EXECUTION_CONTEXT_TYPE Type
    )

/*++

Routine Description:

    Switches to an execution context of type Type.

Arguments:

    Type - The requested context type.

Return Value:

    None.

--*/

{
    PEXECUTION_CONTEXT NewContext;

    if (Type == ExecutionContextFirmware) {
        NewContext = &FirmwareExecutionContext;
    } else {
        NewContext = &ApplicationExecutionContext;
    }

    if (CurrentExecutionContext->Type == NewContext->Type) {
        return;
    }

    ArchSwitchContext(NewContext, CurrentExecutionContext);
    CurrentExecutionContext = NewContext;
}

NTSTATUS
BlpArchInitialize (
    IN ULONG Stage
    )

/*++

Routine Description:

    Internal routine to perform architecture-specific initialization.

Arguments:

    Stage - 0.

Return Value:

    STATUS_SUCCESS.

--*/

{
    NTSTATUS Status;

    if (Stage == 0) {
        ApplicationExecutionContext.Type = ExecutionContextApplication;
        FirmwareExecutionContext.Type = ExecutionContextFirmware;
        CurrentExecutionContext = NULL;

        //
        // Initialize and use application context.
        //
        Status = ArchInitializeContext(&ApplicationExecutionContext);
        if (NT_SUCCESS(Status)) {
            CurrentExecutionContext = &ApplicationExecutionContext;
        }

        //
        // Initialize firmware context if supported.
        //
        if (BlPlatformFlags & FIRMWARE_FLAG_EXECUTION_CONTEXT_SUPPORTED) {
            Status = ArchInitializeContext(&FirmwareExecutionContext);
            if (!NT_SUCCESS(Status)) {
                // TODO: Implement ArchInitializeProcessorFeatures()?
                // ArchInitializeProcessorFeatures();
                return Status;
            }

            //
            // Use firmware execution context if
            // the application context is not
            // supported.
            //
            if (CurrentExecutionContext == NULL) {
                CurrentExecutionContext = &FirmwareExecutionContext;
            }
        }

        //
        // Switch to the correct context.
        //
        ArchSwitchContext(CurrentExecutionContext, NULL);
    }

    return STATUS_SUCCESS;
}