Merge pull request #572 from alessandrod/reloc-fixes

Add tests for relocations + fixes

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/maps.rs

@@ -2,7 +2,10 @@
 
 use core::mem;
 
-use crate::thiserror::{self, Error};
+use crate::{
+    thiserror::{self, Error},
+    BpfSectionKind,
+};
 use alloc::vec::Vec;
 
 /// Invalid map type encontered
@@ -139,33 +142,6 @@ pub struct bpf_map_def {
 /// The first five __u32 of `bpf_map_def` must be defined.
 pub(crate) const MINIMUM_MAP_SIZE: usize = mem::size_of::<u32>() * 5;
 
-/// Kinds of maps
-#[derive(Debug, Copy, Clone, PartialEq, Eq)]
-pub enum MapKind {
-    /// A map holding `.bss` section data
-    Bss,
-    /// A map holding `.data` section data
-    Data,
-    /// A map holding `.rodata` section data
-    Rodata,
-    /// Other maps
-    Other,
-}
-
-impl From<&str> for MapKind {
-    fn from(s: &str) -> Self {
-        if s == ".bss" {
-            MapKind::Bss
-        } else if s.starts_with(".data") {
-            MapKind::Data
-        } else if s.starts_with(".rodata") {
-            MapKind::Rodata
-        } else {
-            MapKind::Other
-        }
-    }
-}
-
 /// Map data defined in `maps` or `.maps` sections
 #[derive(Debug, Clone)]
 pub enum Map {
@@ -248,14 +224,6 @@ impl Map {
         }
     }
 
-    /// Returns the map kind
-    pub fn kind(&self) -> MapKind {
-        match self {
-            Map::Legacy(m) => m.kind,
-            Map::Btf(m) => m.kind,
-        }
-    }
-
     /// Returns the section index
     pub fn section_index(&self) -> usize {
         match self {
@@ -264,11 +232,22 @@ impl Map {
         }
     }
 
-    /// Returns the symbol index
-    pub fn symbol_index(&self) -> usize {
+    /// Returns the section kind.
+    pub fn section_kind(&self) -> BpfSectionKind {
+        match self {
+            Map::Legacy(m) => m.section_kind,
+            Map::Btf(_) => BpfSectionKind::BtfMaps,
+        }
+    }
+
+    /// Returns the symbol index.
+    ///
+    /// This is `None` for data maps (.bss, .data and .rodata) since those don't
+    /// need symbols in order to be relocated.
+    pub fn symbol_index(&self) -> Option<usize> {
         match self {
             Map::Legacy(m) => m.symbol_index,
-            Map::Btf(m) => m.symbol_index,
+            Map::Btf(m) => Some(m.symbol_index),
         }
     }
 }
@@ -283,12 +262,16 @@ pub struct LegacyMap {
     pub def: bpf_map_def,
     /// The section index
     pub section_index: usize,
-    /// The symbol index
-    pub symbol_index: usize,
+    /// The section kind
+    pub section_kind: BpfSectionKind,
+    /// The symbol index.
+    ///
+    /// This is None for data maps (.bss .data and .rodata).  We don't need
+    /// symbols to relocate those since they don't contain multiple maps, but
+    /// are just a flat array of bytes.
+    pub symbol_index: Option<usize>,
     /// The map data
     pub data: Vec<u8>,
-    /// The map kind
-    pub kind: MapKind,
 }
 
 /// A BTF-defined map, most likely from a `.maps` section.
@@ -298,6 +281,5 @@ pub struct BtfMap {
     pub def: BtfMapDef,
     pub(crate) section_index: usize,
     pub(crate) symbol_index: usize,
-    pub(crate) kind: MapKind,
     pub(crate) data: Vec<u8>,
 }

aya-obj/src/obj.rs

@@ -15,7 +15,7 @@ use object::{
 };
 
 use crate::{
-    maps::{BtfMap, LegacyMap, Map, MapKind, MINIMUM_MAP_SIZE},
+    maps::{BtfMap, LegacyMap, Map, MINIMUM_MAP_SIZE},
     relocation::*,
     thiserror::{self, Error},
     util::HashMap,
@@ -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(),
@@ -753,7 +749,7 @@ impl Object {
                 section.index,
                 section
                     .relocations
-                    .drain(..)
+                    .into_iter()
                     .map(|rel| (rel.offset, rel))
                     .collect(),
             );
@@ -762,37 +758,6 @@ impl Object {
         Ok(())
     }
 
-    fn parse_map_section(
-        &mut self,
-        section: &Section,
-        symbols: Vec<Symbol>,
-    ) -> Result<(), ParseError> {
-        if symbols.is_empty() {
-            return Err(ParseError::NoSymbolsInMapSection {});
-        }
-        for (i, sym) in symbols.iter().enumerate() {
-            let start = sym.address as usize;
-            let end = start + sym.size as usize;
-            let data = &section.data[start..end];
-            let name = sym
-                .name
-                .as_ref()
-                .ok_or(ParseError::MapSymbolNameNotFound { i })?;
-            let def = parse_map_def(name, data)?;
-            self.maps.insert(
-                name.to_string(),
-                Map::Legacy(LegacyMap {
-                    section_index: section.index.0,
-                    symbol_index: sym.index,
-                    def,
-                    data: Vec::new(),
-                    kind: MapKind::Other,
-                }),
-            );
-        }
-        Ok(())
-    }
-
     fn parse_btf_maps(
         &mut self,
         section: &Section,
@@ -825,7 +790,6 @@ impl Object {
                                 def,
                                 section_index: section.index.0,
                                 symbol_index,
-                                kind: MapKind::Other,
                                 data: Vec::new(),
                             }),
                         );
@@ -836,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();
 
@@ -851,16 +815,16 @@ impl Object {
         self.section_sizes
             .insert(section.name.to_owned(), section.size);
         match section.kind {
-            BpfSectionKind::Data => {
+            BpfSectionKind::Data | BpfSectionKind::Rodata | BpfSectionKind::Bss => {
                 self.maps
-                    .insert(section.name.to_string(), parse_map(&section, section.name)?);
+                    .insert(section.name.to_string(), parse_data_map_section(&section)?);
             }
             BpfSectionKind::Text => self.parse_text_section(section)?,
             BpfSectionKind::Btf => self.parse_btf(&section)?,
             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 {
@@ -875,19 +839,24 @@ impl Object {
                 self.parse_btf_maps(&section, symbols)?
             }
             BpfSectionKind::Maps => {
-                let symbols: Vec<Symbol> = self
-                    .symbols_by_index
-                    .values()
-                    .filter(|s| {
-                        if let Some(idx) = s.section_index {
-                            idx == section.index.0
-                        } else {
-                            false
-                        }
-                    })
-                    .cloned()
-                    .collect();
-                self.parse_map_section(&section, symbols)?
+                // take out self.maps so we can borrow the iterator below
+                // without cloning or collecting
+                let mut maps = mem::take(&mut self.maps);
+
+                // extract the symbols for the .maps section, we'll need them
+                // during parsing
+                let symbols = self.symbol_table.values().filter(|s| {
+                    s.section_index
+                        .map(|idx| idx == section.index.0)
+                        .unwrap_or(false)
+                });
+
+                let res = parse_maps_section(&mut maps, &section, symbols);
+
+                // put the maps back
+                self.maps = maps;
+
+                res?
             }
             BpfSectionKind::Program => {
                 let program = self.parse_program(&section)?;
@@ -898,7 +867,7 @@ impl Object {
                         section.index,
                         section
                             .relocations
-                            .drain(..)
+                            .into_iter()
                             .map(|rel| (rel.offset, rel))
                             .collect(),
                     );
@@ -911,6 +880,45 @@ impl Object {
     }
 }
 
+// Parses multiple map definition contained in a single `maps` section (which is
+// different from `.maps` which is used for BTF). We can tell where each map is
+// based on the symbol table.
+fn parse_maps_section<'a, I: Iterator<Item = &'a Symbol>>(
+    maps: &mut HashMap<String, Map>,
+    section: &Section,
+    symbols: I,
+) -> Result<(), ParseError> {
+    let mut have_symbols = false;
+
+    // each symbol in the section is a  separate map
+    for (i, sym) in symbols.enumerate() {
+        let start = sym.address as usize;
+        let end = start + sym.size as usize;
+        let data = &section.data[start..end];
+        let name = sym
+            .name
+            .as_ref()
+            .ok_or(ParseError::MapSymbolNameNotFound { i })?;
+        let def = parse_map_def(name, data)?;
+        maps.insert(
+            name.to_string(),
+            Map::Legacy(LegacyMap {
+                section_index: section.index.0,
+                section_kind: section.kind,
+                symbol_index: Some(sym.index),
+                def,
+                data: Vec::new(),
+            }),
+        );
+        have_symbols = true;
+    }
+    if !have_symbols {
+        return Err(ParseError::NoSymbolsForMapsSection);
+    }
+
+    Ok(())
+}
+
 /// Errors caught during parsing the object file
 #[derive(Debug, Error)]
 #[allow(missing_docs)]
@@ -974,25 +982,40 @@ pub enum ParseError {
     #[error("the map number {i} in the `maps` section doesn't have a symbol name")]
     MapSymbolNameNotFound { i: usize },
 
-    #[error("no symbols found for the maps included in the maps section")]
-    NoSymbolsInMapSection {},
+    #[error("no symbols for `maps` section, can't parse maps")]
+    NoSymbolsForMapsSection,
 
     /// No BTF parsed for object
     #[error("no BTF parsed for object")]
     NoBTF,
 }
 
-#[derive(Debug)]
-enum BpfSectionKind {
+/// The kind of an ELF section.
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub enum BpfSectionKind {
+    /// Undefined
     Undefined,
+    /// `maps`
     Maps,
+    /// `.maps`
     BtfMaps,
+    /// A program section
     Program,
+    /// `.data`
     Data,
+    /// `.rodata`
+    Rodata,
+    /// `.bss`
+    Bss,
+    /// `.text`
     Text,
+    /// `.BTF`
     Btf,
+    /// `.BTF.ext`
     BtfExt,
+    /// `license`
     License,
+    /// `version`
     Version,
 }
 
@@ -1070,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<_>, _>>()?,
@@ -1159,10 +1183,11 @@ impl From<KernelVersion> for u32 {
     }
 }
 
-fn parse_map(section: &Section, name: &str) -> Result<Map, ParseError> {
-    let kind = MapKind::from(name);
-    let (def, data) = match kind {
-        MapKind::Bss | MapKind::Data | MapKind::Rodata => {
+// Parsed '.bss' '.data' and '.rodata' sections. These sections are arrays of
+// bytes and are relocated based on their section index.
+fn parse_data_map_section(section: &Section) -> Result<Map, ParseError> {
+    let (def, data) = match section.kind {
+        BpfSectionKind::Bss | BpfSectionKind::Data | BpfSectionKind::Rodata => {
             let def = bpf_map_def {
                 map_type: BPF_MAP_TYPE_ARRAY as u32,
                 key_size: mem::size_of::<u32>() as u32,
@@ -1170,7 +1195,7 @@ fn parse_map(section: &Section, name: &str) -> Result<Map, ParseError> {
                 // .bss will always have data.len() == 0
                 value_size: section.size as u32,
                 max_entries: 1,
-                map_flags: if kind == MapKind::Rodata {
+                map_flags: if section.kind == BpfSectionKind::Rodata {
                     BPF_F_RDONLY_PROG
                 } else {
                     0
@@ -1179,14 +1204,15 @@ fn parse_map(section: &Section, name: &str) -> Result<Map, ParseError> {
             };
             (def, section.data.to_vec())
         }
-        MapKind::Other => (parse_map_def(name, section.data)?, Vec::new()),
+        _ => unreachable!(),
     };
     Ok(Map::Legacy(LegacyMap {
         section_index: section.index.0,
-        symbol_index: 0,
+        section_kind: section.kind,
+        // Data maps don't require symbols to be relocated
+        symbol_index: None,
         def,
         data,
-        kind,
     }))
 }
 
@@ -1306,8 +1332,6 @@ pub fn parse_map_info(info: bpf_map_info, pinned: PinningType) -> Map {
             section_index: 0,
             symbol_index: 0,
             data: Vec::new(),
-            // We should never be loading the .bss or .data or .rodata FDs
-            kind: MapKind::Other,
         })
     } else {
         Map::Legacy(LegacyMap {
@@ -1321,10 +1345,9 @@ pub fn parse_map_info(info: bpf_map_info, pinned: PinningType) -> Map {
                 id: info.id,
             },
             section_index: 0,
-            symbol_index: 0,
+            symbol_index: None,
+            section_kind: BpfSectionKind::Undefined,
             data: Vec::new(),
-            // We should never be loading the .bss or .data or .rodata FDs
-            kind: MapKind::Other,
         })
     }
 }
@@ -1373,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,
@@ -1510,65 +1533,21 @@ mod tests {
         assert_eq!(parse_map_def("foo", &buf).unwrap(), def);
     }
 
-    #[test]
-    fn test_parse_map_error() {
-        assert!(matches!(
-            parse_map(&fake_section(BpfSectionKind::Maps, "maps/foo", &[]), "foo",),
-            Err(ParseError::InvalidMapDefinition { .. })
-        ));
-    }
-
-    #[test]
-    fn test_parse_map() {
-        assert!(matches!(
-            parse_map(
-                &fake_section(
-                    BpfSectionKind::Maps,
-                    "maps/foo",
-                    bytes_of(&bpf_map_def {
-                        map_type: 1,
-                        key_size: 2,
-                        value_size: 3,
-                        max_entries: 4,
-                        map_flags: 5,
-                        id: 0,
-                        pinning: PinningType::None,
-                    })
-                ),
-                "foo"
-            ),
-            Ok(Map::Legacy(LegacyMap{
-                section_index: 0,
-                def: bpf_map_def {
-                    map_type: 1,
-                    key_size: 2,
-                    value_size: 3,
-                    max_entries: 4,
-                    map_flags: 5,
-                    id: 0,
-                    pinning: PinningType::None,
-                },
-                data,
-                ..
-            })) if data.is_empty()
-        ))
-    }
-
     #[test]
     fn test_parse_map_data() {
         let map_data = b"map data";
         assert!(matches!(
-            parse_map(
+            parse_data_map_section(
                 &fake_section(
                     BpfSectionKind::Data,
                     ".bss",
                     map_data,
                 ),
-                ".bss"
             ),
             Ok(Map::Legacy(LegacyMap {
                 section_index: 0,
-                symbol_index: 0,
+                section_kind: BpfSectionKind::Data,
+                symbol_index: None,
                 def: bpf_map_def {
                     map_type: _map_type,
                     key_size: 4,
@@ -1579,8 +1558,7 @@ mod tests {
                     pinning: PinningType::None,
                 },
                 data,
-                kind
-            })) if data == map_data && value_size == map_data.len() as u32 && kind == MapKind::Bss
+            })) if data == map_data && value_size == map_data.len() as u32
         ))
     }
 
@@ -2227,12 +2205,12 @@ mod tests {
                     pinning: PinningType::None,
                 },
                 section_index: 1,
-                symbol_index: 1,
+                section_kind: BpfSectionKind::Rodata,
+                symbol_index: Some(1),
                 data: vec![0, 0, 0],
-                kind: MapKind::Rodata,
             }),
         );
-        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;
@@ -15,6 +16,7 @@ use crate::{
     obj::{Function, Object, Program},
     thiserror::{self, Error},
     util::HashMap,
+    BpfSectionKind,
 };
 
 pub(crate) const INS_SIZE: usize = mem::size_of::<bpf_insn>();
@@ -84,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,
 }
@@ -104,12 +107,15 @@ 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();
         for (name, fd, map) in maps {
             maps_by_section.insert(map.section_index(), (name, fd, map));
-            maps_by_symbol.insert(map.symbol_index(), (name, fd, map));
+            if let Some(index) = map.symbol_index() {
+                maps_by_symbol.insert(index, (name, fd, map));
+            }
         }
 
         let functions = self
@@ -125,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(),
@@ -139,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(),
@@ -163,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;
@@ -193,34 +202,54 @@ fn relocate_maps<'a, I: Iterator<Item = &'a Relocation>>(
                 index: rel.symbol_index,
             })?;
 
-        let section_index = match sym.section_index {
-            Some(index) => index,
+        let Some(section_index) = sym.section_index else {
             // this is not a map relocation
-            None => continue,
+            continue;
         };
 
         // 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 maps_by_symbol.contains_key(&rel.symbol_index) {
-            maps_by_symbol
-                .get(&rel.symbol_index)
-                .ok_or(RelocationError::SectionNotFound {
-                    symbol_index: rel.symbol_index,
-                    symbol_name: sym.name.clone(),
-                    section_index,
-                })?
+        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 {:?} 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
         } else {
-            maps_by_section
-                .get(&section_index)
-                .ok_or(RelocationError::SectionNotFound {
+            let Some(m) = maps_by_section.get(&section_index) else {
+                debug!(
+                    "failed relocating map by section index {}",
+                    section_index
+                );
+                return Err(RelocationError::SectionNotFound {
                     symbol_index: rel.symbol_index,
                     symbol_name: sym.name.clone(),
                     section_index,
-                })?
+                });
+            };
+            let map = &m.2;
+            debug!(
+                "relocating map by section index {}, kind {:?} at insn {ins_index} in section {}",
+                map.section_index(),
+                map.section_kind(),
+                fun.section_index.0,
+            );
+
+            debug_assert_eq!(map.symbol_index(), None);
+            debug_assert!(matches!(
+                map.section_kind(),
+                BpfSectionKind::Bss | BpfSectionKind::Data | BpfSectionKind::Rodata
+            ),);
+            m
         };
+        debug_assert_eq!(map.section_index(), section_index);
 
         let map_fd = fd.ok_or_else(|| RelocationError::MapNotCreated {
             name: (*name).into(),
@@ -240,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,
         }
     }
 
@@ -289,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
@@ -305,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
-            };
-
-            // some_fun() or let x = &some_fun trigger linking, everything else
-            // can be ignored here
+                .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)
+                });
+
+            // 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
         );
 
