Fix GetNextVariableName implementation

The NameSize has to be set if EFI_BUFFER_TOO_SMALL is returned,
according to the UEFI specs [1]. Moved the NameSize assignment to
align with this.

[1] https://uefi.org/specs/UEFI/2.10/08_Services_Runtime_Services.html#getnextvariablename

Signed-off-by: Bence Balogh <bence.balogh@arm.com>
Change-Id: I5b1238cfb88062778d4bba047ee1ba77e86fd3b1
diff --git a/components/service/uefi/smm_variable/backend/test/variable_store_tests.cpp b/components/service/uefi/smm_variable/backend/test/variable_store_tests.cpp
index 0d3450e..fd48f13 100644
--- a/components/service/uefi/smm_variable/backend/test/variable_store_tests.cpp
+++ b/components/service/uefi/smm_variable/backend/test/variable_store_tests.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2021-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -591,9 +591,15 @@
 	/* Enumerate store contents */
 	next_name->NameSize = sizeof(int16_t);
 	next_name->Name[0] = 0;
-
+	/* Check if the correct NameSize is returned if max_name_len is too small */
 	status = uefi_variable_store_get_next_variable_name(&m_uefi_variable_store, next_name,
-							    max_name_len, &total_len);
+							    0, &total_len);
+	UNSIGNED_LONGLONGS_EQUAL(EFI_BUFFER_TOO_SMALL, status);
+	UNSIGNED_LONGLONGS_EQUAL(sizeof(var_name_1), next_name->NameSize);
+
+	/* And then used the previously received next_name->NameSize as max_name_len */
+	status = uefi_variable_store_get_next_variable_name(&m_uefi_variable_store, next_name,
+							    next_name->NameSize, &total_len);
 	UNSIGNED_LONGLONGS_EQUAL(EFI_SUCCESS, status);
 	CHECK_TRUE(compare_variable_name(var_name_1, next_name->Name, next_name->NameSize));
 
diff --git a/components/service/uefi/smm_variable/backend/uefi_variable_store.c b/components/service/uefi/smm_variable/backend/uefi_variable_store.c
index 8dad12d..5b46c13 100644
--- a/components/service/uefi/smm_variable/backend/uefi_variable_store.c
+++ b/components/service/uefi/smm_variable/backend/uefi_variable_store.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -418,9 +418,15 @@
 			&context->variable_index, &cur->Guid, cur->NameSize, cur->Name, &status);
 
 		if (info && (status == EFI_SUCCESS)) {
+			/* The NameSize has to be set in every case according to the UEFI specs.
+			 * In case of EFI_BUFFER_TOO_SMALL it has to reflect the size of buffer
+			 * needed.
+			 */
+			cur->NameSize = info->metadata.name_size;
+			*total_length = sizeof(SMM_VARIABLE_COMMUNICATE_GET_NEXT_VARIABLE_NAME);
+
 			if (info->metadata.name_size <= max_name_len) {
 				cur->Guid = info->metadata.guid;
-				cur->NameSize = info->metadata.name_size;
 				memcpy(cur->Name, info->metadata.name, info->metadata.name_size);
 
 				*total_length =
@@ -452,7 +458,8 @@
 	if (status != EFI_SUCCESS) {
 		memset(cur->Name, 0, max_name_len);
 		memset(&cur->Guid, 0, sizeof(EFI_GUID));
-		cur->NameSize = 0;
+		if (status != EFI_BUFFER_TOO_SMALL)
+			cur->NameSize = 0;
 	}
 
 	return status;