From a5a3ce88b3a3186b2216f60985f8fc634a3d9043 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 9 Feb 2023 13:27:46 +0400 Subject: [PATCH] Command injection and Path traversal security fixes (#888) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- ultralytics/yolo/data/dataloaders/v5loader.py | 2 +- ultralytics/yolo/data/utils.py | 2 +- ultralytics/yolo/engine/trainer.py | 18 +++++++++++------ ultralytics/yolo/utils/dist.py | 20 +++++++------------ 4 files changed, 21 insertions(+), 21 deletions(-) diff --git a/ultralytics/yolo/data/dataloaders/v5loader.py b/ultralytics/yolo/data/dataloaders/v5loader.py index c9f2e15..2b75285 100644 --- a/ultralytics/yolo/data/dataloaders/v5loader.py +++ b/ultralytics/yolo/data/dataloaders/v5loader.py @@ -54,7 +54,7 @@ for orientation in ExifTags.TAGS.keys(): def get_hash(paths): # Returns a single hash value of a list of paths (files or dirs) size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) # sizes - h = hashlib.md5(str(size).encode()) # hash sizes + h = hashlib.sha256(str(size).encode()) # hash sizes h.update(''.join(paths).encode()) # hash paths return h.hexdigest() # return hash diff --git a/ultralytics/yolo/data/utils.py b/ultralytics/yolo/data/utils.py index e9ec668..b1163f6 100644 --- a/ultralytics/yolo/data/utils.py +++ b/ultralytics/yolo/data/utils.py @@ -44,7 +44,7 @@ def img2label_paths(img_paths): def get_hash(paths): # Returns a single hash value of a list of paths (files or dirs) size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) # sizes - h = hashlib.md5(str(size).encode()) # hash sizes + h = hashlib.sha256(str(size).encode()) # hash sizes h.update("".join(paths).encode()) # hash paths return h.hexdigest() # return hash diff --git a/ultralytics/yolo/engine/trainer.py b/ultralytics/yolo/engine/trainer.py index ce7c344..31561d1 100644 --- a/ultralytics/yolo/engine/trainer.py +++ b/ultralytics/yolo/engine/trainer.py @@ -5,6 +5,7 @@ Simple training loop; Boilerplate that could apply to any arbitrary neural netwo import os import subprocess +import sys import time from collections import defaultdict from copy import deepcopy @@ -28,10 +29,10 @@ from ultralytics.yolo.utils import (DEFAULT_CFG, LOGGER, RANK, SETTINGS, TQDM_BA yaml_save) from ultralytics.yolo.utils.autobatch import check_train_batch_size from ultralytics.yolo.utils.checks import check_file, check_imgsz, print_args -from ultralytics.yolo.utils.dist import ddp_cleanup, generate_ddp_command +from ultralytics.yolo.utils.dist import ddp_cleanup, generate_ddp_file, find_free_network_port from ultralytics.yolo.utils.files import get_latest_run, increment_path from ultralytics.yolo.utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, init_seeds, one_cycle, - select_device, strip_optimizer) + select_device, strip_optimizer, TORCH_1_9) class BaseTrainer: @@ -174,13 +175,18 @@ class BaseTrainer: # Run subprocess if DDP training, else train normally if world_size > 1 and "LOCAL_RANK" not in os.environ: - command = generate_ddp_command(world_size, self) + # cmd, file = generate_ddp_command(world_size, self) # security vulnerability in Snyk scans + file = generate_ddp_file(self) if sys.argv[0].endswith('yolo') else os.path.abspath(sys.argv[0]) + torch_distributed_cmd = "torch.distributed.run" if TORCH_1_9 else "torch.distributed.launch" + cmd = [ + sys.executable, "-m", torch_distributed_cmd, "--nproc_per_node", f"{world_size}", "--master_port", + f"{find_free_network_port()}", file] + sys.argv[1:] try: - subprocess.run(command) + subprocess.run(cmd, check=True) except Exception as e: - self.console(e) + self.console.warning(e) finally: - ddp_cleanup(command, self) + ddp_cleanup(self, file) else: self._do_train(int(os.getenv("RANK", -1)), world_size) diff --git a/ultralytics/yolo/utils/dist.py b/ultralytics/yolo/utils/dist.py index 3e0681d..2da0e80 100644 --- a/ultralytics/yolo/utils/dist.py +++ b/ultralytics/yolo/utils/dist.py @@ -44,21 +44,15 @@ def generate_ddp_file(trainer): def generate_ddp_command(world_size, trainer): import __main__ # noqa local import to avoid https://github.com/Lightning-AI/lightning/issues/15218 - file_name = os.path.abspath(sys.argv[0]) - using_cli = not file_name.endswith(".py") - if using_cli: - file_name = generate_ddp_file(trainer) + file = generate_ddp_file(trainer) if sys.argv[0].endswith('yolo') else os.path.abspath(sys.argv[0]) torch_distributed_cmd = "torch.distributed.run" if TORCH_1_9 else "torch.distributed.launch" - return [ + cmd = [ sys.executable, "-m", torch_distributed_cmd, "--nproc_per_node", f"{world_size}", "--master_port", - f"{find_free_network_port()}", file_name] + sys.argv[1:] + f"{find_free_network_port()}", file] + sys.argv[1:] + return cmd, file -def ddp_cleanup(command, trainer): +def ddp_cleanup(trainer, file): # delete temp file if created - tempfile_suffix = f"{id(trainer)}.py" - if tempfile_suffix in "".join(command): - for chunk in command: - if tempfile_suffix in chunk: - os.remove(chunk) - break + if f"{id(trainer)}.py" in file: # if temp_file suffix in file + os.remove(file)