aya: support relocations across multiple text sections + fixes

Fix R_BPF_64_64 text relocations in sections other than .text (for
instance .text.unlikely). Also fix misc bugs triggered by integration
tests.

aya-obj/src/btf/btf.rs

@@ -1045,9 +1045,7 @@ mod tests {
         let name_offset = btf.add_string("&mut int".to_string());
         let ptr_type_id = btf.add_type(BtfType::Ptr(Ptr::new(name_offset, int_type_id)));
 
-        let features = BtfFeatures {
-            ..Default::default()
-        };
+        let features = Default::default();
 
         btf.fixup_and_sanitize(&HashMap::new(), &HashMap::new(), &features)
             .unwrap();

aya-obj/src/lib.rs

@@ -37,8 +37,9 @@
 //! let bytes = std::fs::read("program.o").unwrap();
 //! let mut object = Object::parse(&bytes).unwrap();
 //! // Relocate the programs
-//! object.relocate_calls().unwrap();
-//! object.relocate_maps(std::iter::empty()).unwrap();
+//! let text_sections = std::collections::HashSet::new();
+//! object.relocate_calls(&text_sections).unwrap();
+//! object.relocate_maps(std::iter::empty(), &text_sections).unwrap();
 //!
 //! // Run with rbpf
 //! let instructions = &object.programs["prog_name"].function.instructions;

aya-obj/src/obj.rs

@@ -52,14 +52,13 @@ pub struct Object {
     /// in [ProgramSection]s as keys.
     pub programs: HashMap<String, Program>,
     /// Functions
-    pub functions: HashMap<u64, Function>,
+    pub functions: HashMap<(usize, u64), Function>,
     pub(crate) relocations: HashMap<SectionIndex, HashMap<u64, Relocation>>,
-    pub(crate) symbols_by_index: HashMap<usize, Symbol>,
+    pub(crate) symbol_table: HashMap<usize, Symbol>,
     pub(crate) section_sizes: HashMap<String, 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>,
-    pub(crate) text_section_index: Option<usize>,
 }
 
 /// An eBPF program
@@ -522,7 +521,7 @@ impl Object {
                     is_definition: symbol.is_definition(),
                     kind: symbol.kind(),
                 };
-                bpf_obj.symbols_by_index.insert(symbol.index().0, sym);
+                bpf_obj.symbol_table.insert(symbol.index().0, sym);
 
                 if symbol.is_global() || symbol.kind() == SymbolKind::Data {
                     bpf_obj.symbol_offset_by_name.insert(name, symbol.address());
@@ -564,17 +563,16 @@ impl Object {
             programs: HashMap::new(),
             functions: HashMap::new(),
             relocations: HashMap::new(),
-            symbols_by_index: HashMap::new(),
+            symbol_table: HashMap::new(),
             section_sizes: HashMap::new(),
             symbol_offset_by_name: HashMap::new(),
-            text_section_index: None,
         }
     }
 
     /// Patches map data
     pub fn patch_map_data(&mut self, globals: HashMap<&str, &[u8]>) -> Result<(), ParseError> {
         let symbols: HashMap<String, &Symbol> = self
-            .symbols_by_index
+            .symbol_table
             .iter()
             .filter(|(_, s)| s.name.is_some())
             .map(|(_, s)| (s.name.as_ref().unwrap().clone(), s))
@@ -668,12 +666,10 @@ impl Object {
         })
     }
 
