feat: Add typed union validation

MetaMask · Jul 3, 2024 · 47648d4 · 47648d4
1 parent b7580a7
commit 47648d4
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 6 deletions.
diff --git a/packages/snaps-sdk/src/internals/structs.test.ts b/packages/snaps-sdk/src/internals/structs.test.ts
@@ -1,6 +1,8 @@
 import { is, validate } from 'superstruct';
 
-import { enumValue, literal, union } from './structs';
+import { Text } from '../jsx';
+import { BoxStruct, FieldStruct, TextStruct } from '../jsx/validation';
+import { enumValue, literal, typedUnion, union } from './structs';
 
 describe('enumValue', () => {
   it('validates an enum value', () => {
@@ -38,3 +40,30 @@ describe('literal', () => {
     );
   });
 });
+
+describe('typedUnion', () => {
+  const unionStruct = typedUnion([BoxStruct, TextStruct, FieldStruct]);
+  it('validates strictly the part of the union that matches the type', () => {
+    const result = validate(Text({} as any), unionStruct);
+
+    expect(result[0]?.message).toBe(
+      'At path: props.children -- Expected the value to satisfy a union of `union | array`, but received: undefined',
+    );
+  });
+
+  it('returns an error if the value has no type', () => {
+    const result = validate({}, unionStruct);
+
+    expect(result[0]?.message).toBe(
+      'Expected type to be one of: "Box", "Text", "Field", but received: undefined',
+    );
+  });
+
+  it('returns an error if the type doesnt exist in the union', () => {
+    const result = validate({ type: 'foo' }, unionStruct);
+
+    expect(result[0]?.message).toBe(
+      'Expected type to be one of: "Box", "Text", "Field", but received: "foo"',
+    );
+  });
+});
diff --git a/packages/snaps-sdk/src/internals/structs.ts b/packages/snaps-sdk/src/internals/structs.ts
@@ -1,7 +1,9 @@
+import { assert, hasProperty, isPlainObject } from '@metamask/utils';
 import type { Infer } from 'superstruct';
 import {
   Struct,
   define,
+  is,
   literal as superstructLiteral,
   union as superstructUnion,
 } from 'superstruct';
@@ -76,3 +78,59 @@ export function enumValue<Type extends string>(
 ): Struct<EnumToUnion<Type>, null> {
   return literal(constant as EnumToUnion<Type>);
 }
+
+/**
+ * Create a custom union struct that validates exclusively based on a `type` field.
+ *
+ * This should improve error messaging for unions with many structs in them.
+ *
+ * @param structs - The structs to union.
+ * @returns The `superstruct` struct, which validates that the value satisfies
+ * one of the structs.
+ */
+export function typedUnion<Head extends AnyStruct, Tail extends AnyStruct[]>(
+  structs: [head: Head, ...tail: Tail],
+): Struct<Infer<Head> | InferStructTuple<Tail>[number], null> {
+  return new Struct({
+    type: 'typedUnion',
+    schema: null,
+    *entries(value, context) {
+      if (isPlainObject(value) && hasProperty(value, 'type')) {
+        const { type } = value;
+        const struct = structs.find(({ schema }) => is(type, schema.type));
+
+        assert(struct, 'Expected at least one struct to match');
+
+        for (const entry of struct.entries(value, context)) {
+          yield entry;
+        }
+      }
+    },
+    validator(value, context) {
+      const types = structs.map(({ schema }) => schema.type.type);
+
+      if (
+        !isPlainObject(value) ||
+        !hasProperty(value, 'type') ||
+        typeof value.type !== 'string'
+      ) {
+        return `Expected type to be one of: ${types.join(
+          ', ',
+        )}, but received: undefined`;
+      }
+
+      const { type } = value;
+
+      const struct = structs.find(({ schema }) => is(type, schema.type));
+
+      if (struct) {
+        // This only validates the root of the struct, entries does the rest of the work.
+        return struct.validator(value, context);
+      }
+
+      return `Expected type to be one of: ${types.join(
+        ', ',
+      )}, but received: "${type}"`;
+    },
+  }) as unknown as Struct<Infer<Head> | InferStructTuple<Tail>[number], null>;
+}
diff --git a/packages/snaps-sdk/src/jsx/validation.ts b/packages/snaps-sdk/src/jsx/validation.ts
@@ -25,7 +25,7 @@ import type {
 } from 'superstruct/dist/utils';
 
 import type { Describe } from '../internals';
-import { literal, nullUnion, svg } from '../internals';
+import { literal, nullUnion, svg, typedUnion } from '../internals';
 import type { EmptyObject } from '../types';
 import type {
   GenericSnapElement,
@@ -480,7 +480,7 @@ export const SpinnerStruct: Describe<SpinnerElement> = element('Spinner');
  * This set includes all components, except components that need to be nested in
  * another component (e.g., Field must be contained in a Form).
  */
-export const BoxChildStruct = nullUnion([
+export const BoxChildStruct = typedUnion([
   AddressStruct,
   BoldStruct,
   BoxStruct,
@@ -515,7 +515,7 @@ export const RootJSXElementStruct = nullUnion([
 /**
  * A struct for the {@link JSXElement} type.
  */
-export const JSXElementStruct: Describe<JSXElement> = nullUnion([
+export const JSXElementStruct: Describe<JSXElement> = typedUnion([
   ButtonStruct,
   InputStruct,
   FileInputStruct,

diff --git a/packages/snaps-sdk/src/ui/components/panel.ts b/packages/snaps-sdk/src/ui/components/panel.ts
@@ -1,6 +1,7 @@
 import type { Infer, Struct } from 'superstruct';
-import { array, assign, lazy, literal, object, union } from 'superstruct';
+import { array, assign, lazy, literal, object } from 'superstruct';
 
+import { typedUnion } from '../../internals';
 import { createBuilder } from '../builder';
 import { NodeStruct, NodeType } from '../nodes';
 import { AddressStruct } from './address';
@@ -86,7 +87,7 @@ export type Panel = {
 export const panel = createBuilder(NodeType.Panel, PanelStruct, ['children']);
 
 // This is defined separately from `Component` to avoid circular dependencies.
-export const ComponentStruct = union([
+export const ComponentStruct = typedUnion([
   CopyableStruct,
   DividerStruct,
   HeadingStruct,