From 9b4937ba34250e459f3fa3db21ad4aea24282e78 Mon Sep 17 00:00:00 2001
From: Baitinq <manuelpalenzuelamerino@gmail.com>
Date: Thu, 13 Jun 2024 00:00:06 +0200
Subject: fs-tracer: cleanup syscall handler

---
 fs-tracer/src/main.rs            |  5 ++-
 fs-tracer/src/syscall_handler.rs | 79 ++++++++++++++++++++--------------------
 2 files changed, 44 insertions(+), 40 deletions(-)

diff --git a/fs-tracer/src/main.rs b/fs-tracer/src/main.rs
index ba5f77c..815a9e0 100644
--- a/fs-tracer/src/main.rs
+++ b/fs-tracer/src/main.rs
@@ -5,6 +5,7 @@ use aya::{include_bytes_aligned, Ebpf};
 use aya::{maps::AsyncPerfEventArray, programs::TracePoint};
 use aya_log::EbpfLogger;
 use bytes::BytesMut;
+use core::panic;
 use fs_tracer_common::SyscallInfo;
 use log::{debug, info, warn};
 use std::env;
@@ -97,7 +98,9 @@ async fn main() -> Result<(), anyhow::Error> {
                     }
                     let ptr = buf.as_ptr() as *const SyscallInfo;
                     let data = unsafe { ptr.read_unaligned() };
-                    syscall_handler.handle_syscall(data);
+                    if let Err(_) = syscall_handler.handle_syscall(data) {
+                        panic!("Error handling syscall!");
+                    }
                 }
             }
         }));
diff --git a/fs-tracer/src/syscall_handler.rs b/fs-tracer/src/syscall_handler.rs
index 99b839e..a6ca102 100644
--- a/fs-tracer/src/syscall_handler.rs
+++ b/fs-tracer/src/syscall_handler.rs
@@ -3,12 +3,11 @@ use std::ffi::CStr;
 use crossbeam_channel::Sender;
 use delay_map::HashMapDelay;
 
-use fs_tracer_common::SyscallInfo;
+use fs_tracer_common::{OpenSyscallBPF, SyscallInfo, WriteSyscallBPF};
 
 pub struct SyscallHandler {
     resolved_files: Sender<String>,
     open_files: HashMapDelay<i32, String>,
-    total: u64,
 }
 
 impl SyscallHandler {
@@ -16,52 +15,54 @@ impl SyscallHandler {
         Self {
             resolved_files,
             open_files: HashMapDelay::new(std::time::Duration::from_secs(400)),
-            total: 0,
         }
     }
 
-    pub fn handle_syscall(&mut self, data: SyscallInfo) -> u64 {
+    pub fn handle_syscall(&mut self, data: SyscallInfo) -> Result<(), ()> {
         match data {
-            SyscallInfo::Write(x) => {
-                let filename = self.open_files.get(&x.fd);
-                let filename = match filename {
-                    None => {
-                        println!("DIDNT FIND AN OPEN FILE FOR THE WRITE SYSCALL");
-                        return 0;
-                    }
-                    Some(x) => x,
-                };
-                let contents = CStr::from_bytes_until_nul(&x.buf)
-                    .unwrap_or_default()
-                    .to_str()
-                    .unwrap_or_default();
-                println!("WRITE KERNEL: DATA {:?}", x);
-                let _ = self.resolved_files.send(format!(
-                    r#"
+            SyscallInfo::Write(write_syscall) => self.handle_write(write_syscall),
+            SyscallInfo::Open(open_syscall) => self.handle_open(open_syscall),
+            //TODO: SyscallInfo::Close
+        }
+    }
+
+    fn handle_write(&self, write_syscall: WriteSyscallBPF) -> Result<(), ()> {
+        let filename = match self.open_files.get(&write_syscall.fd) {
+            None => {
+                println!("DIDNT FIND AN OPEN FILE FOR THE WRITE SYSCALL");
+                return Ok(());
+            }
+            Some(str) => str,
+        };
+        let contents = CStr::from_bytes_until_nul(&write_syscall.buf)
+            .unwrap_or_default()
+            .to_str()
+            .unwrap_or_default();
+        println!("WRITE KERNEL: DATA {:?}", write_syscall);
+        let _ = self.resolved_files.send(format!(
+            r#"
                 {{
                     "timestamp": "{}",
                     "absolute_path": "{}",
                     "contents": "{}"
                 }}
                 "#,
-                    chrono::Utc::now().to_rfc3339(),
-                    filename,
-                    contents,
-                ));
-                self.total += 1;
-                println!("Total: {:?}", self.total);
-                return 0;
-            }
-            SyscallInfo::Open(x) => {
-                let filename = CStr::from_bytes_until_nul(&x.filename)
-                    .unwrap_or_default()
-                    .to_str()
-                    .unwrap_or_default();
-                println!("OPEN KERNEL DATA: {:?}", x);
-                println!("OPEN FILENAME: {:?}", filename);
-                self.open_files.insert(x.ret, filename.to_string());
-                return 0;
-            } //SyscallInfo::Close
-        }
+            chrono::Utc::now().to_rfc3339(),
+            filename,
+            contents,
+        ));
+        Ok(())
+    }
+
+    fn handle_open(&mut self, open_syscall: OpenSyscallBPF) -> Result<(), ()> {
+        let filename = CStr::from_bytes_until_nul(&open_syscall.filename)
+            .unwrap_or_default()
+            .to_str()
+            .unwrap_or_default();
+        println!("OPEN KERNEL DATA: {:?}", open_syscall);
+        println!("OPEN FILENAME: {:?}", filename);
+        let fd = open_syscall.ret;
+        self.open_files.insert(fd, filename.to_string());
+        Ok(())
     }
 }
-- 
cgit 1.4.1