#include <gtest/gtest.h>

#include <stdint.h>

extern "C" {
#include "config.h"
#include "task_access.h"
#include "access_control.h"
#include "authentication.h"
#include "api.h"
#include "tlb.h"
}

static const TaskInfo kInitTask = {0};
static const ProcessAddress kCheckAddr = 0x60065;

static uint8_t *g_init_buffer;
static size_t g_init_buffer_size = PAGE_SIZE;
static size_t g_get_buffer_size = PAGE_SIZE;
static size_t g_offset = 0;
static VmaInfo g_vma_task_find_return = {0};
static uint8_t g_not_readable_task_memory = 0;
static uint8_t g_not_writable_task_memory = 0;
static uint8_t g_exec_task_memory = 0;

static PaTzResult g_task_address_to_physical_return = PA_TZ_SUCCESS;
static PaTzResult g_platform_sys_map_return = PA_TZ_SUCCESS;
static PaTzResult g_platform_sys_unmap_return = PA_TZ_SUCCESS;
static PaTzResult g_task_find_return = PA_TZ_SUCCESS;

/*
 * Mocked functions
 */
extern "C" PaTzResult TaskFindVmaWithProcessAddress(const TaskInfo *task,
                                                    ProcessAddress address, 
                                                    VmaInfo *out_vma) {
  *out_vma = g_vma_task_find_return;
  return g_task_find_return;
}

extern "C" PaTzResult TaskAddressToPhysical(TlbConverterInfo* context, 
                                            KernelAddress pgd_address,
                                            ProcessAddress user_virt,
                                            PhysicalAddress *phys,
                                            AccessPermissionFlags* flags) {
  if (g_not_readable_task_memory) {
    flags->read = 0;
  } else {
    flags->read = 1;
  }

  if (g_not_writable_task_memory) {
    flags->write = 0;
  } else {
    flags->write = 1;
  }

  if (g_exec_task_memory) {
    flags->exec = 1;
  }

  *phys = (PhysicalAddress) user_virt;

  return g_task_address_to_physical_return;
}

extern "C" PaTzResult PlatformSysMap(PhysicalAddress phys, size_t size,
                                     MemoryAccessType type, void **virt) {
  *virt = (void *)phys;

  return g_platform_sys_map_return;
}

extern "C" PaTzResult PlatformSysUnmap(void *virt, size_t size) {
  return g_platform_sys_unmap_return;
}
/* End */

static void InitBuffer() {
  if (g_init_buffer) {
    free(g_init_buffer);
  }

  g_init_buffer = (uint8_t *)aligned_alloc(PAGE_SIZE, g_init_buffer_size);

  for (size_t i = 0; i < g_init_buffer_size; i++) {
    g_init_buffer[i] = (int)i % 256;
  }
}

class TaskAccessGetBytesTest : public ::testing::Test {
protected:
  virtual void SetUp() {
    g_task_address_to_physical_return = PA_TZ_SUCCESS;
    g_platform_sys_map_return = PA_TZ_SUCCESS;
    g_platform_sys_unmap_return = PA_TZ_SUCCESS;
    g_not_readable_task_memory = 0;
    g_not_writable_task_memory = 0;
    g_exec_task_memory = 0;

    g_init_buffer_size = PAGE_SIZE;
    g_get_buffer_size = g_init_buffer_size;
    g_offset = 0;
    InitBuffer();
  }
  virtual void TearDown() {
    if (g_init_buffer) {
      free(g_init_buffer);
      g_init_buffer = NULL;
    }
  }
};

class TaskAccessPutBytesTest : public ::testing::Test {
protected:
  virtual void SetUp() {
    g_task_address_to_physical_return = PA_TZ_SUCCESS;
    g_platform_sys_map_return = PA_TZ_SUCCESS;
    g_platform_sys_unmap_return = PA_TZ_SUCCESS;
    g_not_readable_task_memory = 0;
    g_not_writable_task_memory = 0;
    g_exec_task_memory = 0;

    g_init_buffer_size = 3 * PAGE_SIZE;
    g_offset = 0;
    InitBuffer();
  }
  virtual void TearDown() {
    if (g_init_buffer) {
      free(g_init_buffer);
      g_init_buffer = NULL;
    }
  }
};

