From aee72be7ac343d874832f82c8cbf50bdfc641a35 Mon Sep 17 00:00:00 2001
From: howydev <132113803+howydev@users.noreply.github.com>
Date: Thu, 19 Dec 2024 20:38:44 -0500
Subject: [PATCH] chore: fix signing implementation

---
 .../src/ma-v2/account/nativeSMASigner.ts      | 52 +++++++++-------
 .../src/ma-v2/client/client.test.ts           | 21 ++++---
 .../single-signer-validation/signer.ts        | 60 ++++++++++++-------
 .../smart-contracts/src/ma-v2/utils.ts        |  1 +
 4 files changed, 80 insertions(+), 54 deletions(-)

diff --git a/account-kit/smart-contracts/src/ma-v2/account/nativeSMASigner.ts b/account-kit/smart-contracts/src/ma-v2/account/nativeSMASigner.ts
index ee113c3779..99f4110819 100644
--- a/account-kit/smart-contracts/src/ma-v2/account/nativeSMASigner.ts
+++ b/account-kit/smart-contracts/src/ma-v2/account/nativeSMASigner.ts
@@ -41,22 +41,6 @@ export const nativeSMASigner = <TSigner extends SmartAccountSigner>(
   accountAddress: Address,
   entityId: number
 ) => {
-  const apply712MessageWrap = async (digest: Hex): Promise<Hex> => {
-    return hashTypedData({
-      domain: {
-        chainId: Number(chain.id),
-        verifyingContract: accountAddress,
-      },
-      types: {
-        ReplaySafeHash: [{ name: "digest", type: "bytes32" }],
-      },
-      message: {
-        digest,
-      },
-      primaryType: "ReplaySafeHash",
-    });
-  };
-
   return {
     getDummySignature: (): Hex => {
       const dummyEcdsaSignature =
@@ -83,10 +67,22 @@ export const nativeSMASigner = <TSigner extends SmartAccountSigner>(
     }: {
       message: SignableMessage;
     }): Promise<`0x${string}`> {
-      const digest = await apply712MessageWrap(hashMessage(message));
+      const digest = await hashMessage(message);
 
       return pack1271Signature({
-        validationSignature: await signer.signMessage({ raw: digest }),
+        validationSignature: await signer.signTypedData({
+          domain: {
+            chainId: Number(chain.id),
+            verifyingContract: accountAddress,
+          },
+          types: {
+            ReplaySafeHash: [{ name: "digest", type: "bytes32" }],
+          },
+          message: {
+            digest,
+          },
+          primaryType: "ReplaySafeHash",
+        }),
         entityId,
       });
     },
@@ -102,10 +98,22 @@ export const nativeSMASigner = <TSigner extends SmartAccountSigner>(
       const isDeferredAction =
         typedDataDefinition?.primaryType === "DeferredAction" &&
         typedDataDefinition?.domain?.verifyingContract === accountAddress;
-      const digest = isDeferredAction
-        ? hashTypedData(typedDataDefinition)
-        : await apply712MessageWrap(hashTypedData(typedDataDefinition));
-      return signer.signMessage({ raw: digest });
+
+      return isDeferredAction
+        ? signer.signTypedData(typedDataDefinition)
+        : signer.signTypedData({
+            domain: {
+              chainId: Number(chain.id),
+              verifyingContract: accountAddress,
+            },
+            types: {
+              ReplaySafeHash: [{ name: "digest", type: "bytes32" }],
+            },
+            message: {
+              digest: await hashTypedData(typedDataDefinition),
+            },
+            primaryType: "ReplaySafeHash",
+          });
     },
   };
 };
diff --git a/account-kit/smart-contracts/src/ma-v2/client/client.test.ts b/account-kit/smart-contracts/src/ma-v2/client/client.test.ts
index 3bd3097602..d9976fadc0 100644
--- a/account-kit/smart-contracts/src/ma-v2/client/client.test.ts
+++ b/account-kit/smart-contracts/src/ma-v2/client/client.test.ts
@@ -5,6 +5,7 @@ import {
   getContract,
   keccak256,
   toHex,
+  hashMessage,
 } from "viem";
 import { LocalAccountSigner, type SmartAccountSigner } from "@aa-sdk/core";
 import { createSMAV2AccountClient, type SMAV2AccountClient } from "./client.js";
