#include "kg_rpmb.h"

qsee_stor_device_id_type gDeviceID = QSEE_STOR_EMMC_RPMB;
qsee_stor_device_handle_t gDevHandle = 0;
qsee_stor_client_handle_t gClientHandle = 0;
qsee_stor_client_info_t gClientInfo;

uint32_t kg_rpmb_init() {    

    uint32_t num_sectors = KG_RPMB_SECTOR_MAX;

    int32_t ret = qsee_stor_device_init(gDeviceID, 0, &gDevHandle);
    if (QSEE_STOR_SUCCESS != ret) {
        KG_LOG("KG RPMB is unavailable\n");
        return KG_RPMB_UNAVAILABLE;
    }

    ret = qsee_stor_open_partition(&gDevHandle, KG_RPMB_PARTITION_ID, &gClientHandle);
    if (QSEE_STOR_PARTI_NOT_FOUND_ERROR == ret) {
        KG_LOG("rpmb debug\n");
        qsee_stor_device_info_t device_info;

        KG_LOG("KG RPMB: partition not found, try to add partition");
        ret = qsee_stor_device_get_info(&gDevHandle, &device_info);
        if (QSEE_STOR_SUCCESS != ret) {
            KG_LOG("KG RPMB: Failed to get device info\n");
            return KG_RPMB_UNAVAILABLE;
        }
        KG_LOG("KG RPMB: sectors total (%u) - available (%u)\n", device_info.total_sectors, device_info.available_sectors);
        num_sectors = (KG_RPMB_SECTOR_MAX < device_info.available_sectors) ? KG_RPMB_SECTOR_MAX : device_info.available_sectors;

        ret = qsee_stor_add_partition(&gDevHandle, KG_RPMB_PARTITION_ID, num_sectors);
        if (QSEE_STOR_SUCCESS != ret) {
            KG_LOG("KG RPMB: Failed to make new partition - sector %d", num_sectors);
            return KG_RPMB_UNAVAILABLE;
        }
        ret = qsee_stor_open_partition(&gDevHandle, KG_RPMB_PARTITION_ID, &gClientHandle);
        if ((QSEE_STOR_SUCCESS != ret) || (gClientHandle == 0)) {
            KG_LOG("KG RPMB: open failed\n");
            ret = KG_RPMB_UNAVAILABLE;
        }
    } else if (QSEE_STOR_SUCCESS != ret) {
        KG_LOG("KG RPMB: first open failed ret = %d\n", ret);
        return KG_RPMB_UNAVAILABLE;
    }
    KG_LOG("KG RPMB init success\n");

    return KG_SUCCESS;
}

static uint32_t get_bytes_per_sector(void) {
    int ret;

    ret = qsee_stor_client_get_info(&gClientHandle, &gClientInfo);
    if (ret != QSEE_STOR_SUCCESS) {
            KG_LOG("Error retrieving bytes per sector");
            return 0;
    }
    //KG_LOG_DBG("total_sectors : %d", gClientInfo.total_sectors);
    return gClientInfo.bytes_per_sector;
}

