Merge pull request #97 from amosproj/88-vfs_write

88 vfs write

rust/backend/common/src/lib.rs

@@ -15,9 +15,26 @@ pub enum KProbeTypes {
 
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
-pub struct KProbeData {
-    pub pid: u32,
-    pub tid: u32,
-    pub probe_type: KProbeTypes,
-    pub ret: bool,
+pub struct VfsWriteCall {
+    pid: u32,
+    tid: u32,
+    begin_time_stamp: u64,
+    fd: i32,
+    bytes_written: usize,
 }
+
+impl VfsWriteCall {
+    pub fn new(pid: u32, tid: u32, begin_time_stamp: u64, fd: i32, bytes_written: usize) -> Self {
+        Self { pid, tid, begin_time_stamp, fd, bytes_written}
+    }
+}
+
+#[inline(always)]
+pub fn generate_id(pid: u32, tgid: u32) -> u64{
+    let pid_u64 = pid as u64;
+    let tgid_u64 = tgid as u64;
+
+    (pid_u64 << 32) | tgid_u64
+}
+
+

rust/backend/ebpf/src/lib.rs

@@ -10,4 +10,4 @@
 
 mod vfs_tracing;
 
-pub use vfs_tracing::{vfs_write, KPROBES};
+pub use vfs_tracing::{vfs_write, VFS_WRITE_MAP};

rust/backend/ebpf/src/main.rs

@@ -14,7 +14,7 @@ use aya_ebpf::{
     maps::{PerCpuArray, RingBuf},
     programs::XdpContext,
 };
-pub use backend_ebpf::{vfs_write, KPROBES};
+pub use backend_ebpf::{vfs_write, VFS_WRITE_MAP};
 
 #[map(name = "COUNTER")]
 static PACKET_COUNTER: PerCpuArray<u32> = PerCpuArray::with_max_entries(1, 0);

rust/backend/ebpf/src/vfs_tracing.rs

@@ -2,35 +2,76 @@
 //
 // SPDX-License-Identifier: MIT
 
+
+
+const TIME_LIMIT_NS: u64 = 100_000_000;
+
 use aya_ebpf::{
-    macros::{kprobe, map},
-    maps::RingBuf,
-    programs::ProbeContext,
+    macros::{kprobe, map, kretprobe},
+    maps::{HashMap, RingBuf},
+    programs::{ProbeContext, RetProbeContext},
     EbpfContext,
+    helpers::gen::bpf_ktime_get_ns,
 };
-use backend_common::{KProbeData, KProbeTypes};
+use aya_log_ebpf::info;
+use backend_common::{generate_id, VfsWriteCall};
+
+
+
+#[map(name = "VFS_WRITE_MAP")]
+pub static VFS_WRITE_MAP: RingBuf = RingBuf::with_byte_size(1024, 0);
 
-#[map(name = "Kprobes")]
-pub static KPROBES: RingBuf = RingBuf::with_byte_size(1024, 0);
+
+#[map(name = "VfsWriteIntern")]
+static VFS_WRITE_TIMESTAMPS: HashMap<u64, VfsWriteIntern> = HashMap::with_max_entries(1024, 0);
+
+
+struct VfsWriteIntern {
+    begin_time_stamp: u64,
+    fd: i32,
+    bytes_written: usize,
+}
 
 #[kprobe]
 pub fn vfs_write(ctx: ProbeContext) -> Result<(), u32> {
-    let pid = ctx.pid();
-    let tid = ctx.tgid();
-
-    let data = KProbeData {
-        pid,
-        tid,
-        probe_type: KProbeTypes::VfsWrite,
-        ret: false,
+    let id = generate_id(ctx.pid(), ctx.tgid());
+    let data = VfsWriteIntern {
+        begin_time_stamp: unsafe {bpf_ktime_get_ns()},
+        fd: ctx.arg(0).unwrap_or(-1),
+        bytes_written: ctx.arg(2).unwrap_or(usize::MAX) as usize,
     };
-    let mut entry = match KPROBES.reserve::<KProbeData>(0) {
-        Some(entry) => entry,
-        None => return Err(0),
+
+    match VFS_WRITE_TIMESTAMPS.insert(&id, &data, 0) {
+        Ok(_) => Ok(()),
+        Err(_) => Err(0),
+    }
+
+}
+
+
+#[kretprobe]
+pub fn vfs_write_ret(ctx: RetProbeContext) -> Result<(), u32> {
+    let probe_end = unsafe { bpf_ktime_get_ns() };
+
+    let pid = ctx.pid();
+    let tgid = ctx.tgid();
+    let call_id = generate_id(pid, tgid);
+    let data = match unsafe { VFS_WRITE_TIMESTAMPS.get(&call_id) } {
+        None => {return Err(0)}
+        Some(entry) => {entry}
     };
 
-    entry.write(data);
-    entry.submit(0);
+    if  probe_end - data.begin_time_stamp > TIME_LIMIT_NS {
+        let data = VfsWriteCall::new(pid, tgid, data.begin_time_stamp, data.fd, data.bytes_written);
+
+        let mut entry = match VFS_WRITE_MAP.reserve::<VfsWriteCall>(0) {
+            Some(entry) => entry,
+            None => return Err(0),
+        };
+
+        entry.write(data);
+        entry.submit(0);
+    }
 
     Ok(())
-}
+}