Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion packages/x-markdown/src/XMarkdown/__test__/Renderer.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import DOMPurify from 'dompurify';
import React from 'react';
import Renderer from '../core/Renderer';
import { getDOMPurify } from '../utils';

// Mock React components for testing
const MockComponent: React.FC<any> = (props) => {
Expand Down Expand Up @@ -776,6 +776,7 @@ describe('Renderer', () => {
const renderer = new Renderer({ components });

// Spy on DOMPurify.sanitize
const DOMPurify = getDOMPurify();
const sanitizeSpy = jest.spyOn(DOMPurify, 'sanitize');

const html = '<custom-tag>content</custom-tag><script>alert("xss")</script>';
Expand Down Expand Up @@ -807,6 +808,7 @@ describe('Renderer', () => {
});

// Spy on DOMPurify.sanitize
const DOMPurify = getDOMPurify();
const sanitizeSpy = jest.spyOn(DOMPurify, 'sanitize');

const html = '<custom-tag class="test" id="test-id">content</custom-tag>';
Expand Down
4 changes: 3 additions & 1 deletion packages/x-markdown/src/XMarkdown/core/Renderer.ts
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
import type { Config as DOMPurifyConfig } from 'dompurify';
import DOMPurify from 'dompurify';
import type { DOMNode, Element } from 'html-react-parser';
import parseHtml, { domToReact } from 'html-react-parser';
import React, { ReactNode } from 'react';
import AnimationText from '../AnimationText';
import type { ComponentProps, XMarkdownProps } from '../interface';
import { getDOMPurify } from '../utils';

interface RendererOptions {
components?: XMarkdownProps['components'];
Expand Down Expand Up @@ -143,6 +143,8 @@ class Renderer {

// Use DOMPurify to clean HTML while preserving custom components and target attributes
const purifyConfig = this.configureDOMPurify();

const DOMPurify = getDOMPurify();
const cleanHtml = DOMPurify.sanitize(htmlString, purifyConfig);

return parseHtml(cleanHtml, {
Expand Down
181 changes: 116 additions & 65 deletions packages/x-markdown/src/XMarkdown/hooks/useStreaming.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import { useCallback, useEffect, useRef, useState } from 'react';
import { StreamCacheTokenType, XMarkdownProps } from '../interface';
import { StreamCacheTokenType, StreamingOption, XMarkdownProps } from '../interface';

/* ------------ Type ------------ */

Expand Down Expand Up @@ -196,47 +196,124 @@ const safeEncodeURIComponent = (str: string): string => {
}
};

/* ------------ Pure Processing Functions ------------ */
const getIncompleteMarkdownPlaceholder = (
cache: StreamCache,
incompleteMarkdownComponentMap?: StreamingOption['incompleteMarkdownComponentMap'],
components?: XMarkdownProps['components'],
): string | undefined => {
const { token, pending } = cache;
if (token === StreamCacheTokenType.Text) return;
/**
* An image tag starts with '!', if it's the only character, it's incomplete and should be stripped.
* !
* ^
*/
if (token === StreamCacheTokenType.Image && pending === '!') return undefined;

/**
* If a table has more than two lines (header, separator, and at least one row),
* it's considered complete enough to not be replaced by a placeholder.
* | column1 | column2 |\n| -- | --|\n
* ^
*/
if (token === StreamCacheTokenType.Table && pending.split('\n').length > 2) {
return pending;
}

const componentMap = incompleteMarkdownComponentMap || {};
const componentName = componentMap[token] || `incomplete-${token}`;
const encodedPending = safeEncodeURIComponent(pending);

return components?.[componentName]
? `<${componentName} data-raw="${encodedPending}" />`
: undefined;
};

const computeStreamingOutput = (
text: string,
cache: StreamCache,
incompleteMarkdownComponentMap?: StreamingOption['incompleteMarkdownComponentMap'],
components?: XMarkdownProps['components'],
): string => {
if (!text) {
return '';
}

const expectedPrefix = cache.completeMarkdown + cache.pending;
// Reset cache if input doesn't continue from previous state
if (!text.startsWith(expectedPrefix)) {
Object.assign(cache, getInitialCache());
}

const chunk = text.slice(cache.processedLength);
if (!chunk) {
const incompletePlaceholder = getIncompleteMarkdownPlaceholder(
cache,
incompleteMarkdownComponentMap,
components,
);
return cache.completeMarkdown + (incompletePlaceholder || '');
}

cache.processedLength += chunk.length;
const isTextInBlock = isInCodeBlock(text);
for (const char of chunk) {
cache.pending += char;
// Skip processing if inside code block
if (isTextInBlock) {
commitCache(cache);
continue;
}

if (cache.token === StreamCacheTokenType.Text) {
for (const handler of recognizeHandlers) handler.recognize(cache);
} else {
const handler = recognizeHandlers.find((handler) => handler.tokenType === cache.token);
handler?.recognize(cache);
}

if (cache.token === StreamCacheTokenType.Text) {
commitCache(cache);
}
}

const incompletePlaceholder = getIncompleteMarkdownPlaceholder(
cache,
incompleteMarkdownComponentMap,
components,
);
return cache.completeMarkdown + (incompletePlaceholder || '');
};

/* ------------ Main Hook ------------ */
const useStreaming = (
input: string,
config?: { streaming: XMarkdownProps['streaming']; components?: XMarkdownProps['components'] },
) => {
const { streaming, components = {} } = config || {};
const { hasNextChunk: enableCache = false, incompleteMarkdownComponentMap } = streaming || {};
const [output, setOutput] = useState('');

const cacheRef = useRef<StreamCache>(getInitialCache());

const handleIncompleteMarkdown = useCallback(
(cache: StreamCache): string | undefined => {
const { token, pending } = cache;
if (token === StreamCacheTokenType.Text) return;
/**
* An image tag starts with '!', if it's the only character, it's incomplete and should be stripped.
* !
* ^
*/
if (token === StreamCacheTokenType.Image && pending === '!') return undefined;

/**
* If a table has more than two lines (header, separator, and at least one row),
* it's considered complete enough to not be replaced by a placeholder.
* | column1 | column2 |\n| -- | --|\n
* ^
*/
if (token === StreamCacheTokenType.Table && pending.split('\n').length > 2) {
return pending;
}
// Use lazy initializer to compute initial output synchronously on first render
const [output, setOutput] = useState(() => {
if (typeof input !== 'string') {
return '';
}

const componentMap = incompleteMarkdownComponentMap || {};
const componentName = componentMap[token] || `incomplete-${token}`;
const encodedPending = safeEncodeURIComponent(pending);
if (!enableCache) {
return input;
}

return components?.[componentName]
? `<${componentName} data-raw="${encodedPending}" />`
: undefined;
},
[incompleteMarkdownComponentMap, components],
);
// For streaming mode, compute initial output synchronously to avoid empty content on mount
return computeStreamingOutput(
input,
cacheRef.current,
incompleteMarkdownComponentMap,
components,
);
});

const processStreaming = useCallback(
(text: string): void => {
Expand All @@ -246,42 +323,16 @@ const useStreaming = (
return;
}

const expectedPrefix = cacheRef.current.completeMarkdown + cacheRef.current.pending;
// Reset cache if input doesn't continue from previous state
if (!text.startsWith(expectedPrefix)) {
cacheRef.current = getInitialCache();
}

const cache = cacheRef.current;
const chunk = text.slice(cache.processedLength);
if (!chunk) return;

cache.processedLength += chunk.length;
const isTextInBlock = isInCodeBlock(text);
for (const char of chunk) {
cache.pending += char;
// Skip processing if inside code block
if (isTextInBlock) {
commitCache(cache);
continue;
}

if (cache.token === StreamCacheTokenType.Text) {
for (const handler of recognizeHandlers) handler.recognize(cache);
} else {
const handler = recognizeHandlers.find((handler) => handler.tokenType === cache.token);
handler?.recognize(cache);
}

if (cache.token === StreamCacheTokenType.Text) {
commitCache(cache);
}
}
const result = computeStreamingOutput(
text,
cacheRef.current,
incompleteMarkdownComponentMap,
components,
);

const incompletePlaceholder = handleIncompleteMarkdown(cache);
setOutput(cache.completeMarkdown + (incompletePlaceholder || ''));
setOutput(result);
},
[handleIncompleteMarkdown],
[incompleteMarkdownComponentMap, components],
);

useEffect(() => {
Expand Down
18 changes: 18 additions & 0 deletions packages/x-markdown/src/XMarkdown/utils/getDOMPurify.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
import DOMPurify from 'dompurify';
import { JSDOM } from 'jsdom';

/**
* Returns the DOMPurify instance, compatible with both server-side (Node.js) and client-side (browser) environments.
*
* On the server, it creates a DOMPurify instance with a jsdom window; on the client, it returns the browser's DOMPurify.
*
* @see https://github.com/cure53/DOMPurify?tab=readme-ov-file#running-dompurify-on-the-server
*/
export function getDOMPurify() {
if (typeof window === 'undefined') {
const jsWindow = new JSDOM('').window;
return DOMPurify(jsWindow);
}

return DOMPurify;
}
Comment on lines +1 to +43
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Creating a new JSDOM instance on every call to getDOMPurify on the server can be inefficient. It's better to create the server-side DOMPurify instance once and cache it for subsequent calls. This will improve performance in an SSR environment by avoiding the overhead of creating a new JSDOM window repeatedly.

import DOMPurify from 'dompurify';
import { JSDOM } from 'jsdom';

let serverDOMPurify: ReturnType<typeof DOMPurify>;

/**
 * Returns the DOMPurify instance, compatible with both server-side (Node.js) and client-side (browser) environments.
 *
 * On the server, it creates a DOMPurify instance with a jsdom window; on the client, it returns the browser's DOMPurify.
 *
 * @see https://github.com/cure53/DOMPurify?tab=readme-ov-file#running-dompurify-on-the-server
 */
export function getDOMPurify() {
  if (typeof window === 'undefined') {
    if (!serverDOMPurify) {
      const jsWindow = new JSDOM('').window;
      serverDOMPurify = DOMPurify(jsWindow as any);
    }
    return serverDOMPurify;
  }

  return DOMPurify;
}

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will this cause a memory leak for the serverDOMPurify variable?

Comment on lines +34 to +43
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

服务端初始化失败时的回退逻辑存在问题。

第 39 行的回退逻辑 return serverDOMPurify || DOMPurify 在服务端环境下可能导致运行时错误:

  • 如果 initializeServerDOMPurify() 失败(例如 jsdom 加载失败),serverDOMPurifynull
  • 此时直接返回 DOMPurify 在服务端无法工作,因为 DOMPurify 需要一个 window 对象
  • 用户会收到关于缺少 window 的隐晦错误信息

建议改进回退策略:

 export function getDOMPurify(): ReturnType<typeof DOMPurify> {
   if (typeof window === 'undefined') {
     if (!serverDOMPurify) {
       initializeServerDOMPurify();
     }
-    return serverDOMPurify || DOMPurify;
+    if (!serverDOMPurify) {
+      throw new Error(
+        'Failed to initialize DOMPurify for SSR. Please ensure jsdom is installed.'
+      );
+    }
+    return serverDOMPurify;
   }
 
   return DOMPurify;
 }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
export function getDOMPurify(): ReturnType<typeof DOMPurify> {
if (typeof window === 'undefined') {
if (!serverDOMPurify) {
initializeServerDOMPurify();
}
return serverDOMPurify || DOMPurify;
}
return DOMPurify;
}
export function getDOMPurify(): ReturnType<typeof DOMPurify> {
if (typeof window === 'undefined') {
if (!serverDOMPurify) {
initializeServerDOMPurify();
}
if (!serverDOMPurify) {
throw new Error(
'Failed to initialize DOMPurify for SSR. Please ensure jsdom is installed.'
);
}
return serverDOMPurify;
}
return DOMPurify;
}

1 change: 1 addition & 0 deletions packages/x-markdown/src/XMarkdown/utils/index.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
export { getDOMPurify } from './getDOMPurify';