feat(aya): Add SchedClassifier:query_tcx() API

This adds SchedClassifier::query_tcx() API that is able to return an ordered list of ProgramInfo (and revision) of an interface that has TCX programs attached. This is useful to verify the ordering of the programs in our integration tests, but also to for aya-rs#1032. Additionally tidies up the macro used for TCX testing and adds assertions using this new API. Signed-off-by: Dave Tucker <[email protected]>
astoycos · Sep 19, 2024 · c5bc4ba · c5bc4ba
1 parent f5d4fde
commit c5bc4ba
Show file tree

Hide file tree

Showing 7 changed files with 164 additions and 120 deletions.
diff --git a/aya/src/programs/cgroup_device.rs b/aya/src/programs/cgroup_device.rs
@@ -8,7 +8,7 @@ use crate::{
         bpf_prog_get_fd_by_id, define_link_wrapper, load_program, query, CgroupAttachMode, FdLink,
         Link, ProgAttachLink, ProgramData, ProgramError, ProgramFd,
     },
-    sys::{bpf_link_create, LinkTarget, SyscallError},
+    sys::{bpf_link_create, LinkTarget, ProgQueryTarget, SyscallError},
     util::KernelVersion,
 };
 
@@ -120,7 +120,12 @@ impl CgroupDevice {
     /// Queries the cgroup for attached programs.
     pub fn query<T: AsFd>(target_fd: T) -> Result<Vec<CgroupDeviceLink>, ProgramError> {
         let target_fd = target_fd.as_fd();
-        let prog_ids = query(target_fd, BPF_CGROUP_DEVICE, 0, &mut None)?;
+        let (_, prog_ids) = query(
+            ProgQueryTarget::Fd(target_fd),
+            BPF_CGROUP_DEVICE,
+            0,
+            &mut None,
+        )?;
 
         prog_ids
             .into_iter()

diff --git a/aya/src/programs/lirc_mode2.rs b/aya/src/programs/lirc_mode2.rs
@@ -7,7 +7,7 @@ use crate::{
         load_program, query, CgroupAttachMode, Link, ProgramData, ProgramError, ProgramFd,
         ProgramInfo,
     },
-    sys::{bpf_prog_attach, bpf_prog_detach, bpf_prog_get_fd_by_id},
+    sys::{bpf_prog_attach, bpf_prog_detach, bpf_prog_get_fd_by_id, ProgQueryTarget},
 };
 
 /// A program used to decode IR into key events for a lirc device.
@@ -100,7 +100,7 @@ impl LircMode2 {
     /// Queries the lirc device for attached programs.
     pub fn query<T: AsFd>(target_fd: T) -> Result<Vec<LircLink>, ProgramError> {
         let target_fd = target_fd.as_fd();
-        let prog_ids = query(target_fd, BPF_LIRC_MODE2, 0, &mut None)?;
+        let (_, prog_ids) = query(ProgQueryTarget::Fd(target_fd), BPF_LIRC_MODE2, 0, &mut None)?;
 
         prog_ids
             .into_iter()

diff --git a/aya/src/programs/mod.rs b/aya/src/programs/mod.rs
@@ -73,7 +73,7 @@ use std::{
     ffi::CString,
     io,
     num::TryFromIntError,
-    os::fd::{AsFd, AsRawFd, BorrowedFd},
+    os::fd::{AsFd, BorrowedFd},
     path::{Path, PathBuf},
     sync::Arc,
 };
@@ -122,7 +122,7 @@ use crate::{
     sys::{
         bpf_btf_get_fd_by_id, bpf_get_object, bpf_link_get_fd_by_id, bpf_link_get_info_by_fd,
         bpf_load_program, bpf_pin_object, bpf_prog_get_fd_by_id, bpf_prog_query, iter_link_ids,
-        retry_with_verifier_logs, EbpfLoadProgramAttrs, SyscallError,
+        retry_with_verifier_logs, EbpfLoadProgramAttrs, ProgQueryTarget, SyscallError,
     },
     util::KernelVersion,
     VerifierLogLevel,
@@ -688,28 +688,30 @@ fn load_program<T: Link>(
 }
 
 pub(crate) fn query(
-    target_fd: BorrowedFd<'_>,
+    target: ProgQueryTarget<'_>,
     attach_type: bpf_attach_type,
     query_flags: u32,
     attach_flags: &mut Option<u32>,
-) -> Result<Vec<u32>, ProgramError> {
+) -> Result<(u64, Vec<u32>), ProgramError> {
     let mut prog_ids = vec![0u32; 64];
     let mut prog_cnt = prog_ids.len() as u32;
+    let mut revision = 0;
 
     let mut retries = 0;
 
     loop {
         match bpf_prog_query(
-            target_fd.as_fd().as_raw_fd(),
+            &target,
             attach_type,
             query_flags,
             attach_flags.as_mut(),
             &mut prog_ids,
             &mut prog_cnt,
+            &mut revision,
         ) {
             Ok(_) => {
                 prog_ids.resize(prog_cnt as usize, 0);
-                return Ok(prog_ids);
+                return Ok((revision, prog_ids));
             }
             Err((_, io_error)) => {
                 if retries == 0 && io_error.raw_os_error() == Some(ENOSPC) {

diff --git a/aya/src/programs/tc.rs b/aya/src/programs/tc.rs
@@ -8,7 +8,7 @@ use std::{
 
 use thiserror::Error;
 
-use super::FdLink;
+use super::{FdLink, ProgramInfo};
 use crate::{
     generated::{
         bpf_attach_type::{self, BPF_TCX_EGRESS, BPF_TCX_INGRESS},
@@ -17,12 +17,13 @@ use crate::{
         TC_H_CLSACT, TC_H_MIN_EGRESS, TC_H_MIN_INGRESS,
     },
     programs::{
-        define_link_wrapper, load_program, Link, LinkError, LinkOrder, ProgramData, ProgramError,
+        define_link_wrapper, load_program, query, Link, LinkError, LinkOrder, ProgramData,
+        ProgramError,
     },
     sys::{
-        bpf_link_create, bpf_link_get_info_by_fd, bpf_link_update, netlink_find_filter_with_name,
-        netlink_qdisc_add_clsact, netlink_qdisc_attach, netlink_qdisc_detach, LinkTarget,
-        SyscallError,
+        bpf_link_create, bpf_link_get_info_by_fd, bpf_link_update, bpf_prog_get_fd_by_id,
+        netlink_find_filter_with_name, netlink_qdisc_add_clsact, netlink_qdisc_attach,
+        netlink_qdisc_detach, LinkTarget, ProgQueryTarget, SyscallError,
     },
     util::{ifindex_from_ifname, tc_handler_make, KernelVersion},
     VerifierLogLevel,
@@ -340,6 +341,46 @@ impl SchedClassifier {
         let data = ProgramData::from_pinned_path(path, VerifierLogLevel::default())?;
         Ok(Self { data })
     }
+
+    /// Queries a given interface for attached TCX programs.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// # use aya::programs::tc::{TcAttachType, SchedClassifier};
+    /// # #[derive(Debug, thiserror::Error)]
+    /// # enum Error {
+    /// #     #[error(transparent)]
+    /// #     Program(#[from] aya::programs::ProgramError),
+    /// # }
+    /// let (revision, programs) = SchedClassifier::query_tcx("eth0", TcAttachType::Ingress)?;
+    /// # Ok::<(), Error>(())
+    /// ```
+    pub fn query_tcx(
+        interface: &str,
+        attach_type: TcAttachType,
+    ) -> Result<(u64, Vec<ProgramInfo>), ProgramError> {
+        let if_index = ifindex_from_ifname(interface)
+            .map_err(|io_error| TcError::NetlinkError { io_error })?;
+
+        let (revision, prog_ids) = query(
+            ProgQueryTarget::IfIndex(if_index),
+            attach_type.tcx_attach_type()?,
+            0,
+            &mut None,
+        )?;
+
+        let prog_infos = prog_ids
+            .into_iter()
+            .map(|prog_id| {
+                let prog_fd = bpf_prog_get_fd_by_id(prog_id)?;
+                let prog_info = ProgramInfo::new_from_fd(prog_fd.as_fd())?;
+                Ok::<ProgramInfo, ProgramError>(prog_info)
+            })
+            .collect::<Result<_, _>>()?;
+
+        Ok((revision, prog_infos))
+    }
 }
 
 #[derive(Debug, Hash, Eq, PartialEq)]

diff --git a/aya/src/sys/bpf.rs b/aya/src/sys/bpf.rs
@@ -498,25 +498,39 @@ pub(crate) fn bpf_prog_detach(
     Ok(())
 }
 
+#[derive(Debug)]
+pub(crate) enum ProgQueryTarget<'a> {
+    Fd(BorrowedFd<'a>),
+    IfIndex(u32),
+}
+
 pub(crate) fn bpf_prog_query(
-    target_fd: RawFd,
+    target: &ProgQueryTarget<'_>,
     attach_type: bpf_attach_type,
     query_flags: u32,
     attach_flags: Option<&mut u32>,
     prog_ids: &mut [u32],
     prog_cnt: &mut u32,
+    revision: &mut u64,
 ) -> SysResult<i64> {
     let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
 
-    attr.query.__bindgen_anon_1.target_fd = target_fd as u32;
+    match target {
+        ProgQueryTarget::Fd(fd) => {
+            attr.query.__bindgen_anon_1.target_fd = fd.as_raw_fd() as u32;
+        }
+        ProgQueryTarget::IfIndex(ifindex) => {
+            attr.query.__bindgen_anon_1.target_ifindex = *ifindex;
+        }
+    }
     attr.query.attach_type = attach_type as u32;
     attr.query.query_flags = query_flags;
     attr.query.__bindgen_anon_2.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, &mut attr);
 
     *prog_cnt = unsafe { attr.query.__bindgen_anon_2.prog_cnt };
+    *revision = unsafe { attr.query.revision };
 
     if let Some(attach_flags) = attach_flags {
         *attach_flags = unsafe { attr.query.attach_flags };

diff --git a/macro.patch b/macro.patch