refactor: use rehype-sanitize for html tags

package.json

@@ -183,6 +183,7 @@
     "rehype-katex": "^7.0.1",
     "rehype-mathjax": "^7.0.0",
     "rehype-raw": "^7.0.0",
+    "rehype-sanitize": "^6.0.0",
     "remark-cjk-friendly": "^1.1.0",
     "remark-gfm": "^4.0.0",
     "remark-math": "^6.0.0",

src/renderer/src/pages/home/Markdown/Markdown.tsx

@@ -8,6 +8,7 @@ import type { Message } from '@renderer/types'
 import { parseJSON } from '@renderer/utils'
 import { escapeBrackets, removeSvgEmptyLines, withGeminiGrounding } from '@renderer/utils/formats'
 import { findCitationInChildren } from '@renderer/utils/markdown'
+import { sanitizeSchema } from '@renderer/utils/markdown'
 import { isEmpty } from 'lodash'
 import { type FC, useMemo } from 'react'
 import { useTranslation } from 'react-i18next'
@@ -16,6 +17,7 @@ import rehypeKatex from 'rehype-katex'
 // @ts-ignore next-line
 import rehypeMathjax from 'rehype-mathjax'
 import rehypeRaw from 'rehype-raw'
+import rehypeSanitize from 'rehype-sanitize'
 import remarkCjkFriendly from 'remark-cjk-friendly'
 import remarkGfm from 'remark-gfm'
 import remarkMath from 'remark-math'
@@ -24,15 +26,12 @@ import CodeBlock from './CodeBlock'
 import ImagePreview from './ImagePreview'
 import Link from './Link'
 
-const ALLOWED_ELEMENTS =
-  /<(style|p|div|span|b|i|strong|em|ul|ol|li|table|tr|td|th|thead|tbody|h[1-6]|blockquote|pre|code|br|hr|svg|path|circle|rect|line|polyline|polygon|text|g|defs|title|desc|tspan|sub|sup)/i
-
 interface Props {
   message: Message
 }
 
 const remarkPlugins = [remarkMath, remarkGfm, remarkCjkFriendly]
-const disallowedElements = ['iframe']
+
 const Markdown: FC<Props> = ({ message }) => {
   const { t } = useTranslation()
   const { renderInputMessageAsMarkdown, mathEngine } = useSettings()
@@ -47,9 +46,8 @@ const Markdown: FC<Props> = ({ message }) => {
   }, [message, t])
 
   const rehypePlugins = useMemo(() => {
-    const hasElements = ALLOWED_ELEMENTS.test(messageContent)
-    return hasElements ? [rehypeRaw, rehypeMath] : [rehypeMath]
-  }, [messageContent, rehypeMath])
+    return [rehypeRaw, [rehypeSanitize, sanitizeSchema], rehypeMath]
+  }, [rehypeMath])
 
   const components = useMemo(() => {
     const baseComponents = {
@@ -75,7 +73,6 @@ const Markdown: FC<Props> = ({ message }) => {
       remarkPlugins={remarkPlugins}
       className="markdown"
       components={components}
-      disallowedElements={disallowedElements}
       remarkRehypeOptions={{
         footnoteLabel: t('common.footnotes'),
         footnoteLabelTagName: 'h4',

src/renderer/src/utils/markdown.ts

@@ -17,3 +17,67 @@ export const findCitationInChildren = (children) => {
 
   return null
 }
+
+export const MARKDOWN_ALLOWED_TAGS = [
+  'style',
+  'p',
+  'div',
+  'span',
+  'b',
+  'i',
+  'strong',
+  'em',
+  'ul',
+  'ol',
+  'li',
+  'table',
+  'tr',
+  'td',
+  'th',
+  'thead',
+  'tbody',
+  'h1',
+  'h2',
+  'h3',
+  'h4',
+  'h5',
+  'h6',
+  'blockquote',
+  'pre',
+  'code',
+  'br',
+  'hr',
+  'svg',
+  'path',
+  'circle',
+  'rect',
+  'line',
+  'polyline',
+  'polygon',
+  'text',
+  'g',
+  'defs',
+  'title',
+  'desc',
+  'tspan',
+  'sub',
+  'sup'
+]
+
+// rehype-sanitize配置
+export const sanitizeSchema = {
+  tagNames: MARKDOWN_ALLOWED_TAGS,
+  attributes: {
+    '*': ['className', 'style', 'id', 'title'],
+    svg: ['viewBox', 'width', 'height', 'xmlns', 'fill', 'stroke'],
+    path: ['d', 'fill', 'stroke', 'strokeWidth', 'strokeLinecap', 'strokeLinejoin'],
+    circle: ['cx', 'cy', 'r', 'fill', 'stroke'],
+    rect: ['x', 'y', 'width', 'height', 'fill', 'stroke'],
+    line: ['x1', 'y1', 'x2', 'y2', 'stroke'],
+    polyline: ['points', 'fill', 'stroke'],
+    polygon: ['points', 'fill', 'stroke'],
+    text: ['x', 'y', 'fill', 'textAnchor', 'dominantBaseline'],
+    g: ['transform', 'fill', 'stroke'],
+    a: ['href', 'target', 'rel']
+  }
+}

yarn.lock

@@ -4265,6 +4265,7 @@ __metadata:
     rehype-katex: "npm:^7.0.1"
     rehype-mathjax: "npm:^7.0.0"
     rehype-raw: "npm:^7.0.0"
+    rehype-sanitize: "npm:^6.0.0"
     remark-cjk-friendly: "npm:^1.1.0"
     remark-gfm: "npm:^4.0.0"
     remark-math: "npm:^6.0.0"
@@ -8932,6 +8933,17 @@ __metadata:
   languageName: node
   linkType: hard
 
+"hast-util-sanitize@npm:^5.0.0":
+  version: 5.0.2
+  resolution: "hast-util-sanitize@npm:5.0.2"
+  dependencies:
+    "@types/hast": "npm:^3.0.0"
+    "@ungap/structured-clone": "npm:^1.0.0"
+    unist-util-position: "npm:^5.0.0"
+  checksum: 10c0/20951652078a8c21341c1c9a84f90015b2ba01cc41fa16772f122c65cda26a7adb0501fdeba5c8e37e40e2632447e8fe455d0dd2dc27d39663baacca76f2ecb6
+  languageName: node
+  linkType: hard
+
 "hast-util-to-html@npm:^9.0.5":
   version: 9.0.5
   resolution: "hast-util-to-html@npm:9.0.5"
@@ -14481,6 +14493,16 @@ __metadata:
   languageName: node
   linkType: hard
 
+"rehype-sanitize@npm:^6.0.0":
+  version: 6.0.0
+  resolution: "rehype-sanitize@npm:6.0.0"
+  dependencies:
+    "@types/hast": "npm:^3.0.0"
+    hast-util-sanitize: "npm:^5.0.0"
+  checksum: 10c0/43d6c056e63c994cf56e5ee0e157052d2030dc5ac160845ee494af9a26e5906bf5ec5af56c7d90c99f9c4dc0091e45a48a168618135fb6c64a76481ad3c449e9
+  languageName: node
+  linkType: hard
+
 "remark-cjk-friendly@npm:^1.1.0":
   version: 1.1.0
   resolution: "remark-cjk-friendly@npm:1.1.0"