Browse Source

Merge pull request #496 from dave-tucker/program-from-pinned3

aya: Add from_pin for Programs
Dave Tucker 2 years ago
parent
commit
811ab299de

+ 18 - 0
aya/src/programs/cgroup_skb.rs

@@ -2,6 +2,7 @@
 use std::{
     hash::Hash,
     os::unix::prelude::{AsRawFd, RawFd},
+    path::Path,
 };
 
 use crate::{
@@ -138,6 +139,23 @@ impl CgroupSkb {
     pub fn detach(&mut self, link_id: CgroupSkbLinkId) -> Result<(), ProgramError> {
         self.data.links.remove(link_id)
     }
+
+    /// Creates a program from a pinned entry on a bpffs.
+    ///
+    /// Existing links will not be populated. To work with existing links you should use [`crate::programs::links::PinnedLink`].
+    ///
+    /// On drop, any managed links are detached and the program is unloaded. This will not result in
+    /// the program being unloaded from the kernel if it is still pinned.
+    pub fn from_pin<P: AsRef<Path>>(
+        path: P,
+        expected_attach_type: CgroupSkbAttachType,
+    ) -> Result<Self, ProgramError> {
+        let data = ProgramData::from_pinned_path(path)?;
+        Ok(Self {
+            data,
+            expected_attach_type: Some(expected_attach_type),
+        })
+    }
 }
 
 #[derive(Debug, Hash, Eq, PartialEq)]

+ 15 - 0
aya/src/programs/cgroup_sock.rs

@@ -4,6 +4,7 @@ pub use aya_obj::programs::CgroupSockAttachType;
 use std::{
     hash::Hash,
     os::unix::prelude::{AsRawFd, RawFd},
+    path::Path,
 };
 
 use crate::{
@@ -113,6 +114,20 @@ impl CgroupSock {
     pub fn detach(&mut self, link_id: CgroupSockLinkId) -> Result<(), ProgramError> {
         self.data.links.remove(link_id)
     }
+
+    /// Creates a program from a pinned entry on a bpffs.
+    ///
+    /// Existing links will not be populated. To work with existing links you should use [`crate::programs::links::PinnedLink`].
+    ///
+    /// On drop, any managed links are detached and the program is unloaded. This will not result in
+    /// the program being unloaded from the kernel if it is still pinned.
+    pub fn from_pin<P: AsRef<Path>>(
+        path: P,
+        attach_type: CgroupSockAttachType,
+    ) -> Result<Self, ProgramError> {
+        let data = ProgramData::from_pinned_path(path)?;
+        Ok(Self { data, attach_type })
+    }
 }
 
 #[derive(Debug, Hash, Eq, PartialEq)]

+ 15 - 0
aya/src/programs/cgroup_sock_addr.rs

@@ -4,6 +4,7 @@ pub use aya_obj::programs::CgroupSockAddrAttachType;
 use std::{
     hash::Hash,
     os::unix::prelude::{AsRawFd, RawFd},
+    path::Path,
 };
 
 use crate::{
@@ -119,6 +120,20 @@ impl CgroupSockAddr {
     pub fn detach(&mut self, link_id: CgroupSockAddrLinkId) -> Result<(), ProgramError> {
         self.data.links.remove(link_id)
     }
+
+    /// Creates a program from a pinned entry on a bpffs.
+    ///
+    /// Existing links will not be populated. To work with existing links you should use [`crate::programs::links::PinnedLink`].
+    ///
+    /// On drop, any managed links are detached and the program is unloaded. This will not result in
+    /// the program being unloaded from the kernel if it is still pinned.
+    pub fn from_pin<P: AsRef<Path>>(
+        path: P,
+        attach_type: CgroupSockAddrAttachType,
+    ) -> Result<Self, ProgramError> {
+        let data = ProgramData::from_pinned_path(path)?;
+        Ok(Self { data, attach_type })
+    }
 }
 
 #[derive(Debug, Hash, Eq, PartialEq)]

+ 15 - 0
aya/src/programs/cgroup_sockopt.rs

@@ -4,6 +4,7 @@ pub use aya_obj::programs::CgroupSockoptAttachType;
 use std::{
     hash::Hash,
     os::unix::prelude::{AsRawFd, RawFd},
+    path::Path,
 };
 
 use crate::{
@@ -114,6 +115,20 @@ impl CgroupSockopt {
     pub fn detach(&mut self, link_id: CgroupSockoptLinkId) -> Result<(), ProgramError> {
         self.data.links.remove(link_id)
     }
+
+    /// Creates a program from a pinned entry on a bpffs.
+    ///
+    /// Existing links will not be populated. To work with existing links you should use [`crate::programs::links::PinnedLink`].
+    ///
+    /// On drop, any managed links are detached and the program is unloaded. This will not result in
+    /// the program being unloaded from the kernel if it is still pinned.
+    pub fn from_pin<P: AsRef<Path>>(
+        path: P,
+        attach_type: CgroupSockoptAttachType,
+    ) -> Result<Self, ProgramError> {
+        let data = ProgramData::from_pinned_path(path)?;
+        Ok(Self { data, attach_type })
+    }
 }
 
 #[derive(Debug, Hash, Eq, PartialEq)]

+ 12 - 1
aya/src/programs/kprobe.rs

@@ -1,5 +1,5 @@
 //! Kernel space probes.
-use std::io;
+use std::{io, path::Path};
 use thiserror::Error;
 
 use crate::{
@@ -83,6 +83,17 @@ impl KProbe {
     pub fn take_link(&mut self, link_id: KProbeLinkId) -> Result<KProbeLink, ProgramError> {
         self.data.take_link(link_id)
     }
+
+    /// Creates a program from a pinned entry on a bpffs.
+    ///
+    /// Existing links will not be populated. To work with existing links you should use [`crate::programs::links::PinnedLink`].
+    ///
+    /// On drop, any managed links are detached and the program is unloaded. This will not result in
+    /// the program being unloaded from the kernel if it is still pinned.
+    pub fn from_pin<P: AsRef<Path>>(path: P, kind: ProbeKind) -> Result<Self, ProgramError> {
+        let data = ProgramData::from_pinned_path(path)?;
+        Ok(Self { data, kind })
+    }
 }
 
 define_link_wrapper!(

+ 1 - 1
aya/src/programs/links.rs

@@ -181,7 +181,7 @@ impl PinnedLink {
         PinnedLink { inner: link, path }
     }
 
-    /// Creates a [`PinnedLink`] from a valid path on bpffs.
+    /// Creates a [`crate::programs::links::PinnedLink`] from a valid path on bpffs.
     pub fn from_pin<P: AsRef<Path>>(path: P) -> Result<Self, LinkError> {
         let path_string = CString::new(path.as_ref().to_string_lossy().to_string()).unwrap();
         let fd =

+ 1 - 0
aya/src/programs/lsm.rs

@@ -1,4 +1,5 @@
 //! LSM probes.
+
 use crate::{
     generated::{bpf_attach_type::BPF_LSM_MAC, bpf_prog_type::BPF_PROG_TYPE_LSM},
     obj::btf::{Btf, BtfKind},

+ 152 - 38
aya/src/programs/mod.rs

@@ -69,7 +69,7 @@ use std::{
     ffi::CString,
     io,
     os::unix::io::{AsRawFd, RawFd},
-    path::Path,
+    path::{Path, PathBuf},
 };
 use thiserror::Error;
 
@@ -108,8 +108,9 @@ use crate::{
     obj::{self, btf::BtfError, Function, KernelVersion},
     pin::PinError,
     sys::{
-        bpf_get_object, bpf_load_program, bpf_pin_object, bpf_prog_get_fd_by_id,
-        bpf_prog_get_info_by_fd, bpf_prog_query, retry_with_verifier_logs, BpfLoadProgramAttrs,
+        bpf_btf_get_fd_by_id, bpf_get_object, bpf_load_program, bpf_pin_object,
+        bpf_prog_get_fd_by_id, bpf_prog_get_info_by_fd, bpf_prog_query, retry_with_verifier_logs,
+        BpfLoadProgramAttrs,
     },
     util::VerifierLog,
 };
@@ -337,33 +338,33 @@ impl Program {
         }
     }
 
-    /// Unload the program
-    fn unload(&mut self) -> Result<(), ProgramError> {
+    /// Unloads the program from the kernel.
+    pub fn unload(self) -> Result<(), ProgramError> {
         match self {
-            Program::KProbe(p) => p.unload(),
-            Program::UProbe(p) => p.unload(),
-            Program::TracePoint(p) => p.unload(),
-            Program::SocketFilter(p) => p.unload(),
-            Program::Xdp(p) => p.unload(),
-            Program::SkMsg(p) => p.unload(),
-            Program::SkSkb(p) => p.unload(),
-            Program::SockOps(p) => p.unload(),
-            Program::SchedClassifier(p) => p.unload(),
-            Program::CgroupSkb(p) => p.unload(),
-            Program::CgroupSysctl(p) => p.unload(),
-            Program::CgroupSockopt(p) => p.unload(),
-            Program::LircMode2(p) => p.unload(),
-            Program::PerfEvent(p) => p.unload(),
-            Program::RawTracePoint(p) => p.unload(),
-            Program::Lsm(p) => p.unload(),
-            Program::BtfTracePoint(p) => p.unload(),
-            Program::FEntry(p) => p.unload(),
-            Program::FExit(p) => p.unload(),
-            Program::Extension(p) => p.unload(),
-            Program::CgroupSockAddr(p) => p.unload(),
-            Program::SkLookup(p) => p.unload(),
-            Program::CgroupSock(p) => p.unload(),
-            Program::CgroupDevice(p) => p.unload(),
+            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(),
+            Program::CgroupSock(mut p) => p.unload(),
+            Program::CgroupDevice(mut p) => p.unload(),
         }
     }
 
@@ -401,16 +402,10 @@ impl Program {
     }
 }
 
-impl Drop for Program {
-    fn drop(&mut self) {
-        let _ = self.unload();
-    }
-}
-
 #[derive(Debug)]
 pub(crate) struct ProgramData<T: Link> {
     pub(crate) name: Option<String>,
-    pub(crate) obj: obj::Program,
+    pub(crate) obj: Option<obj::Program>,
     pub(crate) fd: Option<RawFd>,
     pub(crate) links: LinkMap<T>,
     pub(crate) expected_attach_type: Option<bpf_attach_type>,
@@ -419,6 +414,7 @@ pub(crate) struct ProgramData<T: Link> {
     pub(crate) attach_prog_fd: Option<RawFd>,
     pub(crate) btf_fd: Option<RawFd>,
     pub(crate) verifier_log_level: u32,
+    pub(crate) path: Option<PathBuf>,
 }
 
 impl<T: Link> ProgramData<T> {
@@ -430,7 +426,7 @@ impl<T: Link> ProgramData<T> {
     ) -> ProgramData<T> {
         ProgramData {
             name,
-            obj,
+            obj: Some(obj),
             fd: None,
             links: LinkMap::new(),
             expected_attach_type: None,
@@ -439,8 +435,68 @@ impl<T: Link> ProgramData<T> {
             attach_prog_fd: None,
             btf_fd,
             verifier_log_level,
+            path: None,
         }
     }
+
+    pub(crate) fn from_bpf_prog_info(
+        name: Option<String>,
+        fd: RawFd,
+        path: &Path,
+        info: bpf_prog_info,
+    ) -> Result<ProgramData<T>, ProgramError> {
+        let attach_btf_id = if info.attach_btf_id > 0 {
+            Some(info.attach_btf_id)
+        } else {
+            None
+        };
+        let attach_btf_obj_fd = if info.attach_btf_obj_id > 0 {
+            let fd = bpf_btf_get_fd_by_id(info.attach_btf_obj_id).map_err(|io_error| {
+                ProgramError::SyscallError {
+                    call: "bpf_btf_get_fd_by_id".to_string(),
+                    io_error,
+                }
+            })?;
+            Some(fd as u32)
+        } else {
+            None
+        };
+
+        Ok(ProgramData {
+            name,
+            obj: None,
+            fd: Some(fd),
+            links: LinkMap::new(),
+            expected_attach_type: None,
+            attach_btf_obj_fd,
+            attach_btf_id,
+            attach_prog_fd: None,
+            btf_fd: None,
+            verifier_log_level: 0,
+            path: Some(path.to_path_buf()),
+        })
+    }
+
+    pub(crate) fn from_pinned_path<P: AsRef<Path>>(
+        path: P,
+    ) -> Result<ProgramData<T>, ProgramError> {
+        let path_string =
+            CString::new(path.as_ref().as_os_str().to_string_lossy().as_bytes()).unwrap();
+        let fd =
+            bpf_get_object(&path_string).map_err(|(_, io_error)| ProgramError::SyscallError {
+                call: "bpf_obj_get".to_owned(),
+                io_error,
+            })? as RawFd;
+
+        let info = bpf_prog_get_info_by_fd(fd).map_err(|io_error| ProgramError::SyscallError {
+            call: "bpf_prog_get_info_by_fd".to_owned(),
+            io_error,
+        })?;
+
+        let info = ProgramInfo(info);
+        let name = info.name_as_str().map(|s| s.to_string());
+        ProgramData::from_bpf_prog_info(name, fd, path.as_ref(), info.0)
+    }
 }
 
 impl<T: Link> ProgramData<T> {
@@ -493,6 +549,11 @@ fn load_program<T: Link>(
     if fd.is_some() {
         return Err(ProgramError::AlreadyLoaded);
     }
+    if obj.is_none() {
+        // This program was loaded from a pin in bpffs
+        return Err(ProgramError::AlreadyLoaded);
+    }
+    let obj = obj.as_ref().unwrap();
     let crate::obj::Program {
         function:
             Function {
@@ -621,6 +682,12 @@ macro_rules! impl_program_unload {
                     unload_program(&mut self.data)
                 }
             }
+
+            impl Drop for $struct_name {
+                fn drop(&mut self) {
+                    let _ = self.unload();
+                }
+            }
         )+
     }
 }
@@ -703,8 +770,17 @@ macro_rules! impl_program_pin{
                 /// To remove the program, the file on the BPF filesystem must be removed.
                 /// Any directories in the the path provided should have been created by the caller.
                 pub fn pin<P: AsRef<Path>>(&mut self, path: P) -> Result<(), PinError> {
+                    self.data.path = Some(path.as_ref().to_path_buf());
                     pin_program(&mut self.data, path)
                 }
+
+                /// Removes the pinned link from the filesystem.
+                pub fn unpin(mut self) -> Result<(), io::Error> {
+                    if let Some(path) = self.data.path.take() {
+                        std::fs::remove_file(path)?;
+                    }
+                    Ok(())
+                }
             }
         )+
     }
@@ -737,6 +813,45 @@ impl_program_pin!(
     CgroupDevice,
 );
 
+macro_rules! impl_from_pin {
+    ($($struct_name:ident),+ $(,)?) => {
+        $(
+            impl $struct_name {
+                /// Creates a program from a pinned entry on a bpffs.
+                ///
+                /// Existing links will not be populated. To work with existing links you should use [`crate::programs::links::PinnedLink`].
+                ///
+                /// On drop, any managed links are detached and the program is unloaded. This will not result in
+                /// the program being unloaded from the kernel if it is still pinned.
+                pub fn from_pin<P: AsRef<Path>>(path: P) -> Result<Self, ProgramError> {
+                    let data = ProgramData::from_pinned_path(path)?;
+                    Ok(Self { data })
+                }
+            }
+        )+
+    }
+}
+
+// Use impl_from_pin if the program doesn't require additional data
+impl_from_pin!(
+    TracePoint,
+    SocketFilter,
+    Xdp,
+    SkMsg,
+    CgroupSysctl,
+    LircMode2,
+    PerfEvent,
+    Lsm,
+    RawTracePoint,
+    BtfTracePoint,
+    FEntry,
+    FExit,
+    Extension,
+    SkLookup,
+    SockOps,
+    CgroupDevice,
+);
+
 macro_rules! impl_try_from_program {
     ($($ty:ident),+ $(,)?) => {
         $(
@@ -849,7 +964,6 @@ impl ProgramInfo {
         unsafe {
             libc::close(fd);
         }
-
         Ok(ProgramInfo(info))
     }
 }

+ 1 - 0
aya/src/programs/perf_event.rs

@@ -1,4 +1,5 @@
 //! Perf event programs.
+
 pub use crate::generated::{
     perf_hw_cache_id, perf_hw_cache_op_id, perf_hw_cache_op_result_id, perf_hw_id, perf_sw_ids,
 };

+ 12 - 1
aya/src/programs/sk_skb.rs

@@ -1,6 +1,6 @@
 //! Skskb programs.
 
-use std::os::unix::io::AsRawFd;
+use std::{os::unix::io::AsRawFd, path::Path};
 
 use crate::{
     generated::{
@@ -105,6 +105,17 @@ impl SkSkb {
     pub fn take_link(&mut self, link_id: SkSkbLinkId) -> Result<SkSkbLink, ProgramError> {
         self.data.take_link(link_id)
     }
+
+    /// Creates a program from a pinned entry on a bpffs.
+    ///
+    /// Existing links will not be populated. To work with existing links you should use [`crate::programs::links::PinnedLink`].
+    ///
+    /// On drop, any managed links are detached and the program is unloaded. This will not result in
+    /// the program being unloaded from the kernel if it is still pinned.
+    pub fn from_pin<P: AsRef<Path>>(path: P, kind: SkSkbKind) -> Result<Self, ProgramError> {
+        let data = ProgramData::from_pinned_path(path)?;
+        Ok(Self { data, kind })
+    }
 }
 
 define_link_wrapper!(

+ 15 - 0
aya/src/programs/tc.rs

@@ -4,6 +4,7 @@ use thiserror::Error;
 use std::{
     ffi::{CStr, CString},
     io,
+    path::Path,
 };
 
 use crate::{
@@ -190,6 +191,20 @@ impl SchedClassifier {
     ) -> Result<SchedClassifierLink, ProgramError> {
         self.data.take_link(link_id)
     }
+
+    /// Creates a program from a pinned entry on a bpffs.
+    ///
+    /// Existing links will not be populated. To work with existing links you should use [`crate::programs::links::PinnedLink`].
+    ///
+    /// On drop, any managed links are detached and the program is unloaded. This will not result in
+    /// the program being unloaded from the kernel if it is still pinned.
+    pub fn from_pin<P: AsRef<Path>>(path: P) -> Result<Self, ProgramError> {
+        let data = ProgramData::from_pinned_path(path)?;
+        let cname = CString::new(data.name.clone().unwrap_or_default())
+            .unwrap()
+            .into_boxed_c_str();
+        Ok(Self { data, name: cname })
+    }
 }
 
 #[derive(Debug, Hash, Eq, PartialEq)]

+ 1 - 0
aya/src/programs/tp_btf.rs

@@ -1,4 +1,5 @@
 //! BTF-enabled raw tracepoints.
+
 use crate::{
     generated::{bpf_attach_type::BPF_TRACE_RAW_TP, bpf_prog_type::BPF_PROG_TYPE_TRACING},
     obj::btf::{Btf, BtfKind},

+ 11 - 0
aya/src/programs/uprobe.rs

@@ -136,6 +136,17 @@ impl UProbe {
     pub fn take_link(&mut self, link_id: UProbeLinkId) -> Result<UProbeLink, ProgramError> {
         self.data.take_link(link_id)
     }
+
+    /// Creates a program from a pinned entry on a bpffs.
+    ///
+    /// Existing links will not be populated. To work with existing links you should use [`crate::programs::links::PinnedLink`].
+    ///
+    /// On drop, any managed links are detached and the program is unloaded. This will not result in
+    /// the program being unloaded from the kernel if it is still pinned.
+    pub fn from_pin<P: AsRef<Path>>(path: P, kind: ProbeKind) -> Result<Self, ProgramError> {
+        let data = ProgramData::from_pinned_path(path)?;
+        Ok(Self { data, kind })
+    }
 }
 
 define_link_wrapper!(

+ 22 - 2
test/integration-test/src/tests/load.rs

@@ -1,4 +1,4 @@
-use std::{process::Command, thread, time};
+use std::{convert::TryInto, process::Command, thread, time};
 
 use aya::{
     include_bytes_aligned,
@@ -192,16 +192,36 @@ fn pin_lifecycle() {
         let mut bpf = Bpf::load(bytes).unwrap();
         let prog: &mut Xdp = bpf.program_mut("pass").unwrap().try_into().unwrap();
         prog.load().unwrap();
+        prog.pin("/sys/fs/bpf/aya-xdp-test-prog").unwrap();
+    }
+
+    // should still be loaded since prog was pinned
+    assert_loaded!("pass", true);
+
+    // 2. Load program from bpffs but don't attach it
+    {
+        let _ = Xdp::from_pin("/sys/fs/bpf/aya-xdp-test-prog").unwrap();
+    }
+
+    // should still be loaded since prog was pinned
+    assert_loaded!("pass", true);
+
+    // 3. Load program from bpffs and attach
+    {
+        let mut prog = Xdp::from_pin("/sys/fs/bpf/aya-xdp-test-prog").unwrap();
         let link_id = prog.attach("lo", XdpFlags::default()).unwrap();
         let link = prog.take_link(link_id).unwrap();
         let fd_link: FdLink = link.try_into().unwrap();
         fd_link.pin("/sys/fs/bpf/aya-xdp-test-lo").unwrap();
+
+        // Unpin the program. It will stay attached since its links were pinned.
+        prog.unpin().unwrap();
     }
 
     // should still be loaded since link was pinned
     assert_loaded!("pass", true);
 
-    // 2. Load a new version of the program, unpin link, and atomically replace old program
+    // 4. Load a new version of the program, unpin link, and atomically replace old program
     {
         let mut bpf = Bpf::load(bytes).unwrap();
         let prog: &mut Xdp = bpf.program_mut("pass").unwrap().try_into().unwrap();