bugfix: Added additional checks when saving/deleting histories.

GaiZhenbiao · Sep 18, 2024 · 526c615 · 526c615
1 parent 226a9b2
commit 526c615
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/modules/models/base_model.py b/modules/models/base_model.py
@@ -1109,6 +1109,10 @@ def delete_chat_history(self, filename):
         else:
             history_file_path = filename
         md_history_file_path = history_file_path[:-5] + ".md"
+        # check if history file path matches user_name
+        # if user access control is not enabled, user_name is empty, don't check
+        assert os.path.dirname(history_file_path) == self.user_name or self.user_name == ""
+        assert os.path.dirname(md_history_file_path) == self.user_name or self.user_name == ""
         try:
             os.remove(history_file_path)
             os.remove(md_history_file_path)

diff --git a/modules/utils.py b/modules/utils.py
@@ -418,6 +418,9 @@ def save_file(filename, model, chatbot):
     else:
         history_file_path = os.path.join(HISTORY_DIR, user_name, filename)
 
+    # check if history file path matches user_name
+    # if user access control is not enabled, user_name is empty, don't check
+    assert os.path.dirname(history_file_path) == model.user_name or model.user_name == ""
     with open(history_file_path, "w", encoding="utf-8") as f:
         json.dump(json_s, f, ensure_ascii=False, indent=4)