Program unload API (#264)

aya: add `program.unload()` API

aya/src/programs/links.rs

@@ -79,6 +79,13 @@ impl<T: Link> LinkMap<T> {
             .detach()
     }
 
+    pub(crate) fn remove_all(&mut self) -> Result<(), ProgramError> {
+        for (_, link) in self.links.drain() {
+            link.detach()?;
+        }
+        Ok(())
+    }
+
     pub(crate) fn forget(&mut self, link_id: T::Id) -> Result<T, ProgramError> {
         self.links.remove(&link_id).ok_or(ProgramError::NotAttached)
     }
@@ -86,9 +93,7 @@ impl<T: Link> LinkMap<T> {
 
 impl<T: Link> Drop for LinkMap<T> {
     fn drop(&mut self) {
-        for (_, link) in self.links.drain() {
+        let _ = self.remove_all();
-            let _ = link.detach();
-        }
     }
 }
 

aya/src/programs/mod.rs

@@ -386,6 +386,15 @@ impl<T: Link> ProgramData<T> {
     }
 }
 
+fn unload_program<T: Link>(data: &mut ProgramData<T>) -> Result<(), ProgramError> {
+    data.links.remove_all()?;
+    let fd = data.fd.take().ok_or(ProgramError::NotLoaded)?;
+    unsafe {
+        libc::close(fd);
+    }
+    Ok(())
+}
+
 fn load_program<T: Link>(
     prog_type: bpf_prog_type,
     data: &mut ProgramData<T>,
@@ -540,6 +549,83 @@ impl<'a, P: ProgramFd> ProgramFd for &'a P {
     }
 }
 
+macro_rules! impl_program_unload {
+    ($($struct_name:ident),+ $(,)?) => {
+        $(
+            impl $struct_name {
+                /// Unloads the program from the kernel.
+                ///
+                /// Links will be detached before unloading the program.  Note
+                /// that owned links obtained using `forget_link()` will not be
+                /// detached.
+                pub fn unload(&mut self) -> Result<(), ProgramError> {
+                    unload_program(&mut self.data)
+                }
+            }
+        )+
+    }
+}
+
+impl_program_unload!(
+    KProbe,
+    UProbe,
+    TracePoint,
+    SocketFilter,
+    Xdp,
+    SkMsg,
+    SkSkb,
+    SchedClassifier,
+    CgroupSkb,
+    CgroupSysctl,
+    CgroupSockopt,
+    LircMode2,
+    PerfEvent,
+    Lsm,
+    RawTracePoint,
+    BtfTracePoint,
+    FEntry,
+    FExit,
+    Extension,
+    CgroupSockAddr,
+    SkLookup,
+    SockOps
+);
+
+#[cfg(test)]
+mod tests {
+    use super::Program;
+
+    #[allow(dead_code)]
+    // When a new program is added, this fn will break as a reminder: consider adding unload()
+    // See [impl_program_unload!()]
+    fn program_implements_unload(a: Program) {
+        let _ = match a {
+            Program::KProbe(mut p) => p.unload(),
+            Program::UProbe(mut p) => p.unload(),
+            Program::TracePoint(mut p) => p.unload(),
+            Program::SocketFilter(mut p) => p.unload(),
+            Program::Xdp(mut p) => p.unload(),
+            Program::SkMsg(mut p) => p.unload(),
+            Program::SkSkb(mut p) => p.unload(),
+            Program::SockOps(mut p) => p.unload(),
+            Program::SchedClassifier(mut p) => p.unload(),
+            Program::CgroupSkb(mut p) => p.unload(),
+            Program::CgroupSysctl(mut p) => p.unload(),
+            Program::CgroupSockopt(mut p) => p.unload(),
+            Program::LircMode2(mut p) => p.unload(),
+            Program::PerfEvent(mut p) => p.unload(),
+            Program::RawTracePoint(mut p) => p.unload(),
+            Program::Lsm(mut p) => p.unload(),
+            Program::BtfTracePoint(mut p) => p.unload(),
+            Program::FEntry(mut p) => p.unload(),
+            Program::FExit(mut p) => p.unload(),
+            Program::Extension(mut p) => p.unload(),
+            Program::CgroupSockAddr(mut p) => p.unload(),
+            Program::SkLookup(mut p) => p.unload(),
+        };
+    }
+}
+
 macro_rules! impl_program_fd {
     ($($struct_name:ident),+ $(,)?) => {
         $(

test/cases/010_load/030_unload/test.ebpf.rs

@@ -0,0 +1,26 @@
+//! ```cargo
+//! [dependencies]
+//! aya-bpf = { path = "../../../../bpf/aya-bpf" }
+//! ```
+
+#![no_std]
+#![no_main]
+
+use aya_bpf::{bindings::xdp_action, macros::xdp, programs::XdpContext};
+
+#[xdp(name = "test_unload")]
+pub fn pass(ctx: XdpContext) -> u32 {
+    match unsafe { try_pass(ctx) } {
+        Ok(ret) => ret,
+        Err(_) => xdp_action::XDP_ABORTED,
+    }
+}
+
+unsafe fn try_pass(_ctx: XdpContext) -> Result<u32, u32> {
+    Ok(xdp_action::XDP_PASS)
+}
+
+#[panic_handler]
+fn panic(_info: &core::panic::PanicInfo) -> ! {
+    unsafe { core::hint::unreachable_unchecked() }
+}

test/cases/010_load/030_unload/test.rs

@@ -0,0 +1,57 @@
+//! ```cargo
+//! [dependencies]
+//! aya = { path = "../../../../aya" }
+//! ```
+
+use aya::{
+    programs::{Xdp, XdpFlags},
+    Bpf,
+};
+use std::convert::TryInto;
+use std::process::Command;
+
+fn is_loaded() -> bool {
+    let output = Command::new("bpftool").args(&["prog"]).output().unwrap();
+    let stdout = String::from_utf8(output.stdout).unwrap();
+    stdout.contains("test_unload")
+}
+
+fn assert_loaded(loaded: bool) {
+    let state = is_loaded();
+    if state == loaded {
+        return;
+    }
+    panic!("Expected loaded: {} but was loaded: {}", loaded, state);
+}
+
+fn main() {
+    println!("Loading XDP program");
+    let mut bpf = Bpf::load_file("test.o").unwrap();
+    let dispatcher: &mut Xdp = bpf.program_mut("test_unload").unwrap().try_into().unwrap();
+
+    dispatcher.load().unwrap();
+
+    let link = dispatcher.attach("eth0", XdpFlags::default()).unwrap();
+
+    {
+        let link_owned = dispatcher.forget_link(link);
+
+        dispatcher.unload().unwrap();
+
+        assert_loaded(true);
+    };
+
+    assert_loaded(false);
+
+    dispatcher.load().unwrap();
+
+    assert_loaded(true);
+
+    dispatcher.attach("eth0", XdpFlags::default()).unwrap();
+
+    assert_loaded(true);
+
+    dispatcher.unload().unwrap();
+
+    assert_loaded(false);
+}

test/cases/010_load/030_unload/test.sh

@@ -0,0 +1,27 @@
+#!/bin/sh
+# SUMMARY: Check that the program can be unloaded
+# LABELS:
+
+set -ex
+
+# Source libraries. Uncomment if needed/defined
+#. "${RT_LIB}"
+. "${RT_PROJECT_ROOT}/_lib/lib.sh"
+
+NAME=test
+
+clean_up() {
+    rm -rf ebpf user ${NAME}.o ${NAME}
+    exec_vm rm ${NAME} ${NAME}.o
+}
+
+trap clean_up EXIT
+
+# Test code goes here
+compile_ebpf ${NAME}.ebpf.rs
+compile_user ${NAME}.rs
+
+scp_vm ${NAME}.o
+scp_vm ${NAME}
+
+exec_vm sudo ./${NAME}