LIT-3956 - wrapped-keys - Allow batch of keys to be stored atomically (storeEncryptedKeyBatch)

local-tests/test.ts

@@ -105,6 +105,7 @@ import { testExportWrappedKey } from './tests/wrapped-keys/testExportWrappedKey'
 import { testSignMessageWithSolanaEncryptedKey } from './tests/wrapped-keys/testSignMessageWithSolanaEncryptedKey';
 import { testSignTransactionWithSolanaEncryptedKey } from './tests/wrapped-keys/testSignTransactionWithSolanaEncryptedKey';
 import { testBatchGeneratePrivateKeys } from './tests/wrapped-keys/testBatchGeneratePrivateKeys';
+import { testFailBatchGeneratePrivateKeysAtomic } from './tests/wrapped-keys/testFailStoreEncryptedKeyBatchIsAtomic';
 
 import { setLitActionsCodeToLocal } from './tests/wrapped-keys/util';
 import { testUseEoaSessionSigsToRequestSingleResponse } from './tests/testUseEoaSessionSigsToRequestSingleResponse';
@@ -151,6 +152,7 @@ setLitActionsCodeToLocal();
     testFailEthereumSignTransactionWrappedKeyWithMissingParam,
     testFailEthereumSignTransactionWrappedKeyWithInvalidParam,
     testFailEthereumSignTransactionWrappedKeyInvalidDecryption,
+    testFailBatchGeneratePrivateKeysAtomic,
 
     // -- import wrapped keys
     testFailImportWrappedKeysWithSamePrivateKey,

local-tests/tests/wrapped-keys/testBatchGeneratePrivateKeys.ts

@@ -7,7 +7,7 @@ import bs58 from 'bs58';
 import { ethers } from 'ethers';
 import { BatchGeneratePrivateKeysActionResult } from '../../../packages/wrapped-keys/src/lib/types';
 
