oobabooga/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2026-03-20 12:24:38 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

Security: server-side file save roots, image URL SSRF protection, extension allowlist

Browse source
This commit is contained in:
oobabooga 2026-03-17 22:24:36 -07:00
parent 88a318894c
commit 256431f258
6 changed files with 60 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
modules/chat.py
View file

@ -2612,19 +2612,23 @@ def handle_load_template_click(instruction_template):
def handle_save_template_click(instruction_template_str):
import gradio as gr
contents = generate_instruction_template_yaml(instruction_template_str)
root = str(shared.user_data_dir / 'instruction-templates') + '/'
return [
"My Template.yaml",
str(shared.user_data_dir / 'instruction-templates') + '/',
root,
contents,
root,
gr.update(visible=True)
]
def handle_delete_template_click(template):
import gradio as gr
root = str(shared.user_data_dir / 'instruction-templates') + '/'
return [
f"{template}.yaml",
str(shared.user_data_dir / 'instruction-templates') + '/',
root,
root,
gr.update(visible=False)
]

13
modules/image_utils.py
View file

@ -77,7 +77,18 @@ def process_message_content(content: Any) -> Tuple[str, List[Image.Image]]:
# Support external URLs
try:
import requests
response = requests.get(image_url, timeout=10)
from urllib.parse import urljoin
from modules.web_search import _validate_url
_validate_url(image_url)
url = image_url
for _ in range(5):
response = requests.get(url, timeout=10, allow_redirects=False)
if response.is_redirect and 'Location' in response.headers:
url = urljoin(url, response.headers['Location'])
_validate_url(url)
else:
break
response.raise_for_status()
image_data = response.content
image = Image.open(io.BytesIO(image_data))

4
modules/ui_chat.py
View file

