/*
 * Copyright (c) 2016 Samsung Electronics Co., Ltd. All rights reserved.
 *
 * Created in Samsung Ukraine R&D Center (SRK) under a contract between
 * LLC "Samsung Electronics Ukraine Company" (Kiev, Ukraine)
 * and "Samsung Electronics Co", Ltd (Seoul, Republic of Korea)
 *
 * Created on: Jul 1, 2016
 * Author: Oleksii Kachkan <o.kachkan@samsung.com>
 * Brief: RPMB access point.
 */

#include "TzwRpmb.h"

#include "TigerLogging.h"
#include "TigerMacros.h"
#include "TigerUtils.h"

// QSEE includes
#if defined(TZ_MODEL_QCOM)
    #include <qsee_stor.h>
    #include <gpTypes.h>
#elif defined(TIGER_TZ_MODEL_BLOWFISH)
    #include <rpmb.h>
    #include <driver/rpmb/rpmb.h>
#elif defined(TZ_MODEL_Kinibi) /*TZ_MODEL_Kinibi*/
    #include "TzwString.h"
    #include "TzwMemory.h"
#endif

// From 80-P4904-1_WECHAT SOTER INTEGRATION GUIDE.pdf
// To make good use of RPMB, the maximum number of ATTKs is 5,
// so the number of copies could be customized between 0 and 4, which defaults to 3.
#define ATTK_COPIES_NUM         (3)

// QSEE implementation
#if defined(TZ_MODEL_QCOM)

// Because QSEE share partition IDs beetween all TA we have to generate unique Id for Tiger
// 0x7E544650 it is HEX presentation of '~TFP'
#define TIGER_ATTK_PARTITION_ID (0x7E544650)

typedef struct {
    qsee_stor_client_handle_t handle;
    uint32_t                  bytesPerSector;
    uint32_t                  totalSectors;
} TigerPartition;

// Function create new partition for tigerfp and open it.
static TEE_Result initTigerPartition(const uint32_t dataSize, IN OUT TigerPartition* partition);

// Function try to open tigerfp partition. If not found returns TEEC_ERROR_ITEM_NOT_FOUND
static TEE_Result openTigerPartition(OUT TigerPartition* partition);

// Normalize (attkhash + len + attk) size by RPMB sector size
static uint32_t normalizeSize(const uint32_t dataSize, const uint32_t bytesPerSector) {
    const uint32_t neededSectors = (dataSize + (bytesPerSector - 1 )) / bytesPerSector;
    return neededSectors * bytesPerSector;
}

TEE_Result rpmbWrite(IN const uint8_t* data, const uint32_t dataSize) {
    LOG_FUNC_BEGIN;
    TIGER_CHECK_FUNCTION_ARGUMENT_RETURN(data != NULL);

    TigerPartition partition;
    TEE_Result status = openTigerPartition(&partition);
    // if Attk save for the first time we need make provisioning
    if (TEEC_ERROR_ITEM_NOT_FOUND == status) {
        LOG_D("Provision is running...");
        status = initTigerPartition(dataSize, &partition);
    }
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("Open tiger partition");

    uint32_t normalizedFullSize = normalizeSize(dataSize, partition.bytesPerSector);
    uint32_t sectorsCount = normalizedFullSize / partition.bytesPerSector;

    if ((sectorsCount * ATTK_COPIES_NUM) > partition.totalSectors) {
        LOG_E("Partition does not have enough sectors.");
        return TEE_ERROR_BAD_STATE;
    }

    // prepare Attk ant it's hash
    uint8_t* normalizedBuf = tzwMalloc(normalizedFullSize);
    TIGER_CHECK_BUFFER_ALLOCATED_RETURN(normalizedBuf);

    // copy data to normalized buffer
    tzwMemMove(normalizedBuf, data, dataSize);

    uint32_t attkCopyNum = 0;
    int ret = QSEE_STOR_SUCCESS;
    for (; attkCopyNum < ATTK_COPIES_NUM; ++attkCopyNum) {
        uint32_t startSector = attkCopyNum * sectorsCount;
        LOG_D("Save Attk to partition (%u), sector %u", partition, startSector);

        ret = qsee_stor_write_sectors(&(partition.handle), startSector, sectorsCount, normalizedBuf);
        if (QSEE_STOR_SUCCESS != ret) {
            LOG_E("Failed to qsee_stor_write_sectors %d", ret);
            continue;
        }
        LOG_D("Attk copy #%u was successfully saved!", attkCopyNum);
    }

    if ((QSEE_STOR_SUCCESS == ret) && (attkCopyNum == ATTK_COPIES_NUM)) {
        LOG_D("Attestation key was successfully saved.");
        status = TEE_SUCCESS;
    } else {
        LOG_E("Failed to save attk key-pair, only %u copies saved", attkCopyNum);
        status = TEE_ERROR_GENERIC;
    }

    tzwFree(normalizedBuf);

    LOG_FUNC_END;
    return status;
}