class IsProcessAddressReadableTest : public ::testing::Test {
protected:
  virtual void SetUp() {
    g_vma_task_find_return.vm_flags = VM_READ;
    g_task_find_return = PA_TZ_SUCCESS;
  }
  virtual void TearDown() {
  }
};

class IsProcessAddressWritableTest : public ::testing::Test {
protected:
  virtual void SetUp() {
    g_vma_task_find_return.vm_flags = VM_WRITE;
    g_task_find_return = PA_TZ_SUCCESS;
  }
  virtual void TearDown() {
  }
};


TEST_F(TaskAccessGetBytesTest, AddrAtStartOfPage) {
  PaTzResult result;
  uint8_t out_buffer[g_get_buffer_size];
  
  result = TaskAccessGetBytes(&kInitTask, (ProcessAddress)g_init_buffer, 
                              g_get_buffer_size, out_buffer);

  EXPECT_TRUE(PA_TZ_SUCCESS == result);
  EXPECT_TRUE(memcmp(g_init_buffer, &out_buffer, sizeof(out_buffer)) == 0);
}

TEST_F(TaskAccessGetBytesTest, AddrAtMiddleOfPage) {
  PaTzResult result;
  g_init_buffer_size = 3 * PAGE_SIZE;
  g_get_buffer_size = 2 * PAGE_SIZE;
  g_offset = PAGE_SIZE / 2;
  uint8_t out_buffer[g_get_buffer_size];

  InitBuffer();
  
  result = TaskAccessGetBytes(&kInitTask, (ProcessAddress)(g_init_buffer + g_offset),
                              g_get_buffer_size, out_buffer);

  EXPECT_TRUE(PA_TZ_SUCCESS == result);
  EXPECT_TRUE(memcmp(g_init_buffer + g_offset, &out_buffer, sizeof(out_buffer)) == 0);
}

TEST_F(TaskAccessGetBytesTest, AddrAtEndOfPage) {
  PaTzResult result;
  g_init_buffer_size = 3 * PAGE_SIZE;
  g_get_buffer_size = 2 * PAGE_SIZE;
  g_offset = PAGE_SIZE - 1;
  uint8_t out_buffer[g_get_buffer_size];

  InitBuffer();
  
  result = TaskAccessGetBytes(&kInitTask, (ProcessAddress)(g_init_buffer + g_offset),
                              g_get_buffer_size, out_buffer);

  EXPECT_TRUE(PA_TZ_SUCCESS == result);
  EXPECT_TRUE(memcmp(g_init_buffer + g_offset, &out_buffer, sizeof(out_buffer)) == 0);
}

TEST_F(TaskAccessGetBytesTest, FailedTaskAddressToPhysical) {
  PaTzResult result;
  uint8_t out_buffer[g_get_buffer_size];
  g_task_address_to_physical_return = PA_TZ_GENERAL_ERROR;

  result = TaskAccessGetBytes(&kInitTask, (ProcessAddress)g_init_buffer,
                              g_get_buffer_size, out_buffer);

  EXPECT_TRUE(PA_TZ_VIRT_TO_PHYS_ERROR == result);
}

TEST_F(TaskAccessGetBytesTest, FailedPlatformSysMap) {
  PaTzResult result;
  uint8_t out_buffer[g_get_buffer_size];
  g_platform_sys_map_return = PA_TZ_GENERAL_ERROR;

  result = TaskAccessGetBytes(&kInitTask, (ProcessAddress)g_init_buffer,
                              g_get_buffer_size, out_buffer);

  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAccessGetBytesTest, FailedPlatformSysUnmap) {
  PaTzResult result;
  uint8_t out_buffer[g_get_buffer_size];
  g_platform_sys_unmap_return = PA_TZ_GENERAL_ERROR;

  result = TaskAccessGetBytes(&kInitTask, (ProcessAddress)(g_init_buffer + g_offset),
                              g_get_buffer_size, out_buffer);

  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAccessGetBytesTest, NonReadbleMemory_Failed) {
  g_not_readable_task_memory = 1;
  PaTzResult result = PA_TZ_SUCCESS;
  uint8_t out_buffer[g_get_buffer_size];

  result = TaskAccessGetBytes(&kInitTask, (ProcessAddress)g_init_buffer,
                              g_get_buffer_size, out_buffer);
  EXPECT_EQ(PA_TZ_MEMORY_PERM_ERROR, result);
}

