From 7a82d43eb8fd78bfebd2a6fb4816dd7ba13965fe Mon Sep 17 00:00:00 2001
From: Matt Schoch <matt@narval.xyz>
Date: Thu, 28 Mar 2024 10:42:35 -0400
Subject: [PATCH] Moving core Action types to Zod schema as source of truth
 (#197)

---
 packages/policy-engine-shared/src/index.ts    |  1 -
 .../src/lib/type/action.type.ts               | 91 +++++++++++++++----
 .../src/lib/type/domain.type.ts               | 19 ----
 3 files changed, 71 insertions(+), 40 deletions(-)

diff --git a/packages/policy-engine-shared/src/index.ts b/packages/policy-engine-shared/src/index.ts
index 9b53dee13..eaac8f8fd 100644
--- a/packages/policy-engine-shared/src/index.ts
+++ b/packages/policy-engine-shared/src/index.ts
@@ -1,4 +1,3 @@
-export * from './lib/schema/action.schema'
 export * from './lib/schema/address.schema'
 export * from './lib/schema/data-store.schema'
 export * from './lib/schema/domain.schema'
diff --git a/packages/policy-engine-shared/src/lib/type/action.type.ts b/packages/policy-engine-shared/src/lib/type/action.type.ts
index 5698bb7f6..e08d45b45 100644
--- a/packages/policy-engine-shared/src/lib/type/action.type.ts
+++ b/packages/policy-engine-shared/src/lib/type/action.type.ts
@@ -1,5 +1,7 @@
-import { Hex } from 'viem'
-import { Address, JwtString, TransactionRequest } from './domain.type'
+import { z } from 'zod'
+import { addressSchema } from '../schema/address.schema'
+import { hexSchema } from '../schema/hex.schema'
+import { Address, JwtString } from './domain.type'
 import {
   AccountClassification,
   AccountType,
@@ -67,26 +69,75 @@ export type BaseAdminRequest = {
   approvals: JwtString[]
 }
 
-export type SignTransactionAction = BaseAction & {
-  action: typeof Action.SIGN_TRANSACTION
-  resourceId: string
-  transactionRequest: TransactionRequest
-}
+export const AccessList = z.array(
+  z.object({
+    address: addressSchema,
+    storageKeys: z.array(hexSchema)
+  })
+)
+export type AccessList = z.infer<typeof AccessList>
+
+export const ActionSchema = z.nativeEnum(Action)
+
+export const BaseActionSchema = z.object({
+  action: ActionSchema,
+  nonce: z.string()
+})
+export type BaseActionSchema = z.infer<typeof BaseActionSchema>
+
+export const TransactionRequest = z.object({
+  chainId: z.number(),
+  from: addressSchema,
+  nonce: z.number().optional(),
+  accessList: AccessList.optional(),
+  data: hexSchema.optional(),
+  gas: z.coerce.bigint().optional(),
+  maxFeePerGas: z.coerce.bigint().optional(),
+  maxPriorityFeePerGas: z.coerce.bigint().optional(),
+  to: addressSchema.nullable().optional(),
+  type: z.literal('2').optional(),
+  value: hexSchema.optional()
+})
+export type TransactionRequest = z.infer<typeof TransactionRequest>
+
+export const SignTransactionAction = z.intersection(
+  BaseActionSchema,
+  z.object({
+    action: z.literal(Action.SIGN_TRANSACTION),
+    resourceId: z.string(),
+    transactionRequest: TransactionRequest
+  })
+)
+export type SignTransactionAction = z.infer<typeof SignTransactionAction>
 
 // Matching viem's SignableMessage options https://viem.sh/docs/actions/wallet/signMessage#message
-export type SignableMessage = string | { raw: Hex }
-
-export type SignMessageAction = BaseAction & {
-  action: typeof Action.SIGN_MESSAGE
-  resourceId: string
-  message: SignableMessage
-}
-
-export type SignTypedDataAction = BaseAction & {
-  action: typeof Action.SIGN_TYPED_DATA
-  resourceId: string
-  typedData: string
-}
+export const SignableMessage = z.union([
+  z.string(),
+  z.object({
+    raw: hexSchema
+  })
+])
+export type SignableMessage = z.infer<typeof SignableMessage>
+
+export const SignMessageAction = z.intersection(
+  BaseActionSchema,
+  z.object({
+    action: z.literal(Action.SIGN_MESSAGE),
+    resourceId: z.string(),
+    message: SignableMessage
+  })
+)
+export type SignMessageAction = z.infer<typeof SignMessageAction>
+
+export const SignTypedDataAction = z.intersection(
+  BaseActionSchema,
+  z.object({
+    action: z.literal(Action.SIGN_TYPED_DATA),
+    resourceId: z.string(),
+    typedData: z.string()
+  })
+)
+export type SignTypedDataAction = z.infer<typeof SignTypedDataAction>
 
 export type CreateOrganizationAction = BaseAction & {
   action: typeof Action.CREATE_ORGANIZATION
diff --git a/packages/policy-engine-shared/src/lib/type/domain.type.ts b/packages/policy-engine-shared/src/lib/type/domain.type.ts
index 1f7a7206f..622333397 100644
--- a/packages/policy-engine-shared/src/lib/type/domain.type.ts
+++ b/packages/policy-engine-shared/src/lib/type/domain.type.ts
@@ -165,22 +165,3 @@ export type EvaluationResponse = {
 export type Hex = `0x${string}` // DOMAIN
 
 export type Address = `0x${string}` // DOMAIN
-
-export type AccessList = {
-  address: Address
-  storageKeys: Hex[]
-}[] // DOMAIN
-
-export type TransactionRequest = {
-  chainId: number
-  from: Address
-  nonce?: number
-  accessList?: AccessList
-  data?: Hex
-  gas?: bigint
-  maxFeePerGas?: bigint
-  maxPriorityFeePerGas?: bigint
-  to?: Address | null
-  type?: '2'
-  value?: Hex
-}