@@ -438,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 {
@@ -476,7 +499,10 @@ fn insn_is_call(ins: &bpf_insn) -> bool {
 mod test {
     use alloc::{string::ToString, vec, vec::Vec};
 
-    use crate::maps::{bpf_map_def, BtfMap, BtfMapDef, LegacyMap, Map, MapKind};
+    use crate::{
+        maps::{BtfMap, LegacyMap, Map},
+        BpfSectionKind,
+    };
 
     use super::*;
 
@@ -498,25 +524,20 @@ mod test {
 
     fn fake_legacy_map(symbol_index: usize) -> Map {
         Map::Legacy(LegacyMap {
-            def: bpf_map_def {
-                ..Default::default()
-            },
+            def: Default::default(),
             section_index: 0,
-            symbol_index,
+            section_kind: BpfSectionKind::Undefined,
+            symbol_index: Some(symbol_index),
             data: Vec::new(),
-            kind: MapKind::Other,
         })
     }
 
     fn fake_btf_map(symbol_index: usize) -> Map {
         Map::Btf(BtfMap {
-            def: BtfMapDef {
-                ..Default::default()
-            },
+            def: Default::default(),
             section_index: 0,
             symbol_index,
             data: Vec::new(),
-            kind: MapKind::Other,
         })
     }
 
@@ -549,6 +570,7 @@ mod test {
         let relocations = vec![Relocation {
             offset: 0x0,
             symbol_index: 1,
+            size: 64,
         }];
         let maps_by_section = HashMap::new();
 
@@ -561,7 +583,7 @@ mod test {
             &maps_by_section,
             &maps_by_symbol,
             &symbol_table,
-            None,
+            &HashSet::new(),
         )
         .unwrap();
 
@@ -596,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();
@@ -617,7 +641,7 @@ mod test {
             &maps_by_section,
             &maps_by_symbol,
             &symbol_table,
-            None,
+            &HashSet::new(),
         )
         .unwrap();
 
@@ -646,6 +670,7 @@ mod test {
         let relocations = vec![Relocation {
             offset: 0x0,
             symbol_index: 1,
+            size: 64,
         }];
         let maps_by_section = HashMap::new();
 
@@ -658,7 +683,7 @@ mod test {
             &maps_by_section,
             &maps_by_symbol,
             &symbol_table,
-            None,
+            &HashSet::new(),
         )
         .unwrap();
 
@@ -693,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();
@@ -714,7 +741,7 @@ mod test {
             &maps_by_section,
             &maps_by_symbol,
             &symbol_table,
-            None,
+            &HashSet::new(),
         )
         .unwrap();
 

aya/src/bpf.rs

@@ -11,6 +11,7 @@ use aya_obj::{
     btf::{BtfFeatures, BtfRelocationError},
     generated::BPF_F_XDP_HAS_FRAGS,
     relocation::BpfRelocationError,
+    BpfSectionKind,
 };
 use log::debug;
 use thiserror::Error;
@@ -23,7 +24,6 @@ use crate::{
     maps::{Map, MapData, MapError},
     obj::{
         btf::{Btf, BtfError},
-        maps::MapKind,
         Object, ParseError, ProgramSection,
     },
     programs::{
@@ -415,14 +415,14 @@ impl<'a> BpfLoader<'a> {
                 }
                 PinningType::None => map.create(&name)?,
             };
-            if !map.obj.data().is_empty() && map.obj.kind() != MapKind::Bss {
+            if !map.obj.data().is_empty() && map.obj.section_kind() != BpfSectionKind::Bss {
                 bpf_map_update_elem_ptr(fd, &0 as *const _, map.obj.data_mut().as_mut_ptr(), 0)
                     .map_err(|(_, io_error)| MapError::SyscallError {
                         call: "bpf_map_update_elem".to_owned(),
                         io_error,
                     })?;
             }
-            if map.obj.kind() == MapKind::Rodata {
+            if map.obj.section_kind() == BpfSectionKind::Rodata {
                 bpf_map_freeze(fd).map_err(|(_, io_error)| MapError::SyscallError {
                     call: "bpf_map_freeze".to_owned(),
                     io_error,
@@ -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

aya/src/maps/bloom_filter.rs

@@ -84,10 +84,7 @@ mod tests {
             bpf_map_type::{BPF_MAP_TYPE_BLOOM_FILTER, BPF_MAP_TYPE_PERF_EVENT_ARRAY},
         },
         maps::{Map, MapData},
-        obj::{
-            self,
-            maps::{LegacyMap, MapKind},
-        },
+        obj::{self, maps::LegacyMap, BpfSectionKind},
         sys::{override_syscall, SysResult, Syscall},
     };
     use libc::{EFAULT, ENOENT};
@@ -103,9 +100,9 @@ mod tests {
                 ..Default::default()
             },
             section_index: 0,
-            symbol_index: 0,
+            section_kind: BpfSectionKind::Maps,
+            symbol_index: None,
             data: Vec::new(),
-            kind: MapKind::Other,
         })
     }
 
@@ -142,9 +139,9 @@ mod tests {
                     ..Default::default()
                 },
                 section_index: 0,
-                symbol_index: 0,
+                section_kind: BpfSectionKind::Maps,
+                symbol_index: None,
                 data: Vec::new(),
-                kind: MapKind::Other,
             }),
             fd: None,
             pinned: false,

aya/src/maps/hash_map/hash_map.rs

@@ -117,10 +117,7 @@ mod tests {
             bpf_map_type::{BPF_MAP_TYPE_HASH, BPF_MAP_TYPE_LRU_HASH},
         },
         maps::{Map, MapData},
-        obj::{
-            self,
-            maps::{LegacyMap, MapKind},
-        },
+        obj::{self, maps::LegacyMap, BpfSectionKind},
         sys::{override_syscall, SysResult, Syscall},
     };
 
@@ -136,9 +133,9 @@ mod tests {
                 ..Default::default()
             },
             section_index: 0,
+            section_kind: BpfSectionKind::Maps,
             data: Vec::new(),
-            kind: MapKind::Other,
-            symbol_index: 0,
+            symbol_index: None,
         })
     }
 
@@ -267,9 +264,9 @@ mod tests {
                     ..Default::default()
                 },
                 section_index: 0,
-                symbol_index: 0,
+                section_kind: BpfSectionKind::Maps,
+                symbol_index: None,
                 data: Vec::new(),
-                kind: MapKind::Other,
             }),
             fd: Some(42),
             pinned: false,