TEE_Result rpmbRead(IN OUT uint8_t* buffer, const uint32_t bufferSize) {
    LOG_FUNC_BEGIN;

    TigerPartition partition;

    TEE_Result status = openTigerPartition(&partition);
    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("openTigerPartition()");

    // in one partition we have ATTK_COPIES_NUM copies of ATTk;
    uint32_t attkSectorscount = partition.totalSectors / ATTK_COPIES_NUM;

    if (bufferSize < (partition.bytesPerSector * attkSectorscount)) {
        LOG_E("buffer is too small.");
        return TEE_ERROR_BAD_PARAMETERS;
    }

    for (uint32_t readAttempts = 0; readAttempts < ATTK_COPIES_NUM; ++readAttempts) {

        LOG_D("Current attk Index %u ", readAttempts);

        const uint32_t startSector = readAttempts * attkSectorscount;
        int ret = qsee_stor_read_sectors(&(partition.handle), startSector, attkSectorscount, buffer);
        if (QSEE_STOR_SUCCESS != ret) {
            LOG_E("Failed to read attk %d", ret);
            continue;
        }
        break;
    }

    LOG_FUNC_END;
    return status;
}


TEE_Result rpmbClear() {
    LOG_FUNC_BEGIN;

    TigerPartition partition;

    TEEC_Result status = openTigerPartition(&partition);

    TIGER_CHECK_TEE_STATUS_SUCCESS_RETURN("openTigerPartition()");
    uint8_t* buffer = tzwMalloc(partition.totalSectors * partition.bytesPerSector);
    TIGER_CHECK_BUFFER_ALLOCATED_RETURN(buffer);

    int ret = qsee_stor_write_sectors(&(partition.handle), 0, partition.totalSectors, buffer);
    if (QSEE_STOR_SUCCESS != ret) {
        LOG_E("Failed to qsee_stor_write_sectors %d", ret);
        status = TEE_ERROR_GENERIC;
    }
    tzwFree(buffer);

    LOG_FUNC_END;
    return status;
}

static TEE_Result openTigerPartition(OUT TigerPartition* partition) {
    qsee_stor_device_handle_t handle = 0;
    qsee_stor_device_info_t deviceInfo;

    int ret = qsee_stor_device_init(QSEE_STOR_EMMC_RPMB, NULL, &handle);
    if (QSEE_STOR_SUCCESS != ret) {
        LOG_E("Failed to init device, error %d", ret);
        return TEE_ERROR_ACCESS_DENIED;
    }
    ret  = qsee_stor_device_get_info(&handle, &deviceInfo);
    if (QSEE_STOR_SUCCESS != ret) {
        LOG_E("Failed to get device info, error code: %d", ret);
        return TEE_ERROR_GENERIC;
    }

    ret = qsee_stor_open_partition(&handle, TIGER_ATTK_PARTITION_ID, &(partition->handle));
    if (QSEE_STOR_SUCCESS != ret) {
        LOG_E("Failed to open partition with ID(%x), code %d", TIGER_ATTK_PARTITION_ID, ret);
        return TEE_ERROR_ITEM_NOT_FOUND;
    }

    qsee_stor_client_info_t partitionInfo;
    ret = qsee_stor_client_get_info(&(partition->handle), &partitionInfo);
    if (QSEE_STOR_SUCCESS != ret) {
        LOG_E("Failed to qsee_stor_client_get_info %d", ret);
        return TEE_ERROR_GENERIC;
    }

    // update partition info
    partition->totalSectors = partitionInfo.total_sectors;
    partition->bytesPerSector = partitionInfo.bytes_per_sector;

    return TEE_SUCCESS;
}

