commit 21dcffbefdd97a0aeefaf4535d945460e177a288 Author: Matt Ehrnschwender Date: Fri May 30 15:57:20 2025 -0400 Initial release diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..719f698 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,4 @@ +[alias] +xtask = "run -p xtask --" +xtask-dist = "run -p xtask --features dist --" +dist = "xtask-dist dist" diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..9915c0c --- /dev/null +++ b/.editorconfig @@ -0,0 +1,16 @@ +root = true + +[*] +insert_final_newline = true +trim_trailing_whitespace = true + +[*.{rs,c}] +indent_size = 4 +intent_style = space + +[*.{toml,def,yaml,yml}] +indent_size = 2 +indent_style = space + +[Makefile] +indent_style = tab diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml new file mode 100644 index 0000000..647ec86 --- /dev/null +++ b/.github/workflows/ci.yaml @@ -0,0 +1,27 @@ +name: CI + +on: + push: + pull_request: + +env: + CARGO_TERM_COLOR: always + +jobs: + ci: + name: Run CI tasks + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, windows-latest] + + runs-on: ${{ matrix.os }} + steps: + - uses: actions/checkout@v4 + + - name: Install clippy and rustfmt + run: | + rustup update + rustup component add clippy rustfmt + + - run: cargo xtask ci diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml new file mode 100644 index 0000000..98240c5 --- /dev/null +++ b/.github/workflows/release.yaml @@ -0,0 +1,99 @@ +name: Release + +on: + push: + tags: + - 'v[0-9]+.[0-9]+.[0-9]+' + +env: + CARGO_TERM_COLOR: always + CARGO_PROFILE_RELEASE_LTO: fat + CARGO_PROFILE_RELEASE_STRIP: symbols + CARGO_BUILD_DEP_INFO_BASEDIR: "." + +concurrency: + group: "release" + cancel-in-progress: true + +jobs: + build-linux: + name: Build Linux dist artifacts + permissions: + actions: write + + env: + CARGO_TARGET_AARCH64_UNKNOWN_LINUX_GNU_LINKER: aarch64-linux-gnu-gcc + CARGO_TARGET_X86_64_UNKNOWN_LINUX_MUSL_RUSTFLAGS: -C target-feature=+crt-static + + + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Install system packages + run: | + sudo apt-get update -y + sudo apt-get install -y musl musl-dev musl-tools gcc-aarch64-linux-gnu libc6-arm64-cross + + - name: Install rust targets + run: rustup target add x86_64-unknown-linux-gnu x86_64-unknown-linux-musl aarch64-unknown-linux-gnu + + - name: Build dist artifacts + run: cargo dist --target x86_64-unknown-linux-gnu --target x86_64-unknown-linux-musl --target aarch64-unknown-linux-gnu + + - name: Upload dist artifacts + uses: actions/upload-artifact@v4 + with: + name: boflink-dist-linux + path: | + target/dist/*.tar.gz + target/dist/*.tar.gz.sha256 + retention-days: 1 + + build-windows: + name: Build Windows dist artifacts + permissions: + actions: write + + env: + RUSTFLAGS: -C target-feature=+crt-static + + runs-on: windows-latest + steps: + - uses: actions/checkout@v4 + + - name: Build dist artifacts + run: cargo dist + + - name: Upload dist artifacts + uses: actions/upload-artifact@v4 + with: + name: boflink-dist-windows + path: | + target/dist/*.zip + target/dist/*.zip.sha256 + retention-days: 1 + + release: + name: Create Release + needs: [build-linux, build-windows] + + permissions: + contents: write + + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Download dist artifacts + uses: actions/download-artifact@v4 + with: + path: dist + pattern: boflink-dist-* + merge-multiple: true + + - name: Create a new Github release + env: + GIT_TAG: ${{ github.ref_name }} + GH_TOKEN: ${{ github.token }} + run: gh release create $GIT_TAG -d -t "Boflink $GIT_TAG" -n "" --verify-tag ./dist/* diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..19bc651 --- /dev/null +++ b/.gitignore @@ -0,0 +1,13 @@ +/target +Cargo.lock +*.o +*.obj +*.d +*.bof +*.dot +*.svg +*.a +*.lib +*.dll +*.exe +*.exp diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..88b00fe --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,15 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased] + +## [0.1.0] - 2025-05-30 + +- Initial release + +[unreleased]: https://github.com/MEhrn00/boflink/compare/v0.1.0...HEAD +[0.1.0]: https://github.com/MEhrn00/boflink/releases/tag/v0.1.0 diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..b5b28a3 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,69 @@ +[workspace] +members = [ + "crates/*", + "crates/jamcrc/cli", + "xtask", +] + +[package] +name = "boflink" +version = "0.1.0" +authors = ["Matt Ehrnschwender "] +edition = "2024" +rust-version = "1.85" +description = """ +Linker for Beacon Object Files. +""" +readme = "README.md" +homepage = "https://github.com/MEhrn00/boflink" +repository = "https://github.com/MEhrn00/boflink" +license = "BSD-3-Clause" +publish = false + + +[dependencies] +anyhow = "1.0.92" +argfile = "0.2.1" +bitflags = "2.9.0" +bumpalo = "3.17.0" +clap-verbosity-flag = "3.0.2" +indexmap = "2.7.1" +jamcrc = { path = "crates/jamcrc" } +log = { version = "0.4.26", features = ["std"] } +num_enum = "0.7.3" +termcolor = "1.4.1" +thiserror = "2.0.11" +typed-arena = "2.0.2" + +[dependencies.clap] +version = "4.5.24" +default-features = false +features = ["std", "color", "help", "suggestions", "usage", "derive"] + +[dependencies.object] +version = "0.36.7" +default-features = false +features = ["archive", "coff", "write", "read"] + +[dev-dependencies] +coffyaml = { path = "crates/coffyaml" } +serde = "1" +serde_yml = "0.0.12" + +[lints.rust] +unsafe_code = "forbid" + +[profile.release] +debug = 1 + +[profile.release-lto] +inherits = "release" +opt-level = 3 +debug = "none" +strip = "symbols" +debug-assertions = false +overflow-checks = false +lto = "fat" +panic = "abort" +incremental = false +codegen-units = 1 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..2354b63 --- /dev/null +++ b/LICENSE @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2025, Matt Ehrnschwender + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + diff --git a/README.md b/README.md new file mode 100644 index 0000000..d37bf55 --- /dev/null +++ b/README.md @@ -0,0 +1,60 @@ +# Boflink + +[![GitHub License](https://img.shields.io/github/license/MEhrn00/boflink)](https://github.com/MEhrn00/boflink/blob/main/LICENSE) +[![GitHub Release](https://img.shields.io/github/v/release/MEhrn00/boflink)](https://github.com/MEhrn00/boflink/releases/latest) + +Linker for Beacon Object Files. + +- [Installation](#installation) +- [Usage](#usage) +- [Examples](#examples) + +## Installation +Requires: [Rust](https://www.rust-lang.org/tools/install) >=1.85.0 +```bash +rustc --version +``` + +### From Source +```bash +git clone https://github.com/MEhrn00/boflink.git +cd boflink +cargo xtask install + +# For an LTO build +cargo xtask install -p release-lto +``` + +## Usage +### Standalone +```bash +boflink [-o ] [options] ... +boflink -o mybof.bof -L/path/to/windows/libs -lkernel32 -ladvapi32 source.c object.o +``` + +### Using MinGW GCC on Linux +MinGW GCC can be used to invoke boflink using its configured link libraries and library search paths. + +```bash +x86_64-w64-mingw32-gcc -B ~/.local/libexec/boflink -fno-lto -nostartfiles ... +x86_64-w64-mingw32-gcc -B ~/.local/libexec/boflink -fno-lto -nostartfiles -o mybof.bof source.c object.o +``` + +### Using Clang on Linux +Clang can be used to invoke boflink using its configured link libraries and library search paths. + +```bash +clang --ld-path=/path/to/boflink --target=x86_64-windows-gnu -nostartfiles ... +clang --ld-path=/path/to/boflink --target=x86_64-windows-gnu -nostartfiles -o mybof.bof source.c object.o +``` + +### Using MSVC on Windows +Windows requires running the boflink executable in a Visual Studio Developer Console. + +```powershell +boflink ... +boflink -o mybof.bof object1.o object2.o -lkernel32 -ladvapi32 +``` + +## Examples +Additional examples can be found in the [examples/](examples/) directory. diff --git a/crates/coffyaml/Cargo.toml b/crates/coffyaml/Cargo.toml new file mode 100644 index 0000000..75f3a99 --- /dev/null +++ b/crates/coffyaml/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "coffyaml" +version = "0.1.0" +authors = ["Matt Ehrnschwender "] +edition = "2024" +description = """ +Library for generating COFFs from LLVM's yaml2obj YAML description files. +""" +readme = "README.md" +homepage = "https://github.com/MEhrn00/boflink/tree/main/crates/coffyaml" +repository = "https://github.com/MEhrn00/boflink" +publish = false + +[dependencies] +hex = { version = "0.4.3", features = ["serde"] } +indexmap = "2.8.0" +serde = { version = "1.0.219", features = ["derive"] } +serde_yml = "0.0.12" +thiserror = "2.0.12" +typed-arena = "2.0.2" + +[dependencies.object] +version = "0.36.7" +default-features = false +features = ["coff", "archive", "write", "read"] + +[lints.rust] +unsafe_code = "forbid" diff --git a/crates/coffyaml/README.md b/crates/coffyaml/README.md new file mode 100644 index 0000000..c8a7919 --- /dev/null +++ b/crates/coffyaml/README.md @@ -0,0 +1,7 @@ +# coffyaml +Library for converting COFFs and import libraries into yaml description files. + +These yaml files follow the LLVM [yaml2obj](https://llvm.org/docs/yaml2obj.html) format +with the addition of import libraries. + +Used for building linker test data. diff --git a/crates/coffyaml/src/archive/builder/armap.rs b/crates/coffyaml/src/archive/builder/armap.rs new file mode 100644 index 0000000..45bda2d --- /dev/null +++ b/crates/coffyaml/src/archive/builder/armap.rs @@ -0,0 +1,176 @@ +use std::collections::HashMap; + +use typed_arena::Arena; + +use super::{ + ArchiveMemberIndex, ArchiveMemberMetadata, MemberSize, longnames::ArchiveMemberName, + make_archive_member_buffer, +}; + +#[derive(Default)] +pub struct ArchiveMapBuilder { + /// The archive member indicies + indicies: Arena, + + /// The string table + string_table: Arena, +} + +impl ArchiveMapBuilder { + /// Adds a symbol to the archive map + pub fn add_symbol(&self, index: ArchiveMemberIndex, symbol: impl AsRef) { + self.indicies.alloc(index); + self.string_table.alloc_str(symbol.as_ref()); + self.string_table.alloc(0); + } + + /// Builds the armap with the specified offsets + pub fn build(mut self, archive_map: &HashMap) -> Vec { + let mut buffer = make_archive_member_buffer( + &ArchiveMemberName::Value("/".to_string()), + &ArchiveMemberMetadata { + date: Some(0), + uid: Some(0), + gid: Some(0), + mode: Some(0), + }, + &self, + ); + + // Number of symbols + buffer.extend((self.indicies.len() as u32).to_be_bytes()); + + // Offsets + for archive_index in self.indicies.iter_mut() { + let offset = *archive_map + .get(archive_index) + .unwrap_or_else(|| unreachable!()); + buffer.extend((offset as u32).to_be_bytes()); + } + + // String table + buffer.append(&mut self.string_table.into_vec()); + + // Padding + if buffer.len() % 2 != 0 { + buffer.push(b'\n'); + } + + buffer + } +} + +impl MemberSize for ArchiveMapBuilder { + fn member_data_size(&self) -> usize { + 4 + 4 * self.indicies.len() + self.string_table.len() + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use crate::archive::builder::ArchiveMemberIndex; + + use super::ArchiveMapBuilder; + + #[test] + fn correct_symbol_count() { + const TEST_COUNTS: &[usize] = &[5, 10, 20, 100, 1024]; + + for &test_count in TEST_COUNTS { + let symbol_map = ArchiveMapBuilder::default(); + + for v in 0..test_count { + symbol_map.add_symbol(ArchiveMemberIndex(0), format!("{v}")); + } + + let built = symbol_map.build(&HashMap::from([(ArchiveMemberIndex(0), 0)])); + let found_symbol_count = u32::from_be_bytes(built[60..60 + 4].try_into().unwrap()); + + assert_eq!( + test_count as u32, found_symbol_count, + "expected number of symbols in the built archive map does not match the found value" + ); + } + } + + #[test] + fn table_entries_remapped() { + const TEST_VALUES: &[(ArchiveMemberIndex, usize)] = &[ + (ArchiveMemberIndex(0), 10), + (ArchiveMemberIndex(1), 20), + (ArchiveMemberIndex(0), 10), + (ArchiveMemberIndex(2), 100), + (ArchiveMemberIndex(3), 200), + ]; + + let test_symbols = TEST_VALUES + .iter() + .enumerate() + .map(|(idx, _)| format!("symbol{idx}")) + .collect::>(); + + let symbol_map = ArchiveMapBuilder::default(); + + for ((member_idx, _), symbol) in TEST_VALUES.iter().zip(test_symbols.iter()) { + symbol_map.add_symbol(*member_idx, symbol); + } + + let symbol_remap: HashMap = + HashMap::from_iter(TEST_VALUES.iter().copied()); + + let built = symbol_map.build(&symbol_remap); + let armap_data = &built[60..]; + + let symbol_count = u32::from_be_bytes(armap_data[..4].try_into().unwrap()); + assert_eq!( + TEST_VALUES.len(), + symbol_count as usize, + "not all symbols were added to the archive map" + ); + + let mut offset_iter = armap_data[4..].chunks_exact(4); + for (_, expected_offset) in TEST_VALUES { + let found_offset = u32::from_be_bytes(offset_iter.next().unwrap().try_into().unwrap()); + assert_eq!( + *expected_offset, found_offset as usize, + "found offset value for built symbol map does not match expected value" + ); + } + } + + #[test] + fn aligned() { + const TEST_VALUES: &[(ArchiveMemberIndex, usize)] = &[ + (ArchiveMemberIndex(0), 10), + (ArchiveMemberIndex(1), 100), + (ArchiveMemberIndex(0), 1000), + (ArchiveMemberIndex(0), 2000), + (ArchiveMemberIndex(2), 10000), + (ArchiveMemberIndex(3), 100000), + ]; + + let test_symbols = TEST_VALUES + .iter() + .map(|(_, idx)| format!("symbol{idx}")) + .collect::>(); + + for i in 1..=TEST_VALUES.len() { + let symbol_map = ArchiveMapBuilder::default(); + for ((member_idx, _), symbol) in TEST_VALUES.iter().take(i).zip(test_symbols.iter()) { + symbol_map.add_symbol(*member_idx, symbol); + } + + let archive_map: HashMap = + HashMap::from_iter(TEST_VALUES.iter().take(i).copied()); + + let built = symbol_map.build(&archive_map); + assert!( + built.len() % 2 == 0, + "built archive map member with {} symbols should be 2-byte aligned", + i, + ); + } + } +} diff --git a/crates/coffyaml/src/archive/builder/gnu.rs b/crates/coffyaml/src/archive/builder/gnu.rs new file mode 100644 index 0000000..e636265 --- /dev/null +++ b/crates/coffyaml/src/archive/builder/gnu.rs @@ -0,0 +1,45 @@ +use std::collections::HashMap; + +use super::{ + ArchiveMemberIndex, ArchiveVariant, ByteSize, MemberSize, + armap::ArchiveMapBuilder, + longnames::{ArchiveLongNamesBuilder, ArchiveMemberName}, +}; + +#[derive(Default)] +pub struct GnuArchiveVariant { + armap: ArchiveMapBuilder, + longnames: ArchiveLongNamesBuilder, +} + +impl ByteSize for GnuArchiveVariant { + fn byte_size(&self) -> usize { + let mut build_size = 0; + + // Archive map + build_size += self.armap.member_size(); + + // Long names + build_size += self.longnames.member_size(); + + build_size + } +} + +impl ArchiveVariant for GnuArchiveVariant { + fn add_exported_symbol(&mut self, member: ArchiveMemberIndex, symbol: impl AsRef) { + self.armap.add_symbol(member, symbol); + } + + fn add_long_name(&mut self, name: impl Into) -> ArchiveMemberName { + self.longnames.add_name(name) + } + + fn build(self, archive_map: HashMap) -> Vec { + let mut buffer = Vec::with_capacity(self.byte_size()); + buffer.append(&mut self.armap.build(&archive_map)); + buffer.append(&mut self.longnames.build()); + + buffer + } +} diff --git a/crates/coffyaml/src/archive/builder/longnames.rs b/crates/coffyaml/src/archive/builder/longnames.rs new file mode 100644 index 0000000..8606cfd --- /dev/null +++ b/crates/coffyaml/src/archive/builder/longnames.rs @@ -0,0 +1,273 @@ +use std::collections::{HashMap, hash_map}; + +use typed_arena::Arena; + +use super::{ArchiveMemberMetadata, MemberSize, make_archive_member_buffer}; + +#[derive(Debug, PartialEq, Eq)] +pub enum ArchiveMemberName { + LongNameOffset(usize), + Value(String), +} + +impl ArchiveMemberName { + pub fn to_name_array(&self) -> [u8; 16] { + match self { + Self::LongNameOffset(offset) => { + let s = format!("/{offset}"); + let mut buffer = [b' '; 16]; + buffer[..s.len()].copy_from_slice(s.as_bytes()); + buffer + } + Self::Value(s) => { + debug_assert!(s.len() < std::mem::offset_of!(object::archive::Header, date)); + debug_assert!(s.chars().last().is_some_and(|v| v == '/')); + + let mut buffer = [b' '; 16]; + buffer[..s.len()].copy_from_slice(s.as_bytes()); + buffer + } + } + } +} + +#[derive(Default)] +pub struct ArchiveLongNamesBuilder { + offset_map: HashMap, + longnames: Arena, +} + +impl ArchiveLongNamesBuilder { + pub fn add_name(&mut self, name: impl Into) -> ArchiveMemberName { + let mut name = name.into(); + name.push('/'); + + if name.len() < 16 { + ArchiveMemberName::Value(name) + } else { + ArchiveMemberName::LongNameOffset(match self.offset_map.entry(name) { + hash_map::Entry::Occupied(entry) => *entry.get(), + hash_map::Entry::Vacant(entry) => { + let offset = self.longnames.len(); + + self.longnames.alloc_str(entry.key().as_str()); + self.longnames.alloc(DELIM); + + entry.insert(offset); + offset + } + }) + } + } + + fn is_empty(&self) -> bool { + self.longnames.len() == 0 + } + + pub fn build(self) -> Vec { + // Do not build the long names member if it is empty + if self.is_empty() { + return Vec::new(); + } + + let mut buffer = make_archive_member_buffer( + &ArchiveMemberName::Value("//".to_string()), + &ArchiveMemberMetadata { + ..Default::default() + }, + &self, + ); + + buffer.append(&mut self.longnames.into_vec()); + + // Padding + if buffer.len() % 2 != 0 { + buffer.push(b'\n'); + } + + buffer + } +} + +impl MemberSize for ArchiveLongNamesBuilder { + fn member_data_size(&self) -> usize { + self.longnames.len() + } + + fn member_size(&self) -> usize { + // Return 0 if the longnames are empty + if self.member_data_size() == 0 { + return 0; + } + + let size = std::mem::size_of::() + self.member_data_size(); + if size % 2 != 0 { size + 1 } else { size } + } +} + +#[cfg(test)] +mod tests { + use crate::archive::builder::MemberSize; + + use super::{ArchiveLongNamesBuilder, ArchiveMemberName}; + + #[test] + fn empty_build() { + let longnames_member = ArchiveLongNamesBuilder::::default(); + + assert_eq!( + longnames_member.member_size(), + 0, + "calculated member size should be 0 for empty long names members" + ); + + let built = longnames_member.build(); + assert!( + built.is_empty(), + "long names member without any entries should be empty when built" + ); + } + + #[test] + fn trailing_slash() { + let mut longnames_member = ArchiveLongNamesBuilder::::default(); + assert_eq!( + longnames_member.add_name("hello"), + ArchiveMemberName::Value("hello/".to_string()), + "trailing forward slash should have been added to the long names return value" + ); + } + + #[test] + fn offsets_valid() { + let mut longnames_member = ArchiveLongNamesBuilder::::default(); + let firstname = "abcdefghijklmnopqrstuvwxyz"; + + assert_eq!( + longnames_member.add_name(firstname), + ArchiveMemberName::LongNameOffset(0), + "this entry should be the first entry in the long names member" + ); + + let secondname = "testnametestnametestname"; + assert_eq!( + longnames_member.add_name(secondname), + ArchiveMemberName::LongNameOffset(firstname.len() + 2), + "offset for the second entry in the long names member should be after the first", + ); + + let thirdname = "testingtestingtestingtest"; + + assert_eq!( + longnames_member.add_name(thirdname), + ArchiveMemberName::LongNameOffset(firstname.len() + 2 + secondname.len() + 2) + ); + } + + #[test] + fn duplicates_same_offset() { + let mut longnames_member = ArchiveLongNamesBuilder::::default(); + + let testvalue = "abcdefghijklmnopqrstuvwxyz"; + + let first_value = longnames_member.add_name(testvalue); + longnames_member.add_name("testnametestnametestname"); + + let duplicate_value = longnames_member.add_name(testvalue); + + assert_eq!( + first_value, duplicate_value, + "duplicate strings inserted in the long names member should return the same offset" + ); + } + + #[test] + fn built_data_valid() { + let mut longnames_member = ArchiveLongNamesBuilder::::default(); + + let first_string = "abcdefghijklmnopqrstuvwxyz"; + let first_name = longnames_member.add_name(first_string); + + let second_string = "testnametestnametestname"; + let second_name = longnames_member.add_name(second_string); + + let third_string = "testingtestingtestingtest"; + let third_name = longnames_member.add_name(third_string); + + let built_data = longnames_member.longnames.into_vec(); + + let mut check_offset = 0; + for (expected_string, expected_name) in [first_string, second_string, third_string] + .iter() + .zip([first_name, second_name, third_name]) + { + assert_eq!( + expected_name, + ArchiveMemberName::LongNameOffset(check_offset), + "found offset value for string '{expected_string}' in the long names data does not match the value returned when the string was inserted" + ); + + let check_str = std::ffi::CStr::from_bytes_until_nul(&built_data[check_offset..]) + .unwrap() + .to_str() + .unwrap(); + + assert_eq!( + format!("{expected_string}/"), + check_str, + "string value at offset {check_offset} does not match the expected value" + ); + + check_offset += check_str.len() + 1; + } + } + + #[test] + fn aligned() { + let name = "abcdefghijklmnopqrstuvwxyz"; + + // The test name should have an even length + assert!( + name.len() % 2 == 0, + "test string '{name}' length should be even" + ); + + let mut longnames_member = ArchiveLongNamesBuilder::::default(); + longnames_member.add_name(name); + + let built = longnames_member.build(); + assert!( + !built.is_empty(), + "built long names member should not be empty" + ); + + assert!( + built.len() % 2 == 0, + "built long names member for an even data length is not 2 byte aligned. length = {}", + built.len() + ); + + let name = "abcdefghijklmnopqrstuvwxy"; + + // The test name should have an odd length + assert!( + name.len() % 2 != 0, + "test string '{name}' length should be odd" + ); + + let mut longnames_member = ArchiveLongNamesBuilder::::default(); + longnames_member.add_name(name); + + let built = longnames_member.build(); + assert!( + !built.is_empty(), + "built long names member should not be empty" + ); + + assert!( + built.len() % 2 == 0, + "built long names member for an odd data length is not 2 byte aligned. length = {}", + built.len() + ); + } +} diff --git a/crates/coffyaml/src/archive/builder/mod.rs b/crates/coffyaml/src/archive/builder/mod.rs new file mode 100644 index 0000000..4b2139e --- /dev/null +++ b/crates/coffyaml/src/archive/builder/mod.rs @@ -0,0 +1,305 @@ +use std::collections::HashMap; + +use longnames::ArchiveMemberName; +use typed_arena::Arena; + +mod armap; +mod gnu; +mod longnames; +mod msvc; +mod sorted_armap; + +pub use gnu::GnuArchiveVariant; +pub use msvc::MsvcArchiveVariant; + +pub trait ByteSize { + fn byte_size(&self) -> usize; +} + +trait MemberSize { + fn member_data_size(&self) -> usize; + + fn member_size(&self) -> usize { + let size = std::mem::size_of::() + self.member_data_size(); + if size % 2 != 0 { size + 1 } else { size } + } +} + +fn make_archive_member_buffer( + name: &ArchiveMemberName, + metadata: &ArchiveMemberMetadata, + member: &impl MemberSize, +) -> Vec { + use object::archive::Header; + use std::mem::offset_of; + + let mut buffer = Vec::with_capacity(member.member_size()); + buffer.extend(name.to_name_array()); + + // Pad to the date field + buffer.resize(offset_of!(Header, date), b' '); + + let mut strbuf = String::with_capacity(12); + + if let Some(date_val) = metadata.date { + make_ascii_base10(&mut strbuf, date_val); + buffer.extend(strbuf.as_bytes()); + } + + // Pad to the uid field + buffer.resize(offset_of!(Header, uid), b' '); + + if let Some(uid_val) = metadata.uid { + make_ascii_base10(&mut strbuf, uid_val); + buffer.extend(strbuf.as_bytes()); + } + + // Pad to the gid field + buffer.resize(offset_of!(Header, gid), b' '); + + if let Some(gid_val) = metadata.gid { + make_ascii_base10(&mut strbuf, gid_val); + buffer.extend(strbuf.as_bytes()); + } + + // Pad to the mode field + buffer.resize(offset_of!(Header, mode), b' '); + + if let Some(mode_val) = metadata.mode { + make_ascii_base10(&mut strbuf, mode_val); + buffer.extend(strbuf.as_bytes()); + } + + // Pad to the size field + buffer.resize(offset_of!(Header, size), b' '); + + make_ascii_base10(&mut strbuf, member.member_data_size() as u64); + buffer.extend(strbuf.as_bytes()); + + // Pad to the terminator field + buffer.resize(offset_of!(Header, terminator), b' '); + + buffer.extend(object::archive::TERMINATOR); + + buffer +} + +pub trait ArchiveVariant: Default + ByteSize { + fn add_exported_symbol(&mut self, member: ArchiveMemberIndex, symbol: impl AsRef); + fn add_long_name(&mut self, name: impl Into) -> ArchiveMemberName; + fn build(self, archive_map: HashMap) -> Vec; +} + +#[derive(Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)] +#[repr(transparent)] +pub struct ArchiveMemberIndex(pub(self) usize); + +#[derive(Default)] +struct ArchiveMemberMetadata { + date: Option, + uid: Option, + gid: Option, + mode: Option, +} + +struct ArchiveMemberBuilder { + name: ArchiveMemberName, + meta: ArchiveMemberMetadata, + data: Vec, +} + +impl ArchiveMemberBuilder { + fn new(name: ArchiveMemberName, data: Vec) -> ArchiveMemberBuilder { + Self { + name, + data, + meta: ArchiveMemberMetadata::default(), + } + } + + // Builds this member + fn build(mut self) -> Vec { + let mut buffer = make_archive_member_buffer(&self.name, &self.meta, &self); + buffer.append(&mut self.data); + if buffer.len() % 2 != 0 { + buffer.push(b'\n'); + } + + buffer + } +} + +impl MemberSize for ArchiveMemberBuilder { + fn member_data_size(&self) -> usize { + self.data.len() + } +} + +/// Converts an integer value to base 10 ascii using the specified string buffer +fn make_ascii_base10(s: &mut String, v: impl Into) { + s.clear(); + + let v = v.into(); + + let mut n = 0; + + while v / 10u64.pow(n + 1) != 0 { + n += 1; + } + + for e in (0..=n).rev() { + let num = ((v / 10u64.pow(e)) % 10) as u32; + s.push(char::from_digit(num, 10).unwrap()); + } +} + +pub struct ArchiveBuilder { + /// The archive variant + variant: V, + + /// The members in the archive + members: Arena, +} + +impl ArchiveBuilder { + /// Create a new [`ArchiveBuilder`] for the specified number of members. + /// + /// This excludes linker members (armap, longnames, etc.). + pub fn with_capacity(members: usize) -> ArchiveBuilder { + Self { + variant: V::default(), + members: Arena::with_capacity(members), + } + } + + /// Adds a member to the archive. + /// + /// Returns an [`ArchiveMemberAccessor`] for modifying the inserted member + /// metadata. + pub fn add_member( + &mut self, + name: impl Into, + data: impl Into>, + ) -> ArchiveMemberAccessor<'_, V> { + let archive_index = self.members.len(); + + ArchiveMemberAccessor { + index: ArchiveMemberIndex(archive_index), + member: self.members.alloc(ArchiveMemberBuilder::new( + self.variant.add_long_name(name), + data.into(), + )), + variant: &mut self.variant, + } + } + + /// Build the archive + pub fn build(mut self) -> Vec { + // Calculate the buffer size needed for building the archive + let mut buffer_size = object::archive::MAGIC.len() + self.variant.byte_size(); + + // Create the map of archive members to their file offsets + let mut archive_map = HashMap::with_capacity(self.members.len()); + for (member_idx, member) in self.members.iter_mut().enumerate() { + archive_map.insert(ArchiveMemberIndex(member_idx), buffer_size); + buffer_size += member.member_size(); + } + + // Build the archive + + let mut buffer = Vec::with_capacity(buffer_size); + // Add the signature + buffer.extend(object::archive::MAGIC); + // Add the variant members + buffer.append(&mut self.variant.build(archive_map)); + + // Add the rest of the members + let members = self.members.into_vec(); + for member in members { + buffer.append(&mut member.build()); + } + + buffer + } +} + +impl ArchiveBuilder { + pub fn msvc_archive_with_capacity(members: usize) -> ArchiveBuilder { + Self::with_capacity(members) + } +} + +impl ArchiveBuilder { + pub fn gnu_archive_with_capacity(members: usize) -> ArchiveBuilder { + Self::with_capacity(members) + } +} + +/// Accessor for modifying metadata of an archive member. +pub struct ArchiveMemberAccessor<'a, V: ArchiveVariant> { + /// The archive index + index: ArchiveMemberIndex, + + /// The archive member + member: &'a mut ArchiveMemberBuilder, + + /// The archive variant + variant: &'a mut V, +} + +impl ArchiveMemberAccessor<'_, V> { + /// Adds a symbol export to this archive member in the symbol table + pub fn export(&mut self, symbol: impl AsRef) { + self.variant.add_exported_symbol(self.index, symbol); + } + + /// Adds the list of exports to the archive symbol table for this member + pub fn exports(&mut self, symbols: I) + where + S: AsRef, + I: IntoIterator, + { + for symbol in symbols { + self.variant.add_exported_symbol(self.index, symbol); + } + } + + /// Sets the date timestamp in the archive header to the specified value + pub fn date(&mut self, ts: u64) { + self.member.meta.date = Some(ts); + } + + /// Sets the user id value to the specified value + pub fn uid(&mut self, uid: u32) { + self.member.meta.uid = Some(uid); + } + + /// Sets the group id value to the specified value + pub fn gid(&mut self, gid: u32) { + self.member.meta.gid = Some(gid); + } + + /// Sets the mode to the specified value + pub fn mode(&mut self, mode: u32) { + self.member.meta.mode = Some(mode); + } +} + +#[cfg(test)] +mod tests { + use super::make_ascii_base10; + + #[test] + fn make_ascii_int() { + const TESTS: &[(u64, &str)] = &[(123, "123"), (848193, "848193"), (0, "0")]; + + for (v, expected) in TESTS { + let mut value = String::with_capacity(expected.len()); + make_ascii_base10(&mut value, *v); + assert_eq!( + value, *expected, + "integer value should have been converted into ASCII" + ); + } + } +} diff --git a/crates/coffyaml/src/archive/builder/msvc.rs b/crates/coffyaml/src/archive/builder/msvc.rs new file mode 100644 index 0000000..a6536b2 --- /dev/null +++ b/crates/coffyaml/src/archive/builder/msvc.rs @@ -0,0 +1,58 @@ +use std::collections::HashMap; + +use super::{ + ArchiveMemberIndex, ArchiveVariant, ByteSize, MemberSize, + armap::ArchiveMapBuilder, + longnames::{ArchiveLongNamesBuilder, ArchiveMemberName}, + sorted_armap::SortedArchiveMapBuilder, +}; + +#[derive(Default)] +pub struct MsvcArchiveVariant { + /// The archive symbol map + armap: ArchiveMapBuilder, + + sorted_armap: SortedArchiveMapBuilder, + + /// The long names + longnames: ArchiveLongNamesBuilder, +} + +impl ByteSize for MsvcArchiveVariant { + fn byte_size(&self) -> usize { + let mut build_size = 0; + + // Archive map + build_size += self.armap.member_size(); + + // Sorted archive map + build_size += self.sorted_armap.member_size(); + + // Long names if it is not empty + build_size += self.longnames.member_size(); + + build_size + } +} + +impl ArchiveVariant for MsvcArchiveVariant { + fn add_exported_symbol(&mut self, member: ArchiveMemberIndex, symbol: impl AsRef) { + self.armap.add_symbol(member, &symbol); + self.sorted_armap.add_symbol(member, symbol); + } + + fn add_long_name(&mut self, name: impl Into) -> ArchiveMemberName { + self.longnames.add_name(name) + } + + fn build(self, archive_map: HashMap) -> Vec { + let mut buffer = Vec::with_capacity(self.byte_size()); + + // Each member `.build()` method should add padding + buffer.append(&mut self.armap.build(&archive_map)); + buffer.append(&mut self.sorted_armap.build(&archive_map)); + buffer.append(&mut self.longnames.build()); + + buffer + } +} diff --git a/crates/coffyaml/src/archive/builder/sorted_armap.rs b/crates/coffyaml/src/archive/builder/sorted_armap.rs new file mode 100644 index 0000000..101fffd --- /dev/null +++ b/crates/coffyaml/src/archive/builder/sorted_armap.rs @@ -0,0 +1,105 @@ +use std::collections::{HashMap, HashSet}; + +use indexmap::IndexSet; + +use super::{ + ArchiveMemberIndex, ArchiveMemberMetadata, MemberSize, longnames::ArchiveMemberName, + make_archive_member_buffer, +}; + +#[derive(Default)] +pub struct SortedArchiveMapBuilder { + member_indicies: HashSet, + string_indicies: HashSet<(String, ArchiveMemberIndex)>, + string_table_size: usize, +} + +impl SortedArchiveMapBuilder { + pub fn add_symbol(&mut self, index: ArchiveMemberIndex, symbol: impl AsRef) { + let symbol_len = symbol.as_ref().len(); + if !self + .string_indicies + .insert((symbol.as_ref().to_string(), index)) + { + return; + } + + self.member_indicies.insert(index); + self.string_table_size += symbol_len + 1; + } + + pub fn build(self, archive_map: &HashMap) -> Vec { + let mut buffer = make_archive_member_buffer( + &ArchiveMemberName::Value("/".to_string()), + &ArchiveMemberMetadata { + date: Some(0), + uid: Some(0), + gid: Some(0), + mode: Some(0), + }, + &self, + ); + + // Create the sorted member offsets + let mut member_offsets: IndexSet = + IndexSet::from_iter(self.member_indicies.into_iter().map(|member_id| { + *archive_map + .get(&member_id) + .unwrap_or_else(|| unreachable!()) + })); + + member_offsets.sort(); + + // Add the number of members + buffer.extend((member_offsets.len() as u32).to_le_bytes()); + + // Add the member offsets + buffer.extend( + member_offsets + .iter() + .flat_map(|offset| (*offset as u32).to_le_bytes()), + ); + + // Create the sorted strings + let mut string_table = self.string_indicies.into_iter().collect::>(); + string_table.sort_by(|(a, _), (b, _)| a.cmp(b)); + + // Add the number of symbols + buffer.extend((string_table.len() as u32).to_le_bytes()); + + // Add the symbol indices + for (_, archive_idx) in &string_table { + let member_offset = *archive_map + .get(archive_idx) + .unwrap_or_else(|| unreachable!()); + + let table_index = member_offsets + .get_index_of(&member_offset) + .unwrap_or_else(|| unreachable!()); + + buffer.extend(u16::try_from(table_index + 1).unwrap().to_le_bytes()); + } + + // Add the strings + for (symbol, _) in string_table { + buffer.extend(symbol.as_bytes()); + buffer.push(0); + } + + // Padding + if buffer.len() % 2 != 0 { + buffer.push(b'\n'); + } + + buffer + } +} + +impl MemberSize for SortedArchiveMapBuilder { + fn member_data_size(&self) -> usize { + 4 + 4 * self.member_indicies.len() + + 4 + + 2 * self.string_indicies.len() + + self.string_table_size + } +} diff --git a/crates/coffyaml/src/archive/mod.rs b/crates/coffyaml/src/archive/mod.rs new file mode 100644 index 0000000..9b5530b --- /dev/null +++ b/crates/coffyaml/src/archive/mod.rs @@ -0,0 +1 @@ +pub(crate) mod builder; diff --git a/crates/coffyaml/src/coff/errors.rs b/crates/coffyaml/src/coff/errors.rs new file mode 100644 index 0000000..309150a --- /dev/null +++ b/crates/coffyaml/src/coff/errors.rs @@ -0,0 +1,16 @@ +use std::num::TryFromIntError; + +#[derive(Debug, thiserror::Error)] +pub enum CoffYamlCoffBuildError { + #[error("{0}")] + IntegerConversion(#[from] TryFromIntError), + + #[error("relocation target symbol {0} does not exist")] + MissingSymbol(String), + + #[error("alignment value of {align} for section index {index} is not valid")] + SectionAlign { index: usize, align: usize }, + + #[error("{0}")] + ObjectWrite(#[from] object::write::Error), +} diff --git a/crates/coffyaml/src/coff/header.rs b/crates/coffyaml/src/coff/header.rs new file mode 100644 index 0000000..d2b57d4 --- /dev/null +++ b/crates/coffyaml/src/coff/header.rs @@ -0,0 +1,333 @@ +use object::pe::{ + IMAGE_FILE_32BIT_MACHINE, IMAGE_FILE_AGGRESIVE_WS_TRIM, IMAGE_FILE_BYTES_REVERSED_HI, + IMAGE_FILE_BYTES_REVERSED_LO, IMAGE_FILE_DEBUG_STRIPPED, IMAGE_FILE_DLL, + IMAGE_FILE_EXECUTABLE_IMAGE, IMAGE_FILE_LARGE_ADDRESS_AWARE, IMAGE_FILE_LINE_NUMS_STRIPPED, + IMAGE_FILE_LOCAL_SYMS_STRIPPED, IMAGE_FILE_MACHINE_AM33, IMAGE_FILE_MACHINE_AMD64, + IMAGE_FILE_MACHINE_ARM, IMAGE_FILE_MACHINE_ARM64, IMAGE_FILE_MACHINE_ARMNT, + IMAGE_FILE_MACHINE_EBC, IMAGE_FILE_MACHINE_I386, IMAGE_FILE_MACHINE_IA64, + IMAGE_FILE_MACHINE_M32R, IMAGE_FILE_MACHINE_MIPS16, IMAGE_FILE_MACHINE_MIPSFPU, + IMAGE_FILE_MACHINE_MIPSFPU16, IMAGE_FILE_MACHINE_POWERPC, IMAGE_FILE_MACHINE_POWERPCFP, + IMAGE_FILE_MACHINE_R4000, IMAGE_FILE_MACHINE_SH3, IMAGE_FILE_MACHINE_SH3DSP, + IMAGE_FILE_MACHINE_SH4, IMAGE_FILE_MACHINE_SH5, IMAGE_FILE_MACHINE_THUMB, + IMAGE_FILE_MACHINE_UNKNOWN, IMAGE_FILE_MACHINE_WCEMIPSV2, IMAGE_FILE_NET_RUN_FROM_SWAP, + IMAGE_FILE_RELOCS_STRIPPED, IMAGE_FILE_REMOVABLE_RUN_FROM_SWAP, IMAGE_FILE_SYSTEM, + IMAGE_FILE_UP_SYSTEM_ONLY, +}; +use serde::{ + Deserialize, Deserializer, Serialize, Serializer, + de::{self, Visitor}, + ser::SerializeSeq, +}; + +#[derive(Debug, Default, Clone, PartialEq, Eq, Deserialize, Serialize)] +#[serde(rename_all = "PascalCase")] +pub struct CoffYamlHeader { + #[serde( + deserialize_with = "machine_deserializer", + serialize_with = "machine_serializer" + )] + pub machine: u16, + + #[serde( + deserialize_with = "characteristics_deserializer", + serialize_with = "characteristics_serializer" + )] + pub characteristics: u16, +} + +fn machine_deserializer<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + struct MachineVisitor; + + impl Visitor<'_> for MachineVisitor { + type Value = u16; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("'IMAGE_FILE_MACHINE_*' string or integer") + } + + fn visit_u64(self, v: u64) -> Result + where + E: de::Error, + { + u16::try_from(v).map_err(serde::de::Error::custom) + } + + fn visit_str(self, v: &str) -> Result + where + E: de::Error, + { + Ok(match v { + "IMAGE_FILE_MACHINE_UNKNOWN" => IMAGE_FILE_MACHINE_UNKNOWN, + "IMAGE_FILE_MACHINE_AM33" => IMAGE_FILE_MACHINE_AM33, + "IMAGE_FILE_MACHINE_AMD64" => IMAGE_FILE_MACHINE_AMD64, + "IMAGE_FILE_MACHINE_ARM" => IMAGE_FILE_MACHINE_ARM, + "IMAGE_FILE_MACHINE_ARMNT" => IMAGE_FILE_MACHINE_ARMNT, + "IMAGE_FILE_MACHINE_ARM64" => IMAGE_FILE_MACHINE_ARM64, + "IMAGE_FILE_MACHINE_EBC" => IMAGE_FILE_MACHINE_EBC, + "IMAGE_FILE_MACHINE_I386" => IMAGE_FILE_MACHINE_I386, + "IMAGE_FILE_MACHINE_IA64" => IMAGE_FILE_MACHINE_IA64, + "IMAGE_FILE_MACHINE_M32R" => IMAGE_FILE_MACHINE_M32R, + "IMAGE_FILE_MACHINE_MIPS16" => IMAGE_FILE_MACHINE_MIPS16, + "IMAGE_FILE_MACHINE_MIPSFPU" => IMAGE_FILE_MACHINE_MIPSFPU, + "IMAGE_FILE_MACHINE_MIPSFPU16" => IMAGE_FILE_MACHINE_MIPSFPU16, + "IMAGE_FILE_MACHINE_POWERPC" => IMAGE_FILE_MACHINE_POWERPC, + "IMAGE_FILE_MACHINE_POWERPCFP" => IMAGE_FILE_MACHINE_POWERPCFP, + "IMAGE_FILE_MACHINE_R4000" => IMAGE_FILE_MACHINE_R4000, + "IMAGE_FILE_MACHINE_SH3" => IMAGE_FILE_MACHINE_SH3, + "IMAGE_FILE_MACHINE_SH3DSP" => IMAGE_FILE_MACHINE_SH3DSP, + "IMAGE_FILE_MACHINE_SH4" => IMAGE_FILE_MACHINE_SH4, + "IMAGE_FILE_MACHINE_SH5" => IMAGE_FILE_MACHINE_SH5, + "IMAGE_FILE_MACHINE_THUMB" => IMAGE_FILE_MACHINE_THUMB, + "IMAGE_FILE_MACHINE_WCEMIPSV2" => IMAGE_FILE_MACHINE_WCEMIPSV2, + _ => { + return Err(serde::de::Error::custom(format!( + "invalid machine type {v}" + ))); + } + }) + } + } + + deserializer.deserialize_any(MachineVisitor) +} + +fn machine_serializer(machine: &u16, serializer: S) -> Result +where + S: Serializer, +{ + match *machine { + IMAGE_FILE_MACHINE_UNKNOWN => serializer.serialize_str("IMAGE_FILE_MACHINE_UNKNOWN"), + IMAGE_FILE_MACHINE_AM33 => serializer.serialize_str("IMAGE_FILE_MACHINE_AM33"), + IMAGE_FILE_MACHINE_AMD64 => serializer.serialize_str("IMAGE_FILE_MACHINE_AMD64"), + IMAGE_FILE_MACHINE_ARM => serializer.serialize_str("IMAGE_FILE_MACHINE_ARM"), + IMAGE_FILE_MACHINE_ARMNT => serializer.serialize_str("IMAGE_FILE_MACHINE_ARMNT"), + IMAGE_FILE_MACHINE_ARM64 => serializer.serialize_str("IMAGE_FILE_MACHINE_ARM64"), + IMAGE_FILE_MACHINE_EBC => serializer.serialize_str("IMAGE_FILE_MACHINE_EBC"), + IMAGE_FILE_MACHINE_I386 => serializer.serialize_str("IMAGE_FILE_MACHINE_I386"), + IMAGE_FILE_MACHINE_IA64 => serializer.serialize_str("IMAGE_FILE_MACHINE_IA64"), + IMAGE_FILE_MACHINE_M32R => serializer.serialize_str("IMAGE_FILE_MACHINE_M32R"), + IMAGE_FILE_MACHINE_MIPS16 => serializer.serialize_str("IMAGE_FILE_MACHINE_MIPS16"), + IMAGE_FILE_MACHINE_MIPSFPU => serializer.serialize_str("IMAGE_FILE_MACHINE_MIPSFPU"), + IMAGE_FILE_MACHINE_MIPSFPU16 => serializer.serialize_str("IMAGE_FILE_MACHINE_MIPSFPU16"), + IMAGE_FILE_MACHINE_POWERPC => serializer.serialize_str("IMAGE_FILE_MACHINE_POWERPC"), + IMAGE_FILE_MACHINE_POWERPCFP => serializer.serialize_str("IMAGE_FILE_MACHINE_POWERPCFP"), + IMAGE_FILE_MACHINE_R4000 => serializer.serialize_str("IMAGE_FILE_MACHINE_R4000"), + IMAGE_FILE_MACHINE_SH3 => serializer.serialize_str("IMAGE_FILE_MACHINE_SH3"), + IMAGE_FILE_MACHINE_SH3DSP => serializer.serialize_str("IMAGE_FILE_MACHINE_SH3DSP"), + IMAGE_FILE_MACHINE_SH4 => serializer.serialize_str("IMAGE_FILE_MACHINE_SH4"), + IMAGE_FILE_MACHINE_SH5 => serializer.serialize_str("IMAGE_FILE_MACHINE_SH5"), + IMAGE_FILE_MACHINE_THUMB => serializer.serialize_str("IMAGE_FILE_MACHINE_THUMB"), + IMAGE_FILE_MACHINE_WCEMIPSV2 => serializer.serialize_str("IMAGE_FILE_MACHINE_WCEMIPSV2"), + o => serializer.serialize_u16(o), + } +} + +fn characteristics_deserializer<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + struct CharacteristicsVisitor; + + impl<'de> Visitor<'de> for CharacteristicsVisitor { + type Value = u16; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("'IMAGE_FILE_*' characteristics string list or integer") + } + + fn visit_u64(self, v: u64) -> Result + where + E: de::Error, + { + u16::try_from(v).map_err(serde::de::Error::custom) + } + + fn visit_seq(self, mut seq: A) -> Result + where + A: de::SeqAccess<'de>, + { + let mut flags = 0; + + while let Some(val) = seq.next_element::<&str>()? { + flags |= match val { + "IMAGE_FILE_RELOCS_STRIPPED" => IMAGE_FILE_RELOCS_STRIPPED, + "IMAGE_FILE_EXECUTABLE_IMAGE" => IMAGE_FILE_EXECUTABLE_IMAGE, + "IMAGE_FILE_LINE_NUMS_STRIPPED" => IMAGE_FILE_LINE_NUMS_STRIPPED, + "IMAGE_FILE_LOCAL_SYMS_STRIPPED" => IMAGE_FILE_LOCAL_SYMS_STRIPPED, + "IMAGE_FILE_AGGRESSIVE_WS_TRIM" => IMAGE_FILE_AGGRESIVE_WS_TRIM, + "IMAGE_FILE_LARGE_ADDRESS_AWARE" => IMAGE_FILE_LARGE_ADDRESS_AWARE, + "IMAGE_FILE_BYTES_REVERSED_LO" => IMAGE_FILE_BYTES_REVERSED_LO, + "IMAGE_FILE_32BIT_MACHINE" => IMAGE_FILE_32BIT_MACHINE, + "IMAGE_FILE_DEBUG_STRIPPED" => IMAGE_FILE_DEBUG_STRIPPED, + "IMAGE_FILE_REMOVABLE_RUN_FROM_SWAP" => IMAGE_FILE_REMOVABLE_RUN_FROM_SWAP, + "IMAGE_FILE_NET_RUN_FROM_SWAP" => IMAGE_FILE_NET_RUN_FROM_SWAP, + "IMAGE_FILE_SYSTEM" => IMAGE_FILE_SYSTEM, + "IMAGE_FILE_DLL" => IMAGE_FILE_DLL, + "IMAGE_FILE_UP_SYSTEM_ONLY" => IMAGE_FILE_UP_SYSTEM_ONLY, + "IMAGE_FILE_BYTES_REVERSED_HI" => IMAGE_FILE_BYTES_REVERSED_HI, + _ => { + return Err(serde::de::Error::custom(format!( + "invalid header characteristic value {val}" + ))); + } + } + } + + Ok(flags) + } + } + + deserializer.deserialize_any(CharacteristicsVisitor) +} + +fn characteristics_serializer(characteristics: &u16, serializer: S) -> Result +where + S: Serializer, +{ + const FLAGS: [(u16, &str); 15] = [ + (IMAGE_FILE_RELOCS_STRIPPED, "IMAGE_FILE_RELOCS_STRIPPED"), + (IMAGE_FILE_EXECUTABLE_IMAGE, "IMAGE_FILE_EXECUTABLE_IMAGE"), + ( + IMAGE_FILE_LINE_NUMS_STRIPPED, + "IMAGE_FILE_LINE_NUMS_STRIPPED", + ), + ( + IMAGE_FILE_LOCAL_SYMS_STRIPPED, + "IMAGE_FILE_LOCAL_SYMS_STRIPPED", + ), + ( + IMAGE_FILE_AGGRESIVE_WS_TRIM, + "IMAGE_FILE_AGGRESSIVE_WS_TRIM", + ), + ( + IMAGE_FILE_LARGE_ADDRESS_AWARE, + "IMAGE_FILE_LARGE_ADDRESS_AWARE", + ), + (IMAGE_FILE_BYTES_REVERSED_LO, "IMAGE_FILE_BYTES_REVERSED_LO"), + (IMAGE_FILE_32BIT_MACHINE, "IMAGE_FILE_32BIT_MACHINE"), + (IMAGE_FILE_DEBUG_STRIPPED, "IMAGE_FILE_DEBUG_STRIPPED"), + ( + IMAGE_FILE_REMOVABLE_RUN_FROM_SWAP, + "IMAGE_FILE_REMOVABLE_RUN_FROM_SWAP", + ), + (IMAGE_FILE_NET_RUN_FROM_SWAP, "IMAGE_FILE_NET_RUN_FROM_SWAP"), + (IMAGE_FILE_SYSTEM, "IMAGE_FILE_SYSTEM"), + (IMAGE_FILE_DLL, "IMAGE_FILE_DLL"), + (IMAGE_FILE_UP_SYSTEM_ONLY, "IMAGE_FILE_UP_SYSTEM_ONLY"), + (IMAGE_FILE_BYTES_REVERSED_HI, "IMAGE_FILE_BYTES_REVERSED_HI"), + ]; + + let contains_flag = |bits: u16, flag: u16| bits & flag != 0; + + let flagcount = FLAGS + .iter() + .filter(|(flag, _)| contains_flag(*characteristics, *flag)) + .count(); + + let mut seq = serializer.serialize_seq(Some(flagcount))?; + for (flag, val) in FLAGS { + if contains_flag(*characteristics, flag) { + seq.serialize_element(val)?; + } + } + + seq.end() +} + +#[cfg(test)] +mod tests { + use super::{CoffYamlHeader, characteristics_deserializer, machine_deserializer}; + use crate::testutils; + use object::pe::{ + IMAGE_FILE_LINE_NUMS_STRIPPED, IMAGE_FILE_MACHINE_AMD64, IMAGE_FILE_MACHINE_I386, + IMAGE_FILE_MACHINE_UNKNOWN, IMAGE_FILE_RELOCS_STRIPPED, + }; + use serde::Deserialize; + + #[test] + fn machine_scalar_deserialize() { + testutils::run_deserializer_tests( + machine_deserializer, + [ + ("34404", IMAGE_FILE_MACHINE_AMD64), + ("0x8664", IMAGE_FILE_MACHINE_AMD64), + ("332", IMAGE_FILE_MACHINE_I386), + ("0x14c", IMAGE_FILE_MACHINE_I386), + ("0", IMAGE_FILE_MACHINE_UNKNOWN), + ], + ); + } + + #[test] + fn machine_string_deserialize() { + testutils::run_deserializer_tests( + machine_deserializer, + [ + ("IMAGE_FILE_MACHINE_AMD64", IMAGE_FILE_MACHINE_AMD64), + ("IMAGE_FILE_MACHINE_I386", IMAGE_FILE_MACHINE_I386), + ("IMAGE_FILE_MACHINE_UNKNOWN", IMAGE_FILE_MACHINE_UNKNOWN), + ], + ); + } + + #[test] + fn characteristics_scalar_deserialize() { + testutils::run_deserializer_tests( + characteristics_deserializer, + [ + ("1", IMAGE_FILE_RELOCS_STRIPPED), + ( + "5", + IMAGE_FILE_RELOCS_STRIPPED | IMAGE_FILE_LINE_NUMS_STRIPPED, + ), + ("0x1", IMAGE_FILE_RELOCS_STRIPPED), + ], + ); + } + + #[test] + fn characteristics_string_deserialize() { + testutils::run_deserializer_tests( + characteristics_deserializer, + [ + ("[ IMAGE_FILE_RELOCS_STRIPPED ]", IMAGE_FILE_RELOCS_STRIPPED), + ( + "[ IMAGE_FILE_RELOCS_STRIPPED, IMAGE_FILE_LINE_NUMS_STRIPPED ]", + IMAGE_FILE_RELOCS_STRIPPED | IMAGE_FILE_LINE_NUMS_STRIPPED, + ), + ], + ); + } + + #[test] + fn header_deserialize() { + testutils::run_deserializer_tests( + CoffYamlHeader::deserialize, + [ + ( + r#" + Machine: IMAGE_FILE_MACHINE_UNKNOWN + Characteristics: [ IMAGE_FILE_RELOCS_STRIPPED ] + "#, + CoffYamlHeader { + machine: IMAGE_FILE_MACHINE_UNKNOWN, + characteristics: IMAGE_FILE_RELOCS_STRIPPED, + }, + ), + ( + r#" + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_RELOCS_STRIPPED, IMAGE_FILE_LINE_NUMS_STRIPPED ] + "#, + CoffYamlHeader { + machine: IMAGE_FILE_MACHINE_AMD64, + characteristics: IMAGE_FILE_RELOCS_STRIPPED | IMAGE_FILE_LINE_NUMS_STRIPPED, + }, + ), + ], + ); + } +} diff --git a/crates/coffyaml/src/coff/mod.rs b/crates/coffyaml/src/coff/mod.rs new file mode 100644 index 0000000..932f554 --- /dev/null +++ b/crates/coffyaml/src/coff/mod.rs @@ -0,0 +1,169 @@ +use std::collections::HashMap; + +use errors::CoffYamlCoffBuildError; +use object::{ + pe::{IMAGE_SYM_ABSOLUTE, IMAGE_SYM_DEBUG, IMAGE_SYM_DTYPE_SHIFT, IMAGE_SYM_UNDEFINED}, + write::coff::{AuxSymbolSection, FileHeader, Relocation, SectionHeader, Symbol, Writer}, +}; +use serde::{Deserialize, Serialize}; + +pub mod errors; +mod header; +mod sections; +mod symbols; + +pub use header::CoffYamlHeader; +pub use sections::{CoffYamlSection, CoffYamlSectionRelocation}; +pub use symbols::{CoffYamlAuxFunctionDefinition, CoffYamlAuxSectionDefinition, CoffYamlSymbol}; + +const SECTION_ALIGN_SHIFT: u32 = 20; + +#[derive(Debug, Default, Clone, Deserialize, Serialize)] +pub struct CoffYaml { + pub header: CoffYamlHeader, + pub sections: Vec, + pub symbols: Vec, +} + +impl CoffYaml { + pub fn build(mut self) -> Result, CoffYamlCoffBuildError> { + let mut buffer = Vec::new(); + + let mut writer = Writer::new(&mut buffer); + writer.reserve_file_header(); + writer.reserve_section_headers(self.sections.len().try_into()?); + + let mut section_headers = Vec::with_capacity(self.sections.len()); + for (idx, section) in self.sections.iter().enumerate() { + let alignment_flag = if let Some(alignment) = section.alignment { + if alignment == 0 || alignment > 8192 || (alignment != 1 && alignment % 2 != 0) { + return Err(CoffYamlCoffBuildError::SectionAlign { + index: idx, + align: alignment, + }); + } + + ((alignment as u32).ilog2() + 1) << SECTION_ALIGN_SHIFT + } else { + 0 + }; + + section_headers.push(SectionHeader { + name: writer.add_name(section.name.as_bytes()), + size_of_raw_data: if let Some(size) = section.size_of_raw_data { + size + } else { + section.section_data.len().try_into()? + }, + pointer_to_raw_data: writer.reserve_section(section.section_data.len()), + pointer_to_relocations: 0, + pointer_to_linenumbers: 0, + number_of_relocations: section.relocations.len().try_into()?, + number_of_linenumbers: 0, + characteristics: section.characteristics | alignment_flag, + }); + } + + for (section_header, section) in section_headers.iter_mut().zip(self.sections.iter()) { + section_header.pointer_to_relocations = + writer.reserve_relocations(section.relocations.len()); + } + + let mut symbol_map = HashMap::with_capacity(self.symbols.len()); + let mut symbol_names = Vec::with_capacity(self.symbols.len()); + + for symbol in self.symbols.iter_mut() { + if let Some(aux_file) = symbol.file.as_mut() { + aux_file.truncate(18); + } + } + + for symbol in &self.symbols { + let idx = writer.reserve_symbol_index(); + symbol_names.push(writer.add_name(symbol.name.as_bytes())); + symbol_map.insert(&symbol.name, idx); + + if let Some(aux_file) = symbol.file.as_ref() { + writer.reserve_aux_file_name(aux_file.as_bytes()); + } + + if symbol.section_definition.as_ref().is_some() { + writer.reserve_aux_section(); + } + } + + writer.reserve_symtab_strtab(); + + writer.write_file_header(FileHeader { + machine: self.header.machine, + time_date_stamp: 0, + characteristics: self.header.characteristics, + })?; + + for header in section_headers { + writer.write_section_header(header); + } + + for section in &self.sections { + writer.write_section(§ion.section_data); + } + + for section in &self.sections { + if section.relocations.len() > u16::MAX as usize { + writer.write_relocations_count(section.relocations.len()); + } + + for reloc in §ion.relocations { + writer.write_relocation(Relocation { + virtual_address: reloc.virtual_address, + symbol: symbol_map.get(&reloc.symbol_name).copied().ok_or_else(|| { + CoffYamlCoffBuildError::MissingSymbol(reloc.symbol_name.clone()) + })?, + typ: reloc.typ, + }); + } + } + + for (symbol_name, symbol) in symbol_names.into_iter().zip(self.symbols.iter()) { + let aux_count = if symbol.file.is_some() { 1 } else { 0 } + + if symbol.section_definition.is_some() { + 1 + } else { + 0 + }; + + writer.write_symbol(Symbol { + name: symbol_name, + value: symbol.value, + section_number: match symbol.section_number { + IMAGE_SYM_UNDEFINED => 0, + IMAGE_SYM_ABSOLUTE => u16::MAX, + IMAGE_SYM_DEBUG => u16::MAX - 1, + _ => symbol.section_number.try_into()?, + }, + typ: (symbol.complex_type << IMAGE_SYM_DTYPE_SHIFT) | (symbol.simple_type & 0xff), + number_of_aux_symbols: aux_count, + storage_class: symbol.storage_class, + }); + + if let Some(aux_file) = symbol.file.as_ref() { + writer.write_aux_file_name(aux_file.as_bytes(), 1); + } + + if let Some(aux_section) = symbol.section_definition.as_ref() { + writer.write_aux_section(AuxSymbolSection { + length: aux_section.length, + number_of_relocations: aux_section.number_of_relocations.into(), + number_of_linenumbers: aux_section.number_of_linenumbers, + check_sum: aux_section.check_sum, + number: aux_section.number.into(), + selection: aux_section.selection, + }); + } + } + + writer.write_strtab(); + + Ok(buffer) + } +} diff --git a/crates/coffyaml/src/coff/sections.rs b/crates/coffyaml/src/coff/sections.rs new file mode 100644 index 0000000..384e470 --- /dev/null +++ b/crates/coffyaml/src/coff/sections.rs @@ -0,0 +1,279 @@ +use object::pe::{ + IMAGE_REL_AMD64_ABSOLUTE, IMAGE_REL_AMD64_ADDR32, IMAGE_REL_AMD64_ADDR32NB, + IMAGE_REL_AMD64_ADDR64, IMAGE_REL_AMD64_PAIR, IMAGE_REL_AMD64_REL32, IMAGE_REL_AMD64_REL32_1, + IMAGE_REL_AMD64_REL32_2, IMAGE_REL_AMD64_REL32_3, IMAGE_REL_AMD64_REL32_4, + IMAGE_REL_AMD64_REL32_5, IMAGE_REL_AMD64_SECREL, IMAGE_REL_AMD64_SECREL7, + IMAGE_REL_AMD64_SECTION, IMAGE_REL_AMD64_SREL32, IMAGE_REL_AMD64_SSPAN32, + IMAGE_REL_AMD64_TOKEN, IMAGE_REL_I386_ABSOLUTE, IMAGE_REL_I386_DIR16, IMAGE_REL_I386_DIR32, + IMAGE_REL_I386_DIR32NB, IMAGE_REL_I386_REL16, IMAGE_REL_I386_REL32, IMAGE_REL_I386_SECREL, + IMAGE_REL_I386_SECREL7, IMAGE_REL_I386_SECTION, IMAGE_REL_I386_SEG12, IMAGE_REL_I386_TOKEN, + IMAGE_SCN_CNT_CODE, IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_CNT_UNINITIALIZED_DATA, + IMAGE_SCN_GPREL, IMAGE_SCN_LNK_COMDAT, IMAGE_SCN_LNK_INFO, IMAGE_SCN_LNK_NRELOC_OVFL, + IMAGE_SCN_LNK_OTHER, IMAGE_SCN_LNK_REMOVE, IMAGE_SCN_MEM_DISCARDABLE, IMAGE_SCN_MEM_EXECUTE, + IMAGE_SCN_MEM_LOCKED, IMAGE_SCN_MEM_NOT_CACHED, IMAGE_SCN_MEM_NOT_PAGED, IMAGE_SCN_MEM_PRELOAD, + IMAGE_SCN_MEM_PURGEABLE, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_SHARED, IMAGE_SCN_MEM_WRITE, + IMAGE_SCN_TYPE_NO_PAD, +}; +use serde::{ + Deserialize, Deserializer, Serialize, Serializer, + de::{self, Visitor}, + ser::SerializeSeq, +}; + +#[derive(Debug, Clone, Default, Deserialize, Serialize)] +#[serde(rename_all = "PascalCase")] +pub struct CoffYamlSection { + pub name: String, + + #[serde( + deserialize_with = "characteristics_deserializer", + serialize_with = "characteristics_serializer" + )] + pub characteristics: u32, + + #[serde(skip_serializing_if = "Option::is_none")] + pub alignment: Option, + + #[serde( + deserialize_with = "hex::serde::deserialize", + serialize_with = "hex::serde::serialize_upper" + )] + pub section_data: Vec, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub size_of_raw_data: Option, + + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub relocations: Vec, +} + +fn characteristics_deserializer<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + struct CharacteristicsVisitor; + + impl<'de> Visitor<'de> for CharacteristicsVisitor { + type Value = u32; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("'IMAGE_SCN_*' string list or integer") + } + + fn visit_u64(self, v: u64) -> Result + where + E: de::Error, + { + u32::try_from(v).map_err(serde::de::Error::custom) + } + + fn visit_seq(self, mut seq: A) -> Result + where + A: de::SeqAccess<'de>, + { + let mut flags = 0; + + while let Some(val) = seq.next_element::<&str>()? { + flags |= match val { + "IMAGE_SCN_TYPE_NO_PAD" => IMAGE_SCN_TYPE_NO_PAD, + "IMAGE_SCN_CNT_CODE" => IMAGE_SCN_CNT_CODE, + "IMAGE_SCN_CNT_INITIALIZED_DATA" => IMAGE_SCN_CNT_INITIALIZED_DATA, + "IMAGE_SCN_CNT_UNINITIALIZED_DATA" => IMAGE_SCN_CNT_UNINITIALIZED_DATA, + "IMAGE_SCN_LNK_OTHER" => IMAGE_SCN_LNK_OTHER, + "IMAGE_SCN_LNK_INFO" => IMAGE_SCN_LNK_INFO, + "IMAGE_SCN_LNK_REMOVE" => IMAGE_SCN_LNK_REMOVE, + "IMAGE_SCN_LNK_COMDAT" => IMAGE_SCN_LNK_COMDAT, + "IMAGE_SCN_GPREL" => IMAGE_SCN_GPREL, + "IMAGE_SCN_MEM_PURGEABLE" => IMAGE_SCN_MEM_PURGEABLE, + "IMAGE_SCN_MEM_LOCKED" => IMAGE_SCN_MEM_LOCKED, + "IMAGE_SCN_MEM_PRELOAD" => IMAGE_SCN_MEM_PRELOAD, + "IMAGE_SCN_LNK_NRELOC_OVFL" => IMAGE_SCN_LNK_NRELOC_OVFL, + "IMAGE_SCN_MEM_DISCARDABLE" => IMAGE_SCN_MEM_DISCARDABLE, + "IMAGE_SCN_MEM_NOT_CACHED" => IMAGE_SCN_MEM_NOT_CACHED, + "IMAGE_SCN_MEM_NOT_PAGED" => IMAGE_SCN_MEM_NOT_PAGED, + "IMAGE_SCN_MEM_SHARED" => IMAGE_SCN_MEM_SHARED, + "IMAGE_SCN_MEM_EXECUTE" => IMAGE_SCN_MEM_EXECUTE, + "IMAGE_SCN_MEM_READ" => IMAGE_SCN_MEM_READ, + "IMAGE_SCN_MEM_WRITE" => IMAGE_SCN_MEM_WRITE, + _ => { + return Err(serde::de::Error::custom(format!( + "invalid section characteristic value {val}" + ))); + } + } + } + + Ok(flags) + } + } + + deserializer.deserialize_any(CharacteristicsVisitor) +} + +fn characteristics_serializer(characteristics: &u32, serializer: S) -> Result +where + S: Serializer, +{ + const FLAGS: [(u32, &str); 20] = [ + (IMAGE_SCN_TYPE_NO_PAD, "IMAGE_SCN_TYPE_NO_PAD"), + (IMAGE_SCN_CNT_CODE, "IMAGE_SCN_CNT_CODE"), + ( + IMAGE_SCN_CNT_INITIALIZED_DATA, + "IMAGE_SCN_CNT_INITIALIZED_DATA", + ), + ( + IMAGE_SCN_CNT_UNINITIALIZED_DATA, + "IMAGE_SCN_CNT_UNINITIALIZED_DATA", + ), + (IMAGE_SCN_LNK_OTHER, "IMAGE_SCN_LNK_OTHER"), + (IMAGE_SCN_LNK_INFO, "IMAGE_SCN_LNK_INFO"), + (IMAGE_SCN_LNK_REMOVE, "IMAGE_SCN_LNK_REMOVE"), + (IMAGE_SCN_LNK_COMDAT, "IMAGE_SCN_LNK_COMDAT"), + (IMAGE_SCN_GPREL, "IMAGE_SCN_GPREL"), + (IMAGE_SCN_MEM_PURGEABLE, "IMAGE_SCN_MEM_PURGEABLE"), + (IMAGE_SCN_MEM_LOCKED, "IMAGE_SCN_MEM_LOCKED"), + (IMAGE_SCN_MEM_PRELOAD, "IMAGE_SCN_MEM_PRELOAD"), + (IMAGE_SCN_LNK_NRELOC_OVFL, "IMAGE_SCN_LNK_NRELOC_OVFL"), + (IMAGE_SCN_MEM_DISCARDABLE, "IMAGE_SCN_MEM_DISCARDABLE"), + (IMAGE_SCN_MEM_NOT_CACHED, "IMAGE_SCN_MEM_NOT_CACHED"), + (IMAGE_SCN_MEM_NOT_PAGED, "IMAGE_SCN_MEM_NOT_PAGED"), + (IMAGE_SCN_MEM_SHARED, "IMAGE_SCN_MEM_SHARED"), + (IMAGE_SCN_MEM_EXECUTE, "IMAGE_SCN_MEM_EXECUTE"), + (IMAGE_SCN_MEM_READ, "IMAGE_SCN_MEM_READ"), + (IMAGE_SCN_MEM_WRITE, "IMAGE_SCN_MEM_WRITE"), + ]; + + let contains_flag = |bits: u32, flag: u32| bits & flag != 0; + + let flagcount = FLAGS + .iter() + .filter(|(flag, _)| contains_flag(*characteristics, *flag)) + .count(); + + let mut seq = serializer.serialize_seq(Some(flagcount))?; + for (flag, val) in FLAGS { + if contains_flag(*characteristics, flag) { + seq.serialize_element(val)?; + } + } + + seq.end() +} + +#[derive(Debug, Clone, Default, Deserialize, Serialize)] +#[serde(rename_all = "PascalCase")] +pub struct CoffYamlSectionRelocation { + pub virtual_address: u32, + pub symbol_name: String, + + #[serde(rename = "Type", deserialize_with = "relocation_type_deserializer")] + pub typ: u16, +} + +fn relocation_type_deserializer<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + struct RelocationTypeVisitor; + + impl Visitor<'_> for RelocationTypeVisitor { + type Value = u16; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("'IMAGE_REL_*' string or integer") + } + + fn visit_u64(self, v: u64) -> Result + where + E: de::Error, + { + u16::try_from(v).map_err(serde::de::Error::custom) + } + + fn visit_str(self, v: &str) -> Result + where + E: de::Error, + { + Ok(match v { + "IMAGE_REL_AMD64_ABSOLUTE" => IMAGE_REL_AMD64_ABSOLUTE, + "IMAGE_REL_AMD64_ADDR64" => IMAGE_REL_AMD64_ADDR64, + "IMAGE_REL_AMD64_ADDR32" => IMAGE_REL_AMD64_ADDR32, + "IMAGE_REL_AMD64_ADDR32NB" => IMAGE_REL_AMD64_ADDR32NB, + "IMAGE_REL_AMD64_REL32" => IMAGE_REL_AMD64_REL32, + "IMAGE_REL_AMD64_REL32_1" => IMAGE_REL_AMD64_REL32_1, + "IMAGE_REL_AMD64_REL32_2" => IMAGE_REL_AMD64_REL32_2, + "IMAGE_REL_AMD64_REL32_3" => IMAGE_REL_AMD64_REL32_3, + "IMAGE_REL_AMD64_REL32_4" => IMAGE_REL_AMD64_REL32_4, + "IMAGE_REL_AMD64_REL32_5" => IMAGE_REL_AMD64_REL32_5, + "IMAGE_REL_AMD64_SECTION" => IMAGE_REL_AMD64_SECTION, + "IMAGE_REL_AMD64_SECREL" => IMAGE_REL_AMD64_SECREL, + "IMAGE_REL_AMD64_SECREL7" => IMAGE_REL_AMD64_SECREL7, + "IMAGE_REL_AMD64_TOKEN" => IMAGE_REL_AMD64_TOKEN, + "IMAGE_REL_AMD64_SREL32" => IMAGE_REL_AMD64_SREL32, + "IMAGE_REL_AMD64_PAIR" => IMAGE_REL_AMD64_PAIR, + "IMAGE_REL_AMD64_SSPAN32" => IMAGE_REL_AMD64_SSPAN32, + "IMAGE_REL_I386_ABSOLUTE" => IMAGE_REL_I386_ABSOLUTE, + "IMAGE_REL_I386_DIR16" => IMAGE_REL_I386_DIR16, + "IMAGE_REL_I386_DIR32" => IMAGE_REL_I386_DIR32, + "IMAGE_REL_I386_DIR32NB" => IMAGE_REL_I386_DIR32NB, + "IMAGE_REL_I386_SEG12" => IMAGE_REL_I386_SEG12, + "IMAGE_REL_I386_SECTION" => IMAGE_REL_I386_SECTION, + "IMAGE_REL_I386_SECREL" => IMAGE_REL_I386_SECREL, + "IMAGE_REL_I386_TOKEN" => IMAGE_REL_I386_TOKEN, + "IMAGE_REL_I386_SECREL7" => IMAGE_REL_I386_SECREL7, + "IMAGE_REL_I386_REL16" => IMAGE_REL_I386_REL16, + "IMAGE_REL_I386_REL32" => IMAGE_REL_I386_REL32, + _ => { + return Err(serde::de::Error::custom(format!( + "invalid relocation type {v}" + ))); + } + }) + } + } + + deserializer.deserialize_any(RelocationTypeVisitor) +} + +#[cfg(test)] +mod tests { + use super::{characteristics_deserializer, relocation_type_deserializer}; + use crate::testutils; + use object::pe::{ + IMAGE_REL_AMD64_ADDR32, IMAGE_REL_I386_DIR32, IMAGE_SCN_CNT_CODE, + IMAGE_SCN_CNT_INITIALIZED_DATA, + }; + + #[test] + fn characteristic_scalar_deserialize() { + testutils::run_deserializer_tests( + characteristics_deserializer, + [("32", IMAGE_SCN_CNT_CODE)], + ); + } + + #[test] + fn characteristic_string_deserialize() { + testutils::run_deserializer_tests( + characteristics_deserializer, + [ + ("[ IMAGE_SCN_CNT_CODE ]", IMAGE_SCN_CNT_CODE), + ( + "[ IMAGE_SCN_CNT_CODE, IMAGE_SCN_CNT_INITIALIZED_DATA ]", + IMAGE_SCN_CNT_CODE | IMAGE_SCN_CNT_INITIALIZED_DATA, + ), + ], + ); + } + + #[test] + fn relocation_type_string_deserialize() { + testutils::run_deserializer_tests( + relocation_type_deserializer, + [ + ("IMAGE_REL_AMD64_ADDR32", IMAGE_REL_AMD64_ADDR32), + ("IMAGE_REL_I386_DIR32", IMAGE_REL_I386_DIR32), + ], + ); + } +} diff --git a/crates/coffyaml/src/coff/symbols.rs b/crates/coffyaml/src/coff/symbols.rs new file mode 100644 index 0000000..c5f7910 --- /dev/null +++ b/crates/coffyaml/src/coff/symbols.rs @@ -0,0 +1,773 @@ +use object::pe::{ + IMAGE_COMDAT_SELECT_ANY, IMAGE_COMDAT_SELECT_ASSOCIATIVE, IMAGE_COMDAT_SELECT_EXACT_MATCH, + IMAGE_COMDAT_SELECT_LARGEST, IMAGE_COMDAT_SELECT_NODUPLICATES, IMAGE_COMDAT_SELECT_SAME_SIZE, + IMAGE_SYM_ABSOLUTE, IMAGE_SYM_CLASS_ARGUMENT, IMAGE_SYM_CLASS_AUTOMATIC, + IMAGE_SYM_CLASS_BIT_FIELD, IMAGE_SYM_CLASS_BLOCK, IMAGE_SYM_CLASS_CLR_TOKEN, + IMAGE_SYM_CLASS_END_OF_FUNCTION, IMAGE_SYM_CLASS_END_OF_STRUCT, IMAGE_SYM_CLASS_ENUM_TAG, + IMAGE_SYM_CLASS_EXTERNAL, IMAGE_SYM_CLASS_EXTERNAL_DEF, IMAGE_SYM_CLASS_FILE, + IMAGE_SYM_CLASS_FUNCTION, IMAGE_SYM_CLASS_LABEL, IMAGE_SYM_CLASS_MEMBER_OF_ENUM, + IMAGE_SYM_CLASS_MEMBER_OF_STRUCT, IMAGE_SYM_CLASS_MEMBER_OF_UNION, IMAGE_SYM_CLASS_NULL, + IMAGE_SYM_CLASS_REGISTER, IMAGE_SYM_CLASS_REGISTER_PARAM, IMAGE_SYM_CLASS_SECTION, + IMAGE_SYM_CLASS_STATIC, IMAGE_SYM_CLASS_STRUCT_TAG, IMAGE_SYM_CLASS_TYPE_DEFINITION, + IMAGE_SYM_CLASS_UNDEFINED_LABEL, IMAGE_SYM_CLASS_UNDEFINED_STATIC, IMAGE_SYM_CLASS_UNION_TAG, + IMAGE_SYM_CLASS_WEAK_EXTERNAL, IMAGE_SYM_DEBUG, IMAGE_SYM_DTYPE_ARRAY, + IMAGE_SYM_DTYPE_FUNCTION, IMAGE_SYM_DTYPE_NULL, IMAGE_SYM_DTYPE_POINTER, IMAGE_SYM_TYPE_BYTE, + IMAGE_SYM_TYPE_CHAR, IMAGE_SYM_TYPE_DOUBLE, IMAGE_SYM_TYPE_DWORD, IMAGE_SYM_TYPE_ENUM, + IMAGE_SYM_TYPE_FLOAT, IMAGE_SYM_TYPE_INT, IMAGE_SYM_TYPE_LONG, IMAGE_SYM_TYPE_MOE, + IMAGE_SYM_TYPE_NULL, IMAGE_SYM_TYPE_SHORT, IMAGE_SYM_TYPE_STRUCT, IMAGE_SYM_TYPE_UINT, + IMAGE_SYM_TYPE_UNION, IMAGE_SYM_TYPE_VOID, IMAGE_SYM_TYPE_WORD, IMAGE_SYM_UNDEFINED, +}; +use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Visitor}; +use serde_yml::with::singleton_map_optional; + +#[derive(Debug, Clone, Default, Deserialize, PartialEq, Eq, Serialize)] +#[serde(rename_all = "PascalCase")] +pub struct CoffYamlSymbol { + pub name: String, + pub value: u32, + + #[serde( + deserialize_with = "section_number_deserializer", + serialize_with = "section_number_serializer" + )] + pub section_number: i32, + + #[serde( + deserialize_with = "simple_type_deserializer", + serialize_with = "simple_type_serializer" + )] + pub simple_type: u16, + + #[serde( + deserialize_with = "complex_type_deserializer", + serialize_with = "complex_type_serializer" + )] + pub complex_type: u16, + + #[serde( + deserialize_with = "storage_class_deserializer", + serialize_with = "storage_class_serializer" + )] + pub storage_class: u8, + + #[serde( + default, + with = "singleton_map_optional", + skip_serializing_if = "Option::is_none" + )] + pub section_definition: Option, + + #[serde( + default, + with = "singleton_map_optional", + skip_serializing_if = "Option::is_none" + )] + pub function_definition: Option, + + #[serde(default, skip_serializing_if = "Option::is_none")] + pub file: Option, +} + +fn section_number_deserializer<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + struct SectionNumberVisitor; + + impl Visitor<'_> for SectionNumberVisitor { + type Value = i32; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str( + "IMAGE_SYM_UNDEFINED, IMAGE_SYM_ABSOLUTE, IMAGE_SYM_DEBUG string or integer", + ) + } + + fn visit_u64(self, v: u64) -> Result + where + E: serde::de::Error, + { + i32::try_from(v).map_err(serde::de::Error::custom) + } + + fn visit_i64(self, v: i64) -> Result + where + E: serde::de::Error, + { + i32::try_from(v).map_err(serde::de::Error::custom) + } + + fn visit_str(self, v: &str) -> Result + where + E: serde::de::Error, + { + Ok(match v { + "IMAGE_SYM_UNDEFINED" => IMAGE_SYM_UNDEFINED, + "IMAGE_SYM_ABSOLUTE" => IMAGE_SYM_ABSOLUTE, + "IMAGE_SYM_DEBUG" => IMAGE_SYM_DEBUG, + _ => { + return Err(serde::de::Error::custom(format!( + "invalid section number string {v}" + ))); + } + }) + } + } + + deserializer.deserialize_any(SectionNumberVisitor) +} + +fn section_number_serializer(section_number: &i32, serializer: S) -> Result +where + S: Serializer, +{ + match *section_number { + IMAGE_SYM_UNDEFINED => serializer.serialize_str("IMAGE_SYM_UNDEFINED"), + IMAGE_SYM_ABSOLUTE => serializer.serialize_str("IMAGE_SYM_ABSOLUTE"), + IMAGE_SYM_DEBUG => serializer.serialize_str("IMAGE_SYM_DEBUG"), + o if o >= 1 => serializer.serialize_i32(o), + _ => Err(serde::ser::Error::custom(format!( + "invalid section number {section_number}" + ))), + } +} + +fn simple_type_deserializer<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + struct SimpleTypeVisitor; + + impl Visitor<'_> for SimpleTypeVisitor { + type Value = u16; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("'IMAGE_SYM_TYPE_*' string or integer") + } + + fn visit_u64(self, v: u64) -> Result + where + E: serde::de::Error, + { + u16::try_from(v).map_err(serde::de::Error::custom) + } + + fn visit_str(self, v: &str) -> Result + where + E: serde::de::Error, + { + Ok(match v { + "IMAGE_SYM_TYPE_NULL" => IMAGE_SYM_TYPE_NULL, + "IMAGE_SYM_TYPE_VOID" => IMAGE_SYM_TYPE_VOID, + "IMAGE_SYM_TYPE_CHAR" => IMAGE_SYM_TYPE_CHAR, + "IMAGE_SYM_TYPE_SHORT" => IMAGE_SYM_TYPE_SHORT, + "IMAGE_SYM_TYPE_INT" => IMAGE_SYM_TYPE_INT, + "IMAGE_SYM_TYPE_LONG" => IMAGE_SYM_TYPE_LONG, + "IMAGE_SYM_TYPE_FLOAT" => IMAGE_SYM_TYPE_FLOAT, + "IMAGE_SYM_TYPE_DOUBLE" => IMAGE_SYM_TYPE_DOUBLE, + "IMAGE_SYM_TYPE_STRUCT" => IMAGE_SYM_TYPE_STRUCT, + "IMAGE_SYM_TYPE_UNION" => IMAGE_SYM_TYPE_UNION, + "IMAGE_SYM_TYPE_ENUM" => IMAGE_SYM_TYPE_ENUM, + "IMAGE_SYM_TYPE_MOE" => IMAGE_SYM_TYPE_MOE, + "IMAGE_SYM_TYPE_BYTE" => IMAGE_SYM_TYPE_BYTE, + "IMAGE_SYM_TYPE_WORD" => IMAGE_SYM_TYPE_WORD, + "IMAGE_SYM_TYPE_UINT" => IMAGE_SYM_TYPE_UINT, + "IMAGE_SYM_TYPE_DWORD" => IMAGE_SYM_TYPE_DWORD, + _ => { + return Err(serde::de::Error::custom(format!( + "invalid symbol simple type {v}" + ))); + } + }) + } + } + + deserializer.deserialize_any(SimpleTypeVisitor) +} + +fn simple_type_serializer(simple_type: &u16, serializer: S) -> Result +where + S: Serializer, +{ + match *simple_type { + IMAGE_SYM_TYPE_NULL => serializer.serialize_str("IMAGE_SYM_TYPE_NULL"), + IMAGE_SYM_TYPE_VOID => serializer.serialize_str("IMAGE_SYM_TYPE_VOID"), + IMAGE_SYM_TYPE_CHAR => serializer.serialize_str("IMAGE_SYM_TYPE_CHAR"), + IMAGE_SYM_TYPE_SHORT => serializer.serialize_str("IMAGE_SYM_TYPE_SHORT"), + IMAGE_SYM_TYPE_INT => serializer.serialize_str("IMAGE_SYM_TYPE_INT"), + IMAGE_SYM_TYPE_LONG => serializer.serialize_str("IMAGE_SYM_TYPE_LONG"), + IMAGE_SYM_TYPE_FLOAT => serializer.serialize_str("IMAGE_SYM_TYPE_FLOAT"), + IMAGE_SYM_TYPE_DOUBLE => serializer.serialize_str("IMAGE_SYM_TYPE_DOUBLE"), + IMAGE_SYM_TYPE_STRUCT => serializer.serialize_str("IMAGE_SYM_TYPE_STRUCT"), + IMAGE_SYM_TYPE_UNION => serializer.serialize_str("IMAGE_SYM_TYPE_UNION"), + IMAGE_SYM_TYPE_ENUM => serializer.serialize_str("IMAGE_SYM_TYPE_ENUM"), + IMAGE_SYM_TYPE_MOE => serializer.serialize_str("IMAGE_SYM_TYPE_MOE"), + IMAGE_SYM_TYPE_BYTE => serializer.serialize_str("IMAGE_SYM_TYPE_BYTE"), + IMAGE_SYM_TYPE_WORD => serializer.serialize_str("IMAGE_SYM_TYPE_WORD"), + IMAGE_SYM_TYPE_UINT => serializer.serialize_str("IMAGE_SYM_TYPE_UINT"), + IMAGE_SYM_TYPE_DWORD => serializer.serialize_str("IMAGE_SYM_TYPE_DWORD"), + o => serializer.serialize_u16(o), + } +} + +fn complex_type_deserializer<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + struct ComplexTypeVisitor; + + impl Visitor<'_> for ComplexTypeVisitor { + type Value = u16; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("'IMAGE_SYM_DTYPE_*' string or integer") + } + + fn visit_u64(self, v: u64) -> Result + where + E: serde::de::Error, + { + u16::try_from(v).map_err(serde::de::Error::custom) + } + + fn visit_str(self, v: &str) -> Result + where + E: serde::de::Error, + { + Ok(match v { + "IMAGE_SYM_DTYPE_NULL" => IMAGE_SYM_DTYPE_NULL, + "IMAGE_SYM_DTYPE_POINTER" => IMAGE_SYM_DTYPE_POINTER, + "IMAGE_SYM_DTYPE_FUNCTION" => IMAGE_SYM_DTYPE_FUNCTION, + "IMAGE_SYM_DTYPE_ARRAY" => IMAGE_SYM_DTYPE_ARRAY, + _ => { + return Err(serde::de::Error::custom(format!( + "invalid symbol complex type {v}" + ))); + } + }) + } + } + + deserializer.deserialize_any(ComplexTypeVisitor) +} + +fn complex_type_serializer(complex_type: &u16, serializer: S) -> Result +where + S: Serializer, +{ + match *complex_type { + IMAGE_SYM_DTYPE_NULL => serializer.serialize_str("IMAGE_SYM_DTYPE_NULL"), + IMAGE_SYM_DTYPE_POINTER => serializer.serialize_str("IMAGE_SYM_DTYPE_POINTER"), + IMAGE_SYM_DTYPE_FUNCTION => serializer.serialize_str("IMAGE_SYM_DTYPE_FUNCTION"), + IMAGE_SYM_DTYPE_ARRAY => serializer.serialize_str("IMAGE_SYM_DTYPE_ARRAY"), + o => serializer.serialize_u16(o), + } +} + +fn storage_class_deserializer<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + struct StorageClassVisitor; + + impl Visitor<'_> for StorageClassVisitor { + type Value = u8; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("'IMAGE_SYM_CLASS_*' string or integer") + } + + fn visit_u64(self, v: u64) -> Result + where + E: serde::de::Error, + { + u8::try_from(v).map_err(serde::de::Error::custom) + } + + fn visit_i64(self, v: i64) -> Result + where + E: serde::de::Error, + { + Ok(i8::try_from(v).map_err(serde::de::Error::custom)? as u8) + } + + fn visit_str(self, v: &str) -> Result + where + E: serde::de::Error, + { + Ok(match v { + "IMAGE_SYM_CLASS_END_OF_FUNCTION" => IMAGE_SYM_CLASS_END_OF_FUNCTION, + "IMAGE_SYM_CLASS_NULL" => IMAGE_SYM_CLASS_NULL, + "IMAGE_SYM_CLASS_AUTOMATIC" => IMAGE_SYM_CLASS_AUTOMATIC, + "IMAGE_SYM_CLASS_EXTERNAL" => IMAGE_SYM_CLASS_EXTERNAL, + "IMAGE_SYM_CLASS_STATIC" => IMAGE_SYM_CLASS_STATIC, + "IMAGE_SYM_CLASS_REGISTER" => IMAGE_SYM_CLASS_REGISTER, + "IMAGE_SYM_CLASS_EXTERNAL_DEF" => IMAGE_SYM_CLASS_EXTERNAL_DEF, + "IMAGE_SYM_CLASS_LABEL" => IMAGE_SYM_CLASS_LABEL, + "IMAGE_SYM_CLASS_UNDEFINED_LABEL" => IMAGE_SYM_CLASS_UNDEFINED_LABEL, + "IMAGE_SYM_CLASS_MEMBER_OF_STRUCT" => IMAGE_SYM_CLASS_MEMBER_OF_STRUCT, + "IMAGE_SYM_CLASS_ARGUMENT" => IMAGE_SYM_CLASS_ARGUMENT, + "IMAGE_SYM_CLASS_STRUCT_TAG" => IMAGE_SYM_CLASS_STRUCT_TAG, + "IMAGE_SYM_CLASS_MEMBER_OF_UNION" => IMAGE_SYM_CLASS_MEMBER_OF_UNION, + "IMAGE_SYM_CLASS_UNION_TAG" => IMAGE_SYM_CLASS_UNION_TAG, + "IMAGE_SYM_CLASS_TYPE_DEFINITION" => IMAGE_SYM_CLASS_TYPE_DEFINITION, + "IMAGE_SYM_CLASS_UNDEFINED_STATIC" => IMAGE_SYM_CLASS_UNDEFINED_STATIC, + "IMAGE_SYM_CLASS_ENUM_TAG" => IMAGE_SYM_CLASS_ENUM_TAG, + "IMAGE_SYM_CLASS_MEMBER_OF_ENUM" => IMAGE_SYM_CLASS_MEMBER_OF_ENUM, + "IMAGE_SYM_CLASS_REGISTER_PARAM" => IMAGE_SYM_CLASS_REGISTER_PARAM, + "IMAGE_SYM_CLASS_BIT_FIELD" => IMAGE_SYM_CLASS_BIT_FIELD, + "IMAGE_SYM_CLASS_BLOCK" => IMAGE_SYM_CLASS_BLOCK, + "IMAGE_SYM_CLASS_FUNCTION" => IMAGE_SYM_CLASS_FUNCTION, + "IMAGE_SYM_CLASS_END_OF_STRUCT" => IMAGE_SYM_CLASS_END_OF_STRUCT, + "IMAGE_SYM_CLASS_FILE" => IMAGE_SYM_CLASS_FILE, + "IMAGE_SYM_CLASS_SECTION" => IMAGE_SYM_CLASS_SECTION, + "IMAGE_SYM_CLASS_WEAK_EXTERNAL" => IMAGE_SYM_CLASS_WEAK_EXTERNAL, + "IMAGE_SYM_CLASS_CLR_TOKEN" => IMAGE_SYM_CLASS_CLR_TOKEN, + _ => { + return Err(serde::de::Error::custom(format!( + "invalid symbol complex type {v}" + ))); + } + }) + } + } + + deserializer.deserialize_any(StorageClassVisitor) +} + +fn storage_class_serializer(storage_class: &u8, serializer: S) -> Result +where + S: Serializer, +{ + match *storage_class { + IMAGE_SYM_CLASS_END_OF_FUNCTION => { + serializer.serialize_str("IMAGE_SYM_CLASS_END_OF_FUNCTION") + } + IMAGE_SYM_CLASS_NULL => serializer.serialize_str("IMAGE_SYM_CLASS_NULL"), + IMAGE_SYM_CLASS_AUTOMATIC => serializer.serialize_str("IMAGE_SYM_CLASS_AUTOMATIC"), + IMAGE_SYM_CLASS_EXTERNAL => serializer.serialize_str("IMAGE_SYM_CLASS_EXTERNAL"), + IMAGE_SYM_CLASS_STATIC => serializer.serialize_str("IMAGE_SYM_CLASS_STATIC"), + IMAGE_SYM_CLASS_REGISTER => serializer.serialize_str("IMAGE_SYM_CLASS_REGISTER"), + IMAGE_SYM_CLASS_EXTERNAL_DEF => serializer.serialize_str("IMAGE_SYM_CLASS_EXTERNAL_DEF"), + IMAGE_SYM_CLASS_LABEL => serializer.serialize_str("IMAGE_SYM_CLASS_LABEL"), + IMAGE_SYM_CLASS_UNDEFINED_LABEL => { + serializer.serialize_str("IMAGE_SYM_CLASS_UNDEFINED_LABEL") + } + IMAGE_SYM_CLASS_MEMBER_OF_STRUCT => { + serializer.serialize_str("IMAGE_SYM_CLASS_MEMBER_OF_STRUCT") + } + IMAGE_SYM_CLASS_ARGUMENT => serializer.serialize_str("IMAGE_SYM_CLASS_ARGUMENT"), + IMAGE_SYM_CLASS_STRUCT_TAG => serializer.serialize_str("IMAGE_SYM_CLASS_STRUCT_TAG"), + IMAGE_SYM_CLASS_MEMBER_OF_UNION => { + serializer.serialize_str("IMAGE_SYM_CLASS_MEMBER_OF_UNION") + } + IMAGE_SYM_CLASS_UNION_TAG => serializer.serialize_str("IMAGE_SYM_CLASS_UNION_TAG"), + IMAGE_SYM_CLASS_TYPE_DEFINITION => { + serializer.serialize_str("IMAGE_SYM_CLASS_TYPE_DEFINITION") + } + IMAGE_SYM_CLASS_UNDEFINED_STATIC => { + serializer.serialize_str("IMAGE_SYM_CLASS_UNDEFINED_STATIC") + } + IMAGE_SYM_CLASS_ENUM_TAG => serializer.serialize_str("IMAGE_SYM_CLASS_ENUM_TAG"), + IMAGE_SYM_CLASS_MEMBER_OF_ENUM => { + serializer.serialize_str("IMAGE_SYM_CLASS_MEMBER_OF_ENUM") + } + IMAGE_SYM_CLASS_REGISTER_PARAM => { + serializer.serialize_str("IMAGE_SYM_CLASS_REGISTER_PARAM") + } + IMAGE_SYM_CLASS_BIT_FIELD => serializer.serialize_str("IMAGE_SYM_CLASS_BIT_FIELD"), + IMAGE_SYM_CLASS_BLOCK => serializer.serialize_str("IMAGE_SYM_CLASS_BLOCK"), + IMAGE_SYM_CLASS_FUNCTION => serializer.serialize_str("IMAGE_SYM_CLASS_FUNCTION"), + IMAGE_SYM_CLASS_END_OF_STRUCT => serializer.serialize_str("IMAGE_SYM_CLASS_END_OF_STRUCT"), + IMAGE_SYM_CLASS_FILE => serializer.serialize_str("IMAGE_SYM_CLASS_FILE"), + IMAGE_SYM_CLASS_SECTION => serializer.serialize_str("IMAGE_SYM_CLASS_SECTION"), + IMAGE_SYM_CLASS_WEAK_EXTERNAL => serializer.serialize_str("IMAGE_SYM_CLASS_WEAK_EXTERNAL"), + IMAGE_SYM_CLASS_CLR_TOKEN => serializer.serialize_str("IMAGE_SYM_CLASS_CLR_TOKEN"), + o => serializer.serialize_u8(o), + } +} + +#[derive(Debug, Default, Deserialize, Serialize, Clone, PartialEq, Eq)] +#[serde(rename_all = "PascalCase")] +pub struct CoffYamlAuxFunctionDefinition { + pub tag_index: u32, + pub total_size: u32, + pub pointer_to_linenumber: u32, + pub pointer_to_next_function: u32, +} + +#[derive(Debug, Clone, Default, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "PascalCase")] +pub struct CoffYamlAuxSectionDefinition { + pub length: u32, + pub number_of_relocations: u16, + pub number_of_linenumbers: u16, + pub check_sum: u32, + pub number: u16, + + #[serde( + default, + deserialize_with = "aux_section_comdat_selection_deserializer", + serialize_with = "aux_section_comdat_selection_serializer" + )] + pub selection: u8, +} + +fn aux_section_comdat_selection_deserializer<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + struct ComdatSelectionVisitor; + + impl<'de> Visitor<'de> for ComdatSelectionVisitor { + type Value = u8; + + fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result { + formatter.write_str("'IMAGE_COMDAT_SELECT_*' string or integer") + } + + fn visit_u64(self, v: u64) -> Result + where + E: serde::de::Error, + { + u8::try_from(v).map_err(serde::de::Error::custom) + } + + fn visit_none(self) -> Result + where + E: serde::de::Error, + { + Ok(0) + } + + fn visit_some(self, deserializer: D) -> Result + where + D: Deserializer<'de>, + { + deserializer.deserialize_any(Self) + } + + fn visit_str(self, v: &str) -> Result + where + E: serde::de::Error, + { + Ok(match v { + "IMAGE_COMDAT_SELECT_NODUPLICATES" => IMAGE_COMDAT_SELECT_NODUPLICATES, + "IMAGE_COMDAT_SELECT_ANY" => IMAGE_COMDAT_SELECT_ANY, + "IMAGE_COMDAT_SELECT_SAME_SIZE" => IMAGE_COMDAT_SELECT_SAME_SIZE, + "IMAGE_COMDAT_SELECT_EXACT_MATCH" => IMAGE_COMDAT_SELECT_EXACT_MATCH, + "IMAGE_COMDAT_SELECT_ASSOCIATIVE" => IMAGE_COMDAT_SELECT_ASSOCIATIVE, + "IMAGE_COMDAT_SELECT_LARGEST" => IMAGE_COMDAT_SELECT_LARGEST, + _ => { + return Err(serde::de::Error::custom(format!( + "invalid COMDAT selection type {v}" + ))); + } + }) + } + } + + deserializer.deserialize_any(ComdatSelectionVisitor) +} + +fn aux_section_comdat_selection_serializer( + selection: &u8, + serializer: S, +) -> Result +where + S: Serializer, +{ + match *selection { + IMAGE_COMDAT_SELECT_NODUPLICATES => { + serializer.serialize_str("IMAGE_COMDAT_SELECT_NODUPLICATES") + } + IMAGE_COMDAT_SELECT_ANY => serializer.serialize_str("IMAGE_COMDAT_SELECT_ANY"), + IMAGE_COMDAT_SELECT_SAME_SIZE => serializer.serialize_str("IMAGE_COMDAT_SELECT_SAME_SIZE"), + IMAGE_COMDAT_SELECT_EXACT_MATCH => { + serializer.serialize_str("IMAGE_COMDAT_SELECT_EXACT_MATCH") + } + IMAGE_COMDAT_SELECT_ASSOCIATIVE => { + serializer.serialize_str("IMAGE_COMDAT_SELECT_ASSOCIATIVE") + } + IMAGE_COMDAT_SELECT_LARGEST => serializer.serialize_str("IMAGE_COMDAT_SELECT_LARGEST"), + o => serializer.serialize_u8(o), + } +} + +#[cfg(test)] +mod tests { + use super::{ + CoffYamlAuxFunctionDefinition, CoffYamlSymbol, aux_section_comdat_selection_deserializer, + complex_type_deserializer, section_number_deserializer, simple_type_deserializer, + storage_class_deserializer, + }; + use crate::testutils; + use object::pe::{ + IMAGE_COMDAT_SELECT_ANY, IMAGE_COMDAT_SELECT_LARGEST, IMAGE_COMDAT_SELECT_NODUPLICATES, + IMAGE_SYM_ABSOLUTE, IMAGE_SYM_CLASS_END_OF_FUNCTION, IMAGE_SYM_CLASS_EXTERNAL, + IMAGE_SYM_CLASS_NULL, IMAGE_SYM_CLASS_STATIC, IMAGE_SYM_DEBUG, IMAGE_SYM_DTYPE_FUNCTION, + IMAGE_SYM_DTYPE_NULL, IMAGE_SYM_TYPE_NULL, IMAGE_SYM_TYPE_VOID, IMAGE_SYM_UNDEFINED, + }; + use serde::Deserialize; + + #[test] + fn section_number_scalar_deserialize() { + testutils::run_deserializer_tests( + section_number_deserializer, + [("0", 0), ("1", 1), ("-1", -1), ("-2", -2)], + ); + } + + #[test] + fn section_number_string_deserialize() { + testutils::run_deserializer_tests( + section_number_deserializer, + [ + ("IMAGE_SYM_UNDEFINED", IMAGE_SYM_UNDEFINED), + ("IMAGE_SYM_ABSOLUTE", IMAGE_SYM_ABSOLUTE), + ("IMAGE_SYM_DEBUG", IMAGE_SYM_DEBUG), + ], + ); + } + + #[test] + fn simple_type_scalar_deserialize() { + testutils::run_deserializer_tests( + simple_type_deserializer, + [("0", IMAGE_SYM_TYPE_NULL), ("1", IMAGE_SYM_TYPE_VOID)], + ); + } + + #[test] + fn simple_type_string_deserialize() { + testutils::run_deserializer_tests( + simple_type_deserializer, + [ + ("IMAGE_SYM_TYPE_NULL", IMAGE_SYM_TYPE_NULL), + ("IMAGE_SYM_TYPE_VOID", IMAGE_SYM_TYPE_VOID), + ], + ); + } + + #[test] + fn complex_type_scalar_deserialize() { + testutils::run_deserializer_tests( + complex_type_deserializer, + [("0", IMAGE_SYM_DTYPE_NULL), ("2", IMAGE_SYM_DTYPE_FUNCTION)], + ); + } + + #[test] + fn complex_type_string_deserialize() { + testutils::run_deserializer_tests( + complex_type_deserializer, + [ + ("IMAGE_SYM_DTYPE_NULL", IMAGE_SYM_DTYPE_NULL), + ("IMAGE_SYM_DTYPE_FUNCTION", IMAGE_SYM_DTYPE_FUNCTION), + ], + ); + } + + #[test] + fn storage_class_scalar_deserialize() { + testutils::run_deserializer_tests( + storage_class_deserializer, + [ + ("-1", IMAGE_SYM_CLASS_END_OF_FUNCTION), + ("0", IMAGE_SYM_CLASS_NULL), + ("2", IMAGE_SYM_CLASS_EXTERNAL), + ("3", IMAGE_SYM_CLASS_STATIC), + ], + ); + } + + #[test] + fn storage_class_string_deserialize() { + testutils::run_deserializer_tests( + storage_class_deserializer, + [ + ( + "IMAGE_SYM_CLASS_END_OF_FUNCTION", + IMAGE_SYM_CLASS_END_OF_FUNCTION, + ), + ("IMAGE_SYM_CLASS_NULL", IMAGE_SYM_CLASS_NULL), + ("IMAGE_SYM_CLASS_EXTERNAL", IMAGE_SYM_CLASS_EXTERNAL), + ("IMAGE_SYM_CLASS_STATIC", IMAGE_SYM_CLASS_STATIC), + ], + ); + } + + #[test] + fn aux_section_comdat_selection_scalar_deserialize() { + testutils::run_deserializer_tests( + aux_section_comdat_selection_deserializer, + [ + ("1", IMAGE_COMDAT_SELECT_NODUPLICATES), + ("2", IMAGE_COMDAT_SELECT_ANY), + ("6", IMAGE_COMDAT_SELECT_LARGEST), + ], + ); + } + + #[test] + fn aux_section_comdat_selection_string_deserialize() { + testutils::run_deserializer_tests( + aux_section_comdat_selection_deserializer, + [ + ( + "IMAGE_COMDAT_SELECT_NODUPLICATES", + IMAGE_COMDAT_SELECT_NODUPLICATES, + ), + ("IMAGE_COMDAT_SELECT_ANY", IMAGE_COMDAT_SELECT_ANY), + ("IMAGE_COMDAT_SELECT_LARGEST", IMAGE_COMDAT_SELECT_LARGEST), + ], + ); + } + + #[test] + fn symbol_no_aux_deserialize() { + testutils::run_deserializer_tests( + CoffYamlSymbol::deserialize, + [ + ( + r#" + Name: all_scalars + Value: 0 + SectionNumber: 1 + SimpleType: 2 + ComplexType: 3 + StorageClass: 4 + "#, + CoffYamlSymbol { + name: "all_scalars".into(), + value: 0, + section_number: 1, + simple_type: 2, + complex_type: 3, + storage_class: 4, + function_definition: None, + section_definition: None, + file: None, + }, + ), + ( + r#" + Name: all_strings + Value: 0 + SectionNumber: IMAGE_SYM_DEBUG + SimpleType: IMAGE_SYM_TYPE_VOID + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + "#, + CoffYamlSymbol { + name: "all_strings".into(), + value: 0, + section_number: IMAGE_SYM_DEBUG, + simple_type: IMAGE_SYM_TYPE_VOID, + complex_type: IMAGE_SYM_DTYPE_FUNCTION, + storage_class: IMAGE_SYM_CLASS_EXTERNAL, + function_definition: None, + section_definition: None, + file: None, + }, + ), + ], + ); + } + + #[test] + fn symbol_with_aux_deserialize() { + testutils::run_deserializer_tests( + CoffYamlSymbol::deserialize, + [ + ( + r#" + Name: aux_section + Value: 0 + SectionNumber: 1 + SimpleType: 2 + ComplexType: 3 + StorageClass: 3 + SectionDefinition: + Length: 123 + NumberOfRelocations: 2 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + "#, + CoffYamlSymbol { + name: "aux_section".into(), + value: 0, + section_number: 1, + simple_type: 2, + complex_type: 3, + storage_class: 3, + section_definition: Some(super::CoffYamlAuxSectionDefinition { + length: 123, + number_of_relocations: 2, + number_of_linenumbers: 0, + check_sum: 0, + number: 0, + selection: 0, + }), + function_definition: None, + file: None, + }, + ), + ( + r#" + Name: aux_function + Value: 2 + SectionNumber: 23 + SimpleType: 4 + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: 2 + FunctionDefinition: + TagIndex: 0 + TotalSize: 1 + PointerToLinenumber: 2 + PointerToNextFunction: 3 + "#, + CoffYamlSymbol { + name: "aux_function".into(), + value: 2, + section_number: 23, + simple_type: 4, + complex_type: IMAGE_SYM_DTYPE_FUNCTION, + storage_class: 2, + function_definition: Some(CoffYamlAuxFunctionDefinition { + tag_index: 0, + total_size: 1, + pointer_to_linenumber: 2, + pointer_to_next_function: 3, + }), + section_definition: None, + file: None, + }, + ), + ( + r#" + Name: aux_file + Value: 3 + SectionNumber: 11 + SimpleType: IMAGE_SYM_TYPE_VOID + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: 5 + File: test.c + "#, + CoffYamlSymbol { + name: "aux_file".into(), + value: 3, + section_number: 11, + simple_type: IMAGE_SYM_TYPE_VOID, + complex_type: IMAGE_SYM_DTYPE_NULL, + storage_class: 5, + file: Some("test.c".into()), + section_definition: None, + function_definition: None, + }, + ), + ], + ); + } +} diff --git a/crates/coffyaml/src/importlib/archconfig.rs b/crates/coffyaml/src/importlib/archconfig.rs new file mode 100644 index 0000000..5250f0a --- /dev/null +++ b/crates/coffyaml/src/importlib/archconfig.rs @@ -0,0 +1,30 @@ +use object::pe::{IMAGE_FILE_MACHINE_AMD64, IMAGE_REL_AMD64_ADDR32NB}; + +use super::{Architecture, errors::ImportlibYamlBuildError}; + +pub(super) struct ArchitectureConfig { + machine: u16, + reloc_type: u16, +} + +impl ArchitectureConfig { + pub fn new(arch: Architecture) -> Result { + Ok(match arch { + Architecture::X86_64 => Self { + machine: IMAGE_FILE_MACHINE_AMD64, + reloc_type: IMAGE_REL_AMD64_ADDR32NB, + }, + _ => return Err(ImportlibYamlBuildError::UnsupportArchitecture(arch)), + }) + } + + #[inline] + pub fn machine(&self) -> u16 { + self.machine + } + + #[inline] + pub fn reloc_type(&self) -> u16 { + self.reloc_type + } +} diff --git a/crates/coffyaml/src/importlib/build.rs b/crates/coffyaml/src/importlib/build.rs new file mode 100644 index 0000000..4567a62 --- /dev/null +++ b/crates/coffyaml/src/importlib/build.rs @@ -0,0 +1,265 @@ +use object::pe::{ + IMAGE_FILE_MACHINE_UNKNOWN, IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ, + IMAGE_SCN_MEM_WRITE, IMAGE_SYM_CLASS_EXTERNAL, IMAGE_SYM_CLASS_SECTION, IMAGE_SYM_CLASS_STATIC, + IMPORT_OBJECT_CODE, IMPORT_OBJECT_NAME, IMPORT_OBJECT_NAME_MASK, IMPORT_OBJECT_NAME_SHIFT, + IMPORT_OBJECT_TYPE_MASK, +}; + +use crate::{ + archive::builder::ArchiveBuilder, + coff::{CoffYaml, CoffYamlHeader, CoffYamlSection, CoffYamlSectionRelocation, CoffYamlSymbol}, +}; + +use super::{Architecture, ArchitectureConfig, ImportlibYaml, errors::ImportlibYamlBuildError}; + +impl ImportlibYaml { + pub fn build(self, arch: Architecture) -> Result, ImportlibYamlBuildError> { + let cfg = ArchitectureConfig::new(arch)?; + + // Import descriptor, NULL import descriptor, NULL thunk data, import members + let member_count = 3 + self.exports.len(); + + let mut archive_builder = ArchiveBuilder::msvc_archive_with_capacity(member_count); + + // The library name for the import descriptor symbols + let library_name = self + .library + .rsplit_once('.') + .and_then(|(prefix, suffix)| { + suffix + .eq_ignore_ascii_case("dll") + .then(|| prefix.to_string()) + }) + .unwrap_or_else(|| self.library.clone()); + + let import_descriptor_name = format!("__IMPORT_DESCRIPTOR_{library_name}"); + let null_import_descriptor_name = "__NULL_IMPORT_DESCRIPTOR"; + let null_thunk_data_name = format!("\x7f{library_name}_NULL_THUNK_DATA"); + + // Add the import descriptor member + let mut member = archive_builder.add_member( + &self.library, + CoffYaml { + header: CoffYamlHeader { + machine: cfg.machine(), + characteristics: 0, + }, + sections: vec![ + CoffYamlSection { + name: ".idata$2".to_string(), + characteristics: IMAGE_SCN_CNT_INITIALIZED_DATA + | IMAGE_SCN_MEM_READ + | IMAGE_SCN_MEM_WRITE, + alignment: Some(4), + section_data: vec![0u8; 20], + size_of_raw_data: None, + relocations: vec![ + CoffYamlSectionRelocation { + virtual_address: 12, + symbol_name: ".idata$6".to_string(), + typ: cfg.reloc_type(), + }, + CoffYamlSectionRelocation { + virtual_address: 0, + symbol_name: ".idata$4".to_string(), + typ: cfg.reloc_type(), + }, + CoffYamlSectionRelocation { + virtual_address: 16, + symbol_name: ".idata$5".to_string(), + typ: cfg.reloc_type(), + }, + ], + }, + CoffYamlSection { + name: ".idata$6".to_string(), + characteristics: IMAGE_SCN_CNT_INITIALIZED_DATA + | IMAGE_SCN_MEM_READ + | IMAGE_SCN_MEM_WRITE, + alignment: Some(2), + section_data: format!("{}\0", &self.library).as_bytes().to_vec(), + ..Default::default() + }, + ], + symbols: vec![ + CoffYamlSymbol { + name: import_descriptor_name.clone(), + section_number: 1, + storage_class: IMAGE_SYM_CLASS_EXTERNAL, + ..Default::default() + }, + CoffYamlSymbol { + name: ".idata$2".to_string(), + section_number: 1, + storage_class: IMAGE_SYM_CLASS_SECTION, + ..Default::default() + }, + CoffYamlSymbol { + name: ".idata$6".to_string(), + section_number: 2, + storage_class: IMAGE_SYM_CLASS_STATIC, + ..Default::default() + }, + CoffYamlSymbol { + name: ".idata$4".to_string(), + storage_class: IMAGE_SYM_CLASS_SECTION, + ..Default::default() + }, + CoffYamlSymbol { + name: ".idata$5".to_string(), + storage_class: IMAGE_SYM_CLASS_SECTION, + ..Default::default() + }, + CoffYamlSymbol { + name: null_import_descriptor_name.to_string(), + storage_class: IMAGE_SYM_CLASS_EXTERNAL, + ..Default::default() + }, + CoffYamlSymbol { + name: null_thunk_data_name.clone(), + storage_class: IMAGE_SYM_CLASS_EXTERNAL, + ..Default::default() + }, + ], + } + .build() + .unwrap(), + ); + member.date(0); + member.uid(0); + member.gid(0); + member.mode(644); + member.export(&import_descriptor_name); + + // Add the NULL import descriptor member + let mut member = archive_builder.add_member( + &self.library, + CoffYaml { + header: CoffYamlHeader { + machine: cfg.machine(), + characteristics: 0, + }, + sections: vec![CoffYamlSection { + name: ".idata$3".to_string(), + characteristics: IMAGE_SCN_CNT_INITIALIZED_DATA + | IMAGE_SCN_MEM_READ + | IMAGE_SCN_MEM_WRITE, + alignment: Some(4), + section_data: vec![0u8; 20], + ..Default::default() + }], + symbols: vec![CoffYamlSymbol { + name: null_import_descriptor_name.to_string(), + section_number: 1, + storage_class: IMAGE_SYM_CLASS_EXTERNAL, + ..Default::default() + }], + } + .build() + .unwrap(), + ); + member.date(0); + member.uid(0); + member.gid(0); + member.mode(644); + member.export(null_import_descriptor_name); + + // Add the NULL thunk data member + let mut member = archive_builder.add_member( + &self.library, + CoffYaml { + header: CoffYamlHeader { + machine: cfg.machine(), + characteristics: 0, + }, + sections: vec![ + CoffYamlSection { + name: ".idata$5".to_string(), + characteristics: IMAGE_SCN_CNT_INITIALIZED_DATA + | IMAGE_SCN_MEM_READ + | IMAGE_SCN_MEM_WRITE, + alignment: Some(8), + section_data: vec![0u8; 8], + ..Default::default() + }, + CoffYamlSection { + name: ".idata$4".to_string(), + characteristics: IMAGE_SCN_CNT_INITIALIZED_DATA + | IMAGE_SCN_MEM_READ + | IMAGE_SCN_MEM_WRITE, + alignment: Some(8), + section_data: vec![0u8; 8], + ..Default::default() + }, + ], + symbols: vec![CoffYamlSymbol { + name: null_thunk_data_name.clone(), + section_number: 1, + storage_class: IMAGE_SYM_CLASS_EXTERNAL, + ..Default::default() + }], + } + .build() + .unwrap(), + ); + member.date(0); + member.uid(0); + member.gid(0); + member.mode(644); + member.export(null_thunk_data_name); + + // Add each import COFF + for export in self.exports { + let mut member = archive_builder.add_member( + &self.library, + build_import_coff(cfg.machine(), &export, &self.library), + ); + member.date(0); + member.uid(0); + member.gid(0); + member.mode(644); + member.exports([format!("__imp_{}", &export), export]); + } + + Ok(archive_builder.build()) + } +} + +fn build_import_coff(machine: u16, export: impl AsRef, dllname: impl AsRef) -> Vec { + let mut buffer = + Vec::with_capacity(20 + export.as_ref().len() + 1 + dllname.as_ref().len() + 1); + + // Header + buffer.extend(IMAGE_FILE_MACHINE_UNKNOWN.to_le_bytes()); // Sig1 + buffer.extend(0xffffu16.to_le_bytes()); // Sig2 + buffer.extend(0u16.to_le_bytes()); // Version + buffer.extend(machine.to_le_bytes()); // Machine + buffer.extend(0u32.to_le_bytes()); // Time-Date Stamp + + // Size of data + buffer.extend(((export.as_ref().len() + 1 + dllname.as_ref().len() + 1) as u32).to_le_bytes()); + + // Ordinal/Hint + buffer.extend(0u16.to_le_bytes()); + + // Import metadata + let mut meta: u16 = 0; + + // Code import type + meta |= IMPORT_OBJECT_CODE & IMPORT_OBJECT_TYPE_MASK; + + // Import by name + meta |= (IMPORT_OBJECT_NAME & IMPORT_OBJECT_NAME_MASK) << IMPORT_OBJECT_NAME_SHIFT; + + buffer.extend(meta.to_le_bytes()); + + // Import name string + buffer.extend(export.as_ref().as_bytes()); + buffer.push(0); + + // DLL name string + buffer.extend(dllname.as_ref().as_bytes()); + buffer.push(0); + + buffer +} diff --git a/crates/coffyaml/src/importlib/errors.rs b/crates/coffyaml/src/importlib/errors.rs new file mode 100644 index 0000000..2455ca4 --- /dev/null +++ b/crates/coffyaml/src/importlib/errors.rs @@ -0,0 +1,7 @@ +use super::Architecture; + +#[derive(Debug, thiserror::Error)] +pub enum ImportlibYamlBuildError { + #[error("architecture {0:?} is not supported")] + UnsupportArchitecture(Architecture), +} diff --git a/crates/coffyaml/src/importlib/legacy_build.rs b/crates/coffyaml/src/importlib/legacy_build.rs new file mode 100644 index 0000000..1662634 --- /dev/null +++ b/crates/coffyaml/src/importlib/legacy_build.rs @@ -0,0 +1,12 @@ +use crate::archive::builder::ArchiveBuilder; + +use super::{Architecture, ArchitectureConfig, ImportlibYaml, errors::ImportlibYamlBuildError}; + +impl ImportlibYaml { + pub fn build_legacy(self, arch: Architecture) -> Result, ImportlibYamlBuildError> { + let _cfg = ArchitectureConfig::new(arch)?; + + let _archive_builder = ArchiveBuilder::gnu_archive_with_capacity(0); + todo!() + } +} diff --git a/crates/coffyaml/src/importlib/mod.rs b/crates/coffyaml/src/importlib/mod.rs new file mode 100644 index 0000000..b5d64ad --- /dev/null +++ b/crates/coffyaml/src/importlib/mod.rs @@ -0,0 +1,16 @@ +use archconfig::ArchitectureConfig; +use serde::{Deserialize, Serialize}; + +pub use object::Architecture; + +mod archconfig; +mod build; +pub mod errors; +mod legacy_build; + +#[derive(Debug, Deserialize, Serialize)] +#[serde(rename_all = "PascalCase")] +pub struct ImportlibYaml { + pub library: String, + pub exports: Vec, +} diff --git a/crates/coffyaml/src/lib.rs b/crates/coffyaml/src/lib.rs new file mode 100644 index 0000000..f1439ad --- /dev/null +++ b/crates/coffyaml/src/lib.rs @@ -0,0 +1,24 @@ +pub mod archive; +pub mod coff; +pub mod importlib; + +#[cfg(test)] +mod testutils { + use serde::Deserializer; + + pub(crate) fn run_deserializer_tests<'a, 'de, U, T, F>(deserializer: F, tests: T) + where + T: IntoIterator, + U: 'a + std::fmt::Debug + std::cmp::PartialEq, + 'a: 'de, + F: Fn( + serde_yml::Deserializer<'de>, + ) -> Result as Deserializer<'de>>::Error>, + { + for (val, expected) in tests { + let de = serde_yml::Deserializer::from_str(val); + let parsed = deserializer(de).unwrap_or_else(|e| panic!("{val}\nError: {e}")); + assert_eq!(parsed, expected); + } + } +} diff --git a/crates/coffyaml/tests/coff.rs b/crates/coffyaml/tests/coff.rs new file mode 100644 index 0000000..5064363 --- /dev/null +++ b/crates/coffyaml/tests/coff.rs @@ -0,0 +1,22 @@ +use coffyaml::coff::CoffYaml; +use object::{coff::CoffFile, pe}; + +const COFF_YAML: &str = include_str!("coff.yaml"); + +#[test] +fn coff_sanity_parse() { + assert!(serde_yml::from_str::(COFF_YAML).is_ok()); +} + +#[test] +fn coff_sanity_build() { + let parsed: CoffYaml = serde_yml::from_str(COFF_YAML).unwrap(); + assert!(parsed.build().is_ok()); +} + +#[test] +fn coff_sanity_object_can_parse() { + let parsed_yaml: CoffYaml = serde_yml::from_str(COFF_YAML).unwrap(); + let built = parsed_yaml.build().unwrap(); + assert!(CoffFile::<_, pe::ImageFileHeader>::parse(built.as_slice()).is_ok()); +} diff --git a/crates/coffyaml/tests/coff.yaml b/crates/coffyaml/tests/coff.yaml new file mode 100644 index 0000000..1871027 --- /dev/null +++ b/crates/coffyaml/tests/coff.yaml @@ -0,0 +1,146 @@ +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_LINE_NUMS_STRIPPED ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: 554889E54883EC20E800000000488D05000000004889C1E800000000B8000000004883C4205DC3909090909090909090 + SizeOfRawData: 48 + Relocations: + - VirtualAddress: 16 + SymbolName: .rdata + Type: IMAGE_REL_AMD64_REL32 + - VirtualAddress: 24 + SymbolName: puts + Type: IMAGE_REL_AMD64_REL32 + - Name: .data + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 16 + SectionData: '' + - Name: .bss + Characteristics: [ IMAGE_SCN_CNT_UNINITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 16 + SectionData: '' + - Name: .rdata + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: 48656C6C6F20576F726C640000000000 + SizeOfRawData: 16 + - Name: .xdata + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 4 + SectionData: '010803050832040301500000' + SizeOfRawData: 12 + - Name: .pdata + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 4 + SectionData: '000000002700000000000000' + SizeOfRawData: 12 + Relocations: + - VirtualAddress: 0 + SymbolName: .text + Type: IMAGE_REL_AMD64_ADDR32NB + - VirtualAddress: 4 + SymbolName: .text + Type: IMAGE_REL_AMD64_ADDR32NB + - VirtualAddress: 8 + SymbolName: .xdata + Type: IMAGE_REL_AMD64_ADDR32NB +symbols: + - Name: .file + Value: 0 + SectionNumber: -2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_FILE + File: hello.c + - Name: main + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + FunctionDefinition: + TagIndex: 0 + TotalSize: 0 + PointerToLinenumber: 0 + PointerToNextFunction: 0 + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 39 + NumberOfRelocations: 3 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .data + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .bss + Value: 0 + SectionNumber: 3 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .rdata + Value: 0 + SectionNumber: 4 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 12 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .xdata + Value: 0 + SectionNumber: 5 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 12 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .pdata + Value: 0 + SectionNumber: 6 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 12 + NumberOfRelocations: 3 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: puts + Value: 0 + SectionNumber: 0 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL diff --git a/crates/coffyaml/tests/importlib.rs b/crates/coffyaml/tests/importlib.rs new file mode 100644 index 0000000..dd34e20 --- /dev/null +++ b/crates/coffyaml/tests/importlib.rs @@ -0,0 +1,75 @@ +use coffyaml::importlib::{Architecture, ImportlibYaml}; +use object::{ + coff::{ImportFile, ImportName}, + read::archive::ArchiveFile, +}; + +const IMPORTLIB_YAML: &str = include_str!("importlib.yaml"); + +#[test] +fn importlib_sanity_parse() { + assert!(serde_yml::from_str::(IMPORTLIB_YAML).is_ok()); +} + +#[test] +fn importlib_sanity_build() { + let parsed: ImportlibYaml = serde_yml::from_str(IMPORTLIB_YAML).unwrap(); + assert!(parsed.build(Architecture::X86_64).is_ok()); +} + +#[test] +fn importlib_sanity_object_can_parse() { + let parsed_yaml: ImportlibYaml = serde_yml::from_str(IMPORTLIB_YAML).unwrap(); + let built = parsed_yaml.build(Architecture::X86_64).unwrap(); + assert!(ArchiveFile::parse(built.as_slice()).is_ok()); +} + +#[test] +fn importlib_symbol_table_exports() { + let parsed_yaml: ImportlibYaml = serde_yml::from_str(IMPORTLIB_YAML).unwrap(); + + let exports_list = parsed_yaml.exports.clone(); + + let built = parsed_yaml.build(Architecture::X86_64).unwrap(); + let parsed_archive = ArchiveFile::parse(built.as_slice()).unwrap(); + + let archive_symbols = parsed_archive.symbols().unwrap().unwrap(); + + for export in exports_list { + assert!( + archive_symbols + .clone() + .any(|symbol| std::str::from_utf8(symbol.unwrap().name()).unwrap() == export), + "could not find '{export}' in symbol table" + ); + } +} + +#[test] +fn importlib_extract_member() { + let parsed_yaml: ImportlibYaml = serde_yml::from_str(IMPORTLIB_YAML).unwrap(); + + let built = parsed_yaml.build(Architecture::X86_64).unwrap(); + let parsed_archive = ArchiveFile::parse(built.as_slice()).unwrap(); + + let mut archive_symbols = parsed_archive.symbols().unwrap().unwrap(); + let symbol = archive_symbols + .find(|symbol| std::str::from_utf8(symbol.unwrap().name()).unwrap() == "ExportedSymbol") + .unwrap() + .unwrap(); + + let extracted_member = parsed_archive.member(symbol.offset()).unwrap(); + let extracted_data = extracted_member.data(built.as_slice()).unwrap(); + + let import_file = ImportFile::parse(extracted_data).unwrap(); + match import_file.import() { + ImportName::Name(s) => { + let name = std::str::from_utf8(s).unwrap(); + assert_eq!( + name, "ExportedSymbol", + "extracted import member public import name does not match name in the symbol map" + ); + } + ImportName::Ordinal(_) => panic!("import value should not be an ordinal"), + } +} diff --git a/crates/coffyaml/tests/importlib.yaml b/crates/coffyaml/tests/importlib.yaml new file mode 100644 index 0000000..afd7709 --- /dev/null +++ b/crates/coffyaml/tests/importlib.yaml @@ -0,0 +1,6 @@ +Library: KERNEL32.dll +Exports: + - GetLastError + - aaa + - ExportedSymbol + - other_export diff --git a/crates/jamcrc/Cargo.toml b/crates/jamcrc/Cargo.toml new file mode 100644 index 0000000..5789edc --- /dev/null +++ b/crates/jamcrc/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "jamcrc" +version = "0.1.0" +authors = ["Matt Ehrnschwender "] +edition = "2024" +description = "JamCRC checksum calculation" +readme = "README.md" +homepage = "https://github.com/MEhrn00/boflink/tree/main/crates/jamcrc" +repository = "https://github.com/MEhrn00/boflink" +publish = false + +[dependencies] +crc32fast = "1.4.2" diff --git a/crates/jamcrc/README.md b/crates/jamcrc/README.md new file mode 100644 index 0000000..fc2bef5 --- /dev/null +++ b/crates/jamcrc/README.md @@ -0,0 +1,26 @@ +# jamcrc +JamCRC wrapper around the [`crc32fast`](https://github.com/srijs/rust-crc32fast) crate. +Used for calculating COFF auxiliary section checksums. + +This is split up into a separate crate to expose it as a command line tool [`jamcrc-cli`](cli). + +## Command line usage +```shell +cargo r -p jamcrc-cli -- [arguments] +cargo r -p jamcrc-cli -- -h +``` + +### Usage +``` +Usage: jamcrc-cli [OPTIONS] [FILE] + +Arguments: + [FILE] Input file to calculate the checksum for. Use "-" to read from stdin + +Options: + -s, --string Input string to calculate the checksum for instead of a file + -i, --init Init value for the calcuation [default: 0] + --ihex Decode the passed in input as hex + --hex Print the calculated checksum as hex + -h, --help Print help +``` diff --git a/crates/jamcrc/cli/Cargo.toml b/crates/jamcrc/cli/Cargo.toml new file mode 100644 index 0000000..313dddc --- /dev/null +++ b/crates/jamcrc/cli/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "jamcrc-cli" +version = "0.1.0" +authors = ["Matt Ehrnschwender "] +edition = "2024" +description = """ +Command line tool for calculating JamCRC checksums. +""" +homepage = "https://github.com/MEhrn00/boflink/tree/main/crates/jamcrc/cli" +repository = "https://github.com/MEhrn00/boflink" +publish = false + +[dependencies] +anyhow = "1.0.98" +hex = "0.4.3" +jamcrc = { path = ".." } + +[dependencies.clap] +version = "4.5.38" +default-features = false +features = ["std", "help", "usage", "derive"] + +[lints.rust] +unsafe_code = "forbid" diff --git a/crates/jamcrc/cli/src/hexstream.rs b/crates/jamcrc/cli/src/hexstream.rs new file mode 100644 index 0000000..1a0a7f4 --- /dev/null +++ b/crates/jamcrc/cli/src/hexstream.rs @@ -0,0 +1,55 @@ +use std::io::Read; + +const DEFAULT_BUFFER_SIZE: usize = (8 * 1024) / 2; + +pub struct HexDecodeStream { + buffer: Vec, + reader: R, +} + +impl HexDecodeStream { + pub fn new(reader: R) -> HexDecodeStream { + Self { + buffer: Vec::with_capacity(DEFAULT_BUFFER_SIZE), + reader, + } + } +} + +impl From for HexDecodeStream { + fn from(value: R) -> Self { + Self::new(value) + } +} + +impl std::io::Read for HexDecodeStream { + fn read(&mut self, buf: &mut [u8]) -> std::io::Result { + self.buffer.resize(buf.len() * 2, 0); + let read_in = self.reader.read(&mut self.buffer)?; + hex::decode_to_slice(&self.buffer[..read_in], &mut buf[..read_in / 2]) + .map_err(std::io::Error::other)?; + Ok(read_in / 2) + } +} + +#[cfg(test)] +mod tests { + use std::io::Read; + + use super::HexDecodeStream; + + #[test] + fn decode_stream_round_trip() { + let input = "hello world"; + + let encoded = hex::encode(input); + let mut stream = HexDecodeStream::new(encoded.as_bytes()); + + let mut decoded = Vec::new(); + stream + .read_to_end(&mut decoded) + .expect("Could not read stream"); + + assert_eq!(decoded, input.as_bytes()); + } +} diff --git a/crates/jamcrc/cli/src/main.rs b/crates/jamcrc/cli/src/main.rs new file mode 100644 index 0000000..a5ed1a6 --- /dev/null +++ b/crates/jamcrc/cli/src/main.rs @@ -0,0 +1,120 @@ +use std::{ + io::{BufRead, BufReader}, + path::PathBuf, +}; + +use clap::Parser; +use hexstream::HexDecodeStream; + +mod hexstream; + +#[derive(Parser, Debug)] +#[command(about)] +struct CliArgs { + /// Input file to calculate the checksum for. Use "-" to read from stdin + #[arg(value_parser = parse_stdin_or_filepath)] + file: Option, + + /// Input string to calculate the checksum for instead of a file + #[arg( + id = "string", + long, + short, + conflicts_with = "file", + default_value = "", + hide_default_value = true + )] + input_string: String, + + /// Init value for the calcuation + #[arg(long, short, default_value = "0", allow_negative_numbers = true)] + init: i32, + + /// Decode the passed in input as hex + #[arg(long)] + ihex: bool, + + /// Print the calculated checksum as hex + #[arg(long)] + hex: bool, +} + +#[derive(Clone, Debug)] +enum StdinOrFilePath { + Stdin, + FilePath(PathBuf), +} + +fn parse_stdin_or_filepath( + argval: &str, +) -> Result> { + match argval { + "-" => Ok(StdinOrFilePath::Stdin), + _ => Ok(StdinOrFilePath::FilePath(PathBuf::from(argval))), + } +} + +fn calculate_buffered( + mut hasher: jamcrc::Hasher, + mut reader: R, +) -> anyhow::Result { + loop { + let buffer = reader.fill_buf()?; + if buffer.is_empty() { + return Ok(hasher.finalize()); + } + + let consumed = buffer.len(); + hasher.update(buffer); + reader.consume(consumed); + } +} + +fn calculate_full(mut hasher: jamcrc::Hasher, data: impl AsRef<[u8]>) -> u32 { + hasher.update(data.as_ref()); + hasher.finalize() +} + +fn main() -> anyhow::Result<()> { + let args = CliArgs::parse(); + + let hasher = jamcrc::Hasher::new_with_initial(args.init.cast_unsigned()); + + let checksum = if let Some(file) = args.file.as_ref() { + match file { + StdinOrFilePath::Stdin => { + if args.ihex { + calculate_buffered( + hasher, + BufReader::new(HexDecodeStream::new(std::io::stdin().lock())), + )? + } else { + calculate_buffered(hasher, std::io::stdin().lock())? + } + } + StdinOrFilePath::FilePath(path) => { + let f = std::fs::File::open(path)?; + if args.ihex { + calculate_buffered(hasher, BufReader::new(HexDecodeStream::new(f)))? + } else { + calculate_buffered(hasher, BufReader::new(f))? + } + } + } + } else if args.ihex { + calculate_buffered( + hasher, + BufReader::new(HexDecodeStream::new(args.input_string.as_bytes())), + )? + } else { + calculate_full(hasher, args.input_string.as_bytes()) + }; + + if args.hex { + println!("{checksum:#x}"); + } else { + println!("{checksum}"); + } + + Ok(()) +} diff --git a/crates/jamcrc/src/lib.rs b/crates/jamcrc/src/lib.rs new file mode 100644 index 0000000..df4a06a --- /dev/null +++ b/crates/jamcrc/src/lib.rs @@ -0,0 +1,40 @@ +/// JamCRC hasher. +pub struct Hasher { + inner: crc32fast::Hasher, +} + +impl Hasher { + /// Creates a new [`Hasher`]. + #[inline] + pub fn new() -> Self { + Self { + inner: crc32fast::Hasher::new(), + } + } + + /// Creates a new [`Hasher`] with an initial state. + #[inline] + pub fn new_with_initial(init: u32) -> Self { + Self { + inner: crc32fast::Hasher::new_with_initial(init), + } + } + + /// Updates the hasher with the specified data. + #[inline] + pub fn update(&mut self, data: &[u8]) { + self.inner.update(data); + } + + /// Returns the JamCRC value. + #[inline] + pub fn finalize(self) -> u32 { + !self.inner.finalize() + } +} + +impl std::default::Default for Hasher { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/objs2yaml/Cargo.toml b/crates/objs2yaml/Cargo.toml new file mode 100644 index 0000000..619b7eb --- /dev/null +++ b/crates/objs2yaml/Cargo.toml @@ -0,0 +1,31 @@ +[package] +name = "objs2yaml" +version = "0.1.0" +authors = ["Matt Ehrnschwender "] +edition = "2024" +description = """ +Generates test data yaml files from a collection of COFFs and import libraries. +""" +readme = "README.md" +homepage = "https://github.com/MEhrn00/boflink/tree/main/crates/objs2yaml" +repository = "https://github.com/MEhrn00/boflink" +publish = false + +[dependencies] +anyhow = "1.0.92" +coffyaml = { path = "../coffyaml" } +serde = "1" +serde_yml = "0.0.12" + +[dependencies.clap] +version = "4.5.24" +default-features = false +features = ["std", "help", "usage", "derive"] + +[dependencies.object] +version = "0.36.7" +default-features = false +features = ["archive", "coff", "write", "read"] + +[lints.rust] +unsafe_code = "forbid" diff --git a/crates/objs2yaml/README.md b/crates/objs2yaml/README.md new file mode 100644 index 0000000..f40a3cf --- /dev/null +++ b/crates/objs2yaml/README.md @@ -0,0 +1,19 @@ +# objs2yaml +Generates test data yaml files from a collection of COFFs and import libraries. + +```shell +cargo r -p objs2yaml -- [arguments] +cargo r -p objs2yaml -- -h +``` + +## Usage +``` +Usage: objs2yaml [OPTIONS] ... + +Arguments: + ... Input files + +Options: + -o, --output Output file. Defaults to stdout + -h, --help Print help +``` diff --git a/crates/objs2yaml/src/main.rs b/crates/objs2yaml/src/main.rs new file mode 100644 index 0000000..1942d2d --- /dev/null +++ b/crates/objs2yaml/src/main.rs @@ -0,0 +1,216 @@ +use std::{io::BufWriter, path::PathBuf}; + +use anyhow::Context; +use clap::Parser; +use coffyaml::{ + coff::{ + CoffYaml, CoffYamlAuxFunctionDefinition, CoffYamlAuxSectionDefinition, CoffYamlHeader, + CoffYamlSection, CoffYamlSectionRelocation, CoffYamlSymbol, + }, + importlib::ImportlibYaml, +}; +use object::{ + Object, ObjectSection, ObjectSymbol, + coff::{CoffFile, ImageSymbol, ImportFile}, + pe::{IMAGE_SYM_ABSOLUTE, IMAGE_SYM_DEBUG}, + read::archive::ArchiveFile, +}; +use serde::Serialize; + +#[derive(Parser, Debug)] +#[command(about)] +struct CliArgs { + /// Input files. + #[arg(required = true, value_name = "files", value_hint = clap::ValueHint::FilePath)] + files: Vec, + + /// Output file. Defaults to stdout. + #[arg(short, long, value_name = "file", value_hint = clap::ValueHint::FilePath)] + output: Option, +} + +#[derive(Debug, Serialize)] +enum ParsedInput { + #[serde(rename = "COFF")] + Coff(CoffYaml), + + #[serde(rename = "IMPORTLIB")] + Importlib(ImportlibYaml), +} + +fn main() -> anyhow::Result<()> { + let args = CliArgs::parse(); + + let mut parsed_inputs = Vec::with_capacity(args.files.len()); + + for file in args.files { + let data = + std::fs::read(&file).with_context(|| format!("could not read {}.", file.display()))?; + + if data + .get(..object::archive::MAGIC.len()) + .is_some_and(|magic| magic == object::archive::MAGIC) + { + parsed_inputs + .push(ParsedInput::Importlib(parse_importlib(data).with_context( + || format!("could not parse {}.", file.display()), + )?)); + } else { + parsed_inputs.push(ParsedInput::Coff( + parse_coff(data).with_context(|| format!("could not parse {}.", file.display()))?, + )); + } + } + + let mut output: Box = if let Some(filepath) = args.output { + Box::new(BufWriter::new( + std::fs::File::create(&filepath) + .with_context(|| format!("could not open {}.", filepath.display()))?, + )) + } else { + Box::new(BufWriter::new(std::io::stdout().lock())) + }; + + write!(output, "--- ")?; + let mut ser = serde_yml::Serializer::new(&mut output); + for parsed in parsed_inputs { + parsed.serialize(&mut ser)?; + } + + Ok(()) +} + +fn parse_importlib(data: Vec) -> anyhow::Result { + let archive = ArchiveFile::parse(data.as_slice())?; + + let first_member = archive.members().nth(3).unwrap()?; + let member_data = first_member.data(data.as_slice())?; + let import_file = ImportFile::parse(member_data)?; + + let library = std::str::from_utf8(import_file.dll())?.to_string(); + + let mut symbols = Vec::with_capacity(archive.members().count() - 3); + + for member in archive.members().skip(3) { + let member = member?; + let member_data = member.data(data.as_slice())?; + + let import_file = ImportFile::parse(member_data)?; + symbols.push(std::str::from_utf8(import_file.symbol())?.to_string()); + } + + Ok(ImportlibYaml { + library, + exports: symbols, + }) +} + +fn parse_coff(data: Vec) -> anyhow::Result { + let coff: CoffFile = CoffFile::parse(data.as_slice())?; + + let coff_header = coff.coff_header(); + + let header = CoffYamlHeader { + machine: coff_header.machine.get(object::LittleEndian), + characteristics: coff_header.characteristics.get(object::LittleEndian), + }; + + let mut sections = Vec::with_capacity(coff.coff_section_table().len()); + for section in coff.sections() { + let coff_section = section.coff_section(); + + let mut characteristics = coff_section.characteristics.get(object::LittleEndian); + let alignment = (characteristics & (0xfu32 << 20) != 0) + .then(|| 2usize.pow((characteristics >> 20 & 0xf) - 1)); + characteristics &= !(0xfu32 << 20); + + let mut relocations = Vec::with_capacity( + coff_section.number_of_relocations.get(object::LittleEndian) as usize, + ); + for reloc in section.coff_relocations()? { + let symbol = coff.symbol_by_index(reloc.symbol())?; + + relocations.push(CoffYamlSectionRelocation { + symbol_name: symbol.name()?.to_string(), + virtual_address: reloc.virtual_address.get(object::LittleEndian), + typ: reloc.typ.get(object::LittleEndian), + }); + } + + sections.push(CoffYamlSection { + name: section.name()?.to_string(), + characteristics, + alignment, + section_data: section.data()?.to_vec(), + size_of_raw_data: Some(coff_section.size_of_raw_data.get(object::LittleEndian)), + relocations, + }); + } + + let symbol_table = coff.coff_symbol_table(); + let mut symbols = Vec::with_capacity(symbol_table.len()); + + for symbol in coff.symbols() { + let coff_symbol = symbol.coff_symbol(); + + let section_definition = if coff_symbol.has_aux_section() { + let aux_section = symbol_table.aux_section(symbol.index())?; + Some(CoffYamlAuxSectionDefinition { + length: aux_section.length.get(object::LittleEndian), + number_of_relocations: aux_section.number_of_relocations.get(object::LittleEndian), + number_of_linenumbers: aux_section.number_of_linenumbers.get(object::LittleEndian), + check_sum: aux_section.check_sum.get(object::LittleEndian), + number: aux_section.number.get(object::LittleEndian), + selection: aux_section.selection, + }) + } else { + None + }; + + let function_definition = if coff_symbol.has_aux_function() { + let aux_function = symbol_table.aux_function(symbol.index())?; + Some(CoffYamlAuxFunctionDefinition { + tag_index: aux_function.tag_index.get(object::LittleEndian), + total_size: aux_function.total_size.get(object::LittleEndian), + pointer_to_linenumber: aux_function.pointer_to_linenumber.get(object::LittleEndian), + pointer_to_next_function: aux_function + .pointer_to_next_function + .get(object::LittleEndian), + }) + } else { + None + }; + + let file = if coff_symbol.has_aux_file_name() { + Some(symbol.name()?.to_string()) + } else { + None + }; + + symbols.push(CoffYamlSymbol { + name: if coff_symbol.has_aux_file_name() { + ".file".to_string() + } else { + symbol.name()?.to_string() + }, + value: coff_symbol.value.get(object::LittleEndian), + section_number: match coff_symbol.section_number.get(object::LittleEndian) { + 0xffff => IMAGE_SYM_ABSOLUTE, + 0xfffe => IMAGE_SYM_DEBUG, + o => o.into(), + }, + simple_type: coff_symbol.base_type(), + complex_type: coff_symbol.derived_type(), + storage_class: coff_symbol.storage_class, + section_definition, + function_definition, + file, + }); + } + + Ok(CoffYaml { + header, + sections, + symbols, + }) +} diff --git a/examples/basic/README.md b/examples/basic/README.md new file mode 100644 index 0000000..c444b1e --- /dev/null +++ b/examples/basic/README.md @@ -0,0 +1,23 @@ +# Basic +Basic example of compiling and linking a single source file. + +## Compiling +A minimal compile script is provided for compiling using MinGW GCC, Clang or MSVC. + +These compile scripts are intended to be used as a reference for the command line flags +needed to run boflink correctly and not as a project build script template. + +MinGW GCC (Linux) +```bash +./compile-mingw.sh +``` + +Clang (Linux) +```bash +./compile-clang.sh +``` + +MSVC (Windows) +```shell +.\compile-msvc.ps1 +``` diff --git a/examples/basic/basic.c b/examples/basic/basic.c new file mode 100644 index 0000000..9a281e8 --- /dev/null +++ b/examples/basic/basic.c @@ -0,0 +1,17 @@ +#include + +#include + +#include "beacon.h" + +void go(void) { + BeaconPrintf(CALLBACK_OUTPUT, "Hello, World!"); + + DWORD pid = GetCurrentProcessId(); + BeaconPrintf(CALLBACK_OUTPUT, "Current process id is %lu", pid); + + char username[UNLEN + 1] = {0}; + if (GetUserNameA(username, &(DWORD){sizeof(username)}) != 0) { + BeaconPrintf(CALLBACK_OUTPUT, "Your username is %s", username); + } +} diff --git a/examples/basic/beacon.h b/examples/basic/beacon.h new file mode 100644 index 0000000..3b9c96e --- /dev/null +++ b/examples/basic/beacon.h @@ -0,0 +1,370 @@ +/* + * Beacon Object Files (BOF) + * ------------------------- + * A Beacon Object File is a light-weight post exploitation tool that runs + * with Beacon's inline-execute command. + * + * Additional BOF resources are available here: + * - https://github.com/Cobalt-Strike/bof_template + * + * Cobalt Strike 4.x + * ChangeLog: + * 1/25/2022: updated for 4.5 + * 7/18/2023: Added BeaconInformation API for 4.9 + * 7/31/2023: Added Key/Value store APIs for 4.9 + * BeaconAddValue, BeaconGetValue, and BeaconRemoveValue + * 8/31/2023: Added Data store APIs for 4.9 + * BeaconDataStoreGetItem, BeaconDataStoreProtectItem, + * BeaconDataStoreUnprotectItem, and BeaconDataStoreMaxEntries + * 9/01/2023: Added BeaconGetCustomUserData API for 4.9 + * 3/21/2024: Updated BeaconInformation API for 4.10 to return a BOOL + * Updated the BEACON_INFO data structure to add new parameters + * 4/19/2024: Added BeaconGetSyscallInformation API for 4.10 + * 4/25/2024: Added APIs to call Beacon's system call implementation + */ +#ifndef _BEACON_H_ +#define _BEACON_H_ +#include + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +/* data API */ +typedef struct { + char * original; /* the original buffer [so we can free it] */ + char * buffer; /* current pointer into our buffer */ + int length; /* remaining length of data */ + int size; /* total size of this buffer */ +} datap; + +DECLSPEC_IMPORT void BeaconDataParse(datap * parser, char * buffer, int size); +DECLSPEC_IMPORT char * BeaconDataPtr(datap * parser, int size); +DECLSPEC_IMPORT int BeaconDataInt(datap * parser); +DECLSPEC_IMPORT short BeaconDataShort(datap * parser); +DECLSPEC_IMPORT int BeaconDataLength(datap * parser); +DECLSPEC_IMPORT char * BeaconDataExtract(datap * parser, int * size); + +/* format API */ +typedef struct { + char * original; /* the original buffer [so we can free it] */ + char * buffer; /* current pointer into our buffer */ + int length; /* remaining length of data */ + int size; /* total size of this buffer */ +} formatp; + +DECLSPEC_IMPORT void BeaconFormatAlloc(formatp * format, int maxsz); +DECLSPEC_IMPORT void BeaconFormatReset(formatp * format); +DECLSPEC_IMPORT void BeaconFormatAppend(formatp * format, const char * text, int len); +DECLSPEC_IMPORT void BeaconFormatPrintf(formatp * format, const char * fmt, ...); +DECLSPEC_IMPORT char * BeaconFormatToString(formatp * format, int * size); +DECLSPEC_IMPORT void BeaconFormatFree(formatp * format); +DECLSPEC_IMPORT void BeaconFormatInt(formatp * format, int value); + +/* Output Functions */ +#define CALLBACK_OUTPUT 0x0 +#define CALLBACK_OUTPUT_OEM 0x1e +#define CALLBACK_OUTPUT_UTF8 0x20 +#define CALLBACK_ERROR 0x0d +#define CALLBACK_CUSTOM 0x1000 +#define CALLBACK_CUSTOM_LAST 0x13ff + + +DECLSPEC_IMPORT void BeaconOutput(int type, const char * data, int len); +DECLSPEC_IMPORT void BeaconPrintf(int type, const char * fmt, ...); + + +/* Token Functions */ +DECLSPEC_IMPORT BOOL BeaconUseToken(HANDLE token); +DECLSPEC_IMPORT void BeaconRevertToken(); +DECLSPEC_IMPORT BOOL BeaconIsAdmin(); + +/* Spawn+Inject Functions */ +DECLSPEC_IMPORT void BeaconGetSpawnTo(BOOL x86, char * buffer, int length); +DECLSPEC_IMPORT void BeaconInjectProcess(HANDLE hProc, int pid, char * payload, int p_len, int p_offset, char * arg, int a_len); +DECLSPEC_IMPORT void BeaconInjectTemporaryProcess(PROCESS_INFORMATION * pInfo, char * payload, int p_len, int p_offset, char * arg, int a_len); +DECLSPEC_IMPORT BOOL BeaconSpawnTemporaryProcess(BOOL x86, BOOL ignoreToken, STARTUPINFO * si, PROCESS_INFORMATION * pInfo); +DECLSPEC_IMPORT void BeaconCleanupProcess(PROCESS_INFORMATION * pInfo); + +/* Utility Functions */ +DECLSPEC_IMPORT BOOL toWideChar(char * src, wchar_t * dst, int max); + +/* Beacon Information */ +/* + * ptr - pointer to the base address of the allocated memory. + * size - the number of bytes allocated for the ptr. + */ +typedef struct { + char * ptr; + size_t size; +} HEAP_RECORD; +#define MASK_SIZE 13 + +/* Information the user can set in the USER_DATA via a UDRL */ +typedef enum { + PURPOSE_EMPTY, + PURPOSE_GENERIC_BUFFER, + PURPOSE_BEACON_MEMORY, + PURPOSE_SLEEPMASK_MEMORY, + PURPOSE_BOF_MEMORY, + PURPOSE_USER_DEFINED_MEMORY = 1000 +} ALLOCATED_MEMORY_PURPOSE; + +typedef enum { + LABEL_EMPTY, + LABEL_BUFFER, + LABEL_PEHEADER, + LABEL_TEXT, + LABEL_RDATA, + LABEL_DATA, + LABEL_PDATA, + LABEL_RELOC, + LABEL_USER_DEFINED = 1000 +} ALLOCATED_MEMORY_LABEL; + +typedef enum { + METHOD_UNKNOWN, + METHOD_VIRTUALALLOC, + METHOD_HEAPALLOC, + METHOD_MODULESTOMP, + METHOD_NTMAPVIEW, + METHOD_USER_DEFINED = 1000, +} ALLOCATED_MEMORY_ALLOCATION_METHOD; + +/** +* This structure allows the user to provide additional information +* about the allocated heap for cleanup. It is mandatory to provide +* the HeapHandle but the DestroyHeap Boolean can be used to indicate +* whether the clean up code should destroy the heap or simply free the pages. +* This is useful in situations where a loader allocates memory in the +* processes current heap. +*/ +typedef struct _HEAPALLOC_INFO { + PVOID HeapHandle; + BOOL DestroyHeap; +} HEAPALLOC_INFO, *PHEAPALLOC_INFO; + +typedef struct _MODULESTOMP_INFO { + HMODULE ModuleHandle; +} MODULESTOMP_INFO, *PMODULESTOMP_INFO; + +typedef union _ALLOCATED_MEMORY_ADDITIONAL_CLEANUP_INFORMATION { + HEAPALLOC_INFO HeapAllocInfo; + MODULESTOMP_INFO ModuleStompInfo; + PVOID Custom; +} ALLOCATED_MEMORY_ADDITIONAL_CLEANUP_INFORMATION, *PALLOCATED_MEMORY_ADDITIONAL_CLEANUP_INFORMATION; + +typedef struct _ALLOCATED_MEMORY_CLEANUP_INFORMATION { + BOOL Cleanup; + ALLOCATED_MEMORY_ALLOCATION_METHOD AllocationMethod; + ALLOCATED_MEMORY_ADDITIONAL_CLEANUP_INFORMATION AdditionalCleanupInformation; +} ALLOCATED_MEMORY_CLEANUP_INFORMATION, *PALLOCATED_MEMORY_CLEANUP_INFORMATION; + +typedef struct _ALLOCATED_MEMORY_SECTION { + ALLOCATED_MEMORY_LABEL Label; // A label to simplify Sleepmask development + PVOID BaseAddress; // Pointer to virtual address of section + SIZE_T VirtualSize; // Virtual size of the section + DWORD CurrentProtect; // Current memory protection of the section + DWORD PreviousProtect; // The previous memory protection of the section (prior to masking/unmasking) + BOOL MaskSection; // A boolean to indicate whether the section should be masked +} ALLOCATED_MEMORY_SECTION, *PALLOCATED_MEMORY_SECTION; + +typedef struct _ALLOCATED_MEMORY_REGION { + ALLOCATED_MEMORY_PURPOSE Purpose; // A label to indicate the purpose of the allocated memory + PVOID AllocationBase; // The base address of the allocated memory block + SIZE_T RegionSize; // The size of the allocated memory block + DWORD Type; // The type of memory allocated + ALLOCATED_MEMORY_SECTION Sections[8]; // An array of section information structures + ALLOCATED_MEMORY_CLEANUP_INFORMATION CleanupInformation; // Information required to cleanup the allocation +} ALLOCATED_MEMORY_REGION, *PALLOCATED_MEMORY_REGION; + +typedef struct { + ALLOCATED_MEMORY_REGION AllocatedMemoryRegions[6]; +} ALLOCATED_MEMORY, *PALLOCATED_MEMORY; + +/* + * version - The version of the beacon dll was added for release 4.10 + * version format: 0xMMmmPP, where MM = Major, mm = Minor, and PP = Patch + * e.g. 0x040900 -> CS 4.9 + * 0x041000 -> CS 4.10 + * + * sleep_mask_ptr - pointer to the sleep mask base address + * sleep_mask_text_size - the sleep mask text section size + * sleep_mask_total_size - the sleep mask total memory size + * + * beacon_ptr - pointer to beacon's base address + * The stage.obfuscate flag affects this value when using CS default loader. + * true: beacon_ptr = allocated_buffer - 0x1000 (Not a valid address) + * false: beacon_ptr = allocated_buffer (A valid address) + * For a UDRL the beacon_ptr will be set to the 1st argument to DllMain + * when the 2nd argument is set to DLL_PROCESS_ATTACH. + * heap_records - list of memory addresses on the heap beacon wants to mask. + * The list is terminated by the HEAP_RECORD.ptr set to NULL. + * mask - the mask that beacon randomly generated to apply + * + * Added in version 4.10 + * allocatedMemory - An ALLOCATED_MEMORY structure that can be set in the USER_DATA + * via a UDRL. + */ +typedef struct { + unsigned int version; + char * sleep_mask_ptr; + DWORD sleep_mask_text_size; + DWORD sleep_mask_total_size; + + char * beacon_ptr; + HEAP_RECORD * heap_records; + char mask[MASK_SIZE]; + + ALLOCATED_MEMORY allocatedMemory; +} BEACON_INFO, *PBEACON_INFO; + +DECLSPEC_IMPORT BOOL BeaconInformation(PBEACON_INFO info); + +/* Key/Value store functions + * These functions are used to associate a key to a memory address and save + * that information into beacon. These memory addresses can then be + * retrieved in a subsequent execution of a BOF. + * + * key - the key will be converted to a hash which is used to locate the + * memory address. + * + * ptr - a memory address to save. + * + * Considerations: + * - The contents at the memory address is not masked by beacon. + * - The contents at the memory address is not released by beacon. + * + */ +DECLSPEC_IMPORT BOOL BeaconAddValue(const char * key, void * ptr); +DECLSPEC_IMPORT void * BeaconGetValue(const char * key); +DECLSPEC_IMPORT BOOL BeaconRemoveValue(const char * key); + +/* Beacon Data Store functions + * These functions are used to access items in Beacon's Data Store. + * BeaconDataStoreGetItem returns NULL if the index does not exist. + * + * The contents are masked by default, and BOFs must unprotect the entry + * before accessing the data buffer. BOFs must also protect the entry + * after the data is not used anymore. + * + */ + +#define DATA_STORE_TYPE_EMPTY 0 +#define DATA_STORE_TYPE_GENERAL_FILE 1 + +typedef struct { + int type; + DWORD64 hash; + BOOL masked; + char* buffer; + size_t length; +} DATA_STORE_OBJECT, *PDATA_STORE_OBJECT; + +DECLSPEC_IMPORT PDATA_STORE_OBJECT BeaconDataStoreGetItem(size_t index); +DECLSPEC_IMPORT void BeaconDataStoreProtectItem(size_t index); +DECLSPEC_IMPORT void BeaconDataStoreUnprotectItem(size_t index); +DECLSPEC_IMPORT size_t BeaconDataStoreMaxEntries(); + +/* Beacon User Data functions */ +DECLSPEC_IMPORT char * BeaconGetCustomUserData(); + +/* Beacon System call */ +/* Syscalls API */ +typedef struct +{ + PVOID fnAddr; + PVOID jmpAddr; + DWORD sysnum; +} SYSCALL_API_ENTRY, *PSYSCALL_API_ENTRY; + +typedef struct +{ + SYSCALL_API_ENTRY ntAllocateVirtualMemory; + SYSCALL_API_ENTRY ntProtectVirtualMemory; + SYSCALL_API_ENTRY ntFreeVirtualMemory; + SYSCALL_API_ENTRY ntGetContextThread; + SYSCALL_API_ENTRY ntSetContextThread; + SYSCALL_API_ENTRY ntResumeThread; + SYSCALL_API_ENTRY ntCreateThreadEx; + SYSCALL_API_ENTRY ntOpenProcess; + SYSCALL_API_ENTRY ntOpenThread; + SYSCALL_API_ENTRY ntClose; + SYSCALL_API_ENTRY ntCreateSection; + SYSCALL_API_ENTRY ntMapViewOfSection; + SYSCALL_API_ENTRY ntUnmapViewOfSection; + SYSCALL_API_ENTRY ntQueryVirtualMemory; + SYSCALL_API_ENTRY ntDuplicateObject; + SYSCALL_API_ENTRY ntReadVirtualMemory; + SYSCALL_API_ENTRY ntWriteVirtualMemory; + SYSCALL_API_ENTRY ntReadFile; + SYSCALL_API_ENTRY ntWriteFile; + SYSCALL_API_ENTRY ntCreateFile; +} SYSCALL_API, *PSYSCALL_API; + +/* Additional Run Time Library (RTL) addresses used to support system calls. + * If they are not set then system calls that require them will fall back + * to the Standard Windows API. + * + * Required to support the following system calls: + * ntCreateFile + */ +typedef struct +{ + PVOID rtlDosPathNameToNtPathNameUWithStatusAddr; + PVOID rtlFreeHeapAddr; + PVOID rtlGetProcessHeapAddr; +} RTL_API, *PRTL_API; + +typedef struct +{ + PSYSCALL_API syscalls; + PRTL_API rtls; +} BEACON_SYSCALLS, *PBEACON_SYSCALLS; + +DECLSPEC_IMPORT BOOL BeaconGetSyscallInformation(PBEACON_SYSCALLS info, BOOL resolveIfNotInitialized); + +/* Beacon System call functions which will use the current system call method */ +DECLSPEC_IMPORT LPVOID BeaconVirtualAlloc(LPVOID lpAddress, SIZE_T dwSize, DWORD flAllocationType, DWORD flProtect); +DECLSPEC_IMPORT LPVOID BeaconVirtualAllocEx(HANDLE processHandle, LPVOID lpAddress, SIZE_T dwSize, DWORD flAllocationType, DWORD flProtect); +DECLSPEC_IMPORT BOOL BeaconVirtualProtect(LPVOID lpAddress, SIZE_T dwSize, DWORD flNewProtect, PDWORD lpflOldProtect); +DECLSPEC_IMPORT BOOL BeaconVirtualProtectEx(HANDLE processHandle, LPVOID lpAddress, SIZE_T dwSize, DWORD flNewProtect, PDWORD lpflOldProtect); +DECLSPEC_IMPORT BOOL BeaconVirtualFree(LPVOID lpAddress, SIZE_T dwSize, DWORD dwFreeType); +DECLSPEC_IMPORT BOOL BeaconGetThreadContext(HANDLE threadHandle, PCONTEXT threadContext); +DECLSPEC_IMPORT BOOL BeaconSetThreadContext(HANDLE threadHandle, PCONTEXT threadContext); +DECLSPEC_IMPORT DWORD BeaconResumeThread(HANDLE threadHandle); +DECLSPEC_IMPORT HANDLE BeaconOpenProcess(DWORD desiredAccess, BOOL inheritHandle, DWORD processId); +DECLSPEC_IMPORT HANDLE BeaconOpenThread(DWORD desiredAccess, BOOL inheritHandle, DWORD threadId); +DECLSPEC_IMPORT BOOL BeaconCloseHandle(HANDLE object); +DECLSPEC_IMPORT BOOL BeaconUnmapViewOfFile(LPCVOID baseAddress); +DECLSPEC_IMPORT SIZE_T BeaconVirtualQuery(LPCVOID address, PMEMORY_BASIC_INFORMATION buffer, SIZE_T length); +DECLSPEC_IMPORT BOOL BeaconDuplicateHandle(HANDLE hSourceProcessHandle, HANDLE hSourceHandle, HANDLE hTargetProcessHandle, LPHANDLE lpTargetHandle, DWORD dwDesiredAccess, BOOL bInheritHandle, DWORD dwOptions); +DECLSPEC_IMPORT BOOL BeaconReadProcessMemory(HANDLE hProcess, LPCVOID lpBaseAddress, LPVOID lpBuffer, SIZE_T nSize, SIZE_T *lpNumberOfBytesRead); +DECLSPEC_IMPORT BOOL BeaconWriteProcessMemory(HANDLE hProcess, LPVOID lpBaseAddress, LPCVOID lpBuffer, SIZE_T nSize, SIZE_T *lpNumberOfBytesWritten); + +/* Beacon Gate APIs */ +DECLSPEC_IMPORT VOID BeaconDisableBeaconGate(); +DECLSPEC_IMPORT VOID BeaconEnableBeaconGate(); + +/* Beacon User Data + * + * version format: 0xMMmmPP, where MM = Major, mm = Minor, and PP = Patch + * e.g. 0x040900 -> CS 4.9 + * 0x041000 -> CS 4.10 +*/ + +#define DLL_BEACON_USER_DATA 0x0d +#define BEACON_USER_DATA_CUSTOM_SIZE 32 +typedef struct +{ + unsigned int version; + PSYSCALL_API syscalls; + char custom[BEACON_USER_DATA_CUSTOM_SIZE]; + PRTL_API rtls; + PALLOCATED_MEMORY allocatedMemory; +} USER_DATA, * PUSER_DATA; + +#ifdef __cplusplus +} +#endif // __cplusplus +#endif // _BEACON_H_ diff --git a/examples/basic/compile-clang.sh b/examples/basic/compile-clang.sh new file mode 100755 index 0000000..8941008 --- /dev/null +++ b/examples/basic/compile-clang.sh @@ -0,0 +1,9 @@ +#!/bin/sh +# Assumes that the boflink executable exists and is located in the user's PATH. + +# Get the path to the 'boflink' executable +boflink=$(which boflink) + +command="clang --ld-path=$boflink --target=x86_64-windows-gnu -nostartfiles -o basic.bof basic.c" +echo $command +eval $command diff --git a/examples/basic/compile-mingw.sh b/examples/basic/compile-mingw.sh new file mode 100755 index 0000000..b3d9881 --- /dev/null +++ b/examples/basic/compile-mingw.sh @@ -0,0 +1,8 @@ +#!/bin/sh +# Assumes that the program was installed using `cargo xtask install`. + +libexec="~/.local/libexec/boflink" + +command="x86_64-w64-mingw32-gcc -B $libexec -fno-lto -nostartfiles -o basic.bof basic.c" +echo $command +eval $command diff --git a/examples/basic/compile-msvc.ps1 b/examples/basic/compile-msvc.ps1 new file mode 100644 index 0000000..442b615 --- /dev/null +++ b/examples/basic/compile-msvc.ps1 @@ -0,0 +1,2 @@ +cl /GS- /c /Fo:basic.obj basic.c +boflink -o basic.bof basic.obj -lkernel32 -ladvapi32 diff --git a/examples/custom-api/README.md b/examples/custom-api/README.md new file mode 100644 index 0000000..21fd048 --- /dev/null +++ b/examples/custom-api/README.md @@ -0,0 +1,23 @@ +# Custom API +Example utilizing a set of custom API exports in-place of the default Cobalt Strike Beacon API. + +## Compiling +A minimal compile script is provided for compiling using MinGW GCC, Clang or MSVC. + +These compile scripts are intended to be used as a reference for the command line flags +needed to run boflink correctly and not as a project build script template. + +MinGW GCC (Linux) +```bash +./compile-mingw.sh +``` + +Clang (Linux) +```bash +./compile-clang.sh +``` + +MSVC (Windows) +```shell +.\compile-msvc.ps1 +``` diff --git a/examples/custom-api/compile-clang.sh b/examples/custom-api/compile-clang.sh new file mode 100755 index 0000000..bf00247 --- /dev/null +++ b/examples/custom-api/compile-clang.sh @@ -0,0 +1,13 @@ +#!/bin/sh +# Assumes that the boflink executable exists and is located in the user's PATH. + +# Get the path to the 'boflink' executable +boflink=$(which boflink) + +command="llvm-dlltool -l libmyapi.a -d myapi.def" +echo $command +eval $command + +command="clang --ld-path=$boflink --target=x86_64-windows-gnu -nostartfiles -Wl,--custom-api=libmyapi.a -o custom-api.bof custom-api.c" +echo $command +eval $command diff --git a/examples/custom-api/compile-mingw.sh b/examples/custom-api/compile-mingw.sh new file mode 100755 index 0000000..d2a2ad5 --- /dev/null +++ b/examples/custom-api/compile-mingw.sh @@ -0,0 +1,12 @@ +#!/bin/sh +# Assumes that the program was installed using `cargo xtask install`. + +libexec="~/.local/libexec/boflink" + +command="x86_64-w64-mingw32-dlltool -l libmyapi.a -d myapi.def" +echo $command +eval $command + +command="x86_64-w64-mingw32-gcc -B $libexec -fno-lto -nostartfiles -Wl,--custom-api=libmyapi.a -o custom-api.bof custom-api.c" +echo $command +eval $command diff --git a/examples/custom-api/compile-msvc.ps1 b/examples/custom-api/compile-msvc.ps1 new file mode 100644 index 0000000..1a2dec8 --- /dev/null +++ b/examples/custom-api/compile-msvc.ps1 @@ -0,0 +1,3 @@ +lib /machine:x64 /def:myapi.def /out:myapi.lib +cl /GS- /c /Fo:custom-api.obj custom-api.c +boflink --custom-api myapi.lib -o custom-api.bof custom-api.obj diff --git a/examples/custom-api/custom-api.c b/examples/custom-api/custom-api.c new file mode 100644 index 0000000..3373962 --- /dev/null +++ b/examples/custom-api/custom-api.c @@ -0,0 +1,13 @@ +#include "myapi.h" + +void go(void) { + int version = MyApiVersion(); + MyApiPrintf("MyApiVersion: %d", version); + + int *value = MyApiAlloc(sizeof(int)); + + *value = 123; + MyApiPrintf("value: %d", *value); + + MyApiFree(value); +} diff --git a/examples/custom-api/myapi.def b/examples/custom-api/myapi.def new file mode 100644 index 0000000..4e55532 --- /dev/null +++ b/examples/custom-api/myapi.def @@ -0,0 +1,6 @@ +LIBRARY MyApi +EXPORTS + MyApiVersion + MyApiPrintf + MyApiAlloc + MyApiFree diff --git a/examples/custom-api/myapi.h b/examples/custom-api/myapi.h new file mode 100644 index 0000000..beb2004 --- /dev/null +++ b/examples/custom-api/myapi.h @@ -0,0 +1,19 @@ +#ifndef MYAPI_H +#define MYAPI_H + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +#include + +__declspec(dllimport) int MyApiVersion(void); +__declspec(dllimport) void MyApiPrintf(const char *format, ...); +__declspec(dllimport) void *MyApiAlloc(size_t size); +__declspec(dllimport) void MyApiFree(void *ptr); + +#ifdef __cplusplus +}; +#endif // __cplusplus + +#endif // MYAPI_H diff --git a/examples/multiple-sources/README.md b/examples/multiple-sources/README.md new file mode 100644 index 0000000..a9af0e4 --- /dev/null +++ b/examples/multiple-sources/README.md @@ -0,0 +1,23 @@ +# Multiple Sources +Example for compiling and linking multiple source files. + +## Compiling +A minimal compile script is provided for compiling using MinGW GCC, Clang or MSVC. + +These compile scripts are intended to be used as a reference for the command line flags +needed to run boflink correctly and not as a project build script template. + +MinGW GCC (Linux) +```bash +./compile-mingw.sh +``` + +Clang (Linux) +```bash +./compile-clang.sh +``` + +MSVC (Windows) +```shell +.\compile-msvc.ps1 +``` diff --git a/examples/multiple-sources/beacon.h b/examples/multiple-sources/beacon.h new file mode 100644 index 0000000..3b9c96e --- /dev/null +++ b/examples/multiple-sources/beacon.h @@ -0,0 +1,370 @@ +/* + * Beacon Object Files (BOF) + * ------------------------- + * A Beacon Object File is a light-weight post exploitation tool that runs + * with Beacon's inline-execute command. + * + * Additional BOF resources are available here: + * - https://github.com/Cobalt-Strike/bof_template + * + * Cobalt Strike 4.x + * ChangeLog: + * 1/25/2022: updated for 4.5 + * 7/18/2023: Added BeaconInformation API for 4.9 + * 7/31/2023: Added Key/Value store APIs for 4.9 + * BeaconAddValue, BeaconGetValue, and BeaconRemoveValue + * 8/31/2023: Added Data store APIs for 4.9 + * BeaconDataStoreGetItem, BeaconDataStoreProtectItem, + * BeaconDataStoreUnprotectItem, and BeaconDataStoreMaxEntries + * 9/01/2023: Added BeaconGetCustomUserData API for 4.9 + * 3/21/2024: Updated BeaconInformation API for 4.10 to return a BOOL + * Updated the BEACON_INFO data structure to add new parameters + * 4/19/2024: Added BeaconGetSyscallInformation API for 4.10 + * 4/25/2024: Added APIs to call Beacon's system call implementation + */ +#ifndef _BEACON_H_ +#define _BEACON_H_ +#include + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +/* data API */ +typedef struct { + char * original; /* the original buffer [so we can free it] */ + char * buffer; /* current pointer into our buffer */ + int length; /* remaining length of data */ + int size; /* total size of this buffer */ +} datap; + +DECLSPEC_IMPORT void BeaconDataParse(datap * parser, char * buffer, int size); +DECLSPEC_IMPORT char * BeaconDataPtr(datap * parser, int size); +DECLSPEC_IMPORT int BeaconDataInt(datap * parser); +DECLSPEC_IMPORT short BeaconDataShort(datap * parser); +DECLSPEC_IMPORT int BeaconDataLength(datap * parser); +DECLSPEC_IMPORT char * BeaconDataExtract(datap * parser, int * size); + +/* format API */ +typedef struct { + char * original; /* the original buffer [so we can free it] */ + char * buffer; /* current pointer into our buffer */ + int length; /* remaining length of data */ + int size; /* total size of this buffer */ +} formatp; + +DECLSPEC_IMPORT void BeaconFormatAlloc(formatp * format, int maxsz); +DECLSPEC_IMPORT void BeaconFormatReset(formatp * format); +DECLSPEC_IMPORT void BeaconFormatAppend(formatp * format, const char * text, int len); +DECLSPEC_IMPORT void BeaconFormatPrintf(formatp * format, const char * fmt, ...); +DECLSPEC_IMPORT char * BeaconFormatToString(formatp * format, int * size); +DECLSPEC_IMPORT void BeaconFormatFree(formatp * format); +DECLSPEC_IMPORT void BeaconFormatInt(formatp * format, int value); + +/* Output Functions */ +#define CALLBACK_OUTPUT 0x0 +#define CALLBACK_OUTPUT_OEM 0x1e +#define CALLBACK_OUTPUT_UTF8 0x20 +#define CALLBACK_ERROR 0x0d +#define CALLBACK_CUSTOM 0x1000 +#define CALLBACK_CUSTOM_LAST 0x13ff + + +DECLSPEC_IMPORT void BeaconOutput(int type, const char * data, int len); +DECLSPEC_IMPORT void BeaconPrintf(int type, const char * fmt, ...); + + +/* Token Functions */ +DECLSPEC_IMPORT BOOL BeaconUseToken(HANDLE token); +DECLSPEC_IMPORT void BeaconRevertToken(); +DECLSPEC_IMPORT BOOL BeaconIsAdmin(); + +/* Spawn+Inject Functions */ +DECLSPEC_IMPORT void BeaconGetSpawnTo(BOOL x86, char * buffer, int length); +DECLSPEC_IMPORT void BeaconInjectProcess(HANDLE hProc, int pid, char * payload, int p_len, int p_offset, char * arg, int a_len); +DECLSPEC_IMPORT void BeaconInjectTemporaryProcess(PROCESS_INFORMATION * pInfo, char * payload, int p_len, int p_offset, char * arg, int a_len); +DECLSPEC_IMPORT BOOL BeaconSpawnTemporaryProcess(BOOL x86, BOOL ignoreToken, STARTUPINFO * si, PROCESS_INFORMATION * pInfo); +DECLSPEC_IMPORT void BeaconCleanupProcess(PROCESS_INFORMATION * pInfo); + +/* Utility Functions */ +DECLSPEC_IMPORT BOOL toWideChar(char * src, wchar_t * dst, int max); + +/* Beacon Information */ +/* + * ptr - pointer to the base address of the allocated memory. + * size - the number of bytes allocated for the ptr. + */ +typedef struct { + char * ptr; + size_t size; +} HEAP_RECORD; +#define MASK_SIZE 13 + +/* Information the user can set in the USER_DATA via a UDRL */ +typedef enum { + PURPOSE_EMPTY, + PURPOSE_GENERIC_BUFFER, + PURPOSE_BEACON_MEMORY, + PURPOSE_SLEEPMASK_MEMORY, + PURPOSE_BOF_MEMORY, + PURPOSE_USER_DEFINED_MEMORY = 1000 +} ALLOCATED_MEMORY_PURPOSE; + +typedef enum { + LABEL_EMPTY, + LABEL_BUFFER, + LABEL_PEHEADER, + LABEL_TEXT, + LABEL_RDATA, + LABEL_DATA, + LABEL_PDATA, + LABEL_RELOC, + LABEL_USER_DEFINED = 1000 +} ALLOCATED_MEMORY_LABEL; + +typedef enum { + METHOD_UNKNOWN, + METHOD_VIRTUALALLOC, + METHOD_HEAPALLOC, + METHOD_MODULESTOMP, + METHOD_NTMAPVIEW, + METHOD_USER_DEFINED = 1000, +} ALLOCATED_MEMORY_ALLOCATION_METHOD; + +/** +* This structure allows the user to provide additional information +* about the allocated heap for cleanup. It is mandatory to provide +* the HeapHandle but the DestroyHeap Boolean can be used to indicate +* whether the clean up code should destroy the heap or simply free the pages. +* This is useful in situations where a loader allocates memory in the +* processes current heap. +*/ +typedef struct _HEAPALLOC_INFO { + PVOID HeapHandle; + BOOL DestroyHeap; +} HEAPALLOC_INFO, *PHEAPALLOC_INFO; + +typedef struct _MODULESTOMP_INFO { + HMODULE ModuleHandle; +} MODULESTOMP_INFO, *PMODULESTOMP_INFO; + +typedef union _ALLOCATED_MEMORY_ADDITIONAL_CLEANUP_INFORMATION { + HEAPALLOC_INFO HeapAllocInfo; + MODULESTOMP_INFO ModuleStompInfo; + PVOID Custom; +} ALLOCATED_MEMORY_ADDITIONAL_CLEANUP_INFORMATION, *PALLOCATED_MEMORY_ADDITIONAL_CLEANUP_INFORMATION; + +typedef struct _ALLOCATED_MEMORY_CLEANUP_INFORMATION { + BOOL Cleanup; + ALLOCATED_MEMORY_ALLOCATION_METHOD AllocationMethod; + ALLOCATED_MEMORY_ADDITIONAL_CLEANUP_INFORMATION AdditionalCleanupInformation; +} ALLOCATED_MEMORY_CLEANUP_INFORMATION, *PALLOCATED_MEMORY_CLEANUP_INFORMATION; + +typedef struct _ALLOCATED_MEMORY_SECTION { + ALLOCATED_MEMORY_LABEL Label; // A label to simplify Sleepmask development + PVOID BaseAddress; // Pointer to virtual address of section + SIZE_T VirtualSize; // Virtual size of the section + DWORD CurrentProtect; // Current memory protection of the section + DWORD PreviousProtect; // The previous memory protection of the section (prior to masking/unmasking) + BOOL MaskSection; // A boolean to indicate whether the section should be masked +} ALLOCATED_MEMORY_SECTION, *PALLOCATED_MEMORY_SECTION; + +typedef struct _ALLOCATED_MEMORY_REGION { + ALLOCATED_MEMORY_PURPOSE Purpose; // A label to indicate the purpose of the allocated memory + PVOID AllocationBase; // The base address of the allocated memory block + SIZE_T RegionSize; // The size of the allocated memory block + DWORD Type; // The type of memory allocated + ALLOCATED_MEMORY_SECTION Sections[8]; // An array of section information structures + ALLOCATED_MEMORY_CLEANUP_INFORMATION CleanupInformation; // Information required to cleanup the allocation +} ALLOCATED_MEMORY_REGION, *PALLOCATED_MEMORY_REGION; + +typedef struct { + ALLOCATED_MEMORY_REGION AllocatedMemoryRegions[6]; +} ALLOCATED_MEMORY, *PALLOCATED_MEMORY; + +/* + * version - The version of the beacon dll was added for release 4.10 + * version format: 0xMMmmPP, where MM = Major, mm = Minor, and PP = Patch + * e.g. 0x040900 -> CS 4.9 + * 0x041000 -> CS 4.10 + * + * sleep_mask_ptr - pointer to the sleep mask base address + * sleep_mask_text_size - the sleep mask text section size + * sleep_mask_total_size - the sleep mask total memory size + * + * beacon_ptr - pointer to beacon's base address + * The stage.obfuscate flag affects this value when using CS default loader. + * true: beacon_ptr = allocated_buffer - 0x1000 (Not a valid address) + * false: beacon_ptr = allocated_buffer (A valid address) + * For a UDRL the beacon_ptr will be set to the 1st argument to DllMain + * when the 2nd argument is set to DLL_PROCESS_ATTACH. + * heap_records - list of memory addresses on the heap beacon wants to mask. + * The list is terminated by the HEAP_RECORD.ptr set to NULL. + * mask - the mask that beacon randomly generated to apply + * + * Added in version 4.10 + * allocatedMemory - An ALLOCATED_MEMORY structure that can be set in the USER_DATA + * via a UDRL. + */ +typedef struct { + unsigned int version; + char * sleep_mask_ptr; + DWORD sleep_mask_text_size; + DWORD sleep_mask_total_size; + + char * beacon_ptr; + HEAP_RECORD * heap_records; + char mask[MASK_SIZE]; + + ALLOCATED_MEMORY allocatedMemory; +} BEACON_INFO, *PBEACON_INFO; + +DECLSPEC_IMPORT BOOL BeaconInformation(PBEACON_INFO info); + +/* Key/Value store functions + * These functions are used to associate a key to a memory address and save + * that information into beacon. These memory addresses can then be + * retrieved in a subsequent execution of a BOF. + * + * key - the key will be converted to a hash which is used to locate the + * memory address. + * + * ptr - a memory address to save. + * + * Considerations: + * - The contents at the memory address is not masked by beacon. + * - The contents at the memory address is not released by beacon. + * + */ +DECLSPEC_IMPORT BOOL BeaconAddValue(const char * key, void * ptr); +DECLSPEC_IMPORT void * BeaconGetValue(const char * key); +DECLSPEC_IMPORT BOOL BeaconRemoveValue(const char * key); + +/* Beacon Data Store functions + * These functions are used to access items in Beacon's Data Store. + * BeaconDataStoreGetItem returns NULL if the index does not exist. + * + * The contents are masked by default, and BOFs must unprotect the entry + * before accessing the data buffer. BOFs must also protect the entry + * after the data is not used anymore. + * + */ + +#define DATA_STORE_TYPE_EMPTY 0 +#define DATA_STORE_TYPE_GENERAL_FILE 1 + +typedef struct { + int type; + DWORD64 hash; + BOOL masked; + char* buffer; + size_t length; +} DATA_STORE_OBJECT, *PDATA_STORE_OBJECT; + +DECLSPEC_IMPORT PDATA_STORE_OBJECT BeaconDataStoreGetItem(size_t index); +DECLSPEC_IMPORT void BeaconDataStoreProtectItem(size_t index); +DECLSPEC_IMPORT void BeaconDataStoreUnprotectItem(size_t index); +DECLSPEC_IMPORT size_t BeaconDataStoreMaxEntries(); + +/* Beacon User Data functions */ +DECLSPEC_IMPORT char * BeaconGetCustomUserData(); + +/* Beacon System call */ +/* Syscalls API */ +typedef struct +{ + PVOID fnAddr; + PVOID jmpAddr; + DWORD sysnum; +} SYSCALL_API_ENTRY, *PSYSCALL_API_ENTRY; + +typedef struct +{ + SYSCALL_API_ENTRY ntAllocateVirtualMemory; + SYSCALL_API_ENTRY ntProtectVirtualMemory; + SYSCALL_API_ENTRY ntFreeVirtualMemory; + SYSCALL_API_ENTRY ntGetContextThread; + SYSCALL_API_ENTRY ntSetContextThread; + SYSCALL_API_ENTRY ntResumeThread; + SYSCALL_API_ENTRY ntCreateThreadEx; + SYSCALL_API_ENTRY ntOpenProcess; + SYSCALL_API_ENTRY ntOpenThread; + SYSCALL_API_ENTRY ntClose; + SYSCALL_API_ENTRY ntCreateSection; + SYSCALL_API_ENTRY ntMapViewOfSection; + SYSCALL_API_ENTRY ntUnmapViewOfSection; + SYSCALL_API_ENTRY ntQueryVirtualMemory; + SYSCALL_API_ENTRY ntDuplicateObject; + SYSCALL_API_ENTRY ntReadVirtualMemory; + SYSCALL_API_ENTRY ntWriteVirtualMemory; + SYSCALL_API_ENTRY ntReadFile; + SYSCALL_API_ENTRY ntWriteFile; + SYSCALL_API_ENTRY ntCreateFile; +} SYSCALL_API, *PSYSCALL_API; + +/* Additional Run Time Library (RTL) addresses used to support system calls. + * If they are not set then system calls that require them will fall back + * to the Standard Windows API. + * + * Required to support the following system calls: + * ntCreateFile + */ +typedef struct +{ + PVOID rtlDosPathNameToNtPathNameUWithStatusAddr; + PVOID rtlFreeHeapAddr; + PVOID rtlGetProcessHeapAddr; +} RTL_API, *PRTL_API; + +typedef struct +{ + PSYSCALL_API syscalls; + PRTL_API rtls; +} BEACON_SYSCALLS, *PBEACON_SYSCALLS; + +DECLSPEC_IMPORT BOOL BeaconGetSyscallInformation(PBEACON_SYSCALLS info, BOOL resolveIfNotInitialized); + +/* Beacon System call functions which will use the current system call method */ +DECLSPEC_IMPORT LPVOID BeaconVirtualAlloc(LPVOID lpAddress, SIZE_T dwSize, DWORD flAllocationType, DWORD flProtect); +DECLSPEC_IMPORT LPVOID BeaconVirtualAllocEx(HANDLE processHandle, LPVOID lpAddress, SIZE_T dwSize, DWORD flAllocationType, DWORD flProtect); +DECLSPEC_IMPORT BOOL BeaconVirtualProtect(LPVOID lpAddress, SIZE_T dwSize, DWORD flNewProtect, PDWORD lpflOldProtect); +DECLSPEC_IMPORT BOOL BeaconVirtualProtectEx(HANDLE processHandle, LPVOID lpAddress, SIZE_T dwSize, DWORD flNewProtect, PDWORD lpflOldProtect); +DECLSPEC_IMPORT BOOL BeaconVirtualFree(LPVOID lpAddress, SIZE_T dwSize, DWORD dwFreeType); +DECLSPEC_IMPORT BOOL BeaconGetThreadContext(HANDLE threadHandle, PCONTEXT threadContext); +DECLSPEC_IMPORT BOOL BeaconSetThreadContext(HANDLE threadHandle, PCONTEXT threadContext); +DECLSPEC_IMPORT DWORD BeaconResumeThread(HANDLE threadHandle); +DECLSPEC_IMPORT HANDLE BeaconOpenProcess(DWORD desiredAccess, BOOL inheritHandle, DWORD processId); +DECLSPEC_IMPORT HANDLE BeaconOpenThread(DWORD desiredAccess, BOOL inheritHandle, DWORD threadId); +DECLSPEC_IMPORT BOOL BeaconCloseHandle(HANDLE object); +DECLSPEC_IMPORT BOOL BeaconUnmapViewOfFile(LPCVOID baseAddress); +DECLSPEC_IMPORT SIZE_T BeaconVirtualQuery(LPCVOID address, PMEMORY_BASIC_INFORMATION buffer, SIZE_T length); +DECLSPEC_IMPORT BOOL BeaconDuplicateHandle(HANDLE hSourceProcessHandle, HANDLE hSourceHandle, HANDLE hTargetProcessHandle, LPHANDLE lpTargetHandle, DWORD dwDesiredAccess, BOOL bInheritHandle, DWORD dwOptions); +DECLSPEC_IMPORT BOOL BeaconReadProcessMemory(HANDLE hProcess, LPCVOID lpBaseAddress, LPVOID lpBuffer, SIZE_T nSize, SIZE_T *lpNumberOfBytesRead); +DECLSPEC_IMPORT BOOL BeaconWriteProcessMemory(HANDLE hProcess, LPVOID lpBaseAddress, LPCVOID lpBuffer, SIZE_T nSize, SIZE_T *lpNumberOfBytesWritten); + +/* Beacon Gate APIs */ +DECLSPEC_IMPORT VOID BeaconDisableBeaconGate(); +DECLSPEC_IMPORT VOID BeaconEnableBeaconGate(); + +/* Beacon User Data + * + * version format: 0xMMmmPP, where MM = Major, mm = Minor, and PP = Patch + * e.g. 0x040900 -> CS 4.9 + * 0x041000 -> CS 4.10 +*/ + +#define DLL_BEACON_USER_DATA 0x0d +#define BEACON_USER_DATA_CUSTOM_SIZE 32 +typedef struct +{ + unsigned int version; + PSYSCALL_API syscalls; + char custom[BEACON_USER_DATA_CUSTOM_SIZE]; + PRTL_API rtls; + PALLOCATED_MEMORY allocatedMemory; +} USER_DATA, * PUSER_DATA; + +#ifdef __cplusplus +} +#endif // __cplusplus +#endif // _BEACON_H_ diff --git a/examples/multiple-sources/compile-clang.sh b/examples/multiple-sources/compile-clang.sh new file mode 100755 index 0000000..bde5872 --- /dev/null +++ b/examples/multiple-sources/compile-clang.sh @@ -0,0 +1,9 @@ +#!/bin/sh +# Assumes that the boflink executable exists and is located in the user's PATH. + +# Get the path to the 'boflink' executable +boflink=$(which boflink) + +command="clang --ld-path=$boflink --target=x86_64-windows-gnu -nostartfiles -o multiple-sources.bof go.c other.c" +echo $command +eval $command diff --git a/examples/multiple-sources/compile-mingw.sh b/examples/multiple-sources/compile-mingw.sh new file mode 100755 index 0000000..2804684 --- /dev/null +++ b/examples/multiple-sources/compile-mingw.sh @@ -0,0 +1,8 @@ +#!/bin/sh +# Assumes that the program was installed using `cargo xtask install`. + +libexec="~/.local/libexec/boflink" + +command="x86_64-w64-mingw32-gcc -B $libexec -fno-lto -nostartfiles -o multiple-sources.o go.c other.c" +echo $command +eval $command diff --git a/examples/multiple-sources/compile-msvc.ps1 b/examples/multiple-sources/compile-msvc.ps1 new file mode 100644 index 0000000..5238543 --- /dev/null +++ b/examples/multiple-sources/compile-msvc.ps1 @@ -0,0 +1,2 @@ +cl /GS- /c go.c other.c +boflink -o multiple-sources.bof go.obj other.obj diff --git a/examples/multiple-sources/go.c b/examples/multiple-sources/go.c new file mode 100644 index 0000000..ba2e2a1 --- /dev/null +++ b/examples/multiple-sources/go.c @@ -0,0 +1,9 @@ +#include "beacon.h" + +#include "other.h" + +void go(void) { + BeaconPrintf(CALLBACK_OUTPUT, "Hello world from the go() function"); + + other_function(); +} diff --git a/examples/multiple-sources/other.c b/examples/multiple-sources/other.c new file mode 100644 index 0000000..ecb80fb --- /dev/null +++ b/examples/multiple-sources/other.c @@ -0,0 +1,7 @@ +#include "other.h" + +#include "beacon.h" + +void other_function() { + BeaconPrintf(CALLBACK_OUTPUT, "Hello world from other_function()"); +} diff --git a/examples/multiple-sources/other.h b/examples/multiple-sources/other.h new file mode 100644 index 0000000..46fca72 --- /dev/null +++ b/examples/multiple-sources/other.h @@ -0,0 +1,14 @@ +#ifndef MULTIPLE_SOURCES_OTHER_H +#define MULTIPLE_SOURCES_OTHER_H + +#ifdef __cplusplus +extern "C" { +#endif + +void other_function(void); + +#ifdef __cplusplus +}; +#endif + +#endif // MULTIPLE_SOURCES_OTHER_H diff --git a/src/api/beaconapi.rs b/src/api/beaconapi.rs new file mode 100644 index 0000000..d137b79 --- /dev/null +++ b/src/api/beaconapi.rs @@ -0,0 +1,127 @@ +use std::path::Path; + +use object::Architecture; + +use crate::{ + libsearch::LibraryFind, + linker::{ApiInit, ApiInitCtx, error::ApiInitError}, + linkobject::import::{ImportMember, ImportName, ImportType}, +}; + +use super::{ApiSymbolError, ApiSymbolSource}; + +/// The Beacon API symbol string values. +/// +/// Symbols are sorted based on commonality. +const BEACONAPI_SYMBOLS: [&str; 52] = [ + "BeaconPrintf", + "BeaconDataParse", + "BeaconOutput", + "BeaconDataExtract", + "BeaconDataInt", + "BeaconGetSpawnTo", + "BeaconCleanupProcess", + "BeaconSpawnTemporaryProcess", + "BeaconDataShort", + "toWideChar", + "BeaconUseToken", + "BeaconGetValue", + "BeaconRemoveValue", + "BeaconInjectProcess", + "BeaconDataLength", + "BeaconAddValue", + "BeaconRevertToken", + "BeaconOpenThread", + "BeaconUnmapViewOfFile", + "BeaconFormatInt", + "BeaconGetSyscallInformation", + "BeaconDataStoreProtectItem", + "BeaconFormatFree", + "BeaconDataStoreUnprotectItem", + "BeaconInformation", + "BeaconDataStoreMaxEntries", + "BeaconDuplicateHandle", + "BeaconOpenProcess", + "BeaconDataStoreGetItem", + "BeaconEnableBeaconGate:", + "BeaconVirtualQuery", + "BeaconWriteProcessMemory", + "BeaconSetThreadContext", + "BeaconVirtualProtect", + "BeaconFormatAppend", + "BeaconDisableBeaconGate", + "BeaconResumeThread", + "BeaconDataPtr", + "BeaconGetThreadContext", + "BeaconIsAdmin", + "BeaconVirtualAlloc", + "BeaconCloseHandle", + "BeaconReadProcessMemory", + "BeaconFormatReset", + "BeaconVirtualAllocEx", + "BeaconFormatPrintf", + "BeaconFormatToString", + "BeaconInjectTemporaryProcess", + "BeaconVirtualFree", + "BeaconGetCustomUserData", + "BeaconVirtualProtectEx", + "BeaconFormatAlloc", +]; + +/// Container for looking up Beacon API symbols. +pub struct BeaconApiSymbols { + architecture: Architecture, + symbols: [&'static str; 52], +} + +impl BeaconApiSymbols { + /// Returns the Beacon API symbols for the specified architecture. + pub fn new(arch: Architecture) -> BeaconApiSymbols { + Self { + architecture: arch, + symbols: BEACONAPI_SYMBOLS, + } + } +} + +impl<'a> ApiSymbolSource<'a> for BeaconApiSymbols { + fn api_path(&self) -> &std::path::Path { + Path::new("Beacon API") + } + + fn extract_api_symbol(&self, symbol: &'a str) -> Result, ApiSymbolError> { + let unprefixed_symbol = if self.architecture == Architecture::I386 { + symbol.trim_start_matches("__imp__") + } else { + symbol.trim_start_matches("__imp_") + }; + + self.symbols + .iter() + .copied() + .find_map(|contained_symbol| { + (contained_symbol == unprefixed_symbol).then_some(ImportMember { + architecture: self.architecture, + symbol: contained_symbol, + dll: "Beacon API", + import: ImportName::Name(contained_symbol), + typ: ImportType::Code, + }) + }) + .ok_or(ApiSymbolError::NotFound) + } +} + +pub struct BeaconApiInit; + +impl ApiInit for BeaconApiInit { + type Output<'a> = BeaconApiSymbols; + + #[inline] + fn initialize_api<'a, L: LibraryFind>( + &self, + ctx: &ApiInitCtx<'_, 'a, L>, + ) -> Result, ApiInitError> { + Ok(BeaconApiSymbols::new(ctx.target_arch.into())) + } +} diff --git a/src/api/error.rs b/src/api/error.rs new file mode 100644 index 0000000..e971856 --- /dev/null +++ b/src/api/error.rs @@ -0,0 +1,26 @@ +use crate::linkobject::archive::{ArchiveParseError, ExtractMemberError, MemberParseError}; + +#[derive(Debug, thiserror::Error)] +pub enum ApiSymbolError { + #[error("member for symbol does not exist")] + NotFound, + + #[error("{0}")] + ArchiveParse(ArchiveParseError), + + #[error("{0}")] + MemberParse(MemberParseError), + + #[error("invalid COFF import library member")] + ImportMember, +} + +impl From for ApiSymbolError { + fn from(value: ExtractMemberError) -> Self { + match value { + ExtractMemberError::NotFound => Self::NotFound, + ExtractMemberError::ArchiveParse(e) => Self::ArchiveParse(e), + ExtractMemberError::MemberParse(e) => Self::MemberParse(e), + } + } +} diff --git a/src/api/mod.rs b/src/api/mod.rs new file mode 100644 index 0000000..b24c0c3 --- /dev/null +++ b/src/api/mod.rs @@ -0,0 +1,7 @@ +mod beaconapi; +mod error; +mod traits; + +pub use beaconapi::*; +pub use error::*; +pub use traits::*; diff --git a/src/api/traits.rs b/src/api/traits.rs new file mode 100644 index 0000000..41a6858 --- /dev/null +++ b/src/api/traits.rs @@ -0,0 +1,13 @@ +use std::path::Path; + +use crate::linkobject::import::ImportMember; + +use super::error::ApiSymbolError; + +pub trait ApiSymbolSource<'a> { + fn extract_api_symbol(&self, symbol: &'a str) -> Result, ApiSymbolError>; + + fn api_path(&self) -> &Path { + Path::new("API") + } +} diff --git a/src/bin/boflink/arguments.rs b/src/bin/boflink/arguments.rs new file mode 100644 index 0000000..2321d4a --- /dev/null +++ b/src/bin/boflink/arguments.rs @@ -0,0 +1,165 @@ +use std::path::PathBuf; + +use boflink::linker::LinkerTargetArch; +use clap::{Parser, ValueEnum}; +use clap_verbosity_flag::{InfoLevel, Verbosity}; + +#[derive(Parser, Debug)] +#[command(version, about)] +pub struct CliArgs { + /// Set the output file name + #[arg( + short, + long, + default_value = "a.bof", + value_name = "file", + value_hint = clap::ValueHint::FilePath + )] + pub output: PathBuf, + + /// Files to link + #[arg( + value_name = "files", + value_hint = clap::ValueHint::FilePath + )] + pub files: Vec, + + /// Add the specified library to search for symbols + #[arg(id = "library", short, long, value_name = "libname")] + pub libraries: Vec, + + /// Add the directory to the library search path + #[arg( + id = "library-path", + short = 'L', + long, + value_name = "directory", + value_hint = clap::ValueHint::DirPath + )] + pub library_paths: Vec, + + /// Set the sysroot path + #[arg( + long, + value_name = "directory", + value_hint = clap::ValueHint::DirPath + )] + pub sysroot: Option, + + /// Set the target machine emulation + #[arg(short, long, value_name = "emulation")] + pub machine: Option, + + /// Name of the entrypoint + #[arg(short, long, value_name = "entry", default_value = "go")] + pub entry: String, + + /// Dump the link graph to the specified file + #[arg(long, value_name = "file", value_hint = clap::ValueHint::FilePath)] + pub dump_link_graph: Option, + + /// Custom API to use instead of the Beacon API + #[arg(long, value_name = "library", visible_alias = "api")] + pub custom_api: Option, + + /// Initialize the .bss section and merge it with the .data section + #[arg(long)] + pub merge_bss: bool, + + /// Print colored output + #[arg(long, value_name = "color", default_value_t = ColorOption::Auto)] + pub color: ColorOption, + + #[command(flatten)] + pub verbose: Verbosity, + + /// Print timing information + #[arg(long)] + pub print_timing: bool, +} + +#[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)] +pub enum TargetEmulation { + #[value(name = "i386pep")] + I386Pep, + + #[value(name = "i386pe")] + I386Pe, +} + +impl From for LinkerTargetArch { + fn from(value: TargetEmulation) -> Self { + match value { + TargetEmulation::I386Pep => LinkerTargetArch::Amd64, + TargetEmulation::I386Pe => LinkerTargetArch::I386, + } + } +} + +#[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)] +pub enum DumpGraphState { + #[value(name = "linked")] + Linked, + + #[value(name = "merged")] + Merged, +} + +impl std::fmt::Display for DumpGraphState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some(v) = self.to_possible_value() { + write!(f, "{}", v.get_name())? + } + + Ok(()) + } +} + +#[derive(ValueEnum, Clone, Copy, Debug, PartialEq, Eq)] +pub enum ColorOption { + #[value(name = "never")] + Never, + + #[value(name = "auto")] + Auto, + + #[value(name = "always")] + Always, + + #[value(name = "ansi")] + AlwaysAnsi, +} + +impl std::fmt::Display for ColorOption { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some(v) = self.to_possible_value() { + write!(f, "{}", v.get_name())?; + } + + Ok(()) + } +} + +impl From for termcolor::ColorChoice { + fn from(val: ColorOption) -> Self { + match val { + ColorOption::Never => termcolor::ColorChoice::Never, + ColorOption::Auto => termcolor::ColorChoice::Auto, + ColorOption::Always => termcolor::ColorChoice::Always, + ColorOption::AlwaysAnsi => termcolor::ColorChoice::AlwaysAnsi, + } + } +} + +/// Parses the command line arguments into the [`CliArgs`]. +pub fn parse_arguments() -> anyhow::Result { + let args = CliArgs::parse_from(argfile::expand_args_from( + std::env::args_os().filter(|arg| arg != "-Bdynamic"), + argfile::parse_fromfile, + argfile::PREFIX, + )?); + + crate::logging::setup_logger(&args)?; + + Ok(args) +} diff --git a/src/bin/boflink/logging.rs b/src/bin/boflink/logging.rs new file mode 100644 index 0000000..8cb4723 --- /dev/null +++ b/src/bin/boflink/logging.rs @@ -0,0 +1,100 @@ +use std::io::{IsTerminal, Write}; + +use log::Level; +use termcolor::{BufferWriter, Color, ColorChoice, ColorSpec, WriteColor}; + +use crate::arguments::{CliArgs, ColorOption}; + +struct CliLogger { + stdout: BufferWriter, + stderr: BufferWriter, +} + +impl log::Log for CliLogger { + #[inline] + fn enabled(&self, _metadata: &log::Metadata) -> bool { + true + } + + fn log(&self, record: &log::Record) { + if record.args().as_str().is_some_and(|args| args.is_empty()) { + return; + } + + let writer = if record.level() <= Level::Warn { + &self.stderr + } else { + &self.stdout + }; + + let mut buffer = writer.buffer(); + write!(buffer, "{}: ", env!("CARGO_BIN_NAME")).unwrap(); + + match record.level() { + Level::Error => { + let _ = buffer.set_color(ColorSpec::new().set_fg(Some(Color::Red)).set_bold(true)); + write!(buffer, "error:").unwrap(); + } + Level::Warn => { + let _ = + buffer.set_color(ColorSpec::new().set_fg(Some(Color::Yellow)).set_bold(true)); + write!(buffer, "warn:").unwrap(); + } + Level::Info => { + let _ = + buffer.set_color(ColorSpec::new().set_fg(Some(Color::Green)).set_bold(true)); + write!(buffer, "info:").unwrap(); + } + Level::Debug => { + let _ = + buffer.set_color(ColorSpec::new().set_fg(Some(Color::White)).set_bold(true)); + write!(buffer, "debug:").unwrap(); + } + Level::Trace => { + let _ = buffer.set_color(ColorSpec::new().set_fg(Some(Color::Blue)).set_bold(true)); + write!(buffer, "trace:").unwrap(); + } + } + + buffer.reset().unwrap(); + writeln!(buffer, " {}", record.args()).unwrap(); + + writer.print(&buffer).unwrap(); + } + + fn flush(&self) {} +} + +/// Sets up logging for the cli +pub fn setup_logger(args: &CliArgs) -> anyhow::Result<()> { + let color_option = if args.color == ColorOption::Auto + && std::env::var("TERM") + .ok() + .is_none_or(|term| !term.eq_ignore_ascii_case("dumb")) + && std::env::var_os("NO_COLOR").is_none() + { + args.color.into() + } else { + ColorChoice::Never + }; + + log::set_boxed_logger(Box::from(CliLogger { + stdout: BufferWriter::stdout( + if color_option != ColorChoice::Never && std::io::stdout().is_terminal() { + color_option + } else { + ColorChoice::Never + }, + ), + stderr: BufferWriter::stderr( + if color_option != ColorChoice::Never && std::io::stderr().is_terminal() { + color_option + } else { + ColorChoice::Never + }, + ), + })) + .map(|()| log::set_max_level(args.verbose.log_level_filter()))?; + + Ok(()) +} diff --git a/src/bin/boflink/main.rs b/src/bin/boflink/main.rs new file mode 100644 index 0000000..71880cd --- /dev/null +++ b/src/bin/boflink/main.rs @@ -0,0 +1,140 @@ +use anyhow::{Result, anyhow, bail}; +use arguments::CliArgs; +use log::{error, info}; + +use boflink::{ + libsearch::LibrarySearcher, + linker::{LinkerBuilder, error::LinkError}, + pathed_item::PathedItem, +}; + +mod arguments; +mod logging; + +#[derive(Debug)] +struct EmptyError; + +impl std::fmt::Display for EmptyError { + fn fmt(&self, _f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + Ok(()) + } +} + +impl std::error::Error for EmptyError {} + +/// cli entrypoint +fn main() { + if let Err(e) = try_main() { + if let Some(link_error) = e.downcast_ref::() { + match link_error { + LinkError::Setup(setup_errors) => { + for setup_error in setup_errors.errors() { + error!("{setup_error}"); + } + } + LinkError::Symbol(symbol_errors) => { + let error_count = symbol_errors.errors().len(); + let mut error_iter = symbol_errors.errors().iter(); + for symbol_error in error_iter.by_ref().take(error_count.saturating_sub(1)) { + error!("{symbol_error}\n"); + } + + if let Some(last_error) = error_iter.next() { + error!("{last_error}"); + } + } + _ => { + error!("{e}"); + } + } + } else if !e.is::() { + error!("{e}"); + } + + std::process::exit(1); + } +} + +/// Main program entrypoint +fn try_main() -> Result<()> { + let mut args = arguments::parse_arguments()?; + + let it = std::time::Instant::now(); + + let link_res = run_linker(&mut args); + + let elapsed = std::time::Instant::now() - it; + if args.print_timing { + info!("link time: {}ms", elapsed.as_micros() as f64 / 1000f64); + } + + link_res +} + +fn run_linker(args: &mut CliArgs) -> anyhow::Result<()> { + let mut library_searcher = LibrarySearcher::new(); + library_searcher.extend_search_paths(std::mem::take(&mut args.library_paths)); + + if cfg!(windows) { + if let Some(libenv) = std::env::var_os("LIB") { + library_searcher.extend_search_paths(std::env::split_paths(&libenv)); + } + } + + let linker = LinkerBuilder::new().library_searcher(library_searcher); + + let linker = if let Some(target_arch) = args.machine.take() { + linker.architecture(target_arch.into()) + } else { + linker + }; + + let linker = if let Some(graph_path) = args.dump_link_graph.take() { + linker.link_graph_path(graph_path) + } else { + linker + }; + + let linker = if let Some(custom_api) = args.custom_api.take() { + linker.custom_api(custom_api) + } else { + linker + }; + + let linker = linker.merge_bss(args.merge_bss); + + let mut error_flag = false; + let inputs = std::mem::take(&mut args.files) + .into_iter() + .filter_map(|file| match std::fs::read(&file) { + Ok(buffer) => Some(PathedItem::new(file, buffer)), + Err(e) => { + error!("could not open {}: {e}", file.display()); + error_flag = true; + None + } + }) + .collect::>(); + + let linker = linker.add_inputs(inputs); + + if error_flag { + bail!(EmptyError); + } + + let linker = linker.add_libraries(std::mem::take(&mut args.libraries)); + + let mut linker = linker.build(); + + match linker.link() { + Ok(built) => { + std::fs::write(&args.output, built) + .map_err(|e| anyhow!("could not write output file: {e}"))?; + } + Err(e) => { + return Err(anyhow!(e)); + } + } + + Ok(()) +} diff --git a/src/drectve/mod.rs b/src/drectve/mod.rs new file mode 100644 index 0000000..c3a7233 --- /dev/null +++ b/src/drectve/mod.rs @@ -0,0 +1,126 @@ +use object::{Object, ObjectSection, coff::CoffFile, pe::IMAGE_SCN_LNK_INFO}; + +use parsers::{Parser, many0, many1, not_token, token}; + +mod parsers; + +pub struct DrectveLibraries<'a> { + section_data: &'a str, +} + +impl<'a> DrectveLibraries<'a> { + fn parse(data: &'a str) -> DrectveLibraries<'a> { + Self { section_data: data } + } +} + +impl<'a> Iterator for DrectveLibraries<'a> { + type Item = &'a str; + + fn next(&mut self) -> Option { + loop { + let ((flag, value), remaining) = many0(token(" ")) + .preceeds(token("-").or(token("/"))) + .preceeds( + many1(not_token(":")).terminated_by(token(":")).then( + many1(not_token("\"")) + .surrounded_by(token("\"")) + .or(many1(not_token(" "))) + .terminated_by(token(" ")), + ), + ) + .parse(self.section_data) + .ok()?; + + self.section_data = remaining; + + if flag.eq_ignore_ascii_case("DEFAULTLIB") { + return Some(value); + } + } + } +} + +pub fn parse_drectve_libraries<'a>(coff: &CoffFile<'a>) -> Option> { + let drectve_section = coff.section_by_name(".drectve")?; + if drectve_section + .coff_section() + .characteristics + .get(object::LittleEndian) + & IMAGE_SCN_LNK_INFO + == 0 + { + return None; + } + + let section_data = drectve_section.data().ok()?; + if section_data + .get(..3) + .is_some_and(|prefix| prefix == [0xef, 0xbb, 0xbf]) + { + Some(DrectveLibraries::parse( + std::str::from_utf8(section_data.get(3..)?).ok()?, + )) + } else { + Some(DrectveLibraries::parse( + std::str::from_utf8(section_data).ok()?, + )) + } +} + +#[cfg(test)] +mod tests { + use super::DrectveLibraries; + + #[test] + fn quoted() { + const INPUT: &str = + " /DEFAULTLIB:\"uuid.lib\" /DEFAULTLIB:\"advapi32.lib\" /DEFAULTLIB:\"OLDNAMES\" "; + + const LIBRARIES: [&str; 3] = ["uuid.lib", "advapi32.lib", "OLDNAMES"]; + + let parsed = DrectveLibraries::parse(INPUT).collect::>(); + for library in LIBRARIES { + assert!( + parsed.contains(&library), + "Could not find {} in {:?}", + library, + parsed + ); + } + } + + #[test] + fn unquoted() { + const INPUT: &str = " /DEFAULTLIB:uuid.lib /DEFAULTLIB:advapi32.lib /DEFAULTLIB:OLDNAMES "; + + const LIBRARIES: [&str; 3] = ["uuid.lib", "advapi32.lib", "OLDNAMES"]; + + let parsed = DrectveLibraries::parse(INPUT).collect::>(); + for library in LIBRARIES { + assert!( + parsed.contains(&library), + "Could not find {} in {:?}", + library, + parsed + ); + } + } + + #[test] + fn mixed() { + const INPUT: &str = + " /DEFAULTLIB:uuid.lib /DEFAULTLIB:\"advapi32.lib\" /DEFAULTLIB:OLDNAMES "; + + const LIBRARIES: [&str; 3] = ["uuid.lib", "advapi32.lib", "OLDNAMES"]; + let parsed = DrectveLibraries::parse(INPUT).collect::>(); + for library in LIBRARIES { + assert!( + parsed.contains(&library), + "Could not find {} in {:?}", + library, + parsed + ); + } + } +} diff --git a/src/drectve/parsers.rs b/src/drectve/parsers.rs new file mode 100644 index 0000000..c4de71f --- /dev/null +++ b/src/drectve/parsers.rs @@ -0,0 +1,337 @@ +use std::marker::PhantomData; + +#[derive(Debug)] +pub struct ParseError<'a> { + #[allow(unused)] + data: &'a str, + + #[allow(unused)] + kind: ParseErrorKind, +} + +#[derive(Debug, PartialEq, Eq)] +pub enum ParseErrorKind { + Or, + Token, + EndOfData, +} + +pub type ParseResult<'a, O> = Result<(O, &'a str), ParseError<'a>>; + +pub trait Parser<'a> { + type Output; + + fn parse(self, data: &'a str) -> ParseResult<'a, Self::Output>; + + #[inline] + fn or

(self, other: P) -> Or<'a, Self::Output, Self, P> + where + Self: Sized, + P: Parser<'a, Output = Self::Output>, + { + Or { + first: self, + second: other, + _m: PhantomData, + _o: PhantomData, + } + } + + #[inline] + fn then

(self, other: P) -> Then<'a, Self::Output, P::Output, Self, P> + where + Self: Sized, + P: Parser<'a>, + { + Then { + first: self, + second: other, + _o: PhantomData, + _m: PhantomData, + } + } + + #[inline] + fn preceeds

(self, parser: P) -> Preceeds<'a, P::Output, Self, P> + where + Self: Sized, + P: Parser<'a>, + { + Preceeds { + first: self, + parser, + _o: PhantomData, + _m: PhantomData, + } + } + + #[inline] + fn surrounded_by

(self, value: P) -> SurroundedBy<'a, Self, P> + where + Self: Sized, + P: Parser<'a> + Clone, + { + SurroundedBy { + surround: value, + parser: self, + _m: PhantomData, + } + } + + #[inline] + fn terminated_by

(self, value: P) -> TerminatedBy<'a, Self, P> + where + Self: Sized, + P: Parser<'a>, + { + TerminatedBy { + terminator: value, + parser: self, + _m: PhantomData, + } + } +} + +impl<'a, T, O> Parser<'a> for T +where + T: FnMut(&'a str) -> ParseResult<'a, O>, +{ + type Output = O; + + #[inline] + fn parse(mut self, data: &'a str) -> Result<(Self::Output, &'a str), ParseError<'a>> { + self(data) + } +} + +#[inline] +pub fn token<'a>(tok: &str) -> impl Parser<'a, Output = &'a str> + Clone { + move |data: &'a str| { + let subdata = data.get(0..tok.len()).ok_or(ParseError { + data, + kind: ParseErrorKind::EndOfData, + })?; + + if subdata == tok { + Ok((subdata, &data[tok.len()..])) + } else { + Err(ParseError { + data, + kind: ParseErrorKind::Token, + }) + } + } +} + +#[inline] +pub fn not_token<'a>(tok: &str) -> impl Parser<'a, Output = &'a str> + Clone { + move |data: &'a str| { + let subdata = data.get(0..tok.len()).ok_or(ParseError { + data, + kind: ParseErrorKind::EndOfData, + })?; + + if subdata != tok { + Ok((subdata, &data[tok.len()..])) + } else { + Err(ParseError { + data, + kind: ParseErrorKind::Token, + }) + } + } +} + +#[inline] +pub fn many0<'a>( + parser: impl Parser<'a, Output = &'a str> + Clone, +) -> impl Parser<'a, Output = &'a str> + Clone { + move |data| { + let mut offset = 0; + let mut parse_data = data; + + while let Ok((parsed, remaining)) = parser.clone().parse(parse_data) { + offset += parsed.len(); + parse_data = remaining; + } + + Ok((&data[..offset], &data[offset..])) + } +} + +#[inline] +pub fn many1<'a>( + parser: impl Parser<'a, Output = &'a str> + Clone, +) -> impl Parser<'a, Output = &'a str> + Clone { + move |data| { + let mut offset = 0; + let mut parse_data = data; + + let (parsed, remaining) = parser.clone().parse(parse_data)?; + offset += parsed.len(); + parse_data = remaining; + + while let Ok((parsed, remaining)) = parser.clone().parse(parse_data) { + offset += parsed.len(); + parse_data = remaining; + } + + Ok((&data[..offset], &data[offset..])) + } +} + +pub struct Or<'a, O, F: Parser<'a, Output = O>, S: Parser<'a, Output = O>> { + first: F, + second: S, + _o: PhantomData, + _m: PhantomData<&'a str>, +} + +impl<'a, F: Parser<'a, Output = &'a str>, S: Parser<'a, Output = &'a str>> Parser<'a> + for Or<'a, &'a str, F, S> +{ + type Output = &'a str; + + fn parse(self, data: &'a str) -> Result<(Self::Output, &'a str), ParseError<'a>> { + if let Ok((parsed, remaining)) = self.first.parse(data) { + Ok((parsed, remaining)) + } else if let Ok((parsed, remaining)) = self.second.parse(data) { + Ok((parsed, remaining)) + } else { + Err(ParseError { + data, + kind: ParseErrorKind::Or, + }) + } + } +} + +pub struct Then<'a, O1, O2, F: Parser<'a, Output = O1>, S: Parser<'a, Output = O2>> { + first: F, + second: S, + _o: PhantomData<(O1, O2)>, + _m: PhantomData<&'a str>, +} + +impl<'a, O1, O2, F: Parser<'a, Output = O1>, S: Parser<'a, Output = O2>> Parser<'a> + for Then<'a, O1, O2, F, S> +{ + type Output = (O1, O2); + + fn parse(self, data: &'a str) -> Result<(Self::Output, &'a str), ParseError<'a>> { + let (p1, remaining) = self.first.parse(data)?; + let (p2, remaining) = self.second.parse(remaining)?; + Ok(((p1, p2), remaining)) + } +} + +pub struct Preceeds<'a, O, F: Parser<'a>, S: Parser<'a, Output = O>> { + first: F, + parser: S, + _o: PhantomData, + _m: PhantomData<&'a str>, +} + +impl<'a, O, F: Parser<'a>, S: Parser<'a, Output = O>> Parser<'a> for Preceeds<'a, O, F, S> { + type Output = O; + + fn parse(self, data: &'a str) -> Result<(Self::Output, &'a str), ParseError<'a>> { + let (_, remaining) = self.first.parse(data)?; + self.parser.parse(remaining) + } +} + +pub struct SurroundedBy<'a, P: Parser<'a>, S: Parser<'a> + Clone> { + surround: S, + parser: P, + _m: PhantomData<&'a str>, +} + +impl<'a, P: Parser<'a>, S: Parser<'a> + Clone> Parser<'a> for SurroundedBy<'a, P, S> { + type Output = P::Output; + + fn parse(self, data: &'a str) -> Result<(Self::Output, &'a str), ParseError<'a>> { + let (_, remaining) = self.surround.clone().parse(data)?; + let (parsed, remaining) = self.parser.parse(remaining)?; + let (_, remaining) = self.surround.parse(remaining)?; + Ok((parsed, remaining)) + } +} + +pub struct TerminatedBy<'a, P: Parser<'a>, T: Parser<'a>> { + terminator: T, + parser: P, + _m: PhantomData<&'a str>, +} + +impl<'a, P: Parser<'a>, T: Parser<'a>> Parser<'a> for TerminatedBy<'a, P, T> { + type Output = P::Output; + + fn parse(self, data: &'a str) -> Result<(Self::Output, &'a str), ParseError<'a>> { + let (parsed, remaining) = self.parser.parse(data)?; + let (_, remaining) = self.terminator.parse(remaining)?; + Ok((parsed, remaining)) + } +} + +#[cfg(test)] +mod tests { + use super::{Parser, many0, many1, not_token, token}; + + #[test] + fn token_parser() { + let (parsed, remaining) = token("hello").parse("hello, world").unwrap(); + assert_eq!(parsed, "hello"); + assert_eq!(remaining, ", world"); + } + + #[test] + fn not_token_parser() { + let (parsed, remaining) = not_token("a").parse("b").unwrap(); + assert_eq!(parsed, "b"); + assert_eq!(remaining, ""); + } + + #[test] + fn many0_combinator() { + let (parsed, remaining) = many0(token("1")).parse("111000").unwrap(); + assert_eq!(parsed, "111"); + assert_eq!(remaining, "000"); + + let (parsed, remaining) = many0(token("0")).parse("111").unwrap(); + assert_eq!(parsed, ""); + assert_eq!(remaining, "111"); + } + + #[test] + fn many1_combinator() { + let (parsed, remaining) = many1(token("1")).parse("111000").unwrap(); + assert_eq!(parsed, "111"); + assert_eq!(remaining, "000"); + + let result = many1(token("0")).parse("111"); + assert!(result.is_err(), "Expected error: {:?}", result); + } + + #[test] + fn or_combinator() { + let (parsed, remaining) = token("1").or(token("2")).parse("2").unwrap(); + assert_eq!(parsed, "2"); + assert!(remaining.is_empty()); + } + + #[test] + fn then_combinator() { + let (parsed, remaining) = token("1").then(token("2")).parse("12").unwrap(); + assert_eq!(parsed.0, "1"); + assert_eq!(parsed.1, "2"); + assert!(remaining.is_empty()); + } + + #[test] + fn preceeds_combinator() { + let (parsed, remaining) = token("0").preceeds(token("1")).parse("01").unwrap(); + assert_eq!(parsed, "1"); + assert_eq!(remaining, ""); + } +} diff --git a/src/graph/built.rs b/src/graph/built.rs new file mode 100644 index 0000000..b3b9d15 --- /dev/null +++ b/src/graph/built.rs @@ -0,0 +1,1026 @@ +use std::{ + cell::OnceCell, + collections::{BTreeMap, LinkedList}, +}; + +use indexmap::IndexMap; +use log::debug; +use object::{ + pe::{ + IMAGE_FILE_LINE_NUMS_STRIPPED, IMAGE_REL_AMD64_REL32, IMAGE_REL_I386_DIR32, + IMAGE_SCN_CNT_CODE, IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_CNT_UNINITIALIZED_DATA, + IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE, IMAGE_SYM_CLASS_EXTERNAL, IMAGE_SYM_CLASS_STATIC, + IMAGE_SYM_TYPE_NULL, + }, + write::coff::{Relocation, SectionHeader, Writer}, +}; + +use crate::linker::LinkerTargetArch; + +use super::{ + edge::{ComdatSelection, DefinitionEdgeWeight, Edge, RelocationEdgeWeight}, + link::{LinkGraph, LinkGraphArena}, + node::{ + CoffNode, LibraryNode, SectionNode, SectionNodeCharacteristics, SectionNodeData, + SymbolName, SymbolNode, SymbolNodeStorageClass, SymbolNodeType, + }, +}; + +const SECTION_ALIGN_SHIFT: u32 = 20; + +#[derive(Debug, thiserror::Error)] +pub enum LinkGraphLinkError { + #[error("{coff_name}: {reference} references symbol '{symbol}' defined in discarded section.")] + DiscardedSection { + coff_name: String, + reference: String, + symbol: String, + }, + + #[error( + "{coff_name}: {section}+{address:#x} relocation is outside section bounds (size = {size:#x})." + )] + RelocationBounds { + coff_name: String, + section: String, + address: u32, + size: u32, + }, + + #[error("{coff_name}: relocation adjustment at '{section}+{address:#x}' overflowed.")] + RelocationOverflow { + coff_name: String, + section: String, + address: u32, + }, +} + +/// An output section with the header and contained sections. +#[derive(Default)] +pub(super) struct OutputSection<'arena, 'data> { + /// The section header + header: SectionHeader, + + /// The list of nodes contained in this output section. + pub nodes: Vec<&'arena SectionNode<'arena, 'data>>, +} + +/// The built link graph with all of the processed inputs. +/// +/// This graph does not allow adding any more inputs and is only used for +/// post-processing and building the linked output file. +pub struct BuiltLinkGraph<'arena, 'data> { + /// Machine value for the output COFF. + machine: LinkerTargetArch, + + /// The sections. + sections: IndexMap<&'arena str, OutputSection<'arena, 'data>>, + + /// The node for the COMMON section. + common_section: OnceCell<&'arena SectionNode<'arena, 'data>>, + + /// Pseudo-COFF for holding metadata sections. + root_coff: &'arena CoffNode<'data>, + + /// The library nodes in the graph. + library_nodes: IndexMap<&'data str, &'arena LibraryNode<'arena, 'data>>, + + /// The API node if it exists. + api_node: Option<&'arena LibraryNode<'arena, 'data>>, + + /// The symbol with external storage class. + external_symbols: IndexMap<&'data str, &'arena SymbolNode<'arena, 'data>>, + + /// Graph arena allocator. + arena: &'arena LinkGraphArena, +} + +impl<'arena, 'data> BuiltLinkGraph<'arena, 'data> { + pub(super) fn new(link_graph: LinkGraph<'arena, 'data>) -> BuiltLinkGraph<'arena, 'data> { + // Partition the sections by name and discard LnkRemove section + let mut sections: IndexMap<&str, OutputSection> = link_graph + .section_nodes + .into_iter() + .filter(|section| { + if section + .characteristics() + .contains(SectionNodeCharacteristics::LnkRemove) + { + debug!( + "{}: discarding 'IMAGE_SCN_LNK_REMOVE' section {}", + section.coff(), + section.name() + ); + section.discard(); + false + } else if section.is_debug() { + debug!( + "{}: discarding debug section {}", + section.coff(), + section.name() + ); + section.discard(); + false + } else { + true + } + }) + .fold(IndexMap::new(), |mut outputs, section_node| { + let section_entry = outputs.entry(section_node.name().group_name()).or_default(); + section_entry.nodes.push(section_node); + outputs + }); + + // Sort grouped sections + sections + .values_mut() + .for_each(|section| section.nodes.sort_by_key(|section| section.name().as_str())); + + // Dedup equivalent .rdata$zzz sections + if let Some(section) = sections.get_mut(".rdata") { + section.nodes.dedup_by(|first, second| { + first + .name() + .group_ordering() + .is_some_and(|ordering| ordering == "zzz") + && second + .name() + .group_ordering() + .is_some_and(|ordering| ordering == "zzz") + // Only dedup them if they have no incoming relocations + && first.relocations().is_empty() + && second.relocations().is_empty() + && first.checksum() == second.checksum() + }); + } + + // Create the built link graph + Self { + machine: link_graph.machine, + sections, + common_section: link_graph.common_section, + library_nodes: link_graph.library_nodes, + root_coff: link_graph.root_coff, + api_node: link_graph.api_node, + external_symbols: link_graph.external_symbols, + arena: link_graph.arena, + } + } + + /// Merge the .bss section with the .data section. + pub fn merge_bss(&mut self) { + self.allocate_commons(); + + let bss_section = self.sections.entry(".bss").or_default(); + let mut bss_nodes = std::mem::take(&mut bss_section.nodes); + + let data_section = self + .sections + .entry(".data") + .or_insert_with(|| OutputSection { + // Manually set the characteristics for the output section to + // match what is expected if the .data section does not already + // exist + header: SectionHeader { + characteristics: IMAGE_SCN_CNT_INITIALIZED_DATA + | IMAGE_SCN_MEM_READ + | IMAGE_SCN_MEM_WRITE, + ..Default::default() + }, + nodes: Vec::with_capacity(bss_nodes.len()), + }); + + data_section.nodes.append(&mut bss_nodes); + debug!("'.bss' output section merged with '.data' section"); + } + + /// Allocate space for COMMON symbols at the end of the .bss + fn allocate_commons(&mut self) { + // Take the value out of the OnceCell to make the function idempotent. + // This function should only run once but may be called multiple times + let common_section = match self.common_section.take() { + Some(section) => section, + None => return, + }; + + // Get the COMMON symbols along with the maximum definition value for + // each symbol. + let mut common_symbols = IndexMap::<&str, &SymbolNode>::from_iter( + common_section + .definitions() + .iter() + .map(|definition| (definition.source().name().as_str(), definition.source())), + ) + .into_values() + .map(|symbol| { + let max_value = symbol + .definitions() + .iter() + .max_by_key(|definition| definition.weight().address()) + .unwrap_or_else(|| { + unreachable!("COMMON symbol should have at least 1 definition associated to it") + }) + .weight() + .address(); + + (symbol, max_value) + }) + .collect::>(); + + // Sort the symbols by size. + common_symbols.sort_by_key(|(_, value)| *value); + + let align = common_section + .characteristics() + .alignment() + .unwrap_or_else(|| { + unreachable!("COMMON section characteristics should have the alignment flag set") + }) as u32; + + // Assign addresses to each symbol. + let mut symbol_addr: u32 = 0; + + for (symbol, symbol_size) in &common_symbols { + symbol_addr = symbol_addr.next_multiple_of(align); + + // Get the first definition edge from the symbol's edge list. + // This will be re-used as the real definition edge with the symbol + // address + let common_def = symbol.definitions().pop_front().unwrap_or_else(|| { + unreachable!("COMMON symbol should have at least 1 definition associated to it") + }); + + // Set the address of the symbol definition to the new value + common_def.weight().set_address(symbol_addr); + + // Clear out the remaining definitions connected to the symbol + symbol.definitions().clear(); + + // Re insert the definition back into the symbol's edge list + symbol.definitions().push_back(common_def); + + // Increment the address for the next symbol + symbol_addr += symbol_size; + } + + // At this point, all of the COMMON symbols should have a single + // definition edge with the address for the symbol. + // The COMMON section still has all of the old definition edges linked + // to its edge list. The COMMON section's edge list needs to be cleared + // and updated with the real definition edges from the symbols. + + // Clear the list of edges associated with the COMMON section + common_section.definitions().clear(); + + // Re insert the definition edges from the COMMON symbols into the + // COMMON section edge list + for (symbol, _) in common_symbols { + let definition = symbol.definitions().front().unwrap_or_else(|| { + unreachable!("COMMON symbol should have at least 1 definition associated to it") + }); + common_section.definitions().push_back(definition); + } + + // Set the size of the COMMON section + common_section.set_uninitialized_size(symbol_addr); + + // Add the COMMON section to the end of the .bss output section + let bss_entry = self + .sections + .entry(".bss") + .or_insert_with(|| OutputSection { + header: SectionHeader { + characteristics: common_section.characteristics().bits(), + ..Default::default() + }, + nodes: Vec::with_capacity(1), + }); + + bss_entry.nodes.push(common_section); + } + + fn apply_import_thunks(&mut self) { + let mut thunk_symbols: LinkedList<(&SymbolNode, SymbolName)> = LinkedList::new(); + + for library_node in self.api_node.iter().chain(self.library_nodes.values()) { + for import_edge in library_node.imports() { + let import_name = import_edge.weight().import_name(); + let symbol = import_edge.source(); + if symbol + .name() + .strip_dllimport() + .is_none_or(|unprefixed| unprefixed != import_name.as_str()) + && !symbol.is_unreferenced() + { + thunk_symbols.push_back((symbol, import_name)); + } + } + } + + if !thunk_symbols.is_empty() { + // jmp [rip + $] + const CODE_THUNK: [u8; 8] = [0xff, 0x25, 0x00, 0x00, 0x00, 0x00, 0x90, 0x90]; + + let code_section_data: &mut [u8] = self + .arena + .alloc_slice_fill_default(CODE_THUNK.len() * thunk_symbols.len()); + + for data_chunk in code_section_data.chunks_mut(CODE_THUNK.len()) { + data_chunk.copy_from_slice(&CODE_THUNK); + } + + let code_section = self.arena.alloc_with(|| { + SectionNode::new( + ".text$zzz", + SectionNodeCharacteristics::CntCode + | SectionNodeCharacteristics::MemExecute + | SectionNodeCharacteristics::MemRead + | SectionNodeCharacteristics::Align8Bytes, + SectionNodeData::Initialized(code_section_data), + 0, + self.root_coff, + ) + }); + + let thunk_reloc = match self.machine { + LinkerTargetArch::Amd64 => RelocationEdgeWeight::new(2, IMAGE_REL_AMD64_REL32), + LinkerTargetArch::I386 => RelocationEdgeWeight::new(2, IMAGE_REL_I386_DIR32), + }; + + for (symbol_num, (symbol_node, import_name)) in thunk_symbols.iter().enumerate() { + let symbol_addr = symbol_num as u32 * 8; + + // Add a definition edge for the existing symbol + let definition_edge = self.arena.alloc_with(|| { + Edge::new( + *symbol_node, + code_section, + DefinitionEdgeWeight::new(symbol_addr, None), + ) + }); + + symbol_node.definitions().push_back(definition_edge); + code_section.definitions().push_back(definition_edge); + + // Add a new thunk import symbol for this symbol + let thunk_import_symbol = self.arena.alloc_with(|| { + SymbolNode::new( + &*self + .arena + .alloc_str(&format!("__imp_{}", import_name.as_str())), + SymbolNodeStorageClass::External, + false, + SymbolNodeType::Value(0), + ) + }); + + // Add a relocation to the thunk import symbol + let relocation_edge = self.arena.alloc_with(|| { + Edge::new( + code_section, + thunk_import_symbol, + RelocationEdgeWeight::new( + thunk_reloc.address() + symbol_addr, + thunk_reloc.typ(), + ), + ) + }); + + code_section.relocations().push_back(relocation_edge); + thunk_import_symbol.references().push_back(relocation_edge); + + // Unlink the import edge from the existing symbol + let removed_import_edge = symbol_node.imports().pop_front().unwrap(); + // Set the source node for the edge to the new thunk import + // symbol + removed_import_edge.replace_source(thunk_import_symbol); + + // Link the edge to the thunk symbol + thunk_import_symbol.imports().push_back(removed_import_edge); + } + + // Add the new code section to the list of sections + self.sections + .entry(code_section.name().group_name()) + .or_default() + .nodes + .push(code_section); + } + } + + /// Handles discarding/keeping sections for COMDAT symbols + fn handle_comdats(&self) { + for symbol in self.external_symbols.values() { + let mut definition_iter = symbol.definitions().iter().peekable(); + + let first_definition = match definition_iter.peek() { + Some(definition) => definition, + None => continue, + }; + + let selection = match first_definition.weight().selection { + Some(sel) => sel, + None => continue, + }; + + if selection == ComdatSelection::Any + || selection == ComdatSelection::SameSize + || selection == ComdatSelection::ExactMatch + { + // Keep the first section but discard the rest. + let _ = definition_iter.next(); + + for remaining in definition_iter { + let section = remaining.target(); + debug!( + "{}: discarding COMDAT {} ({selection:?})", + section.coff(), + section.name(), + ); + section.discard(); + } + } else if selection == ComdatSelection::Largest { + // Find the largest size and discard the rest. + let mut largest_section: Option<&'arena SectionNode<'arena, 'data>> = None; + + for definition in definition_iter { + let section = definition.target(); + + if let Some(largest) = &mut largest_section { + if largest.data().len() < section.data().len() { + debug!( + "{}: discarding COMDAT {} ({selection:?})", + largest.coff(), + largest.name() + ); + largest.discard(); + *largest = section; + } + } else { + largest_section = Some(section); + } + } + } else if selection == ComdatSelection::Associative { + // Associative COMDAT symbols are handled by traversing the + // root of the COMDAT chain. + continue; + } + + for definition in symbol.definitions() { + let root_section = definition.target(); + + // Discard or keep the associated sections depending on if + // the root section was kept or discarded. + let root_discarded = root_section.is_discarded(); + + for associative_section in root_section.associative_bfs() { + if !associative_section.is_discarded() && root_discarded { + debug!( + "{}: discarding COMDAT {}. associative to discarded root ({}:{})", + associative_section.coff(), + associative_section.name(), + root_section.coff().short_name(), + root_section.name() + ); + } + + associative_section.set_discarded(root_discarded); + } + } + } + } + + /// Links the graph components together and builds the final COFF. + pub fn link(mut self) -> Result, LinkGraphLinkError> { + self.apply_import_thunks(); + self.handle_comdats(); + self.allocate_commons(); + + // Remove discarded section nodes. + // Discard output sections which no longer have any input sections. + self.sections.retain(|section_name, section| { + section.nodes.retain(|node| !node.is_discarded()); + if section.nodes.is_empty() { + debug!("discarding output section '{section_name}'"); + false + } else { + true + } + }); + + let mut built_coff = Vec::new(); + let mut coff_writer = Writer::new(&mut built_coff); + + coff_writer.reserve_file_header(); + + for (section_name, section) in self.sections.iter_mut() { + section.header.name = coff_writer.add_name(section_name.as_bytes()); + let mut section_alignment: u32 = 0; + + let mut section_nodes_iter = section.nodes.iter().peekable(); + + // Get the characteristics from the first node and use them if not + // already set + if section.header.characteristics == 0 { + if let Some(first_node) = section_nodes_iter.peek() { + let mut flags = first_node.characteristics().zero_align(); + + // Remove the COMDAT flag + flags.remove(SectionNodeCharacteristics::LnkComdat); + section.header.characteristics = flags.bits(); + } + } + + // Assign virtual addresses to each section + for node in section_nodes_iter { + // Include alignment needed to satisfy input section node + // alignment + if let Some(align) = node.characteristics().alignment() { + let align = align as u32; + section.header.size_of_raw_data = + section.header.size_of_raw_data.next_multiple_of(align); + section_alignment = section_alignment.max(align); + } + + debug!( + "{}: mapping section '{}' to '{}' at address {:#x} with size {:#x}", + node.coff(), + node.name(), + section_name, + section.header.size_of_raw_data, + node.data().len(), + ); + + node.assign_virtual_address(section.header.size_of_raw_data); + section.header.size_of_raw_data += node.data().len() as u32; + } + + // Set the alignment needed for this section + if section_alignment != 0 { + section.header.characteristics |= + (section_alignment.ilog2() + 1) << SECTION_ALIGN_SHIFT; + } + } + + // Reserve section headers + coff_writer.reserve_section_headers(self.sections.len().try_into().unwrap()); + + // Reserve section data only if the data is initialized + for section in self.sections.values_mut() { + if section.header.characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA == 0 { + section.header.pointer_to_raw_data = + coff_writer.reserve_section(section.header.size_of_raw_data as usize); + } + } + + // Reserve relocations skipping relocations to the same output section + for (section_name, section) in self.sections.iter_mut() { + let mut reloc_count = 0usize; + + for section_node in §ion.nodes { + for reloc in section_node.relocations() { + let symbol = reloc.target(); + + if let Some(definition) = symbol + .definitions() + .iter() + .find(|definition| !definition.target().is_discarded()) + { + if definition.target().name().group_name() == *section_name { + continue; + } + } else if symbol.imports().is_empty() { + // Symbol has no imports and all definitions are in + // discarded sections. Return an error. + + let coff_name = section_node.coff().to_string(); + + let symbol_defs = BTreeMap::from_iter( + section_node.definitions().iter().filter_map(|definition| { + let ref_symbol = definition.source(); + if ref_symbol.is_section_symbol() || ref_symbol.is_label() { + None + } else { + Some((definition.weight().address(), ref_symbol.name())) + } + }), + ); + + if let Some(reference_symbol) = + symbol_defs.range(0..=reloc.weight().address()).next_back() + { + return Err(LinkGraphLinkError::DiscardedSection { + coff_name, + reference: reference_symbol.1.demangle().to_string(), + symbol: symbol.name().demangle().to_string(), + }); + } else { + return Err(LinkGraphLinkError::DiscardedSection { + coff_name, + reference: format!( + "{}+{:#x}", + section_node.name(), + reloc.weight().address() + ), + symbol: symbol.name().demangle().to_string(), + }); + } + } + + reloc_count += 1; + } + } + + section.header.number_of_relocations = reloc_count.try_into().unwrap(); + section.header.pointer_to_relocations = coff_writer.reserve_relocations(reloc_count); + } + + // Reserve symbols defined in sections + for section in self.sections.values() { + // Reserve the section symbol + let section_symbol_index = coff_writer.reserve_symbol_index(); + let _ = coff_writer.reserve_aux_section(); + + for section_node in §ion.nodes { + // Assign table indicies to defined symbols + for definition in section_node.definitions() { + let symbol = definition.source(); + + // Section symbol already reserved. Set the index to the + // existing one + if symbol.is_section_symbol() { + symbol + .assign_table_index(section_symbol_index) + .unwrap_or_else(|v| { + panic!( + "symbol {} already assigned to symbol table index {v}", + symbol.name().demangle() + ) + }); + } else if symbol.is_label() { + // Associate labels with the section symbol + symbol + .assign_table_index(section_symbol_index) + .unwrap_or_else(|v| { + panic!( + "symbol {} already assigned to symbol table index {v}", + symbol.name().demangle() + ) + }); + } else { + let _ = symbol.output_name().get_or_init(|| { + coff_writer.add_name(symbol.name().as_str().as_bytes()) + }); + + // Reserve an index for this symbol + symbol + .assign_table_index(coff_writer.reserve_symbol_index()) + .unwrap_or_else(|v| { + panic!( + "symbol {} already assigned to symbol table index {v}", + symbol.name().demangle() + ) + }); + } + } + } + } + + // Reserve API imported symbols + if let Some(api_node) = self.api_node { + for import in api_node.imports() { + let symbol = import.source(); + + let _ = symbol + .output_name() + .get_or_init(|| coff_writer.add_name(symbol.name().as_str().as_bytes())); + + symbol + .assign_table_index(coff_writer.reserve_symbol_index()) + .unwrap_or_else(|v| { + panic!( + "symbol {} already assigned to symbol table index {v}", + symbol.name().demangle() + ) + }); + } + } + + // Reserve library imported symbols + for library in self.library_nodes.values() { + for import in library.imports() { + let symbol = import.source(); + + let name = self.arena.alloc_str(&format!( + "__imp_{}${}", + library.name().trim_dll_suffix(), + import.weight().import_name() + )); + + let _ = symbol + .output_name() + .get_or_init(|| coff_writer.add_name(name.as_bytes())); + + symbol + .assign_table_index(coff_writer.reserve_symbol_index()) + .unwrap_or_else(|v| { + panic!( + "symbol {} already assigned to symbol table index {v}", + symbol.name().demangle() + ) + }); + } + } + + // Finish reserving COFF data + coff_writer.reserve_symtab_strtab(); + + // Write out the file header + coff_writer + .write_file_header(object::write::coff::FileHeader { + machine: self.machine.into(), + time_date_stamp: 0, + characteristics: IMAGE_FILE_LINE_NUMS_STRIPPED, + }) + .unwrap(); + + // Write out the section headers + for section in self.sections.values() { + coff_writer.write_section_header(section.header.clone()); + } + + // Write out the section data + for section in self.sections.values() { + if section.header.size_of_raw_data > 0 + && section.header.characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA == 0 + { + coff_writer.write_section_align(); + + let alignment_byte = if (section.header.characteristics & IMAGE_SCN_CNT_CODE) != 0 { + 0x90u8 + } else { + 0x00u8 + }; + + let mut data_written = 0; + let mut alignment_buffer = vec![alignment_byte; 16]; + + for node in section.nodes.iter() { + // Write alignment padding + let needed = node.virtual_address() - data_written; + if needed > 0 { + alignment_buffer.resize(needed as usize, alignment_byte); + coff_writer.write(&alignment_buffer); + data_written += needed; + } + + let section_data = match node.data() { + SectionNodeData::Initialized(data) => data, + SectionNodeData::Uninitialized(size) => { + // This node contains uninitialized data but the + // output section should be initialized. + // Write out padding bytes to satisfy the size + // requested + alignment_buffer.resize(size as usize, alignment_byte); + alignment_buffer.as_slice() + } + }; + + coff_writer.write(section_data); + data_written += section_data.len() as u32; + } + } + } + + // Write out the relocations skipping relocations to the same section + for (section_name, section) in self.sections.iter() { + for section_node in §ion.nodes { + for reloc in section_node.relocations() { + let target_symbol = reloc.target(); + + if let Some(symbol_definition) = target_symbol + .definitions() + .iter() + .find(|definition| !definition.target().is_discarded()) + { + if symbol_definition.target().name().group_name() == *section_name { + continue; + } + } + + coff_writer.write_relocation(Relocation { + virtual_address: section_node.virtual_address() + reloc.weight().address(), + symbol: target_symbol.table_index().unwrap_or_else(|| { + panic!( + "symbol {} was never assigned a symbol table index", + target_symbol.name().demangle() + ) + }), + typ: reloc.weight().typ(), + }); + } + } + } + + // Write out symbols defined in sections + for (section_index, section) in self.sections.values().enumerate() { + // Write the section symbol + coff_writer.write_symbol(object::write::coff::Symbol { + name: section.header.name, + value: 0, + section_number: (section_index + 1).try_into().unwrap(), + typ: IMAGE_SYM_TYPE_NULL, + storage_class: IMAGE_SYM_CLASS_STATIC, + number_of_aux_symbols: 1, + }); + + coff_writer.write_aux_section(object::write::coff::AuxSymbolSection { + length: section.header.size_of_raw_data, + number_of_relocations: section.header.number_of_relocations, + number_of_linenumbers: 0, + // The object crate will calculate the checksum + check_sum: 0, + number: (section_index + 1).try_into().unwrap(), + selection: 0, + }); + + for section_node in §ion.nodes { + for definition in section_node.definitions() { + let symbol = definition.source(); + + // Skip labels and section symbols + if !symbol.is_section_symbol() && !symbol.is_label() { + coff_writer.write_symbol(object::write::coff::Symbol { + name: symbol.output_name().get().copied().unwrap_or_else(|| { + panic!( + "symbol {} never had the name reserved in the output COFF", + symbol.name().demangle() + ) + }), + value: definition.weight().address() + section_node.virtual_address(), + section_number: (section_index + 1).try_into().unwrap(), + typ: match symbol.typ() { + SymbolNodeType::Value(typ) => typ, + _ => unreachable!(), + }, + storage_class: symbol.storage_class().into(), + number_of_aux_symbols: 0, + }); + } + } + } + } + + // Write out API imported symbols + if let Some(api_node) = self.api_node { + for import in api_node.imports() { + let symbol = import.source(); + coff_writer.write_symbol(object::write::coff::Symbol { + name: symbol.output_name().get().copied().unwrap_or_else(|| { + panic!( + "symbol {} never had the name reserved in the output COFF", + symbol.name().demangle() + ) + }), + value: 0, + section_number: 0, + typ: 0, + storage_class: IMAGE_SYM_CLASS_EXTERNAL, + number_of_aux_symbols: 0, + }); + } + } + + // Write out library imported symbols + for library in self.library_nodes.values() { + for import in library.imports() { + let symbol = import.source(); + coff_writer.write_symbol(object::write::coff::Symbol { + name: symbol.output_name().get().copied().unwrap_or_else(|| { + panic!( + "symbol {} never had the name reserved in the output COFF", + symbol.name().demangle() + ) + }), + value: 0, + section_number: 0, + typ: 0, + storage_class: IMAGE_SYM_CLASS_EXTERNAL, + number_of_aux_symbols: 0, + }); + } + } + + // Finish writing the COFF + coff_writer.write_strtab(); + + // Fixup relocations + for section in self.sections.values() { + let section_data_base = section.header.pointer_to_raw_data as usize; + for section_node in §ion.nodes { + let section_data_ptr = section_data_base + section_node.virtual_address() as usize; + + let section_data = + &mut built_coff[section_data_ptr..section_data_ptr + section_node.data().len()]; + + for reloc_edge in section_node.relocations() { + let target_symbol = reloc_edge.target(); + + let symbol_definition = match target_symbol + .definitions() + .iter() + .find(|definition| !definition.target().is_discarded()) + { + Some(definition) => definition, + None => continue, + }; + + let target_section = symbol_definition.target(); + let reloc = reloc_edge.weight(); + + // Return an error if the relocation is out of bounds. + if reloc.virtual_address + 4 > section_node.data().len() as u32 { + return Err(LinkGraphLinkError::RelocationBounds { + coff_name: section_node.coff().to_string(), + section: section_node.name().to_string(), + address: reloc.virtual_address, + size: section_node.data().len() as u32, + }); + } + + // The relocation bounds check above checks the relocation + // in the graph. This indexes into the built COFF after + // everything is merged. The slice index should always + // be in bounds but if it is not, there is some logic + // error above. Panic with a verbose error message if that + // is the case. + + let reloc_data: [u8; 4] = section_data + .get(reloc.address() as usize..reloc.address() as usize + 4) + .map(|data| data.try_into().unwrap_or_else(|_| unreachable!())) + .unwrap_or_else(|| { + unreachable!( + "relocation in section '{}' is out of bounds", + section_node.name() + ) + }); + + // Update relocations + let relocated_val = if target_symbol.is_section_symbol() { + // Target symbol is a section symbol. Relocations need to + // be adjusted to account for the section shift. + let reloc_val = u32::from_le_bytes(reloc_data); + + reloc_val + .checked_add(target_section.virtual_address()) + .ok_or_else(|| LinkGraphLinkError::RelocationOverflow { + coff_name: section_node.coff().to_string(), + section: section_node.name().to_string(), + address: reloc.address(), + })? + } else if section_node.name().group_name() == target_section.name().group_name() + { + // Relocation targets a symbol defined in the same section. + // Apply the relocation to the symbol address. + + let reloc_addr = reloc.address() + section_node.virtual_address(); + let symbol_addr = + symbol_definition.weight().address() + target_section.virtual_address(); + + let reloc_val = u32::from_be_bytes(reloc_data); + let delta = symbol_addr.wrapping_sub(reloc_addr + 4); + reloc_val.wrapping_add(delta) + } else if target_symbol.is_label() { + // Old relocation target symbol is a label. The current + // relocation points to the section symbol and the label + // was discarded. + // Handle this like a section symbol relocation but + // shift it to point to the label's virtual address in + // the section. + let reloc_val = u32::from_le_bytes(reloc_data); + let symbol_addr = symbol_definition.weight().address(); + + reloc_val + .checked_add(target_section.virtual_address()) + .and_then(|reloc_val| reloc_val.checked_add(symbol_addr)) + .ok_or_else(|| LinkGraphLinkError::RelocationOverflow { + coff_name: section_node.coff().to_string(), + section: section_node.name().to_string(), + address: reloc.address(), + })? + } else { + // Relocation target is symbolic and does not need + // updating + continue; + }; + + // Write the new reloc + section_data[reloc.address() as usize..reloc.address() as usize + 4] + .copy_from_slice(&relocated_val.to_le_bytes()); + } + } + } + + Ok(built_coff) + } +} diff --git a/src/graph/cache.rs b/src/graph/cache.rs new file mode 100644 index 0000000..1b231d8 --- /dev/null +++ b/src/graph/cache.rs @@ -0,0 +1,96 @@ +use std::collections::HashMap; + +use object::{SectionIndex, SymbolIndex}; + +use super::{ + edge::ComdatSelection, + node::{SectionNode, SymbolNode}, +}; + +/// Cache for inserting COFFs into the graph. +pub struct LinkGraphCache<'arena, 'data> { + /// Cached symbols. + symbols: HashMap>, + + /// Cached sections. + sections: HashMap>, + + /// Cached selection and section symbol values for COMDAT symbols. + comdat_selections: HashMap, +} + +impl<'arena, 'data> LinkGraphCache<'arena, 'data> { + #[inline] + pub fn new() -> LinkGraphCache<'arena, 'data> { + Self { + symbols: HashMap::new(), + sections: HashMap::new(), + comdat_selections: HashMap::new(), + } + } + + #[inline] + pub fn with_capacity(symbols: usize, sections: usize) -> LinkGraphCache<'arena, 'data> { + Self { + symbols: HashMap::with_capacity(symbols), + sections: HashMap::with_capacity(sections), + comdat_selections: HashMap::new(), + } + } + + #[inline] + pub fn clear(&mut self) { + self.symbols.clear(); + self.sections.clear(); + self.comdat_selections.clear(); + } + + #[inline] + pub fn reserve_symbols(&mut self, additional: usize) { + self.symbols.reserve(additional); + } + + #[inline] + pub fn reserve_sections(&mut self, additional: usize) { + self.sections.reserve(additional); + } + + #[inline] + pub fn reserve_comdat_selections(&mut self, additional: usize) { + self.comdat_selections.reserve(additional); + } + + #[inline] + pub fn insert_section( + &mut self, + idx: SectionIndex, + section: &'arena SectionNode<'arena, 'data>, + ) { + let _ = self.sections.insert(idx, section); + } + + #[inline] + pub fn insert_symbol(&mut self, idx: SymbolIndex, symbol: &'arena SymbolNode<'arena, 'data>) { + let _ = self.symbols.insert(idx, symbol); + } + + #[inline] + pub fn insert_comdat_selection(&mut self, idx: SectionIndex, selection: ComdatSelection) { + let _ = self.comdat_selections.insert(idx, selection); + } + + #[inline] + pub fn get_symbol(&self, idx: SymbolIndex) -> Option<&'arena SymbolNode<'arena, 'data>> { + self.symbols.get(&idx).copied() + } + + #[inline] + pub fn get_section(&self, idx: SectionIndex) -> Option<&'arena SectionNode<'arena, 'data>> { + self.sections.get(&idx).copied() + } + + #[inline] + pub fn get_comdat_selection(&self, idx: SectionIndex) -> Option { + self.comdat_selections.get(&idx).copied() + } +} diff --git a/src/graph/edge.rs b/src/graph/edge.rs new file mode 100644 index 0000000..0fa147d --- /dev/null +++ b/src/graph/edge.rs @@ -0,0 +1,402 @@ +use std::{cell::Cell, marker::PhantomData}; + +use super::node::SymbolName; + +use __private::SealedTrait; +use num_enum::{IntoPrimitive, TryFromPrimitive}; +use object::pe::{ + IMAGE_COMDAT_SELECT_ANY, IMAGE_COMDAT_SELECT_ASSOCIATIVE, IMAGE_COMDAT_SELECT_EXACT_MATCH, + IMAGE_COMDAT_SELECT_LARGEST, IMAGE_COMDAT_SELECT_NODUPLICATES, IMAGE_COMDAT_SELECT_SAME_SIZE, +}; + +pub trait EdgeListTraversal: SealedTrait {} + +pub struct OutgoingEdges; +impl SealedTrait for OutgoingEdges {} +impl EdgeListTraversal for OutgoingEdges {} + +pub struct IncomingEdges; +impl SealedTrait for IncomingEdges {} +impl EdgeListTraversal for IncomingEdges {} + +pub trait EdgeListEntry<'arena, Source, Target, Weight, Tr: EdgeListTraversal>: + SealedTrait +{ + fn next_node(&self) -> &Cell>>; +} + +#[derive(Debug, Copy, Clone, thiserror::Error)] +#[error("invalid COMDAT selection ({0})")] +pub struct TryFromComdatSelectionError(u8); + +/// An adjacency list for a node's adjacent edges. +pub struct EdgeList<'arena, Source, Target, Weight, Tr: EdgeListTraversal> +where + Edge<'arena, Source, Target, Weight>: EdgeListEntry<'arena, Source, Target, Weight, Tr>, +{ + /// The head edge in the list. + head: Cell>>, + + /// The tail edge in the list. + tail: Cell>>, + + /// The number of edges in the list. + size: Cell, + + /// The traversal type for this edge list. + _traversal: PhantomData, +} + +impl<'arena, Source, Target, Weight, Tr: EdgeListTraversal> + EdgeList<'arena, Source, Target, Weight, Tr> +where + Edge<'arena, Source, Target, Weight>: EdgeListEntry<'arena, Source, Target, Weight, Tr>, +{ + /// Creates a new empty [`EdgeList`]. + pub(super) fn new() -> EdgeList<'arena, Source, Target, Weight, Tr> { + Self { + head: Cell::new(None), + tail: Cell::new(None), + size: Cell::new(0), + _traversal: PhantomData, + } + } + + /// Returns the number of entries in this [`EdgeList`]. + #[inline] + pub fn len(&self) -> usize { + self.size.get() + } + + /// Returns `true` if the list is empty. + #[inline] + pub fn is_empty(&self) -> bool { + self.head.get().is_none() + } + + /// Returns the first edge in the edge list if the list is non-empty. + #[inline] + pub fn front(&self) -> Option<&'arena Edge<'arena, Source, Target, Weight>> { + self.head.get() + } + + /// Returns the last edge in the edge list if the list is non-empty. + #[inline] + pub fn back(&self) -> Option<&'arena Edge<'arena, Source, Target, Weight>> { + self.tail.get() + } + + /// Returns an [`EdgeListIter`] for iterating over the list of edges. + #[inline] + pub fn iter(&self) -> EdgeListIter<'arena, Source, Target, Weight, Tr> { + EdgeListIter((self.head.get(), PhantomData)) + } + + /// Adds an edge to this list with the specified weight and linked to the + /// target node. + pub(super) fn push_back(&self, edge: &'arena Edge<'arena, Source, Target, Weight>) { + if let Some(tail_node) = self.tail.get() { + tail_node.next_node().set(Some(edge)); + self.tail.set(Some(edge)); + } else { + self.head.set(Some(edge)); + self.tail.set(Some(edge)); + } + + self.size.set(self.size.get() + 1); + } + + /// Removes the first item from the edge list and returns it. + /// + /// # Note + /// This will leak the removed edge. + pub(super) fn pop_front(&self) -> Option<&'arena Edge<'arena, Source, Target, Weight>> { + let removed_edge = self.head.get()?; + let size = self.size.get().saturating_sub(1); + + self.head.set(removed_edge.next_node().take()); + if size == 0 { + self.tail.take(); + } + + self.size.set(size); + Some(removed_edge) + } + + /// Removes all of the nodes from the edge list. + /// + /// # Note + /// This does not deallocate the edges since they are handled by the arena. + pub(super) fn clear(&self) { + self.head.set(None); + self.tail.set(None); + self.size.set(0); + } +} + +impl<'arena, Source, Target, Weight, T: EdgeListTraversal> IntoIterator + for EdgeList<'arena, Source, Target, Weight, T> +where + Edge<'arena, Source, Target, Weight>: EdgeListEntry<'arena, Source, Target, Weight, T>, + EdgeListIter<'arena, Source, Target, Weight, T>: + Iterator>, +{ + type Item = ::Item; + type IntoIter = EdgeListIter<'arena, Source, Target, Weight, T>; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +impl<'arena, Source, Target, Weight, T: EdgeListTraversal> IntoIterator + for &EdgeList<'arena, Source, Target, Weight, T> +where + Edge<'arena, Source, Target, Weight>: EdgeListEntry<'arena, Source, Target, Weight, T>, + EdgeListIter<'arena, Source, Target, Weight, T>: + Iterator>, +{ + type Item = ::Item; + type IntoIter = EdgeListIter<'arena, Source, Target, Weight, T>; + + #[inline] + fn into_iter(self) -> Self::IntoIter { + self.iter() + } +} + +/// Iterator for iterating over the edges of an [`EdgeList`]. +pub struct EdgeListIter<'arena, Source, Target, Weight, T: EdgeListTraversal>( + ( + Option<&'arena Edge<'arena, Source, Target, Weight>>, + PhantomData, + ), +); + +impl Clone + for EdgeListIter<'_, Source, Target, Weight, T> +{ + #[inline] + fn clone(&self) -> Self { + Self(self.0) + } +} + +impl<'arena, Source, Target, Weight, T: EdgeListTraversal> Iterator + for EdgeListIter<'arena, Source, Target, Weight, T> +where + Edge<'arena, Source, Target, Weight>: EdgeListEntry<'arena, Source, Target, Weight, T>, +{ + type Item = &'arena Edge<'arena, Source, Target, Weight>; + + fn next(&mut self) -> Option { + let curr = self.0.0?; + self.0.0 = EdgeListEntry::<_, _, _, T>::next_node(curr).get(); + Some(curr) + } +} + +/// A graph edge. +pub struct Edge<'arena, S, T, W> { + /// The next outgoing edge in the list of outgoing edges for the source + /// node + next_outgoing: Cell>>, + + /// The next incoming edge in the list of incoming edges for the target + /// node + next_incoming: Cell>>, + + /// Reference to the source node for this edge. + source_node: Cell<&'arena S>, + + /// Reference to the target node for this edge. + target_node: Cell<&'arena T>, + + /// The edge weight + weight: W, +} + +impl<'arena, S, T, W> Edge<'arena, S, T, W> { + #[inline] + pub(super) fn new( + source_node: &'arena S, + target_node: &'arena T, + weight: W, + ) -> Edge<'arena, S, T, W> { + Self { + next_outgoing: Cell::new(None), + next_incoming: Cell::new(None), + source_node: Cell::new(source_node), + target_node: Cell::new(target_node), + weight, + } + } + + /// Replaces the source node joined to this edge. The edge must be removed + /// from the source node before it can be replaced. + #[inline] + pub(super) fn replace_source(&self, source_node: &'arena S) { + debug_assert!(self.next_outgoing.get().is_none()); + self.source_node.replace(source_node); + } + + /// Returns a reference to the source node joined to this edge. + #[inline] + pub fn source(&self) -> &'arena S { + self.source_node.get() + } + + /// Returns a reference to the target node joined to this edge. + #[inline] + pub fn target(&self) -> &'arena T { + self.target_node.get() + } + + /// Returns a reference to the edge weight + #[inline] + pub fn weight(&self) -> &W { + &self.weight + } + + /// Returns a mutable reference to the edge weight + #[inline] + pub fn weight_mut(&mut self) -> &mut W { + &mut self.weight + } +} + +impl SealedTrait for Edge<'_, S, T, W> {} + +impl<'arena, Source, Target, Weight> EdgeListEntry<'arena, Source, Target, Weight, OutgoingEdges> + for Edge<'arena, Source, Target, Weight> +{ + #[inline] + fn next_node(&self) -> &Cell>> { + &self.next_outgoing + } +} + +impl<'arena, Source, Target, Weight> EdgeListEntry<'arena, Source, Target, Weight, IncomingEdges> + for Edge<'arena, Source, Target, Weight> +{ + #[inline] + fn next_node(&self) -> &Cell>> { + &self.next_incoming + } +} + +/// The weight for a definition edge. +pub struct DefinitionEdgeWeight { + /// The virtual address for the definition. + virtual_address: Cell, + + /// The COMDAT selection if the symbol is a COMDAT symbol. + pub(super) selection: Option, +} + +impl DefinitionEdgeWeight { + #[inline] + pub(super) fn new( + virtual_address: u32, + selection: Option, + ) -> DefinitionEdgeWeight { + Self { + virtual_address: Cell::new(virtual_address), + selection, + } + } + + /// Returns the address of the symbol + #[inline] + pub fn address(&self) -> u32 { + self.virtual_address.get() + } + + /// Sets the virtual address for the symbol. + /// + /// Used for assigning addresses to COMMON symbols. + #[inline] + pub fn set_address(&self, val: u32) { + self.virtual_address.set(val); + } + + /// Returns the COMDAT selection for the symbol if this is a COMDAT symbol. + #[inline] + pub fn selection(&self) -> Option { + self.selection + } +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq, TryFromPrimitive, IntoPrimitive)] +#[num_enum( + error_type( + name = TryFromComdatSelectionError, + constructor = TryFromComdatSelectionError, + ) +)] +#[repr(u8)] +pub enum ComdatSelection { + NoDuplicates = IMAGE_COMDAT_SELECT_NODUPLICATES, + Any = IMAGE_COMDAT_SELECT_ANY, + SameSize = IMAGE_COMDAT_SELECT_SAME_SIZE, + ExactMatch = IMAGE_COMDAT_SELECT_EXACT_MATCH, + Associative = IMAGE_COMDAT_SELECT_ASSOCIATIVE, + Largest = IMAGE_COMDAT_SELECT_LARGEST, +} + +/// The weight for a relocation edge. +pub struct RelocationEdgeWeight { + /// The virtual address of the relocation. + pub(super) virtual_address: u32, + + /// The relocation type. + typ: u16, +} + +impl RelocationEdgeWeight { + #[inline] + pub(super) fn new(virtual_address: u32, typ: u16) -> RelocationEdgeWeight { + Self { + virtual_address, + typ, + } + } + + #[inline] + pub fn address(&self) -> u32 { + self.virtual_address + } + + #[inline] + pub fn typ(&self) -> u16 { + self.typ + } +} + +/// The weight for an import edge. +pub struct ImportEdgeWeight<'data> { + /// The name to import the symbol as. + import_name: SymbolName<'data>, +} + +impl<'data> ImportEdgeWeight<'data> { + #[inline] + pub(super) fn new(import_name: impl Into>) -> ImportEdgeWeight<'data> { + Self { + import_name: import_name.into(), + } + } + + #[inline] + pub fn import_name(&self) -> SymbolName<'data> { + self.import_name + } +} + +/// The weight for a COMDAT associative section edge. +pub struct AssociativeSectionEdgeWeight; + +mod __private { + pub trait SealedTrait {} +} diff --git a/src/graph/link.rs b/src/graph/link.rs new file mode 100644 index 0000000..cb57237 --- /dev/null +++ b/src/graph/link.rs @@ -0,0 +1,926 @@ +use std::{ + cell::OnceCell, + collections::{BTreeMap, HashMap, LinkedList, hash_map}, + hash::{DefaultHasher, Hasher}, + path::Path, + sync::LazyLock, +}; + +use indexmap::{IndexMap, IndexSet}; +use log::warn; +use object::{ + Architecture, Object, ObjectSection, ObjectSymbol, SectionIndex, SymbolIndex, + coff::{CoffFile, CoffHeader, ImageSymbol}, +}; + +use crate::{ + linker::LinkerTargetArch, + linkobject::import::{ImportMember, ImportName}, +}; + +use super::{ + BuiltLinkGraph, SpecLinkGraph, + cache::LinkGraphCache, + edge::{ + AssociativeSectionEdgeWeight, ComdatSelection, DefinitionEdgeWeight, Edge, + ImportEdgeWeight, RelocationEdgeWeight, TryFromComdatSelectionError, + }, + node::{ + CoffNode, LibraryNode, LibraryNodeWeight, SectionNode, SectionNodeCharacteristics, + SectionNodeData, SymbolNode, SymbolNodeStorageClass, SymbolNodeType, TryFromSymbolError, + }, +}; + +pub type LinkGraphArena = bumpalo::Bump; + +pub(super) static ROOT_COFF: LazyLock = + LazyLock::new(|| CoffNode::new(Path::new(""), None)); + +#[derive(Debug, thiserror::Error)] +pub enum LinkGraphAddError { + #[error("invalid architecture '{found:?}', expected '{expected:?}'")] + ArchitectureMismatch { + expected: Architecture, + found: Architecture, + }, + + #[error("could not parse symbol '{name}' at table index {index}: {error}")] + Symbol { + name: String, + index: SymbolIndex, + error: TryFromSymbolError, + }, + + #[error( + "symbol '{symbol_name}' at table index {symbol_index} references invalid section number {section_num}" + )] + SymbolSectionIndex { + symbol_name: String, + symbol_index: SymbolIndex, + section_num: SectionIndex, + }, + + #[error( + "{section}+{address:#x} relocation references invalid target symbol index {symbol_index}" + )] + RelocationTarget { + section: String, + address: u32, + symbol_index: SymbolIndex, + }, + + #[error("could not parse symbol '{name}' at table index {index}: {error}")] + ComdatSymbol { + name: String, + index: SymbolIndex, + error: TryFromComdatSelectionError, + }, + + #[error("COMDAT symbol '{0}' is missing a section symbol")] + MissingComdatSectionSymbol(String), + + #[error("COMDAT section symbol '{symbol}' is missing associative section {associative_index}")] + MissingComdatAssociativeSection { + symbol: String, + associative_index: SectionIndex, + }, + + #[error("{0}")] + Object(#[from] object::read::Error), +} + +#[derive(Debug, thiserror::Error)] +pub enum SymbolError<'arena, 'data> { + #[error("{0}")] + Duplicate(DuplicateSymbolError<'arena, 'data>), + + #[error("{0}")] + Undefined(UndefinedSymbolError<'arena, 'data>), + + #[error("{0}")] + MultiplyDefined(MultiplyDefinedSymbolError<'arena, 'data>), +} + +#[derive(Debug, thiserror::Error)] +pub struct DuplicateSymbolError<'arena, 'data>(&'arena SymbolNode<'arena, 'data>); + +impl std::fmt::Display for DuplicateSymbolError<'_, '_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "duplicate symbol: {}", self.0.name().demangle())?; + + let mut definition_iter = self.0.definitions().iter(); + + for definition in definition_iter.by_ref().take(5) { + write!(f, "\n>>> defined at {}", definition.target().coff())?; + } + + let remaining = definition_iter.count(); + if remaining > 0 { + write!(f, "\n>>> defined {remaining} more times")?; + } + + Ok(()) + } +} + +#[derive(Debug, thiserror::Error)] +pub struct UndefinedSymbolError<'arena, 'data>(&'arena SymbolNode<'arena, 'data>); + +impl std::fmt::Display for UndefinedSymbolError<'_, '_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "undefined symbol: {}", self.0.name().demangle())?; + + let mut reference_iter = self.0.references().iter(); + + for reference in reference_iter.by_ref().take(5) { + let section = reference.source(); + let coff = section.coff(); + + let symbol_defs = + BTreeMap::from_iter(section.definitions().iter().filter_map(|definition| { + let ref_symbol = definition.source(); + if ref_symbol.is_section_symbol() || ref_symbol.is_label() { + None + } else { + Some((definition.weight().address(), ref_symbol.name())) + } + })); + + if let Some(reference_symbol) = symbol_defs + .range(0..=reference.weight().address()) + .next_back() + { + write!( + f, + "\n>>> referenced by {coff}:({})", + reference_symbol.1.demangle() + )?; + } else { + write!( + f, + "\n>>> referenced by {coff}:({}+{:#x})", + section.name(), + reference.weight().address() + )?; + } + } + + let remaining = reference_iter.count(); + if remaining > 0 { + write!(f, "\n>>> referenced {remaining} more times")?; + } + + Ok(()) + } +} + +#[derive(Debug, thiserror::Error)] +pub struct MultiplyDefinedSymbolError<'arena, 'data>(&'arena SymbolNode<'arena, 'data>); + +impl std::fmt::Display for MultiplyDefinedSymbolError<'_, '_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "multiply defined symbol: {}", self.0.name().demangle())?; + + let mut definition_iter = self.0.definitions().iter(); + + for definition in definition_iter.by_ref().take(5) { + write!(f, "\n>>> defined at {}", definition.target().coff())?; + } + + let remaining = definition_iter.count(); + if remaining > 0 { + write!(f, "\n>>> defined {remaining} more times")?; + } + + Ok(()) + } +} + +/// The link graph. +pub struct LinkGraph<'arena, 'data> { + /// Target architecture. + pub(super) machine: LinkerTargetArch, + + /// List of section nodes in the graph. + pub(super) section_nodes: Vec<&'arena SectionNode<'arena, 'data>>, + + /// The node for the COMMON section. + pub(super) common_section: OnceCell<&'arena SectionNode<'arena, 'data>>, + + /// Pseudo-COFF for holding metadata sections. + pub(super) root_coff: &'arena CoffNode<'data>, + + /// List of library nodes in the graph. + pub(super) library_nodes: IndexMap<&'data str, &'arena LibraryNode<'arena, 'data>>, + + /// List of COFF nodes in the graph. + pub(super) coff_nodes: IndexSet<&'arena CoffNode<'data>>, + + /// The API node if it exists. + pub(super) api_node: Option<&'arena LibraryNode<'arena, 'data>>, + + /// List of symbols with external storage class. + pub(super) external_symbols: IndexMap<&'data str, &'arena SymbolNode<'arena, 'data>>, + + /// Local symbols without any definition (absolute/debug symbols) + pub(super) extraneous_symbols: LinkedList<&'arena SymbolNode<'arena, 'data>>, + + /// Number of nodes in the graph. + pub(super) node_count: usize, + + /// COFF insertion cache. + pub(super) cache: LinkGraphCache<'arena, 'data>, + + /// Graph arena allocator. + pub(super) arena: &'arena LinkGraphArena, +} + +impl<'arena, 'data> LinkGraph<'arena, 'data> { + /// Constructs a new empty [`LinkGraph`] using the specified `arena` for + /// holding graph components. + pub fn new( + arena: &'arena LinkGraphArena, + machine: LinkerTargetArch, + ) -> LinkGraph<'arena, 'data> { + Self { + machine, + section_nodes: Vec::new(), + common_section: OnceCell::new(), + library_nodes: IndexMap::new(), + coff_nodes: IndexSet::new(), + root_coff: &*ROOT_COFF, + api_node: None, + external_symbols: IndexMap::new(), + extraneous_symbols: LinkedList::new(), + node_count: 0, + cache: LinkGraphCache::new(), + arena, + } + } + + /// Creates a new [`super::SpecLinkGraph`] for pre-calculating the memory + /// needed for building the graph. + pub fn spec() -> SpecLinkGraph { + Default::default() + } + + /// Returns the number of bytes allocated internally for holding the graph + /// components. + #[allow(unused)] + #[inline] + pub fn allocated_bytes(&self) -> usize { + self.arena.allocated_bytes() + } + + /// Adds a COFF to the graph. + pub fn add_coff( + &mut self, + file_path: &'data Path, + member_path: Option<&'data Path>, + coff: &CoffFile<'data, &'data [u8], C>, + ) -> Result<(), LinkGraphAddError> { + if Architecture::from(self.machine) != coff.architecture() { + return Err(LinkGraphAddError::ArchitectureMismatch { + expected: self.machine.into(), + found: coff.architecture(), + }); + } + + let coff_node = CoffNode::new(file_path, member_path); + + if self.coff_nodes.contains(&coff_node) { + return Ok(()); + } + + let coff_node = self.arena.alloc(coff_node); + self.node_count += 1; + self.coff_nodes.insert(coff_node); + + let symbol_table = coff.coff_symbol_table(); + + self.cache.clear(); + + self.cache.reserve_sections(coff.coff_section_table().len()); + self.cache.reserve_symbols(coff.coff_symbol_table().len()); + let mut comdat_count = 0; + + for section in coff.sections() { + let section_name = section.name()?; + let coff_section = section.coff_section(); + + let characteristics = SectionNodeCharacteristics::from_bits_truncate( + coff_section.characteristics.get(object::LittleEndian), + ); + + let section_data = + if characteristics.contains(SectionNodeCharacteristics::CntUninitializedData) { + SectionNodeData::Uninitialized( + coff_section.size_of_raw_data.get(object::LittleEndian), + ) + } else { + SectionNodeData::Initialized(section.data()?) + }; + + if characteristics.contains(SectionNodeCharacteristics::LnkComdat) { + comdat_count += 1; + } + + let section_node = self.arena.alloc_with(|| { + SectionNode::new(section_name, characteristics, section_data, 0, coff_node) + }); + + self.node_count += 1; + + self.cache.insert_section(section.index(), section_node); + self.section_nodes.push(section_node); + } + + self.cache.reserve_comdat_selections(comdat_count); + + for symbol in coff.symbols() { + let symbol_name = symbol.name()?; + let coff_symbol = symbol.coff_symbol(); + + let graph_symbol = + SymbolNode::try_from_symbol::(symbol_name, coff_symbol).map_err(|e| { + LinkGraphAddError::Symbol { + name: symbol_name.to_string(), + index: symbol.index(), + error: e, + } + })?; + + let graph_symbol = if symbol.is_global() { + *self + .external_symbols + .entry(symbol_name) + .and_modify(|existing| { + if symbol.is_definition() { + existing.set_type(coff_symbol.typ()); + } + }) + .or_insert_with(|| { + self.arena.alloc_with(|| { + self.node_count += 1; + graph_symbol + }) + }) + } else { + self.arena.alloc_with(|| { + self.node_count += 1; + graph_symbol + }) + }; + + self.cache.insert_symbol(symbol.index(), graph_symbol); + + let section_idx = match symbol.section_index() { + Some(idx) => idx, + None => { + if symbol.is_common() { + // Add a definition link for COMMON symbols to hold the + // symbol value + let common_section = *self.common_section.get_or_init(|| { + self.arena.alloc_with(|| { + SectionNode::new( + "COMMON data", + SectionNodeCharacteristics::CntUninitializedData + | SectionNodeCharacteristics::MemRead + | SectionNodeCharacteristics::MemWrite + | match self.machine { + LinkerTargetArch::Amd64 => { + SectionNodeCharacteristics::Align8Bytes + } + LinkerTargetArch::I386 => { + SectionNodeCharacteristics::Align4Bytes + } + }, + SectionNodeData::Uninitialized(0), + 0, + self.root_coff, + ) + }) + }); + + let definition_edge = self.arena.alloc_with(|| { + Edge::new( + graph_symbol, + common_section, + DefinitionEdgeWeight::new(coff_symbol.value(), None), + ) + }); + + graph_symbol.definitions().push_back(definition_edge); + common_section.definitions().push_back(definition_edge); + } else if symbol.is_local() { + self.extraneous_symbols.push_back(graph_symbol); + } + + continue; + } + }; + + let graph_section = self.cache.get_section(section_idx).ok_or_else(|| { + LinkGraphAddError::SymbolSectionIndex { + symbol_name: symbol_name.to_string(), + symbol_index: symbol.index(), + section_num: section_idx, + } + })?; + + let mut definition_edge = Edge::new( + graph_symbol, + graph_section, + DefinitionEdgeWeight::new(coff_symbol.value(), None), + ); + + if coff_symbol.has_aux_section() { + let aux_section = symbol_table.aux_section(symbol.index())?; + let mut checksum = aux_section.check_sum.get(object::LittleEndian); + + if graph_section.is_comdat() { + let selection = + ComdatSelection::try_from(aux_section.selection).map_err(|e| { + LinkGraphAddError::ComdatSymbol { + name: symbol_name.to_string(), + index: symbol.index(), + error: e, + } + })?; + + // If this is a COMDAT to an associative section, add an + // associative edge from the section specified in the + // selection number to this section + if selection == ComdatSelection::Associative { + let aux_section_number = aux_section.number.get(object::LittleEndian); + + let associative_section_index = SectionIndex(aux_section_number as usize); + let associative_section = self + .cache + .get_section(associative_section_index) + .ok_or_else(|| LinkGraphAddError::MissingComdatAssociativeSection { + symbol: symbol_name.to_string(), + associative_index: associative_section_index, + })?; + + associative_section + .associative_edges() + .push_back(self.arena.alloc_with(|| { + Edge::new( + associative_section, + graph_section, + AssociativeSectionEdgeWeight, + ) + })); + } + + // Store the selection value for this section. + self.cache.insert_comdat_selection(section_idx, selection); + } + + // Calculate the checksum value for the .rdata$zzz section since + // it is used to dedup those sections in the output + if graph_section.name().as_str() == ".rdata$zzz" && checksum == 0 { + // Calculate and set the auxiliary checksum value if needed + if let SectionNodeData::Initialized(data) = graph_section.data() { + let mut h = jamcrc::Hasher::new_with_initial(!0); + h.update(data); + checksum = !h.finalize(); + } + } + + graph_section.replace_checksum(checksum); + } else if graph_section.is_comdat() { + // Add the selection to the COMDAT symbol definition. + let selection = self + .cache + .get_comdat_selection(section_idx) + .ok_or_else(|| { + LinkGraphAddError::MissingComdatSectionSymbol(symbol_name.to_string()) + })?; + + definition_edge.weight_mut().selection = Some(selection); + } + + let definition_edge = self.arena.alloc(definition_edge); + + // Link the definition to the symbol + graph_symbol.definitions().push_back(definition_edge); + + // Link the definition to the section + graph_section.definitions().push_back(definition_edge); + } + + for section in coff.sections() { + let graph_section = self + .cache + .get_section(section.index()) + .unwrap_or_else(|| unreachable!()); + + for reloc in section.coff_relocations()? { + let target_symbol = self.cache.get_symbol(reloc.symbol()).ok_or_else(|| { + LinkGraphAddError::RelocationTarget { + section: graph_section.name().to_string(), + address: reloc.virtual_address.get(object::LittleEndian), + symbol_index: reloc.symbol(), + } + })?; + + let reloc_edge = self.arena.alloc_with(|| { + Edge::new( + graph_section, + target_symbol, + RelocationEdgeWeight::new( + reloc.virtual_address.get(object::LittleEndian), + reloc.typ.get(object::LittleEndian), + ), + ) + }); + + graph_section.relocations().push_back(reloc_edge); + target_symbol.references().push_back(reloc_edge); + } + } + + Ok(()) + } + + /// Adds an external symbol to the graph if it does not exist. + /// + /// The newly added symbol will be undefined. + pub fn add_external_symbol(&mut self, name: &'data str) { + self.external_symbols.entry(name).or_insert_with(|| { + self.arena.alloc_with(|| { + SymbolNode::new( + name, + SymbolNodeStorageClass::External, + false, + SymbolNodeType::Value(0), + ) + }) + }); + } + + /// Returns an iterator over the names of the undefined symbols + pub fn undefined_symbols(&self) -> impl Iterator + use<'_, 'data, 'arena> { + self.external_symbols + .iter() + .filter_map(|(name, symbol)| symbol.is_undefined().then_some(*name)) + } + + /// Associates `symbol` as an API imported symbol with metadata from the + /// specified [`ImportMember`]. + /// + /// # Panics + /// Panics if `symbol` does not exist. + #[inline] + pub fn add_api_import( + &mut self, + symbol: &str, + import: &ImportMember<'data>, + ) -> Result<(), LinkGraphAddError> { + let api_node = *self.api_node.get_or_insert_with(|| { + self.arena + .alloc_with(|| LibraryNode::new(LibraryNodeWeight::new(import.dll))) + }); + + self.add_import_edge(symbol, api_node, import) + } + + /// Associates `symbol` with the specified [`ImportMember`]. + /// + /// # Panics + /// Panics if `symbol` does not exist. + #[inline] + pub fn add_library_import( + &mut self, + symbol: &str, + import: &ImportMember<'data>, + ) -> Result<(), LinkGraphAddError> { + let library_node = *self.library_nodes.entry(import.dll).or_insert_with(|| { + self.arena + .alloc_with(|| LibraryNode::new(LibraryNodeWeight::new(import.dll))) + }); + + self.add_import_edge(symbol, library_node, import) + } + + fn add_import_edge( + &mut self, + symbol: &str, + library: &'arena LibraryNode<'arena, 'data>, + import: &ImportMember<'data>, + ) -> Result<(), LinkGraphAddError> { + let symbol_node = self + .external_symbols + .get(symbol) + .copied() + .unwrap_or_else(|| panic!("symbol {symbol} does not exist")); + + if import.architecture != self.machine.into() { + return Err(LinkGraphAddError::ArchitectureMismatch { + expected: self.machine.into(), + found: import.architecture, + }); + } + + let import_name = match import.import { + ImportName::Name(name) => name, + ImportName::Ordinal(o) => { + warn!( + "found ordinal import value '{o}' for symbol \"{symbol}\". Linking public symbol name." + ); + import.symbol + } + }; + + let import_edge = self + .arena + .alloc_with(|| Edge::new(symbol_node, library, ImportEdgeWeight::new(import_name))); + + symbol_node.imports().push_back(import_edge); + library.imports().push_back(import_edge); + + Ok(()) + } + + /// Finishes building the link graph. + pub fn finish(self) -> Result, Vec>> { + let mut symbol_errors = Vec::new(); + + for symbol in self.external_symbols.values().copied() { + if symbol.is_undefined() { + symbol_errors.push(SymbolError::Undefined(UndefinedSymbolError(symbol))); + } else if symbol.is_duplicate() { + symbol_errors.push(SymbolError::Duplicate(DuplicateSymbolError(symbol))); + } else if symbol.is_multiply_defined() { + symbol_errors.push(SymbolError::MultiplyDefined(MultiplyDefinedSymbolError( + symbol, + ))); + } + } + + if !symbol_errors.is_empty() { + return Err(symbol_errors); + } + + Ok(BuiltLinkGraph::new(self)) + } + + /// Writes out the GraphViz dot representation of this graph to the specified + /// [`std::io::Write`]er. + pub fn write_dot_graph(&self, mut w: impl std::io::Write) -> std::io::Result<()> { + writeln!(w, "digraph {{")?; + + let pad = " "; + + let mut idx_val = 0; + let mut node_ids: HashMap = HashMap::with_capacity(self.node_count); + + let mut section_flags = String::new(); + + // Write out the section nodes and the neighboring symbol nodes. + for section in self + .section_nodes + .iter() + .copied() + .chain(self.common_section.get().into_iter().copied()) + { + let section_idx = idx_val; + + let mut h = DefaultHasher::new(); + std::ptr::hash(section, &mut h); + let hid = h.finish(); + node_ids.insert(hid, section_idx); + + idx_val += 1; + + section_flags.clear(); + bitflags::parser::to_writer( + §ion.characteristics().zero_align(), + &mut section_flags, + ) + .unwrap(); + + writeln!( + w, + "{pad}{section_idx} [ label=\"{{ {} | {} | {{ Size: {:#x}\\l | Align: {:#x}\\l | Checksum: {:#x}\\l }} | {{ {} }} }}\" shape=record ]", + section.name(), + section.coff().short_name(), + section.data().len(), + section.characteristics().alignment().unwrap_or(0), + section.checksum(), + section_flags, + )?; + + for reloc in section.relocations().iter() { + let target_symbol = reloc.target(); + let mut h = DefaultHasher::new(); + std::ptr::hash(target_symbol, &mut h); + let hid = h.finish(); + + if let hash_map::Entry::Vacant(e) = node_ids.entry(hid) { + let symbol_idx = idx_val; + idx_val += 1; + + write!(w, "{pad}{symbol_idx} [ label=\"{}\"", target_symbol.name())?; + + if target_symbol.is_undefined() || target_symbol.is_duplicate() { + write!(w, " color=red")?; + } + + writeln!(w, " ]")?; + + e.insert(symbol_idx); + } + } + + for definition in section.definitions().iter() { + let target_symbol = definition.source(); + let mut h = DefaultHasher::new(); + std::ptr::hash(target_symbol, &mut h); + let hid = h.finish(); + if let hash_map::Entry::Vacant(e) = node_ids.entry(hid) { + let symbol_idx = idx_val; + idx_val += 1; + + write!(w, "{pad}{symbol_idx} [ label=\"{}\"", target_symbol.name())?; + + if target_symbol.is_undefined() + || target_symbol.is_duplicate() + || target_symbol.is_multiply_defined() + { + write!(w, " color=red")?; + } + + writeln!(w, " ]")?; + + e.insert(symbol_idx); + } + } + } + + // Write out unreferenced extraneous symbols. + for symbol in self.extraneous_symbols.iter().copied() { + if !symbol.references().is_empty() { + continue; + } + + let symbol_idx = idx_val; + + let mut h = DefaultHasher::new(); + std::ptr::hash(symbol, &mut h); + let hid = h.finish(); + node_ids.insert(hid, symbol_idx); + + writeln!(w, "{pad}{symbol_idx} [ label=\"{}\" ]", symbol.name())?; + + idx_val += 1; + } + + // Write out the library nodes. + for library in self.library_nodes.values().copied() { + let mut h = DefaultHasher::new(); + std::ptr::hash(library, &mut h); + + if let hash_map::Entry::Vacant(e) = node_ids.entry(h.finish()) { + let library_idx = idx_val; + idx_val += 1; + + writeln!( + w, + "{pad}{library_idx} [ label=\"{}\" shape=diamond ]", + library.name(), + )?; + + e.insert(library_idx); + } + } + + let mut api_idx = None; + + // Write out the API node if it exists. + if let Some(api_node) = self.api_node { + let api_idx_val = idx_val; + + writeln!( + w, + "{pad}{api_idx_val} [ label=\"{}\" shape=triangle ]", + api_node.name().trim_dll_suffix() + )?; + + api_idx = Some(api_idx_val); + } + + // Write out relocations, definitions and COMDAT associations. + for section in self + .section_nodes + .iter() + .copied() + .chain(self.common_section.get().into_iter().copied()) + { + let mut h = DefaultHasher::new(); + std::ptr::hash(section, &mut h); + let hid = h.finish(); + let section_idx = node_ids.get(&hid).copied().unwrap(); + + // Relocations + for reloc in section.relocations().iter() { + let target_symbol = reloc.target(); + let mut h = DefaultHasher::new(); + std::ptr::hash(target_symbol, &mut h); + let hid = h.finish(); + let symbol_idx = node_ids.get(&hid).copied().unwrap(); + + writeln!( + w, + "{pad}{section_idx} -> {symbol_idx} [ label=\"relocation (addr {:#x})\" ]", + reloc.weight().address(), + )?; + } + + // Definitions + for definition in section.definitions().iter() { + let target_symbol = definition.source(); + let mut h = DefaultHasher::new(); + std::ptr::hash(target_symbol, &mut h); + let hid = h.finish(); + let symbol_idx = node_ids.get(&hid).copied().unwrap(); + + write!( + w, + "{pad}{symbol_idx} -> {section_idx} [ label=\"defined at {:#x}", + definition.weight().address() + )?; + + if let Some(selection) = definition.weight().selection() { + write!(w, " ({selection:?})")?; + } + + write!(w, "\"")?; + + if target_symbol.is_duplicate() || target_symbol.is_multiply_defined() { + write!(w, " color=red")?; + } + + writeln!(w, " ]")?; + } + + // COMDAT associations + for assocation in section.associative_edges().iter() { + let target_section = assocation.target(); + let mut h = DefaultHasher::new(); + std::ptr::hash(target_section, &mut h); + let hid = h.finish(); + let target_idx = node_ids.get(&hid).copied().unwrap(); + + writeln!( + w, + "{pad}{section_idx} -> {target_idx} [ label=\"associative\" ]" + )?; + } + } + + // Write out API import edges. + if let Some(api_node) = self.api_node { + let api_node_idx = api_idx.unwrap(); + + for import in api_node.imports().iter() { + let target_symbol = import.source(); + let mut h = DefaultHasher::new(); + std::ptr::hash(target_symbol, &mut h); + let symbol_idx = node_ids.get(&h.finish()).copied().unwrap(); + + writeln!( + w, + "{pad}{symbol_idx} -> {api_node_idx} [ label=\"import \\\"{}\\\"\" ]", + import.weight().import_name() + )?; + } + } + + // Write out import edges. + for library in self.library_nodes.values().copied() { + let mut h = DefaultHasher::new(); + std::ptr::hash(library, &mut h); + let library_idx = node_ids.get(&h.finish()).copied().unwrap(); + + for import in library.imports().iter() { + let target_symbol = import.source(); + let mut h = DefaultHasher::new(); + std::ptr::hash(target_symbol, &mut h); + let symbol_idx = node_ids.get(&h.finish()).copied().unwrap(); + + writeln!( + w, + "{pad}{symbol_idx} -> {library_idx} [ label=\"import \\\"{}\\\"\" ]", + import.weight().import_name(), + )?; + } + } + + writeln!(w, "}}")?; + + Ok(()) + } +} diff --git a/src/graph/mod.rs b/src/graph/mod.rs new file mode 100644 index 0000000..31306fd --- /dev/null +++ b/src/graph/mod.rs @@ -0,0 +1,10 @@ +mod built; +mod cache; +pub mod edge; +mod link; +pub mod node; +mod spec; + +pub use built::*; +pub use link::*; +pub use spec::*; diff --git a/src/graph/node/coff.rs b/src/graph/node/coff.rs new file mode 100644 index 0000000..5ec4d6d --- /dev/null +++ b/src/graph/node/coff.rs @@ -0,0 +1,61 @@ +use std::path::Path; + +/// A COFF node. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct CoffNode<'data> { + /// The path on disk. + file_path: &'data Path, + + /// The member path. + member_path: Option<&'data Path>, +} + +impl<'data> CoffNode<'data> { + #[inline] + pub const fn new(file_path: &'data Path, member_path: Option<&'data Path>) -> CoffNode<'data> { + Self { + file_path, + member_path, + } + } + + /// Returns a [`CoffNodeShortName`] for displaying a shortened version of + /// the COFF name. + #[inline] + pub fn short_name(&self) -> CoffNodeShortName<'_, 'data> { + CoffNodeShortName(self) + } +} + +impl std::fmt::Display for CoffNode<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some(member_path) = self.member_path { + write!(f, "{}({})", self.file_path.display(), member_path.display()) + } else { + write!(f, "{}", self.file_path.display()) + } + } +} + +/// Used for writing out a shortened version of a [`CoffNode`]. +#[derive(Debug)] +pub struct CoffNodeShortName<'b, 'data>(&'b CoffNode<'data>); + +impl std::fmt::Display for CoffNodeShortName<'_, '_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some(member_path) = self.0.member_path { + write!( + f, + "{}({})", + self.0.file_path.file_name().unwrap().to_string_lossy(), + member_path.file_name().unwrap().to_string_lossy(), + ) + } else { + write!( + f, + "{}", + self.0.file_path.file_name().unwrap().to_string_lossy() + ) + } + } +} diff --git a/src/graph/node/library.rs b/src/graph/node/library.rs new file mode 100644 index 0000000..3a189c9 --- /dev/null +++ b/src/graph/node/library.rs @@ -0,0 +1,74 @@ +use crate::graph::edge::{EdgeList, ImportEdgeWeight, IncomingEdges}; + +use super::SymbolNode; + +/// A library node in the graph. +pub struct LibraryNode<'arena, 'data> { + /// The list of incoming import edges for this library node. + import_edges: + EdgeList<'arena, SymbolNode<'arena, 'data>, Self, ImportEdgeWeight<'data>, IncomingEdges>, + + /// The node weight. + weight: LibraryNodeWeight<'data>, +} + +impl<'arena, 'data> LibraryNode<'arena, 'data> { + #[inline] + pub fn new(weight: LibraryNodeWeight<'data>) -> LibraryNode<'arena, 'data> { + Self { + import_edges: EdgeList::new(), + weight, + } + } + + #[inline] + pub fn imports( + &self, + ) -> &EdgeList<'arena, SymbolNode<'arena, 'data>, Self, ImportEdgeWeight<'data>, IncomingEdges> + { + &self.import_edges + } + + #[inline] + pub fn name(&self) -> LibraryName<'data> { + self.weight.name + } +} + +/// The weight for a library node. +pub struct LibraryNodeWeight<'data> { + /// The name of the library. + name: LibraryName<'data>, +} + +impl<'data> LibraryNodeWeight<'data> { + #[inline] + pub fn new(name: impl Into>) -> LibraryNodeWeight<'data> { + Self { name: name.into() } + } +} + +/// A library name. +#[derive(Debug, Clone, Copy)] +pub struct LibraryName<'data>(&'data str); + +impl LibraryName<'_> { + pub fn trim_dll_suffix(&self) -> &str { + self.0 + .rsplit_once('.') + .and_then(|(prefix, suffix)| suffix.eq_ignore_ascii_case("dll").then_some(prefix)) + .unwrap_or(self.0) + } +} + +impl<'data> From<&'data str> for LibraryName<'data> { + fn from(value: &'data str) -> Self { + Self(value) + } +} + +impl std::fmt::Display for LibraryName<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} diff --git a/src/graph/node/mod.rs b/src/graph/node/mod.rs new file mode 100644 index 0000000..f9c7454 --- /dev/null +++ b/src/graph/node/mod.rs @@ -0,0 +1,9 @@ +mod coff; +mod library; +mod section; +mod symbol; + +pub use coff::*; +pub use library::*; +pub use section::*; +pub use symbol::*; diff --git a/src/graph/node/section.rs b/src/graph/node/section.rs new file mode 100644 index 0000000..d432f5c --- /dev/null +++ b/src/graph/node/section.rs @@ -0,0 +1,378 @@ +use std::{ + cell::Cell, + collections::{HashSet, VecDeque}, + hash::{DefaultHasher, Hasher}, +}; + +use object::pe::{ + IMAGE_SCN_ALIGN_1BYTES, IMAGE_SCN_ALIGN_2BYTES, IMAGE_SCN_ALIGN_4BYTES, IMAGE_SCN_ALIGN_8BYTES, + IMAGE_SCN_ALIGN_16BYTES, IMAGE_SCN_ALIGN_32BYTES, IMAGE_SCN_ALIGN_64BYTES, + IMAGE_SCN_ALIGN_128BYTES, IMAGE_SCN_ALIGN_256BYTES, IMAGE_SCN_ALIGN_512BYTES, + IMAGE_SCN_ALIGN_1024BYTES, IMAGE_SCN_ALIGN_2048BYTES, IMAGE_SCN_ALIGN_4096BYTES, + IMAGE_SCN_ALIGN_8192BYTES, IMAGE_SCN_CNT_CODE, IMAGE_SCN_CNT_INITIALIZED_DATA, + IMAGE_SCN_CNT_UNINITIALIZED_DATA, IMAGE_SCN_GPREL, IMAGE_SCN_LNK_COMDAT, IMAGE_SCN_LNK_INFO, + IMAGE_SCN_LNK_NRELOC_OVFL, IMAGE_SCN_LNK_OTHER, IMAGE_SCN_LNK_REMOVE, + IMAGE_SCN_MEM_DISCARDABLE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_LOCKED, + IMAGE_SCN_MEM_NOT_CACHED, IMAGE_SCN_MEM_NOT_PAGED, IMAGE_SCN_MEM_PRELOAD, + IMAGE_SCN_MEM_PURGEABLE, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_SHARED, IMAGE_SCN_MEM_WRITE, + IMAGE_SCN_TYPE_NO_PAD, +}; + +use crate::graph::edge::{ + AssociativeSectionEdgeWeight, DefinitionEdgeWeight, EdgeList, IncomingEdges, OutgoingEdges, + RelocationEdgeWeight, +}; + +use super::{CoffNode, SymbolNode}; + +/// A section node in the graph. +pub struct SectionNode<'arena, 'data> { + /// The list of outgoing relocation edges for this section. + relocation_edges: + EdgeList<'arena, Self, SymbolNode<'arena, 'data>, RelocationEdgeWeight, OutgoingEdges>, + + /// The list of incoming definition edges for this section. + definition_edges: + EdgeList<'arena, SymbolNode<'arena, 'data>, Self, DefinitionEdgeWeight, IncomingEdges>, + + /// The list of outgoing COMDAT associative edges for this section. + associative_edges: EdgeList< + 'arena, + Self, + SectionNode<'arena, 'data>, + AssociativeSectionEdgeWeight, + OutgoingEdges, + >, + + /// The COFF this section is from. + coff: &'arena CoffNode<'data>, + + /// The rebased virtual address of the section. + virtual_address: Cell, + + /// If this section is to be discarded. + discarded: Cell, + + /// The name of the section. + name: SectionName<'data>, + + /// The characteristics of the section. + characteristics: SectionNodeCharacteristics, + + /// The section data. + data: Cell>, + + /// The data checksum + checksum: Cell, +} + +impl<'arena, 'data> SectionNode<'arena, 'data> { + #[inline] + pub fn new( + name: impl Into>, + characteristics: SectionNodeCharacteristics, + data: SectionNodeData<'arena>, + checksum: u32, + coff: &'arena CoffNode<'data>, + ) -> SectionNode<'arena, 'data> { + Self { + relocation_edges: EdgeList::new(), + definition_edges: EdgeList::new(), + associative_edges: EdgeList::new(), + virtual_address: Cell::new(0), + discarded: Cell::new(false), + coff, + data: Cell::new(data), + characteristics, + checksum: Cell::from(checksum), + name: name.into(), + } + } + + /// Returns the list of outgoing relocation edges for this section. + #[inline] + pub fn relocations( + &self, + ) -> &EdgeList<'arena, Self, SymbolNode<'arena, 'data>, RelocationEdgeWeight, OutgoingEdges> + { + &self.relocation_edges + } + + /// Returns the list of incoming relocation edges for this section. + #[inline] + pub fn definitions( + &self, + ) -> &EdgeList<'arena, SymbolNode<'arena, 'data>, Self, DefinitionEdgeWeight, IncomingEdges> + { + &self.definition_edges + } + + /// Returns the list of output associative section edges for this section. + /// If this section is linked, the adjacent sections must also be linked. + #[inline] + pub fn associative_edges( + &self, + ) -> &EdgeList< + 'arena, + Self, + SectionNode<'arena, 'data>, + AssociativeSectionEdgeWeight, + OutgoingEdges, + > { + &self.associative_edges + } + + /// Perform a BFS traversal over the associative section edges starting + /// from this section. + pub fn associative_bfs(&'arena self) -> AssociativeBfs<'arena, 'data> { + let queue = VecDeque::from([self]); + let mut h = DefaultHasher::new(); + std::ptr::hash(self, &mut h); + let visited = HashSet::from([h.finish()]); + AssociativeBfs { queue, visited } + } + + /// Returns the COFF associated with this section. + /// + /// This is the COFF where the section node was sourced from. + #[inline] + pub fn coff(&self) -> &'arena CoffNode<'data> { + self.coff + } + + /// Marks this section as being discarded. + #[inline] + pub fn discard(&self) { + self.discarded.set(true); + } + + /// Sets the discarded value for the section. + #[inline] + pub fn set_discarded(&self, val: bool) { + self.discarded.set(val); + } + + /// Keeps this section if it was previously discarded. + #[inline] + #[allow(unused)] + pub(super) fn keep(&self) { + self.discarded.set(false); + } + + /// Returns `true` if this section was discarded. + #[inline] + pub fn is_discarded(&self) -> bool { + self.discarded.get() + } + + /// Returns `true` if this is a debug section. + #[inline] + pub fn is_debug(&self) -> bool { + self.name().group_name() == ".debug" + && self + .name() + .group_ordering() + .is_some_and(|val| val == "S" || val == "T" || val == "P" || val == "F") + } + + /// Returns `true` if this is a COMDAT section. + #[inline] + pub fn is_comdat(&self) -> bool { + self.characteristics() + .contains(SectionNodeCharacteristics::LnkComdat) + } + + /// Returns the name of the section. + #[inline] + pub fn name(&self) -> SectionName<'data> { + self.name + } + + /// Returns the characteristics flags associated with this section. + #[inline] + pub fn characteristics(&self) -> SectionNodeCharacteristics { + self.characteristics + } + + /// Returns the data associated with this section. + #[inline] + pub fn data(&self) -> SectionNodeData<'arena> { + self.data.get() + } + + /// Sets the size value if this section contains uninitialized data. + #[inline] + pub fn set_uninitialized_size(&self, val: u32) { + if matches!(self.data(), SectionNodeData::Uninitialized(_)) { + self.data.set(SectionNodeData::Uninitialized(val)); + } + } + + /// Returns the checksum value for the section data. + #[inline] + pub fn checksum(&self) -> u32 { + self.checksum.get() + } + + /// Replaces the checksum value for the section data. + #[inline] + pub fn replace_checksum(&self, val: u32) { + self.checksum.set(val); + } + + /// Returns the assigned virtual address of the section. + #[inline] + pub fn virtual_address(&self) -> u32 { + self.virtual_address.get() + } + + /// Assigns a virtual address for the section. + #[inline] + pub fn assign_virtual_address(&self, val: u32) { + self.virtual_address.set(val); + } +} + +/// A section name. +#[derive(Debug, Clone, Copy)] +pub struct SectionName<'data>(&'data str); + +impl<'data> SectionName<'data> { + #[inline] + pub fn as_str(&self) -> &'data str { + self.0 + } + + /// Returns the `group name` value (`$`) from + /// the section name. + #[inline] + pub fn group_name(&self) -> &'data str { + self.0 + .split_once('$') + .map(|(group_name, _)| group_name) + .unwrap_or(self.0) + } + + /// Returns the `group ordering` value (`$`) + /// from the section name if this is a grouped section. + #[inline] + pub fn group_ordering(&self) -> Option<&str> { + self.0 + .split_once('$') + .map(|(_, group_ordering)| group_ordering) + } +} + +impl<'data> From<&'data str> for SectionName<'data> { + fn from(value: &'data str) -> Self { + Self(value) + } +} + +impl std::fmt::Display for SectionName<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} + +/// Section node characteristic flags +#[derive(Debug, Copy, Clone)] +pub struct SectionNodeCharacteristics(u32); + +bitflags::bitflags! { + impl SectionNodeCharacteristics: u32 { + const TypeNoPad = IMAGE_SCN_TYPE_NO_PAD; + const CntCode = IMAGE_SCN_CNT_CODE; + const CntInitializedData = IMAGE_SCN_CNT_INITIALIZED_DATA; + const CntUninitializedData = IMAGE_SCN_CNT_UNINITIALIZED_DATA; + const LnkOther = IMAGE_SCN_LNK_OTHER; + const LnkInfo = IMAGE_SCN_LNK_INFO; + const LnkRemove = IMAGE_SCN_LNK_REMOVE; + const LnkComdat = IMAGE_SCN_LNK_COMDAT; + const GPRel = IMAGE_SCN_GPREL; + const MemPurgeable = IMAGE_SCN_MEM_PURGEABLE; + const MemLocked = IMAGE_SCN_MEM_LOCKED; + const MemPreload = IMAGE_SCN_MEM_PRELOAD; + const Align1Bytes = IMAGE_SCN_ALIGN_1BYTES; + const Align2Bytes = IMAGE_SCN_ALIGN_2BYTES; + const Align4Bytes = IMAGE_SCN_ALIGN_4BYTES; + const Align8Bytes = IMAGE_SCN_ALIGN_8BYTES; + const Align16Bytes = IMAGE_SCN_ALIGN_16BYTES; + const Align32Bytes = IMAGE_SCN_ALIGN_32BYTES; + const Align64Bytes = IMAGE_SCN_ALIGN_64BYTES; + const Align128Bytes = IMAGE_SCN_ALIGN_128BYTES; + const Align256Bytes = IMAGE_SCN_ALIGN_256BYTES; + const Align512Bytes = IMAGE_SCN_ALIGN_512BYTES; + const Align1024Bytes = IMAGE_SCN_ALIGN_1024BYTES; + const Align2048Bytes = IMAGE_SCN_ALIGN_2048BYTES; + const Align4096Bytes = IMAGE_SCN_ALIGN_4096BYTES; + const Align8192Bytes = IMAGE_SCN_ALIGN_8192BYTES; + const LnkNRelocOvfl = IMAGE_SCN_LNK_NRELOC_OVFL; + const MemDiscardable = IMAGE_SCN_MEM_DISCARDABLE; + const MemNotCached = IMAGE_SCN_MEM_NOT_CACHED; + const MemNotPaged = IMAGE_SCN_MEM_NOT_PAGED; + const MemShared = IMAGE_SCN_MEM_SHARED; + const MemExecute = IMAGE_SCN_MEM_EXECUTE; + const MemRead = IMAGE_SCN_MEM_READ; + const MemWrite = IMAGE_SCN_MEM_WRITE; + const _ = !0; + } +} + +impl SectionNodeCharacteristics { + /// Returns the alignment value if it exists + pub fn alignment(&self) -> Option { + (self.0 & (0xfu32 << 20) != 0).then(|| 2usize.pow(((self.0 >> 20) & 0xf) - 1)) + } + + /// Returns a new [`SectionNodeCharacteristics`] without the alignment + /// bits set + pub fn zero_align(&self) -> SectionNodeCharacteristics { + Self(self.0 & !(0xfu32 << 20)) + } +} + +/// The section data. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum SectionNodeData<'arena> { + Initialized(&'arena [u8]), + Uninitialized(u32), +} + +impl SectionNodeData<'_> { + pub fn len(&self) -> usize { + match self { + Self::Initialized(data) => data.len(), + Self::Uninitialized(size) => *size as usize, + } + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +/// BFS traversal over sections with associative edges +pub struct AssociativeBfs<'arena, 'data> { + queue: VecDeque<&'arena SectionNode<'arena, 'data>>, + visited: HashSet, +} + +impl<'arena, 'data> Iterator for AssociativeBfs<'arena, 'data> { + type Item = &'arena SectionNode<'arena, 'data>; + + fn next(&mut self) -> Option { + let next_section = self.queue.pop_front()?; + + for edge in next_section.associative_edges() { + let target = edge.target(); + let mut h = DefaultHasher::new(); + std::ptr::hash(target, &mut h); + if self.visited.insert(h.finish()) { + self.queue.push_back(target); + } + } + + Some(next_section) + } +} diff --git a/src/graph/node/symbol.rs b/src/graph/node/symbol.rs new file mode 100644 index 0000000..67e587a --- /dev/null +++ b/src/graph/node/symbol.rs @@ -0,0 +1,399 @@ +use std::{ + cell::{Cell, OnceCell}, + collections::HashSet, +}; + +use num_enum::{IntoPrimitive, TryFromPrimitive}; +use object::{ + coff::{CoffHeader, ImageSymbol}, + pe::{ + IMAGE_SYM_ABSOLUTE, IMAGE_SYM_CLASS_ARGUMENT, IMAGE_SYM_CLASS_AUTOMATIC, + IMAGE_SYM_CLASS_BIT_FIELD, IMAGE_SYM_CLASS_BLOCK, IMAGE_SYM_CLASS_CLR_TOKEN, + IMAGE_SYM_CLASS_END_OF_FUNCTION, IMAGE_SYM_CLASS_END_OF_STRUCT, IMAGE_SYM_CLASS_ENUM_TAG, + IMAGE_SYM_CLASS_EXTERNAL, IMAGE_SYM_CLASS_EXTERNAL_DEF, IMAGE_SYM_CLASS_FILE, + IMAGE_SYM_CLASS_FUNCTION, IMAGE_SYM_CLASS_LABEL, IMAGE_SYM_CLASS_MEMBER_OF_ENUM, + IMAGE_SYM_CLASS_MEMBER_OF_STRUCT, IMAGE_SYM_CLASS_MEMBER_OF_UNION, IMAGE_SYM_CLASS_NULL, + IMAGE_SYM_CLASS_REGISTER, IMAGE_SYM_CLASS_REGISTER_PARAM, IMAGE_SYM_CLASS_SECTION, + IMAGE_SYM_CLASS_STATIC, IMAGE_SYM_CLASS_STRUCT_TAG, IMAGE_SYM_CLASS_TYPE_DEFINITION, + IMAGE_SYM_CLASS_UNDEFINED_LABEL, IMAGE_SYM_CLASS_UNDEFINED_STATIC, + IMAGE_SYM_CLASS_UNION_TAG, IMAGE_SYM_CLASS_WEAK_EXTERNAL, IMAGE_SYM_DEBUG, + }, +}; + +use crate::graph::edge::{ + ComdatSelection, DefinitionEdgeWeight, EdgeList, ImportEdgeWeight, IncomingEdges, + OutgoingEdges, RelocationEdgeWeight, +}; + +use super::{LibraryNode, SectionNode}; + +#[derive(Debug, Copy, Clone, thiserror::Error)] +#[error("unknown storage class value ({0})")] +pub struct TryFromStorageClassError(u8); + +#[derive(Debug, thiserror::Error)] +pub enum TryFromSymbolError { + #[error("{0}")] + StorageClass(#[from] TryFromStorageClassError), +} + +/// A symbol node in the graph. +pub struct SymbolNode<'arena, 'data> { + /// The list of outgoing definition edges for this symbol. + definition_edges: + EdgeList<'arena, Self, SectionNode<'arena, 'data>, DefinitionEdgeWeight, OutgoingEdges>, + + /// The list of outgoing import edges for this symbol. + import_edges: + EdgeList<'arena, Self, LibraryNode<'arena, 'data>, ImportEdgeWeight<'data>, OutgoingEdges>, + + /// The incoming relocation edges for this symbol. + relocation_edges: + EdgeList<'arena, SectionNode<'arena, 'data>, Self, RelocationEdgeWeight, IncomingEdges>, + + /// The symbol table index when inserted into the output COFF. + table_index: OnceCell, + + /// The symbol name for the output COFF. + output_name: OnceCell, + + /// Cached flag for checking if this is an MSVC label symbol. + msvc_label: OnceCell, + + /// The name of the symbol. + name: SymbolName<'arena>, + + /// The storage class of the symbol. + storage_class: SymbolNodeStorageClass, + + /// If this is a section symbol. + section: bool, + + /// The type of symbol. + typ: Cell, +} + +impl<'arena, 'data> SymbolNode<'arena, 'data> { + #[inline] + pub fn new( + name: impl Into>, + storage_class: SymbolNodeStorageClass, + section: bool, + typ: SymbolNodeType, + ) -> SymbolNode<'arena, 'data> { + Self { + definition_edges: EdgeList::new(), + import_edges: EdgeList::new(), + relocation_edges: EdgeList::new(), + table_index: OnceCell::new(), + output_name: OnceCell::new(), + msvc_label: OnceCell::new(), + name: name.into(), + storage_class, + section, + typ: Cell::new(typ), + } + } + + pub fn try_from_symbol<'file, C: CoffHeader>( + name: impl Into>, + coff_symbol: &'arena C::ImageSymbol, + ) -> Result, TryFromSymbolError> { + Ok(Self { + definition_edges: EdgeList::new(), + import_edges: EdgeList::new(), + relocation_edges: EdgeList::new(), + table_index: OnceCell::new(), + output_name: OnceCell::new(), + msvc_label: OnceCell::new(), + name: name.into(), + storage_class: coff_symbol.storage_class().try_into()?, + section: coff_symbol.has_aux_section(), + typ: Cell::new(match coff_symbol.section_number() { + IMAGE_SYM_ABSOLUTE => SymbolNodeType::Absolute(coff_symbol.value()), + IMAGE_SYM_DEBUG => SymbolNodeType::Debug, + _ => SymbolNodeType::Value(coff_symbol.typ()), + }), + }) + } + + /// Returns the list of adjacent outgoing definition edges for this symbol + /// node. + #[inline] + pub fn definitions( + &self, + ) -> &EdgeList<'arena, Self, SectionNode<'arena, 'data>, DefinitionEdgeWeight, OutgoingEdges> + { + &self.definition_edges + } + + /// Returns the list of adjacent incoming relocation edges for this symbol + /// node. + #[inline] + pub fn references( + &self, + ) -> &EdgeList<'arena, SectionNode<'arena, 'data>, Self, RelocationEdgeWeight, IncomingEdges> + { + &self.relocation_edges + } + + /// Returns the list of adjacent outgoing import edges for this symbol + /// node. + #[inline] + pub fn imports( + &self, + ) -> &EdgeList<'arena, Self, LibraryNode<'arena, 'data>, ImportEdgeWeight<'data>, OutgoingEdges> + { + &self.import_edges + } + + /// Returns the name of the symbol. + #[inline] + pub fn name(&self) -> SymbolName<'arena> { + self.name + } + + /// Returns the storage class of the symbol. + #[inline] + pub fn storage_class(&self) -> SymbolNodeStorageClass { + self.storage_class + } + + /// Returns `true` if this is a section symbol. + #[inline] + pub fn is_section_symbol(&self) -> bool { + self.section + } + + /// Returns `true` if this symbol is a label. + pub fn is_label(&self) -> bool { + self.storage_class == SymbolNodeStorageClass::Label || self.is_msvc_label() + } + + /// Returns `true` if this is an MSVC .data label. + /// + /// These are symbols with static storage class, have a name format of + /// `$SG` and are defined in a data section. + pub fn is_msvc_label(&self) -> bool { + *self.msvc_label.get_or_init(|| { + self.storage_class() == SymbolNodeStorageClass::Static + && self + .name() + .as_str() + .strip_prefix("$SG") + .is_some_and(|unprefixed| unprefixed.parse::().is_ok()) + && self + .definitions() + .front() + .is_some_and(|definition| definition.target().name().group_name() == ".data") + }) + } + + /// Returns `true` if this symbol has no references or all sections + /// referencing this symbol have been discarded. + pub fn is_unreferenced(&self) -> bool { + self.references().is_empty() + || self + .references() + .iter() + .all(|reloc| reloc.source().is_discarded()) + } + + /// Returns `true` if this symbol is undefined. + #[inline] + pub fn is_undefined(&self) -> bool { + self.imports().is_empty() && self.definitions().is_empty() + } + + /// Returns `true` if this symbol has multiple non-COMDAT definitions. + pub fn is_duplicate(&self) -> bool { + self.definitions() + .iter() + .filter(|definition| definition.weight().selection().is_none()) + .count() + > 1 + } + + /// Returns `true` if this symbol is multiply defined. + pub fn is_multiply_defined(&self) -> bool { + let mut noduplicates = false; + let mut samesize = false; + let mut exact_match = false; + + let mut sizes = HashSet::with_capacity(self.definitions().len()); + let mut checksums = HashSet::with_capacity(self.definitions().len()); + + for definition in self.definitions().iter() { + let selection = match definition.weight().selection() { + Some(sel) => sel, + None => continue, + }; + + match selection { + ComdatSelection::NoDuplicates => { + noduplicates = true; + } + ComdatSelection::SameSize => { + sizes.insert(definition.target().data().len()); + samesize = true; + } + ComdatSelection::ExactMatch => { + // TODO: This will just check if the section data matches. + // Also need to check that the relocations and definitions + // match. + checksums.insert(definition.target().checksum()); + exact_match = true; + } + _ => (), + } + } + + (noduplicates && self.definitions().len() > 1) + || (samesize && sizes.len() > 1) + || (exact_match && checksums.len() > 1) + } + + /// Returns the type associated with this symbol. + #[inline] + pub fn typ(&self) -> SymbolNodeType { + self.typ.get() + } + + /// Sets the type to the specified value for this symbol. + #[inline] + pub fn set_type(&self, val: u16) { + self.typ.set(SymbolNodeType::Value(val)); + } + + /// Sets the symbol table index for this symbol. + /// + /// This can only be set once. + #[inline] + pub fn assign_table_index(&self, value: u32) -> Result<(), u32> { + self.table_index.set(value) + } + + /// Gets the assigned symbol table index for this symbol. + /// + /// Returns `None` if this symbol has not been assigned an index. + #[inline] + pub fn table_index(&self) -> Option { + self.table_index.get().copied() + } + + /// Gets the name of the symbol for the output COFF. + /// + /// Returns `None` if the name of the symbol was never added to the COFF. + #[inline] + pub fn output_name(&self) -> &OnceCell { + &self.output_name + } +} + +impl std::fmt::Debug for SymbolNode<'_, '_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SymbolNode") + .field("name", &self.name) + .field("storage_class", &self.storage_class) + .field("section", &self.section) + .field("typ", &self.typ) + .finish_non_exhaustive() + } +} + +/// A symbol name. +#[derive(Debug, Clone, Copy)] +pub struct SymbolName<'data>(&'data str); + +impl<'data> SymbolName<'data> { + #[inline] + pub fn as_str(&self) -> &'data str { + self.0 + } + + /// Returns a [`SymbolNameDemangler`] for demangling the name of the symbol. + #[inline] + pub fn demangle(&self) -> SymbolNameDemangler<'_, 'data> { + SymbolNameDemangler(self) + } + + /// Returns the symbol name but without the `__declspeci(dllimport)` prefix + /// if it exists. + #[inline] + pub fn strip_dllimport(&self) -> Option<&'data str> { + self.0.strip_prefix("__imp_") + } +} + +impl<'data> From<&'data str> for SymbolName<'data> { + fn from(value: &'data str) -> Self { + Self(value) + } +} + +impl std::fmt::Display for SymbolName<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} + +/// Wrapper around a [`SymbolName`] for demangling the name string. +#[derive(Debug, Clone, Copy)] +pub struct SymbolNameDemangler<'a, 'data>(&'a SymbolName<'data>); + +impl std::fmt::Display for SymbolNameDemangler<'_, '_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some(unprefixed) = self.0.0.strip_prefix("__imp_") { + write!(f, "__declspec(dllimport) {unprefixed}") + } else { + write!(f, "{}", self.0.0) + } + } +} + +/// The storage class of a symbol. +#[derive(Debug, Copy, Clone, PartialEq, Eq, TryFromPrimitive, IntoPrimitive)] +#[num_enum(error_type(name = TryFromStorageClassError, constructor = TryFromStorageClassError))] +#[repr(u8)] +pub enum SymbolNodeStorageClass { + EndOfFunction = IMAGE_SYM_CLASS_END_OF_FUNCTION, + Null = IMAGE_SYM_CLASS_NULL, + Automatic = IMAGE_SYM_CLASS_AUTOMATIC, + External = IMAGE_SYM_CLASS_EXTERNAL, + Static = IMAGE_SYM_CLASS_STATIC, + Register = IMAGE_SYM_CLASS_REGISTER, + ExternalDef = IMAGE_SYM_CLASS_EXTERNAL_DEF, + Label = IMAGE_SYM_CLASS_LABEL, + UndefinedLabel = IMAGE_SYM_CLASS_UNDEFINED_LABEL, + MemberOfStruct = IMAGE_SYM_CLASS_MEMBER_OF_STRUCT, + Argument = IMAGE_SYM_CLASS_ARGUMENT, + StructTag = IMAGE_SYM_CLASS_STRUCT_TAG, + MemberOfUnion = IMAGE_SYM_CLASS_MEMBER_OF_UNION, + UnionTag = IMAGE_SYM_CLASS_UNION_TAG, + TypeDefinition = IMAGE_SYM_CLASS_TYPE_DEFINITION, + UndefinedStatic = IMAGE_SYM_CLASS_UNDEFINED_STATIC, + EnumTag = IMAGE_SYM_CLASS_ENUM_TAG, + MemberOfEnum = IMAGE_SYM_CLASS_MEMBER_OF_ENUM, + RegisterParam = IMAGE_SYM_CLASS_REGISTER_PARAM, + BitField = IMAGE_SYM_CLASS_BIT_FIELD, + Block = IMAGE_SYM_CLASS_BLOCK, + Function = IMAGE_SYM_CLASS_FUNCTION, + EndOfStruct = IMAGE_SYM_CLASS_END_OF_STRUCT, + File = IMAGE_SYM_CLASS_FILE, + Section = IMAGE_SYM_CLASS_SECTION, + WeakExternal = IMAGE_SYM_CLASS_WEAK_EXTERNAL, + ClrToken = IMAGE_SYM_CLASS_CLR_TOKEN, +} + +/// The type of symbol. +#[derive(Debug, Copy, Clone)] +pub enum SymbolNodeType { + /// A debug symbol. + Debug, + + /// An absolute symbol. + Absolute(#[allow(unused)] u32), + + /// A defined symbol type value. + Value(u16), +} diff --git a/src/graph/spec.rs b/src/graph/spec.rs new file mode 100644 index 0000000..cddb275 --- /dev/null +++ b/src/graph/spec.rs @@ -0,0 +1,130 @@ +use std::{cell::OnceCell, collections::LinkedList}; + +use indexmap::{IndexMap, IndexSet}; +use object::{ + Object, ObjectSection, ObjectSymbol, + coff::{CoffFile, CoffHeader, ImageSymbol}, +}; + +use crate::linker::LinkerTargetArch; + +use super::{ + LinkGraph, ROOT_COFF, + cache::LinkGraphCache, + edge::{DefinitionEdgeWeight, Edge, RelocationEdgeWeight}, + link::LinkGraphArena, + node::{CoffNode, SectionNode, SymbolNode}, +}; + +#[cfg(target_pointer_width = "16")] +const SYSTEM_ALIGNMENT: usize = 2; + +#[cfg(target_pointer_width = "32")] +const SYSTEM_ALIGNMENT: usize = 4; + +#[cfg(target_pointer_width = "64")] +const SYSTEM_ALIGNMENT: usize = 8; + +/// Structure for speculatively calculating the amount of memory needed for +/// building the main link graph. +pub struct SpecLinkGraph { + coffs: usize, + externals: usize, + sections: usize, + max_sections: usize, + max_symbols: usize, + alloc_size: usize, +} + +impl SpecLinkGraph { + /// Creates a new [`SpecLinkGraph`]. + pub fn new() -> SpecLinkGraph { + Self { + coffs: 0, + externals: 0, + sections: 0, + max_sections: 0, + max_symbols: 0, + alloc_size: 0usize, + } + } + + /// Returns the calculated number of bytes needed to hold graph components. + #[inline] + pub fn byte_capacity(&self) -> usize { + self.alloc_size + } + + /// Adds a COFF to the [`LinkGraph`] allocation calculation. + pub fn add_coff<'a, C: CoffHeader>(&mut self, coff: &CoffFile<'a, &'a [u8], C>) { + self.coffs += 1; + + self.max_sections = self.max_sections.max(coff.coff_section_table().len()); + self.max_symbols = self.max_symbols.max(coff.coff_symbol_table().len()); + self.sections += coff.coff_section_table().len(); + + self.alloc_size += std::mem::size_of::(); + self.alloc_size = self.alloc_size.next_multiple_of(SYSTEM_ALIGNMENT); + + for _ in coff.sections() { + self.alloc_size += std::mem::size_of::(); + self.alloc_size = self.alloc_size.next_multiple_of(SYSTEM_ALIGNMENT); + } + + for symbol in coff.symbols() { + self.alloc_size += std::mem::size_of::(); + self.alloc_size = self.alloc_size.next_multiple_of(SYSTEM_ALIGNMENT); + + if symbol.is_global() { + self.externals += 1; + } + + if symbol.is_definition() || symbol.coff_symbol().has_aux_section() { + self.alloc_size += + std::mem::size_of::>(); + self.alloc_size = self.alloc_size.next_multiple_of(SYSTEM_ALIGNMENT); + } + } + + for section in coff.sections() { + for _ in section.relocations() { + self.alloc_size += + std::mem::size_of::>(); + self.alloc_size = self.alloc_size.next_multiple_of(SYSTEM_ALIGNMENT); + } + } + } + + /// Allocates the arena for the [`LinkGraph`]. + pub fn alloc_arena(&self) -> LinkGraphArena { + LinkGraphArena::with_capacity(self.byte_capacity()) + } + + /// Allocates the [`LinkGraph`] using the specified `arena`. + pub fn alloc_graph<'data>( + self, + arena: &LinkGraphArena, + machine: LinkerTargetArch, + ) -> LinkGraph<'_, 'data> { + LinkGraph { + machine, + section_nodes: Vec::with_capacity(self.sections), + common_section: OnceCell::new(), + library_nodes: IndexMap::new(), + coff_nodes: IndexSet::with_capacity(self.coffs), + root_coff: &ROOT_COFF, + api_node: None, + external_symbols: IndexMap::with_capacity(self.externals), + extraneous_symbols: LinkedList::new(), + cache: LinkGraphCache::with_capacity(self.max_symbols, self.max_sections), + node_count: 0, + arena, + } + } +} + +impl Default for SpecLinkGraph { + fn default() -> Self { + Self::new() + } +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..17a40fd --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,7 @@ +mod api; +mod drectve; +pub mod graph; +pub mod libsearch; +pub mod linker; +pub mod linkobject; +pub mod pathed_item; diff --git a/src/libsearch.rs b/src/libsearch.rs new file mode 100644 index 0000000..d4e6133 --- /dev/null +++ b/src/libsearch.rs @@ -0,0 +1,133 @@ +use std::{borrow::Cow, io::ErrorKind, path::PathBuf}; + +use indexmap::IndexSet; +use log::debug; + +use crate::pathed_item::PathedItem; + +pub trait LibraryFind { + fn find_library(&self, name: impl AsRef) -> Result; +} + +#[derive(Debug, thiserror::Error)] +pub enum LibsearchError { + #[error("unable to find library -l{0}")] + NotFound(String), + + #[error("could not open link library {}: {error}", .path.display())] + Io { + path: PathBuf, + error: std::io::Error, + }, +} + +/// A search library name +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +pub struct SearchLibraryName<'a>(&'a str); + +impl SearchLibraryName<'_> { + pub fn value(&self) -> &str { + self.0.trim_start_matches(':') + } + + pub fn is_filename(&self) -> bool { + self.0.starts_with(':') + } +} + +impl<'a> From<&'a str> for SearchLibraryName<'a> { + fn from(value: &'a str) -> Self { + Self(value) + } +} + +impl std::fmt::Display for SearchLibraryName<'_> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +/// A read in link library found from the [`LibrarySearcher`]. +pub type FoundLibrary = PathedItem>; + +impl std::hash::Hash for FoundLibrary { + fn hash(&self, state: &mut H) { + self.path().hash(state); + } +} + +impl std::cmp::PartialEq for FoundLibrary { + fn eq(&self, other: &Self) -> bool { + self.path().eq(other.path()) + } +} + +impl std::cmp::Eq for FoundLibrary {} + +/// Used for finding link libraries. +#[derive(Default)] +pub struct LibrarySearcher { + search_paths: IndexSet, +} + +impl LibrarySearcher { + pub fn new() -> LibrarySearcher { + Default::default() + } + + pub fn extend_search_paths(&mut self, search_paths: I) + where + I: IntoIterator, + P: Into, + { + self.search_paths + .extend(search_paths.into_iter().map(|v| v.into())); + } +} + +impl LibraryFind for LibrarySearcher { + fn find_library(&self, name: impl AsRef) -> Result { + if self.search_paths.is_empty() { + return Err(LibsearchError::NotFound(name.as_ref().to_string())); + } + + let library = SearchLibraryName::from(name.as_ref()); + + let library_filenames: Vec> = if !library.is_filename() { + let name = library.value(); + // Create a vec with the library file names to check. + vec![ + format!("lib{name}.dll.a").into(), + format!("{name}.dll.a").into(), + format!("lib{name}.a").into(), + format!("{name}.lib").into(), + format!("lib{name}.lib").into(), + format!("{name}.a").into(), + ] + } else { + vec![Cow::Borrowed(library.value())] + }; + + for search_path in &self.search_paths { + for filename in &library_filenames { + let full_path = search_path.join(filename.as_ref()); + match std::fs::read(&full_path) { + Ok(data) => { + return Ok(FoundLibrary::new(full_path, data)); + } + Err(e) if e.kind() != ErrorKind::NotFound => { + return Err(LibsearchError::Io { + path: full_path, + error: e, + }); + } + Err(e) => { + debug!("attempt to open {} failed ({})", full_path.display(), e); + } + }; + } + } + + Err(LibsearchError::NotFound(name.as_ref().to_string())) + } +} diff --git a/src/linker/builder.rs b/src/linker/builder.rs new file mode 100644 index 0000000..e933d51 --- /dev/null +++ b/src/linker/builder.rs @@ -0,0 +1,146 @@ +use std::path::PathBuf; + +use indexmap::IndexSet; + +use crate::{ + api::BeaconApiInit, + libsearch::{LibraryFind, LibrarySearcher}, + pathed_item::PathedItem, +}; + +use super::{ConfiguredLinker, CustomApiInit, LinkImpl, LinkerTargetArch}; + +/// Sets up inputs and configures a [`super::Linker`]. +#[derive(Default)] +pub struct LinkerBuilder { + /// The target architecture. + pub(super) target_arch: Option, + + /// The input files to link. + pub(super) inputs: Vec>>, + + /// Link libraries. + pub(super) libraries: IndexSet, + + /// The name of the entrypoint symbol. + pub(super) entrypoint: Option, + + /// Custom BOF API to use. + pub(super) custom_api: Option, + + /// Whether to merge the .bss section with the .data section. + pub(super) merge_bss: bool, + + /// Searcher for finding link libraries. + pub(super) library_searcher: Option, + + /// Output path for dumping the link graph. + pub(super) link_graph_output: Option, +} + +impl LinkerBuilder { + /// Creates a new [`LinkerBuilder`] with the defaults. + pub fn new() -> Self { + Self { + target_arch: Default::default(), + inputs: Default::default(), + libraries: Default::default(), + entrypoint: Default::default(), + custom_api: Default::default(), + merge_bss: false, + library_searcher: None, + link_graph_output: None, + } + } + + /// Sets the target architecture for the linker. + /// + /// This is not needed if the linker can parse the target architecture + /// from the input files. + pub fn architecture(mut self, arch: LinkerTargetArch) -> Self { + self.target_arch = Some(arch); + self + } + + /// Set the output path for dumping the link graph. + pub fn link_graph_path(mut self, path: impl Into) -> Self { + self.link_graph_output = Some(path.into()); + self + } + + /// Merge the .bss section with the .data section. + pub fn merge_bss(mut self, val: bool) -> Self { + self.merge_bss = val; + self + } + + /// Custom BOF API to use instead of the Beacon API. + pub fn custom_api(mut self, api: impl Into) -> Self { + self.custom_api = Some(api.into()); + self + } + + /// Set the library searcher to use for finding link libraries. + pub fn library_searcher(mut self, searcher: L) -> Self { + self.library_searcher = Some(searcher); + self + } + + /// Add an input file to the linker. + pub fn add_input(mut self, input: PathedItem>) -> Self { + self.inputs.push(input); + self + } + + /// Add a set of input files to the linker. + pub fn add_inputs( + mut self, + inputs: impl IntoIterator>>, + ) -> Self { + self.inputs.extend(inputs); + self + } + + /// Add a link library to the linker. + pub fn add_library(mut self, name: impl Into) -> Self { + self.libraries.insert(name.into()); + self + } + + /// Add a set of link libraries to the linker. + pub fn add_libraries, I: IntoIterator>(mut self, names: I) -> Self { + self.libraries.extend(names.into_iter().map(Into::into)); + self + } + + /// Finishes configuring the linker. + pub fn build(mut self) -> Box { + if let Some(library_searcher) = self.library_searcher.take() { + if let Some(custom_api) = self.custom_api.take() { + Box::new(ConfiguredLinker::with_opts( + self, + library_searcher, + CustomApiInit::from(custom_api), + )) + } else { + Box::new(ConfiguredLinker::with_opts( + self, + library_searcher, + BeaconApiInit, + )) + } + } else if let Some(custom_api) = self.custom_api.take() { + Box::new(ConfiguredLinker::with_opts( + self, + LibrarySearcher::new(), + CustomApiInit::from(custom_api), + )) + } else { + Box::new(ConfiguredLinker::with_opts( + self, + LibrarySearcher::new(), + BeaconApiInit, + )) + } + } +} diff --git a/src/linker/configured.rs b/src/linker/configured.rs new file mode 100644 index 0000000..5accb22 --- /dev/null +++ b/src/linker/configured.rs @@ -0,0 +1,469 @@ +use std::{ + collections::VecDeque, + io::BufWriter, + path::{Path, PathBuf}, +}; + +use indexmap::{IndexMap, IndexSet}; +use log::warn; +use object::{Object, coff::CoffFile}; +use typed_arena::Arena; + +use crate::{ + api::{ApiSymbolError, ApiSymbolSource}, + drectve, + graph::LinkGraph, + libsearch::LibraryFind, + linker::error::{DrectveLibsearchError, LinkerSymbolErrors}, + linkobject::archive::{ExtractMemberError, ExtractedMemberContents, LinkArchive}, + pathed_item::PathedItem, +}; + +use super::{ + ApiInit, ApiInitCtx, LinkImpl, LinkerBuilder, LinkerTargetArch, + error::{LinkError, LinkerSetupError, LinkerSetupErrors, LinkerSetupPathError}, +}; + +/// A configured linker. +pub struct ConfiguredLinker { + /// The target architecture. + target_arch: Option, + + /// The unparsed linker inputs + inputs: Vec>>, + + /// The names of the link libraries. + library_names: IndexSet, + + /// The custom API. + custom_api: Api, + + /// The link library searcher. + library_searcher: L, + + /// The name of the entrypoint symbol. + entrypoint: Option, + + /// Whether to merge the .bss section with the .data section. + merge_bss: bool, + + /// Output path for dumping the link graph. + link_graph_output: Option, +} + +impl ConfiguredLinker { + /// Returns a [`LinkerBuilder`] for configuring a linker. + pub fn builder() -> LinkerBuilder { + LinkerBuilder::new() + } + + pub(super) fn with_opts( + builder: LinkerBuilder, + library_searcher: L, + custom_api: Api, + ) -> ConfiguredLinker { + Self { + target_arch: builder.target_arch, + inputs: builder.inputs, + library_names: builder.libraries, + custom_api, + library_searcher, + entrypoint: builder.entrypoint, + merge_bss: builder.merge_bss, + link_graph_output: builder.link_graph_output, + } + } +} + +impl LinkImpl for ConfiguredLinker { + fn link(&mut self) -> Result, LinkError> { + // Parsed input COFFs + let mut parsed_inputs = Vec::with_capacity(self.inputs.len()); + + // Errors during setup + let mut setup_errors = Vec::new(); + + // Parsed link libraries + let mut link_libraries = + IndexMap::with_capacity(self.inputs.len() + self.library_names.len()); + + // The opened link library names including .drectve libraries + let mut library_names: IndexSet<&str> = + IndexSet::from_iter(self.library_names.iter().map(|v| v.as_str())); + + // Spec graph for calculating memory needed for allocating the link + // graph. + let mut spec = LinkGraph::spec(); + + // Queue of .drectve libraries to open + let mut drectve_queue = VecDeque::with_capacity(self.inputs.len()); + + // Parse the command line input files + for input in &self.inputs { + // Check if this is an archive file passed in the command line + if input + .get(..object::archive::MAGIC.len()) + .is_some_and(|magic| magic == object::archive::MAGIC) + { + match LinkArchive::parse(input.as_slice()) + .map_err(|e| LinkerSetupPathError::nomember(input.path(), e)) + { + Ok(parsed) => { + link_libraries.insert(input.path().as_path(), parsed); + } + Err(e) => { + setup_errors.push(LinkerSetupError::Path(e)); + } + }; + } else { + match CoffFile::<_>::parse(input.as_slice()) + .map_err(|e| LinkerSetupPathError::nomember(input.path(), e)) + { + Ok(parsed) => { + // Add .drectve libraries to the drectve_queue. + for library_name in drectve::parse_drectve_libraries(&parsed) + .into_iter() + .flatten() + { + let library_name = library_name.trim_end_matches(".lib"); + if library_names.insert(library_name) { + drectve_queue.push_back((input.path().as_path(), library_name)); + } + } + + spec.add_coff(&parsed); + + // Add the COFF to the list of parsed inputs. + parsed_inputs.push(PathedItem::new(input.path().as_path(), parsed)); + } + Err(e) => { + setup_errors.push(LinkerSetupError::Path(e)); + } + } + } + } + + let library_arena = Arena::with_capacity(library_names.len() + 1); + + // Open link libraries + for link_library in &self.library_names { + let found = match self.library_searcher.find_library(link_library) { + Ok(found) => { + if link_libraries.contains_key(found.path().as_path()) { + continue; + } + + library_arena.alloc(found) + } + Err(e) => { + setup_errors.push(LinkerSetupError::Library(e)); + continue; + } + }; + + let parsed = match LinkArchive::parse(found.as_slice()) { + Ok(parsed) => parsed, + Err(e) => { + setup_errors.push(LinkerSetupError::Path(LinkerSetupPathError::nomember( + found.path(), + e, + ))); + continue; + } + }; + + link_libraries.insert(found.path().as_path(), parsed); + } + + // Open drectve link libraries + while let Some((coff_path, drectve_library)) = drectve_queue.pop_front() { + let found = match self.library_searcher.find_library(drectve_library) { + Ok(found) => { + if link_libraries.contains_key(found.path().as_path()) { + continue; + } + + library_arena.alloc(found) + } + Err(e) => { + setup_errors.push(LinkerSetupError::Path(LinkerSetupPathError::nomember( + coff_path, + DrectveLibsearchError::from(e), + ))); + continue; + } + }; + + let parsed = match LinkArchive::parse(found.as_slice()) { + Ok(parsed) => parsed, + Err(e) => { + setup_errors.push(LinkerSetupError::Path(LinkerSetupPathError::nomember( + found.path(), + e, + ))); + continue; + } + }; + + link_libraries.insert(found.path().as_path(), parsed); + } + + let target_arch = self.target_arch.take().or_else(|| { + parsed_inputs + .iter() + .find_map(|coff| LinkerTargetArch::try_from(coff.architecture()).ok()) + }); + + let target_arch = match target_arch { + Some(target_arch) => target_arch, + None => { + if !setup_errors.is_empty() { + return Err(LinkError::Setup(LinkerSetupErrors(setup_errors))); + } + + if self.inputs.is_empty() { + return Err(LinkError::NoInput); + } + + return Err(LinkError::ArchitectureDetect); + } + }; + + // Initialize the custom API + let api_resolver = match self.custom_api.initialize_api(&ApiInitCtx { + target_arch, + library_searcher: &self.library_searcher, + arena: &library_arena, + }) { + Ok(resolver) => resolver, + Err(e) => { + setup_errors.push(LinkerSetupError::ApiInit(e)); + return Err(LinkError::Setup(LinkerSetupErrors(setup_errors))); + } + }; + + // Check errors + if !setup_errors.is_empty() { + return Err(LinkError::Setup(LinkerSetupErrors(setup_errors))); + } + + if self.inputs.is_empty() { + return Err(LinkError::NoInput); + } + + // Build the graph + let graph_arena = spec.alloc_arena(); + let mut graph = spec.alloc_graph(&graph_arena, target_arch); + + // Add COFFs + for coff in parsed_inputs { + if let Err(e) = graph.add_coff(coff.path(), None, &coff) { + setup_errors.push(LinkerSetupError::Path(LinkerSetupPathError::nomember( + coff.path(), + e, + ))); + } + } + + // Return any errors + if !setup_errors.is_empty() { + return Err(LinkError::Setup(LinkerSetupErrors(setup_errors))); + } + + // Add the entrypoint symbol so that it can be linked in from archives + if let Some(entrypoint) = &mut self.entrypoint { + if target_arch == LinkerTargetArch::I386 { + *entrypoint = format!("_{entrypoint}"); + } + + graph.add_external_symbol(entrypoint); + } + + let mut drectve_queue: VecDeque<((&Path, &Path), &str)> = VecDeque::new(); + + let undefined_count = graph.undefined_symbols().count(); + let mut symbol_search_buffer = VecDeque::with_capacity(undefined_count); + let mut undefined_symbols: IndexSet<&str> = IndexSet::with_capacity(undefined_count); + + // Resolve symbols + loop { + // Get the list of undefined symbols to search for + symbol_search_buffer.extend( + graph + .undefined_symbols() + .filter(|symbol| !undefined_symbols.contains(symbol)), + ); + + // If the search list is empty, finished resolving + if symbol_search_buffer.is_empty() { + break; + } + + // Attempt to resolve each symbol in the search list + 'symbol: while let Some(symbol_name) = symbol_search_buffer.pop_front() { + // Try resolving it as an API import first + match api_resolver.extract_api_symbol(symbol_name) { + Ok(api_import) => { + if let Err(e) = graph.add_api_import(symbol_name, &api_import) { + setup_errors.push(LinkerSetupError::Path( + LinkerSetupPathError::nomember(api_resolver.api_path(), e), + )); + } else { + continue; + } + } + Err(ApiSymbolError::NotFound) => (), + Err(e) => { + setup_errors.push(LinkerSetupError::Path(LinkerSetupPathError::nomember( + api_resolver.api_path(), + e, + ))); + } + } + + // Open any pending libraries in the .drectve queue + while let Some(((library_path, coff_path), drectve_library)) = + drectve_queue.pop_front() + { + match self.library_searcher.find_library(drectve_library) { + Ok(found) => { + if library_names.insert(drectve_library) { + let found = library_arena.alloc(found); + + match LinkArchive::parse(found.as_slice()) { + Ok(parsed) => { + link_libraries.insert(found.path().as_path(), parsed); + } + Err(e) => { + setup_errors.push(LinkerSetupError::Path( + LinkerSetupPathError::new( + library_path, + Some(coff_path), + e, + ), + )); + } + } + } + } + Err(e) => { + setup_errors.push(LinkerSetupError::Path(LinkerSetupPathError::new( + library_path, + Some(coff_path), + DrectveLibsearchError::from(e), + ))); + } + } + } + + // Attempt to resolve the symbol using the opened link libraries + for (library_path, library) in &link_libraries { + let extracted = + match library.extract_symbol(symbol_name) { + Ok(extracted) => extracted, + Err(ExtractMemberError::NotFound) => { + continue; + } + Err(ExtractMemberError::ArchiveParse(e)) => { + setup_errors.push(LinkerSetupError::Path( + LinkerSetupPathError::nomember(library_path, e), + )); + continue; + } + Err(ExtractMemberError::MemberParse(e)) => { + setup_errors.push(LinkerSetupError::Path( + LinkerSetupPathError::new(library_path, Some(e.path), e.kind), + )); + continue; + } + }; + + match extracted.contents() { + ExtractedMemberContents::Coff(coff) => { + // Add any .drectve link libraries from linked in COFFs + // to the drectve queue + for drectve_library in + drectve::parse_drectve_libraries(coff).into_iter().flatten() + { + let drectve_library_name = drectve_library.trim_end_matches(".lib"); + if library_names.contains(drectve_library) { + drectve_queue.push_back(( + (library_path, extracted.path()), + drectve_library_name, + )); + } + } + + if let Err(e) = + graph.add_coff(library_path, Some(extracted.path()), coff) + { + setup_errors.push(LinkerSetupError::Path( + LinkerSetupPathError::new( + library_path, + Some(extracted.path()), + e, + ), + )); + continue; + } + + continue 'symbol; + } + ExtractedMemberContents::Import(import_member) => { + if let Err(e) = graph.add_library_import(symbol_name, import_member) { + setup_errors.push(LinkerSetupError::Path( + LinkerSetupPathError::new( + library_path, + Some(extracted.path()), + e, + ), + )); + continue; + } + + continue 'symbol; + } + } + } + + // Symbol could not be found in any of the link libraries + undefined_symbols.insert(symbol_name); + } + } + + // Write out the link graph + if let Some(graph_path) = self.link_graph_output.as_ref() { + match std::fs::File::create(graph_path) { + Ok(f) => { + if let Err(e) = graph.write_dot_graph(BufWriter::new(f)) { + warn!("could not write link graph: {e}"); + } + } + Err(e) => { + warn!("could not open {}: {e}", graph_path.display()); + } + } + } + + // Return errors + if !setup_errors.is_empty() { + return Err(LinkError::Setup(LinkerSetupErrors(setup_errors))); + } + + // Finish building the link graph + let mut graph = match graph.finish() { + Ok(graph) => graph, + Err(e) => { + return Err(LinkError::Symbol(LinkerSymbolErrors( + e.into_iter().map(|v| v.to_string()).collect(), + ))); + } + }; + + if self.merge_bss { + graph.merge_bss(); + } + + Ok(graph.link()?) + } +} diff --git a/src/linker/error.rs b/src/linker/error.rs new file mode 100644 index 0000000..a86ecd1 --- /dev/null +++ b/src/linker/error.rs @@ -0,0 +1,192 @@ +use std::path::PathBuf; + +use crate::{ + api::ApiSymbolError, + graph::{LinkGraphAddError, LinkGraphLinkError}, + libsearch::LibsearchError, + linkobject::archive::{ArchiveParseError, LinkArchiveParseError, MemberParseErrorKind}, +}; + +#[derive(Debug, thiserror::Error)] +pub enum LinkError { + #[error("{0}")] + Setup(LinkerSetupErrors), + + #[error("{0}")] + Symbol(LinkerSymbolErrors), + + #[error("{0}")] + Graph(#[from] LinkGraphLinkError), + + #[error("no input files")] + NoInput, + + #[error("could not detect architecture")] + ArchitectureDetect, +} + +#[derive(Debug, thiserror::Error)] +#[error("{}", display_vec(.0))] +pub struct LinkerSetupErrors(pub(super) Vec); + +impl LinkerSetupErrors { + pub fn errors(&self) -> &[LinkerSetupError] { + &self.0 + } +} + +#[derive(Debug, thiserror::Error)] +pub enum LinkerSetupError { + #[error("{0}")] + Path(LinkerSetupPathError), + + #[error("{0}")] + Library(LibsearchError), + + #[error("{0}")] + ApiInit(ApiInitError), +} + +#[derive(Debug, thiserror::Error)] +pub enum ApiInitError { + #[error("{}: could not open custom API: {error}", .path.display())] + Io { + path: PathBuf, + error: std::io::Error, + }, + + #[error("unable to find custom API '{0}'")] + NotFound(String), + + #[error("{}: {error}", .path.display())] + Parse { + path: PathBuf, + error: LinkArchiveParseError, + }, +} + +impl From for ApiInitError { + fn from(value: LibsearchError) -> Self { + match value { + LibsearchError::NotFound(name) => Self::NotFound(name), + LibsearchError::Io { path, error } => Self::Io { path, error }, + } + } +} + +#[derive(Debug, thiserror::Error)] +#[error( + "{}{}: {error}", + .path.display(), + .member.as_ref().map(|p| format!("({})", p.display())).unwrap_or_default() +)] +pub struct LinkerSetupPathError { + pub path: PathBuf, + pub member: Option, + pub error: LinkerPathErrorKind, +} + +impl LinkerSetupPathError { + pub fn new>( + path: impl Into, + member: Option

, + error: impl Into, + ) -> LinkerSetupPathError { + Self { + path: path.into(), + member: member.map(Into::into), + error: error.into(), + } + } + + pub fn nomember( + path: impl Into, + error: impl Into, + ) -> LinkerSetupPathError { + Self { + path: path.into(), + member: None, + error: error.into(), + } + } +} + +#[derive(Debug, thiserror::Error)] +pub enum LinkerPathErrorKind { + #[error("{0}")] + DrectveLibrary(#[from] DrectveLibsearchError), + + #[error("{0}")] + ArchiveParse(#[from] LinkArchiveParseError), + + #[error("{0}")] + ArchiveExtract(#[from] ArchiveParseError), + + #[error("{0}")] + MemberExtract(#[from] MemberParseErrorKind), + + #[error("{0}")] + GraphAdd(#[from] LinkGraphAddError), + + #[error("{0}")] + ApiSymbol(#[from] ApiSymbolError), + + #[error("{0}")] + Object(#[from] object::Error), +} + +#[derive(Debug, thiserror::Error)] +pub enum DrectveLibsearchError { + #[error("unable to find library {0}")] + NotFound(String), + + #[error("could not open link library {}: {error}", .path.display())] + Io { + path: PathBuf, + error: std::io::Error, + }, +} + +impl From for DrectveLibsearchError { + fn from(value: LibsearchError) -> Self { + match value { + LibsearchError::Io { path, error } => Self::Io { path, error }, + LibsearchError::NotFound(name) => Self::NotFound(name), + } + } +} + +#[derive(Debug, thiserror::Error)] +#[error("{}", display_vec(.0))] +pub struct LinkerSymbolErrors(pub(super) Vec); + +impl LinkerSymbolErrors { + pub fn errors(&self) -> &[String] { + &self.0 + } +} + +struct DisplayVec<'a, T: std::fmt::Display>(&'a Vec); + +impl<'a, T: std::fmt::Display> std::fmt::Display for DisplayVec<'a, T> { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut value_iter = self.0.iter(); + + let first_value = match value_iter.next() { + Some(v) => v, + None => return Ok(()), + }; + + first_value.fmt(f)?; + + for val in value_iter { + write!(f, "\n{val}")?; + } + + Ok(()) + } +} + +fn display_vec(errors: &Vec) -> DisplayVec<'_, T> { + DisplayVec(errors) +} diff --git a/src/linker/mod.rs b/src/linker/mod.rs new file mode 100644 index 0000000..5a5c4d6 --- /dev/null +++ b/src/linker/mod.rs @@ -0,0 +1,110 @@ +use std::path::{Path, PathBuf}; + +use num_enum::{IntoPrimitive, TryFromPrimitive}; +use object::pe::{IMAGE_FILE_MACHINE_AMD64, IMAGE_FILE_MACHINE_I386}; +use typed_arena::Arena; + +use crate::{ + api::ApiSymbolSource, libsearch::LibraryFind, linkobject::archive::LinkArchive, + pathed_item::PathedItem, +}; +use error::{ApiInitError, LinkError}; + +mod builder; +mod configured; +pub mod error; + +pub use self::configured::*; +pub use builder::*; + +pub trait LinkImpl { + fn link(&mut self) -> Result, LinkError>; +} + +#[derive(Clone, Copy, PartialEq, Eq, TryFromPrimitive, IntoPrimitive)] +#[repr(u16)] +pub enum LinkerTargetArch { + Amd64 = IMAGE_FILE_MACHINE_AMD64, + I386 = IMAGE_FILE_MACHINE_I386, +} + +impl From for object::Architecture { + fn from(value: LinkerTargetArch) -> Self { + match value { + LinkerTargetArch::Amd64 => object::Architecture::X86_64, + LinkerTargetArch::I386 => object::Architecture::I386, + } + } +} + +impl TryFrom for LinkerTargetArch { + type Error = object::Architecture; + + fn try_from(value: object::Architecture) -> Result { + Ok(match value { + object::Architecture::X86_64 => Self::Amd64, + object::Architecture::I386 => Self::I386, + _ => return Err(value), + }) + } +} + +pub struct ApiInitCtx<'b, 'a, L: LibraryFind> { + pub(super) target_arch: LinkerTargetArch, + pub(super) library_searcher: &'b L, + pub(super) arena: &'a Arena>>, +} + +pub trait ApiInit { + type Output<'a>: ApiSymbolSource<'a>; + + fn initialize_api<'a, L: LibraryFind>( + &self, + ctx: &ApiInitCtx<'_, 'a, L>, + ) -> Result, ApiInitError>; +} + +struct CustomApiInit(String); + +impl From for CustomApiInit { + fn from(v: String) -> CustomApiInit { + Self(v) + } +} + +impl ApiInit for CustomApiInit { + type Output<'a> = PathedItem<&'a Path, LinkArchive<'a>>; + + fn initialize_api<'a, L: LibraryFind>( + &self, + ctx: &ApiInitCtx<'_, 'a, L>, + ) -> Result, ApiInitError> { + let custom_api = match std::fs::read(&self.0) { + Ok(buffer) => ctx + .arena + .alloc(PathedItem::new(PathBuf::from(&self.0), buffer)), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + match ctx.library_searcher.find_library(&self.0) { + Ok(found) => ctx.arena.alloc(found), + Err(e) => { + return Err(e.into()); + } + } + } + Err(e) => { + return Err(ApiInitError::Io { + path: PathBuf::from(&self.0), + error: e, + }); + } + }; + + let parsed = + LinkArchive::parse(custom_api.as_slice()).map_err(|e| ApiInitError::Parse { + path: custom_api.path().to_path_buf(), + error: e, + })?; + + Ok(PathedItem::new(custom_api.path().as_path(), parsed)) + } +} diff --git a/src/linkobject/archive/error.rs b/src/linkobject/archive/error.rs new file mode 100644 index 0000000..99ef490 --- /dev/null +++ b/src/linkobject/archive/error.rs @@ -0,0 +1,136 @@ +use std::path::PathBuf; + +use crate::linkobject::import::TryFromImportFileError; + +#[derive(Debug, thiserror::Error)] +pub enum LinkArchiveParseError { + #[error("thin archives are not supported")] + ThinArchive, + + #[error("archive is missing a symbol table")] + NoSymbolMap, + + #[error("{0}")] + Object(#[from] object::read::Error), +} + +#[derive(Debug, thiserror::Error)] +pub enum ExtractMemberError { + #[error("member for symbol does not exist")] + NotFound, + + #[error("{0}")] + ArchiveParse(#[from] ArchiveParseError), + + #[error("{0}")] + MemberParse(#[from] MemberParseError), +} + +#[derive(Debug, thiserror::Error)] +pub enum ArchiveParseError { + #[error("archive member name is invalid: {0}")] + MemberName(std::str::Utf8Error), + + #[error("failed parsing archive file: {0}")] + Object(#[from] object::read::Error), +} + +#[derive(Debug, thiserror::Error)] +#[error("could not parse member {}: {kind}", .path.display())] +pub struct MemberParseError { + pub path: PathBuf, + pub kind: MemberParseErrorKind, +} + +impl MemberParseError { + pub fn new( + path: impl Into, + kind: impl Into, + ) -> MemberParseError { + Self { + path: path.into(), + kind: kind.into(), + } + } +} + +#[derive(Debug, thiserror::Error)] +pub enum MemberParseErrorKind { + #[error("failed parsing legacy import library symbol member: {0}")] + LegacyImportLibrarySymbolMember(#[from] LegacyImportSymbolMemberParseError), + + #[error("failed parsing legacy import library head member: {0}")] + LegacyImportLibraryHeadMember(#[from] LegacyImportHeadMemberParseError), + + #[error("failed parsing legacy import library tail member: {0}")] + LegacyImportLibraryTailMember(#[from] LegacyImportTailMemberParseError), + + #[error("legacy import library is missing symbol '{0}'")] + LegacyImportLibraryMissingSymbol(String), + + #[error("import library member is invalid: {0}")] + ImportFile(#[from] TryFromImportFileError), + + #[error("{0}")] + Object(#[from] object::read::Error), +} + +#[derive(Debug, PartialEq, thiserror::Error)] +pub enum LegacyImportSymbolMemberParseError { + #[error("COFF is not a valid legacy import library symbol COFF")] + Invalid, + + #[error("public symbol is missing")] + MissingPublicSymbol, + + #[error("'_head_*' symbol is missing")] + MissingHeadSymbol, + + #[error("import lookup table is missing")] + IltMissing, + + #[error("import lookup table data is malformed")] + IltMalformed, + + #[error("import lookup table is missing the name table section")] + MissingIltNameSection, + + #[error("import lookup table name section is malformed")] + IltNameMalformed, + + #[error("name string from the import lookup table name table could not be parsed: {0}")] + ImportName(std::str::Utf8Error), + + #[error("{0}")] + Object(#[from] object::read::Error), +} + +#[derive(Debug, thiserror::Error)] +pub enum LegacyImportHeadMemberParseError { + #[error("invalid legacy import library head member")] + Invalid, + + #[error("'*_iname' symbol for the linked tail member is missing")] + MissingInameSymbol, + + #[error("{0}")] + Object(#[from] object::read::Error), +} + +#[derive(Debug, thiserror::Error)] +pub enum LegacyImportTailMemberParseError { + #[error("invalid legacy import library tail member COFF")] + Invalid, + + #[error("'*_iname' symbol is missing")] + MissingInameSymbol, + + #[error("section with the '*_iname' symbol is not valid")] + InameSectionInvalid, + + #[error("could not parse DLL name: {0}")] + DllName(std::str::Utf8Error), + + #[error("{0}")] + Object(#[from] object::read::Error), +} diff --git a/src/linkobject/archive/legacy_importlib.rs b/src/linkobject/archive/legacy_importlib.rs new file mode 100644 index 0000000..4c7f80c --- /dev/null +++ b/src/linkobject/archive/legacy_importlib.rs @@ -0,0 +1,224 @@ +use std::ffi::CStr; + +use object::{ + Object, ObjectSection, ObjectSymbol, + coff::{CoffFile, ImageSymbol}, + pe::{IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE, IMAGE_SYM_CLASS_EXTERNAL}, +}; + +use crate::linkobject::import::{ImportName, ImportType}; + +use super::error::{ + LegacyImportHeadMemberParseError, LegacyImportSymbolMemberParseError, + LegacyImportTailMemberParseError, +}; + +const ILT64_ORDINAL_BIT_SHIFT: u64 = 63; +const ILT32_ORDINAL_BIT_SHIFT: u32 = 31; + +const ILT_ORDINAL_NUMBER_MASK: u64 = 0xffff; + +/// A parsed legacy import library member for a symbol. +pub struct LegacyImportSymbolMember<'a> { + /// The public symbol name. + pub public_symbol: &'a str, + + /// The name to import the symbol as. + pub import_name: ImportName<'a>, + + /// The type of import. + pub typ: ImportType, + + /// The name of the head symbol. + pub head_symbol: &'a str, +} + +impl<'a> LegacyImportSymbolMember<'a> { + pub fn parse( + coff: &CoffFile<'a>, + ) -> Result, LegacyImportSymbolMemberParseError> { + if coff.coff_section_table().len() != 7 { + return Err(LegacyImportSymbolMemberParseError::Invalid); + } + + // This is the first '.idata' section after the .text, .data and .bss + // sections. Use this as a smoke test to check if the COFF is valid. + if coff.section_by_name(".idata$7").is_none() { + return Err(LegacyImportSymbolMemberParseError::Invalid); + } + + let public_symbol = coff + .symbols() + .filter(|symbol| symbol.coff_symbol().storage_class() == IMAGE_SYM_CLASS_EXTERNAL) + .find_map(|symbol| { + symbol + .section_index() + .and_then(|section_idx| coff.section_by_index(section_idx).ok()) + .and_then(|section| section.name().ok()) + .and_then(|section_name| { + (!section_name.starts_with(".idata$")).then_some(symbol) + }) + }) + .ok_or(LegacyImportSymbolMemberParseError::MissingPublicSymbol)?; + + let public_section = coff.section_by_index(public_symbol.section_index().unwrap())?; + + let characteristics = public_section + .coff_section() + .characteristics + .get(object::LittleEndian); + + let typ = if characteristics & IMAGE_SCN_CNT_CODE != 0 { + ImportType::Code + } else if characteristics & IMAGE_SCN_MEM_READ != 0 + && characteristics & IMAGE_SCN_MEM_WRITE == 0 + { + ImportType::Const + } else { + ImportType::Data + }; + + let head_symbol = coff + .symbols() + .filter(|symbol| symbol.coff_symbol().storage_class() == IMAGE_SYM_CLASS_EXTERNAL) + .find_map(|symbol| symbol.is_undefined().then(|| symbol.name().ok()).flatten()) + .ok_or(LegacyImportSymbolMemberParseError::MissingHeadSymbol)?; + + let ilt_section = coff + .section_by_name(".idata$4") + .ok_or(LegacyImportSymbolMemberParseError::IltMissing)?; + + let ilt_data = ilt_section.data()?; + + let ilt = if coff.is_64() { + u64::from_le_bytes( + ilt_data[..8] + .try_into() + .map_err(|_| LegacyImportSymbolMemberParseError::IltMalformed)?, + ) + } else { + u32::from_le_bytes( + ilt_data[..4] + .try_into() + .map_err(|_| LegacyImportSymbolMemberParseError::IltMalformed)?, + ) + .into() + }; + + let import_name = if (coff.is_64() && (ilt & (1 << ILT64_ORDINAL_BIT_SHIFT) != 0)) + || (!coff.is_64() && (ilt & (1 << ILT32_ORDINAL_BIT_SHIFT) != 0)) + { + ImportName::Ordinal((ilt & ILT_ORDINAL_NUMBER_MASK) as u16) + } else { + let name_section = coff + .section_by_name(".idata$6") + .ok_or(LegacyImportSymbolMemberParseError::MissingIltNameSection)?; + let name_data = name_section.data()?; + + let name_bytes = name_data + .get(2..) + .ok_or(LegacyImportSymbolMemberParseError::IltNameMalformed)?; + + let name = CStr::from_bytes_until_nul(name_bytes) + .map(|name| name.to_str()) + .unwrap_or_else(|_| std::str::from_utf8(name_bytes)) + .map_err(LegacyImportSymbolMemberParseError::ImportName)?; + + ImportName::Name(name) + }; + + Ok(LegacyImportSymbolMember { + public_symbol: public_symbol.name()?, + typ, + import_name, + head_symbol, + }) + } +} + +/// The parsed head member for a legacy import library. +pub struct LegacyImportHeadMember<'a> { + /// The name of the '*_iname' symbol for the tail COFF. + pub tail_symbol: &'a str, +} + +impl<'a> LegacyImportHeadMember<'a> { + pub fn parse( + coff: &CoffFile<'a>, + ) -> Result, LegacyImportHeadMemberParseError> { + if coff.coff_section_table().len() != 6 { + return Err(LegacyImportHeadMemberParseError::Invalid); + } + + // This is the first '.idata' section after the .text, .data and .bss + // sections. Use this as a smoke test to check if the COFF is valid. + if coff.section_by_name(".idata$2").is_none() { + return Err(LegacyImportHeadMemberParseError::Invalid); + } + + for symbol in coff + .symbols() + .filter(|symbol| symbol.is_global() && symbol.is_undefined()) + { + let symbol_name = symbol.name()?; + if symbol_name.ends_with("_iname") { + return Ok(LegacyImportHeadMember { + tail_symbol: symbol_name, + }); + } + } + + Err(LegacyImportHeadMemberParseError::MissingInameSymbol) + } +} + +/// The tail member for a legacy import library. +pub struct LegacyImportTailMember<'a> { + /// The DLL name contained in the COFF. + pub dll: &'a str, +} + +impl<'a> LegacyImportTailMember<'a> { + pub fn parse( + coff: &CoffFile<'a>, + ) -> Result, LegacyImportTailMemberParseError> { + if coff.coff_section_table().len() != 6 { + return Err(LegacyImportTailMemberParseError::Invalid); + } + + // This is the first '.idata' section after the .text, .data and .bss + // sections. Use it as a smoke test to check if the COFF is valid. + if coff.section_by_name(".idata$4").is_none() { + return Err(LegacyImportTailMemberParseError::Invalid); + } + + for symbol in coff + .symbols() + .filter(|symbol| symbol.is_global() && symbol.is_definition()) + { + let symbol_name = symbol.name()?; + if symbol_name.ends_with("_iname") { + let iname_section = match symbol.section() { + object::SymbolSection::Section(section_idx) => { + coff.section_by_index(section_idx)? + } + _ => return Err(LegacyImportTailMemberParseError::InameSectionInvalid), + }; + + if iname_section.name()? != ".idata$7" { + return Err(LegacyImportTailMemberParseError::InameSectionInvalid); + } + + let iname_data = iname_section.data()?; + let dll = CStr::from_bytes_until_nul(iname_data) + .map(|name| name.to_str()) + .unwrap_or_else(|_| std::str::from_utf8(iname_data)) + .map_err(LegacyImportTailMemberParseError::DllName)?; + + return Ok(LegacyImportTailMember { dll }); + } + } + + Err(LegacyImportTailMemberParseError::MissingInameSymbol) + } +} diff --git a/src/linkobject/archive/mod.rs b/src/linkobject/archive/mod.rs new file mode 100644 index 0000000..eac1ba1 --- /dev/null +++ b/src/linkobject/archive/mod.rs @@ -0,0 +1,368 @@ +use std::{ + cell::RefCell, + collections::{BTreeMap, HashMap}, + ops::Deref, + path::{Path, PathBuf}, +}; + +use object::{ + Object, + coff::{CoffFile, ImportFile}, + pe::IMAGE_FILE_MACHINE_UNKNOWN, + read::archive::{ArchiveFile, ArchiveMember, ArchiveOffset, ArchiveSymbolIterator}, +}; + +use crate::{ + api::{ApiSymbolError, ApiSymbolSource}, + pathed_item::PathedItem, +}; + +use super::import::ImportMember; + +pub use error::*; + +use legacy_importlib::{LegacyImportHeadMember, LegacyImportSymbolMember, LegacyImportTailMember}; + +pub mod error; +mod legacy_importlib; + +pub struct ExtractedMember<'a> { + path: &'a Path, + contents: ExtractedMemberContents<'a>, +} + +impl<'a> ExtractedMember<'a> { + pub fn new( + path: &'a Path, + contents: impl Into>, + ) -> ExtractedMember<'a> { + Self { + path, + contents: contents.into(), + } + } + + pub fn path(&self) -> &'a Path { + self.path + } + + pub fn contents(&self) -> &ExtractedMemberContents<'a> { + &self.contents + } +} + +pub enum ExtractedMemberContents<'a> { + Coff(CoffFile<'a>), + Import(ImportMember<'a>), +} + +impl<'a> From> for ExtractedMemberContents<'a> { + fn from(value: CoffFile<'a>) -> Self { + Self::Coff(value) + } +} + +impl<'a> From> for ExtractedMemberContents<'a> { + fn from(value: ImportMember<'a>) -> Self { + Self::Import(value) + } +} + +struct CachedSymbolMap<'a> { + cache: HashMap<&'a str, ArchiveOffset>, + iter: Option>, +} + +impl CachedSymbolMap<'_> { + fn find_symbol(&mut self, symbol: &str) -> Option { + if let Some(found) = self.cache.get(symbol).copied() { + return Some(found); + } + + for archive_symbol in self.iter.iter_mut().flatten().flatten() { + let archive_symbol_name = match std::str::from_utf8(archive_symbol.name()) { + Ok(name) => name, + Err(_) => continue, + }; + + self.cache + .insert(archive_symbol_name, archive_symbol.offset()); + if archive_symbol_name == symbol { + return Some(archive_symbol.offset()); + } + } + + None + } +} + +/// A parsed archive file for linking. +pub struct LinkArchive<'a> { + /// The parsed archive file + archive_file: ArchiveFile<'a>, + + /// The cached archive symbol table. + symbol_cache: RefCell>, + + /// Map of legacy import member '_head_*' symbols to the associated + /// library names. + legacy_imports: RefCell>, + + /// The archive file data. + archive_data: &'a [u8], +} + +impl<'a> LinkArchive<'a> { + /// Parses the data. + pub fn parse(data: &'a [u8]) -> Result, LinkArchiveParseError> { + let archive_file = ArchiveFile::parse(data)?; + + if archive_file.is_thin() { + return Err(LinkArchiveParseError::ThinArchive); + } + + let symbols = archive_file + .symbols()? + .ok_or(LinkArchiveParseError::NoSymbolMap)?; + let symbol_count = symbols + .size_hint() + .1 + .unwrap_or_else(|| symbols.clone().count()); + + Ok(Self { + archive_file, + symbol_cache: RefCell::new(CachedSymbolMap { + cache: HashMap::with_capacity(symbol_count), + iter: Some(symbols), + }), + legacy_imports: RefCell::new(BTreeMap::new()), + archive_data: data, + }) + } + + pub fn extract_symbol( + &self, + symbol: &'a str, + ) -> Result, ExtractMemberError> { + let extracted = self.extract_archive_member(symbol)?; + let member_name = std::str::from_utf8(extracted.name()) + .map_err(|e| ExtractMemberError::ArchiveParse(ArchiveParseError::MemberName(e)))?; + + self.parse_member(&extracted, member_name) + .map_err(ExtractMemberError::MemberParse) + } + + fn parse_member( + &self, + member: &ArchiveMember<'a>, + member_name: &'a str, + ) -> Result, MemberParseError> { + let member_data = member + .data(self.archive_data) + .map_err(|e| MemberParseError::new(PathBuf::from(member_name), e))?; + + let member_path = Path::new(member_name); + + if member_data + .get(..2) + .is_some_and(|magic| magic == IMAGE_FILE_MACHINE_UNKNOWN.to_le_bytes()) + { + Ok(ExtractedMember { + path: member_path, + contents: ExtractedMemberContents::Import( + ImportFile::parse(member_data) + .map_err(|e| MemberParseError::new(member_path, e))? + .try_into() + .map_err(|e| MemberParseError::new(member_path, e))?, + ), + }) + } else { + let coff = CoffFile::<&[u8]>::parse(member_data) + .map_err(|e| MemberParseError::new(member_path, e))?; + + match self.parse_legacy_import_member(member_name, &coff) { + Ok(import) => Ok(ExtractedMember::new(member_path, import)), + Err(e) + if matches!( + e.kind, + MemberParseErrorKind::LegacyImportLibrarySymbolMember( + LegacyImportSymbolMemberParseError::Invalid + ) + ) => + { + Ok(ExtractedMember::new(member_path, coff)) + } + Err(e) => Err(e), + } + } + } + + fn parse_legacy_import_member( + &self, + member_name: &str, + coff: &CoffFile<'a>, + ) -> Result, MemberParseError> { + let member_path = Path::new(member_name); + + let symbol_member = LegacyImportSymbolMember::parse(coff) + .map_err(|e| MemberParseError::new(member_path, e))?; + + let mut imports_cache = self.legacy_imports.borrow_mut(); + let dll = match imports_cache.entry(symbol_member.head_symbol) { + std::collections::btree_map::Entry::Occupied(dll_entry) => *dll_entry.get(), + std::collections::btree_map::Entry::Vacant(dll_entry) => { + // Get the head COFF for this symbol import member + let head_coff_member = self + .extract_archive_member(symbol_member.head_symbol) + .map_err(|_| { + MemberParseError::new( + member_path, + MemberParseErrorKind::LegacyImportLibraryMissingSymbol( + symbol_member.head_symbol.to_string(), + ), + ) + })?; + + let head_coff_data = head_coff_member.data(self.archive_data).map_err(|_| { + MemberParseError::new( + member_path, + MemberParseErrorKind::LegacyImportLibraryMissingSymbol( + symbol_member.head_symbol.to_string(), + ), + ) + })?; + + let head_coff = CoffFile::<&[u8]>::parse(head_coff_data).map_err(|e| { + let path = std::str::from_utf8(head_coff_member.name()).unwrap_or(member_name); + MemberParseError::new(Path::new(path), e) + })?; + + let legacy_head_member = + LegacyImportHeadMember::parse(&head_coff).map_err(|e| { + let path = + std::str::from_utf8(head_coff_member.name()).unwrap_or(member_name); + MemberParseError::new(Path::new(path), e) + })?; + + // Get the tail COFF for the head member. + let tail_coff_member = self + .extract_archive_member(legacy_head_member.tail_symbol) + .map_err(|_| { + let path = + std::str::from_utf8(head_coff_member.name()).unwrap_or(member_name); + MemberParseError::new( + Path::new(path), + MemberParseErrorKind::LegacyImportLibraryMissingSymbol( + legacy_head_member.tail_symbol.to_string(), + ), + ) + })?; + + let tail_coff_data = tail_coff_member.data(self.archive_data).map_err(|_| { + let path = std::str::from_utf8(tail_coff_member.name()).unwrap_or(member_name); + MemberParseError::new( + Path::new(path), + MemberParseErrorKind::LegacyImportLibraryMissingSymbol( + symbol_member.head_symbol.to_string(), + ), + ) + })?; + + let tail_coff = CoffFile::<&[u8]>::parse(tail_coff_data).map_err(|e| { + let path = std::str::from_utf8(tail_coff_member.name()).unwrap_or(member_name); + MemberParseError::new(Path::new(path), e) + })?; + + let legacy_tail_member = + LegacyImportTailMember::parse(&tail_coff).map_err(|e| { + let path = + std::str::from_utf8(tail_coff_member.name()).unwrap_or(member_name); + MemberParseError::new(Path::new(path), e) + })?; + + // Store the mapping from the '_head' symbol found + // in the symbol COFF to the DLL name found in the + // '_iname' tail COFF. + dll_entry.insert(legacy_tail_member.dll) + } + }; + + Ok(ImportMember { + architecture: coff.architecture(), + symbol: symbol_member.public_symbol, + dll, + import: symbol_member.import_name, + typ: symbol_member.typ, + }) + } + + fn extract_archive_member( + &self, + symbol: &'a str, + ) -> Result, ExtractMemberError> { + let mut symbol_map = self.symbol_cache.borrow_mut(); + let member_idx = symbol_map + .find_symbol(symbol) + .ok_or(ExtractMemberError::NotFound)?; + + self.archive_file + .member(member_idx) + .map_err(|e| ExtractMemberError::ArchiveParse(ArchiveParseError::Object(e))) + } +} + +impl<'a> ApiSymbolSource<'a> for PathedItem<&Path, LinkArchive<'a>> { + fn api_path(&self) -> &Path { + self.path() + } + + fn extract_api_symbol(&self, symbol: &'a str) -> Result, ApiSymbolError> { + self.deref().extract_api_symbol(symbol) + } +} + +impl<'a> ApiSymbolSource<'a> for LinkArchive<'a> { + fn extract_api_symbol(&self, symbol: &'a str) -> Result, ApiSymbolError> { + let member = match self.extract_archive_member(symbol) { + Ok(member) => member, + Err(e) => return Err(e.into()), + }; + + let member_name = std::str::from_utf8(member.name()) + .map_err(|e| ApiSymbolError::ArchiveParse(ArchiveParseError::MemberName(e)))?; + + let member_path = Path::new(member_name); + + let member_data = member + .data(self.archive_data) + .map_err(|e| ApiSymbolError::MemberParse(MemberParseError::new(member_path, e)))?; + + if member_data + .get(..2) + .is_some_and(|magic| magic == IMAGE_FILE_MACHINE_UNKNOWN.to_le_bytes()) + { + Ok(ImportFile::parse(member_data) + .map_err(|e| ApiSymbolError::MemberParse(MemberParseError::new(member_path, e)))? + .try_into() + .map_err(|e| ApiSymbolError::MemberParse(MemberParseError::new(member_path, e)))?) + } else { + let coff = CoffFile::<&[u8]>::parse(member_data) + .map_err(|e| ApiSymbolError::MemberParse(MemberParseError::new(member_path, e)))?; + + match self.parse_legacy_import_member(member_name, &coff) { + Ok(import) => Ok(import), + Err(e) + if matches!( + e.kind, + MemberParseErrorKind::LegacyImportLibrarySymbolMember( + LegacyImportSymbolMemberParseError::Invalid + ) + ) => + { + Err(ApiSymbolError::ImportMember) + } + Err(e) => Err(ApiSymbolError::MemberParse(e)), + } + } + } +} diff --git a/src/linkobject/import.rs b/src/linkobject/import.rs new file mode 100644 index 0000000..8b1af8e --- /dev/null +++ b/src/linkobject/import.rs @@ -0,0 +1,95 @@ +use std::str::Utf8Error; + +use object::Architecture; + +#[derive(Debug, thiserror::Error)] +pub enum TryFromImportFileError { + #[error("symbol field value could not be parsed: {0}")] + Symbol(Utf8Error), + + #[error("dll field value could not be parsed: {0}")] + Dll(Utf8Error), + + #[error("import field value could not be parsed: {0}")] + ImportName(Utf8Error), +} + +/// An exported name from a DLL. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ImportName<'a> { + /// The symbol is exported by the ordinal value. + Ordinal(u16), + + /// The symbol is exported by the symbol name. + Name(&'a str), +} + +impl<'a> TryFrom> for ImportName<'a> { + type Error = Utf8Error; + + fn try_from(value: object::read::coff::ImportName<'a>) -> Result { + Ok(match value { + object::coff::ImportName::Ordinal(o) => Self::Ordinal(o), + object::coff::ImportName::Name(name) => Self::Name(std::str::from_utf8(name)?), + }) + } +} + +/// The type of symbol being imported. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ImportType { + /// The symbol is for executable code. + Code, + + /// The symbol is for misc data. + Data, + + /// The symbol is a constant value. + Const, +} + +impl From for ImportType { + fn from(value: object::read::coff::ImportType) -> Self { + match value { + object::coff::ImportType::Code => Self::Const, + object::coff::ImportType::Data => Self::Data, + object::coff::ImportType::Const => Self::Const, + } + } +} + +/// A short import COFF member from import libraries. +pub struct ImportMember<'a> { + /// The architecture for the import. + pub(crate) architecture: Architecture, + + /// The public symbol name. + pub(crate) symbol: &'a str, + + /// The name of the DLL the symbol is from. + pub(crate) dll: &'a str, + + /// The name exported from the DLL. + pub(crate) import: ImportName<'a>, + + /// The type of import. + #[allow(unused)] + pub(crate) typ: ImportType, +} + +impl<'a> TryFrom> for ImportMember<'a> { + type Error = TryFromImportFileError; + + fn try_from(value: object::read::coff::ImportFile<'a>) -> Result { + Ok(Self { + architecture: value.architecture(), + symbol: std::str::from_utf8(value.symbol()).map_err(TryFromImportFileError::Symbol)?, + dll: std::str::from_utf8(value.dll()).map_err(TryFromImportFileError::Dll)?, + import: value + .import() + .try_into() + .map_err(TryFromImportFileError::ImportName)?, + typ: value.import_type().into(), + }) + } +} diff --git a/src/linkobject/mod.rs b/src/linkobject/mod.rs new file mode 100644 index 0000000..0397ed7 --- /dev/null +++ b/src/linkobject/mod.rs @@ -0,0 +1,2 @@ +pub mod archive; +pub mod import; diff --git a/src/pathed_item.rs b/src/pathed_item.rs new file mode 100644 index 0000000..ea3b474 --- /dev/null +++ b/src/pathed_item.rs @@ -0,0 +1,46 @@ +use std::path::Path; + +/// An item with an associated path. +pub struct PathedItem, T> { + path: P, + item: T, +} + +impl, T> PathedItem { + /// Creates a new [`PathedItem`] with the specified values. + pub fn new(path: P, item: T) -> PathedItem { + Self { path, item } + } + + /// Returns the [`Path`] associated with the item. + pub fn path(&self) -> &P { + &self.path + } + + /// Returns a mutable reference to the [`Path`] associated with the item. + pub fn path_mut(&mut self) -> &mut P { + &mut self.path + } + + /// Converts the item into a `Box`. + pub fn into_boxed_item(self) -> PathedItem> { + PathedItem { + path: self.path, + item: Box::new(self.item), + } + } +} + +impl, T> std::ops::Deref for PathedItem { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.item + } +} + +impl, T> std::ops::DerefMut for PathedItem { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.item + } +} diff --git a/tests/bss/commons.yaml b/tests/bss/commons.yaml new file mode 100644 index 0000000..df228da --- /dev/null +++ b/tests/bss/commons.yaml @@ -0,0 +1,27 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_RELOCS_STRIPPED ] +sections: [] +symbols: + # This symbol should cause a .bss section to be created + - Name: common_symbol + Value: 4 + SectionNumber: 0 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_RELOCS_STRIPPED ] +sections: [] +symbols: + # This symbol should cause a .bss section to be created + - Name: other_common + Value: 8 + SectionNumber: 0 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_EXTERNAL diff --git a/tests/bss/merged.yaml b/tests/bss/merged.yaml new file mode 100644 index 0000000..7f752ee --- /dev/null +++ b/tests/bss/merged.yaml @@ -0,0 +1,40 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_RELOCS_STRIPPED ] +sections: + - Name: .data + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 16 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 + - Name: .bss + Characteristics: [ IMAGE_SCN_CNT_UNINITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 16 + SectionData: '' + SizeOfRawData: 16 +symbols: + - Name: .data + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .bss + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 diff --git a/tests/bss/mod.rs b/tests/bss/mod.rs new file mode 100644 index 0000000..7082781 --- /dev/null +++ b/tests/bss/mod.rs @@ -0,0 +1,96 @@ +use crate::{link_yaml, setup_linker}; +use boflink::linker::LinkerTargetArch; +use object::{Object, ObjectSection, ObjectSymbol, coff::CoffFile}; + +#[test] +fn resized() { + let linked = link_yaml!("resized.yaml", LinkerTargetArch::Amd64); + let parsed: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + let bss_section = parsed + .section_by_name(".bss") + .expect("Could not find .bss section"); + + assert_eq!( + bss_section + .coff_section() + .size_of_raw_data + .get(object::LittleEndian), + 48, + ".bss section size should be 64" + ); +} + +#[test] +fn common_symbols() { + let linked = link_yaml!("commons.yaml", LinkerTargetArch::Amd64); + let coff: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + + const TEST_SYMBOLS: [(&str, u32); 2] = [("common_symbol", 0), ("other_common", 8)]; + + for (symbol_name, symbol_value) in TEST_SYMBOLS { + let symbol = coff + .symbol_by_name(symbol_name) + .unwrap_or_else(|| panic!("Could not find symbol '{symbol_name}'")); + + let section_idx = symbol + .section_index() + .unwrap_or_else(|| panic!("Could not get section index for symbol '{symbol_name}'")); + + let section = coff + .section_by_index(section_idx) + .unwrap_or_else(|e| panic!("Could not get section '{symbol_name}' is defined in: {e}")); + + let section_name = section.name().expect("Could not get section name"); + assert_eq!( + section_name, ".bss", + "'{symbol_name}' is not defined in the .bss section" + ); + + let value = symbol.coff_symbol().value.get(object::LittleEndian); + assert_eq!( + value, symbol_value, + "'{symbol_name}' should be defined at address {symbol_value}" + ); + } +} + +#[test] +fn merged_bss_data() { + let linked = setup_linker!("merged.yaml", LinkerTargetArch::Amd64) + .merge_bss(true) + .build() + .link() + .expect("Could not link files"); + + let parsed: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + + assert!( + parsed.section_by_name(".bss").is_none_or(|section| section + .coff_section() + .size_of_raw_data + .get(object::LittleEndian) + == 0), + "Output COFF should have an empty .bss section or none at all" + ); + + let data_section = parsed + .section_by_name(".data") + .expect("Could not find .data section"); + assert_eq!( + data_section + .coff_section() + .size_of_raw_data + .get(object::LittleEndian), + 32, + ".data section size is not correct" + ); + + let data_section_data = data_section + .data() + .expect("Could not get .data section data"); + assert_eq!( + data_section_data.len(), + 32, + ".data section should have 32 bytes of initialized data" + ); +} diff --git a/tests/bss/resized.yaml b/tests/bss/resized.yaml new file mode 100644 index 0000000..7b5ed74 --- /dev/null +++ b/tests/bss/resized.yaml @@ -0,0 +1,47 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_RELOCS_STRIPPED ] +sections: + - Name: .bss + Characteristics: [ IMAGE_SCN_CNT_UNINITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 16 + SectionData: '' + SizeOfRawData: 32 +symbols: + - Name: .bss + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 32 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_RELOCS_STRIPPED ] +sections: + - Name: .bss + Characteristics: [ IMAGE_SCN_CNT_UNINITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 16 + SectionData: '' + SizeOfRawData: 16 +symbols: + - Name: .bss + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 diff --git a/tests/comdats/any.yaml b/tests/comdats/any.yaml new file mode 100644 index 0000000..8d13895 --- /dev/null +++ b/tests/comdats/any.yaml @@ -0,0 +1,98 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: '4883EC28488D150000000031C9E800000000904883C428C3' + SizeOfRawData: 24 + Relocations: + - VirtualAddress: 7 + SymbolName: '??_C@_0M@KPLPPDAC@Hello?5World?$AA@' + Type: IMAGE_REL_AMD64_REL32 + - VirtualAddress: 14 + SymbolName: BeaconPrintf + Type: IMAGE_REL_AMD64_REL32 + - Name: .rdata + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_LNK_COMDAT, IMAGE_SCN_MEM_READ ] + Alignment: 1 + SectionData: 48656C6C6F20576F726C6400 + SizeOfRawData: 12 +symbols: + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 24 + NumberOfRelocations: 2 + NumberOfLinenumbers: 0 + CheckSum: 3582930156 + Number: 1 + - Name: .rdata + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 12 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 731237010 + Number: 2 + Selection: IMAGE_COMDAT_SELECT_ANY + - Name: '??_C@_0M@KPLPPDAC@Hello?5World?$AA@' + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + - Name: go + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + - Name: BeaconPrintf + Value: 0 + SectionNumber: 0 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ ] +sections: + # Only one of these sections should be kept + - Name: .rdata + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_LNK_COMDAT, IMAGE_SCN_MEM_READ ] + Alignment: 1 + SectionData: 48656C6C6F20576F726C6400 + SizeOfRawData: 12 +symbols: + - Name: .rdata + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 12 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 731237010 + Number: 2 + Selection: IMAGE_COMDAT_SELECT_ANY + - Name: '??_C@_0M@KPLPPDAC@Hello?5World?$AA@' + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_EXTERNAL diff --git a/tests/comdats/associative.yaml b/tests/comdats/associative.yaml new file mode 100644 index 0000000..1580efa --- /dev/null +++ b/tests/comdats/associative.yaml @@ -0,0 +1,109 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ ] +sections: + - Name: '.root$discard' + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_LNK_COMDAT, IMAGE_SCN_MEM_READ ] + Alignment: 1 + SectionData: '6469736361726400' + SizeOfRawData: 8 + - Name: '.assoc$discard' + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_LNK_COMDAT, IMAGE_SCN_MEM_READ ] + Alignment: 1 + SectionData: '6469736361726400' + SizeOfRawData: 8 +symbols: + - Name: '.root$discard' + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 8 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 3436533630 + Number: 1 + Selection: IMAGE_COMDAT_SELECT_LARGEST + - Name: root_comdat + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + - Name: '.assoc$discard' + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 8 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 3436533630 + Number: 1 + Selection: IMAGE_COMDAT_SELECT_ASSOCIATIVE + - Name: discarded + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ ] +sections: + - Name: '.root$keep' + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_LNK_COMDAT, IMAGE_SCN_MEM_READ ] + Alignment: 1 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 + - Name: '.assoc$keep' + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_LNK_COMDAT, IMAGE_SCN_MEM_READ ] + Alignment: 1 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 +symbols: + - Name: '.root$keep' + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 1 + Selection: IMAGE_COMDAT_SELECT_LARGEST + - Name: root_comdat + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + - Name: '.assoc$keep' + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 2 + Selection: IMAGE_COMDAT_SELECT_ASSOCIATIVE + - Name: discarded + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC diff --git a/tests/comdats/mod.rs b/tests/comdats/mod.rs new file mode 100644 index 0000000..44c678f --- /dev/null +++ b/tests/comdats/mod.rs @@ -0,0 +1,76 @@ +use boflink::linker::LinkerTargetArch; +use object::{Object, ObjectSection, coff::CoffFile, pe::IMAGE_SCN_LNK_COMDAT}; + +use crate::link_yaml; + +#[test] +fn any() { + let linked = link_yaml!("any.yaml", LinkerTargetArch::Amd64); + + let coff: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + + let rdata_section = coff + .section_by_name(".rdata") + .expect("Could not find .rdata section"); + + let flags = rdata_section + .coff_section() + .characteristics + .get(object::LittleEndian); + assert_eq!( + flags & IMAGE_SCN_LNK_COMDAT, + 0, + "IMAGE_SCN_LNK_COMDAT flag should have been removed from the section characteristics" + ); + + let rdata_size = rdata_section + .coff_section() + .size_of_raw_data + .get(object::LittleEndian); + + assert_eq!( + rdata_size, 12, + ".rdata section size should be 12. Section was not deduplicated" + ); +} + +#[test] +fn associative() { + let linked = link_yaml!("associative.yaml", LinkerTargetArch::Amd64); + + let coff: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + + let root_section = coff + .section_by_name(".root") + .expect("Could not find .root section"); + + let root_data = root_section + .data() + .expect("Could not get .root section data"); + + assert_eq!(root_data.len(), 16, ".root section data should be 16 bytes"); + + assert!( + root_data.iter().all(|b| *b == 0), + ".root section data should be all zeros" + ); + + let assoc_section = coff + .section_by_name(".assoc") + .expect("Could not find .assoc section"); + + let assoc_data = assoc_section + .data() + .expect("Could not get .assoc section data"); + + assert_eq!( + assoc_data.len(), + 16, + ".assoc section data should be 16 bytes" + ); + + assert!( + assoc_data.iter().all(|b| *b == 0), + ".assoc section data should be all zeros" + ); +} diff --git a/tests/compilers/clang/amd64_gnuabi_empty.yaml b/tests/compilers/clang/amd64_gnuabi_empty.yaml new file mode 100644 index 0000000..8b01811 --- /dev/null +++ b/tests/compilers/clang/amd64_gnuabi_empty.yaml @@ -0,0 +1,91 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: C3 + SizeOfRawData: 1 + - Name: .data + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 4 + SectionData: '' + - Name: .bss + Characteristics: [ IMAGE_SCN_CNT_UNINITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 4 + SectionData: '' + - Name: .llvm_addrsig + Characteristics: [ IMAGE_SCN_LNK_REMOVE ] + Alignment: 1 + SectionData: '' +symbols: + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 1 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 40735498 + Number: 1 + - Name: .data + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 2 + - Name: .bss + Value: 0 + SectionNumber: 3 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 3 + - Name: .llvm_addrsig + Value: 0 + SectionNumber: 4 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 4 + - Name: '@feat.00' + Value: 0 + SectionNumber: -1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + - Name: go + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + - Name: .file + Value: 0 + SectionNumber: -2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_FILE + File: empty.c +... diff --git a/tests/compilers/clang/i386_gnuabi_empty.yaml b/tests/compilers/clang/i386_gnuabi_empty.yaml new file mode 100644 index 0000000..f4fd6ab --- /dev/null +++ b/tests/compilers/clang/i386_gnuabi_empty.yaml @@ -0,0 +1,91 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_I386 + Characteristics: [ ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: 5589E55DC3 + SizeOfRawData: 5 + - Name: .data + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 4 + SectionData: '' + - Name: .bss + Characteristics: [ IMAGE_SCN_CNT_UNINITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 4 + SectionData: '' + - Name: .llvm_addrsig + Characteristics: [ IMAGE_SCN_LNK_REMOVE ] + Alignment: 1 + SectionData: '' +symbols: + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 5 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 3270712146 + Number: 1 + - Name: .data + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 2 + - Name: .bss + Value: 0 + SectionNumber: 3 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 3 + - Name: .llvm_addrsig + Value: 0 + SectionNumber: 4 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 4 + - Name: '@feat.00' + Value: 1 + SectionNumber: -1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + - Name: _go + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + - Name: .file + Value: 0 + SectionNumber: -2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_FILE + File: empty.c +... diff --git a/tests/compilers/clang/mod.rs b/tests/compilers/clang/mod.rs new file mode 100644 index 0000000..f433e2f --- /dev/null +++ b/tests/compilers/clang/mod.rs @@ -0,0 +1,26 @@ +use boflink::linker::LinkerTargetArch; +use object::{Object, coff::CoffFile}; + +use crate::link_yaml; + +#[test] +fn amd64_gnuabi_empty() { + let linked = link_yaml!("amd64_gnuabi_empty.yaml", LinkerTargetArch::Amd64); + + let parsed: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + assert!( + parsed.symbol_by_name("go").is_some(), + "Could not find go symbol in linked COFF" + ); +} + +#[test] +fn i386_gnuabi_empty() { + let linked = link_yaml!("i386_gnuabi_empty.yaml", LinkerTargetArch::I386); + + let parsed: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + assert!( + parsed.symbol_by_name("_go").is_some(), + "Could not find _go symbol in linked COFF" + ); +} diff --git a/tests/compilers/mingw/amd64_empty.yaml b/tests/compilers/mingw/amd64_empty.yaml new file mode 100644 index 0000000..d2963c0 --- /dev/null +++ b/tests/compilers/mingw/amd64_empty.yaml @@ -0,0 +1,135 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_LINE_NUMS_STRIPPED ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: 554889E5905DC3909090909090909090 + SizeOfRawData: 16 + - Name: .data + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 16 + SectionData: '' + - Name: .bss + Characteristics: [ IMAGE_SCN_CNT_UNINITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 16 + SectionData: '' + - Name: .xdata + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 4 + SectionData: '0104020504030150' + SizeOfRawData: 8 + - Name: .pdata + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 4 + SectionData: '000000000700000000000000' + SizeOfRawData: 12 + Relocations: + - VirtualAddress: 0 + SymbolName: .text + Type: IMAGE_REL_AMD64_ADDR32NB + - VirtualAddress: 4 + SymbolName: .text + Type: IMAGE_REL_AMD64_ADDR32NB + - VirtualAddress: 8 + SymbolName: .xdata + Type: IMAGE_REL_AMD64_ADDR32NB + - Name: '.rdata$zzz' + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: 4743433A2028474E55292031342E322E3120323032343038303120284665646F7261204D696E47572031342E322E312D332E6663343129000000000000000000 + SizeOfRawData: 64 +symbols: + - Name: .file + Value: 0 + SectionNumber: -2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_FILE + File: empty.c + - Name: go + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + FunctionDefinition: + TagIndex: 0 + TotalSize: 0 + PointerToLinenumber: 0 + PointerToNextFunction: 0 + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 7 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .data + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .bss + Value: 0 + SectionNumber: 3 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .xdata + Value: 0 + SectionNumber: 4 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 8 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .pdata + Value: 0 + SectionNumber: 5 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 12 + NumberOfRelocations: 3 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: '.rdata$zzz' + Value: 0 + SectionNumber: 6 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 56 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 +... diff --git a/tests/compilers/mingw/i386_empty.yaml b/tests/compilers/mingw/i386_empty.yaml new file mode 100644 index 0000000..cce3ade --- /dev/null +++ b/tests/compilers/mingw/i386_empty.yaml @@ -0,0 +1,112 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_I386 + Characteristics: [ IMAGE_FILE_LINE_NUMS_STRIPPED, IMAGE_FILE_32BIT_MACHINE ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 4 + SectionData: 5589E5905DC39090 + SizeOfRawData: 8 + - Name: .data + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 4 + SectionData: '' + - Name: .bss + Characteristics: [ IMAGE_SCN_CNT_UNINITIALIZED_DATA, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE ] + Alignment: 4 + SectionData: '' + - Name: '.rdata$zzz' + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 4 + SectionData: 4743433A2028474E55292031342E322E3120323032343038303120284665646F7261204D696E47572031342E322E312D332E666334312900 + SizeOfRawData: 56 + - Name: .eh_frame + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 4 + SectionData: 1400000000000000017A5200017C08011B0C0404880100001C0000001C000000040000000600000000410E088502420D0542C50C04040000 + SizeOfRawData: 56 + Relocations: + - VirtualAddress: 32 + SymbolName: .text + Type: IMAGE_REL_I386_REL32 +symbols: + - Name: .file + Value: 0 + SectionNumber: -2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_FILE + File: empty.c + - Name: _go + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + FunctionDefinition: + TagIndex: 0 + TotalSize: 0 + PointerToLinenumber: 0 + PointerToNextFunction: 0 + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 6 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .data + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .bss + Value: 0 + SectionNumber: 3 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 0 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: '.rdata$zzz' + Value: 0 + SectionNumber: 4 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 56 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .eh_frame + Value: 0 + SectionNumber: 5 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 56 + NumberOfRelocations: 1 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 +... diff --git a/tests/compilers/mingw/mod.rs b/tests/compilers/mingw/mod.rs new file mode 100644 index 0000000..83acb4b --- /dev/null +++ b/tests/compilers/mingw/mod.rs @@ -0,0 +1,26 @@ +use boflink::linker::LinkerTargetArch; +use object::{Object, coff::CoffFile}; + +use crate::link_yaml; + +#[test] +fn amd64_empty() { + let linked = link_yaml!("amd64_empty.yaml", LinkerTargetArch::Amd64); + + let parsed: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + assert!( + parsed.symbol_by_name("go").is_some(), + "Could not find go symbol in linked COFF" + ); +} + +#[test] +fn i386_empty() { + let linked = link_yaml!("i386_empty.yaml", LinkerTargetArch::I386); + + let parsed: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + assert!( + parsed.symbol_by_name("_go").is_some(), + "Could not find _go symbol in linked COFF" + ); +} diff --git a/tests/compilers/mod.rs b/tests/compilers/mod.rs new file mode 100644 index 0000000..277cec6 --- /dev/null +++ b/tests/compilers/mod.rs @@ -0,0 +1,3 @@ +mod clang; +mod mingw; +mod msvc; diff --git a/tests/compilers/msvc/amd64_empty.yaml b/tests/compilers/msvc/amd64_empty.yaml new file mode 100644 index 0000000..b2945c5 --- /dev/null +++ b/tests/compilers/msvc/amd64_empty.yaml @@ -0,0 +1,101 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ ] +sections: + - Name: '.debug$S' + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_DISCARDABLE, IMAGE_SCN_MEM_READ ] + Alignment: 1 + SectionData: 04000000F10000008A0000004C00011100000000433A626F666C696E6B5C74657374735C736F75726365735C656D7074792E6F626A003A003C1100420000D00013002B00FA87000013002B00FA8700004D6963726F736F667420285229204F7074696D697A696E6720436F6D70696C6572000000 + Subsections: + - !Symbols + Records: + - Kind: S_OBJNAME + ObjNameSym: + Signature: 0 + ObjectName: 'C:\boflink\tests\sources\empty.obj' + - Kind: S_COMPILE3 + Compile3Sym: + Flags: [ NoDbgInfo, HotPatch ] + Machine: X64 + FrontendMajor: 19 + FrontendMinor: 43 + FrontendBuild: 34810 + FrontendQFE: 0 + BackendMajor: 19 + BackendMinor: 43 + BackendBuild: 34810 + BackendQFE: 0 + Version: 'Microsoft (R) Optimizing Compiler' + SizeOfRawData: 116 + - Name: '.text$mn' + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: C20000 + SizeOfRawData: 3 + - Name: .chks64 + Characteristics: [ IMAGE_SCN_LNK_INFO, IMAGE_SCN_LNK_REMOVE ] + SectionData: 23076615271ABF1A3F3249567ABA655AB6CBECECF1B07E740000000000000000 + SizeOfRawData: 32 +symbols: + - Name: '@comp.id' + Value: 17074170 + SectionNumber: -1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + - Name: '@feat.00' + Value: 2147549328 + SectionNumber: -1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + - Name: '@vol.md' + Value: 3 + SectionNumber: -1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + - Name: '.debug$S' + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 152 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: '.text$mn' + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 3 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 2452308526 + Number: 0 + - Name: go + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + - Name: .chks64 + Value: 0 + SectionNumber: 3 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 32 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 +... diff --git a/tests/compilers/msvc/i386_empty.yaml b/tests/compilers/msvc/i386_empty.yaml new file mode 100644 index 0000000..38bb399 --- /dev/null +++ b/tests/compilers/msvc/i386_empty.yaml @@ -0,0 +1,101 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_I386 + Characteristics: [ ] +sections: + - Name: '.debug$S' + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_DISCARDABLE, IMAGE_SCN_MEM_READ ] + Alignment: 1 + SectionData: 04000000F10000008A0000004C00011100000000433A626F666C696E6B5C74657374735C736F75726365735C656D7074792E6F626A003A003C1100020000070013002B00FA87000013002B00FA8700004D6963726F736F667420285229204F7074696D697A696E6720436F6D70696C6572000000 + Subsections: + - !Symbols + Records: + - Kind: S_OBJNAME + ObjNameSym: + Signature: 0 + ObjectName: 'C:\boflink\tests\sources\empty.obj' + - Kind: S_COMPILE3 + Compile3Sym: + Flags: [ NoDbgInfo ] + Machine: Pentium3 + FrontendMajor: 19 + FrontendMinor: 43 + FrontendBuild: 34810 + FrontendQFE: 0 + BackendMajor: 19 + BackendMinor: 43 + BackendBuild: 34810 + BackendQFE: 0 + Version: 'Microsoft (R) Optimizing Compiler' + SizeOfRawData: 116 + - Name: '.text$mn' + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: 558BEC5DC3 + SizeOfRawData: 5 + - Name: .chks64 + Characteristics: [ IMAGE_SCN_LNK_INFO, IMAGE_SCN_LNK_REMOVE ] + SectionData: 23076615271ABF1A053150A3D353E00D7178BC05F1E5707E0000000000000000 + SizeOfRawData: 32 +symbols: + - Name: '@comp.id' + Value: 17074170 + SectionNumber: -1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + - Name: '@feat.00' + Value: 2147549329 + SectionNumber: -1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + - Name: '@vol.md' + Value: 3 + SectionNumber: -1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + - Name: '.debug$S' + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 152 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: '.text$mn' + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 5 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 1730930774 + Number: 0 + - Name: _go + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + - Name: .chks64 + Value: 0 + SectionNumber: 3 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 32 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 +... diff --git a/tests/compilers/msvc/mod.rs b/tests/compilers/msvc/mod.rs new file mode 100644 index 0000000..83acb4b --- /dev/null +++ b/tests/compilers/msvc/mod.rs @@ -0,0 +1,26 @@ +use boflink::linker::LinkerTargetArch; +use object::{Object, coff::CoffFile}; + +use crate::link_yaml; + +#[test] +fn amd64_empty() { + let linked = link_yaml!("amd64_empty.yaml", LinkerTargetArch::Amd64); + + let parsed: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + assert!( + parsed.symbol_by_name("go").is_some(), + "Could not find go symbol in linked COFF" + ); +} + +#[test] +fn i386_empty() { + let linked = link_yaml!("i386_empty.yaml", LinkerTargetArch::I386); + + let parsed: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + assert!( + parsed.symbol_by_name("_go").is_some(), + "Could not find _go symbol in linked COFF" + ); +} diff --git a/tests/imports/import_thunks.yaml b/tests/imports/import_thunks.yaml new file mode 100644 index 0000000..c3dcfb7 --- /dev/null +++ b/tests/imports/import_thunks.yaml @@ -0,0 +1,38 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_LINE_NUMS_STRIPPED ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 8 + SectionData: 0000000000000000 + SizeOfRawData: 8 + Relocations: + - VirtualAddress: 2 + SymbolName: import + Type: IMAGE_REL_AMD64_REL32 +symbols: + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 8 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: import + Value: 0 + SectionNumber: 0 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + +--- !IMPORTLIB +Library: LIBRARY +Exports: + - import diff --git a/tests/imports/library_prefix.yaml b/tests/imports/library_prefix.yaml new file mode 100644 index 0000000..1f32f9e --- /dev/null +++ b/tests/imports/library_prefix.yaml @@ -0,0 +1,17 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_RELOCS_STRIPPED ] +sections: [] +symbols: + - Name: __imp_imported_symbol + Value: 0 + SectionNumber: 0 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + +--- !IMPORTLIB +Library: LIBRARY +Exports: + - imported_symbol diff --git a/tests/imports/mod.rs b/tests/imports/mod.rs new file mode 100644 index 0000000..7cf702d --- /dev/null +++ b/tests/imports/mod.rs @@ -0,0 +1,66 @@ +use crate::link_yaml; +use boflink::linker::LinkerTargetArch; +use object::{Object, ObjectSymbol, coff::CoffFile}; + +#[test] +fn library_prefix() { + let linked = link_yaml!("library_prefix.yaml", LinkerTargetArch::Amd64); + let parsed: CoffFile = + CoffFile::parse(linked.as_slice()).expect("Could not parse linked output"); + + assert!( + parsed + .symbol_by_name("__imp_LIBRARY$imported_symbol") + .is_some(), + "Could not find symbol '__imp_LIBRARY$imported_symbol' in linked output" + ); +} + +#[test] +fn import_thunks() { + let linked = link_yaml!("import_thunks.yaml", LinkerTargetArch::Amd64); + let parsed: CoffFile = + CoffFile::parse(linked.as_slice()).expect("Could not parse linked output"); + + let thunk_symbol = parsed + .symbol_by_name("import") + .expect("Could not find symbol 'import'"); + + assert!( + thunk_symbol.is_definition(), + "thunk symbol should be defined" + ); + + let thunk_addr = thunk_symbol.coff_symbol().value.get(object::LittleEndian); + + let text_section = parsed + .section_by_name(".text") + .expect("Could not find .text section"); + + let thunk_reloc = text_section + .coff_relocations() + .unwrap() + .iter() + .next() + .expect(".text section should have a relocation"); + + let reloc_addr = thunk_reloc.virtual_address.get(object::LittleEndian); + + assert_eq!( + thunk_addr + 2, + reloc_addr, + "Thunk relocation address and thunk symbol address do not line up" + ); + + let reloc_target = parsed + .symbol_by_index(thunk_reloc.symbol()) + .expect("Could not get thunk reloc target symbol"); + + let target_name = reloc_target + .name() + .expect("Could not get thunk reloc target name"); + assert_eq!( + target_name, "__imp_LIBRARY$import", + "Thunk relocation target does not point to import symbol" + ); +} diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs new file mode 100644 index 0000000..14b2d06 --- /dev/null +++ b/tests/integration_tests.rs @@ -0,0 +1,6 @@ +mod bss; +mod comdats; +mod compilers; +mod imports; +mod relocations; +mod utils; diff --git a/tests/relocations/defined_symbol_target_no_shift.yaml b/tests/relocations/defined_symbol_target_no_shift.yaml new file mode 100644 index 0000000..3308a80 --- /dev/null +++ b/tests/relocations/defined_symbol_target_no_shift.yaml @@ -0,0 +1,93 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_LINE_NUMS_STRIPPED ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 + - Name: .rdata + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 +symbols: + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .rdata + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_LINE_NUMS_STRIPPED ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 + Relocations: + # This relocation targets a symbol defined in a section and not the section symbol + # itself. The relocation value should be left untouched. + - VirtualAddress: 0 + SymbolName: target_symbol + Type: IMAGE_REL_AMD64_REL32 + - Name: .rdata + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 +symbols: + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .rdata + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: target_symbol + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC diff --git a/tests/relocations/mod.rs b/tests/relocations/mod.rs new file mode 100644 index 0000000..6aba7b3 --- /dev/null +++ b/tests/relocations/mod.rs @@ -0,0 +1,133 @@ +use boflink::linker::LinkerTargetArch; +use object::{Object, ObjectSection, ObjectSymbol, coff::CoffFile}; + +use crate::link_yaml; + +#[test] +fn same_section_flattened() { + let linked = link_yaml!("same_section_flattened.yaml", LinkerTargetArch::Amd64); + let coff: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + + let text_section = coff + .section_by_name(".text") + .expect("Could not find .text section in linked COFF"); + + assert_eq!( + text_section + .coff_section() + .number_of_relocations + .get(object::LittleEndian), + 0, + ".text section header should have 0 for the number of relocations" + ); + + let reloc_count = text_section + .coff_relocations() + .expect("Could not get COFF relocations") + .len(); + assert_eq!(reloc_count, 0, ".text section should have 0 relocations"); + + // Check the relocation to see if it was applied properly so that it points + // to the target symbol + let target_symbol = coff + .symbol_by_name("external_function") + .expect("Could not get external_function symbol"); + + let symbol_addr = target_symbol.coff_symbol().value.get(object::LittleEndian); + + let section_data = text_section + .data() + .expect("Could not get .text section data"); + + let found_reloc_val = u32::from_le_bytes(section_data[2..6].try_into().unwrap()); + let expected_reloc_val = symbol_addr - 2 - 4; + + assert_eq!( + found_reloc_val, expected_reloc_val, + "Flattened relocation value does not point to the target symbol" + ); +} + +#[test] +fn section_target_shifted() { + let linked = link_yaml!("section_target_shifted.yaml", LinkerTargetArch::Amd64); + let coff: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + + let text_section = coff + .section_by_name(".text") + .expect("Could not find .text section in linked COFF"); + + let reloc = text_section + .coff_relocations() + .expect("Could not get .text section relocation") + .iter() + .next() + .expect(".text section should have a relocation"); + + let reloc_addr = reloc.virtual_address.get(object::LittleEndian); + + let section_data = text_section + .data() + .expect("Could not get .text section data"); + + let found_reloc_val = u32::from_le_bytes( + section_data[reloc_addr as usize..reloc_addr as usize + 4] + .try_into() + .unwrap(), + ); + + assert_eq!( + found_reloc_val, 16, + "Relocation value should point to virtual address of shifted section" + ); +} + +#[test] +fn defined_symbol_target_no_shift() { + let linked = link_yaml!( + "defined_symbol_target_no_shift.yaml", + LinkerTargetArch::Amd64 + ); + let coff: CoffFile = CoffFile::parse(linked.as_slice()).expect("Could not parse linked COFF"); + + let text_section = coff + .section_by_name(".text") + .expect("Could not find .text section in linked COFF"); + + let reloc = text_section + .coff_relocations() + .expect("Could not get .text section relocation") + .iter() + .next() + .expect(".text section should have a relocation"); + + let target_symbol = coff + .symbol_by_index(reloc.symbol()) + .expect("Could not get relocation target symbol"); + + let target_name = target_symbol + .name() + .expect("Could not get target symbol name"); + + assert_eq!( + target_name, "target_symbol", + "Relocation target symbol name should be 'target_symbol'" + ); + + let reloc_addr = reloc.virtual_address.get(object::LittleEndian); + + let section_data = text_section + .data() + .expect("Could not get .text section data"); + + let found_reloc_val = u32::from_le_bytes( + section_data[reloc_addr as usize..reloc_addr as usize + 4] + .try_into() + .unwrap(), + ); + + assert_eq!( + found_reloc_val, 0, + "Relocation value should not have shifted" + ); +} diff --git a/tests/relocations/same_section_flattened.yaml b/tests/relocations/same_section_flattened.yaml new file mode 100644 index 0000000..d00f342 --- /dev/null +++ b/tests/relocations/same_section_flattened.yaml @@ -0,0 +1,70 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_LINE_NUMS_STRIPPED ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 + Relocations: + # The relocation here should be applied and removed + # when the other .text section is merged + - VirtualAddress: 2 + SymbolName: external_function + Type: IMAGE_REL_AMD64_REL32 +symbols: + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: external_function + Value: 0 + SectionNumber: 0 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_LINE_NUMS_STRIPPED ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 +symbols: + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: external_function + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_FUNCTION + StorageClass: IMAGE_SYM_CLASS_EXTERNAL + FunctionDefinition: + TagIndex: 0 + TotalSize: 0 + PointerToLinenumber: 0 + PointerToNextFunction: 0 diff --git a/tests/relocations/section_target_shifted.yaml b/tests/relocations/section_target_shifted.yaml new file mode 100644 index 0000000..3e46013 --- /dev/null +++ b/tests/relocations/section_target_shifted.yaml @@ -0,0 +1,88 @@ +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_LINE_NUMS_STRIPPED ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 + - Name: .rdata + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 +symbols: + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .rdata + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + +--- !COFF +header: + Machine: IMAGE_FILE_MACHINE_AMD64 + Characteristics: [ IMAGE_FILE_LINE_NUMS_STRIPPED ] +sections: + - Name: .text + Characteristics: [ IMAGE_SCN_CNT_CODE, IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 + Relocations: + # The address of the .rdata symbol will shift after it has been merged with the + # .rdata section above. The current relocation value needs to be adjusted to account + # for this shift + - VirtualAddress: 0 + SymbolName: .rdata + Type: IMAGE_REL_AMD64_REL32 + - Name: .rdata + Characteristics: [ IMAGE_SCN_CNT_INITIALIZED_DATA, IMAGE_SCN_MEM_READ ] + Alignment: 16 + SectionData: '00000000000000000000000000000000' + SizeOfRawData: 16 +symbols: + - Name: .text + Value: 0 + SectionNumber: 1 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 + - Name: .rdata + Value: 0 + SectionNumber: 2 + SimpleType: IMAGE_SYM_TYPE_NULL + ComplexType: IMAGE_SYM_DTYPE_NULL + StorageClass: IMAGE_SYM_CLASS_STATIC + SectionDefinition: + Length: 16 + NumberOfRelocations: 0 + NumberOfLinenumbers: 0 + CheckSum: 0 + Number: 0 diff --git a/tests/utils/archive_searcher.rs b/tests/utils/archive_searcher.rs new file mode 100644 index 0000000..a9d2ea6 --- /dev/null +++ b/tests/utils/archive_searcher.rs @@ -0,0 +1,30 @@ +use std::{collections::HashMap, path::PathBuf}; + +use boflink::libsearch::{FoundLibrary, LibraryFind, LibsearchError}; + +pub struct MemoryArchiveSearcher { + files: HashMap>, +} + +impl MemoryArchiveSearcher { + pub fn new() -> MemoryArchiveSearcher { + Self { + files: HashMap::new(), + } + } + + pub fn add_library(&mut self, name: impl Into, data: Vec) { + self.files.insert(name.into(), data); + } +} + +impl LibraryFind for MemoryArchiveSearcher { + fn find_library(&self, name: impl AsRef) -> Result { + self.files + .get(name.as_ref()) + .map(|data| FoundLibrary::new(PathBuf::from(name.as_ref()), data.clone())) + .ok_or(boflink::libsearch::LibsearchError::NotFound( + name.as_ref().to_string(), + )) + } +} diff --git a/tests/utils/build.rs b/tests/utils/build.rs new file mode 100644 index 0000000..2ff077c --- /dev/null +++ b/tests/utils/build.rs @@ -0,0 +1,11 @@ +use coffyaml::{coff::CoffYaml, importlib::ImportlibYaml}; +use serde::Deserialize; + +#[derive(Debug, Deserialize)] +pub enum YamlInput { + #[serde(rename = "COFF")] + Coff(CoffYaml), + + #[serde(rename = "IMPORTLIB")] + Importlib(ImportlibYaml), +} diff --git a/tests/utils/macros.rs b/tests/utils/macros.rs new file mode 100644 index 0000000..37da3bc --- /dev/null +++ b/tests/utils/macros.rs @@ -0,0 +1,52 @@ +#[macro_export] +macro_rules! link_yaml { + ($input:literal, $arch:expr) => {{ + const __INPUT_DOC: &str = include_str!($input); + link_yaml!(__INPUT_DOC, $arch) + }}; + + ($input:ident, $arch:expr) => {{ + $crate::setup_linker!($input, $arch) + .build() + .link() + .expect("Could not link files") + }}; +} + +#[macro_export] +macro_rules! setup_linker { + ($input:literal, $arch:expr) => {{ + const __INPUT_DOC: &str = include_str!($input); + setup_linker!(__INPUT_DOC, $arch) + }}; + + ($input:ident, $arch:expr) => {{ + use serde::Deserialize; + let mut __searcher = $crate::utils::archive_searcher::MemoryArchiveSearcher::new(); + let mut __input_libraries = Vec::new(); + let mut __input_coffs = Vec::new(); + + for (idx, document) in serde_yml::Deserializer::from_str($input).enumerate() { + let yaml_input = $crate::utils::build::YamlInput::deserialize(document).unwrap(); + match yaml_input { + $crate::utils::build::YamlInput::Coff(c) => { + __input_coffs.push(boflink::pathed_item::PathedItem::new( + format!("file{}", idx + 1).into(), + c.build().unwrap(), + )); + } + $crate::utils::build::YamlInput::Importlib(c) => { + let library_name = format!("file{}", idx + 1); + __searcher.add_library(library_name.clone(), c.build($arch.into()).unwrap()); + __input_libraries.push(library_name); + } + }; + } + + boflink::linker::LinkerBuilder::new() + .architecture($arch) + .library_searcher(__searcher) + .add_inputs(__input_coffs) + .add_libraries(__input_libraries) + }}; +} diff --git a/tests/utils/mod.rs b/tests/utils/mod.rs new file mode 100644 index 0000000..0e9cd3d --- /dev/null +++ b/tests/utils/mod.rs @@ -0,0 +1,3 @@ +pub mod archive_searcher; +pub mod build; +pub mod macros; diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml new file mode 100644 index 0000000..5160ce3 --- /dev/null +++ b/xtask/Cargo.toml @@ -0,0 +1,45 @@ +[package] +name = "xtask" +version = "0.1.0" +edition = "2024" +publish = false + +[[bin]] +path = "src/main.rs" +name = "xtask-util" + +[dependencies] +flate2 = { version = "1.1.1", optional = true } +hex = { version = "0.4.3", optional = true } +serde = { version = "1.0.219", features = ["derive"], optional = true } +serde_json = { version = "1.0.140", optional = true } +sha2 = { version = "0.10.9", optional = true } +tar = { version = "0.4.44", optional = true } + +[dependencies.zip] +version = "4.0.0" +optional = true +default-features = false +features = ["deflate", "chrono"] + +[dependencies.chrono] +version = "0.4.41" +optional = true +default-features = false +features = ["std"] + +[features] +default = [] +metadata = [ + "dep:serde", + "dep:serde_json" +] +dist = [ + "metadata", + "dep:flate2", + "dep:chrono", + "dep:hex", + "dep:sha2", + "dep:tar", + "dep:zip", +] diff --git a/xtask/src/dist/mod.rs b/xtask/src/dist/mod.rs new file mode 100644 index 0000000..f83cc54 --- /dev/null +++ b/xtask/src/dist/mod.rs @@ -0,0 +1,2 @@ +mod sha256; +pub mod task; diff --git a/xtask/src/dist/sha256.rs b/xtask/src/dist/sha256.rs new file mode 100644 index 0000000..b251c85 --- /dev/null +++ b/xtask/src/dist/sha256.rs @@ -0,0 +1,32 @@ +use std::io::Write; + +use sha2::Digest; + +pub struct Sha256Writer { + hasher: sha2::Sha256, + writer: W, +} + +impl Sha256Writer { + pub fn new(writer: W) -> Sha256Writer { + Self { + hasher: sha2::Sha256::new(), + writer, + } + } + + pub fn finalize(self) -> [u8; 32] { + self.hasher.finalize().into() + } +} + +impl std::io::Write for Sha256Writer { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.hasher.update(buf); + self.writer.write(buf) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.writer.flush() + } +} diff --git a/xtask/src/dist/task.rs b/xtask/src/dist/task.rs new file mode 100644 index 0000000..5b7c164 --- /dev/null +++ b/xtask/src/dist/task.rs @@ -0,0 +1,213 @@ +use std::{ + collections::VecDeque, + error::Error, + io::{BufWriter, Cursor, Read, Write}, + path::PathBuf, + sync::Mutex, + time::UNIX_EPOCH, +}; + +use chrono::Utc; +use flate2::{Compression, write::GzEncoder}; +use sha2::Digest; +use tar::Header; +use zip::{ZipWriter, write::SimpleFileOptions}; + +use crate::utils; + +use super::sha256::Sha256Writer; + +const INSTALL_SCRIPT: &str = r#"#!/bin/sh +/usr/bin/install -m 755 -v boflink ~/.local/bin/boflink +/usr/bin/install -v -D -d ~/.local/libexec/boflink +/usr/bin/ln -svf ~/.local/bin/boflink ~/.local/libexec/boflink/ld +"#; + +pub fn dist() -> Result<(), Box> { + if !std::env::args_os().skip(2).any(|arg| arg == "--skip-tests") { + let _ = utils::shell::run_cargo(["test"]); + } + + let mut build_targets = Vec::new(); + let mut target_flag_found = false; + for arg in std::env::args().skip(2) { + if let Some(target) = arg.strip_prefix("--target=") { + build_targets.extend(target.split(',').map(String::from)); + } else if arg == "-t" || arg == "--target" { + target_flag_found = true; + } else if target_flag_found { + target_flag_found = false; + build_targets.extend(arg.split(',').map(String::from)); + } + } + + if build_targets.is_empty() { + build_targets.push(utils::env::rustc_host()?); + } + + build_targets.sort(); + build_targets.dedup(); + + let version = utils::metadata::package_version("boflink")?; + let workspaceroot = PathBuf::from(utils::metadata::workspace_root()?); + let targetdir = PathBuf::from(utils::metadata::target_directory()?); + let distdir = targetdir.join("dist"); + std::fs::create_dir_all(&distdir)?; + + utils::shell::run_cargo( + ["build", "--release"].into_iter().chain( + build_targets + .iter() + .flat_map(|target| ["--target", target.as_str()]), + ), + )?; + + let thread_count = std::thread::available_parallelism() + .map(|v| v.get()) + .unwrap_or(1) + .min(build_targets.len()); + + let tasks = Mutex::new(VecDeque::from_iter(build_targets)); + + std::thread::scope(|scope| -> Result<(), std::io::Error> { + let mut thrds = VecDeque::with_capacity(thread_count); + + for _ in 0..thread_count { + thrds.push_back(scope.spawn(|| -> Result<(), std::io::Error> { + while let Some(target) = tasks.lock().unwrap().pop_front() { + let builddir = targetdir.join(&target).join("release"); + + let windows_build = target.contains("windows"); + let linux_build = target.contains("linux"); + + let binext = windows_build.then_some("exe"); + let distext = if windows_build { "zip" } else { "tar.gz" }; + + let distname = format!("boflink-v{version}-{target}"); + let distfilename = PathBuf::from(format!("{distname}.{distext}")); + let distsumfilename = PathBuf::from(format!("{distname}.{distext}.sha256")); + let distname = PathBuf::from(distname); + + if linux_build { + let mut tarbuilder = tar::Builder::new(GzEncoder::new( + Sha256Writer::new(BufWriter::new(std::fs::File::create( + distdir.join(&distfilename), + )?)), + Compression::default(), + )); + + tarbuilder.append_path_with_name( + builddir + .join("boflink") + .with_extension(binext.unwrap_or_default()), + distname + .join("boflink") + .with_extension(binext.unwrap_or_default()), + )?; + + tarbuilder.append_path_with_name( + workspaceroot.join("LICENSE"), + distname.join("LICENSE"), + )?; + + let mut header = Header::new_gnu(); + header.set_size(INSTALL_SCRIPT.len() as u64); + header.set_mode(0o755); + header.set_mtime( + std::time::SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(), + ); + tarbuilder.append_data( + &mut header, + distname.join("install"), + INSTALL_SCRIPT.as_bytes(), + )?; + + let sha256sum = tarbuilder.into_inner()?.finish()?.finalize(); + std::fs::write( + distdir.join(&distsumfilename), + format!("{} {}\n", hex::encode(sha256sum), distfilename.display()), + )?; + } else if windows_build { + let mut zipbuilder = ZipWriter::new(Cursor::new(Vec::new())); + + let filepath = builddir + .join("boflink") + .with_extension(binext.unwrap_or_default()); + + let mut file = std::fs::File::open(&filepath)?; + + let modified = + chrono::DateTime::::from(file.metadata()?.modified()?).naive_utc(); + + zipbuilder.start_file_from_path( + distname + .join("boflink") + .with_extension(binext.unwrap_or_default()), + SimpleFileOptions::default() + .unix_permissions(0o755) + .last_modified_time( + zip::DateTime::try_from(modified) + .map_err(std::io::Error::other)?, + ), + )?; + + let mut buffer = Vec::new(); + file.read_to_end(&mut buffer)?; + + zipbuilder.write_all(buffer.as_slice())?; + + buffer.clear(); + + let filepath = workspaceroot.join("LICENSE"); + let mut file = std::fs::File::open(&filepath)?; + + let modified = + chrono::DateTime::::from(file.metadata()?.modified()?).naive_utc(); + + zipbuilder.start_file_from_path( + distname.join("LICENSE"), + SimpleFileOptions::default() + .unix_permissions(0o644) + .last_modified_time( + zip::DateTime::try_from(modified) + .map_err(std::io::Error::other)?, + ), + )?; + + file.read_to_end(&mut buffer)?; + zipbuilder.write_all(buffer.as_slice())?; + + buffer.clear(); + + let result = zipbuilder.finish()?.into_inner(); + let sha256sum: [u8; 32] = sha2::Sha256::digest(result.as_slice()).into(); + + std::fs::write(distdir.join(&distfilename), result.as_slice())?; + + std::fs::write( + distdir.join(&distsumfilename), + format!( + "SHA256 hash of .\\{}:\n{}\nCertUtil: -hashfile command completed successfully.\n", + distfilename.display(), + hex::encode(sha256sum) + ), + )?; + } + } + + Ok(()) + })); + } + + while let Some(thr) = thrds.pop_front() { + thr.join().unwrap()?; + } + + Ok(()) + })?; + + Ok(()) +} diff --git a/xtask/src/main.rs b/xtask/src/main.rs new file mode 100644 index 0000000..8e27c72 --- /dev/null +++ b/xtask/src/main.rs @@ -0,0 +1,43 @@ +use tasks::TASKLIST; + +mod tasks; +mod utils; + +#[cfg(feature = "dist")] +mod dist; + +fn main() { + if let Err(e) = try_main() { + eprintln!("{e}"); + std::process::exit(1); + } +} + +fn try_main() -> Result<(), Box> { + if std::env::args().len() == 1 { + tasks::print_help()?; + return Ok(()); + } + + #[cfg(feature = "dist")] + if std::env::args().nth(1).is_some_and(|task| task == "dist") { + tasks::dist()?; + return Ok(()); + } + + for task in std::env::args().skip(1) { + if !TASKLIST.iter().any(|defined| defined.name == task) { + return Err(format!("unknown task '{task}'").into()); + } + } + + for task in std::env::args().skip(1) { + if let Some(defined) = TASKLIST.iter().find(|defined| defined.name == task) { + (defined.run)()?; + } else { + return Err(format!("unknown task '{task}'").into()); + } + } + + Ok(()) +} diff --git a/xtask/src/tasks.rs b/xtask/src/tasks.rs new file mode 100644 index 0000000..e950f82 --- /dev/null +++ b/xtask/src/tasks.rs @@ -0,0 +1,158 @@ +use std::{error::Error, ffi::OsString}; + +use crate::utils; + +pub struct Task { + pub name: &'static str, + pub help: &'static str, + pub run: fn() -> Result<(), Box>, +} + +pub const TASKLIST: &[Task] = &[ + Task { + name: "install", + help: "Install boflink", + run: install, + }, + Task { + name: "uninstall", + help: "Uninstall boflink", + run: uninstall, + }, + Task { + name: "lint", + help: "Lint with clippy", + run: lint, + }, + Task { + name: "test", + help: "Run tests", + run: test, + }, + Task { + name: "checkfmt", + help: "Check formatting", + run: checkfmt, + }, + Task { + name: "ci", + help: "Run ci workflow", + run: ci, + }, + #[cfg(feature = "dist")] + Task { + name: "dist", + help: "Build a binary release dist archive", + run: dist, + }, + Task { + name: "list", + help: "List tasks", + run: print_help, + }, + Task { + name: "help", + help: "Print help", + run: print_help, + }, +]; + +pub fn install() -> Result<(), Box> { + utils::shell::run_cargo( + ["install", "--path", ".", "--bin", "boflink"] + .into_iter() + .map(OsString::from) + .chain(std::env::args_os().skip(2)), + )?; + + if cfg!(unix) { + let exe_install_path = utils::env::cargo_home()?.join("bin").join("boflink"); + + let libexec = utils::env::install_libexec().join("boflink"); + println!("mkdir -p {}", libexec.display()); + std::fs::create_dir_all(&libexec)?; + + let ldsymlink = libexec.join("ld"); + println!( + "ln -sf {} {}", + exe_install_path.display(), + ldsymlink.display() + ); + + let _ = std::fs::remove_file(&ldsymlink); + + #[cfg(unix)] + std::os::unix::fs::symlink(exe_install_path, ldsymlink)?; + } + + Ok(()) +} + +pub fn uninstall() -> Result<(), Box> { + let _ = utils::shell::run_cargo(["uninstall", "boflink"]); + + if cfg!(unix) { + let libexec = utils::env::install_libexec().join("boflink"); + + let symlink = libexec.join("ld"); + println!("rm -f {}", symlink.display()); + let _ = std::fs::remove_file(symlink); + + println!("rmdir {}", libexec.display()); + let _ = std::fs::remove_dir(libexec); + } + + Ok(()) +} + +pub fn lint() -> Result<(), Box> { + utils::shell::run_cargo([ + "clippy", + "--workspace", + "--all-features", + "--all-targets", + "--", + "-D", + "warnings", + ])?; + Ok(()) +} + +pub fn test() -> Result<(), Box> { + utils::shell::run_cargo(["test", "--workspace"])?; + Ok(()) +} + +pub fn checkfmt() -> Result<(), Box> { + utils::shell::run_cargo(["fmt", "--all", "--check"])?; + Ok(()) +} + +pub fn ci() -> Result<(), Box> { + checkfmt()?; + lint()?; + test()?; + Ok(()) +} + +#[cfg(feature = "dist")] +pub fn dist() -> Result<(), Box> { + crate::dist::task::dist() +} + +pub fn print_help() -> Result<(), Box> { + let descwidth = TASKLIST + .iter() + .max_by_key(|task| task.name.len()) + .map(|task| task.name.len()) + .unwrap_or_default() + + 2; + + println!("tasks:"); + + for task in TASKLIST { + println!(" {: OsString { + static CARGO_EXE: LazyLock = + LazyLock::new(|| std::env::var_os("CARGO").unwrap_or_else(|| "cargo".into())); + + LazyLock::force(&CARGO_EXE).to_os_string() +} + +pub fn rustc() -> OsString { + static RUSTC_EXE: LazyLock = + LazyLock::new(|| std::env::var_os("RUSTC").unwrap_or_else(|| "rustc".into())); + + LazyLock::force(&RUSTC_EXE).to_os_string() +} diff --git a/xtask/src/utils/env/mod.rs b/xtask/src/utils/env/mod.rs new file mode 100644 index 0000000..dcf63e1 --- /dev/null +++ b/xtask/src/utils/env/mod.rs @@ -0,0 +1,78 @@ +use std::{ + error::Error, + path::{Path, PathBuf}, + sync::LazyLock, +}; + +use exe::rustc; + +use super::shell::check_output_projdir; + +pub mod exe; + +pub fn project_root() -> std::io::Result<&'static Path> { + static MANIFEST_DIR: LazyLock> = LazyLock::new(|| { + let path = Path::new(&env!("CARGO_MANIFEST_DIR")).ancestors().nth(1); + + if let Some(path) = path { + if !path.join("Cargo.toml").is_file() { + return None; + } + + if !path.join(".git").is_dir() { + return None; + } + } + + path + }); + + LazyLock::force(&MANIFEST_DIR) + .ok_or_else(|| std::io::Error::other("could not detect project root")) +} + +pub fn cargo_home() -> std::io::Result { + static CARGO_HOME: LazyLock> = LazyLock::new(|| { + std::env::var_os("CARGO_HOME") + .map(PathBuf::from) + .or_else(|| { + if cfg!(unix) { + std::env::var_os("HOME") + } else if cfg!(windows) { + std::env::var_os("USERPROFILE") + } else { + None + } + .map(|home| Path::new(&home).join(".cargo")) + }) + }); + + LazyLock::force(&CARGO_HOME) + .clone() + .ok_or_else(|| std::io::Error::other("could not find CARGO_HOME")) +} + +pub fn install_prefix() -> PathBuf { + static INSTALL_PREFIX: LazyLock = LazyLock::new(|| { + std::env::var_os("PREFIX") + .map(PathBuf::from) + .or_else(|| std::env::var_os("HOME").map(|home| Path::new(&home).join(".local"))) + .unwrap_or_else(|| PathBuf::from("/usr/local")) + }); + + LazyLock::force(&INSTALL_PREFIX).clone() +} + +pub fn install_libexec() -> PathBuf { + install_prefix().join("libexec") +} + +pub fn rustc_host() -> Result> { + let output = check_output_projdir(rustc(), ["--version", "--verbose"])?; + + output + .lines() + .find_map(|line| line.strip_prefix("host: ")) + .map(String::from) + .ok_or_else(|| "could not find rustc host".into()) +} diff --git a/xtask/src/utils/metadata.rs b/xtask/src/utils/metadata.rs new file mode 100644 index 0000000..582cac8 --- /dev/null +++ b/xtask/src/utils/metadata.rs @@ -0,0 +1,87 @@ +use std::sync::LazyLock; + +use serde::Deserialize; + +use super::{env::exe::cargo, shell::check_output_projdir}; + +static CARGO_METADATA: LazyLock> = LazyLock::new(|| { + let output = check_output_projdir(cargo(), ["metadata", "--no-deps", "--format-version=1"])?; + serde_json::from_str::(&output).map_err(std::io::Error::other) +}); + +#[derive(Debug, Deserialize)] +struct CargoMetadata { + packages: Vec, + target_directory: String, + workspace_root: String, +} + +#[derive(Debug, Deserialize)] +#[allow(unused)] +struct CargoMetadataPackage { + name: String, + version: String, + id: String, + license: Option, + license_file: Option, + description: Option, + source: Option, + dependencies: Vec, + targets: Vec, + manifest_path: String, + authors: Vec, +} + +#[derive(Debug, Deserialize)] +#[allow(unused)] +struct CargoMetadataPackageDependency { + name: String, + source: Option, + req: String, + kind: Option, + rename: Option, + optional: bool, + uses_default_features: bool, + features: Vec, +} + +#[derive(Debug, Deserialize)] +#[allow(unused)] +struct CargoMetadataPackageTarget { + kind: Vec, + crate_types: Vec, + name: String, + src_path: String, + edition: String, + doc: bool, + doctest: bool, + test: bool, +} + +pub fn workspace_root() -> std::io::Result { + CARGO_METADATA + .as_ref() + .map_err(std::io::Error::other) + .map(|meta| meta.workspace_root.clone()) +} + +pub fn target_directory() -> std::io::Result { + CARGO_METADATA + .as_ref() + .map_err(std::io::Error::other) + .map(|meta| meta.target_directory.clone()) +} + +pub fn package_version(name: impl AsRef) -> std::io::Result { + CARGO_METADATA + .as_ref() + .map_err(std::io::Error::other) + .and_then(|meta| { + meta.packages + .iter() + .find_map(|package| { + (package.name == name.as_ref()).then(|| package.version.clone()) + }) + .ok_or_else(|| std::io::Error::other("could not find package in metadata")) + }) +} diff --git a/xtask/src/utils/mod.rs b/xtask/src/utils/mod.rs new file mode 100644 index 0000000..258f3a1 --- /dev/null +++ b/xtask/src/utils/mod.rs @@ -0,0 +1,7 @@ +#![allow(unused)] + +pub mod env; +pub mod shell; + +#[cfg(feature = "metadata")] +pub mod metadata; diff --git a/xtask/src/utils/shell.rs b/xtask/src/utils/shell.rs new file mode 100644 index 0000000..93172e2 --- /dev/null +++ b/xtask/src/utils/shell.rs @@ -0,0 +1,84 @@ +use std::{ffi::OsStr, path::Path, process::Command}; + +use crate::utils::env::exe::cargo; + +use super::env::project_root; + +pub fn check_output, S: AsRef>( + path: impl AsRef, + prog: impl AsRef, + args: I, +) -> std::io::Result { + let output = Command::new(prog).current_dir(path).args(args).output()?; + + if !output.status.success() { + Err(std::io::Error::other("command returned non-zero exit code")) + } else { + Ok(std::str::from_utf8(&output.stdout) + .map_err(std::io::Error::other)? + .to_string()) + } +} + +pub fn check_output_projdir, S: AsRef>( + prog: impl AsRef, + args: I, +) -> std::io::Result { + check_output(project_root()?, prog, args) +} + +pub fn run_command, S: AsRef>( + path: impl AsRef, + prog: impl AsRef, + args: I, +) -> std::io::Result<()> { + let status = Command::new(prog).current_dir(path).args(args).status()?; + + if !status.success() { + Err(std::io::Error::other("command returned non-zero exit code")) + } else { + Ok(()) + } +} + +#[allow(unused)] +pub fn run_command_projdir, S: AsRef>( + prog: impl AsRef, + args: I, +) -> std::io::Result<()> { + run_command(project_root()?, prog, args) +} + +pub fn run_echo, S: AsRef>( + path: impl AsRef, + prog: impl AsRef, + args: I, +) -> std::io::Result<()> { + let args = args.into_iter().collect::>(); + + let prog_exe = Path::new(prog.as_ref()) + .file_name() + .unwrap_or_else(|| prog.as_ref()); + + print!("{}", prog_exe.to_string_lossy()); + + if !args.is_empty() { + for arg in &args { + print!(" {}", arg.as_ref().to_string_lossy()); + } + } + + println!(); + run_command(path, prog, args) +} + +pub fn run_echo_projdir, S: AsRef>( + prog: impl AsRef, + args: I, +) -> std::io::Result<()> { + run_echo(project_root()?, prog, args) +} + +pub fn run_cargo, S: AsRef>(args: I) -> std::io::Result<()> { + run_echo_projdir(cargo(), args) +}