fix(builder): validate AI response before deserializing into craft state
Adds an AI-boundary validation/sanitization pass in apply-ai-response.ts so malformed or malicious Sitesmith output can never crash addNodeTree or corrupt live craft state: - resolvedName allowlist: any node whose type.resolvedName is not a registered component in componentResolver is dropped (with its subtree) and a console.warn is logged; buildNodeTree throws before touching Craft.js if the tree ROOT itself is invalid, which existing call sites already catch and warn on (fail soft, never throws into addNodeTree). - update_props protected-key list (node_id) — an AI-supplied node_id can no longer overwrite the target node's real id; style is merged shallowly instead of replaced wholesale so unrelated existing style keys survive. - id policy: regenerate, never skip. Any AI-supplied node id that is 'ROOT', empty/non-string, or collides with an existing/already-used id is replaced with a fresh ai-auto-N id; the node itself is kept. Extends apply-ai-response.test.ts with coverage for all of the above plus regression tests proving fully-valid AI responses still apply exactly as before. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1,5 +1,5 @@
|
||||
import { describe, test, expect } from 'vitest';
|
||||
import { serializeTreeForCraft, __test } from './apply-ai-response';
|
||||
import { describe, test, expect, vi } from 'vitest';
|
||||
import { serializeTreeForCraft, sanitizeAiTree, __test } from './apply-ai-response';
|
||||
|
||||
describe('serializeTreeForCraft', () => {
|
||||
test('flattens nested tree', () => {
|
||||
@@ -88,3 +88,180 @@ describe('findNodeIdByAiNodeId', () => {
|
||||
expect(__test.findNodeIdByAiNodeId({ getNodes: () => ({}) }, 'missing')).toBe(null);
|
||||
});
|
||||
});
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers for the tests below: fake @craftjs/core `query` / `actions` objects
|
||||
// good enough to exercise buildNodeTree/applyPatch without a real Editor.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function makeFakeQuery(existingIds: string[] = ['ROOT']) {
|
||||
return {
|
||||
getNodes: () => Object.fromEntries(existingIds.map((id) => [id, { data: { props: {} } }])),
|
||||
parseFreshNode: (input: any) => ({ toNode: () => input }),
|
||||
node: (id: string) => ({
|
||||
get: () => ({ data: { parent: 'ROOT' } }),
|
||||
childNodes: () => [id],
|
||||
}),
|
||||
};
|
||||
}
|
||||
|
||||
/** A fake query+actions pair backed by a single existing node, wired so
|
||||
* findNodeIdByAiNodeId can resolve `aiNodeId` back to `craftId`. */
|
||||
function makeSingleNodeHarness(craftId: string, aiNodeId: string, initialProps: Record<string, unknown>) {
|
||||
const propsById: Record<string, any> = { [craftId]: { node_id: aiNodeId, ...initialProps } };
|
||||
const query = {
|
||||
getNodes: () =>
|
||||
Object.fromEntries(Object.entries(propsById).map(([id, p]) => [id, { data: { props: p } }])),
|
||||
parseFreshNode: (input: any) => ({ toNode: () => input }),
|
||||
node: (id: string) => ({
|
||||
get: () => ({ data: { parent: 'ROOT' } }),
|
||||
childNodes: () => [craftId],
|
||||
}),
|
||||
};
|
||||
const actions = {
|
||||
setProp: (id: string, fn: (p: any) => void) => {
|
||||
if (propsById[id]) fn(propsById[id]);
|
||||
},
|
||||
delete: (id: string) => {
|
||||
delete propsById[id];
|
||||
},
|
||||
addNodeTree: () => {},
|
||||
};
|
||||
return { query, actions, propsById };
|
||||
}
|
||||
|
||||
describe('sanitizeAiTree', () => {
|
||||
test('unknown resolvedName at the root is rejected (returns null)', () => {
|
||||
const tree = { type: { resolvedName: 'NotARealComponent' }, props: {}, nodes: [] };
|
||||
expect(sanitizeAiTree(tree, new Set())).toBeNull();
|
||||
});
|
||||
|
||||
test('unknown resolvedName on a nested child drops only that subtree; valid siblings survive', () => {
|
||||
const tree = {
|
||||
type: { resolvedName: 'Section' },
|
||||
props: { node_id: 's1' },
|
||||
nodes: [
|
||||
{ type: { resolvedName: 'NotARealComponent' }, props: {}, nodes: [] },
|
||||
{ type: { resolvedName: 'Heading' }, props: { node_id: 'h1' }, nodes: [] },
|
||||
],
|
||||
};
|
||||
const out = sanitizeAiTree(tree, new Set());
|
||||
expect(out).not.toBeNull();
|
||||
expect(out!.nodes!.length).toBe(1);
|
||||
expect(out!.nodes![0].props.node_id).toBe('h1');
|
||||
});
|
||||
|
||||
test('a node id of "ROOT" is regenerated, never left as ROOT', () => {
|
||||
const tree = { type: { resolvedName: 'Heading' }, props: { node_id: 'ROOT' }, nodes: [] };
|
||||
const out = sanitizeAiTree(tree, new Set(['ROOT']));
|
||||
expect(out).not.toBeNull();
|
||||
expect(out!.props.node_id).not.toBe('ROOT');
|
||||
});
|
||||
|
||||
test('an empty/non-string id is regenerated to a non-empty string', () => {
|
||||
const tree = { type: { resolvedName: 'Heading' }, props: { node_id: '' }, nodes: [] };
|
||||
const out = sanitizeAiTree(tree, new Set());
|
||||
expect(typeof out!.props.node_id).toBe('string');
|
||||
expect(out!.props.node_id).toBeTruthy();
|
||||
});
|
||||
|
||||
test('an id colliding with an existing node is regenerated, not reused', () => {
|
||||
const tree = { type: { resolvedName: 'Heading' }, props: { node_id: 'existing-1' }, nodes: [] };
|
||||
const out = sanitizeAiTree(tree, new Set(['existing-1']));
|
||||
expect(out!.props.node_id).not.toBe('existing-1');
|
||||
});
|
||||
|
||||
test('duplicate ids within the same tree are de-duplicated (no two nodes share an id)', () => {
|
||||
const tree = {
|
||||
type: { resolvedName: 'Section' },
|
||||
props: { node_id: 'dup-1' },
|
||||
nodes: [
|
||||
{ type: { resolvedName: 'Heading' }, props: { node_id: 'dup-1' }, nodes: [] },
|
||||
],
|
||||
};
|
||||
const out = sanitizeAiTree(tree, new Set());
|
||||
expect(out!.props.node_id).toBe('dup-1');
|
||||
expect(out!.nodes![0].props.node_id).not.toBe('dup-1');
|
||||
});
|
||||
|
||||
test('a fully-valid tree passes through with structure and ids preserved (regression)', () => {
|
||||
const tree = {
|
||||
type: { resolvedName: 'Section' },
|
||||
props: { node_id: 'sec-1', aiName: 'Hero' },
|
||||
nodes: [
|
||||
{ type: { resolvedName: 'Heading' }, props: { node_id: 'h-1', text: 'Welcome' }, nodes: [] },
|
||||
],
|
||||
};
|
||||
const out = sanitizeAiTree(tree, new Set());
|
||||
expect(out!.props.node_id).toBe('sec-1');
|
||||
expect(out!.nodes!.length).toBe(1);
|
||||
expect(out!.nodes![0].props.node_id).toBe('h-1');
|
||||
expect(out!.nodes![0].props.text).toBe('Welcome');
|
||||
});
|
||||
});
|
||||
|
||||
describe('buildNodeTree', () => {
|
||||
test('builds a NodeTree preserving structure for a fully valid AI tree (regression)', () => {
|
||||
const tree = {
|
||||
type: { resolvedName: 'Section' },
|
||||
props: { node_id: 'sec-1' },
|
||||
nodes: [
|
||||
{ type: { resolvedName: 'Heading' }, props: { node_id: 'h-1', text: 'Hi' }, nodes: [] },
|
||||
],
|
||||
};
|
||||
const result = __test.buildNodeTree(makeFakeQuery(), tree);
|
||||
expect(result.rootNodeId).toBe('sec-1');
|
||||
expect(result.nodes['sec-1']).toBeDefined();
|
||||
// Section wraps direct children in a pre-created section-inner linkedNode.
|
||||
const innerId = result.nodes['sec-1'].data.linkedNodes['section-inner'];
|
||||
expect(innerId).toBeDefined();
|
||||
expect(result.nodes[innerId].data.nodes).toEqual(['h-1']);
|
||||
expect(result.nodes['h-1'].data.parent).toBe(innerId);
|
||||
});
|
||||
|
||||
test('throws when the tree root has an unknown resolvedName (caught by callers, never reaches addNodeTree)', () => {
|
||||
const tree = { type: { resolvedName: 'NotARealComponent' }, props: {}, nodes: [] };
|
||||
expect(() => __test.buildNodeTree(makeFakeQuery(), tree)).toThrow();
|
||||
});
|
||||
});
|
||||
|
||||
describe('applyPatch', () => {
|
||||
test('update_props: does not overwrite node_id, applies a legit prop, merges style shallowly', () => {
|
||||
const { query, actions, propsById } = makeSingleNodeHarness('craft-1', 'ai-node-1', {
|
||||
text: 'Old',
|
||||
style: { color: 'red', fontSize: '16px' },
|
||||
});
|
||||
const ops = [
|
||||
{
|
||||
op: 'update_props',
|
||||
node_id: 'ai-node-1',
|
||||
props: { node_id: 'ai-hacked', text: 'New', style: { color: 'blue' } },
|
||||
},
|
||||
];
|
||||
const result = __test.applyPatch(actions, query, ops);
|
||||
expect(result.ok).toBe(true);
|
||||
expect(propsById['craft-1'].node_id).toBe('ai-node-1'); // protected, unchanged
|
||||
expect(propsById['craft-1'].text).toBe('New'); // legit prop applied
|
||||
// style merges shallowly: color overwritten, fontSize preserved
|
||||
expect(propsById['craft-1'].style).toEqual({ color: 'blue', fontSize: '16px' });
|
||||
});
|
||||
|
||||
test('an op with an unknown resolvedName tree is skipped without throwing; other ops in the batch still apply', () => {
|
||||
const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
|
||||
const { query, actions, propsById } = makeSingleNodeHarness('craft-1', 'ai-node-1', {});
|
||||
const ops = [
|
||||
{
|
||||
op: 'insert_after',
|
||||
node_id: 'ai-node-1',
|
||||
tree: { type: { resolvedName: 'NotARealComponent' }, props: {}, nodes: [] },
|
||||
},
|
||||
{ op: 'delete_node', node_id: 'ai-node-1' },
|
||||
];
|
||||
let result: { ok: boolean; message?: string } | undefined;
|
||||
expect(() => { result = __test.applyPatch(actions, query, ops); }).not.toThrow();
|
||||
expect(result!.ok).toBe(true);
|
||||
expect(propsById['craft-1']).toBeUndefined(); // delete_node (op 2) still applied
|
||||
expect(warnSpy).toHaveBeenCalled(); // op 1's failure was logged, not thrown
|
||||
warnSpy.mockRestore();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -2,6 +2,7 @@ import { useEditor } from '@craftjs/core';
|
||||
import type { NodeTree } from '@craftjs/core';
|
||||
import { usePages } from '../state/PageContext';
|
||||
import { SitesmithResponse, SerializedTreeNode } from '../types/sitesmith';
|
||||
import { componentResolver } from '../components/resolver';
|
||||
|
||||
/** Only Container is a "real" Craft.js canvas in serialized state. Layout
|
||||
* shells (Section/HeroSimple/ColumnLayout/etc) use <Element canvas> linkedNodes
|
||||
@@ -16,6 +17,81 @@ const SHELL_INNER: Record<string, string> = {
|
||||
FormContainer: 'form-inner',
|
||||
};
|
||||
|
||||
/**
|
||||
* update_props may never touch these keys, regardless of what the AI sends.
|
||||
* `node_id` is the stable identifier other ops (and findNodeIdByAiNodeId)
|
||||
* rely on to re-target a node in later turns of the conversation — letting
|
||||
* the AI clobber it would silently orphan future patches.
|
||||
*/
|
||||
const PROTECTED_UPDATE_PROPS_KEYS = new Set<string>(['node_id']);
|
||||
|
||||
/** True if `name` is a registered component in the resolver (i.e. safe to
|
||||
* hand to `actions.addNodeTree` without it throwing "does not exist in
|
||||
* resolver" at render time). */
|
||||
function isKnownResolvedName(name: unknown): name is string {
|
||||
return typeof name === 'string' && Object.prototype.hasOwnProperty.call(componentResolver, name);
|
||||
}
|
||||
|
||||
/**
|
||||
* Validate + sanitize an AI-supplied tree before it is handed to
|
||||
* `buildNodeTree`/`query.parseFreshNode`. Fails SOFT — never throws:
|
||||
*
|
||||
* - **resolvedName allowlist**: any node whose `type.resolvedName` is not a
|
||||
* key of `componentResolver` is dropped (it and its subtree), with a
|
||||
* `console.warn`. If the tree's ROOT is itself invalid, this returns
|
||||
* `null` and the caller must skip the whole op/tree.
|
||||
* - **id policy — REGENERATE (not skip)**: a node id that is `'ROOT'`,
|
||||
* empty/non-string, or collides with an id already in `usedIds` is
|
||||
* replaced with a fresh `ai-auto-N` id (and a warning is logged). The
|
||||
* node itself is kept — only structurally invalid *types* are dropped;
|
||||
* a merely-bad *id* is repaired so we never lose otherwise-valid AI
|
||||
* content over an id collision.
|
||||
*
|
||||
* `usedIds` should be seeded with the live document's existing node ids
|
||||
* (see `buildNodeTree`) so AI-supplied ids can never collide with content
|
||||
* already on the canvas; it is also used internally to keep ids unique
|
||||
* within the tree being sanitized.
|
||||
*/
|
||||
export function sanitizeAiTree(
|
||||
tree: SerializedTreeNode,
|
||||
usedIds: Set<string> = new Set(),
|
||||
idCounter: { n: number } = { n: 0 },
|
||||
): SerializedTreeNode | null {
|
||||
const resolvedName = tree?.type?.resolvedName;
|
||||
if (!isKnownResolvedName(resolvedName)) {
|
||||
console.warn(`sitesmith: dropping node with unknown component "${String(resolvedName)}"`);
|
||||
return null;
|
||||
}
|
||||
|
||||
const rawId = tree.props?.node_id;
|
||||
let id = typeof rawId === 'string' ? rawId.trim() : '';
|
||||
if (!id) {
|
||||
id = `ai-auto-${idCounter.n++}`;
|
||||
} else if (id === 'ROOT') {
|
||||
console.warn('sitesmith: rejecting AI-supplied node id "ROOT" (reserved); regenerating');
|
||||
id = `ai-auto-${idCounter.n++}`;
|
||||
} else if (usedIds.has(id)) {
|
||||
console.warn(`sitesmith: AI-supplied node id "${id}" collides with an existing id; regenerating`);
|
||||
id = `ai-auto-${idCounter.n++}`;
|
||||
}
|
||||
// Belt-and-suspenders: keep drawing fresh ids in the (pathological) case
|
||||
// the freshly generated id itself collides.
|
||||
while (usedIds.has(id)) id = `ai-auto-${idCounter.n++}`;
|
||||
usedIds.add(id);
|
||||
|
||||
const children: SerializedTreeNode[] = [];
|
||||
for (const child of tree.nodes ?? []) {
|
||||
const sanitizedChild = sanitizeAiTree(child, usedIds, idCounter);
|
||||
if (sanitizedChild) children.push(sanitizedChild);
|
||||
}
|
||||
|
||||
return {
|
||||
type: tree.type,
|
||||
props: { ...tree.props, node_id: id },
|
||||
nodes: children,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Flatten a SerializedTreeNode tree into a Craft.js node map ready for
|
||||
* `actions.deserialize()`.
|
||||
@@ -68,11 +144,25 @@ export function serializeTreeForCraft(tree: SerializedTreeNode): { rootNodeId: s
|
||||
* inserting/replacing sections or individual nodes.
|
||||
*/
|
||||
function buildNodeTree(query: any, tree: SerializedTreeNode): NodeTree {
|
||||
const idCounter = { n: 0 };
|
||||
// Validate + sanitize at the AI boundary before touching Craft.js at all:
|
||||
// unknown resolvedNames are dropped (never reach parseFreshNode/addNodeTree,
|
||||
// which would throw), and ids are repaired ('ROOT'/empty/colliding →
|
||||
// regenerated) against the live document's existing node ids so a new
|
||||
// AI-supplied id can never clobber or duplicate one already on the canvas.
|
||||
const existingIds = new Set<string>(
|
||||
typeof query?.getNodes === 'function' ? Object.keys(query.getNodes()) : [],
|
||||
);
|
||||
const sanitized = sanitizeAiTree(tree, existingIds);
|
||||
if (!sanitized) {
|
||||
throw new Error('sitesmith: AI tree root has an unknown/invalid component type; nothing to build');
|
||||
}
|
||||
|
||||
const craftNodes: Record<string, any> = {};
|
||||
|
||||
const walk = (node: SerializedTreeNode, parent: string | null): string => {
|
||||
const id = (node.props.node_id as string | undefined) || `ai-auto-${idCounter.n++}`;
|
||||
// sanitizeAiTree() has already guaranteed every node here has a valid,
|
||||
// unique, non-'ROOT' node_id — no fallback/collision handling needed.
|
||||
const id = node.props.node_id as string;
|
||||
const craftNode = (query.parseFreshNode({
|
||||
id,
|
||||
data: {
|
||||
@@ -138,7 +228,7 @@ function buildNodeTree(query: any, tree: SerializedTreeNode): NodeTree {
|
||||
return id;
|
||||
};
|
||||
|
||||
const rootId = walk(tree, null);
|
||||
const rootId = walk(sanitized, null);
|
||||
return { rootNodeId: rootId, nodes: craftNodes };
|
||||
}
|
||||
|
||||
@@ -156,7 +246,7 @@ export function findNodeIdByAiNodeId(query: any, aiNodeId: string): string | nul
|
||||
}
|
||||
|
||||
/** Exported for unit tests */
|
||||
export const __test = { findNodeIdByAiNodeId };
|
||||
export const __test = { findNodeIdByAiNodeId, buildNodeTree, applyPatch };
|
||||
|
||||
/**
|
||||
* React hook that returns an `apply` function.
|
||||
@@ -243,9 +333,32 @@ function applyPatch(
|
||||
}
|
||||
|
||||
switch (op.op) {
|
||||
case 'update_props':
|
||||
actions.setProp(id, (p: any) => { Object.assign(p, op.props); });
|
||||
case 'update_props': {
|
||||
const patchProps = op.props;
|
||||
if (!patchProps || typeof patchProps !== 'object' || Array.isArray(patchProps)) {
|
||||
console.warn('sitesmith patch: update_props with invalid props payload, skipping', op.node_id);
|
||||
break;
|
||||
}
|
||||
actions.setProp(id, (p: any) => {
|
||||
for (const [key, value] of Object.entries(patchProps)) {
|
||||
if (PROTECTED_UPDATE_PROPS_KEYS.has(key)) {
|
||||
console.warn(`sitesmith patch: refusing to overwrite protected prop "${key}"`);
|
||||
continue;
|
||||
}
|
||||
// Merge `style` shallowly instead of blindly replacing it — the
|
||||
// AI usually only means to change one or two style keys, and a
|
||||
// blind Object.assign would clobber every other style the user
|
||||
// (or a previous AI turn) had already set.
|
||||
if (key === 'style' && value && typeof value === 'object' && !Array.isArray(value)) {
|
||||
const existingStyle = p.style && typeof p.style === 'object' ? p.style : {};
|
||||
p.style = { ...existingStyle, ...(value as Record<string, unknown>) };
|
||||
} else {
|
||||
p[key] = value;
|
||||
}
|
||||
}
|
||||
});
|
||||
break;
|
||||
}
|
||||
|
||||
case 'replace_node': {
|
||||
try {
|
||||
|
||||
Reference in New Issue
Block a user