Skip to content

Commit

Permalink
bugfix: Fix the XSS issue when uploading history files
Browse files Browse the repository at this point in the history
  • Loading branch information
GaiZhenbiao committed Sep 18, 2024
1 parent 71cb89c commit 868767e
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 4 deletions.
8 changes: 5 additions & 3 deletions modules/models/base_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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} 加载对话历史中……")
Expand Down Expand Up @@ -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)
Expand Down
18 changes: 17 additions & 1 deletion modules/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit 868767e

Please sign in to comment.