2 files changed, 9 insertions, 2 deletions
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index b8aa9e1769..2248a8c5ee 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -337,7 +337,7 @@ This function supports reading a CPU Save State register in SMBase relocation ha
 
 @retval EFI_SUCCESS           The register was read from Save State.
 @retval EFI_NOT_FOUND         The register is not defined for the Save State of Processor.
-@retval EFI_INVALID_PARAMETER  This or Buffer is NULL.
+@retval EFI_INVALID_PARAMETER Buffer is NULL, or Width does not meet requirement per Register type.
 
 **/
 EFI_STATUS
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmramSaveState.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmramSaveState.c
index 661cc51f36..fc418c2500 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmramSaveState.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmramSaveState.c
@@ -343,7 +343,7 @@ ReadSaveStateRegisterByIndex (
 
   @retval EFI_SUCCESS           The register was read from Save State.
   @retval EFI_NOT_FOUND         The register is not defined for the Save State of Processor.
-  @retval EFI_INVALID_PARAMETER  This or Buffer is NULL.
+  @retval EFI_INVALID_PARAMETER Buffer is NULL, or Width does not meet requirement per Register type.
 
 **/
 EFI_STATUS
@@ -419,6 +419,13 @@ ReadSaveStateRegister (
     }
 
     //
+    // Make sure the incoming buffer is large enough to hold IoInfo before accessing
+    //
+    if (Width < sizeof (EFI_SMM_SAVE_STATE_IO_INFO)) {
+      return EFI_INVALID_PARAMETER;
+    }
+
+    //
     // Zero the IoInfo structure that will be returned in Buffer
     //
     IoInfo = (EFI_SMM_SAVE_STATE_IO_INFO *)Buffer;