diff --git a/aya/src/programs/mod.rs b/aya/src/programs/mod.rs index 0947c8fd..20eac41f 100644 --- a/aya/src/programs/mod.rs +++ b/aya/src/programs/mod.rs @@ -110,7 +110,7 @@ use crate::{ pin::PinError, sys::{ 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_get_next_id, bpf_prog_query, + bpf_prog_get_fd_by_id, bpf_prog_get_info_by_fd, bpf_prog_query, iter_prog_ids, retry_with_verifier_logs, BpfLoadProgramAttrs, SyscallError, }, util::KernelVersion, @@ -971,56 +971,6 @@ impl ProgramInfo { } } -/// ProgramsIter is an Iterator over loaded eBPF programs. -pub struct ProgramsIter { - current: u32, - error: bool, -} - -impl Iterator for ProgramsIter { - type Item = Result; - - fn next(&mut self) -> Option { - if self.error { - return None; - } - let current = self.current; - - match bpf_prog_get_next_id(current) { - Ok(Some(next)) => { - self.current = next; - Some( - bpf_prog_get_fd_by_id(next) - .map_err(|io_error| SyscallError { - call: "bpf_prog_get_fd_by_id", - io_error, - }) - .and_then(|fd| { - bpf_prog_get_info_by_fd(fd.as_raw_fd()) - .map_err(|io_error| SyscallError { - call: "bpf_prog_get_info_by_fd", - io_error, - }) - .map(ProgramInfo) - }) - .map_err(Into::into), - ) - } - Ok(None) => None, - Err((_, io_error)) => { - // If getting the next program failed, we have to yield None in our next - // iteration to avoid an infinite loop. - self.error = true; - Some(Err(SyscallError { - call: "bpf_prog_get_fd_by_id", - io_error, - } - .into())) - } - } - } -} - /// Returns an iterator over all loaded bpf programs. /// /// This differs from [`crate::Bpf::programs`] since it will return all programs @@ -1044,9 +994,21 @@ impl Iterator for ProgramsIter { /// next program id, get the program fd, or the [`ProgramInfo`] fail. In cases where /// iteration can't be performed, for example the caller does not have the necessary privileges, /// a single item will be yielded containing the error that occurred. -pub fn loaded_programs() -> ProgramsIter { - ProgramsIter { - current: 0, - error: false, - } +pub fn loaded_programs() -> impl Iterator> { + iter_prog_ids() + .map(|id| { + let id = id?; + bpf_prog_get_fd_by_id(id).map_err(|io_error| SyscallError { + call: "bpf_prog_get_fd_by_id", + io_error, + }) + }) + .map(|fd| { + let fd = fd?; + bpf_prog_get_info_by_fd(fd.as_raw_fd()).map_err(|io_error| SyscallError { + call: "bpf_prog_get_info_by_fd", + io_error, + }) + }) + .map(|result| result.map(ProgramInfo).map_err(Into::into)) } diff --git a/aya/src/sys/bpf.rs b/aya/src/sys/bpf.rs index ab8de01c..e4a4becc 100644 --- a/aya/src/sys/bpf.rs +++ b/aya/src/sys/bpf.rs @@ -1,7 +1,7 @@ use std::{ cmp::{self, min}, ffi::{CStr, CString}, - io, + io, iter, mem::{self, MaybeUninit}, os::fd::{AsRawFd, BorrowedFd, FromRawFd as _, OwnedFd, RawFd}, slice, @@ -28,7 +28,7 @@ use crate::{ }, copy_instructions, }, - sys::{syscall, SysResult, Syscall}, + sys::{syscall, SysResult, Syscall, SyscallError}, Btf, Pod, VerifierLogLevel, BPF_OBJ_NAME_LEN, }; @@ -994,17 +994,44 @@ fn sys_bpf(cmd: bpf_cmd, attr: &mut bpf_attr) -> SysResult { syscall(Syscall::Bpf { cmd, attr }) } -pub(crate) fn bpf_prog_get_next_id(id: u32) -> Result, (c_long, io::Error)> { +fn bpf_prog_get_next_id(id: u32) -> Result, SyscallError> { let mut attr = unsafe { mem::zeroed::() }; let u = unsafe { &mut attr.__bindgen_anon_6 }; u.__bindgen_anon_1.start_id = id; match sys_bpf(bpf_cmd::BPF_PROG_GET_NEXT_ID, &mut attr) { - Ok(_) => Ok(Some(unsafe { attr.__bindgen_anon_6.next_id })), - Err((_, io_error)) if io_error.raw_os_error() == Some(ENOENT) => Ok(None), - Err(e) => Err(e), + Ok(code) => { + assert_eq!(code, 0); + Ok(Some(unsafe { attr.__bindgen_anon_6.next_id })) + } + Err((code, io_error)) => { + assert_eq!(code, -1); + if io_error.raw_os_error() == Some(ENOENT) { + Ok(None) + } else { + Err(SyscallError { + call: "bpf_prog_get_next_id", + io_error, + }) + } + } } } +pub(crate) fn iter_prog_ids() -> impl Iterator> { + let mut current_id = Some(0); + iter::from_fn(move || { + let next_id = { + let current_id = current_id?; + bpf_prog_get_next_id(current_id).transpose() + }; + current_id = next_id.as_ref().and_then(|next_id| match next_id { + Ok(next_id) => Some(*next_id), + Err(SyscallError { .. }) => None, + }); + next_id + }) +} + pub(crate) fn retry_with_verifier_logs( max_retries: usize, f: impl Fn(&mut [u8]) -> SysResult,