Allow sp_discovery to return more than one SP

sp_discovery_partition_info_get() could only return a single SP when
passing a UUID. Change it to enable returning multiple SPs when there
are multiple SPs with the same UUID.

Signed-off-by: Jelle Sels <jelle.sels@arm.com>
Signed-off-by: Imre Kis <imre.kis@arm.com>
Change-Id: I167432152395a62db766b18d60c3487bb814b990
diff --git a/components/messaging/ffa/libsp/include/sp_discovery.h b/components/messaging/ffa/libsp/include/sp_discovery.h
index a85e946..977c875 100644
--- a/components/messaging/ffa/libsp/include/sp_discovery.h
+++ b/components/messaging/ffa/libsp/include/sp_discovery.h
@@ -1,6 +1,6 @@
 /* SPDX-License-Identifier: BSD-3-Clause */
 /*
- * Copyright (c) 2020, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2020-2023, Arm Limited and Contributors. All rights reserved.
  */
 
 #ifndef LIBSP_INCLUDE_SP_DISCOVERY_H_
@@ -60,13 +60,18 @@
 /**
  * @brief       Queries the information about a partition by its UUID.
  *
- * @param[in]   uuid  The UUID of the partition
- * @param[out]  info  The partition information
+ * @param[in]      uuid  The UUID of the partition
+ * @param[out]     info  The partition information
+ * @param[in,out]  count As an input value it specifies the count of partition
+ *                       info structures that would fit into the output buffer.
+ *                       As an output it indicates the count of the valid
+ *                       entries in the buffer.
  *
- * @return      The SP API result
+ * @return         The SP API result
  */
 sp_result sp_discovery_partition_info_get(const struct sp_uuid *uuid,
-					  struct sp_partition_info *info);
+					  struct sp_partition_info *info,
+					  uint32_t *count);
 
 /**
  * @brief          Queries partition information of all partitions.
diff --git a/components/messaging/ffa/libsp/mock/mock_sp_discovery.cpp b/components/messaging/ffa/libsp/mock/mock_sp_discovery.cpp
index 47f4ef7..6f18300 100644
--- a/components/messaging/ffa/libsp/mock/mock_sp_discovery.cpp
+++ b/components/messaging/ffa/libsp/mock/mock_sp_discovery.cpp
@@ -66,6 +66,8 @@
 
 void expect_sp_discovery_partition_info_get(const struct sp_uuid *uuid,
 					  const struct sp_partition_info *info,
+					  uint32_t in_count,
+					  const uint32_t *out_count,
 					  sp_result result)
 {
 	mock()
@@ -73,17 +75,22 @@
 		.withMemoryBufferParameter("uuid", (const unsigned char *)uuid,
 					   sizeof(*uuid))
 		.withOutputParameterReturning("info", info, sizeof(*info))
+		.withUnsignedIntParameter("in_count", in_count)
+		.withOutputParameterReturning("out_count", out_count, sizeof(*out_count))
 		.andReturnValue(result);
 }
 
 sp_result sp_discovery_partition_info_get(const struct sp_uuid *uuid,
-					  struct sp_partition_info *info)
+					  struct sp_partition_info *info,
+					  uint32_t *count)
 {
 	return mock()
 		.actualCall("sp_discovery_partition_info_get")
 		.withMemoryBufferParameter("uuid", (const unsigned char *)uuid,
 					   sizeof(*uuid))
 		.withOutputParameter("info", info)
+		.withUnsignedIntParameter("in_count", *count)
+		.withOutputParameter("out_count", count)
 		.returnIntValue();
 }
 
diff --git a/components/messaging/ffa/libsp/mock/mock_sp_discovery.h b/components/messaging/ffa/libsp/mock/mock_sp_discovery.h
index a71ce18..67cb118 100644
--- a/components/messaging/ffa/libsp/mock/mock_sp_discovery.h
+++ b/components/messaging/ffa/libsp/mock/mock_sp_discovery.h
@@ -24,6 +24,8 @@
 
 void expect_sp_discovery_partition_info_get(const struct sp_uuid *uuid,
 					    const struct sp_partition_info *info,
+					    uint32_t in_count,
+					    const uint32_t *out_count,
 					    sp_result result);
 
 void expect_sp_discovery_partition_info_get_all(const struct sp_partition_info info[],
diff --git a/components/messaging/ffa/libsp/mock/test/test_mock_sp_discovery.cpp b/components/messaging/ffa/libsp/mock/test/test_mock_sp_discovery.cpp
index bb4bf07..5b6c538 100644
--- a/components/messaging/ffa/libsp/mock/test/test_mock_sp_discovery.cpp
+++ b/components/messaging/ffa/libsp/mock/test/test_mock_sp_discovery.cpp
@@ -76,11 +76,17 @@
 
 	struct sp_uuid uuid = expected_uuid;
 	struct sp_partition_info info = {0};
+	uint32_t in_count = 16;
+	uint32_t expected_out_count = 8;
+	uint32_t out_count = in_count;
+
 
 	expect_sp_discovery_partition_info_get(&expected_uuid, &expected_info,
+					       in_count, &expected_out_count,
 					       result);
-	LONGS_EQUAL(result, sp_discovery_partition_info_get(&uuid, &info));
+	LONGS_EQUAL(result, sp_discovery_partition_info_get(&uuid, &info, &out_count));
 	MEMCMP_EQUAL(&expected_info, &info, sizeof(&expected_info));
+	UNSIGNED_LONGS_EQUAL(expected_out_count, out_count);
 }
 
 TEST(mock_sp_discovery, sp_discovery_partition_info_get_all)
diff --git a/components/messaging/ffa/libsp/sp_discovery.c b/components/messaging/ffa/libsp/sp_discovery.c
index 15cb0be..7f92b6b 100644
--- a/components/messaging/ffa/libsp/sp_discovery.c
+++ b/components/messaging/ffa/libsp/sp_discovery.c
@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: BSD-3-Clause
 /*
- * Copyright (c) 2020-2021, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2020-2023, Arm Limited and Contributors. All rights reserved.
  */
 
 #include "sp_discovery.h"
