fix(common): kms correctly serializes transactions (#2721)

.changeset/perfect-actors-flow.md

@@ -2,13 +2,13 @@
 "@latticexyz/common": patch
 ---
 
-Added `createKmsAccount`, a [viem custom account](https://viem.sh/docs/accounts/custom#custom-account) that signs transactions using AWS KMS.
+Added `kmsKeyToAccount`, a [viem custom account](https://viem.sh/docs/accounts/custom#custom-account) that signs transactions using AWS KMS.
 
 To use it, you must first install `@aws-sdk/[email protected]` and `[email protected]` dependencies into your project. Then create a KMS account with:
 
 ```ts
-import { createKmsAccount } from "@latticexyz/common/kms";
-const account = createKmsAccount({ keyId: ... });
+import { kmsKeyToAccount } from "@latticexyz/common/kms";
+const account = kmsKeyToAccount({ keyId: ... });
 ```
 
 By default, a `KMSClient` will be created, but you can also pass one in via the `client` option. The default KMS client will use [your environment's AWS SDK configuration](https://docs.aws.amazon.com/sdk-for-javascript/v3/developer-guide/configuring-the-jssdk.html).

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/README.md

@@ -1,11 +1,11 @@
 # KMS Custom Account
 
-`createKmsAccount` is a [viem custom account](https://viem.sh/docs/accounts/custom#custom-account) that signs transactions using AWS KMS.
+`kmsKeyToAccount` is a [viem custom account](https://viem.sh/docs/accounts/custom#custom-account) that signs transactions using AWS KMS.
 
 To use it, you must first install `@aws-sdk/[email protected]` and `[email protected]` dependencies into your project. Then create a KMS account with:
 
 ```ts
-const account = createKmsAccount({ keyId: ... });
+const account = kmsKeyToAccount({ keyId: ... });
 ```
 
 By default, a `KMSClient` will be created, but you can also pass one in via the `client` option. The default KMS client will use [your environment's AWS SDK configuration](https://docs.aws.amazon.com/sdk-for-javascript/v3/developer-guide/configuring-the-jssdk.html).

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

@@ -1,10 +1,14 @@
 import { describe, it, expect, beforeAll } from "vitest";
-import { KMSAccount, createKmsAccount } from "./createKmsAccount";
+import { KmsAccount, kmsKeyToAccount } from "./kmsKeyToAccount";
 import { CreateKeyCommand, KMSClient } from "@aws-sdk/client-kms";
-import { parseGwei, verifyMessage, verifyTypedData } from "viem";
+import { parseGwei, http, verifyMessage, verifyTypedData, createClient, parseEther } from "viem";
+import { foundry } from "viem/chains";
+import { anvilRpcUrl, testClient } from "../../../test/common";
+import { waitForTransaction } from "../../test/waitForTransaction";
+import { getTransactionReceipt, sendTransaction } from "viem/actions";
 
-describe("createKmsAccount", () => {
-  let account: KMSAccount;
+describe("kmsKeyToAccount", () => {
+  let account: KmsAccount;
   let keyId: string;
 
   beforeAll(async () => {
@@ -30,7 +34,7 @@ describe("createKmsAccount", () => {
 
     keyId = createResponse.KeyMetadata.KeyId;
 
-    account = await createKmsAccount({ keyId, client });
+    account = await kmsKeyToAccount({ keyId, client });
   });
 
   it("signMessage", async () => {
@@ -102,4 +106,22 @@ describe("createKmsAccount", () => {
 
     expect(valid).toBeTruthy();
   });
+
+  it("can execute transactions", async () => {
+    // Fund the KMS account
+    testClient.setBalance({ address: account.address, value: parseEther("1") });
+
+    const kmsClient = createClient({
+      chain: foundry,
+      transport: http(anvilRpcUrl),
+      account,
+    });
+
+    // Check that the KMS account can execute transactions
+    const tx = await sendTransaction(kmsClient, {});
+    await waitForTransaction(tx);
+
+    const receipt = await getTransactionReceipt(kmsClient, { hash: tx });
+    expect(receipt.status).toEqual("success");
+  });
 });

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

@@ -0,0 +1,81 @@
+import { KMSClient } from "@aws-sdk/client-kms";
+import { LocalAccount, hashMessage, hashTypedData, keccak256, serializeTransaction, signatureToHex } from "viem";
+import { toAccount } from "viem/accounts";
+import { signWithKms } from "./signWithKms";
+import { getAddressFromKms } from "./getAddressFromKms";
+
+export type KmsKeyToAccountOptions = {
+  keyId: string;
+  client?: KMSClient;
+};
+
+export type KmsAccount = LocalAccount<"aws-kms"> & {
+  getKeyId(): string;
+};
+
+/**
+ * @description Creates an Account from a KMS key.
+ *
+ * @returns A Local Account.
+ */
+
+export async function kmsKeyToAccount({
+  keyId,
+  client = new KMSClient(),
+}: KmsKeyToAccountOptions): Promise<KmsAccount> {
+  const address = await getAddressFromKms({ keyId, client });
+
+  const account = toAccount({
+    address,
+    async signMessage({ message }) {
+      const signature = await signWithKms({
+        client,
+        keyId,
+        hash: hashMessage(message),
+        address,
+      });
+
+      return signatureToHex(signature);
+    },
+    // The logic of this function should be align with viem's signTransaction
+    // https://github.com/wevm/viem/blob/main/src/accounts/utils/signTransaction.ts
+    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: keccak256(serializer(signableTransaction)),
+        address,
+      });
+
+      return serializer(transaction, signature);
+    },
+    async signTypedData(typedData) {
+      const signature = await signWithKms({
+        client,
+        keyId,
+        hash: hashTypedData(typedData),
+        address,
+      });
+
+      return signatureToHex(signature);
+    },
+  });
+
+  return {
+    ...account,
+    source: "aws-kms",
+    getKeyId: () => keyId,
+  };
+}

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/exports/kms.ts

@@ -1 +1 @@
-export { createKmsAccount, type CreateKmsAccountOptions, type KMSAccount } from "../account/kms/createKmsAccount";
+export { kmsKeyToAccount, type KmsKeyToAccountOptions, type KmsAccount } from "../account/kms/kmsKeyToAccount";

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,
+  },
+});