lib/stabilizer/src/code.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 const config = @import("config.zig");
4 const Marsaglia = @import("rng.zig").Marsaglia;
5
6 const Allocator = std.mem.Allocator;
7 const CodeConfig = config.CodeConfig;
8
9 pub const FunctionId = u32;
10 pub const LocationId = u32;
11
12 pub const FunctionOptions = struct {
13 code_size: usize,
14 table_size: usize = 0,
15 table_adjacent: bool = false,
16 stack_pad_unit: ?u8 = null,
17 };
18
19 pub const FunctionImageOptions = struct {
20 code: []const u8,
21 table: []const u8 = &.{},
22 table_adjacent: bool = false,
23 stack_pad_unit: ?u8 = null,
24 stack_pad_table_offsets: []const usize = &.{},
25 };
26
27 pub const FunctionRangeOptions = struct {
28 code_base: [*]const u8,
29 code_limit: [*]const u8,
30 table_base: ?[*]const u8 = null,
31 table_size: usize = 0,
32 table_adjacent: bool = false,
33 stack_pad: ?*const u8 = null,
34 };
35
36 pub const FunctionEntryState = enum {
37 trap,
38 forwarding,
39 };
40
41 pub const FunctionStackPad = struct {
42 unit: u8,
43 bytes: usize,
44 };
45
46 pub const FunctionLocation = struct {
47 id: LocationId,
48 function: FunctionId,
49 base: usize,
50 size: usize,
51 contents: []u8 = &.{},
52 defunct: bool = false,
53 marked: bool = false,
54 };
55
56 const FunctionImage = struct {
57 code: []u8,
58 table: []u8,
59 stack_pad_table_offsets: []usize,
60 stack_pad_address: ?usize,
61 table_adjacent: bool,
62
63 fn deinit(self: *FunctionImage, allocator: Allocator) void {
64 if (self.code.len > 0) allocator.free(self.code);
65 if (self.table.len > 0) allocator.free(self.table);
66 if (self.stack_pad_table_offsets.len > 0) allocator.free(self.stack_pad_table_offsets);
67 self.* = undefined;
68 }
69 };
70
71 const CodeFunction = struct {
72 id: FunctionId,
73 options: FunctionOptions,
74 image: ?FunctionImage = null,
75 entry: FunctionEntryState = .trap,
76 stack_pad_unit: ?u8 = null,
77 current: ?LocationId = null,
78 original_base: usize,
79 };
80
81 pub const CodeRandomizer = struct {
82 allocator: Allocator,
83 config: CodeConfig,
84 rng: Marsaglia,
85 functions: std.ArrayListUnmanaged(CodeFunction) = .empty,
86 locations: std.ArrayListUnmanaged(FunctionLocation) = .empty,
87 live: std.ArrayListUnmanaged(FunctionId) = .empty,
88 next_function_id: FunctionId = 1,
89 next_location_id: LocationId = 1,
90 next_original_base: usize = 0x0800_0000,
91 next_base: usize = 0x1000_0000,
92 next_stack_pad_address: usize = 0x3000_0000,
93 rerandomizing: bool = false,
94
95 pub fn init(allocator: Allocator, code_config: CodeConfig, seed: u64) CodeRandomizer {
96 return .{
97 .allocator = allocator,
98 .config = normalizeCodeConfig(code_config),
99 .rng = Marsaglia.init(seed),
100 };
101 }
102
103 pub fn deinit(self: *CodeRandomizer) void {
104 for (self.functions.items) |*function| {
105 if (function.image) |*image| image.deinit(self.allocator);
106 }
107 for (self.locations.items) |location| self.destroyLocation(location);
108 self.functions.deinit(self.allocator);
109 self.locations.deinit(self.allocator);
110 self.live.deinit(self.allocator);
111 self.* = undefined;
112 }
113
114 pub fn registerFunction(self: *CodeRandomizer, options: FunctionOptions) !FunctionId {
115 if (!self.config.enabled) return error.CodeRandomizationDisabled;
116 const normalized = normalizeFunctionOptions(options);
117 const id = self.next_function_id;
118 self.next_function_id += 1;
119 errdefer self.next_function_id -= 1;
120 const previous_original_base = self.next_original_base;
121 const original_base = self.allocateOriginalBase(normalized.code_size);
122 errdefer self.next_original_base = previous_original_base;
123 try self.functions.append(self.allocator, .{
124 .id = id,
125 .options = normalized,
126 .stack_pad_unit = normalized.stack_pad_unit,
127 .original_base = original_base,
128 });
129 return id;
130 }
131
132 pub fn registerFunctionImage(self: *CodeRandomizer, image_options: FunctionImageOptions) !FunctionId {
133 if (!self.config.enabled) return error.CodeRandomizationDisabled;
134 if (image_options.code.len == 0) return error.EmptyFunctionImage;
135 for (image_options.stack_pad_table_offsets) |offset| {
136 if (offset % @sizeOf(usize) != 0) return error.InvalidStackPadTableOffset;
137 if (offset > image_options.table.len or image_options.table.len - offset < @sizeOf(usize)) {
138 return error.InvalidStackPadTableOffset;
139 }
140 }
141
142 const code = try self.allocator.dupe(u8, image_options.code);
143 errdefer self.allocator.free(code);
144 const table = try self.allocator.dupe(u8, image_options.table);
145 errdefer self.allocator.free(table);
146 const stack_pad_table_offsets = try self.allocator.dupe(usize, image_options.stack_pad_table_offsets);
147 errdefer self.allocator.free(stack_pad_table_offsets);
148
149 const id = self.next_function_id;
150 self.next_function_id += 1;
151 errdefer self.next_function_id -= 1;
152 const previous_original_base = self.next_original_base;
153 const original_base = self.allocateOriginalBase(image_options.code.len);
154 errdefer self.next_original_base = previous_original_base;
155 try self.functions.append(self.allocator, .{
156 .id = id,
157 .options = .{
158 .code_size = image_options.code.len,
159 .table_size = image_options.table.len,
160 .table_adjacent = image_options.table_adjacent,
161 .stack_pad_unit = image_options.stack_pad_unit,
162 },
163 .image = .{
164 .code = code,
165 .table = table,
166 .stack_pad_table_offsets = stack_pad_table_offsets,
167 .stack_pad_address = null,
168 .table_adjacent = image_options.table_adjacent,
169 },
170 .stack_pad_unit = image_options.stack_pad_unit,
171 .original_base = original_base,
172 });
173 return id;
174 }
175
176 pub fn registerFunctionRange(self: *CodeRandomizer, range_options: FunctionRangeOptions) !FunctionId {
177 if (!self.config.enabled) return error.CodeRandomizationDisabled;
178 const code_start = @intFromPtr(range_options.code_base);
179 const code_end = @intFromPtr(range_options.code_limit);
180 if (code_end <= code_start) return error.InvalidFunctionRange;
181 const code_bytes = range_options.code_base[0 .. code_end - code_start];
182 const table_bytes = if (range_options.table_size == 0)
183 &.{}
184 else if (range_options.table_base) |table_base|
185 table_base[0..range_options.table_size]
186 else
187 return error.InvalidFunctionTableRange;
188
189 var stack_pad_table_offsets: std.ArrayListUnmanaged(usize) = .empty;
190 defer stack_pad_table_offsets.deinit(self.allocator);
191 const stack_pad_unit = if (range_options.stack_pad) |stack_pad| blk: {
192 try collectStackPadTableOffsets(
193 self.allocator,
194 table_bytes,
195 @intFromPtr(stack_pad),
196 &stack_pad_table_offsets,
197 );
198 break :blk stack_pad.*;
199 } else null;
200
201 return self.registerFunctionImage(.{
202 .code = code_bytes,
203 .table = table_bytes,
204 .table_adjacent = range_options.table_adjacent,
205 .stack_pad_unit = stack_pad_unit,
206 .stack_pad_table_offsets = stack_pad_table_offsets.items,
207 });
208 }
209
210 pub fn allocationSize(self: *const CodeRandomizer, id: FunctionId) ?usize {
211 const function = self.findFunctionConst(id) orelse return null;
212 return functionAllocationSize(function.options);
213 }
214
215 pub fn entryState(self: *const CodeRandomizer, id: FunctionId) ?FunctionEntryState {
216 const function = self.findFunctionConst(id) orelse return null;
217 return function.entry;
218 }
219
220 pub fn functionCount(self: *const CodeRandomizer) usize {
221 return self.functions.items.len;
222 }
223
224 pub fn stackPadUnit(self: *const CodeRandomizer, id: FunctionId) ?u8 {
225 const function = self.findFunctionConst(id) orelse return null;
226 return function.stack_pad_unit;
227 }
228
229 pub fn stackPad(self: *const CodeRandomizer, id: FunctionId) ?FunctionStackPad {
230 const unit = self.stackPadUnit(id) orelse return null;
231 return .{
232 .unit = unit,
233 .bytes = @as(usize, unit) * config.stack_alignment,
234 };
235 }
236
237 pub fn stackPadAddress(self: *const CodeRandomizer, id: FunctionId) ?usize {
238 const function = self.findFunctionConst(id) orelse return null;
239 const image = function.image orelse return null;
240 return image.stack_pad_address;
241 }
242
243 pub fn originalBase(self: *const CodeRandomizer, id: FunctionId) ?usize {
244 const function = self.findFunctionConst(id) orelse return null;
245 return function.original_base;
246 }
247
248 pub fn adjustAddress(self: *const CodeRandomizer, address: usize) usize {
249 for (self.locations.items) |location| {
250 if (address >= location.base and address < location.base + location.size) {
251 const function = self.findFunctionConst(location.function) orelse return address;
252 return function.original_base + (address - location.base);
253 }
254 }
255 return address;
256 }
257
258 pub fn enterFunction(self: *CodeRandomizer, id: FunctionId, roots: []const usize) !?*FunctionLocation {
259 const function = self.findFunction(id) orelse return error.UnknownFunction;
260 return switch (function.entry) {
261 .trap => try self.trap(id, roots),
262 .forwarding => self.currentLocation(id),
263 };
264 }
265
266 pub fn relocate(self: *CodeRandomizer, id: FunctionId) !*FunctionLocation {
267 if (!self.config.enabled) return error.CodeRandomizationDisabled;
268 const function = self.findFunction(id) orelse return error.UnknownFunction;
269 const old_current = function.current;
270 const allocation_size = functionAllocationSize(function.options);
271 var contents = try self.copyFunctionContents(function, allocation_size);
272 errdefer if (contents.len > 0) self.allocator.free(contents);
273 const base = self.allocateCodeBase(allocation_size);
274 const location_id = self.next_location_id;
275 self.next_location_id += 1;
276 var location_id_committed = false;
277 errdefer {
278 if (!location_id_committed) self.next_location_id -= 1;
279 }
280 try self.locations.append(self.allocator, .{
281 .id = location_id,
282 .function = id,
283 .base = base,
284 .size = allocation_size,
285 .contents = contents,
286 });
287 contents = &.{};
288 location_id_committed = true;
289 if (old_current) |current_id| {
290 if (self.findLocation(current_id)) |location| location.defunct = true;
291 }
292 function.current = location_id;
293 function.entry = .forwarding;
294 if (function.stack_pad_unit != null) {
295 function.stack_pad_unit = self.rng.nextByte();
296 }
297 try self.markLive(id);
298 return &self.locations.items[self.locations.items.len - 1];
299 }
300
301 pub fn beginRerandomization(self: *CodeRandomizer) void {
302 _ = self.beginRerandomizationAt(0);
303 }
304
305 pub fn trapRegisteredFunctions(self: *CodeRandomizer) void {
306 if (!self.config.enabled) return;
307 for (self.functions.items) |*function| {
308 function.entry = .trap;
309 }
310 self.live.clearRetainingCapacity();
311 self.rerandomizing = false;
312 }
313
314 pub fn beginRerandomizationAt(self: *CodeRandomizer, instruction_pointer: usize) usize {
315 var forwarded_instruction_pointer = instruction_pointer;
316 if (!self.config.enabled) return forwarded_instruction_pointer;
317 for (self.live.items) |id| {
318 if (self.findFunction(id)) |function| {
319 if (forwarded_instruction_pointer == function.original_base) {
320 if (function.current) |location_id| {
321 if (self.findLocation(location_id)) |location| {
322 forwarded_instruction_pointer = location.base;
323 }
324 }
325 }
326 function.entry = .trap;
327 }
328 }
329 self.live.clearRetainingCapacity();
330 self.rerandomizing = true;
331 return forwarded_instruction_pointer;
332 }
333
334 pub fn trap(self: *CodeRandomizer, id: FunctionId, roots: []const usize) !*FunctionLocation {
335 if (!self.config.enabled) return error.CodeRandomizationDisabled;
336 if (self.rerandomizing) {
337 for (roots) |root| self.markAddress(root);
338 self.sweep();
339 self.rerandomizing = false;
340 }
341 return self.relocate(id);
342 }
343
344 pub fn markAddress(self: *CodeRandomizer, address: usize) void {
345 if (!self.config.enabled) return;
346 for (self.locations.items) |*location| {
347 if (address >= location.base and address < location.base + location.size) {
348 location.marked = true;
349 }
350 }
351 }
352
353 pub fn sweep(self: *CodeRandomizer) void {
354 if (!self.config.enabled) return;
355 var index: usize = 0;
356 while (index < self.locations.items.len) {
357 const location = &self.locations.items[index];
358 if (location.defunct and !location.marked) {
359 self.destroyLocation(location.*);
360 _ = self.locations.swapRemove(index);
361 } else {
362 location.marked = false;
363 index += 1;
364 }
365 }
366 }
367
368 pub fn currentLocation(self: *CodeRandomizer, id: FunctionId) ?*FunctionLocation {
369 const function = self.findFunction(id) orelse return null;
370 const location_id = function.current orelse return null;
371 return self.findLocation(location_id);
372 }
373
374 pub fn locationCount(self: *const CodeRandomizer) usize {
375 return self.locations.items.len;
376 }
377
378 fn allocateCodeBase(self: *CodeRandomizer, size: usize) usize {
379 const jitter = self.rng.bounded(self.config.shuffle_slots) * self.config.alignment;
380 const base = std.mem.alignForward(usize, self.next_base + jitter, self.config.alignment);
381 self.next_base = base + std.mem.alignForward(usize, size, self.config.alignment) + self.config.alignment;
382 return base;
383 }
384
385 fn allocateOriginalBase(self: *CodeRandomizer, size: usize) usize {
386 const base = std.mem.alignForward(usize, self.next_original_base, self.config.alignment);
387 self.next_original_base = base + std.mem.alignForward(usize, size, self.config.alignment) + self.config.alignment;
388 return base;
389 }
390
391 fn copyFunctionContents(self: *CodeRandomizer, function: *CodeFunction, allocation_size: usize) ![]u8 {
392 const image = if (function.image) |*image| image else return &.{};
393 const contents = try self.allocator.alloc(u8, allocation_size);
394 errdefer self.allocator.free(contents);
395 @memset(contents, 0);
396
397 if (function.current) |current_id| {
398 if (self.findLocation(current_id)) |location| {
399 if (location.contents.len > 0) {
400 const copy_len = @min(contents.len, location.contents.len);
401 @memcpy(contents[0..copy_len], location.contents[0..copy_len]);
402 return contents;
403 }
404 }
405 }
406
407 self.patchStackPadTable(function, image);
408 @memcpy(contents[0..image.code.len], image.code);
409 if (image.table_adjacent and image.table.len > 0) {
410 @memcpy(contents[image.code.len..][0..image.table.len], image.table);
411 }
412 return contents;
413 }
414
415 fn destroyLocation(self: *CodeRandomizer, location: FunctionLocation) void {
416 if (location.contents.len > 0) self.allocator.free(location.contents);
417 }
418
419 fn patchStackPadTable(self: *CodeRandomizer, function: *CodeFunction, image: *FunctionImage) void {
420 if (function.stack_pad_unit == null or image.stack_pad_table_offsets.len == 0) return;
421 if (image.stack_pad_address == null) {
422 image.stack_pad_address = self.allocateStackPadAddress();
423 }
424 const address = image.stack_pad_address.?;
425 for (image.stack_pad_table_offsets) |offset| {
426 std.mem.writeInt(usize, image.table[offset..][0..@sizeOf(usize)], address, .little);
427 }
428 }
429
430 fn allocateStackPadAddress(self: *CodeRandomizer) usize {
431 const address = self.next_stack_pad_address;
432 self.next_stack_pad_address += @sizeOf(usize);
433 return address;
434 }
435
436 fn markLive(self: *CodeRandomizer, id: FunctionId) !void {
437 for (self.live.items) |live_id| if (live_id == id) return;
438 try self.live.append(self.allocator, id);
439 }
440
441 fn findFunction(self: *CodeRandomizer, id: FunctionId) ?*CodeFunction {
442 for (self.functions.items) |*function| if (function.id == id) return function;
443 return null;
444 }
445
446 fn findFunctionConst(self: *const CodeRandomizer, id: FunctionId) ?*const CodeFunction {
447 for (self.functions.items) |*function| if (function.id == id) return function;
448 return null;
449 }
450
451 fn findLocation(self: *CodeRandomizer, id: LocationId) ?*FunctionLocation {
452 for (self.locations.items) |*location| if (location.id == id) return location;
453 return null;
454 }
455 };
456
457 fn normalizeCodeConfig(code_config: CodeConfig) CodeConfig {
458 var out = code_config;
459 if (out.shuffle_slots == 0) out.shuffle_slots = 1;
460 if (out.alignment == 0) out.alignment = config.code_alignment;
461 out.alignment = std.math.ceilPowerOfTwoAssert(usize, out.alignment);
462 if (out.interval_ms == 0) out.interval_ms = config.rerandomize_interval_ms;
463 return out;
464 }
465
466 fn normalizeFunctionOptions(options: FunctionOptions) FunctionOptions {
467 var out = options;
468 if (out.code_size == 0) out.code_size = 1;
469 return out;
470 }
471
472 fn functionAllocationSize(options: FunctionOptions) usize {
473 return if (options.table_adjacent) options.code_size + options.table_size else options.code_size;
474 }
475
476 fn collectStackPadTableOffsets(
477 allocator: Allocator,
478 table: []const u8,
479 stack_pad_address: usize,
480 offsets: *std.ArrayListUnmanaged(usize),
481 ) !void {
482 var offset: usize = 0;
483 while (offset + @sizeOf(usize) <= table.len) : (offset += @sizeOf(usize)) {
484 const value = std.mem.readInt(usize, table[offset..][0..@sizeOf(usize)], .little);
485 if (value == stack_pad_address) try offsets.append(allocator, offset);
486 }
487 }
488
489 test "code randomizer registers upstream function ranges" {
490 var code = CodeRandomizer.init(std.testing.allocator, .{ .shuffle_slots = 4 }, 55);
491 defer code.deinit();
492
493 const image = [_]u8{ 0xaa, 0xbb, 0xcc };
494 var stack_pad: u8 = 7;
495 const stack_pad_address = @intFromPtr(&stack_pad);
496 var table = @as([(@sizeOf(usize) * 3)]u8, @splat(0));
497 std.mem.writeInt(usize, table[0..@sizeOf(usize)], 0x1111, .little);
498 std.mem.writeInt(usize, table[@sizeOf(usize)..][0..@sizeOf(usize)], stack_pad_address, .little);
499 std.mem.writeInt(usize, table[@sizeOf(usize) * 2 ..][0..@sizeOf(usize)], stack_pad_address, .little);
500
501 const function = try code.registerFunctionRange(.{
502 .code_base = image[0..].ptr,
503 .code_limit = image[0..].ptr + image.len,
504 .table_base = table[0..].ptr,
505 .table_size = table.len,
506 .table_adjacent = true,
507 .stack_pad = &stack_pad,
508 });
509 try std.testing.expectEqual(@as(u8, 7), code.stackPadUnit(function).?);
510
511 const location = try code.relocate(function);
512 const relocated_stack_pad_address = code.stackPadAddress(function).?;
513 try std.testing.expectEqualSlices(u8, &image, location.contents[0..image.len]);
514 try std.testing.expectEqual(@as(usize, 0x1111), std.mem.readInt(
515 usize,
516 location.contents[image.len..][0..@sizeOf(usize)],
517 .little,
518 ));
519 try std.testing.expectEqual(relocated_stack_pad_address, std.mem.readInt(
520 usize,
521 location.contents[image.len + @sizeOf(usize) ..][0..@sizeOf(usize)],
522 .little,
523 ));
524 try std.testing.expectEqual(relocated_stack_pad_address, std.mem.readInt(
525 usize,
526 location.contents[image.len + @sizeOf(usize) * 2 ..][0..@sizeOf(usize)],
527 .little,
528 ));
529 }
530
531 test "code randomizer rejects invalid function ranges" {
532 var code = CodeRandomizer.init(std.testing.allocator, .{ .shuffle_slots = 4 }, 55);
533 defer code.deinit();
534
535 const image = [_]u8{0xaa};
536 try std.testing.expectError(error.InvalidFunctionRange, code.registerFunctionRange(.{
537 .code_base = image[0..].ptr,
538 .code_limit = image[0..].ptr,
539 }));
540 try std.testing.expectError(error.InvalidFunctionTableRange, code.registerFunctionRange(.{
541 .code_base = image[0..].ptr,
542 .code_limit = image[0..].ptr + image.len,
543 .table_size = @sizeOf(usize),
544 }));
545 }