257 lines
8.0 KiB
Python
257 lines
8.0 KiB
Python
"""
|
|
Code Execution Sandbox Server
|
|
Runs inside a Docker container, accepts code execution requests.
|
|
Supports: Python, JavaScript
|
|
"""
|
|
import asyncio
|
|
import subprocess
|
|
import tempfile
|
|
import os
|
|
import time
|
|
import resource
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel, Field
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
app = FastAPI(title="Code Execution Sandbox", version="1.0.0")
|
|
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
|
|
|
|
# Use /tmp for temp files (writable by executor user)
|
|
# /workspace is a tmpfs mount owned by root
|
|
WORKSPACE = Path("/tmp")
|
|
WORKSPACE.mkdir(exist_ok=True)
|
|
|
|
# Max code length: 50KB
|
|
MAX_CODE_LENGTH = 50 * 1024
|
|
# Default timeout: 30s, max: 120s
|
|
DEFAULT_TIMEOUT = 30
|
|
MAX_TIMEOUT = 120
|
|
|
|
|
|
class ExecuteRequest(BaseModel):
|
|
language: str = Field(..., description="Programming language: python, javascript")
|
|
code: str = Field(..., description="Source code to execute")
|
|
stdin: str = Field(default="", description="Standard input for the program")
|
|
timeout: int = Field(default=DEFAULT_TIMEOUT, ge=1, le=MAX_TIMEOUT, description="Execution timeout in seconds")
|
|
|
|
|
|
class ExecuteResponse(BaseModel):
|
|
success: bool
|
|
stdout: str = ""
|
|
stderr: str = ""
|
|
exit_code: int = 0
|
|
execution_time_ms: float = 0
|
|
error: Optional[str] = None
|
|
truncated: bool = False
|
|
|
|
|
|
@app.get("/health")
|
|
async def health():
|
|
return {"status": "healthy", "languages": ["python", "javascript"]}
|
|
|
|
|
|
@app.get("/info")
|
|
async def info():
|
|
"""Return sandbox runtime info."""
|
|
py_ver = subprocess.run(["python3", "--version"], capture_output=True, text=True).stdout.strip()
|
|
try:
|
|
node_ver = subprocess.run(["node", "--version"], capture_output=True, text=True).stdout.strip()
|
|
except Exception:
|
|
node_ver = "not installed"
|
|
return {
|
|
"python": py_ver,
|
|
"nodejs": node_ver,
|
|
"workspace": str(WORKSPACE),
|
|
"max_code_length": MAX_CODE_LENGTH,
|
|
"max_timeout": MAX_TIMEOUT,
|
|
}
|
|
|
|
|
|
@app.post("/execute", response_model=ExecuteResponse)
|
|
async def execute_code(req: ExecuteRequest):
|
|
"""Execute code in a sandboxed subprocess."""
|
|
language = req.language.lower().strip()
|
|
code = req.code
|
|
stdin_input = req.stdin
|
|
timeout = min(req.timeout, MAX_TIMEOUT)
|
|
|
|
# Validate language
|
|
if language not in ("python", "py", "python3", "javascript", "js", "node"):
|
|
raise HTTPException(status_code=400, detail=f"Unsupported language: {language}. Supported: python, javascript")
|
|
|
|
# Normalize language
|
|
if language in ("py", "python3"):
|
|
language = "python"
|
|
if language in ("js", "node"):
|
|
language = "javascript"
|
|
|
|
# Validate code length
|
|
if len(code.encode("utf-8")) > MAX_CODE_LENGTH:
|
|
raise HTTPException(status_code=400, detail=f"Code too long. Max {MAX_CODE_LENGTH // 1024}KB")
|
|
|
|
start_time = time.time()
|
|
|
|
try:
|
|
if language == "python":
|
|
result = await _run_python(code, stdin_input, timeout)
|
|
elif language == "javascript":
|
|
result = await _run_javascript(code, stdin_input, timeout)
|
|
else:
|
|
raise HTTPException(status_code=400, detail=f"Unsupported language: {language}")
|
|
except asyncio.TimeoutError:
|
|
execution_time = (time.time() - start_time) * 1000
|
|
return ExecuteResponse(
|
|
success=False,
|
|
stdout="",
|
|
stderr="",
|
|
exit_code=-1,
|
|
execution_time_ms=execution_time,
|
|
error=f"Execution timed out after {timeout}s",
|
|
)
|
|
|
|
execution_time_ms = (time.time() - start_time) * 1000
|
|
result["execution_time_ms"] = round(execution_time_ms, 2)
|
|
return result
|
|
|
|
|
|
async def _run_python(code: str, stdin_input: str, timeout: int) -> dict:
|
|
"""Execute Python code."""
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="w", suffix=".py", dir=WORKSPACE, delete=False, encoding="utf-8"
|
|
) as f:
|
|
f.write(code)
|
|
script_path = f.name
|
|
|
|
try:
|
|
# Use resource limits
|
|
def set_limits():
|
|
# Memory limit: 256MB
|
|
resource.setrlimit(resource.RLIMIT_AS, (256 * 1024 * 1024, 256 * 1024 * 1024))
|
|
# CPU time limit: timeout + 5s grace
|
|
resource.setrlimit(resource.RLIMIT_CPU, (timeout + 5, timeout + 5))
|
|
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"python3",
|
|
"-u", # unbuffered
|
|
script_path,
|
|
stdin=asyncio.subprocess.PIPE,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
preexec_fn=set_limits,
|
|
)
|
|
|
|
try:
|
|
stdout, stderr = await asyncio.wait_for(
|
|
proc.communicate(input=stdin_input.encode("utf-8") if stdin_input else None),
|
|
timeout=timeout,
|
|
)
|
|
except asyncio.TimeoutError:
|
|
proc.kill()
|
|
await proc.wait()
|
|
raise
|
|
|
|
stdout_str = stdout.decode("utf-8", errors="replace")
|
|
stderr_str = stderr.decode("utf-8", errors="replace")
|
|
|
|
# Truncate output to 500KB
|
|
truncated = False
|
|
if len(stdout_str.encode("utf-8")) > 500 * 1024:
|
|
stdout_str = stdout_str[:500 * 1024]
|
|
truncated = True
|
|
if len(stderr_str.encode("utf-8")) > 500 * 1024:
|
|
stderr_str = stderr_str[:500 * 1024]
|
|
truncated = True
|
|
|
|
return {
|
|
"success": proc.returncode == 0,
|
|
"stdout": stdout_str,
|
|
"stderr": stderr_str,
|
|
"exit_code": proc.returncode or 0,
|
|
"error": None,
|
|
"truncated": truncated,
|
|
}
|
|
|
|
finally:
|
|
# Cleanup temp file
|
|
try:
|
|
os.unlink(script_path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
async def _run_javascript(code: str, stdin_input: str, timeout: int) -> dict:
|
|
"""Execute JavaScript code via Node.js."""
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="w", suffix=".js", dir=WORKSPACE, delete=False, encoding="utf-8"
|
|
) as f:
|
|
f.write(code)
|
|
script_path = f.name
|
|
|
|
try:
|
|
def set_limits():
|
|
# Only limit CPU time for JS — V8 needs large virtual address space
|
|
# for CodeRange; rely on Docker --memory for actual RSS limit.
|
|
resource.setrlimit(resource.RLIMIT_CPU, (timeout + 5, timeout + 5))
|
|
|
|
# Use Node.js to run the script with tight heap limits
|
|
env = os.environ.copy()
|
|
node_flags = "--max-old-space-size=32 --max-semi-space-size=8"
|
|
env["NODE_OPTIONS"] = node_flags
|
|
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"node",
|
|
"--max-old-space-size=32",
|
|
"--max-semi-space-size=8",
|
|
script_path,
|
|
stdin=asyncio.subprocess.PIPE,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
env=env,
|
|
preexec_fn=set_limits,
|
|
)
|
|
|
|
try:
|
|
stdout, stderr = await asyncio.wait_for(
|
|
proc.communicate(input=stdin_input.encode("utf-8") if stdin_input else None),
|
|
timeout=timeout,
|
|
)
|
|
except asyncio.TimeoutError:
|
|
proc.kill()
|
|
await proc.wait()
|
|
raise
|
|
|
|
stdout_str = stdout.decode("utf-8", errors="replace")
|
|
stderr_str = stderr.decode("utf-8", errors="replace")
|
|
|
|
truncated = False
|
|
if len(stdout_str.encode("utf-8")) > 500 * 1024:
|
|
stdout_str = stdout_str[:500 * 1024]
|
|
truncated = True
|
|
if len(stderr_str.encode("utf-8")) > 500 * 1024:
|
|
stderr_str = stderr_str[:500 * 1024]
|
|
truncated = True
|
|
|
|
return {
|
|
"success": proc.returncode == 0,
|
|
"stdout": stdout_str,
|
|
"stderr": stderr_str,
|
|
"exit_code": proc.returncode or 0,
|
|
"error": None,
|
|
"truncated": truncated,
|
|
}
|
|
|
|
finally:
|
|
try:
|
|
os.unlink(script_path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|