Commit f1e25cf43e

Jakub Konka <kubkon@jakubkonka.com>
2023-03-16 20:41:46
macho: add hot-code swapping poc
1 parent 266c813
lib/std/macho.zig
@@ -656,6 +656,10 @@ pub const segment_command_64 = extern struct {
     pub fn segName(seg: *const segment_command_64) []const u8 {
         return parseName(&seg.segname);
     }
+
+    pub fn isWriteable(seg: segment_command_64) bool {
+        return seg.initprot & PROT.WRITE != 0;
+    }
 };
 
 pub const PROT = struct {
src/link/MachO/Atom.zig
@@ -183,19 +183,11 @@ pub fn addLazyBinding(macho_file: *MachO, atom_index: Index, binding: Binding) !
     try gop.value_ptr.append(gpa, binding);
 }
 
-pub fn resolveRelocations(macho_file: *MachO, atom_index: Index) !void {
-    const atom = macho_file.getAtom(atom_index);
-    const relocs = macho_file.relocs.get(atom_index) orelse return;
-    const source_sym = atom.getSymbol(macho_file);
-    const source_section = macho_file.sections.get(source_sym.n_sect - 1).header;
-    const file_offset = source_section.offset + source_sym.n_value - source_section.addr;
-
-    log.debug("relocating '{s}'", .{atom.getName(macho_file)});
-
-    for (relocs.items) |*reloc| {
+pub fn resolveRelocations(macho_file: *MachO, atom_index: Index, relocs: []Relocation, code: []u8) !void {
+    log.debug("relocating '{s}'", .{macho_file.getAtom(atom_index).getName(macho_file)});
+    for (relocs) |*reloc| {
         if (!reloc.dirty) continue;
-
-        try reloc.resolve(macho_file, atom_index, file_offset);
+        try reloc.resolve(macho_file, atom_index, code);
         reloc.dirty = false;
     }
 }
src/link/MachO/Relocation.zig
@@ -50,7 +50,7 @@ pub fn getTargetAtomIndex(self: Relocation, macho_file: *MachO) ?Atom.Index {
     return macho_file.getAtomIndexForSymbol(self.target);
 }
 
-pub fn resolve(self: Relocation, macho_file: *MachO, atom_index: Atom.Index, base_offset: u64) !void {
+pub fn resolve(self: Relocation, macho_file: *MachO, atom_index: Atom.Index, code: []u8) !void {
     const arch = macho_file.base.options.target.cpu.arch;
     const atom = macho_file.getAtom(atom_index);
     const source_sym = atom.getSymbol(macho_file);
@@ -68,42 +68,28 @@ pub fn resolve(self: Relocation, macho_file: *MachO, atom_index: Atom.Index, bas
     });
 
     switch (arch) {
-        .aarch64 => return self.resolveAarch64(macho_file, source_addr, target_addr, base_offset),
-        .x86_64 => return self.resolveX8664(macho_file, source_addr, target_addr, base_offset),
+        .aarch64 => return self.resolveAarch64(source_addr, target_addr, code),
+        .x86_64 => return self.resolveX8664(source_addr, target_addr, code),
         else => unreachable,
     }
 }
 
 fn resolveAarch64(
     self: Relocation,
-    macho_file: *MachO,
     source_addr: u64,
     target_addr: i64,
-    base_offset: u64,
+    code: []u8,
 ) !void {
     const rel_type = @intToEnum(macho.reloc_type_arm64, self.type);
     if (rel_type == .ARM64_RELOC_UNSIGNED) {
-        var buffer: [@sizeOf(u64)]u8 = undefined;
-        const code = blk: {
-            switch (self.length) {
-                2 => {
-                    mem.writeIntLittle(u32, buffer[0..4], @truncate(u32, @bitCast(u64, target_addr)));
-                    break :blk buffer[0..4];
-                },
-                3 => {
-                    mem.writeIntLittle(u64, &buffer, @bitCast(u64, target_addr));
-                    break :blk &buffer;
-                },
-                else => unreachable,
-            }
+        return switch (self.length) {
+            2 => mem.writeIntLittle(u32, code[self.offset..][0..4], @truncate(u32, @bitCast(u64, target_addr))),
+            3 => mem.writeIntLittle(u64, code[self.offset..][0..8], @bitCast(u64, target_addr)),
+            else => unreachable,
         };
-        return macho_file.base.file.?.pwriteAll(code, base_offset + self.offset);
     }
 
-    var buffer: [@sizeOf(u32)]u8 = undefined;
-    const amt = try macho_file.base.file.?.preadAll(&buffer, base_offset + self.offset);
-    if (amt != buffer.len) return error.InputOutput;
-
+    var buffer = code[self.offset..][0..4];
     switch (rel_type) {
         .ARM64_RELOC_BRANCH26 => {
             const displacement = math.cast(
@@ -114,10 +100,10 @@ fn resolveAarch64(
                 .unconditional_branch_immediate = mem.bytesToValue(meta.TagPayload(
                     aarch64.Instruction,
                     aarch64.Instruction.unconditional_branch_immediate,
-                ), &buffer),
+                ), buffer),
             };
             inst.unconditional_branch_immediate.imm26 = @truncate(u26, @bitCast(u28, displacement >> 2));
-            mem.writeIntLittle(u32, &buffer, inst.toU32());
+            mem.writeIntLittle(u32, buffer, inst.toU32());
         },
         .ARM64_RELOC_PAGE21,
         .ARM64_RELOC_GOT_LOAD_PAGE21,
@@ -130,31 +116,31 @@ fn resolveAarch64(
                 .pc_relative_address = mem.bytesToValue(meta.TagPayload(
                     aarch64.Instruction,
                     aarch64.Instruction.pc_relative_address,
-                ), &buffer),
+                ), buffer),
             };
             inst.pc_relative_address.immhi = @truncate(u19, pages >> 2);
             inst.pc_relative_address.immlo = @truncate(u2, pages);
-            mem.writeIntLittle(u32, &buffer, inst.toU32());
+            mem.writeIntLittle(u32, buffer, inst.toU32());
         },
         .ARM64_RELOC_PAGEOFF12,
         .ARM64_RELOC_GOT_LOAD_PAGEOFF12,
         => {
             const narrowed = @truncate(u12, @intCast(u64, target_addr));
-            if (isArithmeticOp(&buffer)) {
+            if (isArithmeticOp(buffer)) {
                 var inst = aarch64.Instruction{
                     .add_subtract_immediate = mem.bytesToValue(meta.TagPayload(
                         aarch64.Instruction,
                         aarch64.Instruction.add_subtract_immediate,
-                    ), &buffer),
+                    ), buffer),
                 };
                 inst.add_subtract_immediate.imm12 = narrowed;
-                mem.writeIntLittle(u32, &buffer, inst.toU32());
+                mem.writeIntLittle(u32, buffer, inst.toU32());
             } else {
                 var inst = aarch64.Instruction{
                     .load_store_register = mem.bytesToValue(meta.TagPayload(
                         aarch64.Instruction,
                         aarch64.Instruction.load_store_register,
-                    ), &buffer),
+                    ), buffer),
                 };
                 const offset: u12 = blk: {
                     if (inst.load_store_register.size == 0) {
@@ -170,7 +156,7 @@ fn resolveAarch64(
                     }
                 };
                 inst.load_store_register.offset = offset;
-                mem.writeIntLittle(u32, &buffer, inst.toU32());
+                mem.writeIntLittle(u32, buffer, inst.toU32());
             }
         },
         .ARM64_RELOC_TLVP_LOAD_PAGEOFF12 => {
@@ -180,11 +166,11 @@ fn resolveAarch64(
                 size: u2,
             };
             const reg_info: RegInfo = blk: {
-                if (isArithmeticOp(&buffer)) {
+                if (isArithmeticOp(buffer)) {
                     const inst = mem.bytesToValue(meta.TagPayload(
                         aarch64.Instruction,
                         aarch64.Instruction.add_subtract_immediate,
-                    ), &buffer);
+                    ), buffer);
                     break :blk .{
                         .rd = inst.rd,
                         .rn = inst.rn,
@@ -194,7 +180,7 @@ fn resolveAarch64(
                     const inst = mem.bytesToValue(meta.TagPayload(
                         aarch64.Instruction,
                         aarch64.Instruction.load_store_register,
-                    ), &buffer);
+                    ), buffer);
                     break :blk .{
                         .rd = inst.rt,
                         .rn = inst.rn,
@@ -214,72 +200,62 @@ fn resolveAarch64(
                     .sf = @truncate(u1, reg_info.size),
                 },
             };
-            mem.writeIntLittle(u32, &buffer, inst.toU32());
+            mem.writeIntLittle(u32, buffer, inst.toU32());
         },
         .ARM64_RELOC_POINTER_TO_GOT => {
             const result = @intCast(i32, @intCast(i64, target_addr) - @intCast(i64, source_addr));
-            mem.writeIntLittle(i32, &buffer, result);
+            mem.writeIntLittle(i32, buffer, result);
         },
         .ARM64_RELOC_SUBTRACTOR => unreachable,
         .ARM64_RELOC_ADDEND => unreachable,
         .ARM64_RELOC_UNSIGNED => unreachable,
     }