static TEE_Result initTigerPartition(const uint32_t dataSize, IN OUT TigerPartition* partition) {

    qsee_stor_device_handle_t handle = 0;
    qsee_stor_device_info_t deviceInfo;

    int ret = qsee_stor_device_init(QSEE_STOR_EMMC_RPMB, NULL, &handle);
    if (QSEE_STOR_SUCCESS != ret) {
        LOG_E("Failed to init device, error %d", ret);
        return TEE_ERROR_ACCESS_DENIED;
    }
    ret  = qsee_stor_device_get_info(&handle, &deviceInfo);
    if (QSEE_STOR_SUCCESS != ret) {
        LOG_E("Failed to get device info, error code: %d", ret);
        return TEE_ERROR_GENERIC;
    }

    uint32_t fullDataSize = normalizeSize(dataSize, deviceInfo.bytes_per_sector);
    uint32_t sectorscount = (fullDataSize / deviceInfo.bytes_per_sector) * ATTK_COPIES_NUM;

    ret = qsee_stor_add_partition(&handle, TIGER_ATTK_PARTITION_ID, sectorscount);
    if (QSEE_STOR_SUCCESS != ret) {
        LOG_E("Failed to create partition with ID(%x), code %d", TIGER_ATTK_PARTITION_ID, ret);
        return TEE_ERROR_ACCESS_CONFLICT;
    }
    LOG_D("Tiger partition was successfully created, open it.");

    ret = qsee_stor_open_partition(&handle, TIGER_ATTK_PARTITION_ID, &(partition->handle));
    if (QSEE_STOR_SUCCESS != ret) {
        LOG_E("Failed to open partition with ID(%x), code %d", TIGER_ATTK_PARTITION_ID, ret);
        return TEE_ERROR_ACCESS_DENIED;
    }

    qsee_stor_client_info_t partitionInfo;
    ret = qsee_stor_client_get_info(&(partition->handle), &partitionInfo);
    if (QSEE_STOR_SUCCESS != ret) {
        LOG_E("Failed to qsee_stor_client_get_info %d", ret);
        return TEE_ERROR_GENERIC;
    }

    // update secotrs num
    partition->totalSectors = partitionInfo.total_sectors;
    partition->bytesPerSector = partitionInfo.bytes_per_sector;

    return TEE_SUCCESS;
}

#elif defined(TZ_MODEL_BLOWFISH) /*TZ_MODEL_BLOWFISH*/

#define TYPE_FLAG_VAL_OF_RPMB  RPMB_TYPE_BYTE
#define SOTER_RPMB_PARTITION_ID      17

TEE_Result rpmbWrite(IN const uint8_t* data, const uint32_t dataSize) {

    TEE_Result status = TEE_SUCCESS;

    if(TEE_SUCCESS != (status = TEES_RPMBCheckEnable())){
        LOG_E("RPMB is not supported yet");
        return status;
    }

    status = TEES_RPMBWrite(SOTER_RPMB_PARTITION_ID, 0, (uint8_t*)data, (uint32_t)dataSize, TYPE_FLAG_VAL_OF_RPMB);
    if(TEE_SUCCESS != status){
        LOG_E("Failed to TEES_RPMBWrite");
    }

    return status;
}