aya/src/maps/lpm_trie.rs

@@ -247,10 +247,7 @@ mod tests {
             bpf_map_type::{BPF_MAP_TYPE_LPM_TRIE, BPF_MAP_TYPE_PERF_EVENT_ARRAY},
         },
         maps::{Map, MapData},
-        obj::{
-            self,
-            maps::{LegacyMap, MapKind},
-        },
+        obj::{self, maps::LegacyMap, BpfSectionKind},
         sys::{override_syscall, SysResult, Syscall},
     };
     use libc::{EFAULT, ENOENT};
@@ -266,9 +263,9 @@ mod tests {
                 ..Default::default()
             },
             section_index: 0,
-            symbol_index: 0,
+            section_kind: BpfSectionKind::Maps,
+            symbol_index: None,
             data: Vec::new(),
-            kind: MapKind::Other,
         })
     }
 
@@ -322,9 +319,9 @@ mod tests {
                     ..Default::default()
                 },
                 section_index: 0,
-                symbol_index: 0,
+                section_kind: BpfSectionKind::Maps,
+                symbol_index: None,
                 data: Vec::new(),
-                kind: MapKind::Other,
             }),
             fd: None,
             btf_fd: None,

aya/src/maps/mod.rs

@@ -845,7 +845,7 @@ mod tests {
         bpf_map_def,
         generated::{bpf_cmd, bpf_map_type::BPF_MAP_TYPE_HASH},
         maps::MapData,
-        obj::maps::{LegacyMap, MapKind},
+        obj::{maps::LegacyMap, BpfSectionKind},
         sys::{override_syscall, Syscall},
     };
 
