feat(common): kms custom account (#2688)

Co-authored-by: Kevin Ingersoll <[email protected]>

.changeset/perfect-actors-flow.md

@@ -0,0 +1,13 @@
+---
+"@latticexyz/common": patch
+---
+
+Added `createKmsAccount`, 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: ... });
+```
+
+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).

.github/workflows/test.yml

@@ -17,6 +17,10 @@ jobs:
           - 5432:5432
         # needed because the postgres container does not provide a healthcheck
         options: --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
+      local-kms:
+        image: nsmithuk/local-kms
+        ports:
+          - 8080:8080
     steps:
       - name: Checkout
         uses: actions/checkout@v3
@@ -40,6 +44,7 @@ jobs:
         if: steps.check_changes.outputs.changes_outside_docs
         env:
           DATABASE_URL: "postgres://postgres@localhost:5432/postgres"
+          KMS_ENDPOINT: "http://localhost:8080"
         run: pnpm test:ci
 
       - name: Generate gas reports

packages/common/package.json

@@ -72,5 +72,17 @@
     "@types/node": "^18.15.11",
     "tsup": "^6.7.0",
     "vitest": "0.34.6"
+  },
+  "peerDependencies": {
+    "@aws-sdk/client-kms": "3.x",
+    "asn1.js": "5.x"
+  },
+  "peerDependenciesMeta": {
+    "@aws-sdk/client-kms": {
+      "optional": true
+    },
+    "asn1.js": {
+      "optional": true
+    }
   }
 }

packages/common/src/account/kms/README.md

@@ -0,0 +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.
+
+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: ... });
+```
+
+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/commands/getPublicKey.ts

@@ -0,0 +1,18 @@
+import {
+  GetPublicKeyCommand,
+  GetPublicKeyCommandInput,
+  GetPublicKeyCommandOutput,
+  KMSClient,
+} from "@aws-sdk/client-kms";
+
+export function getPublicKey({
+  keyId,
+  client,
+}: {
+  keyId: GetPublicKeyCommandInput["KeyId"];
+  client: KMSClient;
+}): Promise<GetPublicKeyCommandOutput> {
+  const command = new GetPublicKeyCommand({ KeyId: keyId });
+
+  return client.send(command);
+}

packages/common/src/account/kms/commands/sign.ts

@@ -0,0 +1,23 @@
+import { KMSClient, SignCommand, SignCommandInput, SignCommandOutput } from "@aws-sdk/client-kms";
+import { Hex, fromHex } from "viem";
+
+export async function sign({
+  keyId,
+  hash,
+  client,
+}: {
+  hash: Hex;
+  keyId: SignCommandInput["KeyId"];
+  client: KMSClient;
+}): Promise<SignCommandOutput> {
+  const formatted = Buffer.from(fromHex(hash, "bytes"));
+
+  const command = new SignCommand({
+    KeyId: keyId,
+    Message: formatted,
+    SigningAlgorithm: "ECDSA_SHA_256",
+    MessageType: "DIGEST",
+  });
+
+  return client.send(command);
+}

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

@@ -0,0 +1,34 @@
+import { Address, toHex } from "viem";
+import { publicKeyToAddress } from "viem/utils";
+import { KMSClient, SignCommandInput } from "@aws-sdk/client-kms";
+import { getPublicKey } from "./commands/getPublicKey";
+// @ts-expect-error Could not find a declaration file for module 'asn1.js'.
+import asn1 from "asn1.js";
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+const EcdsaPubKey = asn1.define("EcdsaPubKey", function (this: any) {
+  this.seq().obj(this.key("algo").seq().obj(this.key("a").objid(), this.key("b").objid()), this.key("pubKey").bitstr());
+});
+
+function publicKeyKmsToAddress(publicKey: Uint8Array): Address {
+  const res = EcdsaPubKey.decode(Buffer.from(publicKey));
+
+  const publicKeyBuffer: Buffer = res.pubKey.data;
+
+  const publicKeyHex = toHex(publicKeyBuffer);
+  const address = publicKeyToAddress(publicKeyHex);
+
+  return address;
+}
+
+export async function getAddressFromKMS({
+  keyId,
+  client,
+}: {
+  keyId: SignCommandInput["KeyId"];
+  client: KMSClient;
+}): Promise<Address> {
+  const KMSKey = await getPublicKey({ keyId, client });
+
+  return publicKeyKmsToAddress(KMSKey.PublicKey as Uint8Array);
+}

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

