aya: add support BPF_PROG_TYPE_SK_SKB programs and SockMaps

aya/src/bpf.rs

@@ -19,8 +19,8 @@ use crate::{
         Object, ParseError, ProgramKind,
     },
     programs::{
-        KProbe, ProbeKind, Program, ProgramData, ProgramError, SocketFilter, TracePoint, UProbe,
+        KProbe, ProbeKind, Program, ProgramData, ProgramError, SkSkb, SkSkbKind, SocketFilter,
-        Xdp,
+        TracePoint, UProbe, Xdp,
     },
     sys::bpf_map_update_elem_ptr,
     util::{possible_cpus, POSSIBLE_CPUS},
@@ -173,6 +173,14 @@ impl Bpf {
                     ProgramKind::TracePoint => Program::TracePoint(TracePoint { data }),
                     ProgramKind::SocketFilter => Program::SocketFilter(SocketFilter { data }),
                     ProgramKind::Xdp => Program::Xdp(Xdp { data }),
+                    ProgramKind::SkSkbStreamParser => Program::SkSkb(SkSkb {
+                        data,
+                        kind: SkSkbKind::StreamParser,
+                    }),
+                    ProgramKind::SkSkbStreamVerdict => Program::SkSkb(SkSkb {
+                        data,
+                        kind: SkSkbKind::StreamVerdict,
+                    }),
                 };
 
                 (name, program)

aya/src/maps/mod.rs

@@ -48,12 +48,14 @@ mod map_lock;
 pub mod array;
 pub mod hash_map;
 pub mod perf;
+pub mod sock_map;
 pub mod stack_trace;
 
 pub use array::{Array, PerCpuArray, ProgramArray};
 pub use hash_map::{HashMap, PerCpuHashMap};
 pub use map_lock::*;
 pub use perf::PerfEventArray;
+pub use sock_map::SockMap;
 pub use stack_trace::StackTraceMap;
 
 #[derive(Error, Debug)]

aya/src/maps/sock_map.rs

@@ -0,0 +1,102 @@
+//! An array of eBPF program file descriptors used as a jump table.
+
+use std::{
+    convert::TryFrom,
+    mem,
+    ops::{Deref, DerefMut},
+    os::unix::prelude::RawFd,
+};
+
+use crate::{
+    generated::bpf_map_type::BPF_MAP_TYPE_SOCKMAP,
+    maps::{Map, MapError, MapKeys, MapRef, MapRefMut},
+    sys::{bpf_map_delete_elem, bpf_map_update_elem},
+};
+
+pub struct SockMap<T: Deref<Target = Map>> {
+    pub(crate) inner: T,
+}
+
+impl<T: Deref<Target = Map>> SockMap<T> {
+    fn new(map: T) -> Result<SockMap<T>, MapError> {
+        let map_type = map.obj.def.map_type;
+        if map_type != BPF_MAP_TYPE_SOCKMAP as u32 {
+            return Err(MapError::InvalidMapType {
+                map_type: map_type as u32,
+            })?;
+        }
+        let expected = mem::size_of::<u32>();
+        let size = map.obj.def.key_size as usize;
+        if size != expected {
+            return Err(MapError::InvalidKeySize { size, expected });
+        }
+
+        let expected = mem::size_of::<RawFd>();
+        let size = map.obj.def.value_size as usize;
+        if size != expected {
+            return Err(MapError::InvalidValueSize { size, expected });
+        }
+        let _fd = map.fd_or_err()?;
+
+        Ok(SockMap { inner: map })
+    }
+
+    /// An iterator over the indices of the array that point to a program. The iterator item type
+    /// is `Result<u32, MapError>`.
+    pub unsafe fn indices(&self) -> MapKeys<'_, u32> {
+        MapKeys::new(&self.inner)
+    }
+
+    fn check_bounds(&self, index: u32) -> Result<(), MapError> {
+        let max_entries = self.inner.obj.def.max_entries;
+        if index >= self.inner.obj.def.max_entries {
+            Err(MapError::OutOfBounds { index, max_entries })
+        } else {
+            Ok(())
+        }
+    }
+}
+
+impl<T: Deref<Target = Map> + DerefMut<Target = Map>> SockMap<T> {
+    pub fn set(&mut self, index: u32, tcp_fd: RawFd, flags: u64) -> Result<(), MapError> {
+        let fd = self.inner.fd_or_err()?;
+        self.check_bounds(index)?;
+        bpf_map_update_elem(fd, &index, &tcp_fd, flags).map_err(|(code, io_error)| {
+            MapError::SyscallError {
+                call: "bpf_map_update_elem".to_owned(),
+                code,
+                io_error,
+            }
+        })?;
+        Ok(())
+    }
+
+    /// Clears the value at index in the jump table.
+    pub fn clear_index(&mut self, index: &u32) -> Result<(), MapError> {
+        let fd = self.inner.fd_or_err()?;
+        self.check_bounds(*index)?;
+        bpf_map_delete_elem(fd, index)
+            .map(|_| ())
+            .map_err(|(code, io_error)| MapError::SyscallError {
+                call: "bpf_map_delete_elem".to_owned(),
+                code,
+                io_error,
+            })
+    }
+}
+
+impl TryFrom<MapRef> for SockMap<MapRef> {
+    type Error = MapError;
+
+    fn try_from(a: MapRef) -> Result<SockMap<MapRef>, MapError> {
+        SockMap::new(a)
+    }
+}
+
+impl TryFrom<MapRefMut> for SockMap<MapRefMut> {
+    type Error = MapError;
+
+    fn try_from(a: MapRefMut) -> Result<SockMap<MapRefMut>, MapError> {
+        SockMap::new(a)
+    }
+}

