aya: Apply BTF relocations to all functions

This fix aya wrong logic causing non entrypoint functions to not have
any BTF relocations working.

Also fix missing section_offset computation for instruction offset in
multiple spots.

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::{
@@ -409,7 +409,7 @@ impl Btf {
 
     pub(crate) fn fixup_and_sanitize(
         &mut self,
-        section_sizes: &HashMap<String, u64>,
+        section_infos: &HashMap<String, (SectionIndex, u64)>,
         symbol_offsets: &HashMap<String, u64>,
         features: &BtfFeatures,
     ) -> Result<(), BtfError> {
@@ -484,8 +484,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(),
                             }
@@ -631,7 +631,7 @@ impl Object {
         if let Some(ref mut obj_btf) = self.btf {
             // fixup btf
             obj_btf.fixup_and_sanitize(
-                &self.section_sizes,
+                &self.section_infos,
                 &self.symbol_offset_by_name,
                 features,
             )?;
@@ -1212,7 +1212,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,
         )

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,
-    Function, Object, 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,23 +233,17 @@ impl Object {
                         error: RelocationError::BtfError(e),
                     })?;
 
-            let program_section = match ProgramSection::from_str(&section_name) {
-                Ok(program) => program,
-                Err(_) => continue,
-            };
-            let section_name = program_section.name();
-
-            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,
+            let (section_index, _) = self
+                .section_infos
+                .get(&section_name.to_string())
+                .ok_or_else(|| BtfRelocationError {
+                    section: section_name.to_string(),
+                    error: RelocationError::SectionNotFound,
                 })?;
 
-            match relocate_btf_function(
-                function,
+            match relocate_btf_functions(
+                section_index,
+                &mut self.functions,
                 relos,
                 local_btf,
                 target_btf,
@@ -252,7 +252,7 @@ impl Object {
                 Ok(_) => {}
                 Err(error) => {
                     return Err(BtfRelocationError {
-                        section: section_name.to_owned(),
+                        section: section_name.to_string(),
                         error,
                     })
                 }
@@ -263,16 +263,55 @@ impl Object {
     }
 }
 
-fn relocate_btf_function<'target>(
-    function: &mut Function,
+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::<bpf_insn>();
+    let func_offset = func.section_offset / mem::size_of::<bpf_insn>();
+    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<u32, Vec<Candidate<'target>>>,
 ) -> Result<(), RelocationError> {
+    let mut last_function_opt: Option<&mut Function> = None;
+
     for rel in relos {
+        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 / mem::size_of::<bpf_insn>();
+        let ins_index = (rel.ins_offset - function.section_offset) / mem::size_of::<bpf_insn>();
         if ins_index >= instructions.len() {
             return Err(RelocationError::InvalidInstructionIndex {
                 index: ins_index,
@@ -345,6 +384,8 @@ fn relocate_btf_function<'target>(
         };
 
         comp_rel.apply(function, rel, local_btf, target_btf)?;
+
+        last_function_opt = Some(function);
     }
 
     Ok(())
@@ -861,7 +902,7 @@ impl ComputedRelocation {
     ) -> Result<(), RelocationError> {
         let instructions = &mut function.instructions;
         let num_instructions = instructions.len();
-        let ins_index = rel.ins_offset / mem::size_of::<bpf_insn>();
+        let ins_index = (rel.ins_offset - function.section_offset) / mem::size_of::<bpf_insn>();
         let ins =
             instructions
                 .get_mut(ins_index)

aya-obj/src/obj.rs

@@ -2,6 +2,7 @@
 
 use alloc::{
     borrow::ToOwned,
+    collections::BTreeMap,
     ffi::CString,
     string::{String, ToString},
     vec::Vec,
@@ -66,10 +67,10 @@ pub struct Object {
     /// in [ProgramSection]s as keys.
     pub programs: HashMap<String, Program>,
     /// Functions
-    pub functions: HashMap<(usize, u64), Function>,
+    pub functions: BTreeMap<(usize, u64), Function>,
     pub(crate) relocations: HashMap<SectionIndex, HashMap<u64, Relocation>>,
     pub(crate) symbol_table: HashMap<usize, Symbol>,
-    pub(crate) section_sizes: HashMap<String, u64>,
+    pub(crate) section_infos: HashMap<String, (SectionIndex, u64)>,
     // 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<String, u64>,
@@ -586,10 +587,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(),
         }
     }
@@ -843,8 +844,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

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};
 
@@ -154,7 +154,7 @@ impl Object {
         text_sections: &HashSet<usize>,
     ) -> Result<(), BpfRelocationError> {
         for (name, program) in self.programs.iter() {
-            let mut linker = FunctionLinker::new(
+            let linker = FunctionLinker::new(
                 &self.functions,
                 &self.relocations,
                 &self.symbol_table,
@@ -172,14 +172,10 @@ impl Object {
                         },
                     })?;
 
-            let mut func = func_orig.clone();
-
-            linker
-                .relocate(&mut func, func_orig)
-                .map_err(|error| BpfRelocationError {
-                    function: name.to_owned(),
-                    error,
-                })?;
+            let func = linker.link(func_orig).map_err(|error| BpfRelocationError {
+                function: name.to_owned(),
+                error,
+            })?;
 
             self.functions.insert(program.function_key(), func);
         }
@@ -291,7 +287,7 @@ fn relocate_maps<'a, I: Iterator<Item = &'a Relocation>>(
 }
 
 struct FunctionLinker<'a> {
-    functions: &'a HashMap<(usize, u64), Function>,
+    functions: &'a BTreeMap<(usize, u64), Function>,
     linked_functions: HashMap<u64, usize>,
     relocations: &'a HashMap<SectionIndex, HashMap<u64, Relocation>>,
     symbol_table: &'a HashMap<usize, Symbol>,
@@ -300,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<SectionIndex, HashMap<u64, Relocation>>,
         symbol_table: &'a HashMap<usize, Symbol>,
         text_sections: &'a HashSet<usize>,
@@ -314,6 +310,17 @@ impl<'a> FunctionLinker<'a> {
         }
     }
 
+    fn link(mut self, program_function: &Function) -> Result<Function, 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
+        Ok(fun)
+    }
+
     fn link_function(
         &mut self,
         program: &mut Function,

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;
                 }}