[CHERRY-PICK] Global: fix ArmFfaLibRun() caller couldn't get ret-args

When ArmFfaLibDirectMsgReq(2) is preempted, caller of these functions
should resume it works via ArmFfaLibRun() and the secure partition
will be return with FFA_DIRECT_MSG_RESP(2) with return arguments.

However, since ArmFfaLibRun() gets its return in its stack variable,
So caller of ArmFfaLibRun() doesn't get the return arguments from
secure partition.

To resolve this, add output parameter to ArmFfaLibRun() to
receive return arguments.

Continuous-integration-options: PatchCheck.ignore-multi-package
Fixes: 5d1b38d ("ArmPkg: Add ArmFfaLib used in Dxe driver")
Reported-by: Mariam Elshakfy <[email protected]>
Signed-off-by: Yeoreum Yun <[email protected]>
(cherry picked from commit d8e875e62505e15a48f82ccb49db42be9fda519f)

Author: LeviYeoReum

MdeModulePkg/Library/ArmFfaLib/ArmFfaCommon.c

@@ -553,6 +553,7 @@ ArmFfaLibPartitionInfoGet (
 
   @param [in]   PartId       Partition id
   @param [in]   CpuNumber    Cpu number in partition
+  @param [out]  DirectMsgArg return arguments for direct msg resp/resp2
 
   @retval EFI_SUCCESS
   @retval Other              Error
@@ -561,10 +562,12 @@ ArmFfaLibPartitionInfoGet (
 EFI_STATUS
 EFIAPI
 ArmFfaLibRun (
-  IN  UINT16  PartId,
-  IN  UINT16  CpuNumber
+  IN  UINT16           PartId,
+  IN  UINT16           CpuNumber,
+  OUT DIRECT_MSG_ARGS  *DirectMsgArg OPTIONAL
   )
 {
+  EFI_STATUS    Status;
   ARM_FFA_ARGS  FfaArgs;
 
   ZeroMem (&FfaArgs, sizeof (ARM_FFA_ARGS));
@@ -574,7 +577,39 @@ ArmFfaLibRun (
 
   ArmCallFfa (&FfaArgs);
 
-  return FfaArgsToEfiStatus (&FfaArgs);
+  Status = FfaArgsToEfiStatus (&FfaArgs);
+  if (EFI_ERROR (Status)) {
+    return Status;
+  }
+
+  if (DirectMsgArg != NULL) {
+    ZeroMem (DirectMsgArg, sizeof (DIRECT_MSG_ARGS));
+
+    if (FfaArgs.Arg0 == ARM_FID_FFA_MSG_SEND_DIRECT_RESP) {
+      DirectMsgArg->Arg0 = FfaArgs.Arg3;
+      DirectMsgArg->Arg1 = FfaArgs.Arg4;
+      DirectMsgArg->Arg2 = FfaArgs.Arg5;
+      DirectMsgArg->Arg3 = FfaArgs.Arg6;
+      DirectMsgArg->Arg4 = FfaArgs.Arg7;
+    } else if (FfaArgs.Arg0 == ARM_FID_FFA_MSG_SEND_DIRECT_RESP2) {
+      DirectMsgArg->Arg0  = FfaArgs.Arg4;
+      DirectMsgArg->Arg1  = FfaArgs.Arg5;
+      DirectMsgArg->Arg2  = FfaArgs.Arg6;
+      DirectMsgArg->Arg3  = FfaArgs.Arg7;
+      DirectMsgArg->Arg4  = FfaArgs.Arg8;
+      DirectMsgArg->Arg5  = FfaArgs.Arg9;
+      DirectMsgArg->Arg6  = FfaArgs.Arg10;
+      DirectMsgArg->Arg7  = FfaArgs.Arg11;
+      DirectMsgArg->Arg8  = FfaArgs.Arg12;
+      DirectMsgArg->Arg9  = FfaArgs.Arg13;
+      DirectMsgArg->Arg10 = FfaArgs.Arg14;
+      DirectMsgArg->Arg11 = FfaArgs.Arg15;
+      DirectMsgArg->Arg12 = FfaArgs.Arg16;
+      DirectMsgArg->Arg13 = FfaArgs.Arg17;
+    }
+  }
+
+  return Status;
 }
 
 /**

MdePkg/Include/Library/ArmFfaLib.h

@@ -301,6 +301,7 @@ ArmFfaLibSpmIdGet (
 
   @param [in]   PartId       Partition id
   @param [in]   CpuNumber    Cpu number in partition
+  @param [out]  DirectMsgArg return arguments for direct msg resp/resp2
 
   @retval EFI_SUCCESS
   @retval Other              Error
@@ -309,8 +310,9 @@ ArmFfaLibSpmIdGet (
 EFI_STATUS
 EFIAPI
 ArmFfaLibRun (
-  IN  UINT16  PartId,
-  IN  UINT16  CpuNumber
+  IN  UINT16           PartId,
+  IN  UINT16           CpuNumber,
+  OUT DIRECT_MSG_ARGS  *DirectMsgArg OPTIONAL
   );
 
 /**