Merge pull request #282 from dave-tucker/bpfd

Add atomic replacement of XDP progs and multihoming of Extension progs

aya/src/maps/array/program_array.rs

@@ -4,7 +4,7 @@ use std::{
     convert::TryFrom,
     mem,
     ops::{Deref, DerefMut},
-    os::unix::prelude::RawFd,
+    os::unix::prelude::{AsRawFd, RawFd},
 };
 
 use crate::{
@@ -26,26 +26,28 @@ use crate::{
 ///
 /// # Examples
 /// ```no_run
-/// # let bpf = aya::Bpf::load(&[])?;
+/// # let mut bpf = aya::Bpf::load(&[])?;
 /// use aya::maps::ProgramArray;
-/// use aya::programs::{CgroupSkb, ProgramFd};
+/// use aya::programs::CgroupSkb;
 /// use std::convert::{TryFrom, TryInto};
 ///
 /// let mut prog_array = ProgramArray::try_from(bpf.map_mut("JUMP_TABLE")?)?;
 /// let prog_0: &CgroupSkb = bpf.program("example_prog_0").unwrap().try_into()?;
+/// let prog_0_fd =  prog_0.fd().unwrap();
 /// let prog_1: &CgroupSkb = bpf.program("example_prog_1").unwrap().try_into()?;
+/// let prog_1_fd = prog_1.fd().unwrap();
 /// let prog_2: &CgroupSkb = bpf.program("example_prog_2").unwrap().try_into()?;
-///
+/// let prog_2_fd = prog_2.fd().unwrap();
 /// let flags = 0;
 ///
 /// // bpf_tail_call(ctx, JUMP_TABLE, 0) will jump to prog_0
-/// prog_array.set(0, prog_0, flags);
+/// prog_array.set(0, prog_0_fd, flags);
 ///
 /// // bpf_tail_call(ctx, JUMP_TABLE, 1) will jump to prog_1
-/// prog_array.set(1, prog_1, flags);
+/// prog_array.set(1, prog_1_fd, flags);
 ///
 /// // bpf_tail_call(ctx, JUMP_TABLE, 2) will jump to prog_2
-/// prog_array.set(2, prog_2, flags);
+/// prog_array.set(2, prog_2_fd, flags);
 /// # Ok::<(), aya::BpfError>(())
 /// ```
 #[doc(alias = "BPF_MAP_TYPE_PROG_ARRAY")]
@@ -98,10 +100,10 @@ impl<T: Deref<Target = Map> + DerefMut<Target = Map>> ProgramArray<T> {
     ///
     /// When an eBPF program calls `bpf_tail_call(ctx, prog_array, index)`, control
     /// flow will jump to `program`.
-    pub fn set(&mut self, index: u32, program: impl ProgramFd, flags: u64) -> Result<(), MapError> {
+    pub fn set(&mut self, index: u32, program: ProgramFd, flags: u64) -> Result<(), MapError> {
         let fd = self.inner.fd_or_err()?;
         self.check_bounds(index)?;
-        let prog_fd = program.fd().ok_or(MapError::ProgramNotLoaded)?;
+        let prog_fd = program.as_raw_fd();
 
         bpf_map_update_elem(fd, &index, &prog_fd, flags).map_err(|(code, io_error)| {
             MapError::SyscallError {

aya/src/programs/extension.rs

@@ -9,20 +9,21 @@ use crate::{
     obj::btf::BtfKind,
     programs::{
         define_link_wrapper, load_program, FdLink, FdLinkId, OwnedLink, ProgramData, ProgramError,
+        ProgramFd,
     },
     sys::{self, bpf_link_create},
     Btf,
 };
 
-/// The type returned when loading or attaching an [`Extension`] fails
+/// The type returned when loading or attaching an [`Extension`] fails.
 #[derive(Debug, Error)]
 pub enum ExtensionError {
-    /// target BPF program does not have BTF loaded to the kernel
+    /// Target BPF program does not have BTF loaded to the kernel.
     #[error("target BPF program does not have BTF loaded to the kernel")]
     NoBTF,
 }
 
-/// A program used to extend existing BPF programs
+/// A program used to extend existing BPF programs.
 ///
 /// [`Extension`] programs can be loaded to replace a global
 /// function in a program that has already been loaded.
@@ -34,7 +35,7 @@ pub enum ExtensionError {
 /// # Examples
 ///
 /// ```no_run
-/// use aya::{BpfLoader, programs::{Xdp, XdpFlags, Extension, ProgramFd}};
+/// use aya::{BpfLoader, programs::{Xdp, XdpFlags, Extension}};
 /// use std::convert::TryInto;
 ///
 /// let mut bpf = BpfLoader::new().extension("extension").load_file("app.o")?;
@@ -60,62 +61,18 @@ impl Extension {
     /// Prepares the code included in the extension to replace the code of the function
     /// `func_name` within the eBPF program represented by the `program` file descriptor.
     /// This requires that both the [`Extension`] and `program` have had their BTF
-    /// loaded into the kernel as the verifier must check that the function signatures
-    /// match.
+    /// loaded into the kernel.
+    ///
+    /// The BPF verifier requires that we specify the target program and function name
+    /// at load time, so it can identify that the program and target are BTF compatible
+    /// and to enforce this constraint when programs are attached.
     ///
     /// The extension code will be loaded but inactive until it's attached.
     /// There are no restrictions on what functions may be replaced, so you could replace
     /// the main entry point of your program with an extension.
-    pub fn load<T: AsRawFd>(&mut self, program: T, func_name: &str) -> Result<(), ProgramError> {
+    pub fn load(&mut self, program: ProgramFd, func_name: &str) -> Result<(), ProgramError> {
         let target_prog_fd = program.as_raw_fd();
-
-        let info = sys::bpf_obj_get_info_by_fd(target_prog_fd).map_err(|io_error| {
-            ProgramError::SyscallError {
-                call: "bpf_obj_get_info_by_fd".to_owned(),
-                io_error,
-            }
-        })?;
-
-        if info.btf_id == 0 {
-            return Err(ProgramError::ExtensionError(ExtensionError::NoBTF));
-        }
-
-        let btf_fd = sys::bpf_btf_get_fd_by_id(info.btf_id).map_err(|io_error| {
-            ProgramError::SyscallError {
-                call: "bpf_btf_get_fd_by_id".to_owned(),
-                io_error,
-            }
-        })?;
-
-        let mut buf = vec![0u8; 4096];
-        let btf_info = match sys::btf_obj_get_info_by_fd(btf_fd, &mut buf) {
-            Ok(info) => {
-                if info.btf_size > buf.len() as u32 {
-                    buf.resize(info.btf_size as usize, 0u8);
-                    let btf_info =
-                        sys::btf_obj_get_info_by_fd(btf_fd, &mut buf).map_err(|io_error| {
-                            ProgramError::SyscallError {
-                                call: "bpf_obj_get_info_by_fd".to_owned(),
-                                io_error,
-                            }
-                        })?;
-                    Ok(btf_info)
-                } else {
-                    Ok(info)
-                }
-            }
-            Err(io_error) => Err(ProgramError::SyscallError {
-                call: "bpf_obj_get_info_by_fd".to_owned(),
-                io_error,
-            }),
-        }?;
-
-        let btf = Btf::parse(&buf[0..btf_info.btf_size as usize], Endianness::default())
-            .map_err(ProgramError::Btf)?;
-
-        let btf_id = btf
-            .id_by_type_name_kind(func_name, BtfKind::Func)
-            .map_err(ProgramError::Btf)?;
+        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_prog_fd = Some(target_prog_fd);
@@ -125,7 +82,8 @@ impl Extension {
 
     /// Attaches the extension.
     ///
-    /// Attaches the extension effectively replacing the original target function.
+    /// Attaches the extension to the program and function name specified at load time,
+    /// effectively replacing the original target function.
     ///
     /// The returned value can be used to detach the extension and restore the
     /// original function, see [Extension::detach].
@@ -142,6 +100,34 @@ impl Extension {
         self.data.links.insert(ExtensionLink(FdLink::new(link_fd)))
     }
 
+    /// Attaches the extension to another program.
+    ///
+    /// Attaches the extension to a program and/or function other than the one provided
+    /// at load time. You may only attach to another program/function if the BTF
+    /// type signature is identical to that which was verified on load. Attempting to
+    /// attach to an invalid program/function will result in an error.
+    ///
+    /// Once attached, the extension effectively replaces the original target function.
+    ///
+    /// The returned value can be used to detach the extension and restore the
+    /// original function, see [Extension::detach].
+    pub fn attach_to_program(
+        &mut self,
+        program: ProgramFd,
+        func_name: &str,
+    ) -> Result<ExtensionLinkId, ProgramError> {
+        let target_fd = program.as_raw_fd();
+        let (_, btf_id) = get_btf_info(target_fd, func_name)?;
+        let prog_fd = self.data.fd_or_err()?;
+        // the attach type must be set as 0, which is bpf_attach_type::BPF_CGROUP_INET_INGRESS
+        let link_fd = bpf_link_create(prog_fd, target_fd, BPF_CGROUP_INET_INGRESS, Some(btf_id), 0)
+            .map_err(|(_, io_error)| ProgramError::SyscallError {
+                call: "bpf_link_create".to_owned(),
+                io_error,
+            })? as RawFd;
+        self.data.links.insert(ExtensionLink(FdLink::new(link_fd)))
+    }
+
     /// Detaches the extension.
     ///
     /// Detaching restores the original code overridden by the extension program.
@@ -162,6 +148,63 @@ 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> {
+    // retrieve program information
+    let info =
+        sys::bpf_obj_get_info_by_fd(prog_fd).map_err(|io_error| ProgramError::SyscallError {
+            call: "bpf_obj_get_info_by_fd".to_owned(),
+            io_error,
+        })?;
+
+    // btf_id refers to the ID of the program btf that was loaded with bpf(BPF_BTF_LOAD)
+    if info.btf_id == 0 {
+        return Err(ProgramError::ExtensionError(ExtensionError::NoBTF));
+    }
+
+    // the bpf fd of the BTF object
+    let btf_fd =
+        sys::bpf_btf_get_fd_by_id(info.btf_id).map_err(|io_error| ProgramError::SyscallError {
+            call: "bpf_btf_get_fd_by_id".to_owned(),
+            io_error,
+        })?;
+
+    // we need to read the btf bytes into a buffer but we don't know the size ahead of time.
+    // assume 4kb. if this is too small we can resize based on the size obtained in the response.
+    let mut buf = vec![0u8; 4096];
+    let btf_info = match sys::btf_obj_get_info_by_fd(btf_fd, &mut buf) {
+        Ok(info) => {
+            if info.btf_size > buf.len() as u32 {
+                buf.resize(info.btf_size as usize, 0u8);
+                let btf_info =
+                    sys::btf_obj_get_info_by_fd(btf_fd, &mut buf).map_err(|io_error| {
+                        ProgramError::SyscallError {
+                            call: "bpf_obj_get_info_by_fd".to_owned(),
+                            io_error,
+                        }
+                    })?;
+                Ok(btf_info)
+            } else {
+                Ok(info)
+            }
+        }
+        Err(io_error) => Err(ProgramError::SyscallError {
+            call: "bpf_obj_get_info_by_fd".to_owned(),
+            io_error,
+        }),
+    }?;
+
+    let btf = Btf::parse(&buf[0..btf_info.btf_size as usize], Endianness::default())
+        .map_err(ProgramError::Btf)?;
+
+    let btf_id = btf
+        .id_by_type_name_kind(func_name, BtfKind::Func)
+        .map_err(ProgramError::Btf)?;
+
+    Ok((btf_fd, btf_id))
+}
+
 define_link_wrapper!(
     /// The link used by [Extension] programs.
     ExtensionLink,

aya/src/programs/mod.rs

@@ -214,10 +214,14 @@ pub enum ProgramError {
     },
 }
 
-/// Allows the Fd of a loaded [`Program`] to be retrieved
-pub trait ProgramFd {
-    /// Returns the [`RawFd`] of the program if it has been loaded, or `None`
-    fn fd(&self) -> Option<RawFd>;
+/// A [`Program`] file descriptor.
+#[derive(Copy, Clone)]
+pub struct ProgramFd(RawFd);
+
+impl AsRawFd for ProgramFd {
+    fn as_raw_fd(&self) -> RawFd {
+        self.0
+    }
 }
 
 /// eBPF program type.
@@ -359,6 +363,38 @@ impl Program {
             Program::CgroupSock(p) => p.unload(),
         }
     }
+
+    /// Returns the file descriptor of a program.
+    ///
+    /// Can be used to add a program to a [`crate::maps::ProgramArray`] or attach an [`Extension`] program.
+    /// Can be converted to [`RawFd`] using [`AsRawFd`].
+    pub fn fd(&self) -> Option<ProgramFd> {
+        match self {
+            Program::KProbe(p) => p.fd(),
+            Program::UProbe(p) => p.fd(),
+            Program::TracePoint(p) => p.fd(),
+            Program::SocketFilter(p) => p.fd(),
+            Program::Xdp(p) => p.fd(),
+            Program::SkMsg(p) => p.fd(),
+            Program::SkSkb(p) => p.fd(),
+            Program::SockOps(p) => p.fd(),
+            Program::SchedClassifier(p) => p.fd(),
+            Program::CgroupSkb(p) => p.fd(),
+            Program::CgroupSysctl(p) => p.fd(),
+            Program::CgroupSockopt(p) => p.fd(),
+            Program::LircMode2(p) => p.fd(),
+            Program::PerfEvent(p) => p.fd(),
+            Program::RawTracePoint(p) => p.fd(),
+            Program::Lsm(p) => p.fd(),
+            Program::BtfTracePoint(p) => p.fd(),
+            Program::FEntry(p) => p.fd(),
+            Program::FExit(p) => p.fd(),
+            Program::Extension(p) => p.fd(),
+            Program::CgroupSockAddr(p) => p.fd(),
+            Program::SkLookup(p) => p.fd(),
+            Program::CgroupSock(p) => p.fd(),
+        }
+    }
 }
 
 impl Drop for Program {
@@ -555,42 +591,6 @@ pub(crate) fn query<T: AsRawFd>(
     }
 }
 
-impl ProgramFd for Program {
-    fn fd(&self) -> Option<RawFd> {
-        match self {
-            Program::KProbe(p) => p.data.fd,
-            Program::UProbe(p) => p.data.fd,
-            Program::TracePoint(p) => p.data.fd,
-            Program::SocketFilter(p) => p.data.fd,
-            Program::Xdp(p) => p.data.fd,
-            Program::SkMsg(p) => p.data.fd,
-            Program::SkSkb(p) => p.data.fd,
-            Program::SockOps(p) => p.data.fd,
-            Program::SchedClassifier(p) => p.data.fd,
-            Program::CgroupSkb(p) => p.data.fd,
-            Program::CgroupSysctl(p) => p.data.fd,
-            Program::CgroupSockopt(p) => p.data.fd,
-            Program::LircMode2(p) => p.data.fd,
-            Program::PerfEvent(p) => p.data.fd,
-            Program::RawTracePoint(p) => p.data.fd,
-            Program::Lsm(p) => p.data.fd,
-            Program::BtfTracePoint(p) => p.data.fd,
-            Program::FEntry(p) => p.data.fd,
-            Program::FExit(p) => p.data.fd,
-            Program::Extension(p) => p.data.fd,
-            Program::CgroupSockAddr(p) => p.data.fd,
-            Program::SkLookup(p) => p.data.fd,
-            Program::CgroupSock(p) => p.data.fd,
-        }
-    }
-}
-
-impl<'a, P: ProgramFd> ProgramFd for &'a P {
-    fn fd(&self) -> Option<RawFd> {
-        (*self).fd()
-    }
-}
-
 macro_rules! impl_program_unload {
     ($($struct_name:ident),+ $(,)?) => {
         $(
@@ -634,25 +634,20 @@ impl_program_unload!(
     CgroupSock,
 );
 
-macro_rules! impl_program_fd {
+macro_rules! impl_fd {
     ($($struct_name:ident),+ $(,)?) => {
         $(
-            impl ProgramFd for $struct_name {
-                fn fd(&self) -> Option<RawFd> {
-                    self.data.fd
-                }
-            }
-
-            impl ProgramFd for &mut $struct_name {
-                fn fd(&self) -> Option<RawFd> {
-                    self.data.fd
+            impl $struct_name {
+                /// Returns the file descriptor of this Program.
+                pub fn fd(&self) -> Option<ProgramFd> {
+                    self.data.fd.map(|fd| ProgramFd(fd))
                 }
             }
         )+
     }
 }
 
-impl_program_fd!(
+impl_fd!(
     KProbe,
     UProbe,
     TracePoint,
@@ -674,6 +669,7 @@ impl_program_fd!(
     Extension,
     CgroupSockAddr,
     SkLookup,
+    SockOps,
     CgroupSock,
 );
 

aya/src/programs/xdp.rs

@@ -1,7 +1,7 @@
 //! eXpress Data Path (XDP) programs.
 use bitflags;
 use libc::if_nametoindex;
-use std::{ffi::CString, hash::Hash, io, os::unix::io::RawFd};
+use std::{ffi::CString, hash::Hash, io, mem, os::unix::io::RawFd};
 use thiserror::Error;
 
 use crate::{
@@ -14,7 +14,7 @@ use crate::{
     programs::{
         define_link_wrapper, load_program, FdLink, Link, OwnedLink, ProgramData, ProgramError,
     },
-    sys::{bpf_link_create, kernel_version, netlink_set_xdp_fd},
+    sys::{bpf_link_create, bpf_link_update, kernel_version, netlink_set_xdp_fd},
 };
 
 /// The type returned when attaching an [`Xdp`] program fails on kernels `< 5.9`.
@@ -141,6 +141,46 @@ impl Xdp {
     pub fn take_link(&mut self, link_id: XdpLinkId) -> Result<OwnedLink<XdpLink>, ProgramError> {
         Ok(OwnedLink::new(self.data.take_link(link_id)?))
     }
+
+    /// Atomically replaces the program referenced by the provided link.
+    ///
+    /// Ownership of the link will transfer to this program.
+    pub fn attach_to_link(&mut self, link: OwnedLink<XdpLink>) -> Result<XdpLinkId, ProgramError> {
+        let prog_fd = self.data.fd_or_err()?;
+        match &link.0 {
+            XdpLinkInner::FdLink(fd_link) => {
+                let link_fd = fd_link.fd;
+                bpf_link_update(link_fd, prog_fd, None, 0).map_err(|(_, io_error)| {
+                    ProgramError::SyscallError {
+                        call: "bpf_link_update".to_string(),
+                        io_error,
+                    }
+                })?;
+                // dispose of link and avoid detach on drop
+                mem::forget(link);
+                self.data
+                    .links
+                    .insert(XdpLink(XdpLinkInner::FdLink(FdLink::new(link_fd))))
+            }
+            XdpLinkInner::NlLink(nl_link) => {
+                let if_index = nl_link.if_index;
+                let old_prog_fd = nl_link.prog_fd;
+                let flags = nl_link.flags;
+                let replace_flags = flags | XdpFlags::REPLACE;
+                unsafe {
+                    netlink_set_xdp_fd(if_index, prog_fd, Some(old_prog_fd), replace_flags.bits())
+                        .map_err(|io_error| XdpError::NetlinkError { io_error })?;
+                }
+                // dispose of link and avoid detach on drop
+                mem::forget(link);
+                self.data.links.insert(XdpLink(XdpLinkInner::NlLink(NlLink {
+                    if_index,
+                    prog_fd,
+                    flags,
+                })))
+            }
+        }
+    }
 }
 
 #[derive(Debug)]

aya/src/sys/bpf.rs

@@ -1,5 +1,7 @@
 use crate::{
-    generated::{btf_func_linkage, btf_param, btf_var_secinfo, BTF_INT_SIGNED, BTF_VAR_STATIC},
+    generated::{
+        btf_func_linkage, btf_param, btf_var_secinfo, BPF_F_REPLACE, BTF_INT_SIGNED, BTF_VAR_STATIC,
+    },
     obj::{btf::BtfType, copy_instructions},
     Btf,
 };
@@ -333,6 +335,27 @@ pub(crate) fn bpf_link_create(
     sys_bpf(bpf_cmd::BPF_LINK_CREATE, &attr)
 }
 
+// since kernel 5.7
+pub(crate) fn bpf_link_update(
+    link_fd: RawFd,
+    new_prog_fd: RawFd,
+    old_prog_fd: Option<RawFd>,
+    flags: u32,
+) -> SysResult {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+
+    attr.link_update.link_fd = link_fd as u32;
+    attr.link_update.new_prog_fd = new_prog_fd as u32;
+    if let Some(fd) = old_prog_fd {
+        attr.link_update.old_prog_fd = fd as u32;
+        attr.link_update.flags = flags | BPF_F_REPLACE;
+    } else {
+        attr.link_update.flags = flags;
+    }
+
+    sys_bpf(bpf_cmd::BPF_LINK_UPDATE, &attr)
+}
+
 pub(crate) fn bpf_prog_attach(
     prog_fd: RawFd,
     target_fd: RawFd,