forked from mirrors/gotosocial
170 lines
3.6 KiB
Go
170 lines
3.6 KiB
Go
|
package link
|
||
|
|
||
|
import (
|
||
|
"errors"
|
||
|
"fmt"
|
||
|
"os"
|
||
|
|
||
|
"github.com/cilium/ebpf"
|
||
|
)
|
||
|
|
||
|
type cgroupAttachFlags uint32
|
||
|
|
||
|
// cgroup attach flags
|
||
|
const (
|
||
|
flagAllowOverride cgroupAttachFlags = 1 << iota
|
||
|
flagAllowMulti
|
||
|
flagReplace
|
||
|
)
|
||
|
|
||
|
type CgroupOptions struct {
|
||
|
// Path to a cgroupv2 folder.
|
||
|
Path string
|
||
|
// One of the AttachCgroup* constants
|
||
|
Attach ebpf.AttachType
|
||
|
// Program must be of type CGroup*, and the attach type must match Attach.
|
||
|
Program *ebpf.Program
|
||
|
}
|
||
|
|
||
|
// AttachCgroup links a BPF program to a cgroup.
|
||
|
func AttachCgroup(opts CgroupOptions) (Link, error) {
|
||
|
cgroup, err := os.Open(opts.Path)
|
||
|
if err != nil {
|
||
|
return nil, fmt.Errorf("can't open cgroup: %s", err)
|
||
|
}
|
||
|
|
||
|
clone, err := opts.Program.Clone()
|
||
|
if err != nil {
|
||
|
cgroup.Close()
|
||
|
return nil, err
|
||
|
}
|
||
|
|
||
|
var cg Link
|
||
|
cg, err = newLinkCgroup(cgroup, opts.Attach, clone)
|
||
|
if errors.Is(err, ErrNotSupported) {
|
||
|
cg, err = newProgAttachCgroup(cgroup, opts.Attach, clone, flagAllowMulti)
|
||
|
}
|
||
|
if errors.Is(err, ErrNotSupported) {
|
||
|
cg, err = newProgAttachCgroup(cgroup, opts.Attach, clone, flagAllowOverride)
|
||
|
}
|
||
|
if err != nil {
|
||
|
cgroup.Close()
|
||
|
clone.Close()
|
||
|
return nil, err
|
||
|
}
|
||
|
|
||
|
return cg, nil
|
||
|
}
|
||
|
|
||
|
// LoadPinnedCgroup loads a pinned cgroup from a bpffs.
|
||
|
func LoadPinnedCgroup(fileName string) (Link, error) {
|
||
|
link, err := LoadPinnedRawLink(fileName)
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
|
||
|
return &linkCgroup{link}, nil
|
||
|
}
|
||
|
|
||
|
type progAttachCgroup struct {
|
||
|
cgroup *os.File
|
||
|
current *ebpf.Program
|
||
|
attachType ebpf.AttachType
|
||
|
flags cgroupAttachFlags
|
||
|
}
|
||
|
|
||
|
var _ Link = (*progAttachCgroup)(nil)
|
||
|
|
||
|
func (cg *progAttachCgroup) isLink() {}
|
||
|
|
||
|
func newProgAttachCgroup(cgroup *os.File, attach ebpf.AttachType, prog *ebpf.Program, flags cgroupAttachFlags) (*progAttachCgroup, error) {
|
||
|
if flags&flagAllowMulti > 0 {
|
||
|
if err := haveProgAttachReplace(); err != nil {
|
||
|
return nil, fmt.Errorf("can't support multiple programs: %w", err)
|
||
|
}
|
||
|
}
|
||
|
|
||
|
err := RawAttachProgram(RawAttachProgramOptions{
|
||
|
Target: int(cgroup.Fd()),
|
||
|
Program: prog,
|
||
|
Flags: uint32(flags),
|
||
|
Attach: attach,
|
||
|
})
|
||
|
if err != nil {
|
||
|
return nil, fmt.Errorf("cgroup: %w", err)
|
||
|
}
|
||
|
|
||
|
return &progAttachCgroup{cgroup, prog, attach, flags}, nil
|
||
|
}
|
||
|
|
||
|
func (cg *progAttachCgroup) Close() error {
|
||
|
defer cg.cgroup.Close()
|
||
|
defer cg.current.Close()
|
||
|
|
||
|
err := RawDetachProgram(RawDetachProgramOptions{
|
||
|
Target: int(cg.cgroup.Fd()),
|
||
|
Program: cg.current,
|
||
|
Attach: cg.attachType,
|
||
|
})
|
||
|
if err != nil {
|
||
|
return fmt.Errorf("close cgroup: %s", err)
|
||
|
}
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
func (cg *progAttachCgroup) Update(prog *ebpf.Program) error {
|
||
|
new, err := prog.Clone()
|
||
|
if err != nil {
|
||
|
return err
|
||
|
}
|
||
|
|
||
|
args := RawAttachProgramOptions{
|
||
|
Target: int(cg.cgroup.Fd()),
|
||
|
Program: prog,
|
||
|
Attach: cg.attachType,
|
||
|
Flags: uint32(cg.flags),
|
||
|
}
|
||
|
|
||
|
if cg.flags&flagAllowMulti > 0 {
|
||
|
// Atomically replacing multiple programs requires at least
|
||
|
// 5.5 (commit 7dd68b3279f17921 "bpf: Support replacing cgroup-bpf
|
||
|
// program in MULTI mode")
|
||
|
args.Flags |= uint32(flagReplace)
|
||
|
args.Replace = cg.current
|
||
|
}
|
||
|
|
||
|
if err := RawAttachProgram(args); err != nil {
|
||
|
new.Close()
|
||
|
return fmt.Errorf("can't update cgroup: %s", err)
|
||
|
}
|
||
|
|
||
|
cg.current.Close()
|
||
|
cg.current = new
|
||
|
return nil
|
||
|
}
|
||
|
|
||
|
func (cg *progAttachCgroup) Pin(string) error {
|
||
|
return fmt.Errorf("can't pin cgroup: %w", ErrNotSupported)
|
||
|
}
|
||
|
|
||
|
type linkCgroup struct {
|
||
|
*RawLink
|
||
|
}
|
||
|
|
||
|
var _ Link = (*linkCgroup)(nil)
|
||
|
|
||
|
func (cg *linkCgroup) isLink() {}
|
||
|
|
||
|
func newLinkCgroup(cgroup *os.File, attach ebpf.AttachType, prog *ebpf.Program) (*linkCgroup, error) {
|
||
|
link, err := AttachRawLink(RawLinkOptions{
|
||
|
Target: int(cgroup.Fd()),
|
||
|
Program: prog,
|
||
|
Attach: attach,
|
||
|
})
|
||
|
if err != nil {
|
||
|
return nil, err
|
||
|
}
|
||
|
|
||
|
return &linkCgroup{link}, err
|
||
|
}
|