reworked the way kprobe ctx is saved and restored

kunai-common/src/kunai-macros/src/lib.rs

@@ -117,6 +117,7 @@ pub fn str_enum_derive(item: TokenStream) -> TokenStream {
     };
 
     let mut as_str_arms = vec![];
+    let mut c_str_arms = vec![];
     let mut from_str_arms = vec![];
     let mut try_from_uint_arms = vec![];
     let mut variants = vec![];
@@ -143,6 +144,7 @@ pub fn str_enum_derive(item: TokenStream) -> TokenStream {
         // we generate a match arm delivering the good error name
         if v.fields.is_empty() {
             as_str_arms.push(quote!(Self::#name => #args,));
+            c_str_arms.push(quote!(Self::#name => concat!(#args, '\0'),));
             from_str_arms.push(quote!(#args => Ok(Self::#name),));
             variants.push(quote!(Self::#name));
             try_from_uint_arms.push(quote!(ty if Self::#name as u64 == ty => Ok(Self::#name),));
@@ -195,6 +197,13 @@ pub fn str_enum_derive(item: TokenStream) -> TokenStream {
                     #(#as_str_arms)*
                 }
             }
+
+            #[inline(always)]
+            pub const fn as_str_with_null(&self) -> &'static str{
+                match self {
+                    #(#c_str_arms)*
+                }
+            }
         }
     )
     .into()

kunai-ebpf/src/error.rs

@@ -1,5 +1,6 @@
 use crate::alloc;
 use crate::maps;
+use crate::util;
 
 use aya_log_ebpf::error;
 use kunai_common::bpf_events;
@@ -39,8 +40,6 @@ pub enum ProbeError {
     CoReFieldMissing,
     #[error("failed to get kprobe arg")]
     KProbeArgFailure,
-    #[error("failed to restore kprobe context")]
-    KProbeCtxRestoreFailure,
     #[wrap]
     BpfMapError(MapError),
     #[wrap]
@@ -61,6 +60,8 @@ pub enum ProbeError {
     EventError(bpf_events::Error),
     #[wrap]
     CgroupError(cgroup::Error),
+    #[wrap]
+    KprobeCtxError(util::Error),
 }
 
 /// log a ProbeError using Aya log
@@ -141,4 +142,10 @@ impl From<cgroup::Error> for ProbeError {
     }
 }
 
+impl From<util::Error> for ProbeError {
+    fn from(value: util::Error) -> Self {
+        Self::KprobeCtxError(value)
+    }
+}
+
 pub type ProbeResult<T> = core::result::Result<T, ProbeError>;

kunai-ebpf/src/probes.rs

@@ -32,7 +32,6 @@ mod mmap;
 mod mount;
 mod mprotect;
 mod prctl;
-mod save;
 mod schedule;
 mod send_data;
 

kunai-ebpf/src/probes/bpf_socket.rs

@@ -3,11 +3,18 @@ use kunai_common::{co_re::sock_fprog_kern, net::SocketInfo};
 
 use super::*;
 
+#[kprobe(name = "sk.enter.__sk_attach_prog")]
+pub fn enter_sk_attach_prog(ctx: ProbeContext) -> u32 {
+    unsafe { ignore_result!(ProbeFn::sk_sk_attach_prog.save_ctx(&ctx)) }
+    0
+}
+
 #[kretprobe(name = "sk.exit.__sk_attach_prog")]
 pub fn exit_sk_attach_prog(exit_ctx: ProbeContext) -> u32 {
-    match unsafe {
-        restore_entry_ctx(ProbeFn::__sk_attach_prog)
-            .ok_or(ProbeError::KProbeCtxRestoreFailure)
+    let rc = match unsafe {
+        ProbeFn::sk_sk_attach_prog
+            .restore_ctx()
+            .map_err(ProbeError::from)
             .and_then(|entry_ctx| {
                 let entry_ctx = &entry_ctx.probe_context();
 
@@ -22,14 +29,24 @@ pub fn exit_sk_attach_prog(exit_ctx: ProbeContext) -> u32 {
             log_err!(&exit_ctx, s);
             error::BPF_PROG_FAILURE
         }
-    }
+    };
+    // we cleanup entry context
+    ignore_result!(unsafe { ProbeFn::sk_sk_attach_prog.clean_ctx() });
+    rc
+}
+
+#[kprobe(name = "sk.enter.reuseport_attach_prog")]
+pub fn enter_reuseport_attach_prog(ctx: ProbeContext) -> u32 {
+    unsafe { ignore_result!(ProbeFn::sk_reuseport_attach_prog.save_ctx(&ctx)) }
+    0
 }
 
 #[kretprobe(name = "sk.exit.reuseport_attach_prog")]
 pub fn exit_reuseport_attach_prog(exit_ctx: ProbeContext) -> u32 {
-    match unsafe {
-        restore_entry_ctx(ProbeFn::reuseport_attach_prog)
-            .ok_or(ProbeError::KProbeCtxRestoreFailure)
+    let rc = match unsafe {
+        ProbeFn::sk_reuseport_attach_prog
+            .restore_ctx()
+            .map_err(ProbeError::from)
             .and_then(|entry_ctx| {
                 let entry_ctx = &entry_ctx.probe_context();
 
@@ -44,7 +61,10 @@ pub fn exit_reuseport_attach_prog(exit_ctx: ProbeContext) -> u32 {
             log_err!(&exit_ctx, s);
             error::BPF_PROG_FAILURE
         }
-    }
+    };
+    // we cleanup entry context
+    ignore_result!(unsafe { ProbeFn::sk_reuseport_attach_prog.clean_ctx() });
+    rc
 }
 
 #[inline(always)]

kunai-ebpf/src/probes/clone.rs

@@ -3,21 +3,33 @@ use kunai_common::co_re::{kernel_clone_args, task_struct};
 
 use super::*;
 
-#[kprobe(name = "enter.wake_up_new_task")]
+#[kprobe(name = "clone.enter.kernel_clone")]
+pub fn enter_kernel_clone(ctx: ProbeContext) -> u32 {
+    unsafe {
+        ignore_result!(ProbeFn::clone_kernel_clone.save_ctx(&ctx));
+    }
+    0
+}
+
+#[kprobe(name = "clone.enter.wake_up_new_task")]
 pub fn enter_wake_up_new_task(ctx: ProbeContext) -> u32 {
-    match unsafe { try_enter_wake_up_new_task(&ctx) } {
+    let rc = match unsafe { try_enter_wake_up_new_task(&ctx) } {
         Ok(_) => error::BPF_PROG_SUCCESS,
         Err(s) => {
             log_err!(&ctx, s);
             error::BPF_PROG_FAILURE
         }
-    }
+    };
+    // we cleanup saved context
+    ignore_result!(unsafe { ProbeFn::clone_kernel_clone.clean_ctx() });
+    rc
 }
 
 unsafe fn try_enter_wake_up_new_task(ctx: &ProbeContext) -> ProbeResult<()> {
     // makes sure we are inside kernel_clone
-    if let Ok(entry_ctx) = restore_entry_ctx(ProbeFn::kernel_clone)
-        .ok_or(ProbeError::KProbeCtxRestoreFailure)
+    if let Ok(entry_ctx) = ProbeFn::clone_kernel_clone
+        .restore_ctx()
+        .map_err(ProbeError::from)
         .and_then(|c| Ok(c.probe_context()))
     {
         let clone_args = kernel_clone_args::from_ptr(kprobe_arg!(&entry_ctx, 0)?);

kunai-ebpf/src/probes/connect.rs

@@ -1,33 +1,30 @@
 use super::*;
 
-use aya_bpf::{helpers::bpf_ktime_get_ns, programs::ProbeContext};
+use aya_bpf::programs::ProbeContext;
 use kunai_common::net::IpPort;
 
 #[kprobe(name = "net.enter.__sys_connect")]
 pub fn enter_sys_connect(ctx: ProbeContext) -> u32 {
-    unsafe {
-        ignore_result!(save_context(
-            ProbeFn::__sys_connect,
-            bpf_ktime_get_ns(),
-            &ctx
-        ))
-    }
+    unsafe { ignore_result!(ProbeFn::net_sys_connect.save_ctx(&ctx)) }
     0
 }
 
 #[kretprobe(name = "net.exit.__sys_connect")]
 pub fn exit_sys_connect(ctx: ProbeContext) -> u32 {
-    match unsafe {
-        restore_entry_ctx(ProbeFn::__sys_connect)
-            .ok_or(ProbeError::KProbeCtxRestoreFailure)
+    let rc = match unsafe {
+        ProbeFn::net_sys_connect
+            .restore_ctx()
+            .map_err(ProbeError::from)
             .and_then(|ent_ctx| try_exit_connect(ent_ctx, &ctx))
     } {
         Ok(_) => error::BPF_PROG_SUCCESS,
         Err(s) => {
             log_err!(&ctx, s);
             error::BPF_PROG_FAILURE
         }
-    }
+    };
+    ignore_result!(unsafe { ProbeFn::net_sys_connect.clean_ctx() });
+    rc
 }
 
 const EINPROGRESS: i32 = 115;

kunai-ebpf/src/probes/dns.rs

@@ -79,23 +79,33 @@ impl SockHelper {
     }
 }
 
+#[kprobe(name = "net.dns.enter.vfs_read")]
+pub fn enter_vfs_read(ctx: ProbeContext) -> u32 {
+    unsafe { ignore_result!(ProbeFn::dns_vfs_read.save_ctx(&ctx)) };
+    0
+}
+
 #[kretprobe(name = "net.dns.exit.vfs_read")]
 pub fn exit_vfs_read(ctx: ProbeContext) -> u32 {
-    match unsafe { try_exit_vfs_read(&ctx) } {
+    let rc = match unsafe { try_exit_vfs_read(&ctx) } {
         Ok(_) => error::BPF_PROG_SUCCESS,
         Err(s) => {
             log_err!(&ctx, s);
             error::BPF_PROG_FAILURE
         }
-    }
+    };
+    // we cleanup saved entry context
+    ignore_result!(unsafe { ProbeFn::dns_vfs_read.clean_ctx() });
+    rc
 }
 
 #[inline(always)]
 unsafe fn try_exit_vfs_read(ctx: &ProbeContext) -> ProbeResult<()> {
     let rc = ctx.ret().unwrap_or(-1);
 
-    let entry_ctx =
-        restore_entry_ctx(ProbeFn::vfs_read).ok_or(ProbeError::KProbeCtxRestoreFailure)?;
+    let entry_ctx = ProbeFn::dns_vfs_read
+        .restore_ctx()
+        .map_err(ProbeError::from)?;
     let saved_ctx = entry_ctx.probe_context();
 
     let file = co_re::file::from_ptr(kprobe_arg!(&saved_ctx, 0)?);
@@ -116,19 +126,28 @@ unsafe fn try_exit_vfs_read(ctx: &ProbeContext) -> ProbeResult<()> {
     Ok(())
 }
 
+#[kprobe(name = "net.dns.enter.__sys_recvfrom")]
+pub fn enter_recv(ctx: ProbeContext) -> u32 {
+    unsafe { ignore_result!(ProbeFn::dns_sys_recv_from.save_ctx(&ctx)) }
+    0
+}
+
 #[kretprobe(name = "net.dns.exit.__sys_recvfrom")]
 pub fn exit_recv(ctx: ProbeContext) -> u32 {
-    match unsafe {
-        restore_entry_ctx(ProbeFn::__sys_recvfrom)
-            .ok_or(ProbeError::KProbeCtxRestoreFailure)
+    let rc = match unsafe {
+        ProbeFn::dns_sys_recv_from
+            .restore_ctx()
+            .map_err(ProbeError::from)
             .and_then(|ent_ctx| try_exit_recv(ent_ctx, &ctx))
     } {
         Ok(_) => error::BPF_PROG_SUCCESS,
         Err(s) => {
             log_err!(&ctx, s);
             error::BPF_PROG_FAILURE
         }
-    }
+    };
+    ignore_result!(unsafe { ProbeFn::dns_sys_recv_from.clean_ctx() });
+    rc
 }
 
 #[inline(always)]
@@ -164,19 +183,29 @@ unsafe fn try_exit_recv(
     Ok(())
 }
 
+#[kprobe(name = "net.dns.enter.__sys_recvmsg")]
+pub fn enter_sys_recvmsg(ctx: ProbeContext) -> u32 {
+    unsafe { ignore_result!(ProbeFn::net_dns_sys_recvmsg.save_ctx(&ctx)) }
+    0
+}
+
 #[kretprobe(name = "net.dns.exit.__sys_recvmsg")]
 pub fn exit_sys_recvmsg(ctx: ProbeContext) -> u32 {
-    match unsafe {
-        restore_entry_ctx(ProbeFn::__sys_recvmsg)
-            .ok_or(ProbeError::KProbeCtxRestoreFailure)
+    let rc = match unsafe {
+        ProbeFn::net_dns_sys_recvmsg
+            .restore_ctx()
+            .map_err(ProbeError::from)
             .and_then(|ent_ctx| try_exit_recvmsg(ent_ctx, &ctx))
     } {
         Ok(_) => error::BPF_PROG_SUCCESS,
         Err(s) => {
             log_err!(&ctx, s);
             error::BPF_PROG_FAILURE
         }
-    }
+    };
+    // we cleanup saved entry context
+    ignore_result!(unsafe { ProbeFn::net_dns_sys_recvmsg.clean_ctx() });
+    rc
 }
 
 #[inline(always)]

kunai-ebpf/src/probes/fs.rs

@@ -221,7 +221,7 @@ unsafe fn try_security_path_unlink(ctx: &ProbeContext) -> ProbeResult<()> {
     // as vfs_unlink can be reached without security_path_unlink being called
     // we report error when insertion is failing
     PATHS
-        .insert(&ProbeFn::security_path_unlink.uuid(), &p, 0)
+        .insert(&ProbeFn::security_path_unlink.depth_key(), &p, 0)
         .map_err(|_| MapError::InsertFailure)?;
 
     Ok(())
@@ -250,7 +250,7 @@ unsafe fn try_vfs_unlink(ctx: &ProbeContext) -> ProbeResult<()> {
 
     e.init_from_current_task(Type::FileUnlink)?;
 
-    let path_key = ProbeFn::security_path_unlink.uuid();
+    let path_key = ProbeFn::security_path_unlink.depth_key();
     if let Some(p) = PATHS.get(&path_key) {
         e.data.path.copy_from(&p);
         // make some room in the cache

kunai-ebpf/src/probes/mount.rs

@@ -1,17 +1,27 @@
 use super::*;
 
 use aya_bpf::cty::c_int;
+
 use aya_bpf::maps::LruHashMap;
 use aya_bpf::programs::ProbeContext;
 
 #[map]
 static mut MOUNT_EVENTS: LruHashMap<u128, MountEvent> = LruHashMap::with_max_entries(1024, 0);
 
+#[kprobe(name = "fs.enter.security_sb_mount")]
+pub fn enter_path_mount(ctx: ProbeContext) -> u32 {
+    unsafe {
+        ignore_result!(ProbeFn::fs_security_sb_mount.save_ctx(&ctx));
+    }
+    0
+}
+
 #[kretprobe(name = "fs.exit.security_sb_mount")]
 pub fn exit_security_sb_mount(ctx: ProbeContext) -> u32 {
     match unsafe {
-        restore_entry_ctx(ProbeFn::security_sb_mount)
-            .ok_or(ProbeError::KProbeCtxRestoreFailure)
+        ProbeFn::fs_security_sb_mount
+            .restore_ctx()
+            .map_err(ProbeError::from)
             .and_then(|ent_ctx| try_exit_security_sb_mount(ent_ctx, &ctx))
     } {
         Ok(_) => error::BPF_PROG_SUCCESS,
@@ -60,17 +70,22 @@ unsafe fn try_exit_security_sb_mount(
 // path_mount is available only since 5.9 before that do_mount must be hooked
 #[kretprobe(name = "fs.exit.path_mount")]
 pub fn exit_path_mount(ctx: ProbeContext) -> u32 {
-    match unsafe {
-        restore_entry_ctx(ProbeFn::security_sb_mount)
-            .ok_or(ProbeError::KProbeCtxRestoreFailure)
+    let rc = match unsafe {
+        ProbeFn::fs_security_sb_mount
+            .restore_ctx()
+            .map_err(ProbeError::from)
             .and_then(|ent_ctx| try_exit_path_mount(ent_ctx, &ctx))
     } {
         Ok(_) => error::BPF_PROG_SUCCESS,
         Err(s) => {
             log_err!(&ctx, s);
             error::BPF_PROG_FAILURE
         }
-    }
+    };
+    // we cleanup only at the end of path_mount to let security_sb_mount available
+    // for exit_path_mount probe
+    ignore_result!(unsafe { ProbeFn::fs_security_sb_mount.clean_ctx() });
+    rc
 }
 
 #[inline(always)]