object_rewrite_test.c (9921B)
1 #include <kit/object.h> 2 3 #include <stdint.h> 4 #include <stdio.h> 5 #include <string.h> 6 7 #include "lib/kit_unit.h" 8 9 #define PE_SIZE 0x900u 10 #define PE_OFF 0x80u 11 #define FILE_HDR (PE_OFF + 4u) 12 #define OPT_HDR (FILE_HDR + 20u) 13 #define SEC_HDR (OPT_HDR + 240u) 14 #define DEBUG_RAW 0x400u 15 #define RDATA_RAW 0x600u 16 #define CERT_RAW 0x800u 17 #define SYMTAB_RAW 0x880u 18 19 static void put16(uint8_t* p, uint16_t v) { 20 p[0] = (uint8_t)v; 21 p[1] = (uint8_t)(v >> 8); 22 } 23 24 static void put32(uint8_t* p, uint32_t v) { 25 p[0] = (uint8_t)v; 26 p[1] = (uint8_t)(v >> 8); 27 p[2] = (uint8_t)(v >> 16); 28 p[3] = (uint8_t)(v >> 24); 29 } 30 31 static void put64(uint8_t* p, uint64_t v) { 32 put32(p, (uint32_t)v); 33 put32(p + 4, (uint32_t)(v >> 32)); 34 } 35 36 static uint32_t get32(const uint8_t* p) { 37 return (uint32_t)p[0] | ((uint32_t)p[1] << 8) | 38 ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24); 39 } 40 41 static void put_dir(uint8_t* pe, uint32_t index, uint32_t addr, 42 uint32_t size) { 43 uint8_t* d = pe + OPT_HDR + 112u + index * 8u; 44 put32(d, addr); 45 put32(d + 4, size); 46 } 47 48 static void put_section(uint8_t* pe, uint32_t index, const char* name, 49 uint32_t virtual_size, uint32_t va, 50 uint32_t raw_size, uint32_t raw_off, 51 uint32_t flags) { 52 uint8_t* s = pe + SEC_HDR + index * 40u; 53 size_t n = strlen(name); 54 if (n > 8u) n = 8u; 55 memcpy(s, name, n); 56 put32(s + 8, virtual_size); 57 put32(s + 12, va); 58 put32(s + 16, raw_size); 59 put32(s + 20, raw_off); 60 put32(s + 36, flags); 61 } 62 63 static void build_signed_debug_pe(uint8_t pe[PE_SIZE]) { 64 uint8_t* debug_entry; 65 uint8_t* string_table; 66 memset(pe, 0, PE_SIZE); 67 put16(pe, 0x5a4du); 68 put32(pe + 0x3c, PE_OFF); 69 put32(pe + PE_OFF, 0x00004550u); 70 put16(pe + FILE_HDR, 0x8664u); 71 put16(pe + FILE_HDR + 2, 3u); 72 put32(pe + FILE_HDR + 8, SYMTAB_RAW); 73 put32(pe + FILE_HDR + 12, 1u); 74 put16(pe + FILE_HDR + 16, 240u); 75 put16(pe + FILE_HDR + 18, 0x22u); 76 put16(pe + OPT_HDR, 0x020bu); 77 put32(pe + OPT_HDR + 16, 0x1000u); 78 put32(pe + OPT_HDR + 20, 0x1000u); 79 put32(pe + OPT_HDR + 60, 0x200u); 80 put32(pe + OPT_HDR + 108, 16u); 81 put_dir(pe, 1u, 0x33330000u, 0x44u); /* import metadata sentinel */ 82 put_dir(pe, 4u, CERT_RAW, 0x80u); 83 put_dir(pe, 6u, 0x3000u, 28u); 84 put_dir(pe, 9u, 0x55550000u, 0x66u); /* TLS metadata sentinel */ 85 put_section(pe, 0u, ".text", 0x80u, 0x1000u, 0x200u, 0x200u, 86 0x60000020u); 87 put_section(pe, 1u, "/4", 0x80u, 0x2000u, 0x200u, DEBUG_RAW, 88 0x42000040u); 89 put_section(pe, 2u, "/16", 0x80u, 0x3000u, 0x200u, RDATA_RAW, 90 0x40000040u); 91 memset(pe + 0x200u, 0xa5, 0x200u); 92 memset(pe + DEBUG_RAW, 0xdb, 0x200u); 93 debug_entry = pe + RDATA_RAW; 94 put32(debug_entry + 12, 2u); /* CodeView */ 95 put32(debug_entry + 16, 16u); 96 put32(debug_entry + 20, 0x2020u); 97 put32(debug_entry + 24, 0x420u); 98 memset(pe + CERT_RAW, 0xce, 0x80u); 99 memset(pe + SYMTAB_RAW, 0x5a, 18u); 100 string_table = pe + SYMTAB_RAW + 18u; 101 put32(string_table, 31u); 102 memcpy(string_table + 4u, ".debug_info", 12u); 103 memcpy(string_table + 16u, ".readonly_data", 15u); 104 } 105 106 static int all_byte(const uint8_t* p, size_t n, uint8_t value) { 107 size_t i; 108 for (i = 0; i < n; ++i) 109 if (p[i] != value) return 0; 110 return 1; 111 } 112 113 static void check_signed_transaction(KitUnit* u, const uint8_t* pe) { 114 uint8_t snapshot[PE_SIZE]; 115 KitSlice input; 116 KitLinkedRewriteOptions opts; 117 KitWriter* out = NULL; 118 size_t out_len = 99u; 119 KitStatus st; 120 memcpy(snapshot, pe, PE_SIZE); 121 input.data = pe; 122 input.len = PE_SIZE; 123 memset(&opts, 0, sizeof opts); 124 opts.strip_level = KIT_LINKED_STRIP_DEBUG; 125 CU_EXPECT(u, kit_writer_mem(&u->heap, &out) == KIT_OK && out, 126 "signed transaction memory writer"); 127 st = kit_obj_rewrite_linked(&u->ctx, KIT_SLICE_LIT("signed.exe"), &input, 128 &opts, out, NULL); 129 CU_EXPECT(u, st == KIT_UNSUPPORTED, 130 "signed PE without opt-in is unsupported (%d)", (int)st); 131 (void)kit_writer_mem_bytes(out, &out_len); 132 CU_EXPECT(u, out_len == 0u, "signed failure emits no bytes"); 133 CU_EXPECT(u, memcmp(pe, snapshot, PE_SIZE) == 0, 134 "signed failure leaves input unchanged"); 135 CU_EXPECT(u, strstr(u->last_diag, "--remove-signature") != NULL, 136 "signed failure explains explicit opt-in"); 137 kit_writer_close(out); 138 } 139 140 static void check_strip_debug(KitUnit* u, const uint8_t* pe) { 141 KitSlice input; 142 KitLinkedRewriteOptions opts; 143 KitLinkedRewriteReport report; 144 KitWriter* out = NULL; 145 const uint8_t* bytes; 146 size_t len = 0; 147 memset(&opts, 0, sizeof opts); 148 input.data = pe; 149 input.len = PE_SIZE; 150 memset(&report, 0, sizeof report); 151 opts.strip_level = KIT_LINKED_STRIP_DEBUG; 152 opts.remove_signature = true; 153 CU_EXPECT(u, kit_writer_mem(&u->heap, &out) == KIT_OK && out, 154 "debug rewrite memory writer"); 155 CU_EXPECT(u, 156 kit_obj_rewrite_linked(&u->ctx, KIT_SLICE_LIT("debug.exe"), 157 &input, &opts, out, &report) == KIT_OK, 158 "debug rewrite succeeds"); 159 bytes = kit_writer_mem_bytes(out, &len); 160 CU_EXPECT(u, len == PE_SIZE, "PE rewrite preserves file extent"); 161 CU_EXPECT(u, get32(bytes + OPT_HDR + 16) == 0x1000u, 162 "entry point preserved"); 163 CU_EXPECT(u, get32(bytes + OPT_HDR + 112u + 8u) == 0x33330000u && 164 get32(bytes + OPT_HDR + 112u + 9u * 8u) == 0x55550000u, 165 "import and TLS directories preserved"); 166 CU_EXPECT(u, all_byte(bytes + 0x200u, 0x200u, 0xa5), 167 "mapped text bytes preserved"); 168 CU_EXPECT(u, all_byte(bytes + DEBUG_RAW, 0x200u, 0), 169 "debug section bytes removed"); 170 CU_EXPECT(u, memcmp(bytes + SEC_HDR + 40u, ".discard", 8u) == 0, 171 "long-name debug section header neutralized"); 172 CU_EXPECT(u, all_byte(bytes + CERT_RAW, 0x80u, 0), 173 "certificate bytes removed"); 174 CU_EXPECT(u, get32(bytes + OPT_HDR + 112u + 4u * 8u) == 0u && 175 get32(bytes + OPT_HDR + 112u + 6u * 8u) == 0u, 176 "certificate and debug directories cleared"); 177 CU_EXPECT(u, get32(bytes + FILE_HDR + 8) == SYMTAB_RAW && 178 get32(bytes + FILE_HDR + 12) == 1u, 179 "--strip-debug retains COFF symbol table"); 180 CU_EXPECT(u, report.had_signature && report.removed_signature, 181 "signature removal reported"); 182 CU_EXPECT(u, report.removed_debug_sections == 2u, 183 "debug directory and section reported"); 184 kit_writer_close(out); 185 } 186 187 static void check_strip_all(KitUnit* u, const uint8_t* pe) { 188 KitSlice input; 189 KitLinkedRewriteOptions opts; 190 KitLinkedRewriteReport report; 191 KitWriter* out = NULL; 192 const uint8_t* bytes; 193 size_t len = 0; 194 memset(&opts, 0, sizeof opts); 195 input.data = pe; 196 input.len = PE_SIZE; 197 memset(&report, 0, sizeof report); 198 opts.strip_level = KIT_LINKED_STRIP_ALL; 199 opts.remove_signature = true; 200 CU_EXPECT(u, kit_writer_mem(&u->heap, &out) == KIT_OK && out, 201 "all rewrite memory writer"); 202 CU_EXPECT(u, 203 kit_obj_rewrite_linked(&u->ctx, KIT_SLICE_LIT("all.exe"), &input, 204 &opts, out, &report) == KIT_OK, 205 "strip-all rewrite succeeds"); 206 bytes = kit_writer_mem_bytes(out, &len); 207 CU_EXPECT(u, len == PE_SIZE, "strip-all preserves PE extent"); 208 CU_EXPECT(u, get32(bytes + FILE_HDR + 8) == SYMTAB_RAW && 209 get32(bytes + FILE_HDR + 12) == 0u, 210 "strip-all retains locator for section-name string table"); 211 CU_EXPECT(u, all_byte(bytes + DEBUG_RAW, 0x200u, 0), 212 "strip-all removes long-name debug section"); 213 CU_EXPECT(u, memcmp(bytes + SEC_HDR + 40u, ".discard", 8u) == 0, 214 "strip-all neutralizes long-name debug section header"); 215 CU_EXPECT(u, get32(bytes + SYMTAB_RAW) == 31u && 216 memcmp(bytes + SYMTAB_RAW + 16u, ".readonly_data", 15u) == 217 0, 218 "strip-all preserves long non-debug section name"); 219 CU_EXPECT(u, all_byte(bytes + SYMTAB_RAW + 31u, 18u, 0), 220 "strip-all clears removed COFF symbol records"); 221 CU_EXPECT(u, report.removed_symbols == 1u, 222 "strip-all reports removed symbols"); 223 kit_writer_close(out); 224 } 225 226 static void check_malformed_macho_section(KitUnit* u) { 227 uint8_t macho[184]; 228 KitSlice input; 229 KitLinkedRewriteOptions opts; 230 KitWriter* out = NULL; 231 size_t out_len = 99u; 232 KitStatus st; 233 memset(macho, 0, sizeof macho); 234 put32(macho, 0xfeedfacfu); /* MH_MAGIC_64 */ 235 put32(macho + 12, 2u); /* MH_EXECUTE */ 236 put32(macho + 16, 1u); /* ncmds */ 237 put32(macho + 20, 152u); /* sizeofcmds */ 238 put32(macho + 32, 0x19u); /* LC_SEGMENT_64 */ 239 put32(macho + 36, 152u); 240 memcpy(macho + 40, "__DWARF", 7u); 241 put32(macho + 96, 1u); /* nsects */ 242 memcpy(macho + 104, "__debug_info", 12u); 243 memcpy(macho + 120, "__DWARF", 7u); 244 put64(macho + 144, 16u); /* section size */ 245 put32(macho + 152, 4096u); /* impossible file offset */ 246 put32(macho + 168, 0x02000000u); /* S_ATTR_DEBUG */ 247 input.data = macho; 248 input.len = sizeof macho; 249 memset(&opts, 0, sizeof opts); 250 opts.strip_level = KIT_LINKED_STRIP_DEBUG; 251 CU_EXPECT(u, kit_writer_mem(&u->heap, &out) == KIT_OK && out, 252 "malformed Mach-O memory writer"); 253 st = kit_obj_rewrite_linked(&u->ctx, KIT_SLICE_LIT("bad-macho"), &input, 254 &opts, out, NULL); 255 CU_EXPECT(u, st == KIT_MALFORMED, 256 "out-of-range Mach-O debug section is malformed (%d)", (int)st); 257 (void)kit_writer_mem_bytes(out, &out_len); 258 CU_EXPECT(u, out_len == 0u, "malformed Mach-O emits no bytes"); 259 kit_writer_close(out); 260 } 261 262 int main(void) { 263 KitUnit u; 264 uint8_t pe[PE_SIZE]; 265 kit_unit_init(&u); 266 build_signed_debug_pe(pe); 267 check_signed_transaction(&u, pe); 268 check_strip_debug(&u, pe); 269 check_strip_all(&u, pe); 270 check_malformed_macho_section(&u); 271 kit_unit_summary(&u, "object_rewrite_test"); 272 return kit_unit_status(&u); 273 }