Merge pull request #534 from LIT-Protocol/feat-update-wrapped-keys-ex…

…port

Feat: Update Wrapped keys Export to use Lit Actions

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

@@ -52,6 +52,7 @@ export const testExportWrappedKey = async (devEnv: TinnyEnvironment) => {
   const { decryptedPrivateKey } = await exportPrivateKey({
     pkpSessionSigs: pkpSessionSigsExport,
     litNodeClient: devEnv.litNodeClient,
+    network: 'solana',
   });
 
   if (decryptedPrivateKey !== privateKey) {

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

@@ -49,10 +49,10 @@ export const testGenerateEthereumWrappedKey = async (
 
   console.log(pkpSessionSigsExport);
 
-  // FIXME: Export broken as we can't decrypt data encrypted inside a Lit Action
   const { decryptedPrivateKey } = await exportPrivateKey({
     pkpSessionSigs: pkpSessionSigsExport,
     litNodeClient: devEnv.litNodeClient,
+    network: 'evm',
   });
 
   const wallet = new ethers.Wallet(decryptedPrivateKey);

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

@@ -4,7 +4,7 @@ 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 { Keypair } from '@solana/web3.js';
 
 const { generatePrivateKey, signMessageWithEncryptedKey, exportPrivateKey } =
   api;
@@ -83,16 +83,18 @@ export const testGenerateSolanaWrappedKey = async (
     new Date(Date.now() + 1000 * 60 * 10).toISOString()
   ); // 10 mins expiry
 
-  // FIXME: Export broken as we can't decrypt data encrypted inside a Lit Action
   const { decryptedPrivateKey } = await exportPrivateKey({
     pkpSessionSigs: pkpSessionSigsExport,
     litNodeClient: devEnv.litNodeClient,
+    network: 'solana',
   });
 
-  const wallet = new ethers.Wallet(decryptedPrivateKey);
-  const decryptedPublicKey = wallet.publicKey;
+  const solanaKeyPair = Keypair.fromSecretKey(
+    Buffer.from(decryptedPrivateKey, 'hex')
+  );
+  const decryptedPublicKey = solanaKeyPair.publicKey;
 
-  if (decryptedPublicKey !== generatedPublicKey) {
+  if (decryptedPublicKey.toString() !== generatedPublicKey) {
     throw new Error(
       `Decrypted decryptedPublicKey: ${decryptedPublicKey} doesn't match with the original generatedPublicKey: ${generatedPublicKey}`
     );

packages/wrapped-keys/esbuild.config.js

@@ -25,4 +25,12 @@ const esbuild = require('esbuild');
     outdir: './src/lib/litActions/ethereum/dist',
     inject: ['./buffer.shim.js'],
   });
+  await esbuild.build({
+    entryPoints: ['./src/lib/litActions/common/src/exportPrivateKey.js'],
+    bundle: true,
+    minify: true,
+    sourcemap: false,
+    outdir: './src/lib/litActions/common/dist',
+    inject: ['./buffer.shim.js'],
+  });
 })();

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

@@ -1,55 +1,39 @@
-import { decryptToString } from '@lit-protocol/encryption';
+import { exportPrivateKeyWithLitAction } from '../lit-actions-client';
 
-import { CHAIN_ETHEREUM, LIT_PREFIX } from '../constants';
 import { fetchPrivateKeyMetadata } from '../service-client';
 import { ExportPrivateKeyParams, ExportPrivateKeyResult } from '../types';
-import {
-  getFirstSessionSig,
-  getPkpAccessControlCondition,
-  getPkpAddressFromSessionSig,
-} from '../utils';
+import { getFirstSessionSig, getPkpAccessControlCondition } from '../utils';
+import { getLitActionCid } from '../lit-actions-client/utils';
 
-/** Exports a previously persisted private key from the wrapped keys service for direct use by the caller, along with the keys metadata
+/**
+ * Exports a previously persisted private key from the wrapped keys service for direct use by the caller, along with the keys metadata.
+ * This method fetches the encrypted key from the wrapped keys service, then executes a Lit Action that decrypts the key inside the LIT action and
+ * removes the salt from the decrypted key.
  *
  * @param { ExportPrivateKeyParams } params Parameters required to export the private key
+ *
+ * @returns { Promise<ExportPrivateKeyResult> } - The decrypted private key of the Wrapped Key along with all the associated key info and LIT PKP Address associated with the Wrapped Key
  */
 export async function exportPrivateKey(
   params: ExportPrivateKeyParams
 ): Promise<ExportPrivateKeyResult> {
-  const { pkpSessionSigs, litNodeClient } = params;
+  const { litNodeClient, network, pkpSessionSigs } = params;
 
   const sessionSig = getFirstSessionSig(pkpSessionSigs);
-  const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
-  const allowPkpAddressToDecrypt = getPkpAccessControlCondition(pkpAddress);
-
-  const privateKeyMetadata = await fetchPrivateKeyMetadata({
+  const storedKeyMetadata = await fetchPrivateKeyMetadata({
     sessionSig,
     litNetwork: litNodeClient.config.litNetwork,
   });
 
-  const { ciphertext, dataToEncryptHash, ...privateKeyMetadataMinusEncrypted } =
-    privateKeyMetadata;
-
-  const decryptedPrivateKey = await decryptToString(
-    {
-      accessControlConditions: [allowPkpAddressToDecrypt],
-      chain: CHAIN_ETHEREUM,
-      ciphertext,
-      dataToEncryptHash,
-      sessionSigs: pkpSessionSigs,
-    },
-    litNodeClient
+  const allowPkpAddressToDecrypt = getPkpAccessControlCondition(
+    storedKeyMetadata.pkpAddress
   );
 
-  // It will be of the form lit_<privateKey>
-  if (!decryptedPrivateKey.startsWith(LIT_PREFIX)) {
-    throw new Error(
-      `PKey was not encrypted with salt; all wrapped keys must be prefixed with '${LIT_PREFIX}'`
-    );
-  }
-
-  return {
-    decryptedPrivateKey: decryptedPrivateKey.slice(LIT_PREFIX.length),
-    ...privateKeyMetadataMinusEncrypted,
-  };
+  return exportPrivateKeyWithLitAction({
+    ...params,
+    litActionIpfsCid: getLitActionCid(network, 'exportPrivateKey'),
+    accessControlConditions: [allowPkpAddressToDecrypt],
+    pkpSessionSigs,
+    storedKeyMetadata,
+  });
 }

packages/wrapped-keys/src/lib/lit-actions-client/constants.ts

@@ -2,17 +2,21 @@ import { LitCidRepository } from './types';
 
 const LIT_ACTION_CID_REPOSITORY: LitCidRepository = Object.freeze({
   signTransaction: Object.freeze({
-    evm: 'QmdYUhPCCK5hpDWMK1NiDLNLG6RZQy61QE4J7dBm1Y2nbA',
-    solana: 'QmSi9GL2weCFEP1SMAUw5PDpZRr436Zt3tLUNrSECPA5dT',
+    evm: 'QmSsgmy1N1zFZ5yNPCY7QWQZwrYRLuYUJF1VDygJc7L26o',
+    solana: 'QmdkcMmrtqWSQ8VrPr8KwzuzZnAxGJoVDeZP3NKTWCMZCg',
   }),
   signMessage: Object.freeze({
-    evm: 'QmTMGcyp77NeppGaqF2DmE1F8GXTSxQYzXCrbE7hNudUWx',
-    solana: 'QmUxnWS8VU9QwZRdyfwsEtvhJZcvcdrannjokgZoA1sesy',
+    evm: 'QmWVW51FBH5j3wwaMVy8MR1QyzJgEjuaPh1yqwSGXRCENx',
+    solana: 'QmSPcfFhLofjhNDd5ZdVbhw53zmbR8oV4C2585Bm7C8izH',
   }),
   generateEncryptedKey: Object.freeze({
     evm: 'QmaoPMSqcze3NW3KSA75ecWSkcmWT1J7kVr8LyJPCKRvHd',
     solana: 'QmdRBXYLYvcNHrChmsZ2jFDY8dA99CcSdqHo3p1ES3UThL',
   }),
+  exportPrivateKey: Object.freeze({
+    evm: 'Qmb5ZAm1EZRL7dYTtyYxkPxx4kBmoCjjzcgdrJH9cKMXxR',
+    solana: 'Qmb5ZAm1EZRL7dYTtyYxkPxx4kBmoCjjzcgdrJH9cKMXxR',
+  }),
 });
 
 export { LIT_ACTION_CID_REPOSITORY };

packages/wrapped-keys/src/lib/lit-actions-client/export-private-key.ts

@@ -0,0 +1,47 @@
+import { AccessControlConditions } from '@lit-protocol/types';
+
+import { postLitActionValidation } from './utils';
+import { ExportPrivateKeyParams, StoredKeyMetadata } from '../types';
+
+interface SignMessageWithLitActionParams extends ExportPrivateKeyParams {
+  accessControlConditions: AccessControlConditions;
+  storedKeyMetadata: StoredKeyMetadata;
+  litActionIpfsCid: string;
+}
+
+export async function exportPrivateKeyWithLitAction(
+  args: SignMessageWithLitActionParams
+) {
+  const {
+    accessControlConditions,
+    litNodeClient,
+    pkpSessionSigs,
+    litActionIpfsCid,
+    storedKeyMetadata,
+  } = args;
+
+  const {
+    pkpAddress,
+    ciphertext,
+    dataToEncryptHash,
+    ...storeKeyMetadataMinusEncryptedAndPkp
+  } = storedKeyMetadata;
+  const result = await litNodeClient.executeJs({
+    sessionSigs: pkpSessionSigs,
+    ipfsId: litActionIpfsCid,
+    jsParams: {
+      pkpAddress,
+      ciphertext,
+      dataToEncryptHash,
+      accessControlConditions,
+    },
+  });
+
+  const decryptedPrivateKey = postLitActionValidation(result);
+
+  return {
+    decryptedPrivateKey,
+    pkpAddress,
+    ...storeKeyMetadataMinusEncryptedAndPkp,
+  };
+}

packages/wrapped-keys/src/lib/lit-actions-client/index.ts

@@ -1,9 +1,11 @@
 import { generateKeyWithLitAction } from './generate-key';
 import { signMessageWithLitAction } from './sign-message';
 import { signTransactionWithLitAction } from './sign-transaction';
+import { exportPrivateKeyWithLitAction } from './export-private-key';
 
 export {
   generateKeyWithLitAction,
   signTransactionWithLitAction,
   signMessageWithLitAction,
+  exportPrivateKeyWithLitAction,
 };

packages/wrapped-keys/src/lib/lit-actions-client/types.ts

@@ -3,7 +3,8 @@ import { Network } from '../types';
 export type LitActionType =
   | 'signTransaction'
   | 'signMessage'
-  | 'generateEncryptedKey';
+  | 'generateEncryptedKey'
+  | 'exportPrivateKey';
 
 export type LitCidRepositoryEntry = Readonly<Record<Network, string>>;
 

packages/wrapped-keys/src/lib/litActions/common/src/exportPrivateKey.js

@@ -0,0 +1,43 @@
+const { removeSaltFromDecryptedKey } = require('../../utils');
+
+/**
+ *
+ * Exports the private key after decrypting and removing the salt from it.
+ *
+ * @jsParam pkpAddress - The Eth address of the PKP which is associated with the Wrapped Key
+ * @jsParam ciphertext - For the encrypted Wrapped Key
+ * @jsParam dataToEncryptHash - For the encrypted Wrapped Key
+ * @jsParam accessControlConditions - The access control condition that allows only the pkpAddress to decrypt the Wrapped Key
+ *
+ * @returns { Promise<string> } - Returns a decrypted private key.
+ */
+
+(async () => {
+  let decryptedPrivateKey;
+  try {
+    decryptedPrivateKey = await Lit.Actions.decryptToSingleNode({
+      accessControlConditions,
+      ciphertext,
+      dataToEncryptHash,
+      chain: 'ethereum',
+      authSig: null,
+    });
+  } catch (err) {
+    const errorMessage =
+      'Error: When decrypting to a single node- ' + err.message;
+    Lit.Actions.setResponse({ response: errorMessage });
+    return;
+  }
+
+  if (!decryptedPrivateKey) {
+    // Exit the nodes which don't have the decryptedData
+    return;
+  }
+
+  try {
+    const privateKey = removeSaltFromDecryptedKey(decryptedPrivateKey);
+    Lit.Actions.setResponse({ response: privateKey });
+  } catch (err) {
+    Lit.Actions.setResponse({ response: err.message });
+  }
+})();

packages/wrapped-keys/src/lib/litActions/ethereum/src/signMessageWithEthereumEncryptedKey.js

@@ -35,7 +35,15 @@ const { removeSaltFromDecryptedKey } = require('../../utils');
     return;
   }
 
-  const privateKey = removeSaltFromDecryptedKey(decryptedPrivateKey);
+  let privateKey;
+  try {
+    privateKey = removeSaltFromDecryptedKey(decryptedPrivateKey);
+    Lit.Actions.setResponse({ response: privateKey });
+  } catch (err) {
+    Lit.Actions.setResponse({ response: err.message });
+    return;
+  }
+  const wallet = new ethers.Wallet(privateKey);
 
   try {
     const signature = await wallet.signMessage(messageToSign);

packages/wrapped-keys/src/lib/litActions/ethereum/src/signTransactionWithEthereumEncryptedKey.js

@@ -63,7 +63,14 @@ const { removeSaltFromDecryptedKey } = require('../../utils');
     return;
   }
 
-  const privateKey = removeSaltFromDecryptedKey(decryptedPrivateKey);
+  let privateKey;
+  try {
+    privateKey = removeSaltFromDecryptedKey(decryptedPrivateKey);
+    Lit.Actions.setResponse({ response: privateKey });
+  } catch (err) {
+    Lit.Actions.setResponse({ response: err.message });
+    return;
+  }
   const wallet = new ethers.Wallet(privateKey);
 
   let nonce;

packages/wrapped-keys/src/lib/litActions/solana/src/signMessageWithSolanaEncryptedKey.js

@@ -39,7 +39,14 @@ const { removeSaltFromDecryptedKey } = require('../../utils');
     return;
   }
 
-  const privateKey = removeSaltFromDecryptedKey(decryptedPrivateKey);
+  let privateKey;
+  try {
+    privateKey = removeSaltFromDecryptedKey(decryptedPrivateKey);
+    Lit.Actions.setResponse({ response: privateKey });
+  } catch (err) {
+    Lit.Actions.setResponse({ response: err.message });
+    return;
+  }
   const solanaKeyPair = Keypair.fromSecretKey(Buffer.from(privateKey, 'hex'));
 
   let signature;

packages/wrapped-keys/src/lib/litActions/solana/src/signTransactionWithSolanaEncryptedKey.js

@@ -62,7 +62,14 @@ const { removeSaltFromDecryptedKey } = require('../../utils');
     return;
   }
 
-  const privateKey = removeSaltFromDecryptedKey(decryptedPrivateKey);
+  let privateKey;
+  try {
+    privateKey = removeSaltFromDecryptedKey(decryptedPrivateKey);
+    Lit.Actions.setResponse({ response: privateKey });
+  } catch (err) {
+    Lit.Actions.setResponse({ response: err.message });
+    return;
+  }
 
   const solanaKeyPair = Keypair.fromSecretKey(
     Uint8Array.from(Buffer.from(privateKey, 'hex'))

packages/wrapped-keys/src/lib/litActions/utils.js

@@ -3,7 +3,7 @@ import { LIT_PREFIX } from '../constants';
 export function removeSaltFromDecryptedKey(decryptedPrivateKey) {
   if (!decryptedPrivateKey.startsWith(LIT_PREFIX)) {
     throw new Error(
-      `PKey was not encrypted with salt; all wrapped keys must be prefixed with '${LIT_PREFIX}'`
+      `Error: PKey was not encrypted with salt; all wrapped keys must be prefixed with '${LIT_PREFIX}'`
     );
   }
 

packages/wrapped-keys/src/lib/types.ts

@@ -70,7 +70,7 @@ export type StoreEncryptedKeyMetadataParams = BaseApiParams &
  * @extends BaseApiParams
  *
  */
-export type ExportPrivateKeyParams = BaseApiParams;
+export type ExportPrivateKeyParams = BaseApiParams & ApiParamsSupportedNetworks;
 
 /** Includes the decrypted private key and metadata that was stored alongside it in the wrapped keys service
  *