/*++

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

Module Name:

   bootlib.h

Abstract:

   Boot library definitions.

--*/

#ifndef _BOOTLIB_H
#define _BOOTLIB_H

#include <nt.h>

#if defined(_EFI)
#include "efi.h"

#define PAGE_SIZE  EFI_PAGE_SIZE
#define PAGE_SHIFT EFI_PAGE_SHIFT
#endif /* _EFI */

//
// Address translation types.
//
#define TRANSLATION_TYPE_NONE 0
#define TRANSLATION_TYPE_MAX  1

//
// Set machine type.
//
#if defined(__x86_64__)
#define BOOT_MACHINE_TYPE IMAGE_FILE_MACHINE_AMD64
#elif defined(__i386__)
#define BOOT_MACHINE_TYPE IMAGE_FILE_MACHINE_I386
#endif

#if defined(__i386__) || defined(__x86_64__)
typedef struct __attribute__((packed)) {
    USHORT Limit;
    ULONG_PTR Base;
} DESCRIPTOR_TABLE_REGISTER, *PDESCRIPTOR_TABLE_REGISTER;

typedef struct __attribute__((packed)) {
    DESCRIPTOR_TABLE_REGISTER Gdt;
    DESCRIPTOR_TABLE_REGISTER Idt;
    USHORT LdtSelector;
    USHORT CS, DS, ES, FS, GS, SS;
} DESCRIPTOR_TABLE_CONTEXT, *PDESCRIPTOR_TABLE_CONTEXT;
#endif

//
// Firmware platform flags.
//

#define FIRMWARE_FLAG_EXECUTION_CONTEXT_SUPPORTED 0x00100000

//
// Execution contexts represent the current state
// of the processor/system.
//

typedef enum {
    ExecutionContextApplication,
    ExecutionContextFirmware,
    ExecutionContextMax
} EXECUTION_CONTEXT_TYPE;

#define EXECUTION_CONTEXT_INTERRUPTS_ENABLED 0x02
#define EXECUTION_CONTEXT_PAGING_ENABLED     0x04

typedef struct {
    EXECUTION_CONTEXT_TYPE Type;
#if defined(__i386__) || defined(__x86_64__)
    ULONG_PTR Cr3;
#endif
    ULONG Flags;
    DESCRIPTOR_TABLE_CONTEXT DescriptorTableContext;
} EXECUTION_CONTEXT, *PEXECUTION_CONTEXT;

//
// Application entry option.
// Used to represent options passed to a boot
// application by whatever started it.
//

typedef struct {
    ULONG Type;
    ULONG DataOffset;
    ULONG DataSize;
    ULONG OtherOptionsOffset;
    ULONG NextOptionOffset;
    BOOLEAN IsInvalid;
    UCHAR Reserved[3];
} BOOT_ENTRY_OPTION, *PBOOT_ENTRY_OPTION;

//
// Application entry.
// Used for any boot environment application,
// including the boot manager itself.
//

#define BOOT_APPLICATION_ENTRY_NO_BCD_IDENTIFIER 0x01

typedef struct {
    ULONG Attributes;
    GUID BcdIdentifier;
    PBOOT_ENTRY_OPTION Options;
} BOOT_APPLICATION_ENTRY, *PBOOT_APPLICATION_ENTRY;

//
// Initial application entry.
// Special form of BOOT_APPLICATION_ENTRY used
// when firmware first loads the boot manager.
//

#define BOOT_INIT_APPLICATION_ENTRY_SIGNATURE 0x544e4550415442 /* "BTAPENT" */

typedef struct {
    ULONGLONG Signature;
    ULONG Attributes;
    GUID BcdIdentifier;
    UCHAR Reserved[16];
    BOOT_ENTRY_OPTION Options;
} BOOT_INIT_APPLICATION_ENTRY, *PBOOT_INIT_APPLICATION_ENTRY;

typedef ULONG MEMORY_TYPE;

//
// Memory type classes.
//
#define MEMORY_CLASS_APPLICATION 0xD
#define MEMORY_CLASS_LIBRARY     0xE
#define MEMORY_CLASS_SYSTEM      0xF

