From c7a19bcefba25455279d9e718f6430dee7a84b74 Mon Sep 17 00:00:00 2001
From: Tamir Duberstein <tamird@gmail.com>
Date: Thu, 10 Aug 2023 18:19:24 -0400
Subject: [PATCH] sys: add map_ids to bpf_prog_get_info_by_fd

Allows the caller to pass a slice which the kernel will populate with
map ids used by the program.
---
 aya/src/programs/extension.rs  |  2 +-
 aya/src/programs/lirc_mode2.rs |  2 +-
 aya/src/programs/mod.rs        |  6 +++---
 aya/src/sys/bpf.rs             | 26 ++++++++++++++++++--------
 4 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/aya/src/programs/extension.rs b/aya/src/programs/extension.rs
index 143c909f..00f8b02a 100644
--- a/aya/src/programs/extension.rs
+++ b/aya/src/programs/extension.rs
@@ -151,7 +151,7 @@ impl Extension {
 /// with the name `func_name` within that BTF object.
 fn get_btf_info(prog_fd: i32, func_name: &str) -> Result<(RawFd, u32), ProgramError> {
     // retrieve program information
-    let info = sys::bpf_prog_get_info_by_fd(prog_fd)?;
+    let info = sys::bpf_prog_get_info_by_fd(prog_fd, &mut [])?;
 
     // btf_id refers to the ID of the program btf that was loaded with bpf(BPF_BTF_LOAD)
     if info.btf_id == 0 {
diff --git a/aya/src/programs/lirc_mode2.rs b/aya/src/programs/lirc_mode2.rs
index e39bbf1a..39e1d820 100644
--- a/aya/src/programs/lirc_mode2.rs
+++ b/aya/src/programs/lirc_mode2.rs
@@ -131,7 +131,7 @@ impl LircLink {
 
     /// Get ProgramInfo from this link
     pub fn info(&self) -> Result<ProgramInfo, ProgramError> {
-        bpf_prog_get_info_by_fd(self.prog_fd)
+        bpf_prog_get_info_by_fd(self.prog_fd, &mut [])
             .map(ProgramInfo)
             .map_err(Into::into)
     }
diff --git a/aya/src/programs/mod.rs b/aya/src/programs/mod.rs
index eab30abe..ff643fa5 100644
--- a/aya/src/programs/mod.rs
+++ b/aya/src/programs/mod.rs
@@ -488,7 +488,7 @@ impl<T: Link> ProgramData<T> {
             io_error,
         })?;
 
-        let info = bpf_prog_get_info_by_fd(fd.as_raw_fd())?;
+        let info = bpf_prog_get_info_by_fd(fd.as_raw_fd(), &mut [])?;
         let name = ProgramInfo(info).name_as_str().map(|s| s.to_string());
         ProgramData::from_bpf_prog_info(name, fd, path.as_ref(), info, verifier_log_level)
     }
@@ -955,7 +955,7 @@ impl ProgramInfo {
             io_error,
         })?;
 
-        let info = bpf_prog_get_info_by_fd(fd.as_raw_fd())?;
+        let info = bpf_prog_get_info_by_fd(fd.as_raw_fd(), &mut [])?;
         Ok(ProgramInfo(info))
     }
 }
@@ -991,7 +991,7 @@ pub fn loaded_programs() -> impl Iterator<Item = Result<ProgramInfo, ProgramErro
         })
         .map(|fd| {
             let fd = fd?;
-            bpf_prog_get_info_by_fd(fd.as_raw_fd())
+            bpf_prog_get_info_by_fd(fd.as_raw_fd(), &mut [])
         })
         .map(|result| result.map(ProgramInfo).map_err(Into::into))
 }
diff --git a/aya/src/sys/bpf.rs b/aya/src/sys/bpf.rs
index 597edc76..db70fe5a 100644
--- a/aya/src/sys/bpf.rs
+++ b/aya/src/sys/bpf.rs
@@ -477,10 +477,14 @@ pub(crate) fn bpf_prog_get_fd_by_id(prog_id: u32) -> Result<OwnedFd, SyscallErro
     })
 }
 
-fn bpf_obj_get_info_by_fd<T>(fd: BorrowedFd<'_>) -> Result<T, SyscallError> {
+fn bpf_obj_get_info_by_fd<T, F: FnOnce(&mut T)>(
+    fd: BorrowedFd<'_>,
+    init: F,
+) -> Result<T, SyscallError> {
     let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
-    // info gets entirely populated by the kernel
-    let info = MaybeUninit::zeroed();
+    let mut info = unsafe { mem::zeroed() };
+
+    init(&mut info);
 
     attr.info.bpf_fd = fd.as_raw_fd() as u32;
     attr.info.info = &info as *const _ as u64;
@@ -489,7 +493,7 @@ fn bpf_obj_get_info_by_fd<T>(fd: BorrowedFd<'_>) -> Result<T, SyscallError> {
     match sys_bpf(bpf_cmd::BPF_OBJ_GET_INFO_BY_FD, &mut attr) {
         Ok(code) => {
             assert_eq!(code, 0);
-            Ok(unsafe { info.assume_init() })
+            Ok(info)
         }
         Err((code, io_error)) => {
             assert_eq!(code, -1);
@@ -501,13 +505,19 @@ fn bpf_obj_get_info_by_fd<T>(fd: BorrowedFd<'_>) -> Result<T, SyscallError> {
     }
 }
 
-pub(crate) fn bpf_prog_get_info_by_fd(fd: RawFd) -> Result<bpf_prog_info, SyscallError> {
+pub(crate) fn bpf_prog_get_info_by_fd(
+    fd: RawFd,
+    map_ids: &mut [u32],
+) -> Result<bpf_prog_info, SyscallError> {
     let fd = unsafe { BorrowedFd::borrow_raw(fd) };
-    bpf_obj_get_info_by_fd::<bpf_prog_info>(fd)
+    bpf_obj_get_info_by_fd(fd, |info: &mut bpf_prog_info| {
+        info.nr_map_ids = map_ids.len() as _;
+        info.map_ids = map_ids.as_mut_ptr() as _;
+    })
 }
 
 pub(crate) fn bpf_map_get_info_by_fd(fd: BorrowedFd<'_>) -> Result<bpf_map_info, SyscallError> {
-    bpf_obj_get_info_by_fd::<bpf_map_info>(fd)
+    bpf_obj_get_info_by_fd(fd, |_| {})
 }
 
 pub(crate) fn bpf_link_get_fd_by_id(link_id: u32) -> Result<OwnedFd, SyscallError> {
@@ -525,7 +535,7 @@ pub(crate) fn bpf_link_get_fd_by_id(link_id: u32) -> Result<OwnedFd, SyscallErro
 }
 
 pub(crate) fn bpf_link_get_info_by_fd(fd: BorrowedFd<'_>) -> Result<bpf_link_info, SyscallError> {
-    bpf_obj_get_info_by_fd::<bpf_link_info>(fd)
+    bpf_obj_get_info_by_fd(fd, |_| {})
 }
 
 pub(crate) fn btf_obj_get_info_by_fd(