LIT-3920 - wrapped-keys - Batch generate private keys & optionally sign messages in a single operation

local-tests/setup/tinny-environment.ts

@@ -181,7 +181,6 @@ export class TinnyEnvironment {
         return { privateKey: this.processEnvs.PRIVATE_KEYS[index], index }; // Return the key and its index
       } else {
         console.log('[𐬺🧪 Tinny Environment𐬺] No available keys. Waiting...', {
-          privateKeys: this.processEnvs.PRIVATE_KEYS,
           keysInUse: this.processEnvs.KEY_IN_USE,
         }); // Log a message indicating that we are waiting
         // Wait for the specified interval before checking again
@@ -447,8 +446,9 @@ export class TinnyEnvironment {
    * @throws If there is an error sending the funds.
    */
   getFunds = async (walletAddress: string, amount = '0.001') => {
+    const privateKey = await this.getAvailablePrivateKey();
+
     try {
-      const privateKey = await this.getAvailablePrivateKey();
       const provider = new ethers.providers.JsonRpcBatchProvider(this.rpc);
       const wallet = new ethers.Wallet(privateKey.privateKey, provider);
 
@@ -460,6 +460,9 @@ export class TinnyEnvironment {
       await tx.wait();
     } catch (e) {
       throw new Error(`Failed to send funds to ${walletAddress}: ${e}`);
+    } finally {
+      // @ts-expect-error We don't have a user, but this works
+      this.releasePrivateKeyFromUser({ privateKey });
     }
   };
 

local-tests/test.ts

@@ -104,6 +104,13 @@ import { testFailImportWrappedKeysWithExpiredSessionSig } from './tests/wrapped-
 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 { setLitActionsCodeToLocal } from './tests/wrapped-keys/util';
+import { testUseEoaSessionSigsToRequestSingleResponse } from './tests/testUseEoaSessionSigsToRequestSingleResponse';
+
+// Use the current LIT action code to test against
+setLitActionsCodeToLocal();
 
 (async () => {
   console.log('[𐬺🧪 Tinny𐬺] Running tests...');
@@ -118,6 +125,7 @@ import { testSignTransactionWithSolanaEncryptedKey } from './tests/wrapped-keys/
   // --filter=WrappedKey
   const wrappedKeysTests = {
     // -- valid cases
+    testBatchGeneratePrivateKeys,
     testEthereumSignMessageGeneratedKey,
     testEthereumBroadcastTransactionGeneratedKey,
     testEthereumSignMessageWrappedKey,
@@ -163,6 +171,7 @@ import { testSignTransactionWithSolanaEncryptedKey } from './tests/wrapped-keys/
     testUseEoaSessionSigsToEncryptDecryptString,
     testUseEoaSessionSigsToEncryptDecryptFile,
     testUseEoaSessionSigsToEncryptDecryptZip,
+    testUseEoaSessionSigsToRequestSingleResponse,
   };
 
   const pkpSessionSigsTests = {
@@ -282,22 +291,22 @@ import { testSignTransactionWithSolanaEncryptedKey } from './tests/wrapped-keys/
     tests: {
       // testExample,
       // testBundleSpeed,
-      ...eoaSessionSigsTests,
-      ...pkpSessionSigsTests,
-      ...litActionSessionSigsTests,
-      ...litActionIpfsIdSessionSigsTests,
-      ...capacityDelegationTests,
-      ...bareAuthSigTests,
-
-      ...pkpEthersTest.eoaSessionSigs,
-      ...pkpEthersTest.pkpSessionSigs,
-      ...pkpEthersTest.litActionSessionSigs,
-
-      ...litActionCombiningTests.broadcastAndCombine,
-      ...litActionCombiningTests.decryptAndCombine,
-      ...litActionCombiningTests.ecdsaSignAndCombine,
-
-      ...relayerTests,
+      // ...eoaSessionSigsTests,
+      // ...pkpSessionSigsTests,
+      // ...litActionSessionSigsTests,
+      // ...litActionIpfsIdSessionSigsTests,
+      // ...capacityDelegationTests,
+      // ...bareAuthSigTests,
+      //
+      // ...pkpEthersTest.eoaSessionSigs,
+      // ...pkpEthersTest.pkpSessionSigs,
+      // ...pkpEthersTest.litActionSessionSigs,
+      //
+      // ...litActionCombiningTests.broadcastAndCombine,
+      // ...litActionCombiningTests.decryptAndCombine,
+      // ...litActionCombiningTests.ecdsaSignAndCombine,
+      //
+      // ...relayerTests,
       ...wrappedKeysTests,
     },
     devEnv,

local-tests/tests/testUseEoaSessionSigsToRequestSingleResponse.ts

@@ -0,0 +1,56 @@
+import { getEoaSessionSigs } from 'local-tests/setup/session-sigs/get-eoa-session-sigs';
+import { TinnyEnvironment } from 'local-tests/setup/tinny-environment';
+
+/**
+ * Test Commands:
+ * ✅ NETWORK=datil-dev yarn test:local --filter=testUseEoaSessionSigsToRequestSingleResponse
+ * ✅ NETWORK=datil-test yarn test:local --filter=testUseEoaSessionSigsToRequestSingleResponse
+ * ✅ NETWORK=datil yarn test:local --filter=testUseEoaSessionSigsToRequestSingleResponse
+ */
+export const testUseEoaSessionSigsToRequestSingleResponse = async (
+  devEnv: TinnyEnvironment
+) => {
+  const alice = await devEnv.createRandomPerson();
+
+  try {
+    const eoaSessionSigs = await getEoaSessionSigs(devEnv, alice);
+
+    const res = await devEnv.litNodeClient.executeJs({
+      sessionSigs: eoaSessionSigs,
+      code: `(async () => {
+          console.log('hello world')
+        })();`,
+      numResponsesRequired: 1,
+    });
+  } finally {
+    devEnv.releasePrivateKeyFromUser(alice);
+  }
+  console.log('res:', res);
+
+  // Expected output:
+  // {
+  //   success: true,
+  //   signedData: {},
+  //   decryptedData: {},
+  //   claimData: {},
+  //   response: "",
+  //   logs: "hello world\n",
+  // }
+
+  // -- assertions
+  if (res.response) {
+    throw new Error(`Expected "response" to be falsy`);
+  }
+
+  if (!res.logs) {
+    throw new Error(`Expected "logs" in res`);
+  }
+
+  if (!res.logs.includes('hello world')) {
+    throw new Error(`Expected "logs" to include 'hello world'`);
+  }
+
+  if (!res.success) {
+    throw new Error(`Expected "success" in res`);
+  }
+};

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

@@ -0,0 +1,154 @@
+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';
+
+const { batchGeneratePrivateKeys, exportPrivateKey } = api;
+
+async function verifySolanaSignature(
+  solanaResult: BatchGeneratePrivateKeysActionResult,
+  solanaMessageToSign
+) {
+  const {
+    signMessage: { signature },
+    generateEncryptedPrivateKey: { generatedPublicKey },
+  } = solanaResult;
+  const signatureIsValidForPublicKey = nacl.sign.detached.verify(
+    Buffer.from(solanaMessageToSign),
+    bs58.decode(signature),
+    bs58.decode(generatedPublicKey)
+  );
+
+  console.log({ signatureIsValidForPublicKey, signature });
+  if (!signatureIsValidForPublicKey) {
+    throw new Error(
+      `signature: ${signature} doesn't validate for the Solana public key: ${generatedPublicKey}`
+    );
+  }
+}
+async function verifyEvmSignature(
+  litNodeClient,
+  evmResult,
+  messageToSign,
+  pkpSessionSigs
+) {
+  function verifyMessageSignature() {
+    try {
+      return ethers.utils.verifyMessage(
+        messageToSign,
+        evmResult.signMessage.signature
+      );
+    } catch (err) {
+      throw new Error(
+        `When validating signed Ethereum message is valid: ${err.message}`
+      );
+    }
+  }
+
+  const { decryptedPrivateKey } = await exportPrivateKey({
+    litNodeClient,
+    network: 'evm',
+    id: evmResult.generateEncryptedPrivateKey.id,
+    pkpSessionSigs,
+  });
+
+  const recoveredAddress = verifyMessageSignature();
+
+  const wallet = new ethers.Wallet(decryptedPrivateKey);
+
+  console.log({
+    recoveredAddress,
+    walletAddress: wallet.address,
+    signature: evmResult.signMessage.signature,
+  });
+  if (recoveredAddress !== wallet.address) {
+    throw new Error(
+      "Recovered address from verifyMessage doesn't match the wallet address"
+    );
+  }
+}
+
+/**
+ * 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 testBatchGeneratePrivateKeys = 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 { results } = await batchGeneratePrivateKeys({
+      pkpSessionSigs: pkpSessionSigsSigning,
+      actions: [
+        {
+          network: 'evm',
+          signMessageParams: { messageToSign: evmMessageToSign },
+          generateKeyParams: { memo: 'Test evm key' },
+        },
+        {
+          network: 'solana',
+          signMessageParams: { messageToSign: solanaMessageToSign },
+          generateKeyParams: { memo: 'Test solana key' },
+        },
+      ],
+      litNodeClient: devEnv.litNodeClient,
+    });
+
+    if (results.length !== 2) {
+      throw new Error(
+        `Incorrect # of results; expected 2, got ${results.length}`
+      );
+    }
+
+    if (
+      results[0].generateEncryptedPrivateKey.memo !== 'Test evm key' ||
+      results[1].generateEncryptedPrivateKey.memo !== 'Test solana key'
+    ) {
+      throw new Error(
+        'Results not in order sent; expected evm as first result, solana as second'
+      );
+    }
+
+    if (
+      !results[0].signMessage.signature ||
+      !results[1].signMessage.signature
+    ) {
+      throw new Error('Missing message signature in response');
+    }
+
+    console.log('solana verify sig');
+    await verifySolanaSignature(results[1], solanaMessageToSign);
+
+    console.log('evm verify sig');
+    await verifyEvmSignature(
+      devEnv.litNodeClient,
+      results[0],
+      evmMessageToSign,
+      pkpSessionSigsSigning
+    );
+    console.log('results', results);
+
+    log('✅ testBatchGenerateEncryptedKeys');
+  } catch (err) {
+    console.log(err.message, err, err.stack);
+    throw err;
+  } finally {
+    devEnv.releasePrivateKeyFromUser(alice);
+  }
+};

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

@@ -70,11 +70,7 @@ export const testFailEthereumSignTransactionWrappedKeyWithInvalidParam = async (
         id,
       });
     } catch (e: any) {
-      if (
-        e.message.includes(
-          'Error executing the Signing Lit Action: Error: When signing transaction- invalid hexlify value'
-        )
-      ) {
+      if (e.message.includes('invalid hexlify value')) {
         console.log('✅ THIS IS EXPECTED: ', e);
         console.log(e.message);
         console.log(

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

@@ -66,11 +66,7 @@ export const testFailEthereumSignTransactionWrappedKeyWithMissingParam = async (
         id,
       });
     } catch (e: any) {
-      if (
-        e.message.includes(
-          'Error executing the Signing Lit Action: Error: Missing required field: toAddress'
-        )
-      ) {
+      if (e.message.includes('Missing required field: toAddress')) {
         console.log('✅ THIS IS EXPECTED: ', e);
         console.log(e.message);
         console.log(

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

@@ -11,6 +11,7 @@ import {
   clusterApiUrl,
 } from '@solana/web3.js';
 import { getPkpSessionSigs } from 'local-tests/setup/session-sigs/get-pkp-session-sigs';
+import { ethers } from 'ethers';
 
 const { importPrivateKey, signTransactionWithEncryptedKey } = api;
 
@@ -125,7 +126,7 @@ export const testSignTransactionWithSolanaEncryptedKey = async (
     // const confirmation = await solanaConnection.confirmTransaction(signedTx);
     // console.log(confirmation); // { context: { slot: 321490379 }, value: { err: null } }
 
-    const signatureBuffer = Buffer.from(signedTx, 'base64');
+    const signatureBuffer = Buffer.from(ethers.utils.base58.decode(signedTx));
     solanaTransaction.addSignature(solanaKeypair.publicKey, signatureBuffer);
 
     if (!solanaTransaction.verifySignatures()) {