diff --git a/aya-obj/src/btf/btf.rs b/aya-obj/src/btf/btf.rs index e94b67be..360079f0 100644 --- a/aya-obj/src/btf/btf.rs +++ b/aya-obj/src/btf/btf.rs @@ -11,7 +11,7 @@ use alloc::{ use bytes::BufMut; use log::debug; -use object::Endianness; +use object::{Endianness, SectionIndex}; use crate::{ btf::{ @@ -414,7 +414,7 @@ impl Btf { pub(crate) fn fixup_and_sanitize( &mut self, - section_sizes: &HashMap, + section_infos: &HashMap, symbol_offsets: &HashMap, features: &BtfFeatures, ) -> Result<(), BtfError> { @@ -489,8 +489,8 @@ impl Btf { } else { // We need to get the size of the section from the ELF file // Fortunately, we cached these when parsing it initially - // and we can this up by name in section_sizes - let size = section_sizes.get(&name).ok_or_else(|| { + // and we can this up by name in section_infos + let (_, size) = section_infos.get(&name).ok_or_else(|| { BtfError::UnknownSectionSize { section_name: name.clone(), } @@ -639,7 +639,7 @@ impl Object { } // fixup btf obj_btf.fixup_and_sanitize( - &self.section_sizes, + &self.section_infos, &self.symbol_offset_by_name, features, )?; @@ -1277,7 +1277,7 @@ mod tests { }; btf.fixup_and_sanitize( - &HashMap::from([(".data/foo".to_string(), 32u64)]), + &HashMap::from([(".data/foo".to_string(), (SectionIndex(0), 32u64))]), &HashMap::from([("foo".to_string(), 64u64)]), &features, ) diff --git a/aya-obj/src/btf/relocation.rs b/aya-obj/src/btf/relocation.rs index 3098ab7f..38452446 100644 --- a/aya-obj/src/btf/relocation.rs +++ b/aya-obj/src/btf/relocation.rs @@ -1,12 +1,14 @@ -use core::{mem, ptr, str::FromStr}; +use core::{mem, ops::Bound::Included, ptr}; use alloc::{ borrow::ToOwned, + collections::BTreeMap, format, string::{String, ToString}, vec, vec::Vec, }; +use object::SectionIndex; use crate::{ btf::{ @@ -18,7 +20,7 @@ use crate::{ BPF_K, BPF_LD, BPF_LDX, BPF_ST, BPF_STX, BPF_W, BTF_INT_SIGNED, }, util::HashMap, - Object, Program, ProgramSection, + Function, Object, }; #[cfg(not(feature = "std"))] @@ -43,9 +45,13 @@ enum RelocationError { #[error(transparent)] IOError(#[from] std::io::Error), - /// Program not found - #[error("program not found")] - ProgramNotFound, + /// Section not found + #[error("section not found")] + SectionNotFound, + + /// Function not found + #[error("function not found")] + FunctionNotFound, /// Invalid relocation access string #[error("invalid relocation access string {access_str}")] @@ -227,25 +233,26 @@ impl Object { error: RelocationError::BtfError(e), })?; - let program_section = match ProgramSection::from_str(§ion_name) { - Ok(program) => program, - Err(_) => continue, - }; - let section_name = program_section.name(); - - let program = self - .programs - .get_mut(section_name) - .ok_or(BtfRelocationError { - section: section_name.to_owned(), - error: RelocationError::ProgramNotFound, + let (section_index, _) = self + .section_infos + .get(§ion_name.to_string()) + .ok_or_else(|| BtfRelocationError { + section: section_name.to_string(), + error: RelocationError::SectionNotFound, })?; - match relocate_btf_program(program, relos, local_btf, target_btf, &mut candidates_cache) - { + + match relocate_btf_functions( + section_index, + &mut self.functions, + relos, + local_btf, + target_btf, + &mut candidates_cache, + ) { Ok(_) => {} Err(error) => { return Err(BtfRelocationError { - section: section_name.to_owned(), + section: section_name.to_string(), error, }) } @@ -256,16 +263,55 @@ impl Object { } } -fn relocate_btf_program<'target>( - program: &mut Program, +fn is_relocation_inside_function( + section_index: &SectionIndex, + func: &Function, + rel: &Relocation, +) -> bool { + if section_index.0 != func.section_index.0 { + return false; + } + + let ins_offset = rel.ins_offset / mem::size_of::(); + let func_offset = func.section_offset / mem::size_of::(); + let func_size = func.instructions.len(); + + (func_offset..func_offset + func_size).contains(&ins_offset) +} + +fn function_by_relocation<'a>( + section_index: &SectionIndex, + functions: &'a mut BTreeMap<(usize, u64), Function>, + rel: &Relocation, +) -> Option<&'a mut Function> { + functions + .range_mut(( + Included(&(section_index.0, 0)), + Included(&(section_index.0, u64::MAX)), + )) + .map(|(_, func)| func) + .find(|func| is_relocation_inside_function(section_index, func, rel)) +} + +fn relocate_btf_functions<'target>( + section_index: &SectionIndex, + functions: &mut BTreeMap<(usize, u64), Function>, relos: &[Relocation], local_btf: &Btf, target_btf: &'target Btf, candidates_cache: &mut HashMap>>, ) -> Result<(), RelocationError> { + let mut last_function_opt: Option<&mut Function> = None; + for rel in relos { - let instructions = &mut program.function.instructions; - let ins_index = rel.ins_offset / mem::size_of::(); + let function = match last_function_opt.take() { + Some(func) if is_relocation_inside_function(section_index, func, rel) => func, + _ => function_by_relocation(section_index, functions, rel) + .ok_or(RelocationError::FunctionNotFound)?, + }; + + let instructions = &mut function.instructions; + let ins_index = (rel.ins_offset - function.section_offset) / mem::size_of::(); if ins_index >= instructions.len() { return Err(RelocationError::InvalidInstructionIndex { index: ins_index, @@ -337,7 +383,9 @@ fn relocate_btf_program<'target>( ComputedRelocation::new(rel, &local_spec, None)? }; - comp_rel.apply(program, rel, local_btf, target_btf)?; + comp_rel.apply(function, rel, local_btf, target_btf)?; + + last_function_opt = Some(function); } Ok(()) @@ -847,14 +895,14 @@ impl ComputedRelocation { fn apply( &self, - program: &mut Program, + function: &mut Function, rel: &Relocation, local_btf: &Btf, target_btf: &Btf, ) -> Result<(), RelocationError> { - let instructions = &mut program.function.instructions; + let instructions = &mut function.instructions; let num_instructions = instructions.len(); - let ins_index = rel.ins_offset / mem::size_of::(); + let ins_index = (rel.ins_offset - function.section_offset) / mem::size_of::(); let ins = instructions .get_mut(ins_index) diff --git a/aya-obj/src/lib.rs b/aya-obj/src/lib.rs index a72a8eef..d11cdaf6 100644 --- a/aya-obj/src/lib.rs +++ b/aya-obj/src/lib.rs @@ -42,7 +42,8 @@ //! object.relocate_maps(std::iter::empty(), &text_sections).unwrap(); //! //! // Run with rbpf -//! let instructions = &object.programs["prog_name"].function.instructions; +//! let function = object.functions.get(&object.programs["prog_name"].function_key()).unwrap(); +//! let instructions = &function.instructions; //! let data = unsafe { //! core::slice::from_raw_parts( //! instructions.as_ptr() as *const u8, diff --git a/aya-obj/src/obj.rs b/aya-obj/src/obj.rs index 7536dda8..0a7235fd 100644 --- a/aya-obj/src/obj.rs +++ b/aya-obj/src/obj.rs @@ -2,6 +2,7 @@ use alloc::{ borrow::ToOwned, + collections::BTreeMap, ffi::CString, string::{String, ToString}, vec::Vec, @@ -67,10 +68,10 @@ pub struct Object { /// in [ProgramSection]s as keys. pub programs: HashMap, /// Functions - pub functions: HashMap<(usize, u64), Function>, + pub functions: BTreeMap<(usize, u64), Function>, pub(crate) relocations: HashMap>, pub(crate) symbol_table: HashMap, - pub(crate) section_sizes: HashMap, + pub(crate) section_infos: HashMap, // symbol_offset_by_name caches symbols that could be referenced from a // BTF VAR type so the offsets can be fixed up pub(crate) symbol_offset_by_name: HashMap, @@ -85,8 +86,17 @@ pub struct Program { pub kernel_version: KernelVersion, /// The section containing the program pub section: ProgramSection, - /// The function - pub function: Function, + /// The section index of the program + pub section_index: usize, + /// The address of the program + pub address: u64, +} + +impl Program { + /// The key used by [Object::functions] + pub fn function_key(&self) -> (usize, u64) { + (self.section_index, self.address) + } } /// An eBPF function @@ -578,10 +588,10 @@ impl Object { btf_ext: None, maps: HashMap::new(), programs: HashMap::new(), - functions: HashMap::new(), + functions: BTreeMap::new(), relocations: HashMap::new(), symbol_table: HashMap::new(), - section_sizes: HashMap::new(), + section_infos: HashMap::new(), symbol_offset_by_name: HashMap::new(), } } @@ -647,7 +657,7 @@ impl Object { Ok(()) } - fn parse_program(&self, section: &Section) -> Result { + fn parse_program(&self, section: &Section) -> Result<(Program, Function), ParseError> { let prog_sec = ProgramSection::from_str(section.name)?; let name = prog_sec.name().to_owned(); @@ -665,22 +675,28 @@ impl Object { (FuncSecInfo::default(), LineSecInfo::default(), 0, 0) }; - Ok(Program { - license: self.license.clone(), - kernel_version: self.kernel_version, - section: prog_sec, - function: Function { - name, - address: section.address, - section_index: section.index, - section_offset: 0, - instructions: copy_instructions(section.data)?, - func_info, - line_info, - func_info_rec_size, - line_info_rec_size, + let function = Function { + name, + address: section.address, + section_index: section.index, + section_offset: 0, + instructions: copy_instructions(section.data)?, + func_info, + line_info, + func_info_rec_size, + line_info_rec_size, + }; + + Ok(( + Program { + license: self.license.clone(), + kernel_version: self.kernel_version, + section: prog_sec, + section_index: function.section_index.0, + address: function.address, }, - }) + function, + )) } fn parse_text_section(&mut self, section: Section) -> Result<(), ParseError> { @@ -829,8 +845,8 @@ impl Object { { parts.push(parts[0]); } - self.section_sizes - .insert(section.name.to_owned(), section.size); + self.section_infos + .insert(section.name.to_owned(), (section.index, section.size)); match section.kind { BpfSectionKind::Data | BpfSectionKind::Rodata | BpfSectionKind::Bss => { self.maps @@ -876,7 +892,8 @@ impl Object { res? } BpfSectionKind::Program => { - let program = self.parse_program(§ion)?; + let (program, function) = self.parse_program(§ion)?; + self.functions.insert(program.function_key(), function); self.programs .insert(program.section.name().to_owned(), program); if !section.relocations.is_empty() { @@ -896,10 +913,10 @@ impl Object { Ok(()) } - /// Sanitize BPF programs. - pub fn sanitize_programs(&mut self, features: &Features) { - for program in self.programs.values_mut() { - program.sanitize(features); + /// Sanitize BPF functions. + pub fn sanitize_functions(&mut self, features: &Features) { + for function in self.functions.values_mut() { + function.sanitize(features); } } } @@ -919,9 +936,9 @@ const BPF_FUNC_PROBE_READ_KERNEL: i32 = 113; const BPF_FUNC_PROBE_READ_USER_STR: i32 = 114; const BPF_FUNC_PROBE_READ_KERNEL_STR: i32 = 115; -impl Program { +impl Function { fn sanitize(&mut self, features: &Features) { - for inst in &mut self.function.instructions { + for inst in &mut self.instructions { if !insn_is_helper_call(inst) { continue; } @@ -1679,17 +1696,17 @@ mod tests { assert_matches!( obj.parse_program(&fake_section(BpfSectionKind::Program,"kprobe/foo", bytes_of(&fake_ins()))), - Ok(Program { + Ok((Program { license, kernel_version: KernelVersion::Any, section: ProgramSection::KProbe { .. }, - function: Function { + .. }, Function { name, address: 0, section_index: SectionIndex(0), section_offset: 0, instructions, - ..} }) if license.to_string_lossy() == "GPL" && name == "foo" && instructions.len() == 1 + ..})) if license.to_string_lossy() == "GPL" && name == "foo" && instructions.len() == 1 ); } diff --git a/aya-obj/src/relocation.rs b/aya-obj/src/relocation.rs index 6422ed38..7fb3bc0d 100644 --- a/aya-obj/src/relocation.rs +++ b/aya-obj/src/relocation.rs @@ -2,7 +2,7 @@ use core::mem; -use alloc::{borrow::ToOwned, string::String}; +use alloc::{borrow::ToOwned, collections::BTreeMap, string::String}; use log::debug; use object::{SectionIndex, SymbolKind}; @@ -12,7 +12,7 @@ use crate::{ BPF_PSEUDO_MAP_VALUE, }, maps::Map, - obj::{Function, Object, Program}, + obj::{Function, Object}, util::{HashMap, HashSet}, BpfSectionKind, }; @@ -64,6 +64,15 @@ pub enum RelocationError { caller_name: String, }, + /// Unknown function + #[error("program at section {section_index} and address {address:#x} was not found while relocating")] + UnknownProgram { + /// The function section index + section_index: usize, + /// The function address + address: u64, + }, + /// Referenced map not created yet #[error("the map `{name}` at section `{section_index}` has not been created")] MapNotCreated { @@ -119,13 +128,7 @@ impl Object { } } - let functions = self - .programs - .values_mut() - .map(|p| &mut p.function) - .chain(self.functions.values_mut()); - - for function in functions { + for function in self.functions.values_mut() { if let Some(relocations) = self.relocations.get(&function.section_index) { relocate_maps( function, @@ -150,17 +153,31 @@ impl Object { &mut self, text_sections: &HashSet, ) -> Result<(), BpfRelocationError> { - for (name, program) in self.programs.iter_mut() { + for (name, program) in self.programs.iter() { let linker = FunctionLinker::new( &self.functions, &self.relocations, &self.symbol_table, text_sections, ); - linker.link(program).map_err(|error| BpfRelocationError { + + let func_orig = + self.functions + .get(&program.function_key()) + .ok_or_else(|| BpfRelocationError { + function: name.clone(), + error: RelocationError::UnknownProgram { + section_index: program.section_index, + address: program.address, + }, + })?; + + let func = linker.link(func_orig).map_err(|error| BpfRelocationError { function: name.to_owned(), error, })?; + + self.functions.insert(program.function_key(), func); } Ok(()) @@ -270,7 +287,7 @@ fn relocate_maps<'a, I: Iterator>( } struct FunctionLinker<'a> { - functions: &'a HashMap<(usize, u64), Function>, + functions: &'a BTreeMap<(usize, u64), Function>, linked_functions: HashMap, relocations: &'a HashMap>, symbol_table: &'a HashMap, @@ -279,7 +296,7 @@ struct FunctionLinker<'a> { impl<'a> FunctionLinker<'a> { fn new( - functions: &'a HashMap<(usize, u64), Function>, + functions: &'a BTreeMap<(usize, u64), Function>, relocations: &'a HashMap>, symbol_table: &'a HashMap, text_sections: &'a HashSet, @@ -293,17 +310,15 @@ impl<'a> FunctionLinker<'a> { } } - fn link(mut self, program: &mut Program) -> Result<(), RelocationError> { - let mut fun = program.function.clone(); + fn link(mut self, program_function: &Function) -> Result { + let mut fun = program_function.clone(); // relocate calls in the program's main function. As relocation happens, // it will trigger linking in all the callees. - self.relocate(&mut fun, &program.function)?; + self.relocate(&mut fun, program_function)?; // this now includes the program function plus all the other functions called during // execution - program.function = fun; - - Ok(()) + Ok(fun) } fn link_function( diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs index 938629d0..2b499c12 100644 --- a/aya/src/bpf.rs +++ b/aya/src/bpf.rs @@ -441,18 +441,21 @@ impl<'a> BpfLoader<'a> { &text_sections, )?; obj.relocate_calls(&text_sections)?; - obj.sanitize_programs(&FEATURES); + obj.sanitize_functions(&FEATURES); let programs = obj .programs .drain() - .map(|(name, obj)| { + .map(|(name, prog_obj)| { + let function_obj = obj.functions.get(&prog_obj.function_key()).unwrap().clone(); + let prog_name = if FEATURES.bpf_name { Some(name.clone()) } else { None }; - let section = obj.section.clone(); + let section = prog_obj.section.clone(); + let obj = (prog_obj, function_obj); let program = if self.extensions.contains(name.as_str()) { Program::Extension(Extension { diff --git a/aya/src/programs/mod.rs b/aya/src/programs/mod.rs index 9ee02141..d903cca7 100644 --- a/aya/src/programs/mod.rs +++ b/aya/src/programs/mod.rs @@ -405,7 +405,7 @@ impl Program { #[derive(Debug)] pub(crate) struct ProgramData { pub(crate) name: Option, - pub(crate) obj: Option, + pub(crate) obj: Option<(obj::Program, obj::Function)>, pub(crate) fd: Option, pub(crate) links: LinkMap, pub(crate) expected_attach_type: Option, @@ -421,7 +421,7 @@ pub(crate) struct ProgramData { impl ProgramData { pub(crate) fn new( name: Option, - obj: obj::Program, + obj: (obj::Program, obj::Function), btf_fd: Option, verifier_log_level: u32, ) -> ProgramData { @@ -557,20 +557,21 @@ fn load_program( return Err(ProgramError::AlreadyLoaded); } let obj = obj.as_ref().unwrap(); - let crate::obj::Program { - function: - Function { - instructions, - func_info, - line_info, - func_info_rec_size, - line_info_rec_size, - .. - }, - license, - kernel_version, - .. - } = obj; + let ( + crate::obj::Program { + license, + kernel_version, + .. + }, + Function { + instructions, + func_info, + line_info, + func_info_rec_size, + line_info_rec_size, + .. + }, + ) = obj; let target_kernel_version = match *kernel_version { KernelVersion::Any => { diff --git a/test/integration-test/src/tests/btf_relocations.rs b/test/integration-test/src/tests/btf_relocations.rs index 85291ef4..96df45b8 100644 --- a/test/integration-test/src/tests/btf_relocations.rs +++ b/test/integration-test/src/tests/btf_relocations.rs @@ -191,12 +191,17 @@ impl RelocationTest { }} output_map __attribute__((section(".maps"), used)); + __attribute__ ((noinline)) int bpf_func() {{ + __u32 key = 0; + __u64 value = 0; + {relocation_code} + bpf_map_update_elem(&output_map, &key, &value, BPF_ANY); + return 0; + }} + __attribute__((section("tracepoint/bpf_prog"), used)) int bpf_prog(void *ctx) {{ - __u32 key = 0; - __u64 value = 0; - {relocation_code} - bpf_map_update_elem(&output_map, &key, &value, BPF_ANY); + bpf_func(); return 0; }} diff --git a/test/integration-test/src/tests/rbpf.rs b/test/integration-test/src/tests/rbpf.rs index e896d262..2f46381d 100644 --- a/test/integration-test/src/tests/rbpf.rs +++ b/test/integration-test/src/tests/rbpf.rs @@ -18,7 +18,11 @@ fn run_with_rbpf() { )); assert_eq!(object.programs["pass"].section.name(), "pass"); - let instructions = &object.programs["pass"].function.instructions; + let instructions = &object + .functions + .get(&object.programs["pass"].function_key()) + .unwrap() + .instructions; let data = unsafe { from_raw_parts( instructions.as_ptr() as *const u8, @@ -86,7 +90,11 @@ fn use_map_with_rbpf() { // Executes the program assert_eq!(object.programs.len(), 1); - let instructions = &object.programs["tracepoint"].function.instructions; + let instructions = &object + .functions + .get(&object.programs["tracepoint"].function_key()) + .unwrap() + .instructions; let data = unsafe { from_raw_parts( instructions.as_ptr() as *const u8,