278 lines
7.5 KiB
Go
278 lines
7.5 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 = "bash"
|
|
code = toolCode
|
|
if len(s.Args) > 0 {
|
|
code = injectArgs(language, s.Tool, s.Args, toolCode)
|
|
}
|
|
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+".*\$`),
|
|
}
|
|
|
|
var languagePatterns = map[string][]*regexp.Regexp{
|
|
"bash": bashPatterns,
|
|
"": bashPatterns,
|
|
}
|
|
|
|
// 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()
|
|
}
|
|
|
|
switch name {
|
|
case "execute_code":
|
|
return dispatchExecuteCode(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)
|
|
}
|
|
|
|
|
|
|
|
// 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
|
|
}
|
|
|
|
|
|
|
|
// 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
|
|
}
|
|
|
|
|