feat(ff-a): introduce count flag to ffa_partition_info_get

The partition count flag controls whether to return partition
info descriptors on a call to FFA_PARTITION_INFO_GET. This
patch introduces this functionality.

Signed-off-by: Daniel Boulby <daniel.boulby@arm.com>
Change-Id: Ib4a13c10e107daf726e5d6ec0ce6bd0d451c1b0f
diff --git a/src/api.c b/src/api.c
index de2bf5b..60dc65d 100644
--- a/src/api.c
+++ b/src/api.c
@@ -356,35 +356,57 @@
 }
 
 struct ffa_value api_ffa_partition_info_get(struct vcpu *current,
-					    const struct ffa_uuid *uuid)
+					    const struct ffa_uuid *uuid,
+					    const uint32_t flags)
 {
 	struct vm *current_vm = current->vm;
 	struct vm_locked current_vm_locked;
 	ffa_vm_count_t vm_count = 0;
+	bool count_flag = (flags && FFA_PARTITION_COUNT_FLAG_MASK) ==
+			  FFA_PARTITION_COUNT_FLAG;
 	bool uuid_is_null = ffa_uuid_is_null(uuid);
 	struct ffa_value ret;
 	uint32_t size;
 	struct ffa_partition_info partitions[2 * MAX_VMS];
 
+	/* Bits 31:1 Must Be Zero */
+	if ((flags & ~FFA_PARTITION_COUNT_FLAG) != 0) {
+		return ffa_error(FFA_INVALID_PARAMETERS);
+	}
+
 	/*
-	 * Iterate through the VMs to find the ones with a matching UUID.
-	 * A Null UUID retrieves information for all VMs.
+	 * No need to count if we are returning the number of paritions as we
+	 * already know this.
 	 */
-	for (uint16_t index = 0; index < vm_get_count(); ++index) {
-		struct vm *vm = vm_find_index(index);
+	if (uuid_is_null && count_flag) {
+		vm_count = vm_get_count();
+	} else {
+		/*
+		 * Iterate through the VMs to find the ones with a matching
+		 * UUID. A Null UUID retrieves information for all VMs.
+		 */
+		for (uint16_t index = 0; index < vm_get_count(); ++index) {
+			struct vm *vm = vm_find_index(index);
 
-		if (uuid_is_null || ffa_uuid_equal(uuid, &vm->uuid)) {
-			partitions[vm_count].vm_id = vm->id;
-			partitions[vm_count].vcpu_count = vm->vcpu_count;
-			partitions[vm_count].properties =
-				plat_ffa_partition_properties(current_vm->id,
-							      vm);
-			partitions[vm_count].properties |=
-				vm_are_notifications_enabled(vm)
-					? FFA_PARTITION_NOTIFICATION
-					: 0;
+			if (uuid_is_null || ffa_uuid_equal(uuid, &vm->uuid)) {
+				uint16_t array_index = vm_count;
 
-			++vm_count;
+				++vm_count;
+				if (count_flag) {
+					continue;
+				}
+
+				partitions[array_index].vm_id = vm->id;
+				partitions[array_index].vcpu_count =
+					vm->vcpu_count;
+				partitions[array_index].properties =
+					plat_ffa_partition_properties(
+						current_vm->id, vm);
+				partitions[array_index].properties |=
+					vm_are_notifications_enabled(vm)
+						? FFA_PARTITION_NOTIFICATION
+						: 0;
+			}
 		}
 	}
 
@@ -395,13 +417,13 @@
 	 * When running the Hypervisor:
 	 * - If UUID is Null the Hypervisor forwards the query to the SPMC for
 	 * it to fill with secure partitions information.
-	 * - If UUID is non-Null vm_count may be zero because the UUID  matches
+	 * - If UUID is non-Null vm_count may be zero because the UUID matches
 	 * a secure partition and the query is forwarded to the SPMC.
 	 * When running the SPMC:
 	 * - If UUID is non-Null and vm_count is zero it means there is no such
 	 * partition identified in the system.
 	 */
-	plat_ffa_partition_info_get_forward(uuid, partitions, &vm_count);
+	plat_ffa_partition_info_get_forward(uuid, flags, partitions, &vm_count);
 
 	/*
 	 * Unrecognized UUID: does not match any of the VMs (or SPs)
@@ -411,6 +433,15 @@
 		return ffa_error(FFA_INVALID_PARAMETERS);
 	}
 
+	/*
+	 * If the count flag is set we don't need to return the partition info
+	 * descriptors.
+	 */
+	if (count_flag) {
+		return (struct ffa_value){.func = FFA_SUCCESS_32,
+					  .arg2 = vm_count};
+	}
+
 	size = vm_count * sizeof(partitions[0]);
 	if (size > FFA_MSG_PAYLOAD_MAX) {
 		dlog_error(
diff --git a/src/arch/aarch64/hypervisor/handler.c b/src/arch/aarch64/hypervisor/handler.c
index 9bedee8..e3dcf2b 100644
--- a/src/arch/aarch64/hypervisor/handler.c
+++ b/src/arch/aarch64/hypervisor/handler.c
@@ -492,7 +492,7 @@
 
 		ffa_uuid_init(args->arg1, args->arg2, args->arg3, args->arg4,
 			      &uuid);
-		*args = api_ffa_partition_info_get(current, &uuid);
+		*args = api_ffa_partition_info_get(current, &uuid, args->arg5);
 		return true;
 	}
 	case FFA_ID_GET_32:
diff --git a/src/arch/aarch64/plat/ffa/absent.c b/src/arch/aarch64/plat/ffa/absent.c
index 7750cac..87a552e 100644
--- a/src/arch/aarch64/plat/ffa/absent.c
+++ b/src/arch/aarch64/plat/ffa/absent.c
@@ -298,9 +298,11 @@
 }
 
 bool plat_ffa_partition_info_get_forward(const struct ffa_uuid *uuid,
+					 const uint32_t flags,
 					 const ffa_vm_count_t *ret_count)
 {
 	(void)uuid;
+	(void)flags;
 	(void)ret_count;
 
 	return false;
diff --git a/src/arch/aarch64/plat/ffa/hypervisor.c b/src/arch/aarch64/plat/ffa/hypervisor.c
index b8489e3..d38a6b0 100644
--- a/src/arch/aarch64/plat/ffa/hypervisor.c
+++ b/src/arch/aarch64/plat/ffa/hypervisor.c
@@ -580,6 +580,7 @@
  * Emits FFA_PARTITION_INFO_GET from Hypervisor to SPMC if allowed.
  */
 void plat_ffa_partition_info_get_forward(const struct ffa_uuid *uuid,
+					 const uint32_t flags,
 					 struct ffa_partition_info *partitions,
 					 ffa_vm_count_t *ret_count)
 {
@@ -605,7 +606,8 @@
 				   .arg1 = uuid->uuid[0],
 				   .arg2 = uuid->uuid[1],
 				   .arg3 = uuid->uuid[2],
-				   .arg4 = uuid->uuid[3]});
+				   .arg4 = uuid->uuid[3],
+				   .arg5 = flags});
 	if (ffa_func_id(ret) != FFA_SUCCESS_32) {
 		dlog_verbose(
 			"Failed forwarding FFA_PARTITION_INFO_GET to "
@@ -619,16 +621,22 @@
 		return;
 	}
 
-	tee_partitions = (struct ffa_partition_info *)tee->mailbox.send;
-	for (ffa_vm_count_t index = 0; index < tee_partitions_count; index++) {
-		partitions[vm_count] = tee_partitions[index];
-		++vm_count;
-	}
+	if ((flags && FFA_PARTITION_COUNT_FLAG_MASK) ==
+	    FFA_PARTITION_COUNT_FLAG) {
+		vm_count += tee_partitions_count;
+	} else {
+		tee_partitions = (struct ffa_partition_info *)tee->mailbox.send;
+		for (ffa_vm_count_t index = 0; index < tee_partitions_count;
+		     index++) {
+			partitions[vm_count] = tee_partitions[index];
+			++vm_count;
+		}
 
-	/* Release the RX buffer. */
-	ret = arch_other_world_call(
-		(struct ffa_value){.func = FFA_RX_RELEASE_32});
-	CHECK(ret.func == FFA_SUCCESS_32);
+		/* Release the RX buffer. */
+		ret = arch_other_world_call(
+			(struct ffa_value){.func = FFA_RX_RELEASE_32});
+		CHECK(ret.func == FFA_SUCCESS_32);
+	}
 
 	*ret_count = vm_count;
 }
diff --git a/src/arch/aarch64/plat/ffa/spmc.c b/src/arch/aarch64/plat/ffa/spmc.c
index ec75228..08d8825 100644
--- a/src/arch/aarch64/plat/ffa/spmc.c
+++ b/src/arch/aarch64/plat/ffa/spmc.c
@@ -1220,11 +1220,13 @@
 /** Forward helper for FFA_PARTITION_INFO_GET. */
 void plat_ffa_partition_info_get_forward(  // NOLINTNEXTLINE
 	const struct ffa_uuid *uuid,	   // NOLINTNEXTLINE
+	const uint32_t flags,		   // NOLINTNEXTLINE
 	struct ffa_partition_info *partitions, ffa_vm_count_t *ret_count)
 {
 	/* The SPMC does not forward FFA_PARTITION_INFO_GET. */
 
 	(void)uuid;
+	(void)flags;
 	(void)partitions;
 	(void)ret_count;
 }
diff --git a/src/arch/fake/hypervisor/ffa.c b/src/arch/fake/hypervisor/ffa.c
index cadebee..13c67e6 100644
--- a/src/arch/fake/hypervisor/ffa.c
+++ b/src/arch/fake/hypervisor/ffa.c
@@ -286,9 +286,11 @@
 
 void plat_ffa_partition_info_get_forward(  // NOLINTNEXTLINE
 	const struct ffa_uuid *uuid,	   // NOLINTNEXTLINE
+	const uint32_t flags,		   // NOLINTNEXTLINE
 	struct ffa_partition_info *partitions, ffa_vm_count_t *ret_count)
 {
 	(void)uuid;
+	(void)flags;
 	(void)partitions;
 	(void)ret_count;
 }