@@ -38,74 +38,7 @@
 	return SP_RESULT_FFA(ffa_res);
 }
 
-static sp_result
-partition_info_get(const struct sp_uuid *uuid,
-		   const struct ffa_partition_information **info,
-		   uint32_t *count)
-{
-	const void *buffer = NULL;
-	size_t buffer_size = 0;
-	struct ffa_uuid ffa_uuid = { 0 };
-	sp_result sp_res = SP_RESULT_OK;
-	ffa_result ffa_res = FFA_OK;
-
-	sp_res = sp_rxtx_buffer_rx_get(&buffer, &buffer_size);
-	if (sp_res != SP_RESULT_OK) {
-		*count = UINT32_C(0);
-		return sp_res;
-	}
-
-	/* Safely convert to FF-A UUID format */
-	memcpy(&ffa_uuid.uuid, uuid->uuid, sizeof(ffa_uuid.uuid));
-
-	ffa_res = ffa_partition_info_get(&ffa_uuid, count);
-	if (ffa_res != FFA_OK) {
-		*count = UINT32_C(0);
-		return SP_RESULT_FFA(ffa_res);
-	}
-
-	if ((*count * sizeof(struct ffa_partition_information)) > buffer_size) {
-		/*
-		 * The indicated amount of info structures doesn't fit into the
-		 * RX buffer.
-		 */
-		*count = UINT32_C(0);
-		return SP_RESULT_INTERNAL_ERROR;
-	}
-
-	*info = (const struct ffa_partition_information *)buffer;
-
-	return SP_RESULT_OK;
-}
-
-static sp_result
-partition_info_get_single(const struct sp_uuid *uuid,
-			  const struct ffa_partition_information **info)
-{
-	uint32_t count = 0;
-	sp_result sp_res = SP_RESULT_OK;
-
-	if (uuid == NULL)
-		return SP_RESULT_INVALID_PARAMETERS;
-
-	/*
-	 * Nil UUID means querying all partitions which is handled by a separate
-	 * function.
-	 */
-	if (memcmp(&uuid_nil, uuid, sizeof(struct sp_uuid)) == 0)
-		return SP_RESULT_INVALID_PARAMETERS;
-
-	sp_res = partition_info_get(uuid, info, &count);
-	if (sp_res != SP_RESULT_OK)
-		return sp_res;
-
-	if (count == 0)
-		return SP_RESULT_NOT_FOUND;
-
-	return SP_RESULT_OK;
-}
-
-static void unpack_ffa_info(const struct ffa_partition_information ffa_info[],
+static void unpack_ffa_info(const struct ffa_partition_information *ffa_info,
 			    struct sp_partition_info *sp_info)
 {
 	uint32_t props = ffa_info->partition_properties;
@@ -120,53 +53,20 @@
 		props & FFA_PARTITION_SUPPORTS_INDIRECT_REQUESTS;
 }
 
-sp_result sp_discovery_partition_id_get(const struct sp_uuid *uuid,
-					uint16_t *id)
-{
-	const struct ffa_partition_information *ffa_info = NULL;
-	sp_result sp_res = SP_RESULT_OK;
-
-	if (id == NULL)
-		return SP_RESULT_INVALID_PARAMETERS;
-
-	sp_res = partition_info_get_single(uuid, &ffa_info);
-	if (sp_res != SP_RESULT_OK) {
-		*id = FFA_ID_GET_ID_MASK;
-		return sp_res;
-	}
-
-	*id = ffa_info->partition_id;
-
-	return SP_RESULT_OK;
-}
-
-sp_result sp_discovery_partition_info_get(const struct sp_uuid *uuid,
-					  struct sp_partition_info *info)
-{
-	const struct ffa_partition_information *ffa_info = NULL;
-	sp_result sp_res = SP_RESULT_OK;
-
-	if (info == NULL)
-		return SP_RESULT_INVALID_PARAMETERS;
-
-	sp_res = partition_info_get_single(uuid, &ffa_info);
-	if (sp_res != SP_RESULT_OK) {
-		*info = (struct sp_partition_info){ 0 };
-		return sp_res;
-	}
-
-	unpack_ffa_info(ffa_info, info);
-
-	return SP_RESULT_OK;
-}
-
-sp_result sp_discovery_partition_info_get_all(struct sp_partition_info info[],
-					      uint32_t *count)
+static sp_result
+partition_info_get(const struct sp_uuid *uuid,
+		   struct sp_partition_info info[],
+		   uint32_t *count,
+		   bool allow_nil_uuid)
 {
 	const struct ffa_partition_information *ffa_info = NULL;
 	uint32_t ffa_count = 0;
 	uint32_t i = 0;
 	sp_result sp_res = SP_RESULT_OK;
+	const void *buffer = NULL;
+	size_t buffer_size = 0;
+	struct ffa_uuid ffa_uuid = { 0 };
+	ffa_result ffa_res = FFA_OK;
 
 	if (count == NULL)
 		return SP_RESULT_INVALID_PARAMETERS;
@@ -176,10 +76,40 @@
 		return SP_RESULT_INVALID_PARAMETERS;
 	}
 
-	sp_res = partition_info_get(&uuid_nil, &ffa_info, &ffa_count);
+	if (uuid == NULL || (!allow_nil_uuid &&
+	    memcmp(&uuid_nil, uuid, sizeof(struct sp_uuid)) == 0)) {
+		sp_res = SP_RESULT_INVALID_PARAMETERS;
+		goto out;
+	}
+
+	sp_res = sp_rxtx_buffer_rx_get(&buffer, &buffer_size);
 	if (sp_res != SP_RESULT_OK) {
-		*count = UINT32_C(0);
-		return sp_res;
+		goto out;
+	}
+
+	/* Safely convert to FF-A UUID format */
+	memcpy(&ffa_uuid.uuid, uuid->uuid, sizeof(ffa_uuid.uuid));
+
+	ffa_res = ffa_partition_info_get(&ffa_uuid, &ffa_count);
+	if (ffa_res != FFA_OK) {
+		sp_res = SP_RESULT_FFA(ffa_res);
+		goto out;
+	}
+
+	if ((ffa_count * sizeof(struct ffa_partition_information)) > buffer_size) {
+		/*
+		 * The indicated amount of info structures doesn't fit into the
+		 * RX buffer.
+		 */
+		sp_res = SP_RESULT_INTERNAL_ERROR;
+		goto out;
+	}
+
+	ffa_info = (const struct ffa_partition_information *)buffer;
+
+	if (ffa_count == 0) {
+		sp_res = SP_RESULT_NOT_FOUND;
+		goto out;
 	}
 
 	*count = MIN(*count, ffa_count);
@@ -187,4 +117,48 @@
 		unpack_ffa_info(&ffa_info[i], &info[i]);
 
 	return SP_RESULT_OK;
+
+out:
+	for (i = 0; i < *count; i++)
+		info[i] =  (struct sp_partition_info){ 0 };
+	*count = UINT32_C(0);
+
+	return sp_res;
+}
+
+sp_result sp_discovery_partition_id_get(const struct sp_uuid *uuid,
+					uint16_t *id)
+{
+	struct sp_partition_info sp_info = { 0 };
+	sp_result sp_res = SP_RESULT_OK;
+	uint32_t count = 1;
+
+	if (id == NULL)
+		return SP_RESULT_INVALID_PARAMETERS;
+
+	*id = FFA_ID_GET_ID_MASK;
+
+	if (uuid == NULL || memcmp(&uuid_nil, uuid, sizeof(struct sp_uuid)) == 0)
+		return SP_RESULT_INVALID_PARAMETERS;
+
+	sp_res = partition_info_get(uuid, &sp_info, &count, false);
+	if (sp_res != SP_RESULT_OK)
+		return sp_res;
+
+	*id = sp_info.partition_id;
+
+	return SP_RESULT_OK;
+}
+
+sp_result sp_discovery_partition_info_get(const struct sp_uuid *uuid,
+					  struct sp_partition_info info[],
+					  uint32_t *count)
+{
+	return partition_info_get(uuid, info, count, false);
+}
+
+sp_result sp_discovery_partition_info_get_all(struct sp_partition_info info[],
+					      uint32_t *count)
+{
+	return partition_info_get(&uuid_nil, info, count, true);
 }
