[v2] efi_loader: Enable run-time variable support for tee based variables

diff --git a/include/efi_variable.h b/include/efi_variable.h
index 2c629e4dca92..60491cb64025 100644
--- a/include/efi_variable.h
+++ b/include/efi_variable.h
@@ -142,6 +142,22 @@  struct efi_var_file {
  */
 efi_status_t efi_var_to_file(void);
 
+/**
+ * efi_var_collect() - collect variables in buffer
+ *
+ * A buffer is allocated and filled with variables in a format ready to be
+ * written to disk.
+ *
+ * @bufp:		pointer to pointer of buffer with collected variables
+ * @lenp:		pointer to length of buffer
+ * @check_attr_mask:	bitmask with required attributes of variables to be collected.
+ *                      variables are only collected if all of the required
+ *                      attributes are set.
+ * Return:		status code
+ */
+efi_status_t __maybe_unused efi_var_collect(struct efi_var_file **bufp, loff_t *lenp,
+					    u32 check_attr_mask);
+
 /**
  * efi_var_restore() - restore EFI variables from buffer
  *
@@ -233,4 +249,62 @@  efi_status_t efi_init_secure_state(void);
  */
 enum efi_auth_var_type efi_auth_var_get_type(u16 *name, const efi_guid_t *guid);
 
+/**
+ * efi_get_next_variable_name_mem() - Runtime common code across efi variable
+ *                                    implementations for GetNextVariable()
+ *                                    from the cached memory copy
+ * @variable_name_size:	size of variable_name buffer in byte
+ * @variable_name:	name of uefi variable's name in u16
+ * @vendor:		vendor's guid
+ *
+ * Return: status code
+ */
+efi_status_t __efi_runtime
+efi_get_next_variable_name_mem(efi_uintn_t *variable_name_size, u16 *variable_name,
+			       efi_guid_t *vendor);
+/**
+ * efi_get_variable_mem() - Runtime common code across efi variable
+ *                          implementations for GetVariable() from
+ *                          the cached memory copy
+ *
+ * @variable_name:	name of the variable
+ * @vendor:		vendor GUID
+ * @attributes:		attributes of the variable
+ * @data_size:		size of the buffer to which the variable value is copied
+ * @data:		buffer to which the variable value is copied
+ * @timep:		authentication time (seconds since start of epoch)
+ * Return:		status code
+
+ */
+efi_status_t __efi_runtime
+efi_get_variable_mem(u16 *variable_name, const efi_guid_t *vendor, u32 *attributes,
+		     efi_uintn_t *data_size, void *data, u64 *timep);
+
+/**
+ * efi_get_variable_runtime() - runtime implementation of GetVariable()
+ *
+ * @variable_name:	name of the variable
+ * @guid:		vendor GUID
+ * @attributes:		attributes of the variable
+ * @data_size:		size of the buffer to which the variable value is copied
+ * @data:		buffer to which the variable value is copied
+ * Return:		status code
+ */
+efi_status_t __efi_runtime EFIAPI
+efi_get_variable_runtime(u16 *variable_name, const efi_guid_t *guid,
+			 u32 *attributes, efi_uintn_t *data_size, void *data);
+
+/**
+ * efi_get_next_variable_name_runtime() - runtime implementation of
+ *					  GetNextVariable()
+ *
+ * @variable_name_size:	size of variable_name buffer in byte
+ * @variable_name:	name of uefi variable's name in u16
+ * @guid:		vendor's guid
+ * Return:              status code
+ */
+efi_status_t __efi_runtime EFIAPI
+efi_get_next_variable_name_runtime(efi_uintn_t *variable_name_size,
+				   u16 *variable_name, efi_guid_t *guid);
+
 #endif
diff --git a/lib/efi_loader/Makefile b/lib/efi_loader/Makefile
index 441ac9432e99..9bad1d159b03 100644
--- a/lib/efi_loader/Makefile
+++ b/lib/efi_loader/Makefile
@@ -37,11 +37,11 @@  obj-y += efi_setup.o
 obj-$(CONFIG_EFI_UNICODE_COLLATION_PROTOCOL2) += efi_unicode_collation.o
 obj-y += efi_var_common.o
 obj-y += efi_var_mem.o
+obj-y += efi_var_file.o
 ifeq ($(CONFIG_EFI_MM_COMM_TEE),y)
 obj-y += efi_variable_tee.o
 else
 obj-y += efi_variable.o
-obj-y += efi_var_file.o
 obj-$(CONFIG_EFI_VARIABLES_PRESEED) += efi_var_seed.o
 endif
 obj-y += efi_watchdog.o
diff --git a/lib/efi_loader/efi_var_common.c b/lib/efi_loader/efi_var_common.c
index ee2e67bc8cbf..453cbce5c83a 100644
--- a/lib/efi_loader/efi_var_common.c
+++ b/lib/efi_loader/efi_var_common.c
@@ -166,6 +166,28 @@  efi_status_t EFIAPI efi_query_variable_info(
 	return EFI_EXIT(ret);
 }
 
+efi_status_t __efi_runtime EFIAPI
+efi_get_variable_runtime(u16 *variable_name, const efi_guid_t *guid,
+			 u32 *attributes, efi_uintn_t *data_size, void *data)
+{
+	efi_status_t ret;
+
+	ret = efi_get_variable_mem(variable_name, guid, attributes, data_size, data, NULL);
+
+	/* Remove EFI_VARIABLE_READ_ONLY flag */
+	if (attributes)
+		*attributes &= EFI_VARIABLE_MASK;
+
+	return ret;
+}
+
+efi_status_t __efi_runtime EFIAPI
+efi_get_next_variable_name_runtime(efi_uintn_t *variable_name_size,
+				   u16 *variable_name, efi_guid_t *guid)
+{
+	return efi_get_next_variable_name_mem(variable_name_size, variable_name, guid);
+}
+
 /**
  * efi_set_secure_state - modify secure boot state variables
  * @secure_boot:	value of SecureBoot
diff --git a/lib/efi_loader/efi_var_file.c b/lib/efi_loader/efi_var_file.c
index 6f9d76f2a2d5..b171d2d1a8f7 100644
--- a/lib/efi_loader/efi_var_file.c
+++ b/lib/efi_loader/efi_var_file.c
@@ -46,18 +46,8 @@  static efi_status_t __maybe_unused efi_set_blk_dev_to_system_partition(void)
 	return EFI_SUCCESS;
 }
 
-/**
- * efi_var_collect() - collect non-volatile variables in buffer
- *
- * A buffer is allocated and filled with all non-volatile variables in a
- * format ready to be written to disk.
- *
- * @bufp:	pointer to pointer of buffer with collected variables
- * @lenp:	pointer to length of buffer
- * Return:	status code
- */
-static efi_status_t __maybe_unused efi_var_collect(struct efi_var_file **bufp,
-						   loff_t *lenp)
+efi_status_t __maybe_unused efi_var_collect(struct efi_var_file **bufp, loff_t *lenp,
+					    u32 check_attr_mask)
 {
 	size_t len = EFI_VAR_BUF_SIZE;
 	struct efi_var_file *buf;
@@ -102,11 +92,10 @@  static efi_status_t __maybe_unused efi_var_collect(struct efi_var_file **bufp,
 			free(buf);
 			return ret;
 		}
-		if (!(var->attr & EFI_VARIABLE_NON_VOLATILE))
-			continue;
-		var->length = data_length;
-		var = (struct efi_var_entry *)
-		      ALIGN((uintptr_t)data + data_length, 8);
+		if ((var->attr & check_attr_mask) == check_attr_mask) {
+			var->length = data_length;
+			var = (struct efi_var_entry *)ALIGN((uintptr_t)data + data_length, 8);
+		}
 	}
 
 	buf->reserved = 0;
@@ -137,7 +126,7 @@  efi_status_t efi_var_to_file(void)
 	loff_t actlen;
 	int r;
 
-	ret = efi_var_collect(&buf, &len);
+	ret = efi_var_collect(&buf, &len, EFI_VARIABLE_NON_VOLATILE);
 	if (ret != EFI_SUCCESS)
 		goto error;
 
diff --git a/lib/efi_loader/efi_var_mem.c b/lib/efi_loader/efi_var_mem.c
index 7a2dba7dc263..fd97d5b56300 100644
--- a/lib/efi_loader/efi_var_mem.c
+++ b/lib/efi_loader/efi_var_mem.c
@@ -10,7 +10,7 @@ 
 #include <efi_variable.h>
 #include <u-boot/crc.h>
 
-static struct efi_var_file __efi_runtime_data *efi_var_buf;
+struct efi_var_file __efi_runtime_data *efi_var_buf;
 static struct efi_var_entry __efi_runtime_data *efi_current_var;
 
 /**
@@ -264,3 +264,71 @@  efi_status_t efi_var_mem_init(void)
 		return ret;
 	return ret;
 }
+
+efi_status_t __efi_runtime
+efi_get_variable_mem(u16 *variable_name, const efi_guid_t *vendor, u32 *attributes,
+		     efi_uintn_t *data_size, void *data, u64 *timep)
+{
+	efi_uintn_t old_size;
+	struct efi_var_entry *var;
+	u16 *pdata;
+
+	if (!variable_name || !vendor || !data_size)
+		return EFI_INVALID_PARAMETER;
+	var = efi_var_mem_find(vendor, variable_name, NULL);
+	if (!var)
+		return EFI_NOT_FOUND;
+
+	if (attributes)
+		*attributes = var->attr;
+	if (timep)
+		*timep = var->time;
+
+	old_size = *data_size;
+	*data_size = var->length;
+	if (old_size < var->length)
+		return EFI_BUFFER_TOO_SMALL;
+
+	if (!data)
+		return EFI_INVALID_PARAMETER;
+
+	for (pdata = var->name; *pdata; ++pdata)
+		;
+	++pdata;
+
+	efi_memcpy_runtime(data, pdata, var->length);
+
+	return EFI_SUCCESS;
+}
+
+efi_status_t __efi_runtime
+efi_get_next_variable_name_mem(efi_uintn_t *variable_name_size, u16 *variable_name,
+			       efi_guid_t *vendor)
+{
+	struct efi_var_entry *var;
+	efi_uintn_t old_size;
+	u16 *pdata;
+
+	if (!variable_name_size || !variable_name || !vendor)
+		return EFI_INVALID_PARAMETER;
+
+	efi_var_mem_find(vendor, variable_name, &var);
+
+	if (!var)
+		return EFI_NOT_FOUND;
+
+	for (pdata = var->name; *pdata; ++pdata)
+		;
+	++pdata;
+
+	old_size = *variable_name_size;
+	*variable_name_size = (uintptr_t)pdata - (uintptr_t)var->name;
+
+	if (old_size < *variable_name_size)
+		return EFI_BUFFER_TOO_SMALL;
+
+	efi_memcpy_runtime(variable_name, var->name, *variable_name_size);
+	efi_memcpy_runtime(vendor, &var->guid, sizeof(efi_guid_t));
+
+	return EFI_SUCCESS;
+}
diff --git a/lib/efi_loader/efi_variable.c b/lib/efi_loader/efi_variable.c
index 39a848290380..e509d6dbf0cf 100644
--- a/lib/efi_loader/efi_variable.c
+++ b/lib/efi_loader/efi_variable.c
@@ -282,68 +282,14 @@  efi_get_variable_int(u16 *variable_name, const efi_guid_t *vendor,
 		     u32 *attributes, efi_uintn_t *data_size, void *data,
 		     u64 *timep)
 {
-	efi_uintn_t old_size;
-	struct efi_var_entry *var;
-	u16 *pdata;
-
-	if (!variable_name || !vendor || !data_size)
-		return EFI_INVALID_PARAMETER;
-	var = efi_var_mem_find(vendor, variable_name, NULL);
-	if (!var)
-		return EFI_NOT_FOUND;
-
-	if (attributes)
-		*attributes = var->attr;
-	if (timep)
-		*timep = var->time;
-
-	old_size = *data_size;
-	*data_size = var->length;
-	if (old_size < var->length)
-		return EFI_BUFFER_TOO_SMALL;
-
-	if (!data)
-		return EFI_INVALID_PARAMETER;
-
-	for (pdata = var->name; *pdata; ++pdata)
-		;
-	++pdata;
-
-	efi_memcpy_runtime(data, pdata, var->length);
-
-	return EFI_SUCCESS;
+	return efi_get_variable_mem(variable_name, vendor, attributes, data_size, data, timep);
 }
 
 efi_status_t __efi_runtime
 efi_get_next_variable_name_int(efi_uintn_t *variable_name_size,
 			       u16 *variable_name, efi_guid_t *vendor)
 {
-	struct efi_var_entry *var;
-	efi_uintn_t old_size;
-	u16 *pdata;
-
-	if (!variable_name_size || !variable_name || !vendor)
-		return EFI_INVALID_PARAMETER;
-
-	efi_var_mem_find(vendor, variable_name, &var);
-
-	if (!var)
-		return EFI_NOT_FOUND;
-
-	for (pdata = var->name; *pdata; ++pdata)
-		;
-	++pdata;
-
-	old_size = *variable_name_size;
-	*variable_name_size = (uintptr_t)pdata - (uintptr_t)var->name;
-
-	if (old_size < *variable_name_size)
-		return EFI_BUFFER_TOO_SMALL;
-
-	efi_memcpy_runtime(variable_name, var->name, *variable_name_size);
-	efi_memcpy_runtime(vendor, &var->guid, sizeof(efi_guid_t));
-
-	return EFI_SUCCESS;
+	return efi_get_next_variable_name_mem(variable_name_size, variable_name, vendor);
 }
 
 efi_status_t efi_set_variable_int(u16 *variable_name, const efi_guid_t *vendor,
@@ -504,49 +450,6 @@  efi_status_t __efi_runtime EFIAPI efi_query_variable_info_runtime(
 	return EFI_UNSUPPORTED;
 }
 
-/**
- * efi_get_variable_runtime() - runtime implementation of GetVariable()
- *
- * @variable_name:	name of the variable
- * @vendor:		vendor GUID
- * @attributes:		attributes of the variable
- * @data_size:		size of the buffer to which the variable value is copied
- * @data:		buffer to which the variable value is copied
- * Return:		status code
- */
-static efi_status_t __efi_runtime EFIAPI
-efi_get_variable_runtime(u16 *variable_name, const efi_guid_t *vendor,
-			 u32 *attributes, efi_uintn_t *data_size, void *data)
-{
-	efi_status_t ret;
-
-	ret = efi_get_variable_int(variable_name, vendor, attributes,
-				   data_size, data, NULL);
-
-	/* Remove EFI_VARIABLE_READ_ONLY flag */
-	if (attributes)
-		*attributes &= EFI_VARIABLE_MASK;
-
-	return ret;
-}
-
-/**
- * efi_get_next_variable_name_runtime() - runtime implementation of
- *					  GetNextVariable()
- *
- * @variable_name_size:	size of variable_name buffer in byte
- * @variable_name:	name of uefi variable's name in u16
- * @vendor:		vendor's guid
- * Return: status code
- */
-static efi_status_t __efi_runtime EFIAPI
-efi_get_next_variable_name_runtime(efi_uintn_t *variable_name_size,
-				   u16 *variable_name, efi_guid_t *vendor)
-{
-	return efi_get_next_variable_name_int(variable_name_size, variable_name,
-					      vendor);
-}
-
 /**
  * efi_set_variable_runtime() - runtime implementation of SetVariable()
  *
diff --git a/lib/efi_loader/efi_variable_tee.c b/lib/efi_loader/efi_variable_tee.c
index 9920351a403e..8a3ff8c59386 100644
--- a/lib/efi_loader/efi_variable_tee.c
+++ b/lib/efi_loader/efi_variable_tee.c
@@ -15,6 +15,8 @@ 
 #include <malloc.h>
 #include <mm_communication.h>
 
+#define OPTEE_PAGE_SIZE BIT(12)
+extern struct efi_var_file __efi_runtime_data *efi_var_buf;
 static efi_uintn_t max_buffer_size;	/* comm + var + func + data */
 static efi_uintn_t max_payload_size;	/* func + data */
 
@@ -237,8 +239,31 @@  efi_status_t EFIAPI get_max_payload(efi_uintn_t *size)
 	if (ret != EFI_SUCCESS)
 		goto out;
 
+	/* Make sure the buffer is big enough for storing variables i.e en-US. */
+	if (var_payload->size < MM_VARIABLE_ACCESS_HEADER_SIZE + 0x20) {
+		ret = EFI_DEVICE_ERROR;
+		goto out;
+	}
 	*size = var_payload->size;
-
+	/*
+	 * Although the max payload is configurable on StMM, we only share a single
+	 * page from OP-TEE for the non-secure buffer used to communicate with StMM.
+	 * Since OP-TEE will reject to map anything bigger than that, make sure we are
+	 * in bounds.
+	 */
+	if (*size > OPTEE_PAGE_SIZE)
+		*size = OPTEE_PAGE_SIZE - MM_COMMUNICATE_HEADER_SIZE  -
+			MM_VARIABLE_COMMUNICATE_SIZE;
+	/*
+	 * There seems to be a bug in EDK2 miscalculating the boundaries and size
+	 * checks, so deduct 2 more bytes to fulfill this requirement. Fix it up here
+	 * to ensure backwards compatibility with older versions.
+	 * Ref: StandaloneMmPkg/Drivers/StandaloneMmCpu/AArch64/EventHandle.c
+	 * sizeof (EFI_MM_COMMUNICATE_HEADER) instead the size minus the flexible
+	 * array member
+	 * This is guaranteed to be > 2 due to checks on the beggining.
+	 */
+	*size -= 2;
 out:
 	free(comm_buf);
 	return ret;
@@ -592,39 +617,6 @@  out:
 	return ret;
 }
 
