427 lines
13 KiB
Go
427 lines
13 KiB
Go
package execute
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"path/filepath"
|
|
"sync"
|
|
|
|
"ollie/internal/sandbox"
|
|
"ollie/pkg/paths"
|
|
|
|
"regexp"
|
|
"strings"
|
|
)
|
|
|
|
// CodeStep is one stage in an execution pipeline.
|
|
// Set Code/Language for inline code, Tool/Args for a named script, or Parallel for concurrent fan-out.
|
|
// Set Elevated to run the step outside the sandbox via the configured elevation backend.
|
|
// Set Detach to immediately background the process and return its PID.
|
|
type CodeStep struct {
|
|
Code string `json:"code"`
|
|
Language string `json:"language"`
|
|
Tool string `json:"tool"`
|
|
Args []string `json:"args"`
|
|
Parallel []CodeStep `json:"parallel"`
|
|
Elevated bool `json:"elevated"`
|
|
Detach bool `json:"detach"`
|
|
}
|
|
|
|
// resolveCodeStep loads a CodeStep into executable (code, language, trusted).
|
|
// Tool steps are read from the tools directory and treated as trusted.
|
|
// Inline code steps default to bash and are validated by the caller.
|
|
func resolveCodeStep(s CodeStep) (code, language string, trusted bool, err error) {
|
|
if s.Tool != "" {
|
|
toolCode, terr := ReadTool(s.Tool)
|
|
if terr != nil {
|
|
return "", "", false, terr
|
|
}
|
|
language = detectLanguage(toolCode)
|
|
code = toolCode
|
|
if len(s.Args) > 0 {
|
|
code = injectArgs(language, s.Tool, s.Args, toolCode)
|
|
switch language {
|
|
case "awk", "sed", "jq", "ed", "expect", "bc":
|
|
language = "bash"
|
|
}
|
|
}
|
|
return code, language, true, nil
|
|
}
|
|
language = s.Language
|
|
if language == "" {
|
|
language = "bash"
|
|
}
|
|
return s.Code, language, false, nil
|
|
}
|
|
|
|
// universalPatterns apply to all general-purpose languages.
|
|
var universalPatterns = []*regexp.Regexp{
|
|
regexp.MustCompile(`\bmkfs\b`),
|
|
regexp.MustCompile(`\bdd\b.*\bif=/dev/`),
|
|
regexp.MustCompile(`\b(sudo|su)\s`),
|
|
regexp.MustCompile(`/etc/(shadow|sudoers)`),
|
|
}
|
|
|
|
// bashPatterns apply only to bash (flag syntax, redirects, shell-specific constructs).
|
|
var bashPatterns = []*regexp.Regexp{
|
|
regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s*/(home|var|usr|etc|boot|root|bin|sbin|lib|opt|srv)?`),
|
|
regexp.MustCompile(`rm\s+(-[a-z]*f[a-z]*\s+)*-[a-z]*r[a-z]*\s*/(home|var|usr|etc|boot|root|bin|sbin|lib|opt|srv)?`),
|
|
regexp.MustCompile(`rm\s+.*--recursive.*--force`),
|
|
regexp.MustCompile(`rm\s+.*--force.*--recursive`),
|
|
regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+\.\.(/|\s|$)`),
|
|
regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+~`),
|
|
regexp.MustCompile(`rm\s+(-[a-z]*r[a-z]*\s+)*-[a-z]*f[a-z]*\s+\*`),
|
|
regexp.MustCompile(`:\s*\(\s*\)\s*\{\s*:\s*\|\s*:\s*&`), // fork bomb
|
|
regexp.MustCompile(`>\s*/dev/sd`),
|
|
regexp.MustCompile(`\beval\s+".*\$`),
|
|
}
|
|
|
|
// pythonPatterns apply to python3/python.
|
|
var pythonPatterns = []*regexp.Regexp{
|
|
regexp.MustCompile(`shutil\.rmtree\s*\(\s*['"/]`),
|
|
regexp.MustCompile(`(os\.system|subprocess\.(call|run|popen))\s*\(.*\brm\s+-[a-z]*r[a-z]*f`),
|
|
regexp.MustCompile(`os\.(remove|unlink)\s*\(\s*['"]/(etc|usr|bin|sbin|lib|boot)`),
|
|
}
|
|
|
|
// perlPatterns apply to perl.
|
|
var perlPatterns = []*regexp.Regexp{
|
|
regexp.MustCompile(`(system|exec)\s*\(.*\brm\s+-[a-z]*r[a-z]*f`),
|
|
regexp.MustCompile("`.+rm\\s+-[a-z]*r[a-z]*f"), // backtick execution
|
|
regexp.MustCompile(`unlink\s+glob\s*\(['"]/(etc|usr|bin|sbin|lib|boot)`),
|
|
}
|
|
|
|
// luaPatterns apply to lua.
|
|
var luaPatterns = []*regexp.Regexp{
|
|
regexp.MustCompile(`(os\.execute|io\.popen)\s*\(.*\brm\s+-[a-z]*r[a-z]*f`),
|
|
regexp.MustCompile(`os\.remove\s*\(\s*['"]/(etc|usr|bin|sbin|lib|boot)`),
|
|
}
|
|
|
|
var languagePatterns = map[string][]*regexp.Regexp{
|
|
"bash": bashPatterns,
|
|
"": bashPatterns,
|
|
"python3": pythonPatterns,
|
|
"python": pythonPatterns,
|
|
"perl": perlPatterns,
|
|
"lua": luaPatterns,
|
|
}
|
|
|
|
// Dispatch routes a named execute tool call.
|
|
func (e *Server) Dispatch(ctx context.Context, name string, args json.RawMessage) (string, error) {
|
|
if e.OnPreDispatch != nil {
|
|
e.OnPreDispatch()
|
|
}
|
|
if len(e.allowExecutors) > 0 && !e.allowExecutors[name] {
|
|
return "", fmt.Errorf("executor %q not permitted in this agent context", name)
|
|
}
|
|
switch name {
|
|
case "execute_code":
|
|
return dispatchExecuteCode(ctx, e, args)
|
|
case "call_tool":
|
|
return dispatchCallTool(ctx, e, args)
|
|
case "pipe":
|
|
return dispatchPipeCall(ctx, e, args)
|
|
default:
|
|
return "", fmt.Errorf("unknown execute tool: %s", name)
|
|
}
|
|
}
|
|
|
|
// ValidateCode checks code against dangerous patterns for the given language.
|
|
func (e *Server) ValidateCode(code, language string) error {
|
|
if err := e.checkRateLimit(); err != nil {
|
|
return err
|
|
}
|
|
|
|
normalized := strings.ToLower(code)
|
|
normalized = whitespacePattern.ReplaceAllString(normalized, " ")
|
|
|
|
patterns := append(universalPatterns, languagePatterns[language]...)
|
|
for _, pattern := range patterns {
|
|
if pattern.MatchString(normalized) {
|
|
e.recordValidationFailure()
|
|
return fmt.Errorf("dangerous pattern detected")
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func loadSandboxConfig(name string) (*sandbox.Config, error) {
|
|
if name == "" {
|
|
name = "default"
|
|
}
|
|
path := filepath.Join(paths.CfgDir(), "sandbox", name+".yaml")
|
|
f, err := os.Open(path)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("sandbox %q not found: %w", name, err)
|
|
}
|
|
defer f.Close()
|
|
return sandbox.LoadSandbox(f)
|
|
}
|
|
|
|
type limitedWriter struct {
|
|
mu sync.Mutex
|
|
w io.Writer
|
|
written int
|
|
limit int
|
|
truncated bool
|
|
stream func(string) // if non-nil, called with each chunk of output
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
remaining := lw.limit - lw.written
|
|
toWrite := p
|
|
if len(p) > remaining {
|
|
toWrite = p[:remaining]
|
|
lw.truncated = true
|
|
}
|
|
|
|
written, err := lw.w.Write(toWrite)
|
|
lw.written += written
|
|
if lw.stream != nil {
|
|
lw.stream(string(toWrite[:written]))
|
|
}
|
|
if err != nil {
|
|
return written, err
|
|
}
|
|
return len(p), nil
|
|
}
|
|
|
|
// dispatchExecuteCode handles the execute_code tool: inline code steps only.
|
|
// The legacy "pipe" field is no longer accepted here; use the pipe tool instead.
|
|
func dispatchExecuteCode(ctx context.Context, e *Server, args json.RawMessage) (string, error) {
|
|
steps, timeout, sandboxName, err := execCodeOnlyArgs(args)
|
|
if err != nil {
|
|
return "", fmt.Errorf("execute_code: bad args: %w", err)
|
|
}
|
|
if len(steps) == 0 {
|
|
return "", fmt.Errorf("execute_code: steps is required. Format: {\"steps\": [{\"code\": \"your code here\"}]}")
|
|
}
|
|
return e.dispatchSteps(ctx, steps, timeout, sandboxName)
|
|
}
|
|
|
|
// dispatchCallTool handles the call_tool tool: named tool scripts only.
|
|
// Inline code fields are rejected. Fan-out behaviour is preserved.
|
|
func dispatchCallTool(ctx context.Context, e *Server, args json.RawMessage) (string, error) {
|
|
calls, timeout, sandboxName, err := execCallToolArgs(args)
|
|
if err != nil {
|
|
return "", fmt.Errorf("call_tool: bad args: %w. Expected format: {\"calls\": [{\"tool\": \"tool_name\", \"args\": [\"arg1\", ...]}]}", err)
|
|
}
|
|
if len(calls) == 0 {
|
|
return "", fmt.Errorf("call_tool: calls is required. Format: {\"calls\": [{\"tool\": \"tool_name\", \"args\": [\"arg1\", ...]}]}")
|
|
}
|
|
// Validate: every top-level call must name a tool (or be a parallel group of tools).
|
|
for i, c := range calls {
|
|
if len(c.Parallel) > 0 {
|
|
for j, p := range c.Parallel {
|
|
if p.Tool == "" {
|
|
return "", fmt.Errorf("call_tool: calls[%d].parallel[%d]: tool name is required", i, j)
|
|
}
|
|
if len(e.allowTools) > 0 && !e.allowTools[p.Tool] {
|
|
return "", fmt.Errorf("call_tool: tool %q not permitted in this agent context", p.Tool)
|
|
}
|
|
}
|
|
continue
|
|
}
|
|
if c.Tool == "" {
|
|
return "", fmt.Errorf("call_tool: calls[%d]: tool name is required", i)
|
|
}
|
|
if len(e.allowTools) > 0 && !e.allowTools[c.Tool] {
|
|
return "", fmt.Errorf("call_tool: tool %q not permitted in this agent context", c.Tool)
|
|
}
|
|
}
|
|
return e.dispatchSteps(ctx, calls, timeout, sandboxName)
|
|
}
|
|
|
|
// dispatchPipeCall handles the pipe tool: sequential pipeline of code and/or tool stages.
|
|
func dispatchPipeCall(ctx context.Context, e *Server, args json.RawMessage) (string, error) {
|
|
stages, timeout, sandboxName, err := execPipeArgs(args)
|
|
if err != nil {
|
|
return "", fmt.Errorf("pipe: bad args: %w", err)
|
|
}
|
|
if len(stages) == 0 {
|
|
return "", fmt.Errorf("pipe: stages is required. Format: {\"stages\": [{\"code\": \"...\"} or {\"tool\": \"name\", \"args\": [...]}]}")
|
|
}
|
|
return e.dispatchPipe(ctx, stages, timeout, sandboxName)
|
|
}
|
|
|
|
// dispatchSteps runs steps in parallel when safe (annotation-based), serially otherwise.
|
|
func (e *Server) dispatchSteps(ctx context.Context, steps []CodeStep, timeout int, sandboxName string) (string, error) {
|
|
if e.Strict {
|
|
if err := enforceStrict(steps); err != nil {
|
|
return "", err
|
|
}
|
|
}
|
|
|
|
// Single step: run directly.
|
|
if len(steps) == 1 && len(steps[0].Parallel) == 0 {
|
|
if steps[0].Elevated {
|
|
e.wdMu.RLock()
|
|
dir := e.cwd
|
|
e.wdMu.RUnlock()
|
|
return e.executeElevated(ctx, steps[0].Code, dir, timeout, steps[0].Detach)
|
|
}
|
|
code, lang, trusted, err := resolveCodeStep(steps[0])
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return e.executeWithStdin(ctx, code, lang, timeout, sandboxName, trusted, "", steps[0].Detach)
|
|
}
|
|
|
|
// Multiple steps: auto-batch consecutive read-safe steps in parallel;
|
|
// write/global steps run serially between batches.
|
|
var output string
|
|
for i := 0; i < len(steps); {
|
|
if classifyStep(steps[i]) == lockClassRead {
|
|
j := i + 1
|
|
for j < len(steps) && classifyStep(steps[j]) == lockClassRead {
|
|
j++
|
|
}
|
|
out, err := e.runReadBatch(ctx, i, steps[i:j], timeout, sandboxName, "")
|
|
if err != nil {
|
|
return out, err
|
|
}
|
|
output += out
|
|
i = j
|
|
} else {
|
|
lf, err := acquireFlock(e.lockDir, "rw", true)
|
|
if err != nil {
|
|
return "", fmt.Errorf("step %d: lock: %w", i, err)
|
|
}
|
|
out, runErr := e.runStage(ctx, i, steps[i], timeout, sandboxName, "")
|
|
if lf != nil {
|
|
lf.Close()
|
|
}
|
|
if runErr != nil {
|
|
return output + out, fmt.Errorf("step %d: %w", i, runErr)
|
|
}
|
|
output += out
|
|
i++
|
|
}
|
|
}
|
|
return output, nil
|
|
}
|
|
|
|
// dispatchPipe runs stages as an explicit pipeline: sequential, stdout of each
|
|
// stage feeds stdin of the next.
|
|
func (e *Server) dispatchPipe(ctx context.Context, pipe []CodeStep, timeout int, sandboxName string) (string, error) {
|
|
if e.Strict {
|
|
if err := enforceStrict(pipe); err != nil {
|
|
return "", err
|
|
}
|
|
}
|
|
|
|
if len(e.allowTools) > 0 {
|
|
for i, s := range pipe {
|
|
if s.Tool != "" && !e.allowTools[s.Tool] {
|
|
return "", fmt.Errorf("pipe: stage %d: tool %q not permitted in this agent context", i, s.Tool)
|
|
}
|
|
for j, p := range s.Parallel {
|
|
if p.Tool != "" && !e.allowTools[p.Tool] {
|
|
return "", fmt.Errorf("pipe: stage %d.parallel[%d]: tool %q not permitted in this agent context", i, j, p.Tool)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
var input string
|
|
for i, stage := range pipe {
|
|
out, err := e.runStage(ctx, i, stage, timeout, sandboxName, input)
|
|
if err != nil {
|
|
return out, fmt.Errorf("pipe stage %d: %w", i, err)
|
|
}
|
|
input = out
|
|
}
|
|
return input, nil
|
|
}
|
|
|
|
// enforceStrict rejects any step that uses inline code rather than a named tool.
|
|
func enforceStrict(stages []CodeStep) error {
|
|
for i, s := range stages {
|
|
if len(s.Parallel) > 0 {
|
|
if err := enforceStrict(s.Parallel); err != nil {
|
|
return err
|
|
}
|
|
continue
|
|
}
|
|
if s.Tool == "" {
|
|
return fmt.Errorf("execute_code: step %d rejected: inline code not allowed in strict mode", i)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// execCodeOnlyArgs parses args for the execute_code tool (steps only).
|
|
func execCodeOnlyArgs(args json.RawMessage) (steps []CodeStep, timeout int, sandboxName string, err error) {
|
|
var a struct {
|
|
Steps []CodeStep `json:"steps"`
|
|
Timeout int `json:"timeout"`
|
|
Sandbox string `json:"sandbox"`
|
|
}
|
|
if err = json.Unmarshal(args, &a); err != nil {
|
|
return
|
|
}
|
|
steps = a.Steps
|
|
timeout = a.Timeout
|
|
if timeout <= 0 {
|
|
timeout = 30
|
|
}
|
|
sandboxName = a.Sandbox
|
|
if sandboxName == "" {
|
|
sandboxName = "default"
|
|
}
|
|
return
|
|
}
|
|
|
|
// execCallToolArgs parses args for the call_tool tool (calls array).
|
|
func execCallToolArgs(args json.RawMessage) (calls []CodeStep, timeout int, sandboxName string, err error) {
|
|
var a struct {
|
|
Calls []CodeStep `json:"calls"`
|
|
Timeout int `json:"timeout"`
|
|
Sandbox string `json:"sandbox"`
|
|
}
|
|
if err = json.Unmarshal(args, &a); err != nil {
|
|
return
|
|
}
|
|
calls = a.Calls
|
|
timeout = a.Timeout
|
|
if timeout <= 0 {
|
|
timeout = 30
|
|
}
|
|
sandboxName = a.Sandbox
|
|
if sandboxName == "" {
|
|
sandboxName = "default"
|
|
}
|
|
return
|
|
}
|
|
|
|
// execPipeArgs parses args for the pipe tool (stages array).
|
|
func execPipeArgs(args json.RawMessage) (stages []CodeStep, timeout int, sandboxName string, err error) {
|
|
var a struct {
|
|
Stages []CodeStep `json:"stages"`
|
|
Timeout int `json:"timeout"`
|
|
Sandbox string `json:"sandbox"`
|
|
}
|
|
if err = json.Unmarshal(args, &a); err != nil {
|
|
return
|
|
}
|
|
stages = a.Stages
|
|
timeout = a.Timeout
|
|
if timeout <= 0 {
|
|
timeout = 30
|
|
}
|
|
sandboxName = a.Sandbox
|
|
if sandboxName == "" {
|
|
sandboxName = "default"
|
|
}
|
|
return
|
|
}
|