feat(memory share): support sending fragmented messages.

Updates `memory_init_and_send` to allow sending memory share messages
across multiple fragments.
Tests for this functionality will be added in the next commit.

Change-Id: I6a47735415a8fed1d15322499f1d08cc3e1fd6a2
Signed-off-by: Karl Meakin <karl.meakin@arm.com>
diff --git a/include/runtime_services/ffa_helpers.h b/include/runtime_services/ffa_helpers.h
index f265fe8..3f760fa 100644
--- a/include/runtime_services/ffa_helpers.h
+++ b/include/runtime_services/ffa_helpers.h
@@ -433,10 +433,10 @@
 _Static_assert(sizeof(ffa_memory_attributes_t) == sizeof(uint16_t),
 	       "ffa_memory_attributes_t must be 2 bytes wide");
 
-#define FFA_MEMORY_HANDLE_ALLOCATOR_MASK \
-	((ffa_memory_handle_t)(UINT64_C(1) << 63))
-#define FFA_MEMORY_HANDLE_ALLOCATOR_HYPERVISOR \
-	((ffa_memory_handle_t)(UINT64_C(1) << 63))
+#define FFA_MEMORY_HANDLE_ALLOCATOR_MASK UINT64_C(1)
+#define FFA_MEMORY_HANDLE_ALLOCATOR_SHIFT 63U
+#define FFA_MEMORY_HANDLE_ALLOCATOR_HYPERVISOR UINT64_C(1)
+#define FFA_MEMORY_HANDLE_ALLOCATOR_SPMC UINT64_C(0)
 #define FFA_MEMORY_HANDLE_INVALID (~UINT64_C(0))
 
 /**
@@ -608,6 +608,16 @@
 	return ffa_assemble_handle(r.arg2, r.arg3);
 }
 
+static inline ffa_memory_handle_t ffa_frag_handle(struct ffa_value r)
+{
+	return ffa_assemble_handle(r.arg1, r.arg2);
+}
+
+static inline ffa_id_t ffa_frag_sender(struct ffa_value args)
+{
+	return (args.arg4 >> 16) & 0xffff;
+}
+
 /**
  * Gets the `ffa_composite_memory_region` for the given receiver from an
  * `ffa_memory_region`, or NULL if it is not valid.
@@ -660,6 +670,12 @@
 	enum ffa_memory_shareability shareability, uint32_t *total_length,
 	uint32_t *fragment_length);
 
+uint32_t ffa_memory_fragment_init(
+	struct ffa_memory_region_constituent *fragment,
+	size_t fragment_max_size,
+	const struct ffa_memory_region_constituent constituents[],
+	uint32_t constituent_count, uint32_t *fragment_length);
+
 static inline ffa_id_t ffa_dir_msg_dest(struct ffa_value val) {
 	return (ffa_id_t)val.arg1 & U(0xFFFF);
 }
@@ -711,6 +727,10 @@
 				      uint32_t fragment_length);
 struct ffa_value ffa_mem_relinquish(void);
 struct ffa_value ffa_mem_reclaim(uint64_t handle, uint32_t flags);
+struct ffa_value ffa_mem_frag_rx(ffa_memory_handle_t handle,
+				 uint32_t fragment_length);
+struct ffa_value ffa_mem_frag_tx(ffa_memory_handle_t handle,
+				 uint32_t fragment_length);
 struct ffa_value ffa_notification_bitmap_create(ffa_id_t vm_id,
 						ffa_vcpu_count_t vcpu_count);
 struct ffa_value ffa_notification_bitmap_destroy(ffa_id_t vm_id);
diff --git a/include/runtime_services/ffa_svc.h b/include/runtime_services/ffa_svc.h
index 3abd21f..226bbd1 100644
--- a/include/runtime_services/ffa_svc.h
+++ b/include/runtime_services/ffa_svc.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -86,6 +86,8 @@
 #define FFA_FNUM_MEM_RETRIEVE_RESP		U(0x75)
 #define FFA_FNUM_MEM_RELINQUISH			U(0x76)
 #define FFA_FNUM_MEM_RECLAIM			U(0x77)
+#define FFA_FNUM_MEM_FRAG_RX 			U(0x7A)
+#define FFA_FNUM_MEM_FRAG_TX 			U(0x7B)
 #define FFA_FNUM_NORMAL_WORLD_RESUME		U(0x7C)
 
 /* FF-A v1.1 */
