diff --git a/extension/jest.config.cjs b/extension/jest.config.cjs index 888f955..4ac9a1d 100644 --- a/extension/jest.config.cjs +++ b/extension/jest.config.cjs @@ -29,12 +29,16 @@ module.exports = { ], }, moduleFileExtensions: ["ts", "tsx", "js", "jsx", "json"], - // `webext-storage` is published as ESM-only and our ts-jest config emits - // CommonJS, so importing it directly trips a `SyntaxError: Unexpected token - // 'export'`. None of our tests exercise the real storage flow (they mock the - // storage modules wholesale), so route the import to a small CJS stub. + // `webext-storage` and `abort-utils` are published as ESM-only and our + // ts-jest config emits CommonJS, so importing them directly trips a + // `SyntaxError: Unexpected token 'export'`. None of our tests exercise + // the real storage flow (they mock the storage modules wholesale), and + // the abort-utils stub is a faithful enough subset for the lifecycle + // semantics the rules depend on (signal flip on abortAndReset, onAbort + // dispatch). moduleNameMapper: { "^webext-storage$": "/src/__test-mocks__/webext-storage.ts", + "^abort-utils$": "/src/__test-mocks__/abort-utils.ts", }, // - jest-webextension-mock: installs globalThis.chrome + browser with // jest.fn() stubs for MV2/MV3 APIs the extension uses (runtime, storage, diff --git a/extension/src/__test-mocks__/abort-utils.ts b/extension/src/__test-mocks__/abort-utils.ts new file mode 100644 index 0000000..07be7b2 --- /dev/null +++ b/extension/src/__test-mocks__/abort-utils.ts @@ -0,0 +1,62 @@ +// Copyright (c) 2026 PixieBrix, Inc. +// Licensed under PolyForm Shield 1.0.0 — see LICENSE. + +// In-memory stub of the subset of `abort-utils` consumed by the +// extension. The real package ships pure ESM, which ts-jest with +// `useESM: false` can't transform — so route imports here via +// `moduleNameMapper` in jest.config.cjs. The stub mirrors the runtime +// semantics tests actually depend on: ReusableAbortController issues +// fresh signals after abortAndReset, and onAbort attaches a disposable +// listener. + +export class ReusableAbortController { + private controller = new AbortController(); + + get signal(): AbortSignal { + return this.controller.signal; + } + + abort(reason?: unknown): void { + this.controller.abort(reason); + } + + abortAndReset(reason?: unknown): void { + this.controller.abort(reason); + this.controller = new AbortController(); + } +} + +type Handle = + | { disconnect(): void } + | { abort(reason: unknown): void } + | { abortAndReset(reason: unknown): void } + | ((reason: unknown) => void); + +export function onAbort( + signal: AbortController | AbortSignal | undefined, + ...handles: Handle[] +): { [Symbol.dispose](): void } | undefined { + if (!signal) { + return undefined; + } + const target = signal instanceof AbortController ? signal.signal : signal; + const listener = (): void => { + for (const handle of handles) { + if (typeof handle === "function") { + handle(target.reason); + } else if ("abortAndReset" in handle) { + handle.abortAndReset(target.reason); + } else if ("abort" in handle) { + handle.abort(target.reason); + } else { + handle.disconnect(); + } + } + }; + target.addEventListener("abort", listener, { once: true }); + return { + [Symbol.dispose](): void { + target.removeEventListener("abort", listener); + }, + }; +} diff --git a/extension/src/lib/__tests__/yielding-text-walk.property.test.ts b/extension/src/lib/__tests__/yielding-text-walk.property.test.ts new file mode 100644 index 0000000..e0f61a1 --- /dev/null +++ b/extension/src/lib/__tests__/yielding-text-walk.property.test.ts @@ -0,0 +1,209 @@ +// Copyright (c) 2026 PixieBrix, Inc. +// Licensed under PolyForm Shield 1.0.0 — see LICENSE. + +// Property test for the chunked walker's equivalence with +// `walkTextNodes`. For any tree shape, any text content, any chunkSize, +// the concatenation of all chunks plus the order they were delivered +// must equal what `walkTextNodes` would return. Catches: +// - missing the trailing partial chunk +// - dropping nodes at the chunk boundary (the "next prefetched but +// never pushed" class of bug) +// - double-processing nodes across the sync→async transition +// - drift in the filter predicates (NON_CONTENT_TAGS, minLength) +// between the two walkers + +import fc from "fast-check"; + +import { walkTextNodes } from "../dom-utils"; +import { walkTextNodesChunked } from "../yielding-text-walk"; + +interface FlatTree { + size: number; + parents: readonly number[]; + texts: readonly string[]; + tags: readonly string[]; +} + +// Tag alphabet excludes ` + + + + `; + const values: string[] = []; + walkTextNodesChunked(document.body, { + process: (chunk) => { + for (const node of chunk) { + const text = node.nodeValue?.trim(); + if (text) { + values.push(text); + } + } + }, + }); + expect(values).toContain("real"); + expect(values).not.toContain("script"); + expect(values).not.toContain("style"); + expect(values).not.toContain("noscript"); + expect(values).not.toContain("template"); + }); + + it("skips text inside an existing placeholder", () => { + document.body.innerHTML = ` +

