From 14ff66b9fe02cdedcaf4860aee3ed5223b3d5bf0 Mon Sep 17 00:00:00 2001
From: mikolajkolakowski <Mikolaj.Kolakowski@intel.com>
Date: Tue, 4 Jun 2024 14:42:36 +0200
Subject: [PATCH] PBR crash fix for 2_0 code base

---
 DcpmPkg/common/Pbr.c   |  7 +++++++
 DcpmPkg/common/PbrOs.c | 11 +++++++++++
 2 files changed, 18 insertions(+)

diff --git a/DcpmPkg/common/Pbr.c b/DcpmPkg/common/Pbr.c
index a2e78362..939a9dcf 100644
--- a/DcpmPkg/common/Pbr.c
+++ b/DcpmPkg/common/Pbr.c
@@ -977,6 +977,13 @@ PbrComposeSession(
   }
 
   pPbrMainHeader = (PbrHeader *)pContext->PbrMainHeader;
+
+  if (pPbrMainHeader == NULL) 
+  {
+    // bailout if PBR session is not initialized
+    return EFI_NOT_FOUND;
+  }
+  
   ZeroMem(&pPbrMainHeader->PartitionTable, sizeof(PbrPartitionTable));
   BufferSize = sizeof(PbrHeader);
 
diff --git a/DcpmPkg/common/PbrOs.c b/DcpmPkg/common/PbrOs.c
index 1fccd4d6..c1b523f7 100644
--- a/DcpmPkg/common/PbrOs.c
+++ b/DcpmPkg/common/PbrOs.c
@@ -100,6 +100,7 @@ VOID DeserializePbrMode(UINT32 *pMode, UINT32 defaultMode);
   if(NULL == buffer) { \
     NVDIMM_ERR("Failed to allocate memory for deserializing buffer\n"); \
   } \
+  else \
   if (1 != fread(buffer, size, 1, pFile)) \
   { \
     NVDIMM_ERR("Failed to read the PBR file: %s\n", file); \
@@ -222,7 +223,17 @@ EFI_STATUS PbrDeserializeCtx(
     if (PBR_INVALID_SIG != ctx->PartitionContexts[CtxIndex].PartitionSig) {
       AsciiSPrint(pbr_filename, sizeof(pbr_filename), "%x.pbr", ctx->PartitionContexts[CtxIndex].PartitionSig);
       AsciiSPrint(pbr_dir, sizeof(pbr_dir), "%s%s", PBR_TMP_DIR, pbr_filename);
+
+      ctx->PartitionContexts[CtxIndex].PartitionData = NULL; // initialize buffer to NULL to allow proper error handling in case of fail to open pbr_dir file
       DeserializeBufferEx(pbr_dir, ctx->PartitionContexts[CtxIndex].PartitionData, ctx->PartitionContexts[CtxIndex].PartitionSize);
+
+      if(ctx->PartitionContexts[CtxIndex].PartitionData == NULL)
+      {
+          // we do not free already allocated memory because OS will free it after process exit
+          NVDIMM_ERR("PBR context file corrupted, please remove "PBR_TMP_DIR PBR_CTX_FILE_NAME"\n");
+          ReturnCode = EFI_END_OF_FILE;
+          goto Finish;
+      }
     }
   }