[seal] batched key server requests

.changeset/slimy-plants-add.md

@@ -0,0 +1,5 @@
+---
+'@mysten/seal': minor
+---
+
+introduced deduplication of key server object fetching

packages/seal/src/client.ts

@@ -162,17 +162,20 @@ export class SealClient {
 		for (const objectId of this.#serverObjectIds) {
 			serverObjectIdsMap.set(objectId, (serverObjectIdsMap.get(objectId) ?? 0) + 1);
 		}
+
 		const servicesMap = new Map<string, number>();
 		for (const service of services) {
 			servicesMap.set(service, (servicesMap.get(service) ?? 0) + 1);
 		}
+
 		for (const [objectId, count] of serverObjectIdsMap) {
 			if (servicesMap.get(objectId) !== count) {
 				throw new InconsistentKeyServersError(
 					`Client's key servers must be a subset of the encrypted object's key servers`,
 				);
 			}
 		}
+
 		// Check that the threshold can be met with the client's key servers.
 		if (threshold > this.#serverObjectIds.length) {
 			throw new InvalidThresholdError(
@@ -218,7 +221,7 @@ export class SealClient {
 	/**
 	 * Fetch keys from the key servers and update the cache.
 	 *
-	 * It is recommended to call this function once for all ids of all encrypted obejcts if
+	 * It is recommended to call this function once for all ids of all encrypted objects if
 	 * there are multiple, then call decrypt for each object. This avoids calling fetchKey
 	 * individually for each decrypt.
 	 *
@@ -246,7 +249,7 @@ export class SealClient {
 		}
 
 		let completedServerCount = 0;
-		const remainingKeyServers = new Set<KeyServer>();
+		const remainingKeyServers = new Map<string, KeyServer>();
 		const fullIds = ids.map((id) => createFullId(DST, sessionKey.getPackageId(), id));
 
 		// Count a server as completed if it has keys for all fullIds.
@@ -256,10 +259,11 @@ export class SealClient {
 			for (const fullId of fullIds) {
 				if (!this.#cachedKeys.has(`${fullId}:${server.objectId}`)) {
 					hasAllKeys = false;
-					remainingKeyServers.add(server);
+					remainingKeyServers.set(server.objectId, server);
 					break;
 				}
 			}
+
 			if (hasAllKeys) {
 				completedServerCount++;
 			}
@@ -271,7 +275,7 @@ export class SealClient {
 		}
 
 		// Check server validities.
-		for (const server of remainingKeyServers) {
+		for (const server of remainingKeyServers.values()) {
 			if (server.keyType !== KeyServerType.BonehFranklinBLS12381) {
 				throw new InvalidKeyServerError(
 					`Server ${server.objectId} has invalid key type: ${server.keyType}`,
@@ -285,7 +289,7 @@ export class SealClient {
 		const controller = new AbortController();
 		const errors: Error[] = [];
 
-		const keyFetches = [...remainingKeyServers].map(async (server) => {
+		const keyFetches = [...remainingKeyServers.values()].map(async (server) => {
 			try {
 				const allKeys = await fetchKeysForAllIds(
 					server.url,
@@ -296,6 +300,7 @@ export class SealClient {
 					this.#timeout,
 					controller.signal,
 				);
+
 				// Check validity of the keys and add them to the cache.
 				const receivedIds = new Set<string>();
 				for (const { fullId, key } of allKeys) {
@@ -310,6 +315,7 @@ export class SealClient {
 						console.warn('Received invalid key from key server ' + server.objectId);
 						continue;
 					}
+
 					this.#cachedKeys.set(`${fullId}:${server.objectId}`, keyElement);
 					receivedIds.add(fullId);
 				}
@@ -321,9 +327,11 @@ export class SealClient {
 					receivedIds.size === expectedIds.size &&
 					[...receivedIds].every((id) => expectedIds.has(id));
 
-				// Return early if the completed servers is more than threshold.
+				// Count each occurrence of this servers objectId from the original keyServers array.
 				if (hasAllKeys) {
-					completedServerCount++;
+					const occurrences = keyServers.filter((ks) => ks.objectId === server.objectId).length;
+					completedServerCount += occurrences;
+
 					if (completedServerCount >= threshold) {
 						controller.abort();
 					}
@@ -332,6 +340,7 @@ export class SealClient {
 				if (!controller.signal.aborted) {
 					errors.push(error as Error);
 				}
+
 				// If there are too many errors that the threshold is not attainable, return early with error.
 				if (remainingKeyServers.size - errors.length < threshold - completedServerCount) {
 					controller.abort(error);

packages/seal/src/key-server.ts

@@ -57,9 +57,12 @@ export async function retrieveKeyServers({
 	objectIds: string[];
 	client: SealCompatibleClient;
 }): Promise<KeyServer[]> {
-	// todo: do not fetch the same object ID if this is fetched before.
-	return await Promise.all(
-		objectIds.map(async (objectId) => {
+	const uniqueIds = Array.from(new Set(objectIds));
+	const fetchedServers: Record<string, KeyServer> = {};
+
+	// Only fetch key server information for each unique objectId.
+	await Promise.all(
+		uniqueIds.map(async (objectId) => {
 			let res;
 			try {
 				res = await client.core.getObject({
@@ -74,7 +77,7 @@ export async function retrieveKeyServers({
 				throw new UnsupportedFeatureError(`Unsupported key type ${ks.keyType}`);
 			}
 
-			return {
+			fetchedServers[objectId] = {
 				objectId,
 				name: ks.name,
 				url: ks.url,
@@ -83,6 +86,14 @@ export async function retrieveKeyServers({
 			};
 		}),
 	);
+
+	return objectIds.map((objectId) => {
+		if (!fetchedServers[objectId]) {
+			throw new InvalidGetObjectError(`KeyServer ${objectId} not found`);
+		}
+
+		return fetchedServers[objectId];
+	});
 }
 
 /**