TEE_Result rpmbRead(IN OUT uint8_t* buffer, const uint32_t bufferSize) {

    TEE_Result status = TEE_SUCCESS;

    if(TEE_SUCCESS != (status = TEES_RPMBCheckEnable())){
        LOG_E("RPMB is not supported yet");
        return status;
    }

    status = TEES_RPMBRead(SOTER_RPMB_PARTITION_ID, 0, buffer, bufferSize, TYPE_FLAG_VAL_OF_RPMB);
    if(TEE_SUCCESS != status){
        LOG_E("Failed to TEES_RPMBRead");
    }

    return status;
}

TEE_Result rpmbClear() {

    return TEE_SUCCESS;
}

#elif defined(TZ_MODEL_Kinibi) /*TZ_MODEL_Kinibi*/

#if __USE_MT6768__ == 1

#include "mediatek/tlRpmbDriverApi.h"

#define SIZE_OF_PARTITION  (1024 * 512)
#define SIZE_OF_BLOCK  256
#define POWER_OF_BLOCK 8
#define BLK_NUM(__size) \
        ((__size + SIZE_OF_BLOCK -1 ) >> POWER_OF_BLOCK)


TEE_Result rpmbWrite(IN const uint8_t* data, const uint32_t dataSize)
{
    LOG_FUNC_BEGIN

    tlApiResult_t ret = TLAPI_OK;

    uint32_t session_id = tlApiRpmbOpenSession(SOTER_ATTK);

    if(session_id < 1) {
        LOG_E("open session failed");
        return TEE_ERROR_GENERIC;
    }

    int result = 0;
    ret = tlApiRpmbWriteData(session_id, (uint8_t*)data, dataSize, &result);
    LOG_DEV("tlApiRpmbWriteData, result : 0x%08x", result);
    if( TLAPI_OK != ret ){
        LOG_E("tlApiRpmbWriteData, result : 0x%08x", result);
        LOG_E("tlApiRpmbWriteData failed");
    }

    ret = tlApiRpmbCloseSession(session_id);
    if ( ret != TLAPI_OK){
        LOG_E("close session failed, but we can do nothing any more");
    }

    LOG_FUNC_END

    return TEE_SUCCESS;
}

TEE_Result rpmbRead(IN OUT uint8_t* buffer, const uint32_t bufferSize)
{
    LOG_FUNC_BEGIN

    tlApiResult_t ret = TLAPI_OK;

    uint32_t session_id = tlApiRpmbOpenSession(SOTER_ATTK);

    if(session_id < 1) {
        LOG_E("open session failed");
        return TEE_ERROR_GENERIC;
    }

    int result = 0;
    ret = tlApiRpmbReadData(session_id, buffer, bufferSize, &result);
    LOG_DEV("tlApiRpmbReadData, result : 0x%08x", result);
    if( TLAPI_OK != ret ){
        LOG_E("tlApiRpmbReadData, result : 0x%08x", result);
        LOG_E("tlApiRpmbReadData failed");
    }

    ret = tlApiRpmbCloseSession(session_id);
    if ( ret != TLAPI_OK){
        LOG_E("close session failed, but we can do nothing any more");
    }


    LOG_FUNC_END

    return TEE_SUCCESS;
}

TEE_Result rpmbClear()
{

    LOG_FUNC_BEGIN

    LOG_D("dummy");

    LOG_FUNC_END

    return TEE_SUCCESS;
}

#else /* __CHIPSET__ != MT6768*/

///////////////////////////////////////////////////////////////////////////////
#if defined(CONFIG_SRPMB)
/**
 * Write RPMB data
 *
 * @partition : The partition name in RPMB
 *  The RPMB partition is divided into 8 for multi users
 *  Each partition size is 512KB
 *  It can be 1,2,3,4,5,6,7,8
 * @address : The start address of RPMB for write
 *  The start address is '0' each partition
 *  The data size is 256byte per address '1'
 * @*data : The virtual address of input data
 * @blk_num : The block number of input data length in bytes
 *  The data size is 256byte per block '1'
 *  This value must be less than '256' (32KB at once write)
 */
_TLAPI_EXTERN_C tlApiResult_t tlApiRpmbWrite(
    uint32_t partition,
    uint32_t address,
    uint8_t *data,
    uint32_t blk_num);

