fix SIGBUS due to mapping being larger than file

docs/TODO.md

@@ -19,10 +19,10 @@ IvyBridge(2012) and AMD Zen 2 Family 17H(2019) and Linux 5.9(2020).
       the stack, so `ucontext` isn't on top anymore.
 - [x] `/proc/self/exe`: intercept calls to `readlink`/`readlinkat` with that as argument
 - [x] `auxv`: check if that is setup correctly and completely
-- [ ] JIT support: intercept `mmap`, `mprotect` and `mremap` that change pages to be executable
+- [x] JIT support: intercept `mmap`, `mprotect` and `mremap` that change pages to be executable
 - [ ] `SIGILL` patching fallback
 - [x] `vdso` handling
-- [ ] check why the libc tests are flaky
+- [x] check why the libc tests are flaky
 
 ## Minor things
 
@@ -31,6 +31,8 @@ IvyBridge(2012) and AMD Zen 2 Family 17H(2019) and Linux 5.9(2020).
 - [ ] Ghost page edge case: In all patch strategies, if a range spans multiple pages and we `mmap`
     the first one but can't `mmap` the second one we just let the first one mapped. It would be better
     to unmap them
+- [ ] Right now when patching we mmap a page and may not use it, but we still leave it mapped. This
+      leaks memory. If we fix this correctly the Ghost page issue is also fixed
 - [ ] Re-entrancy for `patchRegion`
     - when a signal comes, while we are in that function, and we need to patch something due to the
       signal we will deadlock
@@ -46,3 +48,5 @@ IvyBridge(2012) and AMD Zen 2 Family 17H(2019) and Linux 5.9(2020).
 - [ ] `modify_ldt`: check what we need to intercept and change
 - [ ] `set_tid_address`: check what we need to intercept and change
 - [ ] performance optimizations for patched code? Peephole might be possible
+- [ ] maybe add a way to run something after the client is finished
+    - could be useful for statistics, cleanup(if necessary), or notifying of suppressed warnings

src/Patcher.zig

