aya: Add Extension::attach_to_program()

This allows for Extension programs already loaded to the kernel to be
attached to another program that is BTF-compatible with the one provided
at `load()` time

Signed-off-by: Dave Tucker <[email protected]>

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")?;
@@ -66,56 +67,9 @@ impl Extension {
     /// 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);
@@ -142,6 +96,30 @@ impl Extension {
         self.data.links.insert(ExtensionLink(FdLink::new(link_fd)))
     }
 
+    /// Attaches the extension to another BTF-compatible program.
+    ///
+    /// Attaches the extension effectively replacing the original target function.
+    /// program.
+    ///
+    /// 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 +140,56 @@ impl Extension {
     }
 }
 
+fn get_btf_info(prog_fd: i32, func_name: &str) -> Result<(RawFd, u32), ProgramError> {
+    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,
+        })?;
+
+    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)?;
+
+    Ok((btf_fd, btf_id))
+}
+
 define_link_wrapper!(
     /// The link used by [Extension] programs.
     ExtensionLink,