/**
 * Read RPMB data
 *
 * @partition : The partition name in RPMB
 *  The RPMB partition is divided into 8 for multi users
 *  Each partition size is 512KB
 *  It can be 1,2,3,4,5,6,7,8
 * @address : The start address of RPMB for write
 *  The start address is '0' each partition
 *  The data size is 256byte per address '1'
 * @blk_num : The block number of input data length in bytes
 *  The data size is 256byte per block '1'
 *  This value must be less than '256' (32KB at once write)
 * @*output_data : The virtual address of output buffer
 *  in read RPMB data
 */
_TLAPI_EXTERN_C tlApiResult_t tlApiRpmbRead(
    uint32_t partition,
    uint32_t address,
    uint32_t blk_num,
    uint8_t *output_data);

/**
 * Remove RPMB partition
 *
 * @return : tlApiResult_t
 *  @TLAPI_OK : The information of register TA is removed successfully.
 *  @E_TLAPI_DRV_UNKNOWN : Failed to finalized exynos driver.
 *  @E_RPMB_NOT_REGISTERED_TA : TA is not registered anywhere.
 */
_TLAPI_EXTERN_C tlApiResult_t tlApiRemovePartition(void);

/**
 * Check RPMB provision state
 *
 * @*data : The address of the data to store rpmb status
 * @return : tlApiResult_t
 *  @TLAPI_OK :
 *  @E_TLAPI_DRV_UNKNOWN : Failed to finalized exynos driver.
 *  @E_RPMB_NOT_PROVISIONED : RPMB is not provisioned.
 */
_TLAPI_EXTERN_C tlApiResult_t tlApiRpmbCheckProvision(uint32_t *state);

/**
 * Check RPMB key provision
 *
  * @return : tlApiResult_t
 *  @TLAPI_OK :
 *  @E_TLAPI_DRV_UNKNOWN : Failed to finalized exynos driver.
 *  @E_RPMB_NOT_PROVISIONED : RPMB key provisioned failed.
 */
_TLAPI_EXTERN_C tlApiResult_t tlApiRpmbKeyProvision(void);

#endif /* srpmb */

///////////////////////////////////////////////////////////////////////////////

#define SIZE_OF_PARTITION  (1024 * 512)
#define SIZE_OF_BLOCK  256
#define POWER_OF_BLOCK 8
#define BLK_NUM(__size) \
        ((__size + SIZE_OF_BLOCK -1 ) >> POWER_OF_BLOCK)


TEE_Result rpmbWrite(IN const uint8_t* data, const uint32_t dataSize) {

    LOG_FUNC_BEGIN

    tlApiResult_t   ret;
    uint32_t partition = 0; //dynamic partition ID
    uint32_t address = 0;
    uint32_t blk_num = BLK_NUM(dataSize);
    uint32_t status;

    ret = tlApiRpmbKeyProvision();
    if (ret != TLAPI_OK) {
        LOG_E("Fail to tlApiRpmbKeyProvision, ret = 0x%08x", ret);
        return TEE_ERROR_GENERIC;
    }

    ret = tlApiRpmbCheckProvision(&status);
    LOG_D("tlApiRpmbCheckProvision, status = 0x%08x", status);
    if (ret != TLAPI_OK) {
        LOG_E("Fail to tlApiRpmbCheckProvision, ret = 0x%08x", ret);
        return TEE_ERROR_GENERIC;
    }

    uint8_t *__data = tzwMalloc(blk_num * SIZE_OF_BLOCK);
    if(!__data){
        LOG_E("Fail to tzwMalloc");
        return TEE_ERROR_OUT_OF_MEMORY;
    }

    tzwMemMove(__data, data, dataSize);

    LOG_DEV("blk_num: %d", blk_num);
    ret = tlApiRpmbWrite(partition, address, (uint8_t*)__data, blk_num);
    if (ret != TLAPI_OK) {
        tzwFree(__data);
        LOG_E("Fail to tlApiRpmbWrite  ret = 0x%08x", ret);
        return TEE_ERROR_ACCESS_DENIED;
    }

    tzwFree(__data);

    LOG_FUNC_END

    return TEE_SUCCESS;
}

