diff --git a/internal/pwru/skb_tracker.go b/internal/pwru/skb_tracker.go new file mode 100644 index 00000000..8cf9390f --- /dev/null +++ b/internal/pwru/skb_tracker.go @@ -0,0 +1,60 @@ +// SPDX-License-Identifier: Apache-2.0 +/* Copyright 2024 Authors of Cilium */ + +package pwru + +import ( + "errors" + "log" + "os" + + "github.com/cilium/ebpf" + "github.com/cilium/ebpf/link" +) + +type skbTracker struct { + links []link.Link +} + +func (t *skbTracker) Close() { + for _, l := range t.links { + _ = l.Close() + } + t.links = nil +} + +func TrackSkb(coll *ebpf.Collection, haveFexit bool) *skbTracker { + var t skbTracker + + kp, err := link.Kprobe("kfree_skbmem", coll.Programs["kprobe_skb_lifetime_termination"], nil) + if err != nil { + if !errors.Is(err, os.ErrNotExist) { + log.Fatalf("Opening kprobe kfree_skbmem: %s\n", err) + } else { + log.Printf("Warn: kfree_skbmem not found, pwru is likely to mismatch skb due to lack of skb lifetime management\n") + } + } else { + t.links = append(t.links, kp) + } + + if haveFexit { + progs := []*ebpf.Program{ + coll.Programs["fexit_skb_clone"], + coll.Programs["fexit_skb_copy"], + } + for _, prog := range progs { + fexit, err := link.AttachTracing(link.TracingOptions{ + Program: prog, + }) + if err != nil { + if !errors.Is(err, os.ErrNotExist) { + log.Fatalf("Opening tracing(%s): %s\n", prog, err) + } + } else { + t.links = append(t.links, fexit) + } + } + } + + return &t +} diff --git a/main.go b/main.go index e0be2455..ececcb75 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" @@ -196,36 +195,9 @@ func main() { defer close() } - if flags.FilterTrackSkb { - kp, err := link.Kprobe("kfree_skbmem", coll.Programs["kprobe_skb_lifetime_termination"], nil) - if err != nil { - if !errors.Is(err, os.ErrNotExist) { - log.Fatalf("Opening kprobe kfree_skbmem: %s\n", err) - } else { - log.Printf("Warn: kfree_skbmem not found, pwru is likely to mismatch skb due to lack of skb lifetime management\n") - } - } else { - defer kp.Close() - } - - if haveFexit { - progs := []*ebpf.Program{ - coll.Programs["fexit_skb_clone"], - coll.Programs["fexit_skb_copy"], - } - for _, prog := range progs { - fexit, err := link.AttachTracing(link.TracingOptions{ - Program: prog, - }) - if err != nil { - if !errors.Is(err, os.ErrNotExist) { - log.Fatalf("Opening tracing(%s): %s\n", prog, err) - } - } else { - defer fexit.Close() - } - } - } + if flags.FilterTrackSkb && len(funcs) != 0 { + t := pwru.TrackSkb(coll, haveFexit) + defer t.Close() } if len(funcs) != 0 {