From 25710d39b7e5829a61cfb86a860226b4bdd8d0ca Mon Sep 17 00:00:00 2001
From: Nikita Smith <git.nikita.smith@gmail.com>
Date: Mon, 2 Jun 2025 13:39:32 -0700
Subject: [PATCH] test optional header fields

---
 src/linker/lnk.c      |  6 ++--
 src/pe/pe.c           |  1 +
 src/pe/pe.h           |  1 +
 src/torture/torture.c | 66 +++++++++++++++++++++++++++++++++++++++++--
 4 files changed, 69 insertions(+), 5 deletions(-)

diff --git a/src/linker/lnk.c b/src/linker/lnk.c
index faa8c3e7..17f2bbd3 100644
--- a/src/linker/lnk.c
+++ b/src/linker/lnk.c
@@ -2005,11 +2005,11 @@ lnk_build_win32_header(Arena *arena, LNK_SymbolTable *symtab, LNK_Config *config
     if (sect->flags & COFF_SectionFlag_CntUninitializedData) {
       sizeof_uninited_data += sect->vsize;
     }
-    if (sect->flags & COFF_SectionFlag_CntInitializedData) {
-      sizeof_inited_data += sect->vsize;
+    if ((sect->flags & COFF_SectionFlag_CntInitializedData) || (sect->flags & COFF_SectionFlag_CntCode)) {
+      sizeof_inited_data += sect->fsize;
     }
     if (sect->flags & COFF_SectionFlag_CntCode) { 
-      sizeof_code += sect->vsize;
+      sizeof_code += sect->fsize;
     }
     sizeof_image = Max(sizeof_image, sects.v[sect_idx]->voff + sects.v[sect_idx]->vsize);
   }
diff --git a/src/pe/pe.c b/src/pe/pe.c
index a8e2ea3a..93da7636 100644
--- a/src/pe/pe.c
+++ b/src/pe/pe.c
@@ -637,6 +637,7 @@ pe_bin_info_from_data(Arena *arena, String8 data)
     info.file_section_align  = file_section_align;
     info.section_count       = clamped_sec_count;
     info.symbol_count        = symbol_count;
+    info.optional_header_off = optional_range.min;
     info.section_table_range = rng_1u64(sec_array_off, sec_array_off + sizeof(COFF_SectionHeader) * clamped_sec_count);
     info.symbol_table_range  = rng_1u64(symbol_array_off, symbol_array_off + sizeof(COFF_Symbol16) * symbol_count);
     info.string_table_range  = rng_1u64(string_table_off, data.size);
diff --git a/src/pe/pe.h b/src/pe/pe.h
index cc225ffc..74dce633 100644
--- a/src/pe/pe.h
+++ b/src/pe/pe.h
@@ -1003,6 +1003,7 @@ struct PE_BinInfo
   U64                  file_section_align;
   U64                  section_count;
   U64                  symbol_count;
+  U64                  optional_header_off;
   Rng1U64              section_table_range;
   Rng1U64              symbol_table_range;
   Rng1U64              string_table_range;
diff --git a/src/torture/torture.c b/src/torture/torture.c
index 6c43492b..74746b02 100644
--- a/src/torture/torture.c
+++ b/src/torture/torture.c
@@ -599,7 +599,9 @@ t_simple_link_test(void)
   String8 main_obj;
   {
     COFF_ObjWriter  *obj_writer = coff_obj_writer_alloc(0, COFF_MachineType_X64);
-    COFF_ObjSection *text_sect  = coff_obj_writer_push_section(obj_writer, str8_lit(".text"), COFF_SectionFlag_CntCode, str8_array_fixed(text_payload));
+    COFF_ObjSection *text_sect  = coff_obj_writer_push_section(obj_writer, str8_lit(".text"), PE_TEXT_SECTION_FLAGS, str8_array_fixed(text_payload));
+    coff_obj_writer_push_section(obj_writer, str8_lit(".data"), PE_DATA_SECTION_FLAGS, str8_lit("qwe"));
+    coff_obj_writer_push_section(obj_writer, str8_lit(".bss"), PE_BSS_SECTION_FLAGS, str8(0, 5));
     coff_obj_writer_push_symbol_extern(obj_writer, str8_lit("my_entry"), 0, text_sect);
     main_obj = coff_obj_writer_serialize(scratch.arena, obj_writer);
     coff_obj_writer_release(&obj_writer);
@@ -626,7 +628,7 @@ t_simple_link_test(void)
   if (pe.is_pe32) {
     goto exit;
   }
-  if (pe.section_count != 1) {
+  if (pe.section_count != 3) {
     goto exit;
   }
   if (pe.arch != Arch_x64) {
@@ -670,11 +672,71 @@ t_simple_link_test(void)
     goto exit;
   }
 
+  COFF_SectionHeader *data_section = t_coff_section_header_from_name(string_table, section_table, pe.section_count, str8_lit(".data"));
+  if (data_section == 0) {
+    goto exit;
+  }
+
+  COFF_SectionHeader *bss_section = t_coff_section_header_from_name(string_table, section_table, pe.section_count, str8_lit(".bss"));
+  if (bss_section == 0) {
+    goto exit;
+  }
+
   String8 text_data = str8_substr(exe, rng_1u64(text_section->foff, text_section->foff + text_section->vsize));
   if (!str8_match(text_data, str8_array_fixed(text_payload), 0)) {
     goto exit;
   }
 
+  PE_OptionalHeader32Plus *opt = str8_deserial_get_raw_ptr(exe, pe.optional_header_off, sizeof(*opt));
+  if (opt->sizeof_code != text_section->fsize) {
+    goto exit;
+  }
+  if (opt->sizeof_inited_data != text_section->fsize + data_section->fsize) {
+    goto exit;
+  }
+  if (opt->sizeof_uninited_data != 0) {
+    goto exit;
+  }
+  if (opt->code_base != 0x1000) {
+    goto exit;
+  }
+  if (opt->image_base != 0x140000000) {
+    goto exit;
+  }
+  if (opt->major_os_ver != 6) {
+    goto exit;
+  }
+  if (opt->minor_os_ver != 0) {
+    goto exit;
+  }
+  if (opt->major_img_ver != 0) {
+    goto exit;
+  }
+  if (opt->minor_img_ver != 0) {
+    goto exit;
+  }
+  if (opt->major_subsystem_ver != 6) {
+    goto exit;
+  }
+  if (opt->minor_subsystem_ver != 0) {
+    goto exit;
+  }
+  if (opt->win32_version_value != 0) {
+    goto exit;
+  }
+  if (opt->sizeof_image != 0x3000) {
+    goto exit;
+  }
+  if (opt->sizeof_headers != 0x200) {
+    goto exit;
+  }
+  if (opt->dll_characteristics != 0x8120) {
+    goto exit;
+  }
+  if (opt->loader_flags != 0) {
+    goto exit;
+  }
+
   result = T_Result_Pass;
 exit:;
   scratch_end(scratch);