Match syscall_enter and syscall_exits to get return value

3 files changed, 99 insertions, 33 deletions

diff --git a/fs-tracer-common/src/lib.rs b/fs-tracer-common/src/lib.rs
index 77ee0e5..9be9774 100644
--- a/fs-tracer-common/src/lib.rs
+++ b/fs-tracer-common/src/lib.rs
@@ -8,19 +8,22 @@ use core::str;
 pub struct WriteSyscallBPF {
     pub pid: u32,
     pub fd: u64,
-    pub buf: [u8; 128],
+    pub buf: [u8; 96],
     pub count: u64,
+
+    pub ret: i64,
 }
 
 unsafe impl Sync for WriteSyscallBPF {}
 
 impl fmt::Debug for WriteSyscallBPF {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        f.debug_struct("WriteSyscallArgs")
+        f.debug_struct("WriteSyscallBPF")
             .field("pid", &self.pid)
             .field("fd", &self.fd)
             .field("buf", &str::from_utf8(&self.buf).unwrap_or("") )
             .field("count", &self.count)
+            .field("ret", &self.ret)
             .finish()
     }
 }
\ No newline at end of file
diff --git a/fs-tracer-ebpf/src/main.rs b/fs-tracer-ebpf/src/main.rs
index d2cae0f..3ac89b1 100644
--- a/fs-tracer-ebpf/src/main.rs
+++ b/fs-tracer-ebpf/src/main.rs
@@ -1,9 +1,13 @@
 #![no_std]
 #![no_main]
 
-use core::ffi::c_void;
-use core::str;
+use core::{ffi::c_void, any::Any};
+use core::{str, ptr};
 
