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

