From 6b2186abafec599184ee4a4aed2062d986800490 Mon Sep 17 00:00:00 2001
From: Download-Fritz <download-fritz@outlook.com>
Date: Mon, 18 Mar 2019 19:58:16 +0100
Subject: [PATCH] OcAppleKernelLib and OcDataHubLib: Security enhancements

---
 Library/OcAppleKernelLib/Prelinked.c | 10 ++++++++--
 Library/OcDataHubLib/OcDataHubLib.c  | 17 ++++++++++++++---
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/Library/OcAppleKernelLib/Prelinked.c b/Library/OcAppleKernelLib/Prelinked.c
index 6f6f1dad..7c7fb781 100644
--- a/Library/OcAppleKernelLib/Prelinked.c
+++ b/Library/OcAppleKernelLib/Prelinked.c
@@ -66,6 +66,9 @@ PrelinkedContextInit (
   if (Context->PrelinkedInfoSegment == NULL) {
     return EFI_NOT_FOUND;
   }
+  if (Context->PrelinkedInfoSegment->FileOffset > MAX_UINT32) {
+    return EFI_UNSUPPORTED;
+  }
 
   Context->PrelinkedInfoSection = MachoGetSectionByName64 (
     &Context->PrelinkedMachContext,
@@ -75,6 +78,9 @@ PrelinkedContextInit (
   if (Context->PrelinkedInfoSection == NULL) {
     return EFI_NOT_FOUND;
   }
+  if (Context->PrelinkedInfoSection->Size > MAX_UINT32) {
+    return EFI_UNSUPPORTED;
+  }
 
   Context->PrelinkedTextSegment = MachoGetSegmentByName64 (
     &Context->PrelinkedMachContext,
@@ -101,7 +107,7 @@ PrelinkedContextInit (
     return EFI_OUT_OF_RESOURCES;
   }
 
-  Context->PrelinkedInfoDocument = XmlDocumentParse (Context->PrelinkedInfo, Context->PrelinkedInfoSection->Size, TRUE);
+  Context->PrelinkedInfoDocument = XmlDocumentParse (Context->PrelinkedInfo, (UINT32)Context->PrelinkedInfoSection->Size, TRUE);
   if (Context->PrelinkedInfoDocument == NULL) {
     PrelinkedContextFree (Context);
     return EFI_INVALID_PARAMETER;
@@ -206,7 +212,7 @@ PrelinkedInjectPrepare (
   SegmentEndOffset = Context->PrelinkedInfoSegment->FileOffset + Context->PrelinkedInfoSegment->FileSize;
 
   if (PRELINKED_ALIGN (SegmentEndOffset) == Context->PrelinkedSize) {
-    Context->PrelinkedSize = PRELINKED_ALIGN (Context->PrelinkedInfoSegment->FileOffset);
+    Context->PrelinkedSize = (UINT32)PRELINKED_ALIGN (Context->PrelinkedInfoSegment->FileOffset);
   }
 
   Context->PrelinkedInfoSegment->VirtualAddress = 0;
diff --git a/Library/OcDataHubLib/OcDataHubLib.c b/Library/OcDataHubLib/OcDataHubLib.c
index 6245565e..42c376f4 100755
--- a/Library/OcDataHubLib/OcDataHubLib.c
+++ b/Library/OcDataHubLib/OcDataHubLib.c
@@ -28,6 +28,7 @@
 #include <Library/MemoryAllocationLib.h>
 #include <Library/OcMiscLib.h>
 #include <Library/OcDataHubLib.h>
+#include <Library/OcGuardLib.h>
 #include <Library/OcStringLib.h>
 #include <Library/UefiBootServicesTableLib.h>
 #include <Library/UefiLib.h>
@@ -181,10 +182,20 @@ SetDataHubEntry (
   PLATFORM_DATA_HEADER  *Entry;
   UINT32                KeySize;
   UINT32                TotalSize;
+  BOOLEAN               Result;
 
-  KeySize   = (UINT32) StrSize (Key);
-  TotalSize = sizeof (*Entry) + KeySize + DataSize;
-  Entry     = AllocateZeroPool (TotalSize);
+  KeySize = (UINT32) StrSize (Key);
+  Result  = OcOverflowTriAddU32 (
+              sizeof (*Entry),
+              KeySize,
+              DataSize,
+              &TotalSize
+              );
+  if (Result) {
+    return EFI_INVALID_PARAMETER;
+  }
+
+  Entry = AllocateZeroPool (TotalSize);
 
   if (Entry == NULL) {
     return EFI_OUT_OF_RESOURCES;