feat: add zod runtime validation for base account (#186)

* feat: add zod runtime validation for base account * feat: add zod runtime validation for base account * feat: add zod runtime validation for simple account * refactor: clean up base schemas * refactor: rebase * refactor: rename abitype import
alchemyplatform · Nov 3, 2023 · e9c48ed · e9c48ed
1 parent 1e1883b
commit e9c48ed
Show file tree

Hide file tree

Showing 15 changed files with 179 additions and 80 deletions.
diff --git a/packages/core/src/account/__tests__/simple.test.ts b/packages/core/src/account/__tests__/simple.test.ts
@@ -1,3 +1,4 @@
+import type { Address } from "viem";
 import { polygonMumbai, type Chain } from "viem/chains";
 import { describe, it } from "vitest";
 import { getDefaultSimpleAccountFactoryAddress } from "../../index.js";
@@ -43,6 +44,64 @@ describe("Account Simple Tests", () => {
       '"0x18dfb3c7000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000000a00000000000000000000000000000000000000000000000000000000000000002000000000000000000000000deadbeefdeadbeefdeadbeefdeadbeefdeadbeef0000000000000000000000008ba1f109551bd432803012645ac136ddd64dba720000000000000000000000000000000000000000000000000000000000000002000000000000000000000000000000000000000000000000000000000000004000000000000000000000000000000000000000000000000000000000000000800000000000000000000000000000000000000000000000000000000000000004deadbeef000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000004cafebabe00000000000000000000000000000000000000000000000000000000"'
     );
   });
