From 9e1109b3ce70a3668771bd11a7fda101eec3ab93 Mon Sep 17 00:00:00 2001 From: Mary Date: Tue, 2 May 2023 16:24:13 +0200 Subject: [PATCH] aya: Move program's functions to the same map --- aya-obj/src/btf/relocation.rs | 27 ++++++---- aya-obj/src/lib.rs | 3 +- aya-obj/src/obj.rs | 72 +++++++++++++++---------- aya-obj/src/relocation.rs | 62 +++++++++++---------- aya/src/bpf.rs | 9 ++-- aya/src/programs/mod.rs | 33 ++++++------ test/integration-test/src/tests/rbpf.rs | 12 ++++- 7 files changed, 131 insertions(+), 87 deletions(-) diff --git a/aya-obj/src/btf/relocation.rs b/aya-obj/src/btf/relocation.rs index 3098ab7f..88d9a54d 100644 --- a/aya-obj/src/btf/relocation.rs +++ b/aya-obj/src/btf/relocation.rs @@ -18,7 +18,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, ProgramSection, }; #[cfg(not(feature = "std"))] @@ -233,15 +233,22 @@ impl Object { }; let section_name = program_section.name(); - let program = self + let function = self .programs .get_mut(section_name) + .and_then(|x| self.functions.get_mut(&x.function_key())) .ok_or(BtfRelocationError { section: section_name.to_owned(), error: RelocationError::ProgramNotFound, })?; - match relocate_btf_program(program, relos, local_btf, target_btf, &mut candidates_cache) - { + + match relocate_btf_function( + function, + relos, + local_btf, + target_btf, + &mut candidates_cache, + ) { Ok(_) => {} Err(error) => { return Err(BtfRelocationError { @@ -256,15 +263,15 @@ impl Object { } } -fn relocate_btf_program<'target>( - program: &mut Program, +fn relocate_btf_function<'target>( + function: &mut Function, relos: &[Relocation], local_btf: &Btf, target_btf: &'target Btf, candidates_cache: &mut HashMap>>, ) -> Result<(), RelocationError> { for rel in relos { - let instructions = &mut program.function.instructions; + let instructions = &mut function.instructions; let ins_index = rel.ins_offset / mem::size_of::(); if ins_index >= instructions.len() { return Err(RelocationError::InvalidInstructionIndex { @@ -337,7 +344,7 @@ 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)?; } Ok(()) @@ -847,12 +854,12 @@ 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 = 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 db1c50d9..1a3e74f6 100644 --- a/aya-obj/src/obj.rs +++ b/aya-obj/src/obj.rs @@ -84,8 +84,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 @@ -646,7 +655,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(); @@ -664,22 +673,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> { @@ -875,7 +890,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() { @@ -895,10 +911,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); } } } @@ -918,9 +934,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; } @@ -1652,17 +1668,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 55eccb5a..669d1200 100644 --- a/aya-obj/src/relocation.rs +++ b/aya-obj/src/relocation.rs @@ -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,35 @@ impl Object { &mut self, text_sections: &HashSet, ) -> Result<(), BpfRelocationError> { - for (name, program) in self.programs.iter_mut() { - let linker = FunctionLinker::new( + for (name, program) in self.programs.iter() { + let mut linker = FunctionLinker::new( &self.functions, &self.relocations, &self.symbol_table, text_sections, ); - linker.link(program).map_err(|error| BpfRelocationError { - function: name.to_owned(), - error, - })?; + + 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 mut func = func_orig.clone(); + + linker + .relocate(&mut func, func_orig) + .map_err(|error| BpfRelocationError { + function: name.to_owned(), + error, + })?; + + self.functions.insert(program.function_key(), func); } Ok(()) @@ -293,19 +314,6 @@ impl<'a> FunctionLinker<'a> { } } - fn link(mut self, program: &mut Program) -> Result<(), RelocationError> { - 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)?; - - // this now includes the program function plus all the other functions called during - // execution - program.function = fun; - - Ok(()) - } - fn link_function( &mut self, program: &mut Function, diff --git a/aya/src/bpf.rs b/aya/src/bpf.rs index 84026a4b..6e21e97d 100644 --- a/aya/src/bpf.rs +++ b/aya/src/bpf.rs @@ -433,18 +433,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/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,