mirror of
https://github.com/modelscope/modelscope.git
synced 2026-07-10 04:22:33 +02:00
262 lines
10 KiB
Python
262 lines
10 KiB
Python
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
"""Regression test for Issue #1672.
|
|
|
|
Pins the `trust_remote_code` gate that prevents `Config.from_file` and
|
|
`check_trust_remote_code_for_config` from executing a `.py` config from an
|
|
untrusted model repository.
|
|
|
|
Coverage matrix
|
|
---------------
|
|
The helper / choke point exercised here is the only mechanism guarding every
|
|
sink listed in the issue:
|
|
|
|
* via `Config.from_file('.py')`:
|
|
- `SalientDetection` (models/cv/salient_detection)
|
|
- `AbnormalDetectionModel` (models/cv/abnormal_object_detection)
|
|
- `DetectionModel` (models/cv/object_detection)
|
|
- `ObjectDetection3DPipeline` (pipelines/cv/object_detection_3d_pipeline)
|
|
|
|
* via `check_trust_remote_code_for_config(...)` invoked before
|
|
`mmcv.Config.fromfile` / `mmseg.apis.init_segmentor`:
|
|
- `ClassificationModel` (image_classification/mmcls_model)
|
|
- `SwinLPanopticSegmentation` (image_panoptic_segmentation)
|
|
- `SemanticSegmentation` (image_semantic_segmentation)
|
|
- `ScrfdDetect` / `TinyMogDetect` / `DamoFdDetect` (face_detection/scrfd)
|
|
- `ImageClassificationMmcvPreprocessor` (preprocessors/cv)
|
|
- `SegformerDetector` (controllable_image_generation)
|
|
"""
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from modelscope.utils.config import Config, check_trust_remote_code_for_config
|
|
|
|
|
|
def _make_repo(root, owner, name):
|
|
"""`check_model_from_owner_group` reads the owner from the parent dir."""
|
|
repo = os.path.join(root, owner, name)
|
|
os.makedirs(repo)
|
|
return repo
|
|
|
|
|
|
def _write(path, content):
|
|
with open(path, 'w') as f:
|
|
f.write(content)
|
|
|
|
|
|
# Canary: any caller that fails the gate must NEVER execute this body.
|
|
_CANARY = ('raise SystemExit("this script must not run")\n'
|
|
'CANARY = "should not be reachable"\n')
|
|
|
|
|
|
class ConfigChokePointSecurityTest(unittest.TestCase):
|
|
"""End-to-end gating at `Config.from_file` for `.py` configs.
|
|
|
|
Covers SalientDetection / AbnormalDetectionModel / DetectionModel /
|
|
ObjectDetection3DPipeline, all of which call `Config.from_file(*.py,
|
|
trust_remote_code=self.trust_remote_code)`.
|
|
"""
|
|
|
|
def setUp(self):
|
|
self.tmp_root = tempfile.mkdtemp(prefix='ms_cfg_security_')
|
|
self.untrusted_dir = _make_repo(self.tmp_root, 'attacker', 'badrepo')
|
|
self.trusted_dir = _make_repo(self.tmp_root, 'damo', 'goodrepo')
|
|
|
|
self.untrusted_py = os.path.join(self.untrusted_dir, 'mmcv_config.py')
|
|
self.trusted_py = os.path.join(self.trusted_dir, 'mmcv_config.py')
|
|
for p in (self.untrusted_py, self.trusted_py):
|
|
_write(p, _CANARY)
|
|
|
|
self.untrusted_json = os.path.join(self.untrusted_dir,
|
|
'configuration.json')
|
|
_write(self.untrusted_json, '{"a": 1}')
|
|
|
|
def tearDown(self):
|
|
shutil.rmtree(self.tmp_root, ignore_errors=True)
|
|
|
|
def test_untrusted_py_without_optin_is_blocked(self):
|
|
with self.assertRaises(RuntimeError):
|
|
Config.from_file(self.untrusted_py)
|
|
|
|
def test_untrusted_py_with_optin_is_allowed(self):
|
|
# opt-in must let the load proceed, which means the canary will run
|
|
# and raise SystemExit -- proving the gate did not block.
|
|
with self.assertRaises(SystemExit):
|
|
Config.from_file(self.untrusted_py, trust_remote_code=True)
|
|
|
|
def test_trusted_owner_py_is_allowed_without_optin(self):
|
|
with self.assertRaises(SystemExit):
|
|
Config.from_file(self.trusted_py)
|
|
|
|
def test_json_passive_load_unaffected(self):
|
|
cfg = Config.from_file(self.untrusted_json)
|
|
self.assertEqual(cfg.a, 1)
|
|
|
|
def test_from_string_in_process_is_not_gated(self):
|
|
# `Config.from_string` materializes a caller-supplied in-process
|
|
# string and must keep working without `trust_remote_code`.
|
|
cfg = Config.from_string('a = 42\n', '.py')
|
|
self.assertEqual(cfg.a, 42)
|
|
|
|
def test_explicit_model_dir_overrides_default(self):
|
|
# Some callers pass `model_dir` explicitly; verify trust derives
|
|
# from that path, not from the file's parent directory.
|
|
nested = os.path.join(self.untrusted_dir, 'subdir')
|
|
os.makedirs(nested)
|
|
nested_py = os.path.join(nested, 'mmcv_config.py')
|
|
_write(nested_py, _CANARY)
|
|
with self.assertRaises(RuntimeError):
|
|
Config.from_file(nested_py, model_dir=self.untrusted_dir)
|
|
# Same file, but with a trusted model_dir override -> allowed.
|
|
with self.assertRaises(SystemExit):
|
|
Config.from_file(nested_py, model_dir=self.trusted_dir)
|
|
|
|
|
|
class CheckTrustRemoteCodeForConfigTest(unittest.TestCase):
|
|
"""Unit tests for the helper invoked at every mmcv / init_segmentor sink.
|
|
|
|
Covers ClassificationModel, SwinLPanopticSegmentation, SemanticSegmentation,
|
|
ScrfdDetect (and subclasses TinyMogDetect / DamoFdDetect),
|
|
ImageClassificationMmcvPreprocessor, and SegformerDetector.
|
|
"""
|
|
|
|
def setUp(self):
|
|
self.tmp_root = tempfile.mkdtemp(prefix='ms_cfg_security_helper_')
|
|
self.untrusted_dir = _make_repo(self.tmp_root, 'attacker', 'badrepo')
|
|
self.trusted_dir = _make_repo(self.tmp_root, 'damo', 'goodrepo')
|
|
# Helper only inspects file extensions; no need to create real files.
|
|
self.untrusted_py = os.path.join(self.untrusted_dir, 'config.py')
|
|
self.trusted_py = os.path.join(self.trusted_dir, 'config.py')
|
|
self.untrusted_json = os.path.join(self.untrusted_dir,
|
|
'configuration.json')
|
|
|
|
def tearDown(self):
|
|
shutil.rmtree(self.tmp_root, ignore_errors=True)
|
|
|
|
def test_non_py_path_always_passes(self):
|
|
# JSON / YAML configs are passive data; helper must not block them.
|
|
check_trust_remote_code_for_config(
|
|
self.untrusted_json,
|
|
trust_remote_code=False,
|
|
model_dir=self.untrusted_dir)
|
|
|
|
def test_untrusted_py_without_optin_raises(self):
|
|
with self.assertRaises(RuntimeError):
|
|
check_trust_remote_code_for_config(
|
|
self.untrusted_py,
|
|
trust_remote_code=False,
|
|
model_dir=self.untrusted_dir)
|
|
|
|
def test_untrusted_py_with_optin_passes(self):
|
|
check_trust_remote_code_for_config(
|
|
self.untrusted_py,
|
|
trust_remote_code=True,
|
|
model_dir=self.untrusted_dir)
|
|
|
|
def test_trusted_owner_py_passes(self):
|
|
check_trust_remote_code_for_config(
|
|
self.trusted_py,
|
|
trust_remote_code=False,
|
|
model_dir=self.trusted_dir)
|
|
|
|
def test_default_model_dir_inferred_from_filename(self):
|
|
# When `model_dir` is omitted, helper must infer it from the file's
|
|
# parent directory and still gate untrusted owners correctly.
|
|
with self.assertRaises(RuntimeError):
|
|
check_trust_remote_code_for_config(
|
|
self.untrusted_py, trust_remote_code=False)
|
|
# Trusted parent should pass.
|
|
check_trust_remote_code_for_config(
|
|
self.trusted_py, trust_remote_code=False)
|
|
|
|
def test_helper_uses_raise_not_assert(self):
|
|
# Security gates must survive `python -O`. The helper raises
|
|
# RuntimeError, not AssertionError; verify the failure type so the
|
|
# contract is locked in.
|
|
try:
|
|
check_trust_remote_code_for_config(
|
|
self.untrusted_py,
|
|
trust_remote_code=False,
|
|
model_dir=self.untrusted_dir)
|
|
except RuntimeError:
|
|
return
|
|
except AssertionError:
|
|
self.fail(
|
|
'Gate used `assert`; would be a no-op under `python -O`.')
|
|
self.fail('Gate did not raise on untrusted `.py` config.')
|
|
|
|
|
|
class SinkWiringTest(unittest.TestCase):
|
|
"""Spot-check that a representative mmcv sink calls the helper before
|
|
reaching `mmcv.Config.fromfile`.
|
|
|
|
We patch the helper to raise a sentinel; if the sink is wired correctly,
|
|
the sentinel propagates and the unsafe `mmcv.Config.fromfile` is never
|
|
invoked.
|
|
"""
|
|
|
|
def setUp(self):
|
|
self.tmp_root = tempfile.mkdtemp(prefix='ms_cfg_sink_wiring_')
|
|
self.untrusted_dir = _make_repo(self.tmp_root, 'attacker', 'badrepo')
|
|
|
|
def tearDown(self):
|
|
shutil.rmtree(self.tmp_root, ignore_errors=True)
|
|
|
|
def _import_module_text(self, dotted):
|
|
# Read the source of a module without importing it (avoids dragging
|
|
# in mmcv / mmdet / torch on test machines that lack them).
|
|
import importlib.util
|
|
spec = importlib.util.find_spec(dotted)
|
|
self.assertIsNotNone(spec, f'{dotted} not found')
|
|
with open(spec.origin, 'r') as f:
|
|
return f.read()
|
|
|
|
def _assert_helper_precedes_sink(self, dotted, sink_call):
|
|
src = self._import_module_text(dotted)
|
|
helper_pos = src.find('check_trust_remote_code_for_config(')
|
|
sink_pos = src.find(sink_call)
|
|
self.assertGreater(
|
|
helper_pos, -1,
|
|
f'{dotted} does not call check_trust_remote_code_for_config')
|
|
self.assertGreater(sink_pos, -1,
|
|
f'{dotted} no longer contains `{sink_call}`')
|
|
self.assertLess(helper_pos, sink_pos,
|
|
f'{dotted}: helper call must precede `{sink_call}`')
|
|
|
|
def test_classification_model_wires_helper(self):
|
|
self._assert_helper_precedes_sink(
|
|
'modelscope.models.cv.image_classification.mmcls_model',
|
|
'mmcv.Config.fromfile(')
|
|
|
|
def test_panoptic_segmentation_wires_helper(self):
|
|
self._assert_helper_precedes_sink(
|
|
'modelscope.models.cv.image_panoptic_segmentation.panseg_model',
|
|
'mmcv.Config.fromfile(')
|
|
|
|
def test_semantic_segmentation_wires_helper(self):
|
|
self._assert_helper_precedes_sink(
|
|
'modelscope.models.cv.image_semantic_segmentation.semantic_seg_model',
|
|
'mmcv.Config.fromfile(')
|
|
|
|
def test_scrfd_detect_wires_helper(self):
|
|
self._assert_helper_precedes_sink(
|
|
'modelscope.models.cv.face_detection.scrfd.scrfd_detect',
|
|
'Config.fromfile(')
|
|
|
|
def test_mmcls_preprocessor_wires_helper(self):
|
|
self._assert_helper_precedes_sink(
|
|
'modelscope.preprocessors.cv.mmcls_preprocessor',
|
|
'mmcv.Config.fromfile(')
|
|
|
|
def test_segformer_detector_wires_helper(self):
|
|
self._assert_helper_precedes_sink(
|
|
'modelscope.models.cv.controllable_image_generation.annotator.annotator',
|
|
'init_segmentor(')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|