mirror of
https://github.com/modelscope/modelscope.git
synced 2026-07-10 04:22:33 +02:00
96 lines
3.6 KiB
Python
96 lines
3.6 KiB
Python
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
"""Regression test for Issue #1668.
|
|
|
|
`TinynasClassificationPipeline.postprocess` previously called `eval()` on
|
|
the raw contents of `label_map.txt` shipped from the remote model repo,
|
|
giving attackers RCE the moment the pipeline produced a prediction. The fix
|
|
swaps `eval` for `ast.literal_eval`.
|
|
"""
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
|
|
import torch # noqa: E402
|
|
|
|
from modelscope.pipelines.cv.tinynas_classification_pipeline import \
|
|
TinynasClassificationPipeline # noqa: E402
|
|
|
|
# `tinynas_classification_pipeline` imports torch/torchvision and our own
|
|
# `tinynas_classfication` module at top level. Stub the latter only when
|
|
# missing so we don't drag heavy custom CUDA ops into the test.
|
|
if 'modelscope.models.cv.tinynas_classfication' not in sys.modules:
|
|
try:
|
|
from modelscope.models.cv.tinynas_classfication import \
|
|
get_zennet # noqa: F401
|
|
except Exception:
|
|
stub = types.ModuleType('modelscope.models.cv.tinynas_classfication')
|
|
stub.get_zennet = lambda *a, **kw: None
|
|
sys.modules['modelscope.models.cv.tinynas_classfication'] = stub
|
|
|
|
|
|
class TinynasLabelMapSecurityTest(unittest.TestCase):
|
|
"""Pin the `ast.literal_eval` gate around remote `label_map.txt`."""
|
|
|
|
def setUp(self):
|
|
self.tmp_root = tempfile.mkdtemp(prefix='ms_tinynas_security_')
|
|
self.label_path = os.path.join(self.tmp_root, 'label_map.txt')
|
|
# Bypass __init__ — it would try to load a real checkpoint.
|
|
self.pipe = TinynasClassificationPipeline.__new__(
|
|
TinynasClassificationPipeline)
|
|
self.pipe.path = self.tmp_root
|
|
|
|
def tearDown(self):
|
|
shutil.rmtree(self.tmp_root, ignore_errors=True)
|
|
|
|
def _write(self, content):
|
|
with open(self.label_path, 'w', encoding='utf-8') as f:
|
|
f.write(content)
|
|
|
|
def _fake_inputs(self, argmax_idx):
|
|
# `outputs` must support `.argmax().item()` and softmax.
|
|
outputs = torch.zeros(1, 3)
|
|
outputs[0, argmax_idx] = 10.0
|
|
return {'outputs': outputs}
|
|
|
|
def test_legitimate_label_map_parses(self):
|
|
self._write("{0: 'cat', 1: 'dog', 2: 'wolf'}")
|
|
out = self.pipe.postprocess(self._fake_inputs(2))
|
|
self.assertEqual(out['labels'], ['wolf'])
|
|
self.assertEqual(len(out['scores']), 1)
|
|
|
|
def test_malicious_import_payload_blocked(self):
|
|
canary = '/tmp/ms_tinynas_rce_canary'
|
|
if os.path.exists(canary):
|
|
os.remove(canary)
|
|
self._write(
|
|
"__import__('os').system('echo HACKED && touch {}')".format(
|
|
canary))
|
|
with self.assertRaises((ValueError, SyntaxError)):
|
|
self.pipe.postprocess(self._fake_inputs(0))
|
|
self.assertFalse(
|
|
os.path.exists(canary),
|
|
'os.system was reached -- eval gate is bypassed.')
|
|
|
|
def test_malicious_function_call_blocked(self):
|
|
self._write("print('pwned')")
|
|
with self.assertRaises((ValueError, SyntaxError)):
|
|
self.pipe.postprocess(self._fake_inputs(0))
|
|
|
|
def test_malicious_subclass_escape_blocked(self):
|
|
self._write('().__class__.__bases__[0].__subclasses__()')
|
|
with self.assertRaises((ValueError, SyntaxError)):
|
|
self.pipe.postprocess(self._fake_inputs(0))
|
|
|
|
def test_malformed_label_map_raises_cleanly(self):
|
|
# Empty file is not a literal container; must not silently no-op.
|
|
self._write('')
|
|
with self.assertRaises((ValueError, SyntaxError)):
|
|
self.pipe.postprocess(self._fake_inputs(0))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|