[PATCH 2/2] MdeModulePkg PiSmmIpl: Handle CommSize OPTIONAL case

[Star Zeng <star.zeng@intel.com>]

Handle CommSize OPTIONAL case for SmmCommunicate.
And return EFI_ACCESS_DENIED when CommunicationBuffer
is not valid for SMM to access.

Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Liming Gao <liming.gao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Star Zeng <star.zeng@intel.com>
---
 MdeModulePkg/Core/PiSmmCore/PiSmmCore.c |  2 +-
 MdeModulePkg/Core/PiSmmCore/PiSmmIpl.c  | 61 +++++++++++++++++++++------------
 2 files changed, 40 insertions(+), 23 deletions(-)

diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
index a7467aca2012..4c1e3e719524 100644
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
@@ -544,7 +544,7 @@ SmmEntryPoint (
         // return EFI_INVALID_PARAMETER
         //
         gSmmCorePrivate->CommunicationBuffer = NULL;
-        gSmmCorePrivate->ReturnStatus = EFI_INVALID_PARAMETER;
+        gSmmCorePrivate->ReturnStatus = EFI_ACCESS_DENIED;
       } else {
         CommunicateHeader = (EFI_SMM_COMMUNICATE_HEADER *)CommunicationBuffer;
         BufferSize -= OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data);
diff --git a/MdeModulePkg/Core/PiSmmCore/PiSmmIpl.c b/MdeModulePkg/Core/PiSmmCore/PiSmmIpl.c
index 2601275ab85f..31d2c9e45e1f 100644
--- a/MdeModulePkg/Core/PiSmmCore/PiSmmIpl.c
+++ b/MdeModulePkg/Core/PiSmmCore/PiSmmIpl.c
@@ -440,37 +440,55 @@ SmmBase2GetSmstLocation (
   after SetVirtualAddressMap().
 
   @param[in] This                The EFI_SMM_COMMUNICATION_PROTOCOL instance.
-  @param[in, out] CommBuffer          A pointer to the buffer to convey into SMRAM.
-  @param[in, out] CommSize            The size of the data buffer being passed in.On exit, the size of data
+  @param[in, out] CommBuffer     A pointer to the buffer to convey into SMRAM.
+  @param[in, out] CommSize       The size of the data buffer being passed in. On exit, the size of data
                                  being returned. Zero if the handler does not wish to reply with any data.
+                                 This parameter is optional and may be NULL.
 
   @retval EFI_SUCCESS            The message was successfully posted.
   @retval EFI_INVALID_PARAMETER  The CommBuffer was NULL.
+  @retval EFI_BAD_BUFFER_SIZE    The buffer is too large for the MM implementation.
+                                 If this error is returned, the MessageLength field
+                                 in the CommBuffer header or the integer pointed by
+                                 CommSize, are updated to reflect the maximum payload
+                                 size the implementation can accommodate.
+  @retval EFI_ACCESS_DENIED      The CommunicateBuffer parameter or CommSize parameter,
+                                 if not omitted, are in address range that cannot be
+                                 accessed by the MM environment.
+
 **/
 EFI_STATUS
 EFIAPI
 SmmCommunicationCommunicate (
   IN CONST EFI_SMM_COMMUNICATION_PROTOCOL  *This,
   IN OUT VOID                              *CommBuffer,
-  IN OUT UINTN                             *CommSize
+  IN OUT UINTN                             *CommSize OPTIONAL
   )
 {
   EFI_STATUS                  Status;
   EFI_SMM_COMMUNICATE_HEADER  *CommunicateHeader;
   BOOLEAN                     OldInSmm;
+  UINTN                       TempCommSize;
 
   //
   // Check parameters
   //
-  if ((CommBuffer == NULL) || (CommSize == NULL)) {
+  if (CommBuffer == NULL) {
     return EFI_INVALID_PARAMETER;
   }
 
-  //
-  // CommSize must hold HeaderGuid and MessageLength
-  //
-  if (*CommSize < OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)) {
-    return EFI_INVALID_PARAMETER;
+  CommunicateHeader = (EFI_SMM_COMMUNICATE_HEADER *) CommBuffer;
+
+  if (CommSize == NULL) {
+    TempCommSize = OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data) + CommunicateHeader->MessageLength;
+  } else {
+    TempCommSize = *CommSize;
+    //
+    // CommSize must hold HeaderGuid and MessageLength
+    //
+    if (TempCommSize < OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)) {
+      return EFI_INVALID_PARAMETER;
+    }
   }
 
   //
@@ -481,7 +499,7 @@ SmmCommunicationCommunicate (
     // Put arguments for Software SMI in gSmmCorePrivate
     //
     gSmmCorePrivate->CommunicationBuffer = CommBuffer;
-    gSmmCorePrivate->BufferSize          = *CommSize;
+    gSmmCorePrivate->BufferSize          = TempCommSize;
 
     //
     // Generate Software SMI
@@ -494,15 +512,17 @@ SmmCommunicationCommunicate (
     //
     // Return status from software SMI 
     //
-    *CommSize = gSmmCorePrivate->BufferSize;
+    if (CommSize != NULL) {
+      *CommSize = gSmmCorePrivate->BufferSize;
+    }
     return gSmmCorePrivate->ReturnStatus;
   }
 
   //
   // If we are in SMM, then the execution mode must be physical, which means that
   // OS established virtual addresses can not be used.  If SetVirtualAddressMap()
-  // has been called, then a direct invocation of the Software SMI is not 
-  // not allowed so return EFI_INVALID_PARAMETER.
+  // has been called, then a direct invocation of the Software SMI is not allowed,
+  // so return EFI_INVALID_PARAMETER.
   //
   if (EfiGoneVirtual()) {
     return EFI_INVALID_PARAMETER;
@@ -524,20 +544,17 @@ SmmCommunicationCommunicate (
   //
   // Before SetVirtualAddressMap(), we are in SMM or SMRAM is open and unlocked, call SmiManage() directly.
   //
-  CommunicateHeader = (EFI_SMM_COMMUNICATE_HEADER *)CommBuffer;
-  *CommSize -= OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data);
+  TempCommSize -= OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data);
   Status = gSmmCorePrivate->Smst->SmiManage (
                                     &CommunicateHeader->HeaderGuid, 
                                     NULL, 
                                     CommunicateHeader->Data, 
-                                    CommSize
+                                    &TempCommSize
                                     );
-
-  //
-  // Update CommunicationBuffer, BufferSize and ReturnStatus
-  // Communicate service finished, reset the pointer to CommBuffer to NULL
-  //
-  *CommSize += OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data);
+  TempCommSize += OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data);
+  if (CommSize != NULL) {
+    *CommSize = TempCommSize;
+  }
 
   //
   // Restore original InSmm state
-- 
2.7.0.windows.1