@@ -134,6 +136,8 @@
 #define FFA_MEM_RETRIEVE_RESP	FFA_FID(SMC_32, FFA_FNUM_MEM_RETRIEVE_RESP)
 #define FFA_MEM_RELINQUISH	FFA_FID(SMC_32, FFA_FNUM_MEM_RELINQUISH)
 #define FFA_MEM_RECLAIM		FFA_FID(SMC_32, FFA_FNUM_MEM_RECLAIM)
+#define FFA_MEM_FRAG_RX 	FFA_FID(SMC_32, FFA_FNUM_MEM_FRAG_RX)
+#define FFA_MEM_FRAG_TX 	FFA_FID(SMC_32, FFA_FNUM_MEM_FRAG_TX)
 #define FFA_NOTIFICATION_BITMAP_CREATE	\
 	FFA_FID(SMC_32, FFA_FNUM_NOTIFICATION_BITMAP_CREATE)
 #define FFA_NOTIFICATION_BITMAP_DESTROY	\
diff --git a/include/runtime_services/spm_common.h b/include/runtime_services/spm_common.h
index 6494947..ad2ba08 100644
--- a/include/runtime_services/spm_common.h
+++ b/include/runtime_services/spm_common.h
@@ -110,8 +110,7 @@
 bool memory_retrieve(struct mailbox_buffers *mb,
 		     struct ffa_memory_region **retrieved, uint64_t handle,
 		     ffa_id_t sender, struct ffa_memory_access receivers[],
-		     uint32_t receiver_count, ffa_memory_region_flags_t flags,
-		     uint32_t mem_func);
+		     uint32_t receiver_count, ffa_memory_region_flags_t flags);
 
 bool hypervisor_retrieve_request(struct mailbox_buffers *mb, uint64_t handle,
 				 void *out, uint32_t out_size);
@@ -124,20 +123,22 @@
 		       ffa_id_t id);
 
 ffa_memory_handle_t memory_send(
-	struct ffa_memory_region *memory_region, uint32_t mem_func,
-	uint32_t fragment_length, uint32_t total_length, struct ffa_value *ret);
+	void *send_buffer, uint32_t mem_func,
+	const struct ffa_memory_region_constituent *constituents,
+	uint32_t constituent_count, uint32_t remaining_constituent_count,
+	uint32_t fragment_length, uint32_t total_length,
+	struct ffa_value *ret);
 
 ffa_memory_handle_t memory_init_and_send(
-	struct ffa_memory_region *memory_region, size_t memory_region_max_size,
-	ffa_id_t sender, struct ffa_memory_access receivers[],
-	uint32_t receiver_count,
+	void *send_buffer, size_t memory_region_max_size, ffa_id_t sender,
+	struct ffa_memory_access receivers[], uint32_t receiver_count,
 	const struct ffa_memory_region_constituent *constituents,
 	uint32_t constituents_count, uint32_t mem_func, struct ffa_value *ret);
 
 bool ffa_partition_info_helper(struct mailbox_buffers *mb,
-			const struct ffa_uuid uuid,
-			const struct ffa_partition_info *expected,
-			const uint16_t expected_size);
+			       const struct ffa_uuid uuid,
+			       const struct ffa_partition_info *expected,
+			       const uint16_t expected_size);
 bool enable_trusted_wdog_interrupt(ffa_id_t source, ffa_id_t dest);
 bool disable_trusted_wdog_interrupt(ffa_id_t source, ffa_id_t dest);
 
