edriver-rust: change lost_cnt to atomic

chriskaliX · Nov 30, 2024 · e35311f · e35311f
1 parent bc3712b
commit e35311f
Show file tree

Hide file tree

Showing 4 changed files with 115 additions and 25 deletions.
diff --git a/plugins/edriver-rust/src/bpf/common/common.h b/plugins/edriver-rust/src/bpf/common/common.h
@@ -38,7 +38,6 @@ static inline struct ipv6_pinfo *get_inet_pinet6(struct inet_sock *inet)
     return pinet6_own_impl;
 }
 
-
 static inline struct ipv6_pinfo *inet6_sk_own_impl(struct sock *__sk, struct inet_sock *inet)
 {
     volatile unsigned char sk_state_own_impl;
@@ -246,7 +245,6 @@ static __always_inline int get_sock_v6(struct sock *sk, struct hds_socket_info_v
 
 /* ===== END ===== */
 
-
 static __always_inline struct file *fget_raw(struct task_struct *task, u64 fd_num)
 {
     struct file **fd = BPF_CORE_READ(task, files, fdt, fd);
@@ -356,4 +354,33 @@ static __always_inline void *get_fd(struct task_struct *task, u64 num)
     return get_path(__builtin_preserve_access_index(&path));
 }
 
+static __always_inline bool is_x86_compat(struct task_struct *task)
+{
+#if defined(bpf_target_x86)
+    return READ_KERN(task->thread_info.status) & TS_COMPAT;
+#else
+    return false;
+#endif
+}
+
+static __always_inline bool is_arm64_compat(struct task_struct *task)
+{
+#if defined(bpf_target_arm64)
+    return READ_KERN(task->thread_info.flags) & _TIF_32BIT;
+#else
+    return false;
+#endif
+}
+
+static __always_inline bool is_compat(struct task_struct *task)
+{
+#if defined(bpf_target_x86)
+    return is_x86_compat(task);
+#elif defined(bpf_target_arm64)
+    return is_arm64_compat(task);
+#else
+    return false;
+#endif
+}
+
 #endif
diff --git a/plugins/edriver-rust/src/bpf/common/consts.h b/plugins/edriver-rust/src/bpf/common/consts.h
@@ -58,6 +58,65 @@
 
 /* consts: vmlinux contants fix */
 #define PF_KTHREAD                  0x00200000
+#define TS_COMPAT                   0x0002    /* 32bit syscall active (64BIT)*/
+
+/* system call mapping */
+#if defined(bpf_target_x86)
+#define NR_exit              80
+#define NR_exit_group        231
+/* execve */
+#define NR_execve            59
+#define NR_execveat          322
+/* ptrace & prctl */
+#define NR_ptrace            101
+#define NR_prctl             157
+/* mount */
+#define NR_mount             40
+/* kernel module init */
+#define NR_init_module       175
+#define NR_finit_module      313
+/* memfd */
+#define NR_memfd_create      319
+/* sockets */
+#define NR_connect           42
+#define NR_accept            43
+#define NR_accept4           288
+#define NR_listen            50
+#define NR_bind              49
+#define NR_recvfrom          45
+#define NR_recvmsg           47
+#define NR_recvmmsg          299
+/* bpf */
+#define NR_bpf               321
+#elif defined(bpf_target_arm64)
+#define NR_exit              93
+#define NR_exit_group        94
+/* execve */
+#define NR_execve            221
+#define NR_execveat          281
+/* ptrace & prctl */
+#define NR_ptrace            117
+#define NR_prctl             167
+/* mount */
+#define NR_mount             40
+/* kernel module init */
+#define NR_init_module       105
+#define NR_finit_module      273
+/* memfd */
+#define NR_memfd_create      279
+/* sockets */
+#define NR_connect           203
+#define NR_accept            202
+#define NR_accept4           242
+#define NR_listen            201
+#define NR_bind              200
+#define NR_recvfrom          207
+#define NR_recvmsg           212
+#define NR_recvmmsg          243
+/* bpf */
+#define NR_bpf               280
+#endif
+
 
 enum buf_index
 {

diff --git a/plugins/edriver-rust/src/bpf/common/edriver.h b/plugins/edriver-rust/src/bpf/common/edriver.h
@@ -82,14 +82,23 @@ int rtp__process_exec(struct bpf_raw_tracepoint_args *ctx)
     return report_event(&c);
 }
 
-// SEC("raw_tracepoint/sys_exit")
-// int rtp__sys_exit(struct bpf_raw_tracepoint_args *ctx)
-// {
-//     /* skip failed syscalls */
-//     if (ctx->args[1])
-//         return 0;
-//     return 0;
-// }
+SEC("raw_tracepoint/sys_exit")
+int rtp__sys_exit(struct bpf_raw_tracepoint_args *ctx)
+{
+    /* skip failed syscalls */
+    if (ctx->args[1])
+        return 0;
+
+    /* task bit transform */
+    struct task_struct *task = (struct task_struct *)bpf_get_current_task();
+
+    /* arch */
+    if (is_x86_compat(task)) {
+
+    }
+
+    return 0;
+}
 
 /* proc_info init */
 static struct proc_info *proc_info_init(struct task_struct *task)

diff --git a/plugins/edriver-rust/src/bpfmgr.rs b/plugins/edriver-rust/src/bpfmgr.rs
@@ -8,8 +8,9 @@ use libbpf_rs::{
 };
 use log::*;
 use sdk::{Client, Record};
+use std::sync::atomic::{AtomicU64, Ordering};
 use std::{
-    sync::{Arc, Mutex},
+    sync::Arc,
     thread,
     time::{Duration, SystemTime, UNIX_EPOCH},
 };
@@ -24,7 +25,7 @@ mod hades_skel {
 use hades_skel::*;
 
 lazy_static! {
-    pub static ref LOSS_CNT: Arc<Mutex<u64>> = Arc::new(Mutex::new(0));
+    pub static ref LOSS_CNT: Arc<AtomicU64> = Arc::new(AtomicU64::new(0));
 }
 
 pub struct Bpfmanager {}
@@ -39,7 +40,6 @@ impl Bpfmanager {
 
         skel.attach().context("Skel attach failed")?;
 
-        let loss_cnt_c = LOSS_CNT.clone();
         let mut trans = Transformer::new();
 
         /* event handle wrap */
@@ -48,7 +48,7 @@ impl Bpfmanager {
             println!("{:?}", map);
         };
 
-        Self::start_heartbeat_thread(client, loss_cnt_c)?;
+        Self::start_heartbeat_thread(client)?;
 
         let binding = skel.maps();
         let map = binding.events();
@@ -63,8 +63,7 @@ impl Bpfmanager {
     }
 
     fn handle_lost_events(_cpu: i32, cnt: u64) {
-        let mut loss_count = LOSS_CNT.lock().unwrap();
-        *loss_count += cnt;
+        LOSS_CNT.fetch_add(cnt, Ordering::SeqCst);
     }
 
     fn bump_rlimit() -> Result<()> {
@@ -78,7 +77,7 @@ impl Bpfmanager {
         Ok(())
     }
 
-    fn start_heartbeat_thread(mut client: Client, loss_counter: Arc<Mutex<u64>>) -> Result<()> {
+    fn start_heartbeat_thread(mut client: Client) -> Result<()> {
         thread::Builder::new()
             .name("heartbeat".to_string())
             .spawn(move || loop {
@@ -92,20 +91,16 @@ impl Bpfmanager {
                 rec.data_type = 900;
 
                 let pld = rec.mut_data();
-                if let Ok(loss_count) = loss_counter.lock() {
-                    pld.fields
-                        .insert("loss_cnt".to_string(), loss_count.to_string());
-                } else {
-                    warn!("Failed to lock loss_counter");
-                    continue;
-                }
+                let loss_count = LOSS_CNT.load(Ordering::SeqCst);
+                pld.fields
+                    .insert("loss_cnt".to_string(), loss_count.to_string());
 
                 if let Err(err) = client.send_record(&rec) {
                     warn!("Heartbeat will exit: {}", err);
                     break;
                 }
 
-                *loss_counter.lock().unwrap() = 0;
+                LOSS_CNT.store(0, Ordering::SeqCst); // Reset the loss counter
                 thread::sleep(Duration::from_secs(30));
             })
             .context("Failed to spawn heartbeat thread")?;