Fix RCE issue for plugins (#1415)

1. Add `trust_remote_code` in func `from_pretrained`  with plugins application
This commit is contained in:
Koko-ry
2025-07-16 14:33:31 +08:00
committed by hjh0119
parent c63952e90c
commit 3d11b891ca

View File

@@ -85,6 +85,7 @@ class Model(ABC):
revision: Optional[str] = DEFAULT_MODEL_REVISION, revision: Optional[str] = DEFAULT_MODEL_REVISION,
cfg_dict: Config = None, cfg_dict: Config = None,
device: str = None, device: str = None,
trust_remote_code: Optional[bool] = False,
**kwargs): **kwargs):
"""Instantiate a model from local directory or remote model repo. Note """Instantiate a model from local directory or remote model repo. Note
that when loading from remote, the model revision can be specified. that when loading from remote, the model revision can be specified.
@@ -96,6 +97,7 @@ class Model(ABC):
cfg_dict(Config, `optional`): An optional model config. If provided, it will replace cfg_dict(Config, `optional`): An optional model config. If provided, it will replace
the config read out of the `model_name_or_path` the config read out of the `model_name_or_path`
device(str, `optional`): The device to load the model. device(str, `optional`): The device to load the model.
trust_remote_code(bool, `optional`): Whether to trust and allow execution of remote code. Default is False.
**kwargs: **kwargs:
task(str, `optional`): The `Tasks` enumeration value to replace the task value task(str, `optional`): The `Tasks` enumeration value to replace the task value
read out of config in the `model_name_or_path`. This is useful when the model to be loaded is not read out of config in the `model_name_or_path`. This is useful when the model to be loaded is not
@@ -181,10 +183,21 @@ class Model(ABC):
f'`{ModelFile.CONFIGURATION}` file not found.') f'`{ModelFile.CONFIGURATION}` file not found.')
model_cfg.model_dir = local_model_dir model_cfg.model_dir = local_model_dir
# install and import remote repos before build # Security check: Only allow execution of remote code or plugins if trust_remote_code is True
register_plugins_repo(cfg.safe_get('plugins')) plugins = cfg.safe_get('plugins')
register_modelhub_repo(local_model_dir, cfg.get('allow_remote', False)) if plugins and not trust_remote_code:
raise RuntimeError(
'Detected plugins field in the model configuration file, but '
'trust_remote_code=True was not explicitly set.\n'
'To prevent potential execution of malicious code, loading has been refused.\n'
'If you trust this model repository, please pass trust_remote_code=True to from_pretrained.'
)
if plugins and trust_remote_code:
logger.warning(
'Use trust_remote_code=True. Will invoke codes or install plugins from remote model repo. '
'Please make sure that you can trust the external codes.')
register_modelhub_repo(local_model_dir, allow_remote=trust_remote_code)
register_plugins_repo(plugins)
for k, v in kwargs.items(): for k, v in kwargs.items():
model_cfg[k] = v model_cfg[k] = v
if device is not None: if device is not None: