[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.
 	 *
@@ -260,6 +263,7 @@ export class SealClient {
 					break;
 				}
 			}
+
 			if (hasAllKeys) {
 				completedServerCount++;
 			}
@@ -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);
 				}
@@ -324,6 +330,7 @@ export class SealClient {
 				// Return early if the completed servers is more than threshold.
 				if (hasAllKeys) {
 					completedServerCount++;
+
 					if (completedServerCount >= threshold) {
 						controller.abort();
 					}
@@ -332,6 +339,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,32 +57,54 @@ 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) => {
-			let res;
-			try {
-				res = await client.core.getObject({
-					objectId,
-				});
-			} catch (e) {
-				throw new InvalidGetObjectError(`KeyServer ${objectId} not found; ${(e as Error).message}`);
+	const uniqueIds = Array.from(new Set(objectIds));
+	const { objects } = await client.core.getObjects({
+		objectIds: uniqueIds,
+	});
+
+	// Create a single pass lookup map from objectId to key server data.
+	const serverDataMap = Object.fromEntries(
+		objects.map((res, i) => {
+			const objectId = uniqueIds[i];
+
+			if (res instanceof Error) {
+				throw new InvalidGetObjectError(
+					`KeyServer ${objectId} not found; ${(res as Error).message}`,
+				);
 			}
 
-			const ks = KeyServerMove.parse(res.object.content);
+			const ks = KeyServerMove.parse(res.content);
 			if (ks.keyType !== 0) {
 				throw new UnsupportedFeatureError(`Unsupported key type ${ks.keyType}`);
 			}
 
-			return {
+			return [
 				objectId,
-				name: ks.name,
-				url: ks.url,
-				keyType: KeyServerType.BonehFranklinBLS12381,
-				pk: new Uint8Array(ks.pk),
-			};
+				{
+					objectId,
+					name: ks.name,
+					url: ks.url,
+					pk: new Uint8Array(ks.pk),
+				},
+			];
 		}),
 	);
+
+	// Return the preserved order of the input objectIds, creating a new object for each occurence.
+	return objectIds.map((objectId) => {
+		const data = serverDataMap[objectId];
+		if (!data) {
+			throw new InvalidGetObjectError(`KeyServer ${objectId} not found`);
+		}
+
+		return {
+			objectId,
+			name: data.name,
+			url: data.url,
+			keyType: KeyServerType.BonehFranklinBLS12381,
+			pk: data.pk,
+		};
+	});
 }
 
 /**