refactor: change how the middleware is defined (#6)

This refactors the middleware overriders to be more specific about the fields that they are allowed to override.

Before, the middleware overriders let you do arbitrary transformations on the UO struct using the overriders.

packages/core/src/middleware/alchemy-paymaster.ts

@@ -2,6 +2,7 @@ import type { Address, Hex } from "viem";
 import type { PublicErc4337Client } from "../client/types.js";
 import type { UserOperationRequest, UserOperationStruct } from "../types.js";
 import { deepHexlify, resolveProperties } from "../utils.js";
+import type { ISmartAccountProvider } from "../provider/types.js";
 
 type ClientWithAlchemyMethod = PublicErc4337Client & {
   request: PublicErc4337Client["request"] &
@@ -27,13 +28,14 @@ export interface AlchemyPaymasterConfig {
 
 export const alchemyPaymasterAndDataMiddleware = (
   config: AlchemyPaymasterConfig
-) => ({
-  dummyPaymasterMiddleware: async (struct: UserOperationStruct) => {
-    struct.paymasterAndData =
-      "0xc03aac639bb21233e0139381970328db8bceeb67fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff0000000000000000000000000000000007aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1c";
-    return struct;
+): Parameters<ISmartAccountProvider["withPaymasterMiddleware"]>["0"] => ({
+  dummyPaymasterDataMiddleware: async (_struct: UserOperationStruct) => {
+    return {
+      paymasterAndData:
+        "0xc03aac639bb21233e0139381970328db8bceeb67fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff0000000000000000000000000000000007aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa1c",
+    };
   },
-  getPaymasterAndDataMiddleware: async (struct: UserOperationStruct) => {
+  paymasterDataMiddleware: async (struct: UserOperationStruct) => {
     const { paymasterAndData } = await (
       config.provider as ClientWithAlchemyMethod
     ).request({
@@ -47,7 +49,6 @@ export const alchemyPaymasterAndDataMiddleware = (
       ],
     });
 
-    struct.paymasterAndData = paymasterAndData;
-    return struct;
+    return { paymasterAndData };
   },
 });

packages/core/src/provider/base.ts

@@ -21,7 +21,11 @@ import {
 } from "../utils.js";
 import type {
   AccountMiddlewareFn,
+  AccountMiddlewareOverrideFn,
+  FeeDataMiddleware,
+  GasEstimatorMiddleware,
   ISmartAccountProvider,
+  PaymasterAndDataMiddleware,
   SendUserOperationResult,
 } from "./types.js";
 
@@ -174,7 +178,7 @@ export class SmartAccountProvider<
       this.dummyPaymasterDataMiddleware,
       this.gasEstimator,
       this.feeDataGetter,
-      this.paymasterMiddleware
+      this.paymasterDataMiddleware
     )({
       initCode,
       sender: this.getAddress(),
@@ -222,11 +226,10 @@ export class SmartAccountProvider<
     return struct;
   };
 
-  readonly paymasterMiddleware: AccountMiddlewareFn = async (
+  readonly paymasterDataMiddleware: AccountMiddlewareFn = async (
     struct: UserOperationStruct
   ): Promise<UserOperationStruct> => {
     struct.paymasterAndData = "0x";
-
     return struct;
   };
 
@@ -271,30 +274,37 @@ export class SmartAccountProvider<
   };
 
   withPaymasterMiddleware = (overrides: {
-    dummyPaymasterMiddleware?: AccountMiddlewareFn;
-    getPaymasterAndDataMiddleware?: AccountMiddlewareFn;
+    dummyPaymasterDataMiddleware?: PaymasterAndDataMiddleware;
+    paymasterDataMiddleware?: PaymasterAndDataMiddleware;
   }): this => {
-    defineReadOnly(
-      this,
-      "dummyPaymasterDataMiddleware",
-      overrides.dummyPaymasterMiddleware ?? this.dummyPaymasterDataMiddleware
-    );
-    defineReadOnly(
-      this,
-      "paymasterMiddleware",
-      overrides.getPaymasterAndDataMiddleware ?? this.paymasterMiddleware
-    );
+    const newDummyMiddleware = overrides.dummyPaymasterDataMiddleware
+      ? this.overrideMiddlewareFunction(overrides.dummyPaymasterDataMiddleware)
+      : this.dummyPaymasterDataMiddleware;
+    defineReadOnly(this, "dummyPaymasterDataMiddleware", newDummyMiddleware);
+
+    const newPaymasterMiddleware = overrides.paymasterDataMiddleware
+      ? this.overrideMiddlewareFunction(overrides.paymasterDataMiddleware)
+      : this.paymasterDataMiddleware;
+    defineReadOnly(this, "paymasterDataMiddleware", newPaymasterMiddleware);
 
     return this;
   };
 
-  withGasEstimator = (override: AccountMiddlewareFn): this => {
-    defineReadOnly(this, "gasEstimator", override);
+  withGasEstimator = (override: GasEstimatorMiddleware): this => {
+    defineReadOnly(
+      this,
+      "gasEstimator",
+      this.overrideMiddlewareFunction(override)
+    );
     return this;
   };
 
-  withFeeDataGetter = (override: AccountMiddlewareFn): this => {
-    defineReadOnly(this, "feeDataGetter", override);
+  withFeeDataGetter = (override: FeeDataMiddleware): this => {
+    defineReadOnly(
+      this,
+      "feeDataGetter",
+      this.overrideMiddlewareFunction(override)
+    );
     return this;
   };
 
@@ -305,4 +315,15 @@ export class SmartAccountProvider<
     defineReadOnly(this, "account", account);
     return this as this & { account: typeof account };
   }
+
+  private overrideMiddlewareFunction = (
+    override: AccountMiddlewareOverrideFn
+  ): AccountMiddlewareFn => {
+    return async (struct: UserOperationStruct) => {
+      return {
+        ...struct,
+        ...(await override(struct)),
+      };
+    };
+  };
 }

packages/core/src/provider/types.ts

@@ -16,13 +16,29 @@ export type AccountMiddlewareFn = (
   struct: UserOperationStruct
 ) => Promise<UserOperationStruct>;
 
+export type AccountMiddlewareOverrideFn<
+  K extends keyof UserOperationStruct = never
+> = (
+  struct: UserOperationStruct
+) => Promise<Required<Pick<UserOperationStruct, K>>>;
+
+export type PaymasterAndDataMiddleware =
+  AccountMiddlewareOverrideFn<"paymasterAndData">;
+
+export type GasEstimatorMiddleware = AccountMiddlewareOverrideFn<
+  "callGasLimit" | "preVerificationGas" | "verificationGasLimit"
+>;
+export type FeeDataMiddleware = AccountMiddlewareOverrideFn<
+  "maxFeePerGas" | "maxPriorityFeePerGas"
+>;
+
 // TODO: this also will need to implement EventEmitteer
 export interface ISmartAccountProvider<
   TTransport extends SupportedTransports = Transport
 > {
   readonly rpcClient: PublicErc4337Client<TTransport>;
   readonly dummyPaymasterDataMiddleware: AccountMiddlewareFn;
-  readonly paymasterMiddleware: AccountMiddlewareFn;
+  readonly paymasterDataMiddleware: AccountMiddlewareFn;
   readonly gasEstimator: AccountMiddlewareFn;
   readonly feeDataGetter: AccountMiddlewareFn;
 
@@ -85,8 +101,8 @@ export interface ISmartAccountProvider<
    * @returns an update instance of this, which now uses the new middleware
    */
   withPaymasterMiddleware: (overrides: {
-    dummyPaymasterMiddleware?: AccountMiddlewareFn;
-    getPaymasterAndDataMiddleware?: AccountMiddlewareFn;
+    dummyPaymasterDataMiddleware?: PaymasterAndDataMiddleware;
+    paymasterDataMiddleware?: PaymasterAndDataMiddleware;
   }) => this;
 
   /**
@@ -96,7 +112,7 @@ export interface ISmartAccountProvider<
    * @param override - a function for overriding the default gas estimator middleware
    * @returns
    */
-  withGasEstimator: (override: AccountMiddlewareFn) => this;
+  withGasEstimator: (override: GasEstimatorMiddleware) => this;
 
   /**
    * Overrides the feeDataGetter middleware which is used for setting the fee fields on the UserOperation
@@ -105,7 +121,7 @@ export interface ISmartAccountProvider<
    * @param override - a function for overriding the default feeDataGetter middleware
    * @returns
    */
-  withFeeDataGetter: (override: AccountMiddlewareFn) => this;
+  withFeeDataGetter: (override: FeeDataMiddleware) => this;
 
   /**
    * Sets the current account to the account returned by the given function. The function parameter

packages/core/src/utils.ts

@@ -32,7 +32,7 @@ export const asyncPipe =
     for (const fn of fns) {
       result = await fn(result);
     }
-    return x;
+    return result;
   };
 
 // based on @ethersproject/properties, but pulled in here to minize the dependency on ethers