OcMachoLib: Fix KC init and segment merging

Library/OcMachoLib/Header.c

@@ -265,7 +265,8 @@ MachoInitializeContext (
   //
   if ((MachHeader->CpuType != MachCpuTypeX8664)
    || ((MachHeader->FileType != MachHeaderFileTypeKextBundle)
-    && (MachHeader->FileType != MachHeaderFileTypeExecute))) {
+    && (MachHeader->FileType != MachHeaderFileTypeExecute)
+    && (MachHeader->FileType != MachHeaderFileTypeFileSet))) {
     return FALSE;
   }
 
@@ -1496,10 +1497,6 @@ MachoMergeSegments64 (
   MACH_LOAD_COMMAND       *LoadCommand;
   MACH_SEGMENT_COMMAND_64 *Segment;
   MACH_SEGMENT_COMMAND_64 *FirstSegment;
-  UINT64                  MaxAddress;
-  UINT64                  MaxOffset;
-  MACH_VM_PROTECTION      MaxInitProt;
-  MACH_VM_PROTECTION      MaxMaxProt;
   MACH_HEADER_64          *Header;
   UINTN                   PrefixLength;
   UINTN                   SkipCount;
@@ -1513,10 +1510,6 @@ MachoMergeSegments64 (
   PrefixLength = AsciiStrLen (Prefix);
   FirstSegment = NULL;
 
-  MaxAddress  = 0;
-  MaxOffset   = 0;
-  MaxInitProt = 0;
-  MaxMaxProt  = 0;
   SkipCount   = 0;
 
   LoadCommand = &Header->Commands[0];
@@ -1525,6 +1518,8 @@ MachoMergeSegments64 (
     //
     // Either skip or stop at unrelated commands.
     //
+    Segment = (MACH_SEGMENT_COMMAND_64 *) (VOID *) LoadCommand;
+
     if (LoadCommand->CommandType != MACH_LOAD_COMMAND_SEGMENT_64
       || AsciiStrnCmp (Segment->SegmentName, Prefix, PrefixLength) != 0) {
       if (FirstSegment != NULL) {
@@ -1538,7 +1533,6 @@ MachoMergeSegments64 (
     //
     // We have a segment starting with the prefix.
     //
-    Segment = (MACH_SEGMENT_COMMAND_64 *) (VOID *) LoadCommand;
 
     //
     // Do not support this for now as it will require changes in the file.

User/Include/File.h

@@ -12,4 +12,6 @@
 
 uint8_t *readFile(const char *str, uint32_t *size);
 
+void writeFile(const char *str, void *data, uint32_t size);
+
 #endif // OC_USER_FILE_H

User/Library/File.c

@@ -24,3 +24,15 @@ uint8_t *readFile(const char *str, uint32_t *size) {
 
   return string;
 }
+
+void writeFile(const char *str, void *data, uint32_t size) {
+  FILE *Fh = fopen("out.bin", "wb");
+
+  if (Fh != NULL) {
+    if (fwrite (data, size, 1, Fh) != 1)
+      abort();
+    fclose(Fh);
+  } else {
+    abort();
+  }
+}

Utilities/TestKernelCollection/KernelCollection.c

+/** @file
+  Copyright (C) 2018, vit9696. All rights reserved.
+
+  All rights reserved.
+
+  This program and the accompanying materials
+  are licensed and made available under the terms and conditions of the BSD License
+  which accompanies this distribution.  The full text of the license may be found at
+  http://opensource.org/licenses/bsd-license.php
+
+  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
+  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
+**/
+
+#include <Library/BaseMemoryLib.h>
+#include <Library/MemoryAllocationLib.h>
+#include <Library/OcMachoLib.h>
+#include <Library/OcMiscLib.h>
+#include <Library/DebugLib.h>
+
+#include <string.h>
+#include <sys/time.h>
+#include <stdint.h>
+#include <stdio.h>
+
+#include <File.h>
+
+/*
+ for fuzzing (TODO):
+ clang-mp-7.0 -Dmain=__main -g -fsanitize=undefined,address,fuzzer -I../Include -I../../Include -I../../../MdePkg/Include/ -include ../Include/Base.h -I../../../EfiPkg/Include/ Macho.c  ../../Library/OcMiscLib/Base64Decode.c ../../Library/OcStringLib/OcAsciiLib.c  ../../Library/OcMachoLib/CxxSymbols.c ../../Library/OcMachoLib/Header.c ../../Library/OcMachoLib/Relocations.c ../../Library/OcMachoLib/Symbols.c -o Macho
+ rm -rf DICT fuzz*.log ; mkdir DICT ; cp /System/Library/Kernels/kernel DICT ; ./Macho -rss_limit_mb=4096M -jobs=4 DICT
+
+ rm -rf fuzz*.log ; mkdir -p DICT ; cp /System/Library/Kernels/kernel DICT/kernel ; ./Macho -jobs=4 -rss_limit_mb=4096M DICT
+
+ rm -rf Macho.dSYM DICT fuzz*.log Macho
+*/
+
+static int FeedMacho(void *file, uint32_t size) {
+  OC_MACHO_CONTEXT Context;
+  if (!MachoInitializeContext (&Context, file, size)) {
+    DEBUG ((DEBUG_WARN, "MachoInitializeContext failure\n"));
+    return -1;
+  }
+
+/*
+  For MH_FILELIST Mach-O as determined by checking Mach-O header:
+
+  1. Save original Mach-O header.
+  2. Find __TEXT segment:
+  - make sure its offset is 0.
+  - make sure its memory size and file size equal.
+  3. Calculate DIFF = NEW_SIZE - file size.
+  4. Write NEW_SIZE to memory size and file size of __TEXT.
+  5. For each segment, that is not text:
+  - Increase file_offset by DIFF.
+  - Increase vm_offset by DIFF.
+  6. For each dyld fixup:
+  - Increase target by DIFF, becuase the offset will not change for anything.
+  7. Insert DIFF bytes between __TEXT and other segments.
+
+VRAM            FILE           FILE NEW                    VRAM
+                       
+__HIB          __TEXT         __TEXT         (bigger)
+__TEXT         __TEXT_EXEC    __TEXT_EXEC    (lower)
+__TEXT_EXEC    __DATA_CONST   __DATA_CONST   (lower)
+__DATA_CONST   __PRELINK_INFO <empty space>
+__PRELINK_INFO __DATA         __DATA         (same off)
+__DATA         __HIB          __HIB          (same off)
+__REGION0      __REGION0      __REGION0      (same off)
+...            ...            ...
+__LINKEDIT     __LINKEDIT     __REGIONX      (new kext)
+                              __PRELINK_INFO (bigger)
+                              __LINKEDIT     (bigger?)                  
+
+*/
+
+
+  BOOLEAN shr = MachoMergeSegments64 (&Context, "__REGION");
+  DEBUG ((DEBUG_WARN, "Shrinking __REGION is %d\n", shr));
+
+  writeFile("out.bin", file, size);
+
+  return 0;
+}
+
+int main(int argc, char** argv) {
+  uint32_t f;
+  uint8_t *b;
+  if ((b = readFile(argc > 1 ? argv[1] : "/System/Library/KernelCollections/BootKernelExtensions.kc", &f)) == NULL) {
+    printf("Read fail\n");
+    return -1;
+  }
+
+  FeedMacho (b, f);
+  free(b);
+  return 0;
+}
+
+INT32 LLVMFuzzerTestOneInput(CONST UINT8 *Data, UINTN Size) {
+  if (Size > 0) {
+    VOID *NewData = AllocatePool (Size);
+    if (NewData) {
+      CopyMem (NewData, Data, Size);
+      FeedMacho (NewData, (UINT32) Size);
+      FreePool (NewData);
+    }
+  }
+  return 0;
+}

Utilities/TestKernelCollection/Makefile

+## @file
+# Copyright (c) 2020, PMheart. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+##
+
+PROJECT = KernelCollection
+PRODUCT = $(PROJECT)$(SUFFIX)
+OBJS    = $(PROJECT).o
+include ../../User/Makefile

Utilities/TestKextInject/KextInject.c

@@ -583,17 +583,11 @@ int wrap_main(int argc, char** argv) {
       DEBUG ((DEBUG_WARN, "Prelink inject complete error %r\n", Status));
     }
 
-    FILE *Fh = fopen("out.bin", "wb");
-
-    if (Fh != NULL) {
-      fwrite (Prelinked, Context.PrelinkedSize, 1, Fh);
-      fclose(Fh);
-
-      if (!EFI_ERROR (Status)) {
-        printf("All good\n");
-      }
+    writeFile("out.bin", Prelinked, Context.PrelinkedSize);
+    if (!EFI_ERROR (Status)) {
+      printf("All good\n");
     } else {
-      printf("File error\n");
+      printf("Inject error\n");
     }
 #endif
     PrelinkedContextFree (&Context);

build_oc.tool

@@ -12,6 +12,7 @@ buildutil() {
     "TestDiskImage"
     "TestHelloWorld"
     "TestImg4"
+    "TestKernelCollection"
     "TestKextInject"
     "TestMacho"
     "TestRsaPreprocess"