uint32_t read_block(uint32_t start_block_index, uint32_t data_len, uint8_t **data)
{
    KG_LOG("QSEE read_block start block index : %d, data_len : %d\n", start_block_index, data_len);
    uint32_t ret = KG_SUCCESS;
    uint32_t qsee_single_sector_size;
    uint32_t qsee_start_block_index = start_block_index;

    qsee_single_sector_size = get_bytes_per_sector();
    if (qsee_single_sector_size <= 0) {
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }
    uint32_t read_loop_count = data_len / qsee_single_sector_size;
    if (data_len % qsee_single_sector_size) {
        read_loop_count++;
    }
    KG_LOG("QSEE read_loop_count : %d\n", read_loop_count);

    uint8_t *tempBuffer = TEE_Malloc(read_loop_count * qsee_single_sector_size, 0);
    if (NULL == tempBuffer) {
        KG_LOG("Failed to alloc read_block\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }
    *data = tempBuffer;

    ret = qsee_stor_read_sectors(&gClientHandle, qsee_start_block_index, read_loop_count, tempBuffer);
    if (QSEE_STOR_SUCCESS != ret) {
        KG_LOG("KG RPMB read failed\n");
        ret = KG_RPMB_READ_FAIL;
        goto exit;
    }
exit:
    return ret;
}

uint32_t write_block(uint32_t start_block_index, uint32_t data_len, uint8_t *data) {
    KG_LOG("write_block start block index : %d, data_len : %d\n", start_block_index, data_len);
    uint32_t ret = KG_SUCCESS;
    uint32_t qsee_single_sector_size;
    uint32_t qsee_start_block_index = start_block_index;

    qsee_single_sector_size = get_bytes_per_sector();
    if (qsee_single_sector_size <= 0) {
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }
    uint32_t write_loop_count = data_len / qsee_single_sector_size;
    if (data_len % qsee_single_sector_size) {
        write_loop_count++;
    }
    KG_LOG("write_loop_count : %d\n", write_loop_count);

    uint8_t *tempBuffer = TEE_Malloc(write_loop_count * qsee_single_sector_size, 0);
    TEE_MemMove(tempBuffer, data, data_len);

    ret = qsee_stor_write_sectors(&gClientHandle, qsee_start_block_index, write_loop_count, tempBuffer);
    if (QSEE_STOR_SUCCESS != ret) {
        KG_LOG("KG RPMB write failed\n");
        ret = KG_RPMB_WRITE_FAIL;
        goto exit;
    }

exit:
    if (tempBuffer != NULL) {
        TEE_Free(tempBuffer);
    }
    return ret;
}

uint32_t kg_rpmb_read(uint8_t *data, uint32_t data_len) {
  
    uint32_t ret = KG_SUCCESS;
    uint32_t size;
    uint8_t *rpmb_data;

    size = get_bytes_per_sector();
    if (size <= 0) {
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }
    
    rpmb_data = TEE_Malloc(size * KG_RPMB_BLOCK_SIZE, 0);
    if (rpmb_data == NULL)
    {
        KG_LOG("Error to alloc memory");
        return KG_ALLOC_BUFFER_FAIL;
    }

    ret = qsee_stor_read_sectors(&gClientHandle, 0, KG_RPMB_BLOCK_SIZE, rpmb_data);
    if (ret != QSEE_STOR_SUCCESS) {
        KG_LOG("KG RPMB read failed\n");
        return KG_RPMB_READ_FAIL;
    }

    TEE_MemMove(data, rpmb_data, sizeof(kg_rpmb_data_t));
    ret = KG_SUCCESS;
exit:
    if(size > 0) {
        TEE_Free(rpmb_data);
    }

    return ret;
}

uint32_t kg_rpmb_write(uint8_t *data, uint32_t data_len) {

    uint32_t ret = KG_SUCCESS;
    uint32_t size;
    uint8_t *rpmb_data;

    size = get_bytes_per_sector();
    if (size <= 0) {
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }
    
    rpmb_data = TEE_Malloc(size * KG_RPMB_BLOCK_SIZE, 0);
    if (rpmb_data == NULL)
    {
        KG_LOG("Error to alloc memory");
        return KG_ALLOC_BUFFER_FAIL;
    }
    TEE_MemMove(rpmb_data, data, data_len);

    ret = qsee_stor_write_sectors(&gClientHandle, 0, KG_RPMB_BLOCK_SIZE, rpmb_data);
    if (ret != QSEE_STOR_SUCCESS) {
        KG_LOG("KG RPMB write failed\n");
        return KG_RPMB_WRITE_FAIL;
    }

    ret = KG_SUCCESS;
exit:
    if (size > 0) {
        TEE_Free(rpmb_data);
    }

    return ret;
}

uint32_t kg_rpmb_read_metadata(uint8_t *data, uint32_t data_len) {
    KG_LOG_DBG("kg_rpmb_read_metadata()");
    uint32_t ret = KG_SUCCESS;
    kg_rpmb_data_t *krd = NULL;
    kg_secure_data_t *ksd = NULL;
    uint8_t *unwrap_data = NULL;
    uint32_t unwrap_data_len = KG_BUF_LEN;

    if (data_len != sizeof(kg_metadata_t)) {
        KG_LOG("Received wrong size of input buffer to read kg metadata from rpmb\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    krd = TEE_Malloc(sizeof(kg_rpmb_data_t), 0);
    if (NULL == krd) {
        KG_LOG("KG TA failed to alloc buffer to read from RPMB\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    ksd = TEE_Malloc(sizeof(kg_secure_data_t), 0);
    if (NULL == ksd) {
        KG_LOG("KG TA failed to alloc buffer to read from RPMB\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    unwrap_data = TEE_Malloc(unwrap_data_len, 0);
    if (NULL == unwrap_data) {
        KG_LOG("KG TA failed to alloc buffer to hold unwrap data\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("KG TA accessing RPMB failed\n");
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }
    if (KG_SUCCESS != kg_rpmb_read((uint8_t *)krd, sizeof(kg_rpmb_data_t))) {
        KG_LOG("KG TA read rpmb data from RPMB failed\n");
        ret = KG_RPMB_READ_FAIL;
        goto exit;
    }
    KG_DUMP("krd ", (uint8_t *)krd, 4096);
    if (krd->magic != KG_MAGIC) {
        KG_LOG("KG TA rpmb data magic check failed\n");
        ret = KG_RPMB_MAGIC_FAIL;
        goto exit;
    }

    // TBD, rpmb data structure version check
    if (krd->kg_wrap_data_len > KG_SECURE_DATA_LEN) {
        KG_LOG("KG TA rpmb wrap data size overflow\n");
        ret = KG_RPMB_UNWRAP_FAIL;
        goto exit;
    }
    KG_LOG("KG rpmb before unwrap data len %d\n", krd->kg_wrap_data_len);
    if (krd->kg_wrap_data_len != 0) {
        KG_DUMP("KG rpmb krd before unwrap\n", (uint8_t *)krd->kg_wrap_data, krd->kg_wrap_data_len);
    }

    if (TZ_API_OK != TZ_unwrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        krd->kg_wrap_data, krd->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
        KG_LOG("Failed to unwrap kg metadata structure on QC\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }

    if (unwrap_data_len != sizeof(kg_secure_data_t)) {
        KG_LOG("KG TA recovering unwraped metadata size check failed\n");
        ret = KG_RPMB_UNWRAP_FAIL;
        goto exit;
    }

    TEE_MemMove(ksd, unwrap_data, unwrap_data_len);

    TEE_MemMove(data, (void *)&(ksd->kg_metadata), sizeof(kg_metadata_t));

exit:
    if (krd != NULL) {
        TEE_MemFill(krd, 0, sizeof(*krd));
        TEE_Free(krd);
        krd = NULL;
    }
    if (ksd != NULL) {
        TEE_MemFill(ksd, 0, sizeof(*ksd));
        TEE_Free(ksd);
        ksd = NULL;
    }
    if (unwrap_data != NULL) {
        TEE_MemFill(unwrap_data, 0, unwrap_data_len);
        TEE_Free(unwrap_data);
        unwrap_data = NULL;
    }
    return ret;
}

uint32_t kg_rpmb_write_metadata(uint8_t *data, uint32_t data_len) {
    KG_LOG_DBG("kg_rpmb_write_metadata()");
    uint32_t ret = KG_SUCCESS;
    kg_rpmb_data_t *krd = NULL;
    kg_secure_data_t *ksd = NULL;
    uint8_t *data_buf = NULL;
    uint32_t data_buf_len = KG_SECURE_DATA_LEN;

    if (data_len != sizeof(kg_metadata_t)) {
        KG_LOG("Received wrong size of input buffer to write kg metadata to rpmb\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    krd = TEE_Malloc(sizeof(kg_rpmb_data_t), 0);
    if (NULL == krd) {
        KG_LOG("KG TA failed to alloc buffer to read from RPMB\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    ksd = TEE_Malloc(sizeof(kg_secure_data_t), 0);
    if (NULL == ksd) {
        KG_LOG("KG TA failed to alloc buffer to read from RPMB\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    data_buf = TEE_Malloc(data_buf_len, 0);
    if (NULL == data_buf) {
        KG_LOG("KG TA failed to alloc buffer to hold wrap data\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("KG TA accessing RPMB failed\n");
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }

    if (KG_SUCCESS != kg_rpmb_read((uint8_t *)krd, sizeof(kg_rpmb_data_t))) {
        KG_LOG("KG TA read metadata from RPMB failed\n");
        ret = KG_RPMB_READ_FAIL;
        goto exit;
    }

    if (krd->magic != KG_MAGIC) {
        KG_LOG("KG TA rpmb data magic check failed\n");
        ret = KG_RPMB_MAGIC_FAIL;
        goto exit;
    }

    if (TZ_API_OK != TZ_unwrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        krd->kg_wrap_data, krd->kg_wrap_data_len, data_buf, &data_buf_len)) {
        KG_LOG("Failed to wrap kg secure data structure in QC\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }

    if (data_buf_len != sizeof(kg_secure_data_t)) {
        KG_LOG("Unwraped kg secure data size check failed\n");
        ret = KG_RPMB_UNWRAP_FAIL;
        goto exit;
    }

    TEE_MemMove(ksd, data_buf, data_buf_len);
    TEE_MemMove(&(ksd->kg_metadata), data, data_len);

    TEE_MemFill(data_buf, 0, KG_SECURE_DATA_LEN);
    data_buf_len = KG_SECURE_DATA_LEN;

    if (TZ_API_OK != TZ_wrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        (uint8_t *)ksd, sizeof(*ksd), data_buf, &data_buf_len)) {
        KG_LOG("Failed to wrap kg secure data structure in QC\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }

    if (data_buf_len > KG_SECURE_DATA_LEN) {
        KG_LOG("Wraped kg secure data size overflow\n");
        ret = KG_RPMB_WRAP_FAIL;
        goto exit;
    }

    TEE_MemFill(krd->kg_wrap_data, 0, KG_SECURE_DATA_LEN);
    krd->kg_wrap_data_len = data_buf_len;
    TEE_MemMove(krd->kg_wrap_data, data_buf, data_buf_len);

    if (KG_SUCCESS != kg_rpmb_write((uint8_t *)krd, sizeof(kg_rpmb_data_t))) {
        KG_LOG("KG TA write metadata to RPMB failed\n");
        ret = KG_RPMB_WRITE_FAIL;
        goto exit;
    }

exit:
    if (krd != NULL) {
        TEE_MemFill(krd, 0, sizeof(*krd));
        TEE_Free(krd);
        krd = NULL;
    }
    if (ksd != NULL) {
        TEE_MemFill(ksd, 0, sizeof(*ksd));
        TEE_Free(ksd);
        krd = NULL;
    }
    if (data_buf != NULL) {
        TEE_Free(data_buf);
        data_buf = NULL;
    }
    return ret;
}

uint32_t kg_rpmb_unwrap_securedata(kg_rpmb_data_t *krd, uint8_t *data) {
    KG_LOG_DBG("kg_rpmb_unwrap_metadata()");
    uint32_t ret = KG_SUCCESS;
    uint8_t *unwrap_data = NULL;
    uint32_t unwrap_data_len = KG_BUF_LEN;

    if (krd == NULL) {
        KG_LOG("krd is NULL\n");
        ret = KG_ERR_INVALID_BUFFER;
        goto exit;
    }

    if (data == NULL) {
        KG_LOG("data is NULL\n");
        ret = KG_ERR_INVALID_BUFFER;
        goto exit;
    }

    unwrap_data = TEE_Malloc(unwrap_data_len, 0);
    if (NULL == unwrap_data) {
        KG_LOG("KG TA failed to alloc buffer to hold unwrap data\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    if (krd->magic != KG_MAGIC) {
        KG_LOG("KG TA rpmb data magic check failed\n");
        ret = KG_RPMB_MAGIC_FAIL;
        goto exit;
    }

    // TBD, rpmb data structure version check
    if (krd->kg_wrap_data_len > KG_SECURE_DATA_LEN) {
        KG_LOG("KG TA rpmb wrap data size overflow\n");
        ret = KG_RPMB_UNWRAP_FAIL;
        goto exit;
    }

    if (TZ_API_OK != TZ_unwrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), krd->kg_wrap_data, krd->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
        KG_LOG("Failed to unwrap kg secure data structure on QC\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }

    if (unwrap_data_len != sizeof(kg_secure_data_t)) {
        KG_LOG("KG TA recovering unwraped secure data size check failed\n");
        ret = KG_RPMB_UNWRAP_FAIL;
        goto exit;
    }

    TEE_MemMove(data, unwrap_data, unwrap_data_len);

exit:
    if (unwrap_data != NULL) {
        TEE_MemFill(unwrap_data, 0, sizeof(kg_secure_data_t));
        TEE_Free(unwrap_data);
        unwrap_data = NULL;
    }
    return ret;
}


uint32_t kg_rpmb_wrap_securedata(kg_rpmb_data_t *krd, kg_secure_data_t *ksd) {
    KG_LOG_DBG("kg_rpmb_wrap_metadata()");
    uint32_t ret = KG_SUCCESS;
    uint8_t *data_buf = NULL;
    uint32_t data_buf_len = KG_SECURE_DATA_LEN;

    if (NULL == krd) {
        KG_LOG("krd is NULL\n");
        ret = KG_ERR_INVALID_BUFFER;
        goto exit;
    }

    if (NULL == ksd) {
        KG_LOG("ksd is null\n");
        ret = KG_ERR_INVALID_BUFFER;
        goto exit;
    }

    data_buf = TEE_Malloc(data_buf_len, 0);
    if (NULL == data_buf) {
        KG_LOG("KG TA failed to alloc buffer to hold wrap data\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    if (krd->magic != KG_MAGIC) {
        KG_LOG("KG TA rpmb data magic check failed\n");
        ret = KG_RPMB_MAGIC_FAIL;
        goto exit;
    }

    TEE_MemFill(data_buf, 0, KG_SECURE_DATA_LEN);
    data_buf_len = KG_SECURE_DATA_LEN;

    if (TZ_API_OK != TZ_wrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME),
         (uint8_t *)ksd, sizeof(*ksd), data_buf, &data_buf_len)) {
        KG_LOG("Failed to wrap kg secure data structure in QC\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }

    if (data_buf_len > KG_SECURE_DATA_LEN) {
        KG_LOG("Wraped kg secure data structure size overflow\n");
        ret = KG_RPMB_WRAP_FAIL;
        goto exit;
    }
    

    TEE_MemFill(krd->kg_wrap_data, 0, KG_SECURE_DATA_LEN);
    krd->kg_wrap_data_len = data_buf_len;
    TEE_MemMove(krd->kg_wrap_data, data_buf, data_buf_len);
	
    KG_LOG_DBG("finish calling kg_rpmb_wrap_metadata\n");

exit:
    if (data_buf != NULL) {
        TEE_Free(data_buf);
        data_buf = NULL;
    }
    return ret;
}