TEST_F(TaskAccessGetBytesTest, ExecutableMemory) {
  g_exec_task_memory = 1;
  PaTzResult result = PA_TZ_SUCCESS;
  uint8_t out_buffer[g_get_buffer_size];

  result = TaskAccessGetBytes(&kInitTask, (ProcessAddress)g_init_buffer,
                              g_get_buffer_size, out_buffer);
  EXPECT_EQ(PA_TZ_SUCCESS, result);
}

TEST_F(TaskAccessPutBytesTest, AddrAtStartOfPage) {
  PaTzResult result = PA_TZ_SUCCESS;
  uint8_t *out_buffer = (uint8_t *)aligned_alloc(PAGE_SIZE, g_init_buffer_size + PAGE_SIZE);

  result = TaskAccessPutBytes(&kInitTask, g_init_buffer, g_init_buffer_size,
                              (ProcessAddress)out_buffer);

  EXPECT_TRUE(PA_TZ_SUCCESS == result);
  EXPECT_TRUE(memcmp(g_init_buffer, out_buffer, g_init_buffer_size) == 0);

  free(out_buffer);
}

TEST_F(TaskAccessPutBytesTest, AddrAtMiddleOfPage) {
  PaTzResult result = PA_TZ_SUCCESS;
  g_offset = PAGE_SIZE / 2;
  uint8_t *out_buffer = (uint8_t *)aligned_alloc(PAGE_SIZE, g_init_buffer_size + PAGE_SIZE);

  result = TaskAccessPutBytes(&kInitTask, g_init_buffer, g_init_buffer_size,
                              (ProcessAddress)(out_buffer + g_offset));

  EXPECT_TRUE(PA_TZ_SUCCESS == result);
  EXPECT_TRUE(memcmp(g_init_buffer, (out_buffer + g_offset), g_init_buffer_size) == 0);

  free(out_buffer);
}

TEST_F(TaskAccessPutBytesTest, AddrAtEndOfPage) {
  PaTzResult result = PA_TZ_SUCCESS;
  g_offset = PAGE_SIZE - 1;
  uint8_t *out_buffer = (uint8_t *)aligned_alloc(PAGE_SIZE, g_init_buffer_size + PAGE_SIZE);

  result = TaskAccessPutBytes(&kInitTask, g_init_buffer, g_init_buffer_size,
                              (ProcessAddress)(out_buffer + g_offset));

  EXPECT_TRUE(PA_TZ_SUCCESS == result);
  EXPECT_TRUE(memcmp(g_init_buffer, (out_buffer + g_offset), g_init_buffer_size) == 0);

  free(out_buffer);
}

TEST_F(TaskAccessPutBytesTest, FailedTaskAddressToPhysical) {
  PaTzResult result = PA_TZ_SUCCESS;
  uint8_t out_buffer[g_init_buffer_size];
  g_task_address_to_physical_return = PA_TZ_GENERAL_ERROR;

  result = TaskAccessPutBytes(&kInitTask, g_init_buffer, g_init_buffer_size,
                              (ProcessAddress)out_buffer);

  EXPECT_TRUE(PA_TZ_VIRT_TO_PHYS_ERROR == result);
}

