/**
 * PROJECT:         ExectOS
 * COPYRIGHT:       See COPYING.md in the top level directory
 * FILE:            xtldr/modules/elf/elf.c
 * DESCRIPTION:     ELF executable file format support module
 * DEVELOPERS:      Jozef Nagy <schkwve@gmail.com>
 */

#include <elf.h>

/* ELF module information */
XTBL_MODINFO = L"ELF executable file format support";

/* EFI XT Loader Protocol */
PXTBL_LOADER_PROTOCOL XtLdrProtocol;

/* XTOS ELF Image Protocol */
XTBL_EXECUTABLE_IMAGE_PROTOCOL XtElfProtocol;

/**
 * Returns the address of the entry point.
 *
 * @param Image
 *        A pointer to the ELF context structure representing the loaded image.
 *
 * @param EntryPoint
 *        A pointer to the memory area where address of the image entry point will be stored.
 *
 * @return This routine returns a status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
ElfGetEntryPoint(IN PVOID ImagePointer,
                 OUT PVOID *EntryPoint)
{
    PELF_IMAGE_CONTEXT Image = ImagePointer;

    /* Validate input data */
    if(!Image)
    {
        /* Invalid parameter passed */
        return STATUS_EFI_INVALID_PARAMETER;
    }

    /* Save entry point and return success */
    *EntryPoint = Image->EntryPoint;
    return STATUS_EFI_SUCCESS;
}

XTCDECL
EFI_STATUS
ElfGetMachineType(IN PVOID ImagePointer,
                  OUT PUSHORT MachineType)
{
    PELF_IMAGE_CONTEXT Image = ImagePointer;

    *MachineType = Image->Header32->e_machine;
    return STATUS_EFI_SUCCESS;
}

XTCDECL
EFI_STATUS
ElfGetSubSystem(IN PVOID ImagePointer,
                OUT PUSHORT SubSystem)
{
    PELF_IMAGE_CONTEXT Image = ImagePointer;
    
    *SubSystem = Image->Header32->e_ident[EI_OSABI];
    return STATUS_EFI_SUCCESS;
}

XTCDECL
EFI_STATUS
ElfRelocateImage(IN PVOID ImagePointer,
                 IN EFI_VIRTUAL_ADDRESS Address)
{
    /* No relocation yet */
    return STATUS_EFI_SUCCESS;
}

