ArmPkg: MmCommunicationDxe: Update MM communicate `MessageLength` check

[mirror_edk2.git] / ArmPkg / Drivers / MmCommunicationDxe / MmCommunication.c

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