Merge "QcomModulePkg: Buffer overflow maybe occur when convert string from ASCII to Unicode"
diff --git a/QcomModulePkg/Library/BootLib/PartitionTableUpdate.c b/QcomModulePkg/Library/BootLib/PartitionTableUpdate.c
index 2bd789a..679df07 100644
--- a/QcomModulePkg/Library/BootLib/PartitionTableUpdate.c
+++ b/QcomModulePkg/Library/BootLib/PartitionTableUpdate.c
@@ -991,7 +991,15 @@
   DeviceDensity = (BlockIo->Media->LastBlock + 1) * BlockIo->Media->BlockSize;
   BlkSz = BlockIo->Media->BlockSize;
 
-  /* Verity that passed block has valid GPT primary header */
+  /* Verity that passed block has valid GPT primary header
+     * Sz is from mNumDataBytes and it will check at CmdDownload
+     * if it is mNumDataBytes > MaxDownLoadSize it will fail early and
+     * will not cause any oob
+     */
+  if (Sz <= BlkSz * 2) {
+    DEBUG ((EFI_D_ERROR, "Gpt Image size is invalid!\n"));
+    return FAILURE;
+  }
   PrimaryGptHdr = (Gpt + BlkSz);
   Ret = ParseGptHeader (&GptHeader, PrimaryGptHdr, DeviceDensity, BlkSz);
   if (Ret) {
@@ -1007,6 +1015,10 @@
   /* Back up partition is stored in the reverse order with back GPT, followed by
    * part entries, find the offset to back up GPT */
   Offset = (2 * PartEntryArrSz);
+  if (Sz < (Offset + (BlkSz * 3))) {
+    DEBUG ((EFI_D_ERROR, "Gpt Image size is invalid!!\n"));
+    return FAILURE;
+  }
   SecondaryGptHdr = Offset + BlkSz + PrimaryGptHdr;
   ParseSecondaryGpt = TRUE;
 
@@ -1081,7 +1093,7 @@
     return FAILURE;
   }
   FlashingGpt = 0;
-  gBS->SetMem ((VOID *)PrimaryGptHdr, Sz, 0x0);
+  gBS->SetMem ((VOID *)Gpt, Sz, 0x0);
 
   DEBUG ((EFI_D_ERROR, "Updated Partition Table Successfully\n"));
   return SUCCESS;