gooseek/backend/internal/computer/sandbox.go

package computer

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
)

type SandboxConfig struct {
	Image        string
	Timeout      time.Duration
	MemoryLimit  string
	CPULimit     string
	NetworkMode  string
	WorkDir      string
	MaxFileSize  int64
	AllowNetwork bool
}

func DefaultSandboxConfig() SandboxConfig {
	return SandboxConfig{
		Image:        "gooseek/sandbox:latest",
		Timeout:      5 * time.Minute,
		MemoryLimit:  "512m",
		CPULimit:     "1.0",
		NetworkMode:  "none",
		WorkDir:      "/workspace",
		MaxFileSize:  10 * 1024 * 1024,
		AllowNetwork: false,
	}
}

type Sandbox struct {
	ID          string
	ContainerID string
	WorkDir     string
	Status      string
	TaskID      string
	CreatedAt   time.Time
}

type SandboxManager struct {
	cfg       SandboxConfig
	sandboxes map[string]*Sandbox
	mu        sync.RWMutex
	useDocker bool
}

func NewSandboxManager(cfg SandboxConfig) *SandboxManager {
	if cfg.Timeout == 0 {
		cfg.Timeout = 5 * time.Minute
	}
	if cfg.MemoryLimit == "" {
		cfg.MemoryLimit = "512m"
	}
	if cfg.WorkDir == "" {
		cfg.WorkDir = "/workspace"
	}

	useDocker := isDockerAvailable()

	return &SandboxManager{
		cfg:       cfg,
		sandboxes: make(map[string]*Sandbox),
		useDocker: useDocker,
	}
}

func isDockerAvailable() bool {
	cmd := exec.Command("docker", "version")
	return cmd.Run() == nil
}

func (sm *SandboxManager) Create(ctx context.Context, taskID string) (*Sandbox, error) {
	sandboxID := uuid.New().String()[:8]

	sandbox := &Sandbox{
		ID:        sandboxID,
		TaskID:    taskID,
		Status:    "creating",
		CreatedAt: time.Now(),
	}

	if sm.useDocker {
		workDir, err := os.MkdirTemp("", fmt.Sprintf("sandbox-%s-", sandboxID))
		if err != nil {
			return nil, fmt.Errorf("failed to create temp dir: %w", err)
		}
		sandbox.WorkDir = workDir

		args := []string{
			"create",
			"--name", fmt.Sprintf("gooseek-sandbox-%s", sandboxID),
			"-v", fmt.Sprintf("%s:%s", workDir, sm.cfg.WorkDir),
			"-w", sm.cfg.WorkDir,
			"--memory", sm.cfg.MemoryLimit,
			"--cpus", sm.cfg.CPULimit,
		}

		if !sm.cfg.AllowNetwork {
			args = append(args, "--network", "none")
		}

		args = append(args, sm.cfg.Image, "tail", "-f", "/dev/null")

		cmd := exec.CommandContext(ctx, "docker", args...)
		output, err := cmd.CombinedOutput()
		if err != nil {
			os.RemoveAll(workDir)
			return nil, fmt.Errorf("failed to create container: %w - %s", err, string(output))
		}

		sandbox.ContainerID = strings.TrimSpace(string(output))

		startCmd := exec.CommandContext(ctx, "docker", "start", sandbox.ContainerID)
		if err := startCmd.Run(); err != nil {
			sm.cleanupContainer(sandbox)
			return nil, fmt.Errorf("failed to start container: %w", err)
		}
	} else {
		workDir, err := os.MkdirTemp("", fmt.Sprintf("sandbox-%s-", sandboxID))
		if err != nil {
			return nil, fmt.Errorf("failed to create temp dir: %w", err)
		}
		sandbox.WorkDir = workDir
	}

	sandbox.Status = "running"

	sm.mu.Lock()
	sm.sandboxes[sandboxID] = sandbox
	sm.mu.Unlock()

	return sandbox, nil
}

