Przeglądaj źródła

Merge pull request #9733 from rragundez/patch-ssrf

Prevent SSRF and HTML injection
Timothy Jaeryang Baek 1 rok temu
rodzic
commit
de2026b5cc
1 zmienionych plików z 7 dodań i 6 usunięć
  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>