diff --git a/Cargo.lock b/Cargo.lock index ad3a0b5..26efea5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1748,9 +1748,9 @@ checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b" [[package]] name = "goblin" -version = "0.4.3" +version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32401e89c6446dcd28185931a01b1093726d0356820ac744023e6850689bf926" +checksum = "daa0a64d21a7eb230583b4c5f4e23b7e4e57974f96620f42a7e75e08ae66d745" dependencies = [ "log", "plain", @@ -2818,22 +2818,22 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] name = "scroll" -version = "0.10.2" +version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fda28d4b4830b807a8b43f7b0e6b5df875311b3e7621d84577188c175b6ec1ec" +checksum = "6ab8598aa408498679922eff7fa985c25d58a90771bd6be794434c5277eab1a6" dependencies = [ "scroll_derive", ] [[package]] name = "scroll_derive" -version = "0.10.5" +version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aaaae8f38bb311444cfb7f1979af0bc9240d95795f75f9ceddf6a59b79ceffa0" +checksum = "1783eabc414609e28a5ba76aee5ddd52199f7107a0b24c2e9746a1ecc34a683d" dependencies = [ "proc-macro2", "quote", - "syn 1.0.109", + "syn 2.0.89", ] [[package]] diff --git a/zluda_trace/Cargo.toml b/zluda_trace/Cargo.toml index b679eaf..e709f17 100644 --- a/zluda_trace/Cargo.toml +++ b/zluda_trace/Cargo.toml @@ -18,7 +18,7 @@ regex = "1.4" dynasm = "1.2" dynasmrt = "1.2" # we don't need elf32, but goblin has a bug where elf64 does not build without elf32 -goblin = { version = "0.4", default-features = false, features = ["elf64", "elf32", "archive"] } +goblin = { version = "0.9", default-features = false, features = ["elf64", "elf32", "archive"] } paste = "1.0" cuda_macros = { path = "../cuda_macros" } cuda_types = { path = "../cuda_types" } diff --git a/zluda_trace/src/lib.rs b/zluda_trace/src/lib.rs index 13e4d2a..50603f1 100644 --- a/zluda_trace/src/lib.rs +++ b/zluda_trace/src/lib.rs @@ -1451,11 +1451,13 @@ pub(crate) fn cuLibraryLoadData_Post( _library_option_values: *mut *mut ::core::ffi::c_void, _num_library_options: ::core::ffi::c_uint, state: &mut trace::StateTracker, - _fn_logger: &mut FnCallLog, + fn_logger: &mut FnCallLog, _result: CUresult, ) { - // TODO: this is not great, the lifetime of `code` is not guaranteed to be 'static state .libraries .insert(unsafe { *library }, trace::CodePointer(code)); + // TODO: this is not correct, but it's enough for now, we just want to + // save the binary to disk + state.record_new_module(unsafe { CUmodule((*library).0.cast()) }, code, fn_logger); } diff --git a/zluda_trace/src/trace.rs b/zluda_trace/src/trace.rs index dfd0011..1242df6 100644 --- a/zluda_trace/src/trace.rs +++ b/zluda_trace/src/trace.rs @@ -9,12 +9,14 @@ use cuda_types::{ use dark_api::fatbin::{ decompress_lz4, decompress_zstd, Fatbin, FatbinFileIterator, FatbinSubmodule, }; +use goblin::{elf, elf32, elf64}; use rustc_hash::{FxHashMap, FxHashSet}; use std::{ borrow::Cow, ffi::{c_void, CStr, CString}, fs::{self, File}, io::{self, Read, Write}, + mem, path::PathBuf, }; use unwrap_or::unwrap_some_or; @@ -121,14 +123,20 @@ impl StateTracker { fn_logger: &mut FnCallLog, ) { self.module_counter += 1; - if unsafe { *(raw_image as *const [u8; 4]) } == *goblin::elf64::header::ELFMAG { + if unsafe { *(raw_image as *const [u8; 4]) } == *elf64::header::ELFMAG { self.saved_modules.insert(module); - // TODO: Parse ELF and write it to disk - fn_logger.log(log::ErrorEntry::UnsupportedModule { - module, - raw_image, - kind: "ELF", - }) + match unsafe { get_elf_size(raw_image) } { + Some(len) => { + let elf_image = + unsafe { std::slice::from_raw_parts(raw_image.cast::(), len) }; + self.record_new_submodule(module, elf_image, fn_logger, "elf"); + } + None => fn_logger.log(log::ErrorEntry::UnsupportedModule { + module, + raw_image, + kind: "ELF", + }), + } } else if unsafe { *(raw_image as *const [u8; 8]) } == *goblin::archive::MAGIC { self.saved_modules.insert(module); // TODO: Figure out how to get size of archive module and write it to disk @@ -198,6 +206,146 @@ impl StateTracker { } } +unsafe fn get_elf_size(start: *const c_void) -> Option { + let start = start.cast::(); + let ei_class = start.add(mem::size_of_val(elf::header::ELFMAG)); + let (header, header_size, is_64bit): (elf::header::Header, _, _) = match *ei_class { + elf::header::ELFCLASS32 => ( + (*start.cast::()).into(), + mem::size_of::() as u64, + false, + ), + elf::header::ELFCLASS64 => ( + (*start.cast::()).into(), + mem::size_of::() as u64, + true, + ), + _ => return None, + }; + let mut max_end = header_size; + max_end = max_end.max(get_max_end_for::( + start, header, is_64bit, + )?); + max_end = max_end.max(get_max_end_for::( + start, header, is_64bit, + )?); + Some(max_end as usize) +} + +unsafe fn get_max_end_for( + elf_start: *const u8, + header: elf::header::Header, + is_64bit: bool, +) -> Option { + if is_64bit && T::header_size(&header) as usize != mem::size_of::() { + return None; + } + if !is_64bit && T::header_size(&header) as usize != mem::size_of::() { + return None; + } + if T::is_extended_header(&header) { + // TODO: support extended headers + return None; + } + let headers_start = T::headers_offset(&header); + if headers_start == 0 { + return Some(0); + } + let mut max_end = 0; + for entry_idx in 0..T::headers_count(&header) as u64 { + let header_start = + headers_start.checked_add(entry_idx.checked_mul(T::header_size(&header) as u64)?)?; + let header_end = header_start.checked_add(T::header_size(&header) as u64)?; + max_end = max_end.max(header_end); + let component = if is_64bit { + (*elf_start + .add(header_start as usize) + .cast::()) + .into() + } else { + (*elf_start + .add(header_start as usize) + .cast::()) + .into() + }; + let component_end = component.start().checked_add(component.size())?; + max_end = max_end.max(component_end); + } + Some(max_end) +} + +trait ElfComponent: Sized { + type Component32: Into + Copy; + type Component64: Into + Copy; + + fn is_extended_header(elf_header: &elf::header::Header) -> bool; + fn header_size(elf_header: &elf::header::Header) -> u16; + fn headers_offset(elf_header: &elf::header::Header) -> u64; + fn headers_count(elf_header: &elf::header::Header) -> u64; + fn start(&self) -> u64; + fn size(&self) -> u64; +} + +impl ElfComponent for elf::program_header::ProgramHeader { + type Component32 = elf32::program_header::program_header32::ProgramHeader; + type Component64 = elf64::program_header::program_header64::ProgramHeader; + + fn is_extended_header(elf_header: &elf::header::Header) -> bool { + const PN_XNUM: u16 = 0xffff; + elf_header.e_phnum >= PN_XNUM + } + + fn header_size(elf_header: &elf::header::Header) -> u16 { + elf_header.e_phentsize + } + + fn headers_offset(elf_header: &elf::header::Header) -> u64 { + elf_header.e_phoff + } + + fn headers_count(elf_header: &elf::header::Header) -> u64 { + elf_header.e_phnum as u64 + } + + fn start(&self) -> u64 { + self.p_offset + } + + fn size(&self) -> u64 { + self.p_filesz + } +} + +impl ElfComponent for elf::section_header::SectionHeader { + type Component32 = elf32::section_header::section_header32::SectionHeader; + type Component64 = elf64::section_header::section_header64::SectionHeader; + + fn is_extended_header(elf_header: &elf::header::Header) -> bool { + const SHN_LORESERVE: u16 = 0xff00; + elf_header.e_phnum >= SHN_LORESERVE + } + + fn header_size(elf_header: &elf::header::Header) -> u16 { + elf_header.e_shentsize + } + + fn headers_offset(elf_header: &elf::header::Header) -> u64 { + elf_header.e_shoff + } + + fn headers_count(elf_header: &elf::header::Header) -> u64 { + elf_header.e_shnum as u64 + } + + fn start(&self) -> u64 { + self.sh_offset + } + + fn size(&self) -> u64 { + self.sh_size + } +} + // This structs writes out information about CUDA execution to the trace dir struct DumpWriter { dump_dir: Option,