1 files changed, 16 insertions, 0 deletions
diff --git a/PrmPkg/Samples/PrmSampleContextBufferModule/Library/DxeContextBufferModuleConfigLib/DxeContextBufferModuleConfigLib.c b/PrmPkg/Samples/PrmSampleContextBufferModule/Library/DxeContextBufferModuleConfigLib/DxeContextBufferModuleConfigLib.c
index 3bf5beba7d..b0b12c012a 100644
--- a/PrmPkg/Samples/PrmSampleContextBufferModule/Library/DxeContextBufferModuleConfigLib/DxeContextBufferModuleConfigLib.c
+++ b/PrmPkg/Samples/PrmSampleContextBufferModule/Library/DxeContextBufferModuleConfigLib/DxeContextBufferModuleConfigLib.c
@@ -151,10 +151,26 @@ ContextBufferModuleConfigLibConstructor (
   //
   // Allocate and populate the context buffer
   //
+#ifdef ALLOCATE_CONTEXT_BUFFER_IN_FW
+  //
+  // The context buffer allocated in FW will continue being used at OS runtime so it must
+  // be a runtime services data buffer.
+  //
   // This sample module uses a single context buffer for all the handlers
   // Todo: This can be done more elegantly in the future. Likely though a library service.
   //
   PrmContextBuffer = AllocateRuntimeZeroPool (sizeof (*PrmContextBuffer));
+#else
+  //
+  // This context buffer is not actually used by PRM handler at OS runtime. The OS will allocate
+  // the actual context buffer passed to the PRM handler.
+  //
+  // This context buffer is used internally in the firmware to associate a PRM handler with a
+  // a static data buffer and a runtime MMIO ranges array so those can be placed into the
+  // PRM_HANDLER_INFORMATION_STRUCT and PRM_MODULE_INFORMATION_STRUCT respectively for the PRM handler.
+  //
+  PrmContextBuffer = AllocateZeroPool (sizeof (*PrmContextBuffer));
+#endif
   ASSERT (PrmContextBuffer != NULL);
   if (PrmContextBuffer == NULL) {
     Status = EFI_OUT_OF_RESOURCES;