tiny.stabilizer.CodeRandomizer

1pub const CodeRandomizer = struct {2    allocator: Allocator,3    config: CodeConfig,4    rng: Marsaglia,5    functions: std.ArrayListUnmanaged(CodeFunction) = .empty,6    locations: std.ArrayListUnmanaged(FunctionLocation) = .empty,7    live: std.ArrayListUnmanaged(FunctionId) = .empty,8    next_function_id: FunctionId = 1,9    next_location_id: LocationId = 1,10    next_original_base: usize = 0x0800_0000,11    next_base: usize = 0x1000_0000,12    next_stack_pad_address: usize = 0x3000_0000,13    rerandomizing: bool = false,1415    pub fn init(allocator: Allocator, code_config: CodeConfig, seed: u64) CodeRandomizer {16        return .{17            .allocator = allocator,18            .config = normalizeCodeConfig(code_config),19            .rng = Marsaglia.init(seed),20        };21    }2223    pub fn deinit(self: *CodeRandomizer) void {24        for (self.functions.items) |*function| {25            if (function.image) |*image| image.deinit(self.allocator);26        }27        for (self.locations.items) |location| self.destroyLocation(location);28        self.functions.deinit(self.allocator);29        self.locations.deinit(self.allocator);30        self.live.deinit(self.allocator);31        self.* = undefined;32    }3334    pub fn registerFunction(self: *CodeRandomizer, options: FunctionOptions) !FunctionId {35        if (!self.config.enabled) return error.CodeRandomizationDisabled;36        const normalized = normalizeFunctionOptions(options);37        const id = self.next_function_id;38        self.next_function_id += 1;39        errdefer self.next_function_id -= 1;40        const previous_original_base = self.next_original_base;41        const original_base = self.allocateOriginalBase(normalized.code_size);42        errdefer self.next_original_base = previous_original_base;43        try self.functions.append(self.allocator, .{44            .id = id,45            .options = normalized,46            .stack_pad_unit = normalized.stack_pad_unit,47            .original_base = original_base,48        });49        return id;50    }5152    pub fn registerFunctionImage(self: *CodeRandomizer, image_options: FunctionImageOptions) !FunctionId {53        if (!self.config.enabled) return error.CodeRandomizationDisabled;54        if (image_options.code.len == 0) return error.EmptyFunctionImage;55        for (image_options.stack_pad_table_offsets) |offset| {56            if (offset % @sizeOf(usize) != 0) return error.InvalidStackPadTableOffset;57            if (offset > image_options.table.len or image_options.table.len - offset < @sizeOf(usize)) {58                return error.InvalidStackPadTableOffset;59            }60        }6162        const code = try self.allocator.dupe(u8, image_options.code);63        errdefer self.allocator.free(code);64        const table = try self.allocator.dupe(u8, image_options.table);65        errdefer self.allocator.free(table);66        const stack_pad_table_offsets = try self.allocator.dupe(usize, image_options.stack_pad_table_offsets);67        errdefer self.allocator.free(stack_pad_table_offsets);6869        const id = self.next_function_id;70        self.next_function_id += 1;71        errdefer self.next_function_id -= 1;72        const previous_original_base = self.next_original_base;73        const original_base = self.allocateOriginalBase(image_options.code.len);74        errdefer self.next_original_base = previous_original_base;75        try self.functions.append(self.allocator, .{76            .id = id,77            .options = .{78                .code_size = image_options.code.len,79                .table_size = image_options.table.len,80                .table_adjacent = image_options.table_adjacent,81                .stack_pad_unit = image_options.stack_pad_unit,82            },83            .image = .{84                .code = code,85                .table = table,86                .stack_pad_table_offsets = stack_pad_table_offsets,87                .stack_pad_address = null,88                .table_adjacent = image_options.table_adjacent,89            },90            .stack_pad_unit = image_options.stack_pad_unit,91            .original_base = original_base,92        });93        return id;94    }9596    pub fn registerFunctionRange(self: *CodeRandomizer, range_options: FunctionRangeOptions) !FunctionId {97        if (!self.config.enabled) return error.CodeRandomizationDisabled;98        const code_start = @intFromPtr(range_options.code_base);99        const code_end = @intFromPtr(range_options.code_limit);100        if (code_end <= code_start) return error.InvalidFunctionRange;101        const code_bytes = range_options.code_base[0 .. code_end - code_start];102        const table_bytes = if (range_options.table_size == 0)103            &.{}104        else if (range_options.table_base) |table_base|105            table_base[0..range_options.table_size]106        else107            return error.InvalidFunctionTableRange;108109        var stack_pad_table_offsets: std.ArrayListUnmanaged(usize) = .empty;110        defer stack_pad_table_offsets.deinit(self.allocator);111        const stack_pad_unit = if (range_options.stack_pad) |stack_pad| blk: {112            try collectStackPadTableOffsets(113                self.allocator,114                table_bytes,115                @intFromPtr(stack_pad),116                &stack_pad_table_offsets,117            );118            break :blk stack_pad.*;119        } else null;120121        return self.registerFunctionImage(.{122            .code = code_bytes,123            .table = table_bytes,124            .table_adjacent = range_options.table_adjacent,125            .stack_pad_unit = stack_pad_unit,126            .stack_pad_table_offsets = stack_pad_table_offsets.items,127        });128    }129130    pub fn allocationSize(self: *const CodeRandomizer, id: FunctionId) ?usize {131        const function = self.findFunctionConst(id) orelse return null;132        return functionAllocationSize(function.options);133    }134135    pub fn entryState(self: *const CodeRandomizer, id: FunctionId) ?FunctionEntryState {136        const function = self.findFunctionConst(id) orelse return null;137        return function.entry;138    }139140    pub fn functionCount(self: *const CodeRandomizer) usize {141        return self.functions.items.len;142    }143144    pub fn stackPadUnit(self: *const CodeRandomizer, id: FunctionId) ?u8 {145        const function = self.findFunctionConst(id) orelse return null;146        return function.stack_pad_unit;147    }148149    pub fn stackPad(self: *const CodeRandomizer, id: FunctionId) ?FunctionStackPad {150        const unit = self.stackPadUnit(id) orelse return null;151        return .{152            .unit = unit,153            .bytes = @as(usize, unit) * config.stack_alignment,154        };155    }156157    pub fn stackPadAddress(self: *const CodeRandomizer, id: FunctionId) ?usize {158        const function = self.findFunctionConst(id) orelse return null;159        const image = function.image orelse return null;160        return image.stack_pad_address;161    }162163    pub fn originalBase(self: *const CodeRandomizer, id: FunctionId) ?usize {164        const function = self.findFunctionConst(id) orelse return null;165        return function.original_base;166    }167168    pub fn adjustAddress(self: *const CodeRandomizer, address: usize) usize {169        for (self.locations.items) |location| {170            if (address >= location.base and address < location.base + location.size) {171                const function = self.findFunctionConst(location.function) orelse return address;172                return function.original_base + (address - location.base);173            }174        }175        return address;176    }177178    pub fn enterFunction(self: *CodeRandomizer, id: FunctionId, roots: []const usize) !?*FunctionLocation {179        const function = self.findFunction(id) orelse return error.UnknownFunction;180        return switch (function.entry) {181            .trap => try self.trap(id, roots),182            .forwarding => self.currentLocation(id),183        };184    }185186    pub fn relocate(self: *CodeRandomizer, id: FunctionId) !*FunctionLocation {187        if (!self.config.enabled) return error.CodeRandomizationDisabled;188        const function = self.findFunction(id) orelse return error.UnknownFunction;189        const old_current = function.current;190        const allocation_size = functionAllocationSize(function.options);191        var contents = try self.copyFunctionContents(function, allocation_size);192        errdefer if (contents.len > 0) self.allocator.free(contents);193        const base = self.allocateCodeBase(allocation_size);194        const location_id = self.next_location_id;195        self.next_location_id += 1;196        var location_id_committed = false;197        errdefer {198            if (!location_id_committed) self.next_location_id -= 1;199        }200        try self.locations.append(self.allocator, .{201            .id = location_id,202            .function = id,203            .base = base,204            .size = allocation_size,205            .contents = contents,206        });207        contents = &.{};208        location_id_committed = true;209        if (old_current) |current_id| {210            if (self.findLocation(current_id)) |location| location.defunct = true;211        }212        function.current = location_id;213        function.entry = .forwarding;214        if (function.stack_pad_unit != null) {215            function.stack_pad_unit = self.rng.nextByte();216        }217        try self.markLive(id);218        return &self.locations.items[self.locations.items.len - 1];219    }220221    pub fn beginRerandomization(self: *CodeRandomizer) void {222        _ = self.beginRerandomizationAt(0);223    }224225    pub fn trapRegisteredFunctions(self: *CodeRandomizer) void {226        if (!self.config.enabled) return;227        for (self.functions.items) |*function| {228            function.entry = .trap;229        }230        self.live.clearRetainingCapacity();231        self.rerandomizing = false;232    }233234    pub fn beginRerandomizationAt(self: *CodeRandomizer, instruction_pointer: usize) usize {235        var forwarded_instruction_pointer = instruction_pointer;236        if (!self.config.enabled) return forwarded_instruction_pointer;237        for (self.live.items) |id| {238            if (self.findFunction(id)) |function| {239                if (forwarded_instruction_pointer == function.original_base) {240                    if (function.current) |location_id| {241                        if (self.findLocation(location_id)) |location| {242                            forwarded_instruction_pointer = location.base;243                        }244                    }245                }246                function.entry = .trap;247            }248        }249        self.live.clearRetainingCapacity();250        self.rerandomizing = true;251        return forwarded_instruction_pointer;252    }253254    pub fn trap(self: *CodeRandomizer, id: FunctionId, roots: []const usize) !*FunctionLocation {255        if (!self.config.enabled) return error.CodeRandomizationDisabled;256        if (self.rerandomizing) {257            for (roots) |root| self.markAddress(root);258            self.sweep();259            self.rerandomizing = false;260        }261        return self.relocate(id);262    }263264    pub fn markAddress(self: *CodeRandomizer, address: usize) void {265        if (!self.config.enabled) return;266        for (self.locations.items) |*location| {267            if (address >= location.base and address < location.base + location.size) {268                location.marked = true;269            }270        }271    }272273    pub fn sweep(self: *CodeRandomizer) void {274        if (!self.config.enabled) return;275        var index: usize = 0;276        while (index < self.locations.items.len) {277            const location = &self.locations.items[index];278            if (location.defunct and !location.marked) {279                self.destroyLocation(location.*);280                _ = self.locations.swapRemove(index);281            } else {282                location.marked = false;283                index += 1;284            }285        }286    }287288    pub fn currentLocation(self: *CodeRandomizer, id: FunctionId) ?*FunctionLocation {289        const function = self.findFunction(id) orelse return null;290        const location_id = function.current orelse return null;291        return self.findLocation(location_id);292    }293294    pub fn locationCount(self: *const CodeRandomizer) usize {295        return self.locations.items.len;296    }297298    fn allocateCodeBase(self: *CodeRandomizer, size: usize) usize {299        const jitter = self.rng.bounded(self.config.shuffle_slots) * self.config.alignment;300        const base = std.mem.alignForward(usize, self.next_base + jitter, self.config.alignment);301        self.next_base = base + std.mem.alignForward(usize, size, self.config.alignment) + self.config.alignment;302        return base;303    }304305    fn allocateOriginalBase(self: *CodeRandomizer, size: usize) usize {306        const base = std.mem.alignForward(usize, self.next_original_base, self.config.alignment);307        self.next_original_base = base + std.mem.alignForward(usize, size, self.config.alignment) + self.config.alignment;308        return base;309    }310311    fn copyFunctionContents(self: *CodeRandomizer, function: *CodeFunction, allocation_size: usize) ![]u8 {312        const image = if (function.image) |*image| image else return &.{};313        const contents = try self.allocator.alloc(u8, allocation_size);314        errdefer self.allocator.free(contents);315        @memset(contents, 0);316317        if (function.current) |current_id| {318            if (self.findLocation(current_id)) |location| {319                if (location.contents.len > 0) {320                    const copy_len = @min(contents.len, location.contents.len);321                    @memcpy(contents[0..copy_len], location.contents[0..copy_len]);322                    return contents;323                }324            }325        }326327        self.patchStackPadTable(function, image);328        @memcpy(contents[0..image.code.len], image.code);329        if (image.table_adjacent and image.table.len > 0) {330            @memcpy(contents[image.code.len..][0..image.table.len], image.table);331        }332        return contents;333    }334335    fn destroyLocation(self: *CodeRandomizer, location: FunctionLocation) void {336        if (location.contents.len > 0) self.allocator.free(location.contents);337    }338339    fn patchStackPadTable(self: *CodeRandomizer, function: *CodeFunction, image: *FunctionImage) void {340        if (function.stack_pad_unit == null or image.stack_pad_table_offsets.len == 0) return;341        if (image.stack_pad_address == null) {342            image.stack_pad_address = self.allocateStackPadAddress();343        }344        const address = image.stack_pad_address.?;345        for (image.stack_pad_table_offsets) |offset| {346            std.mem.writeInt(usize, image.table[offset..][0..@sizeOf(usize)], address, .little);347        }348    }349350    fn allocateStackPadAddress(self: *CodeRandomizer) usize {351        const address = self.next_stack_pad_address;352        self.next_stack_pad_address += @sizeOf(usize);353        return address;354    }355356    fn markLive(self: *CodeRandomizer, id: FunctionId) !void {357        for (self.live.items) |live_id| if (live_id == id) return;358        try self.live.append(self.allocator, id);359    }360361    fn findFunction(self: *CodeRandomizer, id: FunctionId) ?*CodeFunction {362        for (self.functions.items) |*function| if (function.id == id) return function;363        return null;364    }365366    fn findFunctionConst(self: *const CodeRandomizer, id: FunctionId) ?*const CodeFunction {367        for (self.functions.items) |*function| if (function.id == id) return function;368        return null;369    }370371    fn findLocation(self: *CodeRandomizer, id: LocationId) ?*FunctionLocation {372        for (self.locations.items) |*location| if (location.id == id) return location;373        return null;374    }375};

1pub const CodeRandomizer = code_mod.CodeRandomizer;