瀏覽代碼

Prevent SSRF and HTML injection

Rodrigo Agundez 8 月之前
父節點
當前提交
167c8bf00d
共有 1 個文件被更改,包括 7 次插入6 次删除
  1. 7 6
      backend/open_webui/utils/pdf_generator.py

+ 7 - 6
backend/open_webui/utils/pdf_generator.py

@@ -2,6 +2,7 @@ from datetime import datetime
 from io import BytesIO
 from io import BytesIO
 from pathlib import Path
 from pathlib import Path
 from typing import Dict, Any, List
 from typing import Dict, Any, List
+from html import escape
 
 
 from markdown import markdown
 from markdown import markdown
 
 
@@ -11,7 +12,6 @@ from fpdf import FPDF
 from open_webui.env import STATIC_DIR, FONTS_DIR
 from open_webui.env import STATIC_DIR, FONTS_DIR
 from open_webui.models.chats import ChatTitleMessagesForm
 from open_webui.models.chats import ChatTitleMessagesForm
 
 
-
 class PDFGenerator:
 class PDFGenerator:
     """
     """
     Description:
     Description:
@@ -41,13 +41,13 @@ class PDFGenerator:
 
 
     def _build_html_message(self, message: Dict[str, Any]) -> str:
     def _build_html_message(self, message: Dict[str, Any]) -> str:
         """Build HTML for a single message."""
         """Build HTML for a single message."""
-        role = message.get("role", "user")
-        content = message.get("content", "")
+        role = escape(message.get("role", "user"))
+        content = escape(message.get("content", ""))
         timestamp = message.get("timestamp")
         timestamp = message.get("timestamp")
 
 
-        model = message.get("model") if role == "assistant" else ""
+        model = escape(message.get("model") if role == "assistant" else "")
 
 
-        date_str = self.format_timestamp(timestamp) if timestamp else ""
+        date_str = escape(self.format_timestamp(timestamp) if timestamp else "")
 
 
         # extends pymdownx extension to convert markdown to html.
         # extends pymdownx extension to convert markdown to html.
         # - https://facelessuser.github.io/pymdown-extensions/usage_notes/
         # - https://facelessuser.github.io/pymdown-extensions/usage_notes/
@@ -76,6 +76,7 @@ class PDFGenerator:
 
 
     def _generate_html_body(self) -> str:
     def _generate_html_body(self) -> str:
         """Generate the full HTML body for the PDF."""
         """Generate the full HTML body for the PDF."""
+        escaped_title = escape(self.form_data.title)
         return f"""
         return f"""
         <html>
         <html>
             <head>
             <head>
@@ -84,7 +85,7 @@ class PDFGenerator:
             <body>
             <body>
             <div>
             <div>
                 <div>
                 <div>
-                    <h2>{self.form_data.title}</h2>
+                    <h2>{escaped_title}</h2>
                     {self.messages_html}
                     {self.messages_html}
                 </div>
                 </div>
             </div>
             </div>