Skip to content

Commit e9d68a3

Browse files
MaximusHaximusMaximusHaximus
authored
Merge pull request #686 from LIT-Protocol/LIT-3956-write-keys-batch
LIT-3956 - wrapped-keys - Allow batch of keys to be stored atomically (storeEncryptedKeyBatch)
2 parents 1595610 + f777a7e commit e9d68a3

24 files changed

+465
-181
lines changed

‎local-tests/test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ import { testExportWrappedKey } from './tests/wrapped-keys/testExportWrappedKey'
105105
import { testSignMessageWithSolanaEncryptedKey } from './tests/wrapped-keys/testSignMessageWithSolanaEncryptedKey';
106106
import { testSignTransactionWithSolanaEncryptedKey } from './tests/wrapped-keys/testSignTransactionWithSolanaEncryptedKey';
107107
import { testBatchGeneratePrivateKeys } from './tests/wrapped-keys/testBatchGeneratePrivateKeys';
108+
import { testFailBatchGeneratePrivateKeysAtomic } from './tests/wrapped-keys/testFailStoreEncryptedKeyBatchIsAtomic';
108109

109110
import { setLitActionsCodeToLocal } from './tests/wrapped-keys/util';
110111
import { testUseEoaSessionSigsToRequestSingleResponse } from './tests/testUseEoaSessionSigsToRequestSingleResponse';
@@ -151,6 +152,7 @@ setLitActionsCodeToLocal();
151152
testFailEthereumSignTransactionWrappedKeyWithMissingParam,
152153
testFailEthereumSignTransactionWrappedKeyWithInvalidParam,
153154
testFailEthereumSignTransactionWrappedKeyInvalidDecryption,
155+
testFailBatchGeneratePrivateKeysAtomic,
154156

155157
// -- import wrapped keys
156158
testFailImportWrappedKeysWithSamePrivateKey,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import bs58 from 'bs58';
77
import { ethers } from 'ethers';
88
import { BatchGeneratePrivateKeysActionResult } from '../../../packages/wrapped-keys/src/lib/types';
99

10-
const { batchGeneratePrivateKeys, exportPrivateKey } = api;
10+
const { batchGeneratePrivateKeys } = api;
1111

