Skip to content

Commit

Permalink
uprobe: refactor target resolution
Browse files Browse the repository at this point in the history
This attempts to do fewer lossy conversions and to avoid some
allocations.
  • Loading branch information
ajwerner committed Aug 2, 2023
1 parent de8604d commit 81fb4e5
Showing 1 changed file with 36 additions and 40 deletions.
76 changes: 36 additions & 40 deletions aya/src/programs/uprobe.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
use libc::pid_t;
use object::{Object, ObjectSection, ObjectSymbol};
use std::{
borrow::Cow,
error::Error,
ffi::CStr,
fs,
Expand Down Expand Up @@ -82,7 +83,7 @@ impl UProbe {
target: T,
pid: Option<pid_t>,
) -> Result<UProbeLinkId, ProgramError> {
let path = resolve_attach_path(target, pid)?;
let path = resolve_attach_path(&target, pid)?;

let sym_offset = if let Some(fn_name) = fn_name {
resolve_symbol(&path, fn_name).map_err(|error| UProbeError::SymbolError {
Expand Down Expand Up @@ -124,41 +125,36 @@ impl UProbe {
}

fn resolve_attach_path<T: AsRef<Path>>(
target: T,
target: &T,
pid: Option<pid_t>,
) -> Result<String, UProbeError> {
// Look up the path for the target. If it there is a pid, and the target is a library name that
// is in the process's memory map, use the path of that library. Otherwise, use the target
// as-is.
) -> Result<Cow<'_, str>, UProbeError> {
// Look up the path for the target. If it there is a pid, and the target is a library name
// that is in the process's memory map, use the path of that library. Otherwise, use the target as-is.
let target = target.as_ref();
let target_str = &*target.as_os_str().to_string_lossy();

let mut path = if let Some(pid) = pid {
find_lib_in_proc_maps(pid, target_str).map_err(|io_error| UProbeError::FileError {
filename: format!("/proc/{pid}/maps"),
io_error,
})?
} else {
None
};

if path.is_none() {
path = if target.is_absolute() {
Some(target_str)
} else {
let cache = LD_SO_CACHE
.as_ref()
.map_err(|error| UProbeError::InvalidLdSoCache {
io_error: error.clone(),
})?;
cache.resolve(target_str)
}
.map(String::from)
};

path.ok_or(UProbeError::InvalidTarget {
let invalid_target = || UProbeError::InvalidTarget {
path: target.to_owned(),
};
let target_str = target.to_str().ok_or_else(invalid_target)?;
pid.and_then(|pid| {
find_lib_in_proc_maps(pid, target_str)
.map_err(|io_error| UProbeError::FileError {
filename: format!("/proc/{pid}/maps"),
io_error,
})
.map(|v| v.map(Cow::Owned))
.transpose()
})
.or_else(|| target.is_absolute().then(|| Ok(Cow::Borrowed(target_str))))
.or_else(|| {
LD_SO_CACHE
.as_ref()
.map_err(|error| UProbeError::InvalidLdSoCache {
io_error: error.clone(),
})
.map(|cache| cache.resolve(target_str).map(Cow::Borrowed))
.transpose()
})
.unwrap_or_else(|| Err(invalid_target()))
}

// Only run this test on linux with glibc because only in that configuration do we know that we'll
Expand All @@ -175,7 +171,7 @@ fn test_resolve_attach_path() {

// Now let's resolve the path to libc. It should exist in the current process's memory map and
// then in the ld.so.cache.
let libc_path = resolve_attach_path("libc", Some(pid)).unwrap();
let libc_path = resolve_attach_path(&"libc", Some(pid)).unwrap();

// Make sure we got a path that contains libc.
assert!(libc_path.contains("libc"), "libc_path: {}", libc_path);
Expand Down Expand Up @@ -276,16 +272,16 @@ fn find_lib_in_proc_maps(pid: pid_t, lib: &str) -> Result<Option<String>, io::Er
let libs = proc_maps_libs(pid)?;

let ret = if lib.contains(".so") {
libs.iter().find(|(k, _)| k.as_str().starts_with(lib))
libs.into_iter().find(|(k, _)| k.as_str().starts_with(lib))
} else {
let lib = lib.to_string();
let lib1 = lib.clone() + ".so";
let lib2 = lib + "-";
libs.iter()
.find(|(k, _)| k.starts_with(&lib1) || k.starts_with(&lib2))
libs.into_iter().find(|(k, _)| {
k.strip_prefix(lib)
.map(|k| k.starts_with(".so") || k.starts_with('-'))
.unwrap_or_default()
})
};

Ok(ret.map(|(_, v)| v.clone()))
Ok(ret.map(|(_, v)| v))
}

#[derive(Debug)]
Expand Down

0 comments on commit 81fb4e5

Please sign in to comment.