/*
 * board/mt53xx_com/mt53xx_com_rpmb.c
 *
 * MT53xx common board function
 *
 * Copyright (c) 2010-2012 MediaTek Inc.
 * $Author: brianpc.huang $
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for
 * more details.
 * http://www.gnu.org/licenses/old-licenses/gpl-2.0.html
 *
 */

/* ============================================================================
    Includes
============================================================================ */
#include "x_bim.h"
#include "x_printf.h"
#include "x_ldr_env.h"
#include "mt53xx_com_rpmb.h"

/* ============================================================================
    Constant definitions
============================================================================ */
#define RPMB_KEY_SIZE (32)
#define SIGNED_DATA_SIZE (284)

#define TIMER_ID 2
#define TIMER_LOW_REG (REG_RW_TIMER0_LOW + (TIMER_ID * 8))

#ifndef UINT_MAX
#define UINT_MAX 4294967295U
#endif

extern INT32 SHA256(const UINT8 *message_array, UINT32 length,
                    UINT8 Message_Digest[SHA256_HASH_SIZE]);
extern int mtk_hmac_sha256(uint8_t *key, uint32_t key_size, uint8_t *msg, uint32_t msg_size,
                           uint8_t result[SHA256_HASH_SIZE]);
extern int emmc_set_blkcnt(unsigned int count, unsigned int write_flag);
extern int emmc_write_msdc(u64 offset, void *addr, size_t cnt);
extern int emmc_read_msdc(u64 offset, void *addr, size_t cnt);
extern uint32_t readMsTicks(void);
/* ============================================================================
    Global Variable
============================================================================ */
static uint8_t savedNonce[RPMB_NONCE_SIZE];
static uint32_t savedWriteCnt;
static uint8_t macKey[RPMB_KEY_SIZE] = {0};
static int init_read_counter = 0;

/* ============================================================================
    Internal Functions
============================================================================ */
static void _ReverseEndian(void *data, size_t size) {
    uint32_t i;
    char tmp;
    char *swp = (char *)data;
    for (i = 0; i < (size / 2); ++i) {
        tmp = swp[i];
        swp[i] = swp[size - 1 - i];
        swp[size - 1 - i] = tmp;
    }
}

static void _rpmbLoadMacKey(void) {
    static int init = 0;
    if (init == 0) {
        LDR_ENV_T *prLdrEnv = (LDR_ENV_T *)CC_LDR_ENV_OFFSET;
        memcpy(macKey, prLdrEnv->RPMB_KEY, RPMB_KEY_SIZE);
        _ReverseEndian(macKey, RPMB_KEY_SIZE);
        init = 1;
    }
}

static inline void generateNonce(void) {
    uint32_t u4Seed = BIM_READ32(TIMER_LOW_REG);
    uint8_t hash[SHA256_HASH_SIZE];
    SHA256((UINT8 *)&u4Seed, sizeof(uint32_t), hash);
    memcpy(savedNonce, hash, RPMB_NONCE_SIZE);
}

static inline uint32_t verifyNonoce(const uint8_t *nonce) {
    return memcmp(&savedNonce[0], nonce, RPMB_NONCE_SIZE) == 0;
}

static inline int sign(const uint8_t *data, uint8_t *signature) {
    uint8_t tmp[SIGNED_DATA_SIZE];

    memcpy(tmp, data, SIGNED_DATA_SIZE);
    _ReverseEndian(tmp, SIGNED_DATA_SIZE);

    _rpmbLoadMacKey();

    if (mtk_hmac_sha256(macKey, RPMB_KEY_SIZE, tmp, SIGNED_DATA_SIZE, signature) == 0) {
        _ReverseEndian(signature, RPMB_KEY_SIZE);
        return 1;
    } else {
        return 0;
    }
}

static inline int verify(const void *data, const void *check) {
    uint8_t signature[SHA256_HASH_SIZE];

    sign((const uint8_t *)data, signature);
    return memcmp(check, signature, sizeof(signature)) == 0;
}

