Merge pull request #602 from marysaka/fix/btf-reloc-all-functions

aya: Apply BTF relocations to all functions

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<String, u64>,
+        section_infos: &HashMap<String, (SectionIndex, u64)>,
         symbol_offsets: &HashMap<String, u64>,
         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,
         )

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(&section_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(&section_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::<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 instructions = &mut program.function.instructions;
-        let ins_index = rel.ins_offset / mem::size_of::<bpf_insn>();
+        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::<bpf_insn>();
         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::<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/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

@@ -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<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>,
@@ -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<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();
 
@@ -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(&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() {
@@ -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
         );
     }
 

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<usize>,
     ) -> 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<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>,
@@ -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<SectionIndex, HashMap<u64, Relocation>>,
         symbol_table: &'a HashMap<usize, Symbol>,
         text_sections: &'a HashSet<usize>,
@@ -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<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)?;
+        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(

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 {

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

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,