cmd/amlogic/partition_encryption.c

// SPDX-License-Identifier: (GPL-2.0+ OR MIT)
/*
 * Copyright (c) 2019 Amlogic, Inc. All rights reserved.
 */

#include <config.h>
#include <common.h>
#include <command.h>
#include <fdt_support.h>
#include <amlogic/aml_crypto.h>
#include <amlogic/aml_mkl.h>
#include <amlogic/store_wrapper.h>
#include <malloc.h>
#include <u-boot/md5.h>
#include <uboot_aes.h>
#include <amlogic/partition_encryption.h>
#include <linux/libfdt.h>

#define DEBUG           (0)
#define ERROR_REPORT    (1)

//Maximum number of enc partitions which kernel driver supported
#define MAX_ENC_PARTS (4)

typedef struct {
    char part_name[32];
    u8 wrapped[AES_KEY_SIZE_128];
    u8 in_used;
} ENC_PART_ENT_T;

ENC_PART_ENT_T enc_parts[MAX_ENC_PARTS];

int32_t find_free_enc_parts_ent(void)
{
    int32_t i = 0;
    for (i = 0; i < MAX_ENC_PARTS; i++)
        if (enc_parts[i].in_used == 0)
            return i;
    return -1;
}

#if DEBUG
void dump_bufs(const char *name, const uint8_t *buf, int32_t len)
{
    char line_buf[128] = {0x0};
    char *line_buf_ptr = line_buf;
    uint32_t i = 0;
    const uint8_t *ptr = buf;
    const int32_t sz = len;

    printf("\n========= Start dump %s(%d)\n", name, sz);
    while (len - 16 >= 0) {
        printf("%02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x %02x\n",
                ptr[0], ptr[i+1], ptr[i+2], ptr[i+3], ptr[4], ptr[i+5], ptr[i+6], ptr[i+7],
                ptr[i+8], ptr[i+9], ptr[i+10], ptr[i+11], ptr[i+12], ptr[i+13], ptr[i+14], ptr[i+15]);
        ptr += 16;
        len -= 16;
    }

    if (len * 3 <= sizeof(line_buf)) {
        while (len) {
            snprintf(line_buf_ptr, sizeof(line_buf), "%02x ", *ptr);
            line_buf_ptr += 3; /*"xx "*/
            len--;
        }
        if (strlen(line_buf))
            printf("%s\n", line_buf);
        printf("========= End dump %s(%d)\n", name, sz);
    } else {
        printf("remaining length exceeds linebuf length\n");
    }
}
#endif

void dump_enc_parts(void)
{
    int32_t i, j;
    for (i = 0; i < MAX_ENC_PARTS; i++) {
        printf("%d: %s(%s)", i, enc_parts[i].part_name, enc_parts[i].in_used? "In used": "Unused");
        printf("\n    wrapped: ");
        for (j = 0; j < sizeof(enc_parts[i].wrapped); j++)
            printf("%02x", enc_parts[i].wrapped[j]);
        printf("\n");
    }
}

static int hex2bin(char *hex, void *bin, size_t binlen)
{
    int i, c, n1, n2, hexlen, k;

    hexlen = strnlen(hex, 64);
    k = 0;
    n1 = -1;
    n2 = -1;
    for (i = 0; i < hexlen; i++) {
        n2 = n1;
        c = hex[i];
        if (c >= '0' && c <= '9') {
            n1 = c - '0';
        } else if (c >= 'a' && c <= 'f') {
            n1 = c - 'a' + 10;
        } else if (c >= 'A' && c <= 'F') {
            n1 = c - 'A' + 10;
        } else if (c == ' ') {
            n1 = -1;
            continue;
        } else {
            return -1;
        }

        if (n1 >= 0 && n2 >= 0) {
            if (k < binlen) {
                // k is index and should be 0 <= k < binlen
                ((u8 *)bin)[k] = (n2 << 4) | n1;
                n1 = -1;
                k++;
            } else {
#if ERROR_REPORT
                printf("Output overflow(%d >= %lu)\n", k, binlen);
#endif
                goto out;
            }
        }
    }
out:
    return k;
}

