fix: address markdown image review feedback

This commit is contained in:
Ayaan Zaidi
2026-03-07 19:26:29 +05:30
committed by Ayaan Zaidi
parent 4bf902de58
commit d25b493c7f
4 changed files with 52 additions and 7 deletions

View File

@@ -40,24 +40,22 @@ enum ChatMarkdownPreprocessor {
if matches.isEmpty { return Result(cleaned: self.normalize(withoutTimestamps), images: []) }
var images: [InlineImage] = []
var cleaned = withoutTimestamps
let cleaned = NSMutableString(string: withoutTimestamps)
for match in matches.reversed() {
guard match.numberOfRanges >= 3 else { continue }
let label = ns.substring(with: match.range(at: 1))
let source = ns.substring(with: match.range(at: 2))
let start = cleaned.index(cleaned.startIndex, offsetBy: match.range.location)
let end = cleaned.index(start, offsetBy: match.range.length)
if let inlineImage = self.inlineImage(label: label, source: source) {
images.append(inlineImage)
cleaned.replaceSubrange(start..<end, with: "")
cleaned.replaceCharacters(in: match.range, with: "")
} else {
cleaned.replaceSubrange(start..<end, with: self.fallbackImageLabel(label))
cleaned.replaceCharacters(in: match.range, with: self.fallbackImageLabel(label))
}
}
return Result(cleaned: self.normalize(cleaned), images: images.reversed())
return Result(cleaned: self.normalize(cleaned as String), images: images.reversed())
}
private static func inlineImage(label: String, source: String) -> InlineImage? {

View File

@@ -42,6 +42,15 @@ struct ChatMarkdownPreprocessorTests {
#expect(result.images.isEmpty)
}
@Test func handlesUnicodeBeforeRemoteMarkdownImages() {
let markdown = "🙂![Leak](https://example.com/image.png)"
let result = ChatMarkdownPreprocessor.preprocess(markdown: markdown)
#expect(result.cleaned == "🙂Leak")
#expect(result.images.isEmpty)
}
@Test func stripsInboundUntrustedContextBlocks() {
let markdown = """
Conversation info (untrusted metadata):

View File

@@ -665,6 +665,10 @@
return div.innerHTML;
}
function escapeHtmlAttr(text) {
return escapeHtml(text).replaceAll('"', "&quot;").replaceAll("'", "&#39;");
}
// Validate image fields before interpolating data URLs.
const SAFE_IMAGE_MIME_RE = /^image\/(png|jpeg|gif|webp|svg\+xml|bmp|tiff|avif)$/i;
const SAFE_BASE64_RE = /^[A-Za-z0-9+/]+={0,2}$/;
@@ -1725,7 +1729,7 @@
if (!INLINE_DATA_IMAGE_RE.test(href)) {
return escapeHtml(label);
}
return `<img src="${escapeHtml(href)}" alt="${escapeHtml(label)}">`;
return `<img src="${escapeHtmlAttr(href)}" alt="${escapeHtmlAttr(label)}">`;
}
// Configure marked with syntax highlighting and HTML escaping for text

View File

@@ -284,4 +284,38 @@ describe("export html security hardening", () => {
expect(messages?.textContent).toContain("exfil");
expect(messages?.querySelector(`img[src="${dataImage}"]`)).toBeTruthy();
});
it("escapes markdown data-image attributes", () => {
const dataImage = "data:image/png;base64,AAAA";
const session: SessionData = {
header: { id: "session-5", timestamp: now() },
entries: [
{
id: "1",
parentId: null,
timestamp: now(),
type: "message",
message: {
role: "assistant",
content: [
{
type: "text",
text: `![x" onerror="alert(1)](${dataImage})`,
},
],
},
},
],
leafId: "1",
systemPrompt: "",
tools: [],
};
const { document } = renderTemplate(session);
const img = document.querySelector("#messages img");
expect(img).toBeTruthy();
expect(img?.getAttribute("onerror")).toBeNull();
expect(img?.getAttribute("alt")).toBe('x" onerror="alert(1)');
expect(img?.getAttribute("src")).toBe(dataImage);
});
});