core: fix fragmented memory retrieve

add spmc_frag_retrieve for fragmented memory retrieve

Signed-off-by: Jianqiang Mei <meijianqiang.mjq@alibaba-inc.com>

core/arch/arm/kernel/thread_spmc.c

@@ -1383,7 +1383,7 @@ static uint16_t spmc_get_id(void)
 	return args.a2;
 }
 
-static struct ffa_mem_transaction *spmc_retrieve_req(uint64_t cookie)
+static struct ffa_mem_transaction *spmc_retrieve_req(uint64_t cookie, uint32_t* total_length, uint32_t* fragment_length)
 {
 	struct ffa_mem_transaction *trans_descr = nw_rxtx.tx;
 	struct ffa_mem_access *acc_descr_array = NULL;
@@ -1423,7 +1423,8 @@ static struct ffa_mem_transaction *spmc_retrieve_req(uint64_t cookie)
 			     cookie, args.a0);
 		return NULL;
 	}
-
+	*total_length = args.a1;
+	*fragment_length = args.a2;
 	return nw_rxtx.rx;
 }
 
@@ -1465,15 +1466,43 @@ static int set_pages(struct ffa_address_range *regions,
 	return 0;
 }
 
+static struct ffa_mem_transaction *spmc_frag_retrieve(uint64_t cookie, uint32_t next_fragment_offset, uint32_t* fragment_length)
+{
+	struct thread_smc_args args = {
+		.a0 = FFA_MEM_FRAG_RX,
+		.a1 = (uint32_t)(cookie),	/* handle high addr */
+		.a2 = (uint32_t)(cookie>>32),	/* handle low addr */
+		.a3 = next_fragment_offset,	/* offset */
+		.a4 = 0x8001,	/* sender_vm_id */
+	};
+	thread_smccc(&args);
+	if (args.a0 != FFA_MEM_FRAG_TX) {
+		if (args.a0 == FFA_ERROR)
+			EMSG("Failed to fetch cookie %#"PRIx64" error code %d",
+			     cookie, (int)args.a2);
+		else
+			EMSG("Failed to fetch cookie %#"PRIx64" a0 %#"PRIx64,
+			     cookie, args.a0);
+		return NULL;
+	}
+	*fragment_length = args.a3;
+	return nw_rxtx.rx;
+}
+
 struct mobj_ffa *thread_spmc_populate_mobj_from_rx(uint64_t cookie)
 {
 	struct mobj_ffa *ret = NULL;
+	struct ffa_mem_transaction *rx_buffer = NULL;
 	struct ffa_mem_transaction *retrieve_desc = NULL;
 	struct ffa_mem_access *descr_array = NULL;
 	struct ffa_mem_region *descr = NULL;
 	struct mobj_ffa *mf = NULL;
 	unsigned int num_pages = 0;
 	unsigned int offs = 0;
+	unsigned int next_fragment_offset = 0;
+	unsigned int length = 0;
+	unsigned int fragment_length = 0;
+	uint64_t handle = 0;
 	struct thread_smc_args ffa_rx_release_args = {
 		.a0 = FFA_RX_RELEASE
 	};
@@ -1482,13 +1511,31 @@ struct mobj_ffa *thread_spmc_populate_mobj_from_rx(uint64_t cookie)
 	 * OP-TEE is only supporting a single mem_region while the
 	 * specification allows for more than one.
 	 */
-	retrieve_desc = spmc_retrieve_req(cookie);
-	if (!retrieve_desc) {
+	rx_buffer = spmc_retrieve_req(cookie, &length, &fragment_length);
+	if (!rx_buffer) {
 		EMSG("Failed to retrieve cookie from rx buffer %#"PRIx64,
 		     cookie);
 		return NULL;
 	}
 
+	/* Allocate the enough size buffer to copy the fragment descriptor*/
+	retrieve_desc = aligned_alloc(SMALL_PAGE_SIZE, SMALL_PAGE_SIZE*16);
+	memcpy(retrieve_desc, rx_buffer, SMALL_PAGE_SIZE);
+
+	/* Retrieve the remaining fragments. */
+	next_fragment_offset = fragment_length;
+	handle = rx_buffer->global_handle;
+	while (next_fragment_offset < length) {
+		DMSG("Retrieve again, frag offset: %x; total: %x.\n",
+			     next_fragment_offset, length);
+		rx_buffer = spmc_frag_retrieve(handle, next_fragment_offset, &fragment_length);
+		if (rx_buffer != NULL) {
+			memcpy((void *)((uint64_t)retrieve_desc + (uint64_t)next_fragment_offset),
+				 rx_buffer, SMALL_PAGE_SIZE);
+		}
+		next_fragment_offset += fragment_length;
+	}
+
 	descr_array = retrieve_desc->mem_access_array;
 	offs = READ_ONCE(descr_array->region_offs);
 	descr = (struct ffa_mem_region *)((vaddr_t)retrieve_desc + offs);
@@ -1507,6 +1554,7 @@ struct mobj_ffa *thread_spmc_populate_mobj_from_rx(uint64_t cookie)
 	ret = mf;
 
 out:
+	free(retrieve_desc);
 	/* Release RX buffer after the mem retrieve request. */
 	thread_smccc(&ffa_rx_release_args);