-const { batchGeneratePrivateKeys, exportPrivateKey } = api;
+const { batchGeneratePrivateKeys } = api;
 
 async function verifySolanaSignature(
   solanaResult: BatchGeneratePrivateKeysActionResult,

local-tests/tests/wrapped-keys/testFailStoreEncryptedKeyBatchIsAtomic.ts

@@ -0,0 +1,134 @@
+import { log } from '@lit-protocol/misc';
+import { TinnyEnvironment } from 'local-tests/setup/tinny-environment';
+import { api } from '@lit-protocol/wrapped-keys';
+import { getPkpSessionSigs } from 'local-tests/setup/session-sigs/get-pkp-session-sigs';
+import nacl from 'tweetnacl';
+import bs58 from 'bs58';
+import { ethers } from 'ethers';
+import { BatchGeneratePrivateKeysActionResult } from '../../../packages/wrapped-keys/src/lib/types';
+import { batchGenerateKeysWithLitAction } from '../../../packages/wrapped-keys/src/lib/lit-actions-client';
+import {
+  getFirstSessionSig,
+  getPkpAccessControlCondition,
+  getPkpAddressFromSessionSig,
+} from '../../../packages/wrapped-keys/src/lib/utils';
+import { getLitActionCodeOrCidCommon } from '../../../packages/wrapped-keys/src/lib/lit-actions-client/utils';
+import { getKeyTypeFromNetwork } from '../../../packages/wrapped-keys/src/lib/api/utils';
+import { listEncryptedKeyMetadata } from '../../../packages/wrapped-keys/src/lib/api';
+
+const { batchGeneratePrivateKeys, storeEncryptedKeyBatch } = api;
+
+/**
+ * Test Commands:
+ * ✅ NETWORK=datil-dev yarn test:local --filter=testSignMessageWithSolanaEncryptedKey
+ * ✅ NETWORK=datil-test yarn test:local --filter=testSignMessageWithSolanaEncryptedKey
+ * ✅ NETWORK=localchain yarn test:local --filter=testSignMessageWithSolanaEncryptedKey
+ */
+export const testFailBatchGeneratePrivateKeysAtomic = async (
+  devEnv: TinnyEnvironment
+) => {
+  const alice = await devEnv.createRandomPerson();
+
+  try {
+    const pkpSessionSigsSigning = await getPkpSessionSigs(
+      devEnv,
+      alice,
+      null,
+      new Date(Date.now() + 1000 * 60 * 10).toISOString()
+    ); // 10 mins expiry
+
+    const solanaMessageToSign = 'This is a test solana message';
+    const evmMessageToSign = 'This is a test evm message';
+
+    const sessionSig = getFirstSessionSig(pkpSessionSigsSigning);
+    const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
+
+    const allowPkpAddressToDecrypt = getPkpAccessControlCondition(pkpAddress);
+
+    const { litActionCode, litActionIpfsCid } = getLitActionCodeOrCidCommon(
+      'batchGenerateEncryptedKeys'
+    );
+
+    const actionResults = await batchGenerateKeysWithLitAction({
+      litNodeClient: devEnv.litNodeClient,
+      litActionIpfsCid: litActionCode ? undefined : litActionIpfsCid,
+      litActionCode: litActionCode ? litActionCode : undefined,
+      accessControlConditions: [allowPkpAddressToDecrypt],
+      actions: [
+        {
+          network: 'evm',
+          signMessageParams: { messageToSign: evmMessageToSign },
+          generateKeyParams: { memo: 'Test evm key' },
+        },
+        {
+          network: 'solana',
+          signMessageParams: { messageToSign: solanaMessageToSign },
+          generateKeyParams: { memo: 'Test solana key' },
+        },
+      ],
+      pkpSessionSigs: pkpSessionSigsSigning,
+    });
+
+    const keyParamsBatch = actionResults.map((keyData) => {
+      const { generateEncryptedPrivateKey, network } = keyData;
+      return {
+        ...generateEncryptedPrivateKey,
+        keyType: getKeyTypeFromNetwork(network),
+      };
+    });
+
+    // Intentional failure to persist due to missing publicKey
+    delete keyParamsBatch[0].publicKey;
+
+    try {
+      await storeEncryptedKeyBatch({
+        pkpSessionSigs: pkpSessionSigsSigning,
+        litNodeClient: devEnv.litNodeClient,
+        keyBatch: keyParamsBatch,
+      });
+
+      throw new Error(
+        'storeEncryptedKeyBatch() succeeded but we expected it to fail!'
+      );
+    } catch (err) {
+      // We expect `storeEncryptedKeyBatch` to fail w/ a specific error
+      if (
+        err.message.includes(
+          'storeEncryptedKeyBatch() succeeded but we expected it to fail!'
+        ) ||
+        !err.message.includes(
+          'keyParamsBatch[0]: Missing "publicKey" parameter in request'
+        )
+      ) {
+        throw err;
+      }
+
+      try {
+        const keys = await listEncryptedKeyMetadata({
+          litNodeClient: devEnv.litNodeClient,
+          pkpSessionSigs: pkpSessionSigsSigning,
+        });
+
+        console.error(
+          'Got a value back we shouldnt have from listEncryptedKeyMetadata()',
+          keys
+        );
+
+        throw new Error(
+          'Expected `listEncryptedKeyMetadata() to fail, but it didnt!`'
+        );
+      } catch (err) {
+        if (err.message.includes('No keys exist for pkpAddress')) {
+          log('✅ testFailBatchGeneratePrivateKeysAtomic');
+        } else {
+          throw err;
+        }
+      }
+    }
+  } catch (err) {
+    console.log(err.message, err, err.stack);
+    throw err;
+  } finally {
+    devEnv.releasePrivateKeyFromUser(alice);
+  }
+};

packages/wrapped-keys/src/index.ts

@@ -8,6 +8,7 @@ import {
   storeEncryptedKey,
   listEncryptedKeyMetadata,
   batchGeneratePrivateKeys,
+  storeEncryptedKeyBatch,
 } from './lib/api';
 import {
   CHAIN_ETHEREUM,
@@ -45,11 +46,13 @@ import type {
   SignTransactionParamsSupportedEvm,
   SignTransactionParamsSupportedSolana,
   StoreEncryptedKeyParams,
+  StoreEncryptedKeyBatchParams,
   StoredKeyData,
   StoredKeyMetadata,
   ListEncryptedKeyMetadataParams,
   StoreEncryptedKeyResult,
   ImportPrivateKeyResult,
+  StoreEncryptedKeyBatchResult,
 } from './lib/types';
 
 export const constants = {
@@ -70,6 +73,7 @@ export const api = {
   signMessageWithEncryptedKey,
   signTransactionWithEncryptedKey,
   storeEncryptedKey,
+  storeEncryptedKeyBatch,
   batchGeneratePrivateKeys,
 };
 
@@ -101,6 +105,8 @@ export {
   SignTransactionWithEncryptedKeyParams,
   StoreEncryptedKeyParams,
   StoreEncryptedKeyResult,
+  StoreEncryptedKeyBatchParams,
+  StoreEncryptedKeyBatchResult,
   StoredKeyData,
   StoredKeyMetadata,
   SupportedNetworks,

packages/wrapped-keys/src/lib/api/batch-generate-private-keys.ts

@@ -1,7 +1,7 @@
 import { getKeyTypeFromNetwork } from './utils';
 import { batchGenerateKeysWithLitAction } from '../lit-actions-client';
 import { getLitActionCodeOrCidCommon } from '../lit-actions-client/utils';
-import { storePrivateKey } from '../service-client';
+import { storePrivateKeyBatch } from '../service-client';
 import {
   BatchGeneratePrivateKeysActionResult,
   BatchGeneratePrivateKeysParams,
@@ -41,34 +41,40 @@ export async function batchGeneratePrivateKeys(
     pkpSessionSigs,
   });
 
-  const results = await Promise.all(
-    actionResults.map(
-      async (result): Promise<BatchGeneratePrivateKeysActionResult> => {
-        const { generateEncryptedPrivateKey, network } = result;
+  const keyParamsBatch = actionResults.map((keyData) => {
+    const { generateEncryptedPrivateKey, network } = keyData;
+    return {
+      ...generateEncryptedPrivateKey,
+      keyType: getKeyTypeFromNetwork(network),
+    };
+  });
+
+  const { ids } = await storePrivateKeyBatch({
+    pkpAddress,
+    sessionSig,
+    storedKeyMetadataBatch: keyParamsBatch,
+    litNetwork: litNodeClient.config.litNetwork,
+  });
 
-        const signature = result.signMessage?.signature;
+  const results = actionResults.map(
+    (actionResult, ndx): BatchGeneratePrivateKeysActionResult => {
+      const {
+        generateEncryptedPrivateKey: { memo, publicKey },
+      } = actionResult;
+      const id = ids[ndx]; // Result of writes is in same order as provided
 
-        const { id } = await storePrivateKey({
-          sessionSig,
-          storedKeyMetadata: {
-            ...generateEncryptedPrivateKey,
-            keyType: getKeyTypeFromNetwork(network),
-            pkpAddress,
-          },
-          litNetwork: litNodeClient.config.litNetwork,
-        });
+      const signature = actionResult.signMessage?.signature;
 
-        return {
-          ...(signature ? { signMessage: { signature } } : {}),
-          generateEncryptedPrivateKey: {
-            memo: generateEncryptedPrivateKey.memo,
-            id,
-            generatedPublicKey: generateEncryptedPrivateKey.publicKey,
-            pkpAddress,
-          },
-        };
-      }
-    )
+      return {
+        ...(signature ? { signMessage: { signature } } : {}),
+        generateEncryptedPrivateKey: {
+          memo: memo,
+          id,
+          generatedPublicKey: publicKey,
+          pkpAddress,
+        },
+      };
+    }
   );
 
   return { pkpAddress, results };

packages/wrapped-keys/src/lib/api/index.ts

@@ -7,6 +7,7 @@ import { listEncryptedKeyMetadata } from './list-encrypted-key-metadata';
 import { signMessageWithEncryptedKey } from './sign-message-with-encrypted-key';
 import { signTransactionWithEncryptedKey } from './sign-transaction-with-encrypted-key';
 import { storeEncryptedKey } from './store-encrypted-key';
+import { storeEncryptedKeyBatch } from './store-encrypted-key-batch';
 
 export {
   listEncryptedKeyMetadata,
@@ -16,6 +17,7 @@ export {
   exportPrivateKey,
   signMessageWithEncryptedKey,
   storeEncryptedKey,
+  storeEncryptedKeyBatch,
   getEncryptedKey,
   batchGeneratePrivateKeys,
 };

packages/wrapped-keys/src/lib/api/store-encrypted-key-batch.ts

@@ -0,0 +1,54 @@
+import { storePrivateKeyBatch } from '../service-client';
+import { StoreKeyBatchParams } from '../service-client/types';
+import {
+  StoredKeyData,
+  StoreEncryptedKeyBatchParams,
+  StoreEncryptedKeyBatchResult,
+} from '../types';
+import { getFirstSessionSig, getPkpAddressFromSessionSig } from '../utils';
+
+/** Stores a batch of encrypted private keys and their metadata to the wrapped keys backend service
+ *
+ * @param { StoreEncryptedKeyParams } params Parameters required to fetch the encrypted private key metadata
+ * @returns { Promise<StoreEncryptedKeyResult> } The encrypted private key and its associated metadata
+ */
+export async function storeEncryptedKeyBatch(
+  params: StoreEncryptedKeyBatchParams
+): Promise<StoreEncryptedKeyBatchResult> {
+  const { pkpSessionSigs, litNodeClient, keyBatch } = params;
+  const sessionSig = getFirstSessionSig(pkpSessionSigs);
+  const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
+
+  const storedKeyMetadataBatch: StoreKeyBatchParams['storedKeyMetadataBatch'] =
+    keyBatch.map(
+      ({
+        keyType,
+        publicKey,
+        memo,
+        dataToEncryptHash,
+        ciphertext,
+      }): Pick<
+        StoredKeyData,
+        | 'pkpAddress'
+        | 'publicKey'
+        | 'keyType'
+        | 'dataToEncryptHash'
+        | 'ciphertext'
+        | 'memo'
+      > => ({
+        pkpAddress,
+        publicKey,
+        memo,
+        dataToEncryptHash,
+        ciphertext,
+        keyType,
+      })
+    );
+
+  return storePrivateKeyBatch({
+    pkpAddress,
+    storedKeyMetadataBatch,
+    sessionSig: getFirstSessionSig(pkpSessionSigs),
+    litNetwork: litNodeClient.config.litNetwork,
+  });
+}

packages/wrapped-keys/src/lib/service-client/client.ts

@@ -1,8 +1,14 @@
-import { FetchKeyParams, ListKeysParams, StoreKeyParams } from './types';
+import {
+  FetchKeyParams,
+  ListKeysParams,
+  StoreKeyBatchParams,
+  StoreKeyParams,
+} from './types';
 import { generateRequestId, getBaseRequestParams, makeRequest } from './utils';
 import {
   StoredKeyData,
   StoredKeyMetadata,
+  StoreEncryptedKeyBatchResult,
   StoreEncryptedKeyResult,
 } from '../types';
 import { getPkpAddressFromSessionSig } from '../utils';
@@ -92,3 +98,33 @@ export async function storePrivateKey(
 
   return { pkpAddress, id };
 }
+
+/** Stores a batch of up to 25 private key metadata into the wrapped keys service backend
+ *
+ * @param { StoreKeyParams } params Parameters required to store the private key metadata
+ * @returns { Promise<StoreEncryptedKeyResult> } `true` on successful write to the service. Otherwise, this method throws an error.
+ */
+export async function storePrivateKeyBatch(
+  params: StoreKeyBatchParams
+): Promise<StoreEncryptedKeyBatchResult> {
+  const { litNetwork, sessionSig, storedKeyMetadataBatch } = params;
+
+  const requestId = generateRequestId();
+  const { baseUrl, initParams } = getBaseRequestParams({
+    litNetwork,
+    sessionSig,
+    method: 'POST',
+    requestId,
+  });
+
+  const { pkpAddress, ids } = await makeRequest<StoreEncryptedKeyBatchResult>({
+    url: `${baseUrl}_batch`,
+    init: {
+      ...initParams,
+      body: JSON.stringify({ keyParamsBatch: storedKeyMetadataBatch }),
+    },
+    requestId,
+  });
+
+  return { pkpAddress, ids };
+}