mirror of
https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git
synced 2025-06-04 11:24:00 +08:00

### What problem does this PR solve? Add code_executor_manager. #4977. ### Type of change - [x] New Feature (non-breaking change which adds functionality)
437 lines
13 KiB
Python
437 lines
13 KiB
Python
#
|
|
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
import base64
|
|
import os
|
|
import textwrap
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from enum import Enum
|
|
from typing import Dict, Optional
|
|
|
|
import requests
|
|
from pydantic import BaseModel
|
|
|
|
API_URL = os.getenv("SANDBOX_API_URL", "http://localhost:9385/run")
|
|
TIMEOUT = 15
|
|
MAX_WORKERS = 5
|
|
|
|
|
|
class ResultStatus(str, Enum):
|
|
SUCCESS = "success"
|
|
PROGRAM_ERROR = "program_error"
|
|
RESOURCE_LIMIT_EXCEEDED = "resource_limit_exceeded"
|
|
UNAUTHORIZED_ACCESS = "unauthorized_access"
|
|
RUNTIME_ERROR = "runtime_error"
|
|
PROGRAM_RUNNER_ERROR = "program_runner_error"
|
|
|
|
|
|
class ResourceLimitType(str, Enum):
|
|
TIME = "time"
|
|
MEMORY = "memory"
|
|
OUTPUT = "output"
|
|
|
|
|
|
class UnauthorizedAccessType(str, Enum):
|
|
DISALLOWED_SYSCALL = "disallowed_syscall"
|
|
FILE_ACCESS = "file_access"
|
|
NETWORK_ACCESS = "network_access"
|
|
|
|
|
|
class RuntimeErrorType(str, Enum):
|
|
SIGNALLED = "signalled"
|
|
NONZERO_EXIT = "nonzero_exit"
|
|
|
|
|
|
class ExecutionResult(BaseModel):
|
|
status: ResultStatus
|
|
stdout: str
|
|
stderr: str
|
|
exit_code: int
|
|
detail: Optional[str] = None
|
|
resource_limit_type: Optional[ResourceLimitType] = None
|
|
unauthorized_access_type: Optional[UnauthorizedAccessType] = None
|
|
runtime_error_type: Optional[RuntimeErrorType] = None
|
|
|
|
|
|
class TestResult(BaseModel):
|
|
name: str
|
|
passed: bool
|
|
duration: float
|
|
expected_failure: bool = False
|
|
result: Optional[ExecutionResult] = None
|
|
error: Optional[str] = None
|
|
validation_error: Optional[str] = None
|
|
|
|
|
|
def encode_code(code: str) -> str:
|
|
return base64.b64encode(code.encode("utf-8")).decode("utf-8")
|
|
|
|
|
|
def execute_single_test(name: str, code: str, language: str, arguments: dict, expect_fail: bool = False) -> TestResult:
|
|
"""Execute a single test case"""
|
|
payload = {
|
|
"code_b64": encode_code(textwrap.dedent(code)),
|
|
"language": language,
|
|
"arguments": arguments,
|
|
}
|
|
|
|
test_result = TestResult(name=name, passed=False, duration=0, expected_failure=expect_fail)
|
|
|
|
really_processed = False
|
|
try:
|
|
while not really_processed:
|
|
start_time = time.perf_counter()
|
|
|
|
resp = requests.post(API_URL, json=payload, timeout=TIMEOUT)
|
|
resp.raise_for_status()
|
|
response_data = resp.json()
|
|
if response_data["exit_code"] == -429: # too many request
|
|
print(f"[{name}] Reached request limit, retring...")
|
|
time.sleep(0.5)
|
|
continue
|
|
really_processed = True
|
|
|
|
print("-------------------")
|
|
print(f"{name}:\n{response_data}")
|
|
print("-------------------")
|
|
|
|
test_result.duration = time.perf_counter() - start_time
|
|
test_result.result = ExecutionResult(**response_data)
|
|
|
|
# Validate test result expectations
|
|
validate_test_result(name, expect_fail, test_result)
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
test_result.duration = time.perf_counter() - start_time
|
|
test_result.error = f"Request failed: {str(e)}"
|
|
test_result.result = ExecutionResult(
|
|
status=ResultStatus.PROGRAM_RUNNER_ERROR,
|
|
stdout="",
|
|
stderr=str(e),
|
|
exit_code=-999,
|
|
detail="request_failed",
|
|
)
|
|
|
|
return test_result
|
|
|
|
|
|
def validate_test_result(name: str, expect_fail: bool, test_result: TestResult):
|
|
"""Validate if the test result meets expectations"""
|
|
if not test_result.result:
|
|
test_result.passed = False
|
|
test_result.validation_error = "No result returned"
|
|
return
|
|
|
|
test_result.passed = test_result.result.status == ResultStatus.SUCCESS
|
|
# General validation logic
|
|
if expect_fail:
|
|
# Tests expected to fail should return a non-success status
|
|
if test_result.passed:
|
|
test_result.validation_error = "Expected failure but actually succeeded"
|
|
else:
|
|
# Tests expected to succeed should return a success status
|
|
if not test_result.passed:
|
|
test_result.validation_error = f"Unexpected failure (status={test_result.result.status})"
|
|
|
|
|
|
def get_test_cases() -> Dict[str, dict]:
|
|
"""Return test cases (code, whether expected to fail)"""
|
|
return {
|
|
"1 Infinite loop: Should be forcibly terminated": {
|
|
"code": """
|
|
def main():
|
|
while True:
|
|
pass
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"2 Infinite loop: Should be forcibly terminated": {
|
|
"code": """
|
|
def main():
|
|
while True:
|
|
pass
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"3 Infinite loop: Should be forcibly terminated": {
|
|
"code": """
|
|
def main():
|
|
while True:
|
|
pass
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"4 Infinite loop: Should be forcibly terminated": {
|
|
"code": """
|
|
def main():
|
|
while True:
|
|
pass
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"5 Infinite loop: Should be forcibly terminated": {
|
|
"code": """
|
|
def main():
|
|
while True:
|
|
pass
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"6 Infinite loop: Should be forcibly terminated": {
|
|
"code": """
|
|
def main():
|
|
while True:
|
|
pass
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"7 Normal test: Python without dependencies": {
|
|
"code": """
|
|
def main():
|
|
return {"data": "hello, world"}
|
|
""",
|
|
"should_fail": False,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"8 Normal test: Python with pandas, should pass without any error": {
|
|
"code": """
|
|
import pandas as pd
|
|
|
|
def main():
|
|
data = {'Name': ['Alice', 'Bob', 'Charlie'],
|
|
'Age': [25, 30, 35]}
|
|
df = pd.DataFrame(data)
|
|
""",
|
|
"should_fail": False,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"9 Normal test: Nodejs without dependencies, should pass without any error": {
|
|
"code": """
|
|
const https = require('https');
|
|
|
|
async function main(args) {
|
|
return new Promise((resolve, reject) => {
|
|
const req = https.get('https://example.com/', (res) => {
|
|
let data = '';
|
|
|
|
res.on('data', (chunk) => {
|
|
data += chunk;
|
|
});
|
|
|
|
res.on('end', () => {
|
|
clearTimeout(timeout);
|
|
console.log('Body:', data);
|
|
resolve(data);
|
|
});
|
|
});
|
|
|
|
const timeout = setTimeout(() => {
|
|
req.destroy(new Error('Request timeout after 10s'));
|
|
}, 10000);
|
|
|
|
req.on('error', (err) => {
|
|
clearTimeout(timeout);
|
|
console.error('Error:', err.message);
|
|
reject(err);
|
|
});
|
|
});
|
|
}
|
|
|
|
module.exports = { main };
|
|
""",
|
|
"should_fail": False,
|
|
"arguments": {},
|
|
"language": "nodejs",
|
|
},
|
|
"10 Normal test: Nodejs with axios, should pass without any error": {
|
|
"code": """
|
|
const axios = require('axios');
|
|
|
|
async function main(args) {
|
|
try {
|
|
const response = await axios.get('https://example.com/', {
|
|
timeout: 10000
|
|
});
|
|
console.log('Body:', response.data);
|
|
} catch (error) {
|
|
console.error('Error:', error.message);
|
|
}
|
|
}
|
|
|
|
module.exports = { main };
|
|
""",
|
|
"should_fail": False,
|
|
"arguments": {},
|
|
"language": "nodejs",
|
|
},
|
|
"11 Dangerous import: Should fail due to os module import": {
|
|
"code": """
|
|
import os
|
|
|
|
def main():
|
|
pass
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"12 Dangerous import from subprocess: Should fail due to subprocess import": {
|
|
"code": """
|
|
from subprocess import Popen
|
|
|
|
def main():
|
|
pass
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"13 Dangerous call: Should fail due to eval function call": {
|
|
"code": """
|
|
def main():
|
|
eval('os.system("echo hello")')
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"14 Dangerous attribute access: Should fail due to shutil.rmtree": {
|
|
"code": """
|
|
import shutil
|
|
|
|
def main():
|
|
shutil.rmtree('/some/path')
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"15 Dangerous binary operation: Should fail due to unsafe concatenation leading to eval": {
|
|
"code": """
|
|
def main():
|
|
dangerous_string = "os." + "system"
|
|
eval(dangerous_string + '("echo hello")')
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"16 Dangerous function definition: Should fail due to user-defined eval function": {
|
|
"code": """
|
|
def eval_function():
|
|
eval('os.system("echo hello")')
|
|
|
|
def main():
|
|
eval_function()
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
"17 Memory exhaustion(256m): Should fail due to exceeding memory limit(try to allocate 300m)": {
|
|
"code": """
|
|
def main():
|
|
x = ['a' * 1024 * 1024] * 300 # 300MB
|
|
""",
|
|
"should_fail": True,
|
|
"arguments": {},
|
|
"language": "python",
|
|
},
|
|
}
|
|
|
|
|
|
def print_test_report(results: Dict[str, TestResult]):
|
|
print("\n=== 🔍 Test Report ===")
|
|
|
|
max_name_len = max(len(name) for name in results)
|
|
|
|
for name, result in results.items():
|
|
status = "✅" if result.passed else "❌"
|
|
if result.expected_failure:
|
|
status = "⚠️" if result.passed else "✓" # Expected failure case
|
|
|
|
print(f"{status} {name.ljust(max_name_len)} {result.duration:.2f}s")
|
|
|
|
if result.error:
|
|
print(f" REQUEST ERROR: {result.error}")
|
|
if result.validation_error:
|
|
print(f" VALIDATION ERROR: {result.validation_error}")
|
|
|
|
if result.result and not result.passed:
|
|
print(f" STATUS: {result.result.status}")
|
|
if result.result.stderr:
|
|
print(f" STDERR: {result.result.stderr[:200]}...")
|
|
if result.result.detail:
|
|
print(f" DETAIL: {result.result.detail}")
|
|
|
|
passed = sum(1 for r in results.values() if ((not r.expected_failure and r.passed) or (r.expected_failure and not r.passed)))
|
|
failed = len(results) - passed
|
|
|
|
print("\n=== 📊 Statistics ===")
|
|
print(f"✅ Passed: {passed}")
|
|
print(f"❌ Failed: {failed}")
|
|
print(f"📌 Total: {len(results)}")
|
|
|
|
|
|
def main():
|
|
print(f"🔐 Starting sandbox security tests (API: {API_URL})")
|
|
print(f"🚀 Concurrent threads: {MAX_WORKERS}")
|
|
|
|
test_cases = get_test_cases()
|
|
results = {}
|
|
|
|
with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
|
|
futures = {}
|
|
for name, detail in test_cases.items():
|
|
# ✅ Log when a task is submitted
|
|
print(f"✅ Task submitted: {name}")
|
|
time.sleep(0.4)
|
|
future = executor.submit(execute_single_test, name, detail["code"], detail["language"], detail["arguments"], detail["should_fail"])
|
|
futures[future] = name
|
|
|
|
print("\n=== 🚦 Test Progress ===")
|
|
for i, future in enumerate(as_completed(futures)):
|
|
name = futures[future]
|
|
print(f" {i + 1}/{len(test_cases)} completed: {name}")
|
|
try:
|
|
results[name] = future.result()
|
|
except Exception as e:
|
|
print(f"⚠️ Test {name} execution exception: {str(e)}")
|
|
results[name] = TestResult(name=name, passed=False, duration=0, error=f"Execution exception: {str(e)}")
|
|
|
|
print_test_report(results)
|
|
|
|
if any(not r.passed and not r.expected_failure for r in results.values()):
|
|
exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|