diff --git a/components/messaging/ffa/libsp/test/test_sp_discovery.cpp b/components/messaging/ffa/libsp/test/test_sp_discovery.cpp
index 1e61581..5c48fe8 100644
--- a/components/messaging/ffa/libsp/test/test_sp_discovery.cpp
+++ b/components/messaging/ffa/libsp/test/test_sp_discovery.cpp
@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: BSD-3-Clause
 /*
- * Copyright (c) 2020-2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2020-2023, Arm Limited. All rights reserved.
  */
 
 #include <CppUTest/TestHarness.h>
@@ -169,18 +169,20 @@
 {
 	struct sp_partition_info info = { 1 };
 	const struct sp_partition_info expected_info = { 0 };
+	uint32_t count = 1;
 
 	LONGS_EQUAL(SP_RESULT_INVALID_PARAMETERS,
-		    sp_discovery_partition_info_get(NULL, &info));
+		    sp_discovery_partition_info_get(NULL, &info, &count));
 	MEMCMP_EQUAL(&expected_info, &info, sizeof(info));
 }
 
 TEST(sp_discovery, sp_discovery_partition_info_get_null_info)
 {
 	struct sp_uuid uuid = { 1 };
+	uint32_t count = 0;
 
 	LONGS_EQUAL(SP_RESULT_INVALID_PARAMETERS,
-		    sp_discovery_partition_info_get(&uuid, NULL));
+		    sp_discovery_partition_info_get(&uuid, NULL, &count));
 }
 
 TEST(sp_discovery, sp_discovery_partition_info_get_nil_uuid)