func (sm *SandboxManager) Execute(ctx context.Context, sandbox *Sandbox, code string, lang string) (*SandboxResult, error) {
	ctx, cancel := context.WithTimeout(ctx, sm.cfg.Timeout)
	defer cancel()

	startTime := time.Now()

	filename, err := sm.writeCodeFile(sandbox, code, lang)
	if err != nil {
		return nil, err
	}

	var cmd *exec.Cmd
	var stdout, stderr bytes.Buffer

	if sm.useDocker {
		runCmd := sm.getRunCommand(lang, filename)
		cmd = exec.CommandContext(ctx, "docker", "exec", sandbox.ContainerID, "sh", "-c", runCmd)
	} else {
		cmd = sm.getLocalCommand(ctx, lang, filepath.Join(sandbox.WorkDir, filename))
	}

	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err = cmd.Run()
	exitCode := 0
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			exitCode = exitErr.ExitCode()
		} else if ctx.Err() == context.DeadlineExceeded {
			return &SandboxResult{
				Stderr:   "Execution timeout exceeded",
				ExitCode: -1,
				Duration: time.Since(startTime),
			}, nil
		}
	}

	files, _ := sm.collectOutputFiles(sandbox)

	return &SandboxResult{
		Stdout:   stdout.String(),
		Stderr:   stderr.String(),
		ExitCode: exitCode,
		Files:    files,
		Duration: time.Since(startTime),
	}, nil
}

func (sm *SandboxManager) RunCommand(ctx context.Context, sandbox *Sandbox, command string) (*SandboxResult, error) {
	ctx, cancel := context.WithTimeout(ctx, sm.cfg.Timeout)
	defer cancel()

	startTime := time.Now()

	var cmd *exec.Cmd
	var stdout, stderr bytes.Buffer

	if sm.useDocker {
		cmd = exec.CommandContext(ctx, "docker", "exec", sandbox.ContainerID, "sh", "-c", command)
	} else {
		cmd = exec.CommandContext(ctx, "sh", "-c", command)
		cmd.Dir = sandbox.WorkDir
	}

	cmd.Stdout = &stdout
	cmd.Stderr = &stderr

	err := cmd.Run()
	exitCode := 0
	if err != nil {
		if exitErr, ok := err.(*exec.ExitError); ok {
			exitCode = exitErr.ExitCode()
		}
	}

	return &SandboxResult{
		Stdout:   stdout.String(),
		Stderr:   stderr.String(),
		ExitCode: exitCode,
		Duration: time.Since(startTime),
	}, nil
}

func (sm *SandboxManager) WriteFile(ctx context.Context, sandbox *Sandbox, path string, content []byte) error {
	if int64(len(content)) > sm.cfg.MaxFileSize {
		return fmt.Errorf("file size exceeds limit: %d > %d", len(content), sm.cfg.MaxFileSize)
	}

	fullPath := filepath.Join(sandbox.WorkDir, path)
	dir := filepath.Dir(fullPath)

	if err := os.MkdirAll(dir, 0755); err != nil {
		return fmt.Errorf("failed to create directory: %w", err)
	}

	return os.WriteFile(fullPath, content, 0644)
}

func (sm *SandboxManager) ReadFile(ctx context.Context, sandbox *Sandbox, path string) ([]byte, error) {
	fullPath := filepath.Join(sandbox.WorkDir, path)
	return os.ReadFile(fullPath)
}

func (sm *SandboxManager) Destroy(ctx context.Context, sandbox *Sandbox) error {
	sm.mu.Lock()
	delete(sm.sandboxes, sandbox.ID)
	sm.mu.Unlock()

	if sm.useDocker && sandbox.ContainerID != "" {
		sm.cleanupContainer(sandbox)
	}

	if sandbox.WorkDir != "" {
		os.RemoveAll(sandbox.WorkDir)
	}

	return nil
}

func (sm *SandboxManager) cleanupContainer(sandbox *Sandbox) {
	exec.Command("docker", "stop", sandbox.ContainerID).Run()
	exec.Command("docker", "rm", "-f", sandbox.ContainerID).Run()
}

func (sm *SandboxManager) writeCodeFile(sandbox *Sandbox, code string, lang string) (string, error) {
	var filename string
	switch lang {
	case "python", "py":
		filename = "main.py"
	case "javascript", "js", "node":
		filename = "main.js"
	case "typescript", "ts":
		filename = "main.ts"
	case "go", "golang":
		filename = "main.go"
	case "bash", "sh", "shell":
		filename = "script.sh"
	case "ruby", "rb":
		filename = "main.rb"
	default:
		filename = "main.txt"
	}

	fullPath := filepath.Join(sandbox.WorkDir, filename)
	if err := os.WriteFile(fullPath, []byte(code), 0755); err != nil {
		return "", fmt.Errorf("failed to write code file: %w", err)
	}

	return filename, nil
}

