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
85 changes: 70 additions & 15 deletions libs/a2a-core/src/react/components/Message/Message.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,61 @@ marked.use({
},
});

interface FencedCodeBlock {
index: number;
endIndex: number;
language: string;
code: string;
}

// Extract fenced code blocks (```lang\n...\n```) using a single linear scan.
// This intentionally avoids a lazy `([\s\S]*?)` regular expression, which runs
// in polynomial time (ReDoS) on untrusted content containing many unterminated
// code fences. `String.indexOf` keeps the whole pass linear in the input length.
// Exported for unit testing of the ReDoS-safe extraction.
export function extractFencedCodeBlocks(content: string): FencedCodeBlock[] {
const blocks: FencedCodeBlock[] = [];
const fence = '```';
let searchFrom = 0;

while (searchFrom < content.length) {
const fenceStart = content.indexOf(fence, searchFrom);
if (fenceStart === -1) {
break;
}

// The opening fence's info string (language) runs up to the next newline.
const infoEnd = content.indexOf('\n', fenceStart + fence.length);
if (infoEnd === -1) {
break;
}

// A valid opening fence only carries word characters as its language.
const language = content.slice(fenceStart + fence.length, infoEnd);
if (!/^\w*$/.test(language)) {
searchFrom = fenceStart + 1;
continue;
}

const codeStart = infoEnd + 1;
const closingFence = content.indexOf(`\n${fence}`, codeStart);
if (closingFence === -1) {
break;
}

const endIndex = closingFence + 1 + fence.length;
blocks.push({
index: fenceStart,
endIndex,
language,
code: content.slice(codeStart, closingFence),
});
searchFrom = endIndex;
}

return blocks;
}

const useStyles = makeStyles({
'@keyframes fadeIn': {
'0%': { opacity: 0, transform: 'translateY(10px)' },
Expand Down Expand Up @@ -446,9 +501,9 @@ function MessageComponent({
let contentToDownload = message.metadata?.rawContent;

if (!contentToDownload) {
const codeBlockMatch = message.content.match(/```[\w]*\n([\s\S]*?)\n```/);
if (codeBlockMatch) {
contentToDownload = codeBlockMatch[1];
const codeBlocks = extractFencedCodeBlocks(message.content);
if (codeBlocks.length > 0) {
contentToDownload = codeBlocks[0].code;
} else {
contentToDownload = message.content.replace(/^\*\*.*?\*\*\n\n/, '');
}
Expand Down Expand Up @@ -527,24 +582,24 @@ function MessageComponent({
// For regular markdown content, we need to parse and render code blocks with headers
const processedContent = useMemo(() => {
// Extract code blocks and render them separately
const codeBlockRegex = /```(\w*)\n([\s\S]*?)\n```/g;
const content = message.content;
const codeBlocks = extractFencedCodeBlocks(content);
let lastIndex = 0;
const elements: React.ReactNode[] = [];
let match;

while ((match = codeBlockRegex.exec(message.content)) !== null) {
for (const block of codeBlocks) {
// Add content before the code block
if (match.index > lastIndex) {
const textContent = message.content.slice(lastIndex, match.index);
if (block.index > lastIndex) {
const textContent = content.slice(lastIndex, block.index);
const html = sanitizeHtml(marked.parse(textContent, { gfm: true, breaks: true }) as string);
elements.push(
<div key={`text-${lastIndex}`} dangerouslySetInnerHTML={{ __html: html }} />
);
}

// Add the code block with header
const language = match[1] || '';
const code = match[2];
const language = block.language;
const code = block.code;
let highlighted = code;

if (language && Prism.languages[language]) {
Expand All @@ -556,7 +611,7 @@ function MessageComponent({
}

elements.push(
<div key={`code-${match.index}`} className={styles.codeBlockWrapper}>
<div key={`code-${block.index}`} className={styles.codeBlockWrapper}>
<CodeBlockHeader language={language} code={code} />
<div className={styles.codeBlockContent}>
<pre>
Expand All @@ -569,19 +624,19 @@ function MessageComponent({
</div>
);

lastIndex = match.index + match[0].length;
lastIndex = block.endIndex;
}

// Add any remaining content after the last code block
if (lastIndex < message.content.length) {
const remainingContent = message.content.slice(lastIndex);
if (lastIndex < content.length) {
const remainingContent = content.slice(lastIndex);
const html = sanitizeHtml(marked.parse(remainingContent, { gfm: true, breaks: true }) as string);
elements.push(<div key={`text-${lastIndex}`} dangerouslySetInnerHTML={{ __html: html }} />);
}

// If no code blocks were found, just return the parsed markdown
if (elements.length === 0) {
const html = sanitizeHtml(marked.parse(message.content, { gfm: true, breaks: true }) as string);
const html = sanitizeHtml(marked.parse(content, { gfm: true, breaks: true }) as string);
return <div dangerouslySetInnerHTML={{ __html: html }} />;
}

Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { render, screen, fireEvent } from '@testing-library/react';
import { vi, describe, beforeEach, it, expect } from 'vitest';
import '@testing-library/jest-dom';
import { Message } from '../Message';
import { Message, extractFencedCodeBlocks } from '../Message';
import type { Message as MessageType } from '../../../types';

// Mock clipboard API
Expand Down Expand Up @@ -138,4 +138,75 @@ greeting = "Hello"
// Check for copy button
expect(screen.getByRole('button', { name: /copy/i })).toBeInTheDocument();
});

it('should render text that follows a code block', () => {
const message: MessageType = {
id: '1',
content: '```javascript\nconst x = 1;\n```\n\nTrailing explanation text.',
sender: 'assistant',
timestamp: new Date(),
status: 'sent',
};

render(<Message message={message} />);

// The code block header is rendered...
expect(screen.getByText('javascript')).toBeInTheDocument();
// ...and the text after the closing fence is preserved.
expect(screen.getByText('Trailing explanation text.')).toBeInTheDocument();
});
});

describe('extractFencedCodeBlocks', () => {
it('extracts a single code block with language', () => {
const blocks = extractFencedCodeBlocks('```javascript\nconst x = 1;\n```');
expect(blocks).toHaveLength(1);
expect(blocks[0].language).toBe('javascript');
expect(blocks[0].code).toBe('const x = 1;');
expect(blocks[0].index).toBe(0);
});

it('extracts a code block without a language', () => {
const blocks = extractFencedCodeBlocks('```\nplain code\n```');
expect(blocks).toHaveLength(1);
expect(blocks[0].language).toBe('');
expect(blocks[0].code).toBe('plain code');
});

it('extracts multiple code blocks with the text between them intact', () => {
const content = 'intro\n```js\na();\n```\nmiddle\n```py\nb()\n```';
const blocks = extractFencedCodeBlocks(content);
expect(blocks.map((b) => b.language)).toEqual(['js', 'py']);
expect(blocks.map((b) => b.code)).toEqual(['a();', 'b()']);
// The gap between the two blocks maps back to the original 'middle' text.
expect(content.slice(blocks[0].endIndex, blocks[1].index)).toBe('\nmiddle\n');
});

it('ignores an unterminated fence', () => {
expect(extractFencedCodeBlocks('```js\nno closing fence here')).toEqual([]);
});

it('does not treat a fence with a non-word info string as an opening', () => {
// The original regex required the language to be word characters only.
expect(extractFencedCodeBlocks('```not a lang\ncode\n```')).toEqual([]);
});

it('stops the first code block at the first closing fence (lazy match parity)', () => {
const blocks = extractFencedCodeBlocks('```\nfirst\n```\nsecond\n```');
expect(blocks[0].code).toBe('first');
});

it('handles pathological unterminated fences in linear time (ReDoS guard)', () => {
// The exact shape flagged by CodeQL js/polynomial-redos: a string starting
// with '```\n' followed by many repetitions of '```\na' and no closing
// fence. A polynomial-time regex scans to the end from every fence start
// (O(n^2)); the linear scanner returns almost instantly.
const content = `\`\`\`\n${'```\na'.repeat(200000)}`;
const start = performance.now();
const blocks = extractFencedCodeBlocks(content);
const elapsed = performance.now() - start;

expect(blocks).toEqual([]);
expect(elapsed).toBeLessThan(1000);

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Good catch — fixed in f5bb10e. I removed the absolute wall-clock assertion (elapsed < 1000ms) and replaced it with a machine-speed-independent scaling-ratio check.

The guard now times a 10x-larger pathological input and asserts largeTime / smallTime stays below a generous 30x ceiling. A linear scan grows ~10x with 10x input; the old O(n²) regex would grow ~100x. Asserting the ratio rather than absolute time keeps the test stable on slow/contended CI runners while still catching polynomial blow-up. It also warms up the JIT and averages over 20 iterations to reduce noise.

Ran it 3x locally — all green (14/14 each run).

});
});
Loading