Implement query for lirc programs (#32)

Signed-off-by: Sean Young <sean@mess.org>

aya/src/programs/lirc_mode2.rs

@@ -2,8 +2,8 @@ use std::os::unix::prelude::{AsRawFd, RawFd};
 
 use crate::{
     generated::{bpf_attach_type::BPF_LIRC_MODE2, bpf_prog_type::BPF_PROG_TYPE_LIRC_MODE2},
-    programs::{load_program, Link, LinkRef, ProgramData, ProgramError},
+    programs::{load_program, query, Link, LinkRef, ProgramData, ProgramError, ProgramInfo},
-    sys::{bpf_prog_attach, bpf_prog_detach},
+    sys::{bpf_obj_get_info_by_fd, bpf_prog_attach, bpf_prog_detach, bpf_prog_get_fd_by_id},
 };
 
 use libc::{close, dup};
@@ -78,10 +78,34 @@ impl LircMode2 {
 
         Ok(self.data.link(LircLink::new(prog_fd, lircdev_fd)))
     }
+
+    /// Query lirc device for attached programs
+    pub fn query<T: AsRawFd>(target_fd: T) -> Result<Vec<LircLink>, ProgramError> {
+        let prog_ids = query(target_fd.as_raw_fd(), BPF_LIRC_MODE2, 0, &mut None)?;
+
+        let mut prog_fds = Vec::with_capacity(prog_ids.len());
+
+        for id in prog_ids {
+            let fd = bpf_prog_get_fd_by_id(id).map_err(|io_error| ProgramError::SyscallError {
+                call: "bpf_prog_get_fd_by_id".to_owned(),
+                io_error,
+            })?;
+
+            prog_fds.push(fd as RawFd);
+        }
+
+        Ok(prog_fds
+            .into_iter()
+            .map(|prog_fd| LircLink {
+                prog_fd: Some(prog_fd),
+                target_fd: Some(unsafe { dup(target_fd.as_raw_fd()) }),
+            })
+            .collect())
+    }
 }
 
 #[derive(Debug)]
-struct LircLink {
+pub struct LircLink {
     prog_fd: Option<RawFd>,
     target_fd: Option<RawFd>,
 }
@@ -93,6 +117,20 @@ impl LircLink {
             target_fd: Some(unsafe { dup(target_fd) }),
         }
     }
+
+    pub fn info(&self) -> Result<ProgramInfo, ProgramError> {
+        if let Some(fd) = self.prog_fd {
+            match bpf_obj_get_info_by_fd(fd) {
+                Ok(info) => Ok(ProgramInfo(info)),
+                Err(io_error) => Err(ProgramError::SyscallError {
+                    call: "bpf_obj_get_info_by_fd".to_owned(),
+                    io_error,
+                }),
+            }
+        } else {
+            Err(ProgramError::AlreadyDetached)
+        }
+    }
 }
 
 impl Link for LircLink {

aya/src/programs/mod.rs

@@ -51,7 +51,15 @@ mod uprobe;
 mod xdp;
 
 use libc::{close, dup, ENOSPC};
-use std::{cell::RefCell, cmp, convert::TryFrom, ffi::CStr, io, os::unix::io::RawFd, rc::Rc};
+use std::{
+    cell::RefCell,
+    cmp,
+    convert::TryFrom,
+    ffi::CStr,
+    io,
+    os::unix::io::{AsRawFd, RawFd},
+    rc::Rc,
+};
 use thiserror::Error;
 
 pub use cgroup_skb::{CgroupSkb, CgroupSkbAttachType};
@@ -69,10 +77,10 @@ pub use uprobe::{UProbe, UProbeError};
 pub use xdp::{Xdp, XdpError, XdpFlags};
 
 use crate::{
-    generated::{bpf_attach_type, bpf_prog_type},
+    generated::{bpf_attach_type, bpf_prog_info, bpf_prog_type},
     maps::MapError,
     obj::{self, Function},
-    sys::{bpf_load_program, bpf_prog_detach},
+    sys::{bpf_load_program, bpf_prog_detach, bpf_prog_query},
 };
 
 /// Error type returned when working with programs.
@@ -373,6 +381,44 @@ fn load_program(prog_type: bpf_prog_type, data: &mut ProgramData) -> Result<(),
     Ok(())
 }
 
