bugfix: Fix the XSS issue when uploading history files

Author: GaiZhenbiao

modules/models/base_model.py

@@ -961,7 +961,7 @@ def rename_chat_history(self, filename, chatbot):
         filename = os.path.basename(full_path)
 
         self.history_file_path = filename
-        save_file(filename, self, chatbot)
+        save_file(filename, self)
         return init_history_list(self.user_name)
 
     def auto_name_chat_history(
@@ -978,14 +978,14 @@ def auto_name_chat_history(
 
     def auto_save(self, chatbot=None):
         if chatbot is not None:
-            save_file(self.history_file_path, self, chatbot)
+            save_file(self.history_file_path, self)
 
     def export_markdown(self, filename, chatbot):
         if filename == "":
             return
         if not filename.endswith(".md"):
             filename += ".md"
-        save_file(filename, self, chatbot)
+        save_file(filename, self)
 
     def load_chat_history(self, new_history_file_path=None):
         logging.debug(f"{self.user_name} 加载对话历史中……")
@@ -1034,6 +1034,8 @@ def load_chat_history(self, new_history_file_path=None):
                     -len(saved_json["chatbot"]) :
                 ]
                 logging.info(f"Trimmed history: {saved_json['history']}")
+            # Sanitize chatbot
+            saved_json["chatbot"] = remove_html_tags(saved_json["chatbot"])
             logging.debug(f"{self.user_name} 加载对话历史完毕")
             self.history = saved_json["history"]
             self.single_turn = saved_json.get("single_turn", self.single_turn)

modules/utils.py

@@ -243,6 +243,21 @@ def convert_mdtext(md_text):  # deprecated
     output += ALREADY_CONVERTED_MARK
     return output
 
+def remove_html_tags(data):
+    def clean_text(text):
+        # Remove all HTML tags
+        cleaned = re.sub(r'<[^>]+>', '', text)
+        # Remove any remaining HTML entities
+        cleaned = re.sub(r'&[#\w]+;', '', cleaned)
+        # Remove extra whitespace and newlines
+        cleaned = re.sub(r'\s+', ' ', cleaned)
+        return cleaned.strip()
+
+    return [
+        [clean_text(item) for item in sublist]
+        for sublist in data
+    ]
+
 
 def clip_rawtext(chat_message, need_escape=True):
     # first, clip hr line
@@ -380,9 +395,10 @@ def construct_assistant(text):
     return construct_text("assistant", text)
 
 
-def save_file(filename, model, chatbot):
+def save_file(filename, model):
     system = model.system_prompt
     history = model.history
+    chatbot = [(history[i]["content"], history[i + 1]["content"]) for i in range(0, len(history), 2)]
     user_name = model.user_name
     os.makedirs(os.path.join(HISTORY_DIR, user_name), exist_ok=True)
     if filename is None: