lib/closure/src/binary/pe.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const binary = @import("root.zig");
2 const read = @import("read.zig");
3 const std = @import("std");
4
5 pub const sections_max: usize = 96;
6
7 const InspectError = error{
8 UnsupportedPe,
9 PeSectionCapacityExceeded,
10 RangeCapacityExceeded,
11 EvidenceCapacityExceeded,
12 RelocationDirectoryUnmapped,
13 };
14
15 pub const Error = read.Error || InspectError;
16
17 const coff_header_bytes: usize = 20;
18 const section_header_bytes: usize = 40;
19 const machine_x86_64: u16 = 0x8664;
20 const optional_pe32_plus: u16 = 0x20b;
21 const section_execute: u32 = 0x2000_0000;
22 const base_relocation_directory: usize = 5;
23
24 pub fn matches(bytes: []const u8) bool {
25 return bytes.len >= 2 and std.mem.eql(u8, bytes[0..2], "MZ");
26 }
27
28 pub fn inspect(
29 bytes: []const u8,
30 scratch: binary.Scratch,
31 ) Error!binary.Report {
32 const header = try headerFor(bytes);
33 var range_count: usize = 0;
34 var evidence_count: usize = 0;
35 var work: u64 = 1;
36 try survey(
37 bytes,
38 header,
39 scratch.ranges.len,
40 scratch.evidence.len,
41 &range_count,
42 &evidence_count,
43 &work,
44 null,
45 );
46 range_count = 0;
47 evidence_count = 0;
48 work = 1;
49 try survey(
50 bytes,
51 header,
52 scratch.ranges.len,
53 scratch.evidence.len,
54 &range_count,
55 &evidence_count,
56 &work,
57 scratch,
58 );
59 return .{
60 .kind = .pe32_plus,
61 .ranges = scratch.ranges[0..range_count],
62 .evidence = scratch.evidence[0..evidence_count],
63 .work = work,
64 };
65 }
66
67 const Header = struct {
68 coff_offset: usize,
69 optional_offset: usize,
70 optional_bytes: usize,
71 section_offset: usize,
72 section_count: usize,
73 symbol_offset: u32,
74 symbol_count: u32,
75 relocation_rva: u32,
76 relocation_bytes: u32,
77 };
78
79 const RelocationDirectory = struct {
80 rva: u32 = 0,
81 byte_length: u32 = 0,
82 };
83
84 fn headerFor(bytes: []const u8) Error!Header {
85 _ = try read.take(bytes, 0, 64);
86 if (!matches(bytes)) return error.UnsupportedPe;
87 const pe_offset = try read.toOffset(try read.u32le(bytes, 0x3c));
88 const signature = try read.take(bytes, pe_offset, 4);
89 if (!std.mem.eql(u8, signature, "PE\x00\x00")) {
90 return error.UnsupportedPe;
91 }
92 const coff_offset = std.math.add(usize, pe_offset, 4) catch
93 return error.BinaryArithmeticOverflow;
94 _ = try read.take(bytes, coff_offset, coff_header_bytes);
95 if (try read.u16le(bytes, coff_offset) != machine_x86_64) {
96 return error.UnsupportedPe;
97 }
98 const section_count = try read.u16le(bytes, coff_offset + 2);
99 if (section_count > sections_max) {
100 return error.PeSectionCapacityExceeded;
101 }
102 const optional_bytes = try read.u16le(bytes, coff_offset + 16);
103 const optional_offset = std.math.add(
104 usize,
105 coff_offset,
106 coff_header_bytes,
107 ) catch return error.BinaryArithmeticOverflow;
108 _ = try read.take(bytes, optional_offset, optional_bytes);
109 if (optional_bytes < 112 or
110 try read.u16le(bytes, optional_offset) != optional_pe32_plus)
111 {
112 return error.UnsupportedPe;
113 }
114 const directory_count = try read.u32le(bytes, optional_offset + 108);
115 const relocation = try relocationDirectory(
116 bytes,
117 optional_offset,
118 optional_bytes,
119 directory_count,
120 );
121 const section_offset = std.math.add(
122 usize,
123 optional_offset,
124 optional_bytes,
125 ) catch return error.BinaryArithmeticOverflow;
126 if (section_count > 0) {
127 const last = try read.indexed(
128 section_offset,
129 section_count - 1,
130 section_header_bytes,
131 );
132 _ = try read.take(bytes, last, section_header_bytes);
133 }
134 return .{
135 .coff_offset = coff_offset,
136 .optional_offset = optional_offset,
137 .optional_bytes = optional_bytes,
138 .section_offset = section_offset,
139 .section_count = section_count,
140 .symbol_offset = try read.u32le(bytes, coff_offset + 8),
141 .symbol_count = try read.u32le(bytes, coff_offset + 12),
142 .relocation_rva = relocation.rva,
143 .relocation_bytes = relocation.byte_length,
144 };
145 }
146
147 fn relocationDirectory(
148 bytes: []const u8,
149 optional_offset: usize,
150 optional_bytes: usize,
151 directory_count: u32,
152 ) Error!RelocationDirectory {
153 if (directory_count <= base_relocation_directory) return .{};
154 const directory_relative = 112 + base_relocation_directory * 8;
155 const directory_end = directory_relative + 8;
156 if (directory_end > optional_bytes) return error.UnsupportedPe;
157 const directory_offset = std.math.add(
158 usize,
159 optional_offset,
160 directory_relative,
161 ) catch return error.BinaryArithmeticOverflow;
162 _ = try read.take(bytes, directory_offset, 8);
163 return .{
164 .rva = try read.u32le(bytes, directory_offset),
165 .byte_length = try read.u32le(bytes, directory_offset + 4),
166 };
167 }
168
169 fn survey(
170 bytes: []const u8,
171 header: Header,
172 ranges_max: usize,
173 evidence_max: usize,
174 range_count: *usize,
175 evidence_count: *usize,
176 work: *u64,
177 output: ?binary.Scratch,
178 ) Error!void {
179 _ = header.coff_offset;
180 _ = header.optional_offset;
181 _ = header.optional_bytes;
182 try surveySections(
183 bytes,
184 header,
185 ranges_max,
186 range_count,
187 work,
188 output,
189 );
190 try surveyEvidence(
191 bytes,
192 header,
193 evidence_max,
194 evidence_count,
195 work,
196 output,
197 );
198 }
199
200 fn surveySections(
201 bytes: []const u8,
202 header: Header,
203 ranges_max: usize,
204 range_count: *usize,
205 work: *u64,
206 output: ?binary.Scratch,
207 ) Error!void {
208 for (0..header.section_count) |index| {
209 work.* += 1;
210 const offset = try read.indexed(
211 header.section_offset,
212 index,
213 section_header_bytes,
214 );
215 const raw_length = try read.u32le(bytes, offset + 16);
216 const raw_offset = try read.u32le(bytes, offset + 20);
217 _ = try read.take(bytes, raw_offset, raw_length);
218 const characteristics = try read.u32le(bytes, offset + 36);
219 if (raw_length == 0 or characteristics & section_execute == 0) {
220 continue;
221 }
222 if (range_count.* == ranges_max) {
223 return error.RangeCapacityExceeded;
224 }
225 if (output) |present| {
226 present.ranges[range_count.*] = .{
227 .offset = raw_offset,
228 .length = raw_length,
229 };
230 }
231 range_count.* += 1;
232 }
233 }
234
235 fn surveyEvidence(
236 bytes: []const u8,
237 header: Header,
238 evidence_max: usize,
239 evidence_count: *usize,
240 work: *u64,
241 output: ?binary.Scratch,
242 ) Error!void {
243 if (header.symbol_offset != 0 and header.symbol_count != 0) {
244 const symbol_bytes = std.math.mul(
245 u32,
246 header.symbol_count,
247 18,
248 ) catch return error.BinaryArithmeticOverflow;
249 try appendEvidence(
250 bytes,
251 .symbols,
252 header.symbol_offset,
253 symbol_bytes,
254 evidence_max,
255 evidence_count,
256 output,
257 );
258 }
259 if (header.relocation_rva != 0 and header.relocation_bytes != 0) {
260 const relocation_offset = try rvaToOffset(
261 bytes,
262 header,
263 header.relocation_rva,
264 header.relocation_bytes,
265 work,
266 );
267 try appendEvidence(
268 bytes,
269 .relocations,
270 relocation_offset,
271 header.relocation_bytes,
272 evidence_max,
273 evidence_count,
274 output,
275 );
276 }
277 }
278
279 fn appendEvidence(
280 bytes: []const u8,
281 kind: binary.EvidenceKind,
282 offset: u32,
283 length: u32,
284 evidence_max: usize,
285 evidence_count: *usize,
286 output: ?binary.Scratch,
287 ) Error!void {
288 _ = try read.take(bytes, offset, length);
289 if (evidence_count.* == evidence_max) {
290 return error.EvidenceCapacityExceeded;
291 }
292 if (output) |present| {
293 present.evidence[evidence_count.*] = .{
294 .kind = kind,
295 .offset = offset,
296 .length = length,
297 };
298 }
299 evidence_count.* += 1;
300 }
301
302 fn rvaToOffset(
303 bytes: []const u8,
304 header: Header,
305 rva: u32,
306 length: u32,
307 work: *u64,
308 ) Error!u32 {
309 for (0..header.section_count) |index| {
310 work.* += 1;
311 const offset = try read.indexed(
312 header.section_offset,
313 index,
314 section_header_bytes,
315 );
316 const virtual_size = try read.u32le(bytes, offset + 8);
317 const virtual_address = try read.u32le(bytes, offset + 12);
318 const raw_length = try read.u32le(bytes, offset + 16);
319 const raw_offset = try read.u32le(bytes, offset + 20);
320 const mapped_length = @max(virtual_size, raw_length);
321 const relative = std.math.sub(
322 u32,
323 rva,
324 virtual_address,
325 ) catch continue;
326 if (relative > mapped_length) continue;
327 const relative_end = std.math.add(u32, relative, length) catch
328 return error.BinaryArithmeticOverflow;
329 if (relative_end > raw_length) continue;
330 return std.math.add(u32, raw_offset, relative) catch
331 return error.BinaryArithmeticOverflow;
332 }
333 return error.RelocationDirectoryUnmapped;
334 }
335
336 fn put16(bytes: []u8, offset: usize, value: u16) void {
337 std.mem.writeInt(u16, bytes[offset..][0..2], value, .little);
338 }
339
340 fn put32(bytes: []u8, offset: usize, value: u32) void {
341 std.mem.writeInt(u32, bytes[offset..][0..4], value, .little);
342 }
343
344 test "PE inspector reports executable section symbols and relocations" {
345 var bytes: [1_024]u8 = @splat(0);
346 @memcpy(bytes[0..2], "MZ");
347 put32(&bytes, 0x3c, 64);
348 @memcpy(bytes[64..68], "PE\x00\x00");
349 const coff = 68;
350 put16(&bytes, coff, machine_x86_64);
351 put16(&bytes, coff + 2, 1);
352 put32(&bytes, coff + 8, 800);
353 put32(&bytes, coff + 12, 2);
354 put16(&bytes, coff + 16, 160);
355 const optional = coff + coff_header_bytes;
356 put16(&bytes, optional, optional_pe32_plus);
357 put32(&bytes, optional + 108, 6);
358 put32(&bytes, optional + 112 + base_relocation_directory * 8, 0x1010);
359 put32(
360 &bytes,
361 optional + 112 + base_relocation_directory * 8 + 4,
362 8,
363 );
364 const section = optional + 160;
365 put32(&bytes, section + 8, 128);
366 put32(&bytes, section + 12, 0x1000);
367 put32(&bytes, section + 16, 128);
368 put32(&bytes, section + 20, 512);
369 put32(&bytes, section + 36, section_execute);
370 var ranges: [1]binary.Range = undefined;
371 var evidence: [2]binary.Evidence = undefined;
372 const report = try inspect(&bytes, .{
373 .ranges = &ranges,
374 .evidence = &evidence,
375 });
376 try std.testing.expectEqual(binary.Kind.pe32_plus, report.kind);
377 try std.testing.expectEqual(@as(u64, 512), report.ranges[0].offset);
378 try std.testing.expectEqual(
379 binary.EvidenceKind.symbols,
380 report.evidence[0].kind,
381 );
382 try std.testing.expectEqual(
383 binary.EvidenceKind.relocations,
384 report.evidence[1].kind,
385 );
386 try std.testing.expectEqual(@as(u64, 528), report.evidence[1].offset);
387
388 const range_sentinel: binary.Range = .{
389 .offset = 91,
390 .length = 93,
391 };
392 const evidence_sentinel: binary.Evidence = .{
393 .kind = .symbols,
394 .offset = 95,
395 .length = 97,
396 };
397 var guarded_ranges = [_]binary.Range{range_sentinel};
398 var guarded_evidence = [_]binary.Evidence{evidence_sentinel};
399 try std.testing.expectError(
400 error.EvidenceCapacityExceeded,
401 inspect(&bytes, .{
402 .ranges = &guarded_ranges,
403 .evidence = &guarded_evidence,
404 }),
405 );
406 try std.testing.expectEqual(range_sentinel, guarded_ranges[0]);
407 try std.testing.expectEqual(evidence_sentinel, guarded_evidence[0]);
408 }
409
410 test "PE inspector rejects a data directory beyond the optional header" {
411 var bytes: [512]u8 = @splat(0);
412 @memcpy(bytes[0..2], "MZ");
413 put32(&bytes, 0x3c, 64);
414 @memcpy(bytes[64..68], "PE\x00\x00");
415 const coff = 68;
416 put16(&bytes, coff, machine_x86_64);
417 put16(&bytes, coff + 2, 1);
418 put16(&bytes, coff + 16, 112);
419 const optional = coff + coff_header_bytes;
420 put16(&bytes, optional, optional_pe32_plus);
421 put32(&bytes, optional + 108, 6);
422 var ranges: [1]binary.Range = undefined;
423 var evidence: [1]binary.Evidence = undefined;
424 try std.testing.expectError(
425 error.UnsupportedPe,
426 inspect(&bytes, .{
427 .ranges = &ranges,
428 .evidence = &evidence,
429 }),
430 );
431 }
432
433 test "PE inspector rejects a relocation directory outside raw bytes" {
434 var bytes: [512]u8 = @splat(0);
435 @memcpy(bytes[0..2], "MZ");
436 put32(&bytes, 0x3c, 64);
437 @memcpy(bytes[64..68], "PE\x00\x00");
438 const coff = 68;
439 put16(&bytes, coff, machine_x86_64);
440 put16(&bytes, coff + 2, 1);
441 put16(&bytes, coff + 16, 160);
442 const optional = coff + coff_header_bytes;
443 put16(&bytes, optional, optional_pe32_plus);
444 put32(&bytes, optional + 108, 6);
445 put32(&bytes, optional + 112 + base_relocation_directory * 8, 0x4000);
446 put32(
447 &bytes,
448 optional + 112 + base_relocation_directory * 8 + 4,
449 8,
450 );
451 const section = optional + 160;
452 put32(&bytes, section + 8, 64);
453 put32(&bytes, section + 12, 0x1000);
454 put32(&bytes, section + 16, 64);
455 put32(&bytes, section + 20, 400);
456 var ranges: [1]binary.Range = undefined;
457 var evidence: [1]binary.Evidence = undefined;
458 try std.testing.expectError(
459 error.RelocationDirectoryUnmapped,
460 inspect(&bytes, .{
461 .ranges = &ranges,
462 .evidence = &evidence,
463 }),
464 );
465 }