#include "kg_state_cmd.h"

uint32_t KG_get_nonce(tz_nonce_payload_t *sendmsg, tz_nonce_payload_t *respmsg) {
    KG_LOG("KG get ta nonce\n");

    uint32_t ret = KG_SUCCESS;

    uint8_t *nonce_buffer = NULL;
    uint32_t nonce_buffer_len = KG_NONCE_LEN;
    uint8_t *server_buffer = NULL;
    uint32_t server_buffer_len = KG_BUF_LEN;
    uint8_t *server_sig = NULL;
    uint32_t server_sig_len = KG_BUF_LEN;
    uint8_t *pub_mod = NULL;
    uint32_t pub_mod_len = KG_BUF_LEN;
    uint8_t *pub_exp = NULL;
    uint32_t pub_exp_len = KG_BUF_LEN;

    uint8_t *payload_buffer = NULL;
    uint32_t payload_buffer_len = KG_BUF_LEN;
    uint8_t *payload_buffer_b64 = NULL;
    uint32_t payload_buffer_b64_len = KG_B64_BUF_LEN;

    kg_rpmb_info_t* info = NULL;
    uint8_t* unwrap_data = NULL;
    uint32_t unwrap_data_len = KG_BUF_LEN;
    uint8_t* wrap_data = NULL;
    uint8_t* rewrap_data = NULL;
    uint32_t rewrap_data_len = KG_SECURE_DATA_LEN;
    
    EVP_PKEY *pkey = NULL;
    RSA *rsakey = NULL;

    int region_flag = REGION_NONE;

    nonce_buffer = TEE_Malloc(nonce_buffer_len, 0);
    if (NULL == nonce_buffer) {
        KG_LOG("Failed to alloc buffer for nonce buffer\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    server_buffer = TEE_Malloc(server_buffer_len, 0);
    if (NULL == server_buffer) {
        KG_LOG("Failed to alloc buffer for server msg\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    server_sig = TEE_Malloc(server_sig_len, 0);
    if (NULL == server_sig) {
        KG_LOG("Failed to alloc buffer for server sig\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    pub_mod = TEE_Malloc(pub_mod_len, 0);
    if (NULL == pub_mod) {
        KG_LOG("Failed to alloc buffer for pub_mod\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    pub_exp = TEE_Malloc(pub_exp_len, 0);
    if (NULL == pub_exp) {
        KG_LOG("Failed to alloc buffer for pub_mod\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    payload_buffer = TEE_Malloc(payload_buffer_len, 0);
    if (NULL == payload_buffer) {
        KG_LOG("Failed to alloc buffer for payload buffer\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    payload_buffer_b64 = TEE_Malloc(payload_buffer_b64_len, 0);
    if (NULL == payload_buffer_b64) {
        KG_LOG("Failed to alloc buffer for payload buffer b64\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    if (sendmsg->payload.cmd.svr_msg_len > KG_BUF_LEN) {
        KG_LOG_DBG("Recived invalid svr msg buffer\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    if (sendmsg->payload.cmd.svr_sig_len > KG_BUF_LEN) {
        KG_LOG_DBG("Recived invalid svr sig buffer\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    if (sendmsg->payload.cmd.svr_sig_len == 0 || sendmsg->payload.cmd.svr_msg_len == 0) {
        KG_LOG("KG reg info is not set");
        ret = KG_REGION_NOT_CONFIG;
        goto exit;
    }

    KG_LOG_DBG("Input svr msg len %d\n", sendmsg->payload.cmd.svr_msg_len);
    KG_LOG_DBG("Input svr msg content %s\n", sendmsg->payload.cmd.svr_msg_buf);
    if (BASE64_OK != base64_decode(sendmsg->payload.cmd.svr_msg_buf,
            sendmsg->payload.cmd.svr_msg_len, server_buffer, &server_buffer_len)) {
        KG_LOG("Failed to decode base64 encoded input server buffer\n");
        ret = KG_BASE64_DECODE_FAIL;
        goto exit;
    }
    if (server_buffer_len > KG_BUF_LEN) {
        KG_LOG("Decoded server msg length check failed\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    KG_LOG_DBG("Input svr signature len %d\n", sendmsg->payload.cmd.svr_sig_len);
    KG_LOG_DBG("Input svr signature content %s\n", sendmsg->payload.cmd.svr_sig_buf);
    if (BASE64_OK != base64_decode(sendmsg->payload.cmd.svr_sig_buf,
            sendmsg->payload.cmd.svr_sig_len, server_sig, &server_sig_len)) {
        KG_LOG("Failed to decode base64 encoded input server signature\n");
        ret = KG_BASE64_DECODE_FAIL;
        goto exit;
    }

    bool verified_eu = false;
    if (KG_SUCCESS != extract_public_keybytes(enrollment_cert_EU, pub_mod, &pub_mod_len, pub_exp, &pub_exp_len)) {
        KG_LOG_DBG("Failed to extract publick key bytes from policy signing cert EU\n");
        ret = KG_CRYPTO_PKEY_PARSE_FAIL;
        goto check_US;
    }
    KG_DUMP_DBG("dump pubkey n:\n", pub_mod, pub_mod_len);
    KG_DUMP_DBG("dump pubkey e:\n", pub_exp, pub_exp_len);

    if (TZ_API_OK != TZ_verify_CKM_SHA256_RSA_PKCS(pub_mod, pub_mod_len, pub_exp, pub_exp_len,
            server_buffer, server_buffer_len, server_sig, KG_SIG_LEN, &verified_eu)) {
        KG_LOG_DBG("TZ API failed when verifying registration info\n");
        ret = KG_TZ_API_FAIL;
        goto check_US;
    }

check_US:
    TEE_MemFill(pub_mod, 0, KG_BUF_LEN);
    pub_mod_len = KG_BUF_LEN;
    TEE_MemFill(pub_exp, 0, KG_BUF_LEN);
    pub_exp_len = KG_BUF_LEN;

    if (KG_SUCCESS != extract_public_keybytes(enrollment_cert_US, pub_mod, &pub_mod_len, pub_exp, &pub_exp_len)) {
        KG_LOG_DBG("Failed to extract publick key bytes from policy signing cert US\n");
        ret = KG_CRYPTO_PKEY_PARSE_FAIL;
        goto exit;
    }
    KG_DUMP_DBG("dump pubkey n:\n", pub_mod, pub_mod_len);
    KG_DUMP_DBG("dump pubkey e:\n", pub_exp, pub_exp_len);
    bool verified_us = false;
    if (TZ_API_OK != TZ_verify_CKM_SHA256_RSA_PKCS(pub_mod, pub_mod_len, pub_exp, pub_exp_len,
            server_buffer, server_buffer_len, server_sig, KG_SIG_LEN, &verified_us)) {
        KG_LOG_DBG("TZ API failed when verifying server signature\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }

    if (false == verified_eu && false == verified_us) {
        KG_LOG_DBG("Failed to verify server signature\n");
        ret = KG_POLICY_VERIFY_FAIL;
        goto exit;
    }
    KG_LOG("Server signature has been verified successfully\n");

    if (verified_eu == true) {
        region_flag = REGION_EU;
    } else if (verified_us == true) {
        region_flag = REGION_US;
    }

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != (ret = kg_rpmb_init())) {
        KG_LOG("KG rpmb is unavailable when verifying policy\n");
        goto exit;
    }
#endif

    if (KG_SUCCESS != (ret = read_info_object(&info))
        || NULL == info) {
        KG_LOG("failed to read info object\n");
        goto exit;
    }

    if(info->kg_wrap_data_len > KG_SECURE_DATA_LEN){
        KG_LOG("Received invalid size for wrap_data\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
/* error: taking address of packed member 'kg_wrap_data_len' of class or structure 'kg_rpmb_info' may result in an unaligned pointer value [-Werror,-Waddress-of-packed-member]*/ 
    uint32_t temp_wrap_data_len = info->kg_wrap_data_len;
    if (KG_SUCCESS != (ret = read_wrap_data(&temp_wrap_data_len, &wrap_data))
        || NULL == wrap_data) {
        KG_LOG("failed to read wrap data\n");
        goto exit;
    }
    info->kg_wrap_data_len = temp_wrap_data_len;
#else
    if (KG_SUCCESS != (ret = read_wrap_data(&(info->kg_wrap_data_len), &wrap_data))
        || NULL == wrap_data) {
        KG_LOG("failed to read wrap data\n");
        goto exit;
    }
#endif

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

#ifdef CONFIG_QSEE
    if (TZ_API_OK != TZ_unwrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        wrap_data, info->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;
    }
#else
    if (TZ_API_OK != TZ_unwrap_data_with_derived_key((uint8_t *)KG_NAME, strlen(KG_NAME), 
        wrap_data, info->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
        KG_LOG("failed to unwrap data\n")
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#endif

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

    kg_secure_data_t* secure_data = (kg_secure_data_t*)unwrap_data;

    if (region_flag != REGION_NONE && secure_data->kg_metadata.reg_info != REGION_NONE) {
        if(secure_data->kg_metadata.reg_info != region_flag) {
            KG_LOG("KG Region info mismatch when generating device nonce\n");
            ret = KG_REGION_MISMATCH;
            goto exit;
        }
    } else if(region_flag == REGION_NONE && secure_data->kg_metadata.reg_info == REGION_NONE) {
        KG_LOG("KG reg info is not set");
        ret = KG_REGION_NOT_CONFIG;
        goto exit;
    } else if(region_flag == REGION_NONE && secure_data->kg_metadata.reg_info != REGION_NONE) {
        region_flag = secure_data->kg_metadata.reg_info ;    
    }

    if (TZ_API_OK != TZ_gen_rand_data(nonce_buffer, &nonce_buffer_len) || nonce_buffer_len != KG_NONCE_LEN) {
        KG_LOG("KG nonce generation failed\n");
        ret = KG_GEN_NONCE_FAIL;
        goto exit;
    }

    KG_DUMP_DBG("Generated KG nonce is ", nonce_buffer, nonce_buffer_len);

    if (region_flag == REGION_EU) {
        pkey = get_public_key(bl_sign_cert_EU);
    } else if (region_flag == REGION_US) {
        pkey = get_public_key(bl_sign_cert_US);
    }
    
    if (NULL == pkey) {
        KG_LOG_DBG("Invalid pem cert to extract the public key\n");
        goto exit;
    }

    rsakey = EVP_PKEY_get1_RSA(pkey);
    if (!rsakey) {
        KG_LOG_DBG("Failed to create RSA key\n");
        goto exit;
    }

    if (KG_SUCCESS != kg_public_encrypt(nonce_buffer, nonce_buffer_len, rsakey, payload_buffer, (int *)&payload_buffer_len)) {
        KG_LOG_DBG("Failed to encrypted ta nonce\n");
        ret = KG_CRYPTO_RSA_ENC_FAIL;
        goto exit;
    }

    if (payload_buffer_len > KG_MAX_PAYLOAD_LEN) {
        KG_LOG("KG generated ta nonce size check failed\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    KG_DUMP_DBG("encrypted ta nonce: \n", payload_buffer, payload_buffer_len);

    if (BASE64_OK != base64_encode(payload_buffer, payload_buffer_len, payload_buffer_b64, &payload_buffer_b64_len)) {
        KG_LOG("KG base64 encoding failed when generating ta nonce\n");
        ret = KG_BASE64_ENCODE_FAIL;
        goto exit;
    }

    if (payload_buffer_b64_len > KG_B64_BUF_LEN) {
        KG_LOG("Encoded string length check failed\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    TEE_MemMove(secure_data->kg_metadata.ta_nonce, nonce_buffer, KG_NONCE_LEN);

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

#ifdef CONFIG_QSEE
    if (TZ_API_OK != TZ_wrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        unwrap_data, sizeof(kg_secure_data_t), rewrap_data, &rewrap_data_len)) {
        KG_LOG("Failed to wrap data\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#else
    if (TZ_API_OK != TZ_wrap_data_with_derived_key((uint8_t *)KG_NAME, strlen(KG_NAME), 
        unwrap_data, sizeof(kg_secure_data_t), rewrap_data, &rewrap_data_len)) {
        KG_LOG("Failed to wrap kg secure data structure\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#endif

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

    info->kg_wrap_data_len = rewrap_data_len;
    if (KG_SUCCESS != (ret = write_wrap_data(info->kg_wrap_data_len, rewrap_data))) {
        KG_LOG("failed to write wrap data\n");
        goto exit;
    }

    if (KG_SUCCESS != (ret = write_info_object(info))) {
        KG_LOG("failed to write info object\n");
        goto exit;
    }

    TEE_MemFill(respmsg->payload.resp.ta_nonce_buf, 0x0, KG_BUF_LEN);
    
    if(payload_buffer_b64_len > KG_BUF_LEN){
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    respmsg->payload.resp.ta_nonce_len = payload_buffer_b64_len;
    TEE_MemMove(respmsg->payload.resp.ta_nonce_buf, (void *)payload_buffer_b64, payload_buffer_b64_len);
    KG_LOG_DBG("Encoded result string is %s\n", payload_buffer_b64);

exit:
    if (nonce_buffer != NULL) {
        TEE_Free(nonce_buffer);
        nonce_buffer = NULL;
    }
    if (server_sig != NULL) {
        TEE_Free(server_sig);
        server_sig = NULL;
    }
    if (server_buffer != NULL) {
        TEE_Free(server_buffer);
        server_buffer = NULL;
    }
    if (pub_exp != NULL) {
        TEE_Free(pub_exp);
        pub_exp = NULL;
    }
    if (pub_mod != NULL) {
        TEE_Free(pub_mod);
        pub_mod = NULL;
    }
    if (payload_buffer != NULL) {
        TEE_Free(payload_buffer);
        payload_buffer = NULL;
    }
    if (payload_buffer_b64 != NULL) {
        TEE_Free(payload_buffer_b64);
        payload_buffer_b64 = NULL;
    }
    if (info != NULL) {
        TEE_Free(info);
        info = NULL;
    }
    if (wrap_data != NULL) {
        TEE_Free(wrap_data);
        wrap_data = NULL;
    }
    if (unwrap_data != NULL) {
        TEE_MemFill(unwrap_data, 0, unwrap_data_len);
        TEE_Free(unwrap_data);
        unwrap_data = NULL;
    }
    if (rewrap_data != NULL) {
        TEE_Free(rewrap_data);
        rewrap_data = NULL;
    }
    if (pkey != NULL) {
        EVP_PKEY_free(pkey);
    }
    if (rsakey != NULL) {
        RSA_free(rsakey);
    }
    return ret;
}

// DEBUG ONLY
void kg_rpmb_reset() {
    KG_LOG_DBG("rpmb reset start\n");
    kg_rpmb_data_t *krd = NULL;
    krd = TEE_Malloc(sizeof(kg_rpmb_data_t), 0);
    if (NULL == krd) {
        KG_LOG_DBG("Failed to alloc buffer for kg rpmb reset\n");
        goto exit;
    }

    if (KG_SUCCESS != kg_rpmb_init()) {
        goto exit;
    }

    if (KG_SUCCESS != kg_rpmb_write((uint8_t *)krd, sizeof(kg_rpmb_data_t))) {
        goto exit;
    }
    
    if (KG_SUCCESS != kg_rpmb_read((uint8_t *)krd, sizeof(kg_rpmb_data_t))) {
        goto exit;
    }

    KG_DUMP_DBG("RPMB data after reset are\n", (uint8_t *)krd, sizeof(*krd));
exit:
    if (krd != NULL) {
        TEE_Free(krd);
        krd = NULL;
    }
    return;
}
// DEBUG_ONLY

uint32_t KGBL_get_ta_status(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    KG_LOG("KGBL get_ta_status\n");
    uint32_t ret = KG_SUCCESS;
    kg_rpmb_info_t *info = NULL;

/* Finally decided to do NOT completetoken veirification in BL */
#if 0 
    uint8_t *sn_num = NULL;
    uint32_t sn_num_len;
    const char *prefix_sn = "SEC_SN=";
    uint8_t digest[SHA256_DIGEST_LENGTH];
    uint32_t digestLen= sizeof(digest);
    uint8_t *sn_buffer_b64 = NULL;
    uint32_t sn_buffer_b64_len = KG_BUF_LEN;

    kg_rpmb_complete_token_t *complete_token = NULL;
#endif

    if (sendmsg->payload.cmd.data_len > KG_MAX_PAYLOAD_LEN) {
        KG_LOG("Received invalid input buffer when get status\n");
        respmsg->payload.resp.result = KG_STATE_BROKEN;
        goto exit;
    }
    
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("KG rpmb is unavailable when retrieving kg state\n");
        respmsg->payload.resp.result = KG_STATE_BROKEN;
        goto exit;
    }

    ret = read_info_object(&info);
    if(KG_SUCCESS != ret) {
        KG_LOG("Failed to read info object\n");
        goto exit;
    }

    if (info->magic != KG_MAGIC) {
        KG_LOG("KG rpmb magic check failed\n");
        respmsg->payload.resp.result = KG_STATE_ALLZERO;
        goto exit;
    }

    //kg_state always >= 0 and = also indicate error
    if (info->kg_state > KG_STATE_ALLZERO) {
        KG_LOG("KG state corrupted\n");
        respmsg->payload.resp.result = KG_STATE_ERROR;
        goto exit;
    }

/* Finally decided to do NOT completetoken veirification in BL */
#if 0 
    if (info->kg_state == KG_STATE_COMPLETE) {
        sn_num_len = sendmsg->payload.cmd.data_len;
        sn_num = TEE_Malloc(sn_num_len, 0);
        memcpy(sn_num, sendmsg->payload.cmd.data_buf, sn_num_len);

        // Check if "SEC_SN=" exist in the sn_num buffer
        uint32_t idx = 0;
        bool check_flag = true;
        while (idx < 7) {
            if (sn_num[idx] != prefix_sn[idx]) {
                check_flag = false;
            }
            idx++;
        }

        if (false == check_flag) {
            KG_LOG_DBG("device serial number check failed\n");
            respmsg->payload.resp.result = 0x1001;
            //ret = KG_COMPLETE_MAGIC_FAIL;
            goto exit;
        }

        sn_num_len = strlen((char *)sn_num);
        ret = TZ_digest_SHA256(sn_num+idx, sn_num_len-idx, digest, &digestLen);
        //sn_num_len = 11;
        //ret = TZ_digest_SHA256(sn_num+7, 11, digest, &digestLen);
        if (ret != TZ_API_OK) {
            KG_LOG("TZ_digest_SHA1() returns error\n");
            //ret = KG_TZ_API_FAIL;
            respmsg->payload.resp.result = 0x1002;
            goto exit;
        }

        sn_buffer_b64 = TEE_Malloc(sn_buffer_b64_len, 0);
        if (NULL == sn_buffer_b64) {
            KG_LOG("Failed to alloc buffer for dh buffer b64\n");
            //ret = KG_ALLOC_BUFFER_FAIL;
            respmsg->payload.resp.result = 0x1003;
            goto exit;
        }

        KG_LOG("digestLen = %d   sn_buffer_b64_len = %d\n", digestLen, sn_buffer_b64_len);
        if (BASE64_OK != base64_encode(digest, digestLen, sn_buffer_b64, &sn_buffer_b64_len)) {
            KG_LOG("KG base64 encoding failed when generating device serial number\n");
            //ret = KG_BASE64_ENCODE_FAIL;
            respmsg->payload.resp.result = 0x1004;
            goto exit;
        }

        if (sn_buffer_b64_len > KG_SN_HASH_LEN) {
            KG_LOG("Encoded string length check failed\n");
            //ret = KG_BUFFER_SIZE_FAIL;
            respmsg->payload.resp.result = 0x1005;
            goto exit;
        }

        ret = read_complete_token(&complete_token);
        if (ret != KG_SUCCESS) {
            KG_LOG("Failed to read completetoken\n");
            respmsg->payload.resp.result = 0x1006;
            goto exit;
        }

        /* verify the token */
        ret = kgbl_verify_completetoken(complete_token->token, sn_buffer_b64);
        if (ret != KG_SUCCESS) {
            //ret = KG_COMPLETE_VERIFY_FAIL;
            if (ret == KG_COMPLETE_VERIFY_FAIL) {
                respmsg->payload.resp.result = KG_STATE_ERROR;
            } else {
                respmsg->payload.resp.result = ret;
            }
            goto exit;
        }
    }
#endif
    respmsg->payload.resp.result = info->kg_state;
    KG_LOG("KG status is %d\n", info->kg_state);
exit:
    if (NULL != info) {
        TEE_Free(info);
        info = NULL;
    }
/* Finally decided to do NOT completetoken veirification in BL */
#if 0 
    if (NULL != complete_token) {
        TEE_Free(complete_token);
        complete_token = NULL;
    }
#endif    
    return KG_SUCCESS;

}

uint32_t kgbl_verify_completetoken(uint8_t *token, uint8_t *sn_hash) {
    KG_LOG("KGBL verify complete token\n");
    uint32_t ret = KG_SUCCESS;
    uint32_t len_check = KG_SIG_LEN + KG_HASH_LEN + KG_COMPLETE_MSG_LEN;
    uint8_t *data_buffer = NULL;
    uint32_t data_buffer_len = KG_BUF_LEN;
    uint8_t *pub_mod = NULL;
    uint32_t pub_mod_len = KG_BUF_LEN;
    uint8_t *pub_exp = NULL;
    uint32_t pub_exp_len = KG_BUF_LEN;
    const char *complete_msg = "Complete";

    data_buffer = TEE_Malloc(data_buffer_len, 0);
    if (NULL == data_buffer) {
        KG_LOG("Failed to alloc buffer for data buffer\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    pub_mod = TEE_Malloc(pub_mod_len, 0);
    if (NULL == pub_mod) {
        KG_LOG("Failed to alloc buffer for pub_mod\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    pub_exp = TEE_Malloc(pub_exp_len, 0);
    if (NULL == pub_exp) {
        KG_LOG("Failed to alloc buffer for pub_mod\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    if (KG_COMPLETE_MSG_LEN + KG_SN_HASH_LEN > data_buffer_len) {
        KG_LOG_DBG("Buffer length check failed\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    TEE_MemFill(data_buffer, 0x0, data_buffer_len);
    data_buffer_len = 0;
    TEE_MemMove(data_buffer, (void *)complete_msg, KG_COMPLETE_MSG_LEN);
    data_buffer_len = KG_COMPLETE_MSG_LEN;
    TEE_MemMove(data_buffer + data_buffer_len, sn_hash, KG_SN_HASH_LEN);
    data_buffer_len += KG_SN_HASH_LEN;

    // first signature verification using EU cert
    if (KG_SUCCESS != extract_public_keybytes(bl_sign_cert_EU, pub_mod, &pub_mod_len, pub_exp, &pub_exp_len)) {
        KG_LOG_DBG("Failed to extract publick key bytes from pem certificate\n");
        ret = KG_CRYPTO_PKEY_PARSE_FAIL;
        goto exit;
    }

    KG_DUMP_DBG("dump pubkey n:\n", pub_mod, pub_mod_len);
    KG_DUMP_DBG("dump pubkey e:\n", pub_exp, pub_exp_len);
    bool verified = false;
    if (TZ_API_OK != TZ_verify_CKM_SHA256_RSA_PKCS(pub_mod, pub_mod_len, pub_exp, pub_exp_len,
            data_buffer, data_buffer_len, token, KG_SIG_LEN, &verified)) {
        KG_LOG_DBG("Failed to verify signed completeToken\n");
        ret = KG_COMPLETE_VERIFY_FAIL;
        goto exit;
    }
    if (false == verified) {
        // second signature verification using US cert
        pub_mod_len = KG_BUF_LEN;
        pub_exp_len = KG_BUF_LEN;
        TEE_MemFill(pub_mod, 0x0, pub_mod_len);
        TEE_MemFill(pub_exp, 0x0, pub_exp_len);

        if (KG_SUCCESS != extract_public_keybytes(bl_sign_cert_US, pub_mod, &pub_mod_len, pub_exp, &pub_exp_len)) {
            KG_LOG_DBG("Failed to extract publick key bytes from pem certificate\n");
            ret = KG_CRYPTO_PKEY_PARSE_FAIL;
            goto exit;
        }
        KG_DUMP_DBG("dump pubkey n:\n", pub_mod, pub_mod_len);
        KG_DUMP_DBG("dump pubkey e:\n", pub_exp, pub_exp_len);

        if (TZ_API_OK != TZ_verify_CKM_SHA256_RSA_PKCS(pub_mod, pub_mod_len, pub_exp, pub_exp_len,
            data_buffer, data_buffer_len, token, KG_SIG_LEN, &verified)) {
            KG_LOG_DBG("Failed to verify signed completeToken\n");
            ret = KG_COMPLETE_VERIFY_FAIL;
            goto exit;
        }

        if (false == verified) {
            KG_LOG_DBG("Failed to verify complete token signature\n");
            ret = KG_COMPLETE_VERIFY_FAIL;
            goto exit;
        }
    }
    KG_LOG("Complete token signature has been verified successfully\n");

    ret = KG_SUCCESS;
    //kmt->kg_state = KG_STATE_COMPLETE;
exit:
    if (data_buffer != NULL) {
        TEE_Free(data_buffer);
        data_buffer = NULL;
    }
    if (pub_mod != NULL) {
        TEE_Free(pub_mod);
        pub_mod = NULL;
    }
    if (pub_exp != NULL) {
        TEE_Free(pub_exp);
        pub_exp = NULL;
    }
    return ret;
}

uint32_t KG_get_ta_status(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    KG_LOG("[TRACE] KG_get_ta_status start\n");
    uint32_t ret = KG_SUCCESS;
    kg_rpmb_info_t *info = NULL;
    kg_secure_data_t* secure_data = NULL;
    kg_rpmb_complete_token_t* complete_token = NULL;

    uint8_t* wrap_data = NULL;
    uint8_t* unwrap_data = NULL;
    uint32_t unwrap_data_len = KG_BUF_LEN;

    if (sendmsg->payload.cmd.data_len > KG_MAX_PAYLOAD_LEN) {
        KG_LOG("Received invalid input buffer when get status\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    ret = kg_init();
    if (ret != KG_SUCCESS) {
        KG_LOG("KG init failed\n");
        ret = KG_INIT_FAIL;
        goto exit;
    }

    ret = read_info_object(&info);
    if(KG_SUCCESS != ret) {
        KG_LOG("Failed to read info object\n");
        goto exit;
    }

    //kg_state always >= 0 and = also indicate error
    if ( info->kg_state >= KG_STATE_ERROR) {
        KG_LOG("KG state corrupted\n");
        respmsg->payload.resp.result = KG_STATE_ERROR;
        ret = KG_STATE_FAIL;
        goto exit;
    }

    if(info->kg_wrap_data_len > KG_SECURE_DATA_LEN){
        KG_LOG("Received invalid size for wrap_data\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    if (info->kg_state == KG_STATE_COMPLETE) {
#ifdef CONFIG_QSEE
        /* error: taking address of packed member 'kg_wrap_data_len' of class or structure 'kg_rpmb_info' may result in an unaligned pointer value [-Werror,-Waddress-of-packed-member]*/ 
        uint32_t temp_wrap_data_len = info->kg_wrap_data_len;
        if (KG_SUCCESS != (ret = read_wrap_data(&temp_wrap_data_len, &wrap_data))
            || NULL == wrap_data) {
            KG_LOG("failed to read wrap data\n");
            goto exit;
        }
        info->kg_wrap_data_len = temp_wrap_data_len;
#else
        if (KG_SUCCESS != (ret = read_wrap_data(&(info->kg_wrap_data_len), &wrap_data))
            || NULL == wrap_data) {
            KG_LOG("failed to read wrap data\n");
            goto exit;
        }
#endif
        unwrap_data = TEE_Malloc(unwrap_data_len, 0);
        if (NULL == unwrap_data) {
            KG_LOG("failed to alloc unwrap data\n");
            ret = KG_ALLOC_BUFFER_FAIL;
            goto exit;
        }

#ifdef CONFIG_QSEE
        if (TZ_API_OK != TZ_unwrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
            wrap_data, info->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;
        }
#else
        if (TZ_API_OK != TZ_unwrap_data_with_derived_key((uint8_t *)KG_NAME, strlen(KG_NAME), 
            wrap_data, info->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
            KG_LOG("failed to unwrap data\n")
            ret = KG_TZ_API_FAIL;
            goto exit;
        }
#endif
        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;
        }

        secure_data = (kg_secure_data_t*)unwrap_data;
        uint8_t* sn_hash = secure_data->kg_metadata.sn_hash;

        if (KG_SUCCESS != (ret = read_complete_token(&complete_token))
            || NULL == complete_token) {
            KG_LOG("failed to read complete token\n");
            goto exit;
        }

        if (KG_SUCCESS != (ret = kgbl_verify_completetoken(complete_token->token, sn_hash))) {
            KG_LOG("failed to verify complete token");
            respmsg->payload.resp.result = KG_STATE_ERROR;
            ret = KG_STATE_FAIL;
            goto exit;
        }
        
    }
    respmsg->payload.resp.result = info->kg_state;
    KG_LOG("KG status is %d\n", info->kg_state);
exit:
    if (NULL != complete_token) {
        TEE_MemFill(complete_token, 0, sizeof(kg_rpmb_complete_token_t));
        TEE_Free(complete_token);
        complete_token = NULL;
    }
    if (NULL != wrap_data) {
        TEE_Free(wrap_data);
        wrap_data = NULL;
    }
    if (NULL != unwrap_data) {
        TEE_MemFill(unwrap_data, 0, unwrap_data_len);
        TEE_Free(unwrap_data);
        unwrap_data = NULL;
    }
    if (NULL != info) {
        TEE_Free(info);
        info = NULL;
    }
    return ret;
}

uint32_t KG_checking(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {

    KG_LOG("KG checking\n");
    uint32_t ret = KG_SUCCESS;
    kg_rpmb_info_t* info = NULL;

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when checking\n");
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }
#endif

    if (KG_SUCCESS != (ret = read_info_object(&info))
        || NULL == info) {
        KG_LOG("failed to read info object\n");
        goto exit;
    }

    if (info->kg_state != KG_STATE_PREENROLL) {
        KG_LOG("KG state transition to checking is not allowed\n");
        respmsg->payload.resp.result = KG_STATE_ERROR;
        ret = KG_STATE_FAIL;
        goto exit;
    }

    info->kg_state = KG_STATE_CHECKING;

    if (KG_SUCCESS != (ret = write_info_object(info))) {
        KG_LOG("failed to write info object\n");
        goto exit;
    }
exit:
    if (info != NULL) {
        TEE_Free(info);
        info = NULL;
    }
    return ret;
}

uint32_t KG_get_kgid(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    KG_LOG("KG get kgid\n");

    uint32_t ret = KG_SUCCESS;
    kg_rpmb_info_t* info = NULL;

    uint8_t* unwrap_data = NULL;
    uint32_t unwrap_data_len = KG_BUF_LEN;
    uint8_t* wrap_data = NULL;

    if (sendmsg->payload.cmd.data_len > KG_MAX_PAYLOAD_LEN) {
        KG_LOG("Received invalid input buffer when get kgid\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when getting kgid\n");
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }
#endif

    if (KG_SUCCESS != (ret = read_info_object(&info))
        || NULL == info) {
        KG_LOG("failed to read info object\n");
        goto exit;
    }

    if(info->kg_wrap_data_len > KG_SECURE_DATA_LEN){
        KG_LOG("Received invalid size for wrap_data\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
/* error: taking address of packed member 'kg_wrap_data_len' of class or structure 'kg_rpmb_info' may result in an unaligned pointer value [-Werror,-Waddress-of-packed-member]*/ 
    uint32_t temp_wrap_data_len = info->kg_wrap_data_len;
    if (KG_SUCCESS != (ret = read_wrap_data(&temp_wrap_data_len, &wrap_data))
        || NULL == wrap_data) {
        KG_LOG("failed to read wrap data\n");
        goto exit;
    }
    info->kg_wrap_data_len = temp_wrap_data_len;
#else
    if (KG_SUCCESS != (ret = read_wrap_data(&(info->kg_wrap_data_len), &wrap_data))
        || NULL == wrap_data) {
        KG_LOG("failed to read wrap data\n");
        goto exit;
    }
#endif

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

#ifdef CONFIG_QSEE
    if (TZ_API_OK != TZ_unwrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        wrap_data, info->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
        KG_LOG("Failed to unwrap data\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#else
    if (TZ_API_OK != TZ_unwrap_data_with_derived_key((uint8_t *)KG_NAME, strlen(KG_NAME), 
        wrap_data, info->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
        KG_LOG("failed to unwrap data\n")
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#endif
    
    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;
    }

    kg_secure_data_t* secure_data = (kg_secure_data_t*)unwrap_data;
    
    if(KG_ID_LEN > KG_MAX_PAYLOAD_LEN){
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    respmsg->payload.resp.data_len = KG_ID_LEN;
    TEE_MemMove(respmsg->payload.resp.data_buf, secure_data->kg_metadata.kg_id, KG_ID_LEN);

exit:
    if (info != NULL) {
        TEE_Free(info);
        info = NULL;
    }
    if (wrap_data != NULL) {
        TEE_Free(wrap_data);
        wrap_data = 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_verify_registration_info(tz_verify_reg_info_payload_t *sendmsg, tz_verify_reg_info_payload_t *respmsg) {
    KG_LOG("[TRACE] KG_verify_registration_info start\n")

    uint32_t ret = KG_SUCCESS;
    uint32_t len_check = KG_IMEI_HASH_LEN + KG_SN_HASH_LEN + KG_ID_LEN + KG_TIMESTAMP_LEN;
    uint8_t *reg_info_buffer = NULL;
    uint32_t reg_info_buffer_len = KG_BUF_LEN;
    uint8_t *reg_info_sig_buffer = NULL;
    uint32_t reg_info_sig_buffer_len = KG_BUF_LEN;
    uint8_t *pub_mod = NULL;
    uint32_t pub_mod_len = KG_BUF_LEN;
    uint8_t *pub_exp = NULL;
    uint32_t pub_exp_len = KG_BUF_LEN;

    kg_rpmb_info_t* info = NULL;
    kg_secure_data_t* secure_data = NULL;
    kg_rpmb_complete_token_t* token = NULL;

    uint8_t* wrap_data = NULL;
    uint8_t *unwrap_data = NULL;
    uint32_t unwrap_data_len = KG_BUF_LEN;
    uint8_t* rewrap_data = NULL;
    uint32_t rewrap_data_len = KG_SECURE_DATA_LEN;

    if (sendmsg->payload.cmd.reg_info_len > KG_REG_INFO_LEN_MAX ||
        sendmsg->payload.cmd.reg_info_sign_len > KG_REG_INFO_SIGN_MAX) {
        KG_LOG("Received invalid input buffer when verifying reg info\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    if (KG_SUCCESS != (ret = kg_init())) {
        KG_LOG("KG init failed\n");
        ret = KG_INIT_FAIL;
        goto exit;
    }

    reg_info_buffer = TEE_Malloc(reg_info_buffer_len, 0);
    if (NULL == reg_info_buffer) {
        KG_LOG("Failed to alloc buffer for registration info\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    reg_info_sig_buffer = TEE_Malloc(reg_info_sig_buffer_len, 0);
    if (NULL == reg_info_sig_buffer) {
        KG_LOG("Failed to alloc buffer for registration info\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    pub_mod = TEE_Malloc(pub_mod_len, 0);
    if (NULL == pub_mod) {
        KG_LOG("Failed to alloc buffer for pub_mod\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    pub_exp = TEE_Malloc(pub_exp_len, 0);
    if (NULL == pub_exp) {
        KG_LOG("Failed to alloc buffer for pub_mod\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    // 1. reg info base64 decode and length check
    if (BASE64_OK != base64_decode(sendmsg->payload.cmd.reg_info,
            sendmsg->payload.cmd.reg_info_len, reg_info_buffer, &reg_info_buffer_len)) {
        KG_LOG("Failed to decode base64 encoded input reg info\n");
        ret = KG_BASE64_DECODE_FAIL;
        goto exit;
    }
    KG_LOG_DBG("reg_info len is %d\n", reg_info_buffer_len);
    if (len_check != reg_info_buffer_len) {
        KG_LOG("Recieved invalid buffer to verify registration info\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    KG_DUMP_DBG("reg_info_buffer: \n", reg_info_buffer, reg_info_buffer_len);

    // 2. reg sign base64 decode and length check
    if (BASE64_OK != base64_decode(sendmsg->payload.cmd.reg_info_sign,
            sendmsg->payload.cmd.reg_info_sign_len, reg_info_sig_buffer, &reg_info_sig_buffer_len)) {
        KG_LOG("Failed to decode base64 encoded input reg info signature\n");
        ret = KG_BASE64_DECODE_FAIL;
        goto exit;
    }
    KG_LOG_DBG("reg_info_sign len is %d\n", reg_info_sig_buffer_len);
    if (reg_info_sig_buffer_len != KG_SIG_LEN) {
        KG_LOG("Decoded reg info signature buffer length check failed\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    KG_DUMP_DBG("reg_info_sig_buffer: \n", reg_info_sig_buffer, reg_info_sig_buffer_len);

    bool verified_eu = false;
    if (KG_SUCCESS != extract_public_keybytes(policy_sign_cert_EU, pub_mod, &pub_mod_len, pub_exp, &pub_exp_len)) {
        KG_LOG_DBG("Failed to extract publick key bytes from policy signing cert\n");
        ret = KG_CRYPTO_PKEY_PARSE_FAIL;
        goto check_US;
    }
    KG_DUMP_DBG("dump pubkey n:\n", pub_mod, pub_mod_len);
    KG_DUMP_DBG("dump pubkey e:\n", pub_exp, pub_exp_len);

    if (TZ_API_OK != TZ_verify_CKM_SHA256_RSA_PKCS(pub_mod, pub_mod_len, pub_exp, pub_exp_len,
            reg_info_buffer, reg_info_buffer_len, reg_info_sig_buffer, KG_SIG_LEN, &verified_eu)) {
        KG_LOG_DBG("TZ API failed when verifying registration info\n");
        ret = KG_TZ_API_FAIL;
        goto check_US;
    }

check_US:
    TEE_MemFill(pub_mod, 0, KG_BUF_LEN);
    pub_mod_len = KG_BUF_LEN;
    TEE_MemFill(pub_exp, 0, KG_BUF_LEN);
    pub_exp_len = KG_BUF_LEN;

    if (KG_SUCCESS != extract_public_keybytes(policy_sign_cert_US, pub_mod, &pub_mod_len, pub_exp, &pub_exp_len)) {
        KG_LOG_DBG("Failed to extract publick key bytes from policy signing cert US\n");
        ret = KG_CRYPTO_PKEY_PARSE_FAIL;
        goto exit;
    }
    KG_DUMP_DBG("dump pubkey n:\n", pub_mod, pub_mod_len);
    KG_DUMP_DBG("dump pubkey e:\n", pub_exp, pub_exp_len);
    bool verified_us = false;
    if (TZ_API_OK != TZ_verify_CKM_SHA256_RSA_PKCS(pub_mod, pub_mod_len, pub_exp, pub_exp_len,
            reg_info_buffer, reg_info_buffer_len, reg_info_sig_buffer, KG_SIG_LEN, &verified_us)) {
        KG_LOG_DBG("TZ API failed when verifying registration info\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }

    if (false == verified_eu && false == verified_us) {
        KG_LOG_DBG("Failed to verify reg info signature\n");
        ret = KG_REGINFO_VERIFY_FAIL;
        goto exit;
    }
    KG_LOG("Reg info signature has been verified successfully\n");

    char *ptr;
    uint64_t server_time = strtoull((char *)(reg_info_buffer + KG_IMEI_HASH_LEN + KG_SN_HASH_LEN + KG_ID_LEN), &ptr, 10);

    //error check
    KG_LOG_DBG("server_time long is %llu\n", server_time);
    if (server_time >= 5000000000000) {
        if (KG_SUCCESS != (ret = clear_rpmb())) {
            KG_LOG("failed to clear rpmb\n");
            goto exit;
        }
        if (KG_SUCCESS != (ret = kg_init())) {
            KG_LOG("KG init failed\n");
            ret = KG_INIT_FAIL;
            goto exit;
        }
        // after rpmb reset, set server time to be server time - 5000000000000 so that policy checking on server time can pass
        server_time = server_time - 5000000000000;
    }

    if (KG_SUCCESS != (ret = read_info_object(&info))
        || NULL == info) {
        KG_LOG("failed to read info object\n");
        goto exit;
    }

    KG_LOG("wrap data len : %d\n", info->kg_wrap_data_len);
    unwrap_data = TEE_Malloc(unwrap_data_len, 0);
    if (NULL == unwrap_data) {
        KG_LOG("failed to alloc unwrap data\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    if(info->kg_wrap_data_len > KG_SECURE_DATA_LEN){
        KG_LOG("Received invalid size for wrap_data\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    if (0 != info->kg_wrap_data_len) {
#ifdef CONFIG_QSEE
        /* error: taking address of packed member 'kg_wrap_data_len' of class or structure 'kg_rpmb_info' may result in an unaligned pointer value [-Werror,-Waddress-of-packed-member]*/        
        uint32_t temp_wrap_data_len = info->kg_wrap_data_len;
        if (KG_SUCCESS != (ret = read_wrap_data(&temp_wrap_data_len, &wrap_data))
            || NULL == wrap_data) {
            KG_LOG("failed to read wrap data\n");
            goto exit;
        }
        info->kg_wrap_data_len = temp_wrap_data_len;
#else        
        if (KG_SUCCESS != (ret = read_wrap_data(&(info->kg_wrap_data_len), &wrap_data))
            || NULL == wrap_data) {
            KG_LOG("failed to read wrap data\n");
            goto exit;
        }
#endif
#ifdef CONFIG_QSEE
        if (TZ_API_OK != TZ_unwrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
            wrap_data, info->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;
        }
#else        
        if (TZ_API_OK != TZ_unwrap_data_with_derived_key((uint8_t *)KG_NAME, strlen(KG_NAME), 
            wrap_data, info->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
            KG_LOG("failed to unwrap data\n")
            ret = KG_TZ_API_FAIL;
            goto exit;
        }
#endif
    }
    secure_data = (kg_secure_data_t *)unwrap_data;

    secure_data->kg_metadata.svr_timestamp = server_time;
    if (verified_eu == true) {
        secure_data->kg_metadata.reg_info = REGION_EU;
    } else if (verified_us == true) {
        secure_data->kg_metadata.reg_info = REGION_US;
    }

    TEE_MemMove(secure_data->kg_metadata.imei_hash, reg_info_buffer, KG_IMEI_HASH_LEN);
    TEE_MemMove(secure_data->kg_metadata.sn_hash, reg_info_buffer + KG_IMEI_HASH_LEN, KG_SN_HASH_LEN);
    TEE_MemMove(secure_data->kg_metadata.kg_id, reg_info_buffer + KG_IMEI_HASH_LEN + KG_SN_HASH_LEN, KG_ID_LEN);
    KG_LOG_DBG("KG ID : %s", (secure_data->kg_metadata.kg_id));

    // clear complete token for new enrollment
    token = TEE_Malloc(sizeof(kg_rpmb_complete_token_t), 0);
    if (NULL == token) {
        KG_LOG("failed to alloc token\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

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

#ifdef CONFIG_QSEE
/* error: incompatible pointer types passing 'kg_secure_data_t *' (aka 'struct kg_secure_data *') to parameter of type 'uint8_t *' (aka 'unsigned char *') [-Werror,-Wincompatible-pointer-types]*/    
    if (TZ_API_OK != TZ_wrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        (uint8_t *)secure_data, sizeof(kg_secure_data_t), rewrap_data, &rewrap_data_len)) {
        KG_LOG("Failed to wrap kg secure data structure in QC\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#else
    if (TZ_API_OK != TZ_wrap_data_with_derived_key((uint8_t *)KG_NAME, strlen(KG_NAME), 
        (uint8_t *)secure_data, sizeof(kg_secure_data_t), rewrap_data, &rewrap_data_len)) {
        KG_LOG("Failed to wrap kg secure data structure in QC\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#endif

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

    info->kg_wrap_data_len = rewrap_data_len;
    if (KG_SUCCESS != (ret = write_wrap_data(info->kg_wrap_data_len, rewrap_data))) {
        KG_LOG("failed to write wrap data\n");
        goto exit;
    }

    // save token
    if (KG_SUCCESS != (ret = write_complete_token(token))) {
        KG_LOG("failed to write complete token\n");
        goto exit;
    }

    if (KG_SUCCESS != (ret = write_info_object(info))) {
        KG_LOG("failed to write info object\n");
        goto exit;
    }

    KG_LOG("[TRACE] KG_verify_registration_info make response\n")
    respmsg->payload.resq.kg_id_len = KG_ID_LEN;
    TEE_MemMove(respmsg->payload.resq.kg_id, (void *)secure_data->kg_metadata.kg_id, KG_ID_LEN);

exit:
    if (reg_info_buffer != NULL) {
        TEE_Free(reg_info_buffer);
        reg_info_buffer = NULL;
    }
    if (reg_info_sig_buffer != NULL) {
        TEE_Free(reg_info_sig_buffer);
        reg_info_sig_buffer = NULL;
    }
    if (pub_mod != NULL) {
        TEE_Free(pub_mod);
        pub_mod = NULL;
    }
    if (pub_exp != NULL) {
        TEE_Free(pub_exp);
        pub_exp = NULL;
    }
    if (info != NULL) {
        TEE_Free(info);
        info = NULL;
    }
    if (wrap_data != NULL) {
        TEE_Free(wrap_data);
        wrap_data = NULL;
    }
    if (unwrap_data != NULL) {
        TEE_MemFill(unwrap_data, 0, unwrap_data_len);
        TEE_Free(unwrap_data);
        unwrap_data = NULL;
    }
    if (rewrap_data != NULL) {
        TEE_Free(rewrap_data);
        rewrap_data = NULL; 
    }
    if (token != NULL) {
        TEE_Free(token);
        token = NULL;
    }
    KG_LOG("[TRACE] KG_verify_registration_info end\n")
    return ret;
}

uint32_t KG_verify_complete_token(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    KG_LOG("KG verify complete token\n");
    uint32_t ret = KG_SUCCESS;
    uint32_t len_check = KG_SIG_LEN + KG_HASH_LEN + KG_COMPLETE_MSG_LEN;
    uint8_t *complete_buffer = NULL;
    uint32_t complete_buffer_len = KG_BUF_LEN;
    uint8_t *data_buffer = NULL;
    uint32_t data_buffer_len = KG_BUF_LEN;
    uint8_t *pub_mod = NULL;
    uint32_t pub_mod_len = KG_BUF_LEN;
    uint8_t *pub_exp = NULL;
    uint32_t pub_exp_len = KG_BUF_LEN;
    uint8_t ta_nonce_buffer[KG_NONCE_LEN] = {0x0};
    uint8_t *ta_nonce_hash_buffer = NULL;
    uint32_t ta_nonce_hash_len = KG_HASH_LEN;
    uint8_t *svr_nonce_hash_buffer = NULL;
    const char *complete_msg = "Complete";

    kg_rpmb_info_t* info = NULL;
    uint8_t* unwrap_data = NULL;
    uint32_t unwrap_data_len = KG_BUF_LEN;
    uint8_t* wrap_data = NULL;
    kg_rpmb_complete_token_t* token = NULL;

    if (sendmsg->payload.cmd.data_len > KG_MAX_PAYLOAD_LEN) {
        KG_LOG("Received invalid input buffer when verifying complete token\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    
    complete_buffer = TEE_Malloc(complete_buffer_len, 0);
    if (NULL == complete_buffer) {
        KG_LOG("Failed to alloc buffer for complete token\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    data_buffer = TEE_Malloc(data_buffer_len, 0);
    if (NULL == data_buffer) {
        KG_LOG("Failed to alloc buffer for data buffer\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    pub_mod = TEE_Malloc(pub_mod_len, 0);
    if (NULL == pub_mod) {
        KG_LOG("Failed to alloc buffer for pub_mod\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    pub_exp = TEE_Malloc(pub_exp_len, 0);
    if (NULL == pub_exp) {
        KG_LOG("Failed to alloc buffer for pub_mod\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    ta_nonce_hash_buffer = TEE_Malloc(ta_nonce_hash_len, 0);
    if (NULL == ta_nonce_hash_buffer) {
        KG_LOG("Failed to alloc buffer for ta_nonce_hash_buffer\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    svr_nonce_hash_buffer = TEE_Malloc(KG_HASH_LEN, 0);
    if (NULL == svr_nonce_hash_buffer) {
        KG_LOG("Failed to alloc buffer for ta_nonce_hash_buffer\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    token = TEE_Malloc(sizeof(kg_rpmb_complete_token_t), 0);
    if (NULL == token) {
        KG_LOG("failed to alloc token\n");
        ret = KG_ALLOC_BUFFER_FAIL;
        goto exit;
    }

    if (BASE64_OK != base64_decode(sendmsg->payload.cmd.data_buf,
            sendmsg->payload.cmd.data_len, complete_buffer, &complete_buffer_len)) {
        KG_LOG("Failed to decode base64 encoded input completetoken buffer\n");
        ret = KG_BASE64_DECODE_FAIL;
        goto exit;
    }
    if (complete_buffer_len != len_check) {
        KG_LOG_DBG("Decoded completetoken buffer length check failed %d %d\n",complete_buffer_len,len_check);
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    KG_LOG_DBG("complete token is %s\n", complete_buffer);

#ifdef CONFIG_QSEE
    if (KG_SUCCESS != kg_rpmb_init()) {
        KG_LOG("RPMB is not available when getting kgid\n");
        ret = KG_RPMB_UNAVAILABLE;
        goto exit;
    }
#endif

    if (KG_SUCCESS != (ret = read_info_object(&info))
        || NULL == info) {
        KG_LOG("failed to read info object\n");
        goto exit;
    }

    if(info->kg_wrap_data_len > KG_SECURE_DATA_LEN){
        KG_LOG("Received invalid size for wrap_data\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

#ifdef CONFIG_QSEE
/* error: taking address of packed member 'kg_wrap_data_len' of class or structure 'kg_rpmb_info' may result in an unaligned pointer value [-Werror,-Waddress-of-packed-member]*/ 
    uint32_t temp_wrap_data_len = info->kg_wrap_data_len;
    if (KG_SUCCESS != (ret = read_wrap_data(&temp_wrap_data_len, &wrap_data))
        || NULL == wrap_data) {
        KG_LOG("failed to read wrap data\n");
        goto exit;
    }
    info->kg_wrap_data_len = temp_wrap_data_len;
#else
    if (KG_SUCCESS != (ret = read_wrap_data(&(info->kg_wrap_data_len), &wrap_data))
        || NULL == wrap_data) {
        KG_LOG("failed to read wrap data\n");
        goto exit;
    }
#endif

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

#ifdef CONFIG_QSEE
    if (TZ_API_OK != TZ_unwrap_persist_data((uint8_t *)KG_NAME, strlen(KG_NAME), 
        wrap_data, info->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
        KG_LOG("Failed to unwrap data\n");
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#else
    if (TZ_API_OK != TZ_unwrap_data_with_derived_key((uint8_t *)KG_NAME, strlen(KG_NAME), 
        wrap_data, info->kg_wrap_data_len, unwrap_data, &unwrap_data_len)) {
        KG_LOG("failed to unwrap data\n")
        ret = KG_TZ_API_FAIL;
        goto exit;
    }
#endif

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

    kg_secure_data_t* secure_data = (kg_secure_data_t*)unwrap_data;

    uint32_t idx = 0;
    bool check_flag = true;
    while (idx < KG_COMPLETE_MSG_LEN) {
        if ((complete_buffer + KG_SIG_LEN + KG_HASH_LEN)[idx] != complete_msg[idx]) {
            check_flag = false;
        }
        idx++;
    }

    if (false == check_flag) {
        KG_LOG_DBG("Complete msg magic check failed\n");
        ret = KG_COMPLETE_MAGIC_FAIL;
        goto exit;
    }

    if (KG_COMPLETE_MSG_LEN + KG_SN_HASH_LEN > complete_buffer_len) {
        KG_LOG_DBG("Complete buffer length check failed\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    TEE_MemFill(data_buffer, 0x0, data_buffer_len);
    data_buffer_len = 0;
    TEE_MemMove(data_buffer, (void *)complete_msg, KG_COMPLETE_MSG_LEN);
    data_buffer_len = KG_COMPLETE_MSG_LEN;
    TEE_MemMove(data_buffer + data_buffer_len, secure_data->kg_metadata.sn_hash, KG_SN_HASH_LEN);
    data_buffer_len += KG_SN_HASH_LEN;

    // TODO, verify with BL cert and test case from server

    if (secure_data->kg_metadata.reg_info == REGION_EU) {
        if (KG_SUCCESS != extract_public_keybytes(bl_sign_cert_EU, pub_mod, &pub_mod_len, pub_exp, &pub_exp_len)) {
            KG_LOG_DBG("Failed to extract publick key bytes from pem certificate\n");
            ret = KG_CRYPTO_PKEY_PARSE_FAIL;
            goto exit;
        }
    } else if (secure_data->kg_metadata.reg_info == REGION_US) {
        if (KG_SUCCESS != extract_public_keybytes(bl_sign_cert_US, pub_mod, &pub_mod_len, pub_exp, &pub_exp_len)) {
            KG_LOG_DBG("Failed to extract publick key bytes from pem certificate\n");
            ret = KG_CRYPTO_PKEY_PARSE_FAIL;
            goto exit;
        }
    }

    KG_DUMP_DBG("dump pubkey n:\n", pub_mod, pub_mod_len);
    KG_DUMP_DBG("dump pubkey e:\n", pub_exp, pub_exp_len);
    bool verified = false;
    if (TZ_API_OK != TZ_verify_CKM_SHA256_RSA_PKCS(pub_mod, pub_mod_len, pub_exp, pub_exp_len,
            data_buffer, data_buffer_len, complete_buffer, KG_SIG_LEN, &verified)) {
        KG_LOG_DBG("Failed to verify signed completeToken\n");
        ret = KG_COMPLETE_VERIFY_FAIL;
        goto exit;
    }
    if (false == verified) {
        KG_LOG_DBG("Failed to verify complete token signature\n");
        ret = KG_COMPLETE_VERIFY_FAIL;
        goto exit;
    }
    KG_LOG("Complete token signature has been verified successfully\n");

    TEE_MemMove(svr_nonce_hash_buffer, complete_buffer + KG_SIG_LEN, KG_HASH_LEN);
    
    if (TZ_API_OK != TZ_digest_SHA256(secure_data->kg_metadata.ta_nonce, KG_NONCE_LEN, ta_nonce_hash_buffer, &ta_nonce_hash_len)) {
        ret = KG_TZ_API_FAIL;
        goto exit;
    }

    if (ta_nonce_hash_len != KG_HASH_LEN) {
        ret = KG_COMPUTE_HASH_FAIL;
        goto exit;
    }

    idx = 0;
    check_flag = true;
    while (idx < KG_HASH_LEN) {
        if (svr_nonce_hash_buffer[idx] != ta_nonce_hash_buffer[idx]) {
            check_flag = false;
        }
        idx++;
    }

    if (false == check_flag) {
        KG_LOG_DBG("Complete msg ta nonce check failed\n");
        ret = KG_COMPLETE_MAGIC_FAIL;
        goto exit;
    }

    KG_DUMP_DBG("Final rpmb state: \n", (uint8_t *)&(secure_data->kg_metadata), sizeof(kg_metadata_t));

    // save token
    if(KG_TOKEN_LEN > sizeof(kg_rpmb_complete_token_t)){
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }
    TEE_MemMove(token, complete_buffer, KG_TOKEN_LEN);
    if (KG_SUCCESS != (ret = write_complete_token(token))) {
        KG_LOG("failed to write complete token\n");
        goto exit;
    }

    info->kg_state = KG_STATE_COMPLETE;
    if (KG_SUCCESS != (ret = write_info_object(info))) {
        KG_LOG("failed to write info object\n");
        goto exit;
    }
exit:
    if (info != NULL) {
        TEE_Free(info);
        info = NULL;
    }
    if (wrap_data != NULL) {
        TEE_Free(wrap_data);
        wrap_data = NULL;
    }
    if (unwrap_data != NULL) {
        TEE_MemFill(unwrap_data, 0, unwrap_data_len);
        TEE_Free(unwrap_data);
        unwrap_data = NULL;
    }
    if (token != NULL) {
        TEE_Free(token);
        token = NULL;
    }
    if (complete_buffer != NULL) {
        TEE_Free(complete_buffer);
        complete_buffer = NULL;
    }
    if (data_buffer != NULL) {
        TEE_Free(data_buffer);
        data_buffer = NULL;
    }
    if (pub_mod != NULL) {
        TEE_Free(pub_mod);
        pub_mod = NULL;
    }
    if (pub_exp != NULL) {
        TEE_Free(pub_exp);
        pub_exp = NULL;
    }
    if (ta_nonce_hash_buffer != NULL) {
        TEE_Free(ta_nonce_hash_buffer);
        ta_nonce_hash_buffer = NULL;
    }
    if (svr_nonce_hash_buffer != NULL) {
        TEE_Free(svr_nonce_hash_buffer);
        svr_nonce_hash_buffer = NULL;
    }
    return ret;
}

uint32_t KG_reset_rpmb(tz_common_payload_t *sendmsg, tz_common_payload_t *respmsg) {
    KG_LOG("KG reset rpmb\n");
    uint32_t ret = KG_SUCCESS;
    kg_rpmb_info_t *info = NULL;

    if (sendmsg->payload.cmd.data_len > KG_MAX_PAYLOAD_LEN) {
        KG_LOG("Received invalid input buffer when reset rpmb\n");
        ret = KG_BUFFER_SIZE_FAIL;
        goto exit;
    }

    ret = kg_rpmb_init();
    if (ret != KG_SUCCESS) {
        KG_LOG("KG init failed\n");
        ret = KG_INIT_FAIL;
        goto exit;
    }

    ret = read_info_object(&info);
    if(KG_SUCCESS != ret) {
        KG_LOG("Failed to read info object\n");
        goto exit;
    }

    if (info->kg_state == KG_STATE_COMPLETE || info->kg_state == KG_STATE_CHECKING) {
        if (KG_SUCCESS != (ret = clear_rpmb())) {
            KG_LOG("failed to clear rpmb\n");
            goto exit;
        }
    } else {
        KG_LOG("Not allowed reset in current state !!!! \n");
    }

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