From 0ea90f0b882657803356f1fb3b2b3325c2adc7b8 Mon Sep 17 00:00:00 2001
From: essence-of-the-soul <essence-of-the-soul@users.noreply.gitee.com>
Date: Wed, 5 Aug 2026 11:16:15 +0000
Subject: [PATCH] =?UTF-8?q?fix(#IK18N8):=20BwrapSandbox=20=E9=99=8D?=
 =?UTF-8?q?=E7=BA=A7=E8=B7=AF=E5=BE=84=E6=94=B9=E7=94=A8=20preexec=5Ffn=20?=
 =?UTF-8?q?=E8=AE=BE=E7=BD=AE=E8=B5=84=E6=BA=90=E9=99=90=E9=A2=9D?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

关联 issue #IK18N8
---
 ocai-service/ocai/ai/sandbox/bwrap.py | 37 +++++++++++++++++++--------
 1 file changed, 26 insertions(+), 11 deletions(-)

diff --git a/ocai-service/ocai/ai/sandbox/bwrap.py b/ocai-service/ocai/ai/sandbox/bwrap.py
index 901645e..b272841 100644
--- a/ocai-service/ocai/ai/sandbox/bwrap.py
+++ b/ocai-service/ocai/ai/sandbox/bwrap.py
@@ -16,6 +16,7 @@ bubblewrap 提供:
 
 import logging
 import os
+import resource
 import shutil
 import subprocess
 import time
@@ -164,7 +165,7 @@ class BwrapSandbox(SandboxExecutor):
 
         if not self._available:
             return self._fallback_execute(
-                command,
+                ["sh", "-c", command],
                 timeout=effective_timeout,
                 working_dir=working_dir,
                 env=env,
@@ -220,7 +221,7 @@ class BwrapSandbox(SandboxExecutor):
         if not self._available:
             cmd_parts = [interpreter, script_path] + effective_args
             return self._fallback_execute(
-                " ".join(cmd_parts),
+                cmd_parts,
                 timeout=effective_timeout,
                 working_dir=working_dir,
                 env=env,
@@ -436,7 +437,7 @@ class BwrapSandbox(SandboxExecutor):
 
     def _fallback_execute(
         self,
-        command: str,
+        cmd: list[str],
         *,
         timeout: float,
         working_dir: str | None = None,
@@ -446,7 +447,8 @@ class BwrapSandbox(SandboxExecutor):
         """无 bwrap 时的降级执行 — 使用 resource limits
 
         注意: 此降级方案安全性远不如 bwrap，仅在 bwrap 不可用时使用。
-        使用 ulimit 限制 CPU 时间和内存，并在子 shell 中执行。
+        通过 preexec_fn 在子进程中以 resource.setrlimit 设置 CPU/内存/文件大小限制，
+        不再将 command 拼接到 shell 字符串中，避免命令注入风险被放大。
         """
         logger.warning(
             "bwrap not available, using fallback execution with resource limits. "
@@ -456,13 +458,25 @@ class BwrapSandbox(SandboxExecutor):
         start_time = time.time()
         killed = False
 
-        # 使用 ulimit 进行基本资源限制
-        # -t: CPU 时间(秒) -v: 虚拟内存(KB) -f: 文件大小(512B blocks)
+        # 资源限制（等价于原 ulimit -t/-v/-f，但通过内核强制生效）
+        # - RLIMIT_CPU: CPU 时间(秒)
+        # - RLIMIT_AS: 虚拟内存(字节)，原 ulimit -v 以 KB 为单位
+        # - RLIMIT_FSIZE: 文件大小(字节)，原 ulimit -f 以 512B 块为单位
         cpu_limit = min(int(timeout), 300)
-        mem_limit = 512 * 1024  # 512MB in KB
-        file_limit = 1024 * 1024  # ~512MB in 512B blocks
-
-        wrapped_cmd = f"ulimit -t {cpu_limit} -v {mem_limit} -f {file_limit} 2>/dev/null; {command}"
+        mem_limit = 512 * 1024 * 1024  # 512MB in bytes
+        file_limit = 512 * 1024 * 1024  # 512MB in bytes
+
+        def _set_limits() -> None:
+            for limit_id, value in (
+                (resource.RLIMIT_CPU, (cpu_limit, cpu_limit)),
+                (resource.RLIMIT_AS, (mem_limit, mem_limit)),
+                (resource.RLIMIT_FSIZE, (file_limit, file_limit)),
+            ):
+                try:
+                    resource.setrlimit(limit_id, value)
+                except (ValueError, OSError):
+                    # 部分环境(如容器)可能不支持设置某类限额，跳过以保持降级可用
+                    pass
 
         proc_env = os.environ.copy()
         if env:
@@ -470,12 +484,13 @@ class BwrapSandbox(SandboxExecutor):
 
         try:
             proc = subprocess.Popen(
-                ["sh", "-c", wrapped_cmd],
+                cmd,
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
                 stdin=subprocess.DEVNULL,
                 cwd=working_dir,
                 env=proc_env,
+                preexec_fn=_set_limits,
             )
 
             try:
-- 
Gitee