🌱 Add AWS KMS support

package.json

@@ -14,6 +14,8 @@
     "@across-protocol/contracts-v2": "2.5.4",
     "@across-protocol/sdk-v2": "0.23.8",
     "@arbitrum/sdk": "^3.1.3",
+    "@aws-sdk/client-kms": "^3.592.0",
+    "@aws-sdk/client-s3": "^3.592.0",
     "@consensys/linea-sdk": "^0.2.1",
     "@defi-wonderland/smock": "^2.3.5",
     "@eth-optimism/sdk": "^3.2.2",

src/utils/AwskmsUtils.ts

@@ -0,0 +1,100 @@
+import { GetObjectCommand, S3Client } from "@aws-sdk/client-s3";
+import * as KMS from "@aws-sdk/client-kms";
+
+import fs from "fs";
+export interface KeyConfig {
+  keyID: string;
+  accessKeyId: string;
+  secretAccessKey: string;
+  region: string;
+  bucketName: string;
+  fileKey: string;
+}
+
+export interface AwskmsConfig {
+  [network: string]: {
+    [keyName: string]: KeyConfig;
+  };
+}
+
+interface AwsS3StorageConfig {
+  bucket: string;
+  key: string;
+}
+
+const { AWS_S3_STORAGE_CONFIG } = process.env;
+const storageConfig: AwsS3StorageConfig = AWS_S3_STORAGE_CONFIG ? JSON.parse(AWS_S3_STORAGE_CONFIG) : undefined;
+
+export function getAwskmsConfig(keys: string[]): KeyConfig[] {
+  let configOverride: AwskmsConfig = {};
+  if (process.env.AWSKMS_CONFIG) {
+    configOverride = JSON.parse(process.env.AWSKMS_CONFIG);
+  } else {
+    const overrideFname = ".AwskmsOverride.js";
+    try {
+      if (fs.existsSync(`${__dirname}/${overrideFname}`)) {
+        configOverride = require(`./${overrideFname}`);
+      }
+    } catch (err) {
+      // eslint-disable-next-line no-console
+      console.error(err);
+    }
+  }
+
+  const keyConfigs = keys.map((keyName: string): KeyConfig => {
+    return (configOverride["mainnet"][keyName] || {}) as KeyConfig; // Hardcode to "mainnet" network. This makes no impact key retrieval.
+  });
+
+  return keyConfigs;
+}
+
+export const downloadEncryptedKey = async (config: KeyConfig): Promise<Uint8Array | undefined> => {
+  const command = new GetObjectCommand({
+    Bucket: storageConfig.bucket,
+    Key: storageConfig.key,
+  });
+  const client = new S3Client({
+    region: config.region,
+    credentials: {
+      accessKeyId: config.accessKeyId,
+      secretAccessKey: config.secretAccessKey,
+    },
+  });
+  try {
+    const response = await client.send(command);
+    // The Body object also has 'transformToByteArray' and 'transformToWebStream' methods.
+    return response.Body?.transformToByteArray();
+  } catch (err) {
+    // eslint-disable-next-line no-console
+    console.error(err);
+  }
+};
+
+export async function retrieveAwskmsKeys(awskmsConfigs: KeyConfig[]): Promise<string[]> {
+  return await Promise.all(
+    awskmsConfigs.map(async (config) => {
+      const encryptedFileContent = await downloadEncryptedKey(config);
+      const input: KMS.DecryptCommandInput = {
+        KeyId: config.keyID,
+        CiphertextBlob: encryptedFileContent,
+        EncryptionAlgorithm: "SYMMETRIC_DEFAULT",
+      };
+      const decryptCommand = new KMS.DecryptCommand(input);
+
+      const client = new KMS.KMS({
+        region: config.region,
+        credentials: {
+          accessKeyId: config.accessKeyId,
+          secretAccessKey: config.secretAccessKey,
+        },
+      });
+
+      const data = await client.send(decryptCommand);
+      if (!(data.Plaintext instanceof Uint8Array)) {
+        throw new Error("result.plaintext wrong type");
+      }
+
+      return "0x" + Buffer.from(data.Plaintext).toString().trim();
+    })
+  );
+}

