#include "config.h"

#include "crypto.h"
#include "driver_log.h"
#include "kaslr.h"
#include "pa_tz_api.h"
#include "validation_key.h"
#include "kernel_access.h"

#include <tee_internal_api.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;

static PaConfig g_pa_config = {0};

static uint32_t g_config_init_status = 0;
static uint32_t g_is_five_enabled = 1;

/**
 * @brief Prints all fields of configuration structure
 * @param config Configuration structure to be dumped
 */
static void ConfigDump(const Config64_t *config);

/**
 * @brief Load PA configuration from kernel data structure
 * Data structure with config is located in kernel .data section.
 * We can't use virtual addresses to get it, so we need to put it's physical
 * address at some known location in memory. We decided to use one of
 * reserved fields in kernel header for this purpose. Kernel header is located
 * at the beginning of kernel image, which always has fixed
 * address + KASLR offset. In this field we put offset of proca config
 * symbol (g_proca_config) from kernel text section start (_text).
 * So as a summary algorithm to load config from kernel is the following:
 *  1. Get platform specific address of reserved field in kernel.
 *  2. Add KASLR offset to it.
 *  3. Load difference between kernel start and g_proca_config symbol.
 *  4. Add this difference to kernel start address + KASLR offset.
 *  5. Load configuration buffer using obtained address.
 * @param [in] kaslr_offset
 * @param [out] config Output configuration structure
 * @return ::PA_TZ_SUCCESS, ::PA_TZ_GENERAL_ERROR
 */
static PaTzResult ConfigLoadFromKernel(uint32_t kaslr_offset,
                                       Config64_t *config);

/**
 * @brief Read PROCA hash tables configuration
 * @param [in] pa_table_addr Address of proca table structure
 * @param [out] config Application name map configuration
 * @return ::PA_TZ_SUCCESS, ::PA_TZ_GENERAL_ERROR
 */
static PaTzResult GetProcaTablesConf(const KernelAddress pa_table_addr,
                                   PaConfig *config);

/**
 * @brief Initialize memory configuration
 * @param [in] conf Configuration buffer from kernel
 * @param [out] mem_conf Memory configuration
 * @return ::PA_TZ_SUCCESS, ::PA_TZ_GENERAL_ERROR
 */
static void InitMemoryConfig(const Config64_t *conf,
                             MemoryConfig *mem_conf);

/**
 * @brief Validate configuration buffer size, version and magic
 * @param [in] conf Configuration buffer from kernel
 * @return ::PA_TZ_SUCCESS, ::PA_TZ_GENERAL_ERROR
 */
static PaTzResult ConfigValidate(const Config64_t *config);

static void InitMemoryConfig(const Config64_t *conf,
                             MemoryConfig *mem_conf) {
  if (!conf || !mem_conf) {
    LOG_E("Args are NULL\n");
    return;
  }

  mem_conf->va_start       = conf->va_start;
  mem_conf->kimage_vaddr   = conf->kimage_vaddr;
  mem_conf->kimage_voffset = conf->kimage_voffset;
  mem_conf->page_offset    = conf->page_offset;
  mem_conf->va_bits        = conf->va_bits;

  TEE_MemMove(&mem_conf->sys_ram_ranges, &conf->sys_ram_ranges,
              sizeof(mem_conf->sys_ram_ranges[0]) * conf->sys_ram_ranges_num);
  mem_conf->sys_ram_ranges_num = conf->sys_ram_ranges_num;
}

static void ConfigDump(const Config64_t *conf) {
  if (!conf) {
    LOG_E("Config is NULL.\n");
    return;
  }

  LOG_V("version:  0x%08x\n", conf->version);
  LOG_V("size:     0x%08x\n", conf->size);
  LOG_V("magic:    0x%08x\n", conf->magic);

  LOGADDR_V("gaf_addr:         ", conf->gaf_addr);
  LOGADDR_V("proca_table_addr: ", conf->proca_table_addr);

  LOGADDR_V("page_offset:    ", conf->page_offset);
  LOGADDR_V("kimage_vaddr:   ", conf->kimage_vaddr);
  LOGADDR_V("kimage_voffset: ", conf->kimage_voffset);
  LOGADDR_V("va_start:       ", conf->va_start);
  LOG_V("va_bits:        %08d\n", (uint32_t)conf->va_bits);

  LOG_V("Discovered %08d system RAM ranges: \n",
      (uint32_t)conf->sys_ram_ranges_num);

  for (size_t i = 0; i < conf->sys_ram_ranges_num; ++i) {
    LOG_V("0x%08x%08x-0x%08x%08x.\n",
          (uint32_t)(conf->sys_ram_ranges[i].start >> 32),
          (uint32_t)conf->sys_ram_ranges[i].start,
          (uint32_t)(conf->sys_ram_ranges[i].end >> 32),
          (uint32_t)conf->sys_ram_ranges[i].end);
  }
}

