Keep JSON response transformations consistent

src/protect/runtime.js

@@ -606,7 +606,11 @@ export async function createProtection(options = {}) {
       const spanRedactors = redactors.filter((r) => !r.jsonPath);
       if (pathRedactors.length) body = applyPathRedactors(body, pathRedactors, mask, screenCap, transform);
       if (!spanRedactors.length) continue;
+      const beforeSpan = body;
       body = applyRedactors(body, spanRedactors, mask, transform);
+      // Span rewrites may change JSON string values, but not keys, containers or other values.
+      // Check each result before it becomes input to another transformation.
+      if (body !== beforeSpan && !preservesJsonStructure(beforeSpan, body)) return { verdict: 'block' };
       if (transform) continue; // encoding is a body/output concern — headers aren't HTML
       for (const name of Object.keys(headers)) {
         const value = headers[name];
@@ -622,6 +626,9 @@ export async function createProtection(options = {}) {
         }
       }
     }
+    // A rewritten JSON document must still be consumable as JSON. Text-span transformations can
+    // cross its escaping or delimiters; withhold that result rather than emit an invalid document.
+    if (body !== text && isJson(text) && !isJson(body)) return { verdict: 'block' };
     for (const rule of headerMutations) applyHeaderMutation(headers, rule);
     return { verdict: 'redact', body, headers };
   };