TEST_F(TaskAccessPutBytesTest, FailedPlatformSysMap) {
  PaTzResult result = PA_TZ_SUCCESS;
  uint8_t out_buffer[g_init_buffer_size];
  g_platform_sys_map_return = PA_TZ_GENERAL_ERROR;

  result = TaskAccessPutBytes(&kInitTask, g_init_buffer, g_init_buffer_size,
                              (ProcessAddress)&out_buffer);

  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAccessPutBytesTest, FailedPlatformSysUnmap) {
  PaTzResult result = PA_TZ_SUCCESS;
  uint8_t out_buffer[g_init_buffer_size];
  g_platform_sys_unmap_return = PA_TZ_GENERAL_ERROR;

  result = TaskAccessPutBytes(&kInitTask, g_init_buffer, g_init_buffer_size,
                              (ProcessAddress)&out_buffer);

  EXPECT_TRUE(PA_TZ_GENERAL_ERROR == result);
}

TEST_F(TaskAccessPutBytesTest, NonWritableMemory_Failed) {
  g_not_writable_task_memory = 1;
  PaTzResult result = PA_TZ_SUCCESS;
  uint8_t *out_buffer = (uint8_t *)aligned_alloc(PAGE_SIZE, g_init_buffer_size + PAGE_SIZE);

  result = TaskAccessPutBytes(&kInitTask, g_init_buffer, g_init_buffer_size,
                              (ProcessAddress)out_buffer);
  EXPECT_EQ(PA_TZ_MEMORY_PERM_ERROR, result);
}

TEST_F(TaskAccessPutBytesTest, ExecutableMemory_Failed) {
  g_exec_task_memory = 1;
  PaTzResult result = PA_TZ_SUCCESS;
  uint8_t *out_buffer = (uint8_t *)aligned_alloc(PAGE_SIZE, g_init_buffer_size + PAGE_SIZE);

  result = TaskAccessPutBytes(&kInitTask, g_init_buffer, g_init_buffer_size,
                              (ProcessAddress)out_buffer);
  EXPECT_EQ(PA_TZ_MEMORY_PERM_ERROR, result);
}

TEST_F(IsProcessAddressReadableTest, AllParamAreValid) {
  int32_t result;

  result = IsProcessAddressReadable(&kInitTask, kCheckAddr);

  EXPECT_TRUE(1 == result);
}

TEST_F(IsProcessAddressReadableTest, TaskInfoNull) {
  int32_t result;
  const TaskInfo *task_info_null = NULL;

  result = IsProcessAddressReadable(task_info_null, kCheckAddr);

  EXPECT_TRUE(-1 == result);
}

TEST_F(IsProcessAddressReadableTest, FailedTaskFind) {
  int32_t result;
  g_task_find_return = PA_TZ_GENERAL_ERROR;

  result = IsProcessAddressReadable(&kInitTask, kCheckAddr);

  EXPECT_TRUE(-1 == result);
}

TEST_F(IsProcessAddressReadableTest, AddressIsNotReadable) {
  int32_t result;
  g_vma_task_find_return.vm_flags = 0x0;

  result = IsProcessAddressReadable(&kInitTask, kCheckAddr);

  EXPECT_TRUE(0 == result);
}


TEST_F(IsProcessAddressWritableTest, AllParamAreValid) {
  int32_t result;

  result = IsProcessAddressWritable(&kInitTask, kCheckAddr);

  EXPECT_TRUE(1 == result);
}

TEST_F(IsProcessAddressWritableTest, TaskInfoNull) {
  int32_t result;
  const TaskInfo *task_info_null = NULL;

  result = IsProcessAddressWritable(task_info_null, kCheckAddr);

  EXPECT_TRUE(-1 == result);
}

TEST_F(IsProcessAddressWritableTest, FailedTaskFind) {
  int32_t result;
  g_task_find_return = PA_TZ_GENERAL_ERROR;

  result = IsProcessAddressWritable(&kInitTask, kCheckAddr);

  EXPECT_TRUE(-1 == result);
}

TEST_F(IsProcessAddressWritableTest, AddressIsNotWritable) {
  int32_t result;
  g_vma_task_find_return.vm_flags = 0x0;

  result = IsProcessAddressWritable(&kInitTask, kCheckAddr);

  EXPECT_TRUE(0 == result);
}