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/testFailEthereumSignTransactionWrappedKeyInvalidDecryption.ts

@@ -1,14 +1,13 @@
 import { log } from '@lit-protocol/misc';
 import { ethers } from 'ethers';
 import { TinnyEnvironment } from 'local-tests/setup/tinny-environment';
-import { EthereumLitTransaction } from '@lit-protocol/wrapped-keys';
 import { getPkpSessionSigs } from 'local-tests/setup/session-sigs/get-pkp-session-sigs';
-import { getPkpAccessControlCondition } from 'packages/wrapped-keys/src/lib/utils';
 import { encryptString } from '@lit-protocol/encryption';
 import { LIT_PREFIX } from 'packages/wrapped-keys/src/lib/constants';
 import { LIT_ACTION_CID_REPOSITORY } from '../../../packages/wrapped-keys/src/lib/lit-actions-client/constants';
 import { getBaseTransactionForNetwork } from './util';
 import { GLOBAL_OVERWRITE_IPFS_CODE_BY_NETWORK } from '@lit-protocol/constants';
+import { getPkpAccessControlCondition } from '../../../packages/wrapped-keys/src/lib/api/utils';
 
 /**
  * Test Commands:

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

@@ -0,0 +1,130 @@
+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 { batchGenerateKeysWithLitAction } from '../../../packages/wrapped-keys/src/lib/lit-actions-client';
+import { getLitActionCodeOrCidCommon } from '../../../packages/wrapped-keys/src/lib/lit-actions-client/utils';
+import {
+  getFirstSessionSig,
+  getKeyTypeFromNetwork,
+  getPkpAccessControlCondition,
+  getPkpAddressFromSessionSig,
+} from '../../../packages/wrapped-keys/src/lib/api/utils';
+import { listEncryptedKeyMetadata } from '../../../packages/wrapped-keys/src/lib/api';
+
+const { 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);
+  }
+};

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

@@ -3,10 +3,8 @@ 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 { randomSolanaPrivateKey } from 'local-tests/setup/tinny-utils';
-import { listPrivateKeyMetadata } from '../../../packages/wrapped-keys/src/lib/service-client';
-import { getFirstSessionSig } from '../../../packages/wrapped-keys/src/lib/utils';
 
-const { importPrivateKey } = api;
+const { importPrivateKey, listEncryptedKeyMetadata } = api;
 
 /**
  * Test Commands:
@@ -44,9 +42,9 @@ export const testImportWrappedKey = async (devEnv: TinnyEnvironment) => {
       );
     }
 
-    const keys = await listPrivateKeyMetadata({
-      sessionSig: getFirstSessionSig(pkpSessionSigs),
-      litNetwork: devEnv.litNodeClient.config.litNetwork,
+    const keys = await listEncryptedKeyMetadata({
+      pkpSessionSigs,
+      litNodeClient: devEnv.litNodeClient,
     });
 
     if (keys.length !== 1 || keys[0].id !== id) {

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,17 +1,17 @@
-import { getKeyTypeFromNetwork } from './utils';
+import {
+  getFirstSessionSig,
+  getKeyTypeFromNetwork,
+  getPkpAccessControlCondition,
+  getPkpAddressFromSessionSig,
+} 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,
   BatchGeneratePrivateKeysResult,
 } from '../types';
-import {
-  getFirstSessionSig,
-  getPkpAccessControlCondition,
-  getPkpAddressFromSessionSig,
-} from '../utils';
 
 /**
  *  TODO: Document batch behaviour
@@ -41,34 +41,39 @@ 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({
+    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/export-private-key.ts

@@ -1,8 +1,12 @@
+import {
+  getFirstSessionSig,
+  getPkpAccessControlCondition,
+  getPkpAddressFromSessionSig,
+} from './utils';
 import { exportPrivateKeyWithLitAction } from '../lit-actions-client';
 import { getLitActionCodeOrCid } from '../lit-actions-client/utils';
 import { fetchPrivateKey } from '../service-client';
 import { ExportPrivateKeyParams, ExportPrivateKeyResult } from '../types';
-import { getFirstSessionSig, getPkpAccessControlCondition } from '../utils';
 
 /**
  * Exports a previously persisted private key from the wrapped keys service for direct use by the caller, along with the keys metadata.
@@ -19,7 +23,10 @@ export async function exportPrivateKey(
   const { litNodeClient, network, pkpSessionSigs, id } = params;
 
   const sessionSig = getFirstSessionSig(pkpSessionSigs);
+  const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
+
   const storedKeyMetadata = await fetchPrivateKey({
+    pkpAddress,
     id,
     sessionSig,
     litNetwork: litNodeClient.config.litNetwork,

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

@@ -1,13 +1,13 @@
-import { getKeyTypeFromNetwork } from './utils';
-import { generateKeyWithLitAction } from '../lit-actions-client';
-import { getLitActionCodeOrCid } from '../lit-actions-client/utils';
-import { storePrivateKey } from '../service-client';
-import { GeneratePrivateKeyParams, GeneratePrivateKeyResult } from '../types';
 import {
   getFirstSessionSig,
+  getKeyTypeFromNetwork,
   getPkpAccessControlCondition,
   getPkpAddressFromSessionSig,
-} from '../utils';
+} from './utils';
+import { generateKeyWithLitAction } from '../lit-actions-client';
+import { getLitActionCodeOrCid } from '../lit-actions-client/utils';
+import { storePrivateKey } from '../service-client';
+import { GeneratePrivateKeyParams, GeneratePrivateKeyResult } from '../types';
 
 /**
  * Generates a random private key inside a Lit Action, and persists the key and its metadata to the wrapped keys service.
@@ -55,7 +55,6 @@ export async function generatePrivateKey(
       publicKey,
       keyType: getKeyTypeFromNetwork(network),
       dataToEncryptHash,
-      pkpAddress,
       memo,
     },
     litNetwork: litNodeClient.config.litNetwork,