#include "../include/vyx.h"
#include "Base.h"
#include <Uefi.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/BaseMemoryLib.h>
#define PAGE_PRESENT (1ULL << 0)
#define PAGE_TABLE_FLAGS (1ULL<<1)
#define PAGE_RW       (1ULL << 1)
#define PAGE_US       (1ULL << 2)
#define PAGE_PWT      (1ULL << 3)
#define PAGE_PCD      (1ULL << 4)
#define PAGE_ACCESSED (1ULL << 5)
#define PAGE_DIRTY    (1ULL << 6)
#define PAGE_PS       (1ULL << 7)
#define PAGE_GLOBAL   (1ULL << 8)
#define PAGE_NX       (1ULL << 63)
static BOOLEAN is_allocator_init=FALSE;
static UINT64 pt_pool_page_offset=0;
UINTN pad_size(UINTN value,UINTN multiple) {
  if (multiple==0) return value;
  UINTN remainder=value%multiple;
  if (remainder==0) return value;
  return value+multiple-remainder;
}
EFI_STATUS bp_vyx_get_header(bp_disk_FILE *file,bp_vyx_VYX_HEADER **header) {
  if (file==NULL || header==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  if (file->file_size_bytes<sizeof(bp_vyx_VYX_HEADER)) {
    return EFI_BUFFER_TOO_SMALL;
  }
  UINT8 vyx_sig[3]={'V','y','X'};
  if (CompareMem(vyx_sig,file->content,3)!=0) {
    return EFI_COMPROMISED_DATA;
  }
  *header=(bp_vyx_VYX_HEADER *)file->content;
  if ((*header)->vyx_ver!=BP_VYX_SUPPORTED_VERSION) {
    return EFI_UNSUPPORTED;
  }
  return EFI_SUCCESS;
}
EFI_STATUS bp_vyx_create_shelter_boot_conf(UINT8 log_level,BOOLEAN test_benchmark,UINT64 bench_iterations,BOOLEAN log_disable_serial_port,BOOLEAN disable_serial_port,UINT16 log_ring_size,bp_vyx_SHELTER_BOOT_CONFIG *boot_config) {
  if (boot_config==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  boot_config->log_level=log_level;
  boot_config->test_benchmark=test_benchmark;
  boot_config->bench_iterations=bench_iterations;
  boot_config->log_disable_serial_port=log_disable_serial_port;
  boot_config->disable_serial_port=disable_serial_port;
  boot_config->log_ring_size=log_ring_size;
  UINT8 sig_start[8]={'S','h','C','f','g','B','e','g'};
  UINT8 sig_end[8]={'S','h','C','f','g','E','n','d'};
  for (UINTN i=0;i<8;i++) {boot_config->sig_start[i]=sig_start[i];boot_config->sig_end[i]=sig_end[i];};
  return EFI_SUCCESS;
}
EFI_STATUS bp_vyx_init_pt_pool_allocator() {
  pt_pool_page_offset=0;
  is_allocator_init=TRUE;
  return EFI_SUCCESS;
}
UINT64 bp_vyx_get_pt_pool_allocator() {
  return pt_pool_page_offset;
}
EFI_PHYSICAL_ADDRESS bp_vyx_pt_pool_alloc_one_page(EFI_PHYSICAL_ADDRESS pt_pool_base) {
  if (!is_allocator_init) return 0;
  if (pt_pool_base==0) {
    return (EFI_PHYSICAL_ADDRESS)0;
  }
  if (pt_pool_page_offset>=BP_VYX_SHELTER_PT_PAGE_COUNT) {
    return (EFI_PHYSICAL_ADDRESS)0;
  }
  EFI_PHYSICAL_ADDRESS out=pt_pool_base+pt_pool_page_offset*4096;
  pt_pool_page_offset++;
  return out;
}
EFI_STATUS bp_vyx_setup_physical_pages(bp_vyx_VYX_HEADER *header,bp_disk_FILE *vyx_file,bp_disk_FILE *keycard_file,bp_vyx_SHELTER_BOOT_CONFIG *shelter_boot_config,bp_vyx_LOADER_SETTING **loader_setting) {
  if (header==NULL || vyx_file==NULL || loader_setting==NULL || keycard_file==NULL || shelter_boot_config==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  UINTN text_size=header->text_size_padded;
  if (text_size==0) {return EFI_COMPROMISED_DATA;};
  if (text_size%4096!=0) {text_size=pad_size(text_size,4096);};
  UINTN data_size=header->data_size_padded;
  if (data_size!=0 && data_size%4096!=0) {data_size=pad_size(data_size,4096);};
  UINTN rodata_size=header->rodata_size_padded;
  if (rodata_size!=0 && rodata_size%4096!=0) {rodata_size=pad_size(rodata_size,4096);};
  UINTN bss_size=header->bss_size_padded;
  if (bss_size!=0 && bss_size%4096!=0) {bss_size=pad_size(bss_size,4096);};
  UINTN keycard_size=keycard_file->file_size_bytes;
  if (keycard_size==0) {return EFI_COMPROMISED_DATA;};
  if (keycard_size%4096!=0) {keycard_size=pad_size(keycard_size,4096);};
  UINTN keycard_pages=(keycard_size+4095)/4096;
  *loader_setting=AllocatePool(sizeof(bp_vyx_LOADER_SETTING));
  UINTN shelter_boot_config_size=sizeof(bp_vyx_SHELTER_BOOT_CONFIG);
  if (shelter_boot_config_size%4096!=0) {shelter_boot_config_size=pad_size(shelter_boot_config_size,4096);};
  UINTN shelter_boot_config_pages=(shelter_boot_config_size+4095)/4096;
  if (!*loader_setting) {
    return EFI_OUT_OF_RESOURCES;
  }
  EFI_STATUS status=gBS->AllocatePages(AllocateAnyPages,EfiLoaderCode,text_size/4096,&(*loader_setting)->text_addr_phys);
  if (EFI_ERROR(status)) {
    FreePool(*loader_setting);
    return status;
  }
  if (data_size!=0) {
    status=gBS->AllocatePages(AllocateAnyPages,EfiLoaderData,data_size/4096,&(*loader_setting)->data_addr_phys);
    if (EFI_ERROR(status)) {
      FreePool(*loader_setting);
      return status;
    }
  } else {
    (*loader_setting)->data_addr_phys=0;
  }
  if (rodata_size!=0) {
    status=gBS->AllocatePages(AllocateAnyPages,EfiLoaderData,rodata_size/4096,&(*loader_setting)->rodata_addr_phys);
    if (EFI_ERROR(status)) {
      FreePool(*loader_setting);
      return status;
    }
  } else {
    (*loader_setting)->rodata_addr_phys=0;
  }
  if (bss_size!=0) {
    status=gBS->AllocatePages(AllocateAnyPages,EfiLoaderData,bss_size/4096,&(*loader_setting)->bss_addr_phys);
    if (EFI_ERROR(status)) {
      FreePool(*loader_setting);
      return status;
    }
  } else {
    (*loader_setting)->bss_addr_phys=0;
  }
  status=gBS->AllocatePages(AllocateAnyPages,EfiLoaderCode,64,&(*loader_setting)->stack_addr_phys);
  if (EFI_ERROR(status)) {
    FreePool(*loader_setting);
    return status;
  }
  status=gBS->AllocatePages(AllocateAnyPages,EfiLoaderCode,keycard_pages,&(*loader_setting)->keycard_addr_phys);
  if (EFI_ERROR(status)) {
    FreePool(*loader_setting);
    return status;
  }
  status=gBS->AllocatePages(AllocateAnyPages,EfiLoaderCode,shelter_boot_config_pages,&(*loader_setting)->shelter_boot_config_addr_phys);
  if (EFI_ERROR(status)) {
    FreePool(*loader_setting);
    return status;
  }
  status=gBS->AllocatePages(AllocateAnyPages,EfiLoaderCode,BP_VYX_SHELTER_MEMORY_MAP_PAGES_COUNT,&(*loader_setting)->shelter_memory_map_addr_phys);
  if (EFI_ERROR(status)) {
    FreePool(*loader_setting);
    return status;
  }
  status=gBS->AllocatePages(AllocateAnyPages,EfiLoaderCode,BP_VYX_SHELTER_PT_PAGE_COUNT,&(*loader_setting)->shelter_page_table_addr_phys);
  if (EFI_ERROR(status)) {
    FreePool(*loader_setting);
    return status;
  }
  status=gBS->AllocatePages(AllocateAnyPages,EfiLoaderCode,MAX(shelter_boot_config->log_ring_size,4),&(*loader_setting)->shelter_logging_ring_buffer_addr_phys);
  if (EFI_ERROR(status)) {
    FreePool(*loader_setting);
    return status;
  }
  (*loader_setting)->keycard_page_count=keycard_pages;
  (*loader_setting)->shelter_boot_config_pages_count=shelter_boot_config_pages;
  if (bss_size!=0) {SetMem((VOID*)(UINTN)(*loader_setting)->bss_addr_phys,bss_size,0);};
  SetMem((VOID*)(UINTN)(*loader_setting)->stack_addr_phys,64*4096,0);
  SetMem((VOID*)(UINTN)(*loader_setting)->shelter_memory_map_addr_phys,16*4096,0);
  SetMem((VOID*)(UINTN)(*loader_setting)->shelter_logging_ring_buffer_addr_phys,shelter_boot_config->log_ring_size*4096,0);
  (*loader_setting)->shelter_logging_ring_buffer_pages_count=MAX(shelter_boot_config->log_ring_size,4);
  CopyMem((VOID*)(UINTN)(*loader_setting)->text_addr_phys,vyx_file->content+sizeof(bp_vyx_VYX_HEADER),header->text_size_padded);
  if (data_size!=0 && rodata_size!=0) {
    CopyMem((VOID*)(UINTN)(*loader_setting)->data_addr_phys,vyx_file->content+sizeof(bp_vyx_VYX_HEADER)+header->text_size_padded,header->data_size_padded);
    CopyMem((VOID*)(UINTN)(*loader_setting)->rodata_addr_phys,vyx_file->content+sizeof(bp_vyx_VYX_HEADER)+header->text_size_padded+header->data_size_padded,header->rodata_size_padded);
  } else if (data_size==0 && rodata_size!=0) {
    CopyMem((VOID*)(UINTN)(*loader_setting)->rodata_addr_phys,vyx_file->content+sizeof(bp_vyx_VYX_HEADER)+header->text_size_padded,header->rodata_size_padded);
  } else if (data_size!=0 && rodata_size==0) {
    CopyMem((VOID*)(UINTN)(*loader_setting)->data_addr_phys,vyx_file->content+sizeof(bp_vyx_VYX_HEADER)+header->text_size_padded,header->data_size_padded);
  }
  if (!(keycard_file->file_size_bytes<=keycard_pages*4096)) return EFI_COMPROMISED_DATA;
  CopyMem((VOID*)(UINTN)(*loader_setting)->keycard_addr_phys,keycard_file->content,keycard_file->file_size_bytes);
  CopyMem((VOID*)(UINTN)(*loader_setting)->shelter_boot_config_addr_phys,(VOID*)shelter_boot_config,sizeof(bp_vyx_SHELTER_BOOT_CONFIG));
  return EFI_SUCCESS;
}
EFI_STATUS map_page(UINT64* pml4,UINT64 va,UINT64 pa,UINT64 flags) {
  UINT64 pml4_i=(va>>39) & 0x1FF;
  UINT64 pdpt_i=(va>>30) & 0x1FF;
  UINT64 pd_i=(va>>21) & 0x1FF;
  UINT64 pt_i=(va>>12) & 0x1FF;
  UINT64* pdpt;
  UINT64* pd;
  UINT64* pt;
  EFI_PHYSICAL_ADDRESS phys_page;
  if (!(pml4[pml4_i] & 1)) {
    phys_page=bp_vyx_pt_pool_alloc_one_page((EFI_PHYSICAL_ADDRESS)pml4);
    if (phys_page==0) return EFI_OUT_OF_RESOURCES;
    pdpt=(UINT64*)(UINTN)phys_page;
    SetMem(pdpt,4096,0);
    pml4[pml4_i]=phys_page | PAGE_TABLE_FLAGS | PAGE_PRESENT;
  } else {
    pdpt=(UINT64*)(UINTN)(pml4[pml4_i] & ~0xFFF);
  }
  if (!(pdpt[pdpt_i] & 1)) {
    phys_page=bp_vyx_pt_pool_alloc_one_page((EFI_PHYSICAL_ADDRESS)pml4);
    if (phys_page==0) return EFI_OUT_OF_RESOURCES;
    pd=(UINT64*)(UINTN)phys_page;
    SetMem(pd,4096,0);
    pdpt[pdpt_i]=phys_page | PAGE_TABLE_FLAGS | PAGE_PRESENT;
  } else {
    pd=(UINT64*)(UINTN)(pdpt[pdpt_i] & ~0xFFF);
  }
  if (!(pd[pd_i] & 1)) {
    phys_page=bp_vyx_pt_pool_alloc_one_page((EFI_PHYSICAL_ADDRESS)pml4);
    if (phys_page==0) return EFI_OUT_OF_RESOURCES;
    pt=(UINT64*)(UINTN)phys_page;
    SetMem(pt,4096,0);
    pd[pd_i]=phys_page | PAGE_TABLE_FLAGS | PAGE_PRESENT;
  } else {
    pt=(UINT64*)(UINTN)(pd[pd_i] & ~0xFFF);
  }
  pt[pt_i]=(pa & ~0xFFF) | flags | PAGE_PRESENT;
  return EFI_SUCCESS;
}
EFI_STATUS bp_vyx_create_page_table(bp_vyx_VYX_HEADER *header,bp_vyx_LOADER_SETTING *settings,EFI_PHYSICAL_ADDRESS *page_table_address,EFI_VIRTUAL_ADDRESS *stack_address,EFI_VIRTUAL_ADDRESS *page_table_pool_va) {
  if (header==NULL || settings==NULL || page_table_address==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  UINTN text_size=header->text_size_padded;
  if (text_size==0) {return EFI_COMPROMISED_DATA;};
  if (text_size%4096!=0) {text_size=pad_size(text_size,4096);};
  UINTN data_size=header->data_size_padded;
  if (data_size!=0 && data_size%4096!=0) {data_size=pad_size(data_size,4096);};
  UINTN rodata_size=header->rodata_size_padded;
  if (rodata_size!=0 && rodata_size%4096!=0) {rodata_size=pad_size(rodata_size,4096);};
  UINTN bss_size=header->bss_size_padded;
  if (bss_size!=0 && bss_size%4096!=0) {bss_size=pad_size(bss_size,4096);};
  UINTN text_pages_count=text_size/4096;
  UINTN data_pages_count=data_size/4096;
  UINTN rodata_pages_count=rodata_size/4096;
  UINTN bss_pages_count=bss_size/4096;
  EFI_VIRTUAL_ADDRESS va=header->text_offset;
  EFI_VIRTUAL_ADDRESS va_pt_pool=(header->text_offset+header->text_size_padded+text_size+data_size+rodata_size+bss_size+0x200000ULL) & ~0xFFFULL;
  if (va_pt_pool+BP_VYX_SHELTER_PT_PAGE_COUNT*4096>header->stack_base) {
    return EFI_COMPROMISED_DATA;
  }
  EFI_STATUS status;
  UINT64 *pml4;
  pml4=(UINT64 *)(UINTN)settings->shelter_page_table_addr_phys;
  SetMem(pml4,4096*BP_VYX_SHELTER_PT_PAGE_COUNT,0);
  bp_vyx_init_pt_pool_allocator();
  for (UINTN i=0;i<BP_VYX_SHELTER_PT_PAGE_COUNT;i++) {
    status=map_page(pml4,va_pt_pool+i*4096,settings->shelter_page_table_addr_phys+i*4096,PAGE_PRESENT | PAGE_NX | PAGE_RW);
    if (EFI_ERROR(status)) return status;
  }
  *page_table_pool_va=va_pt_pool;
  for (UINTN i=0;i<text_pages_count;i++) {
    status=map_page(pml4,va,settings->text_addr_phys+i*4096,PAGE_PRESENT);
    if (EFI_ERROR(status)) return status;
    va+=4096;
  }
  if (settings->data_addr_phys!=0) {
    for (UINTN i=0;i<data_pages_count;i++) {
      status=map_page(pml4,va,settings->data_addr_phys+i*4096,PAGE_PRESENT | PAGE_RW | PAGE_NX);
      if (EFI_ERROR(status)) return status;
      va+=4096;
    }
  }
  if (settings->rodata_addr_phys!=0) {
    for (UINTN i=0;i<rodata_pages_count;i++) {
      status=map_page(pml4,va,settings->rodata_addr_phys+i*4096,PAGE_PRESENT | PAGE_NX);
      if (EFI_ERROR(status)) return status;
      va+=4096;
    }
  }
  if (settings->bss_addr_phys!=0) {
    for (UINTN i=0;i<bss_pages_count;i++) {
      status=map_page(pml4,va,settings->bss_addr_phys+i*4096,PAGE_PRESENT | PAGE_RW | PAGE_NX);
      if (EFI_ERROR(status)) return status;
      va+=4096;
    }
  }
  EFI_VIRTUAL_ADDRESS stack_base=header->stack_base;
  for (UINTN i=1;i<=64;i++) {
    status=map_page(pml4,stack_base+i*4096,settings->stack_addr_phys+(i-1)*4096,PAGE_PRESENT | PAGE_RW | PAGE_NX);
    if (EFI_ERROR(status)) return status;
  }
  *stack_address=stack_base;
  for (UINTN i=0;i<settings->keycard_page_count;i++) {
    status=map_page(pml4,settings->keycard_addr_phys+i*4096,settings->keycard_addr_phys+i*4096,PAGE_PRESENT);
    if (EFI_ERROR(status)) return status;
  }
  for (UINTN i=0;i<settings->shelter_boot_config_pages_count;i++) {
    status=map_page(pml4,BP_VYX_SHELTER_BOOT_CONFIG_VA+i*4096,settings->shelter_boot_config_addr_phys+i*4096,PAGE_PRESENT | PAGE_NX | PAGE_RW);
    if (EFI_ERROR(status)) return status;
  }
  for (UINTN i=0;i<BP_VYX_SHELTER_MEMORY_MAP_PAGES_COUNT;i++) {
    status=map_page(pml4,BP_VYX_SHELTER_MEMORY_MAP_VA+i*4096,settings->shelter_memory_map_addr_phys+i*4096,PAGE_PRESENT | PAGE_NX);
    if (EFI_ERROR(status)) return status;
  }
  for (UINTN i=0;i<settings->shelter_logging_ring_buffer_pages_count;i++) {
    status=map_page(pml4,BP_VYX_SHELTER_LOGGING_RING_BUFFER_VA+i*4096,settings->shelter_logging_ring_buffer_addr_phys+i*4096,PAGE_PRESENT | PAGE_NX | PAGE_RW);
    if (EFI_ERROR(status)) return status;
  }
  *page_table_address=settings->shelter_page_table_addr_phys;
  return EFI_SUCCESS;
}
