feat(wrapped-keys): add versioned key update support and client APIs

Author: Ansonhkg

packages/e2e/src/test-helpers/executeJs/wrappedKeys.ts

@@ -341,6 +341,56 @@ export const registerWrappedKeysTests = () => {
         expect(storedKey.dataToEncryptHash).toBeTruthy();
       });
 
+      test('updateEncryptedKey rotates ciphertext and returns version history', async () => {
+        const pkpSessionSigs = await TestHelper.createPkpSessionSigs({
+          testEnv,
+          alice,
+          delegationAuthSig: aliceDelegationAuthSig,
+        });
+
+        const initialPayload = TestHelper.createStorePayload(
+          TestHelper.randomMemo('update-before')
+        );
+
+        const { id } = await wrappedKeysApi.storeEncryptedKey({
+          pkpSessionSigs,
+          litClient: testEnv.litClient,
+          ...initialPayload,
+        });
+
+        const newCiphertext = TestHelper.randomCiphertext();
+        const newMemo = TestHelper.randomMemo('update-after');
+
+        const updateResult = await wrappedKeysApi.updateEncryptedKey({
+          pkpSessionSigs,
+          litClient: testEnv.litClient,
+          id,
+          ciphertext: newCiphertext,
+          memo: newMemo,
+        });
+
+        expect(updateResult.id).toBe(id);
+        expect(updateResult.pkpAddress).toBe(alice.pkp!.ethAddress);
+        expect(updateResult.updatedAt).toBeTruthy();
+
+        const fetched = await wrappedKeysApi.getEncryptedKey({
+          pkpSessionSigs,
+          litClient: testEnv.litClient,
+          id,
+          includeVersions: true,
+        });
+
+        expect(fetched.ciphertext).toBe(newCiphertext);
+        expect(fetched.memo).toBe(newMemo);
+        expect(fetched.updatedAt).toBeTruthy();
+        expect(fetched.versions).toBeDefined();
+        expect(fetched.versions?.length).toBe(1);
+        expect(fetched.versions?.[0].ciphertext).toBe(
+          initialPayload.ciphertext
+        );
+        expect(fetched.versions?.[0].memo).toBe(initialPayload.memo);
+      });
+
       test('importPrivateKey persists an externally generated key', async () => {
         const pkpSessionSigs = await TestHelper.createPkpSessionSigs({
           testEnv,

packages/wrapped-keys/src/index.ts

@@ -9,6 +9,7 @@ import {
   signTransactionWithEncryptedKey,
   storeEncryptedKey,
   storeEncryptedKeyBatch,
+  updateEncryptedKey,
 } from './lib/api';
 import {
   CHAIN_ETHEREUM,
@@ -43,6 +44,7 @@ export const api = {
   storeEncryptedKey,
   storeEncryptedKeyBatch,
   batchGeneratePrivateKeys,
+  updateEncryptedKey,
 };
 
 export const config = {
@@ -76,6 +78,9 @@ export type {
   StoreEncryptedKeyResult,
   StoredKeyData,
   StoredKeyMetadata,
+  UpdateEncryptedKeyParams,
+  UpdateEncryptedKeyResult,
+  WrappedKeyVersion,
 } from './lib/types';
 
 export type {

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

@@ -15,7 +15,7 @@ import { GetEncryptedKeyDataParams, StoredKeyData } from '../types';
 export async function getEncryptedKey(
   params: GetEncryptedKeyDataParams
 ): Promise<StoredKeyData> {
-  const { pkpSessionSigs, litClient, id } = params;
+  const { pkpSessionSigs, litClient, id, includeVersions } = params;
 
   const sessionSig = getFirstSessionSig(pkpSessionSigs);
   const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
@@ -26,5 +26,6 @@ export async function getEncryptedKey(
     id,
     sessionSig,
     litNetwork,
+    includeVersions,
   });
 }

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

@@ -8,6 +8,7 @@ import { signMessageWithEncryptedKey } from './sign-message-with-encrypted-key';
 import { signTransactionWithEncryptedKey } from './sign-transaction-with-encrypted-key';
 import { storeEncryptedKey } from './store-encrypted-key';
 import { storeEncryptedKeyBatch } from './store-encrypted-key-batch';
+import { updateEncryptedKey } from './update-encrypted-key';
 
 export {
   listEncryptedKeyMetadata,
@@ -20,4 +21,5 @@ export {
   storeEncryptedKeyBatch,
   getEncryptedKey,
   batchGeneratePrivateKeys,
+  updateEncryptedKey,
 };

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

@@ -0,0 +1,35 @@
+import { updatePrivateKey } from '../service-client';
+import { UpdateEncryptedKeyParams, UpdateEncryptedKeyResult } from '../types';
+import {
+  getFirstSessionSig,
+  getLitNetworkFromClient,
+  getPkpAddressFromSessionSig,
+} from './utils';
+
+/** Update an existing wrapped key and append the previous state to versions. */
+export async function updateEncryptedKey(
+  params: UpdateEncryptedKeyParams
+): Promise<UpdateEncryptedKeyResult> {
+  const {
+    pkpSessionSigs,
+    litClient,
+    id,
+    ciphertext,
+    evmContractConditions,
+    memo,
+  } = params;
+
+  const sessionSig = getFirstSessionSig(pkpSessionSigs);
+  const pkpAddress = getPkpAddressFromSessionSig(sessionSig);
+  const litNetwork = getLitNetworkFromClient(litClient);
+
+  return updatePrivateKey({
+    pkpAddress,
+    id,
+    sessionSig,
+    ciphertext,
+    evmContractConditions,
+    memo,
+    litNetwork,
+  });
+}

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

@@ -3,13 +3,15 @@ import {
   ListKeysParams,
   StoreKeyBatchParams,
   StoreKeyParams,
+  UpdateKeyParams,
 } from './types';
 import { generateRequestId, getBaseRequestParams, makeRequest } from './utils';
 import {
   StoredKeyData,
   StoredKeyMetadata,
   StoreEncryptedKeyBatchResult,
   StoreEncryptedKeyResult,
+  UpdateEncryptedKeyResult,
 } from '../types';
 
 /** Fetches previously stored private key metadata from the wrapped keys service.
@@ -48,7 +50,7 @@ export async function listPrivateKeyMetadata(
 export async function fetchPrivateKey(
   params: FetchKeyParams
 ): Promise<StoredKeyData> {
-  const { litNetwork, sessionSig, id, pkpAddress } = params;
+  const { litNetwork, sessionSig, id, pkpAddress, includeVersions } = params;
 
   const requestId = generateRequestId();
   const { baseUrl, initParams } = getBaseRequestParams({
@@ -58,8 +60,9 @@ export async function fetchPrivateKey(
     requestId,
   });
 
+  const query = includeVersions ? '?includeVersions=true' : '';
   return makeRequest<StoredKeyData>({
-    url: `${baseUrl}/${pkpAddress}/${id}`,
+    url: `${baseUrl}/${pkpAddress}/${id}${query}`,
     init: initParams,
     requestId,
   });
@@ -124,3 +127,45 @@ export async function storePrivateKeyBatch(
 
   return { pkpAddress, ids };
 }
+
+/** Updates an existing wrapped key and appends prior state to versions.
+ *
+ * @param { UpdateKeyParams } params Parameters required to update the private key metadata
+ * @returns { Promise<UpdateEncryptedKeyResult> } id/pkpAddress/updatedAt on successful update
+ */
+export async function updatePrivateKey(
+  params: UpdateKeyParams
+): Promise<UpdateEncryptedKeyResult> {
+  const {
+    litNetwork,
+    sessionSig,
+    pkpAddress,
+    id,
+    ciphertext,
+    evmContractConditions,
+    memo,
+  } = params;
+
+  const requestId = generateRequestId();
+  const { baseUrl, initParams } = getBaseRequestParams({
+    litNetwork,
+    sessionSig,
+    method: 'PUT',
+    requestId,
+  });
+
+  return makeRequest<UpdateEncryptedKeyResult>({
+    url: `${baseUrl}/${pkpAddress}/${id}`,
+    init: {
+      ...initParams,
+      body: JSON.stringify({
+        ciphertext,
+        ...(evmContractConditions !== undefined
+          ? { evmContractConditions }
+          : {}),
+        ...(memo !== undefined ? { memo } : {}),
+      }),
+    },
+    requestId,
+  });
+}

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

@@ -3,11 +3,13 @@ import {
   storePrivateKey,
   storePrivateKeyBatch,
   listPrivateKeyMetadata,
+  updatePrivateKey,
 } from './client';
 
 export {
   fetchPrivateKey,
   storePrivateKey,
   storePrivateKeyBatch,
   listPrivateKeyMetadata,
+  updatePrivateKey,
 };

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

@@ -11,6 +11,7 @@ interface BaseApiParams {
 export type FetchKeyParams = BaseApiParams & {
   pkpAddress: string;
   id: string;
+  includeVersions?: boolean;
 };
 
 export type ListKeysParams = BaseApiParams & { pkpAddress: string };
@@ -34,9 +35,17 @@ export interface StoreKeyBatchParams extends BaseApiParams {
   >[];
 }
 
+export interface UpdateKeyParams extends BaseApiParams {
+  pkpAddress: string;
+  id: string;
+  ciphertext: string;
+  evmContractConditions?: string;
+  memo?: string;
+}
+
 export interface BaseRequestParams {
   sessionSig: AuthSig;
-  method: 'GET' | 'POST';
+  method: 'GET' | 'POST' | 'PUT';
   litNetwork: LIT_NETWORKS_KEYS;
   requestId: string;
 }

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

@@ -47,6 +47,7 @@ export type ListEncryptedKeyMetadataParams = BaseApiParams;
  */
 export type GetEncryptedKeyDataParams = BaseApiParams & {
   id: string;
+  includeVersions?: boolean;
 };
 
 /** Metadata for a key that has been stored, encrypted, on the wrapped keys backend service
@@ -68,6 +69,8 @@ export interface StoredKeyMetadata {
   litNetwork: LIT_NETWORKS_KEYS;
   memo: string;
   id: string;
+  updatedAt?: string;
+  versions?: WrappedKeyVersion[];
 }
 
 /** Complete encrypted private key data, including the `ciphertext` and `dataToEncryptHash` necessary to decrypt the key
@@ -81,6 +84,22 @@ export interface StoredKeyData extends StoredKeyMetadata {
   dataToEncryptHash: string;
 }
 
+/** Represents a historical version of a wrapped key after an update operation */
+export interface WrappedKeyVersion
+  extends Pick<
+    StoredKeyData,
+    | 'ciphertext'
+    | 'dataToEncryptHash'
+    | 'keyType'
+    | 'litNetwork'
+    | 'memo'
+    | 'publicKey'
+  > {
+  id: string;
+  evmContractConditions?: string;
+  updatedAt: string;
+}
+
 /** Properties required to persist an encrypted key into the wrapped-keys backend storage service
  *
  * @typedef StoreEncryptedKeyParams
@@ -130,6 +149,21 @@ export interface StoreEncryptedKeyBatchResult {
   pkpAddress: string;
 }
 
+/** Properties required to update an existing encrypted key in the wrapped-keys backend storage service */
+export type UpdateEncryptedKeyParams = BaseApiParams & {
+  id: string;
+  ciphertext: string;
+  evmContractConditions?: string;
+  memo?: string;
+};
+
+/** Result of updating a private key in the wrapped keys backend service */
+export interface UpdateEncryptedKeyResult {
+  id: string;
+  pkpAddress: string;
+  updatedAt: string;
+}
+
 /** Exporting a previously persisted key only requires valid pkpSessionSigs and a LIT Node Client instance configured for the appropriate network.
  *
  * @typedef ExportPrivateKeyParams