fix(common): kms correctly serializes transactions

packages/common/package.json

@@ -74,6 +74,7 @@
   "devDependencies": {
     "@types/debug": "^4.1.7",
     "@types/node": "^18.15.11",
+    "@viem/anvil": "^0.0.7",
     "tsup": "^6.7.0",
     "vitest": "0.34.6"
   },

packages/common/src/account/kms/createKmsAccount.test.ts

@@ -1,7 +1,19 @@
 import { describe, it, expect, beforeAll } from "vitest";
 import { KMSAccount, createKmsAccount } from "./createKmsAccount";
 import { CreateKeyCommand, KMSClient } from "@aws-sdk/client-kms";
-import { parseGwei, verifyMessage, verifyTypedData } from "viem";
+import {
+  createPublicClient,
+  parseGwei,
+  createWalletClient,
+  http,
+  parseEther,
+  verifyMessage,
+  verifyTypedData,
+} from "viem";
+import { foundry } from "viem/chains";
+import { privateKeyToAccount } from "viem/accounts";
+import { anvilRpcUrl } from "../test/common";
+import { waitForTransaction } from "./test/waitForTransaction";
 
 describe("createKmsAccount", () => {
   let account: KMSAccount;
@@ -102,4 +114,38 @@ describe("createKmsAccount", () => {
 
     expect(valid).toBeTruthy();
   });
+
+  it("can execute transactions", async () => {
+    const publicClient = createPublicClient({
+      chain: foundry,
+      transport: http(anvilRpcUrl),
+    });
+
+    let balance = await publicClient.getBalance({ address: account.address });
+    expect(balance).toEqual(0n);
+
+    const kmsClient = createWalletClient({
+      chain: foundry,
+      transport: http(anvilRpcUrl),
+      account,
+    });
+
+    const adminClient = createWalletClient({
+      chain: foundry,
+      transport: http(anvilRpcUrl),
+      account: privateKeyToAccount("0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80"),
+    });
+
+    let tx = await adminClient.sendTransaction({ to: account.address, value: parseEther("1") });
+    await waitForTransaction(tx);
+
+    balance = await publicClient.getBalance({ address: account.address });
+    expect(balance).toEqual(1000000000000000000n);
+
+    tx = await kmsClient.sendTransaction({});
+    await waitForTransaction(tx);
+
+    balance = await publicClient.getBalance({ address: account.address });
+    expect(balance).toBeLessThan(1000000000000000000n);
+  });
 });

packages/common/src/account/kms/createKmsAccount.ts

@@ -1,5 +1,5 @@
 import { KMSClient } from "@aws-sdk/client-kms";
-import { LocalAccount, hashMessage, hashTypedData, keccak256, serializeTransaction } from "viem";
+import { LocalAccount, hashMessage, hashTypedData, keccak256, serializeTransaction, signatureToHex } from "viem";
 import { toAccount } from "viem/accounts";
 import { signWithKms } from "./signWithKms";
 import { getAddressFromKms } from "./getAddressFromKms";
