feat: add session-scoped process isolation
ProcessScope tracks PIDs per session key for namespace-like isolation. Agents can only see and kill their own processes. - Register/Deregister/Owns for session-PID binding - ListPIDs filters dead processes via Unix signal 0 - KillAll sends SIGTERM to all session-owned processes - Cleanup removes all tracking for a session
This commit is contained in:
parent
0b7d3b9198
commit
b3d5dc9e1b
1 changed files with 141 additions and 0 deletions
141
pkg/tools/process_scope.go
Normal file
141
pkg/tools/process_scope.go
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"os"
|
||||
"sync"
|
||||
"syscall"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
)
|
||||
|
||||
// ProcessScope tracks PIDs per session key, providing namespace-like isolation
|
||||
// for exec tool processes. Inspired by Google Kubernetes pod process isolation
|
||||
// and Linux cgroups — agents can only see and kill their own processes.
|
||||
type ProcessScope struct {
|
||||
pids sync.Map // sessionKey -> *pidSet
|
||||
}
|
||||
|
||||
type pidSet struct {
|
||||
mu sync.Mutex
|
||||
pids map[int]bool
|
||||
}
|
||||
|
||||
// NewProcessScope creates a new process scope tracker.
|
||||
func NewProcessScope() *ProcessScope {
|
||||
return &ProcessScope{}
|
||||
}
|
||||
|
||||
// Register adds a PID to a session's scope.
|
||||
func (ps *ProcessScope) Register(sessionKey string, pid int) {
|
||||
set := ps.getOrCreate(sessionKey)
|
||||
set.mu.Lock()
|
||||
defer set.mu.Unlock()
|
||||
set.pids[pid] = true
|
||||
}
|
||||
|
||||
// Deregister removes a PID from a session's scope (process exited normally).
|
||||
func (ps *ProcessScope) Deregister(sessionKey string, pid int) {
|
||||
v, ok := ps.pids.Load(sessionKey)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
set := v.(*pidSet)
|
||||
set.mu.Lock()
|
||||
defer set.mu.Unlock()
|
||||
delete(set.pids, pid)
|
||||
}
|
||||
|
||||
// Owns returns true if the given PID belongs to the session's scope.
|
||||
func (ps *ProcessScope) Owns(sessionKey string, pid int) bool {
|
||||
v, ok := ps.pids.Load(sessionKey)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
set := v.(*pidSet)
|
||||
set.mu.Lock()
|
||||
defer set.mu.Unlock()
|
||||
return set.pids[pid]
|
||||
}
|
||||
|
||||
// ListPIDs returns all live PIDs for a session (filters out already-exited processes).
|
||||
func (ps *ProcessScope) ListPIDs(sessionKey string) []int {
|
||||
v, ok := ps.pids.Load(sessionKey)
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
set := v.(*pidSet)
|
||||
set.mu.Lock()
|
||||
defer set.mu.Unlock()
|
||||
|
||||
var live []int
|
||||
for pid := range set.pids {
|
||||
// Check if process is still running (Unix signal 0).
|
||||
if isProcessAlive(pid) {
|
||||
live = append(live, pid)
|
||||
} else {
|
||||
delete(set.pids, pid)
|
||||
}
|
||||
}
|
||||
return live
|
||||
}
|
||||
|
||||
// KillAll kills all processes owned by a session. Returns number killed.
|
||||
// Used during cascade stop to clean up spawned processes.
|
||||
func (ps *ProcessScope) KillAll(sessionKey string) int {
|
||||
v, ok := ps.pids.Load(sessionKey)
|
||||
if !ok {
|
||||
return 0
|
||||
}
|
||||
set := v.(*pidSet)
|
||||
set.mu.Lock()
|
||||
defer set.mu.Unlock()
|
||||
|
||||
killed := 0
|
||||
for pid := range set.pids {
|
||||
if err := killProcess(pid); err == nil {
|
||||
killed++
|
||||
}
|
||||
delete(set.pids, pid)
|
||||
}
|
||||
|
||||
if killed > 0 {
|
||||
logger.InfoCF("process_scope", "Killed session processes", map[string]interface{}{
|
||||
"session_key": sessionKey,
|
||||
"killed": killed,
|
||||
})
|
||||
}
|
||||
return killed
|
||||
}
|
||||
|
||||
// Cleanup removes all tracking for a session.
|
||||
func (ps *ProcessScope) Cleanup(sessionKey string) {
|
||||
ps.pids.Delete(sessionKey)
|
||||
}
|
||||
|
||||
func (ps *ProcessScope) getOrCreate(sessionKey string) *pidSet {
|
||||
if v, ok := ps.pids.Load(sessionKey); ok {
|
||||
return v.(*pidSet)
|
||||
}
|
||||
set := &pidSet{pids: make(map[int]bool)}
|
||||
actual, _ := ps.pids.LoadOrStore(sessionKey, set)
|
||||
return actual.(*pidSet)
|
||||
}
|
||||
|
||||
// isProcessAlive checks if a process is still running via signal 0 (Unix).
|
||||
func isProcessAlive(pid int) bool {
|
||||
proc, err := os.FindProcess(pid)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
err = proc.Signal(syscall.Signal(0))
|
||||
return err == nil
|
||||
}
|
||||
|
||||
// killProcess sends SIGTERM to a process. Falls back to SIGKILL if needed.
|
||||
func killProcess(pid int) error {
|
||||
proc, err := os.FindProcess(pid)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return proc.Signal(syscall.SIGTERM)
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue