Skip to content

Commit 4adf506

Browse files
AnsonhkgAnsonhkg
authored
Merge pull request #1006 from LIT-Protocol/feature/jss-141-add-update-function-for-wrapped-keys
Feature/jss 141 add update function for wrapped keys
2 parents 2285897 + cd1d347 commit 4adf506

File tree

11 files changed

+281
-5
lines changed

11 files changed

+281
-5
lines changed

‎.github/workflows/release-docker-images.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ jobs:
2828
runs-on: ubuntu-latest
2929
strategy:
3030
matrix:
31-
app: [lit-auth-server, lit-login-server]
31+
app: [lit-auth-server, lit-login-server, explorer]
3232
steps:
3333
- name: Checkout
3434
uses: actions/checkout@v4

‎local-tests/test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ import { testFailImportWrappedKeysWithMaxExpirySessionSig } from './tests/wrappe
106106
import { testFailImportWrappedKeysWithInvalidSessionSig } from './tests/wrapped-keys/testFailImportWrappedKeysWithInvalidSessionSig';
107107
import { testFailImportWrappedKeysWithExpiredSessionSig } from './tests/wrapped-keys/testFailImportWrappedKeysWithExpiredSessionSig';
108108
import { testExportWrappedKey } from './tests/wrapped-keys/testExportWrappedKey';
109+
import { testUpdateWrappedKey } from './tests/wrapped-keys/testUpdateWrappedKey';
109110
import { testSignMessageWithSolanaEncryptedKey } from './tests/wrapped-keys/testSignMessageWithSolanaEncryptedKey';
110111
import { testSignTransactionWithSolanaEncryptedKey } from './tests/wrapped-keys/testSignTransactionWithSolanaEncryptedKey';
111112
import { testBatchGeneratePrivateKeys } from './tests/wrapped-keys/testBatchGeneratePrivateKeys';
@@ -158,6 +159,7 @@ setLitActionsCodeToLocal();
158159

159160
// -- export wrapped keys
160161
testExportWrappedKey,
162+
testUpdateWrappedKey,
161163