@@ -0,0 +1,87 @@
+import { Hex, isAddressEqual, signatureToHex, toHex } from "viem";
+import { recoverAddress } from "viem/utils";
+import { KMSClient, SignCommandInput } from "@aws-sdk/client-kms";
+import { sign } from "./commands/sign";
+// @ts-expect-error Could not find a declaration file for module 'asn1.js'.
+import asn1 from "asn1.js";
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+const EcdsaSigAsnParse = asn1.define("EcdsaSig", function (this: any) {
+  this.seq().obj(this.key("r").int(), this.key("s").int());
+});
+
+async function getRS(signParams: {
+  hash: Hex;
+  keyId: SignCommandInput["KeyId"];
+  client: KMSClient;
+}): Promise<{ r: Hex; s: Hex }> {
+  const signature = await sign(signParams);
+
+  if (signature.Signature === undefined) {
+    throw new Error("Signature is undefined.");
+  }
+
+  const decoded = EcdsaSigAsnParse.decode(Buffer.from(signature.Signature), "der");
+
+  const r = BigInt(decoded.r);
+  let s = BigInt(decoded.s);
+
+  const secp256k1N = BigInt("0xfffffffffffffffffffffffffffffffebaaedce6af48a03bbfd25e8cd0364141");
+  const secp256k1halfN = secp256k1N / 2n;
+
+  if (s > secp256k1halfN) {
+    s = secp256k1N - s;
+  }
+
+  return {
+    r: toHex(r),
+    s: toHex(s),
+  };
+}
+
+async function getRecovery(hash: Hex, r: Hex, s: Hex, expectedAddress: Hex): Promise<number> {
+  let recovery: number;
+  for (recovery = 0; recovery <= 1; recovery++) {
+    const signature = signatureToHex({
+      r,
+      s,
+      v: recovery ? 28n : 27n,
+      yParity: recovery,
+    });
+
+    const address = await recoverAddress({ hash, signature });
+
+    if (isAddressEqual(address, expectedAddress)) {
+      return recovery;
+    }
+  }
+  throw new Error("Failed to calculate recovery param");
+}
+
+type SignParameters = {
+  hash: Hex;
+  keyId: SignCommandInput["KeyId"];
+  client: KMSClient;
+  address: Hex;
+};
+
+type SignReturnType = Hex;
+
+/**
+ * @description Signs a hash with a given KMS key.
+ *
+ * @param hash The hash to sign.
+ *
+ * @returns The signature.
+ */
+export async function signWithKms({ hash, address, keyId, client }: SignParameters): Promise<SignReturnType> {
+  const { r, s } = await getRS({ keyId, hash, client });
+  const recovery = await getRecovery(hash, r, s, address);
+
+  return signatureToHex({
+    r,
+    s,
+    v: recovery ? 28n : 27n,
+    yParity: recovery,
+  });
+}

packages/common/src/createKmsAccount.test.ts

@@ -0,0 +1,105 @@
+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";
+
+describe("createKmsAccount", () => {
+  let account: KMSAccount;
+  let keyId: string;
+
+  beforeAll(async () => {
+    const client = new KMSClient({
+      endpoint: process.env.KMS_ENDPOINT,
+      region: "local",
+      credentials: {
+        accessKeyId: "AKIAXTTRUF7NU7KDMIED",
+        secretAccessKey: "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",
+      },
+    });
+
+    const command = new CreateKeyCommand({
+      KeyUsage: "SIGN_VERIFY",
+      CustomerMasterKeySpec: "ECC_SECG_P256K1",
+    });
+
+    const createResponse = await client.send(command);
+
+    if (!createResponse.KeyMetadata || !createResponse.KeyMetadata.KeyId) {
+      throw new Error("key creation failed");
+    }
+
+    keyId = createResponse.KeyMetadata.KeyId;
+
+    account = await createKmsAccount({ keyId, client });
+  });
+
+  it("signMessage", async () => {
+    const message = "hello world";
+    const signature = await account.signMessage({ message });
+
+    const valid = await verifyMessage({
+      address: account.address,
+      message,
+      signature,
+    });
+
+    expect(valid).toBeTruthy();
+  });
+
+  it("signTransaction", async () => {
+    await account.signTransaction({
+      chainId: 1,
+      maxFeePerGas: parseGwei("20"),
+      maxPriorityFeePerGas: parseGwei("3"),
+      gas: 21000n,
+      nonce: 69,
+      to: "0xf39fd6e51aad88f6f4ce6ab8827279cfffb92266",
+    });
+  });
+
+  it("signTypedData", async () => {
+    const chainId = 1;
+    const verifyingContract = "0x95222290dd7278aa3ddd389cc1e1d165cc4bafe5";
+    const domain = { chainId, verifyingContract } as const;
+    const types = {
+      Person: [
+        { name: "name", type: "string" },
+        { name: "wallet", type: "address" },
+      ],
+      Mail: [
+        { name: "from", type: "Person" },
+        { name: "to", type: "Person" },
+        { name: "contents", type: "string" },
+      ],
+    };
+    const message = {
+      from: {
+        name: "Cow",
+        wallet: "0xCD2a3d9F938E13CD947Ec05AbC7FE734Df8DD826",
+      },
+      to: {
+        name: "Bob",
+        wallet: "0xbBbBBBBbbBBBbbbBbbBbbbbBBbBbbbbBbBbbBBbB",
+      },
+      contents: "Hello, Bob!",
+    };
+
+    const signature = await account.signTypedData({
+      domain,
+      types,
+      primaryType: "Mail",
+      message,
+    });
+
+    const valid = await verifyTypedData({
+      address: account.address,
+      signature,
+      domain,
+      types,
+      primaryType: "Mail",
+      message,
+    });
+
+    expect(valid).toBeTruthy();
+  });
+});