diff --git a/internal/pwru/kprobe.go b/internal/pwru/kprobe.go index 518d2ddd..f59cd45f 100644 --- a/internal/pwru/kprobe.go +++ b/internal/pwru/kprobe.go @@ -198,7 +198,7 @@ func AttachKprobeMulti(ctx context.Context, bar *pb.ProgressBar, kprobes []Kprob return } -func KprobeFuncs(ctx context.Context, funcs Funcs, coll *ebpf.Collection, a2n Addr2Name, useKprobeMulti bool, batch uint) *kprober { +func KprobeSkbFuncs(ctx context.Context, funcs Funcs, coll *ebpf.Collection, a2n Addr2Name, useKprobeMulti bool, batch uint) *kprober { msg := "kprobe" if useKprobeMulti { msg = "kprobe-multi" @@ -243,3 +243,23 @@ func KprobeFuncs(ctx context.Context, funcs Funcs, coll *ebpf.Collection, a2n Ad return &k } + +func KprobeNonSkbFuncs(nonSkbFuncs []string, funcs Funcs, coll *ebpf.Collection) *kprober { + var k kprober + k.kprobeBatch = uint(len(nonSkbFuncs)) + + for _, fn := range nonSkbFuncs { + if _, ok := funcs[fn]; ok { + continue + } + + kp, err := link.Kprobe(fn, coll.Programs["kprobe_skb_by_stackid"], nil) + if err != nil { + log.Fatalf("Opening kprobe %s: %s\n", fn, err) + } + + k.links = append(k.links, kp) + } + + return &k +} diff --git a/main.go b/main.go index 469ab828..e256cb17 100644 --- a/main.go +++ b/main.go @@ -16,7 +16,6 @@ import ( "github.com/cilium/ebpf" "github.com/cilium/ebpf/btf" - "github.com/cilium/ebpf/link" "github.com/cilium/ebpf/rlimit" "golang.org/x/sys/unix" @@ -202,21 +201,13 @@ func main() { defer t.Detach() } - if len(flags.FilterNonSkbFuncs) > 0 { - for _, fn := range flags.FilterNonSkbFuncs { - if _, ok := funcs[fn]; ok { - continue - } - kp, err := link.Kprobe(fn, coll.Programs["kprobe_skb_by_stackid"], nil) - if err != nil { - log.Fatalf("Opening kprobe %s: %s\n", fn, err) - } - defer kp.Close() - } + if nonSkbFuncs := flags.FilterNonSkbFuncs; len(nonSkbFuncs) != 0 { + k := pwru.KprobeNonSkbFuncs(nonSkbFuncs, funcs, coll) + defer k.DetachKprobes() } if len(funcs) != 0 { - k := pwru.KprobeFuncs(ctx, funcs, coll, addr2name, useKprobeMulti, flags.FilterKprobeBatch) + k := pwru.KprobeSkbFuncs(ctx, funcs, coll, addr2name, useKprobeMulti, flags.FilterKprobeBatch) defer k.DetachKprobes() }