@@ -21,6 +22,9 @@ describe("MA v2 Tests", async () => {
   let client: ReturnType<typeof instance.getClient> &
     ReturnType<typeof publicActions>;
 
+  const isValidSigSuccess = "0x1626ba7e";
+  const isValidSigFail = "0xffffffff";
+
   beforeAll(async () => {
     client = instance.getClient().extend(publicActions);
   });
@@ -220,22 +224,21 @@ describe("MA v2 Tests", async () => {
       hooks: [],
     });
 
-    const message = keccak256(toHex("testmessage"));
+    const message = "testmessage";
 
+    console.log("outside; ", hashMessage(message));
     const signature = await provider.signMessage({ message });
 
-    console.log(await provider.account.isAccountDeployed());
-
-    console.log(
-      await accountContract.read.isValidSignature([message, signature])
-    );
+    await expect(
+      accountContract.read.isValidSignature([hashMessage(message), signature])
+    ).resolves.toEqual(isValidSigSuccess);
 
     // await expect(
-    //   accountContract.read.isValidSignature({
+    //   accountContract.read.isValidSignature([
     //     message,
     //     signature,
-    //   })
-    // ).resolves.toBeTruthy();
+    //   ])
+    // ).resolves.toEqual(isValidSigSuccess)
   });
 
   const givenConnectedProvider = async ({
diff --git a/account-kit/smart-contracts/src/ma-v2/modules/single-signer-validation/signer.ts b/account-kit/smart-contracts/src/ma-v2/modules/single-signer-validation/signer.ts
index b61feade1a..8bd29cc031 100644
--- a/account-kit/smart-contracts/src/ma-v2/modules/single-signer-validation/signer.ts
+++ b/account-kit/smart-contracts/src/ma-v2/modules/single-signer-validation/signer.ts
@@ -42,23 +42,6 @@ export const singleSignerMessageSigner = <TSigner extends SmartAccountSigner>(
   accountAddress: Address,
   entityId: number
 ) => {
-  const apply712MessageWrap = async (digest: Hex): Promise<Hex> => {
-    return hashTypedData({
-      domain: {
-        chainId: Number(chain.id),
-        salt: accountAddress,
-        verifyingContract: getDefaultSingleSignerValidationModuleAddress(chain),
-      },
-      types: {
-        ReplaySafeHash: [{ name: "digest", type: "bytes32" }],
-      },
-      message: {
-        digest,
-      },
-      primaryType: "ReplaySafeHash",
-    });
-  };
-
   return {
     getDummySignature: (): Hex => {
       const dummyEcdsaSignature =
@@ -85,10 +68,24 @@ export const singleSignerMessageSigner = <TSigner extends SmartAccountSigner>(
     }: {
       message: SignableMessage;
     }): Promise<`0x${string}`> {
-      const digest = await apply712MessageWrap(hashMessage(message));
+      const digest = await hashMessage(message);
 
       return pack1271Signature({
-        validationSignature: await signer.signMessage({ raw: digest }),
+        validationSignature: await signer.signTypedData({
+          domain: {
+            chainId: Number(chain.id),
+            verifyingContract:
+              getDefaultSingleSignerValidationModuleAddress(chain),
+            salt: accountAddress,
+          },
+          types: {
+            ReplaySafeHash: [{ name: "digest", type: "bytes32" }],
+          },
+          message: {
+            digest,
+          },
+          primaryType: "ReplaySafeHash",
+        }),
         entityId,
       });
     },
@@ -100,11 +97,28 @@ export const singleSignerMessageSigner = <TSigner extends SmartAccountSigner>(
     >(
       typedDataDefinition: TypedDataDefinition<typedData, primaryType>
     ): Promise<Hex> => {
-      const digest = await apply712MessageWrap(
-        hashTypedData(typedDataDefinition)
-      );
+      // the accounts domain already gives replay protection across accounts for deferred actions, so we don't need to apply another wrapping
+      const isDeferredAction =
+        typedDataDefinition?.primaryType === "DeferredAction" &&
+        typedDataDefinition?.domain?.verifyingContract === accountAddress;
 
-      return signer.signMessage({ raw: digest });
+      return isDeferredAction
+        ? signer.signTypedData(typedDataDefinition)
+        : signer.signTypedData({
+            domain: {
+              chainId: Number(chain.id),
+              verifyingContract:
+                getDefaultSingleSignerValidationModuleAddress(chain),
+              salt: accountAddress,
+            },
+            types: {
+              ReplaySafeHash: [{ name: "digest", type: "bytes32" }],
+            },
+            message: {
+              digest: await hashTypedData(typedDataDefinition),
+            },
+            primaryType: "ReplaySafeHash",
+          });
     },
   };
 };
diff --git a/account-kit/smart-contracts/src/ma-v2/utils.ts b/account-kit/smart-contracts/src/ma-v2/utils.ts
index 5bcf18fa76..74d58518be 100644
--- a/account-kit/smart-contracts/src/ma-v2/utils.ts
+++ b/account-kit/smart-contracts/src/ma-v2/utils.ts
@@ -42,6 +42,7 @@ export const pack1271Signature = ({
     "0x00",
     toHex(entityId, { size: 4 }),
     "0xFF",
+    "0x00", // EOA type signature
     validationSignature,
   ]);
 };