aya: Move program's functions to the same map

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<u32, Vec<Candidate<'target>>>,
 ) -> 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::<bpf_insn>();
         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::<bpf_insn>();
         let ins =

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,

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<Program, ParseError> {
+    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(&section)?;
+                let (program, function) = self.parse_program(&section)?;
+                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
         );
     }
 

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<usize>,
     ) -> 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,

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 {

aya/src/programs/mod.rs

@@ -405,7 +405,7 @@ impl Program {
 #[derive(Debug)]
 pub(crate) struct ProgramData<T: Link> {
     pub(crate) name: Option<String>,
-    pub(crate) obj: Option<obj::Program>,
+    pub(crate) obj: Option<(obj::Program, obj::Function)>,
     pub(crate) fd: Option<RawFd>,
     pub(crate) links: LinkMap<T>,
     pub(crate) expected_attach_type: Option<bpf_attach_type>,
@@ -421,7 +421,7 @@ pub(crate) struct ProgramData<T: Link> {
 impl<T: Link> ProgramData<T> {
     pub(crate) fn new(
         name: Option<String>,
-        obj: obj::Program,
+        obj: (obj::Program, obj::Function),
         btf_fd: Option<RawFd>,
         verifier_log_level: u32,
     ) -> ProgramData<T> {
@@ -557,20 +557,21 @@ fn load_program<T: Link>(
         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 => {

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,