OcAppleKernelLib: More progress with 11.0, many thx to @Download-Fritz

Include/Acidanthera/Library/OcAppleKernelLib.h

@@ -51,6 +51,11 @@
 #define KC_LINKEDIT_SEGMENT                       "__LINKEDIT"
 #define KC_MOSCOW_SEGMENT                         "__MOSCOW101"
 
+//
+// As PageCount is UINT16, we can only index 2^16 * 4096 Bytes with one chain.
+//
+#define PRELINKED_KEXTS_MAX_SIZE (BIT16 * MACHO_PAGE_SIZE)
+
 //
 // Failsafe default for plist reserve allocation.
 //
@@ -159,6 +164,14 @@ typedef struct {
   // Whether this kernel is a kernel collection (used by macOS 11.0+).
   //
   BOOLEAN                  IsKernelCollection;
+
+  UINT32                   KextsFileOffset;
+
+  UINT64                   KextsVmAddress;
+
+  MACH_DYLD_CHAINED_STARTS_IN_SEGMENT  *KextsFixupChains;
+
+  UINT64 VirtualBase;
 } PRELINKED_CONTEXT;
 
 //
@@ -304,7 +317,8 @@ PrelinkedDependencyInsert (
 EFI_STATUS
 PrelinkedInjectPrepare (
   IN OUT PRELINKED_CONTEXT  *Context,
-  IN     UINT32             LinkedExpansion
+  IN     UINT32             LinkedExpansion,
+  IN     UINT32             ReservedExeSize
   );
 
 /**
@@ -322,16 +336,18 @@ PrelinkedInjectComplete (
 /**
   Updated required reserve size to inject this kext.
 
-  @param[in,out] ReservedSize    Current reserved size, updated.
-  @param[in]     InfoPlistSize   Kext Info.plist size.
-  @param[in]     Executable      Kext executable, optional.
-  @param[in]     ExecutableSize  Kext executable size, optional.
+  @param[in,out] ReservedInfoSize  Current reserved PLIST size, updated.
+  @param[in,out] ReservedExeSize   Current reserved KEXT size, updated.
+  @param[in]     InfoPlistSize     Kext Info.plist size.
+  @param[in]     Executable        Kext executable, optional.
+  @param[in]     ExecutableSize    Kext executable size, optional.
 
   @return  EFI_SUCCESS on success.
 **/
 EFI_STATUS
 PrelinkedReserveKextSize (
-  IN OUT UINT32       *ReservedSize,
+  IN OUT UINT32       *ReservedInfoSize,
+  IN OUT UINT32       *ReservedExeSize,
   IN     UINT32       InfoPlistSize,
   IN     UINT8        *Executable OPTIONAL,
   IN     UINT32       ExecutableSize OPTIONAL
@@ -361,6 +377,84 @@ PrelinkedInjectKext (
   IN     UINT32             ExecutableSize OPTIONAL
   );
 
+EFI_STATUS
+KcRebuildMachHeader (
+  IN OUT PRELINKED_CONTEXT  *Context
+  );
+
+/*
+  Returns the size required to store a segment's fixup chains information.
+
+  @param[in] SegmentSize  The size, in bytes, of the segment to index.
+
+  @retval 0      The segment is too large to index with a single structure.
+  @retval other  The size, in bytes, required to store a segment's fixup chain
+                 information.
+*/
+UINT32
+KcGetSegmentFixupChainsSize (
+  IN UINT32  SegmentSize
+  );
+
+/*
+  Initialises a structure that stores a segments's fixup chains information.
+
+  @param[out] SegChain      The information structure to initialise.
+  @param[in]  SegChainSize  The size, in bytes, available to SegChain.
+  @param[in]  VmAddress     The virtual address of the segment to index.
+  @param[in]  VmSize        The virtual size of the segment to index.
+*/
+EFI_STATUS
+KcInitKextFixupChains (
+  IN OUT PRELINKED_CONTEXT  *Context,
+  IN     UINT32             SegChainSize,
+  IN     UINT32             ReservedSize
+  );
+
+/*
+  Indexes all relocations of MachContext into the kernel described by Context.
+
+  @param[in,out] Context      Prelinked context.
+  @param[in]     MachContext  The context of the Mach-O to index. It must have
+                              been prelinked by OcAppleKernelLib. The image
+                              must reside in Segment.
+*/
+VOID
+KcKextIndexFixups (
+  IN OUT PRELINKED_CONTEXT  *Context,
+  IN     OC_MACHO_CONTEXT   *MachContext
+  );
+
+/*
+  Retrieves a KC KEXT's virtual size.
+
+  @param[in] Context        Prelinked context.
+  @param[in] SourceAddress  The virtual address within the KC image of the KEXT.
+
+  @retval 0      An error has occured.
+  @retval other  The virtual size, in bytes, of the KEXT at SourceAddress.
+*/
+UINT32
+KcGetKextSize (
+  IN PRELINKED_CONTEXT  *Context,
+  IN UINT64             SourceAddress
+  );
+
+/*
+  Apply the delta from KC header to the file's offsets.
+
+  @param[in,out] Context  The context of the KEXT to rebase.
+  @param[in]     Delta    The offset from KC header the KEXT starts at.
+
+  @retval EFI_SUCCESS  The file has beem rebased successfully.
+  @retval other        An error has occured.
+*/
+EFI_STATUS
+KcKextApplyFileDelta (
+  IN OUT OC_MACHO_CONTEXT  *Context,
+  IN     UINT32            Delta
+  );
+
 /**
   Initialize patcher from prelinked context for kext patching.
 

Include/Apple/IndustryStandard/AppleMachoImage.h

@@ -1807,6 +1807,27 @@ typedef struct {
          IsAuth     :  1;  ///< 0 -> not authenticated.  1 -> authenticated
 } MACH_DYKD_CHAINED_PTR_64_KERNEL_CACHE_REBASE;
 
+// header of the LC_DYLD_CHAINED_FIXUPS payload
+typedef struct {
+  UINT32 FixupsVersion;  ///< 0
+  UINT32 StartsOffset;   ///< offset of dyld_chained_starts_in_image in chain_data
+  UINT32 ImportsOffset;  ///< offset of imports table in chain_data
+  UINT32 SymbolsOffset;  ///< offset of symbol strings in chain_data
+  UINT32 ImportsCount;   ///< number of imported symbol names
+  UINT32 ImportsFormat;  ///< DYLD_CHAINED_IMPORT*
+  UINT32 SymbolsFormat;  ///< 0 => uncompressed, 1 => zlib compressed
+} MACHO_DYLD_CHAINED_FIXUPS_HEADER;
+
+///
+/// This struct is embedded in LC_DYLD_CHAINED_FIXUPS payload
+///
+typedef struct {
+  UINT32 NumSegments;
+  UINT32 SegInfoOffset[];  ///< each entry is offset into this struct for that
+                           ///< segment followed by pool of
+                           ///< dyld_chain_starts_in_segment data
+} MACH_DYLD_CHAINED_STARTS_IN_IMAGE;
+
 typedef struct {
   UINT32 Size;             ///< size of this (amount kernel needs to copy)
   UINT16 PageSize;         ///< 0x1000 or 0x4000
@@ -1995,6 +2016,7 @@ typedef UINT32 MACH_HEADER_FILE_TYPE;
 #define MACH_HEADER_FLAG_APP_EXTENSION_SAFE               BIT25
 #define MACH_HEADER_FLAG_NLIST_OUTOFSYNC_WITH_DYLDINFO    BIT26
 #define MACH_HEADER_FLAG_SIM_SUPPORT                      BIT27
+#define MACH_HEADER_FLAG_DYLIB_IN_CACHE                   BIT31
 
 typedef UINT32 MACH_HEADER_FLAGS;
 

Include/Intel/Protocol/LegacyBios.h

@@ -1164,7 +1164,7 @@ BOOLEAN
   the Stack argument
 
   @param[in] This        The protocol instance pointer.
-  @param[in] Segment     The segemnt of 16-bit mode call.
+  @param[in] Segment     The segment of 16-bit mode call.
   @param[in] Offset      The offset of 16-bit mdoe call.
   @param[in] Reg         Register contexted passed into (and returned) from thunk to
                          16-bit mode.

Legacy/BootPlatform/PciBusNoEnumerationDxe/PciIo.c

@@ -1435,14 +1435,14 @@ Returns:
       if (IS_PCI_GFX(&PciIoDevice->Pci)) {
 
         //
-        //Get the boot VGA on the same segement
+        //Get the boot VGA on the same segment
         //
         Temp = ActiveVGADeviceOnTheSameSegment (PciIoDevice);
 
         if (!Temp) {
 
           //
-          // If there is no VGA device on the segement, set
+          // If there is no VGA device on the segment, set
           // this graphics card to decode the palette range
           //
           DisableCommand  |= EFI_PCI_COMMAND_VGA_PALETTE_SNOOP;

Library/OcAppleKernelLib/PrelinkedContext.c

@@ -25,6 +25,7 @@
 #include <Library/OcStringLib.h>
 
 #include "PrelinkedInternal.h"
+#include "ProcessorBind.h"
 
 STATIC
 UINT64
@@ -196,6 +197,7 @@ PrelinkedContextInit (
   UINT64                   SegmentEndOffset;
   UINT32                   PrelinkedInfoRootIndex;
   UINT32                   PrelinkedInfoRootCount;
+  PRELINKED_KEXT           *PrelinkedKext;
 
   ASSERT (Context != NULL);
   ASSERT (Prelinked != NULL);
@@ -263,10 +265,13 @@ PrelinkedContextInit (
   //
   InitializeListHead (&Context->PrelinkedKexts);
   InitializeListHead (&Context->InjectedKexts);
-  if (InternalCachedPrelinkedKernel (Context) == NULL) {
+  PrelinkedKext = InternalCachedPrelinkedKernel (Context);
+  if (PrelinkedKext == NULL) {
     return EFI_INVALID_PARAMETER;
   }
 
+  Context->VirtualBase = PrelinkedKext->Context.VirtualBase;
+
   Context->PrelinkedLastAddress = MACHO_ALIGN (MachoGetLastAddress64 (&Context->PrelinkedMachContext));
   if (Context->PrelinkedLastAddress == 0) {
     return EFI_INVALID_PARAMETER;
@@ -478,12 +483,34 @@ PrelinkedDependencyInsert (
 EFI_STATUS
 PrelinkedInjectPrepare (
   IN OUT PRELINKED_CONTEXT  *Context,
-  IN     UINT32             LinkedExpansion
+  IN     UINT32             LinkedExpansion,
+  IN     UINT32             ReservedExeSize
   )
 {
-  UINT64  SegmentEndOffset;
+  EFI_STATUS Status;
+  UINT64     SegmentEndOffset;
+  UINT32     AlignedExpansion;
+
+  Context->PrelinkedInfoSegment->VirtualAddress = 0;
+  Context->PrelinkedInfoSegment->Size           = 0;
+  Context->PrelinkedInfoSegment->FileOffset     = 0;
+  Context->PrelinkedInfoSegment->FileSize       = 0;
+  Context->PrelinkedInfoSection->Address        = 0;
+  Context->PrelinkedInfoSection->Size           = 0;
+  Context->PrelinkedInfoSection->Offset         = 0;
 
   if (Context->IsKernelCollection) {
+    //
+    // For newer variant of the prelinkedkernel plist we need to kill it
+    // in both inner and outer images.
+    //
+    Context->InnerInfoSegment->VirtualAddress = 0;
+    Context->InnerInfoSegment->Size           = 0;
+    Context->InnerInfoSegment->FileOffset     = 0;
+    Context->InnerInfoSegment->FileSize       = 0;
+    Context->InnerInfoSection->Address        = 0;
+    Context->InnerInfoSection->Size           = 0;
+    Context->InnerInfoSection->Offset         = 0;
     //
     // For newer variant (KC mode) __LINKEDIT is last, and we need to expand it to enable
     // dyld fixup generation.
@@ -493,16 +520,18 @@ PrelinkedInjectPrepare (
       return EFI_OUT_OF_RESOURCES;
     }
 
-    LinkedExpansion = MACHO_ALIGN (LinkedExpansion);
-
     ASSERT (Context->PrelinkedLastAddress == Context->PrelinkedLastLoadAddress);
 
-    Context->PrelinkedSize              += LinkedExpansion;
-    Context->PrelinkedLastAddress       += LinkedExpansion;
-    Context->PrelinkedLastLoadAddress   += LinkedExpansion;
-    Context->LinkeditSegment->Size      += LinkedExpansion;
-    Context->LinkeditSegment->FileSize  += LinkedExpansion;
-    Context->LinkeditSegment->Size      += LinkedExpansion;
+    Context->KextsFixupChains = (VOID *) (Context->Prelinked +
+      Context->LinkeditSegment->FileOffset + Context->LinkeditSegment->FileSize);
+
+    AlignedExpansion = MACHO_ALIGN (LinkedExpansion);
+
+    Context->PrelinkedSize              += AlignedExpansion;
+    Context->PrelinkedLastAddress       += AlignedExpansion;
+    Context->PrelinkedLastLoadAddress   += AlignedExpansion;
+    Context->LinkeditSegment->Size      += AlignedExpansion;
+    Context->LinkeditSegment->FileSize  += AlignedExpansion;
   } else {
     //
     // For older variant of the prelinkedkernel plist info is normally
@@ -515,48 +544,36 @@ PrelinkedInjectPrepare (
     if (MACHO_ALIGN (SegmentEndOffset) == Context->PrelinkedSize) {
       Context->PrelinkedSize = (UINT32) MACHO_ALIGN (Context->PrelinkedInfoSegment->FileOffset);
     }
-  }
 
-  Context->PrelinkedInfoSegment->VirtualAddress = 0;
-  Context->PrelinkedInfoSegment->Size           = 0;
-  Context->PrelinkedInfoSegment->FileOffset     = 0;
-  Context->PrelinkedInfoSegment->FileSize       = 0;
-  Context->PrelinkedInfoSection->Address        = 0;
-  Context->PrelinkedInfoSection->Size           = 0;
-  Context->PrelinkedInfoSection->Offset         = 0;
+    Context->PrelinkedLastAddress = MACHO_ALIGN (MachoGetLastAddress64 (&Context->PrelinkedMachContext));
+    if (Context->PrelinkedLastAddress == 0) {
+      return EFI_INVALID_PARAMETER;
+    }
 
-  if (Context->IsKernelCollection) {
     //
-    // For newer variant of the prelinkedkernel plist we need to kill it
-    // in both inner and outer images.
+    // Prior to plist there usually is prelinked text. 
     //
-    Context->InnerInfoSegment->VirtualAddress = 0;
-    Context->InnerInfoSegment->Size           = 0;
-    Context->InnerInfoSegment->FileOffset     = 0;
-    Context->InnerInfoSegment->FileSize       = 0;
-    Context->InnerInfoSection->Address        = 0;
-    Context->InnerInfoSection->Size           = 0;
-    Context->InnerInfoSection->Offset         = 0;
+    SegmentEndOffset = Context->PrelinkedTextSegment->FileOffset + Context->PrelinkedTextSegment->FileSize;
 
-    return EFI_SUCCESS;
-  }
-
-  Context->PrelinkedLastAddress = MACHO_ALIGN (MachoGetLastAddress64 (&Context->PrelinkedMachContext));
-  if (Context->PrelinkedLastAddress == 0) {
-    return EFI_INVALID_PARAMETER;
+    if (MACHO_ALIGN (SegmentEndOffset) != Context->PrelinkedSize) {
+      //
+      // TODO: Implement prelinked text relocation when it is not preceding prelinked info
+      // and is not in the end of prelinked info.
+      //
+      return EFI_UNSUPPORTED;
+    }
   }
-
   //
-  // Prior to plist there usually is prelinked text. 
+  // Append the injected KEXTs to the current kernel end.
   //
-  SegmentEndOffset = Context->PrelinkedTextSegment->FileOffset + Context->PrelinkedTextSegment->FileSize;
+  Context->KextsFileOffset = Context->PrelinkedSize;
+  Context->KextsVmAddress  = Context->PrelinkedLastAddress;
 
-  if (MACHO_ALIGN (SegmentEndOffset) != Context->PrelinkedSize) {
-    //
-    // TODO: Implement prelinked text relocation when it is not preceding prelinked info
-    // and is not in the end of prelinked info.
-    //
-    return EFI_UNSUPPORTED;
+  if (Context->IsKernelCollection) {
+    Status = KcInitKextFixupChains (Context, LinkedExpansion, ReservedExeSize);
+    if (EFI_ERROR (Status)) {
+      return Status;
+    }
   }
 
   return EFI_SUCCESS;
@@ -567,9 +584,32 @@ PrelinkedInjectComplete (
   IN OUT PRELINKED_CONTEXT  *Context
   )
 {
+  EFI_STATUS  Status;
   CHAR8       *ExportedInfo;
   UINT32      ExportedInfoSize;
   UINT32      NewSize;
+  UINT32      KextsSize;
+  UINT32      ChainSize;
+
+  if (Context->IsKernelCollection) {
+    //
+    // Fix up the segment fixup chains structure to reflect the actually
+    // injected segment size.
+    //
+    KextsSize = Context->PrelinkedSize - Context->KextsFileOffset;
+
+    ChainSize = KcGetSegmentFixupChainsSize (KextsSize);
+    ASSERT (ChainSize != 0);
+    ASSERT (ChainSize <= Context->KextsFixupChains->Size);
+
+    Context->KextsFixupChains->Size      = ChainSize;
+    Context->KextsFixupChains->PageCount = KextsSize / MACHO_PAGE_SIZE;
+
+    Status = KcRebuildMachHeader (Context);
+    if (EFI_ERROR (Status)) {
+      return Status;
+    }
+  }
 
   ExportedInfo = XmlDocumentExport (Context->PrelinkedInfoDocument, &ExportedInfoSize, 0);
   if (ExportedInfo == NULL) {
@@ -630,7 +670,8 @@ PrelinkedInjectComplete (
 
 EFI_STATUS
 PrelinkedReserveKextSize (
-  IN OUT UINT32       *ReservedSize,
+  IN OUT UINT32       *ReservedInfoSize,
+  IN OUT UINT32       *ReservedExeSize,
   IN     UINT32       InfoPlistSize,
   IN     UINT8        *Executable,
   IN     UINT32       ExecutableSize OPTIONAL
@@ -659,11 +700,14 @@ PrelinkedReserveKextSize (
     }
   }
 
-  if (OcOverflowTriAddU32 (*ReservedSize, InfoPlistSize, ExecutableSize, &ExecutableSize)) {
+  if (OcOverflowAddU32 (*ReservedInfoSize, InfoPlistSize, &InfoPlistSize)
+   || OcOverflowAddU32 (*ReservedExeSize, ExecutableSize, &ExecutableSize)) {
     return EFI_INVALID_PARAMETER;
   }
 
-  *ReservedSize = ExecutableSize;
+  *ReservedInfoSize = InfoPlistSize;
+  *ReservedExeSize  = ExecutableSize;
+
   return EFI_SUCCESS;
 }
 
@@ -699,6 +743,7 @@ PrelinkedInjectKext (
   CHAR8             ExecutableSizeStr[24];
   CHAR8             ExecutableLoadAddrStr[24];
   CHAR8             KmodInfoStr[24];
+  UINT32            KextOffset;
 
   PrelinkedKext = NULL;
 
@@ -716,28 +761,32 @@ PrelinkedInjectKext (
       DEBUG ((DEBUG_INFO, "OCAK: Injected kext %a/%a is not a supported executable\n", BundlePath, ExecutablePath));
       return EFI_INVALID_PARAMETER;
     }
+    //
+    // Append the KEXT to the current prelinked end.
+    //
+    KextOffset = Context->PrelinkedSize;
 
     ExecutableSize = MachoExpandImage64 (
       &ExecutableContext,
-      &Context->Prelinked[Context->PrelinkedSize],
-      Context->PrelinkedAllocSize - Context->PrelinkedSize,
+      &Context->Prelinked[KextOffset],
+      Context->PrelinkedAllocSize - KextOffset,
       TRUE
       );
 
     AlignedExecutableSize = MACHO_ALIGN (ExecutableSize);
 
-    if (OcOverflowAddU32 (Context->PrelinkedSize, AlignedExecutableSize, &NewPrelinkedSize)
+    if (OcOverflowAddU32 (KextOffset, AlignedExecutableSize, &NewPrelinkedSize)
       || NewPrelinkedSize > Context->PrelinkedAllocSize
       || ExecutableSize == 0) {
       return EFI_BUFFER_TOO_SMALL;
     }
 
     ZeroMem (
-      &Context->Prelinked[Context->PrelinkedSize + ExecutableSize],
+      &Context->Prelinked[KextOffset + ExecutableSize],
       AlignedExecutableSize - ExecutableSize
       );
 
-    if (!MachoInitializeContext (&ExecutableContext, &Context->Prelinked[Context->PrelinkedSize], ExecutableSize)) {
+    if (!MachoInitializeContext (&ExecutableContext, &Context->Prelinked[KextOffset], ExecutableSize)) {
       return EFI_INVALID_PARAMETER;
     }
 
@@ -839,9 +888,32 @@ PrelinkedInjectKext (
     Context->PrelinkedSize                  += AlignedExecutableSize;
     Context->PrelinkedLastAddress           += AlignedExecutableSize;
     Context->PrelinkedLastLoadAddress       += AlignedExecutableSize;
-    Context->PrelinkedTextSegment->Size     += AlignedExecutableSize;
-    Context->PrelinkedTextSegment->FileSize += AlignedExecutableSize;
-    Context->PrelinkedTextSection->Size     += AlignedExecutableSize;
+
+    if (Context->IsKernelCollection) {
+      //
+      // For KC, our KEXTs have their own segment - do not mod __PRELINK_INFO.
+      // Integrate the KEXT into KC by indexing its fixups and rebasing.
+      //
+      KcKextIndexFixups (Context, &ExecutableContext);
+      Status = KcKextApplyFileDelta (&ExecutableContext, KextOffset);
+      if (EFI_ERROR (Status)) {
+        DEBUG ((
+          DEBUG_WARN,
+          "Failed to rebase injected kext %a/%a by %u\n",
+          BundlePath,
+          ExecutablePath,
+          KextOffset
+          ));
+        return Status;
+      }
+    } else {
+      //
+      // For legacy prelinkedkernel we append to __PRELINK_TEXT.
+      //
+      Context->PrelinkedTextSegment->Size     += AlignedExecutableSize;
+      Context->PrelinkedTextSegment->FileSize += AlignedExecutableSize;
+      Context->PrelinkedTextSection->Size     += AlignedExecutableSize;
+    }
   }
 
   //

Library/OcAppleKernelLib/PrelinkedKext.c

@@ -129,8 +129,18 @@ InternalCreatePrelinkedKext (
   // BundleLibraries, CompatibleVersion, and KmodInfo are optional and thus not checked.
   //
   if (!Found || KextIdentifier == NULL || SourceBase < VirtualBase
-    || (Prelinked != NULL && (VirtualBase == 0 || SourceBase == 0 || SourceSize == 0 || SourceSize > MAX_UINT32))) {
-    return NULL;
+    || (Prelinked != NULL && (VirtualBase == 0 || SourceBase == 0))) {
+    //
+    // KC expectedly lacks PRELINK_INFO_EXECUTABLE_SIZE_KEY - parse the KEXT at
+    // SourceBase and retrieve its virtual size from the Mach-O header.
+    //
+    if (SourceSize == 0 && Prelinked != NULL && Prelinked->IsKernelCollection) {
+      SourceSize = KcGetKextSize (Prelinked, SourceBase);
+    }
+    
+    if (SourceSize == 0 || SourceSize > MAX_UINT32) {
+      return NULL;
+    }
   }
 
   if (Prelinked != NULL) {

Library/OcMachoLib/Header.c

@@ -651,6 +651,10 @@ MachoGetNextSegment64 (
     return NULL;
   }
 
+  if (NextSegment->VirtualAddress + NextSegment->Size < NextSegment->VirtualAddress) {
+    return NULL;
+  }
+
   return NextSegment;
 }
 

Platform/OpenCore/OpenCoreKernel.c

@@ -12,6 +12,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 
 **/
 
+#include "ProcessorBind.h"
 #include <OpenCore.h>
 
 #include <Library/BaseLib.h>
@@ -168,15 +169,16 @@ OcKernelReadDarwinVersion (
 }
 
 STATIC
-UINT32
+EFI_STATUS
 OcKernelLoadKextsAndReserve (
   IN OC_STORAGE_CONTEXT  *Storage,
-  IN OC_GLOBAL_CONFIG    *Config
+  IN OC_GLOBAL_CONFIG    *Config,
+  OUT UINT32             *ReservedExeSize,
+  OUT UINT32             *ReservedInfoSize
   )
 {
   EFI_STATUS           Status;
   UINT32               Index;
-  UINT32               ReserveSize;
   CHAR8                *BundlePath;
   CHAR8                *Comment;
   CHAR8                *PlistPath;
@@ -184,7 +186,8 @@ OcKernelLoadKextsAndReserve (
   CHAR16               FullPath[OC_STORAGE_SAFE_PATH_MAX];
   OC_KERNEL_ADD_ENTRY  *Kext;
 
-  ReserveSize = PRELINK_INFO_RESERVE_SIZE;
+  *ReservedInfoSize = PRELINK_INFO_RESERVE_SIZE;
+  *ReservedExeSize  = 0;
 
   for (Index = 0; Index < Config->Kernel.Add.Count; ++Index) {
     Kext = Config->Kernel.Add.Values[Index];
@@ -289,17 +292,38 @@ OcKernelLoadKextsAndReserve (
       }
     }
 
-    PrelinkedReserveKextSize (
-      &ReserveSize,
+    Status = PrelinkedReserveKextSize (
+      ReservedInfoSize,
+      ReservedExeSize,
       Kext->PlistDataSize,
       Kext->ImageData,
       Kext->ImageDataSize
       );
+    if (EFI_ERROR (Status)) {
+      DEBUG ((
+        DEBUG_ERROR,
+        "OC: Failed to fit kext %a (%a)\n",
+        BundlePath,
+        Comment
+        ));
+      Kext->Enabled = FALSE;
+      FreePool (Kext->PlistData);
+      Kext->PlistData = NULL;
+      continue;
+    }
   }
 
-  DEBUG ((DEBUG_INFO, "OC: Kext reservation size %u\n", ReserveSize));
+  if (*ReservedExeSize > PRELINKED_KEXTS_MAX_SIZE
+   || *ReservedInfoSize + *ReservedExeSize < *ReservedExeSize) {
+    return EFI_UNSUPPORTED;
+  }
 
-  return ReserveSize;
+  DEBUG ((
+    DEBUG_INFO,
+    "OC: Kext reservation size %u\n",
+    *ReservedInfoSize + *ReservedExeSize
+    ));
+  return EFI_SUCCESS;
 }
 
 STATIC
@@ -584,7 +608,9 @@ OcKernelProcessPrelinked (
   IN     UINT32            DarwinVersion,
   IN OUT UINT8             *Kernel,
   IN     UINT32            *KernelSize,
-  IN     UINT32            AllocatedSize
+  IN     UINT32            AllocatedSize,
+  IN     UINT32            LinkedExpansion,
+  IN     UINT32            ReservedExeSize
   )
 {
   EFI_STATUS           Status;
@@ -605,9 +631,12 @@ OcKernelProcessPrelinked (
 
     OcKernelBlockKexts (Config, DarwinVersion, &Context);
 
-    Status = PrelinkedInjectPrepare (&Context, 0);
+    Status = PrelinkedInjectPrepare (
+      &Context,
+      LinkedExpansion,
+      ReservedExeSize
+      );
     if (!EFI_ERROR (Status)) {
-
       for (Index = 0; Index < Config->Kernel.Add.Count; ++Index) {
         Kext = Config->Kernel.Add.Values[Index];
 
@@ -665,6 +694,8 @@ OcKernelProcessPrelinked (
           ));
       }
 
+      ASSERT (Context.PrelinkedSize - Context.KextsFileOffset <= ReservedExeSize);
+
       Status = PrelinkedInjectComplete (&Context);
       if (EFI_ERROR (Status)) {
         DEBUG ((DEBUG_WARN, "OC: Prelink insertion error - %r\n", Status));
@@ -693,6 +724,7 @@ OcKernelFileOpen (
   )
 {
   EFI_STATUS         Status;
+  BOOLEAN            Result;
   UINT8              *Kernel;
   UINT32             KernelSize;
   UINT32             AllocatedSize;
@@ -701,6 +733,10 @@ OcKernelFileOpen (
   EFI_STATUS         PrelinkedStatus;
   EFI_TIME           ModificationTime;
   UINT32             DarwinVersion;
+  UINT32             ReservedInfoSize;
+  UINT32             ReservedExeSize;
+  UINT32             LinkedExpansion;
+  UINT32             ReservedFullSize;
 
   Status = SafeFileOpen (This, NewHandle, FileName, OpenMode, Attributes);
 
@@ -725,13 +761,35 @@ OcKernelFileOpen (
     && OcStriStr (FileName, L"kernel") != NULL
     && StrCmp (FileName, L"System\\Library\\Kernels\\kernel") != 0) {
 
+    OcKernelLoadKextsAndReserve (
+      mOcStorage,
+      mOcConfiguration,
+      &ReservedInfoSize,
+      &ReservedExeSize
+      );
+
+    LinkedExpansion = KcGetSegmentFixupChainsSize (ReservedExeSize);
+    if (LinkedExpansion == 0) {
+      return EFI_UNSUPPORTED;
+    }
+
+    Result = OcOverflowTriAddU32 (
+      ReservedInfoSize,
+      ReservedExeSize,
+      LinkedExpansion,
+      &ReservedFullSize
+      );
+    if (Result) {
+      return EFI_UNSUPPORTED;
+    }
+
     DEBUG ((DEBUG_INFO, "OC: Trying XNU hook on %s\n", FileName));
     Status = ReadAppleKernel (
       *NewHandle,
       &Kernel,
       &KernelSize,
       &AllocatedSize,
-      OcKernelLoadKextsAndReserve (mOcStorage, mOcConfiguration)
+      ReservedFullSize
       );
     DEBUG ((DEBUG_INFO, "OC: Result of XNU hook on %s is %r\n", FileName, Status));
 
@@ -744,7 +802,9 @@ OcKernelFileOpen (
         DarwinVersion,
         Kernel,
         &KernelSize,
-        AllocatedSize
+        AllocatedSize,
+        LinkedExpansion,
+        ReservedExeSize
         );
 
       DEBUG ((DEBUG_INFO, "OC: Prelinked status - %r\n", PrelinkedStatus));

Utilities/TestKernelCollection/KernelCollection.c → Utilities/TestKernelCollection2/KernelCollection2.c

@@ -19,6 +19,7 @@
 #include <Library/OcAppleKernelLib.h>
 #include <Library/DebugLib.h>
 
+#include <assert.h>
 #include <string.h>
 #include <sys/time.h>
 #include <stdint.h>
@@ -128,8 +129,44 @@ int main(int argc, char** argv) {
     return -1;
   }
 
+  uint32_t ReservedInfoSize = PrelinkedSize + 5*1024*1024;
+  uint32_t ReservedExeSize  = 0;
+  EFI_STATUS Status = PrelinkedReserveKextSize (
+    &ReservedInfoSize,
+    &ReservedExeSize,
+    LiluKextInfoPlistDataSize,
+    LiluKextData,
+    LiluKextDataSize
+    );
+  /*Status |= PrelinkedReserveKextSize (
+    &ReservedInfoSize,
+    &ReservedExeSize,
+    VsmcKextInfoPlistDataSize,
+    VsmcKextData,
+    VsmcKextDataSize
+    );*/
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "PrelinkedReserveKextSize failed - %r.\n", Status));
+    return -1;
+  }
+
+  DEBUG ((DEBUG_WARN, "Kexts size %u\n", ReservedExeSize));
+
+  uint32_t LinkedExpansion = KcGetSegmentFixupChainsSize (ReservedExeSize);
+  if (LinkedExpansion == 0) {
+    DEBUG ((DEBUG_ERROR, "KcGetSegmentFixupChainsSize failed.\n"));
+    return -1;
+  }
 
-  AllocSize = MACHO_ALIGN (PrelinkedSize + 1*1024*1024);
+  uint32_t LinkedExpansionA = MACHO_ALIGN (LinkedExpansion);
+
+  if (ReservedInfoSize + ReservedExeSize < ReservedExeSize
+   || ReservedInfoSize + ReservedExeSize + LinkedExpansionA < LinkedExpansionA) {
+    DEBUG ((DEBUG_ERROR, "Overflow.\n"));
+    return -1;
+  }
+
+  AllocSize = MACHO_ALIGN (ReservedInfoSize + ReservedExeSize + LinkedExpansionA);
 
   Prelinked = realloc (Prelinked, AllocSize);
   if (Prelinked == NULL) {
@@ -142,11 +179,15 @@ int main(int argc, char** argv) {
   ApplyKernelPatches (Prelinked, PrelinkedSize);
 #endif
 
-  EFI_STATUS Status = PrelinkedContextInit (&Context, Prelinked, PrelinkedSize, AllocSize);
+  Status = PrelinkedContextInit (&Context, Prelinked, PrelinkedSize, AllocSize);
 
   if (!EFI_ERROR (Status)) {
 
-    Status = PrelinkedInjectPrepare (&Context, EFI_PAGE_SIZE);
+    Status = PrelinkedInjectPrepare (
+      &Context,
+      LinkedExpansion,
+      ReservedExeSize
+      );
     if (EFI_ERROR (Status)) {
       DEBUG ((DEBUG_WARN, "Prelink inject prepare error %r\n", Status));
       return -1;

Utilities/TestKernelCollection/Makefile → Utilities/TestKernelCollection2/Makefile

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 ##
 
-PROJECT = KernelCollection
+PROJECT = KernelCollection2
 PRODUCT = $(PROJECT)$(SUFFIX)
 OBJS    = $(PROJECT).o \
 	Lilu.o \
@@ -11,6 +11,7 @@ OBJS    = $(PROJECT).o \
 	KextPatcher.o \
 	PrelinkedKext.o \
 	PrelinkedContext.o \
+	KernelCollection.o \
 	Vtables.o \
 	Link.o \
 	KernelReader.o \