@@ -209,9 +209,13 @@ pub const Statistics = struct {
 /// Scans a memory region for instructions that require patching and applies the patches
 /// using a hierarchy of tactics (Direct/Punning -> Successor Eviction -> Neighbor Eviction).
 ///
-/// The region is processed Back-to-Front to ensure that modifications (punning) only
-/// constrain instructions that have already been processed or are locked.
+/// NOTE: This function leaves the region as R|W and the caller is responsible for changing it to
+/// the desired protection
 pub fn patchRegion(region: []align(page_size) u8) !void {
+    log.info(
+        "Patching region: 0x{x} - 0x{x}",
+        .{ @intFromPtr(region.ptr), @intFromPtr(&region[region.len - 1]) },
+    );
     // For now just do a coarse lock.
     // TODO: should we make this more fine grained?
     mutex.lock();
@@ -296,8 +300,6 @@ pub fn patchRegion(region: []align(page_size) u8) !void {
     {
         // Apply patches.
         try posix.mprotect(region, posix.PROT.READ | posix.PROT.WRITE);
-        defer posix.mprotect(region, posix.PROT.READ | posix.PROT.EXEC) catch
-            @panic("patchRegion: mprotect back to R|X failed. Can't continue");
 
         var stats = Statistics.empty;
         // Used to track which bytes have been modified or used for constraints (punning),
@@ -854,7 +856,7 @@ fn ensureRangeWritable(
         const gop = try allocated_pages.getOrPut(gpa, page_addr);
         if (gop.found_existing) {
             const ptr: [*]align(page_size) u8 = @ptrFromInt(page_addr);
-            try posix.mprotect(ptr[0..page_addr], protection);
+            try posix.mprotect(ptr[0..page_size], protection);
         } else {
             const addr = posix.mmap(
                 @ptrFromInt(page_addr),

src/main.zig

@@ -256,11 +256,12 @@ fn patchLoadedElf(base: usize) !void {
 
         const page_start = mem.alignBackward(usize, vaddr, page_size);
         const page_end = mem.alignForward(usize, vaddr + memsz, page_size);
+        const size = page_end - page_start;
 
-        const region = @as([*]align(page_size) u8, @ptrFromInt(page_start))[0 .. page_end - page_start];
+        const region = @as([*]align(page_size) u8, @ptrFromInt(page_start))[0..size];
 
-        log.info("Patching segment: 0x{x} - 0x{x}", .{ page_start, page_end });
         try Patcher.patchRegion(region);
+        try posix.mprotect(region, elfToMmapProt(phdr.p_flags));
     }
 }
 
@@ -332,10 +333,9 @@ test "nolibc_nopie_exit" {
 test "nolibc_pie_exit" {
     try testHelper(&.{ flicker_path, getTestExePath("nolibc_pie_exit") }, "");
 }
-// BUG: This one is flaky
-// test "libc_pie_exit" {
-//     try testHelper(&.{ flicker_path, getTestExePath("libc_pie_exit") }, "");
-// }
+test "libc_pie_exit" {
+    try testHelper(&.{ flicker_path, getTestExePath("libc_pie_exit") }, "");
+}
 
 test "nolibc_nopie_helloWorld" {
     try testHelper(&.{ flicker_path, getTestExePath("nolibc_nopie_helloWorld") }, "Hello World!\n");
@@ -343,10 +343,9 @@ test "nolibc_nopie_helloWorld" {
 test "nolibc_pie_helloWorld" {
     try testHelper(&.{ flicker_path, getTestExePath("nolibc_pie_helloWorld") }, "Hello World!\n");
 }
-// BUG: This one is flaky
-// test "libc_pie_helloWorld" {
-//     try testHelper(&.{ flicker_path, getTestExePath("libc_pie_helloWorld") }, "Hello World!\n");
-// }
+test "libc_pie_helloWorld" {
+    try testHelper(&.{ flicker_path, getTestExePath("libc_pie_helloWorld") }, "Hello World!\n");
+}
 
 test "nolibc_nopie_printArgs" {
     try testPrintArgs("nolibc_nopie_printArgs");
@@ -354,10 +353,9 @@ test "nolibc_nopie_printArgs" {
 test "nolibc_pie_printArgs" {
     try testPrintArgs("nolibc_pie_printArgs");
 }
-// BUG: This one is flaky
-// test "libc_pie_printArgs" {
-//     try testPrintArgs("libc_pie_printArgs");
-// }
+test "libc_pie_printArgs" {
+    try testPrintArgs("libc_pie_printArgs");
+}
 
 test "nolibc_nopie_readlink" {
     try testReadlink("nolibc_nopie_readlink");
@@ -365,10 +363,9 @@ test "nolibc_nopie_readlink" {
 test "nolibc_pie_readlink" {
     try testReadlink("nolibc_pie_readlink");
 }
-// BUG: This one just outputs the path to the flicker executable and is likely also flaky
-// test "libc_pie_readlink" {
-//     try testReadlink("libc_pie_readlink");
-// }
+test "libc_pie_readlink" {
+    try testReadlink("libc_pie_readlink");
+}
 
 test "nolibc_nopie_clone_raw" {
     try testHelper(
@@ -396,10 +393,6 @@ test "nolibc_pie_clone_no_new_stack" {
     );
 }
 
-test "echo" {
-    try testHelper(&.{ "echo", "Hello", "There" }, "Hello There\n");
-}
-
 test "nolibc_nopie_fork" {
     try testHelper(
         &.{ flicker_path, getTestExePath("nolibc_nopie_fork") },
@@ -412,13 +405,12 @@ test "nolibc_pie_fork" {
         "Child: I'm alive!\nParent: Child died.\n",
     );
 }
-// BUG: This one is flaky
-// test "libc_pie_fork" {
-//     try testHelper(
-//         &.{ flicker_path, getTestExePath("libc_pie_fork") },
-//         "Child: I'm alive!\nParent: Child died.\n",
-//     );
-// }
+test "libc_pie_fork" {
+    try testHelper(
+        &.{ flicker_path, getTestExePath("libc_pie_fork") },
+        "Child: I'm alive!\nParent: Child died.\n",
+    );
+}
 
 test "nolibc_nopie_signal_handler" {
     try testHelper(
@@ -445,13 +437,16 @@ test "nolibc_pie_vdso_clock" {
         "Time gotten\n",
     );
 }
-// BUG: This one is flaky
-// test "libc_pie_vdso_clock" {
-//     try testHelper(
-//         &.{ flicker_path, getTestExePath("libc_pie_vdso_clock") },
-//         "Time gotten\n",
-//     );
-// }
+test "libc_pie_vdso_clock" {
+    try testHelper(
+        &.{ flicker_path, getTestExePath("libc_pie_vdso_clock") },
+        "Time gotten\n",
+    );
+}
+
+test "echo" {
+    try testHelper(&.{ "echo", "Hello", "There" }, "Hello There\n");
+}
 
 fn testPrintArgs(comptime name: []const u8) !void {
     const exe_path = getTestExePath(name);

src/syscalls.zig

@@ -1,8 +1,13 @@
 const std = @import("std");
 const linux = std.os.linux;
+const posix = std.posix;
 const Patcher = @import("Patcher.zig");
 const assert = std.debug.assert;
 
+const page_size = std.heap.pageSize();
+
+const log = std.log.scoped(.syscalls);
+
 /// Represents the stack layout pushed by `syscallEntry` before calling the handler.
 pub const SavedContext = extern struct {
     padding: u64, // Result of `sub $8, %rsp` for alignment
@@ -74,38 +79,93 @@ export fn syscall_handler(ctx: *SavedContext) callconv(.c) void {
             asm volatile (
                 \\ mov %[rsp], %%rsp
                 \\ syscall
-                \\ ud2
                 :
                 : [rsp] "r" (rsp_orig),
                   [number] "{rax}" (ctx.rax),
-                : .{ .memory = true }
-            );
+                : .{ .memory = true });
             unreachable;
         },
-        .execve, .execveat => |s| {
-            // TODO: option to persist across new processes
-            std.debug.print("syscall {} called\n", .{s});
+        .mmap => {
+            // mmap(void *addr, size_t length, int prot, int flags, int fd, off_t offset)
+
+            const prot: u32 = @intCast(ctx.rdx);
+            // Execute the syscall first to get the address (rax)
+            ctx.rax = executeSyscall(ctx);
+            const addr = ctx.rax;
+            var len = ctx.rsi;
+            const flags: linux.MAP = @bitCast(@as(u32, @intCast(ctx.r10)));
+            const fd: linux.fd_t = @bitCast(@as(u32, @truncate(ctx.r8)));
+            const offset = ctx.r9;
+
+            const is_error = @as(i64, @bitCast(ctx.rax)) < 0;
+            if (is_error) return;
+            if ((prot & posix.PROT.EXEC) == 0) return;
+
+            // If file-backed (not anonymous), clamp len to file size to avoid SIGBUS
+            if (!flags.ANONYMOUS) {
+                var stat: linux.Stat = undefined;
+                if (0 == linux.fstat(fd, &stat) and linux.S.ISREG(stat.mode)) {
+                    const file_size: u64 = @intCast(stat.size);
+                    len = if (offset >= file_size) 0 else @min(len, file_size - offset);
+                }
+            }
+
+            if (len <= 0) return;
+            // mmap addresses are always page aligned
+            const ptr = @as([*]align(page_size) u8, @ptrFromInt(addr));
+            // Check if we can patch it
+            Patcher.patchRegion(ptr[0..len]) catch |err| {
+                std.log.warn("JIT Patching failed: {}", .{err});
+            };
+
+            // patchRegion leaves it as RW. We need to restore to requested prot.
+            _ = linux.syscall3(.mprotect, addr, len, prot);
+            return;
         },
-        .prctl, .arch_prctl, .set_tid_address => |s| {
+        .mprotect => {
+            // mprotect(void *addr, size_t len, int prot)
+            // TODO: cleanup trampolines, when removing X
+            const prot: u32 = @intCast(ctx.rdx);
+            if ((prot & posix.PROT.EXEC) != 0) {
+                const addr = ctx.rdi;
+                const len = ctx.rsi;
+                // mprotect requires addr to be page aligned.
+                if (len > 0 and std.mem.isAligned(addr, page_size)) {
+                    const ptr = @as([*]align(page_size) u8, @ptrFromInt(addr));
+                    Patcher.patchRegion(ptr[0..len]) catch |err| {
+                        std.log.warn("mprotect Patching failed: {}", .{err});
+                    };
+                    // patchRegion leaves it R|W.
+                }
+            }
+            ctx.rax = executeSyscall(ctx);
+            return;
+        },
+        .execve, .execveat => {
+            // TODO: option to persist across new processes
+            ctx.rax = executeSyscall(ctx);
+            return;
+        },
+        .prctl, .arch_prctl, .set_tid_address => {
             // TODO: what do we need to handle from these?
             // process name
             // fs base(gs?)
             // thread id pointers
-            std.debug.print("syscall {} called\n", .{s});
-        },
-        .mmap, .mprotect => {
-            // TODO: JIT support
-            // TODO: cleanup
+            ctx.rax = executeSyscall(ctx);
+            return;
         },
         .munmap, .mremap => {
             // TODO: cleanup
+            ctx.rax = executeSyscall(ctx);
+            return;
         },
-
-        else => {},
-    }
-
+        else => {
             // Write result back to the saved RAX so it is restored to the application.
             ctx.rax = executeSyscall(ctx);
+            return;
+        },
+    }
+    unreachable;
 }
 
 inline fn executeSyscall(ctx: *SavedContext) u64 {

src/test/clone_no_new_stack.zig

@@ -30,9 +30,6 @@ pub fn main() !void {
         \\ mov $60, %%rax       # SYS_exit
         \\ syscall
         \\
-        \\ # Should not be reached
-        \\ ud2
-        \\
         \\ 1:
         \\ # Parent Path continues
         : [ret] "={rax}" (-> usize),

src/test/clone_raw.zig

@@ -34,9 +34,6 @@ pub fn main() !void {
         \\ mov $60, %%rax       # SYS_exit
         \\ syscall
         \\
-        \\ # Should not be reached
-        \\ ud2
-        \\
         \\ 1:
         \\ # Parent Path continues
         : [ret] "={rax}" (-> usize),