Skip to content

Commit

Permalink
Merge pull request #20 from ChrisDenton/comdat
Browse files Browse the repository at this point in the history 
Add an optional COMDAT record for the `.idata$3` section
  • Loading branch information
@bjorn3
bjorn3 authored Aug 16, 2024
2 parents 504d89a + 32d9046 commit 002e259
Show file tree
Hide file tree
Showing 2 changed files with 119 additions and 18 deletions.
56 changes: 42 additions & 14 deletions src/coff_import_file.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,13 @@ use std::path::PathBuf;
use std::str::from_utf8;

use object::pe::{
ImageFileHeader, ImageImportDescriptor, ImageRelocation, ImageSectionHeader, ImageSymbol,
ImportObjectHeader, IMAGE_FILE_32BIT_MACHINE, IMAGE_REL_AMD64_ADDR32NB,
IMAGE_REL_ARM64_ADDR32NB, IMAGE_REL_ARM_ADDR32NB, IMAGE_REL_I386_DIR32NB,
IMAGE_SCN_ALIGN_2BYTES, IMAGE_SCN_ALIGN_4BYTES, IMAGE_SCN_ALIGN_8BYTES,
IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_LNK_INFO, IMAGE_SCN_LNK_REMOVE, IMAGE_SCN_MEM_READ,
IMAGE_SCN_MEM_WRITE, IMAGE_SYM_CLASS_EXTERNAL, IMAGE_SYM_CLASS_NULL, IMAGE_SYM_CLASS_SECTION,
ImageAuxSymbolSection, ImageFileHeader, ImageImportDescriptor, ImageRelocation,
ImageSectionHeader, ImageSymbol, ImportObjectHeader, IMAGE_COMDAT_SELECT_ANY,
IMAGE_FILE_32BIT_MACHINE, IMAGE_REL_AMD64_ADDR32NB, IMAGE_REL_ARM64_ADDR32NB,
IMAGE_REL_ARM_ADDR32NB, IMAGE_REL_I386_DIR32NB, IMAGE_SCN_ALIGN_2BYTES, IMAGE_SCN_ALIGN_4BYTES,
IMAGE_SCN_ALIGN_8BYTES, IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_LNK_COMDAT,
IMAGE_SCN_LNK_INFO, IMAGE_SCN_LNK_REMOVE, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE,
IMAGE_SYM_CLASS_EXTERNAL, IMAGE_SYM_CLASS_NULL, IMAGE_SYM_CLASS_SECTION,
IMAGE_SYM_CLASS_STATIC, IMAGE_SYM_CLASS_WEAK_EXTERNAL, IMAGE_WEAK_EXTERN_SEARCH_ALIAS,
};
use object::pod::bytes_of;
Expand Down Expand Up @@ -463,11 +464,14 @@ impl<'a> ObjectFactory<'a> {
/// Creates a NULL import descriptor. This is a small object file whcih
/// contains a NULL import descriptor. It is used to terminate the imports
/// from a specific DLL.
fn create_null_import_descriptor(&self) -> Result<NewArchiveMember<'_>> {
///
/// If `comdat` is set to true, the NULL import descriptor's section (`.idata$3`) will be
/// turned into a COMDAT section.
fn create_null_import_descriptor(&self, comdat: bool) -> Result<NewArchiveMember<'_>> {
let mut buffer = Vec::new();

const NUMBER_OF_SECTIONS: usize = 1;
const NUMBER_OF_SYMBOLS: usize = 1;
let number_of_symbols = if comdat { 3 } else { 1 };

// COFF Header
let header = ImageFileHeader {
Expand All @@ -477,7 +481,7 @@ impl<'a> ObjectFactory<'a> {
pointer_to_symbol_table: u32!((size_of::<ImageFileHeader>() + (NUMBER_OF_SECTIONS * size_of::<ImageSectionHeader>()) +
// .idata$3
size_of::<ImageImportDescriptor>()).try_into().unwrap()),
number_of_symbols: u32!(NUMBER_OF_SYMBOLS.try_into().unwrap()),
number_of_symbols: u32!(number_of_symbols.try_into().unwrap()),
size_of_optional_header: u16!(0),
characteristics: u16!(if self.is_64_bit() {
0
Expand Down Expand Up @@ -506,6 +510,7 @@ impl<'a> ObjectFactory<'a> {
| IMAGE_SCN_CNT_INITIALIZED_DATA
| IMAGE_SCN_MEM_READ
| IMAGE_SCN_MEM_WRITE
| if comdat { IMAGE_SCN_LNK_COMDAT } else { 0 }
),
}];
buffer.write_all(bytes_of(&section_table))?;
Expand All @@ -521,16 +526,38 @@ impl<'a> ObjectFactory<'a> {
buffer.write_all(bytes_of(&import_descriptor))?;

// Symbol Table
let mut symbol_table: [_; NUMBER_OF_SYMBOLS] = [ImageSymbol {
if comdat {
let symbol = ImageSymbol {
name: *b".idata$3",
value: u32!(0),
section_number: u16!(1),
typ: u16!(0),
storage_class: IMAGE_SYM_CLASS_STATIC,
number_of_aux_symbols: 1,
};
let aux = ImageAuxSymbolSection {
length: u32!(size_of::<ImageImportDescriptor>().try_into().unwrap()),
number_of_relocations: u16!(0),
number_of_linenumbers: u16!(0),
check_sum: u32!(0),
selection: IMAGE_COMDAT_SELECT_ANY,
number: u16!(0),
reserved: 0,
high_number: u16!(0),
};
buffer.write_all(bytes_of(&symbol))?;
buffer.write_all(bytes_of(&aux))?;
}
let mut null_descriptor_symbol = ImageSymbol {
name: [0; 8],
value: u32!(0),
section_number: u16!(1),
typ: u16!(0),
storage_class: IMAGE_SYM_CLASS_EXTERNAL,
number_of_aux_symbols: 0,
}];
set_name_to_string_table_entry(&mut symbol_table[0], size_of::<u32>());
buffer.write_all(bytes_of(&symbol_table))?;
};
set_name_to_string_table_entry(&mut null_descriptor_symbol, size_of::<u32>());
buffer.write_all(bytes_of(&null_descriptor_symbol))?;

// String Table
write_string_table(&mut buffer, &[NULL_IMPORT_DESCRIPTOR_SYMBOL_NAME])?;
Expand Down Expand Up @@ -826,6 +853,7 @@ pub fn write_import_library<W: Write + Seek>(
exports: &[COFFShortExport],
machine: MachineTypes,
mingw: bool,
comdat: bool,
) -> Result<()> {
let native_machine = if machine == MachineTypes::ARM64EC {
MachineTypes::ARM64
Expand All @@ -838,7 +866,7 @@ pub fn write_import_library<W: Write + Seek>(

members.push(of.create_import_descriptor()?);

members.push(of.create_null_import_descriptor()?);
members.push(of.create_null_import_descriptor(comdat)?);

members.push(of.create_null_thunk()?);

Expand Down
81 changes: 77 additions & 4 deletions tests/import_library.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,10 @@ use std::process::Command;

use ar_archive_writer::{ArchiveKind, COFFShortExport, MachineTypes};
use common::{create_archive_with_ar_archive_writer, create_archive_with_llvm_ar};
use object::coff::CoffHeader;
use object::pe::ImageFileHeader;
use object::read::archive::ArchiveFile;
use object::{Architecture, SubArchitecture};
use object::{bytes_of, Architecture, Object, SubArchitecture, U32Bytes};
use pretty_assertions::assert_eq;

mod common;
Expand Down Expand Up @@ -99,6 +101,7 @@ fn create_import_library_with_ar_archive_writer(
temp_dir: &Path,
machine_type: MachineTypes,
mingw: bool,
comdat: bool,
) -> Vec<u8> {
let mut output_bytes = Cursor::new(Vec::new());
ar_archive_writer::write_import_library(
Expand All @@ -107,6 +110,7 @@ fn create_import_library_with_ar_archive_writer(
&get_members(machine_type),
machine_type,
mingw,
comdat,
)
.unwrap();

Expand All @@ -129,7 +133,7 @@ fn compare_to_lib() {
let temp_dir = common::create_tmp_dir("import_library_compare_to_lib");

let archive_writer_bytes =
create_import_library_with_ar_archive_writer(&temp_dir, machine_type, false);
create_import_library_with_ar_archive_writer(&temp_dir, machine_type, false, false);

let llvm_lib_bytes = {
let machine_arg = match machine_type {
Expand Down Expand Up @@ -164,6 +168,11 @@ fn compare_to_lib() {
llvm_lib_bytes, archive_writer_bytes,
"Import library differs. Machine type: {machine_type:?}",
);

compare_comdat(
&create_import_library_with_ar_archive_writer(&temp_dir, machine_type, false, true),
&llvm_lib_bytes,
);
}
}

Expand All @@ -178,7 +187,7 @@ fn compare_to_dlltool() {
let temp_dir = common::create_tmp_dir("import_library_compare_to_dlltool");

let archive_writer_bytes =
create_import_library_with_ar_archive_writer(&temp_dir, machine_type, true);
create_import_library_with_ar_archive_writer(&temp_dir, machine_type, true, false);

let llvm_lib_bytes = {
let machine_arg = match machine_type {
Expand Down Expand Up @@ -215,6 +224,11 @@ fn compare_to_dlltool() {
llvm_lib_bytes, archive_writer_bytes,
"Import library differs. Machine type: {machine_type:?}",
);

compare_comdat(
&create_import_library_with_ar_archive_writer(&temp_dir, machine_type, true, true),
&llvm_lib_bytes,
);
}
}

Expand All @@ -237,10 +251,11 @@ fn wrap_in_archive() {
let mut import_lib_bytes = Cursor::new(Vec::new());
ar_archive_writer::write_import_library(
&mut import_lib_bytes,
&temp_dir.join("MyLibrary.dll").to_string_lossy().to_string(),
&temp_dir.join("MyLibrary.dll").to_string_lossy(),
&get_members(machine_type),
machine_type,
false,
false,
)
.unwrap();
let import_lib_bytes = import_lib_bytes.into_inner();
Expand Down Expand Up @@ -281,3 +296,61 @@ fn wrap_in_archive() {
);
}
}

fn compare_comdat(archive_writer_bytes: &[u8], llvm_bytes: &[u8]) {
let archive_writer = ArchiveFile::parse(archive_writer_bytes).unwrap();
let llvm = ArchiveFile::parse(llvm_bytes).unwrap();

for (archive_member, llvm_member) in archive_writer.members().zip(llvm.members()) {
let archive_member = archive_member.unwrap();
let llvm_member = llvm_member.unwrap();

if archive_member.size() != llvm_member.size() {
// Ensure that the member header is the same except for the file size.
let mut llvm_file_header = *llvm_member.header().unwrap();
llvm_file_header.size = *b"163 ";
assert_eq!(
bytes_of(archive_member.header().unwrap()),
bytes_of(&llvm_file_header)
);

// Ensure that the LLVM generated object contains .idata$3
object::File::parse(llvm_member.data(llvm_bytes).unwrap())
.unwrap()
.section_by_name_bytes(b".idata$3")
.unwrap();

// Ensure the COFF file headers are the same except for the symbol count.
let llvm_data = llvm_member.data(llvm_bytes).unwrap();
let archive_data = archive_member.data(archive_writer_bytes).unwrap();
let mut offset = 0;
let mut header = *ImageFileHeader::parse(llvm_data, &mut offset).unwrap();
header.number_of_symbols = U32Bytes::from_bytes(3_u32.to_le_bytes());
assert_eq!(
&archive_data[..size_of::<ImageFileHeader>()],
bytes_of(&header)
);

// The rest of the object file will always be the same as `expected`
// for all import files no matter the platform.
let expected = [
46, 105, 100, 97, 116, 97, 36, 51, 0, 0, 0, 0, 0, 0, 0, 0, 20, 0, 0, 0, 60, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 64, 16, 48, 192, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 46, 105, 100, 97, 116, 97, 36, 51, 0, 0, 0, 0, 1,
0, 0, 0, 3, 1, 20, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0,
4, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 2, 0, 29, 0, 0, 0, 95, 95, 78, 85, 76, 76, 95,
73, 77, 80, 79, 82, 84, 95, 68, 69, 83, 67, 82, 73, 80, 84, 79, 82, 0,
];
assert_eq!(&archive_data[size_of::<ImageFileHeader>()..], &expected);
} else {
assert_eq!(
bytes_of(archive_member.header().unwrap()),
bytes_of(llvm_member.header().unwrap())
);
assert_eq!(
archive_member.data(archive_writer_bytes).unwrap(),
llvm_member.data(llvm_bytes).unwrap()
);
}
}
}

0 comments on commit 002e259

Please sign in to comment.