Support tracing by pid

main.go

@@ -23,6 +23,9 @@ func main() {
 			if c.IsSet("event") && c.IsSet("exclude-event") {
 				return fmt.Errorf("'event' and 'exclude-event' can't be used in parallel")
 			}
+			if c.IsSet("container") && c.IsSet("pid") {
+				return fmt.Errorf("'container' and 'pid' can't be used in parallel")
+			}
 			events, err := prepareEventsToTrace(c.StringSlice("event"), c.StringSlice("exclude-event"))
 			if err != nil {
 				return err
@@ -34,6 +37,7 @@ func main() {
 				ShowExecEnv:           c.Bool("show-exec-env"),
 				OutputFormat:          c.String("output"),
 				PerfBufferSize:        c.Int("perf-buffer-size"),
+				PidsToTrace:           c.IntSlice("pid"),
 				BlobPerfBufferSize:    c.Int("blob-perf-buffer-size"),
 				OutputPath:            c.String("output-path"),
 				FilterFileWrite:       c.StringSlice("filter-file-write"),
@@ -100,6 +104,12 @@ func main() {
 				Value:   false,
 				Usage:   "trace only containers",
 			},
+			&cli.IntSliceFlag{
+				Name:  "pid",
+				Aliases: []string{"p"},
+				Value: nil,
+				Usage: "trace only the specified pid. use this flag multiple times to choose multiple pids",
+			},
 			&cli.BoolFlag{
 				Name:  "detect-original-syscall",
 				Value: false,

tracee/consts.go

@@ -26,13 +26,19 @@ const (
 type bpfConfig uint32
 
 const (
-	configContMode          bpfConfig = 0
+	configMode              bpfConfig = 0
 	configDetectOrigSyscall bpfConfig = 1
 	configExecEnv           bpfConfig = 2
 	configCaptureFiles      bpfConfig = 3
 	configExtractDynCode    bpfConfig = 4
 )
 
+const (
+	modeSystem    uint32 = 0
+	modePid       uint32 = 1
+	modeContainer uint32 = 2
+)
+
 const (
 	tailVfsWrite uint32 = 0
 	tailSendBin  uint32 = 1

tracee/event_monitor_ebpf.c

@@ -121,12 +121,16 @@
 #define TAG_OWNER          61UL
 #define TAG_GROUP          62UL
 
-#define CONFIG_CONT_MODE        0
+#define CONFIG_MODE             0
 #define CONFIG_SHOW_SYSCALL     1
 #define CONFIG_EXEC_ENV         2
 #define CONFIG_CAPTURE_FILES    3
 #define CONFIG_EXTRACT_DYN_CODE 4
 
+#define MODE_SYSTEM     0
+#define MODE_PID        1
+#define MODE_CONTAINER  2
+
 #if LINUX_VERSION_CODE < KERNEL_VERSION(4, 14, 0)
 #error Minimal required kernel version is 4.14
 #endif
@@ -777,7 +781,7 @@ static __always_inline int should_trace()
     struct task_struct *task = (struct task_struct *)bpf_get_current_task();
 
     u32 rc = 0;
-    if (get_config(CONFIG_CONT_MODE))
+    if (get_config(CONFIG_MODE) == MODE_CONTAINER)
         rc = lookup_pid_ns(task);
     else
         rc = lookup_pid();
@@ -799,7 +803,7 @@ static __always_inline int init_context(context_t *context)
     struct task_struct *task;
     task = (struct task_struct *)bpf_get_current_task();
 
-    if (get_config(CONFIG_CONT_MODE)) {
+    if (get_config(CONFIG_MODE) == MODE_CONTAINER) {
         context->tid = get_task_ns_pid(task);
         context->pid = get_task_ns_tgid(task);
         context->ppid = get_task_ns_ppid(task);
@@ -1269,7 +1273,7 @@ static __always_inline int trace_ret_generic_fork(struct pt_regs *ctx, u32 id, u
     bool delete_args = true;
     int rc = trace_ret_generic(ctx, id, types, tags, delete_args);
 
-    if (!rc && !get_config(CONFIG_CONT_MODE)) {
+    if (!rc && (get_config(CONFIG_MODE) != MODE_CONTAINER)) {
         u32 pid = PT_REGS_RC(ctx);
         add_pid_fork(pid);
     }
@@ -1586,13 +1590,12 @@ int syscall__execve(struct pt_regs *ctx,
 {
     context_t context = {};
 
-    u32 ret = 0;
-    if (get_config(CONFIG_CONT_MODE))
-        ret = add_pid_ns_if_needed();
-    else
-        ret = add_pid();
+    if (get_config(CONFIG_MODE) == MODE_CONTAINER)
+        add_pid_ns_if_needed();
+    else if (get_config(CONFIG_MODE) == MODE_SYSTEM)
+        add_pid();
 
-    if (ret == 0)
+    if (!should_trace())
         return 0;
 
     init_context(&context);
@@ -1655,13 +1658,12 @@ int syscall__execveat(struct pt_regs *ctx,
 {
     context_t context = {};
 
-    u32 ret = 0;
-    if (get_config(CONFIG_CONT_MODE))
-        ret = add_pid_ns_if_needed();
-    else
-        ret = add_pid();
+    if (get_config(CONFIG_MODE) == MODE_CONTAINER)
+        add_pid_ns_if_needed();
+    else if (get_config(CONFIG_MODE) == MODE_SYSTEM)
+        add_pid();
 
-    if (ret == 0)
+    if (!should_trace())
         return 0;
 
     init_context(&context);
@@ -1736,7 +1738,7 @@ int trace_do_exit(struct pt_regs *ctx, long code)
     context.argnum = 0;
     context.retval = code;
 
-    if (get_config(CONFIG_CONT_MODE))
+    if (get_config(CONFIG_MODE) == MODE_CONTAINER)
         remove_pid_ns_if_needed();
     else
         remove_pid();

tracee/tracee.go

@@ -23,6 +23,7 @@ import (
 type TraceeConfig struct {
 	EventsToTrace         []int32
 	ContainerMode         bool
+	PidsToTrace           []int
 	DetectOriginalSyscall bool
 	ShowExecEnv           bool
 	OutputFormat          string
@@ -281,8 +282,15 @@ func (t *Tracee) initBPF(ebpfProgram string) error {
 
 	bpfConfig := bpf.NewTable(t.bpfModule.TableId("config_map"), t.bpfModule)
 
-	binary.LittleEndian.PutUint32(key, uint32(configContMode))
-	binary.LittleEndian.PutUint32(leaf, boolToUInt32(t.config.ContainerMode))
+	mode := modeSystem
+	if t.config.ContainerMode {
+		mode = modeContainer
+	} else if len(t.config.PidsToTrace) > 0 {
+		mode = modePid
+	}
+
+	binary.LittleEndian.PutUint32(key, uint32(configMode))
+	binary.LittleEndian.PutUint32(leaf, mode)
 	bpfConfig.Set(key, leaf)
 	binary.LittleEndian.PutUint32(key, uint32(configDetectOrigSyscall))
 	binary.LittleEndian.PutUint32(leaf, boolToUInt32(t.config.DetectOriginalSyscall))
@@ -297,6 +305,13 @@ func (t *Tracee) initBPF(ebpfProgram string) error {
 	binary.LittleEndian.PutUint32(leaf, boolToUInt32(t.config.CaptureMem))
 	bpfConfig.Set(key, leaf)
 
+	pidsMap := bpf.NewTable(t.bpfModule.TableId("pids_map"), t.bpfModule)
+	for _, pid := range t.config.PidsToTrace {
+		binary.LittleEndian.PutUint32(key, uint32(pid))
+		binary.LittleEndian.PutUint32(leaf, uint32(pid))
+		pidsMap.Set(key, leaf)
+	}
+
 	// Load send_bin function to prog_array to be used as tail call
 	progArrayBPFTable := bpf.NewTable(t.bpfModule.TableId("prog_array"), t.bpfModule)
 	binary.LittleEndian.PutUint32(key, tailVfsWrite)