static int unwrap_key(u8 *in, u32 in_sz, u8 *out, u32 *out_sz)
{
    int ret = 0;
    struct amlkl_params kl_param;
    uint32_t dgpk_key_slot = 31; //only test slot
    uint8_t dst[AES_KEY_SIZE_128] = {0};

    /* dgpk_func_0_bs_0_tsep_0.test.vector.txt */
    uint8_t dgpk_ek3[16] = {
        0xa6, 0x20, 0x65, 0x7f, 0xa9, 0x6b, 0x19, 0x66,
        0xa1, 0x03, 0x47, 0xb4, 0x23, 0xde, 0x20, 0x25
    };

    uint8_t dgpk_ek2[16] = {
        0xa2, 0xe3, 0x7a, 0xe8, 0x6e, 0xce, 0x64, 0x78,
        0x69, 0x18, 0x51, 0x0c, 0x57, 0xec, 0x68, 0xcf
    };

    uint8_t dgpk_ek1[16] = {
        0xfa, 0x29, 0x3b, 0x43, 0x25, 0x7c, 0x96, 0xbd,
        0x55, 0x2b, 0x60, 0xd5, 0x95, 0x60, 0xcf, 0xcc
    };
    if (!in || !out) {
#if ERROR_REPORT
        printf("Invalid arguments.\n");
#endif
        return -1;
    }
    if (!out_sz || *out_sz < AES_KEY_SIZE_128) {
#if ERROR_REPORT
        printf("output size too small\n");
#endif
        return -1;
    }

    memset(&kl_param, 0, sizeof(kl_param));
    kl_param.kt_handle = dgpk_key_slot;
    kl_param.levels = AML_KL_LEVEL_3;
    kl_param.usage.crypto = AML_KT_FLAG_ENC_DEC;
    kl_param.usage.algo = AML_KT_ALGO_AES;
    kl_param.usage.uid = AML_KT_USER_M2M_0;
    kl_param.kl_algo = AML_KL_ALGO_AES;
    kl_param.kl_mode = AML_KL_MODE_AML;
#if CONFIG_PARTITION_ENCRYPTION_LOCAL
    /* use mrk ACRK */
    kl_param.mrk_cfg_index = AML_KL_MRK_ACRK;
#else
    /* use mrk DGPK1 */
    kl_param.mrk_cfg_index = AML_KL_MRK_DGPK1;
#endif
    /* function id 2 */
    kl_param.func_id = AML_KL_FUNC_AES_2;

    memcpy(kl_param.eks[0], dgpk_ek1, 16);
    memcpy(kl_param.eks[1], dgpk_ek2, 16);
    memcpy(kl_param.eks[2], dgpk_ek3, 16);
    ret = aml_mkl_run(&kl_param);
    if (ret != 0) {
#if ERROR_REPORT
        printf("aml_mkl_run error[ret:%d]\n", ret);
#endif
        return -1;
    }
    ret = aes_ecb_dec_keytbl(AES_KEY_SIZE_128, in, dst,
            NULL, AES_KEY_SIZE_128, dgpk_key_slot);
#if DEBUG
    printf("in:\n");
    printf("%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x\n",
            in[0], in[1], in[2],  in[3],  in[4], in[5], in[6], in[7],
            in[8], in[9], in[10], in[11], in[12],in[13],in[14], in[15]);

    printf("unwrap:");
    printf("%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x\n",
            dst[0], dst[1], dst[2],  dst[3],  dst[4], dst[5], dst[6], dst[7],
            dst[8], dst[9], dst[10], dst[11], dst[12],dst[13],dst[14], dst[15]);
#endif
    memcpy(out, dst, AES_KEY_SIZE_128);
    *out_sz = AES_KEY_SIZE_128;
    return ret;
}