static PaTzResult ConfigValidate(const Config64_t *config) {
  if (!config) {
    LOG_D("Config is NULL.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  if (sizeof(*config) != config->size) {
    LOG_E("Wrong config size obtained from kernel.\n");
    LOG_D("Current config size %d > input config size %d\n",
        sizeof(*config), config->size);
    return PA_TZ_GENERAL_ERROR;
  }

  if (kPaConfigVersion != config->version) {
    LOG_E("Didn't match current config version and supported config version\n");
    LOG_D("Current config version = 0x%08x, supported config version = 0x%08x\n",
          config->version, kPaConfigVersion);
    return PA_TZ_GENERAL_ERROR;
  }

  if (config->magic != kPaConfigMagic) {
    LOG_E("Didn't match current config magic and expected value\n");
    LOG_D("Current config magic = 0x%08x, expected = 0x%08x\n",
          config->magic, kPaConfigMagic);
    return PA_TZ_GENERAL_ERROR;
  }

  if (config->sys_ram_ranges_num > kMaxMemoryRangesNum) {
    LOG_E("Number of system RAM ranges is bigger than maximal possible\n");
    LOG_D("Current config sys_ram_ranges_num = 0x%08x, expected < 0x%08x\n",
          config->sys_ram_ranges_num, kMaxMemoryRangesNum);
    return PA_TZ_GENERAL_ERROR;
  }

  return PA_TZ_SUCCESS;
}

static PaTzResult ConfigLoadFromKernel(uint32_t kaslr_offset,
                                       Config64_t *config) {
  uint64_t config_sym_text_offset = 0;

  PaTzResult result = PhysicalGetBytes(KernelGetPaConfigOffsetSymAddress() +
        kaslr_offset, sizeof(config_sym_text_offset), &config_sym_text_offset);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Can not read config offset.\n");
    LOGADDR_D("Config offset symbol phys: ", config_sym_text_offset);
    return result;
  }

  LOGADDR_D("Proca config struct _text offset: ", config_sym_text_offset);

  if (config_sym_text_offset == kFiveDisabledMagicNumber) {
    LOG_D("FIVE disabled magic number found!\n");
    return PA_TZ_PROCA_NOT_SUPPORTED;
  }

  PhysicalAddress config_phys = KernelGetTextSymAddress() +
                                          config_sym_text_offset;
  LOGADDR_D("Proca config struct phys: ", config_phys);

  result = PhysicalGetBytes(config_phys + kaslr_offset,
                            sizeof(*config), config);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Can not load config struct.\n");
    return result;
  }

  result = ConfigValidate(config);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Config is invalid.\n");
    return result;
  }

  ConfigDump(config);

  return PA_TZ_SUCCESS;
}

PaTzResult ConfigInit() {
  Config64_t config = {0};

  g_config_init_status = 0;

  PaTzResult result;

  while (true) {
    result = KernelNextMemoryConfiguration();
    if (result != PA_TZ_SUCCESS) {
      return PA_TZ_GENERAL_ERROR;
    }

    result = KernelGetKaslrOffset(&g_pa_config.kaslr_offset);
    if (result != PA_TZ_SUCCESS) {
      LOG_E("Obtaining KASLR offset is failed.\n");
      continue;
    }

    result = ConfigLoadFromKernel(g_pa_config.kaslr_offset, &config);
    if (result == PA_TZ_PROCA_NOT_SUPPORTED) {
      LOG_I("This device does not support FIVE\n");
      g_is_five_enabled = 0;
      break;
    } else if (result != PA_TZ_SUCCESS) {
      LOG_E("Can not load configuration structure\n");
      continue;
    }

    InitMemoryConfig(&config, &g_pa_config.memory_conf);

    result = LoadGafInfo(config.gaf_addr,
                         &g_pa_config.gaf);
    if (result == PA_TZ_SUCCESS) {
      LOG_I("GAF is loaded.\n");
      break;
    } else {
      LOG_E("Can not load GAF info.\n");
    }
  }

  if (g_is_five_enabled == 1) {
    result = GetProcaTablesConf(config.proca_table_addr,
                                &g_pa_config);
    if (result != PA_TZ_SUCCESS) {
      LOG_E("Can not load application name map table configuration\n");
      return PA_TZ_GENERAL_ERROR;
    }
  }

  g_config_init_status = 1;

  return PA_TZ_SUCCESS;
}

const PaConfig *GetConfig(void) {
  return &g_pa_config;
}

static PaTzResult GetProcaTablesConf(const KernelAddress pa_table_addr,
                                    PaConfig *config) {
  unsigned int shift;
  LOGADDR_V("g_proca_table ", pa_table_addr);
  PaTzResult result = KernelAccessGetUint32(NULL,
            pa_table_addr + GetConfig()->gaf.proca_table_hash_tables_shift,
            &shift);
  if (result != PA_TZ_SUCCESS) {
    LOG_E("Cannot read PROCA hash tables shift.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  if (shift == 0 || shift > 32) {
    LOG_V("PROCA hash table shift = %u\n", shift);
    LOG_E("PROCA hash table shift is invalid.\n");
    return PA_TZ_GENERAL_ERROR;
  }

  config->app_name_map.shift = config->pid_map.shift = shift;
  config->app_name_map.size = config->pid_map.size = (1ULL << shift);

  LOGADDR_V("config->app_name_map.size = ", config->app_name_map.size);
  LOG_V("config->app_name_map.shift = 0x%08x\n", config->app_name_map.shift);

  config->pid_map.start
            = pa_table_addr + config->gaf.proca_table_pid_map;
  config->app_name_map.start
            = pa_table_addr + config->gaf.proca_table_app_name_map;

  return PA_TZ_SUCCESS;
}

uint32_t ConfigIsInited(void) {
  return g_config_init_status;
}

uint32_t IsFiveEnabled(void) {
  return g_is_five_enabled;
}