1+ # -- encoding: utf-8 --
2+ # Copyright (c) 2024 Huawei Technologies Co., Ltd. All Rights Reserved.
3+ # This file is a part of the ModelEngine Project.
4+ # Licensed under the MIT License. See License.txt in the project root for license information.
5+ # ======================================================================================================================
6+ import asyncio
7+ import importlib
8+ import inspect
9+ import json
10+ import multiprocessing
11+ import platform
12+ import re
13+ from typing import Any , Dict , List , Tuple
14+ from pydantic import BaseModel
15+
16+ if platform .system () == 'Windows' :
17+ from enum import IntEnum
18+
19+ class InternalErrorCode (IntEnum ):
20+ EXCEPTION_FROM_USER_CODE_OCCURRED = 0x7F000105
21+ TIME_OUT_EXCEPTION_FROM_USER_CODE_OCCURRED = 0x7F000106 # java 不存在
22+ else :
23+ from fitframework .core .exception .fit_exception import InternalErrorCode
24+ try :
25+ import resource
26+ except ImportError :
27+ resource = None
28+ try :
29+ from .safe_global import safe_builtins
30+ except ImportError as e :
31+ from safe_global import safe_builtins
32+
33+ _PYTHON_REPL_HEADER = '''
34+ import json
35+ from typing import Any
36+
37+ Output = Any
38+
39+
40+ '''
41+
42+ GLOBAL_CONFIG = \
43+ {
44+ "header" : _PYTHON_REPL_HEADER ,
45+ "header_len" : len (_PYTHON_REPL_HEADER .split ('\n ' )),
46+ "entrypoint" : 'main' ,
47+ "whitelist" : ['json' , 'typing' ],
48+ "blacklist" : ['os' , 'sys' , 'cmd' , 'subprocess' , 'multiprocessing' , 'timeit' , 'platform' ],
49+ "timeout" : 10 ,
50+ "max_pool" : 4 ,
51+ "mem_limit" : 181 * 1024 * 1024 ,
52+ "verbose" : False
53+ }
54+
55+
56+ class Result (BaseModel ):
57+ isOk : bool
58+ value : Any = None
59+ error_code : int
60+ msg : str = None
61+
62+ @staticmethod
63+ def ok (data : Any ) -> 'Result' :
64+ return Result (isOk = True , value = data , error_code = 0 )
65+
66+ @staticmethod
67+ def err (err_code : int , err_msg : str ) -> 'Result' :
68+ return Result (isOk = False , error_code = err_code , msg = err_msg )
69+
70+
71+ # 创建一个安全的执行环境
72+ def _create_restricted_exec_env (config : Dict [str , object ]):
73+ def safer_import (name , my_globals = None , my_locals = None , fromlist = (), level = 0 ):
74+ if name not in config ['whitelist' ] or name in config ['blacklist' ]:
75+ raise NameError (f'model { name } )
76+ return importlib .import_module (name )
77+
78+ safe_globals = {
79+ '__builtins__' : {
80+ ** safe_builtins ,
81+ '__import__' : safer_import ,
82+ 'Args' : Dict
83+ }
84+ }
85+ return safe_globals
86+
87+
88+ # 获取内存使用（单位：kB）
89+ def _get_current_memory_usage ():
90+ with open ('/proc/self/status' ) as f :
91+ mem_usage = f .read ().split ('VmPeak:' )[1 ].split ('\n ' )[0 ].strip ()
92+ return int (mem_usage .split ()[0 ].strip ())
93+
94+
95+ # 执行受限代码
96+ def _execute_code_with_restricted_python (args : Dict [str , object ], code : str , config : Dict [str , object ]):
97+ if resource :
98+ resource .setrlimit (resource .RLIMIT_AS , (GLOBAL_CONFIG ["mem_limit" ], GLOBAL_CONFIG ["mem_limit" ]))
99+ loop = asyncio .new_event_loop ()
100+ asyncio .set_event_loop (loop )
101+ try :
102+ full_python_code = (f"{ config ['header' ]}
103+ f'{ code } \n \n ' )
104+
105+ safer_globals = _create_restricted_exec_env (config )
106+ exec (full_python_code , safer_globals )
107+ entrypoint = config ['entrypoint' ]
108+ if (entrypoint not in safer_globals or
109+ not inspect .isfunction (safer_globals .get (entrypoint ))):
110+ raise NameError ("main function not defined" )
111+ entrypoint = safer_globals .get (entrypoint )
112+ if inspect .iscoroutinefunction (entrypoint ):
113+ ret = loop .run_until_complete (asyncio .wait_for (entrypoint (args ), config ['timeout' ]))
114+ return Result .ok (json .dumps (ret ))
115+ else :
116+ return Result .err (InternalErrorCode .EXCEPTION_FROM_USER_CODE_OCCURRED .value ,
117+ "Unable to execute non-asynchronous function" )
118+ except asyncio .TimeoutError :
119+ return Result .err (InternalErrorCode .TIME_OUT_EXCEPTION_FROM_USER_CODE_OCCURRED .value ,
120+ "[TimeoutError] Execution timed out" )
121+ except Exception as err :
122+ return Result .err (InternalErrorCode .EXCEPTION_FROM_USER_CODE_OCCURRED .value , _get_except_msg (err , config ))
123+ finally :
124+ loop .close ()
125+
126+
127+ def _get_except_msg (error : Any , config : Dict [str , object ]) -> str :
128+ if isinstance (error , SyntaxError ):
129+ error_msg = f"{ error .msg } { error .lineno - config ['header_len' ]} { error .offset } { error .text }
130+ elif isinstance (error , KeyError ):
131+ error_msg = f"key { str (error )}
132+ else :
133+ error_msg = str (error )
134+ return f"[{ error .__class__ .__name__ } { error_msg }
135+
136+
137+ def _get_free_process_pool (pools : List [Tuple [multiprocessing .Lock , multiprocessing .Pool ]], index ):
138+ lock = pools [index ][0 ]
139+ if lock .acquire ():
140+ return lock
141+ raise multiprocessing .TimeoutError ()
142+
143+
144+ def execute_node_impl (pools : List [Tuple [multiprocessing .Lock , multiprocessing .Pool ]], index : int ,
145+ args : Dict [str , object ], code : str , config : Dict [str , object ]):
146+ match = _validate_escape (code )
147+ if match is not None :
148+ return Result .err (InternalErrorCode .EXCEPTION_FROM_USER_CODE_OCCURRED .value ,
149+ f'{ match .group ()} )
150+ lock = _get_free_process_pool (pools , index )
151+ pool = pools [index ][1 ]
152+ try :
153+ result = pool .apply_async (_execute_code_with_restricted_python , args = [args , code , config ])
154+ return result .get (config ['timeout' ])
155+ except multiprocessing .TimeoutError :
156+ index = pools .index ((lock , pool ))
157+ pool .terminate ()
158+ pools [index ] = (lock , multiprocessing .Pool (processes = 1 ))
159+ return Result .err (InternalErrorCode .TIME_OUT_EXCEPTION_FROM_USER_CODE_OCCURRED .value ,
160+ "[TimeoutError] Execution timed out" )
161+ finally :
162+ lock .release ()
163+
164+
165+ def _validate_escape (code : str ) -> bool :
166+ # 校验代码中是否存在获取栈帧的字段，禁用可能用于沙箱逃逸的端
167+ pattern = r'.gi_frame|.tb_frame|__[a-zA-Z]+__'
168+ return re .search (pattern , code )
0 commit comments