real

+ hidden + `; + const values: string[] = []; + walkTextNodesChunked(document.body, { + process: (chunk) => { + for (const node of chunk) { + const text = node.nodeValue?.trim(); + if (text) { + values.push(text); + } + } + }, + }); + expect(values).toContain("real"); + expect(values).not.toContain("hidden"); + }); + + it("respects minLength", () => { + document.body.innerHTML = `

short

a long string of text

`; + const seen: string[] = []; + walkTextNodesChunked(document.body, { + minLength: 10, + process: (chunk) => { + for (const node of chunk) { + seen.push(node.nodeValue ?? ""); + } + }, + }); + expect(seen).toEqual(["a long string of text"]); + }); + + it("respects shouldSkipParent", () => { + document.body.innerHTML = ` +

article-text

+ + `; + const seen: string[] = []; + walkTextNodesChunked(document.body, { + shouldSkipParent: (parent) => parent.closest("aside") !== null, + process: (chunk) => { + for (const node of chunk) { + const text = node.nodeValue?.trim(); + if (text) { + seen.push(text); + } + } + }, + }); + expect(seen).toEqual(["article-text"]); + }); +}); diff --git a/extension/src/lib/yielding-text-walk.ts b/extension/src/lib/yielding-text-walk.ts new file mode 100644 index 0000000..942d329 --- /dev/null +++ b/extension/src/lib/yielding-text-walk.ts @@ -0,0 +1,182 @@ +// Copyright (c) 2026 PixieBrix, Inc. +// Licensed under PolyForm Shield 1.0.0 — see LICENSE. + +// Chunked text-node walk for rules that scan large `innerText` +// (pii-redact, secrets-redact, encoded-payload-redact, +// prompt-injection-redact). Yields between chunks so the regex loops +// don't block the event loop for hundreds of milliseconds on dense +// pages, and surfaces an `AbortSignal` so an SPA route change can +// cancel a scan started against the old tree before it finishes +// mutating the new one. +// +// API note: the walk is synchronous when the tree fits in one chunk +// (the common case for small subtrees + the default chunkSize=100). +// Trees that span multiple chunks continue across scheduler yields, +// with `onComplete` firing in a microtask after the last chunk. The +// sync-first-fast-path lets existing rule tests `apply` and then +// assert without an explicit `await` — only the multi-chunk path +// (large pages, abort tests) needs to flush. +// +// The filter (NON_CONTENT_TAGS skip, isInsidePlaceholder skip, +// minLength, optional shouldSkipParent) mirrors `walkTextNodes` in +// dom-utils — same predicates, different consumer shape. + +import { isInsidePlaceholder, isNonContentTag } from "./dom-utils"; + +export interface WalkTextNodesChunkedOptions { + // Aborts the walk at the next chunk boundary. Already-processed + // chunks stay; partially-processed chunks are not possible because + // `process` runs atomically on the full chunk. + signal?: AbortSignal; + // Minimum nodeValue length before the node is offered to `process`. + // Same semantics as `walkTextNodes`. + minLength?: number; + // Extra parent-element predicate beyond the universal NON_CONTENT_TAGS + // and placeholder skips. + shouldSkipParent?: (parent: Element) => boolean; + // Cap on the per-process batch size. Larger = fewer scheduler yields, + // longer blocking windows. The default trades off well for the + // text-walk rules — most batches fit in one chunk and never yield. + chunkSize?: number; + // Invoked once per filled chunk plus once for the trailing partial + // chunk. Synchronous: any DOM mutations should land before the next + // yield so a route-change abort doesn't leave a half-mutated batch. + process: (chunk: Text[]) => void; + // Called after the last chunk is processed (or skipped due to abort). + // Runs synchronously when the walk fit in one chunk; otherwise in a + // microtask after the final yield. Use for cross-chunk post-passes + // (e.g., prompt-injection-redact's outermost-match dedupe of + // collected container elements). + onComplete?: () => void; + // Override the yield mechanism. Production uses requestIdleCallback + // (yields to layout/paint) with a setTimeout(0) fallback; tests can + // pass a microtask-based yield to avoid driving fake timers. + yieldStrategy?: () => Promise; +} + +function defaultYieldStrategy(): Promise { + if (typeof globalThis.requestIdleCallback === "function") { + return new Promise((resolve) => { + globalThis.requestIdleCallback(() => { + resolve(); + }); + }); + } + return new Promise((resolve) => { + setTimeout(resolve, 0); + }); +} + +export function walkTextNodesChunked( + root: ParentNode, + options: WalkTextNodesChunkedOptions, +): void { + const { + signal, + minLength = 0, + shouldSkipParent, + chunkSize = 100, + process, + onComplete, + yieldStrategy = defaultYieldStrategy, + } = options; + + if (signal?.aborted) { + return; + } + + const walker = document.createTreeWalker(root, NodeFilter.SHOW_TEXT, { + acceptNode(node) { + const parent = node.parentElement; + if (!parent) { + return NodeFilter.FILTER_REJECT; + } + if (isNonContentTag(parent.tagName)) { + return NodeFilter.FILTER_REJECT; + } + if (isInsidePlaceholder(parent)) { + return NodeFilter.FILTER_REJECT; + } + if (shouldSkipParent?.(parent)) { + return NodeFilter.FILTER_REJECT; + } + const value = node.nodeValue; + if (!value || value.length < minLength) { + return NodeFilter.FILTER_REJECT; + } + return NodeFilter.FILTER_ACCEPT; + }, + }); + + let chunk: Text[] = []; + // Holds a node we pre-fetched before yielding (to keep walker.currentNode + // anchored to a still-connected node across the boundary). Consumed + // as the first node on the next walkSync entry — we can't let + // walker.nextNode() consume it because we'd want to revisit it + // ourselves, and the chunk-size invariant breaks if we seed it + // straight into `chunk` and then push another node on top. + let pendingResume: Text | null = null; + + function runFinalize(): void { + if (signal?.aborted) { + return; + } + onComplete?.(); + } + + // Walks until the chunk fills or the tree is exhausted. Returns + // "yield" when a chunk was just flushed and walking should continue + // after a scheduler yield; "done" when the walker has no more nodes + // and the trailing partial chunk (if any) has been processed. + function walkSync(): "yield" | "done" { + let next: Node | null = pendingResume ?? walker.nextNode(); + pendingResume = null; + while (next) { + chunk.push(next as Text); + if (chunk.length >= chunkSize) { + // Pre-fetch the resume point BEFORE process() runs. The + // consumer rules call replaceMatchesInTextNode in `process`, + // which detaches every text node in the chunk — including + // the one currently held in walker.currentNode. nextNode() + // off a detached node returns null and the walk silently + // ends after the first chunk. Grabbing the next node up + // front and re-anchoring walker.currentNode to it after + // process keeps the traversal valid across the boundary. + const resume = walker.nextNode(); + process(chunk); + chunk = []; + if (!resume) { + return "done"; + } + walker.currentNode = resume; + pendingResume = resume as Text; + return "yield"; + } + next = walker.nextNode(); + } + if (chunk.length > 0) { + process(chunk); + chunk = []; + } + return "done"; + } + + function continueAsync(): void { + void yieldStrategy().then(() => { + if (signal?.aborted) { + return; + } + if (walkSync() === "done") { + runFinalize(); + } else { + continueAsync(); + } + }); + } + + if (walkSync() === "done") { + runFinalize(); + } else { + continueAsync(); + } +} diff --git a/extension/src/rules/__tests__/encoded-payload-redact.test.ts b/extension/src/rules/__tests__/encoded-payload-redact.test.ts index cca7f1d..38dfb88 100644 --- a/extension/src/rules/__tests__/encoded-payload-redact.test.ts +++ b/extension/src/rules/__tests__/encoded-payload-redact.test.ts @@ -2,6 +2,7 @@ // Licensed under PolyForm Shield 1.0.0 — see LICENSE. import { PLACEHOLDER_CLASS } from "../../lib/placeholder"; +import { __resetRouteChangeForTesting } from "../../lib/route-change"; import { encodedPayloadRedactRule } from "../encoded-payload-redact"; // Benign English sentences padded to a length whose base64 encoding @@ -53,10 +54,13 @@ async function flushMutations(): Promise { beforeEach(() => { document.body.innerHTML = ""; jest.useFakeTimers(); + history.replaceState(null, "", "/initial"); + __resetRouteChangeForTesting(); }); afterEach(() => { encodedPayloadRedactRule.teardown(); + __resetRouteChangeForTesting(); jest.useRealTimers(); }); @@ -223,4 +227,45 @@ describe("encoded-payload-redact teardown", () => { expect(document.querySelector(`.${PLACEHOLDER_CLASS}`)).toBeNull(); }); + + it("teardown aborts the in-flight chunked walk", () => { + const payload = base64Encode(LONG_PROSE); + document.body.innerHTML = Array.from( + { length: 200 }, + (_, i) => `

