Add multi kprobe attach

Signed-off-by: Vincent Li <[email protected]>

bpf/kprobe_pwru.c

@@ -380,35 +380,35 @@ handle_everything(struct sk_buff *skb, struct pt_regs *ctx) {
 	return 0;
 }
 
-SEC("kprobe/skb-1")
+SEC("kprobe.multi/skb-1")
 int kprobe_skb_1(struct pt_regs *ctx) {
 	struct sk_buff *skb = (struct sk_buff *) PT_REGS_PARM1(ctx);
 
 	return handle_everything(skb, ctx);
 }
 
-SEC("kprobe/skb-2")
+SEC("kprobe.multi/skb-2")
 int kprobe_skb_2(struct pt_regs *ctx) {
 	struct sk_buff *skb = (struct sk_buff *) PT_REGS_PARM2(ctx);
 
 	return handle_everything(skb, ctx);
 }
 
-SEC("kprobe/skb-3")
+SEC("kprobe.multi/skb-3")
 int kprobe_skb_3(struct pt_regs *ctx) {
 	struct sk_buff *skb = (struct sk_buff *) PT_REGS_PARM3(ctx);
 
 	return handle_everything(skb, ctx);
 }
 
-SEC("kprobe/skb-4")
+SEC("kprobe.multi/skb-4")
 int kprobe_skb_4(struct pt_regs *ctx) {
 	struct sk_buff *skb = (struct sk_buff *) PT_REGS_PARM4(ctx);
 
 	return handle_everything(skb, ctx);
 }
 
-SEC("kprobe/skb-5")
+SEC("kprobe.multi/skb-5")
 int kprobe_skb_5(struct pt_regs *ctx) {
 	struct sk_buff *skb = (struct sk_buff *) PT_REGS_PARM5(ctx);
 

main.go

@@ -5,6 +5,7 @@
 package main
 
 import (
+	"bufio"
 	"bytes"
 	"context"
 	"encoding/binary"
@@ -13,6 +14,7 @@ import (
 	"log"
 	"os"
 	"os/signal"
+	"sort"
 	"syscall"
 
 	pb "github.com/cheggaaa/pb/v3"
@@ -79,6 +81,103 @@ func singleKprobe(funcs pwru.Funcs, progs Progs, ctx context.Context) ([]link.Li
 	return kprobes, nil
 }
 
+func multiKprobe(funcs pwru.Funcs, progs Progs, ffuncs []string, ctx context.Context) ([]link.Link, error) {
+
+	var kprobes []link.Link
+	ksyms := make(map[*ebpf.Program][]string)
+	ignored := 0
+
+	for name, pos := range funcs {
+		var p *ebpf.Program
+		switch pos {
+		case 1:
+			p = progs[1]
+			ksyms[p] = append(ksyms[p], name)
+		case 2:
+			p = progs[2]
+			ksyms[p] = append(ksyms[p], name)
+		case 3:
+			p = progs[3]
+			ksyms[p] = append(ksyms[p], name)
+		case 4:
+			p = progs[4]
+			ksyms[p] = append(ksyms[p], name)
+		case 5:
+			p = progs[5]
+			ksyms[p] = append(ksyms[p], name)
+		default:
+			ignored += 1
+			continue
+		}
+
+	}
+
+	log.Println("Attaching multi kprobes...")
+	bar := pb.StartNew(len(funcs))
+
+	for prog, syms := range ksyms {
+		// find functions exist in syms and in /sys/kernel/debug/tracing/available_filter_functions
+		s := interSection(syms, ffuncs)
+		sort.Strings(s)
+		log.Printf("Attached (prog %s ) to (funcs %s)\n", prog, s)
+
+		opts := link.KprobeMultiOptions{Symbols: s}
+
+		lnk, err := link.KprobeMulti(prog, opts)
+		if err != nil {
+			return nil, fmt.Errorf("attaching '%s' failed: %w", prog, err)
+		} else {
+			kprobes = append(kprobes, lnk)
+		}
+
+		select {
+		case <-ctx.Done():
+			return kprobes, nil
+		default:
+		}
+
+		bar.Increment()
+
+	}
+
+	bar.Finish()
+	log.Printf("Attached (ignored %d)\n", ignored)
+
+	return kprobes, nil
+}
+
+// readLines reads a whole file into memory
+// and returns a slice of its lines.
+func readLines(path string) ([]string, error) {
+	file, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	defer file.Close()
+
+	var lines []string
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		lines = append(lines, scanner.Text())
+	}
+	return lines, scanner.Err()
+}
+
+func interSection(s1, s2 []string) (inter []string) {
+	var s []string
+	hash := make(map[string]bool)
+	for _, e := range s1 {
+		hash[e] = true
+	}
+	for _, e := range s2 {
+		// If elements present in the hashmap then append intersection list.
+		if hash[e] {
+			s = append(s, e)
+		}
+	}
+	return s
+}
+
 func main() {
 	var (
 		cfgMap, events, printSkbMap, printStackMap *ebpf.Map
@@ -170,11 +269,24 @@ func main() {
 	log.Printf("Per cpu buffer size: %d bytes\n", flags.PerCPUBuffer)
 	pwru.ConfigBPFMap(&flags, cfgMap)
 
-	kprobes, err := singleKprobe(funcs, kProbeSkb, ctx)
+	var kprobes []link.Link
+
+	ffuncs, err := readLines("/sys/kernel/debug/tracing/available_filter_functions")
 	if err != nil {
-		log.Fatalf("Failed to attach kprobes: %s", err)
+		log.Fatalf("readLines: %s", err)
 	}
 
+	kprobes, err = multiKprobe(funcs, kProbeSkb, ffuncs, ctx)
+
+	if err != nil {
+		log.Printf("Failed to attach multi kprobes: %s", err)
+		/*
+			kprobes, err = singleKprobe(funcs, kProbeSkb, ctx)
+			if err != nil {
+				log.Printf("Failed to attach single kprobe: %s", err)
+			}
+		*/
+	}
 	defer func() {
 		select {
 		case <-ctx.Done():