aya/src/obj/mod.rs

@@ -73,6 +73,8 @@ pub enum ProgramKind {
     TracePoint,
     SocketFilter,
     Xdp,
+    SkSkbStreamParser,
+    SkSkbStreamVerdict,
 }
 
 impl FromStr for ProgramKind {
@@ -88,6 +90,8 @@ impl FromStr for ProgramKind {
             "xdp" => Xdp,
             "trace_point" => TracePoint,
             "socket_filter" => SocketFilter,
+            "sk_skb/stream_parser" => SkSkbStreamParser,
+            "sk_skb/stream_verdict" => SkSkbStreamVerdict,
             _ => {
                 return Err(ParseError::InvalidProgramKind {
                     kind: kind.to_string(),
@@ -247,7 +251,8 @@ impl Object {
     }
 
     fn parse_section(&mut self, mut section: Section) -> Result<(), BpfError> {
-        let parts = section.name.split("/").collect::<Vec<_>>();
+        let mut parts = section.name.rsplitn(2, "/").collect::<Vec<_>>();
+        parts.reverse();
 
         match parts.as_slice() {
             &[name]
@@ -269,7 +274,9 @@ impl Object {
             | &[ty @ "uretprobe", name]
             | &[ty @ "socket_filter", name]
             | &[ty @ "xdp", name]
-            | &[ty @ "trace_point", name] => {
+            | &[ty @ "trace_point", name]
+            | &[ty @ "sk_skb/stream_parser", name]
+            | &[ty @ "sk_skb/stream_verdict", name] => {
                 self.programs
                     .insert(name.to_string(), self.parse_program(&section, ty, name)?);
                 if !section.relocations.is_empty() {

aya/src/programs/mod.rs

@@ -47,6 +47,7 @@
 mod kprobe;
 mod perf_attach;
 mod probe;
+mod sk_skb;
 mod socket_filter;
 mod trace_point;
 mod uprobe;
@@ -59,15 +60,17 @@ use thiserror::Error;
 pub use kprobe::{KProbe, KProbeError};
 use perf_attach::*;
 pub use probe::ProbeKind;
+pub use sk_skb::{SkSkb, SkSkbKind};
 pub use socket_filter::{SocketFilter, SocketFilterError};
 pub use trace_point::{TracePoint, TracePointError};
 pub use uprobe::{UProbe, UProbeError};
 pub use xdp::{Xdp, XdpError, XdpFlags};
 
 use crate::{
-    generated::bpf_prog_type,
+    generated::{bpf_attach_type, bpf_prog_type},
+    maps::MapError,
     obj::{self, Function},
-    sys::bpf_load_program,
+    sys::{bpf_load_program, bpf_prog_detach},
 };
 #[derive(Debug, Error)]
 pub enum ProgramError {
@@ -106,6 +109,9 @@ pub enum ProgramError {
     #[error("unexpected program type")]
     UnexpectedProgramType,
 
+    #[error(transparent)]
+    MapError(#[from] MapError),
+
     #[error(transparent)]
     KProbeError(#[from] KProbeError),
 
@@ -134,6 +140,7 @@ pub enum Program {
     TracePoint(TracePoint),
     SocketFilter(SocketFilter),
     Xdp(Xdp),
+    SkSkb(SkSkb),
 }
 
 impl Program {
@@ -160,6 +167,7 @@ impl Program {
             Program::TracePoint(_) => BPF_PROG_TYPE_TRACEPOINT,
             Program::SocketFilter(_) => BPF_PROG_TYPE_SOCKET_FILTER,
             Program::Xdp(_) => BPF_PROG_TYPE_XDP,
+            Program::SkSkb(_) => BPF_PROG_TYPE_SK_SKB,
         }
     }
 
@@ -175,6 +183,7 @@ impl Program {
             Program::TracePoint(p) => &p.data,
             Program::SocketFilter(p) => &p.data,
             Program::Xdp(p) => &p.data,
+            Program::SkSkb(p) => &p.data,
         }
     }
 
@@ -185,6 +194,7 @@ impl Program {
             Program::TracePoint(p) => &mut p.data,
             Program::SocketFilter(p) => &mut p.data,
             Program::Xdp(p) => &mut p.data,
+            Program::SkSkb(p) => &mut p.data,
         }
     }
 }
@@ -352,6 +362,30 @@ impl Drop for FdLink {
     }
 }
 
+#[derive(Debug)]
+struct ProgAttachLink {
+    prog_fd: Option<RawFd>,
+    map_fd: Option<RawFd>,
+    attach_type: bpf_attach_type,
+}
+
+impl Link for ProgAttachLink {
+    fn detach(&mut self) -> Result<(), ProgramError> {
+        if let Some(prog_fd) = self.prog_fd.take() {
+            let _ = bpf_prog_detach(prog_fd, self.map_fd.take().unwrap(), self.attach_type);
+            Ok(())
+        } else {
+            Err(ProgramError::AlreadyDetached)
+        }
+    }
+}
+
+impl Drop for ProgAttachLink {
+    fn drop(&mut self) {
+        let _ = self.detach();
+    }
+}
+
 impl ProgramFd for Program {
     fn fd(&self) -> Option<RawFd> {
         self.data().fd
@@ -370,7 +404,7 @@ macro_rules! impl_program_fd {
     }
 }
 
-impl_program_fd!(KProbe, UProbe, TracePoint, SocketFilter, Xdp);
+impl_program_fd!(KProbe, UProbe, TracePoint, SocketFilter, Xdp, SkSkb);
 
 macro_rules! impl_try_from_program {
     ($($ty:ident),+ $(,)?) => {
@@ -400,4 +434,4 @@ macro_rules! impl_try_from_program {
     }
 }
 
-impl_try_from_program!(KProbe, UProbe, TracePoint, SocketFilter, Xdp);
+impl_try_from_program!(KProbe, UProbe, TracePoint, SocketFilter, Xdp, SkSkb);

aya/src/programs/sk_skb.rs

@@ -0,0 +1,61 @@
+use std::ops::Deref;
+
+use crate::{
+    generated::{
+        bpf_attach_type::{BPF_SK_SKB_STREAM_PARSER, BPF_SK_SKB_STREAM_VERDICT},
+        bpf_prog_type::BPF_PROG_TYPE_SK_SKB,
+    },
+    maps::{Map, SockMap},
+    programs::{load_program, LinkRef, ProgAttachLink, ProgramData, ProgramError},
+    sys::bpf_prog_attach,
+};
+
+#[derive(Copy, Clone, Debug)]
+pub enum SkSkbKind {
+    StreamParser,
+    StreamVerdict,
+}
+
+#[derive(Debug)]
+pub struct SkSkb {
+    pub(crate) data: ProgramData,
+    pub(crate) kind: SkSkbKind,
+}
+
+impl SkSkb {
+    /// Loads the program inside the kernel.
+    ///
+    /// See also [`Program::load`](crate::programs::Program::load).
+    pub fn load(&mut self) -> Result<(), ProgramError> {
+        load_program(BPF_PROG_TYPE_SK_SKB, &mut self.data)
+    }
+
+    /// Returns the name of the program.
+    pub fn name(&self) -> String {
+        self.data.name.to_string()
+    }
+
+    pub fn attach<T: Deref<Target = Map>>(
+        &mut self,
+        map: &SockMap<T>,
+    ) -> Result<LinkRef, ProgramError> {
+        let prog_fd = self.data.fd_or_err()?;
+        let map_fd = map.inner.fd_or_err()?;
+
+        let attach_type = match self.kind {
+            SkSkbKind::StreamParser => BPF_SK_SKB_STREAM_PARSER,
+            SkSkbKind::StreamVerdict => BPF_SK_SKB_STREAM_VERDICT,
+        };
+        bpf_prog_attach(prog_fd, map_fd, attach_type).map_err(|(_, io_error)| {
+            ProgramError::SyscallError {
+                call: "bpf_link_create".to_owned(),
+                io_error,
+            }
+        })?;
+        Ok(self.data.link(ProgAttachLink {
+            prog_fd: Some(prog_fd),
+            map_fd: Some(map_fd),
+            attach_type,
+        }))
+    }
+}

aya/src/sys/bpf.rs

@@ -222,6 +222,34 @@ pub(crate) fn bpf_link_create(
     sys_bpf(bpf_cmd::BPF_LINK_CREATE, &attr)
 }
 
+pub(crate) fn bpf_prog_attach(
+    prog_fd: RawFd,
+    map_fd: RawFd,
+    attach_type: bpf_attach_type,
+) -> SysResult {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+
+    attr.__bindgen_anon_5.attach_bpf_fd = prog_fd as u32;
+    attr.__bindgen_anon_5.target_fd = map_fd as u32;
+    attr.__bindgen_anon_5.attach_type = attach_type as u32;
+
+    sys_bpf(bpf_cmd::BPF_PROG_ATTACH, &attr)
+}
+
+pub(crate) fn bpf_prog_detach(
+    prog_fd: RawFd,
+    map_fd: RawFd,
+    attach_type: bpf_attach_type,
+) -> SysResult {
+    let mut attr = unsafe { mem::zeroed::<bpf_attr>() };
+
+    attr.__bindgen_anon_5.attach_bpf_fd = prog_fd as u32;
+    attr.__bindgen_anon_5.target_fd = map_fd as u32;
+    attr.__bindgen_anon_5.attach_type = attach_type as u32;
+
+    sys_bpf(bpf_cmd::BPF_PROG_ATTACH, &attr)
+}
+
 fn sys_bpf<'a>(cmd: bpf_cmd, attr: &'a bpf_attr) -> SysResult {
     syscall(Syscall::Bpf { cmd, attr })
 }