blob-${i}: ${payload}

`, + ).join(""); + + encodedPayloadRedactRule.apply(document.body); + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + + encodedPayloadRedactRule.teardown(); + jest.advanceTimersByTime(0); + + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + }); + + it("route change aborts the in-flight chunked walk", () => { + const payload = base64Encode(LONG_PROSE); + document.body.innerHTML = Array.from( + { length: 200 }, + (_, i) => `

blob-${i}: ${payload}

`, + ).join(""); + + encodedPayloadRedactRule.apply(document.body); + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + + history.replaceState(null, "", "/new-route"); + globalThis.dispatchEvent(new Event("popstate")); + jest.advanceTimersByTime(0); + + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + }); }); diff --git a/extension/src/rules/__tests__/pii-redact.test.ts b/extension/src/rules/__tests__/pii-redact.test.ts index 1c3d9e8..4b3a15b 100644 --- a/extension/src/rules/__tests__/pii-redact.test.ts +++ b/extension/src/rules/__tests__/pii-redact.test.ts @@ -1,4 +1,5 @@ import { PLACEHOLDER_CLASS } from "../../lib/placeholder"; +import { __resetRouteChangeForTesting } from "../../lib/route-change"; import { piiRedactRule } from "../pii-redact"; const VALID_CARD = "4111 1111 1111 1111"; // Visa test number, Luhn-valid. @@ -15,10 +16,13 @@ async function flushMutations(): Promise { beforeEach(() => { document.body.innerHTML = ""; jest.useFakeTimers(); + history.replaceState(null, "", "/initial"); + __resetRouteChangeForTesting(); }); afterEach(() => { piiRedactRule.teardown(); + __resetRouteChangeForTesting(); jest.useRealTimers(); }); @@ -138,4 +142,86 @@ describe("pii-redact lazy-loaded subtrees", () => { expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength(1); }); + + it("masks every match when the tree spans multiple chunks", async () => { + // Tree exceeds chunkSize=100 so the walker yields mid-scan. Without + // the resume-anchor in walkSync, the walker's currentNode points to + // the (now-detached) last node of chunk 1 after replaceMatchesInTextNode + // runs — and nextNode() returns null, silently dropping every match + // past the first 100. This test runs the rule end-to-end and + // confirms all 200 SSNs get hidden, not just the first chunk. + document.body.innerHTML = Array.from( + { length: 200 }, + (_, i) => `