static INT32 ubootRpmbReadCounter(void) {
    RPMB_DATAFRAM dFrame;
    RPMB_DATAFRAM dFrameResp;
    UINT64 offset = 0;

    generateNonce();

    if (emmc_set_blkcnt(1, 0) != 0) {
        printf("emmc_set_blkcnt failed.\n");
        return -1;
    }

    /* Send Read Counter Request */
    memset(&dFrame, 0, sizeof(RPMB_DATAFRAM));
    dFrame.type = REQ_GET_CNT;
    memcpy(dFrame.nonce, savedNonce, RPMB_NONCE_SIZE);

    _ReverseEndian(&dFrame, sizeof(RPMB_DATAFRAM));
    if (emmc_write_msdc(offset, (void *)(&dFrame), sizeof(RPMB_DATAFRAM)) != 0) {
        printf("emmc_write_msdc failed.\n");
        return -2;
    }

    if (emmc_set_blkcnt(1, 0) != 0) {
        printf("emmc_set_blkcnt failed.\n");
        return -3;
    }

    /* Read Result */
    memset(&dFrameResp, 0, sizeof(RPMB_DATAFRAM));
    if (emmc_read_msdc(offset, (void *)(&dFrameResp), sizeof(RPMB_DATAFRAM)) != 0) {
        printf("emmc_read_msdc failed.\n");
        return -4;
    }
    _ReverseEndian(&dFrameResp, sizeof(RPMB_DATAFRAM));

    if (dFrameResp.result != RESULT_OK) {
        printf("dFrameResp.result(0x%x) != RESULT_OK\n", dFrameResp.result);
        return -5;
    }

    if (dFrameResp.type != RESP_GET_CNT) {
        printf("dFrameResp.type(0x%x) != RESP_GET_CNT(0x200)\n", dFrameResp.type);
        return -6;
    }

    if (verifyNonoce(dFrameResp.nonce) == 0) {
        printf("verifyNonoce fail\n");
        return -7;
    }

    if (verify(&dFrameResp, dFrameResp.mac) == 0) {
        printf("verify dFrameResp.mac fail\n");
        return -8;
    }

    savedWriteCnt = dFrameResp.writeCnt;
    init_read_counter = 1;

    return (dFrameResp.result & 0xF);
}

/* ============================================================================
    Exported Entry Points
============================================================================ */
INT32 ubootRpmbReadData(uint8_t *data, RPMB_ADDR addr) {
    RPMB_DATAFRAM dFrame;
    RPMB_DATAFRAM dFrameResp;
    UINT64 offset = 0;

    generateNonce();

    if (access_emmc_rpmb(1) != 0) {
        printf("access_emmc_rpmb failed.\n");
        return -1;
    }

    if (init_read_counter == 0) {
        if (ubootRpmbReadCounter() != 0) {
            printf("ubootRpmbReadCounter failed.\n");
            return -2;
        }
    }

    if (emmc_set_blkcnt(1, 0) != 0) {
        printf("emmc_set_blkcnt failed.\n");
        return -3;
    }

    /* Send Read Counter Request */
    memset(&dFrame, 0, sizeof(RPMB_DATAFRAM));
    dFrame.type = REQ_RD_DATA;
    dFrame.addr = (unsigned short)addr;
    memcpy(dFrame.nonce, savedNonce, RPMB_NONCE_SIZE);

    _ReverseEndian(&dFrame, sizeof(RPMB_DATAFRAM));
    if (emmc_write_msdc(offset, (void *)(&dFrame), sizeof(RPMB_DATAFRAM)) != 0) {
        printf("emmc_write_msdc failed.\n");
        return -4;
    }

    if (emmc_set_blkcnt(1, 0) != 0) {
        printf("emmc_set_blkcnt failed.\n");
        return -5;
    }

    /* Read Result */
    memset(&dFrameResp, 0, sizeof(RPMB_DATAFRAM));
    if (emmc_read_msdc(offset, (void *)(&dFrameResp), sizeof(RPMB_DATAFRAM)) != 0) {
        printf("emmc_read_msdc failed.\n");
        return -6;
    }
    _ReverseEndian(&dFrameResp, sizeof(RPMB_DATAFRAM));

    if (access_emmc_rpmb(0) != 0) {
        printf("access_emmc_rpmb failed.\n");
        return -7;
    }

    if (dFrameResp.result != RESULT_OK) {
        printf("dFrameResp.result(0x%x) != RESULT_OK\n", dFrameResp.result);
        return -8;
    }

    if (dFrameResp.type != RESP_RD_DATA) {
        printf("dFrameResp.type(0x%x) != RESP_RD_DATA(0x400)\n", dFrameResp.type);
        return -9;
    }

    if (verifyNonoce(dFrameResp.nonce) == 0) {
        printf("verifyNonoce fail\n");
        return -10;
    }

    if (verify(&dFrameResp, dFrameResp.mac) == 0) {
        printf("verify dFrameResp->mac fail\n");
        return -11;
    }

    memcpy(data, dFrameResp.data, RPMB_DATA_SIZE);
    init_read_counter = 1;

    return (dFrameResp.result & 0xF);
}