@@ -2051,6 +2058,44 @@ function htmlEscape(str) {
   return String(str).replace(/[&<>"']/g, (c) => ({ '&': '&amp;', '<': '&lt;', '>': '&gt;', '"': '&quot;', "'": '&#39;' })[c]);
 }
 
+function isJson(text) {
+  try {
+    JSON.parse(text);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+function preservesJsonStructure(before, after) {
+  if (!isJson(before)) return true;
+  if (!isJson(after)) return false;
+  const expected = jsonStructure(before);
+  const actual = jsonStructure(after);
+  for (;;) {
+    const left = expected.next();
+    const right = actual.next();
+    if (left.done || right.done) return left.done === right.done;
+    if (left.value !== right.value) return false;
+  }
+}
+
+// Called only for validated JSON. String values are the only interchangeable tokens; key names,
+// punctuation and non-string tokens remain exact, including number spellings and repeated keys.
+function* jsonStructure(text) {
+  const tokens = /"(?:[^"\\]|\\[\s\S])*"|-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?|[^\s]/g;
+  for (const match of text.matchAll(tokens)) {
+    const token = match[0];
+    if (token[0] !== '"') {
+      yield 'token:' + token;
+      continue;
+    }
+    let next = match.index + token.length;
+    while (text[next] === ' ' || text[next] === '\t' || text[next] === '\r' || text[next] === '\n') next++;
+    yield text[next] === ':' ? 'key:' + JSON.parse(token) : 'string';
+  }
+}
+
 // `transform` (optional): map a matched span to its replacement (the `encode` action passes
 // htmlEscape). Without it, matches are replaced by the `mask` string (the `redact` action).
 function applyRedactors(body, redactors, mask, transform) {
@@ -2121,10 +2166,13 @@ function applyPathRedactors(text, pathRedactors, mask, cap, transform) {
   // Preserve out-of-safe-range integers across the parse→stringify round-trip: JSON.parse would
   // round e.g. a 20-digit id. We quote such number tokens to a sentinel string before parsing and
   // unquote them after stringifying, so untouched big ints survive losslessly.
-  const preserved = preserveBigInts(text);
+  let preserved;
   let obj;
   try {
-    obj = JSON.parse(preserved);
+    // Only valid JSON reaches tokenization; malformed intermediate text is not a token source.
+    JSON.parse(text);
+    preserved = preserveBigInts(text, mask);
+    obj = JSON.parse(preserved.text);
   } catch {
     return text;
   }
@@ -2133,61 +2181,50 @@ function applyPathRedactors(text, pathRedactors, mask, cap, transform) {
     const pred = conditionPredicate(r.condition);
     walkLeaves(obj, r.jsonPath, (loc) => {
       try {
-        if (pred(loc.value)) {
+        const number = preserved.numbers.get(loc.value);
+        // Detection uses JSON.parse's numeric value. Match that same value, while retaining the
+        // original token for any untouched leaf and for string transformations.
+        if (pred(number === undefined ? loc.value : Number(number))) {
           // `encode`: escape the leaf's own value in place; `redact`: replace it with the mask.
-          loc.parent[loc.key] = transform ? transform(String(loc.value)) : mask;
+          setOwn(loc.parent, loc.key, transform ? transform(number ?? String(loc.value)) : mask);
           changed = true;
         }
       } catch {
         /* skip this leaf */
       }
     });
   }
-  return changed ? restoreBigInts(JSON.stringify(obj)) : text;
+  return changed ? restoreBigInts(JSON.stringify(obj), preserved.numbers) : text;
 }
 
-const BIGINT_OPEN = '__PSBIGINT_9c2f__';
-const BIGINT_CLOSE = '__DNEGIB__';
-
-// Quote every out-of-safe-range integer *value* (a bare number token outside a string) into a
-// sentinel string, so JSON.parse keeps it verbatim. String-aware scan (respects \ escapes) so a
-// number inside a string value is never touched. Plain-ASCII sentinel → survives JSON.stringify.
-function preserveBigInts(text) {
-  let out = '';
-  let inStr = false;
-  for (let i = 0; i < text.length; ) {
-    const ch = text[i];
-    if (inStr) {
-      out += ch;
-      if (ch === '\\') { out += text[i + 1] ?? ''; i += 2; continue; }
-      if (ch === '"') inStr = false;
-      i++;
-      continue;
-    }
-    if (ch === '"') { inStr = true; out += ch; i++; continue; }
-    if (ch === '-' || (ch >= '0' && ch <= '9')) {
-      let j = ch === '-' ? i + 1 : i;
-      let digits = 0;
-      while (j < text.length && text[j] >= '0' && text[j] <= '9') { digits++; j++; }
-      const next = text[j];
-      const isIntToken = digits > 0 && next !== '.' && next !== 'e' && next !== 'E';
-      if (isIntToken && digits >= 16) {
-        out += `"${BIGINT_OPEN}${text.slice(i, j)}${BIGINT_CLOSE}"`;
-      } else {
-        out += text.slice(i, j || i + 1);
-      }
-      i = j > i ? j : i + 1;
-      continue;
-    }
-    out += ch;
-    i++;
+// Preserve whole integer tokens, never a digit sequence inside a string, fraction or exponent.
+// The placeholders are unique to this document and cannot alias a literal string or the mask.
+function preserveBigInts(text, mask) {
+  const occupied = new Set([mask]);
+  const integers = [];
+  const tokens = /"(?:[^"\\]|\\[\s\S])*"|-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?/g;
+  for (const match of text.matchAll(tokens)) {
+    const token = match[0];
+    if (token[0] === '"') occupied.add(JSON.parse(token));
+    else if (/^-?\d{16,}$/.test(token)) integers.push({ start: match.index, token });
   }
-  return out;
+  const numbers = new Map();
+  const chunks = [];
+  let offset = 0;
+  let index = 0;
+  for (const { start, token } of integers) {
+    let marker;
+    do { marker = '__PSNUMBER_' + index++ + '__'; } while (occupied.has(marker));
+    numbers.set(marker, token);
+    chunks.push(text.slice(offset, start), JSON.stringify(marker));
+    offset = start + token.length;
+  }
+  chunks.push(text.slice(offset));
+  return { text: chunks.join(''), numbers };
 }
 
-function restoreBigInts(text) {
-  if (!text.includes(BIGINT_OPEN)) return text;
-  return text.replace(new RegExp(`"${BIGINT_OPEN}(-?\\d+)${BIGINT_CLOSE}"`, 'g'), '$1');
+function restoreBigInts(text, numbers) {
+  return text.replace(/"(__PSNUMBER_\d+__)"/g, (token, marker) => numbers.get(marker) ?? token);
 }
 
 // Response-hardening actions. Mutate the (lowercase-keyed) headers object in place; a `null` value

tests/protect/response-json-encoding.test.ts

@@ -0,0 +1,158 @@
+import { describe, expect, it, vi } from 'vitest';
+import { createServer } from 'node:http';
+import { createProtection } from '../../src/protect/runtime.js';
+
+async function screen(text: string, value: string, mode = 'block', type = 'application/json', action = 'encode') {
+  const protection = await createProtection({
+    mode,
+    rules: { firewall: [], whitelists: [], whitelist_keys: {} },
+    responseRules: [{
+      id: 'text-format', phase: 'response', action,
+      rule_v2: [{ parameter: 'response.body', match: { type: 'contains', value } }],
+    }],
+  });
+  return protection.screenResponse(new Response(text, { headers: { 'content-type': type } }));
+}
+
+describe('span encoding and JSON representation', () => {
+  const text = JSON.stringify({ message: '<label title="sample">text</label>' });
+  const match = '<label title=\\"sample\\">';
+
+  it.each(['encode', 'redact'])('withholds a %s span crossing object fields', async (action) => {
+    const response = await screen('{"label":"first","state":"second"}', 'first","state":"second', 'block', 'application/json', action);
+    expect(response.status).toBe(500);
+    expect(await response.json()).toHaveProperty('error');
+  });
+
+  it('withholds a span that changes an object key', async () => {
+    const response = await screen('{"<label>":"sample"}', '<label>');
+    expect(response.status).toBe(500);
+  });
+
+  it.each([
+    ['{"list":["first","second"]}', 'first","second'],
+    ['{"value":true}', 'true'],
+    ['{"value":123}', '123'],
+    ['{"value":null}', 'null'],
+  ])('preserves containers and non-string values in %s', async (document, value) => {
+    const response = await screen(document, value, 'block', 'application/json', 'redact');
+    expect(response.status).toBe(500);
+  });
+
+  it('allows nested string changes without changing adjacent values', async () => {
+    const response = await screen('{"list":[{"label":"<label>","count":1}],"enabled":true,"empty":null}', '<label>');
+    expect(response.status).toBe(200);
+    expect(await response.json()).toEqual({ list: [{ label: '&lt;label&gt;', count: 1 }], enabled: true, empty: null });
+  });
+
+  it('checks span changes separately from explicit path transformations', async () => {
+    const protection = await createProtection({
+      mode: 'block',
+      rules: { firewall: [], whitelists: [], whitelist_keys: {} },
+      responseRules: [
+        {
+          id: 'field-format', phase: 'response', action: 'redact',
+          rule_v2: [{
+            parameter: 'response.body', mutations: ['json_decode'],
+            match: { type: 'array_key_value', key: 'count', match: { type: 'isset' } },
+          }],
+        },
+        {
+          id: 'text-format', phase: 'response', action: 'encode',
+          rule_v2: [{ parameter: 'response.body', match: { type: 'contains', value: '<label>' } }],
+        },
+      ],
+    });
+    const response = await protection.screenResponse(new Response('{"count":12,"label":"<label>"}'));
+    expect(response.status).toBe(200);
+    expect(await response.json()).toEqual({ count: '[REDACTED]', label: '&lt;label&gt;' });
+  });
+
+  it('does not tokenize an incomplete intermediate document', async () => {
+    const document = JSON.stringify({ value: 12, label: '"'.repeat(128) + 'sample' });
+    const suffix = 'sample"}';
+    const incomplete = document.slice(0, -suffix.length) + '[REDACTED]';
+    const matchAll = String.prototype.matchAll;
+    let scannedIncomplete = false;
+    const scanSpy = vi.spyOn(String.prototype, 'matchAll').mockImplementation(function (regexp) {
+      if (String(this) === incomplete) scannedIncomplete = true;
+      return matchAll.call(this, regexp);
+    });
+    try {
+      const protection = await createProtection({
+        mode: 'block',
+        rules: { firewall: [], whitelists: [], whitelist_keys: {} },
+        responseRules: [
+          {
+            id: 'text-format', phase: 'response', action: 'redact',
+            rule_v2: [{ parameter: 'response.body', match: { type: 'contains', value: suffix } }],
+          },
+          {
+            id: 'field-format', phase: 'response', action: 'redact',
+            rule_v2: [{
+              parameter: 'response.body', mutations: ['json_decode'],
+              match: { type: 'array_key_value', key: 'value', match: { type: 'isset' } },
+            }],
+          },
+        ],
+      });
+      const response = await protection.screenResponse(new Response(document));
+      expect(response.status).toBe(500);
+      expect(scannedIncomplete).toBe(false);
+    } finally {
+      scanSpy.mockRestore();
+    }
+  });
+
+  it.each(['application/json', 'text/plain'])('withholds a rewrite that cannot remain valid JSON (%s)', async (type) => {
+    const response = await screen(text, match, 'block', type);
+    expect(response.status).toBe(500);
+    expect(await response.json()).toEqual({ error: 'Response withheld by Patchstack (sensitive data detected)' });
+  });
+
+  it('keeps a representable span encoded inside JSON', async () => {
+    const response = await screen(JSON.stringify({ message: '<label>text</label>' }), '<label>');
+    expect(response.status).toBe(200);
+    expect(await response.json()).toEqual({ message: '&lt;label&gt;text</label>' });
+  });
+
+  it('leaves dry-run output byte-for-byte unchanged', async () => {
+    const response = await screen(text, match, 'dry-run');
+    expect(response.status).toBe(200);
+    expect(await response.text()).toBe(text);
+  });
+
+  it('still encodes ordinary HTML text', async () => {
+    const response = await screen('<label title="sample">text</label>', '<label title="sample">', 'block', 'text/html');
+    expect(response.status).toBe(200);
+    expect(await response.text()).toBe('&lt;label title=&quot;sample&quot;&gt;text</label>');
+  });
+
+  it('sends a complete withheld response through the Node adapter', async () => {
+    const protection = await createProtection({
+      mode: 'block',
+      rules: { firewall: [], whitelists: [], whitelist_keys: {} },
+      responseRules: [{
+        id: 'text-format', phase: 'response', action: 'encode',
+        rule_v2: [{ parameter: 'response.body', match: { type: 'contains', value: match } }],
+      }],
+    });
+    const middleware = protection.node({ screenResponses: true });
+    const server = createServer((req, res) => middleware(req, res, () => {
+      res.writeHead(200, { 'content-type': 'application/json', 'content-length': Buffer.byteLength(text) });
+      res.end(text);
+    }));
+    await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve));
+    try {
+      const address = server.address() as { port: number };
+      const response = await fetch('http://127.0.0.1:' + address.port);
+      const received = await response.text();
+      expect(response.status).toBe(500);
+      expect(Number(response.headers.get('content-length'))).toBe(Buffer.byteLength(received));
+      expect(JSON.parse(received)).toEqual({ error: 'Response withheld by Patchstack (sensitive data detected)' });
+    } finally {
+      await new Promise<void>((resolve, reject) => server.close((error) => error ? reject(error) : resolve()));
+      protection.stop();
+    }
+  });
+});