返回提交历史
Modified
etc/unittest/mcp.py
+80
-0
Modified
g4f/mcp/pa_provider.py
+167
-22
Modified
g4f/mcp/tools.py
+38
-4
XFEstudio/gpt4free
Harden sandbox security: execution timeout, max recursion depth, output size cap
Agent-Logs-Url: https://github.com/xtekky/gpt4free/sessions/41556926-6205-4207-b36b-e10e22a8b87e Co-authored-by: hlohaus <983577+hlohaus@users.noreply.github.com>
56163491
代码差异
3 个文件
+285
-26
@@ -477,3 +477,83 @@ class TestSafeMode(unittest.IsolatedAsyncioTestCase):
477
477
result = await tool.execute({"code": "import math\nresult = math.factorial(5)"})
478
478
self.assertTrue(result.get("success"))
479
479
self.assertEqual(result.get("result"), 120)
480
481
class TestSecurityHardening(unittest.IsolatedAsyncioTestCase):
482
"""Tests for execution timeout, recursion depth, and output size limits."""
483
484
def test_execution_timeout(self):
485
"""Infinite loop is interrupted by the timeout."""
486
import time
487
start = time.time()
488
r = execute_safe_code("while True: pass", timeout=0.5)
489
elapsed = time.time() - start
490
self.assertFalse(r.success)
491
self.assertIn("timed out", r.error.lower())
492
self.assertLess(elapsed, 3.0, "Should have returned within 3 s")
493
494
def test_execution_continues_after_timeout(self):
495
"""The sandbox is usable again after a previous execution timed out."""
496
execute_safe_code("while True: pass", timeout=0.3)
497
r = execute_safe_code("result = 'ok'", timeout=5.0)
498
self.assertTrue(r.success)
499
self.assertEqual(r.result, "ok")
500
501
def test_recursion_depth_limit(self):
502
"""Deep recursion is blocked by the max_depth parameter."""
503
r = execute_safe_code(
504
"def f(n): return f(n + 1)\nf(0)",
505
max_depth=50,
506
timeout=5.0,
507
)
508
self.assertFalse(r.success)
509
510
def test_output_truncation(self):
511
"""stdout is capped at MAX_OUTPUT_BYTES; truncation notice appears."""
512
from g4f.mcp.pa_provider import MAX_OUTPUT_BYTES
513
# Produce more bytes than the limit
514
r = execute_safe_code(
515
f"print('A' * {MAX_OUTPUT_BYTES + 1000})",
516
timeout=5.0,
517
)
518
self.assertTrue(r.success)
519
self.assertLessEqual(len(r.stdout), MAX_OUTPUT_BYTES + 50)
520
self.assertIn("truncated", r.stderr.lower())
521
522
def test_timeout_none_disables_limit(self):
523
"""Passing timeout=None does not impose a time limit."""
524
r = execute_safe_code("result = sum(range(100))", timeout=None)
525
self.assertTrue(r.success)
526
self.assertEqual(r.result, 4950)
527
528
async def test_tool_respects_timeout_param(self):
529
"""PythonExecuteTool forwards timeout to execute_safe_code."""
530
tool = PythonExecuteTool(safe_mode=False)
531
import time
532
start = time.time()
533
result = await tool.execute({"code": "while True: pass", "timeout": 0.5})
534
elapsed = time.time() - start
535
self.assertFalse(result.get("success"))
536
self.assertLess(elapsed, 3.0)
537
538
async def test_tool_safe_mode_ignores_timeout_param(self):
539
"""In safe mode, timeout parameter is ignored and default is used."""
540
from g4f.mcp.pa_provider import MAX_EXEC_TIMEOUT
541
tool = PythonExecuteTool(safe_mode=True)
542
# Passing a very large timeout in safe mode should be ignored;
543
# the default MAX_EXEC_TIMEOUT is used instead.
544
result = await tool.execute({
545
"code": "result = 1",
546
"timeout": MAX_EXEC_TIMEOUT * 100,
547
})
548
self.assertTrue(result.get("success"))
549
550
async def test_tool_safe_mode_ignores_max_depth_param(self):
551
"""In safe mode, max_depth parameter is ignored."""
552
from g4f.mcp.pa_provider import MAX_RECURSION_DEPTH
553
tool = PythonExecuteTool(safe_mode=True)
554
# Even passing a huge depth, safe-mode always uses MAX_RECURSION_DEPTH
555
result = await tool.execute({
556
"code": "result = 1",
557
"max_depth": MAX_RECURSION_DEPTH * 100,
558
})
559
self.assertTrue(result.get("success"))
@@ -26,11 +26,13 @@ The sandbox mitigates the following vectors:
26
26
* **Code injection** — ``exec``, ``eval``, ``compile``, and ``input`` are removed
27
27
from the sandbox built-ins so code in the sandbox cannot spawn secondary
28
28
execution contexts.
29
30
Known limitations: the sandbox does not enforce CPU/memory limits or wall-clock
31
timeouts. Callers that need to bound execution time should wrap
32
:func:`execute_safe_code` with a ``asyncio.wait_for`` or ``concurrent.futures``
33
timeout.
29
* **Execution timeout** — code runs in a dedicated thread; if it does not
30
complete within :data:`MAX_EXEC_TIMEOUT` seconds the result is returned with
31
an error and the thread is abandoned.
32
* **Runaway recursion** — ``sys.setrecursionlimit`` is reduced to
33
:data:`MAX_RECURSION_DEPTH` for the duration of the sandboxed call.
34
* **Output flooding** — stdout and stderr are each capped at
35
:data:`MAX_OUTPUT_BYTES`; excess output is silently truncated.
34
36
35
37
Typical layout of a ``.pa.py`` file::
36
38
@@ -56,8 +58,9 @@ from __future__ import annotations
56
58
57
59
import io
58
60
import ast
61
import sys
59
62
import json
60
import contextlib
63
import threading
61
64
import traceback
62
65
import builtins as _builtins
63
66
from pathlib import Path
@@ -111,10 +114,85 @@ SAFE_MODULES: FrozenSet[str] = frozenset({
111
114
})
112
115
113
116
117
# ---------------------------------------------------------------------------
118
# Security limits
119
# ---------------------------------------------------------------------------
120
121
#: Wall-clock seconds allowed for a single :func:`execute_safe_code` call.
122
MAX_EXEC_TIMEOUT: float = 30.0
123
124
#: Maximum Python call-stack depth inside the sandbox (passed to
125
#: ``sys.setrecursionlimit``). The default CPython limit is 1 000; using a
126
#: lower value catches infinite-recursion attacks early.
127
MAX_RECURSION_DEPTH: int = 500
128
129
#: Maximum number of UTF-8 bytes captured from *each* of stdout and stderr.
130
#: Writes beyond this limit are silently dropped and a truncation notice is
131
#: appended to stderr.
132
MAX_OUTPUT_BYTES: int = 65_536 # 64 KiB
133
134
114
135
# ---------------------------------------------------------------------------
115
136
# Sandbox helpers
116
137
# ---------------------------------------------------------------------------
117
138
139
class _LimitedStringIO(io.StringIO):
140
"""StringIO that stops accepting writes once *max_bytes* of UTF-8 content
141
have been accumulated. Additional writes are silently discarded and
142
``truncated`` is set to ``True``."""
143
144
def __init__(self, max_bytes: int = MAX_OUTPUT_BYTES) -> None:
145
super().__init__()
146
self._max_bytes = max_bytes
147
self._bytes_written = 0
148
self.truncated = False
149
150
def write(self, s: str) -> int:
151
if self._bytes_written >= self._max_bytes:
152
self.truncated = True
153
return 0
154
encoded = s.encode("utf-8", errors="replace")
155
remaining = self._max_bytes - self._bytes_written
156
if len(encoded) > remaining:
157
s = encoded[:remaining].decode("utf-8", errors="replace")
158
self.truncated = True
159
n = super().write(s)
160
self._bytes_written += len(s.encode("utf-8", errors="replace"))
161
return n
162
163
164
def _exec_in_thread(
165
compiled: Any,
166
safe_globals: Dict[str, Any],
167
local_vars: Dict[str, Any],
168
max_depth: int,
169
exc_box: List,
170
) -> None:
171
"""Run *compiled* code with a bounded recursion depth.
172
173
``sys.setrecursionlimit`` is set to *max_depth* for the lifetime of this
174
call and restored afterwards. stdout / stderr capture is handled by
175
the custom ``print`` injected into the sandbox builtins — no global
176
``sys.stdout`` redirection is performed so an abandoned timeout thread
177
cannot corrupt the caller's output streams.
178
179
Any exception is stored in *exc_box* (a one-element list) so the caller
180
can inspect it without needing to join the thread.
181
182
This function is designed to run in a *daemon* thread so that it is
183
automatically discarded when the process exits, even if the sandboxed
184
code is stuck in an infinite loop.
185
"""
186
prev = sys.getrecursionlimit()
187
sys.setrecursionlimit(max_depth)
188
try:
189
exec(compiled, safe_globals, local_vars) # noqa: S102
190
except Exception: # noqa: BLE001
191
exc_box.append(traceback.format_exc())
192
finally:
193
sys.setrecursionlimit(prev)
194
195
118
196
def _make_restricted_import(allowed: FrozenSet[str]):
119
197
"""Return a ``__import__`` replacement that only allows *allowed* modules."""
120
198
original = _builtins.__import__
@@ -135,7 +213,11 @@ def _make_restricted_import(allowed: FrozenSet[str]):
135
213
return _restricted_import
136
214
137
215
138
def _make_safe_globals(allowed: FrozenSet[str] = SAFE_MODULES) -> Dict[str, Any]:
216
def _make_safe_globals(
217
allowed: FrozenSet[str] = SAFE_MODULES,
218
stdout_buf: Optional[io.StringIO] = None,
219
stderr_buf: Optional[io.StringIO] = None,
220
) -> Dict[str, Any]:
139
221
"""Return a ``globals`` dict suitable for sandboxed ``exec``."""
140
222
workspace = get_workspace_dir()
141
223
@@ -168,6 +250,19 @@ def _make_safe_globals(allowed: FrozenSet[str] = SAFE_MODULES) -> Dict[str, Any]
168
250
safe_builtins["open"] = _safe_open
169
251
safe_builtins["__import__"] = _make_restricted_import(allowed)
170
252
253
# Override print / input so stdout/stderr stay local to this sandbox
254
# execution and are never written to the real sys.stdout/stderr. This
255
# avoids the global-state side-effect that contextlib.redirect_stdout
256
# would cause when the thread is abandoned after a timeout.
257
if stdout_buf is not None:
258
_real_print = _builtins.print
259
260
def _safe_print(*args, **kwargs):
261
kwargs.setdefault("file", stdout_buf)
262
_real_print(*args, **kwargs)
263
264
safe_builtins["print"] = _safe_print
265
171
266
return {
172
267
"__builtins__": safe_builtins,
173
268
"__name__": "__pa_provider__",
@@ -222,51 +317,101 @@ def execute_safe_code(
222
317
code: str,
223
318
extra_globals: Optional[Dict[str, Any]] = None,
224
319
allowed_modules: FrozenSet[str] = SAFE_MODULES,
320
timeout: Optional[float] = MAX_EXEC_TIMEOUT,
321
max_depth: int = MAX_RECURSION_DEPTH,
225
322
) -> SafeExecutionResult:
226
323
"""Execute *code* inside a safe sandbox with whitelisted module imports.
227
324
325
The execution runs in a dedicated thread so that a wall-clock *timeout*
326
can be enforced without blocking the caller's event loop. A custom
327
``sys.setrecursionlimit`` guards against stack-overflow attacks. Both
328
stdout and stderr are capped at :data:`MAX_OUTPUT_BYTES`.
329
228
330
Args:
229
331
code: Python source code to execute.
230
332
extra_globals: Additional names injected into the execution globals.
231
333
allowed_modules: Frozenset of top-level module names that may be imported.
334
timeout: Wall-clock seconds before the execution is abandoned. Pass
335
``None`` to disable. Defaults to :data:`MAX_EXEC_TIMEOUT`.
336
max_depth: Maximum recursion depth inside the sandbox. Defaults to
337
:data:`MAX_RECURSION_DEPTH`.
232
338
233
339
Returns:
234
340
:class:`SafeExecutionResult` containing captured stdout/stderr, any
235
341
``result`` variable assigned in the code, or error information.
236
342
"""
237
stdout_buf = io.StringIO()
238
stderr_buf = io.StringIO()
343
stdout_buf = _LimitedStringIO(MAX_OUTPUT_BYTES)
344
stderr_buf = _LimitedStringIO(MAX_OUTPUT_BYTES)
239
345
240
safe_globals = _make_safe_globals(allowed_modules)
346
safe_globals = _make_safe_globals(allowed_modules, stdout_buf=stdout_buf, stderr_buf=stderr_buf)
241
347
if extra_globals:
242
348
safe_globals.update(extra_globals)
243
349
244
350
local_vars: Dict[str, Any] = {}
245
351
352
# Compile outside the thread so SyntaxErrors surface immediately.
246
353
try:
247
354
compiled = compile(code, "<pa_provider>", "exec")
248
with (
249
contextlib.redirect_stdout(stdout_buf),
250
contextlib.redirect_stderr(stderr_buf),
251
):
252
exec(compiled, safe_globals, local_vars) # noqa: S102
355
except SyntaxError:
356
return SafeExecutionResult(
357
success=False,
358
stdout="",
359
stderr="",
360
error=traceback.format_exc(),
361
)
253
362
363
# Run in a daemon thread with timeout and recursion-depth enforcement.
364
# We use a raw daemon Thread (not ThreadPoolExecutor) so that if the
365
# sandboxed code runs forever the thread is discarded when the process
366
# exits rather than blocking interpreter shutdown.
367
exc_box: List = []
368
thread = threading.Thread(
369
target=_exec_in_thread,
370
args=(compiled, safe_globals, local_vars, max_depth, exc_box),
371
daemon=True,
372
name="g4f-sandbox",
373
)
374
thread.start()
375
thread.join(timeout=timeout)
376
377
if thread.is_alive():
378
# The thread is still running — timeout was hit. We cannot kill it
379
# but as a daemon thread it will be reaped when the process exits.
380
stdout = stdout_buf.getvalue()
381
stderr = stderr_buf.getvalue()
382
if stdout_buf.truncated or stderr_buf.truncated:
383
stderr += "\n[Output truncated: size limit reached]"
254
384
return SafeExecutionResult(
255
success=True,
256
stdout=stdout_buf.getvalue(),
257
stderr=stderr_buf.getvalue(),
258
result=local_vars.get("result"),
259
locals=local_vars,
385
success=False,
386
stdout=stdout,
387
stderr=stderr,
388
error=(
389
f"Execution timed out after {timeout:.1f} s. "
390
"The thread has been abandoned."
391
),
260
392
)
261
393
262
except Exception:
394
if exc_box:
263
395
return SafeExecutionResult(
264
396
success=False,
265
397
stdout=stdout_buf.getvalue(),
266
398
stderr=stderr_buf.getvalue(),
267
error=traceback.format_exc(),
399
error=exc_box[0],
268
400
)
269
401
402
stdout = stdout_buf.getvalue()
403
stderr = stderr_buf.getvalue()
404
if stdout_buf.truncated or stderr_buf.truncated:
405
stderr += "\n[Output truncated: size limit reached]"
406
407
return SafeExecutionResult(
408
success=True,
409
stdout=stdout,
410
stderr=stderr,
411
result=local_vars.get("result"),
412
locals=local_vars,
413
)
414
270
415
271
416
# ---------------------------------------------------------------------------
272
417
# .pa.py provider loader
@@ -511,7 +511,21 @@ class PythonExecuteTool(MCPTool):
511
511
"items": {"type": "string"},
512
512
"description": (
513
513
"Optional list of additional module names to allow "
514
"beyond the default whitelist"
514
"beyond the default whitelist (ignored in safe mode)"
515
),
516
},
517
"timeout": {
518
"type": "number",
519
"description": (
520
"Wall-clock seconds to allow before aborting execution "
521
f"(max {30.0}s; ignored in safe mode)"
522
),
523
},
524
"max_depth": {
525
"type": "integer",
526
"description": (
527
"Maximum Python call-stack depth inside the sandbox "
528
f"(max {500}; ignored in safe mode)"
515
529
),
516
530
},
517
531
},
@@ -519,21 +533,41 @@ class PythonExecuteTool(MCPTool):
519
533
}
520
534
521
535
async def execute(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
522
from .pa_provider import execute_safe_code, SAFE_MODULES
536
from .pa_provider import execute_safe_code, SAFE_MODULES, MAX_EXEC_TIMEOUT, MAX_RECURSION_DEPTH
523
537
524
538
code = arguments.get("code", "")
525
539
if not code:
526
540
return {"error": "code parameter is required"}
527
541
528
542
if self.safe_mode:
529
# In safe mode the caller cannot expand the module allowlist
543
# In safe mode the caller cannot override any security parameters
530
544
allowed = SAFE_MODULES
545
timeout = MAX_EXEC_TIMEOUT
546
max_depth = MAX_RECURSION_DEPTH
531
547
else:
532
548
extra_names = arguments.get("allowed_extra_modules") or []
533
549
allowed = SAFE_MODULES | frozenset(extra_names)
550
# Allow callers to reduce (but not exceed) the defaults
551
requested_timeout = arguments.get("timeout")
552
timeout = (
553
min(float(requested_timeout), MAX_EXEC_TIMEOUT)
554
if requested_timeout is not None
555
else MAX_EXEC_TIMEOUT
556
)
557
requested_depth = arguments.get("max_depth")
558
max_depth = (
559
min(int(requested_depth), MAX_RECURSION_DEPTH)
560
if requested_depth is not None
561
else MAX_RECURSION_DEPTH
562
)
534
563
535
564
try:
536
exec_result = execute_safe_code(code, allowed_modules=allowed)
565
exec_result = execute_safe_code(
566
code,
567
allowed_modules=allowed,
568
timeout=timeout,
569
max_depth=max_depth,
570
)
537
571
return exec_result.to_dict()
538
572
except Exception as exc:
539
573
return {"error": f"Execution error: {exc}"}