| ... | @@ -0,0 +1,436 @@ |
| 1 | //! Represents export trie used in MachO executables and dynamic libraries. |
| 2 | //! The purpose of an export trie is to encode as compactly as possible all |
| 3 | //! export symbols for the loader `dyld`. |
| 4 | //! The export trie encodes offset and other information using ULEB128 |
| 5 | //! encoding, and is part of the __LINKEDIT segment. |
| 6 | //! |
| 7 | //! Description from loader.h: |
| 8 | //! |
| 9 | //! The symbols exported by a dylib are encoded in a trie. This is a compact |
| 10 | //! representation that factors out common prefixes. It also reduces LINKEDIT pages |
| 11 | //! in RAM because it encodes all information (name, address, flags) in one small, |
| 12 | //! contiguous range. The export area is a stream of nodes. The first node sequentially |
| 13 | //! is the start node for the trie. |
| 14 | //! |
| 15 | //! Nodes for a symbol start with a uleb128 that is the length of the exported symbol |
| 16 | //! information for the string so far. If there is no exported symbol, the node starts |
| 17 | //! with a zero byte. If there is exported info, it follows the length. |
| 18 | //! |
| 19 | //! First is a uleb128 containing flags. Normally, it is followed by a uleb128 encoded |
| 20 | //! offset which is location of the content named by the symbol from the mach_header |
| 21 | //! for the image. If the flags is EXPORT_SYMBOL_FLAGS_REEXPORT, then following the flags |
| 22 | //! is a uleb128 encoded library ordinal, then a zero terminated UTF8 string. If the string |
| 23 | //! is zero length, then the symbol is re-export from the specified dylib with the same name. |
| 24 | //! If the flags is EXPORT_SYMBOL_FLAGS_STUB_AND_RESOLVER, then following the flags is two |
| 25 | //! uleb128s: the stub offset and the resolver offset. The stub is used by non-lazy pointers. |
| 26 | //! The resolver is used by lazy pointers and must be called to get the actual address to use. |
| 27 | //! |
| 28 | //! After the optional exported symbol information is a byte of how many edges (0-255) that |
| 29 | //! this node has leaving it, followed by each edge. Each edge is a zero terminated UTF8 of |
| 30 | //! the addition chars in the symbol, followed by a uleb128 offset for the node that edge points to. |
| 31 | const Trie = @This(); |
| 32 | |
| 33 | const std = @import("std"); |
| 34 | const mem = std.mem; |
| 35 | const leb = std.debug.leb; |
| 36 | const log = std.log.scoped(.link); |
| 37 | const testing = std.testing; |
| 38 | const assert = std.debug.assert; |
| 39 | const Allocator = mem.Allocator; |
| 40 | |
| 41 | pub const Symbol = struct { |
| 42 | name: []const u8, |
| 43 | vmaddr_offset: u64, |
| 44 | export_flags: u64, |
| 45 | }; |
| 46 | |
| 47 | const Edge = struct { |
| 48 | from: *Node, |
| 49 | to: *Node, |
| 50 | label: []const u8, |
| 51 | |
| 52 | fn deinit(self: *Edge, alloc: *Allocator) void { |
| 53 | self.to.deinit(alloc); |
| 54 | alloc.destroy(self.to); |
| 55 | self.from = undefined; |
| 56 | self.to = undefined; |
| 57 | } |
| 58 | }; |
| 59 | |
| 60 | const Node = struct { |
| 61 | /// Export flags associated with this exported symbol (if any). |
| 62 | export_flags: ?u64 = null, |
| 63 | /// VM address offset wrt to the section this symbol is defined against (if any). |
| 64 | vmaddr_offset: ?u64 = null, |
| 65 | /// Offset of this node in the trie output byte stream. |
| 66 | trie_offset: ?usize = null, |
| 67 | /// List of all edges originating from this node. |
| 68 | edges: std.ArrayListUnmanaged(Edge) = .{}, |
| 69 | |
| 70 | fn deinit(self: *Node, alloc: *Allocator) void { |
| 71 | for (self.edges.items) |*edge| { |
| 72 | edge.deinit(alloc); |
| 73 | } |
| 74 | self.edges.deinit(alloc); |
| 75 | } |
| 76 | |
| 77 | const PutResult = struct { |
| 78 | /// Node reached at this stage of `put` op. |
| 79 | node: *Node, |
| 80 | /// Count of newly inserted nodes at this stage of `put` op. |
| 81 | node_count: usize, |
| 82 | }; |
| 83 | |
| 84 | /// Inserts a new node starting from `self`. |
| 85 | fn put(self: *Node, alloc: *Allocator, label: []const u8, node_count: usize) !PutResult { |
| 86 | var curr_node_count = node_count; |
| 87 | // Check for match with edges from this node. |
| 88 | for (self.edges.items) |*edge| { |
| 89 | const match = mem.indexOfDiff(u8, edge.label, label) orelse return PutResult{ |
| 90 | .node = edge.to, |
| 91 | .node_count = curr_node_count, |
| 92 | }; |
| 93 | if (match == 0) continue; |
| 94 | if (match == edge.label.len) return edge.to.put(alloc, label[match..], curr_node_count); |
| 95 | |
| 96 | // Found a match, need to splice up nodes. |
| 97 | // From: A -> B |
| 98 | // To: A -> C -> B |
| 99 | const mid = try alloc.create(Node); |
| 100 | mid.* = .{}; |
| 101 | const to_label = edge.label; |
| 102 | const to_node = edge.to; |
| 103 | edge.to = mid; |
| 104 | edge.label = label[0..match]; |
| 105 | curr_node_count += 1; |
| 106 | |
| 107 | try mid.edges.append(alloc, .{ |
| 108 | .from = mid, |
| 109 | .to = to_node, |
| 110 | .label = to_label[match..], |
| 111 | }); |
| 112 | |
| 113 | if (match == label.len) { |
| 114 | return PutResult{ .node = to_node, .node_count = curr_node_count }; |
| 115 | } else { |
| 116 | return mid.put(alloc, label[match..], curr_node_count); |
| 117 | } |
| 118 | } |
| 119 | |
| 120 | // Add a new node. |
| 121 | const node = try alloc.create(Node); |
| 122 | node.* = .{}; |
| 123 | curr_node_count += 1; |
| 124 | |
| 125 | try self.edges.append(alloc, .{ |
| 126 | .from = self, |
| 127 | .to = node, |
| 128 | .label = label, |
| 129 | }); |
| 130 | |
| 131 | return PutResult{ .node = node, .node_count = curr_node_count }; |
| 132 | } |
| 133 | |
| 134 | /// This method should only be called *after* updateOffset has been called! |
| 135 | /// In case this is not upheld, this method will panic. |
| 136 | fn writeULEB128Mem(self: Node, buffer: *std.ArrayListUnmanaged(u8)) !void { |
| 137 | assert(self.trie_offset != null); // You need to call updateOffset first. |
| 138 | if (self.vmaddr_offset) |offset| { |
| 139 | // Terminal node info: encode export flags and vmaddr offset of this symbol. |
| 140 | var info_buf_len: usize = 0; |
| 141 | var info_buf: [@sizeOf(u64) * 2]u8 = undefined; |
| 142 | info_buf_len += try leb.writeULEB128Mem(info_buf[0..], self.export_flags.?); |
| 143 | info_buf_len += try leb.writeULEB128Mem(info_buf[info_buf_len..], offset); |
| 144 | |
| 145 | // Encode the size of the terminal node info. |
| 146 | var size_buf: [@sizeOf(u64)]u8 = undefined; |
| 147 | const size_buf_len = try leb.writeULEB128Mem(size_buf[0..], info_buf_len); |
| 148 | |
| 149 | // Now, write them to the output buffer. |
| 150 | buffer.appendSliceAssumeCapacity(size_buf[0..size_buf_len]); |
| 151 | buffer.appendSliceAssumeCapacity(info_buf[0..info_buf_len]); |
| 152 | } else { |
| 153 | // Non-terminal node is delimited by 0 byte. |
| 154 | buffer.appendAssumeCapacity(0); |
| 155 | } |
| 156 | // Write number of edges (max legal number of edges is 256). |
| 157 | buffer.appendAssumeCapacity(@intCast(u8, self.edges.items.len)); |
| 158 | |
| 159 | for (self.edges.items) |edge| { |
| 160 | // Write edges labels. |
| 161 | buffer.appendSliceAssumeCapacity(edge.label); |
| 162 | buffer.appendAssumeCapacity(0); |
| 163 | |
| 164 | var buf: [@sizeOf(u64)]u8 = undefined; |
| 165 | const buf_len = try leb.writeULEB128Mem(buf[0..], edge.to.trie_offset.?); |
| 166 | buffer.appendSliceAssumeCapacity(buf[0..buf_len]); |
| 167 | } |
| 168 | } |
| 169 | |
| 170 | const UpdateResult = struct { |
| 171 | /// Current size of this node in bytes. |
| 172 | node_size: usize, |
| 173 | /// True if the trie offset of this node in the output byte stream |
| 174 | /// would need updating; false otherwise. |
| 175 | updated: bool, |
| 176 | }; |
| 177 | |
| 178 | /// Updates offset of this node in the output byte stream. |
| 179 | fn updateOffset(self: *Node, offset: usize) UpdateResult { |
| 180 | var node_size: usize = 0; |
| 181 | if (self.vmaddr_offset) |vmaddr| { |
| 182 | node_size += sizeULEB128Mem(self.export_flags.?); |
| 183 | node_size += sizeULEB128Mem(vmaddr); |
| 184 | node_size += sizeULEB128Mem(node_size); |
| 185 | } else { |
| 186 | node_size += 1; // 0x0 for non-terminal nodes |
| 187 | } |
| 188 | node_size += 1; // 1 byte for edge count |
| 189 | |
| 190 | for (self.edges.items) |edge| { |
| 191 | const next_node_offset = edge.to.trie_offset orelse 0; |
| 192 | node_size += edge.label.len + 1 + sizeULEB128Mem(next_node_offset); |
| 193 | } |
| 194 | |
| 195 | const trie_offset = self.trie_offset orelse 0; |
| 196 | const updated = offset != trie_offset; |
| 197 | self.trie_offset = offset; |
| 198 | |
| 199 | return .{ .node_size = node_size, .updated = updated }; |
| 200 | } |
| 201 | |
| 202 | /// Calculates number of bytes in ULEB128 encoding of value. |
| 203 | fn sizeULEB128Mem(value: u64) usize { |
| 204 | var res: usize = 0; |
| 205 | var v = value; |
| 206 | while (true) { |
| 207 | v = v >> 7; |
| 208 | res += 1; |
| 209 | if (v == 0) break; |
| 210 | } |
| 211 | return res; |
| 212 | } |
| 213 | }; |
| 214 | |
| 215 | /// Count of nodes in the trie. |
| 216 | /// The count is updated at every `put` call. |
| 217 | /// The trie always consists of at least a root node, hence |
| 218 | /// the count always starts at 1. |
| 219 | node_count: usize = 1, |
| 220 | /// The root node of the trie. |
| 221 | root: Node = .{}, |
| 222 | |
| 223 | /// Insert a symbol into the trie, updating the prefixes in the process. |
| 224 | /// This operation may change the layout of the trie by splicing edges in |
| 225 | /// certain circumstances. |
| 226 | pub fn put(self: *Trie, alloc: *Allocator, symbol: Symbol) !void { |
| 227 | const res = try self.root.put(alloc, symbol.name, 0); |
| 228 | self.node_count += res.node_count; |
| 229 | res.node.vmaddr_offset = symbol.vmaddr_offset; |
| 230 | res.node.export_flags = symbol.export_flags; |
| 231 | } |
| 232 | |
| 233 | /// Write the trie to a buffer ULEB128 encoded. |
| 234 | pub fn writeULEB128Mem(self: *Trie, alloc: *Allocator, buffer: *std.ArrayListUnmanaged(u8)) !void { |
| 235 | var ordered_nodes: std.ArrayListUnmanaged(*Node) = .{}; |
| 236 | defer ordered_nodes.deinit(alloc); |
| 237 | |
| 238 | try ordered_nodes.ensureCapacity(alloc, self.node_count); |
| 239 | walkInOrder(&self.root, &ordered_nodes); |
| 240 | |
| 241 | var offset: usize = 0; |
| 242 | var more: bool = true; |
| 243 | while (more) { |
| 244 | offset = 0; |
| 245 | more = false; |
| 246 | for (ordered_nodes.items) |node| { |
| 247 | const res = node.updateOffset(offset); |
| 248 | offset += res.node_size; |
| 249 | if (res.updated) more = true; |
| 250 | } |
| 251 | } |
| 252 | |
| 253 | try buffer.ensureCapacity(alloc, buffer.items.len + offset); |
| 254 | for (ordered_nodes.items) |node| { |
| 255 | try node.writeULEB128Mem(buffer); |
| 256 | } |
| 257 | } |
| 258 | |
| 259 | /// Walks the trie in DFS order gathering all nodes into a linear stream of nodes. |
| 260 | fn walkInOrder(node: *Node, list: *std.ArrayListUnmanaged(*Node)) void { |
| 261 | list.appendAssumeCapacity(node); |
| 262 | for (node.edges.items) |*edge| { |
| 263 | walkInOrder(edge.to, list); |
| 264 | } |
| 265 | } |
| 266 | |
| 267 | pub fn deinit(self: *Trie, alloc: *Allocator) void { |
| 268 | self.root.deinit(alloc); |
| 269 | } |
| 270 | |
| 271 | test "Trie node count" { |
| 272 | var gpa = testing.allocator; |
| 273 | var trie: Trie = .{}; |
| 274 | defer trie.deinit(gpa); |
| 275 | |
| 276 | testing.expectEqual(trie.node_count, 1); |
| 277 | |
| 278 | try trie.put(gpa, .{ |
| 279 | .name = "_main", |
| 280 | .vmaddr_offset = 0, |
| 281 | .export_flags = 0, |
| 282 | }); |
| 283 | testing.expectEqual(trie.node_count, 2); |
| 284 | |
| 285 | // Inserting the same node shouldn't update the trie. |
| 286 | try trie.put(gpa, .{ |
| 287 | .name = "_main", |
| 288 | .vmaddr_offset = 0, |
| 289 | .export_flags = 0, |
| 290 | }); |
| 291 | testing.expectEqual(trie.node_count, 2); |
| 292 | |
| 293 | try trie.put(gpa, .{ |
| 294 | .name = "__mh_execute_header", |
| 295 | .vmaddr_offset = 0x1000, |
| 296 | .export_flags = 0, |
| 297 | }); |
| 298 | testing.expectEqual(trie.node_count, 4); |
| 299 | |
| 300 | // Inserting the same node shouldn't update the trie. |
| 301 | try trie.put(gpa, .{ |
| 302 | .name = "__mh_execute_header", |
| 303 | .vmaddr_offset = 0x1000, |
| 304 | .export_flags = 0, |
| 305 | }); |
| 306 | testing.expectEqual(trie.node_count, 4); |
| 307 | try trie.put(gpa, .{ |
| 308 | .name = "_main", |
| 309 | .vmaddr_offset = 0, |
| 310 | .export_flags = 0, |
| 311 | }); |
| 312 | testing.expectEqual(trie.node_count, 4); |
| 313 | } |
| 314 | |
| 315 | test "Trie basic" { |
| 316 | var gpa = testing.allocator; |
| 317 | var trie: Trie = .{}; |
| 318 | defer trie.deinit(gpa); |
| 319 | |
| 320 | // root |
| 321 | testing.expect(trie.root.edges.items.len == 0); |
| 322 | |
| 323 | // root --- _st ---> node |
| 324 | try trie.put(gpa, .{ |
| 325 | .name = "_st", |
| 326 | .vmaddr_offset = 0, |
| 327 | .export_flags = 0, |
| 328 | }); |
| 329 | testing.expect(trie.root.edges.items.len == 1); |
| 330 | testing.expect(mem.eql(u8, trie.root.edges.items[0].label, "_st")); |
| 331 | |
| 332 | { |
| 333 | // root --- _st ---> node --- art ---> node |
| 334 | try trie.put(gpa, .{ |
| 335 | .name = "_start", |
| 336 | .vmaddr_offset = 0, |
| 337 | .export_flags = 0, |
| 338 | }); |
| 339 | testing.expect(trie.root.edges.items.len == 1); |
| 340 | |
| 341 | const nextEdge = &trie.root.edges.items[0]; |
| 342 | testing.expect(mem.eql(u8, nextEdge.label, "_st")); |
| 343 | testing.expect(nextEdge.to.edges.items.len == 1); |
| 344 | testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "art")); |
| 345 | } |
| 346 | { |
| 347 | // root --- _ ---> node --- st ---> node --- art ---> node |
| 348 | // | |
| 349 | // | --- main ---> node |
| 350 | try trie.put(gpa, .{ |
| 351 | .name = "_main", |
| 352 | .vmaddr_offset = 0, |
| 353 | .export_flags = 0, |
| 354 | }); |
| 355 | testing.expect(trie.root.edges.items.len == 1); |
| 356 | |
| 357 | const nextEdge = &trie.root.edges.items[0]; |
| 358 | testing.expect(mem.eql(u8, nextEdge.label, "_")); |
| 359 | testing.expect(nextEdge.to.edges.items.len == 2); |
| 360 | testing.expect(mem.eql(u8, nextEdge.to.edges.items[0].label, "st")); |
| 361 | testing.expect(mem.eql(u8, nextEdge.to.edges.items[1].label, "main")); |
| 362 | |
| 363 | const nextNextEdge = &nextEdge.to.edges.items[0]; |
| 364 | testing.expect(mem.eql(u8, nextNextEdge.to.edges.items[0].label, "art")); |
| 365 | } |
| 366 | } |
| 367 | |
| 368 | test "Trie.writeULEB128Mem" { |
| 369 | var gpa = testing.allocator; |
| 370 | var trie: Trie = .{}; |
| 371 | defer trie.deinit(gpa); |
| 372 | |
| 373 | try trie.put(gpa, .{ |
| 374 | .name = "__mh_execute_header", |
| 375 | .vmaddr_offset = 0, |
| 376 | .export_flags = 0, |
| 377 | }); |
| 378 | try trie.put(gpa, .{ |
| 379 | .name = "_main", |
| 380 | .vmaddr_offset = 0x1000, |
| 381 | .export_flags = 0, |
| 382 | }); |
| 383 | |
| 384 | var buffer: std.ArrayListUnmanaged(u8) = .{}; |
| 385 | defer buffer.deinit(gpa); |
| 386 | |
| 387 | try trie.writeULEB128Mem(gpa, &buffer); |
| 388 | |
| 389 | const exp_buffer = [_]u8{ |
| 390 | 0x0, |
| 391 | 0x1, |
| 392 | 0x5f, |
| 393 | 0x0, |
| 394 | 0x5, |
| 395 | 0x0, |
| 396 | 0x2, |
| 397 | 0x5f, |
| 398 | 0x6d, |
| 399 | 0x68, |
| 400 | 0x5f, |
| 401 | 0x65, |
| 402 | 0x78, |
| 403 | 0x65, |
| 404 | 0x63, |
| 405 | 0x75, |
| 406 | 0x74, |
| 407 | 0x65, |
| 408 | 0x5f, |
| 409 | 0x68, |
| 410 | 0x65, |
| 411 | 0x61, |
| 412 | 0x64, |
| 413 | 0x65, |
| 414 | 0x72, |
| 415 | 0x0, |
| 416 | 0x21, |
| 417 | 0x6d, |
| 418 | 0x61, |
| 419 | 0x69, |
| 420 | 0x6e, |
| 421 | 0x0, |
| 422 | 0x25, |
| 423 | 0x2, |
| 424 | 0x0, |
| 425 | 0x0, |
| 426 | 0x0, |
| 427 | 0x3, |
| 428 | 0x0, |
| 429 | 0x80, |
| 430 | 0x20, |
| 431 | 0x0, |
| 432 | }; |
| 433 | |
| 434 | testing.expect(buffer.items.len == exp_buffer.len); |
| 435 | testing.expect(mem.eql(u8, buffer.items, exp_buffer[0..])); |
| 436 | } |
Jakub Konka <kubkon@jakubkonka.com> 2020-10-09 17:41:52+02:00
GitHub <noreply@github.com> 2020-10-09 17:41:52+02:00
Signed by PGP key 4AEE18F83AFDEB23