Refactor code of attaching kprobe

Hide attaching kprobes progress from `main.go`.

Signed-off-by: Leon Hwang <hffilwlqm@gmail.com>

Author: Asphaltt

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,49 @@ 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 {
+	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
+}

main.go

@@ -14,7 +14,6 @@ import (
 	"syscall"
 	"time"
 
-	pb "github.com/cheggaaa/pb/v3"
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/btf"
 	"github.com/cilium/ebpf/link"
@@ -198,35 +197,16 @@ 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)
+			defer kp.Close()
 		}
 	}
 
@@ -239,15 +219,12 @@ func main() {
 			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)
+				defer fexit.Close()
 			}
 		}
 	}
@@ -258,41 +235,17 @@ func main() {
 				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)
+			defer kp.Close()
 		}
 	}
 
-	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.KprobeFuncs(ctx, funcs, coll, addr2name, useKprobeMulti, flags.FilterKprobeBatch)
+		defer k.DetachKprobes()
 	}
-	log.Printf("Attached (ignored %d)\n", ignored)
 
 	log.Println("Listening for events..")