INT32 ubootRpmbWriteData(const uint8_t *data, RPMB_ADDR addr, unsigned short blkCnt) {
    RPMB_DATAFRAM dFrame;
    RPMB_DATAFRAM dFrameResp;
    uint8_t signature[SHA256_HASH_SIZE];
    UINT64 offset = 0;

    if (access_emmc_rpmb(1) != 0) {
        printf("access_emmc_rpmb failed.\n");
        return -1;
    }

    if (init_read_counter == 0) {
        if (ubootRpmbReadCounter() != 0) {
            printf("ubootRpmbReadCounter failed.\n");
            return -2;
        }
    }

    if (emmc_set_blkcnt(1, (1 << 31)) != 0) {
        printf("emmc_set_blkcnt failed.\n");
        return -3;
    }

    /* Send Read Counter Request */
    memset(&dFrame, 0, sizeof(RPMB_DATAFRAM));
    dFrame.type = REQ_WR_DATA;
    dFrame.blkCnt = blkCnt;
    dFrame.addr = (unsigned short)addr;
    dFrame.writeCnt = savedWriteCnt;
    memcpy(dFrame.data, data, RPMB_DATA_SIZE);

    sign((const uint8_t *)&dFrame, signature);
    memcpy(dFrame.mac, signature, sizeof(dFrame.mac));

    _ReverseEndian(&dFrame, sizeof(RPMB_DATAFRAM));
    if (emmc_write_msdc(offset, (void *)(&dFrame), sizeof(RPMB_DATAFRAM)) != 0) {
        printf("emmc_write_msdc failed.\n");
        return -4;
    }

    if (emmc_set_blkcnt(1, 0) != 0) {
        printf("emmc_set_blkcnt failed.\n");
        return -5;
    }

    memset(&dFrame, 0, sizeof(RPMB_DATAFRAM));
    dFrame.type = REQ_RD_RESULT;

    _ReverseEndian(&dFrame, sizeof(RPMB_DATAFRAM));
    if (emmc_write_msdc(offset, (void *)(&dFrame), sizeof(RPMB_DATAFRAM)) != 0) {
        printf("emmc_write_msdc failed.\n");
        return -6;
    }

    if (emmc_set_blkcnt(1, 0) != 0) {
        printf("emmc_set_blkcnt failed.\n");
        return -7;
    }

    /* Read Result */
    memset(&dFrameResp, 0, sizeof(RPMB_DATAFRAM));
    if (emmc_read_msdc(offset, (void *)(&dFrameResp), sizeof(RPMB_DATAFRAM)) != 0) {
        printf("emmc_read_msdc failed.\n");
        return -8;
    }
    _ReverseEndian(&dFrameResp, sizeof(RPMB_DATAFRAM));

    if (access_emmc_rpmb(0) != 0) {
        printf("access_emmc_rpmb failed.\n");
        return -9;
    }

    if (dFrameResp.result != RESULT_OK) {
        printf("dFrameResp.result(0x%x) != RESULT_OK\n", dFrameResp.result);
        return -10;
    }

    if (dFrameResp.type != RESP_WR_DATA) {
        printf("dFrameResp.type(0x%x) != RESP_WR_DATA(0x300)\n", dFrameResp.type);
        return -11;
    }

    if (verify(&dFrameResp, dFrameResp.mac) == 0) {
        printf("verify dFrameResp->mac fail\n");
        return -12;
    }

    if (dFrameResp.writeCnt != UINT_MAX) {
        if (dFrameResp.writeCnt != (savedWriteCnt + 1)) {
            printf("verify counter fail\n");
            return -13;
        }
    }
    savedWriteCnt = dFrameResp.writeCnt;

    return (dFrameResp.result & 0xF);
}

