ollie/toolsrv/shell.go

package toolsrv

import (
	"bytes"
	"context"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"time"

	"ollie/paths"
	"ollie/sandbox"
)

// universalPatterns apply to all code.
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 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,
}
// ValidateCode checks code against dangerous patterns.
func (s *Server) ValidateCode(code, language string) error {
	if err := s.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) {
			s.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
}
// Connects to the broker socket, sends the request with the current env,
// and streams the framed response back.
func (s *Server) executeBypass(ctx context.Context, cmd, dir string, timeout int, doDetach ...bool) (string, error) {
	return s.executeBypassOpts(ctx, cmd, dir, timeout, false, doDetach...)
}

func (s *Server) executeBypassSudo(ctx context.Context, cmd, dir string, timeout int) (string, error) {
	return s.executeBypassOpts(ctx, cmd, dir, timeout, true)
}

func (s *Server) executeBypassOpts(ctx context.Context, cmd, dir string, timeout int, sudo bool, doDetach ...bool) (string, error) {
	xdg := os.Getenv("XDG_RUNTIME_DIR")
	if xdg == "" {
		return "", fmt.Errorf("bypass not available: no XDG_RUNTIME_DIR")
	}
	sockPath := filepath.Join(xdg, "ollie", "bypass.sock")

	wantDetach := len(doDetach) > 0 && doDetach[0]

	var cancel context.CancelFunc
	if wantDetach || timeout <= 0 {
		ctx, cancel = context.WithCancel(ctx)
	} else {
		ctx, cancel = context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
	}

	// Connect to broker
	conn, err := net.DialTimeout("unix", sockPath, 5*time.Second)
	if err != nil {
		cancel()
		return "", fmt.Errorf("bypass not available: %w", err)
	}

	// Send request — merge process env with session-scoped envExtra.
	envMap := make(map[string]string, len(os.Environ()))
	for _, kv := range os.Environ() {
		if k, v, ok := strings.Cut(kv, "="); ok {
			envMap[k] = v
		}
	}
	s.envMu.RLock()
	for k, v := range s.envExtra {
		envMap[k] = v
	}
	s.envMu.RUnlock()
	reqJSON, _ := json.Marshal(struct {
		Cmd     string            `json:"cmd"`
		Cwd     string            `json:"cwd"`
		Env     map[string]string `json:"env"`
		Session string            `json:"session,omitempty"`
		Sudo    bool              `json:"sudo,omitempty"`
	}{Cmd: cmd, Cwd: dir, Env: envMap, Session: s.sessionID, Sudo: sudo})
	reqJSON = append(reqJSON, '\n')
	if _, err := conn.Write(reqJSON); err != nil {
		conn.Close()
		cancel()
		return "", fmt.Errorf("bypass execution failed: write: %w", err)
	}

	// Set up detach channel
	detachCh := make(chan struct{}, 1)
	s.detachMu.Lock()
	s.detachCh = detachCh
	s.detachMu.Unlock()
	defer func() {
		s.detachMu.Lock()
		s.detachCh = nil
		s.detachMu.Unlock()
	}()

	if wantDetach {
		close(detachCh)
	}

	// readFrames reads from the connection, writing to w and streaming.
	// Returns exit code when 'x' frame arrives or -1 on error.
	readFrames := func(w io.Writer, stream func(string)) int {
		header := make([]byte, 5)
		for {
			conn.SetReadDeadline(time.Now().Add(1 * time.Second)) //nolint:errcheck
			_, err := io.ReadFull(conn, header)
			if err != nil {
				if ne, ok := err.(net.Error); ok && ne.Timeout() {
					// Check for context cancellation
					select {
					case <-ctx.Done():
						return -1
					default:
						continue
					}
				}
				return -1
			}

			frameType := header[0]
			length := binary.BigEndian.Uint32(header[1:5])

			payload := make([]byte, length)
			if length > 0 {
				conn.SetReadDeadline(time.Now().Add(30 * time.Second)) //nolint:errcheck
				if _, err := io.ReadFull(conn, payload); err != nil {
					return -1
				}
			}

			switch frameType {
			case 'd':
				w.Write(payload) //nolint:errcheck
				if stream != nil {
					stream(string(payload))
				}
			case 'x':
				var code int
				fmt.Sscanf(string(payload), "%d", &code)
				return code
			}
		}
	}

	// Check if detach was requested
	select {
	case <-detachCh:
		// Detach: background the socket read into a goroutine
		cmdStr := "bypass: " + cmd
		if len(cmdStr) > 80 {
			cmdStr = cmdStr[:77] + "..."
		}
		ring := NewRingBuffer(RingBufSize)
		pid := int(time.Now().UnixNano() & 0x7FFFFFFF) // synthetic PID

			proc := &DetachedProcess{
			PID:     pid,
				Command: cmdStr,
				Started: time.Now(),
				Ring:    ring,
				Done:    make(chan struct{}),
		}
		s.detachMu.Lock()
		s.detached = append(s.detached, proc)
		s.detachMu.Unlock()

		go func() {
			defer conn.Close()
			defer cancel()
			exitCode := readFrames(ring, nil)
			proc.Mu.Lock()
			proc.Exited = true
			proc.ExitCode = exitCode
			proc.Mu.Unlock()
			close(proc.Done)
			if s.OnExit != nil {
				s.OnExit(proc.PID, proc.ExitCode)
			}
		}()

		if s.OnDetach != nil {
			s.OnDetach(proc.PID, cmdStr)
		}
		return fmt.Sprintf("[detached: pid %d]", pid), nil

	default:
		// Normal (foreground) execution — must also handle manual detach.
		var outputBuf bytes.Buffer
		streamFn := StreamFunc(ctx)
		lw := &limitedWriter{
			w:      &outputBuf,
			limit:  10 * 1024 * 1024,
			stream: streamFn,
		}

		// Run readFrames in a goroutine so we can select on detach.
		type frameResult struct{ exitCode int }
		frameCh := make(chan frameResult, 1)
		go func() {
			code := readFrames(lw, nil)
			frameCh <- frameResult{code}
		}()

		select {
		case fr := <-frameCh:
			// Normal completion
			conn.Close()
			cancel()
			combined := outputBuf.String()
			if fr.exitCode != 0 {
				if combined == "" {
					return "", fmt.Errorf("bypass execution failed (exit %d)", fr.exitCode)
				}
				errOutput := combined
				const maxErrOutput = 8192
				if len(errOutput) > maxErrOutput {
					errOutput = errOutput[:maxErrOutput] + fmt.Sprintf("\n[...truncated %d bytes in error]", len(combined)-maxErrOutput)
				}
				return combined, fmt.Errorf("bypass execution failed (exit %d)\nOutput: %s", fr.exitCode, errOutput)
			}
			return combined, nil

		case <-ctx.Done():
			// Timeout or external cancellation
			conn.Close()
			cancel()
			combined := outputBuf.String()
			if ctx.Err() == context.DeadlineExceeded {
				return combined, fmt.Errorf("bypass execution timeout after %d seconds", timeout)
			}
			return combined, fmt.Errorf("bypass execution interrupted")

		case <-detachCh:
			// Manual detach: background the connection read
			cmdStr := "bypass: " + cmd
			if len(cmdStr) > 80 {
				cmdStr = cmdStr[:77] + "..."
			}
			ring := NewRingBuffer(RingBufSize)
			pid := int(time.Now().UnixNano() & 0x7FFFFFFF)

			// Splice: future output goes to ring buffer, stop streaming
			lw.mu.Lock()
			lw.w = ring
			lw.stream = nil
			lw.mu.Unlock()

			proc := &DetachedProcess{
				PID:     pid,
				Command: cmdStr,
				Started: time.Now(),
				Ring:    ring,
				Done:    make(chan struct{}),
			}
			s.detachMu.Lock()
			s.detached = append(s.detached, proc)
			s.detachMu.Unlock()

			go func() {
				fr := <-frameCh
				conn.Close()
				cancel()
				proc.Mu.Lock()
				proc.Exited = true
				proc.ExitCode = fr.exitCode
				proc.Mu.Unlock()
				close(proc.Done)
				if s.OnExit != nil {
					s.OnExit(proc.PID, proc.ExitCode)
				}
			}()

			if s.OnDetach != nil {
				s.OnDetach(proc.PID, cmdStr)
			}

			partial := outputBuf.String()
			return partial + fmt.Sprintf("\n[detached: pid %d]", pid), nil
		}
	}
}