/**
 * Loads an ELF image file.
 *
 * @param FileHandle
 *        The handle of the opened ELF file.
 *
 * @param MemoryType
 *        Supplies the type of memory to be assigned to the memory descriptor.
 *
 * @param VirtualAddress
 *        Optional virtual address pointing to the memory area where ELF file will be loaded.
 *
 * @param Image
 *        Supplies pointer to the memory area where loaded ELF image context will be stored.
 *
 * @return This routine returns status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
ElfLoadImage(IN PEFI_FILE_HANDLE FileHandle,
             IN LOADER_MEMORY_TYPE MemoryType,
             IN PVOID VirtualAddress,
             OUT PVOID *ImagePointer)
{
    EFI_GUID FileInfoGuid = EFI_FILE_INFO_PROTOCOL_GUID;
    PELF_IMAGE_CONTEXT ImageData;
    EFI_PHYSICAL_ADDRESS Address;
    PEFI_FILE_INFO FileInfo;
    UINT_PTR ReadSize;
    EFI_STATUS Status;
    // UINT SectionSize;
    SIZE_T Pages;
    PUCHAR Data;

    /* Set required size for getting file information */
    ReadSize = sizeof(EFI_FILE_INFO) + 32;

    /* Allocate necessary amount of memory */
    Status = XtLdrProtocol->Memory.AllocatePool(ReadSize, (PVOID *)&FileInfo);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Memory allocation failure */
        XtLdrProtocol->Debug.Print(L"ERROR: Memory pool allocation failure\n");
        return Status;
    }

    /* First attempt to get file information */
    Status = FileHandle->GetInfo(FileHandle, &FileInfoGuid, &ReadSize, FileInfo);
    if(Status == STATUS_EFI_BUFFER_TOO_SMALL)
    {
        /* Buffer it too small, but EFI tells the required size, let's reallocate */
        XtLdrProtocol->Memory.FreePool(&FileInfo);
        Status = XtLdrProtocol->Memory.AllocatePool(ReadSize, (PVOID *)&FileInfo);
        if(Status != STATUS_EFI_SUCCESS)
        {
            /* Memory allocation failure */
            XtLdrProtocol->Debug.Print(L"ERROR: Memory pool allocation failure\n");
            return Status;
        }

        /* Second attempt to get file information */
        Status = FileHandle->GetInfo(FileHandle, &FileInfoGuid, &ReadSize, FileInfo);
    }
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Unable to get file information */
        XtLdrProtocol->Debug.Print(L"ERROR: Failed to get file information\n");
        return Status;
    }

    /* Allocate memory for storing image data */
    Status = XtLdrProtocol->Memory.AllocatePool(sizeof(ELF_IMAGE_CONTEXT), (PVOID *)&ImageData);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Memory allocation failure */
        XtLdrProtocol->Debug.Print(L"ERROR: Memory pool allocation failure\n");
        return Status;
    }

    /* Store file size and memory type, nullify data and free up memory */
    ImageData->Data = NULL;
    ImageData->FileSize = FileInfo->FileSize;
    ImageData->MemoryType = MemoryType;
    XtLdrProtocol->Memory.FreePool(FileInfo);

    /* Calculate number of pages */
    Pages = EFI_SIZE_TO_PAGES(ImageData->FileSize);

    /* Allocate pages */
    Status = XtLdrProtocol->Memory.AllocatePages(Pages, &Address);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Pages allocation failure */
        XtLdrProtocol->Debug.Print(L"ERROR: Pages allocation failure\n");
        XtLdrProtocol->Memory.FreePool(ImageData);
        return Status;
    }

    /* Read ELF image */
    ReadSize = Pages * EFI_PAGE_SIZE;
    Data = (PUCHAR)(UINT_PTR)Address;
    Status = FileHandle->Read(FileHandle, &ReadSize, Data);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to read data */
        XtLdrProtocol->Debug.Print(L"ERROR: Unable to read ELF image file\n");
        XtLdrProtocol->Memory.FreePages(Pages, (EFI_PHYSICAL_ADDRESS)(UINT_PTR)Data);
        XtLdrProtocol->Memory.FreePool(ImageData);
        return Status;
    }

    /* Extract header */
    ImageData->Header32 = (PELF_IMAGE_HEADER32)Data;
    ImageData->Header64 = (PELF_IMAGE_HEADER64)Data;

    /* Set physical and virtual addresses */
    ImageData->PhysicalAddress = (PVOID)(UINT_PTR)Address;
    if(VirtualAddress)
    {
        ImageData->VirtualAddress = VirtualAddress;
    }
    else
    {
        ImageData->VirtualAddress = (PVOID)(UINT_PTR)Address;
    }

    /* Validate ELF file */
    if(ImageData->Header32->e_ident[EI_MAG0] != 0x7F ||
       ImageData->Header32->e_ident[EI_MAG1] != 0x45 ||
       ImageData->Header32->e_ident[EI_MAG2] != 0x4C ||
       ImageData->Header32->e_ident[EI_MAG3] != 0x46)
    {
        /* ELF file header is invalid */
        return STATUS_EFI_INVALID_PARAMETER;
    }

    /* Check architecture */
    if(ImageData->Header32->e_ident[EI_CLASS] == 1)
    {
        /* 32-bit executable */
        XtLdrProtocol->Debug.Print(L"Kernel is a 32-bit executable\n");
        
        /* Set entry point */
        ImageData->EntryPoint = (PVOID)(UINT_PTR)ImageData->Header32->e_entry;

        /* Load individual segments according to program headers */
        PELF_IMAGE_PROGRAM_HEADER32 ProgramHeaders = (PELF_IMAGE_PROGRAM_HEADER32)(Data + ImageData->Header32->e_phoff);
        for (UINT Count = 0; Count < ImageData->Header32->e_phnum; Count++)
        {
            /* PT_DYNAMIC will be used for relocations */
            if(ProgramHeaders[Count].p_type == PT_LOAD)
            {
                UINT32 PageCount = EFI_SIZE_TO_PAGES((UINT32)ProgramHeaders[Count].p_memsz);

                Status = XtLdrProtocol->Memory.AllocatePages(PageCount, (PEFI_PHYSICAL_ADDRESS)&PageCount);
                if(Status != STATUS_EFI_SUCCESS)
                {
                    /* Failed to allocate memory */
                    XtLdrProtocol->Debug.Print(L"ERROR: Page allocation failure\n");
                    return STATUS_EFI_PROTOCOL_ERROR;
                }

                /* Read segment into memory */
                Status = FileHandle->SetPosition(FileHandle, (UINT64)&ProgramHeaders[Count].p_offset);
                Status = FileHandle->Read(FileHandle, (PUINT_PTR)&ProgramHeaders[Count].p_filesz, &ProgramHeaders[Count].p_paddr);
                if(Status != STATUS_EFI_SUCCESS)
                {
                    /* Failed to read data */
                    XtLdrProtocol->Debug.Print(L"ERROR: Unable to load ELF segment into memory!\n");
                    XtLdrProtocol->Memory.FreePages(Pages, (EFI_PHYSICAL_ADDRESS)(UINT_PTR)&ProgramHeaders[Count].p_paddr);
                    return Status;
                }

                /* If memory size is larger than file size, zero the memory out */
                UINT32 ZeroCount = ProgramHeaders[Count].p_memsz - ProgramHeaders[Count].p_filesz;
                if(ZeroCount > 0)
                {
                    UINT_PTR ZeroedMemoryStart = (UINT_PTR)ProgramHeaders[Count].p_paddr + (UINT_PTR)ProgramHeaders[Count].p_filesz;
                    RtlZeroMemory((PVOID)ZeroedMemoryStart, ZeroCount);
                }
            }
		}
    }
    else if(ImageData->Header32->e_ident[EI_CLASS] == 2)
    {
        /* 64-bit executable */
        XtLdrProtocol->Debug.Print(L"Kernel is a 64-bit executable\n");
        
        /* Set entry point */
        ImageData->EntryPoint = (PVOID)(UINT_PTR)ImageData->Header64->e_entry;
        XtLdrProtocol->Debug.Print(L"Entry point: 0x%lx (0x%lx)\n", ImageData->EntryPoint, ImageData->Header64->e_entry);

        /* Load individual segments according to program headers */
        PELF_IMAGE_PROGRAM_HEADER64 ProgramHeaders = (PELF_IMAGE_PROGRAM_HEADER64)(Data + ImageData->Header64->e_phoff);
        for (UINT Count = 0; Count < ImageData->Header64->e_phnum; Count++)
        {
            /* PT_DYNAMIC will be used for relocations */
            if(ProgramHeaders[Count].p_type == PT_LOAD)
            {
                /* Allocate memory for the program headers */
                UINT32 PageCount = EFI_SIZE_TO_PAGES((UINT32)ProgramHeaders[Count].p_memsz);

                Status = XtLdrProtocol->Memory.AllocatePages(PageCount, (PEFI_PHYSICAL_ADDRESS)&PageCount);
                if(Status != STATUS_EFI_SUCCESS)
                {
                    /* Failed to allocate memory */
                    XtLdrProtocol->Debug.Print(L"ERROR: Page allocation failure\n");
                    return STATUS_EFI_PROTOCOL_ERROR;
                }

                /* Read segment into memory */
                Status = FileHandle->SetPosition(FileHandle, (UINT64)&ProgramHeaders[Count].p_offset);
                Status = FileHandle->Read(FileHandle, (PUINT_PTR)&ProgramHeaders[Count].p_filesz, (PVOID)(UINT_PTR)ProgramHeaders[Count].p_paddr);
                if(Status != STATUS_EFI_SUCCESS)
                {
                    /* Failed to read data */
                    XtLdrProtocol->Debug.Print(L"ERROR: Unable to load ELF segment into memory!\n");
                    XtLdrProtocol->Memory.FreePages(Pages, (EFI_PHYSICAL_ADDRESS)(UINT_PTR)&ProgramHeaders[Count].p_paddr);
                    return Status;
                }

                /* If memory size is larger than file size, zero the memory out */
                UINT32 ZeroCount = ProgramHeaders[Count].p_memsz - ProgramHeaders[Count].p_filesz;
                if(ZeroCount > 0)
                {
                    UINT_PTR ZeroedMemoryStart = (UINT_PTR)ProgramHeaders[Count].p_paddr + (UINT_PTR)ProgramHeaders[Count].p_filesz;
                    RtlZeroMemory((PVOID)ZeroedMemoryStart, ZeroCount);
                }
            }
		}
    }
    else
    {
        /* Invalid executable */
        XtLdrProtocol->Debug.Print(L"ERROR: ELF executable has invalid architecture\n");
        return STATUS_EFI_UNSUPPORTED;
    }

    /* Check endianness */
    if(ImageData->Header32->e_ident[EI_DATA] != 1)
    {
        /* Big-endian */
        XtLdrProtocol->Debug.Print(L"ERROR: XTLDR only supports LSB (little-endian) ELF executables\n");
        return STATUS_EFI_UNSUPPORTED;
    }

    /* Store image data */
    *ImagePointer = ImageData;

    /* Return SUCCESS */
    return STATUS_EFI_SUCCESS;
}

