Override libc’s malloc with Zig

The LD_PRELOAD trick allows you load a dynamic library before any other libraries are loaded. You can use this to override functions inside an application at runtime. This is often used to override functions in libc to allow some debugging.

Doing it with Zig is cool because Zig itself does not rely on libc. So if we override the malloc family of functions we don’t have to worry about malloc recursively calling itself.

We can also use the Zig standard library and you don’t need to worry about how to implement malloc because it can be copied and pasted from ziglibc with some minor alterations.

malloc.zig

const builtin = @import("builtin");
const std = @import("std");

const alloc_align = 16;
const alloc_metadata_len = std.mem.alignForward(@sizeOf(usize), alloc_align);
var gpa = std.heap.GeneralPurposeAllocator(.{
    .MutexType = std.Thread.Mutex,
}){};

export fn malloc(size: usize) callconv(.C) ?[*]align(alloc_align) u8 {
    std.debug.assert(size > 0); // TODO: what should we do in this case?
    const full_len = alloc_metadata_len + size;
    const buf = gpa.allocator().alignedAlloc(u8, alloc_align, full_len) catch |err| switch (err) {
        error.OutOfMemory => {
            std.log.info("malloc return null", .{});
            return null;
        },
    };
    @ptrCast(*usize, buf).* = full_len;
    const result = @intToPtr([*]align(alloc_align) u8, @ptrToInt(buf.ptr) + alloc_metadata_len);
    std.log.info("malloc({}) return {*}", .{ size, result });
    return result;
}

fn getGpaBuf(ptr: [*]u8) []align(alloc_align) u8 {
    const start = @ptrToInt(ptr) - alloc_metadata_len;
    const len = @intToPtr(*usize, start).*;
    return @alignCast(alloc_align, @intToPtr([*]u8, start)[0..len]);
}

export fn realloc(ptr: ?[*]align(alloc_align) u8, size: usize) callconv(.C) ?[*]align(alloc_align) u8 {
    std.log.info("realloc {*} {}", .{ ptr, size });
    const gpa_buf = getGpaBuf(ptr orelse {
        const result = malloc(size);
        std.log.info("realloc return {*} (from malloc)", .{result});
        return result;
    });
    if (size == 0) {
        gpa.allocator().free(gpa_buf);
        return null;
    }

    const gpa_size = alloc_metadata_len + size;
    if (gpa.allocator().rawResize(gpa_buf, std.math.log2(alloc_align), gpa_size, @returnAddress())) {
        @ptrCast(*usize, gpa_buf.ptr).* = gpa_size;
        std.log.info("realloc return {*}", .{ptr});
        return ptr;
    }

    const new_buf = gpa.allocator().reallocAdvanced(
        gpa_buf,
        gpa_size,
        @returnAddress(),
    ) catch |e| switch (e) {
        error.OutOfMemory => {
            std.log.info("realloc out-of-mem from {} to {}", .{ gpa_buf.len, gpa_size });
            return null;
        },
    };
    @ptrCast(*usize, new_buf.ptr).* = gpa_size;
    const result = @intToPtr([*]align(alloc_align) u8, @ptrToInt(new_buf.ptr) + alloc_metadata_len);
    std.log.info("realloc return {*}", .{result});
    return result;
}

export fn calloc(nmemb: usize, size: usize) callconv(.C) ?[*]align(alloc_align) u8 {
    const total = std.math.mul(usize, nmemb, size) catch {
        // TODO: set errno
        //errno = c.ENOMEM;
        return null;
    };
    const ptr = malloc(total) orelse return null;
    @memset(ptr[0..total], 0);
    return ptr;
}

export fn free(ptr: ?[*]align(alloc_align) u8) callconv(.C) void {
    std.log.info("free {*}", .{ptr});
    const p = ptr orelse return;
    gpa.allocator().free(getGpaBuf(p));
}

You can see that realloc complicates things a bit. It needs to be included otherwise we could have pointers allocated by libc passed to our free.

Build and run

We can compile the above and run it like so:

$ zig build-lib malloc.zig -dynamic
$ LD_PRELOAD=./libmalloc.so ls
...
info: malloc(128) return u8@7fa33a7a6310
info: realloc [*]align(16) u8@0 20800
info: malloc(20800) return u8@7fa33a488010
info: realloc return [*]align(16) u8@7fa33a488010 (from malloc)
info: malloc(32) return u8@7fa33a4900d0
info: malloc(2) return u8@7fa33a7a4610
info: malloc(32816) return u8@7fa33a47f010
info: malloc(11) return u8@7fa33a7a4630
info: malloc(15) return u8@7fa33a7a4650
info: malloc(13) return u8@7fa33a7a4670
info: free [*]align(16) u8@7fa33a47f010
info: free [*]align(16) u8@0
info: realloc [*]align(16) u8@0 72
info: malloc(72) return u8@7fa33a4a3e90
info: realloc return [*]align(16) u8@7fa33a4a3e90 (from malloc)
info: realloc [*]align(16) u8@0 144
info: malloc(144) return u8@7fa33a7a6410
info: realloc return [*]align(16) u8@7fa33a7a6410 (from malloc)
info: realloc [*]align(16) u8@0 168
info: malloc(168) return u8@7fa33a7a6510
info: realloc return [*]align(16) u8@7fa33a7a6510 (from malloc)
info: malloc(1024) return u8@7fa33a487010
libmalloc.so*  libmalloc.so.o  malloc.zig
info: free [*]align(16) u8@7fa33a7a4610
info: free [*]align(16) u8@0
info: free [*]align(16) u8@7fa33a4900d0
info: free [*]align(16) u8@7fa33a487010

I have removed most of the output.

malloc.zig

Build and run

Related