Refactor code of attaching kprobe to non-skb funcs

Hide attaching kprobe to non-skb funcs from `main.go`.

Signed-off-by: Leon Hwang <[email protected]>

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
+}

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()
 	}