#include <gtest/gtest.h>

extern "C" {
  #include "config.h"
  #include "kaslr.h"
  #include "memory.h"
  #include "kernel_access.h"
}

typedef struct __attribute__ ((packed)) {
  uint32_t version;
  uint32_t size;

  uint64_t va_bits;
  uint64_t va_start;
  uint64_t page_offset;
  uint64_t kimage_vaddr;
  uint64_t kimage_voffset;

  KernelAddress gaf_addr;
  KernelAddress proca_table_addr;

  MemoryRange sys_ram_ranges[kMaxMemoryRangesNum];
  uint64_t sys_ram_ranges_num;

  uint32_t magic;
} Config64_t;

typedef struct __attribute__ ((packed)) {
  char prefix[0x28];
  uint64_t pa_config_text_offset;
} KernelHeader;

static const uint32_t kKaslrOffset = 0x1000;

static Config64_t g_pa_config_buff;
static KernelHeader g_kernel_header;

static const PhysicalAddress kInvalidAddr = NULL;
static const PhysicalAddress kKernelTextSymPa = (PhysicalAddress)&g_kernel_header;
static const PhysicalAddress kPaConfigTextOffsetPa = (PhysicalAddress)&g_kernel_header.pa_config_text_offset;

static const uint32_t kInvalidPaConfigVersion = 0xffffffff;
static const uint32_t kInvalidPaConfigMagic   = 0xffffffff;
static const uint32_t kInvalidPaConfigSize    = 0xffffffff;
static const uint32_t kInvalidSysRamRangesNum = kMaxMemoryRangesNum + 1;
static const uint32_t kDefaultSysRamRangesNum = 3;
static const uint32_t kValidShiftValue = 10;
static const uint32_t kInvalidShiftValue = 42;

static PhysicalAddress g_pa_config_text_offset = kPaConfigTextOffsetPa;

static size_t gPhysOffsetCounter = 0;
static uint32_t g_shift_value;

static uint32_t g_is_kaslr_enabled = 0;
static PaTzResult g_get_kaslr_result = PA_TZ_SUCCESS;

static PaTzResult g_load_gaf_info_return = PA_TZ_SUCCESS;

extern "C" PaTzResult LoadGafInfo(KernelAddress gaf_virt, GafInfo *gaf) {
  memset(gaf, 0, sizeof(*gaf));
  return g_load_gaf_info_return;
}

PaTzResult KernelNextMemoryConfiguration(void) {
  if (gPhysOffsetCounter == 0) {
    gPhysOffsetCounter++;
    return PA_TZ_SUCCESS;
  }

  return PA_TZ_GENERAL_ERROR;
}

extern "C" PaTzResult KernelGetKaslrOffset(uint32_t *kaslr_offset) {
  if (g_is_kaslr_enabled) {
    *kaslr_offset = kKaslrOffset;
  } else {
    *kaslr_offset = 0;
  }

  return g_get_kaslr_result;
}

extern "C" PaTzResult PhysicalGetBytes(PhysicalAddress phys, size_t size, void *out) {
  if (phys == kInvalidAddr) {
    return PA_TZ_MEMORY_PERM_ERROR;
  }
  memcpy(out, (void *)phys, size);
  return PA_TZ_SUCCESS;
}

extern "C" PaTzResult KernelAccessGetUint32(KernelAccessInfo *context,
                                            KernelAddress address, uint32_t *out) {
  if (address == kInvalidAddr) {
    return PA_TZ_MEMORY_PERM_ERROR;
  }
  *out = g_shift_value;
  return PA_TZ_SUCCESS;
}

extern "C" PhysicalAddress KernelGetPaConfigOffsetSymAddress() {
  return g_pa_config_text_offset - g_is_kaslr_enabled * kKaslrOffset;
}

extern "C" PhysicalAddress KernelGetTextSymAddress() {
  return kKernelTextSymPa - g_is_kaslr_enabled * kKaslrOffset;
}


class ConfigInitTest : public ::testing::Test {
protected:
  virtual void SetUp() {
    gPhysOffsetCounter = 0;
    g_is_kaslr_enabled = 0;
    g_get_kaslr_result = PA_TZ_SUCCESS;
    g_pa_config_text_offset = kPaConfigTextOffsetPa;
    g_kernel_header.pa_config_text_offset = (PhysicalAddress)&g_pa_config_buff - (PhysicalAddress)&g_kernel_header;

    g_pa_config_buff.size = sizeof(g_pa_config_buff);
    g_pa_config_buff.magic = kPaConfigMagic;
    g_pa_config_buff.version = kPaConfigVersion;
    g_pa_config_buff.sys_ram_ranges_num = kDefaultSysRamRangesNum;
    g_load_gaf_info_return = PA_TZ_SUCCESS;
    g_pa_config_buff.proca_table_addr = kInvalidAddr + 1;
    g_shift_value = kValidShiftValue;
  }

