/*
 * =====================================================================================
 *
 *  Filename:  kg_rpmb.c
 *
 *  Description:  KG RPMB interfaces
 *
 *  Version:  1.0
 *  Created:  03/13/2020 10:04:00 AM
 *  Revision:  none
 *  Compiler:  gcc
 *
 *  Company:  Samsung Electronics
 *  Copyright (c) 2020 by Samsung Electronics, All rights reserved.
 *
 * =====================================================================================
 */

#include "kg_rpmb.h"

#ifdef TEEGRIS_USE_MTK_RPMB_DRV
#include "tlRpmbDriverApi.h"
#endif

#define KG_SINGLE_BLOCK_SIZE    256
/**
 * @brief
 * kg_rpmb_init
 * KG RPMB partition initialization
 *
 * @return KG status code
 */
uint32_t kg_rpmb_init(void) 
{
    KG_LOG("kg_rpmb_init()");
#ifdef TEEGRIS_USE_MTK_RPMB_DRV
    return KG_SUCCESS;
#else    
    TEE_Result ret = TEES_RPMBCheckEnable();

    if (TEE_SUCCESS != ret) {
        KG_LOG("RPMB is unavailable : 0x%x", ret);
        return KG_RPMB_UNAVAILABLE;
    }
    return KG_SUCCESS;
#endif    
}

#if 0 // rpmb read/write with rpmb byte I/O
/**
 * @brief
 * kg_rpmb_read
 * Read RPMB data
 *
 * @param[out] *data     - data buffer with RPMB content
 * @param[in]  *data_len - read data length
 *
 * @return KG status code
 */
uint32_t kg_rpmb_read(uint8_t *data, uint32_t data_len) 
{
    KG_LOG("kg_rpmb_read()");
    TEE_Result tee_ret = TEE_SUCCESS;

    tee_ret = TEES_RPMBRead(KG_RPMB_PARTITION_ID, 0, data, data_len, RPMB_TYPE_BYTE);
    if (tee_ret != TEE_SUCCESS) {
        KG_LOG("KG TA failed to read RPMB : 0x%x", tee_ret);
        return KG_RPMB_READ_FAIL;
    }

    KG_LOG("Success to read RPMB : 0x%x", tee_ret);
    return KG_SUCCESS;
}

/**
 * @brief
 * kg_rpmb_write
 * Write KG TA data to RPMB
 *
 * @param[out] *data     - data buffer with RPMB content
 * @param[in]  *data_len - read data length
 *
 * @return KG status code
 */
uint32_t kg_rpmb_write(uint8_t *data, uint32_t data_len) 
{
    KG_LOG("kg_rpmb_write() data_len : %d", data_len);
    TEE_Result tee_ret = TEE_SUCCESS;

    tee_ret = TEES_RPMBWrite(KG_RPMB_PARTITION_ID, 0, data, data_len, RPMB_TYPE_BYTE);
    if (tee_ret != TEE_SUCCESS) {
        KG_LOG("KG TA failed to write RPMB : 0x%x", tee_ret);
        return KG_RPMB_WRITE_FAIL;
    }

    KG_LOG("success to write RPMB : 0x%x", tee_ret);
    return KG_SUCCESS;
}
#endif 

