From 65c5eba2e10c5accdcd7d995df8ee4274d51de27 Mon Sep 17 00:00:00 2001
From: Ken Lautner <klautner@microsoft.com>
Date: Wed, 20 Sep 2023 16:30:06 -0700
Subject: [PATCH] Some fixes in the MmTestPointMmiHandlerGetDataByOffset
 function

---
 .../TestPointLib/MmTestPointCommunication.c   | 22 ++++++++++---------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/MinPlatformPkg/Test/Library/TestPointLib/MmTestPointCommunication.c b/MinPlatformPkg/Test/Library/TestPointLib/MmTestPointCommunication.c
index 9f096100b4..87a3309c96 100644
--- a/MinPlatformPkg/Test/Library/TestPointLib/MmTestPointCommunication.c
+++ b/MinPlatformPkg/Test/Library/TestPointLib/MmTestPointCommunication.c
@@ -209,10 +209,9 @@ MmiHandlerTestPointCopyData (
 **/
 VOID
 MmTestPointMmiHandlerGetDataByOffset (
-  IN VOID     *MmiHandlerTestPointParameterGetDataByOffset
+  IN MMI_HANDLER_TEST_POINT_PARAMETER_GET_DATA_BY_OFFSET     *MmiHandlerTestPointParameterGetDataByOffset
   )
 {
-  MMI_HANDLER_TEST_POINT_PARAMETER_GET_DATA_BY_OFFSET    *MmiHandlerTestPointGetDataByOffset;
   VOID                                                   *Data;
   UINTN                                                  DataSize;
   EFI_STATUS                                             Status;
@@ -224,8 +223,6 @@ MmTestPointMmiHandlerGetDataByOffset (
     return;
   }
 
-  MmiHandlerTestPointGetDataByOffset = (MMI_HANDLER_TEST_POINT_PARAMETER_GET_DATA_BY_OFFSET *)MmiHandlerTestPointParameterGetDataByOffset;
-
   //
   // Sanity check
   //
@@ -239,17 +236,22 @@ MmTestPointMmiHandlerGetDataByOffset (
   DataSize = 0;
   Status = GetAllMmTestPointData (&DataSize, NULL);
   if (Status != EFI_BUFFER_TOO_SMALL) {
-    MmiHandlerTestPointGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_NOT_FOUND;
+    MmiHandlerTestPointParameterGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_NOT_FOUND;
     goto Done;
   }
   Data = AllocatePool (DataSize);
   if (Data == NULL) {
-    MmiHandlerTestPointGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_OUT_OF_RESOURCES;
+    MmiHandlerTestPointParameterGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_OUT_OF_RESOURCES;
     goto Done;
   }
   Status = GetAllMmTestPointData (&DataSize, Data);
   if (EFI_ERROR(Status)) {
-    MmiHandlerTestPointGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)Status;
+    MmiHandlerTestPointParameterGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)Status;
+    goto Done;
+  }
+
+  if (DataSize > MmiHandlerTestPointParameterGetDataByOffset->DataSize) {
+    DEBUG((DEBUG_ERROR, "[%a] - The Datasize we are going to copy over is larger than expected.  Aborting.\n", __func__));
     goto Done;
   }
 
@@ -261,12 +263,12 @@ MmTestPointMmiHandlerGetDataByOffset (
   SpeculationBarrier ();
 
   CopyMem (
-    MmiHandlerTestPointGetDataByOffset->Data,
+    MmiHandlerTestPointParameterGetDataByOffset->Data,
     Data,
     DataSize
     );
 
-  MmiHandlerTestPointGetDataByOffset->Header.ReturnStatus = 0;
+  MmiHandlerTestPointParameterGetDataByOffset->Header.ReturnStatus = 0;
 
 Done:
   if (Data != NULL) {
@@ -342,7 +344,7 @@ MmTestPointMmiHandler (
       DEBUG((DEBUG_INFO, "MmTestPointMmiHandler: MM communication buffer size invalid!\n"));
       return EFI_SUCCESS;
     }
-    MmTestPointMmiHandlerGetDataByOffset(CommBuffer);
+    MmTestPointMmiHandlerGetDataByOffset((MMI_HANDLER_TEST_POINT_PARAMETER_GET_DATA_BY_OFFSET *)(UINTN)CommBuffer);
     break;
   default:
     break;