//
// Memory types.
//
#define MEMORY_TYPE_BOOT_APPLICATION      0xD0000002
#define MEMORY_TYPE_BOOT_APPLICATION_2    0xD0000013
#define MEMORY_TYPE_FREE                  0xF0000001
#define MEMORY_TYPE_UNUSABLE              0xF0000002
#define MEMORY_TYPE_RESERVED              0xF0000003
#define MEMORY_TYPE_BOOT_SERVICES         0xF0000004
#define MEMORY_TYPE_FREE_ZEROED           0xF0000005
#define MEMORY_TYPE_RUNTIME_SERVICES_CODE 0xF0000006
#define MEMORY_TYPE_PERSISTENT            0xF0000007
#define MEMORY_TYPE_ACPI_RECLAIM          0xF0000008
#define MEMORY_TYPE_ACPI_NVS              0xF0000009
#define MEMORY_TYPE_MMIO                  0xF000000A
#define MEMORY_TYPE_MMIO_PORT_SPACE       0xF000000B
#define MEMORY_TYPE_PAL_CODE              0xF000000C
#define MEMORY_TYPE_RUNTIME_SERVICES_DATA 0xF000000E

//
// Memory caching attributes.
//
#define MEMORY_ATTRIBUTE_UC  0x00000001
#define MEMORY_ATTRIBUTE_WC  0x00000002
#define MEMORY_ATTRIBUTE_WT  0x00000004
#define MEMORY_ATTRIBUTE_WB  0x00000008
#define MEMORY_ATTRIBUTE_UCE 0x00000010  

//
// Memory protection attributes.
//
#define MEMORY_ATTRIBUTE_WP 0x00000100
#define MEMORY_ATTRIBUTE_RP 0x00000200
#define MEMORY_ATTRIBUTE_XP 0x00000400

//
// Memory location attributes.
//
#define MEMORY_ATTRIBUTE_BELOW_1MIB 0x00080000

//
// Memory runtime mapping attributes.
//
#define MEMORY_ATTRIBUTE_RUNTIME 0x01000000

typedef struct {
    LIST_ENTRY ListEntry;

    ULONGLONG FirstPage;
    ULONGLONG MappedFirstPage;
    ULONG PageCount;

    ULONG Attributes;
    MEMORY_TYPE Type;
} MEMORY_DESCRIPTOR, *PMEMORY_DESCRIPTOR;

typedef enum {
    MDL_TYPE_PHYSICAL,
    MDL_TYPE_VIRTUAL
} MEMORY_DESCRIPTOR_LIST_TYPE;

typedef struct {
    LIST_ENTRY ListEntry;

    PLIST_ENTRY Head;
    PLIST_ENTRY Current;
    MEMORY_DESCRIPTOR_LIST_TYPE Type;
} MEMORY_DESCRIPTOR_LIST, *PMEMORY_DESCRIPTOR_LIST;

#define BOOT_MEMORY_INFO_VERSION 1

typedef struct {
    ULONG Version;
    ULONG MdlOffset;
    ULONG DescriptorCount;
    ULONG DescriptorSize;
    ULONG FirstPageOffset;
} BOOT_MEMORY_INFO, *PBOOT_MEMORY_INFO;

#define BOOT_FIRMWARE_DATA_VERSION 2

typedef struct {
    ULONG Version;
    ULONG Reserved;
    EFI_HANDLE ImageHandle;
    EFI_SYSTEM_TABLE *SystemTable;
    ULONG Reserved2;
#if defined(__i386__) || defined(__x86_64__)
    ULONG_PTR Cr3;
#endif
    DESCRIPTOR_TABLE_CONTEXT DescriptorTableContext;
} BOOT_FIRMWARE_DATA, *PBOOT_FIRMWARE_DATA;

#define BOOT_RETURN_DATA_VERSION 1

typedef struct {
    ULONG Version;
    NTSTATUS Status;
    ULONG Flags;
} BOOT_RETURN_DATA, *PBOOT_RETURN_DATA;

//
// Firmware-independent application parameters.
// Passed to any boot application's entry point.
//

#define BOOT_APPLICATION_PARAMETERS_SIGNATURE 0x50504120544f4f42 /* "BOOT APP" */
#define BOOT_APPLICATION_PARAMETERS_VERSION   2