diff --git a/spm/cactus/cactus_tests/cactus_test_memory_sharing.c b/spm/cactus/cactus_tests/cactus_test_memory_sharing.c
index ef49d21..2fbc63b 100644
--- a/spm/cactus/cactus_tests/cactus_test_memory_sharing.c
+++ b/spm/cactus/cactus_tests/cactus_test_memory_sharing.c
@@ -89,7 +89,7 @@
 		0);
 
 	expect(memory_retrieve(mb, &m, handle, source, &receiver, 1,
-			       retrv_flags, mem_func),
+			       retrv_flags),
 	       true);
 
 	composite = ffa_memory_region_get_composite(m, 0);
@@ -232,10 +232,9 @@
 					 CACTUS_ERROR_TEST);
 	}
 
-	handle = memory_init_and_send((struct ffa_memory_region *)mb->send,
-				      PAGE_SIZE, vm_id, &receiver, 1,
-				      constituents, constituents_count,
-				      mem_func, &ffa_ret);
+	handle = memory_init_and_send(mb->send, PAGE_SIZE, vm_id, &receiver, 1,
+				       constituents, constituents_count,
+				       mem_func, &ffa_ret);
 
 	/*
 	 * If returned an invalid handle, we should break the test.
diff --git a/tftf/tests/misc_tests/test_invalid_access.c b/tftf/tests/misc_tests/test_invalid_access.c
index 353162b..b05d694 100644
--- a/tftf/tests/misc_tests/test_invalid_access.c
+++ b/tftf/tests/misc_tests/test_invalid_access.c
@@ -306,8 +306,7 @@
 
 	GET_TFTF_MAILBOX(mb);
 
-	handle = memory_init_and_send((struct ffa_memory_region *)mb.send,
-				      PAGE_SIZE, SENDER, &receiver, 1,
+	handle = memory_init_and_send(mb.send, PAGE_SIZE, SENDER, &receiver, 1,
 				      constituents, constituents_count,
 				      FFA_MEM_SHARE_SMC32, &ret);
 
diff --git a/tftf/tests/runtime_services/secure_service/ffa_helpers.c b/tftf/tests/runtime_services/secure_service/ffa_helpers.c
index 82413a2..ad63e68 100644
--- a/tftf/tests/runtime_services/secure_service/ffa_helpers.c
+++ b/tftf/tests/runtime_services/secure_service/ffa_helpers.c
@@ -280,6 +280,31 @@
 		constituent_count, total_length, fragment_length);
 }
 
+uint32_t ffa_memory_fragment_init(
+	struct ffa_memory_region_constituent *fragment,
+	size_t fragment_max_size,
+	const struct ffa_memory_region_constituent constituents[],
+	uint32_t constituent_count, uint32_t *fragment_length)
+{
+	const uint32_t fragment_max_constituents =
+		fragment_max_size /
+		sizeof(struct ffa_memory_region_constituent);
+
+	uint32_t count_to_copy =
+		MIN(constituent_count, fragment_max_constituents);
+
+	for (uint32_t i = 0; i < count_to_copy; ++i) {
+		fragment[i] = constituents[i];
+	}
+
+	if (fragment_length != NULL) {
+		*fragment_length = count_to_copy *
+				   sizeof(struct ffa_memory_region_constituent);
+	}
+
+	return constituent_count - count_to_copy;
+}
+
 /**
  * Initialises the given `ffa_memory_region` to be used for an
  * `FFA_MEM_RETRIEVE_REQ` by the receiver of a memory transaction.
@@ -570,6 +595,34 @@
 	return ffa_service_call(&args);
 }
 
+struct ffa_value ffa_mem_frag_rx(ffa_memory_handle_t handle,
+				 uint32_t fragment_offset)
+{
+	/* Note that sender MBZ at virtual instance. */
+	struct ffa_value args = {
+		.fid = FFA_MEM_FRAG_RX,
+		.arg1 = (uint32_t)handle,
+		.arg2 = (uint32_t)(handle >> 32),
+		.arg3 = fragment_offset,
+	};
+
+	return ffa_service_call(&args);
+}
+
+struct ffa_value ffa_mem_frag_tx(ffa_memory_handle_t handle,
+				 uint32_t fragment_length)
+{
+	struct ffa_value args = {
+		.fid = FFA_MEM_FRAG_TX,
+		.arg1 = (uint32_t)handle,
+		.arg2 = (uint32_t)(handle >> 32),
+		.arg3 = fragment_length,
+	};
+
+	/* Note that sender MBZ at virtual instance. */
+	return ffa_service_call(&args);
+}
+
 /** Create Notifications Bitmap for the given VM */
 struct ffa_value ffa_notification_bitmap_create(ffa_id_t vm_id,
 						ffa_vcpu_count_t vcpu_count)
