Merge pull request #748 from aya-rs/btf_obj_fd-owned

programs: Plug attach_btf_obj_fd leak

aya/src/programs/extension.rs

@@ -1,5 +1,5 @@
 //! Extension programs.
-use std::os::fd::{AsRawFd, RawFd};
+use std::os::fd::{AsFd as _, AsRawFd as _, OwnedFd};
 use thiserror::Error;
 
 use object::Endianness;
@@ -72,7 +72,7 @@ impl Extension {
         let target_prog_fd = program.as_raw_fd();
         let (btf_fd, btf_id) = get_btf_info(target_prog_fd, func_name)?;
 
-        self.data.attach_btf_obj_fd = Some(btf_fd as u32);
+        self.data.attach_btf_obj_fd = Some(btf_fd);
         self.data.attach_prog_fd = Some(target_prog_fd);
         self.data.attach_btf_id = Some(btf_id);
         load_program(BPF_PROG_TYPE_EXT, &mut self.data)
@@ -149,7 +149,7 @@ impl Extension {
 
 /// Retrieves the FD of the BTF object for the provided `prog_fd` and the BTF ID of the function
 /// with the name `func_name` within that BTF object.
-fn get_btf_info(prog_fd: i32, func_name: &str) -> Result<(RawFd, u32), ProgramError> {
+fn get_btf_info(prog_fd: i32, func_name: &str) -> Result<(OwnedFd, u32), ProgramError> {
     // retrieve program information
     let info = sys::bpf_prog_get_info_by_fd(prog_fd, &mut [])?;
 
@@ -165,7 +165,7 @@ fn get_btf_info(prog_fd: i32, func_name: &str) -> Result<(RawFd, u32), ProgramEr
     // assume 4kb. if this is too small we can resize based on the size obtained in the response.
     let mut buf = vec![0u8; 4096];
     loop {
-        let info = sys::btf_obj_get_info_by_fd(btf_fd, &mut buf)?;
+        let info = sys::btf_obj_get_info_by_fd(btf_fd.as_fd(), &mut buf)?;
         let btf_size = info.btf_size as usize;
         if btf_size > buf.len() {
             buf.resize(btf_size, 0u8);

aya/src/programs/mod.rs

@@ -406,7 +406,7 @@ pub(crate) struct ProgramData<T: Link> {
     pub(crate) fd: Option<RawFd>,
     pub(crate) links: LinkMap<T>,
     pub(crate) expected_attach_type: Option<bpf_attach_type>,
-    pub(crate) attach_btf_obj_fd: Option<u32>,
+    pub(crate) attach_btf_obj_fd: Option<OwnedFd>,
     pub(crate) attach_btf_id: Option<u32>,
     pub(crate) attach_prog_fd: Option<RawFd>,
     pub(crate) btf_fd: Option<Arc<OwnedFd>>,
@@ -450,16 +450,9 @@ impl<T: Link> ProgramData<T> {
         } 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| SyscallError {
-                    call: "bpf_btf_get_fd_by_id",
-                    io_error,
-                })?;
-            Some(fd as u32)
-        } else {
-            None
-        };
+        let attach_btf_obj_fd = (info.attach_btf_obj_id != 0)
+            .then(|| bpf_btf_get_fd_by_id(info.attach_btf_obj_id))
+            .transpose()?;
 
         Ok(ProgramData {
             name,
@@ -604,7 +597,7 @@ fn load_program<T: Link>(
         kernel_version: target_kernel_version,
         expected_attach_type: *expected_attach_type,
         prog_btf_fd: btf_fd.as_ref().map(|f| f.as_fd()),
-        attach_btf_obj_fd: *attach_btf_obj_fd,
+        attach_btf_obj_fd: attach_btf_obj_fd.as_ref().map(|fd| fd.as_fd()),
         attach_btf_id: *attach_btf_id,
         attach_prog_fd: *attach_prog_fd,
         func_info_rec_size: *func_info_rec_size,

aya/src/sys/bpf.rs

@@ -118,7 +118,7 @@ pub(crate) struct BpfLoadProgramAttrs<'a> {
     pub(crate) kernel_version: u32,
     pub(crate) expected_attach_type: Option<bpf_attach_type>,
     pub(crate) prog_btf_fd: Option<BorrowedFd<'a>>,
-    pub(crate) attach_btf_obj_fd: Option<u32>,
+    pub(crate) attach_btf_obj_fd: Option<BorrowedFd<'a>>,
     pub(crate) attach_btf_id: Option<u32>,
     pub(crate) attach_prog_fd: Option<RawFd>,
     pub(crate) func_info_rec_size: usize,
@@ -181,7 +181,7 @@ pub(crate) fn bpf_load_program(
         u.log_size = log_buf.len() as u32;
     }
     if let Some(v) = aya_attr.attach_btf_obj_fd {
-        u.__bindgen_anon_1.attach_btf_obj_fd = v;
+        u.__bindgen_anon_1.attach_btf_obj_fd = v.as_raw_fd() as _;
     }
     if let Some(v) = aya_attr.attach_prog_fd {
         u.__bindgen_anon_1.attach_prog_fd = v as u32;
@@ -539,10 +539,9 @@ pub(crate) fn bpf_link_get_info_by_fd(fd: BorrowedFd<'_>) -> Result<bpf_link_inf
 }
 
 pub(crate) fn btf_obj_get_info_by_fd(
-    fd: RawFd,
+    fd: BorrowedFd<'_>,
     buf: &mut [u8],
 ) -> Result<bpf_btf_info, SyscallError> {
-    let fd = unsafe { BorrowedFd::borrow_raw(fd) };
     bpf_obj_get_info_by_fd(fd, |info: &mut bpf_btf_info| {
         info.btf = buf.as_mut_ptr() as _;
         info.btf_size = buf.len() as _;
@@ -595,14 +594,18 @@ unsafe fn fd_sys_bpf(cmd: bpf_cmd, attr: &mut bpf_attr) -> SysResult<OwnedFd> {
     Ok(OwnedFd::from_raw_fd(fd))
 }
 
-pub(crate) fn bpf_btf_get_fd_by_id(id: u32) -> Result<RawFd, io::Error> {
+pub(crate) fn bpf_btf_get_fd_by_id(id: u32) -> Result<OwnedFd, SyscallError> {
     let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
     attr.__bindgen_anon_6.__bindgen_anon_1.btf_id = id;
 
-    match sys_bpf(bpf_cmd::BPF_BTF_GET_FD_BY_ID, &mut attr) {
-        Ok(v) => Ok(v as RawFd),
-        Err((_, err)) => Err(err),
-    }
+    // SAFETY: BPF_BTF_GET_FD_BY_ID returns a new file descriptor.
+    unsafe { fd_sys_bpf(bpf_cmd::BPF_BTF_GET_FD_BY_ID, &mut attr) }.map_err(|(code, io_error)| {
+        assert_eq!(code, -1);
+        SyscallError {
+            call: "bpf_btf_get_fd_by_id",
+            io_error,
+        }
+    })
 }
 
 pub(crate) fn is_prog_name_supported() -> bool {