-/**
- * efi_get_variable_runtime() - runtime implementation of GetVariable()
- *
- * @variable_name:	name of the variable
- * @guid:		vendor GUID
- * @attributes:		attributes of the variable
- * @data_size:		size of the buffer to which the variable value is copied
- * @data:		buffer to which the variable value is copied
- * Return:		status code
- */
-static efi_status_t __efi_runtime EFIAPI
-efi_get_variable_runtime(u16 *variable_name, const efi_guid_t *guid,
-			 u32 *attributes, efi_uintn_t *data_size, void *data)
-{
-	return EFI_UNSUPPORTED;
-}
-
-/**
- * efi_get_next_variable_name_runtime() - runtime implementation of
- *					  GetNextVariable()
- *
- * @variable_name_size:	size of variable_name buffer in byte
- * @variable_name:	name of uefi variable's name in u16
- * @guid:		vendor's guid
- * Return:              status code
- */
-static efi_status_t __efi_runtime EFIAPI
-efi_get_next_variable_name_runtime(efi_uintn_t *variable_name_size,
-				   u16 *variable_name, efi_guid_t *guid)
-{
-	return EFI_UNSUPPORTED;
-}
-
 /**
  * efi_query_variable_info() - get information about EFI variables
  *
@@ -674,8 +666,10 @@  efi_set_variable_runtime(u16 *variable_name, const efi_guid_t *guid,
  */
 void efi_variables_boot_exit_notify(void)
 {
-	u8 *comm_buf;
 	efi_status_t ret;
+	u8 *comm_buf;
+	loff_t len;
+	struct efi_var_file *var_buf;
 
 	comm_buf = setup_mm_hdr(NULL, 0,
 				SMM_VARIABLE_FUNCTION_EXIT_BOOT_SERVICE, &ret);
@@ -688,6 +682,18 @@  void efi_variables_boot_exit_notify(void)
 		log_err("Unable to notify StMM for ExitBootServices\n");
 	free(comm_buf);
 
+	/*
+	 * Populate the list for runtime variables.
+	 * asking EFI_VARIABLE_RUNTIME_ACCESS is redundant, since
+	 * efi_var_mem_notify_exit_boot_services will clean those, but that's fine
+	 */
+	ret = efi_var_collect(&var_buf, &len, EFI_VARIABLE_RUNTIME_ACCESS);
+	if (ret != EFI_SUCCESS)
+		log_err("Can't populate EFI variables. No runtime variables will be available\n");
+	else
+		memcpy(efi_var_buf, var_buf, len);
+	free(var_buf);
+
 	/* Update runtime service table */
 	efi_runtime_services.query_variable_info =
 			efi_query_variable_info_runtime;
@@ -707,6 +713,11 @@  efi_status_t efi_init_variables(void)
 {
 	efi_status_t ret;
 
+	/* Create a cached copy of the variables that will be enabled on ExitBootServices() */
+	ret = efi_var_mem_init();
+	if (ret != EFI_SUCCESS)
+		return ret;
+
 	ret = get_max_payload(&max_payload_size);
 	if (ret != EFI_SUCCESS)
 		return ret;