From b2527904e4fc3680afb0ee81bb2d39f31c6dbc8e Mon Sep 17 00:00:00 2001 From: Viet Anh Duong Date: Mon, 18 Dec 2023 14:48:08 +0000 Subject: [PATCH] Init basic project Signed-off-by: Viet Anh Duong --- .gitignore | 4 ++ go.mod | 10 +++++ go.sum | 16 +++++++ module.go | 112 ++++++++++++++++++++++++++++++++++++++++++++++ module_options.go | 20 +++++++++ module_utils.go | 46 +++++++++++++++++++ stack_table.go | 40 +++++++++++++++++ table.go | 23 ++++++++++ 8 files changed, 271 insertions(+) create mode 100644 go.mod create mode 100644 go.sum create mode 100644 module.go create mode 100644 module_options.go create mode 100644 module_utils.go create mode 100644 stack_table.go create mode 100644 table.go diff --git a/.gitignore b/.gitignore index 3b735ec..40e19fd 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,7 @@ # Go workspace file go.work + +/.vscode + +.DS_Store diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..b1dd7a1 --- /dev/null +++ b/go.mod @@ -0,0 +1,10 @@ +module github.com/vietanhduong/wbpf + +go 1.21.4 + +require github.com/cilium/ebpf v0.12.3 + +require ( + golang.org/x/exp v0.0.0-20230224173230-c95f2b4c22f2 // indirect + golang.org/x/sys v0.14.1-0.20231108175955-e4099bfacb8c // indirect +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..44cd9bb --- /dev/null +++ b/go.sum @@ -0,0 +1,16 @@ +github.com/cilium/ebpf v0.12.3 h1:8ht6F9MquybnY97at+VDZb3eQQr8ev79RueWeVaEcG4= +github.com/cilium/ebpf v0.12.3/go.mod h1:TctK1ivibvI3znr66ljgi4hqOT8EYQjz1KWBfb1UVgM= +github.com/frankban/quicktest v1.14.5 h1:dfYrrRyLtiqT9GyKXgdh+k4inNeTvmGbuSgZ3lx3GhA= +github.com/frankban/quicktest v1.14.5/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= +github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38= +github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= +github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= +golang.org/x/exp v0.0.0-20230224173230-c95f2b4c22f2 h1:Jvc7gsqn21cJHCmAWx0LiimpP18LZmUxkT5Mp7EZ1mI= +golang.org/x/exp v0.0.0-20230224173230-c95f2b4c22f2/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc= +golang.org/x/sys v0.14.1-0.20231108175955-e4099bfacb8c h1:3kC/TjQ+xzIblQv39bCOyRk8fbEeJcDHwbyxPUU2BpA= +golang.org/x/sys v0.14.1-0.20231108175955-e4099bfacb8c/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= diff --git a/module.go b/module.go new file mode 100644 index 0000000..8a7ca67 --- /dev/null +++ b/module.go @@ -0,0 +1,112 @@ +package wbpf + +import ( + "bytes" + "fmt" + "os" + "runtime" + + "github.com/cilium/ebpf" + "github.com/cilium/ebpf/link" +) + +var ErrProgNotFound = fmt.Errorf("prog not found") + +type Module struct { + collection *ebpf.Collection + + kprobes map[string]link.Link +} + +func NewModule(opts ...ModuleOption) (*Module, error) { + var modOpts moduleOptions + for _, opt := range opts { + opt(&modOpts) + } + + if (modOpts.file == "" && len(modOpts.content) == 0) || + (modOpts.file != "" && len(modOpts.content) != 0) { + return nil, fmt.Errorf("only one of file or content must be specified") + } + mod, err := newModule(&modOpts) + if err != nil { + return nil, err + } + runtime.SetFinalizer(mod, func(m *Module) { m.Close() }) + return mod, nil +} + +func (m *Module) GetTable(name string) (*Table, error) { + tbl, ok := m.collection.Maps[name] + if !ok || tbl == nil { + return nil, ErrTableNotFound + } + info, err := tbl.Info() + if err != nil { + return nil, fmt.Errorf("map info: %w", err) + } + return &Table{Map: tbl, info: info, mod: m}, nil +} + +func (m *Module) AttackKprobe(sysname, prog string) error { + sysname = GetSyscallName(sysname) + if _, ok := m.kprobes[sysname]; ok { + return nil + } + + p, ok := m.collection.Programs[prog] + if !ok || p == nil { + return ErrProgNotFound + } + kprobe, err := link.Kprobe(sysname, p, nil) + if err != nil { + return fmt.Errorf("link kprobe (%s): %w", sysname, err) + } + m.kprobes[sysname] = kprobe + return nil +} + +func (m *Module) DetachKprobe(sysname string) { + sysname = GetSyscallName(sysname) + if kprobe, ok := m.kprobes[sysname]; ok { + kprobe.Close() + delete(m.kprobes, sysname) + } +} + +func (m *Module) Close() { + if m == nil { + return + } + // Close collection after all probes have been closed + if m.collection != nil { + defer m.collection.Close() + } + // Detach Kprobes + for name, l := range m.kprobes { + l.Close() + delete(m.kprobes, name) + } +} + +func newModule(opts *moduleOptions) (*Module, error) { + if opts.file != "" { + buf, err := os.ReadFile(opts.file) + if err != nil { + return nil, fmt.Errorf("os read file: %w", err) + } + opts.content = buf + } + spec, err := ebpf.LoadCollectionSpecFromReader(bytes.NewReader(opts.content)) + if err != nil { + return nil, fmt.Errorf("ebpf load collection spec: %w", err) + } + + mod := &Module{ + kprobes: make(map[string]link.Link), + } + if mod.collection, err = ebpf.NewCollection(spec); err != nil { + return nil, fmt.Errorf("ebpf new collection: %w", err) + } + return mod, nil +} diff --git a/module_options.go b/module_options.go new file mode 100644 index 0000000..9810f7a --- /dev/null +++ b/module_options.go @@ -0,0 +1,20 @@ +package wbpf + +type moduleOptions struct { + file string + content []byte +} + +type ModuleOption func(*moduleOptions) + +func WithElfFile(path string) ModuleOption { + return func(mo *moduleOptions) { + mo.file = path + } +} + +func WithElfFileContent(content []byte) ModuleOption { + return func(mo *moduleOptions) { + mo.content = content[:] + } +} diff --git a/module_utils.go b/module_utils.go new file mode 100644 index 0000000..8cd484e --- /dev/null +++ b/module_utils.go @@ -0,0 +1,46 @@ +package wbpf + +import ( + "runtime" + "strings" +) + +func GetSyscallPrefix() string { + switch runtime.GOARCH { + case "amd64": + return "__x64_sys_" + case "arm64": + return "__arm64_sys_" + case "s390x": + return "__s390x_sys_" + case "s390": + return "__s390_sys_" + } + return "sys_" +} + +func GetSyscallName(name string) string { + if strings.HasPrefix(name, GetSyscallPrefix()) { + return name + } + return GetSyscallPrefix() + name +} + +var syscallPrefixes = []string{ + "sys_", + "__x64_sys_", + "__x32_compat_sys_", + "__ia32_compat_sys_", + "__arm64_sys_", + "__s390x_sys_", + "__s390_sys_", +} + +func FixSyscallName(name string) string { + for _, p := range syscallPrefixes { + if strings.HasPrefix(name, p) { + return GetSyscallName(name[len(p):]) + } + } + return name +} diff --git a/stack_table.go b/stack_table.go new file mode 100644 index 0000000..b45d8e0 --- /dev/null +++ b/stack_table.go @@ -0,0 +1,40 @@ +package wbpf + +import ( + "fmt" + "unsafe" + + "github.com/cilium/ebpf" +) + +const MAX_STACK_DEPTH = 127 + +type stacktrace struct{ insptr [MAX_STACK_DEPTH]uint64 } + +func (t *Table) GetStackAddr(stackId int, clear bool) ([]uint64, error) { + if t.TableType() != ebpf.StackTrace { + return nil, ErrIncorrectTableType + } + + if stackId < 0 { + return nil, nil + } + + var b []byte + var err error + if b, err = t.LookupBytes(uint32(stackId)); len(b) == 0 { + return nil, fmt.Errorf("lookup key 0x%08x: %w", stackId, err) + } + + stack := (*stacktrace)(unsafe.Pointer(&b[0])) + var addrs []uint64 + for i := 0; i < MAX_STACK_DEPTH && stack.insptr[i] != 0; i++ { + addrs = append(addrs, stack.insptr[i]) + } + if clear { + if err := t.Delete(uint32(stackId)); err != nil { + return nil, fmt.Errorf("delete key 0x%08x: %w", stackId, err) + } + } + return addrs, nil +} diff --git a/table.go b/table.go new file mode 100644 index 0000000..eaafcc4 --- /dev/null +++ b/table.go @@ -0,0 +1,23 @@ +package wbpf + +import ( + "fmt" + + "github.com/cilium/ebpf" +) + +var ( + ErrTableNotFound = fmt.Errorf("table not found") + ErrIncorrectTableType = fmt.Errorf("incorrect table type") +) + +type Table struct { + *ebpf.Map + info *ebpf.MapInfo + + mod *Module +} + +func (t *Table) TableType() ebpf.MapType { return t.info.Type } + +func (t *Table) TableName() string { return t.info.Name }