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

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

 #include <common.h>
 #include <command.h>
 #include <mapmem.h>
 #include <tee.h>

 #define RMA_OS_DEBUG_ENABLE (0)

 #define CMD_RMA_OS_INIT_DATA      (0)
 #define CMD_RMA_OS_GET_CMD_LIST   (1)
 #define CMD_RMA_OS_RUN_CMD        (2)
 #define CMD_RMA_OS_GET_INFO       (3)

 #define TA_RMA_TEE_OS_UUID \
 		{ 0x99a216e5, 0x6718, 0x4dfd, \
 		{ 0xbb, 0xb1, 0x37, 0xcd, 0x0d, 0x00, 0x10, 0xc6 } }

 #define RMA_CMD_MAX_NUM           (64)
 #define RMA_CMD_MAX_LENGTH        (128)
 #define RMA_AUTH_HDR_MAGIC        (0x414D52) // "RMA"
 #define RMA_AUTH_DATA_MAX_SIZE    (32 + 256 + 16 * RMA_CMD_MAX_NUM)

 struct rma_info_output {
 	char name[RMA_CMD_MAX_LENGTH];
 	char info[RMA_CMD_MAX_LENGTH];
 };

 static struct rma_info_output rma_info[RMA_CMD_MAX_NUM + 1];

 static void dump_rma_info(struct rma_info_output *result, uint32_t len,
 		char *title, char *comment)
 {
 	uint32_t i;

 	for (i = 0; i < 10; i++)
 		printf("---------");

 	printf("\n");

 	printf("%-55s | %s\n", title, comment);

 	for (i = 0; i < 10; i++)
 		printf("---------");
 	printf("\n");

 	for (i = 0; i < len; i++)
 		printf("%-55s | %s\n", result[i].name, result[i].info);

 	for (i = 0; i < 10; i++)
 		printf("---------");

 	printf("\n");
 }

 static int do_rma_init(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
 {
 	int ret = CMD_RET_SUCCESS;
 	uint8_t *buf;
 	uint32_t map_addr = 0;
 	uint32_t magic;
 	uint32_t number;
 	uint32_t real_size;
 	char rma_data[RMA_AUTH_DATA_MAX_SIZE];
 	struct tee_param params;
 	struct udevice *dev;
 	struct tee_open_session_arg open_arg;
 	struct tee_invoke_arg invoke_arg;
 	const struct tee_optee_ta_uuid uuid = TA_RMA_TEE_OS_UUID;

 	if (argc != 2) {
 		printf("Usage: rma init ADDR, ADDR is memory address\n");
 		return CMD_RET_USAGE;
 	}

 	/* Get RMA data memory address */
 	map_addr = simple_strtoul(argv[1], NULL, 16);
 	if (map_addr <= 0 || map_addr >= 0xffffffff) {
 		printf("RMA data memory address: %#x invalid!\n", map_addr);
 		return CMD_RET_USAGE;
 	}

 	/* Map RMA data memory */
 	buf = map_sysmem(map_addr, RMA_AUTH_DATA_MAX_SIZE);

 	/* Check RMA data magic */
 	magic = *(uint32_t *)buf;
 	if (magic != RMA_AUTH_HDR_MAGIC) {
 		printf("RMA data magic: %#x invalid, it must be %#x\n",
 			magic, RMA_AUTH_HDR_MAGIC);
 		unmap_sysmem(buf);
 		return -CMD_RET_FAILURE;
 	}

 	/* Check RMA data size */
 	number = *(uint32_t *)(buf + 16);
 	real_size = RMA_AUTH_DATA_MAX_SIZE - (16 * (RMA_CMD_MAX_NUM - number));
 	if (real_size > RMA_AUTH_DATA_MAX_SIZE) {
 		printf("RMA data size %#x larger than %#x\n",
 			real_size, RMA_AUTH_DATA_MAX_SIZE);
 		unmap_sysmem(buf);
 		return -CMD_RET_FAILURE;
 	}

 	/* Get RMA data */
 	memset(rma_data, 0, sizeof(rma_data));
 	memcpy(rma_data, buf, real_size);

 	/* Unmap RMA data memory */
 	unmap_sysmem(buf);

 #if RMA_OS_DEBUG_ENABLE
 	/* Debug print RMA data */
 	int index;

 	printf("[Debug] RMA Data:\n");
 	for (index = 0; index < real_size; index++) {
 		if ((index % 16) == 0)
 			printf("\n");
 		printf("%02x ", rma_data[index]);
 	}
 	printf("\nPrint data finished!\n");
 #endif

 	/* Find TEE device */
 	dev = tee_find_device(NULL, NULL, NULL, NULL);
 	if (!dev) {
 		printf("tee_find_device failed");
 		return -CMD_RET_FAILURE;
 	}

 	/* Open TEE session */
 	memset(&open_arg, 0, sizeof(open_arg));
 	tee_optee_ta_uuid_to_octets(open_arg.uuid, &uuid);
 	ret = tee_open_session(dev, &open_arg, 0, NULL);
 	if (ret) {
 		printf("tee_open_session failed, ret: %#x\n", ret);
 		return CMD_RET_FAILURE;
 	}
 	if (open_arg.ret) {
 		printf("tee_open_session failed, ret: %#x, origin: %#x\n",
 				open_arg.ret, open_arg.ret_origin);
 		return CMD_RET_FAILURE;
 	}

 	/* Alloc share memory for RMA data */
 	memset(&params, 0, sizeof(params));
 	params.attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
 	params.u.memref.size = real_size;
 	ret = tee_shm_alloc(dev, real_size, 0, &params.u.memref.shm);
 	if (ret) {
 		printf("tee_shm_alloc failed, ret: %#x\n", ret);
 		tee_close_session(dev, open_arg.session);
 		return CMD_RET_FAILURE;
 	}
 	memcpy(params.u.memref.shm->addr, rma_data, real_size);

 	/* Invoke TEE command */
 	memset(&invoke_arg, 0, sizeof(invoke_arg));
 	invoke_arg.session = open_arg.session;
 	invoke_arg.func = CMD_RMA_OS_INIT_DATA;
 	ret = tee_invoke_func(dev, &invoke_arg, 1, &params);
 	if (ret) {
 		printf("tee_invoke_func failed, ret = %#x\n", ret);
 		ret = CMD_RET_FAILURE;
 		goto exit;
 	}
 	if (invoke_arg.ret) {
 		printf("RMA OS init failed, ret: %#x, origin: %#x\n",
 				invoke_arg.ret, invoke_arg.ret_origin);
 		ret = CMD_RET_FAILURE;
 		goto exit;
 	}
 	printf("RMA OS init and verify authentication data successfully!\n");

 exit:
 	tee_shm_free(params.u.memref.shm);
 	tee_close_session(dev, open_arg.session);

 	return ret;
 }

 static int do_rma_list(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
 {
 	int ret = CMD_RET_SUCCESS;
 	uint32_t info_size;
 	struct tee_param params;
 	struct udevice *dev;
 	struct tee_open_session_arg open_arg;
 	struct tee_invoke_arg invoke_arg;
 	const struct tee_optee_ta_uuid uuid = TA_RMA_TEE_OS_UUID;

 	/* Find TEE device */
 	dev = tee_find_device(NULL, NULL, NULL, NULL);
 	if (!dev) {
 		printf("call tee_find_device failed");
 		return -CMD_RET_FAILURE;
 	}

 	/* Open TEE session */
 	memset(&open_arg, 0, sizeof(open_arg));
 	tee_optee_ta_uuid_to_octets(open_arg.uuid, &uuid);
 	ret = tee_open_session(dev, &open_arg, 0, NULL);
 	if (ret) {
 		printf("call tee_open_session failed, ret: %#x\n", ret);
 		return CMD_RET_FAILURE;
 	}
 	if (open_arg.ret) {
 		printf("call tee_open_session failed, ret: %#x, origin: %#x\n",
 				open_arg.ret, open_arg.ret_origin);
 		return CMD_RET_FAILURE;
 	}

 	/* Alloc share memory for RMA command list */
 	memset((void *)rma_info, 0, sizeof(rma_info));
 	info_size = sizeof(rma_info);
 	memset(&params, 0, sizeof(params));
 	params.attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
 	params.u.memref.size = info_size;
 	ret = tee_shm_alloc(dev, info_size, 0, &params.u.memref.shm);
 	if (ret) {
 		printf("tee_shm_alloc failed, ret: %#x\n", ret);
 		tee_close_session(dev, open_arg.session);
 		return CMD_RET_FAILURE;
 	}

 	/* Invoke TEE command */
 	memset(&invoke_arg, 0, sizeof(invoke_arg));
 	invoke_arg.session = open_arg.session;
 	invoke_arg.func = CMD_RMA_OS_GET_CMD_LIST;
 	ret = tee_invoke_func(dev, &invoke_arg, 1, &params);
 	if (ret) {
 		printf("tee_invoke_func failed, ret = %#x\n", ret);
 		ret = CMD_RET_FAILURE;
 		goto exit;
 	}
 	if (invoke_arg.ret) {
 		printf("Get RMA Command List failed, ret: %#x, origin: %#x\n",
 					invoke_arg.ret, invoke_arg.ret_origin);
 		ret = CMD_RET_FAILURE;
 		goto exit;
 	}

 	/* Dump RMA command list */
 	info_size = (params.u.memref.size / sizeof(struct rma_info_output));
 	memcpy(rma_info, params.u.memref.shm->addr, params.u.memref.size);
 	dump_rma_info(rma_info, info_size, "Command", "Status");

 exit:
 	tee_shm_free(params.u.memref.shm);
 	tee_close_session(dev, open_arg.session);

 	return ret;
 }

 static int do_rma_cmd(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
 {
 	int ret = CMD_RET_SUCCESS;
 	int cmd = 0;
 	struct tee_param params[4];
 	struct udevice *dev;
 	struct tee_open_session_arg open_arg;
 	struct tee_invoke_arg invoke_arg;
 	const struct tee_optee_ta_uuid uuid = TA_RMA_TEE_OS_UUID;

 	if (argc != 2) {
 		printf("Usage: rma cmd ID, ID is command index\n");
 		return CMD_RET_USAGE;
 	}

 	/* Check RMA command range */
 	cmd = simple_strtoul(argv[1], NULL, 10);
 	if (cmd <= 0 || cmd > RMA_CMD_MAX_NUM) {
 		printf("Command ID: %d invalid! Range: 1~%d\n", cmd, RMA_CMD_MAX_NUM);
 		return CMD_RET_USAGE;
 	}

 	/* Find TEE device */
 	dev = tee_find_device(NULL, NULL, NULL, NULL);
 	if (!dev) {
 		printf("tee_find_device failed");
 		return -CMD_RET_FAILURE;
 	}

 	/* Open TEE session */
 	memset(&open_arg, 0, sizeof(open_arg));
 	tee_optee_ta_uuid_to_octets(open_arg.uuid, &uuid);
 	ret = tee_open_session(dev, &open_arg, 0, NULL);
 	if (ret) {
 		printf("tee_open_session failed, ret: %#x\n", ret);
 		ret = CMD_RET_FAILURE;
 		return ret;
 	}
 	if (open_arg.ret) {
 		printf("tee_open_session failed, ret: %#x, origin: %#x\n",
 				open_arg.ret, open_arg.ret_origin);
 		ret = CMD_RET_FAILURE;
 		return ret;
 	}

 	/* Set TEE invoke params */
 	memset(params, 0, sizeof(params));
 	params[0].attr = TEE_PARAM_ATTR_TYPE_VALUE_INOUT;
 	params[0].u.value.a = cmd;
 	params[1].attr = TEE_PARAM_ATTR_TYPE_VALUE_INOUT;
 	params[2].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
 	params[3].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;

 	/* Invoke TEE command */
 	memset(&invoke_arg, 0, sizeof(invoke_arg));
 	invoke_arg.session = open_arg.session;
 	invoke_arg.func = CMD_RMA_OS_RUN_CMD;
 	ret = tee_invoke_func(dev, &invoke_arg, 4, params);
 	if (ret) {
 		printf("tee_invoke_func failed, ret = %#x\n", ret);
 		ret = CMD_RET_FAILURE;
 		goto exit;
 	}
 	if (invoke_arg.ret) {
 		printf("Run RMA Command %d failed, ret: %#x, origin: %#x\n",
 				cmd, invoke_arg.ret, invoke_arg.ret_origin);
 		ret = CMD_RET_FAILURE;
 		goto exit;
 	}
 	printf("Run RMA Command %d successfully!\n", cmd);

 exit:
 	tee_close_session(dev, open_arg.session);

 	return ret;
 }

 static int do_rma_info(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
 {
 	int ret = CMD_RET_SUCCESS;
 	uint32_t info_size;
 	struct tee_param params;
 	struct udevice *dev;
 	struct tee_open_session_arg open_arg;
 	struct tee_invoke_arg invoke_arg;
 	const struct tee_optee_ta_uuid uuid = TA_RMA_TEE_OS_UUID;

 	/* Find TEE device */
 	dev = tee_find_device(NULL, NULL, NULL, NULL);
 	if (!dev) {
 		printf("call tee_find_device failed");
 		return -CMD_RET_FAILURE;
 	}

 	/* Open TEE session */
 	memset(&open_arg, 0, sizeof(open_arg));
 	tee_optee_ta_uuid_to_octets(open_arg.uuid, &uuid);
 	ret = tee_open_session(dev, &open_arg, 0, NULL);
 	if (ret) {
 		printf("call tee_open_session failed, ret: %#x\n", ret);
 		return CMD_RET_FAILURE;
 	}
 	if (open_arg.ret) {
 		printf("call tee_open_session failed, ret: %#x, origin: %#x\n",
 				open_arg.ret, open_arg.ret_origin);
 		return CMD_RET_FAILURE;
 	}

 	/* Alloc share memory for RMA OS info */
 	memset((void *)rma_info, 0, sizeof(rma_info));
 	info_size = sizeof(rma_info);
 	memset(&params, 0, sizeof(params));
 	params.attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
 	params.u.memref.size = info_size;
 	ret = tee_shm_alloc(dev, info_size, 0, &params.u.memref.shm);
 	if (ret) {
 		printf("tee_shm_alloc failed, ret: %#x\n", ret);
 		tee_close_session(dev, open_arg.session);
 		return CMD_RET_FAILURE;
 	}

 	/* Invoke TEE command */
 	memset(&invoke_arg, 0, sizeof(invoke_arg));
 	invoke_arg.session = open_arg.session;
 	invoke_arg.func = CMD_RMA_OS_GET_INFO;
 	ret = tee_invoke_func(dev, &invoke_arg, 1, &params);
 	if (ret) {
 		printf("tee_invoke_func failed, ret = %#x\n", ret);
 		ret = CMD_RET_FAILURE;
 		goto exit;
 	}
 	if (invoke_arg.ret) {
 		printf("Get RMA OS information failed, ret: %#x, origin: %#x\n",
 					invoke_arg.ret, invoke_arg.ret_origin);
 		ret = CMD_RET_FAILURE;
 		goto exit;
 	}

 	/* Dump RMA OS information */
 	info_size = (params.u.memref.size / sizeof(struct rma_info_output));
 	memcpy(rma_info, params.u.memref.shm->addr, params.u.memref.size);
 	dump_rma_info(rma_info, info_size, "Information", "Status");

 exit:
 	tee_shm_free(params.u.memref.shm);
 	tee_close_session(dev, open_arg.session);

 	return ret;
 }

 static cmd_tbl_t cmd_rma_sub[] = {
 	U_BOOT_CMD_MKENT(init, 2, 0, do_rma_init, "", ""),
 	U_BOOT_CMD_MKENT(list, 1, 0, do_rma_list, "", ""),
 	U_BOOT_CMD_MKENT(cmd, 2, 0, do_rma_cmd, "", ""),
 	U_BOOT_CMD_MKENT(info, 1, 0, do_rma_info, "", ""),
 };

 static int do_rma(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
 {
 	cmd_tbl_t *c;

 	if (argc < 2)
 		return cmd_usage(cmdtp);

 	argc--;
 	argv++;

 	/* Check uboot sub command */
 	c = find_cmd_tbl(argv[0], &cmd_rma_sub[0], ARRAY_SIZE(cmd_rma_sub));

 	if (c)
 		return  c->cmd(cmdtp, flag, argc, argv);
 	else
 		return cmd_usage(cmdtp);
 }

 static char rma_help_text[] =
 	"\n"
 	"[rma info] -Show RMA OS information\n"
 	"[rma init] -Init RMA OS and verify authentication data\n"
 	"            e.g: rma init ADDR; ADDR is memory address, e.g: 0x1080000\n"
 	"[rma list] -Show RMA OS command permission list\n"
 	"[rma cmd]  -Run a RMA OS command\n"
 	"            e.g: rma cmd ID; ID is command index, e.g: 1,2,3...\n";

 U_BOOT_CMD(rma, CONFIG_SYS_MAXARGS, 0, do_rma,
 	"RMA OS Function", rma_help_text);
	// SPDX-License-Identifier: (GPL-2.0+ OR MIT)
	/*
	* Copyright (c) 2022 Amlogic, Inc. All rights reserved.
	*/

	#include <common.h>
	#include <command.h>
	#include <mapmem.h>
	#include <tee.h>

	#define RMA_OS_DEBUG_ENABLE (0)

	#define CMD_RMA_OS_INIT_DATA (0)
	#define CMD_RMA_OS_GET_CMD_LIST (1)
	#define CMD_RMA_OS_RUN_CMD (2)
	#define CMD_RMA_OS_GET_INFO (3)

	#define TA_RMA_TEE_OS_UUID \
	{ 0x99a216e5, 0x6718, 0x4dfd, \
	{ 0xbb, 0xb1, 0x37, 0xcd, 0x0d, 0x00, 0x10, 0xc6 } }

	#define RMA_CMD_MAX_NUM (64)
	#define RMA_CMD_MAX_LENGTH (128)
	#define RMA_AUTH_HDR_MAGIC (0x414D52) // "RMA"
	#define RMA_AUTH_DATA_MAX_SIZE (32 + 256 + 16 * RMA_CMD_MAX_NUM)

	struct rma_info_output {
	char name[RMA_CMD_MAX_LENGTH];
	char info[RMA_CMD_MAX_LENGTH];
	};

	static struct rma_info_output rma_info[RMA_CMD_MAX_NUM + 1];

	static void dump_rma_info(struct rma_info_output *result, uint32_t len,
	char title, char comment)
	{
	uint32_t i;

	for (i = 0; i < 10; i++)
	printf("---------");

	printf("\n");

	printf("%-55s \| %s\n", title, comment);

	for (i = 0; i < 10; i++)
	printf("---------");
	printf("\n");

	for (i = 0; i < len; i++)
	printf("%-55s \| %s\n", result[i].name, result[i].info);

	for (i = 0; i < 10; i++)
	printf("---------");

	printf("\n");
	}

	static int do_rma_init(cmd_tbl_t cmdtp, int flag, int argc, char const argv[])
	{
	int ret = CMD_RET_SUCCESS;
	uint8_t *buf;
	uint32_t map_addr = 0;
	uint32_t magic;
	uint32_t number;
	uint32_t real_size;
	char rma_data[RMA_AUTH_DATA_MAX_SIZE];
	struct tee_param params;
	struct udevice *dev;
	struct tee_open_session_arg open_arg;
	struct tee_invoke_arg invoke_arg;
	const struct tee_optee_ta_uuid uuid = TA_RMA_TEE_OS_UUID;

	if (argc != 2) {
	printf("Usage: rma init ADDR, ADDR is memory address\n");
	return CMD_RET_USAGE;
	}

	/* Get RMA data memory address */
	map_addr = simple_strtoul(argv[1], NULL, 16);
	if (map_addr <= 0 \|\| map_addr >= 0xffffffff) {
	printf("RMA data memory address: %#x invalid!\n", map_addr);
	return CMD_RET_USAGE;
	}

	/* Map RMA data memory */
	buf = map_sysmem(map_addr, RMA_AUTH_DATA_MAX_SIZE);

	/* Check RMA data magic */
	magic = (uint32_t )buf;
	if (magic != RMA_AUTH_HDR_MAGIC) {
	printf("RMA data magic: %#x invalid, it must be %#x\n",
	magic, RMA_AUTH_HDR_MAGIC);
	unmap_sysmem(buf);
	return -CMD_RET_FAILURE;
	}

	/* Check RMA data size */
	number = (uint32_t )(buf + 16);
	real_size = RMA_AUTH_DATA_MAX_SIZE - (16 * (RMA_CMD_MAX_NUM - number));
	if (real_size > RMA_AUTH_DATA_MAX_SIZE) {
	printf("RMA data size %#x larger than %#x\n",
	real_size, RMA_AUTH_DATA_MAX_SIZE);
	unmap_sysmem(buf);
	return -CMD_RET_FAILURE;
	}

	/* Get RMA data */
	memset(rma_data, 0, sizeof(rma_data));
	memcpy(rma_data, buf, real_size);

	/* Unmap RMA data memory */
	unmap_sysmem(buf);

	#if RMA_OS_DEBUG_ENABLE
	/* Debug print RMA data */
	int index;

	printf("[Debug] RMA Data:\n");
	for (index = 0; index < real_size; index++) {
	if ((index % 16) == 0)
	printf("\n");
	printf("%02x ", rma_data[index]);
	}
	printf("\nPrint data finished!\n");
	#endif

	/* Find TEE device */
	dev = tee_find_device(NULL, NULL, NULL, NULL);
	if (!dev) {
	printf("tee_find_device failed");
	return -CMD_RET_FAILURE;
	}

	/* Open TEE session */
	memset(&open_arg, 0, sizeof(open_arg));
	tee_optee_ta_uuid_to_octets(open_arg.uuid, &uuid);
	ret = tee_open_session(dev, &open_arg, 0, NULL);
	if (ret) {
	printf("tee_open_session failed, ret: %#x\n", ret);
	return CMD_RET_FAILURE;
	}
	if (open_arg.ret) {
	printf("tee_open_session failed, ret: %#x, origin: %#x\n",
	open_arg.ret, open_arg.ret_origin);
	return CMD_RET_FAILURE;
	}

	/* Alloc share memory for RMA data */
	memset(&params, 0, sizeof(params));
	params.attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
	params.u.memref.size = real_size;
	ret = tee_shm_alloc(dev, real_size, 0, &params.u.memref.shm);
	if (ret) {
	printf("tee_shm_alloc failed, ret: %#x\n", ret);
	tee_close_session(dev, open_arg.session);
	return CMD_RET_FAILURE;
	}
	memcpy(params.u.memref.shm->addr, rma_data, real_size);

	/* Invoke TEE command */
	memset(&invoke_arg, 0, sizeof(invoke_arg));
	invoke_arg.session = open_arg.session;
	invoke_arg.func = CMD_RMA_OS_INIT_DATA;
	ret = tee_invoke_func(dev, &invoke_arg, 1, &params);
	if (ret) {
	printf("tee_invoke_func failed, ret = %#x\n", ret);
	ret = CMD_RET_FAILURE;
	goto exit;
	}
	if (invoke_arg.ret) {
	printf("RMA OS init failed, ret: %#x, origin: %#x\n",
	invoke_arg.ret, invoke_arg.ret_origin);
	ret = CMD_RET_FAILURE;
	goto exit;
	}
	printf("RMA OS init and verify authentication data successfully!\n");

	exit:
	tee_shm_free(params.u.memref.shm);
	tee_close_session(dev, open_arg.session);

	return ret;
	}

	static int do_rma_list(cmd_tbl_t cmdtp, int flag, int argc, char const argv[])
	{
	int ret = CMD_RET_SUCCESS;
	uint32_t info_size;
	struct tee_param params;
	struct udevice *dev;
	struct tee_open_session_arg open_arg;
	struct tee_invoke_arg invoke_arg;
	const struct tee_optee_ta_uuid uuid = TA_RMA_TEE_OS_UUID;

	/* Find TEE device */
	dev = tee_find_device(NULL, NULL, NULL, NULL);
	if (!dev) {
	printf("call tee_find_device failed");
	return -CMD_RET_FAILURE;
	}

	/* Open TEE session */
	memset(&open_arg, 0, sizeof(open_arg));
	tee_optee_ta_uuid_to_octets(open_arg.uuid, &uuid);
	ret = tee_open_session(dev, &open_arg, 0, NULL);
	if (ret) {
	printf("call tee_open_session failed, ret: %#x\n", ret);
	return CMD_RET_FAILURE;
	}
	if (open_arg.ret) {
	printf("call tee_open_session failed, ret: %#x, origin: %#x\n",
	open_arg.ret, open_arg.ret_origin);
	return CMD_RET_FAILURE;
	}

	/* Alloc share memory for RMA command list */
	memset((void *)rma_info, 0, sizeof(rma_info));
	info_size = sizeof(rma_info);
	memset(&params, 0, sizeof(params));
	params.attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
	params.u.memref.size = info_size;
	ret = tee_shm_alloc(dev, info_size, 0, &params.u.memref.shm);
	if (ret) {
	printf("tee_shm_alloc failed, ret: %#x\n", ret);
	tee_close_session(dev, open_arg.session);
	return CMD_RET_FAILURE;
	}

	/* Invoke TEE command */
	memset(&invoke_arg, 0, sizeof(invoke_arg));
	invoke_arg.session = open_arg.session;
	invoke_arg.func = CMD_RMA_OS_GET_CMD_LIST;
	ret = tee_invoke_func(dev, &invoke_arg, 1, &params);
	if (ret) {
	printf("tee_invoke_func failed, ret = %#x\n", ret);
	ret = CMD_RET_FAILURE;
	goto exit;
	}
	if (invoke_arg.ret) {
	printf("Get RMA Command List failed, ret: %#x, origin: %#x\n",
	invoke_arg.ret, invoke_arg.ret_origin);
	ret = CMD_RET_FAILURE;
	goto exit;
	}

	/* Dump RMA command list */
	info_size = (params.u.memref.size / sizeof(struct rma_info_output));
	memcpy(rma_info, params.u.memref.shm->addr, params.u.memref.size);
	dump_rma_info(rma_info, info_size, "Command", "Status");

	exit:
	tee_shm_free(params.u.memref.shm);
	tee_close_session(dev, open_arg.session);

	return ret;
	}

	static int do_rma_cmd(cmd_tbl_t cmdtp, int flag, int argc, char const argv[])
	{
	int ret = CMD_RET_SUCCESS;
	int cmd = 0;
	struct tee_param params[4];
	struct udevice *dev;
	struct tee_open_session_arg open_arg;
	struct tee_invoke_arg invoke_arg;
	const struct tee_optee_ta_uuid uuid = TA_RMA_TEE_OS_UUID;

	if (argc != 2) {
	printf("Usage: rma cmd ID, ID is command index\n");
	return CMD_RET_USAGE;
	}

	/* Check RMA command range */
	cmd = simple_strtoul(argv[1], NULL, 10);
	if (cmd <= 0 \|\| cmd > RMA_CMD_MAX_NUM) {
	printf("Command ID: %d invalid! Range: 1~%d\n", cmd, RMA_CMD_MAX_NUM);
	return CMD_RET_USAGE;
	}

	/* Find TEE device */
	dev = tee_find_device(NULL, NULL, NULL, NULL);
	if (!dev) {
	printf("tee_find_device failed");
	return -CMD_RET_FAILURE;
	}

	/* Open TEE session */
	memset(&open_arg, 0, sizeof(open_arg));
	tee_optee_ta_uuid_to_octets(open_arg.uuid, &uuid);
	ret = tee_open_session(dev, &open_arg, 0, NULL);
	if (ret) {
	printf("tee_open_session failed, ret: %#x\n", ret);
	ret = CMD_RET_FAILURE;
	return ret;
	}
	if (open_arg.ret) {
	printf("tee_open_session failed, ret: %#x, origin: %#x\n",
	open_arg.ret, open_arg.ret_origin);
	ret = CMD_RET_FAILURE;
	return ret;
	}

	/* Set TEE invoke params */
	memset(params, 0, sizeof(params));
	params[0].attr = TEE_PARAM_ATTR_TYPE_VALUE_INOUT;
	params[0].u.value.a = cmd;
	params[1].attr = TEE_PARAM_ATTR_TYPE_VALUE_INOUT;
	params[2].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
	params[3].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;

	/* Invoke TEE command */
	memset(&invoke_arg, 0, sizeof(invoke_arg));
	invoke_arg.session = open_arg.session;
	invoke_arg.func = CMD_RMA_OS_RUN_CMD;
	ret = tee_invoke_func(dev, &invoke_arg, 4, params);
	if (ret) {
	printf("tee_invoke_func failed, ret = %#x\n", ret);
	ret = CMD_RET_FAILURE;
	goto exit;
	}
	if (invoke_arg.ret) {
	printf("Run RMA Command %d failed, ret: %#x, origin: %#x\n",
	cmd, invoke_arg.ret, invoke_arg.ret_origin);
	ret = CMD_RET_FAILURE;
	goto exit;
	}
	printf("Run RMA Command %d successfully!\n", cmd);

	exit:
	tee_close_session(dev, open_arg.session);

	return ret;
	}

	static int do_rma_info(cmd_tbl_t cmdtp, int flag, int argc, char const argv[])
	{
	int ret = CMD_RET_SUCCESS;
	uint32_t info_size;
	struct tee_param params;
	struct udevice *dev;
	struct tee_open_session_arg open_arg;
	struct tee_invoke_arg invoke_arg;
	const struct tee_optee_ta_uuid uuid = TA_RMA_TEE_OS_UUID;

	/* Find TEE device */
	dev = tee_find_device(NULL, NULL, NULL, NULL);
	if (!dev) {
	printf("call tee_find_device failed");
	return -CMD_RET_FAILURE;
	}

	/* Open TEE session */
	memset(&open_arg, 0, sizeof(open_arg));
	tee_optee_ta_uuid_to_octets(open_arg.uuid, &uuid);
	ret = tee_open_session(dev, &open_arg, 0, NULL);
	if (ret) {
	printf("call tee_open_session failed, ret: %#x\n", ret);
	return CMD_RET_FAILURE;
	}
	if (open_arg.ret) {
	printf("call tee_open_session failed, ret: %#x, origin: %#x\n",
	open_arg.ret, open_arg.ret_origin);
	return CMD_RET_FAILURE;
	}

	/* Alloc share memory for RMA OS info */
	memset((void *)rma_info, 0, sizeof(rma_info));
	info_size = sizeof(rma_info);
	memset(&params, 0, sizeof(params));
	params.attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
	params.u.memref.size = info_size;
	ret = tee_shm_alloc(dev, info_size, 0, &params.u.memref.shm);
	if (ret) {
	printf("tee_shm_alloc failed, ret: %#x\n", ret);
	tee_close_session(dev, open_arg.session);
	return CMD_RET_FAILURE;
	}

	/* Invoke TEE command */
	memset(&invoke_arg, 0, sizeof(invoke_arg));
	invoke_arg.session = open_arg.session;
	invoke_arg.func = CMD_RMA_OS_GET_INFO;
	ret = tee_invoke_func(dev, &invoke_arg, 1, &params);
	if (ret) {
	printf("tee_invoke_func failed, ret = %#x\n", ret);
	ret = CMD_RET_FAILURE;
	goto exit;
	}
	if (invoke_arg.ret) {
	printf("Get RMA OS information failed, ret: %#x, origin: %#x\n",
	invoke_arg.ret, invoke_arg.ret_origin);
	ret = CMD_RET_FAILURE;
	goto exit;
	}

	/* Dump RMA OS information */
	info_size = (params.u.memref.size / sizeof(struct rma_info_output));
	memcpy(rma_info, params.u.memref.shm->addr, params.u.memref.size);
	dump_rma_info(rma_info, info_size, "Information", "Status");

	exit:
	tee_shm_free(params.u.memref.shm);
	tee_close_session(dev, open_arg.session);

	return ret;
	}

	static cmd_tbl_t cmd_rma_sub[] = {
	U_BOOT_CMD_MKENT(init, 2, 0, do_rma_init, "", ""),
	U_BOOT_CMD_MKENT(list, 1, 0, do_rma_list, "", ""),
	U_BOOT_CMD_MKENT(cmd, 2, 0, do_rma_cmd, "", ""),
	U_BOOT_CMD_MKENT(info, 1, 0, do_rma_info, "", ""),
	};

	static int do_rma(cmd_tbl_t cmdtp, int flag, int argc, char const argv[])
	{
	cmd_tbl_t *c;

	if (argc < 2)
	return cmd_usage(cmdtp);

	argc--;
	argv++;

	/* Check uboot sub command */
	c = find_cmd_tbl(argv[0], &cmd_rma_sub[0], ARRAY_SIZE(cmd_rma_sub));

	if (c)
	return c->cmd(cmdtp, flag, argc, argv);
	else
	return cmd_usage(cmdtp);
	}

	static char rma_help_text[] =
	"\n"
	"[rma info] -Show RMA OS information\n"
	"[rma init] -Init RMA OS and verify authentication data\n"
	" e.g: rma init ADDR; ADDR is memory address, e.g: 0x1080000\n"
	"[rma list] -Show RMA OS command permission list\n"
	"[rma cmd] -Run a RMA OS command\n"
	" e.g: rma cmd ID; ID is command index, e.g: 1,2,3...\n";

	U_BOOT_CMD(rma, CONFIG_SYS_MAXARGS, 0, do_rma,
	"RMA OS Function", rma_help_text);