+
+  it("should correctly do base runtime validation when entrypoint are invalid", () => {
+    expect(
+      () =>
+        new SimpleSmartContractAccount({
+          entryPointAddress: 1 as unknown as Address,
+          chain,
+          owner,
+          factoryAddress: "0xdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef",
+          rpcClient: "ALCHEMY_RPC_URL",
+        })
+    ).toThrowErrorMatchingInlineSnapshot(`
+      "[
+        {
+          \\"code\\": \\"invalid_type\\",
+          \\"expected\\": \\"string\\",
+          \\"received\\": \\"number\\",
+          \\"path\\": [
+            \\"entryPointAddress\\"
+          ],
+          \\"message\\": \\"Expected string, received number\\"
+        }
+      ]"
+    `);
+  });
+
+  it("should correctly do base runtime validation when multiple inputs are invalid", () => {
+    expect(
+      () =>
+        new SimpleSmartContractAccount({
+          entryPointAddress: 1 as unknown as Address,
+          chain: "0x1" as unknown as Chain,
+          owner,
+          factoryAddress: "0xdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef",
+          rpcClient: "ALCHEMY_RPC_URL",
+        })
+    ).toThrowErrorMatchingInlineSnapshot(`
+      "[
+        {
+          \\"code\\": \\"invalid_type\\",
+          \\"expected\\": \\"string\\",
+          \\"received\\": \\"number\\",
+          \\"path\\": [
+            \\"entryPointAddress\\"
+          ],
+          \\"message\\": \\"Expected string, received number\\"
+        },
+        {
+          \\"code\\": \\"custom\\",
+          \\"fatal\\": true,
+          \\"path\\": [
+            \\"chain\\"
+          ],
+          \\"message\\": \\"Invalid input\\"
+        }
+      ]"
+    `);
+  });
 });
 
 const givenConnectedProvider = ({

diff --git a/packages/core/src/account/base.ts b/packages/core/src/account/base.ts
@@ -20,39 +20,19 @@ import type { SmartAccountSigner } from "../signer/types.js";
 import { wrapSignatureWith6492 } from "../signer/utils.js";
 import type { BatchUserOperationCallData } from "../types.js";
 import { getDefaultEntryPointAddress } from "../utils/defaults.js";
-import type { ISmartContractAccount, SignTypedDataParams } from "./types.js";
+import { createBaseSmartAccountParamsSchema } from "./schema.js";
+import type {
+  BaseSmartAccountParams,
+  ISmartContractAccount,
+  SignTypedDataParams,
+} from "./types.js";
 
 export enum DeploymentState {
   UNDEFINED = "0x0",
   NOT_DEPLOYED = "0x1",
   DEPLOYED = "0x2",
 }
 
-export interface BaseSmartAccountParams<
-  TTransport extends SupportedTransports = Transport
-> {
-  rpcClient: string | PublicErc4337Client<TTransport>;
-  factoryAddress: Address;
-  chain: Chain;
-
-  /**
-   * The address of the entry point contract.
-   * If not provided, the default entry point contract will be used.
-   * Check out https://docs.alchemy.com/reference/eth-supportedentrypoints for all the supported entrypoints
-   */
-  entryPointAddress?: Address;
-
-  /**
-   * Owner account signer for the account if there is one.
-   */
-  owner?: SmartAccountSigner | undefined;
-
-  /**
-   * The address of the account if it is already deployed.
-   */
-  accountAddress?: Address;
-}
-
 export abstract class BaseSmartContractAccount<
   TTransport extends SupportedTransports = Transport
 > implements ISmartContractAccount
@@ -72,6 +52,8 @@ export abstract class BaseSmartContractAccount<
     | PublicErc4337Client<HttpTransport>;
 
   constructor(params: BaseSmartAccountParams<TTransport>) {
+    createBaseSmartAccountParamsSchema<TTransport>().parse(params);
+
     this.entryPointAddress =
       params.entryPointAddress ?? getDefaultEntryPointAddress(params.chain);
 

diff --git a/packages/core/src/account/schema.ts b/packages/core/src/account/schema.ts
@@ -0,0 +1,22 @@
+import { Address } from "abitype/zod";
+import type { Transport } from "viem";
+import z from "zod";
+import { createPublicErc4337ClientSchema } from "../client/schema.js";
+import type { SupportedTransports } from "../client/types";
+import { SignerSchema } from "../signer/schema.js";
+import { ChainSchema } from "../utils/index.js";
+
+export const createBaseSmartAccountParamsSchema = <
+  TTransport extends SupportedTransports = Transport
+>() =>
+  z.object({
+    rpcClient: z.union([
+      z.string(),
+      createPublicErc4337ClientSchema<TTransport>(),
+    ]),
+    factoryAddress: Address,
+    owner: SignerSchema.optional(),
+    entryPointAddress: Address.optional(),
+    chain: ChainSchema,
+    accountAddress: Address.optional(),
+  });
diff --git a/packages/core/src/account/simple.ts b/packages/core/src/account/simple.ts
@@ -9,12 +9,10 @@ import {
 } from "viem";
 import { SimpleAccountAbi } from "../abis/SimpleAccountAbi.js";
 import { SimpleAccountFactoryAbi } from "../abis/SimpleAccountFactoryAbi.js";
-import type { BatchUserOperationCallData } from "../types.js";
-import {
-  BaseSmartContractAccount,
-  type BaseSmartAccountParams,
-} from "./base.js";
 import type { SmartAccountSigner } from "../signer/types.js";
+import type { BatchUserOperationCallData } from "../types.js";
+import { BaseSmartContractAccount } from "./base.js";
+import type { BaseSmartAccountParams } from "./types.js";
 
 export interface SimpleSmartAccountParams<
   TTransport extends Transport | FallbackTransport = Transport

diff --git a/packages/core/src/account/types.ts b/packages/core/src/account/types.ts
@@ -1,11 +1,18 @@
 import type { Address } from "abitype";
-import type { Hash, Hex } from "viem";
+import type { Hash, Hex, Transport } from "viem";
 import type { SignTypedDataParameters } from "viem/accounts";
+import type { z } from "zod";
+import type { SupportedTransports } from "../client/types";
 import type { SmartAccountSigner } from "../signer/types";
 import type { BatchUserOperationCallData } from "../types";
+import type { createBaseSmartAccountParamsSchema } from "./schema";
 
 export type SignTypedDataParams = Omit<SignTypedDataParameters, "privateKey">;
 
+export type BaseSmartAccountParams<
+  TTransport extends SupportedTransports = Transport
+> = z.infer<ReturnType<typeof createBaseSmartAccountParamsSchema<TTransport>>>;
+
 export interface ISmartContractAccount {
   /**
    * @returns the init code for the account

diff --git a/packages/core/src/client/schema.ts b/packages/core/src/client/schema.ts
@@ -0,0 +1,17 @@
+import type { Transport } from "viem";
+import { z } from "zod";
+import type { PublicErc4337Client, SupportedTransports } from "./types";
+
+export const createPublicErc4337ClientSchema = <
+  TTransport extends SupportedTransports = Transport
+>() =>
+  z.custom<PublicErc4337Client<TTransport>>((provider) => {
+    return (
+      provider != null &&
+      typeof provider === "object" &&
+      "request" in provider &&
+      "type" in provider &&
+      "key" in provider &&
+      "name" in provider
+    );
+  });
diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
@@ -7,11 +7,14 @@ export { SimpleAccountAbi } from "./abis/SimpleAccountAbi.js";
 export { SimpleAccountFactoryAbi } from "./abis/SimpleAccountFactoryAbi.js";
 
 export { BaseSmartContractAccount } from "./account/base.js";
-export type { BaseSmartAccountParams } from "./account/base.js";
+export { createBaseSmartAccountParamsSchema } from "./account/schema.js";
 export { SimpleSmartContractAccount } from "./account/simple.js";
 export type { SimpleSmartAccountParams } from "./account/simple.js";
 export type * from "./account/types.js";
+export type { BaseSmartAccountParams } from "./account/types.js";
+
 export { LocalAccountSigner } from "./signer/local-account.js";
+export { SignerSchema } from "./signer/schema.js";
 export type { SmartAccountSigner } from "./signer/types.js";
 export {
   verifyEIP6492Signature,
@@ -24,6 +27,7 @@ export {
   createPublicErc4337FromClient,
   erc4337ClientActions,
 } from "./client/create-client.js";
+export { createPublicErc4337ClientSchema } from "./client/schema.js";
 export type * from "./client/types.js";
 
 export {
@@ -33,6 +37,10 @@ export {
 } from "./ens/utils.js";
 
 export { SmartAccountProvider, noOpMiddleware } from "./provider/base.js";
+export {
+  createSmartAccountProviderConfigSchema,
+  SmartAccountProviderOptsSchema,
+} from "./provider/schema.js";
 export type * from "./provider/types.js";
 
 export type * from "./types.js";
@@ -48,6 +56,7 @@ export {
   getDefaultSimpleAccountFactoryAddress,
   getUserOperationHash,
   resolveProperties,
+  ChainSchema,
 } from "./utils/index.js";
 
 export { Logger } from "./logger.js";

diff --git a/packages/core/src/provider/__tests__/base.test.ts b/packages/core/src/provider/__tests__/base.test.ts
@@ -244,10 +244,11 @@ describe("Base Tests", () => {
       "[
         {
           \\"code\\": \\"custom\\",
-          \\"message\\": \\"Invalid input\\",
+          \\"fatal\\": true,
           \\"path\\": [
             \\"chain\\"
-          ]
+          ],
+          \\"message\\": \\"Invalid input\\"
         },
         {
           \\"code\\": \\"invalid_type\\",

diff --git a/packages/core/src/provider/base.ts b/packages/core/src/provider/base.ts
@@ -41,7 +41,7 @@ import {
   resolveProperties,
   type Deferrable,
 } from "../utils/index.js";
-import { SmartAccountProviderConfigSchema } from "./schema.js";
+import { createSmartAccountProviderConfigSchema } from "./schema.js";
 import type {
   AccountMiddlewareFn,
   AccountMiddlewareOverrideFn,
@@ -85,7 +85,7 @@ export class SmartAccountProvider<
     | PublicErc4337Client<HttpTransport>;
 
   constructor(config: SmartAccountProviderConfig<TTransport>) {
-    SmartAccountProviderConfigSchema<TTransport>().parse(config);
+    createSmartAccountProviderConfigSchema<TTransport>().parse(config);
 
     const { rpcProvider, entryPointAddress, chain, opts } = config;
 

diff --git a/packages/core/src/provider/schema.ts b/packages/core/src/provider/schema.ts
@@ -1,8 +1,9 @@
-import { Address as zAddress } from "abitype/zod";
-import type { Chain, Transport } from "viem";
+import { Address } from "abitype/zod";
+import type { Transport } from "viem";
 import z from "zod";
-import type { PublicErc4337Client, SupportedTransports } from "../client/types";
-import { getChain } from "../utils/index.js";
+import { createPublicErc4337ClientSchema } from "../client/schema.js";
+import type { SupportedTransports } from "../client/types";
+import { ChainSchema } from "../utils/index.js";
 
 export const SmartAccountProviderOptsSchema = z.object({
   /**
@@ -26,43 +27,15 @@ export const SmartAccountProviderOptsSchema = z.object({
   minPriorityFeePerBid: z.bigint().optional().default(100_000_000n),
 });
 
-export const SmartAccountProviderConfigSchema = <
+export const createSmartAccountProviderConfigSchema = <
   TTransport extends SupportedTransports = Transport
 >() => {
   return z.object({
     rpcProvider: z.union([
       z.string(),
-      z
-        .any()
-        .refine<PublicErc4337Client<TTransport>>(
-          (provider): provider is PublicErc4337Client<TTransport> => {
-            return (
-              typeof provider === "object" &&
-              "request" in provider &&
-              "type" in provider &&
-              "key" in provider &&
-              "name" in provider
-            );
-          }
-        ),
+      createPublicErc4337ClientSchema<TTransport>(),
     ]),
-
-    chain: z.any().refine<Chain>((chain): chain is Chain => {
-      if (
-        !(typeof chain === "object") ||
-        !("id" in chain) ||
-        typeof chain.id !== "number"
-      ) {
-        return false;
-      }
-
-      try {
-        return getChain(chain.id) !== undefined;
-      } catch {
-        return false;
-      }
-    }),
-
+    chain: ChainSchema,
     /**
      * Optional entry point contract address for override if needed.
      * If not provided, the entry point contract address for the provider is the connected account's entry point contract,
@@ -71,8 +44,7 @@ export const SmartAccountProviderConfigSchema = <
      * Refer to https://docs.alchemy.com/reference/eth-supportedentrypoints for all the supported entrypoints
      * when using Alchemy as your RPC provider.
      */
-    entryPointAddress: zAddress.optional(),
-
+    entryPointAddress: Address.optional(),
     opts: SmartAccountProviderOptsSchema.optional(),
   });
 };
diff --git a/packages/core/src/provider/types.ts b/packages/core/src/provider/types.ts
@@ -26,8 +26,8 @@ import type {
 } from "../types.js";
 import type { Deferrable } from "../utils";
 import type {
-  SmartAccountProviderConfigSchema,
   SmartAccountProviderOptsSchema,
+  createSmartAccountProviderConfigSchema,
 } from "./schema.js";
 
 type WithRequired<T, K extends keyof T> = Required<Pick<T, K>>;
@@ -87,7 +87,9 @@ export type SmartAccountProviderOpts = z.infer<
 
 export type SmartAccountProviderConfig<
   TTransport extends SupportedTransports = Transport
-> = z.infer<ReturnType<typeof SmartAccountProviderConfigSchema<TTransport>>>;
+> = z.infer<
+  ReturnType<typeof createSmartAccountProviderConfigSchema<TTransport>>
+>;
 
 // TODO: this also will need to implement EventEmitteer
 export interface ISmartAccountProvider<

diff --git a/packages/core/src/signer/schema.ts b/packages/core/src/signer/schema.ts
@@ -0,0 +1,13 @@
+import { z } from "zod";
+import type { SmartAccountSigner } from "./types";
+
+export const SignerSchema = z.custom<SmartAccountSigner>((signer) => {
+  return (
+    signer != null &&
+    typeof signer === "object" &&
+    "signerType" in signer &&
+    "signMessage" in signer &&
+    "signTypedData" in signer &&
+    "getAddress" in signer
+  );
+});
diff --git a/packages/core/src/utils/index.ts b/packages/core/src/utils/index.ts
@@ -158,4 +158,5 @@ export function defineReadOnly<T, K extends keyof T>(
 
 export * from "./bigint.js";
 export * from "./defaults.js";
+export * from "./schema.js";
 export * from "./userop.js";