macho: unify code signature between stage1 and stage2

3 files changed, 28 insertions, 98 deletions

diff --git a/src/link/MachO.zig b/src/link/MachO.zig
index 020faa88b3..921c49763d 100644
--- a/src/link/MachO.zig
+++ b/src/link/MachO.zig
@@ -773,7 +773,7 @@ fn linkWithLLD(self: *MachO, comp: *Compilation) !void {
             var signature = CodeSignature.init(self.base.allocator);
             defer signature.deinit();
             const emit = self.base.options.emit.?;
-            try signature.calcAdhocSignatureFile(
+            try signature.calcAdhocSignature(
                 out_file,
                 emit.sub_path,
                 text_cmd,
@@ -1741,15 +1741,21 @@ fn writeCodeSignaturePadding(self: *MachO) !void {
 }
 
 fn writeCodeSignature(self: *MachO) !void {
-    const code_sig_cmd = &self.load_commands.items[self.code_signature_cmd_index.?].LinkeditData;
+    const text_segment = self.load_commands.items[self.text_segment_cmd_index.?].Segment;
+    const code_sig_cmd = self.load_commands.items[self.code_signature_cmd_index.?].LinkeditData;
+
     var code_sig = CodeSignature.init(self.base.allocator);
     defer code_sig.deinit();
-
-    try code_sig.calcAdhocSignature(self);
+    try code_sig.calcAdhocSignature(
+        self.base.file.?,
+        self.base.options.emit.?.sub_path,
+        text_segment.inner,
+        code_sig_cmd,
+        self.base.options.output_mode,
+    );
 
     var buffer = try self.base.allocator.alloc(u8, code_sig.size());
     defer self.base.allocator.free(buffer);
-
     code_sig.write(buffer);
 
     log.debug("writing code signature from 0x{x} to 0x{x}\n", .{ code_sig_cmd.dataoff, code_sig_cmd.dataoff + buffer.len });
diff --git a/src/link/MachO/CodeSignature.zig b/src/link/MachO/CodeSignature.zig
index 98189ae219..c7cfdfcb44 100644
--- a/src/link/MachO/CodeSignature.zig
+++ b/src/link/MachO/CodeSignature.zig
@@ -10,8 +10,6 @@ const testing = std.testing;
 const Allocator = mem.Allocator;
 const Sha256 = std.crypto.hash.sha2.Sha256;
 
-const MachO = @import("../MachO.zig");
-
 const hash_size: u8 = 32;
 const page_size: u16 = 0x1000;
 
@@ -52,7 +50,7 @@ const CodeDirectory = struct {
     }
 };
 
-alloc: *Allocator,
+allocator: *Allocator,
 inner: macho.SuperBlob = .{
     .magic = macho.CSMAGIC_EMBEDDED_SIGNATURE,
     .length = @sizeOf(macho.SuperBlob),
@@ -60,13 +58,11 @@ inner: macho.SuperBlob = .{
 },
 cdir: ?CodeDirectory = null,
 
-pub fn init(alloc: *Allocator) CodeSignature {
-    return .{
-        .alloc = alloc,
-    };
+pub fn init(allocator: *Allocator) CodeSignature {
+    return .{ .allocator = allocator };
 }
 
-pub fn calcAdhocSignatureFile(
+pub fn calcAdhocSignature(
     self: *CodeSignature,
     file: fs.File,
     id: []const u8,
@@ -107,10 +103,10 @@ pub fn calcAdhocSignatureFile(
     const total_pages = mem.alignForward(file_size, page_size) / page_size;
 
     var hash: [hash_size]u8 = undefined;
-    var buffer = try self.alloc.alloc(u8, page_size);
-    defer self.alloc.free(buffer);
+    var buffer = try self.allocator.alloc(u8, page_size);
+    defer self.allocator.free(buffer);
 
-    try cdir.data.ensureCapacity(self.alloc, total_pages * hash_size + id.len + 1);
+    try cdir.data.ensureCapacity(self.allocator, total_pages * hash_size + id.len + 1);
 
     // 1. Save the identifier and update offsets
     cdir.inner.identOffset = cdir.inner.length;
@@ -142,80 +138,6 @@ pub fn calcAdhocSignatureFile(
     self.cdir = cdir;
 }
 
-pub fn calcAdhocSignature(self: *CodeSignature, bin_file: *const MachO) !void {
-    const text_segment = bin_file.load_commands.items[bin_file.text_segment_cmd_index.?].Segment;
-    const code_sig_cmd = bin_file.load_commands.items[bin_file.code_signature_cmd_index.?].LinkeditData;
-
-    const execSegBase: u64 = text_segment.inner.fileoff;
-    const execSegLimit: u64 = text_segment.inner.filesize;
-    const execSegFlags: u64 = if (bin_file.base.options.output_mode == .Exe) macho.CS_EXECSEG_MAIN_BINARY else 0;
-    const file_size = code_sig_cmd.dataoff;
-    var cdir = CodeDirectory{
-        .inner = .{
-            .magic = macho.CSMAGIC_CODEDIRECTORY,
-            .length = @sizeOf(macho.CodeDirectory),
-            .version = macho.CS_SUPPORTSEXECSEG,
-            .flags = macho.CS_ADHOC,
-            .hashOffset = 0,
-            .identOffset = 0,
-            .nSpecialSlots = 0,
-            .nCodeSlots = 0,
-            .codeLimit = @intCast(u32, file_size),
-            .hashSize = hash_size,
-            .hashType = macho.CS_HASHTYPE_SHA256,
-            .platform = 0,
-            .pageSize = @truncate(u8, std.math.log2(page_size)),
-            .spare2 = 0,
-            .scatterOffset = 0,
-            .teamOffset = 0,
-            .spare3 = 0,
-            .codeLimit64 = 0,
-            .execSegBase = execSegBase,
-            .execSegLimit = execSegLimit,
-            .execSegFlags = execSegFlags,
-        },
-    };
-
-    const total_pages = mem.alignForward(file_size, page_size) / page_size;
-
-    var hash: [hash_size]u8 = undefined;
-    var buffer = try bin_file.base.allocator.alloc(u8, page_size);
-    defer bin_file.base.allocator.free(buffer);
-    const macho_file = bin_file.base.file.?;
-
-    const id = bin_file.base.options.emit.?.sub_path;
-    try cdir.data.ensureCapacity(self.alloc, total_pages * hash_size + id.len + 1);
-
-    // 1. Save the identifier and update offsets
-    cdir.inner.identOffset = cdir.inner.length;
-    cdir.data.appendSliceAssumeCapacity(id);
-    cdir.data.appendAssumeCapacity(0);
-
-    // 2. Calculate hash for each page (in file) and write it to the buffer
-    // TODO figure out how we can cache several hashes since we won't update
-    // every page during incremental linking
-    cdir.inner.hashOffset = cdir.inner.identOffset + @intCast(u32, id.len) + 1;
-    var i: usize = 0;
-    while (i < total_pages) : (i += 1) {
-        const fstart = i * page_size;
-        const fsize = if (fstart + page_size > file_size) file_size - fstart else page_size;
-        const len = try macho_file.preadAll(buffer, fstart);
-        assert(fsize <= len);
-
-        Sha256.hash(buffer[0..fsize], &hash, .{});
-
-        cdir.data.appendSliceAssumeCapacity(hash[0..]);
-        cdir.inner.nCodeSlots += 1;
-    }
-
-    // 3. Update CodeDirectory length
-    cdir.inner.length += @intCast(u32, cdir.data.items.len);
-
-    self.inner.length += @sizeOf(macho.BlobIndex) + cdir.size();
-    self.inner.count = 1;
-    self.cdir = cdir;
-}
-
 pub fn size(self: CodeSignature) u32 {
     return self.inner.length;
 }
@@ -230,7 +152,7 @@ pub fn write(self: CodeSignature, buffer: []u8) void {
 
 pub fn deinit(self: *CodeSignature) void {
     if (self.cdir) |*cdir| {
-        cdir.data.deinit(self.alloc);
+        cdir.data.deinit(self.allocator);
     }
 }
 
diff --git a/src/link/MachO/Parser.zig b/src/link/MachO/Parser.zig
index 2f2e74933e..de2a466f1f 100644
--- a/src/link/MachO/Parser.zig
+++ b/src/link/MachO/Parser.zig
@@ -18,9 +18,9 @@ header: ?macho.mach_header_64 = null,
 /// Load commands
 load_commands: std.ArrayListUnmanaged(LoadCommand) = .{},
 
-text_cmd_index: ?usize = null,
+text_cmd_index: ?u16 = null,
 
-linkedit_cmd_index: ?usize = null,
+linkedit_cmd_index: ?u16 = null,
 linkedit_cmd_offset: ?u64 = null,
 
 code_sig_cmd_offset: ?u64 = null,
@@ -28,9 +28,7 @@ code_sig_cmd_offset: ?u64 = null,
 end_pos: ?u64 = null,
 
 pub fn init(allocator: *Allocator) Parser {
-    return .{
-        .allocator = allocator,
-    };
+    return .{ .allocator = allocator };
 }
 
 pub fn parse(self: *Parser, reader: anytype) !void {
@@ -46,10 +44,10 @@ pub fn parse(self: *Parser, reader: anytype) !void {
         switch (cmd.cmd()) {
             macho.LC_SEGMENT_64 => {
                 const x = cmd.Segment;
-                if (mem.eql(u8, mem.trimRight(u8, x.inner.segname[0..], &[_]u8{0}), "__LINKEDIT")) {
+                if (mem.eql(u8, parseName(&x.inner.segname), "__LINKEDIT")) {
                     self.linkedit_cmd_index = i;
                     self.linkedit_cmd_offset = off;
-                } else if (mem.eql(u8, mem.trimRight(u8, x.inner.segname[0..], &[_]u8{0}), "__TEXT")) {
+                } else if (mem.eql(u8, parseName(&x.inner.segname), "__TEXT")) {
                     self.text_cmd_index = i;
                 }
             },
@@ -78,3 +76,7 @@ pub fn deinit(self: *Parser) void {
     }
     self.load_commands.deinit(self.allocator);
 }
+
+fn parseName(name: *const [16]u8) []const u8 {
+    return mem.trimRight(u8, name.*[0..], &[_]u8{0});
+}