@@ -188,9 +190,10 @@
 	struct sp_uuid uuid = { 0 };
 	struct sp_partition_info info = { 1 };
 	const struct sp_partition_info expected_info = { 0 };
+	uint32_t count = 1;
 
 	LONGS_EQUAL(SP_RESULT_INVALID_PARAMETERS,
-		    sp_discovery_partition_info_get(&uuid, &info));
+		    sp_discovery_partition_info_get(&uuid, &info, &count));
 	MEMCMP_EQUAL(&expected_info, &info, sizeof(info));
 }
 
@@ -199,9 +202,11 @@
 	struct sp_uuid uuid = { 1 };
 	struct sp_partition_info info = { 1 };
 	const struct sp_partition_info expected_info = { 0 };
+	uint32_t count = 1;
 
 	expect_sp_rxtx_buffer_rx_get(&rx_buffer, &rx_buffer_size, result);
-	LONGS_EQUAL(result, sp_discovery_partition_info_get(&uuid, &info));
+	LONGS_EQUAL(result, sp_discovery_partition_info_get(&uuid, &info,
+							    &count));
 	MEMCMP_EQUAL(&expected_info, &info, sizeof(info));
 }
 
@@ -209,13 +214,14 @@
 {
 	struct sp_uuid uuid = { 1 };
 	struct ffa_uuid ffa_uuid = { 1 };
-	uint32_t count = 0;
+	uint32_t count = 1;
 	struct sp_partition_info info = { 1 };
 	const struct sp_partition_info expected_info = { 0 };
 
 	expect_sp_rxtx_buffer_rx_get(&rx_buffer, &rx_buffer_size, SP_RESULT_OK);
 	expect_ffa_partition_info_get(&ffa_uuid, &count, result);
-	LONGS_EQUAL(result, sp_discovery_partition_info_get(&uuid, &info));
+	LONGS_EQUAL(result, sp_discovery_partition_info_get(&uuid, &info,
+							    &count));
 	MEMCMP_EQUAL(&expected_info, &info, sizeof(info));
 }
 
