From 454655dcc3e3939f84cb75ce7b762c1dd08cb906 Mon Sep 17 00:00:00 2001 From: yihong1120 Date: Sat, 27 Jul 2024 10:04:04 +0800 Subject: [PATCH] Boost security --- examples/YOLOv8_server_api/model_downloader.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/examples/YOLOv8_server_api/model_downloader.py b/examples/YOLOv8_server_api/model_downloader.py index c143112..11cfc73 100644 --- a/examples/YOLOv8_server_api/model_downloader.py +++ b/examples/YOLOv8_server_api/model_downloader.py @@ -1,7 +1,7 @@ from __future__ import annotations import datetime -from pathlib import Path +import os import requests from flask import Blueprint @@ -45,12 +45,21 @@ def download_model(model_name): response.headers['Last-Modified'], '%a, %d %b %Y %H:%M:%S GMT', ) - local_file_path = Path(MODELS_DIRECTORY) / model_name + + # Use os.path.join to safely construct the file path + local_file_path = os.path.join(MODELS_DIRECTORY, model_name) + + # Ensure the constructed path is within the expected directory + common_path = os.path.commonpath( + [local_file_path, MODELS_DIRECTORY], + ) + if common_path != MODELS_DIRECTORY: + return jsonify({'error': 'Invalid model name.'}), 400 # Check local file's last modified time - if local_file_path.exists(): + if os.path.exists(local_file_path): local_last_modified = datetime.datetime.fromtimestamp( - local_file_path.stat().st_mtime, + os.path.getmtime(local_file_path), ) # If local file is up-to-date, return 304 Not Modified