node-${i}: ${SSN}

`, + ).join(""); + + piiRedactRule.apply(document.body); + // Chunk 1 (100 nodes) runs synchronously. + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + + // Drain the chunked walk's yields. Each yield is a setTimeout(0) + // whose `.then()` callback schedules the next chunk via microtask; + // alternating timer + microtask drains the queue. + for (let i = 0; i < 5; i++) { + jest.advanceTimersByTime(0); + await flushMutations(); + } + + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 200, + ); + }); + + it("route change aborts the in-flight chunked walk", () => { + // Same shape as the teardown abort test, but the cancellation + // signal is a route-change event (the rule subscribes to + // subscribeRouteChange on first apply). Confirms the + // route-change → abortAndReset wiring fires. + document.body.innerHTML = Array.from( + { length: 200 }, + (_, i) => `

node-${i}: ${SSN}

`, + ).join(""); + + piiRedactRule.apply(document.body); + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + + history.replaceState(null, "", "/new-route"); + globalThis.dispatchEvent(new Event("popstate")); + jest.advanceTimersByTime(0); + + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + }); + + it("teardown aborts the in-flight chunked walk", () => { + // 200 text nodes — exceeds the 100-node chunkSize default, so the + // walk yields after chunk 1. teardown fires before the yield's + // setTimeout(0) resolves; the continuation sees the aborted signal + // and bails, leaving only the first chunk's matches masked. + document.body.innerHTML = Array.from( + { length: 200 }, + (_, i) => `

node-${i}: ${SSN}

