#include <Uefi.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/DebugLib.h>

// 1. 프로토콜 구조체 정의 (다시 심플하게 1개만!)
typedef struct _VULN_TEST_PROTOCOL {
  VOID* (EFIAPI *ProcessImageSize) (UINT32 X, UINT32 Y);
} VULN_TEST_PROTOCOL;

// =======================================================
// [SINK] 메모리 할당만 전담하는 래퍼 함수 (인자 1개)
// =======================================================
__attribute__((noinline, used, optimize("O0")))
VOID* EFIAPI AllocatePoolWrapper (IN UINT32 Size) {
  VOID* Buffer = NULL;
  
  // 실제 UEFI 서비스를 여기서 호출
  gBS->AllocatePool(EfiBootServicesData, (UINTN)Size, &Buffer);
  
  return Buffer;
}

// =======================================================
// [CORE] 정수 오버플로우 계산 후 래퍼를 호출하는 함수
// =======================================================
__attribute__((noinline, used, optimize("O0")))
VOID* EFIAPI VulnerableProtocolFunction (IN UINT32 X, IN UINT32 Y) {
  UINT32 AllocationSize = X * Y * 4;
  VOID* Buffer = NULL;

  if (AllocationSize > 0) {
    // 직접 호출하는 대신, 인자 하나(Size)만 넣어서 래퍼 함수를 호출!
    Buffer = AllocatePoolWrapper(AllocationSize);
  }
  
  return Buffer;
}

// =======================================================
// [WRAPPER] 눈치 볼 필요 없는 아주 깔끔한 래퍼 함수
// =======================================================
__attribute__((noinline, used, optimize("O0")))
VOID* EFIAPI ProtocolWrapperFunction (IN UINT32 ParamX, IN UINT32 ParamY) {
  // 꼬리 물기 방지? 더미 연산? 이제 다 필요 없어! 바로 리턴해!
  return VulnerableProtocolFunction(ParamX, ParamY);
}

// 2. 구조체 변수에 래퍼 하나만 깔끔하게 등록
VULN_TEST_PROTOCOL mMyVulnProtocol = {
  ProtocolWrapperFunction
};

// -------------------------------------------------------
// [ENTRY] 드라이버 엔트리 포인트
// -------------------------------------------------------
EFI_STATUS EFIAPI test_1EntryPoint (IN EFI_HANDLE ImageHandle, IN EFI_SYSTEM_TABLE *SystemTable) {
  EFI_HANDLE Handle = NULL;
  EFI_GUID gGuid = { 0x11223344, 0x5566, 0x7788, { 0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE, 0xFF, 0x00 } };

  return gBS->InstallMultipleProtocolInterfaces(&Handle, &gGuid, &mMyVulnProtocol, NULL);
}