void apply_enc_parts(void)
{
    int32_t i = 0, rc;
    u8 unwrap[AES_KEY_SIZE_128] = {0};
    uint32_t unwrap_sz = sizeof(unwrap);
    u64 fdt_addr;

    fdt_addr = env_get_ulong("dtb_mem_addr", 16, 0);
    for (i = 0; i < MAX_ENC_PARTS; i++) {
        if (enc_parts[i].in_used != 0) {
            rc = unwrap_key(enc_parts[i].wrapped, AES_KEY_SIZE_128,
                    unwrap, &unwrap_sz);
            if (rc < 0) {
#if ERROR_REPORT
                printf("unwrap key for %s failed.(%d)\n", enc_parts[i].part_name, rc);
#endif
                continue;
            }
            rc = fdt_find_and_setprop((void*)fdt_addr, "/keys", enc_parts[i].part_name,
                    unwrap, unwrap_sz, 1 /* create */);
            if (rc) {
                if (rc == -FDT_ERR_NOSPACE) {
                    /* no space. Try to make extra space for it */
                    int prop_sz = 0x10;
                    rc = fdt_shrink_to_minimum((void*)fdt_addr, prop_sz);
                    if (rc < 0) {
#if ERROR_REPORT
                        printf("Unable to add extra size(%d) for %s\n",
                                prop_sz, enc_parts[i].part_name);
#endif
                    } else {
                        rc = fdt_find_and_setprop((void*)fdt_addr, "/keys", enc_parts[i].part_name,
                                unwrap, unwrap_sz, 1 /* create */);
                        if (rc) {
#if ERROR_REPORT
                            printf("Failed to set partition enc for %s(%d)\n",
                                    enc_parts[i].part_name, rc);
#endif
                        }
                    }
                }
#if ERROR_REPORT
                else {
                    printf("Failed to set partition enc for %s(%d)\n",
                           enc_parts[i].part_name, rc);
                }
#endif
            }
        }
    }
    return;
}

int32_t find_enc_parts(const char* part_name)
{
    int32_t i = 0;

    if (part_name == NULL)
        return -1;

    for (i = 0; i < MAX_ENC_PARTS; i++) {
        if (enc_parts[i].in_used != 0) {
            if (!strcmp(part_name, enc_parts[i].part_name))
                return i;
        }
    }
    return -1;
}

int set_enc_parts(char *name, char *wrapped_hex, int wrapped_hex_sz)
{
    int32_t i = 0, rc;

    i = find_enc_parts(name);
    if (i < 0) {
        i = find_free_enc_parts_ent();
        if (i < 0) {
#if ERROR_REPORT
            printf("Unable to find free slot for %s\n", name);
#endif
            return CMD_RET_FAILURE;
        }
    }
    rc = hex2bin(wrapped_hex, enc_parts[i].wrapped, sizeof(enc_parts[i].wrapped));
    if (rc < 0) {
#if ERROR_REPORT
        printf("Unable to parse data hex2bin for %s\n", wrapped_hex);
#endif
        return CMD_RET_FAILURE;
    }
    strncpy(enc_parts[i].part_name, name, sizeof(enc_parts[i].part_name) -1);
    enc_parts[i].in_used = 1;

    return CMD_RET_SUCCESS;
}

int del_enc_parts(char *name)
{
    if (name) {
        int i = find_enc_parts(name);
        if (i >= 0)
            memset(&enc_parts[i], 0x0, sizeof(ENC_PART_ENT_T));
        return CMD_RET_SUCCESS;
    } else
        return CMD_RET_FAILURE;
}

int clear_enc_parts(void)
{
    memset(enc_parts, 0x0, sizeof(enc_parts));
    return CMD_RET_SUCCESS;
}