TEE_Result rpmbRead(IN OUT uint8_t* buffer, const uint32_t bufferSize)
{

    LOG_FUNC_BEGIN

    tlApiResult_t   ret;
    uint32_t partition = 0; //dynamic partition ID
    uint32_t address = 0;
    uint32_t blk_num = BLK_NUM(bufferSize);
    uint32_t status;

    ret = tlApiRpmbKeyProvision();
    if (ret != TLAPI_OK) {
        LOG_E("Fail to tlApiRpmbKeyProvision, ret = 0x%08x", ret);
        return TEE_ERROR_GENERIC;
    }

    ret = tlApiRpmbCheckProvision(&status);
    LOG_D("tlApiRpmbCheckProvision, status = 0x%08x", status);
    if (ret != TLAPI_OK) {
        LOG_E("Fail to tlApiRpmbCheckProvision, ret = 0x%08x", ret);
        return TEE_ERROR_GENERIC;
    }

    uint8_t *__data = tzwMalloc(blk_num * SIZE_OF_BLOCK);
    if(!__data){
        LOG_E("Fail to tzwMalloc");
        return TEE_ERROR_OUT_OF_MEMORY;
    }

    ret = tlApiRpmbRead(partition, address, blk_num, __data);
    if (ret != TLAPI_OK) {
        tzwFree(__data);
        LOG_E("Fail to tlApiRpmbRead, ret = 0x%08x", ret);
        return TEE_ERROR_ITEM_NOT_FOUND;
    }

    tzwMemMove(buffer, __data, bufferSize);
    tzwFree(__data);

    LOG_FUNC_END

    return TEE_SUCCESS;
}

TEE_Result rpmbClear()
{

    LOG_FUNC_BEGIN

    tlApiResult_t ret;
    ret = tlApiRpmbKeyProvision();
    if (ret != TLAPI_OK) {
        LOG_E("Fail to tlApiRpmbKeyProvision, ret = 0x%08x", ret);
        return TEE_ERROR_GENERIC;
    }

    ret = tlApiRemovePartition();
    if (ret != TLAPI_OK) {
        LOG_E("Fail to tlApiRemovePartition, ret = 0x%08x", ret);
        return TEE_ERROR_GENERIC;
    }

    LOG_FUNC_END

    return TEE_SUCCESS;
}
#endif /* ! MT6768*/

#endif

////////////////////////////////////////////////////////////////////////////////////
//
#if defined (__USE_TESTBED__)
TEE_Result testRpmbWrite(const uint8_t* data, const uint32_t lenData, uint8_t* response, uint32_t* lenRsp){

    LOG_FUNC_BEGIN

    logByteArrayHex(data, lenData, "input data for testRpmbWrite");

    TEE_Result ret = rpmbWrite(data, lenData);
    if(TEE_SUCCESS != ret){
        LOG_E("Fail to rpmbWrite, ret = 0x%08x", ret);
    }

    *lenRsp = 0;

    LOG_FUNC_END

    return ret;
}

TEE_Result testRpmbRead(const uint8_t* data, const uint32_t lenData, uint8_t* response, uint32_t* lenRsp){

    LOG_FUNC_BEGIN

    uint32_t bufferSize = read32(data);

    TEE_Result ret = rpmbRead(response, bufferSize);
    if(TEE_SUCCESS != ret){
        LOG_E("Fail to rpmbRead, ret = 0x%08x", ret);
    }

    *lenRsp = bufferSize;

    logByteArrayHex(response, bufferSize, "read data from testRpmbRead");

    LOG_FUNC_END

    return TEE_SUCCESS;
}

TEE_Result testRpmbClear(const uint8_t* data, const uint32_t lenData, uint8_t* response, uint32_t* lenRsp){

    LOG_FUNC_BEGIN

    rpmbClear();

    LOG_FUNC_END

    return TEE_SUCCESS;
}
#endif