-    fn parse_text_section(&mut self, mut section: Section) -> Result<(), ParseError> {
-        self.text_section_index = Some(section.index.0);
-
+    fn parse_text_section(&mut self, section: Section) -> Result<(), ParseError> {
         let mut symbols_by_address = HashMap::new();
 
-        for sym in self.symbols_by_index.values() {
+        for sym in self.symbol_table.values() {
             if sym.is_definition
                 && sym.kind == SymbolKind::Text
                 && sym.section_index == Some(section.index.0)
@@ -729,7 +725,7 @@ impl Object {
                 };
 
             self.functions.insert(
-                sym.address,
+                (section.index.0, sym.address),
                 Function {
                     address,
                     name: sym.name.clone().unwrap(),
@@ -804,7 +800,7 @@ impl Object {
         Ok(())
     }
 
-    fn parse_section(&mut self, mut section: Section) -> Result<(), ParseError> {
+    fn parse_section(&mut self, section: Section) -> Result<(), ParseError> {
         let mut parts = section.name.rsplitn(2, '/').collect::<Vec<_>>();
         parts.reverse();
 
@@ -828,7 +824,7 @@ impl Object {
             BpfSectionKind::BtfExt => self.parse_btf_ext(&section)?,
             BpfSectionKind::BtfMaps => {
                 let symbols: HashMap<String, Symbol> = self
-                    .symbols_by_index
+                    .symbol_table
                     .values()
                     .filter(|s| {
                         if let Some(idx) = s.section_index {
@@ -849,7 +845,7 @@ impl Object {
 
                 // extract the symbols for the .maps section, we'll need them
                 // during parsing
-                let symbols = self.symbols_by_index.values().filter(|s| {
+                let symbols = self.symbol_table.values().filter(|s| {
                     s.section_index
                         .map(|idx| idx == section.index.0)
                         .unwrap_or(false)
@@ -1097,6 +1093,7 @@ impl<'data, 'file, 'a> TryFrom<&'a ObjSection<'data, 'file>> for Section<'a> {
                             _ => return Err(ParseError::UnsupportedRelocationTarget),
                         },
                         offset,
+                        size: r.size(),
                     })
                 })
                 .collect::<Result<Vec<_>, _>>()?,
@@ -1399,8 +1396,8 @@ mod tests {
     }
 
     fn fake_sym(obj: &mut Object, section_index: usize, address: u64, name: &str, size: u64) {
-        let idx = obj.symbols_by_index.len();
-        obj.symbols_by_index.insert(
+        let idx = obj.symbol_table.len();
+        obj.symbol_table.insert(
             idx + 1,
             Symbol {
                 index: idx + 1,
@@ -2213,7 +2210,7 @@ mod tests {
                 data: vec![0, 0, 0],
             }),
         );
-        obj.symbols_by_index.insert(
+        obj.symbol_table.insert(
             1,
             Symbol {
                 index: 1,

aya-obj/src/relocation.rs

@@ -1,6 +1,7 @@
 //! Program relocation handling.
 
 use core::mem;
+use std::collections::HashSet;
 
 use alloc::{borrow::ToOwned, string::String};
 use log::debug;
@@ -85,6 +86,7 @@ pub enum RelocationError {
 pub(crate) struct Relocation {
     // byte offset of the instruction to be relocated
     pub(crate) offset: u64,
+    pub(crate) size: u8,
     // index of the symbol to relocate to
     pub(crate) symbol_index: usize,
 }
@@ -105,6 +107,7 @@ impl Object {
     pub fn relocate_maps<'a, I: Iterator<Item = (&'a str, Option<i32>, &'a Map)>>(
         &mut self,
         maps: I,
+        text_sections: &HashSet<usize>,
     ) -> Result<(), BpfRelocationError> {
         let mut maps_by_section = HashMap::new();
         let mut maps_by_symbol = HashMap::new();
@@ -128,8 +131,8 @@ impl Object {
                     relocations.values(),
                     &maps_by_section,
                     &maps_by_symbol,
-                    &self.symbols_by_index,
-                    self.text_section_index,
+                    &self.symbol_table,
+                    text_sections,
                 )
                 .map_err(|error| BpfRelocationError {
                     function: function.name.clone(),
@@ -142,13 +145,16 @@ impl Object {
     }
 
     /// Relocates function calls
-    pub fn relocate_calls(&mut self) -> Result<(), BpfRelocationError> {
+    pub fn relocate_calls(
+        &mut self,
+        text_sections: &HashSet<usize>,
+    ) -> Result<(), BpfRelocationError> {
         for (name, program) in self.programs.iter_mut() {
             let linker = FunctionLinker::new(
-                self.text_section_index,
                 &self.functions,
                 &self.relocations,
-                &self.symbols_by_index,
+                &self.symbol_table,
+                text_sections,
             );
             linker.link(program).map_err(|error| BpfRelocationError {
                 function: name.to_owned(),
@@ -166,7 +172,7 @@ fn relocate_maps<'a, I: Iterator<Item = &'a Relocation>>(
     maps_by_section: &HashMap<usize, (&str, Option<i32>, &Map)>,
     maps_by_symbol: &HashMap<usize, (&str, Option<i32>, &Map)>,
     symbol_table: &HashMap<usize, Symbol>,
-    text_section_index: Option<usize>,
+    text_sections: &HashSet<usize>,
 ) -> Result<(), RelocationError> {
     let section_offset = fun.section_offset;
     let instructions = &mut fun.instructions;
@@ -202,16 +208,17 @@ fn relocate_maps<'a, I: Iterator<Item = &'a Relocation>>(
         };
 
         // calls and relocation to .text symbols are handled in a separate step
-        if insn_is_call(&instructions[ins_index]) || sym.section_index == text_section_index {
+        if insn_is_call(&instructions[ins_index]) || text_sections.contains(&section_index) {
             continue;
         }
 
         let (name, fd, map) = if let Some(m) = maps_by_symbol.get(&rel.symbol_index) {
             let map = &m.2;
             debug!(
-                "relocating map by symbol index {}, kind {:?}",
-                map.section_index(),
-                map.section_kind()
+                "relocating map by symbol index {:?}, kind {:?} at insn {ins_index} in section {}",
+                map.symbol_index(),
+                map.section_kind(),
+                fun.section_index.0
             );
             debug_assert_eq!(map.symbol_index().unwrap(), rel.symbol_index);
             m
@@ -229,9 +236,10 @@ fn relocate_maps<'a, I: Iterator<Item = &'a Relocation>>(
             };
             let map = &m.2;
             debug!(
-                "relocating map by section index {}, kind {:?}",
+                "relocating map by section index {}, kind {:?} at insn {ins_index} in section {}",
                 map.section_index(),
-                map.section_kind()
+                map.section_kind(),
+                fun.section_index.0,
             );
 
             debug_assert_eq!(map.symbol_index(), None);
@@ -261,26 +269,26 @@ fn relocate_maps<'a, I: Iterator<Item = &'a Relocation>>(
 }
 
 struct FunctionLinker<'a> {
-    text_section_index: Option<usize>,
-    functions: &'a HashMap<u64, Function>,
+    functions: &'a HashMap<(usize, u64), Function>,
     linked_functions: HashMap<u64, usize>,
     relocations: &'a HashMap<SectionIndex, HashMap<u64, Relocation>>,
     symbol_table: &'a HashMap<usize, Symbol>,
+    text_sections: &'a HashSet<usize>,
 }
 
 impl<'a> FunctionLinker<'a> {
     fn new(
-        text_section_index: Option<usize>,
-        functions: &'a HashMap<u64, Function>,
+        functions: &'a HashMap<(usize, u64), Function>,
         relocations: &'a HashMap<SectionIndex, HashMap<u64, Relocation>>,
         symbol_table: &'a HashMap<usize, Symbol>,
+        text_sections: &'a HashSet<usize>,
     ) -> FunctionLinker<'a> {
         FunctionLinker {
-            text_section_index,
             functions,
             linked_functions: HashMap::new(),
             relocations,
             symbol_table,
+            text_sections,
         }
     }
 
@@ -310,6 +318,10 @@ impl<'a> FunctionLinker<'a> {
         // at `start_ins`. We'll use `start_ins` to do pc-relative calls.
         let start_ins = program.instructions.len();
         program.instructions.extend(&fun.instructions);
+        debug!(
+            "linked function `{}` at instruction {}",
+            fun.name, start_ins
+        );
 
         // link func and line info into the main program
         // the offset needs to be adjusted
@@ -326,101 +338,110 @@ impl<'a> FunctionLinker<'a> {
     fn relocate(&mut self, program: &mut Function, fun: &Function) -> Result<(), RelocationError> {
         let relocations = self.relocations.get(&fun.section_index);
 
-        debug!("relocating program {} function {}", program.name, fun.name);
-
         let n_instructions = fun.instructions.len();
         let start_ins = program.instructions.len() - n_instructions;
 
+        debug!(
+            "relocating program `{}` function `{}` size {}",
+            program.name, fun.name, n_instructions
+        );
+
         // process all the instructions. We can't only loop over relocations since we need to
         // patch pc-relative calls too.
         for ins_index in start_ins..start_ins + n_instructions {
             let ins = program.instructions[ins_index];
             let is_call = insn_is_call(&ins);
 
-            // only resolve relocations for calls or for instructions that
-            // reference symbols in the .text section (eg let callback =
-            // &some_fun)
-            let rel = if let Some(relocations) = relocations {
-                self.text_relocation_info(
-                    relocations,
-                    (fun.section_offset + (ins_index - start_ins) * INS_SIZE) as u64,
-                )?
-                // if not a call and not a .text reference, ignore the
-                // relocation (see relocate_maps())
-                .and_then(|(_, sym)| {
-                    if is_call {
-                        return Some(sym.address);
-                    }
-
-                    match sym.kind {
-                        SymbolKind::Text => Some(sym.address),
-                        SymbolKind::Section if sym.section_index == self.text_section_index => {
-                            Some(sym.address + ins.imm as u64)
-                        }
-                        _ => None,
-                    }
+            let rel = relocations
+                .and_then(|relocations| {
+                    relocations
+                        .get(&((fun.section_offset + (ins_index - start_ins) * INS_SIZE) as u64))
                 })
-            } else {
-                None
-            };
+                .and_then(|rel| {
+                    // get the symbol for the relocation
+                    self.symbol_table
+                        .get(&rel.symbol_index)
+                        .map(|sym| (rel, sym))
+                })
+                .filter(|(_rel, sym)| {
+                    // only consider text relocations, data relocations are
+                    // relocated in relocate_maps()
+                    sym.kind == SymbolKind::Text
+                        || sym
+                            .section_index
+                            .map(|section_index| self.text_sections.contains(&section_index))
+                            .unwrap_or(false)
+                });
 
-            // some_fun() or let x = &some_fun trigger linking, everything else
-            // can be ignored here
+            // not a call and not a text relocation, we don't need to do anything
             if !is_call && rel.is_none() {
                 continue;
             }
 
-            let callee_address = if let Some(address) = rel {
-                // We have a relocation entry for the instruction at `ins_index`, the address of
-                // the callee is the address of the relocation's target symbol.
-                address
+            let (callee_section_index, callee_address) = if let Some((rel, sym)) = rel {
+                let address = match sym.kind {
+                    SymbolKind::Text => sym.address,
+                    // R_BPF_64_32 this is a call
+                    SymbolKind::Section if rel.size == 32 => {
+                        sym.address + (ins.imm + 1) as u64 * INS_SIZE as u64
+                    }
+                    // R_BPF_64_64 this is a ld_imm64 text relocation
+                    SymbolKind::Section if rel.size == 64 => sym.address + ins.imm as u64,
+                    _ => todo!(), // FIXME: return an error here,
+                };
+                (sym.section_index.unwrap(), address)
             } else {
                 // The caller and the callee are in the same ELF section and this is a pc-relative
                 // call. Resolve the pc-relative imm to an absolute address.
                 let ins_size = INS_SIZE as i64;
-                (fun.section_offset as i64
-                    + ((ins_index - start_ins) as i64) * ins_size
-                    + (ins.imm + 1) as i64 * ins_size) as u64
+                (
+                    fun.section_index.0,
+                    (fun.section_offset as i64
+                        + ((ins_index - start_ins) as i64) * ins_size
+                        + (ins.imm + 1) as i64 * ins_size) as u64,
+                )
             };
 
             debug!(
-                "relocating {} to callee address {} ({})",
+                "relocating {} to callee address {:#x} in section {} ({}) at instruction {ins_index}",
                 if is_call { "call" } else { "reference" },
                 callee_address,
+                callee_section_index,
                 if rel.is_some() {
                     "relocation"
                 } else {
-                    "relative"
+                    "pc-relative"
                 },
             );
 
             // lookup and link the callee if it hasn't been linked already. `callee_ins_index` will
             // contain the instruction index of the callee inside the program.
-            let callee =
-                self.functions
-                    .get(&callee_address)
-                    .ok_or(RelocationError::UnknownFunction {
-                        address: callee_address,
-                        caller_name: fun.name.clone(),
-                    })?;
+            let callee = self
+                .functions
+                .get(&(callee_section_index, callee_address))
+                .ok_or(RelocationError::UnknownFunction {
+                    address: callee_address,
+                    caller_name: fun.name.clone(),
+                })?;
 
-            debug!("callee is {}", callee.name);
+            debug!("callee is `{}`", callee.name);
 
-            let callee_ins_index = self.link_function(program, callee)?;
+            let callee_ins_index = self.link_function(program, callee)? as i32;
 
             let mut ins = &mut program.instructions[ins_index];
-            ins.imm = if callee_ins_index < ins_index {
-                -((ins_index - callee_ins_index + 1) as i32)
-            } else {
-                (callee_ins_index - ins_index - 1) as i32
-            };
+            let ins_index = ins_index as i32;
+            ins.imm = callee_ins_index - ins_index - 1;
+            debug!(
+                "callee `{}` is at ins {callee_ins_index}, {} from current instruction {ins_index}",
+                callee.name, ins.imm
+            );
             if !is_call {
                 ins.set_src_reg(BPF_PSEUDO_FUNC as u8);
             }
         }
 
         debug!(
-            "finished relocating program {} function {}",
+            "finished relocating program `{}` function `{}`",
             program.name, fun.name
         );
 
@@ -459,25 +480,6 @@ impl<'a> FunctionLinker<'a> {
         }
         Ok(())
     }
-
-    fn text_relocation_info(
-        &self,
-        relocations: &HashMap<u64, Relocation>,
-        offset: u64,
-    ) -> Result<Option<(Relocation, Symbol)>, RelocationError> {
-        if let Some(rel) = relocations.get(&offset) {
-            let sym =
-                self.symbol_table
-                    .get(&rel.symbol_index)
-                    .ok_or(RelocationError::UnknownSymbol {
-                        index: rel.symbol_index,
-                    })?;
-
-            Ok(Some((*rel, sym.clone())))
-        } else {
-            Ok(None)
-        }
-    }
 }
 
 fn insn_is_call(ins: &bpf_insn) -> bool {
@@ -498,7 +500,7 @@ mod test {
     use alloc::{string::ToString, vec, vec::Vec};
 
     use crate::{
-        maps::{bpf_map_def, BtfMap, BtfMapDef, LegacyMap, Map},
+        maps::{BtfMap, LegacyMap, Map},
         BpfSectionKind,
     };
 
@@ -568,6 +570,7 @@ mod test {
         let relocations = vec![Relocation {
             offset: 0x0,
             symbol_index: 1,
+            size: 64,
         }];
         let maps_by_section = HashMap::new();
 
@@ -580,7 +583,7 @@ mod test {
             &maps_by_section,
             &maps_by_symbol,
             &symbol_table,
-            None,
+            &HashSet::new(),
         )
         .unwrap();
 
@@ -615,10 +618,12 @@ mod test {
             Relocation {
                 offset: 0x0,
                 symbol_index: 1,
+                size: 64,
             },
             Relocation {
                 offset: mem::size_of::<bpf_insn>() as u64,
                 symbol_index: 2,
+                size: 64,
             },
         ];
         let maps_by_section = HashMap::new();
@@ -636,7 +641,7 @@ mod test {
             &maps_by_section,
             &maps_by_symbol,
             &symbol_table,
-            None,
+            &HashSet::new(),
         )
         .unwrap();
 
@@ -665,6 +670,7 @@ mod test {
         let relocations = vec![Relocation {
             offset: 0x0,
             symbol_index: 1,
+            size: 64,
         }];
         let maps_by_section = HashMap::new();
 
@@ -677,7 +683,7 @@ mod test {
             &maps_by_section,
             &maps_by_symbol,
             &symbol_table,
-            None,
+            &HashSet::new(),
         )
         .unwrap();
 
@@ -712,10 +718,12 @@ mod test {
             Relocation {
                 offset: 0x0,
                 symbol_index: 1,
+                size: 64,
             },
             Relocation {
                 offset: mem::size_of::<bpf_insn>() as u64,
                 symbol_index: 2,
+                size: 64,
             },
         ];
         let maps_by_section = HashMap::new();
@@ -733,7 +741,7 @@ mod test {
             &maps_by_section,
             &maps_by_symbol,
             &symbol_table,
-            None,
+            &HashSet::new(),
         )
         .unwrap();
 

aya/src/bpf.rs

@@ -431,11 +431,18 @@ impl<'a> BpfLoader<'a> {
             maps.insert(name, map);
         }
 
+        let text_sections = obj
+            .functions
+            .keys()
+            .map(|(section_index, _)| *section_index)
+            .collect();
+
         obj.relocate_maps(
             maps.iter()
                 .map(|(s, data)| (s.as_str(), data.fd, &data.obj)),
+            &text_sections,
         )?;
-        obj.relocate_calls()?;
+        obj.relocate_calls(&text_sections)?;
 
         let programs = obj
             .programs

test/integration-test/src/tests/rbpf.rs

@@ -69,14 +69,20 @@ fn use_map_with_rbpf() {
         }
     }
 
+    let text_sections = object
+        .functions
+        .iter()
+        .map(|((section_index, _), _)| *section_index)
+        .collect();
     object
         .relocate_maps(
             maps.iter()
                 .map(|(s, (fd, map))| (s.as_ref() as &str, Some(*fd), map)),
+            &text_sections,
         )
         .expect("Relocation failed");
     // Actually there is no local function call involved.
-    object.relocate_calls().unwrap();
+    object.relocate_calls(&text_sections).unwrap();
 
     // Executes the program
     assert_eq!(object.programs.len(), 1);