refactor

2 files changed, 78 insertions, 13 deletions

diff --git a/main.bpf.c b/main.bpf.c
index 7c275a5..8b0b2a2 100644
--- a/main.bpf.c
+++ b/main.bpf.c
@@ -4,10 +4,14 @@
 
 #include <bpf/bpf_helpers.h>
 
-// TODO: Make this configurable via a flag from the userland part.
+// TODO: Split out this file into several *.bpf.c programs.
+
+// TODO: Make UID_FILTER configurable via a flag from the userland part.
 // For now, this is set to my own user for development purposes.
 #define UID_FILTER 1001 
 
+// Helper structs for opening file(s)
+
 struct open_event {
     int fd;
     int syscall_id;
@@ -47,7 +51,7 @@ int handle_enter_open(struct trace_event_raw_sys_enter *ctx) {
 }
 
 SEC("tracepoint/syscalls/sys_exit_open")
-int handle_exit_open(struct trace_event_raw_sys_exit *args) {
+int handle_exit_open(struct trace_event_raw_sys_exit *ctx) {
     if ((bpf_get_current_uid_gid() & 0xFFFFFFFF) != UID_FILTER)
         return 0;
 
@@ -56,14 +60,13 @@ int handle_exit_open(struct trace_event_raw_sys_exit *args) {
     if (!eventp) {
         return 0;
     }
-    eventp->fd = args->ret;
-    bpf_perf_event_output(args, &open_event_map, BPF_F_CURRENT_CPU, eventp, sizeof(struct open_event));
+    eventp->fd = ctx->ret;
+    bpf_perf_event_output(ctx, &open_event_map, BPF_F_CURRENT_CPU, eventp, sizeof(struct open_event));
     bpf_map_delete_elem(&open_event_temp_map, &tid);
 
     return 0;
 }
 
-
 SEC("tracepoint/syscalls/sys_enter_openat")
 int handle_enter_openat(struct trace_event_raw_sys_enter *ctx) {
     if ((bpf_get_current_uid_gid() & 0xFFFFFFFF) != UID_FILTER)
@@ -81,11 +84,40 @@ int handle_enter_openat(struct trace_event_raw_sys_enter *ctx) {
 }
 
 SEC("tracepoint/syscalls/sys_exit_openat")
-int handle_exit_openat(struct trace_event_raw_sys_exit *args) {
+int handle_exit_openat(struct trace_event_raw_sys_exit *ctx) {
     if ((bpf_get_current_uid_gid() & 0xFFFFFFFF) != UID_FILTER)
         return 0;
 
-    return handle_exit_open(args);
+    return handle_exit_open(ctx);
+}
+
+// Helper structs for other syscalls on FDs
+
+struct fd_event {
+    int fd;
+    int syscall_id;
+    u32 tid;
+};
+
+struct {
+    __uint(type, BPF_MAP_TYPE_PERF_EVENT_ARRAY);
+    __uint(key_size, sizeof(u32));
+    __uint(value_size, sizeof(u32));
+} fd_event_map SEC(".maps");
+
+SEC("tracepoint/syscalls/sys_enter_close")
+int handle_enter_close(struct trace_event_raw_sys_enter *ctx) {
+    if ((bpf_get_current_uid_gid() & 0xFFFFFFFF) != UID_FILTER)
+        return 0;
+
+    struct fd_event event = {
+        .fd = (int)ctx->args[0],
+        .syscall_id = ctx->id,
+        .tid = bpf_get_current_pid_tgid(),
+    };
+    bpf_perf_event_output(ctx, &fd_event_map, BPF_F_CURRENT_CPU, &event, sizeof(struct fd_event));
+
+    return 0;
 }
 
 char LICENSE[] SEC("license") = "Dual BSD/GPL";
diff --git a/main.go b/main.go
index 0b330a2..a260017 100644
--- a/main.go
+++ b/main.go
@@ -8,13 +8,14 @@ import (
 	"fmt"
 	"log"
 	"runtime"
+	"sync"
 
 	"ioriotng/internal/tracepoints"
 
 	bpf "github.com/aquasecurity/libbpfgo"
 )
 
-type openatEvent struct {
+type openEvent struct {
 	FD        int32
 	SyscallID int32
 	TID       uint32
@@ -22,13 +23,23 @@ type openatEvent struct {
 	Comm      [16]byte
 }
 
-func (e openatEvent) String() string {
+func (e openEvent) String() string {
 	filename := e.Filename[:]
 	comm := e.Comm[:]
 	return fmt.Sprintf("syscall:%d tid:%v fd:%v filename:%s, comm:%s",
 		e.SyscallID, e.TID, e.FD, string(filename), string(comm))
 }
 
+type fdEvent struct {
+	FD        int32
+	SyscallID int32
+	TID       uint32
+}
+
+func (e fdEvent) String() string {
+	return fmt.Sprintf("syscall:%d tid:%v fd:%v", e.SyscallID, e.TID, e.FD)
+}
+
 func resizeMap(module *bpf.Module, name string, size uint32) error {
 	m, err := module.GetMap("open_event_map")
 	if err != nil {
@@ -66,16 +77,39 @@ func main() {
 		log.Fatal(err)
 	}
 
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		if err := openEvents(bpfModule); err != nil {
+			log.Fatal(err)
+		}
+	}()
+
+	go func() {
+		defer wg.Done()
+	}()
+
+	wg.Wait()
+	log.Println("Good bye")
+}
+
+func openEvents(bpfModule *bpf.Module) error {
 	eventsChannel := make(chan []byte)
 	lostChannel := make(chan uint64)
 	pb, err := bpfModule.InitPerfBuf("open_event_map", eventsChannel, lostChannel, 1)
 	if err != nil {
-		log.Fatal(err)
+		return err
 	}
+	defer func() {
+		pb.Stop()
+		pb.Close()
+	}()
 
 	pb.Poll(300)
 	for ev := range eventsChannel {
-		var e openatEvent
+		var e openEvent
 		if err := binary.Read(bytes.NewReader(ev), binary.LittleEndian, &e); err != nil {
 			log.Fatal(err)
 
@@ -85,8 +119,7 @@ func main() {
 		pb.Poll(300)
 	}
 
-	pb.Stop()
-	pb.Close()
+	return nil
 }
 
 func ksymArch() string {