Isolator MCP Server

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") }