func (sm *SandboxManager) getRunCommand(lang, filename string) string {
	switch lang {
	case "python", "py":
		return fmt.Sprintf("python3 %s/%s", sm.cfg.WorkDir, filename)
	case "javascript", "js", "node":
		return fmt.Sprintf("node %s/%s", sm.cfg.WorkDir, filename)
	case "typescript", "ts":
		return fmt.Sprintf("npx ts-node %s/%s", sm.cfg.WorkDir, filename)
	case "go", "golang":
		return fmt.Sprintf("go run %s/%s", sm.cfg.WorkDir, filename)
	case "bash", "sh", "shell":
		return fmt.Sprintf("bash %s/%s", sm.cfg.WorkDir, filename)
	case "ruby", "rb":
		return fmt.Sprintf("ruby %s/%s", sm.cfg.WorkDir, filename)
	default:
		return fmt.Sprintf("cat %s/%s", sm.cfg.WorkDir, filename)
	}
}

func (sm *SandboxManager) getLocalCommand(ctx context.Context, lang, filepath string) *exec.Cmd {
	switch lang {
	case "python", "py":
		return exec.CommandContext(ctx, "python3", filepath)
	case "javascript", "js", "node":
		return exec.CommandContext(ctx, "node", filepath)
	case "go", "golang":
		return exec.CommandContext(ctx, "go", "run", filepath)
	case "bash", "sh", "shell":
		return exec.CommandContext(ctx, "bash", filepath)
	case "ruby", "rb":
		return exec.CommandContext(ctx, "ruby", filepath)
	default:
		return exec.CommandContext(ctx, "cat", filepath)
	}
}

func (sm *SandboxManager) collectOutputFiles(sandbox *Sandbox) (map[string][]byte, error) {
	files := make(map[string][]byte)

	err := filepath.Walk(sandbox.WorkDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}
		if info.IsDir() {
			return nil
		}

		relPath, err := filepath.Rel(sandbox.WorkDir, path)
		if err != nil {
			return nil
		}

		if strings.HasPrefix(relPath, "main.") || strings.HasPrefix(relPath, "script.") {
			return nil
		}

		if info.Size() > sm.cfg.MaxFileSize {
			return nil
		}

		content, err := os.ReadFile(path)
		if err != nil {
			return nil
		}

		files[relPath] = content
		return nil
	})

	return files, err
}

func (sm *SandboxManager) ListSandboxes() []*Sandbox {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	result := make([]*Sandbox, 0, len(sm.sandboxes))
	for _, s := range sm.sandboxes {
		result = append(result, s)
	}
	return result
}

func (sm *SandboxManager) GetSandbox(id string) (*Sandbox, bool) {
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	s, ok := sm.sandboxes[id]
	return s, ok
}

func (sm *SandboxManager) CopyToContainer(ctx context.Context, sandbox *Sandbox, src string, dst string) error {
	if !sm.useDocker {
		srcData, err := os.ReadFile(src)
		if err != nil {
			return err
		}
		return sm.WriteFile(ctx, sandbox, dst, srcData)
	}

	cmd := exec.CommandContext(ctx, "docker", "cp", src, fmt.Sprintf("%s:%s", sandbox.ContainerID, dst))
	return cmd.Run()
}

func (sm *SandboxManager) CopyFromContainer(ctx context.Context, sandbox *Sandbox, src string, dst string) error {
	if !sm.useDocker {
		srcPath := filepath.Join(sandbox.WorkDir, src)
		srcData, err := os.ReadFile(srcPath)
		if err != nil {
			return err
		}
		return os.WriteFile(dst, srcData, 0644)
	}

	cmd := exec.CommandContext(ctx, "docker", "cp", fmt.Sprintf("%s:%s", sandbox.ContainerID, src), dst)
	return cmd.Run()
}

func (sm *SandboxManager) StreamLogs(ctx context.Context, sandbox *Sandbox) (io.ReadCloser, error) {
	if !sm.useDocker {
		return nil, fmt.Errorf("streaming not supported without Docker")
	}

	cmd := exec.CommandContext(ctx, "docker", "logs", "-f", sandbox.ContainerID)
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, err
	}

	if err := cmd.Start(); err != nil {
		return nil, err
	}

	return stdout, nil
}