@@ -861,9 +861,9 @@ mod tests {
                 ..Default::default()
             },
             section_index: 0,
-            symbol_index: 0,
+            section_kind: BpfSectionKind::Maps,
+            symbol_index: Some(0),
             data: Vec::new(),
-            kind: MapKind::Other,
         })
     }
 

test/integration-ebpf/Cargo.toml

@@ -22,3 +22,7 @@ path = "src/pass.rs"
 [[bin]]
 name = "test"
 path = "src/test.rs"
+
+[[bin]]
+name = "relocations"
+path = "src/relocations.rs"

test/integration-ebpf/src/bpf/text_64_64_reloc.c

@@ -0,0 +1,28 @@
+#include <linux/bpf.h>
+#include <bpf/bpf_helpers.h>
+
+char _license[] SEC("license") = "GPL";
+
+struct {
+        __uint(type, BPF_MAP_TYPE_ARRAY);
+        __type(key, __u32);
+        __type(value, __u64);
+        __uint(max_entries, 2);
+} RESULTS SEC(".maps");
+
+static __u64
+inc_cb(void *map, __u32 *key, void *val,
+	       void *data)
+{
+	__u64 *value = val;
+	*value += 1;
+	return 0;
+}
+
+SEC("uprobe/test_text_64_64_reloc")
+int test_text_64_64_reloc(struct pt_regs *ctx)
+{
+	bpf_for_each_map_elem(&RESULTS, inc_cb, NULL, 0);
+    return 0;
+}
+

test/integration-ebpf/src/relocations.rs

