lib/stabilizer/src/runtime.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 const config = @import("config.zig");
4 const code_mod = @import("code.zig");
5 const CodeRandomizer = code_mod.CodeRandomizer;
6 const ShuffleAllocator = @import("heap.zig").ShuffleAllocator;
7 const Marsaglia = @import("rng.zig").Marsaglia;
8 const stack = @import("stack.zig");
9
10 const Allocator = std.mem.Allocator;
11
12 pub const ConstructorFn = *const fn () callconv(.c) void;
13
14 pub const Runtime = struct {
15 allocator: Allocator,
16 config: config.Config,
17 heap: ShuffleAllocator,
18 stack: stack.StackPads,
19 code: CodeRandomizer,
20 rng: Marsaglia,
21 stack_pad_slots: std.ArrayListUnmanaged(*u8) = .empty,
22 constructors: std.ArrayListUnmanaged(ConstructorFn) = .empty,
23
24 pub fn init(allocator: Allocator, runtime_config: config.Config) !Runtime {
25 var rng = Marsaglia.init(runtime_config.seed);
26 var stack_pads = try stack.StackPads.init(allocator, runtime_config.stack, &rng);
27 errdefer stack_pads.deinit();
28 return .{
29 .allocator = allocator,
30 .config = runtime_config,
31 .heap = ShuffleAllocator.init(allocator, runtime_config.heap, runtime_config.seed ^ 0x4845_4150),
32 .stack = stack_pads,
33 .code = CodeRandomizer.init(allocator, runtime_config.code, runtime_config.seed ^ 0x434f_4445),
34 .rng = rng,
35 };
36 }
37
38 pub fn deinit(self: *Runtime) void {
39 self.constructors.deinit(self.allocator);
40 self.stack_pad_slots.deinit(self.allocator);
41 self.code.deinit();
42 self.stack.deinit();
43 self.heap.deinit();
44 self.* = undefined;
45 }
46
47 pub fn randomizedAllocator(self: *Runtime) Allocator {
48 return self.heap.allocator();
49 }
50
51 pub fn malloc(self: *Runtime, len: usize) ![*]u8 {
52 return self.heap.malloc(len);
53 }
54
55 pub fn calloc(self: *Runtime, count: usize, len: usize) ![*]u8 {
56 return self.heap.calloc(count, len);
57 }
58
59 pub fn realloc(self: *Runtime, ptr: ?[*]u8, new_len: usize) !?[*]u8 {
60 return self.heap.realloc(ptr, new_len);
61 }
62
63 pub fn free(self: *Runtime, ptr: ?[*]u8) void {
64 self.heap.freePointer(ptr);
65 }
66
67 pub fn freeIfOwned(self: *Runtime, ptr: ?[*]u8) bool {
68 return self.heap.freePointerIfOwned(ptr);
69 }
70
71 pub fn ownsPointer(self: *const Runtime, ptr: ?[*]u8) bool {
72 return self.heap.ownsPointer(ptr);
73 }
74
75 pub fn registerStackPad(self: *Runtime, slot: *u8) !void {
76 for (self.stack_pad_slots.items) |existing| {
77 if (existing == slot) return;
78 }
79 try self.stack_pad_slots.append(self.allocator, slot);
80 }
81
82 pub fn registeredStackPadCount(self: *const Runtime) usize {
83 return self.stack_pad_slots.items.len;
84 }
85
86 pub fn registerConstructor(self: *Runtime, constructor: ConstructorFn) !void {
87 try self.constructors.append(self.allocator, constructor);
88 }
89
90 pub fn constructorCount(self: *const Runtime) usize {
91 return self.constructors.items.len;
92 }
93
94 pub fn runConstructors(self: *Runtime) void {
95 for (self.constructors.items) |constructor| constructor();
96 }
97
98 pub fn registerFunctionRange(self: *Runtime, options: code_mod.FunctionRangeOptions) !code_mod.FunctionId {
99 return self.code.registerFunctionRange(options);
100 }
101
102 pub fn enterFunction(self: *Runtime, id: code_mod.FunctionId, roots: []const usize) !?usize {
103 const location = try self.code.enterFunction(id, roots);
104 return if (location) |active| active.base else null;
105 }
106
107 pub fn trapFunction(self: *Runtime, id: code_mod.FunctionId, roots: []const usize) !usize {
108 const location = try self.code.trap(id, roots);
109 return location.base;
110 }
111
112 pub fn adjustAddress(self: *const Runtime, address: usize) usize {
113 return self.code.adjustAddress(address);
114 }
115
116 pub fn adjustBacktrace(self: *const Runtime, adjusted: []usize, real: []const usize) []usize {
117 const count = @min(adjusted.len, real.len);
118 for (real[0..count], adjusted[0..count]) |address, *out| {
119 out.* = self.adjustAddress(address);
120 }
121 return adjusted[0..count];
122 }
123
124 pub fn runProgram(self: *Runtime, main: anytype, args: anytype) @TypeOf(@call(.auto, main, args)) {
125 self.code.trapRegisteredFunctions();
126 self.runConstructors();
127 return @call(.auto, main, args);
128 }
129
130 pub fn nextStackPad(self: *Runtime) stack.StackPad {
131 return self.stack.next(&self.rng);
132 }
133
134 pub fn rerandomize(self: *Runtime) void {
135 _ = self.rerandomizeAt(0);
136 }
137
138 pub fn rerandomizeAt(self: *Runtime, instruction_pointer: usize) usize {
139 if (self.config.code.enabled) {
140 if (self.code.functionCount() == 0) self.rerandomizeStackPads();
141 return self.code.beginRerandomizationAt(instruction_pointer);
142 } else {
143 self.rerandomizeStackPads();
144 return instruction_pointer;
145 }
146 }
147
148 fn rerandomizeStackPads(self: *Runtime) void {
149 if (!self.config.stack.enabled) return;
150 self.stack.refill(&self.rng);
151 for (self.stack_pad_slots.items) |slot| {
152 slot.* = self.rng.nextByte();
153 }
154 }
155 };
156
157 test "runtime rerandomizes registered stack pad slots" {
158 const seed = 9;
159 var expected_rng = Marsaglia.init(seed);
160 _ = expected_rng.nextByte();
161
162 var runtime = try Runtime.init(std.testing.allocator, .{
163 .seed = seed,
164 .stack = .{ .entries = 1 },
165 .code = .{ .enabled = false },
166 });
167 defer runtime.deinit();
168
169 var first: u8 = 0;
170 var second: u8 = 0;
171 try runtime.registerStackPad(&first);
172 try runtime.registerStackPad(&second);
173 try runtime.registerStackPad(&first);
174 try std.testing.expectEqual(@as(usize, 2), runtime.registeredStackPadCount());
175
176 const expected_stack = expected_rng.nextByte();
177 const expected_first = expected_rng.nextByte();
178 const expected_second = expected_rng.nextByte();
179
180 runtime.rerandomize();
181
182 try std.testing.expectEqual(expected_stack, runtime.nextStackPad().unit);
183 try std.testing.expectEqual(expected_first, first);
184 try std.testing.expectEqual(expected_second, second);
185 }
186
187 test "runtime rerandomize forwards interrupts at live function entries" {
188 var runtime = try Runtime.init(std.testing.allocator, .{
189 .seed = 9,
190 .code = .{ .enabled = true },
191 });
192 defer runtime.deinit();
193
194 const active = try runtime.code.registerFunction(.{ .code_size = 64 });
195 const inactive = try runtime.code.registerFunction(.{ .code_size = 64 });
196 const active_base = (try runtime.enterFunction(active, &.{})).?;
197 const active_original = runtime.code.originalBase(active).?;
198 const inactive_original = runtime.code.originalBase(inactive).?;
199
200 try std.testing.expectEqual(active_base, runtime.rerandomizeAt(active_original));
201 try std.testing.expectEqual(code_mod.FunctionEntryState.trap, runtime.code.entryState(active).?);
202 try std.testing.expectEqual(code_mod.FunctionEntryState.trap, runtime.code.entryState(inactive).?);
203 try std.testing.expect(runtime.code.rerandomizing);
204 try std.testing.expectEqual(@as(usize, 0), runtime.code.live.items.len);
205 try std.testing.expectEqual(active_original + 1, runtime.rerandomizeAt(active_original + 1));
206 try std.testing.expectEqual(inactive_original, runtime.rerandomizeAt(inactive_original));
207 }
208
209 test "runtime rerandomizeAt preserves instruction pointer in stack-only mode" {
210 const seed = 9;
211 var expected_rng = Marsaglia.init(seed);
212 _ = expected_rng.nextByte();
213
214 var runtime = try Runtime.init(std.testing.allocator, .{
215 .seed = seed,
216 .stack = .{ .entries = 1 },
217 .code = .{ .enabled = false },
218 });
219 defer runtime.deinit();
220
221 const expected_stack = expected_rng.nextByte();
222 try std.testing.expectEqual(@as(usize, 0x1234), runtime.rerandomizeAt(0x1234));
223 try std.testing.expectEqual(expected_stack, runtime.nextStackPad().unit);
224 }
225
226 test "runtime enterFunction relocates traps and preserves rooted code locations" {
227 var runtime = try Runtime.init(std.testing.allocator, .{
228 .code = .{ .enabled = true },
229 });
230 defer runtime.deinit();
231
232 const first = try runtime.code.registerFunction(.{ .code_size = 64 });
233 const second = try runtime.code.registerFunction(.{ .code_size = 64 });
234 const second_original = runtime.code.originalBase(second).?;
235
236 const first_base = (try runtime.enterFunction(first, &.{})).?;
237 const second_base = (try runtime.enterFunction(second, &.{})).?;
238
239 runtime.rerandomize();
240 const second_next = (try runtime.enterFunction(second, &.{first_base})).?;
241 try std.testing.expect(second_next != second_base);
242 try std.testing.expectEqual(second_original, runtime.adjustAddress(second_base));
243
244 runtime.rerandomize();
245 _ = try runtime.enterFunction(first, &.{second_base});
246 try std.testing.expectEqual(second_original, runtime.adjustAddress(second_base));
247
248 runtime.rerandomize();
249 _ = try runtime.enterFunction(first, &.{});
250 try std.testing.expectEqual(second_base, runtime.adjustAddress(second_base));
251 }
252
253 test "runtime adjusts relocated addresses back to original code" {
254 var runtime = try Runtime.init(std.testing.allocator, .{
255 .code = .{ .shuffle_slots = 4 },
256 });
257 defer runtime.deinit();
258
259 const function = try runtime.code.registerFunction(.{
260 .code_size = 64,
261 .table_size = 16,
262 .table_adjacent = true,
263 });
264 const original = runtime.code.originalBase(function).?;
265 const first = try runtime.code.relocate(function);
266 const first_base = first.base;
267 const first_size = first.size;
268 try std.testing.expectEqual(original + 5, runtime.adjustAddress(first_base + 5));
269 try std.testing.expectEqual(original + 70, runtime.adjustAddress(first_base + 70));
270
271 const unrelated = first_base + first_size + 1024;
272 try std.testing.expectEqual(unrelated, runtime.adjustAddress(unrelated));
273
274 const second = try runtime.code.relocate(function);
275 const second_base = second.base;
276 try std.testing.expectEqual(original + 9, runtime.adjustAddress(first_base + 9));
277 try std.testing.expectEqual(original + 9, runtime.adjustAddress(second_base + 9));
278
279 const real_backtrace = [_]usize{ first_base + 1, unrelated, second_base + 2 };
280 var adjusted_backtrace: [real_backtrace.len + 1]usize = undefined;
281 const adjusted = runtime.adjustBacktrace(adjusted_backtrace[0..], real_backtrace[0..]);
282 try std.testing.expectEqual(@as(usize, real_backtrace.len), adjusted.len);
283 try std.testing.expectEqual(original + 1, adjusted[0]);
284 try std.testing.expectEqual(unrelated, adjusted[1]);
285 try std.testing.expectEqual(original + 2, adjusted[2]);
286
287 const truncated = runtime.adjustBacktrace(adjusted_backtrace[0..1], real_backtrace[0..]);
288 try std.testing.expectEqual(@as(usize, 1), truncated.len);
289 try std.testing.expectEqual(original + 1, truncated[0]);
290
291 runtime.code.sweep();
292 try std.testing.expectEqual(first_base + 9, runtime.adjustAddress(first_base + 9));
293 try std.testing.expectEqual(original + 9, runtime.adjustAddress(second_base + 9));
294 }
295
296 var constructor_trace: [64]u8 = undefined;
297 var constructor_trace_len: usize = 0;
298
299 fn resetConstructorTrace() void {
300 @memset(&constructor_trace, 0);
301 constructor_trace_len = 0;
302 }
303
304 fn appendConstructorTrace(value: u8) void {
305 constructor_trace[constructor_trace_len] = value;
306 constructor_trace_len += 1;
307 }
308
309 fn traceConstructor1() callconv(.c) void {
310 appendConstructorTrace(1);
311 }
312
313 fn traceConstructor2() callconv(.c) void {
314 appendConstructorTrace(2);
315 }
316
317 fn traceConstructor3() callconv(.c) void {
318 appendConstructorTrace(3);
319 }
320
321 fn traceProgramMain() u8 {
322 appendConstructorTrace(9);
323 return 42;
324 }
325
326 test "runtime runs registered constructors in insertion order" {
327 resetConstructorTrace();
328
329 var runtime = try Runtime.init(std.testing.allocator, .{});
330 defer runtime.deinit();
331
332 try runtime.registerConstructor(traceConstructor2);
333 try runtime.registerConstructor(traceConstructor1);
334 try runtime.registerConstructor(traceConstructor2);
335 try std.testing.expectEqual(@as(usize, 3), runtime.constructorCount());
336
337 runtime.runConstructors();
338
339 try std.testing.expectEqual(@as(usize, 3), constructor_trace_len);
340 try std.testing.expectEqualSlices(u8, &.{ 2, 1, 2 }, constructor_trace[0..constructor_trace_len]);
341 }
342
343 test "runtime arms code and runs constructors before program main" {
344 resetConstructorTrace();
345
346 var runtime = try Runtime.init(std.testing.allocator, .{
347 .code = .{ .enabled = true },
348 });
349 defer runtime.deinit();
350
351 const function = try runtime.code.registerFunction(.{ .code_size = 64 });
352 _ = try runtime.enterFunction(function, &.{});
353 try std.testing.expectEqual(code_mod.FunctionEntryState.forwarding, runtime.code.entryState(function).?);
354
355 try runtime.registerConstructor(traceConstructor1);
356 try runtime.registerConstructor(traceConstructor3);
357 const result = runtime.runProgram(traceProgramMain, .{});
358
359 try std.testing.expectEqual(@as(u8, 42), result);
360 try std.testing.expectEqual(code_mod.FunctionEntryState.trap, runtime.code.entryState(function).?);
361 try std.testing.expectEqual(@as(usize, 0), runtime.code.live.items.len);
362 try std.testing.expect(!runtime.code.rerandomizing);
363 try std.testing.expectEqualSlices(u8, &.{ 1, 3, 9 }, constructor_trace[0..constructor_trace_len]);
364 }