mirror of
https://github.com/modelscope/modelscope.git
synced 2026-07-09 20:09:17 +02:00
fix (#1740)
This commit is contained in:
@@ -131,7 +131,7 @@ def check_model_from_owner_group(model_dir: str,
|
||||
"""This checking is for the torch.load, this function may eval malicious code into memory
|
||||
|
||||
Args:
|
||||
model_dir: The local model_dir
|
||||
model_dir: The local model_dir or model_id
|
||||
owner_group: The owner group to trust
|
||||
|
||||
Returns:
|
||||
@@ -142,6 +142,15 @@ def check_model_from_owner_group(model_dir: str,
|
||||
if owner_group is None:
|
||||
owner_group = ['iic', 'damo']
|
||||
model_dir = model_dir.rstrip('/').rstrip('\\')
|
||||
model_dir = os.path.dirname(model_dir)
|
||||
group = os.path.basename(model_dir)
|
||||
return group in owner_group
|
||||
parent_dir = os.path.dirname(model_dir)
|
||||
group = os.path.basename(parent_dir)
|
||||
if group in owner_group:
|
||||
return True
|
||||
# Also check cache path pattern: {cache_root}/{owner}--{model_name}/snapshots/{revision}
|
||||
# Only check the grandparent directory to avoid path traversal attacks
|
||||
grandparent = os.path.basename(os.path.dirname(parent_dir))
|
||||
if '--' in grandparent:
|
||||
prefix = grandparent.split('--', 1)[0]
|
||||
if prefix in owner_group:
|
||||
return True
|
||||
return False
|
||||
|
||||
Reference in New Issue
Block a user