`, + ).join(""); + + piiRedactRule.apply(document.body); + // Sync chunk 1: 100 placeholders. + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + + piiRedactRule.teardown(); + // Fire the yield's setTimeout(0). The continuation checks the + // signal first — aborted — and returns without processing + // chunk 2. + jest.advanceTimersByTime(0); + + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + }); }); diff --git a/extension/src/rules/__tests__/prompt-injection-redact.test.ts b/extension/src/rules/__tests__/prompt-injection-redact.test.ts index 6d6a9f9..dd80cc1 100644 --- a/extension/src/rules/__tests__/prompt-injection-redact.test.ts +++ b/extension/src/rules/__tests__/prompt-injection-redact.test.ts @@ -1,4 +1,5 @@ import { PLACEHOLDER_CLASS } from "../../lib/placeholder"; +import { __resetRouteChangeForTesting } from "../../lib/route-change"; import { promptInjectionRedactRule } from "../prompt-injection-redact"; import { FIXTURES } from "./injection-fixtures"; @@ -267,3 +268,63 @@ describe("prompt-injection-redact", () => { }); }); }); + +describe("prompt-injection-redact abort", () => { + beforeEach(() => { + jest.useFakeTimers(); + history.replaceState(null, "", "/initial"); + __resetRouteChangeForTesting(); + }); + afterEach(() => { + promptInjectionRedactRule.teardown(); + __resetRouteChangeForTesting(); + jest.useRealTimers(); + }); + + it("teardown aborts the in-flight chunked walk before the hide pass runs", () => { + // 200 paragraphs, each with an injection phrase. The chunked walk + // fills chunk 1 (100 nodes), yields, and only schedules the second + // pass (filterToOutermost + hide) via onComplete after the LAST + // chunk. teardown's abortAndReset fires before the yield resolves + // — so the continuation bails, onComplete never runs, and nothing + // gets hidden. The collected first-chunk containers stay in the + // DOM unchanged. + document.body.innerHTML = Array.from( + { length: 200 }, + (_, i) => `

${FIXTURES.IGNORE_HACKED} ${i}

`, + ).join(""); + + promptInjectionRedactRule.apply(document.body); + // Chunk 1 processed synchronously: collects containers but does + // NOT yet hide (onComplete is gated to the end of the walk). + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength(0); + expect(document.querySelectorAll("p")).toHaveLength(200); + + promptInjectionRedactRule.teardown(); + jest.advanceTimersByTime(0); + + // Aborted before onComplete: no placeholders ever installed. + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength(0); + expect(document.querySelectorAll("p")).toHaveLength(200); + }); + + it("route change aborts the in-flight chunked walk before onComplete", () => { + // Mirror of the teardown test, but the cancellation signal is a + // route-change event. Verifies the subscribeRouteChange wiring. + document.body.innerHTML = Array.from( + { length: 200 }, + (_, i) => `

${FIXTURES.IGNORE_HACKED} ${i}

`, + ).join(""); + + promptInjectionRedactRule.apply(document.body); + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength(0); + expect(document.querySelectorAll("p")).toHaveLength(200); + + history.replaceState(null, "", "/new-route"); + globalThis.dispatchEvent(new Event("popstate")); + jest.advanceTimersByTime(0); + + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength(0); + expect(document.querySelectorAll("p")).toHaveLength(200); + }); +}); diff --git a/extension/src/rules/__tests__/secrets-redact.test.ts b/extension/src/rules/__tests__/secrets-redact.test.ts index 5f07168..5e73cb7 100644 --- a/extension/src/rules/__tests__/secrets-redact.test.ts +++ b/extension/src/rules/__tests__/secrets-redact.test.ts @@ -1,4 +1,5 @@ import { PLACEHOLDER_CLASS } from "../../lib/placeholder"; +import { __resetRouteChangeForTesting } from "../../lib/route-change"; import { secretsRedactRule } from "../secrets-redact"; const AWS_KEY = "AKIAIOSFODNN7EXAMPLE"; @@ -27,10 +28,13 @@ async function flushMutations(): Promise { beforeEach(() => { document.body.innerHTML = ""; jest.useFakeTimers(); + history.replaceState(null, "", "/initial"); + __resetRouteChangeForTesting(); }); afterEach(() => { secretsRedactRule.teardown(); + __resetRouteChangeForTesting(); jest.useRealTimers(); }); @@ -229,4 +233,46 @@ describe("secrets-redact lazy-loaded subtrees", () => { expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength(1); }); + + it("teardown aborts the in-flight chunked walk", () => { + // Tree exceeds the 100-node chunkSize so the walk yields after + // chunk 1; teardown's abortAndReset bails the continuation. + document.body.innerHTML = Array.from( + { length: 200 }, + (_, i) => `

key-${i}: ${AWS_KEY}

