A-Tune/fix-collection-train-file-overwriting-through-soft-links.patch

58 lines
2.0 KiB
Diff

From c5e491e5dffab4dda814f2e1ba11c21714cac0c6 Mon Sep 17 00:00:00 2001
From: zhoupengcheng <zhoupengcheng11@huawei.com>
Date: Wed, 1 Nov 2023 11:14:37 +0800
Subject: [PATCH] fix-collection-train-file-overwriting-through-soft-links.patch
---
analysis/atuned/collector.py | 10 +++++++++-
analysis/engine/train.py | 4 +++-
2 files changed, 12 insertions(+), 2 deletions(-)
diff --git a/analysis/atuned/collector.py b/analysis/atuned/collector.py
index 4749284..9a264dd 100755
--- a/analysis/atuned/collector.py
+++ b/analysis/atuned/collector.py
@@ -39,6 +39,15 @@ class Collector(Resource):
args = COLLECTOR_POST_PARSER.parse_args()
current_app.logger.info(args)
n_pipe = get_npipe(args.get("pipe"))
+
+ path = args.get("file")
+ path = os.path.abspath(path)
+ if not path.startswith("/var/atune_data/collection/"):
+ return "Files outside the /var/atune_data/collection/ directory cannot be modified.", 400
+
+ if os.path.exists(path):
+ return "File already exists!", 400
+
monitors = []
mpis = []
field_name = []
@@ -91,7 +100,6 @@ class Collector(Resource):
if n_pipe is not None:
n_pipe.close()
- path = args.get("file")
save_file(path, data, field_name)
result = {}
result["path"] = path
diff --git a/analysis/engine/train.py b/analysis/engine/train.py
index 7608660..462b16c 100644
--- a/analysis/engine/train.py
+++ b/analysis/engine/train.py
@@ -49,8 +49,10 @@ class Training(Resource):
return "Illegal model name provide: {}".format(err), 400
characterization = WorkloadCharacterization(model_path)
+ output_path = TRAINING_MODEL_PATH + model_name
+ if os.path.exists(output_path):
+ return "File already exists!", 400
try:
- output_path = TRAINING_MODEL_PATH + model_name
characterization.retrain(data_path, output_path)
except Exception as err:
LOGGER.error(err)
--
2.33.0