execute: add process detach mechanism
Allows the user to detach a running process from the agent mid-execution. The process continues running; output flows to a 64KB ring buffer. Detached processes can be listed, signaled (TERM/KILL), and dismissed. - DetachedProcess struct with ring buffer, signal, output methods - Server.Detach() signals the currently executing process to detach - Execute() rewritten: cmd.Start() + select on wait/ctx/detach - cleanupDetached() on Close() sends SIGTERM to all - OnDetach/OnExit hooks for integration with D-Bus layer
This commit is contained in:
Levi Neely
parent
aca0250b65
commit
16d057741c
|
|
@ -7,6 +7,7 @@ import (
|
|||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
|
||||
"ollie/internal/sandbox"
|
||||
"ollie/pkg/paths"
|
||||
|
|
@ -158,6 +159,7 @@ func loadSandboxConfig(name string) (*sandbox.Config, error) {
|
|||
}
|
||||
|
||||
type limitedWriter struct {
|
||||
mu sync.Mutex
|
||||
w io.Writer
|
||||
written int
|
||||
limit int
|
||||
|
|
@ -165,6 +167,8 @@ type limitedWriter struct {
|
|||
}
|
||||
|
||||
func (lw *limitedWriter) Write(p []byte) (n int, err error) {
|
||||
lw.mu.Lock()
|
||||
defer lw.mu.Unlock()
|
||||
if lw.written >= lw.limit {
|
||||
lw.truncated = true
|
||||
return len(p), nil
|
||||
|
|
|
|||
|
|
@ -0,0 +1,96 @@
|
|||
package execute
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
)
|
||||
|
||||
const ringBufSize = 64 * 1024 // 64KB ring buffer per detached process
|
||||
|
||||
// DetachedProcess represents a process that the agent has detached from
|
||||
// but which continues running. The user can view its output and signal it.
|
||||
type DetachedProcess struct {
|
||||
PID int
|
||||
Command string
|
||||
Started time.Time
|
||||
Exited bool
|
||||
ExitCode int
|
||||
|
||||
ring *ringBuffer
|
||||
cmd *os.Process
|
||||
done chan struct{}
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// Output returns the current contents of the ring buffer.
|
||||
func (p *DetachedProcess) Output() string {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
return p.ring.String()
|
||||
}
|
||||
|
||||
// Signal sends a signal to the detached process.
|
||||
func (p *DetachedProcess) Signal(sig syscall.Signal) error {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if p.Exited {
|
||||
return fmt.Errorf("process %d already exited", p.PID)
|
||||
}
|
||||
if p.cmd == nil {
|
||||
return fmt.Errorf("process %d has no handle", p.PID)
|
||||
}
|
||||
// Send to process group
|
||||
return syscall.Kill(-p.PID, sig)
|
||||
}
|
||||
|
||||
// ringBuffer is a fixed-size circular byte buffer.
|
||||
type ringBuffer struct {
|
||||
buf []byte
|
||||
size int
|
||||
pos int
|
||||
full bool
|
||||
}
|
||||
|
||||
func newRingBuffer(size int) *ringBuffer {
|
||||
return &ringBuffer{buf: make([]byte, size), size: size}
|
||||
}
|
||||
|
||||
// Write implements io.Writer.
|
||||
func (r *ringBuffer) Write(p []byte) (int, error) {
|
||||
n := len(p)
|
||||
if n >= r.size {
|
||||
// Data larger than buffer: just keep the tail
|
||||
copy(r.buf, p[n-r.size:])
|
||||
r.pos = 0
|
||||
r.full = true
|
||||
return n, nil
|
||||
}
|
||||
if r.pos+n <= r.size {
|
||||
copy(r.buf[r.pos:], p)
|
||||
} else {
|
||||
first := r.size - r.pos
|
||||
copy(r.buf[r.pos:], p[:first])
|
||||
copy(r.buf, p[first:])
|
||||
r.full = true
|
||||
}
|
||||
r.pos = (r.pos + n) % r.size
|
||||
if r.pos == 0 && n > 0 {
|
||||
r.full = true
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// String returns the buffer contents in order.
|
||||
func (r *ringBuffer) String() string {
|
||||
if !r.full {
|
||||
return string(r.buf[:r.pos])
|
||||
}
|
||||
// Buffer has wrapped: data from pos..end + 0..pos
|
||||
out := make([]byte, r.size)
|
||||
copy(out, r.buf[r.pos:])
|
||||
copy(out[r.size-r.pos:], r.buf[:r.pos])
|
||||
return string(out)
|
||||
}
|
||||
|
|
@ -7,6 +7,7 @@ import (
|
|||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestMain points OLLIE_CFG_PATH at testdata so loadSandboxConfig uses the
|
||||
|
|
@ -515,3 +516,61 @@ func TestRunReadBatch_WithLockDir(t *testing.T) {
|
|||
t.Errorf("output = %q; want 'locked'", out)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- detach ----
|
||||
|
||||
func TestDetach(t *testing.T) {
|
||||
s := newServer(t)
|
||||
ctx := context.Background()
|
||||
|
||||
// Start a long-running process, then detach it
|
||||
var detachedPID int
|
||||
s.OnDetach = func(pid int, cmd string) {
|
||||
detachedPID = pid
|
||||
}
|
||||
|
||||
done := make(chan string, 1)
|
||||
go func() {
|
||||
out, _ := s.Execute(ctx, "echo started; sleep 30", "bash", 60, "default", true)
|
||||
done <- out
|
||||
}()
|
||||
|
||||
// Give the process time to start and reach the select
|
||||
for i := 0; i < 50; i++ {
|
||||
if s.Detach() {
|
||||
break
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
|
||||
out := <-done
|
||||
if !strings.Contains(out, "[detached: pid") {
|
||||
t.Fatalf("expected detach marker in output, got: %q", out)
|
||||
}
|
||||
if detachedPID == 0 {
|
||||
t.Fatal("OnDetach not called")
|
||||
}
|
||||
|
||||
// Verify process is in detached list
|
||||
procs := s.ListDetached()
|
||||
if len(procs) != 1 {
|
||||
t.Fatalf("expected 1 detached process, got %d", len(procs))
|
||||
}
|
||||
if procs[0].PID != detachedPID {
|
||||
t.Errorf("PID mismatch: got %d, want %d", procs[0].PID, detachedPID)
|
||||
}
|
||||
if procs[0].Exited {
|
||||
t.Error("process should still be running")
|
||||
}
|
||||
|
||||
// Kill it
|
||||
if err := s.SignalDetached(detachedPID, 15); err != nil { // SIGTERM
|
||||
t.Fatalf("SignalDetached: %v", err)
|
||||
}
|
||||
|
||||
// Wait for exit
|
||||
<-procs[0].done
|
||||
if !procs[0].Exited {
|
||||
t.Error("process should have exited after SIGTERM")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -65,6 +65,13 @@ type Server struct {
|
|||
validationFailures int
|
||||
lastFailure time.Time
|
||||
blockedUntil time.Time
|
||||
|
||||
// Detached process management
|
||||
detachMu sync.Mutex
|
||||
detachCh chan struct{} // signal to detach the currently running process
|
||||
detached []*DetachedProcess
|
||||
OnDetach func(pid int, cmd string) // hook: called when a process is detached
|
||||
OnExit func(pid int, exitCode int) // hook: called when a detached process exits
|
||||
}
|
||||
|
||||
// Option configures a Server.
|
||||
|
|
@ -386,6 +393,7 @@ func (e *Server) SetEnv(key, value string) {
|
|||
}
|
||||
// Close is called when the session ends. Calls OnClose hook if registered.
|
||||
func (e *Server) Close() {
|
||||
e.cleanupDetached()
|
||||
if e.OnClose != nil {
|
||||
e.OnClose()
|
||||
}
|
||||
|
|
@ -505,7 +513,9 @@ func (e *Server) executeWithStdin(ctx context.Context, code, language string, ti
|
|||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
|
||||
defer cancel()
|
||||
// NOTE: do NOT defer cancel() here — the detach path must avoid cancelling
|
||||
// the context (which would kill the detached process via cmd.Cancel).
|
||||
// Each exit path calls cancel() explicitly except detach.
|
||||
|
||||
var cmd *exec.Cmd
|
||||
var interpreter []string
|
||||
|
|
@ -606,19 +616,106 @@ func (e *Server) executeWithStdin(ctx context.Context, code, language string, ti
|
|||
cmd.Stdout = lw
|
||||
cmd.Stderr = lw
|
||||
|
||||
err = cmd.Run()
|
||||
output := outputBuf.Bytes()
|
||||
if lw.truncated {
|
||||
output = append(output, []byte("\n[output truncated at 10MB]")...)
|
||||
// Set up detach channel for this execution
|
||||
detachCh := make(chan struct{}, 1)
|
||||
e.detachMu.Lock()
|
||||
e.detachCh = detachCh
|
||||
e.detachMu.Unlock()
|
||||
defer func() {
|
||||
e.detachMu.Lock()
|
||||
e.detachCh = nil
|
||||
e.detachMu.Unlock()
|
||||
}()
|
||||
|
||||
if err = cmd.Start(); err != nil {
|
||||
return "", fmt.Errorf("execution failed: %v", err)
|
||||
}
|
||||
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
return "", fmt.Errorf("execution timeout after %d seconds", timeout)
|
||||
// Wait for completion, context cancellation, or detach signal
|
||||
waitCh := make(chan error, 1)
|
||||
go func() { waitCh <- cmd.Wait() }()
|
||||
|
||||
select {
|
||||
case err = <-waitCh:
|
||||
// Normal completion
|
||||
cancel()
|
||||
output := outputBuf.Bytes()
|
||||
if lw.truncated {
|
||||
output = append(output, []byte("\n[output truncated at 10MB]")...)
|
||||
}
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
return "", fmt.Errorf("execution timeout after %d seconds", timeout)
|
||||
}
|
||||
if err != nil {
|
||||
return string(output), fmt.Errorf("execution failed: %v\nOutput: %s", err, string(output))
|
||||
}
|
||||
return string(output), nil
|
||||
|
||||
case <-ctx.Done():
|
||||
// Context cancelled (interrupt or timeout) — cancel() already fired via ctx
|
||||
syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
|
||||
<-waitCh // reap
|
||||
output := outputBuf.Bytes()
|
||||
if ctx.Err() == context.DeadlineExceeded {
|
||||
return string(output), fmt.Errorf("execution timeout after %d seconds", timeout)
|
||||
}
|
||||
return string(output), fmt.Errorf("execution interrupted")
|
||||
|
||||
case <-detachCh:
|
||||
// Detach: move process to registry, continue running
|
||||
// Neutralize cmd.Cancel so context cancellation won't SIGKILL the process
|
||||
cmd.Cancel = nil
|
||||
// Cancel the timeout context (harmless now that cmd.Cancel is nil)
|
||||
cancel()
|
||||
|
||||
cmdStr := strings.Join(interpreter, " ")
|
||||
if len(cmdStr) > 80 {
|
||||
cmdStr = cmdStr[:77] + "..."
|
||||
}
|
||||
ring := newRingBuffer(ringBufSize)
|
||||
// Splice: future output goes to ring buffer instead of outputBuf.
|
||||
lw.mu.Lock()
|
||||
lw.w = ring
|
||||
lw.mu.Unlock()
|
||||
|
||||
proc := &DetachedProcess{
|
||||
PID: cmd.Process.Pid,
|
||||
Command: cmdStr,
|
||||
Started: time.Now(),
|
||||
ring: ring,
|
||||
cmd: cmd.Process,
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
e.detachMu.Lock()
|
||||
e.detached = append(e.detached, proc)
|
||||
e.detachMu.Unlock()
|
||||
|
||||
// Monitor for exit in background
|
||||
go func() {
|
||||
waitErr := <-waitCh
|
||||
proc.mu.Lock()
|
||||
proc.Exited = true
|
||||
if waitErr != nil {
|
||||
if exitErr, ok := waitErr.(*exec.ExitError); ok {
|
||||
proc.ExitCode = exitErr.ExitCode()
|
||||
} else {
|
||||
proc.ExitCode = -1
|
||||
}
|
||||
}
|
||||
proc.mu.Unlock()
|
||||
close(proc.done)
|
||||
if e.OnExit != nil {
|
||||
e.OnExit(proc.PID, proc.ExitCode)
|
||||
}
|
||||
}()
|
||||
|
||||
if e.OnDetach != nil {
|
||||
e.OnDetach(proc.PID, cmdStr)
|
||||
}
|
||||
|
||||
partial := outputBuf.String()
|
||||
return partial + fmt.Sprintf("\n[detached: pid %d]", proc.PID), nil
|
||||
}
|
||||
if err != nil {
|
||||
return string(output), fmt.Errorf("execution failed: %v\nOutput: %s", err, string(output))
|
||||
}
|
||||
return string(output), nil
|
||||
}
|
||||
|
||||
// prependOlliePath returns env with $OLLIE_CFG_PATH/scripts/x/ prepended to PATH so wrapper
|
||||
|
|
@ -662,4 +759,81 @@ func plan9Namespace(env map[string]string) string {
|
|||
return "/tmp/ns." + user + "." + disp
|
||||
}
|
||||
|
||||
// Detach signals the currently running process to be detached from the agent.
|
||||
// The process continues running; its output is captured in a ring buffer.
|
||||
// Returns false if no process is currently running.
|
||||
func (e *Server) Detach() bool {
|
||||
e.detachMu.Lock()
|
||||
ch := e.detachCh
|
||||
e.detachMu.Unlock()
|
||||
if ch == nil {
|
||||
return false
|
||||
}
|
||||
select {
|
||||
case ch <- struct{}{}:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// ListDetached returns all detached processes (running and exited).
|
||||
func (e *Server) ListDetached() []*DetachedProcess {
|
||||
e.detachMu.Lock()
|
||||
defer e.detachMu.Unlock()
|
||||
out := make([]*DetachedProcess, len(e.detached))
|
||||
copy(out, e.detached)
|
||||
return out
|
||||
}
|
||||
|
||||
// SignalDetached sends a signal to a detached process by PID.
|
||||
func (e *Server) SignalDetached(pid int, sig syscall.Signal) error {
|
||||
e.detachMu.Lock()
|
||||
defer e.detachMu.Unlock()
|
||||
for _, p := range e.detached {
|
||||
if p.PID == pid {
|
||||
return p.Signal(sig)
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("no detached process with pid %d", pid)
|
||||
}
|
||||
|
||||
// GetDetachedOutput returns the ring buffer contents for a detached process.
|
||||
func (e *Server) GetDetachedOutput(pid int) (string, error) {
|
||||
e.detachMu.Lock()
|
||||
defer e.detachMu.Unlock()
|
||||
for _, p := range e.detached {
|
||||
if p.PID == pid {
|
||||
return p.Output(), nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("no detached process with pid %d", pid)
|
||||
}
|
||||
|
||||
// DismissDetached removes an exited process from the list.
|
||||
func (e *Server) DismissDetached(pid int) bool {
|
||||
e.detachMu.Lock()
|
||||
defer e.detachMu.Unlock()
|
||||
for i, p := range e.detached {
|
||||
if p.PID == pid && p.Exited {
|
||||
e.detached = append(e.detached[:i], e.detached[i+1:]...)
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// cleanupDetached sends SIGTERM to all running detached processes.
|
||||
func (e *Server) cleanupDetached() {
|
||||
e.detachMu.Lock()
|
||||
defer e.detachMu.Unlock()
|
||||
for _, p := range e.detached {
|
||||
p.mu.Lock()
|
||||
if !p.Exited && p.cmd != nil {
|
||||
syscall.Kill(-p.PID, syscall.SIGTERM)
|
||||
}
|
||||
p.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
var _ tools.Server = (*Server)(nil) // compile-time interface check
|
||||
|
|
|
|||
Reference in New Issue