From b94140bc2676a5640c4f20acf0e0ebbb5638dd40 Mon Sep 17 00:00:00 2001
From: Doekin <105162544+Doekin@users.noreply.github.com>
Date: Fri, 6 Jun 2025 00:29:47 +0800
Subject: [PATCH] feat: enable rendering and download of inline base64-encoded
 images (#6669)

This commit introduces support for displaying and downloading
inline base64-encoded images (specifically PNG and JPEG formats)
within Markdown content.

Key changes:
- Modified 'urlTransform' in the Markdown component to allow 'data:image/png'
  and 'data:image/jpeg' URLs, enabling their rendering.
- Updated the 'download' utility to handle 'data:' URLs,
  allowing users to save these inline images.

Signed-off-by: Chan Lee <Leetimemp@gmail.com>
---
 .../src/pages/home/Markdown/Markdown.tsx      |  8 +++-
 src/renderer/src/utils/download.ts            | 37 ++++++++++++-------
 2 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/src/renderer/src/pages/home/Markdown/Markdown.tsx b/src/renderer/src/pages/home/Markdown/Markdown.tsx
index 8e2d64177a..9c9884abac 100644
--- a/src/renderer/src/pages/home/Markdown/Markdown.tsx
+++ b/src/renderer/src/pages/home/Markdown/Markdown.tsx
@@ -12,7 +12,7 @@ import { findCitationInChildren, getCodeBlockId } from '@renderer/utils/markdown
 import { isEmpty } from 'lodash'
 import { type FC, memo, useCallback, useMemo } from 'react'
 import { useTranslation } from 'react-i18next'
-import ReactMarkdown, { type Components } from 'react-markdown'
+import ReactMarkdown, { type Components, defaultUrlTransform } from 'react-markdown'
 import rehypeKatex from 'rehype-katex'
 // @ts-ignore rehype-mathjax is not typed
 import rehypeMathjax from 'rehype-mathjax'
@@ -88,6 +88,11 @@ const Markdown: FC<Props> = ({ block }) => {
     } as Partial<Components>
   }, [onSaveCodeBlock])
 
+  const urlTransform = useCallback((value: string) => {
+    if (value.startsWith('data:image/png') || value.startsWith('data:image/jpeg')) return value
+    return defaultUrlTransform(value)
+  }, [])
+
   // if (role === 'user' && !renderInputMessageAsMarkdown) {
   //   return <p style={{ marginBottom: 5, whiteSpace: 'pre-wrap' }}>{messageContent}</p>
   // }
@@ -103,6 +108,7 @@ const Markdown: FC<Props> = ({ block }) => {
       className="markdown"
       components={components}
       disallowedElements={DISALLOWED_ELEMENTS}
+      urlTransform={urlTransform}
       remarkRehypeOptions={{
         footnoteLabel: t('common.footnotes'),
         footnoteLabelTagName: 'h4',
diff --git a/src/renderer/src/utils/download.ts b/src/renderer/src/utils/download.ts
index cbbbf22e51..5e207eff67 100644
--- a/src/renderer/src/utils/download.ts
+++ b/src/renderer/src/utils/download.ts
@@ -1,20 +1,31 @@
 export const download = (url: string, filename?: string) => {
-  // 处理 file:// 协议
-  if (url.startsWith('file://')) {
+  // 处理可直接通过 <a> 标签下载的 URL:
+  // - 本地文件 ( file:// )
+  // - 对象 URL ( blob: )
+  // - 相对安全的内联数据 ( data:image/png, data:image/jpeg )
+  //   (注: 其他 data 类型，如 data:text/html 或 data:image/svg+xml，
+  //    因其潜在安全风险，不在此处理，将由后续 fetch 逻辑处理或被 CSP 阻止。)
+  const SUPPORTED_PREFIXES = ['file://', 'blob:', 'data:image/png', 'data:image/jpeg']
+  if (SUPPORTED_PREFIXES.some((prefix) => url.startsWith(prefix))) {
     const link = document.createElement('a')
     link.href = url
-    link.download = filename || url.split('/').pop() || 'download'
-    document.body.appendChild(link)
-    link.click()
-    link.remove()
-    return
-  }
 
-  // 处理 Blob URL
-  if (url.startsWith('blob:')) {
-    const link = document.createElement('a')
-    link.href = url
-    link.download = filename || `${Date.now()}_diagram.svg`
+    let resolvedFilename = filename
+    if (!resolvedFilename) {
+      if (url.startsWith('file://')) {
+        const pathname = new URL(url).pathname
+        resolvedFilename = decodeURIComponent(pathname.substring(pathname.lastIndexOf('/') + 1))
+      } else if (url.startsWith('blob:')) {
+        resolvedFilename = `${Date.now()}_diagram.svg`
+      } else if (url.startsWith('data:')) {
+        const mimeMatch = url.match(/^data:([^;,]+)[;,]/)
+        const mimeType = mimeMatch && mimeMatch[1]
+        const extension = getExtensionFromMimeType(mimeType)
+        resolvedFilename = `${Date.now()}_download${extension}`
+      } else resolvedFilename = 'download'
+    }
+    link.download = resolvedFilename
+
     document.body.appendChild(link)
     link.click()
     link.remove()