1212
async function verifySolanaSignature(
1313
solanaResult: BatchGeneratePrivateKeysActionResult,

‎local-tests/tests/wrapped-keys/testFailEthereumSignTransactionWrappedKeyInvalidDecryption.ts‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
11
import { log } from '@lit-protocol/misc';
22
import { ethers } from 'ethers';
33
import { TinnyEnvironment } from 'local-tests/setup/tinny-environment';
4-
import { EthereumLitTransaction } from '@lit-protocol/wrapped-keys';
54
import { getPkpSessionSigs } from 'local-tests/setup/session-sigs/get-pkp-session-sigs';
6-
import { getPkpAccessControlCondition } from 'packages/wrapped-keys/src/lib/utils';
75
import { encryptString } from '@lit-protocol/encryption';
86
import { LIT_PREFIX } from 'packages/wrapped-keys/src/lib/constants';
97
import { LIT_ACTION_CID_REPOSITORY } from '../../../packages/wrapped-keys/src/lib/lit-actions-client/constants';
108
import { getBaseTransactionForNetwork } from './util';
119
import { GLOBAL_OVERWRITE_IPFS_CODE_BY_NETWORK } from '@lit-protocol/constants';
10+
import { getPkpAccessControlCondition } from '../../../packages/wrapped-keys/src/lib/api/utils';
1211

1312
/**
1413
* Test Commands:

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

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
import { log } from '@lit-protocol/misc';
2+
import { TinnyEnvironment } from 'local-tests/setup/tinny-environment';
3+
import { api } from '@lit-protocol/wrapped-keys';
4+
import { getPkpSessionSigs } from 'local-tests/setup/session-sigs/get-pkp-session-sigs';
5+
import { batchGenerateKeysWithLitAction } from '../../../packages/wrapped-keys/src/lib/lit-actions-client';
6+
import { getLitActionCodeOrCidCommon } from '../../../packages/wrapped-keys/src/lib/lit-actions-client/utils';
7+
import {
8+
getFirstSessionSig,
9+
getKeyTypeFromNetwork,
10+
getPkpAccessControlCondition,
11+
getPkpAddressFromSessionSig,
12+
} from '../../../packages/wrapped-keys/src/lib/api/utils';
13+
import { listEncryptedKeyMetadata } from '../../../packages/wrapped-keys/src/lib/api';
14+
15+
const { storeEncryptedKeyBatch } = api;
16+
17+
/**
18+
* Test Commands:
19+
* ✅ NETWORK=datil-dev yarn test:local --filter=testSignMessageWithSolanaEncryptedKey
20+
* ✅ NETWORK=datil-test yarn test:local --filter=testSignMessageWithSolanaEncryptedKey
21+
* ✅ NETWORK=localchain yarn test:local --filter=testSignMessageWithSolanaEncryptedKey
22+
*/
23+
export const testFailBatchGeneratePrivateKeysAtomic = async (
24+
devEnv: TinnyEnvironment
25+
) => {
26+
const alice = await devEnv.createRandomPerson();
27+
28+
try {
29+
const pkpSessionSigsSigning = await getPkpSessionSigs(
30+
devEnv,
31+
alice,
32+
null,
33+
new Date(Date.now() + 1000 * 60 * 10).toISOString()
34+
); // 10 mins expiry
35+
36+
const solanaMessageToSign = 'This is a test solana message';
37+
const evmMessageToSign = 'This is a test evm message';
38+
39+
const sessionSig = getFirstSessionSig(pkpSessionSigsSigning);
40+
const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
41+
42+
const allowPkpAddressToDecrypt = getPkpAccessControlCondition(pkpAddress);
43+
44+
const { litActionCode, litActionIpfsCid } = getLitActionCodeOrCidCommon(
45+
'batchGenerateEncryptedKeys'
46+
);
47+
48+
const actionResults = await batchGenerateKeysWithLitAction({
49+
litNodeClient: devEnv.litNodeClient,
50+
litActionIpfsCid: litActionCode ? undefined : litActionIpfsCid,
51+
litActionCode: litActionCode ? litActionCode : undefined,
52+
accessControlConditions: [allowPkpAddressToDecrypt],
53+
actions: [
54+
{
55+
network: 'evm',
56+
signMessageParams: { messageToSign: evmMessageToSign },
57+
generateKeyParams: { memo: 'Test evm key' },
58+
},
59+
{
60+
network: 'solana',
61+
signMessageParams: { messageToSign: solanaMessageToSign },
62+
generateKeyParams: { memo: 'Test solana key' },
63+
},
64+
],
65+
pkpSessionSigs: pkpSessionSigsSigning,
66+
});
67+
68+
const keyParamsBatch = actionResults.map((keyData) => {
69+
const { generateEncryptedPrivateKey, network } = keyData;
70+
return {
71+
...generateEncryptedPrivateKey,
72+
keyType: getKeyTypeFromNetwork(network),
73+
};
74+
});
75+
76+
// Intentional failure to persist due to missing publicKey
77+
delete keyParamsBatch[0].publicKey;
78+
79+
try {
80+
await storeEncryptedKeyBatch({
81+
pkpSessionSigs: pkpSessionSigsSigning,
82+
litNodeClient: devEnv.litNodeClient,
83+
keyBatch: keyParamsBatch,
84+
});
85+
86+
throw new Error(
87+
'storeEncryptedKeyBatch() succeeded but we expected it to fail!'
88+
);
89+
} catch (err) {
90+
// We expect `storeEncryptedKeyBatch` to fail w/ a specific error
91+
if (
92+
err.message.includes(
93+
'storeEncryptedKeyBatch() succeeded but we expected it to fail!'
94+
) ||
95+
!err.message.includes(
96+
'keyParamsBatch[0]: Missing "publicKey" parameter in request'
97+
)
98+
) {
99+
throw err;
100+
}
101+
102+
try {
103+
const keys = await listEncryptedKeyMetadata({
104+
litNodeClient: devEnv.litNodeClient,
105+
pkpSessionSigs: pkpSessionSigsSigning,
106+
});
107+
108+
console.error(
109+
'Got a value back we shouldnt have from listEncryptedKeyMetadata()',
110+
keys
111+
);
112+
113+
throw new Error(
114+
'Expected `listEncryptedKeyMetadata() to fail, but it didnt!`'
115+
);
116+
} catch (err) {
117+
if (err.message.includes('No keys exist for pkpAddress')) {
118+
log('✅ testFailBatchGeneratePrivateKeysAtomic');
119+
} else {
120+
throw err;
121+
}
122+
}
123+
}
124+
} catch (err) {
125+
console.log(err.message, err, err.stack);
126+
throw err;
127+
} finally {
128+
devEnv.releasePrivateKeyFromUser(alice);
129+
}
130+
};

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,8 @@ import { TinnyEnvironment } from 'local-tests/setup/tinny-environment';
33
import { api } from '@lit-protocol/wrapped-keys';
44
import { getPkpSessionSigs } from 'local-tests/setup/session-sigs/get-pkp-session-sigs';
55
import { randomSolanaPrivateKey } from 'local-tests/setup/tinny-utils';
6-
import { listPrivateKeyMetadata } from '../../../packages/wrapped-keys/src/lib/service-client';
7-
import { getFirstSessionSig } from '../../../packages/wrapped-keys/src/lib/utils';
86

9-
const { importPrivateKey } = api;
7+
const { importPrivateKey, listEncryptedKeyMetadata } = api;
108

119
/**
1210
* Test Commands:
@@ -44,9 +42,9 @@ export const testImportWrappedKey = async (devEnv: TinnyEnvironment) => {
4442
);
4543
}
4644

47-
const keys = await listPrivateKeyMetadata({
48-
sessionSig: getFirstSessionSig(pkpSessionSigs),
49-
litNetwork: devEnv.litNodeClient.config.litNetwork,
45+
const keys = await listEncryptedKeyMetadata({
46+
pkpSessionSigs,
47+
litNodeClient: devEnv.litNodeClient,
5048
});
5149

5250
if (keys.length !== 1 || keys[0].id !== id) {

‎packages/wrapped-keys/src/index.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import {
88
storeEncryptedKey,
99
listEncryptedKeyMetadata,
1010
batchGeneratePrivateKeys,
11+
storeEncryptedKeyBatch,
1112
} from './lib/api';
1213
import {
1314
CHAIN_ETHEREUM,
@@ -45,11 +46,13 @@ import type {
4546
SignTransactionParamsSupportedEvm,
4647
SignTransactionParamsSupportedSolana,
4748
StoreEncryptedKeyParams,
49+
StoreEncryptedKeyBatchParams,
4850
StoredKeyData,
4951
StoredKeyMetadata,
5052
ListEncryptedKeyMetadataParams,
5153
StoreEncryptedKeyResult,
5254
ImportPrivateKeyResult,
55+
StoreEncryptedKeyBatchResult,
5356
} from './lib/types';
5457

5558
export const constants = {
@@ -70,6 +73,7 @@ export const api = {
7073
signMessageWithEncryptedKey,
7174
signTransactionWithEncryptedKey,
7275
storeEncryptedKey,
76+
storeEncryptedKeyBatch,
7377
batchGeneratePrivateKeys,
7478
};
7579

@@ -101,6 +105,8 @@ export {
101105
SignTransactionWithEncryptedKeyParams,
102106
StoreEncryptedKeyParams,
103107
StoreEncryptedKeyResult,
108+
StoreEncryptedKeyBatchParams,
109+
StoreEncryptedKeyBatchResult,
104110
StoredKeyData,
105111
StoredKeyMetadata,
106112
SupportedNetworks,

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

Lines changed: 37 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
1-
import { getKeyTypeFromNetwork } from './utils';
1+
import {
2+
getFirstSessionSig,
3+
getKeyTypeFromNetwork,
4+
getPkpAccessControlCondition,
5+
getPkpAddressFromSessionSig,
6+
} from './utils';
27
import { batchGenerateKeysWithLitAction } from '../lit-actions-client';
38
import { getLitActionCodeOrCidCommon } from '../lit-actions-client/utils';
4-
import { storePrivateKey } from '../service-client';
9+
import { storePrivateKeyBatch } from '../service-client';
510
import {
611
BatchGeneratePrivateKeysActionResult,
712
BatchGeneratePrivateKeysParams,
813
BatchGeneratePrivateKeysResult,
914
} from '../types';
10-
import {
11-
getFirstSessionSig,
12-
getPkpAccessControlCondition,
13-
getPkpAddressFromSessionSig,
14-
} from '../utils';
1515

1616
/**
1717
* TODO: Document batch behaviour
@@ -41,34 +41,39 @@ export async function batchGeneratePrivateKeys(
4141
pkpSessionSigs,
4242
});
4343

44-
const results = await Promise.all(
45-
actionResults.map(
46-
async (result): Promise<BatchGeneratePrivateKeysActionResult> => {
47-
const { generateEncryptedPrivateKey, network } = result;
44+
const keyParamsBatch = actionResults.map((keyData) => {
45+
const { generateEncryptedPrivateKey, network } = keyData;
46+
return {
47+
...generateEncryptedPrivateKey,
48+
keyType: getKeyTypeFromNetwork(network),
49+
};
50+
});
51+
52+
const { ids } = await storePrivateKeyBatch({
53+
sessionSig,
54+
storedKeyMetadataBatch: keyParamsBatch,
55+
litNetwork: litNodeClient.config.litNetwork,
56+
});
4857

49-
const signature = result.signMessage?.signature;
58+
const results = actionResults.map(
59+
(actionResult, ndx): BatchGeneratePrivateKeysActionResult => {
60+
const {
61+
generateEncryptedPrivateKey: { memo, publicKey },
62+
} = actionResult;
63+
const id = ids[ndx]; // Result of writes is in same order as provided
5064

51-
const { id } = await storePrivateKey({
52-
sessionSig,
53-
storedKeyMetadata: {
54-
...generateEncryptedPrivateKey,
55-
keyType: getKeyTypeFromNetwork(network),
56-
pkpAddress,
57-
},
58-
litNetwork: litNodeClient.config.litNetwork,
59-
});
65+
const signature = actionResult.signMessage?.signature;
6066

61-
return {
62-
...(signature ? { signMessage: { signature } } : {}),
63-
generateEncryptedPrivateKey: {
64-
memo: generateEncryptedPrivateKey.memo,
65-
id,
66-
generatedPublicKey: generateEncryptedPrivateKey.publicKey,
67-
pkpAddress,
68-
},
69-
};
70-
}
71-
)
67+
return {
68+
...(signature ? { signMessage: { signature } } : {}),
69+
generateEncryptedPrivateKey: {
70+
memo: memo,
71+
id,
72+
generatedPublicKey: publicKey,
73+
pkpAddress,
74+
},
75+
};
76+
}
7277
);
7378

7479
return { pkpAddress, results };

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
1+
import {
2+
getFirstSessionSig,
3+
getPkpAccessControlCondition,
4+
getPkpAddressFromSessionSig,
5+
} from './utils';
16
import { exportPrivateKeyWithLitAction } from '../lit-actions-client';
27
import { getLitActionCodeOrCid } from '../lit-actions-client/utils';
38
import { fetchPrivateKey } from '../service-client';
49
import { ExportPrivateKeyParams, ExportPrivateKeyResult } from '../types';
5-
import { getFirstSessionSig, getPkpAccessControlCondition } from '../utils';
610

711
/**
812
* 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(
1923
const { litNodeClient, network, pkpSessionSigs, id } = params;
2024

2125
const sessionSig = getFirstSessionSig(pkpSessionSigs);
26+
const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
27+
2228
const storedKeyMetadata = await fetchPrivateKey({
29+
pkpAddress,
2330
id,
2431
sessionSig,
2532
litNetwork: litNodeClient.config.litNetwork,

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

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,13 @@
1-
import { getKeyTypeFromNetwork } from './utils';
2-
import { generateKeyWithLitAction } from '../lit-actions-client';
3-
import { getLitActionCodeOrCid } from '../lit-actions-client/utils';
4-
import { storePrivateKey } from '../service-client';
5-
import { GeneratePrivateKeyParams, GeneratePrivateKeyResult } from '../types';
61
import {
72
getFirstSessionSig,
3+
getKeyTypeFromNetwork,
84
getPkpAccessControlCondition,
95
getPkpAddressFromSessionSig,
10-
} from '../utils';
6+
} from './utils';
7+
import { generateKeyWithLitAction } from '../lit-actions-client';
8+
import { getLitActionCodeOrCid } from '../lit-actions-client/utils';
9+
import { storePrivateKey } from '../service-client';
10+
import { GeneratePrivateKeyParams, GeneratePrivateKeyResult } from '../types';
1111

1212
/**
1313
* 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(
5555
publicKey,
5656
keyType: getKeyTypeFromNetwork(network),
5757
dataToEncryptHash,
58-
pkpAddress,
5958
memo,
6059
},
6160
litNetwork: litNodeClient.config.litNetwork,

0 commit comments

Comments
 (0)