Add multi kprobe attach

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

main.go

@@ -79,6 +79,67 @@ func singleKprobe(funcs pwru.Funcs, progs Progs, ctx context.Context) ([]link.Li
 	return kprobes, nil
 }
 
+func multiKprobe(funcs pwru.Funcs, progs Progs, 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 kprobes...")
+	bar := pb.StartNew(len(funcs))
+
+	for prog, syms := range ksyms {
+
+		opts := link.KprobeMultiOptions{Symbols: syms}
+
+		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
+}
+
 func main() {
 	var (
 		cfgMap, events, printSkbMap, printStackMap *ebpf.Map
@@ -170,11 +231,16 @@ 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
+
+	kprobes, err = multiKprobe(funcs, kProbeSkb, ctx)
 	if err != nil {
-		log.Fatalf("Failed to attach kprobes: %s", err)
+		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():