[PATCH v3 6/6] ArmPkg: MmCommunicationDxe: Update MM communicate `MessageLength` check

["Kun Qin" <kuqin12@gmail.com>]

REF: https://bugzilla.tianocore.org/show_bug.cgi?id=3751

Current MM communicate routine from ArmPkg would conduct few checks prior
to proceeding with SMC calls. However, the inspection step is different
from PI specification.

This patch updated MM communicate input argument inspection routine to
assure that "if the `MessageLength` is zero, or too large for the MM
implementation to manage, the MM implementation must update the
`MessageLength` to reflect the size of the `Data` buffer that it can
tolerate", as described by `EFI_MM_COMMUNICATION_PROTOCOL.Communicate()`
section in PI specification.

Cc: Leif Lindholm <leif@nuviainc.com>
Cc: Ard Biesheuvel <ardb+tianocore@kernel.org>
Cc: Bret Barkelew <Bret.Barkelew@microsoft.com>
Cc: Michael Kubacki <michael.kubacki@microsoft.com>
Cc: Sami Mujawar <sami.mujawar@arm.com>

Signed-off-by: Kun Qin <kuqin12@gmail.com>
Reviewed-by: Sami Mujawar <sami.mujawar@arm.com>
---

Notes:
    v2:
    - Splitting patch into 4 of 4 [Ard]
    - Uncrustify style update
    
    v3:
    - Added cc entry and reviewed-by tag [Sami]

 ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
index 2f89b7c5b6c4..85d9034555f0 100644
--- a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
+++ b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
@@ -92,6 +92,7 @@ MmCommunication2Communicate (
     return EFI_INVALID_PARAMETER;
   }
 
+  Status            = EFI_SUCCESS;
   CommunicateHeader = CommBufferVirtual;
   // CommBuffer is a mandatory parameter. Hence, Rely on
   // MessageLength + Header to ascertain the
@@ -109,28 +110,33 @@ MmCommunication2Communicate (
         (*CommSize > mNsCommBuffMemRegion.Length))
     {
       *CommSize = mNsCommBuffMemRegion.Length;
-      return EFI_BAD_BUFFER_SIZE;
+      Status    = EFI_BAD_BUFFER_SIZE;
     }
 
     //
     // CommSize should cover at least MessageLength + sizeof (EFI_MM_COMMUNICATE_HEADER);
     //
     if (*CommSize < BufferSize) {
-      return EFI_INVALID_PARAMETER;
+      Status = EFI_INVALID_PARAMETER;
     }
   }
 
   //
-  // If the buffer size is 0 or greater than what can be tolerated by the MM
+  // If the message length is 0 or greater than what can be tolerated by the MM
   // environment then return the expected size.
   //
-  if ((BufferSize == 0) ||
+  if ((CommunicateHeader->MessageLength == 0) ||
       (BufferSize > mNsCommBuffMemRegion.Length))
   {
     CommunicateHeader->MessageLength = mNsCommBuffMemRegion.Length -
                                        sizeof (CommunicateHeader->HeaderGuid) -
                                        sizeof (CommunicateHeader->MessageLength);
-    return EFI_BAD_BUFFER_SIZE;
+    Status = EFI_BAD_BUFFER_SIZE;
+  }
+
+  // MessageLength or CommSize check has failed, return here.
+  if (EFI_ERROR (Status)) {
+    return Status;
   }
 
   // SMC Function ID
-- 
2.34.1.windows.1