`, + ).join(""); + + secretsRedactRule.apply(document.body); + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + + secretsRedactRule.teardown(); + jest.advanceTimersByTime(0); + + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + }); + + it("route change aborts the in-flight chunked walk", () => { + // Mirrors the teardown abort test but via the route-change wiring. + document.body.innerHTML = Array.from( + { length: 200 }, + (_, i) => `

key-${i}: ${AWS_KEY}

`, + ).join(""); + + secretsRedactRule.apply(document.body); + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + + history.replaceState(null, "", "/new-route"); + globalThis.dispatchEvent(new Event("popstate")); + jest.advanceTimersByTime(0); + + expect(document.querySelectorAll(`.${PLACEHOLDER_CLASS}`)).toHaveLength( + 100, + ); + }); }); diff --git a/extension/src/rules/encoded-payload-redact.ts b/extension/src/rules/encoded-payload-redact.ts index 1609fc0..e84f64f 100644 --- a/extension/src/rules/encoded-payload-redact.ts +++ b/extension/src/rules/encoded-payload-redact.ts @@ -19,10 +19,12 @@ // Matches are replaced inline with a click-to-reveal placeholder. False // positives cost one click, not lost data. -import { walkTextNodes } from "../lib/dom-utils"; +import { ReusableAbortController } from "abort-utils"; import type { InlineMatch } from "../lib/placeholder"; import { replaceMatchesInTextNode } from "../lib/placeholder"; +import { subscribeRouteChange } from "../lib/route-change"; import { createSubtreeWatcher } from "../lib/subtree-watcher"; +import { walkTextNodesChunked } from "../lib/yielding-text-walk"; import type { Rule } from "./types"; const RULE_ID = "encoded-payload-redact" as const; @@ -296,13 +298,24 @@ function collectMatches(text: string): InlineMatch[] { return merged; } +// See pii-redact for lifecycle rationale. +const lifecycle = new ReusableAbortController(); +let unsubscribeRouteChange: (() => void) | null = null; + function scanAndMask(root: ParentNode): void { - for (const node of walkTextNodes(root, { minLength: MIN_TEXT_LENGTH })) { - const matches = collectMatches(node.nodeValue ?? ""); - if (matches.length > 0) { - replaceMatchesInTextNode(node, matches, RULE_ID); - } - } + const signal = lifecycle.signal; + walkTextNodesChunked(root, { + signal, + minLength: MIN_TEXT_LENGTH, + process: (chunk) => { + for (const node of chunk) { + const matches = collectMatches(node.nodeValue ?? ""); + if (matches.length > 0) { + replaceMatchesInTextNode(node, matches, RULE_ID); + } + } + }, + }); } const watcher = createSubtreeWatcher({ @@ -315,6 +328,9 @@ const watcher = createSubtreeWatcher({ }); function apply(root: ParentNode): void { + unsubscribeRouteChange ??= subscribeRouteChange(() => { + lifecycle.abortAndReset(); + }); scanAndMask(root); watcher.start(root); } @@ -327,5 +343,8 @@ export const encodedPayloadRedactRule = { apply, teardown: () => { watcher.stop(); + lifecycle.abortAndReset(); + unsubscribeRouteChange?.(); + unsubscribeRouteChange = null; }, } satisfies Rule; diff --git a/extension/src/rules/pii-redact.ts b/extension/src/rules/pii-redact.ts index a0d7f1b..ad1814b 100644 --- a/extension/src/rules/pii-redact.ts +++ b/extension/src/rules/pii-redact.ts @@ -1,10 +1,12 @@ // Copyright (c) 2026 PixieBrix, Inc. // Licensed under PolyForm Shield 1.0.0 — see LICENSE. -import { walkTextNodes } from "../lib/dom-utils"; +import { ReusableAbortController } from "abort-utils"; import type { InlineMatch } from "../lib/placeholder"; import { replaceMatchesInTextNode } from "../lib/placeholder"; +import { subscribeRouteChange } from "../lib/route-change"; import { createSubtreeWatcher } from "../lib/subtree-watcher"; +import { walkTextNodesChunked } from "../lib/yielding-text-walk"; import type { Rule } from "./types"; const RULE_ID = "pii-redact" as const; @@ -76,13 +78,29 @@ function collectMatches(text: string): InlineMatch[] { return merged; } +// Lifecycle controller covers every async scan kicked off by apply +// and the watcher. abortAndReset on route change cancels any in-flight +// chunked walk so a scan started against the old tree can't keep +// mutating the new one; abortAndReset on teardown stops everything. +// Incremental subtree-watcher batches do NOT abort — they target their +// own scoped root and don't conflict with the previous scan. +const lifecycle = new ReusableAbortController(); +let unsubscribeRouteChange: (() => void) | null = null; + function scanAndMask(root: ParentNode): void { - for (const node of walkTextNodes(root, { minLength: MIN_TEXT_LENGTH })) { - const matches = collectMatches(node.nodeValue ?? ""); - if (matches.length > 0) { - replaceMatchesInTextNode(node, matches, RULE_ID); - } - } + const signal = lifecycle.signal; + walkTextNodesChunked(root, { + signal, + minLength: MIN_TEXT_LENGTH, + process: (chunk) => { + for (const node of chunk) { + const matches = collectMatches(node.nodeValue ?? ""); + if (matches.length > 0) { + replaceMatchesInTextNode(node, matches, RULE_ID); + } + } + }, + }); } const watcher = createSubtreeWatcher({ @@ -95,6 +113,9 @@ const watcher = createSubtreeWatcher({ }); function apply(root: ParentNode): void { + unsubscribeRouteChange ??= subscribeRouteChange(() => { + lifecycle.abortAndReset(); + }); scanAndMask(root); watcher.start(root); } @@ -106,5 +127,8 @@ export const piiRedactRule = { apply, teardown: () => { watcher.stop(); + lifecycle.abortAndReset(); + unsubscribeRouteChange?.(); + unsubscribeRouteChange = null; }, } satisfies Rule; diff --git a/extension/src/rules/prompt-injection-redact.ts b/extension/src/rules/prompt-injection-redact.ts index c7e1475..92d2924 100644 --- a/extension/src/rules/prompt-injection-redact.ts +++ b/extension/src/rules/prompt-injection-redact.ts @@ -13,13 +13,12 @@ // shipped extension bundle therefore contains plaintext regexes, not // `atob` decoding (matters for Chrome Web Store review). +import { ReusableAbortController } from "abort-utils"; import { REVEALED_ATTR } from "../lib/dom-markers"; -import { - filterToOutermost, - isInsidePlaceholder, - walkTextNodes, -} from "../lib/dom-utils"; +import { filterToOutermost, isInsidePlaceholder } from "../lib/dom-utils"; import { replaceWithBlockPlaceholder } from "../lib/placeholder"; +import { subscribeRouteChange } from "../lib/route-change"; +import { walkTextNodesChunked } from "../lib/yielding-text-walk"; import { INJECTION_PATTERNS } from "./injection-patterns.generated"; import type { Rule } from "./types"; @@ -54,49 +53,75 @@ function findContainer(textNode: Text): HTMLElement | null { return parent; } -function apply(root: ParentNode): void { +// Lifecycle controller cancels in-flight scans when an SPA route swap +// arrives so the post-route apply doesn't compete with the pre-route +// walk on the new tree. Teardown aborts and unsubscribes. +const lifecycle = new ReusableAbortController(); +let unsubscribeRouteChange: (() => void) | null = null; + +function scanAndMask(root: ParentNode): void { + const signal = lifecycle.signal; const containers = new Set(); - for (const node of walkTextNodes(root, { + walkTextNodesChunked(root, { + signal, minLength: MIN_TEXT_LENGTH, - // SVG /<desc>/<text> are the injection carriers inside SVG, and - // `svg-text-strip` handles them by blanking the text in place. Letting - // prompt-injection-redact also fire on those text nodes is destructive: - // the SVG has no p/li/td ancestor, so `findContainer` escalates all the - // way up to the surrounding <article>/<section> and the entire product - // header gets replaced with a single placeholder (#133). + // SVG <title>/<desc>/<text> are the injection carriers inside + // SVG, and `svg-text-strip` handles them by blanking the text in + // place. Letting prompt-injection-redact also fire on those text + // nodes is destructive: the SVG has no p/li/td ancestor, so + // `findContainer` escalates all the way up to the surrounding + // <article>/<section> and the entire product header gets + // replaced with a single placeholder (#133). shouldSkipParent: (parent) => parent.closest("svg") !== null, - })) { - if (!containsInjection(node.nodeValue ?? "")) { - continue; - } - const container = findContainer(node); - if (container) { - containers.add(container); - } - } + process: (chunk) => { + for (const node of chunk) { + if (!containsInjection(node.nodeValue ?? "")) { + continue; + } + const container = findContainer(node); + if (container) { + containers.add(container); + } + } + }, + // Second-pass hide runs after every text node is collected, so + // outermost dedupe sees the full container set. Synchronous when + // the tree fit in one chunk (the common test case); microtask- + // deferred when chunks fired. The helper guards against signal + // abort before calling us. + onComplete: () => { + for (const element of filterToOutermost([...containers])) { + if (!element.isConnected) { + continue; + } + if (isInsidePlaceholder(element)) { + continue; + } + // The container is an ancestor of the matched text node, so the reveal + // stamp lands here — not on the text node we walked from. If a previous + // run hid this container and the user revealed it, a re-apply (rule + // disable→enable, or a future MutationObserver-driven re-scan) would + // otherwise re-hide the same block. Same shape as the disguised-ad-flag + // bug fixed in #160. (Originally landed via #161; preserved here after + // the chunked-walk refactor moved the hide pass into onComplete.) + if (element.closest(REVEALED_SELECTOR)) { + continue; + } + replaceWithBlockPlaceholder( + element, + RULE_ID, + "[possible prompt injection hidden — click to reveal]", + ); + } + }, + }); +} - for (const element of filterToOutermost([...containers])) { - if (!element.isConnected) { - continue; - } - if (isInsidePlaceholder(element)) { - continue; - } - // The container is an ancestor of the matched text node, so the reveal - // stamp lands here — not on the text node we walked from. If a previous - // run hid this container and the user revealed it, a re-apply (rule - // disable→enable, or a future MutationObserver-driven re-scan) would - // otherwise re-hide the same block. Same shape as the disguised-ad-flag - // bug fixed in #160. - if (element.closest(REVEALED_SELECTOR)) { - continue; - } - replaceWithBlockPlaceholder( - element, - RULE_ID, - "[possible prompt injection hidden — click to reveal]", - ); - } +function apply(root: ParentNode): void { + unsubscribeRouteChange ??= subscribeRouteChange(() => { + lifecycle.abortAndReset(); + }); + scanAndMask(root); } export const promptInjectionRedactRule = { @@ -105,4 +130,9 @@ export const promptInjectionRedactRule = { description: "Hide page sections containing phrases common in prompt-injection attacks.", apply, + teardown: () => { + lifecycle.abortAndReset(); + unsubscribeRouteChange?.(); + unsubscribeRouteChange = null; + }, } satisfies Rule; diff --git a/extension/src/rules/secrets-redact.ts b/extension/src/rules/secrets-redact.ts index 8fc297f..e706f25 100644 --- a/extension/src/rules/secrets-redact.ts +++ b/extension/src/rules/secrets-redact.ts @@ -1,10 +1,12 @@ // Copyright (c) 2026 PixieBrix, Inc. // Licensed under PolyForm Shield 1.0.0 — see LICENSE. -import { walkTextNodes } from "../lib/dom-utils"; +import { ReusableAbortController } from "abort-utils"; import type { InlineMatch } from "../lib/placeholder"; import { replaceMatchesInTextNode } from "../lib/placeholder"; +import { subscribeRouteChange } from "../lib/route-change"; import { createSubtreeWatcher } from "../lib/subtree-watcher"; +import { walkTextNodesChunked } from "../lib/yielding-text-walk"; import type { Rule } from "./types"; const RULE_ID = "secrets-redact" as const; @@ -147,13 +149,26 @@ function collectMatches(text: string): InlineMatch[] { return merged; } +// See pii-redact for lifecycle rationale — same pattern: route-change +// aborts in-flight chunked walks; incremental subtree-watcher batches +// do not. +const lifecycle = new ReusableAbortController(); +let unsubscribeRouteChange: (() => void) | null = null; + function scanAndMask(root: ParentNode): void { - for (const node of walkTextNodes(root, { minLength: MIN_TEXT_LENGTH })) { - const matches = collectMatches(node.nodeValue ?? ""); - if (matches.length > 0) { - replaceMatchesInTextNode(node, matches, RULE_ID); - } - } + const signal = lifecycle.signal; + walkTextNodesChunked(root, { + signal, + minLength: MIN_TEXT_LENGTH, + process: (chunk) => { + for (const node of chunk) { + const matches = collectMatches(node.nodeValue ?? ""); + if (matches.length > 0) { + replaceMatchesInTextNode(node, matches, RULE_ID); + } + } + }, + }); } const watcher = createSubtreeWatcher({ @@ -166,6 +181,9 @@ const watcher = createSubtreeWatcher({ }); function apply(root: ParentNode): void { + unsubscribeRouteChange ??= subscribeRouteChange(() => { + lifecycle.abortAndReset(); + }); scanAndMask(root); watcher.start(root); } @@ -178,5 +196,8 @@ export const secretsRedactRule = { apply, teardown: () => { watcher.stop(); + lifecycle.abortAndReset(); + unsubscribeRouteChange?.(); + unsubscribeRouteChange = null; }, } satisfies Rule;