162164
// -- solana wrapped keys
163165
testSignMessageWithSolanaEncryptedKey,

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

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
import { log } from '@lit-protocol/misc';
2+
import { randomBytes } from 'crypto';
3+
4+
import { api } from '@lit-protocol/wrapped-keys';
5+
import { TinnyEnvironment } from 'local-tests/setup/tinny-environment';
6+
import { getPkpSessionSigs } from 'local-tests/setup/session-sigs/get-pkp-session-sigs';
7+
8+
const { generatePrivateKey, getEncryptedKey, updateEncryptedKey } = api;
9+
10+
/**
11+
* Test Commands:
12+
* ✅ NETWORK=datil-dev yarn test:local --filter=testUpdateWrappedKey
13+
* ✅ NETWORK=datil-test yarn test:local --filter=testUpdateWrappedKey
14+
* ✅ NETWORK=custom yarn test:local --filter=testUpdateWrappedKey
15+
*/
16+
export const testUpdateWrappedKey = async (devEnv: TinnyEnvironment) => {
17+
const alice = await devEnv.createRandomPerson();
18+
19+
try {
20+
console.log('1. Fetch PKP session sigs');
21+
const pkpSessionSigs = await getPkpSessionSigs(
22+
devEnv,
23+
alice,
24+
null,
25+
new Date(Date.now() + 1000 * 60 * 10).toISOString()
26+
); // 10 mins expiry
27+
28+
console.log('2. Generate a wrapped key');
29+
const { id, pkpAddress } = await generatePrivateKey({
30+
pkpSessionSigs,
31+
network: 'evm',
32+
litNodeClient: devEnv.litNodeClient,
33+
memo: 'Test update key',
34+
});
35+
console.log({ id, pkpAddress });
36+
37+
console.log('3. Fetch initial encrypted key (without versions)');
38+
const getEncryptedKeyFirst = await getEncryptedKey({
39+
pkpSessionSigs,
40+
litNodeClient: devEnv.litNodeClient,
41+
id,
42+
});
43+
console.log({ getEncryptedKeyFirst });
44+
const newCiphertext1 = randomBytes(48).toString('base64');
45+
const newCiphertext2 = randomBytes(48).toString('base64');
46+
47+
console.log('4. Update encrypted key with new ciphertext/memo');
48+
const updateEncryptedKeyFirst = await updateEncryptedKey({
49+
pkpSessionSigs,
50+
litNodeClient: devEnv.litNodeClient,
51+
id,
52+
ciphertext: newCiphertext1,
53+
memo: 'rotated memo',
54+
});
55+
console.log({ updateEncryptedKeyFirst });
56+
57+
if (updateEncryptedKeyFirst.pkpAddress !== pkpAddress) {
58+
throw new Error('Updated key pkpAddress mismatch');
59+
}
60+
61+
console.log('5. Second update to generate another version');
62+
const updateEncryptedKeySecond = await updateEncryptedKey({
63+
pkpSessionSigs,
64+
litNodeClient: devEnv.litNodeClient,
65+
id,
66+
ciphertext: newCiphertext2,
67+
memo: 'rotated memo v2',
68+
});
69+
console.log({ updateEncryptedKeySecond });
70+
71+
console.log('6. Fetch updated key including versions');
72+
const getEncryptedKeyUpdated = await getEncryptedKey({
73+
pkpSessionSigs,
74+
litNodeClient: devEnv.litNodeClient,
75+
id,
76+
includeVersions: true,
77+
});
78+
console.log({ getEncryptedKeyUpdated });
79+
80+
if (getEncryptedKeyUpdated.ciphertext !== newCiphertext2) {
81+
throw new Error('Ciphertext was not updated');
82+
}
83+
if (
84+
!getEncryptedKeyUpdated.versions ||
85+
getEncryptedKeyUpdated.versions.length !== 2
86+
) {
87+
throw new Error('Versions array missing or incorrect length');
88+
}
89+
if (
90+
getEncryptedKeyUpdated.versions[0].ciphertext !==
91+
getEncryptedKeyFirst.ciphertext
92+
) {
93+
throw new Error('Initial version ciphertext mismatch');
94+
}
95+
if (getEncryptedKeyUpdated.versions[1].ciphertext !== newCiphertext1) {
96+
throw new Error('First update version ciphertext mismatch');
97+
}
98+
if (
99+
!getEncryptedKeyUpdated.updatedAt ||
100+
!getEncryptedKeyUpdated.versions[0].updatedAt ||
101+
!getEncryptedKeyUpdated.versions[1].updatedAt
102+
) {
103+
throw new Error('updatedAt timestamps not set');
104+
}
105+
106+
log('✅ testUpdateWrappedKey');
107+
} finally {
108+
devEnv.releasePrivateKeyFromUser(alice);
109+
}
110+
};

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import {
1010
listEncryptedKeyMetadata,
1111
batchGeneratePrivateKeys,
1212
storeEncryptedKeyBatch,
13+
updateEncryptedKey,
1314
} from './lib/api';
1415
import {
1516
CHAIN_ETHEREUM,
@@ -54,6 +55,9 @@ import type {
5455
StoreEncryptedKeyResult,
5556
ImportPrivateKeyResult,
5657
StoreEncryptedKeyBatchResult,
58+
UpdateEncryptedKeyParams,
59+
UpdateEncryptedKeyResult,
60+
WrappedKeyVersion,
5761
} from './lib/types';
5862

5963
export const constants = {
@@ -77,6 +81,7 @@ export const api = {
7781
storeEncryptedKey,
7882
storeEncryptedKeyBatch,
7983
batchGeneratePrivateKeys,
84+
updateEncryptedKey,
8085
};
8186

8287
export const config = {
@@ -112,4 +117,7 @@ export {
112117
StoredKeyData,
113118
StoredKeyMetadata,
114119
SupportedNetworks,
120+
UpdateEncryptedKeyParams,
121+
UpdateEncryptedKeyResult,
122+
WrappedKeyVersion,
115123
};

‎packages/wrapped-keys/src/lib/api/get-encrypted-key.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import { GetEncryptedKeyDataParams, StoredKeyData } from '../types';
1111
export async function getEncryptedKey(
1212
params: GetEncryptedKeyDataParams
1313
): Promise<StoredKeyData> {
14-
const { pkpSessionSigs, litNodeClient, id } = params;
14+
const { pkpSessionSigs, litNodeClient, id, includeVersions } = params;
1515

1616
const sessionSig = getFirstSessionSig(pkpSessionSigs);
1717
const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
@@ -21,5 +21,6 @@ export async function getEncryptedKey(
2121
id,
2222
sessionSig,
2323
litNetwork: litNodeClient.config.litNetwork,
24+
includeVersions,
2425
});
2526
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import { signTransactionWithEncryptedKey } from './sign-transaction-with-encrypt
99
import { signTypedDataWithEncryptedKey } from './sign-typed-data-with-encrypted-key';
1010
import { storeEncryptedKey } from './store-encrypted-key';
1111
import { storeEncryptedKeyBatch } from './store-encrypted-key-batch';
12+
import { updateEncryptedKey } from './update-encrypted-key';
1213

1314
export {
1415
listEncryptedKeyMetadata,
@@ -22,4 +23,5 @@ export {
2223
storeEncryptedKeyBatch,
2324
getEncryptedKey,
2425
batchGeneratePrivateKeys,
26+
updateEncryptedKey,
2527
};

‎packages/wrapped-keys/src/lib/api/update-encrypted-key.ts‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
import { updatePrivateKey } from '../service-client';
2+
import { UpdateEncryptedKeyParams, UpdateEncryptedKeyResult } from '../types';
3+
import { getFirstSessionSig, getPkpAddressFromSessionSig } from './utils';
4+
5+
/** Update an existing wrapped key and append the previous state to versions.
6+
*
7+
* @param { UpdateEncryptedKeyParams } params Parameters required to update the encrypted private key
8+
* @returns { Promise<UpdateEncryptedKeyResult> } id/pkpAddress/updatedAt on successful update
9+
*/
10+
export async function updateEncryptedKey(
11+
params: UpdateEncryptedKeyParams
12+
): Promise<UpdateEncryptedKeyResult> {
13+
const {
14+
pkpSessionSigs,
15+
litNodeClient,
16+
id,
17+
ciphertext,
18+
evmContractConditions,
19+
memo,
20+
} = params;
21+
22+
const sessionSig = getFirstSessionSig(pkpSessionSigs);
23+
const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
24+
25+
return updatePrivateKey({
26+
pkpAddress,
27+
id,
28+
sessionSig,
29+
ciphertext,
30+
evmContractConditions,
31+
memo,
32+
litNetwork: litNodeClient.config.litNetwork,
33+
});
34+
}

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

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,15 @@ import {
33
ListKeysParams,
44
StoreKeyBatchParams,
55
StoreKeyParams,
6+
UpdateKeyParams,
67
} from './types';
78
import { generateRequestId, getBaseRequestParams, makeRequest } from './utils';
89
import {
910
StoredKeyData,
1011
StoredKeyMetadata,
1112
StoreEncryptedKeyBatchResult,
1213
StoreEncryptedKeyResult,
14+
UpdateEncryptedKeyResult,
1315
} from '../types';
1416

1517
/** Fetches previously stored private key metadata from the wrapped keys service.
@@ -48,7 +50,7 @@ export async function listPrivateKeyMetadata(
4850
export async function fetchPrivateKey(
4951
params: FetchKeyParams
5052
): Promise<StoredKeyData> {
51-
const { litNetwork, sessionSig, id, pkpAddress } = params;
53+
const { litNetwork, sessionSig, id, pkpAddress, includeVersions } = params;
5254

5355
const requestId = generateRequestId();
5456
const { baseUrl, initParams } = getBaseRequestParams({
@@ -58,8 +60,9 @@ export async function fetchPrivateKey(
5860
requestId,
5961
});
6062

63+
const query = includeVersions ? '?includeVersions=true' : '';
6164
return makeRequest<StoredKeyData>({
62-
url: `${baseUrl}/${pkpAddress}/${id}`,
65+
url: `${baseUrl}/${pkpAddress}/${id}${query}`,
6366
init: initParams,
6467
requestId,
6568
});
@@ -124,3 +127,45 @@ export async function storePrivateKeyBatch(
124127

125128
return { pkpAddress, ids };
126129
}
130+
131+
/** Updates an existing wrapped key and appends prior state to versions.
132+
*
133+
* @param { UpdateKeyParams } params Parameters required to update the private key metadata
134+
* @returns { Promise<UpdateEncryptedKeyResult> } id/pkpAddress/updatedAt on successful update
135+
*/
136+
export async function updatePrivateKey(
137+
params: UpdateKeyParams
138+
): Promise<UpdateEncryptedKeyResult> {
139+
const {
140+
litNetwork,
141+
sessionSig,
142+
pkpAddress,
143+
id,
144+
ciphertext,
145+
evmContractConditions,
146+
memo,
147+
} = params;
148+
149+
const requestId = generateRequestId();
150+
const { baseUrl, initParams } = getBaseRequestParams({
151+
litNetwork,
152+
sessionSig,
153+
method: 'PUT',
154+
requestId,
155+
});
156+
157+
return makeRequest<UpdateEncryptedKeyResult>({
158+
url: `${baseUrl}/${pkpAddress}/${id}`,
159+
init: {
160+
...initParams,
161+
body: JSON.stringify({
162+
ciphertext,
163+
...(evmContractConditions !== undefined
164+
? { evmContractConditions }
165+
: {}),
166+
...(memo !== undefined ? { memo } : {}),
167+
}),
168+
},
169+
requestId,
170+
});
171+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@ import {
33
storePrivateKey,
44
storePrivateKeyBatch,
55
listPrivateKeyMetadata,
6+
updatePrivateKey,
67
} from './client';
78

89
export {
910
fetchPrivateKey,
1011
storePrivateKey,
1112
storePrivateKeyBatch,
1213
listPrivateKeyMetadata,
14+
updatePrivateKey,
1315
};

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ interface BaseApiParams {
1111
export type FetchKeyParams = BaseApiParams & {
1212
pkpAddress: string;
1313
id: string;
14+
includeVersions?: boolean;
1415
};
1516

1617
export type ListKeysParams = BaseApiParams & { pkpAddress: string };
@@ -34,9 +35,17 @@ export interface StoreKeyBatchParams extends BaseApiParams {
3435
>[];
3536
}
3637

38+
export interface UpdateKeyParams extends BaseApiParams {
39+
pkpAddress: string;
40+
id: string;
41+
ciphertext: string;
42+
evmContractConditions?: string;
43+
memo?: string;
44+
}
45+
3746
export interface BaseRequestParams {
3847
sessionSig: AuthSig;
39-
method: 'GET' | 'POST';
48+
method: 'GET' | 'POST' | 'PUT';
4049
litNetwork: LIT_NETWORKS_KEYS;
4150
requestId: string;
4251
}

0 commit comments

Comments
 (0)