@ -350,13 +350,13 @@ def create_event_handlers():
shared.gradio['load_template'].click(chat.handle_load_template_click, gradio('instruction_template'), gradio('instruction_template_str', 'instruction_template'), show_progress=False)
shared.gradio['save_template'].click(
ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(
chat.handle_save_template_click, gradio('instruction_template_str'), gradio('save_filename', 'save_root', 'save_contents', 'file_saver'), show_progress=False)
chat.handle_save_template_click, gradio('instruction_template_str'), gradio('save_filename', 'save_root', 'save_contents', 'save_root_state', 'file_saver'), show_progress=False)
shared.gradio['restore_character'].click(
ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(
chat.restore_character_for_ui, gradio('interface_state'), gradio('interface_state', 'name2', 'context', 'greeting', 'character_picture'), show_progress=False)
shared.gradio['delete_template'].click(chat.handle_delete_template_click, gradio('instruction_template'), gradio('delete_filename', 'delete_root', 'file_deleter'), show_progress=False)
shared.gradio['delete_template'].click(chat.handle_delete_template_click, gradio('instruction_template'), gradio('delete_filename', 'delete_root', 'delete_root_state', 'file_deleter'), show_progress=False)
shared.gradio['save_chat_history'].click(
lambda x: json.dumps(x, indent=4), gradio('history'), gradio('temporary_text')).then(
None, gradio('temporary_text', 'character_menu', 'mode'), None, js=f'(hist, char, mode) => {{{ui.save_files_js}; saveHistory(hist, char, mode)}}')

46
modules/ui_file_saving.py
View file

@ -9,6 +9,12 @@ from modules.utils import gradio, sanitize_filename
def create_ui():
mu = shared.args.multi_user
# Server-side per-session root paths for the generic file saver/deleter.
# Set by the handler that opens the dialog, read by the confirm handler.
# Using gr.State so they are session-scoped and safe for multi-user.
shared.gradio['save_root_state'] = gr.State(None)
shared.gradio['delete_root_state'] = gr.State(None)
# Text file saver
with gr.Group(visible=False, elem_classes='file-saver') as shared.gradio['file_saver']:
shared.gradio['save_filename'] = gr.Textbox(lines=1, label='File name')
@ -66,13 +72,13 @@ def create_event_handlers():
ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(
handle_save_preset_click, gradio('interface_state'), gradio('save_preset_contents', 'save_preset_filename', 'preset_saver'), show_progress=False)
shared.gradio['delete_preset'].click(handle_delete_preset_click, gradio('preset_menu'), gradio('delete_filename', 'delete_root', 'file_deleter'), show_progress=False)
shared.gradio['save_grammar'].click(handle_save_grammar_click, gradio('grammar_string'), gradio('save_contents', 'save_filename', 'save_root', 'file_saver'), show_progress=False)
shared.gradio['delete_grammar'].click(handle_delete_grammar_click, gradio('grammar_file'), gradio('delete_filename', 'delete_root', 'file_deleter'), show_progress=False)
shared.gradio['delete_preset'].click(handle_delete_preset_click, gradio('preset_menu'), gradio('delete_filename', 'delete_root', 'delete_root_state', 'file_deleter'), show_progress=False)
shared.gradio['save_grammar'].click(handle_save_grammar_click, gradio('grammar_string'), gradio('save_contents', 'save_filename', 'save_root', 'save_root_state', 'file_saver'), show_progress=False)
shared.gradio['delete_grammar'].click(handle_delete_grammar_click, gradio('grammar_file'), gradio('delete_filename', 'delete_root', 'delete_root_state', 'file_deleter'), show_progress=False)
shared.gradio['save_preset_confirm'].click(handle_save_preset_confirm_click, gradio('save_preset_filename', 'save_preset_contents'), gradio('preset_menu', 'preset_saver'), show_progress=False)
shared.gradio['save_confirm'].click(handle_save_confirm_click, gradio('save_root', 'save_filename', 'save_contents'), gradio('file_saver'), show_progress=False)
shared.gradio['delete_confirm'].click(handle_delete_confirm_click, gradio('delete_root', 'delete_filename'), gradio('file_deleter'), show_progress=False)
shared.gradio['save_confirm'].click(handle_save_confirm_click, gradio('save_root_state', 'save_filename', 'save_contents'), gradio('save_root_state', 'file_saver'), show_progress=False)
shared.gradio['delete_confirm'].click(handle_delete_confirm_click, gradio('delete_root_state', 'delete_filename'), gradio('delete_root_state', 'file_deleter'), show_progress=False)
shared.gradio['save_character_confirm'].click(handle_save_character_confirm_click, gradio('name2', 'greeting', 'context', 'character_picture', 'save_character_filename'), gradio('character_menu', 'character_saver'), show_progress=False)
shared.gradio['delete_character_confirm'].click(handle_delete_character_confirm_click, gradio('character_menu'), gradio('character_menu', 'character_deleter'), show_progress=False)
@ -105,24 +111,30 @@ def handle_save_preset_confirm_click(filename, contents):
]
def handle_save_confirm_click(root, filename, contents):
def handle_save_confirm_click(root_state, filename, contents):
try:
if root_state is None:
return None, gr.update(visible=False)
filename = sanitize_filename(filename)
utils.save_file(root + filename, contents)
utils.save_file(root_state + filename, contents)
except Exception:
traceback.print_exc()
return gr.update(visible=False)
return None, gr.update(visible=False)
def handle_delete_confirm_click(root, filename):
def handle_delete_confirm_click(root_state, filename):
try:
if root_state is None:
return None, gr.update(visible=False)
filename = sanitize_filename(filename)
utils.delete_file(root + filename)
utils.delete_file(root_state + filename)
except Exception:
traceback.print_exc()
return gr.update(visible=False)
return None, gr.update(visible=False)
def handle_save_character_confirm_click(name2, greeting, context, character_picture, filename):
@ -165,26 +177,32 @@ def handle_save_preset_click(state):
def handle_delete_preset_click(preset):
root = str(shared.user_data_dir / "presets") + "/"
return [
f"{preset}.yaml",
str(shared.user_data_dir / "presets") + "/",
root,
root,
gr.update(visible=True)
]
def handle_save_grammar_click(grammar_string):
root = str(shared.user_data_dir / "grammars") + "/"
return [
grammar_string,
"My Fancy Grammar.gbnf",
str(shared.user_data_dir / "grammars") + "/",
root,
root,
gr.update(visible=True)
]
def handle_delete_grammar_click(grammar_file):
root = str(shared.user_data_dir / "grammars") + "/"
return [
grammar_file,
str(shared.user_data_dir / "grammars") + "/",
root,
root,
gr.update(visible=True)
]

6
modules/ui_session.py
View file

@ -30,7 +30,7 @@ def create_ui():
if not mu:
shared.gradio['save_settings'].click(
ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(
handle_save_settings, gradio('interface_state', 'preset_menu', 'extensions_menu', 'show_controls', 'theme_state'), gradio('save_contents', 'save_filename', 'save_root', 'file_saver'), show_progress=False)
handle_save_settings, gradio('interface_state', 'preset_menu', 'extensions_menu', 'show_controls', 'theme_state'), gradio('save_contents', 'save_filename', 'save_root', 'save_root_state', 'file_saver'), show_progress=False)
shared.gradio['toggle_dark_mode'].click(
lambda x: 'dark' if x == 'light' else 'light', gradio('theme_state'), gradio('theme_state')).then(
@ -51,10 +51,12 @@ def create_ui():
def handle_save_settings(state, preset, extensions, show_controls, theme):
contents = ui.save_settings(state, preset, extensions, show_controls, theme, manual_save=True)
root = str(shared.user_data_dir) + "/"
return [
contents,
"settings.yaml",
str(shared.user_data_dir) + "/",
root,
root,
gr.update(visible=True)
]

4
modules/utils.py
View file

@ -47,6 +47,10 @@ def save_file(fname, contents):
logger.error(f'Invalid file path: \"{fname}\"')
return
if Path(abs_path_str).suffix.lower() not in ('.yaml', '.yml', '.json', '.txt', '.gbnf'):
logger.error(f'Refusing to save file with disallowed extension: \"{fname}\"')
return
with open(abs_path_str, 'w', encoding='utf-8') as f:
f.write(contents)