Init basic project

Signed-off-by: Viet Anh Duong <[email protected]>

.gitignore

@@ -19,3 +19,7 @@
 
 # Go workspace file
 go.work
+
+/.vscode
+
+.DS_Store

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

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=

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

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[:]
+	}
+}

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

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

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 }