kit

kit
git clone https://git.ryansepassi.com/git/kit.git
Log | Files | Refs | README

rewrite.c (12703B)

      1 #include "obj/rewrite.h"
      2 
      3 #include <limits.h>
      4 #include <string.h>
      5 
      6 #include "obj/coff/coff.h"
      7 
      8 #define PE_DEBUG_DIRECTORY_SIZE 28u
      9 
     10 typedef struct CoffRewriteView {
     11   ObjRewrite rw;
     12   u8* file_header;
     13   u8* optional_header;
     14   u8* section_headers;
     15   u16 section_count;
     16   u16 optional_size;
     17   u32 directory_count;
     18   u32 size_of_headers;
     19   u32 symbol_table_off;
     20   u32 symbol_count;
     21   u8* string_table;
     22   u32 string_table_size;
     23 } CoffRewriteView;
     24 
     25 static KitStatus coff_bad(CoffRewriteView* v, const char* detail) {
     26   kit_ctx_diagf(v->rw.ctx, "%.*s: malformed linked PE image (%s)",
     27                 KIT_SLICE_ARG(v->rw.name), detail);
     28   return KIT_MALFORMED;
     29 }
     30 
     31 static u8* coff_directory(CoffRewriteView* v, u32 index) {
     32   u64 off = 112u + (u64)index * COFF_DATA_DIRECTORY_SIZE;
     33   if (index >= v->directory_count || off + COFF_DATA_DIRECTORY_SIZE >
     34                                                v->optional_size)
     35     return NULL;
     36   return v->optional_header + (size_t)off;
     37 }
     38 
     39 static u8* coff_section(CoffRewriteView* v, u32 index) {
     40   if (index >= v->section_count) return NULL;
     41   return v->section_headers + (size_t)index * COFF_SECTION_HEADER_SIZE;
     42 }
     43 
     44 /* Translate a PE RVA range to its complete on-disk representation. Header
     45  * RVAs are file offsets; section RVAs are bounded by SizeOfRawData because
     46  * zero-fill tails have no bytes to rewrite. */
     47 static int coff_rva_file_range(CoffRewriteView* v, u32 rva, u32 size,
     48                                u64* off_out) {
     49   u32 i;
     50   if ((u64)rva + size > UINT32_MAX + 1ull) return 0;
     51   if (rva < v->size_of_headers) {
     52     if ((u64)rva + size > v->size_of_headers ||
     53         !obj_rw_range(v->rw.len, rva, size))
     54       return 0;
     55     *off_out = rva;
     56     return 1;
     57   }
     58   for (i = 0; i < v->section_count; ++i) {
     59     const u8* sec = coff_section(v, i);
     60     u32 va = obj_rw_u32(sec + 12, 0);
     61     u32 raw_size = obj_rw_u32(sec + 16, 0);
     62     u32 raw_off = obj_rw_u32(sec + 20, 0);
     63     u64 delta;
     64     if (rva < va) continue;
     65     delta = (u64)rva - va;
     66     if (delta > raw_size || size > (u64)raw_size - delta) continue;
     67     if (!obj_rw_range(v->rw.len, (u64)raw_off + delta, size)) return 0;
     68     *off_out = (u64)raw_off + delta;
     69     return 1;
     70   }
     71   return 0;
     72 }
     73 
     74 static int coff_name_prefix(const u8* name, size_t name_len,
     75                             const char* prefix) {
     76   size_t n = strlen(prefix);
     77   return n <= name_len && memcmp(name, prefix, n) == 0;
     78 }
     79 
     80 static KitStatus coff_section_name(CoffRewriteView* v, const u8* sec,
     81                                    const u8** name_out,
     82                                    size_t* name_len_out) {
     83   const u8* raw = sec;
     84   size_t len = 0;
     85   u32 off = 0;
     86   u32 i;
     87   if (raw[0] != '/') {
     88     while (len < 8u && raw[len]) ++len;
     89     *name_out = raw;
     90     *name_len_out = len;
     91     return KIT_OK;
     92   }
     93 
     94   /* Linked PE images produced by LLVM use COFF's /<decimal> string-table
     95    * convention for DWARF section names even though the PE specification only
     96    * blesses it for relocatable objects. Resolve it here so a real external
     97    * producer's .debug_* sections are not mistaken for ordinary load data. */
     98   if (raw[1] < '0' || raw[1] > '9') return coff_bad(v, "section name");
     99   for (i = 1; i < 8u && raw[i]; ++i) {
    100     u32 digit;
    101     if (raw[i] < '0' || raw[i] > '9') return coff_bad(v, "section name");
    102     digit = (u32)(raw[i] - '0');
    103     if (off > (UINT32_MAX - digit) / 10u)
    104       return coff_bad(v, "section name");
    105     off = off * 10u + digit;
    106   }
    107   if (!v->string_table || off < 4u || off >= v->string_table_size)
    108     return coff_bad(v, "section name");
    109   while (off + len < v->string_table_size && v->string_table[off + len])
    110     ++len;
    111   if (off + len == v->string_table_size)
    112     return coff_bad(v, "section name");
    113   *name_out = v->string_table + off;
    114   *name_len_out = len;
    115   return KIT_OK;
    116 }
    117 
    118 static KitStatus coff_debug_section(CoffRewriteView* v, const u8* sec,
    119                                     int* debug_out) {
    120   const u8* name;
    121   size_t name_len;
    122   KitStatus st = coff_section_name(v, sec, &name, &name_len);
    123   if (st != KIT_OK) return st;
    124   *debug_out = coff_name_prefix(name, name_len, ".debug") ||
    125                coff_name_prefix(name, name_len, ".zdebug") ||
    126                coff_name_prefix(name, name_len, ".stab");
    127   return KIT_OK;
    128 }
    129 
    130 static KitStatus coff_validate(CoffRewriteView* v, const KitContext* ctx,
    131                                KitSlice name, const KitSlice* input) {
    132   KitStatus st;
    133   u32 peoff, i, strings_size;
    134   u64 sections_size, symbols_size, strings_off;
    135   memset(v, 0, sizeof *v);
    136   st = obj_rewrite_init(&v->rw, ctx, name, input);
    137   if (st != KIT_OK) return st;
    138   if (v->rw.len < COFF_DOS_HEADER_SIZE ||
    139       obj_rw_u16(v->rw.bytes, 0) != IMAGE_DOS_SIGNATURE)
    140     return coff_bad(v, "DOS header");
    141   peoff = obj_rw_u32(v->rw.bytes + 0x3cu, 0);
    142   if (!obj_rw_range(v->rw.len, peoff,
    143                     4u + COFF_FILE_HEADER_SIZE) ||
    144       obj_rw_u32(v->rw.bytes + peoff, 0) != IMAGE_NT_SIGNATURE)
    145     return coff_bad(v, "PE header");
    146   v->file_header = v->rw.bytes + peoff + 4u;
    147   v->section_count = obj_rw_u16(v->file_header + 2, 0);
    148   v->optional_size = obj_rw_u16(v->file_header + 16, 0);
    149   if (v->optional_size == 0) {
    150     kit_ctx_diagf(ctx, "%.*s: relocatable object is not a linked image",
    151                   KIT_SLICE_ARG(name));
    152     return KIT_INVALID;
    153   }
    154   if (v->optional_size < 112u ||
    155       !obj_rw_range(v->rw.len, (u64)peoff + 24u, v->optional_size))
    156     return coff_bad(v, "optional header");
    157   v->optional_header = v->file_header + COFF_FILE_HEADER_SIZE;
    158   if (obj_rw_u16(v->optional_header, 0) !=
    159       IMAGE_NT_OPTIONAL_HDR64_MAGIC)
    160     return coff_bad(v, "PE32+ optional header");
    161   v->size_of_headers = obj_rw_u32(v->optional_header + 60, 0);
    162   v->directory_count = obj_rw_u32(v->optional_header + 108, 0);
    163   if (v->directory_count > COFF_NUM_DATA_DIRECTORIES)
    164     v->directory_count = COFF_NUM_DATA_DIRECTORIES;
    165   if ((u64)112u + (u64)v->directory_count * COFF_DATA_DIRECTORY_SIZE >
    166       v->optional_size)
    167     return coff_bad(v, "data directories");
    168   v->section_headers = v->optional_header + v->optional_size;
    169   sections_size = (u64)v->section_count * COFF_SECTION_HEADER_SIZE;
    170   if (!obj_rw_range(v->rw.len,
    171                     (u64)(v->section_headers - v->rw.bytes), sections_size))
    172     return coff_bad(v, "section table");
    173   for (i = 0; i < v->section_count; ++i) {
    174     const u8* sec = coff_section(v, i);
    175     u32 raw_size = obj_rw_u32(sec + 16, 0);
    176     u32 raw_off = obj_rw_u32(sec + 20, 0);
    177     if (raw_size && !obj_rw_range(v->rw.len, raw_off, raw_size))
    178       return coff_bad(v, "section contents");
    179   }
    180   v->symbol_table_off = obj_rw_u32(v->file_header + 8, 0);
    181   v->symbol_count = obj_rw_u32(v->file_header + 12, 0);
    182   if (!v->symbol_table_off && !v->symbol_count) return KIT_OK;
    183   if (!v->symbol_table_off ||
    184       v->symbol_count > UINT32_MAX / COFF_SYMBOL_SIZE)
    185     return coff_bad(v, "COFF symbol table");
    186   symbols_size = (u64)v->symbol_count * COFF_SYMBOL_SIZE;
    187   strings_off = (u64)v->symbol_table_off + symbols_size;
    188   if (!obj_rw_range(v->rw.len, v->symbol_table_off, symbols_size) ||
    189       !obj_rw_range(v->rw.len, strings_off, 4u))
    190     return coff_bad(v, "COFF symbol table");
    191   strings_size = obj_rw_u32(v->rw.bytes + (size_t)strings_off, 0);
    192   if (strings_size < 4u ||
    193       !obj_rw_range(v->rw.len, strings_off, strings_size))
    194     return coff_bad(v, "COFF string table");
    195   v->string_table = v->rw.bytes + (size_t)strings_off;
    196   v->string_table_size = strings_size;
    197   return KIT_OK;
    198 }
    199 
    200 static KitStatus coff_check_signature(CoffRewriteView* v,
    201                                       const KitLinkedRewriteOptions* opts) {
    202   u8* dir = coff_directory(v, IMAGE_DIRECTORY_ENTRY_SECURITY);
    203   u32 off, size;
    204   if (!dir) return KIT_OK;
    205   off = obj_rw_u32(dir, 0);
    206   size = obj_rw_u32(dir + 4, 0);
    207   if (!off && !size) return KIT_OK;
    208   if (!off || !size || !obj_rw_range(v->rw.len, off, size))
    209     return coff_bad(v, "certificate table");
    210   v->rw.report.had_signature = true;
    211   if (!opts->remove_signature) {
    212     kit_ctx_diagf(v->rw.ctx,
    213                   "%.*s: signed PE image is unchanged; pass "
    214                   "--remove-signature and sign the result again",
    215                   KIT_SLICE_ARG(v->rw.name));
    216     return KIT_UNSUPPORTED;
    217   }
    218   memset(v->rw.bytes + off, 0, size);
    219   memset(dir, 0, COFF_DATA_DIRECTORY_SIZE);
    220   v->rw.report.removed_signature = true;
    221   return KIT_OK;
    222 }
    223 
    224 static KitStatus coff_strip_debug_directory(CoffRewriteView* v) {
    225   u8* dir = coff_directory(v, IMAGE_DIRECTORY_ENTRY_DEBUG);
    226   u32 rva, size, i;
    227   u64 off;
    228   if (!dir) return KIT_OK;
    229   rva = obj_rw_u32(dir, 0);
    230   size = obj_rw_u32(dir + 4, 0);
    231   if (!rva && !size) return KIT_OK;
    232   if (!rva || !size || size % PE_DEBUG_DIRECTORY_SIZE != 0 ||
    233       !coff_rva_file_range(v, rva, size, &off))
    234     return coff_bad(v, "debug directory");
    235   for (i = 0; i < size / PE_DEBUG_DIRECTORY_SIZE; ++i) {
    236     const u8* entry = v->rw.input + (size_t)off +
    237                       (size_t)i * PE_DEBUG_DIRECTORY_SIZE;
    238     u32 data_size = obj_rw_u32(entry + 16, 0);
    239     u32 data_off = obj_rw_u32(entry + 24, 0);
    240     if (data_size) {
    241       if (!data_off || !obj_rw_range(v->rw.len, data_off, data_size))
    242         return coff_bad(v, "debug payload");
    243       memset(v->rw.bytes + data_off, 0, data_size);
    244     }
    245   }
    246   memset(v->rw.bytes + (size_t)off, 0, size);
    247   memset(dir, 0, COFF_DATA_DIRECTORY_SIZE);
    248   ++v->rw.report.removed_debug_sections;
    249   return KIT_OK;
    250 }
    251 
    252 static KitStatus coff_strip_debug_sections(CoffRewriteView* v) {
    253   u32 i;
    254   for (i = 0; i < v->section_count; ++i) {
    255     u8* sec = coff_section(v, i);
    256     u32 raw_size, raw_off;
    257     int debug;
    258     KitStatus st = coff_debug_section(v, sec, &debug);
    259     if (st != KIT_OK) return st;
    260     if (!debug) continue;
    261     raw_size = obj_rw_u32(sec + 16, 0);
    262     raw_off = obj_rw_u32(sec + 20, 0);
    263     if (raw_size) memset(v->rw.bytes + raw_off, 0, raw_size);
    264     /* PE loaders validate the complete mapped section table. Removing a
    265      * header in place (or leaving a zero-sized hole before later sections)
    266      * makes otherwise-valid images unloadable. Keep the ordinal, RVA, raw
    267      * span, and permissions intact, but zero the payload and replace any
    268      * short or string-table-backed debug name with an inline neutral name. */
    269     memcpy(sec, ".discard", 8u);
    270     ++v->rw.report.removed_debug_sections;
    271   }
    272   return KIT_OK;
    273 }
    274 
    275 static KitStatus coff_strip_symbol_table(CoffRewriteView* v) {
    276   u32 off = v->symbol_table_off;
    277   u32 count = v->symbol_count;
    278   u64 symbols_size, total;
    279   u32 i;
    280   int needs_strings = 0;
    281   if (!off && !count) return KIT_OK;
    282   symbols_size = (u64)count * COFF_SYMBOL_SIZE;
    283   total = symbols_size + v->string_table_size;
    284   for (i = 0; i < v->section_count; ++i) {
    285     const u8* sec = coff_section(v, i);
    286     if (sec[0] == '/') {
    287       const u8* resolved;
    288       size_t resolved_len;
    289       KitStatus st = coff_section_name(v, sec, &resolved, &resolved_len);
    290       if (st != KIT_OK) return st;
    291       (void)resolved;
    292       (void)resolved_len;
    293       needs_strings = 1;
    294     }
    295   }
    296   if (needs_strings) {
    297     /* NumberOfSymbols controls where readers locate the string table. Move it
    298      * over the removed symbol records so /offset section names keep resolving
    299      * with a zero symbol count. The fixed file extent remains transactional. */
    300     memmove(v->rw.bytes + off, v->string_table, v->string_table_size);
    301     memset(v->rw.bytes + off + v->string_table_size, 0,
    302            (size_t)total - v->string_table_size);
    303     v->string_table = v->rw.bytes + off;
    304   } else {
    305     memset(v->rw.bytes + off, 0, (size_t)total);
    306     obj_rw_put_u32(v->file_header + 8, 0, 0);
    307     v->string_table = NULL;
    308     v->string_table_size = 0;
    309   }
    310   obj_rw_put_u32(v->file_header + 12, 0, 0);
    311   v->rw.report.removed_symbols += count;
    312   return KIT_OK;
    313 }
    314 
    315 KitStatus obj_rewrite_coff(const KitContext* ctx, KitSlice name,
    316                            const KitSlice* input,
    317                            const KitLinkedRewriteOptions* opts,
    318                            KitWriter* out,
    319                            KitLinkedRewriteReport* report_out) {
    320   CoffRewriteView v;
    321   KitStatus st = coff_validate(&v, ctx, name, input);
    322   if (st == KIT_OK) st = coff_check_signature(&v, opts);
    323   if (st == KIT_OK) st = coff_strip_debug_directory(&v);
    324   if (st == KIT_OK) st = coff_strip_debug_sections(&v);
    325   if (st == KIT_OK && opts->strip_level == KIT_LINKED_STRIP_ALL)
    326     st = coff_strip_symbol_table(&v);
    327   if (st == KIT_OK) {
    328     u16 characteristics = obj_rw_u16(v.file_header + 18, 0);
    329     obj_rw_put_u16(v.file_header + 18,
    330                    (u16)(characteristics | IMAGE_FILE_DEBUG_STRIPPED), 0);
    331     st = obj_rewrite_commit(&v.rw, out, report_out);
    332   }
    333   if (st != KIT_OK) obj_rewrite_discard(&v.rw);
    334   return st;
    335 }