Commit 4a2f1e748e

Andrew Kelley <andrew@ziglang.org>
2022-10-23 01:03:59
link.MachO: multi-thread first round of sha256 hashing
when computing the adhoc code signature
1 parent e67c756
Changed files (3)
src/link/MachO/CodeSignature.zig
@@ -1,4 +1,6 @@
 const CodeSignature = @This();
+const Compilation = @import("../../Compilation.zig");
+const WaitGroup = @import("../../WaitGroup.zig");
 
 const std = @import("std");
 const assert = std.debug.assert;
@@ -258,17 +260,19 @@ pub const WriteOpts = struct {
 
 pub fn writeAdhocSignature(
     self: *CodeSignature,
-    allocator: Allocator,
+    comp: *const Compilation,
     opts: WriteOpts,
     writer: anytype,
 ) !void {
+    const gpa = comp.gpa;
+
     var header: macho.SuperBlob = .{
         .magic = macho.CSMAGIC_EMBEDDED_SIGNATURE,
         .length = @sizeOf(macho.SuperBlob),
         .count = 0,
     };
 
-    var blobs = std.ArrayList(Blob).init(allocator);
+    var blobs = std.ArrayList(Blob).init(gpa);
     defer blobs.deinit();
 
     self.code_directory.inner.execSegBase = opts.exec_seg_base;
@@ -276,37 +280,49 @@ pub fn writeAdhocSignature(
     self.code_directory.inner.execSegFlags = if (opts.output_mode == .Exe) macho.CS_EXECSEG_MAIN_BINARY else 0;
     self.code_directory.inner.codeLimit = opts.file_size;
 
-    const total_pages = mem.alignForward(opts.file_size, self.page_size) / self.page_size;
-
-    var buffer = try allocator.alloc(u8, self.page_size);
-    defer allocator.free(buffer);
+    const total_pages = @intCast(u32, mem.alignForward(opts.file_size, self.page_size) / self.page_size);
 
-    try self.code_directory.code_slots.ensureTotalCapacityPrecise(allocator, total_pages);
+    try self.code_directory.code_slots.ensureTotalCapacityPrecise(gpa, total_pages);
+    self.code_directory.code_slots.items.len = total_pages;
+    self.code_directory.inner.nCodeSlots += total_pages;
 
     // Calculate hash for each page (in file) and write it to the buffer
-    var hash: [hash_size]u8 = undefined;
-    var i: usize = 0;
-    while (i < total_pages) : (i += 1) {
-        const fstart = i * self.page_size;
-        const fsize = if (fstart + self.page_size > opts.file_size)
-            opts.file_size - fstart
-        else
-            self.page_size;
-        const len = try opts.file.preadAll(buffer, fstart);
-        assert(fsize <= len);
-
-        Sha256.hash(buffer[0..fsize], &hash, .{});
-
-        self.code_directory.code_slots.appendAssumeCapacity(hash);
-        self.code_directory.inner.nCodeSlots += 1;
+    var wg: WaitGroup = .{};
+    {
+        const results = try gpa.alloc(fs.File.PReadError!usize, total_pages);
+        defer gpa.free(results);
+        {
+            const buffer = try gpa.alloc(u8, self.page_size * total_pages);
+            defer gpa.free(buffer);
+
+            wg.reset();
+            defer wg.wait();
+
+            var i: usize = 0;
+            while (i < total_pages) : (i += 1) {
+                const fstart = i * self.page_size;
+                const fsize = if (fstart + self.page_size > opts.file_size)
+                    opts.file_size - fstart
+                else
+                    self.page_size;
+                const out_hash = &self.code_directory.code_slots.items[i];
+                wg.start();
+                try comp.thread_pool.spawn(workerSha256Hash, .{
+                    opts.file, fstart, buffer[0..fsize], out_hash, &results[i], &wg,
+                });
+            }
+        }
+        for (results) |result| _ = try result;
     }
 
     try blobs.append(.{ .code_directory = &self.code_directory });
     header.length += @sizeOf(macho.BlobIndex);
     header.count += 1;
 
+    var hash: [hash_size]u8 = undefined;
+
     if (self.requirements) |*req| {
-        var buf = std.ArrayList(u8).init(allocator);
+        var buf = std.ArrayList(u8).init(gpa);
         defer buf.deinit();
         try req.write(buf.writer());
         Sha256.hash(buf.items, &hash, .{});
@@ -318,7 +334,7 @@ pub fn writeAdhocSignature(
     }
 
     if (self.entitlements) |*ents| {
-        var buf = std.ArrayList(u8).init(allocator);
+        var buf = std.ArrayList(u8).init(gpa);
         defer buf.deinit();
         try ents.write(buf.writer());
         Sha256.hash(buf.items, &hash, .{});
@@ -356,6 +372,19 @@ pub fn writeAdhocSignature(
     }
 }
 
+fn workerSha256Hash(
+    file: fs.File,
+    fstart: usize,
+    buffer: []u8,
+    hash: *[hash_size]u8,
+    err: *fs.File.PReadError!usize,
+    wg: *WaitGroup,
+) void {
+    defer wg.finish();
+    err.* = file.preadAll(buffer, fstart);
+    Sha256.hash(buffer, hash, .{});
+}
+
 pub fn size(self: CodeSignature) u32 {
     var ssize: u32 = @sizeOf(macho.SuperBlob) + @sizeOf(macho.BlobIndex) + self.code_directory.size();
     if (self.requirements) |req| {
src/link/MachO/zld.zig
@@ -3035,14 +3035,19 @@ pub const Zld = struct {
         return @intCast(u32, offset);
     }
 
-    fn writeCodeSignature(self: *Zld, code_sig: *CodeSignature, offset: u32) !void {
+    fn writeCodeSignature(
+        self: *Zld,
+        comp: *const Compilation,
+        code_sig: *CodeSignature,
+        offset: u32,
+    ) !void {
         const seg_id = self.getSegmentByName("__TEXT").?;
         const seg = self.segments.items[seg_id];
 
         var buffer = std.ArrayList(u8).init(self.gpa);
         defer buffer.deinit();
         try buffer.ensureTotalCapacityPrecise(code_sig.size());
-        try code_sig.writeAdhocSignature(self.gpa, .{
+        try code_sig.writeAdhocSignature(comp, .{
             .file = self.file,
             .exec_seg_base = seg.fileoff,
             .exec_seg_limit = seg.filesize,
@@ -4379,7 +4384,7 @@ pub fn linkWithZld(macho_file: *MachO, comp: *Compilation, prog_node: *std.Progr
         try zld.writeHeader(ncmds, @intCast(u32, lc_buffer.items.len + headers_buf.items.len));
 
         if (codesig) |*csig| {
-            try zld.writeCodeSignature(csig, codesig_offset.?); // code signing always comes last
+            try zld.writeCodeSignature(comp, csig, codesig_offset.?); // code signing always comes last
         }
     }
 
src/link/MachO.zig
@@ -625,7 +625,7 @@ pub fn flushModule(self: *MachO, comp: *Compilation, prog_node: *std.Progress.No
     try self.writeHeader(ncmds, @intCast(u32, lc_buffer.items.len + headers_buf.items.len));
 
     if (codesig) |*csig| {
-        try self.writeCodeSignature(csig, codesig_offset.?); // code signing always comes last
+        try self.writeCodeSignature(comp, csig, codesig_offset.?); // code signing always comes last
     }
 
     if (self.d_sym) |*d_sym| {
@@ -4037,13 +4037,13 @@ fn writeCodeSignaturePadding(
     return @intCast(u32, offset);
 }
 
-fn writeCodeSignature(self: *MachO, code_sig: *CodeSignature, offset: u32) !void {
+fn writeCodeSignature(self: *MachO, comp: *const Compilation, code_sig: *CodeSignature, offset: u32) !void {
     const seg = self.getSegment(self.text_section_index.?);
 
     var buffer = std.ArrayList(u8).init(self.base.allocator);
     defer buffer.deinit();
     try buffer.ensureTotalCapacityPrecise(code_sig.size());
-    try code_sig.writeAdhocSignature(self.base.allocator, .{
+    try code_sig.writeAdhocSignature(comp, .{
         .file = self.base.file.?,
         .exec_seg_base = seg.fileoff,
         .exec_seg_limit = seg.filesize,