aya: add support for function calls

aya/src/bpf.rs

@@ -100,6 +100,7 @@ impl Bpf {
         }
 
         obj.relocate_maps(maps.as_slice())?;
+        obj.relocate_calls()?;
 
         let programs = obj
             .programs
@@ -231,9 +232,9 @@ pub enum BpfError {
     #[error("BTF error")]
     BtfError(#[from] BtfError),
 
-    #[error("error relocating BPF program `{program_name}`: {error}")]
+    #[error("error relocating `{function}`: {error}")]
     RelocationError {
-        program_name: String,
+        function: String,
         error: Box<dyn Error + Send + Sync>,
     },
 

aya/src/generated/linux_bindings_aarch64.rs

@@ -85,12 +85,14 @@ pub const BPF_LDX: u32 = 1;
 pub const BPF_ST: u32 = 2;
 pub const BPF_STX: u32 = 3;
 pub const BPF_ALU: u32 = 4;
+pub const BPF_JMP: u32 = 5;
 pub const BPF_W: u32 = 0;
 pub const BPF_H: u32 = 8;
 pub const BPF_B: u32 = 16;
 pub const BPF_K: u32 = 0;
 pub const BPF_ALU64: u32 = 7;
 pub const BPF_DW: u32 = 24;
+pub const BPF_CALL: u32 = 128;
 pub const BPF_PSEUDO_MAP_FD: u32 = 1;
 pub const BPF_PSEUDO_MAP_VALUE: u32 = 2;
 pub const BPF_PSEUDO_BTF_ID: u32 = 3;

aya/src/generated/linux_bindings_x86_64.rs

@@ -85,12 +85,14 @@ pub const BPF_LDX: u32 = 1;
 pub const BPF_ST: u32 = 2;
 pub const BPF_STX: u32 = 3;
 pub const BPF_ALU: u32 = 4;
+pub const BPF_JMP: u32 = 5;
 pub const BPF_W: u32 = 0;
 pub const BPF_H: u32 = 8;
 pub const BPF_B: u32 = 16;
 pub const BPF_K: u32 = 0;
 pub const BPF_ALU64: u32 = 7;
 pub const BPF_DW: u32 = 24;
+pub const BPF_CALL: u32 = 128;
 pub const BPF_PSEUDO_MAP_FD: u32 = 1;
 pub const BPF_PSEUDO_MAP_VALUE: u32 = 2;
 pub const BPF_PSEUDO_BTF_ID: u32 = 3;

aya/src/obj/btf/relocation.rs

@@ -171,7 +171,7 @@ impl Object {
                 .programs
                 .get_mut(section_name)
                 .ok_or(BpfError::RelocationError {
-                    program_name: section_name.to_owned(),
+                    function: section_name.to_owned(),
                     error: Box::new(RelocationError::ProgramNotFound),
                 })?;
             match relocate_btf_program(
@@ -185,7 +185,7 @@ impl Object {
                 Err(ErrorWrapper::BtfError(e)) => return Err(e)?,
                 Err(ErrorWrapper::RelocationError(error)) => {
                     return Err(BpfError::RelocationError {
-                        program_name: section_name.to_owned(),
+                        function: section_name.to_owned(),
                         error: Box::new(error),
                     })
                 }
@@ -204,7 +204,7 @@ fn relocate_btf_program<'target>(
     candidates_cache: &mut HashMap<u32, Vec<Candidate<'target>>>,
 ) -> Result<(), ErrorWrapper> {
     for rel in relos {
-        let instructions = &mut program.instructions;
+        let instructions = &mut program.function.instructions;
         let ins_index = rel.ins_offset as usize / std::mem::size_of::<bpf_insn>();
         if ins_index >= instructions.len() {
             return Err(RelocationError::InvalidInstructionIndex {
@@ -763,7 +763,7 @@ impl ComputedRelocation {
         local_btf: &Btf,
         target_btf: &Btf,
     ) -> Result<(), ErrorWrapper> {
-        let instructions = &mut program.instructions;
+        let instructions = &mut program.function.instructions;
         let num_instructions = instructions.len();
         let ins_index = rel.ins_offset as usize / std::mem::size_of::<bpf_insn>();
         let mut ins =

aya/src/obj/mod.rs

@@ -3,7 +3,7 @@ mod relocation;
 
 use object::{
     read::{Object as ElfObject, ObjectSection, Section as ObjSection},
-    Endianness, ObjectSymbol, ObjectSymbolTable, SectionIndex, SymbolIndex,
+    Endianness, ObjectSymbol, ObjectSymbolTable, RelocationTarget, SectionIndex,
 };
 use std::{
     collections::HashMap,
@@ -14,7 +14,7 @@ use std::{
 };
 use thiserror::Error;
 
-pub use relocation::*;
+use relocation::*;
 
 use crate::{
     bpf_map_def,
@@ -34,8 +34,9 @@ pub struct Object {
     pub btf_ext: Option<BtfExt>,
     pub(crate) maps: HashMap<String, Map>,
     pub(crate) programs: HashMap<String, Program>,
-    pub(crate) relocations: HashMap<SectionIndex, Vec<Relocation>>,
+    pub(crate) functions: HashMap<u64, Function>,
-    pub(crate) symbols_by_index: HashMap<SymbolIndex, Symbol>,
+    pub(crate) relocations: HashMap<SectionIndex, HashMap<u64, Relocation>>,
+    pub(crate) symbols_by_index: HashMap<usize, Symbol>,
 }
 
 #[derive(Debug, Clone)]
@@ -50,9 +51,17 @@ pub struct Map {
 pub(crate) struct Program {
     pub(crate) license: CString,
     pub(crate) kernel_version: KernelVersion,
-    pub(crate) instructions: Vec<bpf_insn>,
     pub(crate) kind: ProgramKind,
+    pub(crate) function: Function,
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct Function {
+    pub(crate) address: u64,
+    pub(crate) name: String,
     pub(crate) section_index: SectionIndex,
+    pub(crate) section_offset: usize,
+    pub(crate) instructions: Vec<bpf_insn>,
 }
 
 #[derive(Debug, Copy, Clone)]
@@ -107,23 +116,26 @@ impl Object {
 
         let mut bpf_obj = Object::new(endianness, license, kernel_version);
 
-        for s in obj.sections() {
-            bpf_obj.parse_section(Section::try_from(&s)?)?;
-        }
-
         if let Some(symbol_table) = obj.symbol_table() {
             for symbol in symbol_table.symbols() {
-                bpf_obj.symbols_by_index.insert(
+                let sym = Symbol {
-                    symbol.index(),
+                    index: symbol.index().0,
-                    Symbol {
+                    name: symbol.name().ok().map(String::from),
-                        name: symbol.name().ok().map(String::from),
+                    section_index: symbol.section().index(),
-                        section_index: symbol.section().index(),
+                    address: symbol.address(),
-                        address: symbol.address(),
+                    size: symbol.size(),
-                    },
+                    is_definition: symbol.is_definition(),
-                );
+                };
+                bpf_obj
+                    .symbols_by_index
+                    .insert(symbol.index().0, sym.clone());
             }
         }
 
+        for s in obj.sections() {
+            bpf_obj.parse_section(Section::try_from(&s)?)?;
+        }
+
         return Ok(bpf_obj);
     }
 
@@ -136,46 +148,105 @@ impl Object {
             btf_ext: None,
             maps: HashMap::new(),
             programs: HashMap::new(),
+            functions: HashMap::new(),
             relocations: HashMap::new(),
             symbols_by_index: HashMap::new(),
         }
     }
 
-    fn parse_program(&self, section: &Section, ty: &str) -> Result<Program, ParseError> {
+    fn parse_btf(&mut self, section: &Section) -> Result<(), BtfError> {
-        let num_instructions = section.data.len() / mem::size_of::<bpf_insn>();
+        self.btf = Some(Btf::parse(section.data, self.endianness)?);
-        if section.data.len() % mem::size_of::<bpf_insn>() > 0 {
+
-            return Err(ParseError::InvalidProgramCode);
+        Ok(())
-        }
+    }
-        let instructions = (0..num_instructions)
-            .map(|i| unsafe {
-                ptr::read_unaligned(
-                    (section.data.as_ptr() as usize + i * mem::size_of::<bpf_insn>())
-                        as *const bpf_insn,
-                )
-            })
-            .collect::<Vec<_>>();
 
+    fn parse_btf_ext(&mut self, section: &Section) -> Result<(), BtfError> {
+        self.btf_ext = Some(BtfExt::parse(section.data, self.endianness)?);
+        Ok(())
+    }
+
+    fn parse_program(
+        &self,
+        section: &Section,
+        ty: &str,
+        name: &str,
+    ) -> Result<Program, ParseError> {
         Ok(Program {
-            section_index: section.index,
             license: self.license.clone(),
             kernel_version: self.kernel_version,
-            instructions,
             kind: ProgramKind::from_str(ty)?,
+            function: Function {
+                name: name.to_owned(),
+                address: section.address,
+                section_index: section.index,
+                section_offset: 0,
+                instructions: copy_instructions(section.data)?,
+            },
         })
     }
 
-    fn parse_btf(&mut self, section: &Section) -> Result<(), BtfError> {
+    fn parse_text_section(&mut self, mut section: Section) -> Result<(), ParseError> {
-        self.btf = Some(Btf::parse(section.data, self.endianness)?);
+        let mut symbols_by_address = HashMap::new();
 
-        Ok(())
+        for sym in self.symbols_by_index.values() {
-    }
+            if sym.is_definition && sym.section_index == Some(section.index) {
+                if symbols_by_address.contains_key(&sym.address) {
+                    return Err(ParseError::SymbolTableConflict {
+                        section_index: section.index.0,
+                        address: sym.address,
+                    });
+                }
+                symbols_by_address.insert(sym.address, sym);
+            }
+        }
+
+        let mut offset = 0;
+        while offset < section.data.len() {
+            let address = section.address + offset as u64;
+            let sym = symbols_by_address
+                .get(&address)
+                .ok_or(ParseError::UnknownSymbol {
+                    section_index: section.index.0,
+                    address,
+                })?;
+            if sym.size == 0 {
+                return Err(ParseError::InvalidSymbol {
+                    index: sym.index,
+                    name: sym.name.clone(),
+                });
+            }
+
+            self.functions.insert(
+                sym.address,
+                Function {
+                    address,
+                    name: sym.name.clone().unwrap(),
+                    section_index: section.index,
+                    section_offset: offset,
+                    instructions: copy_instructions(
+                        &section.data[offset..offset + sym.size as usize],
+                    )?,
+                },
+            );
+
+            offset += sym.size as usize;
+        }
+
+        if !section.relocations.is_empty() {
+            self.relocations.insert(
+                section.index,
+                section
+                    .relocations
+                    .drain(..)
+                    .map(|rel| (rel.offset, rel))
+                    .collect(),
+            );
+        }
 
-    fn parse_btf_ext(&mut self, section: &Section) -> Result<(), BtfError> {
-        self.btf_ext = Some(BtfExt::parse(section.data, self.endianness)?);
         Ok(())
     }
 
-    fn parse_section(&mut self, section: Section) -> Result<(), BpfError> {
+    fn parse_section(&mut self, mut section: Section) -> Result<(), BpfError> {
         let parts = section.name.split("/").collect::<Vec<_>>();
 
         match parts.as_slice() {
@@ -185,6 +256,7 @@ impl Object {
                 self.maps
                     .insert(name.to_string(), parse_map(&section, name)?);
             }
+            &[".text"] => self.parse_text_section(section)?,
             &[".BTF"] => self.parse_btf(&section)?,
             &[".BTF.ext"] => self.parse_btf_ext(&section)?,
             &["maps", name] => {
@@ -199,9 +271,16 @@ impl Object {
             | &[ty @ "xdp", name]
             | &[ty @ "trace_point", name] => {
                 self.programs
-                    .insert(name.to_string(), self.parse_program(&section, ty)?);
+                    .insert(name.to_string(), self.parse_program(&section, ty, name)?);
                 if !section.relocations.is_empty() {
-                    self.relocations.insert(section.index, section.relocations);
+                    self.relocations.insert(
+                        section.index,
+                        section
+                            .relocations
+                            .drain(..)
+                            .map(|rel| (rel.offset, rel))
+                            .collect(),
+                    );
                 }
             }
 
@@ -233,8 +312,8 @@ pub enum ParseError {
         source: object::read::Error,
     },
 
-    #[error("unsupported relocation")]
+    #[error("unsupported relocation target")]
-    UnsupportedRelocationKind,
+    UnsupportedRelocationTarget,
 
     #[error("invalid program kind `{kind}`")]
     InvalidProgramKind { kind: String },
@@ -244,10 +323,21 @@ pub enum ParseError {
 
     #[error("error parsing map `{name}`")]
     InvalidMapDefinition { name: String },
+
+    #[error("two or more symbols in section `{section_index}` have the same address {address:x}")]
+    SymbolTableConflict { section_index: usize, address: u64 },
+
+    #[error("unknown symbol in section `{section_index}` at address {address:x}")]
+    UnknownSymbol { section_index: usize, address: u64 },
+
+    #[error("invalid symbol, index `{index}` name: {}", .name.as_ref().unwrap_or(&"[unknown]".into()))]
+    InvalidSymbol { index: usize, name: Option<String> },
 }
 
+#[derive(Debug)]
 struct Section<'a> {
     index: SectionIndex,
+    address: u64,
     name: &'a str,
     data: &'a [u8],
     relocations: Vec<Relocation>,
@@ -262,19 +352,24 @@ impl<'data, 'file, 'a> TryFrom<&'a ObjSection<'data, 'file>> for Section<'a> {
             index: index.0,
             source,
         };
+
         Ok(Section {
             index,
+            address: section.address(),
             name: section.name().map_err(map_err)?,
             data: section.data().map_err(map_err)?,
             relocations: section
                 .relocations()
-                .map(|(offset, r)| Relocation {
+                .map(|(offset, r)| {
-                    kind: r.kind(),
+                    Ok(Relocation {
-                    target: r.target(),
+                        symbol_index: match r.target() {
-                    addend: r.addend(),
+                            RelocationTarget::Symbol(index) => index.0,
-                    offset,
+                            _ => return Err(ParseError::UnsupportedRelocationTarget),
+                        },
+                        offset,
+                    })
                 })
-                .collect::<Vec<_>>(),
+                .collect::<Result<Vec<_>, _>>()?,
         })
     }
 }
@@ -367,6 +462,22 @@ fn parse_map_def(name: &str, data: &[u8]) -> Result<bpf_map_def, ParseError> {
     Ok(unsafe { ptr::read_unaligned(data.as_ptr() as *const bpf_map_def) })
 }
 
+fn copy_instructions(data: &[u8]) -> Result<Vec<bpf_insn>, ParseError> {
+    if data.len() % mem::size_of::<bpf_insn>() > 0 {
+        return Err(ParseError::InvalidProgramCode);
+    }
+    let num_instructions = data.len() / mem::size_of::<bpf_insn>();
+    let instructions = (0..num_instructions)
+        .map(|i| unsafe {
+            ptr::read_unaligned(
+                (data.as_ptr() as usize + i * mem::size_of::<bpf_insn>()) as *const bpf_insn,
+            )
+        })
+        .collect::<Vec<_>>();
+
+    Ok(instructions)
+}
+
 #[cfg(test)]
 mod tests {
     use matches::assert_matches;
@@ -378,6 +489,7 @@ mod tests {
     fn fake_section<'a>(name: &'a str, data: &'a [u8]) -> Section<'a> {
         Section {
             index: SectionIndex(0),
+            address: 0,
             name,
             data,
             relocations: Vec::new(),
@@ -563,7 +675,11 @@ mod tests {
         let obj = fake_obj();
 
         assert_matches!(
-            obj.parse_program(&fake_section("kprobe/foo", &42u32.to_ne_bytes(),), "kprobe"),
+            obj.parse_program(
+                &fake_section("kprobe/foo", &42u32.to_ne_bytes(),),
+                "kprobe",
+                "foo"
+            ),
             Err(ParseError::InvalidProgramCode)
         );
     }
@@ -573,14 +689,19 @@ mod tests {
         let obj = fake_obj();
 
         assert_matches!(
-            obj.parse_program(&fake_section("kprobe/foo", bytes_of(&fake_ins())), "kprobe"),
+            obj.parse_program(&fake_section("kprobe/foo", bytes_of(&fake_ins())), "kprobe", "foo"),
             Ok(Program {
                 license,
-                kernel_version,
+                kernel_version: KernelVersion::Any,
                 kind: ProgramKind::KProbe,
-                section_index: SectionIndex(0),
+                function: Function {
-                instructions
+                    name,
-            }) if license.to_string_lossy() == "GPL" && kernel_version == KernelVersion::Any && instructions.len() == 1
+                    address: 0,
+                    section_index: SectionIndex(0),
+                    section_offset: 0,
+                    instructions
+                }
+            }) if license.to_string_lossy() == "GPL" && name == "foo" && instructions.len() == 1
         );
     }
 

aya/src/obj/relocation.rs

@@ -1,25 +1,25 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, mem};
 
-use object::{RelocationKind, RelocationTarget, SectionIndex, SymbolIndex};
+use object::SectionIndex;
 use thiserror::Error;
 
 use crate::{
-    generated::{bpf_insn, BPF_PSEUDO_MAP_FD, BPF_PSEUDO_MAP_VALUE},
+    generated::{
+        bpf_insn, BPF_CALL, BPF_JMP, BPF_K, BPF_PSEUDO_CALL, BPF_PSEUDO_MAP_FD,
+        BPF_PSEUDO_MAP_VALUE,
+    },
     maps::Map,
-    obj::Object,
+    obj::{Function, Object, Program},
     BpfError,
 };
 
-use super::Program;
+const INS_SIZE: usize = mem::size_of::<bpf_insn>();
 
 #[derive(Debug, Error)]
-pub enum RelocationError {
+enum RelocationError {
     #[error("unknown symbol, index `{index}`")]
     UnknownSymbol { index: usize },
 
-    #[error("unknown symbol section, index `{index}`")]
-    UnknownSymbolSection { index: usize },
-
     #[error("section `{section_index}` not found, referenced by symbol `{}`",
             .symbol_name.clone().unwrap_or_else(|| .symbol_index.to_string()))]
     SectionNotFound {
@@ -28,30 +28,35 @@ pub enum RelocationError {
         symbol_name: Option<String>,
     },
 
+    #[error("function {address:#x} not found")]
+    UnknownFunction { address: u64 },
+
     #[error("the map `{name}` at section `{section_index}` has not been created")]
     MapNotCreated { section_index: usize, name: String },
 
-    #[error("invalid instruction index `{index}` referenced by relocation #{relocation_number}")]
+    #[error("invalid offset `{offset}` applying relocation #{relocation_number}")]
-    InvalidInstructionIndex {
+    InvalidRelocationOffset {
-        index: usize,
+        offset: u64,
-        num_instructions: usize,
         relocation_number: usize,
     },
 }
 
 #[derive(Debug, Copy, Clone)]
 pub(crate) struct Relocation {
-    pub(crate) kind: RelocationKind,
+    // byte offset of the instruction to be relocated
-    pub(crate) target: RelocationTarget,
     pub(crate) offset: u64,
-    pub(crate) addend: i64,
+    // index of the symbol to relocate to
+    pub(crate) symbol_index: usize,
 }
 
 #[derive(Debug, Clone)]
 pub(crate) struct Symbol {
+    pub(crate) index: usize,
     pub(crate) section_index: Option<SectionIndex>,
     pub(crate) name: Option<String>,
     pub(crate) address: u64,
+    pub(crate) size: u64,
+    pub(crate) is_definition: bool,
 }
 
 impl Object {
@@ -61,72 +66,242 @@ impl Object {
             .map(|map| (map.obj.section_index, map))
             .collect::<HashMap<_, _>>();
 
-        let symbol_table = &self.symbols_by_index;
+        let functions = self
-        for (program_name, program) in self.programs.iter_mut() {
+            .programs
-            if let Some(relocations) = self.relocations.get(&program.section_index) {
+            .values_mut()
-                relocate_program(program, relocations, &maps_by_section, symbol_table).map_err(
+            .map(|p| &mut p.function)
-                    |error| BpfError::RelocationError {
+            .chain(self.functions.values_mut());
-                        program_name: program_name.clone(),
+
-                        error: Box::new(error),
+        for function in functions {
-                    },
+            if let Some(relocations) = self.relocations.get(&function.section_index) {
-                )?;
+                relocate_maps(
+                    function,
+                    relocations.values(),
+                    &maps_by_section,
+                    &self.symbols_by_index,
+                )
+                .map_err(|error| BpfError::RelocationError {
+                    function: function.name.clone(),
+                    error: Box::new(error),
+                })?;
             }
         }
+
+        Ok(())
+    }
+
+    pub fn relocate_calls(&mut self) -> Result<(), BpfError> {
+        for (name, program) in self.programs.iter_mut() {
+            let linker =
+                FunctionLinker::new(&self.functions, &self.relocations, &self.symbols_by_index);
+            linker
+                .link(program)
+                .map_err(|error| BpfError::RelocationError {
+                    function: name.clone(),
+                    error: Box::new(error),
+                })?;
+        }
+
         Ok(())
     }
 }
 
-fn relocate_program(
+fn relocate_maps<'a, I: Iterator<Item = &'a Relocation>>(
-    program: &mut Program,
+    fun: &mut Function,
-    relocations: &[Relocation],
+    relocations: I,
     maps_by_section: &HashMap<usize, &Map>,
-    symbol_table: &HashMap<SymbolIndex, Symbol>,
+    symbol_table: &HashMap<usize, Symbol>,
 ) -> Result<(), RelocationError> {
-    for (rel_n, rel) in relocations.iter().enumerate() {
+    let section_offset = fun.section_offset;
-        match rel.target {
+    let instructions = &mut fun.instructions;
-            RelocationTarget::Symbol(index) => {
+    let function_size = instructions.len() * INS_SIZE;
-                let sym = symbol_table
+
-                    .get(&index)
+    for (rel_n, rel) in relocations.enumerate() {
-                    .ok_or(RelocationError::UnknownSymbol { index: index.0 })?;
+        let rel_offset = rel.offset as usize;
-
+        if rel_offset < section_offset || rel_offset >= section_offset + function_size {
-                let section_index = sym
+            // the relocation doesn't apply to this function
-                    .section_index
+            continue;
-                    .ok_or(RelocationError::UnknownSymbolSection { index: index.0 })?;
+        }
-
+
-                let map = maps_by_section.get(&section_index.0).ok_or(
+        // make sure that the relocation offset is properly aligned
-                    RelocationError::SectionNotFound {
+        let ins_offset = rel_offset - section_offset;
-                        symbol_index: index.0,
+        if ins_offset % INS_SIZE != 0 {
-                        symbol_name: sym.name.clone(),
+            return Err(RelocationError::InvalidRelocationOffset {
-                        section_index: section_index.0,
+                offset: rel.offset,
-                    },
+                relocation_number: rel_n,
-                )?;
+            });
-
+        }
-                let map_fd = map.fd.ok_or_else(|| RelocationError::MapNotCreated {
+        let ins_index = ins_offset / INS_SIZE;
-                    name: map.obj.name.clone(),
+
+        // calls are relocated in a separate step
+        if is_call(&instructions[ins_index]) {
+            continue;
+        }
+
+        // a map relocation points to the ELF section that contains the map
+        let sym = symbol_table
+            .get(&rel.symbol_index)
+            .ok_or(RelocationError::UnknownSymbol {
+                index: rel.symbol_index,
+            })?;
+
+        let section_index = match sym.section_index {
+            Some(index) => index,
+            // this is not a map relocation
+            None => continue,
+        };
+
+        let map =
+            maps_by_section
+                .get(&section_index.0)
+                .ok_or(RelocationError::SectionNotFound {
+                    symbol_index: rel.symbol_index,
+                    symbol_name: sym.name.clone(),
                     section_index: section_index.0,
                 })?;
 
-                let instructions = &mut program.instructions;
+        let map_fd = map.fd.ok_or_else(|| RelocationError::MapNotCreated {
-                let ins_index = (rel.offset / std::mem::size_of::<bpf_insn>() as u64) as usize;
+            name: map.obj.name.clone(),
-                if ins_index >= instructions.len() {
+            section_index: section_index.0,
-                    return Err(RelocationError::InvalidInstructionIndex {
+        })?;
-                        index: ins_index,
+
-                        num_instructions: instructions.len(),
+        if !map.obj.data.is_empty() {
-                        relocation_number: rel_n,
+            instructions[ins_index].set_src_reg(BPF_PSEUDO_MAP_VALUE as u8);
-                    });
+            instructions[ins_index + 1].imm = instructions[ins_index].imm + sym.address as i32;
-                }
+        } else {
-                if !map.obj.data.is_empty() {
+            instructions[ins_index].set_src_reg(BPF_PSEUDO_MAP_FD as u8);
-                    instructions[ins_index].set_src_reg(BPF_PSEUDO_MAP_VALUE as u8);
-                    instructions[ins_index + 1].imm =
-                        instructions[ins_index].imm + sym.address as i32;
-                } else {
-                    instructions[ins_index].set_src_reg(BPF_PSEUDO_MAP_FD as u8);
-                }
-                instructions[ins_index].imm = map_fd;
-            }
-            RelocationTarget::Section(_index) => {}
-            RelocationTarget::Absolute => todo!(),
         }
+        instructions[ins_index].imm = map_fd;
     }
+
     Ok(())
 }
+
+struct FunctionLinker<'a> {
+    functions: &'a HashMap<u64, Function>,
+    linked_functions: HashMap<u64, usize>,
+    relocations: &'a HashMap<SectionIndex, HashMap<u64, Relocation>>,
+    symbol_table: &'a HashMap<usize, Symbol>,
+}
+
+impl<'a> FunctionLinker<'a> {
+    fn new(
+        functions: &'a HashMap<u64, Function>,
+        relocations: &'a HashMap<SectionIndex, HashMap<u64, Relocation>>,
+        symbol_table: &'a HashMap<usize, Symbol>,
+    ) -> FunctionLinker<'a> {
+        FunctionLinker {
+            functions,
+            linked_functions: HashMap::new(),
+            relocations,
+            symbol_table,
+        }
+    }
+
+    fn link(mut self, program: &mut Program) -> Result<(), RelocationError> {
+        let mut fun = program.function.clone();
+        // relocate calls in the program's main function. As relocation happens,
+        // it will trigger linking in all the callees.
+        self.relocate(&mut fun, &program.function)?;
+
+        // this now includes the program function plus all the other functions called during
+        // execution
+        program.function = fun;
+
+        Ok(())
+    }
+
+    fn link_function(
+        &mut self,
+        program: &mut Function,
+        fun: &Function,
+    ) -> Result<usize, RelocationError> {
+        match self.linked_functions.get(&fun.address) {
+            Some(fun_ins_index) => return Ok(*fun_ins_index), // already linked
+            None => {}
+        };
+
+        // append fun.instructions to the program and record that `fun.address` has been inserted
+        // at `start_ins`. We'll use `start_ins` to do pc-relative calls.
+        let start_ins = program.instructions.len();
+        program.instructions.extend(&fun.instructions);
+        self.linked_functions.insert(fun.address, start_ins);
+
+        // relocate `fun`, recursively linking in all the callees
+        self.relocate(program, fun)?;
+
+        Ok(start_ins)
+    }
+
+    fn relocate(&mut self, program: &mut Function, fun: &Function) -> Result<(), RelocationError> {
+        let relocations = self.relocations.get(&fun.section_index);
+        let rel_info = |offset| relocations.and_then(|rels| rels.get(&offset));
+        let rel_target_address = |rel: &Relocation, symbol_table: &HashMap<usize, Symbol>| {
+            let sym =
+                symbol_table
+                    .get(&rel.symbol_index)
+                    .ok_or(RelocationError::UnknownSymbol {
+                        index: rel.symbol_index,
+                    })?;
+
+            Ok(sym.address)
+        };
+
+        let n_instructions = fun.instructions.len();
+        let start_ins = program.instructions.len() - 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 {
+            if !is_call(&program.instructions[ins_index]) {
+                continue;
+            }
+
+            let callee_address = if let Some(rel) = rel_info((ins_index * INS_SIZE) as u64) {
+                // 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.
+                rel_target_address(rel, self.symbol_table)?
+            } 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
+                    + (program.instructions[ins_index].imm + 1) as i64 * ins_size)
+                    as u64
+            };
+
+            // 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,
+                    })?;
+            let callee_ins_index = self.link_function(program, callee)?;
+
+            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
+            };
+        }
+
+        Ok(())
+    }
+}
+
+fn is_call(ins: &bpf_insn) -> bool {
+    let klass = (ins.code & 0x07) as u32;
+    let op = (ins.code & 0xF0) as u32;
+    let src = (ins.code & 0x08) as u32;
+
+    klass == BPF_JMP
+        && op == BPF_CALL
+        && src == BPF_K
+        && ins.src_reg() as u32 == BPF_PSEUDO_CALL
+        && ins.dst_reg() == 0
+        && ins.off == 0
+}

aya/src/programs/mod.rs

@@ -14,7 +14,11 @@ pub use socket_filter::{SocketFilter, SocketFilterError};
 pub use trace_point::{TracePoint, TracePointError};
 pub use xdp::{Xdp, XdpError};
 
-use crate::{generated::bpf_prog_type, obj, sys::bpf_load_program};
+use crate::{
+    generated::bpf_prog_type,
+    obj::{self, Function},
+    sys::bpf_load_program,
+};
 #[derive(Debug, Error)]
 pub enum ProgramError {
     #[error("the program is already loaded")]
@@ -207,13 +211,13 @@ fn load_program(prog_type: bpf_prog_type, data: &mut ProgramData) -> Result<(),
         return Err(ProgramError::AlreadyLoaded);
     }
     let crate::obj::Program {
-        instructions,
+        function: Function { instructions, .. },
         license,
         kernel_version,
         ..
     } = obj;
 
-    let mut ret = Ok(1);
+    let mut ret = Ok(1); // FIXME
     let mut log_buf = VerifierLog::new();
     for i in 0..3 {
         log_buf.reset();

xtask/src/codegen/aya.rs

@@ -88,6 +88,8 @@ fn codegen_bindings(opts: &Options) -> Result<(), anyhow::Error> {
         "BPF_W",
         "BPF_H",
         "BPF_B",
+        "BPF_JMP",
+        "BPF_CALL",
         "SO_ATTACH_BPF",
         "SO_DETACH_BPF",
         // BTF