kit

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

pe-image-read.c (14022B)


      1 /* PE32+ linked-image reader round-trip (read_coff_image, no external
      2  * toolchain).
      3  *
      4  * Links a tiny PIE executable in memory with kit's own COFF linker — a
      5  * .text entry plus a .data slot that takes an absolute (R_ABS64) reference
      6  * to an imported ExitProcess from KERNEL32.dll (via a short-import shim) —
      7  * then re-opens the emitted bytes through the public kit_obj_open and
      8  * asserts the neutral image view the reader populates:
      9  *   - kind EXEC, nonzero entry / image base
     10  *   - segments + sections (one per PE section, .text executable)
     11  *   - dependency KERNEL32.dll carrying the ExitProcess import
     12  *   - dynamic symbol ExitProcess (undefined import)
     13  *   - base relocation(s) for the absolute .data pointer (PIE)
     14  *   - raw escape hatch: 16 data directories + subsystem + dllchars,
     15  *     IMPORT directory populated
     16  *
     17  * Runs on every host (the reader is ours); covers both x86_64 and aarch64
     18  * Windows targets. */
     19 
     20 #include <kit/core.h>
     21 #include <kit/link.h>
     22 #include <kit/object.h>
     23 #include <setjmp.h>
     24 #include <stdarg.h>
     25 #include <stdio.h>
     26 #include <stdlib.h>
     27 #include <string.h>
     28 
     29 #include "core/core.h"
     30 #include "core/pool.h"
     31 #include "link/link.h"
     32 #include "obj/obj.h"
     33 
     34 /* ---- short-import wire constants (mirror pe-import-smoke.c). ---- */
     35 #define SHIM_HEADER_SIZE 20u
     36 #define SHIM_SYM_CSTR "ExitProcess"
     37 #define SHIM_DLL_CSTR "KERNEL32.dll"
     38 #define SHIM_SYM_NUL_LEN 12u /* "ExitProcess\0" */
     39 #define SHIM_DLL_NUL_LEN 13u /* "KERNEL32.dll\0" */
     40 #define SHIM_DATA_LEN (SHIM_SYM_NUL_LEN + SHIM_DLL_NUL_LEN)
     41 #define SHIM_TOTAL_LEN (SHIM_HEADER_SIZE + SHIM_DATA_LEN)
     42 #define COFF_SHIMP_SIG2 0xFFFFu
     43 /* TypeFlags = Type=CODE(0) | (NameType=NAME(1) << 2) = 0x0004. */
     44 #define COFF_SHIMP_TYPEFLAGS 0x0004u
     45 
     46 /* ---- env vtables --------------------------------------------------- */
     47 
     48 static void* heap_alloc(KitHeap* h, size_t n, size_t a) {
     49   (void)h;
     50   (void)a;
     51   return n ? malloc(n) : NULL;
     52 }
     53 static void* heap_realloc(KitHeap* h, void* p, size_t o, size_t n, size_t a) {
     54   (void)h;
     55   (void)o;
     56   (void)a;
     57   return realloc(p, n);
     58 }
     59 static void heap_free(KitHeap* h, void* p, size_t n) {
     60   (void)h;
     61   (void)n;
     62   free(p);
     63 }
     64 static KitHeap g_heap = {heap_alloc, heap_realloc, heap_free, NULL};
     65 
     66 static void diag_emit(KitDiagSink* s, KitDiagKind k, KitSrcLoc loc,
     67                       const char* fmt, va_list ap) {
     68   static const char* names[] = {"note", "warning", "error", "fatal"};
     69   (void)s;
     70   (void)loc;
     71   fprintf(stderr, "%s: ", names[k]);
     72   vfprintf(stderr, fmt, ap);
     73   fputc('\n', stderr);
     74 }
     75 static KitDiagSink g_diag = {diag_emit, NULL, 0, 0};
     76 static KitContext g_ctx;
     77 
     78 static int g_failures;
     79 static const char* g_case = "?";
     80 #define EXPECT(cond, ...)                                               \
     81   do {                                                                  \
     82     if (!(cond)) {                                                      \
     83       fprintf(stderr, "FAIL [%s] %s:%d: ", g_case, __FILE__, __LINE__); \
     84       fprintf(stderr, __VA_ARGS__);                                     \
     85       fputc('\n', stderr);                                              \
     86       g_failures++;                                                     \
     87     }                                                                   \
     88   } while (0)
     89 
     90 /* ---- target / compiler ------------------------------------------- */
     91 
     92 static void target_windows(KitTargetSpec* t, KitArchKind arch) {
     93   memset(t, 0, sizeof *t);
     94   t->arch = arch;
     95   t->os = KIT_OS_WINDOWS;
     96   t->obj = KIT_OBJ_COFF;
     97   t->ptr_size = 8;
     98   t->ptr_align = 8;
     99   t->big_endian = false;
    100   t->pic = KIT_PIC_PIE;
    101   t->code_model = KIT_CM_SMALL;
    102 }
    103 
    104 static Compiler* make_compiler(const KitTargetSpec* t) {
    105   KitTargetOptions opts;
    106   KitTarget* target = NULL;
    107   KitCompiler* cc = NULL;
    108   memset(&opts, 0, sizeof opts);
    109   opts.spec = *t;
    110   if (kit_target_new(&g_ctx, &opts, &target) != KIT_OK || !target) return NULL;
    111   if (kit_compiler_new(target, &g_ctx, &cc) != KIT_OK || !cc) {
    112     kit_target_free(target);
    113     return NULL;
    114   }
    115   return (Compiler*)cc;
    116 }
    117 
    118 static void free_compiler(Compiler* c) {
    119   const KitTarget* target;
    120   if (!c) return;
    121   target = kit_compiler_target((KitCompiler*)c);
    122   kit_compiler_free((KitCompiler*)c);
    123   kit_target_free((KitTarget*)target);
    124 }
    125 
    126 /* ---- short-import shim builder ------------------------------------ */
    127 
    128 static void build_short_import(uint8_t buf[SHIM_TOTAL_LEN], uint16_t machine) {
    129   memset(buf, 0, SHIM_TOTAL_LEN);
    130   buf[2] = (uint8_t)(COFF_SHIMP_SIG2 & 0xFF); /* Sig2 = 0xFFFF */
    131   buf[3] = (uint8_t)((COFF_SHIMP_SIG2 >> 8) & 0xFF);
    132   buf[6] = (uint8_t)(machine & 0xFF);
    133   buf[7] = (uint8_t)((machine >> 8) & 0xFF);
    134   buf[12] = (uint8_t)(SHIM_DATA_LEN & 0xFFu); /* SizeOfData */
    135   buf[13] = (uint8_t)((SHIM_DATA_LEN >> 8) & 0xFFu);
    136   buf[18] = (uint8_t)(COFF_SHIMP_TYPEFLAGS & 0xFF);
    137   buf[19] = (uint8_t)((COFF_SHIMP_TYPEFLAGS >> 8) & 0xFF);
    138   memcpy(buf + SHIM_HEADER_SIZE, SHIM_SYM_CSTR, SHIM_SYM_NUL_LEN);
    139   memcpy(buf + SHIM_HEADER_SIZE + SHIM_SYM_NUL_LEN, SHIM_DLL_CSTR,
    140          SHIM_DLL_NUL_LEN);
    141 }
    142 
    143 /* ---- program ObjBuilder ------------------------------------------- */
    144 
    145 /* mainCRTStartup body: a single return. The exact encoding is irrelevant
    146  * to the reader; differ per arch only so the linker sees plausible code. */
    147 static const uint8_t TEXT_X64[1] = {0xc3};                    /* ret */
    148 static const uint8_t TEXT_AA64[4] = {0xc0, 0x03, 0x5f, 0xd6}; /* ret */
    149 
    150 static ObjBuilder* build_program(Compiler* c, KitArchKind arch) {
    151   ObjBuilder* ob = obj_new(c);
    152   Pool* p = c->global;
    153   Sym text_name = pool_intern_slice(p, SLICE_LIT(".text"));
    154   Sym data_name = pool_intern_slice(p, SLICE_LIT(".data"));
    155   Sym main_name = pool_intern_slice(p, SLICE_LIT("mainCRTStartup"));
    156   Sym exit_name = pool_intern_slice(p, SLICE_LIT(SHIM_SYM_CSTR));
    157   const uint8_t* text = arch == KIT_ARCH_X86_64 ? TEXT_X64 : TEXT_AA64;
    158   u32 text_len =
    159       arch == KIT_ARCH_X86_64 ? (u32)sizeof TEXT_X64 : (u32)sizeof TEXT_AA64;
    160   ObjSecId tsec = obj_section(ob, text_name, SEC_TEXT, SF_ALLOC | SF_EXEC, 16);
    161   ObjSecId dsec = obj_section(ob, data_name, SEC_DATA, SF_ALLOC | SF_WRITE, 8);
    162   ObjSymId exit_sym;
    163   uint8_t zeros[8] = {0};
    164 
    165   obj_write(ob, tsec, text, text_len);
    166   obj_symbol(ob, main_name, SB_GLOBAL, SK_FUNC, tsec, 0, text_len);
    167 
    168   /* .data: an 8-byte absolute pointer to the imported ExitProcess. The
    169    * R_ABS64 both forces ExitProcess to be imported and (in a PIE) yields a
    170    * base relocation, so the reader's import + base-reloc paths both run. */
    171   exit_sym = obj_symbol(ob, exit_name, SB_GLOBAL, SK_UNDEF, OBJ_SEC_NONE, 0, 0);
    172   obj_write(ob, dsec, zeros, sizeof zeros);
    173   obj_reloc(ob, dsec, 0, R_ABS64, exit_sym, 0);
    174 
    175   obj_finalize(ob);
    176   return ob;
    177 }
    178 
    179 /* Link a PE image and copy the emitted bytes into a fresh malloc buffer
    180  * (so the reader runs fully independent of the producing compiler).
    181  * Returns NULL on failure. */
    182 static uint8_t* link_pe(Compiler* c, KitArchKind arch, uint16_t machine,
    183                         size_t* out_len) {
    184   ObjBuilder* prog = build_program(c, arch);
    185   uint8_t shim[SHIM_TOTAL_LEN];
    186   Linker* l;
    187   LinkImage* img;
    188   KitWriter* w = NULL;
    189   const uint8_t* bytes;
    190   size_t n = 0;
    191   uint8_t* copy = NULL;
    192 
    193   build_short_import(shim, machine);
    194 
    195   l = link_new(c);
    196   if (!l) return NULL;
    197   link_add_obj(l, prog);
    198   (void)link_add_obj_bytes(l, "ExitProcess.lib-member", shim, SHIM_TOTAL_LEN);
    199   link_set_entry(l, KIT_SLICE_LIT("mainCRTStartup"));
    200   link_set_pie(l, 1);
    201   link_set_emit_static_exe(l, 1);
    202 
    203   img = link_resolve(l);
    204   if (!img) {
    205     link_free(l);
    206     return NULL;
    207   }
    208   if (kit_writer_mem(&g_heap, &w) != KIT_OK || !w) {
    209     link_image_free(img);
    210     link_free(l);
    211     return NULL;
    212   }
    213   link_emit_image_writer(img, w);
    214   bytes = kit_writer_mem_bytes(w, &n);
    215   if (bytes && n) {
    216     copy = (uint8_t*)malloc(n);
    217     if (copy) memcpy(copy, bytes, n);
    218   }
    219   *out_len = n;
    220   kit_writer_close(w);
    221   link_image_free(img);
    222   link_free(l);
    223   return copy;
    224 }
    225 
    226 /* ---- the round-trip assertions ------------------------------------ */
    227 
    228 static void run_case(const char* name, KitArchKind arch, uint16_t machine) {
    229   Compiler* c;
    230   uint8_t* pe;
    231   size_t pe_len = 0;
    232   KitTargetSpec t;
    233   KitObjFile* f = NULL;
    234   KitSlice input;
    235   KitObjImageInfo info;
    236   KitStatus st;
    237 
    238   g_case = name;
    239   target_windows(&t, arch);
    240   c = make_compiler(&t);
    241   if (!c) {
    242     EXPECT(0, "make_compiler failed");
    243     return;
    244   }
    245   if (setjmp(c->panic)) {
    246     EXPECT(0, "panic while linking PE");
    247     compiler_run_cleanups(c);
    248     free_compiler(c);
    249     return;
    250   }
    251   pe = link_pe(c, arch, machine, &pe_len);
    252   free_compiler(c);
    253   if (!pe || !pe_len) {
    254     EXPECT(0, "link_pe produced no bytes");
    255     free(pe);
    256     return;
    257   }
    258 
    259   /* Detection should route the image to COFF/Windows. */
    260   EXPECT(kit_detect_fmt(pe, pe_len) == KIT_BIN_PE, "detect_fmt != KIT_BIN_PE");
    261 
    262   input.data = pe;
    263   input.len = pe_len;
    264   st = kit_obj_open(&g_ctx, KIT_SLICE_LIT("image.exe"), &input, &f);
    265   EXPECT(st == KIT_OK && f, "kit_obj_open failed (st=%d)", (int)st);
    266   if (!f) {
    267     free(pe);
    268     return;
    269   }
    270 
    271   EXPECT(kit_obj_kind(f) == KIT_OBJ_KIND_EXEC, "kind != EXEC (%d)",
    272          (int)kit_obj_kind(f));
    273 
    274   st = kit_obj_image_info(f, &info);
    275   EXPECT(st == KIT_OK, "image_info failed");
    276   EXPECT(info.image_base != 0, "image_base == 0");
    277   EXPECT(info.entry > info.image_base, "entry (%llu) not above base (%llu)",
    278          (unsigned long long)info.entry, (unsigned long long)info.image_base);
    279 
    280   /* Sections + a .text section. */
    281   {
    282     KitObjSection idx;
    283     EXPECT(kit_obj_nsections(f) > 0, "no sections");
    284     EXPECT(kit_obj_section_by_name(f, KIT_SLICE_LIT(".text"), &idx) == KIT_OK,
    285            ".text section not found");
    286   }
    287 
    288   /* Segments: at least one, with an executable one present. */
    289   {
    290     KitObjSegIter* it = NULL;
    291     KitObjSegInfo seg;
    292     int nseg = 0, nexec = 0;
    293     EXPECT(kit_obj_segiter_new(f, &it) == KIT_OK, "segiter_new failed");
    294     while (it && kit_obj_segiter_next(it, &seg) == KIT_ITER_ITEM) {
    295       ++nseg;
    296       if (seg.perms & KIT_SEG_X) ++nexec;
    297       EXPECT(seg.vaddr >= info.image_base, "segment vaddr below image base");
    298     }
    299     kit_obj_segiter_free(it);
    300     EXPECT(nseg > 0, "no segments");
    301     EXPECT(nexec > 0, "no executable segment");
    302   }
    303 
    304   /* Dependency KERNEL32.dll carrying the ExitProcess import. */
    305   {
    306     KitObjDepIter* it = NULL;
    307     KitObjDepInfo dep;
    308     int found_dll = 0, found_imp = 0;
    309     EXPECT(kit_obj_depiter_new(f, &it) == KIT_OK, "depiter_new failed");
    310     while (it && kit_obj_depiter_next(it, &dep) == KIT_ITER_ITEM) {
    311       if (kit_slice_eq_cstr(dep.name, SHIM_DLL_CSTR)) {
    312         found_dll = 1;
    313         for (uint32_t i = 0; i < dep.nimports; ++i)
    314           if (kit_slice_eq_cstr(dep.imports[i], SHIM_SYM_CSTR)) found_imp = 1;
    315       }
    316     }
    317     kit_obj_depiter_free(it);
    318     EXPECT(found_dll, "KERNEL32.dll dependency not found");
    319     EXPECT(found_imp, "ExitProcess import not listed under KERNEL32.dll");
    320   }
    321 
    322   /* Dynamic symbol ExitProcess (undefined import). */
    323   {
    324     KitObjSymIter* it = NULL;
    325     KitObjSymInfo sym;
    326     int found = 0;
    327     EXPECT(kit_obj_dynsymiter_new(f, &it) == KIT_OK, "dynsymiter_new failed");
    328     while (it && kit_obj_symiter_next(it, &sym) == KIT_ITER_ITEM)
    329       if (kit_slice_eq_cstr(sym.name, SHIM_SYM_CSTR)) found = 1;
    330     kit_obj_symiter_free(it);
    331     EXPECT(found, "ExitProcess not in dynamic symbols");
    332   }
    333 
    334   /* Raw escape hatch: 16 data dirs + subsystem + dllchars; IMPORT set. */
    335   {
    336     KitObjImageRawIter* it = NULL;
    337     KitObjImageRaw r;
    338     int ndatadir = 0, have_subsys = 0, have_dllchars = 0;
    339     uint64_t import_rva = 0;
    340     EXPECT(kit_obj_image_rawiter_new(f, &it) == KIT_OK, "rawiter_new failed");
    341     while (it && kit_obj_image_rawiter_next(it, &r) == KIT_ITER_ITEM) {
    342       if (r.tag < 16) {
    343         ++ndatadir;
    344         if (r.tag == 1) import_rva = r.value; /* IMAGE_DIRECTORY_ENTRY_IMPORT */
    345       } else if (r.tag == KIT_OBJ_RAW_PE_SUBSYSTEM) {
    346         have_subsys = 1;
    347         EXPECT(r.value == 3, "subsystem != WINDOWS_CUI (%llu)",
    348                (unsigned long long)r.value);
    349       } else if (r.tag == KIT_OBJ_RAW_PE_DLLCHARS) {
    350         have_dllchars = 1;
    351       }
    352     }
    353     kit_obj_image_rawiter_free(it);
    354     EXPECT(ndatadir == 16, "expected 16 data directories, saw %d", ndatadir);
    355     EXPECT(have_subsys, "subsystem raw entry missing");
    356     EXPECT(have_dllchars, "dllcharacteristics raw entry missing");
    357     EXPECT(import_rva != 0, "IMPORT data directory RVA is zero");
    358   }
    359 
    360   /* Base relocations: the PIE .data absolute pointer needs at least one. */
    361   {
    362     KitObjRelocIter* it = NULL;
    363     KitObjReloc rel;
    364     int n = 0;
    365     EXPECT(kit_obj_dynreliter_new(f, &it) == KIT_OK, "dynreliter_new failed");
    366     while (it && kit_obj_reliter_next(it, &rel) == KIT_ITER_ITEM) ++n;
    367     kit_obj_reliter_free(it);
    368     EXPECT(n > 0, "no base relocations for PIE image");
    369   }
    370 
    371   kit_obj_free(f);
    372   free(pe);
    373 }
    374 
    375 int main(int argc, char** argv) {
    376   memset(&g_ctx, 0, sizeof g_ctx);
    377   g_ctx.heap = &g_heap;
    378   g_ctx.diag = &g_diag;
    379   g_ctx.now = -1;
    380 
    381   /* Optional: regenerate the committed x86_64 PE objdump fixture (no
    382    * asserts). Used to produce test/objdump/x86_64-windows/cases/pe-image.exe
    383    * from this same in-memory link, so the non-gated objdump golden is
    384    * reproducible. */
    385   if (argc > 1) {
    386     KitTargetSpec t;
    387     Compiler* c;
    388     target_windows(&t, KIT_ARCH_X86_64);
    389     c = make_compiler(&t);
    390     if (c && setjmp(c->panic) == 0) {
    391       size_t n = 0;
    392       uint8_t* pe = link_pe(c, KIT_ARCH_X86_64, 0x8664u, &n);
    393       if (pe && n) {
    394         FILE* fp = fopen(argv[1], "wb");
    395         if (fp) {
    396           fwrite(pe, 1, n, fp);
    397           fclose(fp);
    398         }
    399         fprintf(stderr, "wrote %zu bytes to %s\n", n, argv[1]);
    400       }
    401       free(pe);
    402     }
    403     free_compiler(c);
    404     return 0;
    405   }
    406 
    407   run_case("x86_64-windows", KIT_ARCH_X86_64, 0x8664u);
    408   run_case("aarch64-windows", KIT_ARCH_ARM_64, 0xAA64u);
    409 
    410   if (g_failures) {
    411     fprintf(stderr, "FAILED %d assertion(s)\n", g_failures);
    412     return 1;
    413   }
    414   fprintf(stderr, "OK pe-image-read\n");
    415   return 0;
    416 }