  virtual void TearDown() {
  }
};

TEST_F(ConfigInitTest, GetPaConfigOffsetFailed) {
  g_pa_config_text_offset = kInvalidAddr;
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_GENERAL_ERROR);
  EXPECT_FALSE(ConfigIsInited()) <<
      "Config SHOULD not be inited if initialization failed";
}

TEST_F(ConfigInitTest, GetPaConfigBufferFailed) {
  g_kernel_header.pa_config_text_offset = kInvalidAddr - (PhysicalAddress)&g_kernel_header;
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_GENERAL_ERROR);
  EXPECT_FALSE(ConfigIsInited()) <<
      "Config SHOULD not be inited if initialization failed";
}

TEST_F(ConfigInitTest, PaConfigInvalidSizeFailed) {
  g_pa_config_buff.size = kInvalidPaConfigSize;
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_GENERAL_ERROR);
  EXPECT_FALSE(ConfigIsInited()) <<
      "Config SHOULD not be inited if initialization failed";
}

TEST_F(ConfigInitTest, PaConfigInvalidSysRamRangesNumFailed) {
  g_pa_config_buff.sys_ram_ranges_num = kInvalidSysRamRangesNum;
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_GENERAL_ERROR);
  EXPECT_FALSE(ConfigIsInited()) <<
      "Config SHOULD not be inited if initialization failed";
}

TEST_F(ConfigInitTest, PaConfigInvalidVersionFailed) {
  g_pa_config_buff.version = kInvalidPaConfigVersion;
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_GENERAL_ERROR);
  EXPECT_FALSE(ConfigIsInited()) <<
      "Config SHOULD not be inited if initialization failed";
}

TEST_F(ConfigInitTest, PaConfigInvalidMagicFailed) {
  g_pa_config_buff.magic = kInvalidPaConfigMagic;
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_GENERAL_ERROR);
  EXPECT_FALSE(ConfigIsInited()) <<
      "Config SHOULD not be inited if initialization failed";
}

TEST_F(ConfigInitTest, LoadGafInfoFailed) {
  g_load_gaf_info_return = PA_TZ_GENERAL_ERROR;
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_GENERAL_ERROR);
  EXPECT_FALSE(ConfigIsInited()) <<
      "Config SHOULD not be inited if initialization failed";
}

TEST_F(ConfigInitTest, HashTablesShiftLoadFailed) {
  g_pa_config_buff.proca_table_addr = kInvalidAddr;
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_GENERAL_ERROR);
  EXPECT_FALSE(ConfigIsInited()) <<
      "Config SHOULD not be inited if initialization failed";
}

TEST_F(ConfigInitTest, HashTablesZeroShiftValueFailed) {
  g_shift_value = 0;
  PaTzResult result = ConfigInit();
  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
  EXPECT_FALSE(ConfigIsInited()) <<
      "Config SHOULD not be inited if initialization failed";
}

TEST_F(ConfigInitTest, HashTablesInvalidShiftValueFailed) {
  g_shift_value = kInvalidShiftValue;
  PaTzResult result = ConfigInit();
  EXPECT_EQ(result, PA_TZ_GENERAL_ERROR);
  EXPECT_FALSE(ConfigIsInited()) <<
      "Config SHOULD not be inited if initialization failed";
}

TEST_F(ConfigInitTest, KaslrIsEnabled) {
  g_is_kaslr_enabled = 1;
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_SUCCESS);
}

TEST_F(ConfigInitTest, KaslrIsDisabled) {
  g_is_kaslr_enabled = 0;
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_SUCCESS);
}

TEST_F(ConfigInitTest, KaslrIsFailed) {
  PaTzResult result;

  g_get_kaslr_result = PA_TZ_GENERAL_ERROR;
  result = ConfigInit();;

  EXPECT_EQ(PA_TZ_GENERAL_ERROR, result) <<
      "Initialization SHOULD fail if getting Kaslr is failed";
  EXPECT_FALSE(ConfigIsInited()) <<
      "Config SHOULD not be inited if initialization failed";
}

TEST_F(ConfigInitTest, FiveIsEnabled) {
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_SUCCESS);
  EXPECT_TRUE(IsFiveEnabled());
}

TEST_F(ConfigInitTest, FiveIsDisabled) {
  g_kernel_header.pa_config_text_offset = kFiveDisabledMagicNumber;
  PaTzResult result = ConfigInit();
  EXPECT_TRUE(result == PA_TZ_SUCCESS);
  EXPECT_FALSE(IsFiveEnabled());
}