cmd/amlogic/partition_encryption.c - yocto/uboot - Git Browser for ODROID

 // 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>
 #include <asm/amlogic/arch/bl31_apis.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];

 s32 find_free_enc_parts_ent(void)
 {
 	s32 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 u8 *buf, s32 len)
 {
 	char line_buf[128] = {0x0};
 	char *line_buf_ptr = line_buf;
 	u32 i = 0;
 	const u8 *ptr = buf;
 	const s32 sz = len;

 	printf("\n========= Start dump %s(%d)\n", name, sz);
 	while (len - 16 >= 0) {
 		printf("%02x %02x %02x %02x %02x %02x %02x %02x",
 		       ptr[0], ptr[i + 1], ptr[i + 2], ptr[i + 3], ptr[i + 4], ptr[i + 5],
 		       ptr[i + 6], ptr[i + 7]);
 		printf("%02x %02x %02x %02x %02x %02x %02x %02x\n",
 		       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)
 {
 	s32 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;
 #if !CONFIG_PARTITION_ENCRYPTION_LOCAL
 	struct amlkl_params kl_param;
 	/* dgpk_func_0_bs_0_tsep_0.test.vector.txt */
 	u8 dgpk_ek3[16] = {
 		0xa6, 0x20, 0x65, 0x7f, 0xa9, 0x6b, 0x19, 0x66,
 		0xa1, 0x03, 0x47, 0xb4, 0x23, 0xde, 0x20, 0x25
 	};

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

 	u8 dgpk_ek1[16] = {
 		0xfa, 0x29, 0x3b, 0x43, 0x25, 0x7c, 0x96, 0xbd,
 		0x55, 0x2b, 0x60, 0xd5, 0x95, 0x60, 0xcf, 0xcc
 	};
 #endif
 	u32 key_slot = 31; //only test slot
 	u8 dst[AES_KEY_SIZE_128] = {0};

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

 #if CONFIG_PARTITION_ENCRYPTION_LOCAL
 	/* when mrk is ACRK, derive it in bl31 */
 	ret = partition_enc_kl_derive_key(key_slot);
 #else
 	memset(&kl_param, 0, sizeof(kl_param));
 	kl_param.kt_handle = 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;
 	/* use mrk DGPK1 */
 	kl_param.mrk_cfg_index = AML_KL_MRK_DGPK1;
 	/* 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);
 #endif
 	if (ret != 0) {
 #if ERROR_REPORT
 		printf("derive key error[ret:%d]\n", ret);
 #endif
 		return -1;
 	}
 	ret = aes_ecb_dec_keytbl(AES_KEY_SIZE_128, in, dst,
 				 NULL, AES_KEY_SIZE_128, 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)
 {
 	s32 i = 0, rc;
 	u8 unwrap[AES_KEY_SIZE_128] = {0};
 	u32 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
 						}
 					}
 				} else {
 #if ERROR_REPORT
 					printf("Failed to set partition enc for %s(%d)\n",
 					       enc_parts[i].part_name, rc);
 #endif
 				}
 			}
 		}
 	}
 	return;
 }

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

 	if (!part_name)
 		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)
 {
 	s32 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;
 	}
 	strlcpy(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 u32 args_len = strlen(PARTITION_ENC_ARGS);
 	char *pair[MAX_ENC_PARTS] = {NULL, NULL, NULL, NULL};
 	char *tmp;
 	u32 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;
 		}
 		strlcpy(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 u64 get_lbn(u64 offset)
 {
 	return (offset >> LBN_SHIFT) + IV_FACTOR_BASE_ADDR;
 }

 s32 calc_aes_128_cbc(u8 *in, u32 in_sz,
 		     u8 *key, u32 key_sz,
 		     u8 *iv, u32 iv_sz,
 		     u8 *out, s32 *out_sz, s32 mode)
 {
 	s32 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 s32 get_iv(u64 offset, u8 *md5, u32 md5_sz, u8 *out, s32 *out_sz)
 {
 	u64 lbn = get_lbn(offset);
 	u8 iv[AES_BLOCK_SIZE] = {};
 	u8 iv_iv[AES_BLOCK_SIZE] = {0xff};
 	s32 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;
 }

 u32 calc_md5(u8 *in, u32 in_sz, u8 *md5_hash, u32 *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)
 {
 	s32 ret = 0;
 	s32 remaining, read_sz, ret_val = 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) {
 #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
 			ret_val = -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
 			ret_val = -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
	// 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>
	#include <asm/amlogic/arch/bl31_apis.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];

	s32 find_free_enc_parts_ent(void)
	{
	s32 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 u8 buf, s32 len)
	{
	char line_buf[128] = {0x0};
	char *line_buf_ptr = line_buf;
	u32 i = 0;
	const u8 *ptr = buf;
	const s32 sz = len;

	printf("\n========= Start dump %s(%d)\n", name, sz);
	while (len - 16 >= 0) {
	printf("%02x %02x %02x %02x %02x %02x %02x %02x",
	ptr[0], ptr[i + 1], ptr[i + 2], ptr[i + 3], ptr[i + 4], ptr[i + 5],
	ptr[i + 6], ptr[i + 7]);
	printf("%02x %02x %02x %02x %02x %02x %02x %02x\n",
	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)
	{
	s32 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;
	#if !CONFIG_PARTITION_ENCRYPTION_LOCAL
	struct amlkl_params kl_param;
	/* dgpk_func_0_bs_0_tsep_0.test.vector.txt */
	u8 dgpk_ek3[16] = {
	0xa6, 0x20, 0x65, 0x7f, 0xa9, 0x6b, 0x19, 0x66,
	0xa1, 0x03, 0x47, 0xb4, 0x23, 0xde, 0x20, 0x25
	};

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

	u8 dgpk_ek1[16] = {
	0xfa, 0x29, 0x3b, 0x43, 0x25, 0x7c, 0x96, 0xbd,
	0x55, 0x2b, 0x60, 0xd5, 0x95, 0x60, 0xcf, 0xcc
	};
	#endif
	u32 key_slot = 31; //only test slot
	u8 dst[AES_KEY_SIZE_128] = {0};

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

	#if CONFIG_PARTITION_ENCRYPTION_LOCAL
	/* when mrk is ACRK, derive it in bl31 */
	ret = partition_enc_kl_derive_key(key_slot);
	#else
	memset(&kl_param, 0, sizeof(kl_param));
	kl_param.kt_handle = 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;
	/* use mrk DGPK1 */
	kl_param.mrk_cfg_index = AML_KL_MRK_DGPK1;
	/* 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);
	#endif
	if (ret != 0) {
	#if ERROR_REPORT
	printf("derive key error[ret:%d]\n", ret);
	#endif
	return -1;
	}
	ret = aes_ecb_dec_keytbl(AES_KEY_SIZE_128, in, dst,
	NULL, AES_KEY_SIZE_128, 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)
	{
	s32 i = 0, rc;
	u8 unwrap[AES_KEY_SIZE_128] = {0};
	u32 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
	}
	}
	} else {
	#if ERROR_REPORT
	printf("Failed to set partition enc for %s(%d)\n",
	enc_parts[i].part_name, rc);
	#endif
	}
	}
	}
	}
	return;
	}

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

	if (!part_name)
	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)
	{
	s32 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;
	}
	strlcpy(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 u32 args_len = strlen(PARTITION_ENC_ARGS);
	char *pair[MAX_ENC_PARTS] = {NULL, NULL, NULL, NULL};
	char *tmp;
	u32 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;
	}
	strlcpy(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 u64 get_lbn(u64 offset)
	{
	return (offset >> LBN_SHIFT) + IV_FACTOR_BASE_ADDR;
	}

	s32 calc_aes_128_cbc(u8 *in, u32 in_sz,
	u8 *key, u32 key_sz,
	u8 *iv, u32 iv_sz,
	u8 out, s32 out_sz, s32 mode)
	{
	s32 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 s32 get_iv(u64 offset, u8 md5, u32 md5_sz, u8 out, s32 *out_sz)
	{
	u64 lbn = get_lbn(offset);
	u8 iv[AES_BLOCK_SIZE] = {};
	u8 iv_iv[AES_BLOCK_SIZE] = {0xff};
	s32 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;
	}

	u32 calc_md5(u8 in, u32 in_sz, u8 md5_hash, u32 *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)
	{
	s32 ret = 0;
	s32 remaining, read_sz, ret_val = 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) {
	#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
	ret_val = -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
	ret_val = -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