src/utils/CLIUtils.ts

@@ -1,10 +1,10 @@
 import minimist from "minimist";
 import { Signer } from "ethers";
-import { constants as sdkConsts } from "@across-protocol/sdk-v2";
+import { constants as sdkConsts } from "@across-protocol/sdk";
 import { SignerOptions, getSigner } from "./SignerUtils";
 import { isDefined } from "./TypeGuards";
 
-const keyTypes = ["secret", "mnemonic", "privateKey", "gckms", "void"];
+const keyTypes = ["secret", "mnemonic", "privateKey", "gckms", "awskms", "void"];
 
 /**
  * Retrieves a signer based on both the CLI args and the env.
@@ -45,6 +45,6 @@ export function retrieveSignerFromCLIArgs(): Promise<Signer> {
  * @param keyType The key type to check.
  * @returns True if the key type is valid, false otherwise.
  */
-function isValidKeyType(keyType: string): keyType is "secret" | "mnemonic" | "privateKey" | "gckms" | "void" {
+function isValidKeyType(keyType: string): keyType is "secret" | "mnemonic" | "privateKey" | "gckms" | "awskms" | "void" {
   return keyTypes.includes(keyType);
 }

src/utils/SignerUtils.ts

@@ -1,8 +1,9 @@
 import { readFile } from "fs/promises";
 import { constants as ethersConsts, VoidSigner } from "ethers";
-import { typeguards } from "@across-protocol/sdk-v2";
+import { typeguards } from "@across-protocol/sdk";
 import { Signer, Wallet, retrieveGckmsKeys, getGckmsConfig, isDefined, assert } from "./";
 import { ArweaveWalletJWKInterface, ArweaveWalletJWKInterfaceSS } from "../interfaces";
+import { getAwskmsConfig, retrieveAwskmsKeys } from "./AwskmsUtils";
 
 /**
  * Signer options for the getSigner function.
@@ -49,6 +50,9 @@ export async function getSigner({ keyType, gckmsKeys, cleanEnv, roAddress }: Sig
     case "gckms":
       signer = await getGckmsSigner(gckmsKeys);
       break;
+    case "awskms":
+      signer = await getAwskmsSigner(gckmsKeys);
+      break;
     case "secret":
       signer = await getSecretSigner();
       break;
@@ -59,7 +63,7 @@ export async function getSigner({ keyType, gckmsKeys, cleanEnv, roAddress }: Sig
       throw new Error(`getSigner: Unsupported signer key type (${keyType})`);
   }
   if (!signer) {
-    throw new Error('Must specify "secret", "mnemonic", "privateKey", "gckms" or "void" for keyType');
+    throw new Error('Must specify "secret", "mnemonic", "privateKey", "gckms", "awskms" or "void" for keyType');
   }
   if (cleanEnv) {
     cleanKeysFromEnvironment();
@@ -92,6 +96,19 @@ async function getGckmsSigner(keys?: string[]): Promise<Signer> {
   return new Wallet(privateKeys[0]); // GCKMS retrieveGckmsKeys returns multiple keys. For now we only support 1.
 }
 
+/**
+ * Retrieves a signer based on the AWSKMS key set in the args.
+ * @returns A signer based on the AWSKMS key set in the args.
+ * @throws If the AWSKMS key is not set.
+ */
+async function getAwskmsSigner(keys?: string[]): Promise<Signer> {
+  if (!isDefined(keys) || keys.length === 0) {
+    throw new Error("Wallet AWSKMS selected but no keys parameter set! Set AWSKMS key (--keys <key>) to use");
+  }
+  const privateKeys = await retrieveAwskmsKeys(getAwskmsConfig(keys));
+  return new Wallet(privateKeys[0]); // AWSKMS retrieveAwskmsKeys returns multiple keys. For now we only support 1.
+}
+
 /**
  * Retrieves a signer based on the mnemonic set in the env.
  * @returns A signer based on the mnemonic set in the env.