Harden chat streaming DOM sinks (#2498)

2026-06-04 20:49:37 +02:00
parent e113c10d01
commit b59bbe80ce
6 changed files with 190 additions and 44 deletions
--- a/static/js/chat.js
+++ b/static/js/chat.js
@@ -844,7 +844,7 @@ import createResearchSynapse from './researchSynapse.js';
      var _charNameInit = presetsModule.getCharacterName ? presetsModule.getCharacterName() : '';
      if (_charNameInit) roleLabel = _charNameInit;
      const roleTs = new Date().toLocaleTimeString([], {hour: '2-digit', minute:'2-digit'});
-      holder.innerHTML = `<div class="role">${roleLabel} <span class="role-timestamp">${roleTs}</span></div><div class="body"></div>`;
+      holder.innerHTML = `<div class="role">${uiModule.esc(roleLabel)} <span class="role-timestamp">${roleTs}</span></div><div class="body"></div>`;
      _applyModelColor(holder.querySelector('.role'), modelName);
      holder.style.position = 'relative';
      
@@ -2002,7 +2002,7 @@ import createResearchSynapse from './researchSynapse.js';
                const node = document.createElement('div')
                node.className = 'agent-thread-node running';
                const cmdHtml = cmd ? `<pre class="agent-thread-cmd">${esc(cmd)}</pre>` : '';
-                node.innerHTML = `<div class="agent-thread-dot"></div><div class="agent-thread-header"><span class="agent-thread-icon">\u25B6</span><span class="agent-thread-tool">${toolLabel}</span><span class="agent-thread-wave">▁▂▃</span></div><div class="agent-thread-content">${cmdHtml}</div>`;
+                node.innerHTML = `<div class="agent-thread-dot"></div><div class="agent-thread-header"><span class="agent-thread-icon">\u25B6</span><span class="agent-thread-tool">${esc(toolLabel)}</span><span class="agent-thread-wave">▁▂▃</span></div><div class="agent-thread-content">${cmdHtml}</div>`;
                // Expand/collapse via delegated click handler (init at module bottom).
                threadWrap.appendChild(node);
                currentToolBubble = node;
@@ -2132,10 +2132,19 @@ import createResearchSynapse from './researchSynapse.js';
                if (json.screenshot && currentToolBubble) {
                  const contentEl = currentToolBubble.querySelector('.agent-thread-content');
                  if (contentEl) {
-                    const details = document.createElement('details');
-                    details.className = 'agent-tool-output';
-                    details.innerHTML = `<summary>Screenshot</summary><img src="${json.screenshot}" style="max-width:100%;border-radius:6px;margin-top:6px;border:1px solid var(--border)" />`;
-                    contentEl.appendChild(details);
+                    const screenshotSrc = chatRenderer.safeToolScreenshotSrc(json.screenshot);
+                    if (screenshotSrc) {
+                      const details = document.createElement('details');
+                      details.className = 'agent-tool-output';
+                      const summary = document.createElement('summary');
+                      summary.textContent = 'Screenshot';
+                      const img = document.createElement('img');
+                      img.src = screenshotSrc;
+                      img.style.cssText = 'max-width:100%;border-radius:6px;margin-top:6px;border:1px solid var(--border)';
+                      details.appendChild(summary);
+                      details.appendChild(img);
+                      contentEl.appendChild(details);
+                    }
                  }
                }
                // --- Reload sessions after manage_session tool (delete, rename, etc.) ---
@@ -3271,7 +3280,7 @@ import createResearchSynapse from './researchSynapse.js';
      var meta = sessionModule.getSessions().find(function(s) { return s.id === sessionId; });
      var roleLabel = _shortModel(meta && meta.model);
      var roleTs = new Date().toLocaleTimeString([], {hour: '2-digit', minute:'2-digit'});
-      holder.innerHTML = '<div class="role">' + roleLabel + ' <span class="role-timestamp">' + roleTs + '</span></div><div class="body"></div>';
+      holder.innerHTML = '<div class="role">' + uiModule.esc(roleLabel) + ' <span class="role-timestamp">' + roleTs + '</span></div><div class="body"></div>';
      _applyModelColor(holder.querySelector('.role'), meta && meta.model);

      var bodyDiv = holder.querySelector('.body');
@@ -4073,7 +4082,7 @@ import createResearchSynapse from './researchSynapse.js';
      const roleTs = new Date().toLocaleTimeString([], {hour: '2-digit', minute:'2-digit'});
      const agentMeta = sessionModule.getSessions().find(s => s.id === sessionModule.getCurrentSessionId());
      const agentModelLabel = _shortModel(agentMeta?.model);
-      holder.innerHTML = `<div class="role">${agentModelLabel} <span class="role-timestamp">${roleTs}</span></div><div class="body"></div>`;
+      holder.innerHTML = `<div class="role">${uiModule.esc(agentModelLabel)} <span class="role-timestamp">${roleTs}</span></div><div class="body"></div>`;
      _applyModelColor(holder.querySelector('.role'), agentMeta?.model);
      box.appendChild(holder);

--- a/static/js/chatRenderer.js
+++ b/static/js/chatRenderer.js
@@ -26,6 +26,29 @@ function _safeHref(url) {
  return '#';
 }

+export function safeToolScreenshotSrc(raw) {
+  const src = String(raw || '').trim();
+  if (/^data:image\/(?:png|jpe?g|gif|webp);base64,[a-z0-9+/=\s]+$/i.test(src)) {
+    return src;
+  }
+  return '';
+}
+
+export function safeDisplayImageSrc(raw) {
+  const src = String(raw || '').trim();
+  if (!src) return '';
+  if (/^data:image\/(?:png|jpe?g|gif|webp);base64,[a-z0-9+/=\s]+$/i.test(src)) {
+    return src;
+  }
+  try {
+    const parsed = new URL(src, window.location.origin);
+    if (parsed.protocol === 'http:' || parsed.protocol === 'https:') {
+      return parsed.href;
+    }
+  } catch (_) {}
+  return '';
+}
+
 function _makeActionBtn(className, title, text, handler) {
  const btn = document.createElement('button');
  btn.className = className;
@@ -1058,12 +1081,19 @@ export function buildImageBubble(imageUrl, prompt, model, size, quality, imageId
  const body = document.createElement('div');
  body.className = 'body';

+  const safeImageUrl = safeDisplayImageSrc(imageUrl);
+  if (!safeImageUrl) {
+    body.textContent = '[Image unavailable]';
+    wrap.appendChild(body);
+    return wrap;
+  }
+
  const img = document.createElement('img');
  img.className = 'generated-image';
  img.alt = prompt || 'Generated image';
  img.title = prompt || 'Generated image';
-  img.src = imageUrl;
-  img.addEventListener('click', () => { window.open(img.src, '_blank'); });
+  img.src = safeImageUrl;
+  img.addEventListener('click', () => { window.open(safeImageUrl, '_blank', 'noopener,noreferrer'); });
  body.appendChild(img);

  if (prompt) {
@@ -1953,8 +1983,9 @@ export function addMessage(role, content, modelName, metadata) {
            if (ev.output && ev.output.trim()) {
              outHtml = `<details class="agent-tool-output"><summary>Output</summary><pre>${esc(ev.output)}</pre></details>`;
            }
-            if (ev.screenshot) {
-              outHtml += `<details class="agent-tool-output"><summary>Screenshot</summary><img src="${esc(ev.screenshot)}" style="max-width:100%;border-radius:6px;margin-top:6px;border:1px solid var(--border)" /></details>`;
+            const screenshotSrc = safeToolScreenshotSrc(ev.screenshot);
+            if (screenshotSrc) {
+              outHtml += `<details class="agent-tool-output"><summary>Screenshot</summary><img src="${esc(screenshotSrc)}" style="max-width:100%;border-radius:6px;margin-top:6px;border:1px solid var(--border)" /></details>`;
            }
            // File-write/edit diff (persisted in the tool event) \u2014 re-render it
            // so it survives reload, matching the live stream.
@@ -2308,6 +2339,8 @@ const chatRenderer = {
  updateSessionCostUI,
  roleTimestamp,
  stripToolBlocks,
+  safeToolScreenshotSrc,
+  safeDisplayImageSrc,
  buildSourcesBox,
  buildFindingsBox,
  appendReportButton,
--- a/static/js/compare/selector.js
+++ b/static/js/compare/selector.js
@@ -1195,7 +1195,7 @@ async function showModelSelector() {
              const row = document.createElement('div');
              row.className = 'compare-probe-row';
              row.dataset.idx = 'p' + i;
-              row.innerHTML = `<span class="compare-probe-spinner">▁▂▃</span><span class="compare-probe-name">${p.label || p.id}</span><span class="compare-probe-status"></span>`;
+              row.innerHTML = `<span class="compare-probe-spinner">▁▂▃</span><span class="compare-probe-name">${escapeHtml(p.label || p.id)}</span><span class="compare-probe-status"></span>`;
              const waveEl = row.querySelector('.compare-probe-spinner');
              const waveFrames = WAVE_FRAMES;
              let wIdx = 0;
--- a/static/js/compare/stream.js
+++ b/static/js/compare/stream.js
@@ -1,7 +1,7 @@
 // compare/stream.js — SSE streaming to panes
 import state from './state.js';
 import { addFinishBadge } from './vote.js';
-import { getModelCost } from '../chatRenderer.js';
+import { getModelCost, safeDisplayImageSrc } from '../chatRenderer.js';
 import markdownModule from '../markdown.js';
 import spinnerModule from '../spinner.js';
 import uiModule from '../ui.js';
@@ -11,6 +11,16 @@ var escapeHtml = uiModule.esc;

 const WAVE_FRAMES = ['▁▂▃', '▂▃▄', '▃▄▅', '▄▅▆', '▅▆▇', '▆▅▄', '▅▄▃', '▄▃▂'];

+function _safeHttpHref(raw) {
+  try {
+    const parsed = new URL(String(raw || '').trim(), window.location.origin);
+    if (parsed.protocol === 'http:' || parsed.protocol === 'https:') {
+      return parsed.href;
+    }
+  } catch (_) {}
+  return '';
+}
+
 // ── Lazy-registered functions from compare.js (avoids circular deps) ──
 let _rerollPane = null;
 let _autoPreviewHtml = null;
@@ -36,9 +46,12 @@ function _renderSearchResults(data) {
    const card = document.createElement('div');
    card.className = 'compare-search-result';
    const titleLink = document.createElement('a');
-    titleLink.href = r.url || '#';
-    titleLink.target = '_blank';
-    titleLink.rel = 'noopener';
+    const safeUrl = _safeHttpHref(r.url);
+    if (safeUrl) {
+      titleLink.href = safeUrl;
+      titleLink.target = '_blank';
+      titleLink.rel = 'noopener noreferrer';
+    }
    titleLink.className = 'search-result-title';
    titleLink.textContent = r.title || 'Untitled';
    card.appendChild(titleLink);
@@ -344,7 +357,7 @@ async function streamToPane(paneIdx, sessionId, message, aiMsgEl, opts) {
              const cmdHtml = cmd ? `<pre class="agent-thread-cmd">${escapeHtml(cmd)}</pre>` : '';
              const node = document.createElement('div');
              node.className = 'agent-thread-node running';
-              node.innerHTML = `<div class="agent-thread-dot"></div><div class="agent-thread-header"><span class="agent-thread-icon">\u25B6</span><span class="agent-thread-tool">${toolLabel}</span><span class="agent-thread-wave">▁▂▃</span></div><div class="agent-thread-content">${cmdHtml}</div>`;
+              node.innerHTML = `<div class="agent-thread-dot"></div><div class="agent-thread-header"><span class="agent-thread-icon">\u25B6</span><span class="agent-thread-tool">${escapeHtml(toolLabel)}</span><span class="agent-thread-wave">▁▂▃</span></div><div class="agent-thread-content">${cmdHtml}</div>`;
              node.querySelector('.agent-thread-header').addEventListener('click', () => node.classList.toggle('open'));
              // Animate wave
              const waveEl = node.querySelector('.agent-thread-wave');
@@ -363,28 +376,33 @@ async function streamToPane(paneIdx, sessionId, message, aiMsgEl, opts) {
            if (json.image_url) {
              // Stop image spinner and render generated image in pane
              if (aiMsgEl._imgSpinner) { aiMsgEl._imgSpinner.destroy(); aiMsgEl._imgSpinner = null; }
+              const safeImageUrl = safeDisplayImageSrc(json.image_url);
              aiBody.innerHTML = '';
-              const img = document.createElement('img');
-              img.className = 'compare-gen-image';
-              img.src = json.image_url;
-              img.alt = json.image_prompt || '';
-              img.title = json.image_prompt || '';
-              img.addEventListener('click', () => window.open(img.src, '_blank'));
-              aiBody.appendChild(img);
-              if (json.image_prompt) {
-                const caption = document.createElement('div');
-                caption.style.cssText = 'font-size:0.82em;color:color-mix(in srgb, var(--fg) 55%, transparent);margin-top:6px;line-height:1.4;';
-                caption.textContent = json.image_prompt;
-                aiBody.appendChild(caption);
+              if (!safeImageUrl) {
+                aiBody.textContent = '[Image unavailable]';
+              } else {
+                const img = document.createElement('img');
+                img.className = 'compare-gen-image';
+                img.src = safeImageUrl;
+                img.alt = json.image_prompt || '';
+                img.title = json.image_prompt || '';
+                img.addEventListener('click', () => window.open(safeImageUrl, '_blank', 'noopener,noreferrer'));
+                aiBody.appendChild(img);
+                if (json.image_prompt) {
+                  const caption = document.createElement('div');
+                  caption.style.cssText = 'font-size:0.82em;color:color-mix(in srgb, var(--fg) 55%, transparent);margin-top:6px;line-height:1.4;';
+                  caption.textContent = json.image_prompt;
+                  aiBody.appendChild(caption);
+                }
+                // Show model name below image (hidden in blind mode until vote)
+                if (json.image_model && !state._blindMode) {
+                  const modelLabel = document.createElement('div');
+                  modelLabel.style.cssText = 'font-size:0.75em;color:color-mix(in srgb, var(--fg) 40%, transparent);margin-top:4px;';
+                  modelLabel.textContent = json.image_model;
+                  aiBody.appendChild(modelLabel);
+                }
+                aiMsgEl._imageData = { url: safeImageUrl, prompt: json.image_prompt, model: json.image_model, size: json.image_size, quality: json.image_quality };
              }
-              // Show model name below image (hidden in blind mode until vote)
-              if (json.image_model && !state._blindMode) {
-                const modelLabel = document.createElement('div');
-                modelLabel.style.cssText = 'font-size:0.75em;color:color-mix(in srgb, var(--fg) 40%, transparent);margin-top:4px;';
-                modelLabel.textContent = json.image_model;
-                aiBody.appendChild(modelLabel);
-              }
-              aiMsgEl._imageData = { url: json.image_url, prompt: json.image_prompt, model: json.image_model, size: json.image_size, quality: json.image_quality };
            } else if (currentToolBlock) {
              // Stop wave animation
              if (currentToolBlock._waveInterval) { clearInterval(currentToolBlock._waveInterval); currentToolBlock._waveInterval = null; }
--- a/static/js/group.js
+++ b/static/js/group.js
@@ -676,7 +676,7 @@ function _createGroupBubble(model, box) {
  // Role label — use character name if assigned, otherwise model name
  const roleLabel = model._groupName || (model.character ? model.character.characterName : chatRenderer.shortModel(model.mid));
  const roleTs = new Date().toLocaleTimeString([], { hour: '2-digit', minute: '2-digit' });
-  wrap.innerHTML = `<div class="role">${roleLabel} <span class="role-timestamp">${roleTs}</span></div><div class="body"></div>`;
+  wrap.innerHTML = `<div class="role">${uiModule.esc(roleLabel)} <span class="role-timestamp">${roleTs}</span></div><div class="body"></div>`;
  chatRenderer.applyModelColor(wrap.querySelector('.role'), model.mid);

  // Spinner — identical to chat.js line 3062
@@ -860,11 +860,14 @@ async function _streamToHolder(modelIdx, sessionId, msg, holderEl, abortCtrl) {
          }
          // Generated image
          else if (json.type === 'generated_image' && json.url) {
-            const img = document.createElement('img');
-            img.src = json.url;
-            img.style.cssText = 'max-width:100%;border-radius:8px;margin:8px 0;';
-            img.loading = 'lazy';
-            bodyEl.appendChild(img);
+            const safeImageUrl = chatRenderer.safeDisplayImageSrc(json.url);
+            if (safeImageUrl) {
+              const img = document.createElement('img');
+              img.src = safeImageUrl;
+              img.style.cssText = 'max-width:100%;border-radius:8px;margin:8px 0;';
+              img.loading = 'lazy';
+              bodyEl.appendChild(img);
+            }
          }
          // Error
          else if (json.error) {
--- a/tests/test_chat_tool_screenshot_xss.py
+++ b/tests/test_chat_tool_screenshot_xss.py
@@ -0,0 +1,83 @@
+"""Regression guards for agent-tool screenshot DOM sinks."""
+
+from pathlib import Path
+
+
+_REPO = Path(__file__).resolve().parent.parent
+
+
+def test_live_tool_screenshot_does_not_template_raw_sse_value():
+    chat = (_REPO / "static" / "js" / "chat.js").read_text(encoding="utf-8")
+
+    assert "safeToolScreenshotSrc(json.screenshot)" in chat
+    assert 'img.src = screenshotSrc' in chat
+    assert 'details.innerHTML = `<summary>Screenshot</summary><img src="${json.screenshot}"' not in chat
+
+
+def test_restored_tool_screenshot_uses_raster_data_url_whitelist():
+    renderer = (_REPO / "static" / "js" / "chatRenderer.js").read_text(encoding="utf-8")
+
+    assert "export function safeToolScreenshotSrc(raw)" in renderer
+    assert "(?:png|jpe?g|gif|webp)" in renderer
+    assert "safeToolScreenshotSrc(ev.screenshot)" in renderer
+    assert 'src="${esc(ev.screenshot)}"' not in renderer
+
+
+def test_streaming_tool_labels_are_escaped_before_inner_html():
+    chat = (_REPO / "static" / "js" / "chat.js").read_text(encoding="utf-8")
+    compare = (_REPO / "static" / "js" / "compare" / "stream.js").read_text(encoding="utf-8")
+
+    assert '<span class="agent-thread-tool">${esc(toolLabel)}</span>' in chat
+    assert '<span class="agent-thread-tool">${toolLabel}</span>' not in chat
+    assert '<span class="agent-thread-tool">${escapeHtml(toolLabel)}</span>' in compare
+    assert '<span class="agent-thread-tool">${toolLabel}</span>' not in compare
+
+
+def test_generated_image_urls_are_vetted_before_assignment_or_open():
+    renderer = (_REPO / "static" / "js" / "chatRenderer.js").read_text(encoding="utf-8")
+    compare = (_REPO / "static" / "js" / "compare" / "stream.js").read_text(encoding="utf-8")
+    group = (_REPO / "static" / "js" / "group.js").read_text(encoding="utf-8")
+
+    assert "export function safeDisplayImageSrc(raw)" in renderer
+    assert "safeDisplayImageSrc(imageUrl)" in renderer
+    assert "img.src = safeImageUrl" in renderer
+    assert "window.open(safeImageUrl, '_blank', 'noopener,noreferrer')" in renderer
+    assert "safeDisplayImageSrc," in renderer
+    assert "safeDisplayImageSrc(json.image_url)" in compare
+    assert "img.src = json.image_url" not in compare
+    assert "chatRenderer.safeDisplayImageSrc(json.url)" in group
+    assert "img.src = json.url" not in group
+
+
+def test_group_chat_role_labels_are_escaped_before_inner_html():
+    group = (_REPO / "static" / "js" / "group.js").read_text(encoding="utf-8")
+
+    assert '<div class="role">${uiModule.esc(roleLabel)}' in group
+    assert '<div class="role">${roleLabel}' not in group
+
+
+def test_main_chat_role_labels_are_escaped_before_inner_html():
+    chat = (_REPO / "static" / "js" / "chat.js").read_text(encoding="utf-8")
+
+    assert '<div class="role">${uiModule.esc(roleLabel)}' in chat
+    assert "'<div class=\"role\">' + uiModule.esc(roleLabel)" in chat
+    assert '<div class="role">${uiModule.esc(agentModelLabel)}' in chat
+    assert '<div class="role">${roleLabel}' not in chat
+    assert "'<div class=\"role\">' + roleLabel" not in chat
+    assert '<div class="role">${agentModelLabel}' not in chat
+
+
+def test_compare_search_result_links_are_http_only():
+    compare = (_REPO / "static" / "js" / "compare" / "stream.js").read_text(encoding="utf-8")
+
+    assert "function _safeHttpHref(raw)" in compare
+    assert "const safeUrl = _safeHttpHref(r.url);" in compare
+    assert "titleLink.href = safeUrl;" in compare
+    assert "titleLink.href = r.url || '#';" not in compare
+
+
+def test_compare_probe_provider_labels_are_escaped():
+    selector = (_REPO / "static" / "js" / "compare" / "selector.js").read_text(encoding="utf-8")
+
+    assert "${escapeHtml(p.label || p.id)}" in selector
+    assert "${p.label || p.id}" not in selector