var whitespacePattern = regexp.MustCompile(`\s+`)

func (s *Server) checkRateLimit() error {
	s.rateLimitMu.Lock()
	defer s.rateLimitMu.Unlock()

	now := time.Now()
	if now.Before(s.blockedUntil) {
		remaining := s.blockedUntil.Sub(now).Round(time.Second)
		return &RateLimitedError{Remaining: remaining}
	}
	return nil
}

func (s *Server) recordValidationFailure() {
	s.rateLimitMu.Lock()
	defer s.rateLimitMu.Unlock()

	now := time.Now()
	if now.Sub(s.lastFailure) > failureWindow {
		s.validationFailures = 0
	}

	s.validationFailures++
	s.lastFailure = now

	if s.validationFailures >= maxFailures {
		s.blockedUntil = now.Add(blockDuration)
		s.validationFailures = 0
	}
}

// executeWithStdin runs code in a sandbox and returns combined stdout+stderr.
// For languages where code is itself passed via stdin (ed, expect, bc), stdinData is ignored.
func (s *Server) executeWithStdin(ctx context.Context, code, language string, timeout int, sandboxName string, trusted bool, stdinData string, doDetach ...bool) (string, error) {
	if timeout < 0 {
		timeout = 30
	}

	if !trusted {
		if err := s.ValidateCode(code, language); err != nil {
			return "", err
		}
	}

	var cfg *sandbox.Config
	var err error
	if !s.Yolo {
		cfg, err = loadSandboxConfig(sandboxName)
		if err != nil {
			return "", err
		}
	}

	s.wdMu.RLock()
	workDir := s.cwd
	s.wdMu.RUnlock()
	if workDir == "" {
		workDir, _ = os.Getwd()
	}

	var cancel context.CancelFunc
	if timeout > 0 {
		ctx, cancel = context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
	} else {
		// timeout=0 means no timeout (run indefinitely, s.g. for daemons)
		ctx, cancel = context.WithCancel(ctx)
	}
	// 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
	// codeStdin: non-empty means the code itself is fed via stdin (ed, expect, bc).
	// In these cases stdinData cannot be used simultaneously.
	var codeStdin string
	switch language {
	case "bash", "":
		interpreter = []string{"bash", "-c", code}
	default:
		cancel()
		return "", fmt.Errorf("unsupported language: %s (only bash is supported)", language)
	}
	s.envMu.RLock()
	envMap := make(map[string]string, len(s.envExtra))
	for k, v := range s.envExtra {
		envMap[k] = v
	}
	s.envMu.RUnlock()
	for _, ev := range os.Environ() {
		k, v, _ := strings.Cut(ev, "=")
		if _, ok := envMap[k]; !ok {
			envMap[k] = v
		}
	}
	if _, ok := envMap["NAMESPACE"]; !ok {
		if ns := plan9Namespace(envMap); ns != "" {
			envMap["NAMESPACE"] = ns
		}
	}
	getenv := func(key string) string { return envMap[key] }

	if s.Yolo {
		cmd = exec.CommandContext(ctx, interpreter[0], interpreter[1:]...)
	} else {
		wrapped, wrapErr := sandbox.WrapCommand(cfg, interpreter, workDir, getenv)
		if wrapErr != nil {
			cancel()
			return "", wrapErr
		}
		cmd = exec.CommandContext(ctx, wrapped[0], wrapped[1:]...)
	}
	cmd.Dir = workDir
	switch {
	case codeStdin != "":
		cmd.Stdin = strings.NewReader(codeStdin)
	case stdinData != "":
		cmd.Stdin = strings.NewReader(stdinData)
	}

	s.envMu.RLock()
	baseEnv := os.Environ()
	// Remove keys that envExtra overrides so duplicates don't leak through.
	filtered := baseEnv[:0]
	for _, ev := range baseEnv {
		k, _, _ := strings.Cut(ev, "=")
		if _, overridden := s.envExtra[k]; !overridden {
			filtered = append(filtered, ev)
		}
	}
	cmd.Env = filtered
	for k, v := range s.envExtra {
		cmd.Env = append(cmd.Env, k+"="+v)
	}
	s.envMu.RUnlock()
	// Export computed NAMESPACE so tools like ollie-9p can find the server.
	if ns := envMap["NAMESPACE"]; ns != "" {
		cmd.Env = append(cmd.Env, "NAMESPACE="+ns)
	}

	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	cmd.Cancel = func() error {
		if cmd.Process != nil {
			syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		}
		return nil
	}
	cmd.WaitDelay = time.Second

	var outputBuf bytes.Buffer
	lw := &limitedWriter{
		w:      &outputBuf,
		limit:  10 * 1024 * 1024,
		stream: StreamFunc(ctx),
	}
	cmd.Stdout = lw
	cmd.Stderr = lw

	// Set up detach channel for this execution
	detachCh := make(chan struct{}, 1)
	s.detachMu.Lock()
	s.detachCh = detachCh
	s.detachMu.Unlock()
	defer func() {
		s.detachMu.Lock()
		s.detachCh = nil
		s.detachMu.Unlock()
	}()

	if err = cmd.Start(); err != nil {
		cancel()
		return "", fmt.Errorf("execution failed: %v", err)
	}

	// If detach requested, immediately signal the detach channel
	if len(doDetach) > 0 && doDetach[0] {
		close(detachCh)
	}

	// 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 {
			// Cap output in error message to avoid flooding the agent context.
			// The full output is still returned as the first return value.
			errOutput := string(output)
			const maxErrOutput = 8192
			if len(errOutput) > maxErrOutput {
				errOutput = errOutput[:maxErrOutput] + fmt.Sprintf("\n[...truncated %d bytes in error]", len(output)-maxErrOutput)
			}
			return string(output), fmt.Errorf("execution failed: %v\nOutput: %s", err, errOutput)
		}
		return string(output), nil

	case <-ctx.Done():
		// Context cancelled (interrupt or timeout)
		cancel()
		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 and WaitDelay so context cancellation
		// won't kill the process after the timeout fires.
		cmd.Cancel = nil
		cmd.WaitDelay = 0
		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.stream = nil
		lw.mu.Unlock()

			proc := &DetachedProcess{
				PID:     cmd.Process.Pid,
				Command: cmdStr,
				Started: time.Now(),
				Ring:    ring,
				Done:    make(chan struct{}),
		}
		s.detachMu.Lock()
		s.detached = append(s.detached, proc)
		s.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 s.OnExit != nil {
				s.OnExit(proc.PID, proc.ExitCode)
			}
		}()

		if s.OnDetach != nil {
			s.OnDetach(proc.PID, cmdStr)
		}

		partial := outputBuf.String()
		return partial + fmt.Sprintf("\n[detached: pid %d]", proc.PID), nil
	}
}

// plan9Namespace computes the plan9port namespace directory from the environment.
// Convention: $NAMESPACE if set, else /tmp/ns.<user>.<display> where display is
// $DISPLAY with trailing .0 stripped and / replaced by _.
func plan9Namespace(env map[string]string) string {
	if ns := env["NAMESPACE"]; ns != "" {
		return ns
	}
	user := env["USER"]
	if user == "" {
		return ""
	}
	disp := env["DISPLAY"]
	if disp == "" {
		return ""
	}
	// Canonicalize: strip trailing .0
	if strings.HasSuffix(disp, ".0") {
		disp = disp[:len(disp)-2]
	}
	// Replace / with _
	disp = strings.ReplaceAll(disp, "/", "_")
	return "/tmp/ns." + user + "." + disp
}