+pub(crate) fn query<T: AsRawFd>(
+    target_fd: T,
+    attach_type: bpf_attach_type,
+    query_flags: u32,
+    attach_flags: &mut Option<u32>,
+) -> Result<Vec<u32>, ProgramError> {
+    let mut prog_ids = vec![0u32; 64];
+    let mut prog_cnt = prog_ids.len() as u32;
+
+    let mut retries = 0;
+
+    loop {
+        match bpf_prog_query(
+            target_fd.as_raw_fd(),
+            attach_type,
+            query_flags,
+            attach_flags.as_mut(),
+            &mut prog_ids,
+            &mut prog_cnt,
+        ) {
+            Ok(_) => {
+                prog_ids.resize(prog_cnt as usize, 0);
+                return Ok(prog_ids);
+            }
+            Err((_, io_error)) if retries == 0 && io_error.raw_os_error() == Some(ENOSPC) => {
+                prog_ids.resize(prog_cnt as usize, 0);
+                retries += 1;
+            }
+            Err((_, io_error)) => {
+                return Err(ProgramError::SyscallError {
+                    call: "bpf_prog_query".to_owned(),
+                    io_error,
+                });
+            }
+        }
+    }
+}
+
 /// Detach an attached program.
 pub trait Link: std::fmt::Debug {
     fn detach(&mut self) -> Result<(), ProgramError>;
@@ -533,3 +579,33 @@ impl_try_from_program!(
     CgroupSkb,
     LircMode2
 );
+
+/// Provides information about a loaded program, like name, id and statistics
+pub struct ProgramInfo(bpf_prog_info);
+
+impl ProgramInfo {
+    /// The name of the program as was provided when it was load. This is limited to 16 bytes
+    pub fn name(&self) -> &[u8] {
+        let length = self
+            .0
+            .name
+            .iter()
+            .rposition(|ch| *ch != 0)
+            .map(|pos| pos + 1)
+            .unwrap_or(0);
+
+        // The name field is defined as [std::os::raw::c_char; 16]. c_char may be signed or
+        // unsigned depending on the platform; that's why we're using from_raw_parts here
+        unsafe { std::slice::from_raw_parts(self.0.name.as_ptr() as *const _, length) }
+    }
+
+    /// The name of the program as a &str. If the name was not valid unicode, None is returned
+    pub fn name_as_str(&self) -> Option<&str> {
+        std::str::from_utf8(self.name()).ok()
+    }
+
+    /// The program id for this program. Each program has a unique id.
+    pub fn id(&self) -> u32 {
+        self.0.id
+    }
+}

aya/src/sys/bpf.rs

@@ -11,7 +11,7 @@ use libc::{c_long, ENOENT};
 
 use crate::{
     bpf_map_def,
-    generated::{bpf_attach_type, bpf_attr, bpf_cmd, bpf_insn, bpf_prog_type},
+    generated::{bpf_attach_type, bpf_attr, bpf_cmd, bpf_insn, bpf_prog_info, bpf_prog_type},
     maps::PerCpuValues,
     programs::VerifierLog,
     sys::SysResult,
@@ -264,6 +264,59 @@ pub(crate) fn bpf_prog_detach(
     sys_bpf(bpf_cmd::BPF_PROG_DETACH, &attr)
 }
 
+pub(crate) fn bpf_prog_query(
+    target_fd: RawFd,
+    attach_type: bpf_attach_type,
+    query_flags: u32,
+    attach_flags: Option<&mut u32>,
+    prog_ids: &mut [u32],
+    prog_cnt: &mut u32,
+) -> SysResult {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+
+    attr.query.target_fd = target_fd as u32;
+    attr.query.attach_type = attach_type as u32;
+    attr.query.query_flags = query_flags;
+    attr.query.prog_cnt = prog_ids.len() as u32;
+    attr.query.prog_ids = prog_ids.as_mut_ptr() as u64;
+
+    let ret = sys_bpf(bpf_cmd::BPF_PROG_QUERY, &attr);
+
+    *prog_cnt = unsafe { attr.query.prog_cnt };
+
+    if let Some(attach_flags) = attach_flags {
+        *attach_flags = unsafe { attr.query.attach_flags };
+    }
+
+    ret
+}
+
+pub(crate) fn bpf_prog_get_fd_by_id(prog_id: u32) -> Result<RawFd, io::Error> {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+
+    attr.__bindgen_anon_6.__bindgen_anon_1.prog_id = prog_id;
+
+    match sys_bpf(bpf_cmd::BPF_PROG_GET_FD_BY_ID, &attr) {
+        Ok(v) => Ok(v as RawFd),
+        Err((_, err)) => Err(err),
+    }
+}
+
+pub(crate) fn bpf_obj_get_info_by_fd(prog_fd: RawFd) -> Result<bpf_prog_info, io::Error> {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+    // info gets entirely populated by the kernel
+    let info = unsafe { MaybeUninit::zeroed().assume_init() };
+
+    attr.info.bpf_fd = prog_fd as u32;
+    attr.info.info = &info as *const _ as u64;
+    attr.info.info_len = mem::size_of::<bpf_prog_info>() as u32;
+
+    match sys_bpf(bpf_cmd::BPF_OBJ_GET_INFO_BY_FD, &attr) {
+        Ok(_) => Ok(info),
+        Err((_, err)) => Err(err),
+    }
+}
+
 fn sys_bpf(cmd: bpf_cmd, attr: &bpf_attr) -> SysResult {
     syscall(Syscall::Bpf { cmd, attr })
 }