fix: Fix app crashing when you upload an incorrect JSON to node editor (#3911)

## What type of PR is this? (check all applicable)

- [x] Bug Fix


## Have you discussed this change with the InvokeAI team?
- [x] Yes, we feel very passionate about this.     

## Description

Uploading an incorrect JSON file to the Node Editor would crash the app.

While this is a much larger problem that we will tackle while refining
the Node Editor, this is a fix that should address 99% of the cases out
there.

When saving an InvokeAI node graph, there are three primary keys.

1. `nodes` - which has all the node related data.
2. `edges` - which has all the edges related data
3. `viewport` - which has all the viewport related data.

So when we load back the JSON, we now check if all three of these keys
exist in the retrieved JSON object. While the `viewport` itself is not a
mandatory key to repopulate the graph, checking for it will allow us to
treat it as an additional check to ensure that the graph was saved from
InvokeAI.

As a result ...

- If you upload an invalid JSON file, the app now warns you that the
JSON is invalid.
- If you upload a JSON of a graph editor that is not InvokeAI, it simply
warns you that you are uploading a non InvokeAI graph.

So effectively, you should not be able to load any graph that is not
generated by ReactFlow.

Here are the edge cases:

- What happens if a user maintains the above key structure but tampers
with the data inside them? Well tested it. Turns out because we validate
and build the graph based on the JSON data, if you tamper with any data
that is needed to rebuild that node, it simply will skip that and load
the rest of the graph with valid data.
- What happens if a user uploads a graph that was made by some other
random ReactFlow app? Well, same as above. Because we do not have to
parse that in our setup, it simply will skip it and only display what
are setup to do.

I think that just about covers 99% of the cases where this could go
wrong. If there's any other edges cases, can add checks if need be. But
can't think of any at the moment.

## Related Tickets & Documents

### Closes
- #3893 
- #3881

## [optional] Are there any post deployment tasks we need to perform?

Yes. Making @psychedelicious a little bit happier. :P
This commit is contained in:
blessedcoolant 2023-07-24 02:15:46 +12:00 committed by GitHub
commit 02618a701d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 189 additions and 102 deletions

View File

@ -102,8 +102,7 @@
"openInNewTab": "Open in New Tab",
"dontAskMeAgain": "Don't ask me again",
"areYouSure": "Are you sure?",
"imagePrompt": "Image Prompt",
"clearNodes": "Are you sure you want to clear all nodes?"
"imagePrompt": "Image Prompt"
},
"gallery": {
"generations": "Generations",
@ -615,6 +614,11 @@
"initialImageNotSetDesc": "Could not load initial image",
"nodesSaved": "Nodes Saved",
"nodesLoaded": "Nodes Loaded",
"nodesNotValidGraph": "Not a valid InvokeAI Node Graph",
"nodesNotValidJSON": "Not a valid JSON",
"nodesCorruptedGraph": "Cannot load. Graph seems to be corrupted.",
"nodesUnrecognizedTypes": "Cannot load. Graph has unrecognized types",
"nodesBrokenConnections": "Cannot load. Some connections are broken.",
"nodesLoadedFailed": "Failed To Load Nodes",
"nodesCleared": "Nodes Cleared"
},
@ -700,9 +704,10 @@
},
"nodes": {
"reloadSchema": "Reload Schema",
"saveNodes": "Save Nodes",
"loadNodes": "Load Nodes",
"clearNodes": "Clear Nodes",
"saveGraph": "Save Graph",
"loadGraph": "Load Graph (saved from Node Editor) (Do not copy-paste metadata)",
"clearGraph": "Clear Graph",
"clearGraphDesc": "Are you sure you want to clear all nodes?",
"zoomInNodes": "Zoom In",
"zoomOutNodes": "Zoom Out",
"fitViewportNodes": "Fit View",

View File

@ -2,11 +2,11 @@ import { HStack } from '@chakra-ui/react';
import CancelButton from 'features/parameters/components/ProcessButtons/CancelButton';
import { memo } from 'react';
import { Panel } from 'reactflow';
import LoadNodesButton from '../ui/LoadNodesButton';
import ClearGraphButton from '../ui/ClearGraphButton';
import LoadGraphButton from '../ui/LoadGraphButton';
import NodeInvokeButton from '../ui/NodeInvokeButton';
import ReloadSchemaButton from '../ui/ReloadSchemaButton';
import SaveNodesButton from '../ui/SaveNodesButton';
import ClearNodesButton from '../ui/ClearNodesButton';
import SaveGraphButton from '../ui/SaveGraphButton';
const TopCenterPanel = () => {
return (
@ -15,9 +15,9 @@ const TopCenterPanel = () => {
<NodeInvokeButton />
<CancelButton />
<ReloadSchemaButton />
<SaveNodesButton />
<LoadNodesButton />
<ClearNodesButton />
<SaveGraphButton />
<LoadGraphButton />
<ClearGraphButton />
</HStack>
</Panel>
);

View File

@ -9,17 +9,17 @@ import {
Text,
useDisclosure,
} from '@chakra-ui/react';
import { makeToast } from 'features/system/util/makeToast';
import { RootState } from 'app/store/store';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import IAIIconButton from 'common/components/IAIIconButton';
import { nodeEditorReset } from 'features/nodes/store/nodesSlice';
import { addToast } from 'features/system/store/systemSlice';
import { memo, useRef, useCallback } from 'react';
import { makeToast } from 'features/system/util/makeToast';
import { memo, useCallback, useRef } from 'react';
import { useTranslation } from 'react-i18next';
import { FaTrash } from 'react-icons/fa';
const ClearNodesButton = () => {
const ClearGraphButton = () => {
const { t } = useTranslation();
const dispatch = useAppDispatch();
const { isOpen, onOpen, onClose } = useDisclosure();
@ -46,8 +46,8 @@ const ClearNodesButton = () => {
<>
<IAIIconButton
icon={<FaTrash />}
tooltip={t('nodes.clearNodes')}
aria-label={t('nodes.clearNodes')}
tooltip={t('nodes.clearGraph')}
aria-label={t('nodes.clearGraph')}
onClick={onOpen}
isDisabled={nodes.length === 0}
/>
@ -62,11 +62,11 @@ const ClearNodesButton = () => {
<AlertDialogContent>
<AlertDialogHeader fontSize="lg" fontWeight="bold">
{t('nodes.clearNodes')}
{t('nodes.clearGraph')}
</AlertDialogHeader>
<AlertDialogBody>
<Text>{t('common.clearNodes')}</Text>
<Text>{t('nodes.clearGraphDesc')}</Text>
</AlertDialogBody>
<AlertDialogFooter>
@ -83,4 +83,4 @@ const ClearNodesButton = () => {
);
};
export default memo(ClearNodesButton);
export default memo(ClearGraphButton);

View File

@ -0,0 +1,161 @@
import { FileButton } from '@mantine/core';
import { useAppDispatch } from 'app/store/storeHooks';
import IAIIconButton from 'common/components/IAIIconButton';
import { loadFileEdges, loadFileNodes } from 'features/nodes/store/nodesSlice';
import { addToast } from 'features/system/store/systemSlice';
import { makeToast } from 'features/system/util/makeToast';
import i18n from 'i18n';
import { memo, useCallback, useRef } from 'react';
import { useTranslation } from 'react-i18next';
import { FaUpload } from 'react-icons/fa';
import { useReactFlow } from 'reactflow';
interface JsonFile {
[key: string]: unknown;
}
function sanityCheckInvokeAIGraph(jsonFile: JsonFile): {
isValid: boolean;
message: string;
} {
// Check if primary keys exist
const keys = ['nodes', 'edges', 'viewport'];
for (const key of keys) {
if (!(key in jsonFile)) {
return {
isValid: false,
message: i18n.t('toast.nodesNotValidGraph'),
};
}
}
// Check if nodes and edges are arrays
if (!Array.isArray(jsonFile.nodes) || !Array.isArray(jsonFile.edges)) {
return {
isValid: false,
message: i18n.t('toast.nodesNotValidGraph'),
};
}
// Check if data is present in nodes
const nodeKeys = ['data', 'type'];
const nodeTypes = ['invocation', 'progress_image'];
if (jsonFile.nodes.length > 0) {
for (const node of jsonFile.nodes) {
for (const nodeKey of nodeKeys) {
if (!(nodeKey in node)) {
return {
isValid: false,
message: i18n.t('toast.nodesNotValidGraph'),
};
}
if (nodeKey === 'type' && !nodeTypes.includes(node[nodeKey])) {
return {
isValid: false,
message: i18n.t('toast.nodesUnrecognizedTypes'),
};
}
}
}
}
// Check Edge Object
const edgeKeys = ['source', 'sourceHandle', 'target', 'targetHandle'];
if (jsonFile.edges.length > 0) {
for (const edge of jsonFile.edges) {
for (const edgeKey of edgeKeys) {
if (!(edgeKey in edge)) {
return {
isValid: false,
message: i18n.t('toast.nodesBrokenConnections'),
};
}
}
}
}
return {
isValid: true,
message: i18n.t('toast.nodesLoaded'),
};
}
const LoadGraphButton = () => {
const { t } = useTranslation();
const dispatch = useAppDispatch();
const { fitView } = useReactFlow();
const uploadedFileRef = useRef<() => void>(null);
const restoreJSONToEditor = useCallback(
(v: File | null) => {
if (!v) return;
const reader = new FileReader();
reader.onload = async () => {
const json = reader.result;
try {
const retrievedNodeTree = await JSON.parse(String(json));
const { isValid, message } =
sanityCheckInvokeAIGraph(retrievedNodeTree);
if (isValid) {
dispatch(loadFileNodes(retrievedNodeTree.nodes));
dispatch(loadFileEdges(retrievedNodeTree.edges));
fitView();
dispatch(
addToast(makeToast({ title: message, status: 'success' }))
);
} else {
dispatch(
addToast(
makeToast({
title: message,
status: 'error',
})
)
);
}
// Cleanup
reader.abort();
} catch (error) {
if (error) {
dispatch(
addToast(
makeToast({
title: t('toast.nodesNotValidJSON'),
status: 'error',
})
)
);
}
}
};
reader.readAsText(v);
// Cleanup
uploadedFileRef.current?.();
},
[fitView, dispatch, t]
);
return (
<FileButton
resetRef={uploadedFileRef}
accept="application/json"
onChange={restoreJSONToEditor}
>
{(props) => (
<IAIIconButton
icon={<FaUpload />}
tooltip={t('nodes.loadGraph')}
aria-label={t('nodes.loadGraph')}
{...props}
/>
)}
</FileButton>
);
};
export default memo(LoadGraphButton);

View File

@ -1,79 +0,0 @@
import { FileButton } from '@mantine/core';
import { makeToast } from 'features/system/util/makeToast';
import { useAppDispatch } from 'app/store/storeHooks';
import IAIIconButton from 'common/components/IAIIconButton';
import { loadFileEdges, loadFileNodes } from 'features/nodes/store/nodesSlice';
import { addToast } from 'features/system/store/systemSlice';
import { memo, useCallback, useRef } from 'react';
import { useTranslation } from 'react-i18next';
import { FaUpload } from 'react-icons/fa';
import { useReactFlow } from 'reactflow';
const LoadNodesButton = () => {
const { t } = useTranslation();
const dispatch = useAppDispatch();
const { fitView } = useReactFlow();
const uploadedFileRef = useRef<() => void>(null);
const restoreJSONToEditor = useCallback(
(v: File | null) => {
if (!v) return;
const reader = new FileReader();
reader.onload = async () => {
const json = reader.result;
const retrievedNodeTree = await JSON.parse(String(json));
if (!retrievedNodeTree) {
dispatch(
addToast(
makeToast({
title: t('toast.nodesLoadedFailed'),
status: 'error',
})
)
);
}
if (retrievedNodeTree) {
dispatch(loadFileNodes(retrievedNodeTree.nodes));
dispatch(loadFileEdges(retrievedNodeTree.edges));
fitView();
dispatch(
addToast(
makeToast({ title: t('toast.nodesLoaded'), status: 'success' })
)
);
}
// Cleanup
reader.abort();
};
reader.readAsText(v);
// Cleanup
uploadedFileRef.current?.();
},
[fitView, dispatch, t]
);
return (
<FileButton
resetRef={uploadedFileRef}
accept="application/json"
onChange={restoreJSONToEditor}
>
{(props) => (
<IAIIconButton
icon={<FaUpload />}
tooltip={t('nodes.loadNodes')}
aria-label={t('nodes.loadNodes')}
{...props}
/>
)}
</FileButton>
);
};
export default memo(LoadNodesButton);

View File

@ -6,7 +6,7 @@ import { memo, useCallback } from 'react';
import { useTranslation } from 'react-i18next';
import { FaSave } from 'react-icons/fa';
const SaveNodesButton = () => {
const SaveGraphButton = () => {
const { t } = useTranslation();
const editorInstance = useAppSelector(
(state: RootState) => state.nodes.editorInstance
@ -37,12 +37,12 @@ const SaveNodesButton = () => {
<IAIIconButton
icon={<FaSave />}
fontSize={18}
tooltip={t('nodes.saveNodes')}
aria-label={t('nodes.saveNodes')}
tooltip={t('nodes.saveGraph')}
aria-label={t('nodes.saveGraph')}
onClick={saveEditorToJSON}
isDisabled={nodes.length === 0}
/>
);
};
export default memo(SaveNodesButton);
export default memo(SaveGraphButton);