/**
 * This routine is the entry point of the XT EFI boot loader module.
 *
 * @param ImageHandle
 *        Firmware-allocated handle that identifies the image.
 *
 * @param SystemTable
 *        Provides the EFI system table.
 *
 * @return This routine returns status code.
 *
 * @since XT 1.0
 */
XTCDECL
EFI_STATUS
XtLdrModuleMain(IN EFI_HANDLE ImageHandle,
                IN PEFI_SYSTEM_TABLE SystemTable)
{
    EFI_GUID Guid = XT_ELF_IMAGE_PROTOCOL_GUID;
    EFI_STATUS Status;

    /* Open the XTLDR protocol */
    Status = BlGetXtLdrProtocol(SystemTable, ImageHandle, &XtLdrProtocol);
    if(Status != STATUS_EFI_SUCCESS)
    {
        /* Failed to open loader protocol */
        return STATUS_EFI_PROTOCOL_ERROR;
    }

    /* Set routines available via ELF image protocol */
    XtElfProtocol.GetEntryPoint = ElfGetEntryPoint;
    XtElfProtocol.GetMachineType = ElfGetMachineType;
    XtElfProtocol.GetSubSystem = ElfGetSubSystem;
    XtElfProtocol.LoadImage = ElfLoadImage;
    XtElfProtocol.RelocateImage = ElfRelocateImage;

    /* Register ELF protocol */
    return XtLdrProtocol->Protocol.Install(&XtElfProtocol, &Guid);
}