From 1689f8c5a6c158a9828e9d894f9e10b296716f1c Mon Sep 17 00:00:00 2001 From: Vincent Li Date: Mon, 26 Sep 2022 03:53:37 +0000 Subject: [PATCH] Add multi kprobe attach Signed-off-by: Vincent Li --- bpf/kprobe_pwru.c | 10 ++--- main.go | 112 +++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 115 insertions(+), 7 deletions(-) diff --git a/bpf/kprobe_pwru.c b/bpf/kprobe_pwru.c index 73e54631..1639fedb 100644 --- a/bpf/kprobe_pwru.c +++ b/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); diff --git a/main.go b/main.go index 3115d5bc..c8806ddc 100644 --- a/main.go +++ b/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,99 @@ 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...") + + for prog, syms := range ksyms { + // find functions exist in syms and in /sys/kernel/debug/tracing/available_filter_functions + s := interSection(ffuncs, syms) + sort.Strings(s) + log.Printf("Attaching (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: + } + + } + + 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() +} + +// took from https://go.dev/play/p/eGGcyIlZD6y, the code is bit unclear to me :) +func interSection(ffuncs, syms []string) (inter []string) { + out := []string{} + bucket := map[string]bool{} + for _, i := range ffuncs { + for _, j := range syms { + if i == j && !bucket[i] { + out = append(out, i) + bucket[i] = true + } + } + } + return out +} + func main() { var ( cfgMap, events, printSkbMap, printStackMap *ebpf.Map @@ -170,11 +265,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():