@@ -0,0 +1,45 @@
+#![no_std]
+#![no_main]
+
+use core::hint;
+
+use aya_bpf::{
+    macros::{map, uprobe},
+    maps::Array,
+    programs::ProbeContext,
+};
+
+#[map]
+static mut RESULTS: Array<u64> = Array::with_max_entries(3, 0);
+
+#[uprobe]
+pub fn test_64_32_call_relocs(_ctx: ProbeContext) {
+    // this will link set_result and do a forward call
+    set_result(0, hint::black_box(1));
+
+    // set_result is already linked, this will just do the forward call
+    set_result(1, hint::black_box(2));
+
+    // this will link set_result_backward after set_result. Then will do a
+    // backward call to set_result.
+    set_result_backward(2, hint::black_box(3));
+}
+
+#[inline(never)]
+fn set_result(index: u32, value: u64) {
+    unsafe {
+        if let Some(v) = RESULTS.get_ptr_mut(index) {
+            *v = value;
+        }
+    }
+}
+
+#[inline(never)]
+fn set_result_backward(index: u32, value: u64) {
+    set_result(index, value);
+}
+
+#[panic_handler]
+fn panic(_info: &core::panic::PanicInfo) -> ! {
+    unsafe { core::hint::unreachable_unchecked() }
+}

test/integration-test-macros/src/lib.rs

