/*
 * 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"
#include "TzwMemory.h"
#include "TzwString.h"

// QSEE includes
#if defined(TIGER_TZ_MODEL_QCOM)
    #include <qsee_stor.h>
    #include <gpTypes.h>
#elif defined(TIGER_TZ_MODEL_BLOWFISH)
    #include "tlRpmbDriverApi.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(TIGER_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)
//got from MTK_RPMB_User_Guide.pdf:
#define RPMB_INVAID_SESSION 0xFFFFFFFF
#define RSA_2048_MAX_SIZE 2048

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

    uint32_t sessionId = TEE_RpmbOpenSession(SOTER_ATTK);
    if (sessionId == RPMB_INVAID_SESSION) {
        LOG_E("Failed to open RpmbSession.");
        return TEE_ERROR_DEVICE_ACCESS_DENIED;
    }

    TEE_Result status = TEE_SUCCESS;
    for (uint32_t copyNum = 0; copyNum < ATTK_COPIES_NUM; ++copyNum) {
        uint32_t writePosition = copyNum * RSA_2048_MAX_SIZE;
        LOG_D("Current Index %u , position(%u)", copyNum, writePosition);

        int ret = 0;
        TEE_RpmbWriteDatabyOffset(sessionId, writePosition, (uint8_t*) data, dataSize, &ret);
        if (0 != ret) {
            LOG_E("Failed to TEE_RpmbWriteDatabyOffset %d", ret);
            status = TEE_ERROR_STORAGE_NOT_AVAILABLE;
            break;
        }
        LOG_D("Attk copy #%u was successfully saved!", copyNum);
    }

    TEE_RpmbCloseSession(sessionId);
    return status;
}

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

    uint32_t sessionId = TEE_RpmbOpenSession(SOTER_ATTK);
    if (sessionId == RPMB_INVAID_SESSION) {
        LOG_E("Failed to open RpmbSession.");
        return TEE_ERROR_DEVICE_ACCESS_DENIED;
    }

    TEE_Result status = TEE_SUCCESS;

    for (uint32_t readAttempts = 0; readAttempts < ATTK_COPIES_NUM; ++readAttempts) {
        uint32_t readPosition = readAttempts * RSA_2048_MAX_SIZE;
        LOG_D("Current Index %u , position(%u)", readAttempts, readPosition);

        int ret = 0;
        TEE_RpmbReadDatabyOffset(sessionId, readPosition, buffer, bufferSize, &ret);
        if (0 != ret) {
            LOG_E("Failed to read %d", ret);
            status = TEE_ERROR_STORAGE_NOT_AVAILABLE;
            continue;
        }

        status = TEE_SUCCESS;
        break;
    }

    TEE_RpmbCloseSession(sessionId);
    return status;
}

TEE_Result rpmbClear() {
    uint8_t* buffer = tzwMalloc(RSA_2048_MAX_SIZE);
    TIGER_CHECK_BUFFER_ALLOCATED_RETURN(buffer);

    uint32_t sessionId = TEE_RpmbOpenSession(SOTER_ATTK);
    if (sessionId == RPMB_INVAID_SESSION) {
        LOG_E("Failed to open RpmbSession.");
        tzwFree(buffer);
        return TEE_ERROR_DEVICE_ACCESS_DENIED;
    }

    for (uint32_t copyNum = 0; copyNum < ATTK_COPIES_NUM; ++copyNum) {
        uint32_t writePosition = copyNum * RSA_2048_MAX_SIZE;
        LOG_D("Current Index %u , position(%u)", copyNum, writePosition);

        int ret = 0;
        TEE_RpmbWriteDatabyOffset(sessionId, writePosition, (uint8_t*) buffer, RSA_2048_MAX_SIZE, &ret);
        if (0 != ret) {
            LOG_E("Failed to TEE_RpmbWriteDatabyOffset %d", ret);
            continue;
        }
    }

    TEE_RpmbCloseSession(sessionId);
    tzwFree(buffer);

    return TEE_SUCCESS;
}
#endif