Add timeout for poll, make event configurable

WYXsb · Aug 30, 2023 · 160d430 · 160d430
1 parent 145befb
commit 160d430
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 7 deletions.
diff --git a/duetector/static/config.toml b/duetector/static/config.toml
@@ -32,12 +32,23 @@ exclude_gid = [
 [tracer]
 disabled = false
 
+[tracer.clonetracer]
+disabled = false
+attach_event = "__x64_sys_clone"
+poll_timeout = 10
+
+[tracer.tcpconnecttracer]
+disabled = false
+poll_timeout = 10
+
 [tracer.unametracer]
 disabled = false
 enable_cache = true
 
 [tracer.opentracer]
 disabled = false
+attach_event = "do_sys_openat2"
+poll_timeout = 10
 
 [collector]
 disabled = false

diff --git a/duetector/tracers/clone.py b/duetector/tracers/clone.py
@@ -10,10 +10,23 @@ class CloneTracer(BccTracer):
     A tracer for clone syscall
     """
 
+    default_config = {
+        **BccTracer.default_config,
+        "attach_event": "__x64_sys_clone",
+        "poll_timeout": 10,
+    }
     attach_type = "kprobe"
-    attatch_args = {"fn_name": "do_trace", "event": "__x64_sys_clone"}
+
+    @property
+    def attatch_args(self):
+        return {"fn_name": "do_trace", "event": self.config.attach_event}
+
     poll_fn = "perf_buffer_poll"
-    poll_args = {}
+
+    @property
+    def poll_args(self):
+        return {"timeout": int(self.config.poll_timeout)}
+
     data_t = namedtuple("CloneTracking", ["pid", "timestamp", "comm"])
     prog = """
     #include <linux/sched.h>
@@ -73,6 +86,6 @@ def print_callback(data: NamedTuple):
     poller = tracer.get_poller(b)
     while True:
         try:
-            poller()
+            poller(**tracer.poll_args)
         except KeyboardInterrupt:
             exit()
diff --git a/duetector/tracers/openat2.py b/duetector/tracers/openat2.py
@@ -10,10 +10,25 @@ class OpenTracer(BccTracer):
     A tracer for openat2 syscall
     """
 
+    default_config = {
+        **BccTracer.default_config,
+        "attach_event": "do_sys_openat2",
+        "poll_timeout": 10,
+    }
+
     attach_type = "kprobe"
+
+    @property
+    def attatch_args(self):
+        return {"fn_name": "do_trace", "event": self.config.attach_event}
+
     attatch_args = {"fn_name": "trace_entry", "event": "do_sys_openat2"}
     poll_fn = "ring_buffer_poll"
-    poll_args = {}
+
+    @property
+    def poll_args(self):
+        return {"timeout": int(self.config.poll_timeout)}
+
     data_t = namedtuple("OpenTracking", ["pid", "uid", "gid", "comm", "fname", "timestamp"])
 
     prog = """
@@ -72,6 +87,6 @@ def print_callback(data: NamedTuple):
     poller = tracer.get_poller(b)
     while True:
         try:
-            poller()
+            poller(**tracer.poll_args)
         except KeyboardInterrupt:
             exit()
diff --git a/duetector/tracers/tcpconnect.py b/duetector/tracers/tcpconnect.py
@@ -10,10 +10,19 @@ class TcpconnectTracer(BccTracer):
     A tracer for openat2 syscall
     """
 
+    default_config = {
+        **BccTracer.default_config,
+        "poll_timeout": 10,
+    }
+
     attach_type = "kprobe"
     attatch_args = {"fn_name": "do_trace", "event": "tcp_v4_connect"}
     poll_fn = "ring_buffer_poll"
-    poll_args = {}
+
+    @property
+    def poll_args(self):
+        return {"timeout": int(self.config.poll_timeout)}
+
     data_t = namedtuple("TcpTracking", ["pid", "comm", "saddr", "daddr", "dport"])
 
     # define BPF program
@@ -135,6 +144,6 @@ def print_callback(data: NamedTuple):
     poller = tracer.get_poller(b)
     while True:
         try:
-            poller()
+            poller(**tracer.poll_args)
         except KeyboardInterrupt:
             exit()