@@ -223,15 +229,16 @@
 {
 	struct sp_uuid uuid = { 1 };
 	struct ffa_uuid ffa_uuid = { 1 };
-	uint32_t count =
+	uint32_t expected_count =
 		(rx_buffer_size / sizeof(struct ffa_partition_information)) + 1;
+	uint32_t count = 1;
 	struct sp_partition_info info = { 1 };
 	const struct sp_partition_info expected_info = { 0 };
 
 	expect_sp_rxtx_buffer_rx_get(&rx_buffer, &rx_buffer_size, SP_RESULT_OK);
-	expect_ffa_partition_info_get(&ffa_uuid, &count, SP_RESULT_OK);
+	expect_ffa_partition_info_get(&ffa_uuid, &expected_count, SP_RESULT_OK);
 	LONGS_EQUAL(SP_RESULT_INTERNAL_ERROR,
-		    sp_discovery_partition_info_get(&uuid, &info));
+		    sp_discovery_partition_info_get(&uuid, &info, &count));
 	MEMCMP_EQUAL(&expected_info, &info, sizeof(info));
 }
 
@@ -239,15 +246,17 @@
 {
 	struct sp_uuid uuid = { 1 };
 	struct ffa_uuid ffa_uuid = { 1 };
-	uint32_t count = 0;
+	uint32_t expected_count = 0;
+	uint32_t count = 1;
 	struct sp_partition_info info = { 1 };
 	const struct sp_partition_info expected_info = { 0 };
 
 	expect_sp_rxtx_buffer_rx_get(&rx_buffer, &rx_buffer_size, SP_RESULT_OK);
-	expect_ffa_partition_info_get(&ffa_uuid, &count, SP_RESULT_OK);
+	expect_ffa_partition_info_get(&ffa_uuid, &expected_count, SP_RESULT_OK);
 	LONGS_EQUAL(SP_RESULT_NOT_FOUND,
-		    sp_discovery_partition_info_get(&uuid, &info));
+		    sp_discovery_partition_info_get(&uuid, &info, &count));
 	MEMCMP_EQUAL(&expected_info, &info, sizeof(info));
+	UNSIGNED_LONGS_EQUAL(expected_count, count);
 }
 
 TEST(sp_discovery, sp_discovery_partition_info)
@@ -258,22 +267,29 @@
 		(rx_buffer_size / sizeof(struct ffa_partition_information));
 	const uint16_t expected_id = 1234;
 	const uint16_t expected_context_count = 23456;
-	struct sp_partition_info info = { 0 };
+	struct sp_partition_info *info =
+		(struct sp_partition_info *)calloc(count, sizeof(struct sp_partition_info));
 