int init_enc_parts(void)
{
    s32 rc = CMD_RET_SUCCESS;
#ifdef PARTITION_ENC_ARGS
    char *args = NULL;
    const uint32_t args_len = strlen(PARTITION_ENC_ARGS);
    char *pair[MAX_ENC_PARTS] = {NULL, NULL, NULL, NULL};
    char *tmp;
    uint32_t i = 0;
    char *part_name;
    char *part_key_hex;
#if DEBUG
    printf("%s:%d: args_len = %d\n", __func__, __LINE__, args_len);
#endif
    if (args_len) {
        args = calloc(args_len + 1, sizeof(char));// include EOS
        if (!args) {
#if ERROR_REPORT
            printf("failed to alloc buf for args.(%u)\n", args_len);
#endif
            return -1;
        }
        strncpy(args, PARTITION_ENC_ARGS, args_len);
        tmp = args;

        for (i = 0; i < MAX_ENC_PARTS; i++) {
            pair[i] = strsep(&tmp, ";");

            if (!pair[i]) {
                break;
            } else if (!strlen(pair[i])) {
                // If delimiter appears at the end of string,
                // drop it and continue the parsing
                pair[i] = NULL;
                i--;
                continue;
            }
        }
        if (i == MAX_ENC_PARTS && strsep(&tmp, ";")) {
#if ERROR_REPORT
            printf("Too many enc parts provided by PARTITION_ENC_ARGS.\n");
#endif
            rc = CMD_RET_FAILURE;
            goto out;
        }
        /* Init enc_parts */
        clear_enc_parts();
        //memset(enc_parts, 0x0, sizeof(enc_parts));
        i = 0;
        while (i < MAX_ENC_PARTS && pair[i]) {
#if DEBUG
            printf("parsing pair[%s] %d ...\n", pair[i], i);
#endif
            tmp = pair[i];
            part_name = strsep(&tmp, ":");
            part_key_hex = tmp;
            if (part_name && part_key_hex) {
                rc = set_enc_parts(part_name, part_key_hex, strnlen(part_key_hex, 32));
            } else {
#if ERROR_REPORT
                printf("parsing error.\n");
#endif
                rc = CMD_RET_FAILURE;
                goto out;
            }
            i++;
        }
        rc = CMD_RET_SUCCESS;
    }
out:
    if (args)
        free(args);
#else
    printf("PARTITION_ENC_ARGS is not defied.\n");
    rc = CMD_RET_FAILURE;
#endif
    return rc;
}

static char partition_enc_help_text[] =
"\n"
"[set] <partition> {<value-hexdump-string>}\n"
"\t\t\tSet partition key to data.  DATA is in continuous\n"
"\t\t\t    hexdump format, e.g. aabb1122\n"
"[del] <partition>\tDelete <partition> in table of partition encryption.\n"
"[init]\t\t\tInitialize table of partition encryption.\n"
"[clear]\t\t\tClear table of partition encryption.\n"
"[dump]\t\t\tDump table of partition encryption.\n"
"[apply]\t\t\tApply table of partition encryption to DTS.\n"
"\n";

