From 95c8b2d06e68bb9997bd01e43c76573f23d61f18 Mon Sep 17 00:00:00 2001
From: Leon Hwang <hffilwlqm@gmail.com>
Date: Thu, 5 Dec 2024 22:12:33 +0800
Subject: [PATCH] bpflbr: Trace only for specified pid by --filter-pid

It should be helpful to trace only for specified process.

Signed-off-by: Leon Hwang <hffilwlqm@gmail.com>
---
 bpf/lbr.c                       | 15 +++++++++++----
 internal/bpflbr/bpf_tracing.go  |  1 +
 internal/bpflbr/flags.go        |  2 ++
 internal/bpflbr/lbr_bpf_prog.go |  3 ++-
 4 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/bpf/lbr.c b/bpf/lbr.c
index 83613ef..1d470a8 100644
--- a/bpf/lbr.c
+++ b/bpf/lbr.c
@@ -10,11 +10,10 @@ struct lbr_config {
     __u32 suppress_lbr:1;
     __u32 output_stack:1;
     __u32 pad:30;
+    __u32 pid;
 };
 
-volatile const struct lbr_config lbr_config = {
-    .suppress_lbr = 0,
-};
+volatile const struct lbr_config lbr_config = {};
 #define cfg (&lbr_config)
 
 #define MAX_STACK_DEPTH 50
@@ -55,11 +54,19 @@ emit_lbr_event(void *ctx)
 
     if (!cfg->suppress_lbr)
         event->nr_bytes = bpf_get_branch_snapshot(event->lbr, sizeof(event->lbr), 0); /* required 5.16 kernel. */
+
+    /* Other filters must be after bpf_get_branch_snapshot() to avoid polluting
+     * LBR entries.
+     */
+
+    event->pid = bpf_get_current_pid_tgid() >> 32;
+    if (cfg->pid && event->pid != cfg->pid)
+        return BPF_OK;
+
     bpf_get_func_ret(ctx, (void *) &retval); /* required 5.17 kernel. */
     event->func_ret = retval;
     event->func_ip = bpf_get_func_ip(ctx); /* required 5.17 kernel. */
     event->cpu = cpu;
-    event->pid = bpf_get_current_pid_tgid() >> 32;
     bpf_get_current_comm(event->comm, sizeof(event->comm));
     event->func_stack_id = -1;
     if (cfg->output_stack)
diff --git a/internal/bpflbr/bpf_tracing.go b/internal/bpflbr/bpf_tracing.go
index 174dd88..5fbd780 100644
--- a/internal/bpflbr/bpf_tracing.go
+++ b/internal/bpflbr/bpf_tracing.go
@@ -28,6 +28,7 @@ func NewBPFTracing(spec *ebpf.CollectionSpec, reusedMaps map[string]*ebpf.Map, i
 	var cfg LbrConfig
 	cfg.SetSuppressLbr(suppressLbr)
 	cfg.SetOutputStack(outputFuncStack)
+	cfg.FilterPid = uint32(filterPid)
 
 	if err := spec.Variables["lbr_config"].Set(cfg); err != nil {
 		return nil, fmt.Errorf("failed to set lbr config: %w", err)
diff --git a/internal/bpflbr/flags.go b/internal/bpflbr/flags.go
index 8bda15b..6dc51cc 100644
--- a/internal/bpflbr/flags.go
+++ b/internal/bpflbr/flags.go
@@ -32,6 +32,7 @@ var (
 
 	suppressLbr     bool
 	outputFuncStack bool
+	filterPid       uint32
 )
 
 type ProgFlag struct {
@@ -133,6 +134,7 @@ func ParseFlags() (*Flags, error) {
 	f.StringVarP(&mode, "mode", "m", TracingModeExit, "mode of lbr tracing, exit or entry")
 	f.BoolVar(&suppressLbr, "suppress-lbr", false, "suppress LBR perf event")
 	f.BoolVar(&outputFuncStack, "output-stack", false, "output function call stack")
+	f.Uint32Var(&filterPid, "filter-pid", 0, "filter pid for tracing")
 
 	return &flags, f.Parse(os.Args)
 }
diff --git a/internal/bpflbr/lbr_bpf_prog.go b/internal/bpflbr/lbr_bpf_prog.go
index b354d13..3e8eb7f 100644
--- a/internal/bpflbr/lbr_bpf_prog.go
+++ b/internal/bpflbr/lbr_bpf_prog.go
@@ -9,7 +9,8 @@ const (
 )
 
 type LbrConfig struct {
-	Flags uint32
+	Flags     uint32
+	FilterPid uint32
 }
 
 func (cfg *LbrConfig) SetSuppressLbr(v bool) {