/*
 * 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 "TigerMacroses.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>*/
#endif

#if defined (__USE_MTK__)
    #include <mediatek/tlRpmbDriverApi.h>	
    #include <mediatek/TlApiError.h>	
	#include "TigerStorageUtils.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;
}

// BLOWFISH FUNCTIONS IMPLEMENTATION
#elif defined(TIGER_TZ_MODEL_BLOWFISH)

#define TYPE_FLAG_VAL_OF_RPMB  RPMB_TYPE_BYTE
#define SOTER_RPMB_PARTITION_ID      17

#define TMP_ATTK_FILE "TPM_ATTK_FILE"
TEE_Result rpmbWrite(IN const uint8_t* data, const uint32_t dataSize) {
#if defined(__USE_MTK__)
    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;
#else
    TEE_Result status = TEE_SUCCESS;

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

    LOG_D("RPMB Write Partition ID[SOTER] is %d", SOTER_RPMB_PARTITION_ID);

    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;
#endif
}

TEE_Result rpmbRead(IN OUT uint8_t* buffer, const uint32_t bufferSize) {
#if defined(__USE_MTK__)
    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;
#else
    TEE_Result status = TEE_SUCCESS;

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

    LOG_D("RPMB Read Partition ID[SOTER] is %d", SOTER_RPMB_PARTITION_ID);

    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;
#endif
}

TEE_Result rpmbClear() {
    LOG_FUNC_BEGIN

    LOG_D("rpmbClear dummy");
    LOG_FUNC_END
    return TEE_SUCCESS;
}
#endif