typedef struct {
    ULONGLONG Signature;
    ULONG Version;
    ULONG Size;

    //
    // Machine information.
    //
    ULONG MachineType;
    ULONG TranslationType;

    //
    // Image information.
    //
    PVOID ImageBase;
    ULONG ImageSize;

    //
    // Offsets to ancillary structures.
    //
    ULONG MemoryInfoOffset;
    ULONG ApplicationEntryOffset;
    ULONG BootDeviceOffset;
    ULONG FirmwareDataOffset;
    ULONG ReturnDataOffset;
    ULONG PlatformDataOffset;
} BOOT_APPLICATION_PARAMETERS, *PBOOT_APPLICATION_PARAMETERS;

//
// Library parameters.
// Specifies how the boot library should be
// set up.
//

typedef struct {
    ULONG Flags;
    ULONG TranslationType;
    ULONG MinimumPageAllocation;
} BOOT_LIBRARY_PARAMETERS, *PBOOT_LIBRARY_PARAMETERS;

typedef struct {
    LARGE_INTEGER ImageBase;
    ULONG ImageSize;
    ULONG ImageOffset;
} BOOT_RAMDISK_IDENTIFIER, *PBOOT_RAMDISK_IDENTIFIER;

#define BOOT_HARDDRIVE_PARTITION_TYPE_GPT 0x00
#define BOOT_HARDDRIVE_PARTITION_TYPE_MBR 0x01
#define BOOT_HARDDRIVE_PARTITION_TYPE_RAW 0x02

typedef struct {
    ULONG PartitionType;

    union {
        struct {
            ULONG Signature;
        } Mbr;

        struct {
            GUID Signature;
        } Gpt;

        struct {
            ULONG DriveNumber;
        } Raw;
    };
} BOOT_HARDDRIVE_IDENTIFIER, *PBOOT_HARDDRIVE_IDENTIFIER;

typedef struct {
    ULONG DriveNumber;
} BOOT_CDROM_IDENTIFIER, *PBOOT_CDROM_IDENTIFIER;

#define BOOT_BLOCK_DEVICE_TYPE_HARDDRIVE 0x00
#define BOOT_BLOCK_DEVICE_TYPE_CDROM     0x02
#define BOOT_BLOCK_DEVICE_TYPE_RAMDISK   0x03

typedef struct {
    ULONG Type;

    union {
        BOOT_RAMDISK_IDENTIFIER Ramdisk;
        BOOT_HARDDRIVE_IDENTIFIER Harddrive;
        BOOT_CDROM_IDENTIFIER Cdrom;
    };
} BOOT_BLOCK_IDENTIFIER, *PBOOT_BLOCK_IDENTIFIER;

typedef struct {
    union {
        struct {
            ULONG PartitionNumber;
        } Mbr;

        struct {
            GUID PartitionIdentifier;
        } Gpt;

        struct {
            ULONG BootEntry;
        } ElTorito;
    };

    BOOT_BLOCK_IDENTIFIER Parent;
} BOOT_PARTITION_IDENTIFIER, *PBOOT_PARTITION_IDENTIFIER;

typedef struct {
    union {
        struct {
            PVOID PartitionOffset;
        } Mbr;

        struct {
            GUID PartitionIdentifier;
        } Gpt;

        struct {
            ULONG BootEntry;
        } ElTorito;
    };

    BOOT_BLOCK_IDENTIFIER Parent;
} BOOT_PARTITION_IDENTIFIER_EX, *PBOOT_PARTITION_IDENTIFIER_EX;

#define BOOT_DEVICE_TYPE_BLOCK        0x00
#define BOOT_DEVICE_TYPE_PARTITION    0x02
#define BOOT_DEVICE_TYPE_PARTITION_EX 0x06

#define BOOT_DEVICE_ATTRIBUTE_NO_PARENT_SIGNATURE 0x04

typedef struct {
    ULONG Type;
    ULONG Attributes;
    ULONG Size;
    ULONG Pad;

    union {
        BOOT_BLOCK_IDENTIFIER Block;
        BOOT_PARTITION_IDENTIFIER Partition;
        BOOT_PARTITION_IDENTIFIER_EX PartitionEx;
    };
} BOOT_DEVICE, *PBOOT_DEVICE;