-    try macho_file.base.file.?.pwriteAll(&buffer, base_offset + self.offset);
 }
 
 fn resolveX8664(
     self: Relocation,
-    macho_file: *MachO,
     source_addr: u64,
     target_addr: i64,
-    base_offset: u64,
+    code: []u8,
 ) !void {
     const rel_type = @intToEnum(macho.reloc_type_x86_64, self.type);
-    var buffer: [@sizeOf(u64)]u8 = undefined;
-    const code = blk: {
-        switch (rel_type) {
-            .X86_64_RELOC_BRANCH,
-            .X86_64_RELOC_GOT,
-            .X86_64_RELOC_GOT_LOAD,
-            .X86_64_RELOC_TLV,
-            => {
-                const displacement = @intCast(i32, @intCast(i64, target_addr) - @intCast(i64, source_addr) - 4);
-                mem.writeIntLittle(u32, buffer[0..4], @bitCast(u32, displacement));
-                break :blk buffer[0..4];
-            },
-            .X86_64_RELOC_SIGNED,
-            .X86_64_RELOC_SIGNED_1,
-            .X86_64_RELOC_SIGNED_2,
-            .X86_64_RELOC_SIGNED_4,
-            => {
-                const correction: u3 = switch (rel_type) {
-                    .X86_64_RELOC_SIGNED => 0,
-                    .X86_64_RELOC_SIGNED_1 => 1,
-                    .X86_64_RELOC_SIGNED_2 => 2,
-                    .X86_64_RELOC_SIGNED_4 => 4,
-                    else => unreachable,
-                };
-                const displacement = @intCast(i32, target_addr - @intCast(i64, source_addr + correction + 4));
-                mem.writeIntLittle(u32, buffer[0..4], @bitCast(u32, displacement));
-                break :blk buffer[0..4];
-            },
-            .X86_64_RELOC_UNSIGNED => {
-                switch (self.length) {
-                    2 => {
-                        mem.writeIntLittle(u32, buffer[0..4], @truncate(u32, @bitCast(u64, target_addr)));
-                        break :blk buffer[0..4];
-                    },
-                    3 => {
-                        mem.writeIntLittle(u64, buffer[0..8], @bitCast(u64, target_addr));
-                        break :blk &buffer;
-                    },
-                    else => unreachable,
-                }
-            },
-            .X86_64_RELOC_SUBTRACTOR => unreachable,
-        }
-    };
-    try macho_file.base.file.?.pwriteAll(code, base_offset + self.offset);
+    switch (rel_type) {
+        .X86_64_RELOC_BRANCH,
+        .X86_64_RELOC_GOT,
+        .X86_64_RELOC_GOT_LOAD,
+        .X86_64_RELOC_TLV,
+        => {
+            const displacement = @intCast(i32, @intCast(i64, target_addr) - @intCast(i64, source_addr) - 4);
+            mem.writeIntLittle(u32, code[self.offset..][0..4], @bitCast(u32, displacement));
+        },
+        .X86_64_RELOC_SIGNED,
+        .X86_64_RELOC_SIGNED_1,
+        .X86_64_RELOC_SIGNED_2,
+        .X86_64_RELOC_SIGNED_4,
+        => {
+            const correction: u3 = switch (rel_type) {
+                .X86_64_RELOC_SIGNED => 0,
+                .X86_64_RELOC_SIGNED_1 => 1,
+                .X86_64_RELOC_SIGNED_2 => 2,
+                .X86_64_RELOC_SIGNED_4 => 4,
+                else => unreachable,
+            };
+            const displacement = @intCast(i32, target_addr - @intCast(i64, source_addr + correction + 4));
+            mem.writeIntLittle(u32, code[self.offset..][0..4], @bitCast(u32, displacement));
+        },
+        .X86_64_RELOC_UNSIGNED => {
+            switch (self.length) {
+                2 => {
+                    mem.writeIntLittle(u32, code[self.offset..][0..4], @truncate(u32, @bitCast(u64, target_addr)));
+                },
+                3 => {
+                    mem.writeIntLittle(u64, code[self.offset..][0..8], @bitCast(u64, target_addr));
+                },
+                else => unreachable,
+            }
+        },
+        .X86_64_RELOC_SUBTRACTOR => unreachable,
+    }
 }
 
 inline fn isArithmeticOp(inst: *const [4]u8) bool {
src/link/MachO.zig
@@ -221,6 +221,9 @@ lazy_bindings: BindingTable = .{},
 /// Table of tracked Decls.
 decls: std.AutoArrayHashMapUnmanaged(Module.Decl.Index, DeclMetadata) = .{},
 
+/// Mach task used when the compiler is in hot-code swapping mode.
+mach_task: ?std.os.darwin.MachTask = null,
+
 const DeclMetadata = struct {
     atom: Atom.Index,
     section: u8,
@@ -584,7 +587,21 @@ pub fn flushModule(self: *MachO, comp: *Compilation, prog_node: *std.Progress.No
     try self.allocateSpecialSymbols();
 
     for (self.relocs.keys()) |atom_index| {
-        try Atom.resolveRelocations(self, atom_index);
+        if (self.relocs.get(atom_index) == null) continue;
+
+        const atom = self.getAtom(atom_index);
+        const sym = atom.getSymbol(self);
+        const section = self.sections.get(sym.n_sect - 1).header;
+        const file_offset = section.offset + sym.n_value - section.addr;
+
+        var code = std.ArrayList(u8).init(self.base.allocator);
+        defer code.deinit();
+        try code.resize(atom.size);
+
+        const amt = try self.base.file.?.preadAll(code.items, file_offset);
+        if (amt != code.items.len) return error.InputOutput;
+
+        try self.writeAtom(atom_index, code.items);
     }
 
     if (build_options.enable_logging) {
@@ -1052,14 +1069,38 @@ pub fn parseDependentLibs(self: *MachO, syslibroot: ?[]const u8, dependent_libs:
     }
 }
 
-pub fn writeAtom(self: *MachO, atom_index: Atom.Index, code: []const u8) !void {
+pub fn writeAtom(self: *MachO, atom_index: Atom.Index, code: []u8) !void {
     const atom = self.getAtom(atom_index);
     const sym = atom.getSymbol(self);
     const section = self.sections.get(sym.n_sect - 1);
     const file_offset = section.header.offset + sym.n_value - section.header.addr;
     log.debug("writing atom for symbol {s} at file offset 0x{x}", .{ atom.getName(self), file_offset });
+
+    if (self.relocs.get(atom_index)) |relocs| {
+        try Atom.resolveRelocations(self, atom_index, relocs.items, code);
+    }
+
+    if (self.base.child_pid) |pid| blk: {
+        const task = self.mach_task orelse {
+            log.warn("cannot hot swap: no Mach task acquired for child process with pid {d}", .{pid});
+            break :blk;
+        };
+        self.writeAtomToMemory(task, section.segment_index, sym.n_value, code) catch |err| {
+            log.warn("cannot hot swap: writing to memory failed: {s}", .{@errorName(err)});
+        };
+    }
+
     try self.base.file.?.pwriteAll(code, file_offset);
-    try Atom.resolveRelocations(self, atom_index);
+}
+
+fn writeAtomToMemory(self: *MachO, task: std.os.darwin.MachTask, segment_index: u8, addr: u64, code: []const u8) !void {
+    const segment = self.segments.items[segment_index];
+    if (!segment.isWriteable()) {
+        try task.setCurrProtection(addr, code.len, macho.PROT.READ | macho.PROT.WRITE | macho.PROT.COPY);
+    }
+    defer if (!segment.isWriteable()) task.setCurrProtection(addr, code.len, segment.initprot) catch {};
+    const nwritten = try task.writeMem(addr, code, self.base.options.target.cpu.arch);
+    if (nwritten != code.len) return error.InputOutput;
 }
 
 fn writePtrWidthAtom(self: *MachO, atom_index: Atom.Index) !void {
@@ -2063,7 +2104,7 @@ pub fn updateFunc(self: *MachO, module: *Module, func: *Module.Fn, air: Air, liv
     else
         try codegen.generateFunction(&self.base, decl.srcLoc(), func, air, liveness, &code_buffer, .none);
 
-    const code = switch (res) {
+    var code = switch (res) {
         .ok => code_buffer.items,
         .fail => |em| {
             decl.analysis = .codegen_failure;
@@ -2115,7 +2156,7 @@ pub fn lowerUnnamedConst(self: *MachO, typed_value: TypedValue, decl_index: Modu
     const res = try codegen.generateSymbol(&self.base, decl.srcLoc(), typed_value, &code_buffer, .none, .{
         .parent_atom_index = self.getAtom(atom_index).getSymbolIndex().?,
     });
-    const code = switch (res) {
+    var code = switch (res) {
         .ok => code_buffer.items,
         .fail => |em| {
             decl.analysis = .codegen_failure;
@@ -2202,7 +2243,7 @@ pub fn updateDecl(self: *MachO, module: *Module, decl_index: Module.Decl.Index)
             .parent_atom_index = atom.getSymbolIndex().?,
         });
 
-    const code = switch (res) {
+    var code = switch (res) {
         .ok => code_buffer.items,
         .fail => |em| {
             decl.analysis = .codegen_failure;
@@ -2375,7 +2416,7 @@ pub fn getOutputSection(self: *MachO, sect: macho.section_64) !?u8 {
     return sect_id;
 }
 
-fn updateDeclCode(self: *MachO, decl_index: Module.Decl.Index, code: []const u8) !u64 {
+fn updateDeclCode(self: *MachO, decl_index: Module.Decl.Index, code: []u8) !u64 {
     const gpa = self.base.allocator;
     const mod = self.base.options.module.?;
     const decl = mod.declPtr(decl_index);
src/link.zig
@@ -392,6 +392,19 @@ pub const File = struct {
                         .linux => std.os.ptrace(std.os.linux.PTRACE.ATTACH, pid, 0, 0) catch |err| {
                             log.warn("ptrace failure: {s}", .{@errorName(err)});
                         },
+                        .macos => {
+                            const macho = base.cast(MachO).?;
+                            if (macho.mach_task == null) {
+                                if (std.os.darwin.machTaskForPid(pid)) |task| {
+                                    macho.mach_task = task;
+                                    std.os.ptrace(std.os.darwin.PT.ATTACHEXC, pid, 0, 0) catch |err| {
+                                        log.warn("ptrace failure: {s}", .{@errorName(err)});
+                                    };
+                                } else |err| {
+                                    log.warn("failed to acquire Mach task for child process: {s}", .{@errorName(err)});
+                                }
+                            }
+                        },
                         else => return error.HotSwapUnavailableOnHostOperatingSystem,
                     }
                 }
@@ -430,6 +443,9 @@ pub const File = struct {
                         .linux => std.os.ptrace(std.os.linux.PTRACE.DETACH, pid, 0, 0) catch |err| {
                             log.warn("ptrace failure: {s}", .{@errorName(err)});
                         },
+                        .macos => std.os.ptrace(std.os.darwin.PT.KILL, pid, 0, 0) catch |err| {
+                            log.warn("ptrace failure: {s}", .{@errorName(err)});
+                        },
                         else => return error.HotSwapUnavailableOnHostOperatingSystem,
                     }
                 }
src/main.zig
@@ -3851,15 +3851,39 @@ fn runOrTestHotSwap(
     if (runtime_args_start) |i| {
         try argv.appendSlice(all_args[i..]);
     }
-    var child = std.ChildProcess.init(argv.items, gpa);
 
-    child.stdin_behavior = .Inherit;
-    child.stdout_behavior = .Inherit;
-    child.stderr_behavior = .Inherit;
+    switch (builtin.target.os.tag) {
+        .macos, .ios, .tvos, .watchos => {
+            const PosixSpawn = std.os.darwin.PosixSpawn;
+            var attr = try PosixSpawn.Attr.init();
+            defer attr.deinit();
+            const flags: u16 = std.os.darwin.POSIX_SPAWN_SETSIGDEF |
+                std.os.darwin.POSIX_SPAWN_SETSIGMASK |
+                std.os.darwin._POSIX_SPAWN_DISABLE_ASLR;
+            try attr.set(flags);
+
+            var arena_allocator = std.heap.ArenaAllocator.init(gpa);
+            defer arena_allocator.deinit();
+            const arena = arena_allocator.allocator();
+
+            const argv_buf = try arena.allocSentinel(?[*:0]u8, argv.items.len, null);
+            for (argv.items, 0..) |arg, i| argv_buf[i] = (try arena.dupeZ(u8, arg)).ptr;
+
+            const pid = try PosixSpawn.spawn(argv.items[0], null, attr, argv_buf, std.c.environ);
+            return pid;
+        },
+        else => {
+            var child = std.ChildProcess.init(argv.items, gpa);
+
+            child.stdin_behavior = .Inherit;
+            child.stdout_behavior = .Inherit;
+            child.stderr_behavior = .Inherit;
 
-    try child.spawn();
+            try child.spawn();
 
-    return child.id;
+            return child.id;
+        },
+    }
 }
 
 const AfterUpdateHook = union(enum) {
build.zig
@@ -152,6 +152,7 @@ pub fn build(b: *std.Build) !void {
     if (only_install_lib_files)
         return;
 
+    const entitlements = b.option([]const u8, "entitlements", "Path to entitlements file for hot-code swapping without sudo on macOS");
     const tracy = b.option([]const u8, "tracy", "Enable Tracy integration. Supply path to Tracy source");
     const tracy_callstack = b.option(bool, "tracy-callstack", "Include callstack information with Tracy data. Does nothing if -Dtracy is not provided") orelse (tracy != null);
     const tracy_allocation = b.option(bool, "tracy-allocation", "Include allocation information with Tracy data. Does nothing if -Dtracy is not provided") orelse (tracy != null);
@@ -173,6 +174,7 @@ pub fn build(b: *std.Build) !void {
     exe.pie = pie;
     exe.sanitize_thread = sanitize_thread;
     exe.build_id = b.option(bool, "build-id", "Include a build id note") orelse false;
+    exe.entitlements = entitlements;
     exe.install();
 
     const compile_step = b.step("compile", "Build the self-hosted compiler");