/* rpmb block I/O */
uint32_t kg_rpmb_read(uint8_t *data, uint32_t data_len)
{
    TEE_Result tee_ret = TEE_SUCCESS;
    //uint8_t *pTemp = NULL;
    int i;
    //int read_cnt = (int)((float)data_len / 256 + 0.999);
    int read_cnt = data_len / KG_SINGLE_BLOCK_SIZE;
    //if (data_len % 256) read_cnt += 1;

    //pTemp = (uint8_t *)data;
    KG_LOG("kg_rpmb_block_read()... %d %d\n",data_len,read_cnt);
    for (i = 0; i < read_cnt; i++) {
        tee_ret = TEES_RPMBRead(KG_RPMB_PARTITION_ID, i, data+(KG_SINGLE_BLOCK_SIZE*i)/*pTemp*/, 1/*KG_RPMB_BLOCK_UNIT*/, RPMB_TYPE_BLOCK);
        if (tee_ret != TEE_SUCCESS) {
            KG_LOG("%s: Something wrong while reading RPMB(%d)\n",__func__,tee_ret);
            return KG_RPMB_READ_FAIL;
        }
        //pTemp += 256;
    }
    if (data_len % KG_SINGLE_BLOCK_SIZE) {
        tee_ret = TEES_RPMBRead(KG_RPMB_PARTITION_ID, read_cnt, data+read_cnt, (data_len % KG_SINGLE_BLOCK_SIZE), RPMB_TYPE_BYTE);
        if (tee_ret != TEE_SUCCESS) {
            KG_LOG("%s: Something wrong while reading RPMB /w Byte I/O(%d)\n",__func__,tee_ret);
            return KG_RPMB_READ_FAIL;
        }
    }
    KG_LOG("Success to block read RPMB : 0x%x",tee_ret);
    return KG_SUCCESS;
}
uint32_t kg_rpmb_write(uint8_t *data, uint32_t data_len)
{
    TEE_Result tee_ret = TEE_SUCCESS;
    //uint8_t *pTemp = NULL;
    int i;
    //int write_cnt = (int)((float)data_len / 256 + 0.999) + 1;
    int write_cnt = data_len / KG_SINGLE_BLOCK_SIZE;
    //if (data_len % 256) write_cnt += 1;

    //pTemp = (uint8_t *)data;
    KG_LOG("kg_rpmb_block_write()... %d %d\n",data_len,write_cnt);
    for (i = 0; i < write_cnt; i++) {
        tee_ret = TEES_RPMBWrite(KG_RPMB_PARTITION_ID, i, data+(KG_SINGLE_BLOCK_SIZE*i)/*pTemp*/, 1/*KG_RPMB_BLOCK_UNIT*/, RPMB_TYPE_BLOCK);
        if (tee_ret != TEE_SUCCESS) {
            KG_LOG("%s: Something wrong while writing RPMB(%d)\n",__func__,tee_ret);
            return KG_RPMB_WRITE_FAIL;
        }
        //pTemp += 256;
    }
    if (data_len % KG_SINGLE_BLOCK_SIZE) {
        tee_ret = TEES_RPMBWrite(KG_RPMB_PARTITION_ID, write_cnt, data+write_cnt, (data_len % KG_SINGLE_BLOCK_SIZE), RPMB_TYPE_BYTE);
        if (tee_ret != TEE_SUCCESS) {
            KG_LOG("%s: Something wrong while writing RPMB /w Byte I/O(%d)\n",__func__,tee_ret);
            return KG_RPMB_WRITE_FAIL;
        }
    }
    KG_LOG("Success to block write RPMB : 0x%x",tee_ret);
    return KG_SUCCESS;
}

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

    // 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_data_with_derived_key((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 Teegris\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(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, sizeof(kg_secure_data_t));
        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;
    }

    KG_DUMP_DBG("read wrapped secure data when writing metadata: \n", krd->kg_wrap_data, krd->kg_wrap_data_len);
    if (TZ_API_OK != TZ_unwrap_data_with_derived_key((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 unwrap kg secure data structure in Teegris\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }

    KG_LOG_DBG("unwrap data len is %d\n", data_buf_len);

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

    TEE_MemMove((uint8_t *)ksd, data_buf, data_buf_len);
    TEE_MemMove(&(ksd->kg_metadata), data, data_len);
    KG_DUMP_DBG("read secure data when writing metadata: \n", (uint8_t *)ksd, sizeof(*ksd));

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

    if (TZ_API_OK != TZ_wrap_data_with_derived_key((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 Teegris\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);

    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;
    }
    KG_LOG_DBG("finish calling write rpmb metadata\n");
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_data_with_derived_key((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 Teegris\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_data_with_derived_key((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 Teegris\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;
}

#ifdef TEEGRIS_USE_MTK_RPMB_DRV
#define KG_SINGLE_BLOCK_SIZE    256
#define MTK_MAX_RPMB_TRANSFER_BLK   32
#define MTK_MAX_RPMB_REQUEST_SIZE   (KG_SINGLE_BLOCK_SIZE*MTK_MAX_RPMB_TRANSFER_BLK)

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

    do
    {
        crSession = tlApiRpmbOpenSession(KNOXGUARD_ID);
        if (crSession == 0xFFFFFFFF) {
            KG_LOG("Faile to get a permission and a valid ID");
            return KG_RPMB_UNAVAILABLE;
        }

        uint32_t i = 0;
        uint32_t mba_round_cnt = 0; // RPMB_MULTI_BLOCK_ACCESS round counter
        uint32_t mtk_start_block_index = start_block_index*2;

        uint32_t read_loop_count = data_len / KG_SINGLE_BLOCK_SIZE;
        if (data_len % KG_SINGLE_BLOCK_SIZE) {
            read_loop_count++;
        }

        KG_LOG("read_loop_count : %d\n", read_loop_count);

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

        if (MTK_MAX_RPMB_TRANSFER_BLK <= read_loop_count) {
            mba_round_cnt = read_loop_count / MTK_MAX_RPMB_TRANSFER_BLK;
            for (i=0; i<mba_round_cnt; i++) {
                __tlApiRpmbReadData(crSession, 
                                    (mtk_start_block_index*KG_SINGLE_BLOCK_SIZE)+i*(MTK_MAX_RPMB_TRANSFER_BLK*KG_SINGLE_BLOCK_SIZE),
                                    tempBuffer+i*(MTK_MAX_RPMB_TRANSFER_BLK*KG_SINGLE_BLOCK_SIZE),
                                    MTK_MAX_RPMB_TRANSFER_BLK*KG_SINGLE_BLOCK_SIZE,
                                    &result);
                if (result) {
                    KG_LOG("Failed to tlApiRpmbReadData1 : (%d)\n", result);
                    ret = KG_RPMB_READ_FAIL;
                    goto exit;
                }
            }
            if (read_loop_count % MTK_MAX_RPMB_TRANSFER_BLK) {
                __tlApiRpmbReadData(crSession,
                                    (mtk_start_block_index*KG_SINGLE_BLOCK_SIZE)+mba_round_cnt*(MTK_MAX_RPMB_TRANSFER_BLK*KG_SINGLE_BLOCK_SIZE),
                                    tempBuffer+mba_round_cnt*(MTK_MAX_RPMB_TRANSFER_BLK*KG_SINGLE_BLOCK_SIZE),
                                    (read_loop_count % MTK_MAX_RPMB_TRANSFER_BLK)*KG_SINGLE_BLOCK_SIZE,
                                    &result);
                if (result) {
                    KG_LOG("Failed to tlApiRpmbReadData2 : (%d)\n", result);
                    ret = KG_RPMB_READ_FAIL;
                    goto exit;
                }
            }
        } else {
            __tlApiRpmbReadData(crSession, mtk_start_block_index*KG_SINGLE_BLOCK_SIZE, tempBuffer, read_loop_count*KG_SINGLE_BLOCK_SIZE, &result);
            if (result) {
                KG_LOG("Failed to tlApiRpmbReadData0 : (%x)\n", result);
                ret = KG_RPMB_READ_FAIL;
                goto exit;
            }
        }

exit:
        tlApiRpmbCloseSession(crSession);
    } while(false);

    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 crSession;
    int result;

    do {
        crSession = tlApiRpmbOpenSession(KNOXGUARD_ID);
        if (crSession == 0xFFFFFFFF) {
            KG_LOG("Failed to get a permission and a valid ID");
            return KG_RPMB_UNAVAILABLE;
        }

        uint32_t i = 0;
        uint32_t mba_round_cnt = 0; // RPMB_MULTI_BLOCK_ACCESS round counter
        uint32_t mtk_start_block_index = start_block_index*2;

        uint32_t write_loop_count = data_len / KG_SINGLE_BLOCK_SIZE;
        if (data_len % KG_SINGLE_BLOCK_SIZE) {
            write_loop_count++;
        }
        KG_LOG("write_loop_count : %d\n", write_loop_count);

        if (MTK_MAX_RPMB_TRANSFER_BLK <= write_loop_count) {
            mba_round_cnt = write_loop_count / MTK_MAX_RPMB_TRANSFER_BLK;
            for (i=0; i<mba_round_cnt; i++) {
                __tlApiRpmbWriteData(crSession,
                                     (mtk_start_block_index*KG_SINGLE_BLOCK_SIZE)+i*(MTK_MAX_RPMB_TRANSFER_BLK*KG_SINGLE_BLOCK_SIZE),
                                     data+i*(MTK_MAX_RPMB_TRANSFER_BLK*KG_SINGLE_BLOCK_SIZE),
                                     MTK_MAX_RPMB_TRANSFER_BLK*KG_SINGLE_BLOCK_SIZE,
                                     &result);
                if (result) {
                    KG_LOG("Failed to tlApiRpmbWriteData1 : (%x)\n", result);
                    ret = KG_RPMB_WRITE_FAIL;
                    goto exit;
                }
            }
            if (write_loop_count % MTK_MAX_RPMB_TRANSFER_BLK) {
                __tlApiRpmbWriteData(crSession,
                                     (mtk_start_block_index*KG_SINGLE_BLOCK_SIZE)+mba_round_cnt*(MTK_MAX_RPMB_TRANSFER_BLK*KG_SINGLE_BLOCK_SIZE),
                                     data+mba_round_cnt*(MTK_MAX_RPMB_TRANSFER_BLK*KG_SINGLE_BLOCK_SIZE),
                                     (write_loop_count % MTK_MAX_RPMB_TRANSFER_BLK)*KG_SINGLE_BLOCK_SIZE,
                                     &result);
                if (result) {
                    KG_LOG("Failed to tlApiRpmbWriteData2 : (%x)\n", result);
                    ret = KG_RPMB_WRITE_FAIL;
                    goto exit;
                }
            }
        } else {
            __tlApiRpmbWriteData(crSession, mtk_start_block_index*KG_SINGLE_BLOCK_SIZE, data, write_loop_count*KG_SINGLE_BLOCK_SIZE, &result);
            if (result) {
                KG_LOG("Failed to tlApiRpmbWriteData0 : (%x)\n", result);
                ret = KG_RPMB_WRITE_FAIL;
                goto exit;
            }
/*    
            if (data_len % KG_SINGLE_BLOCK_SIZE) {
                __tlApiRpmbWriteData(crSession, (mtk_start_block_index+write_loop_count)*KG_SINGLE_BLOCK_SIZE, data+(KG_SINGLE_BLOCK_SIZE*write_loop_count), (data_len % KG_SINGLE_BLOCK_SIZE), &result);
                if (result) {
                    KG_LOG("Failed to tlApiRpmbWriteData2 : (%x)\n", result);
                    ret = KG_RPMB_WRITE_FAIL;
                    goto exit;
                }
            }
*/
        }
exit:
        tlApiRpmbCloseSession(crSession);
    } while(false);

    return ret;
}
#else
#define KG_SINGLE_BLOCK_SIZE    256
uint32_t read_block(uint32_t start_block_index, uint32_t data_len, uint8_t **data) {
    KG_LOG("read_block start block index : %d, data_len : %d\n", start_block_index, data_len);
    uint32_t ret = KG_SUCCESS;
    TEE_Result tee_ret = TEE_SUCCESS;
    uint32_t i = 0;
    uint32_t teegris_start_block_index = start_block_index * 2; // need to change calculated value
    
    uint32_t read_loop_count = data_len / KG_SINGLE_BLOCK_SIZE;
    if (data_len % KG_SINGLE_BLOCK_SIZE) {
        read_loop_count++;
    }

    KG_LOG("read_loop_count : %d\n", read_loop_count);

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

    for (i=0; i<read_loop_count; ++i) {
        tee_ret = TEES_RPMBRead(KG_RPMB_PARTITION_ID, i+teegris_start_block_index, tempBuffer+(KG_SINGLE_BLOCK_SIZE*i), 1, RPMB_TYPE_BLOCK);
        if (TEE_SUCCESS != tee_ret) {
            KG_LOG("Failed to TEES_RPMBRead : %d\n", tee_ret);
            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;
    TEE_Result tee_ret = TEE_SUCCESS;
    uint32_t i = 0;
    uint32_t teegris_start_block_index = start_block_index * 2; // need to change calculated value

    uint32_t write_loop_count = data_len / KG_SINGLE_BLOCK_SIZE;
    KG_LOG("write_loop_count : %d\n", write_loop_count);

    for (i=0; i<write_loop_count; ++i) {
        tee_ret = TEES_RPMBWrite(KG_RPMB_PARTITION_ID, i+teegris_start_block_index, data+(KG_SINGLE_BLOCK_SIZE*i)/*pTemp*/, 1/*KG_RPMB_BLOCK_UNIT*/, RPMB_TYPE_BLOCK);
        if (tee_ret != TEE_SUCCESS) {
            KG_LOG("%s: Something wrong while writing RPMB(%d)\n",__func__,tee_ret);
            return KG_RPMB_WRITE_FAIL;
        }
    }

    if (data_len % KG_SINGLE_BLOCK_SIZE) {
        tee_ret = TEES_RPMBWrite(KG_RPMB_PARTITION_ID, (write_loop_count+teegris_start_block_index)*KG_SINGLE_BLOCK_SIZE, data+(KG_SINGLE_BLOCK_SIZE*write_loop_count)/*pTemp*/, (data_len % KG_SINGLE_BLOCK_SIZE)/*KG_RPMB_BLOCK_UNIT*/, RPMB_TYPE_BYTE);
        if (tee_ret != TEE_SUCCESS) {
            KG_LOG("%s: Something wrong while writing RPMB(%d)\n",__func__,tee_ret);
            return KG_RPMB_WRITE_FAIL;
        }
    }

    return ret;
}
#endif