Skip to content

Commit

Permalink
aya: add Bpf::many_programs_mut
Browse files Browse the repository at this point in the history
Store programs in a `hashbrown::HashMap` and expose `get_many_mut`. We
can revisit this dependency when
rust-lang/rust#97601 is resolved.
  • Loading branch information
tamird committed Aug 11, 2023
1 parent 5b872db commit 8717bf1
Show file tree
Hide file tree
Showing 4 changed files with 92 additions and 13 deletions.
1 change: 1 addition & 0 deletions aya/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ async-io = { workspace = true, optional = true }
aya-obj = { workspace = true, features = ["std"] }
bitflags = { workspace = true }
bytes = { workspace = true }
hashbrown = { workspace = true, default-features = true }
lazy_static = { workspace = true }
libc = { workspace = true }
log = { workspace = true }
Expand Down
29 changes: 28 additions & 1 deletion aya/src/bpf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -768,7 +768,9 @@ impl<'a> Default for BpfLoader<'a> {
#[derive(Debug)]
pub struct Bpf {
maps: HashMap<String, Map>,
programs: HashMap<String, Program>,
// TODO(https://github.com/rust-lang/rust/issues/97601): use `HashMap` when `get_many_mut` is
// stabilized.
programs: hashbrown::HashMap<String, Program>,
}

impl Bpf {
Expand Down Expand Up @@ -914,6 +916,31 @@ impl Bpf {
self.programs.get_mut(name)
}

/// Attempts to get mutable references to `N` programs in the map at once.
///
/// Returns an array of length `N` with the results of each query. For soundness, at most one
/// mutable reference will be returned to any value. `None` will be returned if any of the
/// keys are duplicates or missing.
///
/// # Examples
///
/// ```no_run
/// # let mut bpf = aya::Bpf::load(&[])?;
/// use aya::programs::UProbe;
///
/// let [ssl_read, ssl_write ]= bpf.many_programs_mut(["SSL_read", "SSL_write"]).unwrap();
/// let ssl_read: &mut UProbe = ssl_read.try_into()?;
/// let ssl_write: &mut UProbe = ssl_write.try_into()?;
/// ssl_read.load()?;
/// ssl_read.attach(Some("SSL_read"), 0, "libssl", None)?;
/// ssl_write.load()?;
/// ssl_write.attach(Some("SSL_write"), 0, "libssl", None)?;
/// # Ok::<(), aya::BpfError>(())
/// ```
pub fn many_programs_mut<const N: usize>(&mut self, names: [&str; N]) -> Option<[&mut Program; N]> {
self.programs.get_many_mut(names)
}

/// An iterator over all the programs.
///
/// # Examples
Expand Down
15 changes: 3 additions & 12 deletions aya/src/programs/extension.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,22 +37,13 @@ pub enum ExtensionError {
/// use aya::{BpfLoader, programs::{Xdp, XdpFlags, Extension}};
///
/// let mut bpf = BpfLoader::new().extension("extension").load_file("app.o")?;
/// let mut prog = None;
/// let mut ext = None;
/// for (name, program) in bpf.programs_mut() {
/// match name {
/// "main" => prog = Some(program),
/// "extension" => ext = Some(program),
/// _ => {},
/// }
/// }
///
/// let prog: &mut Xdp = prog.unwrap().try_into()?;
/// let [prog, ext] = bpf.many_programs_mut(["main", "extension"]).unwrap();
/// let prog: &mut Xdp = prog.try_into()?;
/// prog.load()?;
/// prog.attach("eth0", XdpFlags::default())?;
///
/// let prog_fd = prog.fd().unwrap();
/// let ext: &mut Extension = ext.unwrap().try_into()?;
/// let ext: &mut Extension = ext.try_into()?;
/// ext.load(prog_fd, "function_to_replace")?;
/// ext.attach()?;
/// Ok::<(), aya::BpfError>(())
Expand Down
Loading

0 comments on commit 8717bf1

Please sign in to comment.