Skip to content

Commit 122dd16

Browse files
mikelodder7mikelodder7
committed
e2ee
Signed-off-by: Michael Lodder <[email protected]>
1 parent 2afa74c commit 122dd16

File tree

5 files changed

+138
-49
lines changed

5 files changed

+138
-49
lines changed

‎packages/core/src/lib/lit-core.ts‎

Lines changed: 36 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -50,12 +50,14 @@ import {
5050
RejectedNodePromises,
5151
SessionSigsMap,
5252
SuccessNodePromises,
53+
WalletEncryptedPayload,
5354
} from '@lit-protocol/types';
5455

5556
import { composeLitUrl } from './endpoint-version';
5657
import {
5758
CoreNodeConfig,
5859
EpochCache,
60+
GenericResponse,
5961
HandshakeWithNode,
6062
Listener,
6163
NodeCommandServerKeysResponse,
@@ -545,13 +547,30 @@ export class LitCore {
545547
requestId
546548
);
547549

550+
if (!handshakeResult.ok) {
551+
const error = handshakeResult.error ?? "";
552+
const errorObject = handshakeResult.errorObject ?? "";
553+
554+
logErrorWithRequestId(
555+
requestId,
556+
`Error handshaking with the node ${url}: ${error} - ${errorObject}`
557+
);
558+
throw new Error(`${error} - ${errorObject}`);
559+
}
560+
561+
if (!handshakeResult.data) {
562+
throw new Error('Handshake response data is empty');
563+
}
564+
const handshakeResponse = handshakeResult.data;
565+
548566
const keys: JsonHandshakeResponse = {
549-
serverPubKey: handshakeResult.serverPublicKey,
550-
subnetPubKey: handshakeResult.subnetPublicKey,
551-
networkPubKey: handshakeResult.networkPublicKey,
552-
networkPubKeySet: handshakeResult.networkPublicKeySet,
553-
hdRootPubkeys: handshakeResult.hdRootPubkeys,
554-
latestBlockhash: handshakeResult.latestBlockhash,
567+
serverPubKey: handshakeResponse.serverPublicKey,
568+
subnetPubKey: handshakeResponse.subnetPublicKey,
569+
networkPubKey: handshakeResponse.networkPublicKey,
570+
networkPubKeySet: handshakeResponse.networkPublicKeySet,
571+
hdRootPubkeys: handshakeResponse.hdRootPubkeys,
572+
latestBlockhash: handshakeResponse.latestBlockhash,
573+
nodeIdentityKey: handshakeResponse.nodeIdentityKey,
555574
};
556575

557576
// Nodes that have just bootstrapped will not have negotiated their keys, yet
@@ -587,7 +606,7 @@ export class LitCore {
587606
this.config.checkNodeAttestation ||
588607
NETWORKS_REQUIRING_SEV.includes(this.config.litNetwork)
589608
) {
590-
const attestation = handshakeResult.attestation;
609+
const attestation = handshakeResponse.attestation;
591610

592611
if (!attestation) {
593612
throw new InvalidNodeAttestation(
@@ -887,7 +906,7 @@ export class LitCore {
887906
protected _handshakeWithNode = async (
888907
params: HandshakeWithNode,
889908
requestId: string
890-
): Promise<NodeCommandServerKeysResponse> => {
909+
): Promise<GenericResponse<NodeCommandServerKeysResponse>> => {
891910
// -- get properties from params
892911
const { url } = params;
893912

@@ -1029,9 +1048,9 @@ export class LitCore {
10291048
protected _getNodePromises = (
10301049
nodeUrls: string[],
10311050
// eslint-disable-next-line @typescript-eslint/no-explicit-any
1032-
callback: (url: string) => Promise<any>
1051+
callback: (url: string) => { url: string, promise: Promise<any>},
10331052
// eslint-disable-next-line @typescript-eslint/no-explicit-any
1034-
): Promise<any>[] => {
1053+
): { url: string, promise: Promise<any> }[] => {
10351054
// FIXME: Replace <any> usage with explicit, strongly typed handlers
10361055

10371056
const nodePromises = [];
@@ -1090,25 +1109,25 @@ export class LitCore {
10901109
* @returns { Promise<SuccessNodePromises<T> | RejectedNodePromises> }
10911110
*/
10921111
protected _handleNodePromises = async <T>(
1093-
nodePromises: Promise<T>[],
1112+
nodePromises: { url: string, promise: Promise<T> }[],
10941113
requestId: string,
10951114
minNodeCount: number
1096-
): Promise<SuccessNodePromises<T> | RejectedNodePromises> => {
1115+
): Promise<SuccessNodePromises<WalletEncryptedPayload> | RejectedNodePromises> => {
10971116
async function waitForNSuccessesWithErrors<T>(
1098-
promises: Promise<T>[],
1117+
promises: { url: string, promise: Promise<T> }[],
10991118
n: number
11001119
// eslint-disable-next-line @typescript-eslint/no-explicit-any
1101-
): Promise<{ successes: T[]; errors: any[] }> {
1120+
): Promise<{ successes: { url: string, result: T }[]; errors: any[] }> {
11021121
let responses = 0;
1103-
const successes: T[] = [];
1122+
const successes: { url: string, result: T }[] = [];
11041123
// eslint-disable-next-line @typescript-eslint/no-explicit-any
11051124
const errors: any[] = [];
11061125

11071126
return new Promise((resolve) => {
1108-
promises.forEach((promise) => {
1127+
promises.forEach(({ url, promise }) => {
11091128
promise
11101129
.then((result) => {
1111-
successes.push(result);
1130+
successes.push({ url, result });
11121131
if (successes.length >= n) {
11131132
// If we've got enough successful responses to continue, resolve immediately even if some are pending
11141133
resolve({ successes, errors });

‎packages/core/src/lib/types.ts‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,13 @@ export interface SendNodeCommand {
88
requestId: string;
99
}
1010

11+
export interface GenericResponse<T> {
12+
ok: boolean;
13+
error?: string;
14+
errorObject?: string;
15+
data?: T;
16+
}
17+
1118
export interface NodeCommandServerKeysResponse {
1219
serverPublicKey: string;
1320
subnetPublicKey: string;
@@ -16,6 +23,7 @@ export interface NodeCommandServerKeysResponse {
1623
hdRootPubkeys: string[];
1724
attestation?: NodeAttestation;
1825
latestBlockhash?: string;
26+
nodeIdentityKey: string;
1927
}
2028

2129
export interface HandshakeWithNode {

‎packages/crypto/src/lib/crypto.ts‎

Lines changed: 4 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -353,14 +353,11 @@ async function getAmdCert(url: string): Promise<Uint8Array> {
353353
}
354354
}
355355

356-
export const walletEncrypt = async (
356+
export const walletEncrypt = (
357357
myWalletSecretKey: Uint8Array,
358358
theirWalletPublicKey: Uint8Array,
359-
sessionSig: AuthSig,
360359
message: Uint8Array
361-
): Promise<WalletEncryptedPayload> => {
362-
const uint8SessionSig = Buffer.from(JSON.stringify(sessionSig));
363-
360+
): WalletEncryptedPayload => {
364361
const random = new Uint8Array(16);
365362
crypto.getRandomValues(random);
366363
const dateNow = Date.now();
@@ -372,12 +369,10 @@ export const walletEncrypt = async (
372369
nacl.lowlevel.crypto_scalarmult_base(myWalletPublicKey, myWalletSecretKey);
373370

374371
// Construct AAD (Additional Authenticated Data) - data that is authenticated but not encrypted
375-
const sessionSignature = uint8SessionSig; // Replace with actual session signature
376372
const theirPublicKey = Buffer.from(theirWalletPublicKey); // Replace with their public key
377373
const myPublicKey = Buffer.from(myWalletPublicKey); // Replace with your wallet public key
378374

379375
const aad = Buffer.concat([
380-
sessionSignature,
381376
random,
382377
timestamp,
383378
theirPublicKey,
@@ -398,17 +393,16 @@ export const walletEncrypt = async (
398393
V1: {
399394
verification_key: uint8ArrayToHex(myWalletPublicKey),
400395
ciphertext_and_tag: uint8ArrayToHex(ciphertext),
401-
session_signature: uint8ArrayToHex(sessionSignature),
402396
random: uint8ArrayToHex(random),
403397
created_at: new Date(dateNow).toISOString(),
404398
},
405399
};
406400
};
407401

408-
export const walletDecrypt = async (
402+
export const walletDecrypt = (
409403
myWalletSecretKey: Uint8Array,
410404
payload: WalletEncryptedPayload
411-
): Promise<Uint8Array> => {
405+
): Uint8Array => {
412406
const dateSent = new Date(payload.V1.created_at);
413407
const createdAt = Math.floor(dateSent.getTime() / 1000);
414408
const timestamp = Buffer.alloc(8);
@@ -419,15 +413,11 @@ export const walletDecrypt = async (
419413

420414
// Construct AAD
421415
const random = Buffer.from(hexToUint8Array(payload.V1.random));
422-
const sessionSignature = Buffer.from(
423-
hexToUint8Array(payload.V1.session_signature)
424-
); // Replace with actual session signature
425416
const theirPublicKey = hexToUint8Array(payload.V1.verification_key);
426417
const theirPublicKeyBuffer = Buffer.from(theirPublicKey); // Replace with their public key
427418
const myPublicKey = Buffer.from(myWalletPublicKey); // Replace with your wallet public key
428419

429420
const aad = Buffer.concat([
430-
sessionSignature,
431421
random,
432422
timestamp,
433423
theirPublicKeyBuffer,

‎packages/lit-node-client-nodejs/src/lib/lit-node-client-nodejs.ts‎

Lines changed: 78 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,8 @@ import {
5151
hashLitMessage,
5252
verifyAndDecryptWithSignatureShares,
5353
verifySignature,
54+
walletDecrypt,
55+
walletEncrypt,
5456
} from '@lit-protocol/crypto';
5557
import {
5658
defaultMintClaimCallback,
@@ -127,10 +129,12 @@ import type {
127129
JsonPkpSignSdkParams,
128130
JsonSignSessionKeyRequestV1,
129131
JsonSignSessionKeyRequestV2,
132+
JsonSignSessionKeyResponse,
130133
LitNodeClientConfig,
131134
NodeBlsSigningShare,
132135
NodeCommandResponse,
133136
NodeSet,
137+
PKPSignEndpointResponse,
134138
RejectedNodePromises,
135139
SessionKeyPair,
136140
SessionSigningTemplate,
@@ -139,7 +143,10 @@ import type {
139143
SignSessionKeyProp,
140144
SignSessionKeyResponse,
141145
SuccessNodePromises,
146+
WalletEncryptedPayload,
142147
} from '@lit-protocol/types';
148+
import { json } from 'node:stream/consumers';
149+
import { GenericResponse } from 'packages/core/src/lib/types';
143150

144151
export class LitNodeClientNodeJs extends LitCore implements ILitNodeClient {
145152
/** Tracks the total max price a user is willing to pay for each supported product type
@@ -604,7 +611,7 @@ export class LitNodeClientNodeJs extends LitCore implements ILitNodeClient {
604611
url: string,
605612
formattedParams: JsonExecutionSdkParams & { sessionSigs: SessionSigsMap },
606613
requestId: string,
607-
nodeSet: NodeSet[]
614+
nodeSet: { node: NodeSet, nodeIdentityKey: string }[]
608615
) {
609616
// -- choose the right signature
610617
const sessionSig = this._getSessionSigByUrl({
@@ -615,7 +622,7 @@ export class LitNodeClientNodeJs extends LitCore implements ILitNodeClient {
615622
const reqBody: JsonExecutionRequest = {
616623
...formattedParams,
617624
authSig: sessionSig,
618-
nodeSet,
625+
nodeSet: nodeSet,
619626
};
620627

621628
const urlWithPath = composeLitUrl({
@@ -889,6 +896,17 @@ export class LitNodeClientNodeJs extends LitCore implements ILitNodeClient {
889896
// Handle promises for commands sent to Lit nodes
890897

891898
const targetNodeUrls = targetNodePrices.map(({ url }) => url);
899+
let keySets: Record<string, { theirPublicKey: Uint8Array, secretKey: Uint8Array }> = {};
900+
targetNodeUrls.forEach((url) => {
901+
const theirPublicKey = uint8arrayFromString(this.serverKeys[url].nodeIdentityKey);
902+
const keyPair = nacl.box.keyPair.generate();
903+
const secretKey = keyPair.secretKey;
904+
905+
keySets[url] = {
906+
theirPublicKey,
907+
secretKey,
908+
};
909+
});
892910
const nodePromises = this._getNodePromises(
893911
targetNodeUrls,
894912
(url: string) => {
@@ -906,6 +924,8 @@ export class LitNodeClientNodeJs extends LitCore implements ILitNodeClient {
906924
params.messageToSign!
907925
);
908926

927+
const { theirPublicKey, secretKey } = keySets[url];
928+
909929
const reqBody: JsonPkpSignRequest = {
910930
toSign: normalizeArray(toSign),
911931
signingScheme: params.signingScheme,
@@ -929,7 +949,13 @@ export class LitNodeClientNodeJs extends LitCore implements ILitNodeClient {
929949
endpoint: LIT_ENDPOINT.PKP_SIGN,
930950
});
931951

932-
return this.generatePromise(urlWithPath, reqBody, requestId);
952+
const encrypted = walletEncrypt(
953+
theirPublicKey,
954+
secretKey,
955+
uint8arrayFromString(JSON.stringify(reqBody))
956+
);
957+
958+
return { url, promise: this.generatePromise(urlWithPath, encrypted, requestId) };
933959
}
934960
);
935961

@@ -944,7 +970,29 @@ export class LitNodeClientNodeJs extends LitCore implements ILitNodeClient {
944970
return this._throwNodeError(res, requestId);
945971
}
946972

947-
const responseData = res.values;
973+
const responseData = res.values.map((values) => {
974+
const secretKey = keySets[values.url].secretKey;
975+
976+
const decrypted = walletDecrypt(
977+
secretKey,
978+
values.result,
979+
);
980+
981+
const response: GenericResponse<PKPSignEndpointResponse> = JSON.parse(uint8arrayToString(decrypted, 'utf8'));
982+
983+
if (!response.ok) {
984+
const error = response.error ?? "";
985+
const errorObject = response.errorObject ?? {};
986+
987+
throw new Error(`${error} - ${errorObject}`);
988+
}
989+
990+
if (!response.data) {
991+
throw new Error('No data returned from pkp sign');
992+
}
993+
994+
return response.data;
995+
});
948996

949997
logWithRequestId(
950998
requestId,
@@ -1521,6 +1569,7 @@ export class LitNodeClientNodeJs extends LitCore implements ILitNodeClient {
15211569
};
15221570

15231571
getSignSessionKeyShares = async (
1572+
nodeIdentityKey: string,
15241573
url: string,
15251574
params: {
15261575
body: {
@@ -1538,11 +1587,34 @@ export class LitNodeClientNodeJs extends LitCore implements ILitNodeClient {
15381587
url,
15391588
endpoint: LIT_ENDPOINT.SIGN_SESSION_KEY,
15401589
});
1541-
return await this._sendCommandToNode({
1590+
1591+
const theirPublicKey = uint8arrayFromString(nodeIdentityKey, 'base16');
1592+
const keyPair = nacl.box.keyPair();
1593+
const secretKey = keyPair.secretKey;
1594+
1595+
const encrypted = await walletEncrypt(
1596+
secretKey,
1597+
theirPublicKey,
1598+
uint8arrayFromString(JSON.stringify(params.body), 'utf8')
1599+
);
1600+
1601+
const response: WalletEncryptedPayload = await this._sendCommandToNode({
15421602
url: urlWithPath,
1543-
data: params.body,
1603+
data: encrypted,
15441604
requestId,
15451605
});
1606+
1607+
const decrypted = await walletDecrypt(
1608+
secretKey,
1609+
response,
1610+
);
1611+
1612+
const outerResponse: GenericResponse<JsonSignSessionKeyResponse> = JSON.parse(uint8arrayToString(decrypted, 'utf8'));
1613+
if (!outerResponse.ok) {
1614+
throw new Error(`${outerResponse.error} - ${outerResponse.errorObject}`);
1615+
}
1616+
1617+
return outerResponse.data;
15461618
};
15471619

15481620
getMaxPricesForNodeProduct = async ({

0 commit comments

Comments
 (0)