diff --git a/src/linker/lnk_symbol_table.c b/src/linker/lnk_symbol_table.c index 83ab62384..98178cd17 100644 --- a/src/linker/lnk_symbol_table.c +++ b/src/linker/lnk_symbol_table.c @@ -148,6 +148,28 @@ lnk_error_multiply_defined_symbol(LNK_Symbol *dst, LNK_Symbol *src) lnk_error_obj(LNK_Error_MultiplyDefinedSymbol, dst_ref.obj, "symbol \"%S\" (No. %#x) is multiply defined in %S (No. %#x)", dst->name, dst_ref.symbol_idx, src_ref.obj->path, src_ref.symbol_idx); } +internal B32 +lnk_symbol_has_nonzero_prefix(LNK_Obj *obj, COFF_ParsedSymbol symbol) +{ + if (symbol.value == 0) { + return 0; + } + COFF_SectionHeader *section = lnk_coff_section_header_from_section_number(obj, symbol.section_number); + if (section->flags & COFF_SectionFlag_CntUninitializedData) { + return 0; + } + String8 data = lnk_obj_section_data_from_number(obj, symbol.section_number); + if (symbol.value > data.size) { + return 0; + } + for (U32 i = 0; i < symbol.value; i += 1) { + if (data.str[i] != 0) { + return 1; + } + } + return 0; +} + internal B32 lnk_can_replace_symbol(LNK_Symbol *dst, LNK_Symbol *src) { @@ -325,10 +347,17 @@ lnk_can_replace_symbol(LNK_Symbol *dst, LNK_Symbol *src) if (src_select == dst_select) { switch (src_select) { - case COFF_ComdatSelect_Null: - case COFF_ComdatSelect_Any: { - can_replace = lnk_obj_is_before(src_obj, dst_obj); - } break; + case COFF_ComdatSelect_Null: { + can_replace = lnk_obj_is_before(src_obj, dst_obj); + } break; + case COFF_ComdatSelect_Any: { + // LLVM emits prefix data before the function entry point (the symbol + // value). AutoRTFM uses a non-zero prefix to locate the closed function. + // Prefer it over an uninstrumented copy, regardless of input/insert order. + B32 src_has_prefix = lnk_symbol_has_nonzero_prefix(src_obj, src_parsed); + B32 dst_has_prefix = lnk_symbol_has_nonzero_prefix(dst_obj, dst_parsed); + can_replace = src_has_prefix != dst_has_prefix ? src_has_prefix : lnk_obj_is_before(src_obj, dst_obj); + } break; case COFF_ComdatSelect_NoDuplicates: { lnk_error_multiply_defined_symbol(dst, src); } break; diff --git a/src/linker/tests/linker_tests.c b/src/linker/tests/linker_tests.c index 55fae91d6..96f744045 100644 --- a/src/linker/tests/linker_tests.c +++ b/src/linker/tests/linker_tests.c @@ -3253,6 +3253,65 @@ TEST(image_base) T_Ok(str8_match(text_data, str8_array_fixed(expected_text), 0)); } +TEST(comdat_any_nonzero_prefix) +{ + // The symbol value is the prefix length, not the length of the function. + String8 plain = str8_lit("\xB8\x01\x00\x00\x00\xC3"); + String8 zero = str8_lit("\x00\x00\x00\x00\xB8\x02\x00\x00\x00\xC3"); + String8 nonzero = str8_lit("\x00\x00\x00\xA5\xB8\x03\x00\x00\x00\xC3"); + String8 nonzero2 = str8_lit("\x5A\x00\x00\x00\xB8\x04\x00\x00\x00\xC3"); + struct { + String8 data[2]; + U32 value[2]; + COFF_ComdatSelectType selection[2]; + S32 winner; // -1 means the first input must win. + } cases[] = { + {{plain, nonzero}, {0, 4}, {COFF_ComdatSelect_Any, COFF_ComdatSelect_Any}, 1}, + {{zero, nonzero}, {4, 4}, {COFF_ComdatSelect_Any, COFF_ComdatSelect_Any}, 1}, + {{nonzero, nonzero2}, {4, 4}, {COFF_ComdatSelect_Any, COFF_ComdatSelect_Any}, -1}, + {{plain, zero}, {0, 4}, {COFF_ComdatSelect_Any, COFF_ComdatSelect_Any}, -1}, + {{plain, plain}, {0, 0}, {COFF_ComdatSelect_Any, COFF_ComdatSelect_Any}, -1}, + {{zero, zero}, {4, 4}, {COFF_ComdatSelect_Any, COFF_ComdatSelect_Any}, -1}, + // Out-of-range values must not turn function bytes into a valid prefix. + {{plain, nonzero}, {7, 4}, {COFF_ComdatSelect_Any, COFF_ComdatSelect_Any}, 1}, + {{plain, nonzero}, {max_U32, 4}, {COFF_ComdatSelect_Any, COFF_ComdatSelect_Any}, 1}, + // Equal-sized Largest/SameSize and Any promoted to Largest retain input order. + {{zero, nonzero}, {4, 4}, {COFF_ComdatSelect_Largest, COFF_ComdatSelect_Largest}, -1}, + {{zero, nonzero}, {4, 4}, {COFF_ComdatSelect_SameSize, COFF_ComdatSelect_SameSize}, -1}, + {{zero, nonzero}, {4, 4}, {COFF_ComdatSelect_Any, COFF_ComdatSelect_Largest}, -1}, + }; + for (U64 case_idx = 0; case_idx < ArrayCount(cases); case_idx += 1) { + for (U32 obj_idx = 0; obj_idx < 2; obj_idx += 1) { + T_Ok(t_write_def_obj(obj_idx ? "second.obj" : "first.obj", (T_COFF_DefObj){ + .machine = T_COFF_DefSetMachine(X64), + .sections = (T_COFF_DefSection[]){ + {"text", ".text", cases[case_idx].data[obj_idx], .flags = "rx:code", + .raw_flags = COFF_SectionFlag_LnkCOMDAT}, {0}}, + .symbols = (T_COFF_DefSymbol[]){ + T_COFF_DefSymbol_Secdef("text", cases[case_idx].selection[obj_idx]), + T_COFF_DefSymbol_ExternFunc("entry", "text", cases[case_idx].value[obj_idx]), {0}}, + })); + } + for (U32 reverse = 0; reverse < 2; reverse += 1) { + for (U32 workers = 1; workers <= 4; workers += 3) { + t_invoke_linkerf("/subsystem:console /entry:entry /nodefaultlib /opt:ref,noicf /rad_workers:%u /out:prefix.exe %s", + workers, reverse ? "second.obj first.obj" : "first.obj second.obj"); + T_Ok(g_last_exit_code == 0); + String8 image = t_read_file(arena, str8_lit("prefix.exe")); + PE_BinInfo bin = pe_bin_info_from_data(arena, image); + U32 winner = cases[case_idx].winner < 0 ? reverse : (U32)cases[case_idx].winner; + U64 entry_off = pe_foff_from_voff(image, &bin, bin.entry_point); + U32 prefix_size = cases[case_idx].value[winner]; + T_Ok(entry_off >= prefix_size); + if (entry_off >= prefix_size) { + String8 actual = str8_substr(image, r1u64s(entry_off - prefix_size, cases[case_idx].data[winner].size)); + T_Ok(str8_match(actual, cases[case_idx].data[winner], 0)); + } + } + } + } +} + TEST(comdat_any) { T_Ok(t_write_def_obj("1.obj", (T_COFF_DefObj){