diff --git a/tftf/tests/runtime_services/secure_service/spm_common.c b/tftf/tests/runtime_services/secure_service/spm_common.c
index e75be35..dce4b7f 100644
--- a/tftf/tests/runtime_services/secure_service/spm_common.c
+++ b/tftf/tests/runtime_services/secure_service/spm_common.c
@@ -4,7 +4,6 @@
  * SPDX-License-Identifier: BSD-3-Clause
  */
 
-#include "assert.h"
 #include "stdint.h"
 
 #include "ffa_helpers.h"
@@ -212,8 +211,7 @@
 bool memory_retrieve(struct mailbox_buffers *mb,
 		     struct ffa_memory_region **retrieved, uint64_t handle,
 		     ffa_id_t sender, struct ffa_memory_access receivers[],
-		     uint32_t receiver_count, ffa_memory_region_flags_t flags,
-		     uint32_t mem_func)
+		     uint32_t receiver_count, ffa_memory_region_flags_t flags)
 {
 	struct ffa_value ret;
 	uint32_t fragment_size;
@@ -233,8 +231,8 @@
 	ret = ffa_mem_retrieve_req(descriptor_size, descriptor_size);
 
 	if (ffa_func_id(ret) != FFA_MEM_RETRIEVE_RESP) {
-		ERROR("Couldn't retrieve the memory page. Error: %x\n",
-		      ffa_error_code(ret));
+		ERROR("%s: couldn't retrieve the memory page. Error: %d\n",
+		      __func__, ffa_error_code(ret));
 		return false;
 	}
 
@@ -277,8 +275,9 @@
 				 void *out, uint32_t out_size)
 {
 	struct ffa_value ret;
-	uint32_t fragment_size;
 	uint32_t total_size;
+	uint32_t fragment_size;
+	uint32_t fragment_offset;
 	struct ffa_memory_region *region_out = out;
 
 	if (out == NULL || mb == NULL) {
@@ -291,8 +290,8 @@
 				   sizeof(struct ffa_memory_region));
 
 	if (ffa_func_id(ret) != FFA_MEM_RETRIEVE_RESP) {
-		ERROR("Couldn't retrieve the memory page. Error: %x\n",
-		      ffa_error_code(ret));
+		ERROR("%s: couldn't retrieve the memory page. Error: %d\n",
+		      __func__, ffa_error_code(ret));
 		return false;
 	}
 
@@ -301,17 +300,12 @@
 	 * of the state of transaction. When the sum of all fragment_size of all
 	 * fragments is equal to total_size, the memory transaction has been
 	 * completed.
-	 * This is a simple test with only one segment. As such, upon
-	 * successful ffa_mem_retrieve_req, total_size must be equal to
-	 * fragment_size.
 	 */
 	total_size = ret.arg1;
 	fragment_size = ret.arg2;
-
-	if (total_size != fragment_size) {
-		ERROR("Only expect one memory segment to be sent!\n");
-		return false;
-	}
+	fragment_offset = fragment_size;
+	VERBOSE("total_size=%d, fragment_size=%d, fragment_offset=%d\n",
+		total_size, fragment_size, fragment_offset);
 
 	if (fragment_size > PAGE_SIZE) {
 		ERROR("Fragment should be smaller than RX buffer!\n");
@@ -343,6 +337,75 @@
 		return false;
 	}
 
+	while (fragment_offset < total_size) {
+		VERBOSE("Calling again. frag offset: %d; total: %d\n",
+			fragment_offset, total_size);
+		ret = ffa_rx_release();
+		if (ret.fid != FFA_SUCCESS_SMC32) {
+			ERROR("ffa_rx_release() failed: %d\n",
+			      ffa_error_code(ret));
+			return false;
+		}
+
+		ret = ffa_mem_frag_rx(handle, fragment_offset);
+		if (ret.fid != FFA_MEM_FRAG_TX) {
+			ERROR("ffa_mem_frag_rx() failed: %d\n",
+			      ffa_error_code(ret));
+			return false;
+		}
+
+		if (ffa_frag_handle(ret) != handle) {
+			ERROR("%s: fragment handle mismatch: expected %llu, "
+			      "got "
+			      "%llu\n",
+			      __func__, handle, ffa_frag_handle(ret));
+			return false;
+		}
+
+		/* Sender MBZ at physical instance. */
+		if (ffa_frag_sender(ret) != 0) {
+			ERROR("%s: fragment sender mismatch: expected %d, got "
+			      "%d\n",
+			      __func__, 0, ffa_frag_sender(ret));
+			return false;
+		}
+
+		fragment_size = ret.arg3;
+		if (fragment_size == 0) {
+			ERROR("%s: fragment size must not be 0\n", __func__);
+			return false;
+		}
+
+		if (fragment_offset + fragment_size > out_size) {
+			ERROR("%s: fragment is too big to fit in out buffer "
+			      "(%d > %d)\n",
+			      __func__, fragment_offset + fragment_size,
+			      out_size);
+			return false;
+		}
+
+		VERBOSE("copying fragment at offset %d with size %d\n",
+			fragment_offset, fragment_size);
+		memcpy((uint8_t *)out + fragment_offset, mb->recv,
+		       fragment_size);
+
+		fragment_offset += fragment_size;
+	}
+
+	if (fragment_offset != total_size) {
+		ERROR("%s: fragment size mismatch: expected %d, got %d\n",
+		      __func__, total_size, fragment_offset);
+		return false;
+	}
+
+	ret = ffa_rx_release();
+	if (ret.fid != FFA_SUCCESS_SMC32) {
+		ERROR("ffa_rx_release() failed: %d\n", ffa_error_code(ret));
+		return false;
+	}
+
+	VERBOSE("Memory Retrieved!\n");
+
 	return true;
 }
 