@@ -25,38 +25,46 @@ export async function createKmsAccount({
   const account = toAccount({
     address,
     async signMessage({ message }) {
-      const hash = hashMessage(message);
       const signature = await signWithKms({
         client,
         keyId,
-        hash,
+        hash: hashMessage(message),
         address,
       });
 
-      return signature;
+      return signatureToHex(signature);
     },
-    async signTransaction(transaction) {
-      const unsignedTx = serializeTransaction(transaction);
-      const hash = keccak256(unsignedTx);
+    async signTransaction(transaction, { serializer = serializeTransaction } = {}) {
+      // eslint-disable-next-line @typescript-eslint/explicit-function-return-type
+      const signableTransaction = (() => {
+        // For EIP-4844 Transactions, we want to sign the transaction payload body (tx_payload_body) without the sidecars (ie. without the network wrapper).
+        // See: https://github.com/ethereum/EIPs/blob/e00f4daa66bd56e2dbd5f1d36d09fd613811a48b/EIPS/eip-4844.md#networking
+        if (transaction.type === "eip4844")
+          return {
+            ...transaction,
+            sidecars: false,
+          };
+        return transaction;
+      })();
+
       const signature = await signWithKms({
         client,
         keyId,
-        hash,
+        hash: keccak256(serializer(signableTransaction)),
         address,
       });
 
-      return signature;
+      return serializer(transaction, signature);
     },
     async signTypedData(typedData) {
-      const hash = hashTypedData(typedData);
       const signature = await signWithKms({
         client,
         keyId,
-        hash,
+        hash: hashTypedData(typedData),
         address,
       });
 
-      return signature;
+      return signatureToHex(signature);
     },
   });
 

packages/common/src/account/kms/signWithKms.ts

@@ -1,4 +1,4 @@
-import { Hex, isAddressEqual, signatureToHex, toHex } from "viem";
+import { Hex, Signature, isAddressEqual, signatureToHex, toHex } from "viem";
 import { recoverAddress } from "viem/utils";
 import { KMSClient, SignCommandInput } from "@aws-sdk/client-kms";
 import { sign } from "./sign";
@@ -65,7 +65,7 @@ type SignParameters = {
   address: Hex;
 };
 
-type SignReturnType = Hex;
+type SignReturnType = Signature;
 
 /**
  * @description Signs a hash with a given KMS key.
@@ -78,10 +78,10 @@ export async function signWithKms({ hash, address, keyId, client }: SignParamete
   const { r, s } = await getRS({ keyId, hash, client });
   const recovery = await getRecovery(hash, r, s, address);
 
-  return signatureToHex({
+  return {
     r,
     s,
     v: recovery ? 28n : 27n,
     yParity: recovery,
-  });
+  };
 }

packages/common/src/test/minePending.ts

@@ -0,0 +1,9 @@
+import { testClient } from "../../test/common";
+
+export async function minePending(): Promise<void> {
+  const content = await testClient.getTxpoolContent();
+  if (!Object.keys(content.pending).length) return;
+
+  await testClient.mine({ blocks: 1 });
+  await minePending();
+}

packages/common/src/test/waitForTransaction.ts

@@ -0,0 +1,13 @@
+import { Hex } from "viem";
+import { getTransactionReceipt } from "viem/actions";
+import { testClient } from "../../test/common";
+import { minePending } from "./minePending";
+
+export async function waitForTransaction(hash: Hex): Promise<void> {
+  await minePending();
+  const receipt = await getTransactionReceipt(testClient, { hash });
+  if (receipt.status === "reverted") {
+    // TODO: better error
+    throw new Error(`Transaction reverted (${hash})`);
+  }
+}

packages/common/test/common.ts

@@ -0,0 +1,18 @@
+import { createTestClient, http, publicActions, walletActions } from "viem";
+
+export const anvilHost = "127.0.0.1";
+export const anvilPort = 8565;
+
+// ID of the current test worker. Used by the `@viem/anvil` proxy server.
+export const poolId = Number(process.env.VITEST_POOL_ID ?? 1);
+
+export const anvilRpcUrl = `http://${anvilHost}:${anvilPort}/${poolId}`;
+
+export const testClient = createTestClient({
+  mode: "anvil",
+  // TODO: if tests get slow, try switching to websockets?
+  transport: http(anvilRpcUrl),
+  pollingInterval: 10,
+})
+  .extend(publicActions)
+  .extend(walletActions);

packages/common/test/globalSetup.ts

@@ -0,0 +1,13 @@
+import { startProxy as startAnvilProxy } from "@viem/anvil";
+import { anvilHost, anvilPort } from "./common";
+
+export default async function globalSetup(): Promise<() => Promise<void>> {
+  const shutdownAnvilProxy = await startAnvilProxy({
+    host: anvilHost,
+    port: anvilPort,
+  });
+
+  return async () => {
+    await shutdownAnvilProxy();
+  };
+}

packages/common/test/setup.ts

@@ -0,0 +1,18 @@
+import { beforeAll, beforeEach } from "vitest";
+import { testClient } from "./common";
+
+// Some test suites deploy contracts in a `beforeAll` handler, so we restore chain state here.
+beforeAll(async () => {
+  const state = await testClient.dumpState();
+  return async (): Promise<void> => {
+    await testClient.loadState({ state });
+  };
+});
+
+// Some tests execute transactions, so we restore chain state here.
+beforeEach(async () => {
+  const state = await testClient.dumpState();
+  return async (): Promise<void> => {
+    await testClient.loadState({ state });
+  };
+});

packages/common/vitest.config.ts

@@ -0,0 +1,12 @@
+import { defineConfig } from "vitest/config";
+
+export default defineConfig({
+  test: {
+    globalSetup: ["test/globalSetup.ts"],
+    setupFiles: ["test/setup.ts"],
+    // Temporarily set a low teardown timeout because anvil hangs otherwise
+    // Could move this timeout to anvil setup after https://github.com/wevm/anvil.js/pull/46
+    teardownTimeout: 500,
+    hookTimeout: 15000,
+  },
+});