Refactor code of main.go

internal/pwru/kprobe.go

@@ -18,6 +18,13 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
+type kprober struct {
+	links []link.Link
+
+	kprobeMulti bool
+	kprobeBatch uint
+}
+
 type Kprobe struct {
 	hookFunc  string // internal use
 	HookFuncs []string
@@ -115,13 +122,15 @@ func AttachKprobes(ctx context.Context, bar *pb.ProgressBar, kps []Kprobe, batch
 }
 
 // DetachKprobes detaches kprobes concurrently.
-func DetachKprobes(links []link.Link, batch uint) {
+func (k *kprober) DetachKprobes() {
 	log.Println("Detaching kprobes...")
 
+	links := k.links
 	bar := pb.StartNew(len(links))
 	defer bar.Finish()
 
-	if batch == 0 || batch >= uint(len(links)) {
+	batch := k.kprobeBatch
+	if k.kprobeMulti || batch >= uint(len(links)) {
 		for _, l := range links {
 			_ = l.Close()
 			bar.Increment()
@@ -188,3 +197,69 @@ func AttachKprobeMulti(ctx context.Context, bar *pb.ProgressBar, kprobes []Kprob
 
 	return
 }
+
+func KprobeSkbFuncs(ctx context.Context, funcs Funcs, coll *ebpf.Collection, a2n Addr2Name, useKprobeMulti bool, batch uint) *kprober {
+	msg := "kprobe"
+	if useKprobeMulti {
+		msg = "kprobe-multi"
+	}
+	log.Printf("Attaching kprobes (via %s)...\n", msg)
+
+	ignored := 0
+	bar := pb.StartNew(len(funcs))
+
+	pwruKprobes := make([]Kprobe, 0, len(funcs))
+	funcsByPos := GetFuncsByPos(funcs)
+	for pos, fns := range funcsByPos {
+		fn, ok := coll.Programs[fmt.Sprintf("kprobe_skb_%d", pos)]
+		if ok {
+			pwruKprobes = append(pwruKprobes, Kprobe{HookFuncs: fns, Prog: fn})
+		} else {
+			ignored += len(fns)
+			bar.Add(len(fns))
+		}
+	}
+
+	var k kprober
+	k.kprobeMulti = useKprobeMulti
+	k.kprobeBatch = batch
+
+	if !useKprobeMulti {
+		l, i := AttachKprobes(ctx, bar, pwruKprobes, batch)
+		k.links = l
+		ignored += i
+	} else {
+		l, i := AttachKprobeMulti(ctx, bar, pwruKprobes, a2n)
+		k.links = l
+		ignored += i
+	}
+	bar.Finish()
+	select {
+	case <-ctx.Done():
+		return nil
+	default:
+	}
+	log.Printf("Attached (ignored %d)\n", ignored)
+
+	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
+}

internal/pwru/skb_tracker.go

@@ -0,0 +1,61 @@
+// 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) Detach() {
+	for _, l := range t.links {
+		_ = l.Close()
+	}
+	t.links = nil
+}
+
+func TrackSkb(coll *ebpf.Collection, haveFexit, trackSkbClone 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")
+			return &t
+		}
+	} else {
+		t.links = append(t.links, kp)
+	}
+
+	if haveFexit && trackSkbClone {
+		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
+}

main.go

@@ -14,10 +14,8 @@ import (
 	"syscall"
 	"time"
 
-	pb "github.com/cheggaaa/pb/v3"
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/btf"
-	"github.com/cilium/ebpf/link"
 	"github.com/cilium/ebpf/rlimit"
 	"golang.org/x/sys/unix"
 
@@ -197,101 +195,20 @@ func main() {
 		defer close()
 	}
 
-	var kprobes []link.Link
-	defer func() {
-		batch := uint(0)
-		if !useKprobeMulti {
-			batch = flags.FilterKprobeBatch
-		}
-		pwru.DetachKprobes(kprobes, batch)
-	}()
-
-	msg := "kprobe"
-	if useKprobeMulti {
-		msg = "kprobe-multi"
-	}
-	log.Printf("Attaching kprobes (via %s)...\n", msg)
-	ignored := 0
-	bar := pb.StartNew(len(funcs))
-
 	if flags.FilterTrackSkb || flags.FilterTrackSkbByStackid {
-		kp, err := link.Kprobe("kfree_skbmem", coll.Programs["kprobe_skb_lifetime_termination"], nil)
-		bar.Increment()
-		if err != nil {
-			if !errors.Is(err, os.ErrNotExist) {
-				log.Fatalf("Opening kprobe kfree_skbmem: %s\n", err)
-			} else {
-				ignored += 1
-				log.Printf("Warn: kfree_skbmem not found, pwru is likely to mismatch skb due to lack of skb lifetime management\n")
-			}
-		} else {
-			kprobes = append(kprobes, kp)
-		}
+		t := pwru.TrackSkb(coll, haveFexit, flags.FilterTrackSkb)
+		defer t.Detach()
 	}
 
-	if haveFexit && flags.FilterTrackSkb {
-		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,
-			})
-			bar.Increment()
-			if err != nil {
-				if !errors.Is(err, os.ErrNotExist) {
-					log.Fatalf("Opening tracing(%s): %s\n", prog, err)
-				} else {
-					ignored += 1
-				}
-			} else {
-				kprobes = append(kprobes, fexit)
-			}
-		}
+	if nonSkbFuncs := flags.FilterNonSkbFuncs; len(nonSkbFuncs) != 0 {
+		k := pwru.KprobeNonSkbFuncs(nonSkbFuncs, funcs, coll)
+		defer k.DetachKprobes()
 	}
 
-	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)
-			bar.Increment()
-			if err != nil {
-				log.Fatalf("Opening kprobe %s: %s\n", fn, err)
-			}
-			kprobes = append(kprobes, kp)
-		}
-	}
-
-	pwruKprobes := make([]pwru.Kprobe, 0, len(funcs))
-	funcsByPos := pwru.GetFuncsByPos(funcs)
-	for pos, fns := range funcsByPos {
-		fn, ok := coll.Programs[fmt.Sprintf("kprobe_skb_%d", pos)]
-		if ok {
-			pwruKprobes = append(pwruKprobes, pwru.Kprobe{HookFuncs: fns, Prog: fn})
-		} else {
-			ignored += len(fns)
-			bar.Add(len(fns))
-		}
-	}
-	if !useKprobeMulti {
-		l, i := pwru.AttachKprobes(ctx, bar, pwruKprobes, flags.FilterKprobeBatch)
-		kprobes = append(kprobes, l...)
-		ignored += i
-	} else {
-		l, i := pwru.AttachKprobeMulti(ctx, bar, pwruKprobes, addr2name)
-		kprobes = append(kprobes, l...)
-		ignored += i
-	}
-	bar.Finish()
-	select {
-	case <-ctx.Done():
-		return
-	default:
+	if len(funcs) != 0 {
+		k := pwru.KprobeSkbFuncs(ctx, funcs, coll, addr2name, useKprobeMulti, flags.FilterKprobeBatch)
+		defer k.DetachKprobes()
 	}
-	log.Printf("Attached (ignored %d)\n", ignored)
 
 	log.Println("Listening for events..")