[v2,15/15] avb_verify: support using OP-TEE TA AVB

diff --git a/common/avb_verify.c b/common/avb_verify.c
index 20e35ade3029..a50c05bf7847 100644
--- a/common/avb_verify.c
+++ b/common/avb_verify.c
@@ -10,6 +10,8 @@ 
 #include <image.h>
 #include <malloc.h>
 #include <part.h>
+#include <tee.h>
+#include <tee/optee_ta_avb.h>
 
 const unsigned char avb_root_pub[1032] = {
 	0x0, 0x0, 0x10, 0x0, 0x55, 0xd9, 0x4, 0xad, 0xd8, 0x4,
@@ -594,6 +596,79 @@  static AvbIOResult validate_vbmeta_public_key(AvbOps *ops,
 	return AVB_IO_RESULT_OK;
 }
 
+#ifdef CONFIG_OPTEE_TA_AVB
+static void uuid_to_octets(u8 d[TEE_UUID_LEN],
+			   const struct tee_optee_ta_uuid *s)
+{
+	d[0] = s->time_low >> 24;
+	d[1] = s->time_low >> 16;
+	d[2] = s->time_low >> 8;
+	d[3] = s->time_low;
+	d[4] = s->time_mid >> 8;
+	d[5] = s->time_mid;
+	d[6] = s->time_hi_and_version >> 8;
+	d[7] = s->time_hi_and_version;
+	memcpy(d + 8, s->clock_seq_and_node, sizeof(s->clock_seq_and_node));
+}
+
+static int get_open_session(struct AvbOpsData *ops_data)
+{
+	struct udevice *tee = NULL;
+
+	while (!ops_data->tee) {
+		const struct tee_optee_ta_uuid uuid = TA_AVB_UUID;
+		struct tee_open_session_arg arg;
+		int rc;
+
+		tee = tee_find_device(tee, NULL, NULL, NULL);
+		if (!tee)
+			return -ENODEV;
+
+		memset(&arg, 0, sizeof(arg));
+		uuid_to_octets(arg.uuid, &uuid);
+		rc = tee_open_session(tee, &arg, 0, NULL);
+		if (!rc) {
+			ops_data->tee = tee;
+			ops_data->session = arg.session;
+		}
+	}
+
+	return 0;
+}
+
+static AvbIOResult invoke_func(struct AvbOpsData *ops_data, u32 func,
+			       ulong num_param, struct tee_param *param)
+{
+	struct tee_invoke_arg arg;
+
+	if (get_open_session(ops_data))
+		return AVB_IO_RESULT_ERROR_IO;
+
+	memset(&arg, 0, sizeof(arg));
+	arg.func = func;
+	arg.session = ops_data->session;
+
+	if (tee_invoke_func(ops_data->tee, &arg, num_param, param))
+		return AVB_IO_RESULT_ERROR_IO;
+	switch (arg.ret) {
+	case TEE_SUCCESS:
+		return AVB_IO_RESULT_OK;
+	case TEE_ERROR_OUT_OF_MEMORY:
+		return AVB_IO_RESULT_ERROR_OOM;
+	case TEE_ERROR_TARGET_DEAD:
+		/*
+		 * The TA has paniced, close the session to reload the TA
+		 * for the next request.
+		 */
+		tee_close_session(ops_data->tee, ops_data->session);
+		ops_data->tee = NULL;
+		return AVB_IO_RESULT_ERROR_IO;
+	default:
+		return AVB_IO_RESULT_ERROR_IO;
+	}
+}
+#endif
+
 /**
  * read_rollback_index() - gets the rollback index corresponding to the
  * location of given by @out_rollback_index.
@@ -609,6 +684,7 @@  static AvbIOResult read_rollback_index(AvbOps *ops,
 				       size_t rollback_index_slot,
 				       u64 *out_rollback_index)
 {
+#ifndef CONFIG_OPTEE_TA_AVB
 	/* For now we always return 0 as the stored rollback index. */
 	printf("%s not supported yet\n", __func__);
 
@@ -616,6 +692,27 @@  static AvbIOResult read_rollback_index(AvbOps *ops,
 		*out_rollback_index = 0;
 
 	return AVB_IO_RESULT_OK;
+#else
+	AvbIOResult rc;
+	struct tee_param param[2];
+
+	if (rollback_index_slot >= TA_AVB_MAX_ROLLBACK_LOCATIONS)
+		return AVB_IO_RESULT_ERROR_NO_SUCH_VALUE;
+
+	memset(param, 0, sizeof(param));
+	param[0].attr = TEE_PARAM_ATTR_TYPE_VALUE_INPUT;
+	param[0].u.value.a = rollback_index_slot;
+	param[1].attr = TEE_PARAM_ATTR_TYPE_VALUE_OUTPUT;
+
+	rc = invoke_func(ops->user_data, TA_AVB_CMD_READ_ROLLBACK_INDEX,
+			 ARRAY_SIZE(param), param);
+	if (rc)
+		return rc;
+
+	*out_rollback_index = (u64)param[1].u.value.a << 32 |
+			      (u32)param[1].u.value.b;
+	return AVB_IO_RESULT_OK;
+#endif
 }
 
 /**
@@ -633,10 +730,27 @@  static AvbIOResult write_rollback_index(AvbOps *ops,
 					size_t rollback_index_slot,
 					u64 rollback_index)
 {
+#ifndef CONFIG_OPTEE_TA_AVB
 	/* For now this is a no-op. */
 	printf("%s not supported yet\n", __func__);
 
 	return AVB_IO_RESULT_OK;
+#else
+	struct tee_param param[2];
+
+	if (rollback_index_slot >= TA_AVB_MAX_ROLLBACK_LOCATIONS)
+		return AVB_IO_RESULT_ERROR_NO_SUCH_VALUE;
+
+	memset(param, 0, sizeof(param));
+	param[0].attr = TEE_PARAM_ATTR_TYPE_VALUE_INPUT;
+	param[0].u.value.a = rollback_index_slot;
+	param[1].attr = TEE_PARAM_ATTR_TYPE_VALUE_INPUT;
+	param[1].u.value.a = (u32)(rollback_index >> 32);
+	param[1].u.value.b = (u32)rollback_index;
+
+	return invoke_func(ops->user_data, TA_AVB_CMD_WRITE_ROLLBACK_INDEX,
+			   ARRAY_SIZE(param), param);
+#endif
 }
 
 /**
@@ -652,6 +766,7 @@  static AvbIOResult write_rollback_index(AvbOps *ops,
  */
 static AvbIOResult read_is_device_unlocked(AvbOps *ops, bool *out_is_unlocked)
 {
+#ifndef CONFIG_OPTEE_TA_AVB
 	/* For now we always return that the device is unlocked. */
 
 	printf("%s not supported yet\n", __func__);
@@ -659,6 +774,16 @@  static AvbIOResult read_is_device_unlocked(AvbOps *ops, bool *out_is_unlocked)
 	*out_is_unlocked = true;
 
 	return AVB_IO_RESULT_OK;
+#else
+	AvbIOResult rc;
+	struct tee_param param = { .attr = TEE_PARAM_ATTR_TYPE_VALUE_OUTPUT };
+
+	rc = invoke_func(ops->user_data, TA_AVB_CMD_READ_LOCK_STATE, 1, &param);
+	if (rc)
+		return rc;
+	*out_is_unlocked = !param.u.value.a;
+	return AVB_IO_RESULT_OK;
+#endif
 }
 
 /**
@@ -737,6 +862,11 @@  void avb_ops_free(AvbOps *ops)
 
 	ops_data = ops->user_data;
 
-	if (ops_data)
+	if (ops_data) {
+#ifdef CONFIG_OPTEE_TA_AVB
+		if (ops_data->tee)
+			tee_close_session(ops_data->tee, ops_data->session);
+#endif
 		avb_free(ops_data);
+	}
 }
diff --git a/doc/README.avb2 b/doc/README.avb2
index 120279fedbe2..a29cee1b6f50 100644
--- a/doc/README.avb2
+++ b/doc/README.avb2
@@ -18,6 +18,13 @@  Integrity of the bootloader (U-boot BLOB and environment) is out of scope.
 For additional details check:
 https://android.googlesource.com/platform/external/avb/+/master/README.md
 
+1.1. AVB using OP-TEE (optional)
+---------------------------------
+If AVB is configured to use OP-TEE (see 4. below) rollback indexes and
+device lock state are stored in RPMB. The RPMB partition is managed by
+OP-TEE (https://www.op-tee.org/) which is a secure OS leveraging ARM
+TrustZone.
+
 
 2. AVB 2.0 U-BOOT SHELL COMMANDS
 -----------------------------------
@@ -61,6 +68,12 @@  CONFIG_LIBAVB=y
 CONFIG_AVB_VERIFY=y
 CONFIG_CMD_AVB=y
 
+In addtion optionally if storing rollback indexes in RPMB with help of
+OP-TEE:
+CONFIG_TEE=y
+CONFIG_OPTEE=y
+CONFIG_OPTEE_TA_AVB=y
+CONFIG_SUPPORT_EMMC_RPMB=y
 
 Then add `avb verify` invocation to your android boot sequence of commands,
 e.g.:
diff --git a/include/avb_verify.h b/include/avb_verify.h
index eaa60f5393ef..a532a2331aea 100644
--- a/include/avb_verify.h
+++ b/include/avb_verify.h
@@ -27,6 +27,10 @@  struct AvbOpsData {
 	struct AvbOps ops;
 	int mmc_dev;
 	enum avb_boot_state boot_state;
+#ifdef CONFIG_OPTEE_TA_AVB
+	struct udevice *tee;
+	u32 session;
+#endif
 };
 
 struct mmc_part {