ragflow/sandbox/tests/sandbox_security_tests_full.py
Yongteng Lei 2d7c1368f0
Feat: add code_executor_manager (#7814)
### What problem does this PR solve?

Add code_executor_manager. #4977.

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
2025-05-23 16:33:38 +08:00

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()