mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-09-17 23:22:27 +02:00
feat: manage task process lifetimes and preserve turn history
This commit is contained in:
Ed1s0nZ
@@ -0,0 +1,313 @@
|
||||
//go:build linux
|
||||
|
||||
package processguard
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"golang.org/x/sys/unix"
|
||||
)
|
||||
|
||||
var rootLock *os.File // retained until server exit; never inherited by commands
|
||||
|
||||
func configurePlatform(o *Options) error {
|
||||
if o.CgroupRoot == "" {
|
||||
if o.Mode == "required" {
|
||||
return fmt.Errorf("required isolation needs security.process_isolation.cgroup_root")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if o.Mode == "process_group" {
|
||||
return fmt.Errorf("cgroup_root cannot be combined with process_group mode")
|
||||
}
|
||||
if o.CgroupRoot == "auto" {
|
||||
data, err := os.ReadFile("/proc/self/cgroup")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, line := range strings.Split(string(data), "\n") {
|
||||
if strings.HasPrefix(line, "0::") {
|
||||
o.CgroupRoot = filepath.Join("/sys/fs/cgroup", strings.TrimPrefix(line, "0::"))
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
root, err := validateRoot(o.CgroupRoot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
o.CgroupRoot = root
|
||||
// An exclusive host-side lock prevents one server's recovery sweep from
|
||||
// killing tasks owned by another server using the same delegated root.
|
||||
hash := sha256.Sum256([]byte(root))
|
||||
lockPath := filepath.Join(os.TempDir(), fmt.Sprintf("cyberstrike-cgroup-%d-%x.lock", os.Getuid(), hash[:12]))
|
||||
fd, err := unix.Open(lockPath, unix.O_CREAT|unix.O_RDWR|unix.O_CLOEXEC|unix.O_NOFOLLOW, 0600)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
lock := os.NewFile(uintptr(fd), lockPath)
|
||||
if err = unix.Flock(fd, unix.LOCK_EX|unix.LOCK_NB); err != nil {
|
||||
lock.Close()
|
||||
return fmt.Errorf("cgroup root is already owned: %w", err)
|
||||
}
|
||||
success := false
|
||||
defer func() {
|
||||
if !success {
|
||||
lock.Close()
|
||||
}
|
||||
}()
|
||||
// cgroup v2 requires the delegated parent to have no processes before
|
||||
// domain controllers can be enabled. Move only this server, never outsiders.
|
||||
data, err := os.ReadFile(filepath.Join(root, "cgroup.procs"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, pid := range strings.Fields(string(data)) {
|
||||
if pid != strconv.Itoa(os.Getpid()) {
|
||||
return fmt.Errorf("delegated root contains another process %s", pid)
|
||||
}
|
||||
}
|
||||
if len(strings.Fields(string(data))) > 0 {
|
||||
supervisor := filepath.Join(root, "supervisor")
|
||||
if err = os.Mkdir(supervisor, 0700); err != nil && !os.IsExist(err) {
|
||||
return err
|
||||
}
|
||||
if err = os.WriteFile(filepath.Join(supervisor, "cgroup.procs"), []byte(strconv.Itoa(os.Getpid())), 0600); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err = os.WriteFile(filepath.Join(root, "cgroup.subtree_control"), []byte("+cpu +memory +pids"), 0600); err != nil {
|
||||
return fmt.Errorf("delegate cpu, memory and pids controllers: %w", err)
|
||||
}
|
||||
// Recover only our names under the exclusively owned root. No PID replay.
|
||||
entries, err := os.ReadDir(root)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() && validTaskName(entry.Name()) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
err = killAndRemoveCgroup(ctx, filepath.Join(root, entry.Name()))
|
||||
cancel()
|
||||
if err != nil {
|
||||
return fmt.Errorf("recover %s: %w", entry.Name(), err)
|
||||
}
|
||||
}
|
||||
}
|
||||
rootLock = lock
|
||||
success = true
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateRoot(root string) (string, error) {
|
||||
if !filepath.IsAbs(root) {
|
||||
return "", fmt.Errorf("cgroup root must be absolute")
|
||||
}
|
||||
root = filepath.Clean(root)
|
||||
resolved, err := filepath.EvalSymlinks(root)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if root != resolved || root == "/sys/fs/cgroup" || root == "/" {
|
||||
return "", fmt.Errorf("use a dedicated delegated cgroup, not the hierarchy root or a symlink")
|
||||
}
|
||||
var st unix.Statfs_t
|
||||
if err = unix.Statfs(root, &st); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if st.Type != unix.CGROUP2_SUPER_MAGIC {
|
||||
return "", fmt.Errorf("%s is not cgroup v2", root)
|
||||
}
|
||||
return root, nil
|
||||
}
|
||||
func validTaskName(name string) bool {
|
||||
if !strings.HasPrefix(name, "task-") || len(name) != 41 {
|
||||
return false
|
||||
}
|
||||
for _, c := range name[5:] {
|
||||
if !(c >= '0' && c <= '9' || c >= 'a' && c <= 'f' || c == '-') {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
type cgroupGroup struct {
|
||||
mu sync.Mutex
|
||||
path string
|
||||
dir *os.File
|
||||
watcher *watchdog
|
||||
closed bool
|
||||
}
|
||||
|
||||
func newPlatformGroup(id string, o Options) (Group, error) {
|
||||
if o.CgroupRoot == "" {
|
||||
if o.Mode == "required" {
|
||||
return nil, fmt.Errorf("required isolation has no delegated cgroup root")
|
||||
}
|
||||
return newUnixGroup()
|
||||
}
|
||||
root, err := validateRoot(o.CgroupRoot)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
name := "task-" + id
|
||||
if !validTaskName(name) {
|
||||
return nil, fmt.Errorf("invalid task run ID")
|
||||
}
|
||||
path := filepath.Join(root, name)
|
||||
if err = os.Mkdir(path, 0700); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
success := false
|
||||
defer func() {
|
||||
if !success {
|
||||
_ = os.Remove(path)
|
||||
}
|
||||
}()
|
||||
limits := map[string]string{"pids.max": strconv.Itoa(o.MaxProcesses), "memory.max": strconv.FormatInt(o.MemoryMaxBytes, 10), "memory.oom.group": "1"}
|
||||
if o.CPUQuotaMicros > 0 {
|
||||
limits["cpu.max"] = fmt.Sprintf("%d 100000", o.CPUQuotaMicros)
|
||||
}
|
||||
for file, value := range limits {
|
||||
if err = os.WriteFile(filepath.Join(path, file), []byte(value), 0600); err != nil {
|
||||
return nil, fmt.Errorf("set %s: %w", file, err)
|
||||
}
|
||||
}
|
||||
if _, err = os.Stat(filepath.Join(path, "cgroup.kill")); err != nil {
|
||||
return nil, fmt.Errorf("cgroup.kill requires Linux 5.14+: %w", err)
|
||||
}
|
||||
dir, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
g := &cgroupGroup{path: path, dir: dir}
|
||||
w, err := startWatchdog(watchRequest{Name: "cgroup", Path: path}, func() {
|
||||
// A guardian crash is also fail-closed while the owner is still alive.
|
||||
_ = os.WriteFile(filepath.Join(path, "cgroup.kill"), []byte("1"), 0600)
|
||||
})
|
||||
if err != nil {
|
||||
dir.Close()
|
||||
return nil, err
|
||||
}
|
||||
g.watcher = w
|
||||
success = true
|
||||
return g, nil
|
||||
}
|
||||
func (g *cgroupGroup) Name() string { return "cgroup_v2" }
|
||||
func (g *cgroupGroup) Prepare(cmd *exec.Cmd) (*Launch, error) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
if g.closed {
|
||||
return nil, fmt.Errorf("cgroup is closed")
|
||||
}
|
||||
if _, err := g.watcher.send(watchRequest{Op: "ping"}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if cmd.SysProcAttr == nil {
|
||||
cmd.SysProcAttr = &syscall.SysProcAttr{}
|
||||
}
|
||||
// clone3(CLONE_INTO_CGROUP), not a racy write of a newly started PID.
|
||||
cmd.SysProcAttr.UseCgroupFD = true
|
||||
cmd.SysProcAttr.CgroupFD = int(g.dir.Fd())
|
||||
return &Launch{Commit: func() error { return nil }, Dispose: func() {}}, nil
|
||||
}
|
||||
func (g *cgroupGroup) Release(pid int) error { return nil }
|
||||
func (g *cgroupGroup) Close(ctx context.Context) error {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
g.closed = true
|
||||
if g.dir == nil {
|
||||
return nil
|
||||
}
|
||||
if err := killAndRemoveCgroup(ctx, g.path); err != nil {
|
||||
return err
|
||||
}
|
||||
watchErr := g.watcher.close()
|
||||
err := errors.Join(watchErr, g.dir.Close())
|
||||
g.dir = nil
|
||||
return err
|
||||
}
|
||||
func killAndRemoveCgroup(ctx context.Context, path string) error {
|
||||
if err := os.WriteFile(filepath.Join(path, "cgroup.kill"), []byte("1"), 0600); err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
for {
|
||||
data, err := os.ReadFile(filepath.Join(path, "cgroup.events"))
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if strings.Contains(string(data), "populated 0") {
|
||||
return removeCgroupTree(path)
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(10 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
}
|
||||
func removeCgroupTree(path string) error {
|
||||
entries, err := os.ReadDir(path)
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
for _, e := range entries {
|
||||
if e.IsDir() {
|
||||
if err = removeCgroupTree(filepath.Join(path, e.Name())); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
err = os.Remove(path)
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
func guardianMain(dec *json.Decoder, enc *json.Encoder) error {
|
||||
var req watchRequest
|
||||
if err := dec.Decode(&req); err != nil {
|
||||
return err
|
||||
}
|
||||
if req.Name == "process_group" {
|
||||
return groupGuardian(dec, enc)
|
||||
}
|
||||
if req.Name != "cgroup" || !validTaskName(filepath.Base(req.Path)) {
|
||||
return fmt.Errorf("invalid cgroup guardian")
|
||||
}
|
||||
if _, err := validateRoot(req.Path); err != nil {
|
||||
return err
|
||||
}
|
||||
return serveGuardian(dec, enc, func(r watchRequest) error {
|
||||
if r.Op != "ping" {
|
||||
return fmt.Errorf("unknown command")
|
||||
}
|
||||
return nil
|
||||
}, func() error {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
return killAndRemoveCgroup(ctx, req.Path)
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user