Skip to content

Commit

Permalink
Merge pull request #602 from marysaka/fix/btf-reloc-all-functions
Browse files Browse the repository at this point in the history
aya: Apply BTF relocations to all functions
  • Loading branch information
alessandrod committed May 29, 2023
2 parents 3211d2c + c4e721f commit 3a9a54f
Show file tree
Hide file tree
Showing 9 changed files with 211 additions and 113 deletions.
12 changes: 6 additions & 6 deletions aya-obj/src/btf/btf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use alloc::{
use bytes::BufMut;

use log::debug;
use object::Endianness;
use object::{Endianness, SectionIndex};

use crate::{
btf::{
Expand Down Expand Up @@ -414,7 +414,7 @@ impl Btf {

pub(crate) fn fixup_and_sanitize(
&mut self,
section_sizes: &HashMap<String, u64>,
section_infos: &HashMap<String, (SectionIndex, u64)>,
symbol_offsets: &HashMap<String, u64>,
features: &BtfFeatures,
) -> Result<(), BtfError> {
Expand Down Expand Up @@ -489,8 +489,8 @@ impl Btf {
} else {
// We need to get the size of the section from the ELF file
// Fortunately, we cached these when parsing it initially
// and we can this up by name in section_sizes
let size = section_sizes.get(&name).ok_or_else(|| {
// and we can this up by name in section_infos
let (_, size) = section_infos.get(&name).ok_or_else(|| {
BtfError::UnknownSectionSize {
section_name: name.clone(),
}
Expand Down Expand Up @@ -639,7 +639,7 @@ impl Object {
}
// fixup btf
obj_btf.fixup_and_sanitize(
&self.section_sizes,
&self.section_infos,
&self.symbol_offset_by_name,
features,
)?;
Expand Down Expand Up @@ -1277,7 +1277,7 @@ mod tests {
};

btf.fixup_and_sanitize(
&HashMap::from([(".data/foo".to_string(), 32u64)]),
&HashMap::from([(".data/foo".to_string(), (SectionIndex(0), 32u64))]),
&HashMap::from([("foo".to_string(), 64u64)]),
&features,
)
Expand Down
104 changes: 76 additions & 28 deletions aya-obj/src/btf/relocation.rs
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
use core::{mem, ptr, str::FromStr};
use core::{mem, ops::Bound::Included, ptr};

use alloc::{
borrow::ToOwned,
collections::BTreeMap,
format,
string::{String, ToString},
vec,
vec::Vec,
};
use object::SectionIndex;

use crate::{
btf::{
Expand All @@ -18,7 +20,7 @@ use crate::{
BPF_K, BPF_LD, BPF_LDX, BPF_ST, BPF_STX, BPF_W, BTF_INT_SIGNED,
},
util::HashMap,
Object, Program, ProgramSection,
Function, Object,
};

#[cfg(not(feature = "std"))]
Expand All @@ -43,9 +45,13 @@ enum RelocationError {
#[error(transparent)]
IOError(#[from] std::io::Error),

/// Program not found
#[error("program not found")]
ProgramNotFound,
/// Section not found
#[error("section not found")]
SectionNotFound,

/// Function not found
#[error("function not found")]
FunctionNotFound,

/// Invalid relocation access string
#[error("invalid relocation access string {access_str}")]
Expand Down Expand Up @@ -227,25 +233,26 @@ impl Object {
error: RelocationError::BtfError(e),
})?;

let program_section = match ProgramSection::from_str(&section_name) {
Ok(program) => program,
Err(_) => continue,
};
let section_name = program_section.name();

let program = self
.programs
.get_mut(section_name)
.ok_or(BtfRelocationError {
section: section_name.to_owned(),
error: RelocationError::ProgramNotFound,
let (section_index, _) = self
.section_infos
.get(&section_name.to_string())
.ok_or_else(|| BtfRelocationError {
section: section_name.to_string(),
error: RelocationError::SectionNotFound,
})?;
match relocate_btf_program(program, relos, local_btf, target_btf, &mut candidates_cache)
{

match relocate_btf_functions(
section_index,
&mut self.functions,
relos,
local_btf,
target_btf,
&mut candidates_cache,
) {
Ok(_) => {}
Err(error) => {
return Err(BtfRelocationError {
section: section_name.to_owned(),
section: section_name.to_string(),
error,
})
}
Expand All @@ -256,16 +263,55 @@ impl Object {
}
}

fn relocate_btf_program<'target>(
program: &mut Program,
fn is_relocation_inside_function(
section_index: &SectionIndex,
func: &Function,
rel: &Relocation,
) -> bool {
if section_index.0 != func.section_index.0 {
return false;
}

let ins_offset = rel.ins_offset / mem::size_of::<bpf_insn>();
let func_offset = func.section_offset / mem::size_of::<bpf_insn>();
let func_size = func.instructions.len();

(func_offset..func_offset + func_size).contains(&ins_offset)
}

fn function_by_relocation<'a>(
section_index: &SectionIndex,
functions: &'a mut BTreeMap<(usize, u64), Function>,
rel: &Relocation,
) -> Option<&'a mut Function> {
functions
.range_mut((
Included(&(section_index.0, 0)),
Included(&(section_index.0, u64::MAX)),
))
.map(|(_, func)| func)
.find(|func| is_relocation_inside_function(section_index, func, rel))
}

fn relocate_btf_functions<'target>(
section_index: &SectionIndex,
functions: &mut BTreeMap<(usize, u64), Function>,
relos: &[Relocation],
local_btf: &Btf,
target_btf: &'target Btf,
candidates_cache: &mut HashMap<u32, Vec<Candidate<'target>>>,
) -> Result<(), RelocationError> {
let mut last_function_opt: Option<&mut Function> = None;

for rel in relos {
let instructions = &mut program.function.instructions;
let ins_index = rel.ins_offset / mem::size_of::<bpf_insn>();
let function = match last_function_opt.take() {
Some(func) if is_relocation_inside_function(section_index, func, rel) => func,
_ => function_by_relocation(section_index, functions, rel)
.ok_or(RelocationError::FunctionNotFound)?,
};

let instructions = &mut function.instructions;
let ins_index = (rel.ins_offset - function.section_offset) / mem::size_of::<bpf_insn>();
if ins_index >= instructions.len() {
return Err(RelocationError::InvalidInstructionIndex {
index: ins_index,
Expand Down Expand Up @@ -337,7 +383,9 @@ fn relocate_btf_program<'target>(
ComputedRelocation::new(rel, &local_spec, None)?
};

comp_rel.apply(program, rel, local_btf, target_btf)?;
comp_rel.apply(function, rel, local_btf, target_btf)?;

last_function_opt = Some(function);
}

Ok(())
Expand Down Expand Up @@ -847,14 +895,14 @@ impl ComputedRelocation {

fn apply(
&self,
program: &mut Program,
function: &mut Function,
rel: &Relocation,
local_btf: &Btf,
target_btf: &Btf,
) -> Result<(), RelocationError> {
let instructions = &mut program.function.instructions;
let instructions = &mut function.instructions;
let num_instructions = instructions.len();
let ins_index = rel.ins_offset / mem::size_of::<bpf_insn>();
let ins_index = (rel.ins_offset - function.section_offset) / mem::size_of::<bpf_insn>();
let ins =
instructions
.get_mut(ins_index)
Expand Down
3 changes: 2 additions & 1 deletion aya-obj/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,8 @@
//! object.relocate_maps(std::iter::empty(), &text_sections).unwrap();
//!
//! // Run with rbpf
//! let instructions = &object.programs["prog_name"].function.instructions;
//! let function = object.functions.get(&object.programs["prog_name"].function_key()).unwrap();
//! let instructions = &function.instructions;
//! let data = unsafe {
//! core::slice::from_raw_parts(
//! instructions.as_ptr() as *const u8,
Expand Down
Loading

0 comments on commit 3a9a54f

Please sign in to comment.