@@ -10,7 +10,7 @@ pub fn integration_test(_attr: TokenStream, item: TokenStream) -> TokenStream {
     let expanded = quote! {
         #item
 
-        inventory::submit!(IntegrationTest {
+        inventory::submit!(crate::IntegrationTest {
             name: concat!(module_path!(), "::", #name_str),
             test_fn: #name,
         });

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

@@ -0,0 +1,313 @@
+use anyhow::{Context, Result};
+use std::{path::PathBuf, process::Command, thread::sleep, time::Duration};
+use tempfile::TempDir;
+
+use aya::{maps::Array, programs::TracePoint, BpfLoader, Btf, Endianness};
+
+use super::integration_test;
+
+// In the tests below we often use values like 0xAAAAAAAA or -0x7AAAAAAA. Those values have no
+// special meaning, they just have "nice" bit patterns that can be helpful while debugging.
+
+#[integration_test]
+fn relocate_field() {
+    let test = RelocationTest {
+        local_definition: r#"
+            struct foo {
+              __u8 a;
+              __u8 b;
+              __u8 c;
+              __u8 d;
+            };
+        "#,
+        target_btf: r#"
+            struct foo {
+              __u8 a;
+              __u8 c;
+              __u8 b;
+              __u8 d;
+            } s1;
+        "#,
+        relocation_code: r#"
+            __u8 memory[] = {1, 2, 3, 4};
+            struct foo *ptr = (struct foo *) &memory;
+            value = __builtin_preserve_access_index(ptr->c);
+        "#,
+    }
+    .build()
+    .unwrap();
+    assert_eq!(test.run().unwrap(), 2);
+    assert_eq!(test.run_no_btf().unwrap(), 3);
+}
+
+#[integration_test]
+fn relocate_enum() {
+    let test = RelocationTest {
+        local_definition: r#"
+            enum foo { D = 0xAAAAAAAA };
+        "#,
+        target_btf: r#"
+            enum foo { D = 0xBBBBBBBB } e1;
+        "#,
+        relocation_code: r#"
+            #define BPF_ENUMVAL_VALUE 1
+            value = __builtin_preserve_enum_value(*(typeof(enum foo) *)D, BPF_ENUMVAL_VALUE);
+        "#,
+    }
+    .build()
+    .unwrap();
+    assert_eq!(test.run().unwrap(), 0xBBBBBBBB);
+    assert_eq!(test.run_no_btf().unwrap(), 0xAAAAAAAA);
+}
+
+#[integration_test]
+fn relocate_enum_signed() {
+    let test = RelocationTest {
+        local_definition: r#"
+            enum foo { D = -0x7AAAAAAA };
+        "#,
+        target_btf: r#"
+            enum foo { D = -0x7BBBBBBB } e1;
+        "#,
+        relocation_code: r#"
+            #define BPF_ENUMVAL_VALUE 1
+            value = __builtin_preserve_enum_value(*(typeof(enum foo) *)D, BPF_ENUMVAL_VALUE);
+        "#,
+    }
+    .build()
+    .unwrap();
+    assert_eq!(test.run().unwrap() as i64, -0x7BBBBBBBi64);
+    assert_eq!(test.run_no_btf().unwrap() as i64, -0x7AAAAAAAi64);
+}
+
+#[integration_test]
+fn relocate_enum64() {
+    let test = RelocationTest {
+        local_definition: r#"
+            enum foo { D = 0xAAAAAAAABBBBBBBB };
+        "#,
+        target_btf: r#"
+            enum foo { D = 0xCCCCCCCCDDDDDDDD } e1;
+        "#,
+        relocation_code: r#"
+            #define BPF_ENUMVAL_VALUE 1
+            value = __builtin_preserve_enum_value(*(typeof(enum foo) *)D, BPF_ENUMVAL_VALUE);
+        "#,
+    }
+    .build()
+    .unwrap();
+    assert_eq!(test.run().unwrap(), 0xCCCCCCCCDDDDDDDD);
+    assert_eq!(test.run_no_btf().unwrap(), 0xAAAAAAAABBBBBBBB);
+}
+
+#[integration_test]
+fn relocate_enum64_signed() {
+    let test = RelocationTest {
+        local_definition: r#"
+            enum foo { D = -0xAAAAAAABBBBBBBB };
+        "#,
+        target_btf: r#"
+            enum foo { D = -0xCCCCCCCDDDDDDDD } e1;
+        "#,
+        relocation_code: r#"
+            #define BPF_ENUMVAL_VALUE 1
+            value = __builtin_preserve_enum_value(*(typeof(enum foo) *)D, BPF_ENUMVAL_VALUE);
+        "#,
+    }
+    .build()
+    .unwrap();
+    assert_eq!(test.run().unwrap() as i64, -0xCCCCCCCDDDDDDDDi64);
+    assert_eq!(test.run_no_btf().unwrap() as i64, -0xAAAAAAABBBBBBBBi64);
+}
+
+#[integration_test]
+fn relocate_pointer() {
+    let test = RelocationTest {
+        local_definition: r#"
+            struct foo {};
+            struct bar { struct foo *f; };
+        "#,
+        target_btf: r#"
+            struct foo {};
+            struct bar { struct foo *f; };
+        "#,
+        relocation_code: r#"
+            __u8 memory[] = {42, 0, 0, 0, 0, 0, 0, 0};
+            struct bar* ptr = (struct bar *) &memory;
+            value = (__u64) __builtin_preserve_access_index(ptr->f);
+        "#,
+    }
+    .build()
+    .unwrap();
+    assert_eq!(test.run().unwrap(), 42);
+    assert_eq!(test.run_no_btf().unwrap(), 42);
+}
+
+/// Utility code for running relocation tests:
+/// - Generates the eBPF program using probided local definition and relocation code
+/// - Generates the BTF from the target btf code
+struct RelocationTest {
+    /// Data structure definition, local to the eBPF program and embedded in the eBPF bytecode
+    local_definition: &'static str,
+    /// Target data structure definition. What the vmlinux would actually contain.
+    target_btf: &'static str,
+    /// Code executed by the eBPF program to test the relocation.
+    /// The format should be:
+    // __u8 memory[] = { ... };
+    // __u32 value = BPF_CORE_READ((struct foo *)&memory, ...);
+    //
+    // The generated code will be executed by attaching a tracepoint to sched_switch
+    // and emitting `__u32 value` an a map. See the code template below for more details.
+    relocation_code: &'static str,
+}
+
+impl RelocationTest {
+    /// Build a RelocationTestRunner
+    fn build(&self) -> Result<RelocationTestRunner> {
+        Ok(RelocationTestRunner {
+            ebpf: self.build_ebpf()?,
+            btf: self.build_btf()?,
+        })
+    }
+
+    /// - Generate the source eBPF filling a template
+    /// - Compile it with clang
+    fn build_ebpf(&self) -> Result<Vec<u8>> {
+        let local_definition = self.local_definition;
+        let relocation_code = self.relocation_code;
+        let (_tmp_dir, compiled_file) = compile(&format!(
+            r#"
+                #include <linux/bpf.h>
+
+                static long (*bpf_map_update_elem)(void *map, const void *key, const void *value, __u64 flags) = (void *) 2;
+
+                {local_definition}
+
+                struct {{
+                  int (*type)[BPF_MAP_TYPE_ARRAY];
+                  __u32 *key;
+                  __u64 *value;
+                  int (*max_entries)[1];
+                }} output_map
+                __attribute__((section(".maps"), used));
+
+                __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);
+                  return 0;
+                }}
+
+                char _license[] __attribute__((section("license"), used)) = "GPL";
+            "#
+        ))
+        .context("Failed to compile eBPF program")?;
+        let bytecode =
+            std::fs::read(compiled_file).context("Error reading compiled eBPF program")?;
+        Ok(bytecode)
+    }
+
+    /// - Generate the target BTF source with a mock main()
+    /// - Compile it with clang
+    /// - Extract the BTF with llvm-objcopy
+    fn build_btf(&self) -> Result<Btf> {
+        let target_btf = self.target_btf;
+        let relocation_code = self.relocation_code;
+        // BTF files can be generated and inspected with these commands:
+        // $ clang -c -g -O2 -target bpf target.c
+        // $ pahole --btf_encode_detached=target.btf -V target.o
+        // $ bpftool btf dump file ./target.btf  format c
+        let (tmp_dir, compiled_file) = compile(&format!(
+            r#"
+                #include <linux/bpf.h>
+
+                {target_btf}
+                int main() {{
+                    __u64 value = 0;
+                    // This is needed to make sure to emit BTF for the defined types,
+                    // it could be dead code eliminated if we don't.
+                    {relocation_code};
+                    return value;
+                }}
+            "#
+        ))
+        .context("Failed to compile BTF")?;
+        Command::new("llvm-objcopy")
+            .current_dir(tmp_dir.path())
+            .args(["--dump-section", ".BTF=target.btf"])
+            .arg(compiled_file)
+            .status()
+            .context("Failed to run llvm-objcopy")?
+            .success()
+            .then_some(())
+            .context("Failed to extract BTF")?;
+        let btf = Btf::parse_file(tmp_dir.path().join("target.btf"), Endianness::default())
+            .context("Error parsing generated BTF")?;
+        Ok(btf)
+    }
+}
+
+/// Compile an eBPF program and return the path of the compiled object.
+/// Also returns a TempDir handler, dropping it will clear the created dicretory.
+fn compile(source_code: &str) -> Result<(TempDir, PathBuf)> {
+    let tmp_dir = tempfile::tempdir().context("Error making temp dir")?;
+    let source = tmp_dir.path().join("source.c");
+    std::fs::write(&source, source_code).context("Writing bpf program failed")?;
+    Command::new("clang")
+        .current_dir(&tmp_dir)
+        .args(["-c", "-g", "-O2", "-target", "bpf"])
+        .arg(&source)
+        .status()
+        .context("Failed to run clang")?
+        .success()
+        .then_some(())
+        .context("Failed to compile eBPF source")?;
+    Ok((tmp_dir, source.with_extension("o")))
+}
+
+struct RelocationTestRunner {
+    ebpf: Vec<u8>,
+    btf: Btf,
+}
+
+impl RelocationTestRunner {
+    /// Run test and return the output value
+    fn run(&self) -> Result<u64> {
+        self.run_internal(true).context("Error running with BTF")
+    }
+
+    /// Run without loading btf
+    fn run_no_btf(&self) -> Result<u64> {
+        self.run_internal(false)
+            .context("Error running without BTF")
+    }
+
+    fn run_internal(&self, with_relocations: bool) -> Result<u64> {
+        let mut loader = BpfLoader::new();
+        if with_relocations {
+            loader.btf(Some(&self.btf));
+        } else {
+            loader.btf(None);
+        }
+        let mut bpf = loader.load(&self.ebpf).context("Loading eBPF failed")?;
+        let program: &mut TracePoint = bpf
+            .program_mut("bpf_prog")
+            .context("bpf_prog not found")?
+            .try_into()
+            .context("program not a tracepoint")?;
+        program.load().context("Loading tracepoint failed")?;
+        // Attach to sched_switch and wait some time to make sure it executed at least once
+        program
+            .attach("sched", "sched_switch")
+            .context("attach failed")?;
+        sleep(Duration::from_millis(1000));
+        // To inspect the loaded eBPF bytecode, increse the timeout and run:
+        // $ sudo bpftool prog dump xlated name bpf_prog
+
+        let output_map: Array<_, u64> = bpf.take_map("output_map").unwrap().try_into().unwrap();
+        let key = 0;
+        output_map.get(&key, 0).context("Getting key 0 failed")
+    }
+}

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

@@ -1,4 +1,4 @@
-use super::{integration_test, IntegrationTest};
+use super::integration_test;
 
 use aya::include_bytes_aligned;
 use object::{Object, ObjectSymbol};

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

@@ -13,7 +13,7 @@ use log::warn;
 
 use crate::tests::kernel_version;
 
-use super::{integration_test, IntegrationTest};
+use super::integration_test;
 
 const MAX_RETRIES: u32 = 100;
 const RETRY_DURATION_MS: u64 = 10;

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

@@ -4,6 +4,7 @@ use libc::{uname, utsname};
 use regex::Regex;
 use std::{ffi::CStr, mem};
 
+pub mod btf_relocations;
 pub mod elf;
 pub mod load;
 pub mod rbpf;

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

@@ -4,7 +4,7 @@ use std::collections::HashMap;
 use aya::include_bytes_aligned;
 use aya_obj::{generated::bpf_insn, Object, ProgramSection};
 
-use super::{integration_test, IntegrationTest};
+use super::integration_test;
 
 #[integration_test]
 fn run_with_rbpf() {
@@ -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);

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

@@ -1,313 +1,70 @@
-use anyhow::{Context, Result};
-use std::{path::PathBuf, process::Command, thread::sleep, time::Duration};
-use tempfile::TempDir;
+use std::{process::exit, time::Duration};
 
-use aya::{maps::Array, programs::TracePoint, BpfLoader, Btf, Endianness};
-
-use super::{integration_test, IntegrationTest};
-
-// In the tests below we often use values like 0xAAAAAAAA or -0x7AAAAAAA. Those values have no
-// special meaning, they just have "nice" bit patterns that can be helpful while debugging.
+use aya::{
+    include_bytes_aligned,
+    programs::{ProgramError, UProbe},
+    Bpf,
+};
+use integration_test_macros::integration_test;
 
 #[integration_test]
-fn relocate_field() {
-    let test = RelocationTest {
-        local_definition: r#"
-            struct foo {
-              __u8 a;
-              __u8 b;
-              __u8 c;
-              __u8 d;
-            };
-        "#,
-        target_btf: r#"
-            struct foo {
-              __u8 a;
-              __u8 c;
-              __u8 b;
-              __u8 d;
-            } s1;
-        "#,
-        relocation_code: r#"
-            __u8 memory[] = {1, 2, 3, 4};
-            struct foo *ptr = (struct foo *) &memory;
-            value = __builtin_preserve_access_index(ptr->c);
-        "#,
-    }
-    .build()
-    .unwrap();
-    assert_eq!(test.run().unwrap(), 2);
-    assert_eq!(test.run_no_btf().unwrap(), 3);
+fn relocations() {
+    let bpf = load_and_attach(
+        "test_64_32_call_relocs",
+        include_bytes_aligned!("../../../../target/bpfel-unknown-none/release/relocations"),
+    );
+
+    trigger_relocations_program();
+    std::thread::sleep(Duration::from_millis(100));
+
+    let m = aya::maps::Array::<_, u64>::try_from(bpf.map("RESULTS").unwrap()).unwrap();
+    assert_eq!(m.get(&0, 0).unwrap(), 1);
+    assert_eq!(m.get(&1, 0).unwrap(), 2);
+    assert_eq!(m.get(&2, 0).unwrap(), 3);
 }
 
 #[integration_test]
-fn relocate_enum() {
-    let test = RelocationTest {
-        local_definition: r#"
-            enum foo { D = 0xAAAAAAAA };
-        "#,
-        target_btf: r#"
-            enum foo { D = 0xBBBBBBBB } e1;
-        "#,
-        relocation_code: r#"
-            #define BPF_ENUMVAL_VALUE 1
-            value = __builtin_preserve_enum_value(*(typeof(enum foo) *)D, BPF_ENUMVAL_VALUE);
-        "#,
-    }
-    .build()
-    .unwrap();
-    assert_eq!(test.run().unwrap(), 0xBBBBBBBB);
-    assert_eq!(test.run_no_btf().unwrap(), 0xAAAAAAAA);
-}
+fn text_64_64_reloc() {
+    let mut bpf = load_and_attach(
+        "test_text_64_64_reloc",
+        include_bytes_aligned!("../../../../target/bpfel-unknown-none/release/text_64_64_reloc.o"),
+    );
 
-#[integration_test]
-fn relocate_enum_signed() {
-    let test = RelocationTest {
-        local_definition: r#"
-            enum foo { D = -0x7AAAAAAA };
-        "#,
-        target_btf: r#"
-            enum foo { D = -0x7BBBBBBB } e1;
-        "#,
-        relocation_code: r#"
-            #define BPF_ENUMVAL_VALUE 1
-            value = __builtin_preserve_enum_value(*(typeof(enum foo) *)D, BPF_ENUMVAL_VALUE);
-        "#,
-    }
-    .build()
-    .unwrap();
-    assert_eq!(test.run().unwrap() as i64, -0x7BBBBBBBi64);
-    assert_eq!(test.run_no_btf().unwrap() as i64, -0x7AAAAAAAi64);
-}
+    let mut m = aya::maps::Array::<_, u64>::try_from(bpf.map_mut("RESULTS").unwrap()).unwrap();
+    m.set(0, 1, 0).unwrap();
+    m.set(1, 2, 0).unwrap();
 
-#[integration_test]
-fn relocate_enum64() {
-    let test = RelocationTest {
-        local_definition: r#"
-            enum foo { D = 0xAAAAAAAABBBBBBBB };
-        "#,
-        target_btf: r#"
-            enum foo { D = 0xCCCCCCCCDDDDDDDD } e1;
-        "#,
-        relocation_code: r#"
-            #define BPF_ENUMVAL_VALUE 1
-            value = __builtin_preserve_enum_value(*(typeof(enum foo) *)D, BPF_ENUMVAL_VALUE);
-        "#,
-    }
-    .build()
-    .unwrap();
-    assert_eq!(test.run().unwrap(), 0xCCCCCCCCDDDDDDDD);
-    assert_eq!(test.run_no_btf().unwrap(), 0xAAAAAAAABBBBBBBB);
-}
+    trigger_relocations_program();
+    std::thread::sleep(Duration::from_millis(100));
 
-#[integration_test]
-fn relocate_enum64_signed() {
-    let test = RelocationTest {
-        local_definition: r#"
-            enum foo { D = -0xAAAAAAABBBBBBBB };
-        "#,
-        target_btf: r#"
-            enum foo { D = -0xCCCCCCCDDDDDDDD } e1;
-        "#,
-        relocation_code: r#"
-            #define BPF_ENUMVAL_VALUE 1
-            value = __builtin_preserve_enum_value(*(typeof(enum foo) *)D, BPF_ENUMVAL_VALUE);
-        "#,
-    }
-    .build()
-    .unwrap();
-    assert_eq!(test.run().unwrap() as i64, -0xCCCCCCCDDDDDDDDi64);
-    assert_eq!(test.run_no_btf().unwrap() as i64, -0xAAAAAAABBBBBBBBi64);
+    assert_eq!(m.get(&0, 0).unwrap(), 2);
+    assert_eq!(m.get(&1, 0).unwrap(), 3);
 }
 
-#[integration_test]
-fn relocate_pointer() {
-    let test = RelocationTest {
-        local_definition: r#"
-            struct foo {};
-            struct bar { struct foo *f; };
-        "#,
-        target_btf: r#"
-            struct foo {};
-            struct bar { struct foo *f; };
-        "#,
-        relocation_code: r#"
-            __u8 memory[] = {42, 0, 0, 0, 0, 0, 0, 0};
-            struct bar* ptr = (struct bar *) &memory;
-            value = (__u64) __builtin_preserve_access_index(ptr->f);
-        "#,
-    }
-    .build()
+fn load_and_attach(name: &str, bytes: &[u8]) -> Bpf {
+    let mut bpf = Bpf::load(bytes).unwrap();
+
+    let prog: &mut UProbe = bpf.program_mut(name).unwrap().try_into().unwrap();
+    if let Err(ProgramError::LoadError {
+        io_error,
+        verifier_log,
+    }) = prog.load()
+    {
+        println!("Failed to load program `{name}`: {io_error}. Verifier log:\n{verifier_log:#}");
+        exit(1);
+    };
+
+    prog.attach(
+        Some("trigger_relocations_program"),
+        0,
+        "/proc/self/exe",
+        None,
+    )
     .unwrap();
-    assert_eq!(test.run().unwrap(), 42);
-    assert_eq!(test.run_no_btf().unwrap(), 42);
-}
 
-/// Utility code for running relocation tests:
-/// - Generates the eBPF program using probided local definition and relocation code
-/// - Generates the BTF from the target btf code
-struct RelocationTest {
-    /// Data structure definition, local to the eBPF program and embedded in the eBPF bytecode
-    local_definition: &'static str,
-    /// Target data structure definition. What the vmlinux would actually contain.
-    target_btf: &'static str,
-    /// Code executed by the eBPF program to test the relocation.
-    /// The format should be:
-    // __u8 memory[] = { ... };
-    // __u32 value = BPF_CORE_READ((struct foo *)&memory, ...);
-    //
-    // The generated code will be executed by attaching a tracepoint to sched_switch
-    // and emitting `__u32 value` an a map. See the code template below for more details.
-    relocation_code: &'static str,
+    bpf
 }
 
-impl RelocationTest {
-    /// Build a RelocationTestRunner
-    fn build(&self) -> Result<RelocationTestRunner> {
-        Ok(RelocationTestRunner {
-            ebpf: self.build_ebpf()?,
-            btf: self.build_btf()?,
-        })
-    }
-
-    /// - Generate the source eBPF filling a template
-    /// - Compile it with clang
-    fn build_ebpf(&self) -> Result<Vec<u8>> {
-        let local_definition = self.local_definition;
-        let relocation_code = self.relocation_code;
-        let (_tmp_dir, compiled_file) = compile(&format!(
-            r#"
-                #include <linux/bpf.h>
-
-                static long (*bpf_map_update_elem)(void *map, const void *key, const void *value, __u64 flags) = (void *) 2;
-
-                {local_definition}
-
-                struct {{
-                  int (*type)[BPF_MAP_TYPE_ARRAY];
-                  __u32 *key;
-                  __u64 *value;
-                  int (*max_entries)[1];
-                }} output_map
-                __attribute__((section(".maps"), used));
-
-                __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);
-                  return 0;
-                }}
-
-                char _license[] __attribute__((section("license"), used)) = "GPL";
-            "#
-        ))
-        .context("Failed to compile eBPF program")?;
-        let bytecode =
-            std::fs::read(compiled_file).context("Error reading compiled eBPF program")?;
-        Ok(bytecode)
-    }
-
-    /// - Generate the target BTF source with a mock main()
-    /// - Compile it with clang
-    /// - Extract the BTF with llvm-objcopy
-    fn build_btf(&self) -> Result<Btf> {
-        let target_btf = self.target_btf;
-        let relocation_code = self.relocation_code;
-        // BTF files can be generated and inspected with these commands:
-        // $ clang -c -g -O2 -target bpf target.c
-        // $ pahole --btf_encode_detached=target.btf -V target.o
-        // $ bpftool btf dump file ./target.btf  format c
-        let (tmp_dir, compiled_file) = compile(&format!(
-            r#"
-                #include <linux/bpf.h>
-
-                {target_btf}
-                int main() {{
-                    __u64 value = 0;
-                    // This is needed to make sure to emit BTF for the defined types,
-                    // it could be dead code eliminated if we don't.
-                    {relocation_code};
-                    return value;
-                }}
-            "#
-        ))
-        .context("Failed to compile BTF")?;
-        Command::new("llvm-objcopy")
-            .current_dir(tmp_dir.path())
-            .args(["--dump-section", ".BTF=target.btf"])
-            .arg(compiled_file)
-            .status()
-            .context("Failed to run llvm-objcopy")?
-            .success()
-            .then_some(())
-            .context("Failed to extract BTF")?;
-        let btf = Btf::parse_file(tmp_dir.path().join("target.btf"), Endianness::default())
-            .context("Error parsing generated BTF")?;
-        Ok(btf)
-    }
-}
-
-/// Compile an eBPF program and return the path of the compiled object.
-/// Also returns a TempDir handler, dropping it will clear the created dicretory.
-fn compile(source_code: &str) -> Result<(TempDir, PathBuf)> {
-    let tmp_dir = tempfile::tempdir().context("Error making temp dir")?;
-    let source = tmp_dir.path().join("source.c");
-    std::fs::write(&source, source_code).context("Writing bpf program failed")?;
-    Command::new("clang")
-        .current_dir(&tmp_dir)
-        .args(["-c", "-g", "-O2", "-target", "bpf"])
-        .arg(&source)
-        .status()
-        .context("Failed to run clang")?
-        .success()
-        .then_some(())
-        .context("Failed to compile eBPF source")?;
-    Ok((tmp_dir, source.with_extension("o")))
-}
-
-struct RelocationTestRunner {
-    ebpf: Vec<u8>,
-    btf: Btf,
-}
-
-impl RelocationTestRunner {
-    /// Run test and return the output value
-    fn run(&self) -> Result<u64> {
-        self.run_internal(true).context("Error running with BTF")
-    }
-
-    /// Run without loading btf
-    fn run_no_btf(&self) -> Result<u64> {
-        self.run_internal(false)
-            .context("Error running without BTF")
-    }
-
-    fn run_internal(&self, with_relocations: bool) -> Result<u64> {
-        let mut loader = BpfLoader::new();
-        if with_relocations {
-            loader.btf(Some(&self.btf));
-        } else {
-            loader.btf(None);
-        }
-        let mut bpf = loader.load(&self.ebpf).context("Loading eBPF failed")?;
-        let program: &mut TracePoint = bpf
-            .program_mut("bpf_prog")
-            .context("bpf_prog not found")?
-            .try_into()
-            .context("program not a tracepoint")?;
-        program.load().context("Loading tracepoint failed")?;
-        // Attach to sched_switch and wait some time to make sure it executed at least once
-        program
-            .attach("sched", "sched_switch")
-            .context("attach failed")?;
-        sleep(Duration::from_millis(1000));
-        // To inspect the loaded eBPF bytecode, increse the timeout and run:
-        // $ sudo bpftool prog dump xlated name bpf_prog
-
-        let output_map: Array<_, u64> = bpf.take_map("output_map").unwrap().try_into().unwrap();
-        let key = 0;
-        output_map.get(&key, 0).context("Getting key 0 failed")
-    }
-}
+#[no_mangle]
+#[inline(never)]
+pub extern "C" fn trigger_relocations_program() {}

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

@@ -5,7 +5,7 @@ use aya::{
 };
 use log::info;
 
-use super::{integration_test, kernel_version, IntegrationTest};
+use super::{integration_test, kernel_version};
 
 #[integration_test]
 fn xdp() {