-	((struct ffa_partition_information *)rx_buffer)->partition_id =
-		expected_id;
-	((struct ffa_partition_information *)rx_buffer)
-		->execution_context_count = expected_context_count;
+	for (uint32_t i = 0; i < count; i++) {
+		((struct ffa_partition_information *)rx_buffer)[i].partition_id = expected_id + i;
+		((struct ffa_partition_information *)rx_buffer)[i].execution_context_count = expected_context_count + i;
+		((struct ffa_partition_information *)rx_buffer)[i].partition_properties = 0;
+	}
 	expect_sp_rxtx_buffer_rx_get(&rx_buffer, &rx_buffer_size, SP_RESULT_OK);
 	expect_ffa_partition_info_get(&ffa_uuid, &count, SP_RESULT_OK);
 	LONGS_EQUAL(SP_RESULT_OK,
-		    sp_discovery_partition_info_get(&uuid, &info));
-	UNSIGNED_LONGS_EQUAL(expected_id, info.partition_id);
-	UNSIGNED_LONGS_EQUAL(expected_context_count,
-			     info.execution_context_count);
-	CHECK_FALSE(info.supports_direct_requests);
-	CHECK_FALSE(info.can_send_direct_requests);
-	CHECK_FALSE(info.supports_indirect_requests);
+		    sp_discovery_partition_info_get(&uuid, info, &count));
+
+	for (uint32_t i = 0; i < count; i++) {
+		UNSIGNED_LONGS_EQUAL(expected_id + i, info[i].partition_id);
+		UNSIGNED_LONGS_EQUAL(expected_context_count + i,
+				info[i].execution_context_count);
+		CHECK_FALSE(info[i].supports_direct_requests);
+		CHECK_FALSE(info[i].can_send_direct_requests);
+		CHECK_FALSE(info[i].supports_indirect_requests);
+	}
+
+	free(info);
 }
 
 TEST(sp_discovery, sp_discovery_partition_info_support_direct_req)
@@ -288,7 +304,7 @@
 	expect_sp_rxtx_buffer_rx_get(&rx_buffer, &rx_buffer_size, SP_RESULT_OK);
 	expect_ffa_partition_info_get(&ffa_uuid, &count, SP_RESULT_OK);
 	LONGS_EQUAL(SP_RESULT_OK,
-		    sp_discovery_partition_info_get(&uuid, &info));
+		    sp_discovery_partition_info_get(&uuid, &info, &count));
 	CHECK_TRUE(info.supports_direct_requests);
 	CHECK_FALSE(info.can_send_direct_requests);
 	CHECK_FALSE(info.supports_indirect_requests);
@@ -306,7 +322,7 @@
 	expect_sp_rxtx_buffer_rx_get(&rx_buffer, &rx_buffer_size, SP_RESULT_OK);
 	expect_ffa_partition_info_get(&ffa_uuid, &count, SP_RESULT_OK);
 	LONGS_EQUAL(SP_RESULT_OK,
-		    sp_discovery_partition_info_get(&uuid, &info));
+		    sp_discovery_partition_info_get(&uuid, &info, &count));
 	CHECK_FALSE(info.supports_direct_requests);
 	CHECK_TRUE(info.can_send_direct_requests);
 	CHECK_FALSE(info.supports_indirect_requests);
@@ -324,7 +340,7 @@
 	expect_sp_rxtx_buffer_rx_get(&rx_buffer, &rx_buffer_size, SP_RESULT_OK);
 	expect_ffa_partition_info_get(&ffa_uuid, &count, SP_RESULT_OK);
 	LONGS_EQUAL(SP_RESULT_OK,
-		    sp_discovery_partition_info_get(&uuid, &info));
+		    sp_discovery_partition_info_get(&uuid, &info, &count));
 	CHECK_FALSE(info.supports_direct_requests);
 	CHECK_FALSE(info.can_send_direct_requests);
 	CHECK_TRUE(info.supports_indirect_requests);
@@ -394,7 +410,7 @@
 
 	expect_sp_rxtx_buffer_rx_get(&rx_buffer, &rx_buffer_size, SP_RESULT_OK);
 	expect_ffa_partition_info_get(&ffa_uuid, &expected_count, SP_RESULT_OK);
-	LONGS_EQUAL(SP_RESULT_OK,
+	LONGS_EQUAL(SP_RESULT_NOT_FOUND,
 		    sp_discovery_partition_info_get_all(&info, &count));
 	UNSIGNED_LONGS_EQUAL(0, count);
 }