@@ -363,21 +426,103 @@
 	return true;
 }
 
+bool send_fragmented_memory_region(
+	void *send_buffer,
+	const struct ffa_memory_region_constituent constituents[],
+	uint32_t constituent_count, uint32_t remaining_constituent_count,
+	uint32_t sent_length, uint32_t total_length, bool allocator_is_spmc,
+	struct ffa_value ret)
+{
+
+	uint64_t handle;
+	uint64_t handle_mask;
+	uint64_t expected_handle_mask =
+		allocator_is_spmc ? FFA_MEMORY_HANDLE_ALLOCATOR_SPMC
+				  : FFA_MEMORY_HANDLE_ALLOCATOR_HYPERVISOR;
+	ffa_memory_handle_t fragment_handle = FFA_MEMORY_HANDLE_INVALID;
+	uint32_t fragment_length;
+
+	/* Send the remaining fragments. */
+	while (remaining_constituent_count != 0) {
+		VERBOSE("%s: %d constituents left to send.\n", __func__,
+			remaining_constituent_count);
+		if (ret.fid != FFA_MEM_FRAG_RX) {
+			ERROR("ffa_mem_frax_tx() failed: %d\n",
+			      ffa_error_code(ret));
+			return false;
+		}
+
+		if (fragment_handle == FFA_MEMORY_HANDLE_INVALID) {
+			fragment_handle = ffa_frag_handle(ret);
+		} else if (ffa_frag_handle(ret) != fragment_handle) {
+			ERROR("%s: fragment handle mismatch: expected %llu, "
+			      "got %llu\n",
+			      __func__, fragment_handle, ffa_frag_handle(ret));
+			return false;
+		}
+
+		if (ret.arg3 != sent_length) {
+			ERROR("%s: fragment length mismatch: expected %u, got "
+			      "%lu\n",
+			      __func__, sent_length, ret.arg3);
+			return false;
+		}
+
+		remaining_constituent_count = ffa_memory_fragment_init(
+			send_buffer, PAGE_SIZE,
+			constituents + constituent_count -
+				remaining_constituent_count,
+			remaining_constituent_count, &fragment_length);
+
+		ret = ffa_mem_frag_tx(fragment_handle, fragment_length);
+		sent_length += fragment_length;
+	}
+
+	if (sent_length != total_length) {
+		ERROR("%s: fragment length mismatch: expected %u, got %u\n",
+		      __func__, total_length, sent_length);
+		return false;
+	}
+
+	if (ret.fid != FFA_SUCCESS_SMC32) {
+		ERROR("%s: ffa_mem_frax_tx() failed: %d\n", __func__,
+		      ffa_error_code(ret));
+		return false;
+	}
+
+	handle = ffa_mem_success_handle(ret);
+	handle_mask = (handle >> FFA_MEMORY_HANDLE_ALLOCATOR_SHIFT) &
+		      FFA_MEMORY_HANDLE_ALLOCATOR_MASK;
+
+	if (handle_mask != expected_handle_mask) {
+		ERROR("%s: handle mask mismatch: expected %llu, got %llu\n",
+		      __func__, expected_handle_mask, handle_mask);
+		return false;
+	}
+
+	if (fragment_handle != FFA_MEMORY_HANDLE_INVALID && handle != fragment_handle) {
+		ERROR("%s: fragment handle mismatch: expectd %d, got %llu\n",
+		      __func__, fragment_length, handle);
+		return false;
+	}
+
+	return true;
+}
+
 /**
  * Helper to call memory send function whose func id is passed as a parameter.
- * Returns a valid handle in case of successful operation or
- * FFA_MEMORY_HANDLE_INVALID if something goes wrong. Populates *ret with a
- * resulting smc value to handle the error higher in the test chain.
- *
- * TODO: Do memory send with 'ffa_memory_region' taking multiple segments
  */
 ffa_memory_handle_t memory_send(
-	struct ffa_memory_region *memory_region, uint32_t mem_func,
-	uint32_t fragment_length, uint32_t total_length, struct ffa_value *ret)
+	void *send_buffer, uint32_t mem_func,
+	const struct ffa_memory_region_constituent *constituents,
+	uint32_t constituent_count, uint32_t remaining_constituent_count,
+	uint32_t fragment_length, uint32_t total_length,
+	struct ffa_value *ret)
 {
-	if (fragment_length != total_length) {
-		ERROR("For now, fragment_length and total_length need to be"
-		      " equal");
+	if (remaining_constituent_count == 0 && fragment_length != total_length) {
+		ERROR("%s: fragment_length and total_length need "
+		      "to be equal (fragment_length = %d, total_length = %d)\n",
+		      __func__, fragment_length, total_length);
 		return FFA_MEMORY_HANDLE_INVALID;
 	}
 
@@ -392,16 +537,20 @@
 		*ret = ffa_mem_donate(total_length, fragment_length);
 		break;
 	default:
-		*ret = (struct ffa_value){0};
-		ERROR("TFTF - Invalid func id %x!\n", mem_func);
+		ERROR("%s: Invalid func id %d!\n", __func__, mem_func);
 		return FFA_MEMORY_HANDLE_INVALID;
 	}
 
 	if (is_ffa_call_error(*ret)) {
-		VERBOSE("Failed to send memory to: %x\n",
-			memory_region->receivers[0]
-					.receiver_permissions
-					.receiver);
+		VERBOSE("%s: Failed to send memory: %d\n", __func__,
+			ffa_error_code(ret));
+		return FFA_MEMORY_HANDLE_INVALID;
+	}
+
+	if (!send_fragmented_memory_region(
+		    send_buffer, constituents, constituent_count,
+		    remaining_constituent_count, fragment_length, total_length,
+		    true, *ret)) {
 		return FFA_MEMORY_HANDLE_INVALID;
 	}
 
@@ -414,9 +563,8 @@
  * doing it in this file for simplicity and for testing purposes.
  */
 ffa_memory_handle_t memory_init_and_send(
-	struct ffa_memory_region *memory_region, size_t memory_region_max_size,
-	ffa_id_t sender, struct ffa_memory_access receivers[],
-	uint32_t receiver_count,
+	void *send_buffer, size_t memory_region_max_size, ffa_id_t sender,
+	struct ffa_memory_access receivers[], uint32_t receiver_count,
 	const struct ffa_memory_region_constituent *constituents,
 	uint32_t constituents_count, uint32_t mem_func, struct ffa_value *ret)
 {
@@ -430,22 +578,14 @@
 			: FFA_MEMORY_NORMAL_MEM;
 
 	remaining_constituent_count = ffa_memory_region_init(
-		memory_region, memory_region_max_size, sender, receivers,
+		send_buffer, memory_region_max_size, sender, receivers,
 		receiver_count, constituents, constituents_count, 0, 0, type,
 		FFA_MEMORY_CACHE_WRITE_BACK, FFA_MEMORY_INNER_SHAREABLE,
 		&total_length, &fragment_length);
 
-	/*
-	 * For simplicity of the test, and at least for the time being,
-	 * the following condition needs to be true.
-	 */
-	if (remaining_constituent_count != 0U) {
-		ERROR("Remaining constituent should be 0\n");
-		return FFA_MEMORY_HANDLE_INVALID;
-	}
-
-	return memory_send(memory_region, mem_func, fragment_length,
-			   total_length, ret);
+	return memory_send(send_buffer, mem_func, constituents,
+			   constituents_count, remaining_constituent_count,
+			   fragment_length, total_length, ret);
 }
 
 static bool ffa_uuid_equal(const struct ffa_uuid uuid1,
diff --git a/tftf/tests/runtime_services/secure_service/test_ffa_exceptions.c b/tftf/tests/runtime_services/secure_service/test_ffa_exceptions.c
index 6b16da4..71e81dc 100644
--- a/tftf/tests/runtime_services/secure_service/test_ffa_exceptions.c
+++ b/tftf/tests/runtime_services/secure_service/test_ffa_exceptions.c
@@ -64,10 +64,9 @@
 
 	GET_TFTF_MAILBOX(mb);
 
-	handle = memory_init_and_send((struct ffa_memory_region *)mb.send,
-				      PAGE_SIZE, SENDER, &receiver, 1,
-				      constituents, constituents_count,
-				      FFA_MEM_SHARE_SMC32, &ret);
+	handle = memory_init_and_send(mb.send, PAGE_SIZE, SENDER, &receiver, 1,
+				   constituents, constituents_count,
+				   FFA_MEM_SHARE_SMC32, &ret);
 
 	if (handle == FFA_MEMORY_HANDLE_INVALID) {
 		return TEST_RESULT_FAIL;
diff --git a/tftf/tests/runtime_services/secure_service/test_ffa_memory_sharing.c b/tftf/tests/runtime_services/secure_service/test_ffa_memory_sharing.c
index 12ae953..f69c590 100644
--- a/tftf/tests/runtime_services/secure_service/test_ffa_memory_sharing.c
+++ b/tftf/tests/runtime_services/secure_service/test_ffa_memory_sharing.c
@@ -20,12 +20,21 @@
 #define SENDER HYP_ID
 #define RECEIVER SP_ID(1)
 
+/*
+ * A number of pages that is large enough that it must take two fragments to
+ * share.
+ */
+#define FRAGMENTED_SHARE_PAGE_COUNT                                            \
+	(sizeof(struct ffa_memory_region) /                                    \
+	 sizeof(struct ffa_memory_region_constituent))
+
 static const struct ffa_uuid expected_sp_uuids[] = {
 		{PRIMARY_UUID}, {SECONDARY_UUID}, {TERTIARY_UUID}
 	};
 
 /* Memory section to be used for memory share operations */
-static __aligned(PAGE_SIZE) uint8_t share_page[PAGE_SIZE];
+static __aligned(PAGE_SIZE) uint8_t
+	share_page[PAGE_SIZE * FRAGMENTED_SHARE_PAGE_COUNT];
 static __aligned(PAGE_SIZE) uint8_t consecutive_donate_page[PAGE_SIZE];
 static __aligned(PAGE_SIZE) uint8_t four_share_pages[PAGE_SIZE * 4];
 
@@ -69,8 +78,8 @@
 					mem_func, &ret);
 
 	if (handle != FFA_MEMORY_HANDLE_INVALID) {
-		ERROR("Received a valid FF-A memory handle, and that isn't"
-		       " expected.\n");
+		ERROR("Received a valid FF-A memory handle, and that isn't "
+		      "expected.\n");
 		return false;
 	}
 
@@ -433,8 +442,9 @@
 		return TEST_RESULT_FAIL;
 	}
 
-	handle = memory_send(mb.send, FFA_MEM_LEND_SMC32, fragment_length,
-			     total_length, &ret);
+	handle = memory_send(mb.send, FFA_MEM_LEND_SMC32, constituents,
+			     constituents_count, remaining_constituent_count,
+			     fragment_length, total_length, &ret);
 
 	if (handle == FFA_MEMORY_HANDLE_INVALID) {
 		ERROR("Memory Share failed!\n");
@@ -633,12 +643,10 @@
 hypervisor_retrieve_request_test_helper(uint32_t mem_func,
 					bool multiple_receivers)
 {
-	struct ffa_memory_region_constituent sent_constituents[] = {{
-		.address = (void *)share_page,
-		.page_count = 1,
-		.reserved = 0,
-	}};
-	uint32_t sent_constituents_count = ARRAY_SIZE(sent_constituents);
+	static struct ffa_memory_region_constituent
+		sent_constituents[FRAGMENTED_SHARE_PAGE_COUNT];
+	uint32_t sent_constituents_count = 1;
+
 	__aligned(PAGE_SIZE) static uint8_t page[PAGE_SIZE * 2] = {0};
 	struct ffa_memory_region *hypervisor_retrieve_response =
 		(struct ffa_memory_region *)page;
@@ -673,6 +681,16 @@
 	uint32_t receiver_count =
 		multiple_receivers ? ARRAY_SIZE(receivers) : 1;
 
+	/* Add a page per constituent, so that we exhaust the size of a single
+	 * fragment (for testing). In a real world scenario, the whole region
+	 * could be described in a single constituent.
+	 */
+	for (uint32_t i = 0; i < sent_constituents_count; i++) {
+		sent_constituents[i].address = share_page + i * PAGE_SIZE;
+		sent_constituents[i].page_count = 1;
+		sent_constituents[i].reserved = 0;
+	}
+
 	CHECK_SPMC_TESTING_SETUP(1, 2, expected_sp_uuids);
 	GET_TFTF_MAILBOX(mb);
 
@@ -753,12 +771,6 @@
 		return TEST_RESULT_FAIL;
 	}
 
-	ret = ffa_rx_release();
-	if (is_ffa_call_error(ret)) {
-		ERROR("rx release failed: %d\n", ffa_error_code(ret));
-		return TEST_RESULT_FAIL;
-	}
-
 	return TEST_RESULT_SUCCESS;
 }