typedef ULONG BCDE_DATA_TYPE;

#define BCDE_DATA_FORMAT_MASK 0x0F000000

#define BCDE_DATA_FORMAT_DEVICE 0x01000000
#define BCDE_DATA_FORMAT_STRING 0x02000000

#define BCDE_DATA_TYPE_APPLICATION_DEVICE 0x11000001
#define BCDE_DATA_TYPE_APPLICATION_PATH   0x22000002

#define BCDE_DATA_TYPE_BCD_DEVICE 0x21000022
#define BCDE_DATA_TYPE_BCD_PATH   0x22000023

typedef struct {
    GUID Options;
    BOOT_DEVICE Device;
} BCDE_DEVICE, *PBCDE_DEVICE;

//
// Enable/disable debug printing.
//
#if defined(_DEBUG)
#define DebugPrint(String) ConsolePrint(String)
#define DebugPrintf(Format, ...) ConsolePrintf(Format, __VA_ARGS__)
#else
#define DebugPrint(String)
#define DebugPrintf(Format, ...)
#endif

extern PEXECUTION_CONTEXT CurrentExecutionContext;
extern ULONG BlPlatformFlags;
extern PBOOT_DEVICE BlpBootDevice;
extern PBOOT_APPLICATION_PARAMETERS BlpApplicationParameters;
extern BOOT_LIBRARY_PARAMETERS BlpLibraryParameters;
extern BOOT_APPLICATION_ENTRY BlpApplicationEntry;

#if defined(_EFI)
extern PBOOT_FIRMWARE_DATA EfiFirmwareParameters;
extern EFI_HANDLE EfiImageHandle;
extern EFI_SYSTEM_TABLE *EfiST;
extern EFI_BOOT_SERVICES *EfiBS;
extern EFI_RUNTIME_SERVICES *EfiRT;
extern SIMPLE_TEXT_OUTPUT_INTERFACE *EfiConOut;
extern SIMPLE_INPUT_INTERFACE *EfiConIn;
extern EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *EfiConInEx;
#endif

VOID
ConsolePrint (
    IN PWSTR String
    );

VOID
ConsolePrintf (
    IN PWSTR Format,
    ...
    );

VOID
BlpArchGetDescriptorTableContext (
    PDESCRIPTOR_TABLE_CONTEXT Context
    );

BOOLEAN
BlArchIsFiveLevelPagingActive (
    );

VOID
BlpArchSwitchContext (
    IN EXECUTION_CONTEXT_TYPE Type
    );

NTSTATUS
BlpArchInitialize (
    IN ULONG Stage
    );

NTSTATUS
BlpFwInitialize (
    IN ULONG Stage,
    IN PBOOT_FIRMWARE_DATA FirmwareData
    );

NTSTATUS
BlpMmDestroy (
    IN ULONG Stage
    );

NTSTATUS
BlpMmInitializeConstraints (
    VOID
    );

NTSTATUS
BlpMmInitialize (
    IN PBOOT_MEMORY_INFO MemoryInfo,
    IN ULONG TranslationType,
    IN PBOOT_LIBRARY_PARAMETERS LibraryParameters
    );

ULONG
BlGetBootOptionSize (
    IN PBOOT_ENTRY_OPTION Option
    );

ULONG
BlGetBootOptionListSize (
    IN PBOOT_ENTRY_OPTION Options
    );

NTSTATUS
BlGetBootOptionDevice (
    IN PBOOT_ENTRY_OPTION Options,
    IN BCDE_DATA_TYPE Type,
    IN OUT PBOOT_DEVICE *Device,
    IN OUT PBOOT_ENTRY_OPTION *OtherOptions
    );

NTSTATUS
BlGetBootOptionString (
    IN PBOOT_ENTRY_OPTION Options,
    IN BCDE_DATA_TYPE Type,
    IN OUT PWSTR *String
    );

NTSTATUS
BlInitializeLibrary (
    IN PBOOT_APPLICATION_PARAMETERS ApplicationParameters,
    IN PBOOT_LIBRARY_PARAMETERS LibraryParameters
    );

NTSTATUS
BlDestroyLibrary (
    VOID
    );

#endif