From aa76195a971bf03c3db7d6cc60db6d8fde2ba7b9 Mon Sep 17 00:00:00 2001
From: Download-Fritz <download-fritz@outlook.com>
Date: Sat, 6 Apr 2019 13:19:19 +0200
Subject: [PATCH] OcAppleDiskImageLib: Initialize a caller-provided context

---
 Include/Library/OcAppleDiskImageLib.h         |  2 +-
 .../OcAppleDiskImageLib/OcAppleDiskImageLib.c | 23 ++++++-------------
 TestsUser/DiskImage/DiskImage.c               |  8 +++----
 3 files changed, 12 insertions(+), 21 deletions(-)

diff --git a/Include/Library/OcAppleDiskImageLib.h b/Include/Library/OcAppleDiskImageLib.h
index 4511bd79..cd6d13a5 100644
--- a/Include/Library/OcAppleDiskImageLib.h
+++ b/Include/Library/OcAppleDiskImageLib.h
@@ -36,7 +36,7 @@ BOOLEAN
 OcAppleDiskImageInitializeContext (
   IN  VOID                         *Buffer,
   IN  UINTN                        BufferLength,
-  OUT OC_APPLE_DISK_IMAGE_CONTEXT  **Context
+  OUT OC_APPLE_DISK_IMAGE_CONTEXT  *Context
   );
 
 VOID
diff --git a/Library/OcAppleDiskImageLib/OcAppleDiskImageLib.c b/Library/OcAppleDiskImageLib/OcAppleDiskImageLib.c
index 01158a86..51a42500 100644
--- a/Library/OcAppleDiskImageLib/OcAppleDiskImageLib.c
+++ b/Library/OcAppleDiskImageLib/OcAppleDiskImageLib.c
@@ -26,11 +26,10 @@ BOOLEAN
 OcAppleDiskImageInitializeContext (
   IN  VOID                         *Buffer,
   IN  UINTN                        BufferLength,
-  OUT OC_APPLE_DISK_IMAGE_CONTEXT  **Context
+  OUT OC_APPLE_DISK_IMAGE_CONTEXT  *Context
   )
 {
   BOOLEAN                     Result;
-  OC_APPLE_DISK_IMAGE_CONTEXT *DmgContext;
   UINTN                       TrailerOffset;
   UINT8                       *BufferBytes;
   UINT8                       *BufferBytesCurrent;
@@ -162,19 +161,12 @@ OcAppleDiskImageInitializeContext (
     return FALSE;
   }
 
-  DmgContext = AllocateZeroPool (sizeof (*DmgContext));
-  if (DmgContext == NULL) {
-    FreePool (DmgBlocks);
-    return FALSE;
-  }
-
-  DmgContext->Buffer      = BufferBytes;
-  DmgContext->Length      = (TrailerOffset + sizeof (*Trailer));
-  DmgContext->BlockCount  = DmgBlockCount;
-  DmgContext->Blocks      = DmgBlocks;
-  DmgContext->SectorCount = SectorCount;
-
-  *Context = DmgContext;
+  Context->Buffer        = BufferBytes;
+  Context->Length        = (TrailerOffset + sizeof (*Trailer));
+  Context->BlockCount    = DmgBlockCount;
+  Context->Blocks        = DmgBlocks;
+  Context->SectorCount   = SectorCount;
+  Context->BlockIoHandle = NULL;
 
   return TRUE;
 }
@@ -193,7 +185,6 @@ OcAppleDiskImageFreeContext (
   }
 
   FreePool (Context->Blocks);
-  FreePool (Context);
 }
 
 BOOLEAN
diff --git a/TestsUser/DiskImage/DiskImage.c b/TestsUser/DiskImage/DiskImage.c
index 750f0a22..4c82c0dd 100644
--- a/TestsUser/DiskImage/DiskImage.c
+++ b/TestsUser/DiskImage/DiskImage.c
@@ -62,7 +62,7 @@ int main (int argc, char *argv[]) {
 
     BOOLEAN                     Result;
     EFI_STATUS                  Status;
-    OC_APPLE_DISK_IMAGE_CONTEXT *DmgContext;
+    OC_APPLE_DISK_IMAGE_CONTEXT DmgContext;
 
     Result = OcAppleDiskImageInitializeContext (Dmg, DmgSize, &DmgContext);
     if (!Result) {
@@ -70,14 +70,14 @@ int main (int argc, char *argv[]) {
       continue;
     }
 
-    UncompSize = (DmgContext->SectorCount * APPLE_DISK_IMAGE_SECTOR_SIZE);
+    UncompSize = (DmgContext.SectorCount * APPLE_DISK_IMAGE_SECTOR_SIZE);
     UncompDmg  = malloc (UncompSize);
     if (UncompDmg == NULL) {
       printf ("DMG data allocation failed.\n");
       continue;
     }
 
-    Result = OcAppleDiskImageRead (DmgContext, 0, UncompSize, UncompDmg);
+    Result = OcAppleDiskImageRead (&DmgContext, 0, UncompSize, UncompDmg);
     if (!Result) {
       printf ("DMG read error\n");
       continue;
@@ -106,7 +106,7 @@ int main (int argc, char *argv[]) {
 
     printf ("Decompressed the entire DMG...\n");
 
-    OcAppleDiskImageFreeContext (DmgContext);
+    OcAppleDiskImageFreeContext (&DmgContext);
 
     printf ("Success...\n");
   }