cmd/tee.c - u-boot - Git at Google

 #include <common.h>
 #include <command.h>
 #include <dm.h>
 #include <mapmem.h>
 #include <tee/ta_helper.h>
 #include <tee/ta_hello_world.h>
 #include <linux/kernel.h>

 /*
  * Probe and list all compliant TEE devices.
  */
 static int
 do_tee_devices(cmd_tbl_t * cmdtp, int flag, int argc, char *const argv[])
 {
 	struct udevice *dev = NULL;
 	int rc;

 	for (rc = uclass_first_device(UCLASS_TEE, &dev); dev;
 	     rc = uclass_next_device(&dev)) {
 		/* $name@$instance */
 		printf("%s@%llx\n", dev->name, map_to_sysmem(dev));
 	}

 	return rc;
 }

 static int do_tee_hello(cmd_tbl_t * cmdtp, int flag, int argc,
 			char *const argv[])
 {
 	int rc;
 	struct tee_optee_ta_uuid uuid = TA_HELLO_WORLD_UUID;
 	struct tee_param param = {0};
 	param.attr = TEE_PARAM_ATTR_TYPE_VALUE_INOUT;
 	param.u.value.a = 16;

 	printf("Invoking TA to increment value 16.\n");
 	rc = ta_call(&uuid, TA_HELLO_WORLD_CMD_INC_VALUE, 1,
 		     &param);
 	if (rc)
 		return rc;

 	printf("TA incremented my value to %lld.\n", param.u.value.a);
 	if (param.u.value.a != 17)
 		return -EINVAL;

 	return rc;
 }

 static int do_tee_echo(cmd_tbl_t * cmdtp, int flag, int argc,
 			char *const argv[])
 {
 	int rc;
 	struct ta_context context = {0};
 	struct tee_optee_ta_uuid uuid = TA_HELLO_WORLD_UUID;
 	struct tee_param params[4] = {0};
 	char *msg = "Hello from U-Boot!";
 	size_t msg_size;
 	struct tee_shm *shm_msg = NULL;
 	const size_t buf_size = 128;
 	struct tee_shm *shm_buf = NULL;

 	if (argc > 1)
 		msg = argv[1];
 	msg_size = strlen(msg) + 1;

 	/* Connect to TA. */
 	rc = ta_open(&uuid, &context);
 	if (rc)
 		return rc;

 	/* Pass message in via shared memory via params[0]. */
 	rc = tee_shm_alloc(context.tee, msg_size, TEE_SHM_ALLOC, &shm_msg);
 	if (rc)
 		goto out;

 	memcpy(shm_msg->addr, msg, msg_size);

 	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
 	params[0].u.memref.shm = shm_msg;
 	params[0].u.memref.size = msg_size;

 	/* Receive responses via shared memory via params[3]. */
 	rc = tee_shm_alloc(context.tee, buf_size, TEE_SHM_ALLOC, &shm_buf);
 	if (rc)
 		goto out;

 	params[3].attr = TEE_PARAM_ATTR_TYPE_MEMREF_OUTPUT;
 	params[3].u.memref.shm = shm_buf;
 	params[3].u.memref.size = buf_size;

 	/* Invoke TA_HELLO_WORLD_CMD_ECHO exposed by the TA. */
 	printf("Sending \"%s\" to the hello world TA.\n", msg);
 	rc = ta_invoke(&context, TA_HELLO_WORLD_CMD_ECHO,
 			ARRAY_SIZE(params), params);
 	if (rc)
 		goto out;

 	if (params[3].u.memref.size == 0 || params[3].u.memref.size > buf_size) {
 		rc = -ERANGE;
 		goto out;
 	}

 	char *resp = (char *)shm_buf->addr;
 	resp[params[3].u.memref.size - 1] = 0;
 	printf("TA said: \"%s\"\n", resp);

 out:
 	tee_shm_free(shm_msg);
 	tee_shm_free(shm_buf);
 	ta_close(&context);
 	return rc;
 }

 static int write_keyval(const char *key, const void *buf, size_t bytes)
 {
 	int rc;
 	struct ta_context context = {0};
 	struct tee_optee_ta_uuid uuid = TA_HELLO_WORLD_UUID;
 	struct tee_param params[2] = {0};

 	struct tee_shm *shm_key = NULL;
 	size_t key_size = strlen(key) + 1;
 	struct tee_shm *shm_buf = NULL;

 	/* Connect to TA. */
 	rc = ta_open(&uuid, &context);
 	if (rc)
 		return rc;

 	/* Pass in @key via shared memory. */
 	rc = tee_shm_alloc(context.tee, key_size, TEE_SHM_ALLOC, &shm_key);
 	if (rc)
 		goto out;

 	memcpy(shm_key->addr, key, key_size);
 	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
 	params[0].u.memref.shm = shm_key;
 	params[0].u.memref.size = key_size;

 	/* Pass in buf via shared memory. */
 	rc = tee_shm_alloc(context.tee, bytes, TEE_SHM_ALLOC, &shm_buf);
 	if (rc)
 		goto out;

 	memcpy(shm_buf->addr, buf, bytes);
 	params[1].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
 	params[1].u.memref.shm = shm_buf;
 	params[1].u.memref.size = bytes;

 	/* Invoke the TA. */
 	rc = ta_invoke(&context, TA_HELLO_WORLD_CMD_WRITE_PERSISTENT_VALUE,
 			ARRAY_SIZE(params), params);
 	if (rc)
 		goto out;

 out:
 	tee_shm_free(shm_key);
 	tee_shm_free(shm_buf);
 	ta_close(&context);
 	return rc;
 }

 static int read_keyval(const char *key, void *buf, size_t bytes)
 {
 	int rc;
 	struct ta_context context = {0};
 	struct tee_optee_ta_uuid uuid = TA_HELLO_WORLD_UUID;
 	struct tee_param params[2] = {0};

 	struct tee_shm *shm_key = NULL;
 	size_t key_size = strlen(key) + 1;
 	struct tee_shm *shm_buf = NULL;

 	/* Connect to TA. */
 	rc = ta_open(&uuid, &context);
 	if (rc)
 		return rc;

 	/* Pass in @key via shared memory. */
 	rc = tee_shm_alloc(context.tee, key_size, TEE_SHM_ALLOC, &shm_key);
 	if (rc)
 		goto out;

 	memcpy(shm_key->addr, key, key_size);
 	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
 	params[0].u.memref.shm = shm_key;
 	params[0].u.memref.size = key_size;

 	/* Pass in buf via shared memory. */
 	rc = tee_shm_alloc(context.tee, bytes, TEE_SHM_ALLOC, &shm_buf);
 	if (rc)
 		goto out;

 	params[1].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
 	params[1].u.memref.shm = shm_buf;
 	params[1].u.memref.size = bytes;

 	/* Invoke the TA. */
 	rc = ta_invoke(&context, TA_HELLO_WORLD_CMD_READ_PERSISTENT_VALUE,
 			ARRAY_SIZE(params), params);
 	if (rc)
 		goto out;

 	if (params[1].u.memref.size == 0 || params[1].u.memref.size > bytes)
 		return -ERANGE;
 	memcpy(buf, (void *)(params[1].u.memref.shm->addr),
 			params[1].u.memref.size);

 out:
 	tee_shm_free(shm_key);
 	tee_shm_free(shm_buf);
 	ta_close(&context);
 	return rc;
 }

 static int do_tee_set(cmd_tbl_t * cmdtp, int flag, int argc,
 			char *const argv[])
 {
 	if (argc < 3 || !argv[1][0] || !argv[2][0])
 		return -EINVAL;

 	const char *key = argv[1];
 	const char *value = argv[2];

 	return write_keyval(key, value, strlen(value) + 1);
 }

 static int do_tee_get(cmd_tbl_t * cmdtp, int flag, int argc,
 			char *const argv[])
 {
 	if (argc < 2 || !argv[1][0])
 		return -EINVAL;

 	const char *key = argv[1];
 	char value[64];
 	int rc = read_keyval(key, value, sizeof(value));
 	if (rc)
 		return rc;

 	value[sizeof(value) - 1] = 0;
 	printf("\"%s\"\n", value);

 	return 0;
 }

 static cmd_tbl_t tee_commands[] = {
 	U_BOOT_CMD_MKENT(devices, /* maxargs = */ 1, /* repeatable = */ 0,
 			 do_tee_devices, "", ""),
 	U_BOOT_CMD_MKENT(hello, /* maxargs = */ 1, /* repeatable = */ 0,
 			 do_tee_hello, "", ""),
 	U_BOOT_CMD_MKENT(echo, /* maxargs = */ 2, /* repeatable = */ 0,
 			 do_tee_echo, "", ""),
 	U_BOOT_CMD_MKENT(set, /* maxargs = */ 3, /* repeatable = */ 0,
 			 do_tee_set, "", ""),
 	U_BOOT_CMD_MKENT(get, /* maxargs = */ 2, /* repeatable = */ 0,
 			 do_tee_get, "", ""),
 };

 static int do_tee(cmd_tbl_t *cmdtp, int flag, int argc, char *const argv[])
 {
 	cmd_tbl_t *tee_cmd;
 	int ret;

 	if (argc < 2)
 		return CMD_RET_USAGE;

 	argc--;
 	argv++;

 	tee_cmd = find_cmd_tbl(argv[0], tee_commands, ARRAY_SIZE(tee_commands));
 	if (!tee_cmd || argc > tee_cmd->maxargs)
 		return CMD_RET_USAGE;

 	ret = tee_cmd->cmd(tee_cmd, flag, argc, argv);
 	return cmd_process_error(cmdtp, ret);
 }

 U_BOOT_CMD(tee, /* maxargs = */ 4, /* repeatable = */ 0, do_tee,
 	   "Trusted Execution Environment operations.",
 	   "devices - Probes TEE devices.\n"
 	   "tee hello - Pings the hello world TA.\n"
 	   "tee echo [msg] - Sends a message to the hello world TA.\n"
 	   "tee set <key> <value> - Sets a string value to a key in the hello world TA's secure storage.\n"
 	   "tee get <key> - Gets the value to a key in the hello world TA's secure storage.\n"
 	   );
	#include <common.h>
	#include <command.h>
	#include <dm.h>
	#include <mapmem.h>
	#include <tee/ta_helper.h>
	#include <tee/ta_hello_world.h>
	#include <linux/kernel.h>

	/*
	* Probe and list all compliant TEE devices.
	*/
	static int
	do_tee_devices(cmd_tbl_t * cmdtp, int flag, int argc, char *const argv[])
	{
	struct udevice *dev = NULL;
	int rc;

	for (rc = uclass_first_device(UCLASS_TEE, &dev); dev;
	rc = uclass_next_device(&dev)) {
	/* $name@$instance */
	printf("%s@%llx\n", dev->name, map_to_sysmem(dev));
	}

	return rc;
	}

	static int do_tee_hello(cmd_tbl_t * cmdtp, int flag, int argc,
	char *const argv[])
	{
	int rc;
	struct tee_optee_ta_uuid uuid = TA_HELLO_WORLD_UUID;
	struct tee_param param = {0};
	param.attr = TEE_PARAM_ATTR_TYPE_VALUE_INOUT;
	param.u.value.a = 16;

	printf("Invoking TA to increment value 16.\n");
	rc = ta_call(&uuid, TA_HELLO_WORLD_CMD_INC_VALUE, 1,
	&param);
	if (rc)
	return rc;

	printf("TA incremented my value to %lld.\n", param.u.value.a);
	if (param.u.value.a != 17)
	return -EINVAL;

	return rc;
	}

	static int do_tee_echo(cmd_tbl_t * cmdtp, int flag, int argc,
	char *const argv[])
	{
	int rc;
	struct ta_context context = {0};
	struct tee_optee_ta_uuid uuid = TA_HELLO_WORLD_UUID;
	struct tee_param params[4] = {0};
	char *msg = "Hello from U-Boot!";
	size_t msg_size;
	struct tee_shm *shm_msg = NULL;
	const size_t buf_size = 128;
	struct tee_shm *shm_buf = NULL;

	if (argc > 1)
	msg = argv[1];
	msg_size = strlen(msg) + 1;

	/* Connect to TA. */
	rc = ta_open(&uuid, &context);
	if (rc)
	return rc;

	/* Pass message in via shared memory via params[0]. */
	rc = tee_shm_alloc(context.tee, msg_size, TEE_SHM_ALLOC, &shm_msg);
	if (rc)
	goto out;

	memcpy(shm_msg->addr, msg, msg_size);

	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
	params[0].u.memref.shm = shm_msg;
	params[0].u.memref.size = msg_size;

	/* Receive responses via shared memory via params[3]. */
	rc = tee_shm_alloc(context.tee, buf_size, TEE_SHM_ALLOC, &shm_buf);
	if (rc)
	goto out;

	params[3].attr = TEE_PARAM_ATTR_TYPE_MEMREF_OUTPUT;
	params[3].u.memref.shm = shm_buf;
	params[3].u.memref.size = buf_size;

	/* Invoke TA_HELLO_WORLD_CMD_ECHO exposed by the TA. */
	printf("Sending \"%s\" to the hello world TA.\n", msg);
	rc = ta_invoke(&context, TA_HELLO_WORLD_CMD_ECHO,
	ARRAY_SIZE(params), params);
	if (rc)
	goto out;

	if (params[3].u.memref.size == 0 \|\| params[3].u.memref.size > buf_size) {
	rc = -ERANGE;
	goto out;
	}

	char resp = (char )shm_buf->addr;
	resp[params[3].u.memref.size - 1] = 0;
	printf("TA said: \"%s\"\n", resp);

	out:
	tee_shm_free(shm_msg);
	tee_shm_free(shm_buf);
	ta_close(&context);
	return rc;
	}

	static int write_keyval(const char key, const void buf, size_t bytes)
	{
	int rc;
	struct ta_context context = {0};
	struct tee_optee_ta_uuid uuid = TA_HELLO_WORLD_UUID;
	struct tee_param params[2] = {0};

	struct tee_shm *shm_key = NULL;
	size_t key_size = strlen(key) + 1;
	struct tee_shm *shm_buf = NULL;

	/* Connect to TA. */
	rc = ta_open(&uuid, &context);
	if (rc)
	return rc;

	/* Pass in @key via shared memory. */
	rc = tee_shm_alloc(context.tee, key_size, TEE_SHM_ALLOC, &shm_key);
	if (rc)
	goto out;

	memcpy(shm_key->addr, key, key_size);
	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
	params[0].u.memref.shm = shm_key;
	params[0].u.memref.size = key_size;

	/* Pass in buf via shared memory. */
	rc = tee_shm_alloc(context.tee, bytes, TEE_SHM_ALLOC, &shm_buf);
	if (rc)
	goto out;

	memcpy(shm_buf->addr, buf, bytes);
	params[1].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
	params[1].u.memref.shm = shm_buf;
	params[1].u.memref.size = bytes;

	/* Invoke the TA. */
	rc = ta_invoke(&context, TA_HELLO_WORLD_CMD_WRITE_PERSISTENT_VALUE,
	ARRAY_SIZE(params), params);
	if (rc)
	goto out;

	out:
	tee_shm_free(shm_key);
	tee_shm_free(shm_buf);
	ta_close(&context);
	return rc;
	}

	static int read_keyval(const char key, void buf, size_t bytes)
	{
	int rc;
	struct ta_context context = {0};
	struct tee_optee_ta_uuid uuid = TA_HELLO_WORLD_UUID;
	struct tee_param params[2] = {0};

	struct tee_shm *shm_key = NULL;
	size_t key_size = strlen(key) + 1;
	struct tee_shm *shm_buf = NULL;

	/* Connect to TA. */
	rc = ta_open(&uuid, &context);
	if (rc)
	return rc;

	/* Pass in @key via shared memory. */
	rc = tee_shm_alloc(context.tee, key_size, TEE_SHM_ALLOC, &shm_key);
	if (rc)
	goto out;

	memcpy(shm_key->addr, key, key_size);
	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
	params[0].u.memref.shm = shm_key;
	params[0].u.memref.size = key_size;

	/* Pass in buf via shared memory. */
	rc = tee_shm_alloc(context.tee, bytes, TEE_SHM_ALLOC, &shm_buf);
	if (rc)
	goto out;

	params[1].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
	params[1].u.memref.shm = shm_buf;
	params[1].u.memref.size = bytes;

	/* Invoke the TA. */
	rc = ta_invoke(&context, TA_HELLO_WORLD_CMD_READ_PERSISTENT_VALUE,
	ARRAY_SIZE(params), params);
	if (rc)
	goto out;

	if (params[1].u.memref.size == 0 \|\| params[1].u.memref.size > bytes)
	return -ERANGE;
	memcpy(buf, (void *)(params[1].u.memref.shm->addr),
	params[1].u.memref.size);

	out:
	tee_shm_free(shm_key);
	tee_shm_free(shm_buf);
	ta_close(&context);
	return rc;
	}

	static int do_tee_set(cmd_tbl_t * cmdtp, int flag, int argc,
	char *const argv[])
	{
	if (argc < 3 \|\| !argv[1][0] \|\| !argv[2][0])
	return -EINVAL;

	const char *key = argv[1];
	const char *value = argv[2];

	return write_keyval(key, value, strlen(value) + 1);
	}

	static int do_tee_get(cmd_tbl_t * cmdtp, int flag, int argc,
	char *const argv[])
	{
	if (argc < 2 \|\| !argv[1][0])
	return -EINVAL;

	const char *key = argv[1];
	char value[64];
	int rc = read_keyval(key, value, sizeof(value));
	if (rc)
	return rc;

	value[sizeof(value) - 1] = 0;
	printf("\"%s\"\n", value);

	return 0;
	}

	static cmd_tbl_t tee_commands[] = {
	U_BOOT_CMD_MKENT(devices, /* maxargs = / 1, / repeatable = */ 0,
	do_tee_devices, "", ""),
	U_BOOT_CMD_MKENT(hello, /* maxargs = / 1, / repeatable = */ 0,
	do_tee_hello, "", ""),
	U_BOOT_CMD_MKENT(echo, /* maxargs = / 2, / repeatable = */ 0,
	do_tee_echo, "", ""),
	U_BOOT_CMD_MKENT(set, /* maxargs = / 3, / repeatable = */ 0,
	do_tee_set, "", ""),
	U_BOOT_CMD_MKENT(get, /* maxargs = / 2, / repeatable = */ 0,
	do_tee_get, "", ""),
	};

	static int do_tee(cmd_tbl_t cmdtp, int flag, int argc, char const argv[])
	{
	cmd_tbl_t *tee_cmd;
	int ret;

	if (argc < 2)
	return CMD_RET_USAGE;

	argc--;
	argv++;

	tee_cmd = find_cmd_tbl(argv[0], tee_commands, ARRAY_SIZE(tee_commands));
	if (!tee_cmd \|\| argc > tee_cmd->maxargs)
	return CMD_RET_USAGE;

	ret = tee_cmd->cmd(tee_cmd, flag, argc, argv);
	return cmd_process_error(cmdtp, ret);
	}

	U_BOOT_CMD(tee, /* maxargs = / 4, / repeatable = */ 0, do_tee,
	"Trusted Execution Environment operations.",
	"devices - Probes TEE devices.\n"
	"tee hello - Pings the hello world TA.\n"
	"tee echo [msg] - Sends a message to the hello world TA.\n"
	"tee set <key> <value> - Sets a string value to a key in the hello world TA's secure storage.\n"
	"tee get <key> - Gets the value to a key in the hello world TA's secure storage.\n"
	);