package cmd
import (
"archive/tar" // Needed for CopyToContainer
"bytes"
"context"
"encoding/json"
"fmt"
"io" // Needed for CopyToContainer
"os" // Needed for CopyToContainer path manipulation
"path/filepath"
"strings"
"time"
"github.com/docker/docker/api/types/container"
// "github.com/docker/docker/api/types/mount" // REMOVED mount import
"github.com/docker/docker/client"
"github.com/docker/docker/pkg/stdcopy"
"github.com/spf13/cobra"
)
// Result structure for JSON output
type ExecutionResult struct {
Status string `json:"status"` // "success", "timeout", "error"
ExitCode int64 `json:"exitCode"`
Stdout string `json:"stdout"`
Stderr string `json:"stderr"`
DurationMs int64 `json:"durationMs"`
Error string `json:"error,omitempty"` // Only present if status is "error" or "timeout"
ImageTag string `json:"imageTag,omitempty"` // Only if --save-image was used
}
// Flags for the run command
var (
language string
dirPath string
timeout time.Duration
runtime string
readOnlyRootfs bool
networkMode string
capDrop []string
pidsLimit int64
memoryLimit int64 // in MB
cpusLimit float64
saveImageTag string
saveDockerfile string
containerWorkdir string // Will be read from server config later, needed for copy target
entrypointFile string // Added flag for explicit entrypoint filename
envVars []string // New flag for environment variables
)
// runCmd represents the run command
var runCmd = &cobra.Command{
Use: "run",
Short: "Run code in an isolated environment",
Long: `Executes code provided in a specified directory within a secure container.
Supports Python, Go, and JavaScript. Enforces security constraints and resource limits.`,
RunE: func(cmd *cobra.Command, args []string) error {
startTime := time.Now()
// Basic validation
if dirPath == "" {
return fmt.Errorf("--dir flag is required")
}
// No Abs path needed for copy source
// absDirPath, err := filepath.Abs(dirPath)
// if err != nil {
// return fmt.Errorf("failed to get absolute path for --dir: %w", err)
// }
// dirPath = absDirPath
if language == "" {
return fmt.Errorf("--language flag is required")
}
if containerWorkdir == "" {
return fmt.Errorf("--container-workdir flag is required (this should be passed by the caller)")
}
if entrypointFile == "" { // Add check for required entrypoint flag
return fmt.Errorf("--entrypoint flag is required")
}
// --- Log received flags ---
fmt.Fprintf(os.Stderr, "[isolator CLI DEBUG] Received flags:\n")
fmt.Fprintf(os.Stderr, " --language: %s\n", language)
fmt.Fprintf(os.Stderr, " --dir: %s\n", dirPath)
fmt.Fprintf(os.Stderr, " --entrypoint: %s\n", entrypointFile) // Log the crucial flag
fmt.Fprintf(os.Stderr, " --timeout: %s\n", timeout)
fmt.Fprintf(os.Stderr, " --container-workdir: %s\n", containerWorkdir)
// --- End Log ---
var execResult ExecutionResult
execResult.Status = "error" // Default to error
// --- Docker Execution Logic ---
ctx := context.Background()
cli, err := client.NewClientWithOpts(client.FromEnv, client.WithAPIVersionNegotiation())
if err != nil {
execResult.Error = fmt.Sprintf("Failed to create docker client: %s", err.Error())
outputJSON(execResult, startTime) // Attempt to output error
return fmt.Errorf("failed to create docker client: %w", err)
}
defer cli.Close()
// Get base image based on language
image, err := getLangImage(language) // Changed helper function
if err != nil {
execResult.Error = err.Error()
outputJSON(execResult, startTime)
return err
}
// Construct the command using the provided entrypointFile flag
var execCmd string // The actual command string to run
var containerCmd []string // The command string slice
switch strings.ToLower(language) {
case "python":
containerCmd = []string{"python", entrypointFile} // Use relative path
execCmd = fmt.Sprintf("python %s", entrypointFile)
case "go":
containerCmd = []string{"go", "run", entrypointFile} // Keep relative for Go
execCmd = fmt.Sprintf("go run %s", entrypointFile)
case "javascript", "js", "node":
containerCmd = []string{"node", entrypointFile} // Use relative path
execCmd = fmt.Sprintf("node %s", entrypointFile)
default:
return fmt.Errorf("logic error: unsupported language %s reached command construction", language)
}
// Check if entrypoint file actually exists IN THE SOURCE DIR
entrypointPathOnHost := filepath.Join(dirPath, entrypointFile)
if _, err := os.Stat(entrypointPathOnHost); os.IsNotExist(err) {
errMsg := fmt.Sprintf("Entrypoint file '%s' not found in source directory '%s'", entrypointFile, dirPath)
execResult.Error = errMsg
outputJSON(execResult, startTime)
return fmt.Errorf(errMsg)
}
// Wrap the command in a shell to add debugging 'ls /' and 'ls /workspace'
// Keep this debug for now, though ls might fail initially before copy
debugCmd := fmt.Sprintf("echo '--- ls / ---' && ls -al / && echo '--- ls /workspace (before copy) ---' && ls -al %s; echo '--- Running Command ---' && %s", containerWorkdir, execCmd)
containerShellCmd := []string{"sh", "-c", debugCmd}
// Create container config (NO WorkingDir initially, command uses relative path assuming copy destination)
containerCfg := &container.Config{
Image: image,
Cmd: containerShellCmd, // Use shell wrapper for debug ls
WorkingDir: containerWorkdir, // Set working dir where files will be copied
Env: envVars, // Pass environment variables from flag
Tty: false,
AttachStdout: true,
AttachStderr: true,
// User: default non-root user
}
// Create host config (NO Mounts)
hostCfg := &container.HostConfig{
Resources: container.Resources{
Memory: memoryLimit * 1024 * 1024,
NanoCPUs: int64(cpusLimit * 1e9),
PidsLimit: &pidsLimit,
},
NetworkMode: container.NetworkMode(networkMode),
ReadonlyRootfs: readOnlyRootfs, // Keep rootfs read-only if possible
// Mounts: REMOVED
CapDrop: capDrop,
SecurityOpt: []string{"no-new-privileges:true"},
}
// Create execution context with timeout
execCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
// Create container
resp, err := cli.ContainerCreate(execCtx, containerCfg, hostCfg, nil, nil, "")
if err != nil {
execResult.Error = fmt.Sprintf("Failed to create container: %s", err.Error())
outputJSON(execResult, startTime)
return fmt.Errorf("failed to create container: %w", err)
}
containerID := resp.ID
// Ensure container cleanup
defer func() {
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cleanupCancel()
_ = cli.ContainerRemove(cleanupCtx, containerID, container.RemoveOptions{Force: true, RemoveVolumes: true})
}()
// --- Copy files into the STARTING container ---
// Create a tar archive of the source directory contents
tarReader, err := createTarArchive(dirPath)
if err != nil {
execResult.Error = fmt.Sprintf("Failed to create tar archive for copying: %s", err.Error())
outputJSON(execResult, startTime)
// Try to remove container before erroring
cli.ContainerRemove(context.Background(), containerID, container.RemoveOptions{Force: true})
return fmt.Errorf("failed to create tar archive: %w", err)
}
// No need to defer tarReader.Close() because CopyToContainer consumes it
// Copy the archive to the container's working directory
// Use a separate context for copy operation if needed, but execCtx should be okay
err = cli.CopyToContainer(execCtx, containerID, containerWorkdir, tarReader, container.CopyToContainerOptions{
AllowOverwriteDirWithFile: false,
CopyUIDGID: false, // Let ownership be default inside container
})
if err != nil {
execResult.Error = fmt.Sprintf("Failed to copy files to container: %s", err.Error())
outputJSON(execResult, startTime)
// Try to remove container before erroring
cli.ContainerRemove(context.Background(), containerID, container.RemoveOptions{Force: true})
return fmt.Errorf("failed to copy files to container: %w", err)
}
fmt.Fprintf(os.Stderr, "[isolator CLI DEBUG] Copied files from %s to container %s:%s\n", dirPath, containerID, containerWorkdir)
// Start the container (AFTER copying files)
if err := cli.ContainerStart(execCtx, containerID, container.StartOptions{}); err != nil {
execResult.Error = fmt.Sprintf("Failed to start container: %s", err.Error())
outputJSON(execResult, startTime)
return fmt.Errorf("failed to start container: %w", err)
}
// --- Wait for completion and gather results ---
statusCh, errCh := cli.ContainerWait(execCtx, containerID, container.WaitConditionNotRunning)
select {
case err := <-errCh:
if err != nil && !strings.Contains(err.Error(), "context deadline exceeded") {
execResult.Error = fmt.Sprintf("Error waiting for container: %s", err.Error())
}
case status := <-statusCh:
execResult.ExitCode = status.StatusCode
if status.StatusCode == 0 {
execResult.Status = "success"
} else {
execResult.Status = "error"
}
case <-execCtx.Done():
execResult.Status = "timeout"
execResult.Error = fmt.Sprintf("Execution timed out after %s", timeout)
}
// --- Get logs regardless of status ---
logCtx, logCancel := context.WithTimeout(context.Background(), 10*time.Second)
defer logCancel()
logsReader, logErr := cli.ContainerLogs(logCtx, containerID, container.LogsOptions{
ShowStdout: true,
ShowStderr: true,
Timestamps: false,
Follow: false,
})
// Handle log retrieval error AFTER processing main execution status/error
var stdoutStr, stderrStr string
if logErr == nil {
defer logsReader.Close()
stdoutBuf := new(bytes.Buffer)
stderrBuf := new(bytes.Buffer)
_, copyErr := stdcopy.StdCopy(stdoutBuf, stderrBuf, logsReader)
if copyErr != nil {
fmt.Fprintf(os.Stderr, "[isolator CLI DEBUG] Error reading container logs: %s\n", copyErr.Error())
}
stdoutStr = stdoutBuf.String()
stderrStr = stderrBuf.String()
} else {
fmt.Fprintf(os.Stderr, "[isolator CLI DEBUG] Error getting container logs: %s\n", logErr.Error())
}
execResult.Stdout = stdoutStr
execResult.Stderr = stderrStr
// Refine error message if needed
if execResult.Status == "error" && execResult.Error == "" {
trimmedStderr := strings.TrimSpace(stderrStr)
// Check if stderr contains the debug output prefix, ignore it for the main error
debugPrefix := "echo '--- ls / ---'"
if strings.HasPrefix(trimmedStderr, debugPrefix) {
// Find the actual error message after the debug output
parts := strings.SplitN(trimmedStderr, "--- Running Command ---", 2)
if len(parts) > 1 {
actualStderr := strings.TrimSpace(parts[1])
if actualStderr != "" {
execResult.Error = actualStderr
}
}
} else if trimmedStderr != "" {
execResult.Error = trimmedStderr
}
if execResult.Error == "" { // Fallback if stderr was empty or only debug
execResult.Error = fmt.Sprintf("Container exited with status code %d", execResult.ExitCode)
}
}
// --- Handle saving image (optional) ---
if execResult.Status == "success" && saveImageTag != "" {
commitCtx, commitCancel := context.WithTimeout(context.Background(), 30*time.Second)
defer commitCancel()
_, err := cli.ContainerCommit(commitCtx, containerID, container.CommitOptions{Reference: saveImageTag})
if err != nil {
fmt.Fprintf(os.Stderr, "Warning: Failed to commit container image %s: %s\n", saveImageTag, err.Error())
} else {
execResult.ImageTag = saveImageTag
}
}
// --- Handle saving dockerfile (optional) ---
if saveDockerfile != "" {
baseImage, _ := getLangImage(language)
dockerfileContent := fmt.Sprintf("FROM %s\nWORKDIR %s\n# Add COPY commands for files in %s if needed\nCMD [\"%s\"]",
baseImage, containerWorkdir, containerWorkdir, strings.Join(containerCmd, "\", \"")) // Use containerCmd
err := os.WriteFile(saveDockerfile, []byte(dockerfileContent), 0644)
if err != nil {
fmt.Fprintf(os.Stderr, "Warning: Failed to save Dockerfile to %s: %s\n", saveDockerfile, err.Error())
}
}
// --- Finalize and output JSON ---
outputJSON(execResult, startTime)
return nil // Success/failure indicated in JSON output
},
}
// Helper function to get Docker image based on language
func getLangImage(language string) (image string, err error) {
switch strings.ToLower(language) {
case "python":
image = "python:3.11-alpine"
case "go":
image = "golang:1.21-alpine"
case "javascript", "js", "node":
image = "node:20-alpine"
default:
err = fmt.Errorf("unsupported language: %s", language)
}
return
}
// createTarArchive creates a reader for a tar archive containing the files in the source directory.
func createTarArchive(sourceDir string) (io.ReadCloser, error) {
pr, pw := io.Pipe() // Use a pipe to stream tar data
tw := tar.NewWriter(pw)
go func() {
defer pw.Close() // Close writer when done
defer tw.Close() // Close tar writer
// Walk the source directory
err := filepath.Walk(sourceDir, func(file string, fi os.FileInfo, err error) error {
if err != nil {
return err
}
// Skip the root directory itself
if file == sourceDir {
return nil
}
// Create tar header
header, err := tar.FileInfoHeader(fi, fi.Name()) // Use relative path within archive
if err != nil {
return fmt.Errorf("failed to create tar header for %s: %w", file, err)
}
// Adjust header name to be relative to the sourceDir root for tar structure
relPath, err := filepath.Rel(sourceDir, file)
if err != nil {
return fmt.Errorf("failed to get relative path for %s: %w", file, err)
}
header.Name = relPath // Use relative path in tar archive
// Write header
if err := tw.WriteHeader(header); err != nil {
return fmt.Errorf("failed to write tar header for %s: %w", file, err)
}
// If not a directory, write file content
if !fi.IsDir() {
f, err := os.Open(file)
if err != nil {
return fmt.Errorf("failed to open file %s: %w", file, err)
}
defer f.Close()
if _, err := io.Copy(tw, f); err != nil {
return fmt.Errorf("failed to copy file content for %s: %w", file, err)
}
}
return nil
})
if err != nil {
// Close the pipe writer with an error if Walk failed
_ = pw.CloseWithError(fmt.Errorf("error during tar archiving: %w", err))
}
}()
return pr, nil // Return the reader part of the pipe
}
// Helper function to output JSON result
func outputJSON(result ExecutionResult, startTime time.Time) {
result.DurationMs = time.Since(startTime).Milliseconds()
jsonOutput, err := json.Marshal(result)
if err != nil {
fmt.Fprintf(os.Stderr, "Internal Error: Failed to marshal result to JSON: %s\n", err.Error())
fmt.Printf(`{"status": "error", "error": "Internal Error: Failed to marshal result to JSON", "durationMs": %d}`, result.DurationMs)
return
}
fmt.Println(string(jsonOutput))
}
func init() {
rootCmd.AddCommand(runCmd)
// Define flags
runCmd.Flags().StringVarP(&language, "language", "l", "", "Language of the code (python, go, javascript)")
runCmd.Flags().StringVarP(&dirPath, "dir", "d", "", "Path to the host directory containing input files")
runCmd.Flags().DurationVarP(&timeout, "timeout", "t", 60*time.Second, "Execution timeout")
runCmd.Flags().StringVar(&runtime, "runtime", "docker", "Container runtime to use (docker)")
runCmd.Flags().BoolVar(&readOnlyRootfs, "read-only", true, "Mount container root filesystem as read-only")
runCmd.Flags().StringVar(&networkMode, "network", "none", "Container network mode (e.g., none, bridge)")
runCmd.Flags().StringSliceVar(&capDrop, "cap-drop", []string{"ALL"}, "Linux capabilities to drop")
runCmd.Flags().Int64Var(&pidsLimit, "pids-limit", 64, "Maximum number of processes allowed in the container")
runCmd.Flags().Int64VarP(&memoryLimit, "memory", "m", 256, "Memory limit in MB")
runCmd.Flags().Float64VarP(&cpusLimit, "cpus", "c", 0.5, "CPU limit (e.g., 0.5 for half a core)")
runCmd.Flags().StringVar(&containerWorkdir, "container-workdir", "", "Working directory inside the container (target for copy)")
runCmd.Flags().StringVar(&saveImageTag, "save-image", "", "Tag to save the container image as on success (e.g., my-image:v1)")
runCmd.Flags().StringVar(&saveDockerfile, "save-dockerfile", "", "Path to save the base Dockerfile used")
runCmd.Flags().StringVarP(&entrypointFile, "entrypoint", "e", "", "Filename of the entrypoint script within the directory")
runCmd.Flags().StringSliceVar(&envVars, "env", []string{}, "Environment variables to set in the container (e.g., --env KEY1=VALUE1 --env KEY2=VALUE2)") // Add env flag
// Mark required flags
_ = runCmd.MarkFlagRequired("language")
_ = runCmd.MarkFlagRequired("dir")
_ = runCmd.MarkFlagRequired("container-workdir")
_ = runCmd.MarkFlagRequired("entrypoint")
}