Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import time
 from datetime import datetime
 from collections import deque
+import html
 
 # --- GLOBAL SERVER STATE (In-Memory Only) ---
 class ChatServer:
@@ -21,16 +22,24 @@ class ChatServer:
             if u in self.user_rooms:
                 del self.user_rooms[u]
 
+    def _sanitize(self, text):
+        """Prevents HTML injection by escaping user input."""
+        # Escape all HTML characters (&, <, >, ", ')
+        escaped = html.escape(text)
+        # Re-allow line breaks only
+        return escaped.replace("\n", "<br>")
+
     def _add_message_to_queue(self, queue, user, message, profile_pic):
-        html_text = message.replace("\n", "<br>")
+        safe_text = self._sanitize(message)
+        
         # Stacking Logic: Check if last message was same user and same text
-        if queue and queue[-1]['user'] == user and queue[-1]['text'] == html_text:
+        if queue and queue[-1]['user'] == user and queue[-1]['text'] == safe_text:
             queue[-1]['count'] = queue[-1].get('count', 1) + 1
-            queue[-1]['time'] = datetime.now().strftime("%H:%M") # Update time to latest
+            queue[-1]['time'] = datetime.now().strftime("%H:%M") 
         else:
             msg_data = {
                 "user": user,
-                "text": html_text,
+                "text": safe_text,
                 "pic": profile_pic,
                 "time": datetime.now().strftime("%H:%M"),
                 "count": 1
@@ -56,14 +65,14 @@ def get_messages_html(messages):
     if not messages:
         return "<div style='color: gray; text-align: center; margin-top: 20px;'>No messages yet.</div>"
     
-    html = "<div id='chat-flow' style='display: flex; flex-direction: column; gap: 12px; font-family: sans-serif;'>"
+    html_output = "<div id='chat-flow' style='display: flex; flex-direction: column; gap: 12px; font-family: sans-serif;'>"
     for m in messages:
         pic_url = m.get('pic') or "https://huggingface.co/front/assets/huggingface_logo-noborder.svg"
         display_text = m['text']
         if m.get('count', 1) > 1:
             display_text += f" <span style='color: #888; font-size: 0.8em; font-weight: normal;'>(x{m['count']})</span>"
             
-        html += f"""
+        html_output += f"""
         <div class="msg-row" style="display: flex; align-items: flex-start; gap: 10px; margin-bottom: 4px;">
             <img src="{pic_url}" style="width: 36px; height: 36px; border-radius: 50%; border: 1px solid rgba(255,255,255,0.1); flex-shrink: 0;">
             <div style="background: rgba(255, 255, 255, 0.07); padding: 10px 14px; border-radius: 0 16px 16px 16px; max-width: 85%;">
@@ -76,8 +85,9 @@ def get_messages_html(messages):
             </div>
         </div>
         """
-    # Robust Auto-Scroll Script using MutationObserver
-    html += """
+    
+    # SYSTEM SCRIPT (Safe since it's hardcoded by us, not user input)
+    html_output += """
     </div>
     <script>
         (function() {
@@ -89,33 +99,32 @@ def get_messages_html(messages):
                 const observer = new MutationObserver(scrollDown);
                 observer.observe(container, { childList: true, subtree: true });
                 container.dataset.observed = "true";
-                scrollDown(); // Initial scroll
+                scrollDown();
             } else if (container) {
                 container.scrollTo({ top: container.scrollHeight, behavior: 'smooth' });
             }
         })();
     </script>
     """
-    return html
+    return html_output
 
 def get_active_users_html(current_user):
     server.prune_server()
-    html = "<div style='display: flex; flex-direction: column; gap: 8px; font-family: sans-serif;'>"
+    html_list = "<div style='display: flex; flex-direction: column; gap: 8px; font-family: sans-serif;'>"
     for user, data in server.active_users.items():
         pic = data.get('pic') or "https://huggingface.co/front/assets/huggingface_logo-noborder.svg"
         label = f"{user}" + (" (You)" if user == current_user else "")
-        html += f"""
+        html_list += f"""
         <div style="display: flex; align-items: center; gap: 10px; font-size: 0.85em; padding: 4px;">
             <img src="{pic}" style="width: 22px; height: 22px; border-radius: 50%;">
             <span style="color: #ccc;">{label}</span>
             <div style="width: 8px; height: 8px; background: #4ade80; border-radius: 50%; margin-left: auto;"></div>
         </div>
         """
-    return html + "</div>"
+    return html_list + "</div>"
 
 # --- CSS ---
 custom_css = """
-/* Modern Font Stack */
 @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');
 
 body, .gradio-container, .main-title, #chat-container, #side-panel, textarea, button {
@@ -158,8 +167,7 @@ body, .gradio-container, .main-title, #chat-container, #side-panel, textarea, bu
 }
 """
 
-# Updated Blocks constructor for Gradio 6.0+ compatibility
-with gr.Blocks() as demo:
+with gr.Blocks(css=custom_css) as demo:
     user_session = gr.State(None)
     current_room = gr.State("Main")
     dm_target = gr.State(None)
@@ -190,7 +198,7 @@ with gr.Blocks() as demo:
             
             with gr.Row():
                 msg_input = gr.Textbox(
-                    placeholder="Type a message...", 
+                    placeholder="Type a message (HTML is disabled)...", 
                     show_label=False, 
                     scale=10,
                     elem_id="msg-input"
@@ -226,6 +234,7 @@ with gr.Blocks() as demo:
 
     def send_msg(text, session, room, dm_user):
         if not text or not session: return ""
+        # The server class now handles sanitization
         if dm_user:
             server.send_dm(session['name'], dm_user, text, session['pic'])
         else:
@@ -246,8 +255,4 @@ with gr.Blocks() as demo:
     room_list.change(lambda name: (name, None), room_list, [current_room, dm_target])
     start_dm_btn.click(lambda target: ("Main", target), dm_user_input, [current_room, dm_target])
 
-# Pass theme and css to launch() for Gradio 6.0+
-demo.launch(
-    theme=gr.themes.Soft(primary_hue="orange", neutral_hue="slate"), 
-    css=custom_css
-)
+demo.launch()