+use aya_bpf::helpers::gen::bpf_probe_read;
+use aya_bpf::helpers::{bpf_get_current_comm, bpf_probe_read_buf, bpf_probe_read_user, bpf_probe_read_user_str, bpf_probe_read_user_str_bytes};
+use aya_bpf::maps::HashMap;
+//use aya_bpf::helpers::gen::bpf_probe_read;
 use aya_bpf::{
     macros::{tracepoint, map},
     programs::TracePointContext, BpfContext, helpers::{gen::{self, bpf_probe_read_kernel_str}, bpf_probe_read_kernel_str_bytes}, maps::{PerfEventArray, PerCpuArray},
@@ -11,22 +15,31 @@ use aya_bpf::{
 use aya_log_ebpf::info;
 use fs_tracer_common::WriteSyscallBPF;
 
-#[derive(Clone, Copy)]
- struct WriteSyscallArgs {
-     fd: u64,
-     buf: *const u8,
-     count: u64,
-}
-
 #[map]
 static EVENTS: PerfEventArray<WriteSyscallBPF> = PerfEventArray::with_max_entries(1024, 0);
 
+#[map]
+static SYSCALLENTERS: HashMap<u32, WriteSyscallBPF> = HashMap::with_max_entries(1024, 0);
+
+enum SyscallType {
+    ENTER,
+    EXIT,
+}
+
 //#[map]
 //static mut READ_FROM_USERSPACE_BUFFER: PerCpuArray<[u8;2048]> = PerCpuArray::with_max_entries(1, 0);
 
 #[tracepoint]
-pub fn fs_tracer(ctx: TracePointContext) -> u32 {
-    match try_fs_tracer(ctx) {
+pub fn fs_tracer_enter(ctx: TracePointContext) -> u32 {
+    match try_fs_tracer(ctx, SyscallType::ENTER) {
+        Ok(ret) => ret,
+        Err(ret) => ret,
+    }
+}
+
+#[tracepoint]
+pub fn fs_tracer_exit(ctx: TracePointContext) -> u32 {
+    match try_fs_tracer(ctx, SyscallType::EXIT) {
         Ok(ret) => ret,
         Err(ret) => ret,
     }
@@ -34,21 +47,32 @@ pub fn fs_tracer(ctx: TracePointContext) -> u32 {
 
 #[inline(always)]
 fn ptr_at<T>(ctx: &TracePointContext, offset: usize) -> Option<*const T> {
-    let start = ctx.as_ptr();
+    let start = ctx.as_ptr(); //maybe try using the  bpf_probe_read here to see if we can use result of that to know the type of the syscall
 
     Some(unsafe { start.add(offset) } as *const T)
 }
 
-fn try_fs_tracer(ctx: TracePointContext) -> Result<u32, u32> {
+fn try_fs_tracer(ctx: TracePointContext, syscall_type: SyscallType) -> Result<u32, u32> {
     let syscall_nr = unsafe { * ptr_at::<i32>(&ctx, 8).unwrap() };
-    
-    return handle_syscall(ctx, syscall_nr);
+    //info!( &ctx, "syscall_nr: {}", syscall_nr);
+
+    //let typee = unsafe { *(cmd.as_ptr().add(0)  as *const u16)};
+    //let typee = ctx.read_at(0)
+
+   /* let a: [u64; 1] = [0u64; 1];
+    let ret = unsafe { bpf_probe_read(a.as_ptr() as *mut c_void, 8, ctx.as_ptr().add(0) as *const c_void)}; //TODO: Maybe we can try reading some high btis to get the type of the syscall exit or enter
+    info!(&ctx, "ret: {}", ret);
+ info!(&ctx, "x: {}", unsafe {a[0]});
+    info!(&ctx, "syscall_nr: {}", syscall_nr);*/ 
+    //info!(&ctx, "type: {}", typee);
+    //return Ok(1);
+    return handle_syscall(ctx, syscall_nr, syscall_type);
 }
 
-fn handle_syscall(ctx: TracePointContext, syscall_nr: i32) -> Result<u32, u32> {
+fn handle_syscall(ctx: TracePointContext, syscall_nr: i32, syscall_type: SyscallType) -> Result<u32, u32> {
     match syscall_nr {
         1 => {
-            return handle_sys_write(ctx);
+            return handle_sys_write(ctx, syscall_type);
         },
         2 => {
             return Ok(0)
@@ -64,7 +88,7 @@ fn handle_syscall(ctx: TracePointContext, syscall_nr: i32) -> Result<u32, u32> {
         },
         _ => {
             info!(&ctx, "unhandled syscall: {}", syscall_nr);
-            panic!("syscall: {}",syscall_nr);
+            return Err(1)
         }
     }
 }
@@ -77,8 +101,26 @@ static mut write_syscall_bpf: WriteSyscallBPF = WriteSyscallBPF {
     count: 0,
 };*/
 
-fn handle_sys_write(ctx: TracePointContext) -> Result<u32, u32> {
-    //info!(&ctx, "handle_sys_write start");
+#[derive(Clone, Copy)]
+ struct WriteSyscallArgs {
+     fd: u64,
+     buf: *const u8,
+     count: u64,
+}
+
+fn handle_sys_write(ctx: TracePointContext, syscall_type: SyscallType) -> Result<u32, u32> {
+    match syscall_type {
+        SyscallType::ENTER => {
+            return handle_sys_write_enter(ctx);
+        },
+        SyscallType::EXIT => {
+            return handle_sys_write_exit(ctx);
+        },
+    }
+}
+
+fn handle_sys_write_enter(ctx: TracePointContext) -> Result<u32, u32> {
+   // info!(&ctx, "handle_sys_write start");
     let args = unsafe { *ptr_at::<WriteSyscallArgs>(&ctx, 16).unwrap() };
 
     // if fd is stdout, stderr or stdin, ignore
@@ -87,35 +129,51 @@ fn handle_sys_write(ctx: TracePointContext) -> Result<u32, u32> {
     }
 
    // info!(&ctx, "argfs fd: {}", args.fd);
-    let mut buf = [0u8; 128];
+    let mut buf = [0u8; 96]; //we need to make this muuuuuch bigger
     //get_string_from_userspace(args.buf, unsafe { &mut *READ_FROM_USERSPACE_BUFFER.get_ptr_mut(0).unwrap() });
-    get_string_from_userspace(args.buf, &mut buf);
+    //get_string_from_userspace(args.buf, &mut buf);
+    let _ = unsafe { bpf_probe_read_user_str_bytes(args.buf, &mut buf) };
     let buf_ref = &buf;
    // info!(&ctx, "buf: {}", unsafe { str::from_utf8_unchecked(buf_ref) });
     //info!(&ctx, "count: {}", args.count);                                                                                                               ";
 
-    let mut anotherbuf = [0u8; 128];
-    unsafe { bpf_probe_read_kernel_str(anotherbuf.as_mut_ptr() as *mut c_void, 128, buf_ref.as_ptr() as *const c_void) };
+    let mut anotherbuf = [0u8; 96];
+    unsafe { bpf_probe_read_kernel_str(anotherbuf.as_mut_ptr() as *mut c_void, 96, buf_ref.as_ptr() as *const c_void) };
 
-    EVENTS.output(&ctx, &WriteSyscallBPF {
+    let tgid: u32 = ctx.tgid();
+    let _ = SYSCALLENTERS.insert(&tgid, &WriteSyscallBPF {
         pid: ctx.pid(),
         fd: args.fd,
         buf: anotherbuf,
         count: args.count,
+        ret: -9999,
     }, 0);
 
-    //info!(&ctx, "handle_sys_write end");
     return Ok(0)
 }
 
+fn handle_sys_write_exit(ctx: TracePointContext) -> Result<u32, u32> {
+    //info!(&ctx, "handle_sys_write_exit start");
+    let ret = unsafe { *ptr_at::<i64>(&ctx, 16).unwrap() }; //TODO: We cant use unwrap, thats why we couldnt use the aya helper fns
+
+    let tgid = ctx.tgid();
+    if let Some(& syscall) = unsafe { SYSCALLENTERS.get(&tgid) } {
+        let mut newsyscall = syscall.clone();
+        newsyscall.ret = ret;
+        EVENTS.output(&ctx, &newsyscall, 0);
+    }
+    //syscall_enter.ret = ret;
+    //EVENTS.output(&ctx, &syscall_enter, 0);
+
+    return Ok(0)
+}
+
+// thread can only execute 1 syscall at a time <- 
+
 //TODO: How are we going to correlate. We have open of a filename, we need to insert that into (pid, fd) -> filename. on close we remove from map. we need some timeout to remove stale entries
 //TODO: to get the fd from open, we need to know the return value of the syscall. for that we need a tracepoint on end and keep a map of (tgid, pid) -> WriteSyscallBPF). we need to differenciate the syscalls by id
 //TODO: Maybe we can use git itself for the diffs etc.
 
-fn get_string_from_userspace(ptr: *const u8, buf: &mut [u8]) {
-    unsafe { gen::bpf_probe_read_user_str( buf.as_mut_ptr() as *mut c_void, buf.len() as u32, ptr as *const c_void) };
-}
-
 #[panic_handler]
 fn panic(_info: &core::panic::PanicInfo) -> ! {
     unsafe { core::hint::unreachable_unchecked() }
diff --git a/fs-tracer/src/main.rs b/fs-tracer/src/main.rs
index 25edeb7..443e28e 100644
--- a/fs-tracer/src/main.rs
+++ b/fs-tracer/src/main.rs
@@ -41,13 +41,18 @@ async fn main() -> Result<(), anyhow::Error> {
         // This can happen if you remove all log statements from your eBPF program.
         warn!("failed to initialize eBPF logger: {}", e);
     }
-    let program: &mut TracePoint = bpf.program_mut("fs_tracer").unwrap().try_into()?;
+    let program: &mut TracePoint = bpf.program_mut("fs_tracer_enter").unwrap().try_into()?;
     program.load()?;
     //program.attach("syscalls", "sys_enter_open")?;
     program.attach("syscalls", "sys_enter_write")?;
+// program.attach("syscalls", "sys_exit_write")?;
     //program.attach("syscalls", "sys_enter_lseek")?;
     //program.attach("syscalls", "sys_enter_close")?;
 
+    let program2: &mut TracePoint = bpf.program_mut("fs_tracer_exit").unwrap().try_into()?;
+    program2.load()?;
+    program2.attach("syscalls", "sys_exit_write")?;
+
     println!("Num of cpus: {}", online_cpus()?.len());
 
     let mut perf_array =