From 182d70608fce514ae1aec5366a7bbae1dc936844 Mon Sep 17 00:00:00 2001 From: yonada Date: Thu, 25 Apr 2024 13:51:34 +0100 Subject: [PATCH] fix(common): kms correctly serializes transactions (#2721) --- .changeset/perfect-actors-flow.md | 6 +- packages/common/package.json | 1 + packages/common/src/account/kms/README.md | 4 +- .../src/account/kms/createKmsAccount.ts | 67 --------------- ...ccount.test.ts => kmsKeyToAccount.test.ts} | 32 ++++++-- .../common/src/account/kms/kmsKeyToAccount.ts | 81 +++++++++++++++++++ .../common/src/account/kms/signWithKms.ts | 8 +- packages/common/src/exports/kms.ts | 2 +- packages/common/src/test/minePending.ts | 9 +++ .../common/src/test/waitForTransaction.ts | 13 +++ packages/common/test/common.ts | 18 +++++ packages/common/test/globalSetup.ts | 13 +++ packages/common/test/setup.ts | 18 +++++ packages/common/vitest.config.ts | 12 +++ pnpm-lock.yaml | 3 + 15 files changed, 205 insertions(+), 82 deletions(-) delete mode 100644 packages/common/src/account/kms/createKmsAccount.ts rename packages/common/src/account/kms/{createKmsAccount.test.ts => kmsKeyToAccount.test.ts} (69%) create mode 100644 packages/common/src/account/kms/kmsKeyToAccount.ts create mode 100644 packages/common/src/test/minePending.ts create mode 100644 packages/common/src/test/waitForTransaction.ts create mode 100644 packages/common/test/common.ts create mode 100644 packages/common/test/globalSetup.ts create mode 100644 packages/common/test/setup.ts create mode 100644 packages/common/vitest.config.ts diff --git a/.changeset/perfect-actors-flow.md b/.changeset/perfect-actors-flow.md index da4c3ce49b..cf79fb3b00 100644 --- a/.changeset/perfect-actors-flow.md +++ b/.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/client-kms@3.x` and `asn1.js@5.x` 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). diff --git a/packages/common/package.json b/packages/common/package.json index d95e037316..f00c94c85a 100644 --- a/packages/common/package.json +++ b/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" }, diff --git a/packages/common/src/account/kms/README.md b/packages/common/src/account/kms/README.md index 47c4e82eeb..1f9fa8a248 100644 --- a/packages/common/src/account/kms/README.md +++ b/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/client-kms@3.x` and `asn1.js@5.x` 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). diff --git a/packages/common/src/account/kms/createKmsAccount.ts b/packages/common/src/account/kms/createKmsAccount.ts deleted file mode 100644 index 6a78279e0b..0000000000 --- a/packages/common/src/account/kms/createKmsAccount.ts +++ /dev/null @@ -1,67 +0,0 @@ -import { KMSClient } from "@aws-sdk/client-kms"; -import { LocalAccount, hashMessage, hashTypedData, keccak256, serializeTransaction } from "viem"; -import { toAccount } from "viem/accounts"; -import { signWithKms } from "./signWithKms"; -import { getAddressFromKms } from "./getAddressFromKms"; - -export type CreateKmsAccountOptions = { - keyId: string; - client?: KMSClient; -}; - -export type KMSAccount = LocalAccount<"aws-kms">; - -/** - * @description Creates an Account from a KMS key and instance. - * - * @returns A Local Account. - */ -export async function createKmsAccount({ - keyId, - client = new KMSClient(), -}: CreateKmsAccountOptions): Promise { - const address = await getAddressFromKms({ keyId, client }); - - const account = toAccount({ - address, - async signMessage({ message }) { - const hash = hashMessage(message); - const signature = await signWithKms({ - client, - keyId, - hash, - address, - }); - - return signature; - }, - async signTransaction(transaction) { - const unsignedTx = serializeTransaction(transaction); - const hash = keccak256(unsignedTx); - const signature = await signWithKms({ - client, - keyId, - hash, - address, - }); - - return signature; - }, - async signTypedData(typedData) { - const hash = hashTypedData(typedData); - const signature = await signWithKms({ - client, - keyId, - hash, - address, - }); - - return signature; - }, - }); - - return { - ...account, - source: "aws-kms", - }; -} diff --git a/packages/common/src/account/kms/createKmsAccount.test.ts b/packages/common/src/account/kms/kmsKeyToAccount.test.ts similarity index 69% rename from packages/common/src/account/kms/createKmsAccount.test.ts rename to packages/common/src/account/kms/kmsKeyToAccount.test.ts index 022bf5ca30..7d739bad27 100644 --- a/packages/common/src/account/kms/createKmsAccount.test.ts +++ b/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"); + }); }); diff --git a/packages/common/src/account/kms/kmsKeyToAccount.ts b/packages/common/src/account/kms/kmsKeyToAccount.ts new file mode 100644 index 0000000000..c9f82ff00a --- /dev/null +++ b/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 { + 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, + }; +} diff --git a/packages/common/src/account/kms/signWithKms.ts b/packages/common/src/account/kms/signWithKms.ts index f26a1a2299..fac4759d58 100644 --- a/packages/common/src/account/kms/signWithKms.ts +++ b/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, - }); + }; } diff --git a/packages/common/src/exports/kms.ts b/packages/common/src/exports/kms.ts index 04a1b3b293..4759c209e4 100644 --- a/packages/common/src/exports/kms.ts +++ b/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"; diff --git a/packages/common/src/test/minePending.ts b/packages/common/src/test/minePending.ts new file mode 100644 index 0000000000..1ce4743f92 --- /dev/null +++ b/packages/common/src/test/minePending.ts @@ -0,0 +1,9 @@ +import { testClient } from "../../test/common"; + +export async function minePending(): Promise { + const content = await testClient.getTxpoolContent(); + if (!Object.keys(content.pending).length) return; + + await testClient.mine({ blocks: 1 }); + await minePending(); +} diff --git a/packages/common/src/test/waitForTransaction.ts b/packages/common/src/test/waitForTransaction.ts new file mode 100644 index 0000000000..abf5fe8556 --- /dev/null +++ b/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 { + await minePending(); + const receipt = await getTransactionReceipt(testClient, { hash }); + if (receipt.status === "reverted") { + // TODO: better error + throw new Error(`Transaction reverted (${hash})`); + } +} diff --git a/packages/common/test/common.ts b/packages/common/test/common.ts new file mode 100644 index 0000000000..d62e4ab516 --- /dev/null +++ b/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); diff --git a/packages/common/test/globalSetup.ts b/packages/common/test/globalSetup.ts new file mode 100644 index 0000000000..6f9bd68343 --- /dev/null +++ b/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> { + const shutdownAnvilProxy = await startAnvilProxy({ + host: anvilHost, + port: anvilPort, + }); + + return async () => { + await shutdownAnvilProxy(); + }; +} diff --git a/packages/common/test/setup.ts b/packages/common/test/setup.ts new file mode 100644 index 0000000000..35d83f4304 --- /dev/null +++ b/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 => { + await testClient.loadState({ state }); + }; +}); + +// Some tests execute transactions, so we restore chain state here. +beforeEach(async () => { + const state = await testClient.dumpState(); + return async (): Promise => { + await testClient.loadState({ state }); + }; +}); diff --git a/packages/common/vitest.config.ts b/packages/common/vitest.config.ts new file mode 100644 index 0000000000..61f51f284f --- /dev/null +++ b/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, + }, +}); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 1718e499b2..155d6f99e6 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -321,6 +321,9 @@ importers: '@types/node': specifier: ^18.15.11 version: 18.15.11 + '@viem/anvil': + specifier: ^0.0.7 + version: 0.0.7(debug@4.3.4) tsup: specifier: ^6.7.0 version: 6.7.0(postcss@8.4.23)(typescript@5.4.2)