int do_partition_enc(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
{
    int rc = CMD_RET_SUCCESS;

    if (argc < 2) {
        return CMD_RET_USAGE;
    }

    if (argc == 2) {
        if (!strcmp(argv[1], "init"))
            rc = init_enc_parts();
        else if (!strcmp(argv[1], "clear"))
            rc = clear_enc_parts();
        else if (!strcmp(argv[1], "apply"))
            apply_enc_parts();
        else if (!strcmp(argv[1], "dump"))
            dump_enc_parts();
        else
            rc = CMD_RET_USAGE;
    } else if (argc == 3) {
        if (!strcmp(argv[1], "del"))
            rc = del_enc_parts(argv[2]);
        else
            rc = CMD_RET_USAGE;
    } else if (argc == 4) {
        if (!strcmp(argv[1], "set"))
            rc = set_enc_parts(argv[2], argv[3], strnlen(argv[3], 32));
        else
            rc = CMD_RET_USAGE;
    } else
        rc = CMD_RET_USAGE;

#if ERROR_REPORT
    if (rc != CMD_RET_SUCCESS)
        printf("Invalid arguments: %s\n", argv[1]);
#endif
    return rc;
}

U_BOOT_CMD(partition_enc,	4,	0,	do_partition_enc,
        "partition encryption command", partition_enc_help_text);

#if CONFIG_PARTITION_ENCRYPTION_LOCAL

#define LBN_SIZE                (512)
#define AES_BLOCK_SIZE          (16)
#define LBN_SHIFT               (9)
#define IV_FACTOR_BASE_ADDR     (0x800)

/* Get logical block number for encryption.
 * Size of each lbn is 512 bytes. */
static inline uint64_t get_lbn(uint64_t offset)
{
    return (offset >> LBN_SHIFT) + IV_FACTOR_BASE_ADDR;
}

int32_t calc_aes_128_cbc(uint8_t *in, uint32_t in_sz,
        uint8_t *key, uint32_t key_sz,
        uint8_t *iv, uint32_t iv_sz,
        uint8_t *out, int32_t *out_sz, int32_t mode)
{
    int32_t ret = 0;
    u8 key_exp[AES256_EXPAND_KEY_LENGTH];

    /* First we expand the key. */
    aes_expand_key(key, key_sz, key_exp);

    if (mode == DECRYPT)
        aes_cbc_decrypt_blocks(AES128_KEY_LENGTH, key_exp, iv,
                in, out, in_sz / AES_BLOCK_LENGTH);
    else if (mode == ENCRYPT)
        aes_cbc_encrypt_blocks(AES128_KEY_LENGTH, key_exp, iv,
                in, out, in_sz / AES_BLOCK_LENGTH);
#if ERROR_REPORT
    else
        printf("Unsupported mode\n");
#endif
    return ret;
}

static int32_t get_iv(uint64_t offset, uint8_t *md5, uint32_t md5_sz, uint8_t *out, int32_t *out_sz)
{
    uint64_t lbn = get_lbn(offset);
    uint8_t iv[AES_BLOCK_SIZE] = {};
    uint8_t iv_iv[AES_BLOCK_SIZE] = {0xff};
    int32_t ret = 0;

    if (!out || !out_sz || !*out_sz) {
#if ERROR_REPORT
        printf("Invalid input.\n");
#endif
        return -1;
    }

    iv[0] = lbn & 0xFF;
    iv[1] = (lbn >> 8) & 0xFF;
    iv[2] = (lbn >> 16) & 0xFF;
    iv[3] = (lbn >> 24) & 0xFF;

    ret = calc_aes_128_cbc(iv, sizeof(iv),
            md5, md5_sz,          /* key */
            iv_iv, sizeof(iv_iv), /* iv  */
            out, out_sz, 0);

    return ret;
}

uint32_t calc_md5(uint8_t *in, uint32_t in_sz, uint8_t *md5_hash, uint32_t *md5_hash_len)
{
    md5(in, in_sz, md5_hash);
    return 0;
}

int part_crypt(u8 *in, u64 in_sz,
        u8 *key, u32 key_sz,
        u8 *out, u64 out_sz,
        u64 offset, OP_MODE_T mode)
{
    int32_t ret = 0;
    int32_t remaining, read_sz, retVal = 0;
    u8 md5[16];
    u32 md5_sz = sizeof(md5);
    u64 buf_off = 0;
    u8 page_iv[AES_BLOCK_SIZE] = {};
    s32 page_iv_sz = AES_BLOCK_SIZE;
    s32 wr_sz, wr_total = 0;
    u32 padding = 0;
    u8 *temp = NULL;
    u32 temp_sz = 0;
    u32 remain_of_blk = 0;

    if (!in_sz || !key || !key_sz || !out) {
#if ERROR_REPORT
        printf("Invalid input.\n");
#endif
        return -1;
    }
    if (offset & 0x1ff) {
#if ERROR_REPORT
        printf("off should be multiple of 512\n");
#endif
        return -1;
    }
    if (out_sz < in_sz) {
#if ERROR_REPORT
        printf("Output size should >= input size(%#llx)\n", in_sz);
#endif
        return -1;
    }

    if (mode != ENCRYPT && mode != DECRYPT) {
#if ERROR_REPORT
        printf("Invalid mode.\n");
#endif
        return -1;
    }

    remaining = in_sz;
    /* partition key hash */
    calc_md5(key, key_sz, md5, &md5_sz);
    while (remaining > 0) {
        /* Read in a lbn if it is possible */
        read_sz = (remaining >= LBN_SIZE)? LBN_SIZE: remaining;
        remain_of_blk = read_sz & (AES_BLOCK_LENGTH - 1);
        if (remain_of_blk) {
            /* Since LBN_SIZE is multiple of AES_BLOCK_LENGTH, entering
             * this case means it is the last read */
            padding = AES_BLOCK_LENGTH - remain_of_blk;
#if DEBUG
            printf("remain_of_blk = %d, padding = %u\n", remain_of_blk, padding);
#endif
            temp_sz = read_sz + padding;
            temp = (u8*)malloc(temp_sz);
            if (temp == NULL) {
#if ERROR_REPORT
                printf("failed to alloc temp buffer.(%d)\n", temp_sz);
#endif
                ret = -1;
                goto out;
            }
            /* copy input to temp buffer */
            memcpy(temp, &in[buf_off], read_sz);
        }
        /* update iv */
        page_iv_sz = AES_BLOCK_SIZE;
        ret = get_iv(offset, md5, md5_sz, page_iv, &page_iv_sz);
        if (ret < 0) {
#if ERROR_REPORT
            printf("Error to get enc iv\n");
#endif
            retVal = -1;
            goto out;
        }
        if (remain_of_blk) {
            wr_sz = temp_sz;
            ret = calc_aes_128_cbc(temp, temp_sz,   /* input */
                    key, key_sz,                    /* key */
                    page_iv, page_iv_sz,            /* iv */
                    temp, &wr_sz, mode);            /* output */
        } else {
            wr_sz = read_sz;
            ret = calc_aes_128_cbc(&in[buf_off], read_sz,   /* input */
                    key, key_sz,                            /* key */
                    page_iv, page_iv_sz,                    /* iv */
                    &out[buf_off], &wr_sz, mode);           /* output */
        }
        if (ret < 0) {
#if ERROR_REPORT
            printf("Error to calc aes 128 cbc\n");
#endif
            retVal = -1;
            goto out;
        }
        if (remain_of_blk) {
            /* copy result from temp to out buffer */
            memcpy(&out[buf_off], temp, read_sz);
        }

        offset += read_sz;
        buf_off += read_sz;
        wr_total += wr_sz;
        remaining -= read_sz;
    }

out:
    if (temp) free(temp);

    return ret;
}

int part_dec(const char *name, u8 *in, u64 in_sz, u8 *out, u64 out_sz, u64 off)
{
    s32 rc = env_get_yesno("local_dec");

    if (rc == 1) {
        int i = find_enc_parts(name);
        u8 unwrap[AES_KEY_SIZE_128] = {0x0};
        u32 unwrap_sz = AES_KEY_SIZE_128;

        if (i < 0)
            return -1;

        printf("%s is an enc part\n", name);
        rc = unwrap_key(enc_parts[i].wrapped, sizeof(enc_parts[i].wrapped),
                unwrap, &unwrap_sz);
        if (rc) {
#if ERROR_REPORT
            printf("unwrap key failed.(%d)\n", rc);
#endif
            return rc;
        }

        return part_crypt(in, in_sz, unwrap, unwrap_sz, out, out_sz, off, DECRYPT);
    } else
        return -1;
}

#endif //CONFIG_PARTITION_ENCRYPTION_LOCAL