✨ multikas v2: splitPlans

- Adds the ability to share and split DEKs
- Reconstructs keys using share ids
- Does NOT support TDFs that were created with id-less splits. While the old code did support reading them, it did not support creating them

lib/tdf3/index.ts

@@ -10,6 +10,7 @@ import {
   type DecryptKeyMiddleware,
   type DecryptStreamMiddleware,
   EncryptParamsBuilder,
+  type SplitStep,
 } from './src/client/builders.js';
 import { type ClientConfig, createSessionKeys } from './src/client/index.js';
 import {
@@ -56,6 +57,7 @@ export type {
   EncryptStreamMiddleware,
   DecryptKeyMiddleware,
   DecryptStreamMiddleware,
+  SplitStep,
 };
 
 export {

lib/tdf3/src/client/builders.ts

@@ -26,6 +26,11 @@ export type EncryptStreamMiddleware = (
   stream: DecoratedReadableStream
 ) => Promise<DecoratedReadableStream>;
 
+export type SplitStep = {
+  kas: string;
+  split?: string;
+};
+
 export type EncryptParams = {
   source: ReadableStream<Uint8Array>;
   opts?: { keypair: PemKeyPair };
@@ -40,6 +45,7 @@ export type EncryptParams = {
   eo?: EntityObject;
   payloadKey?: Binary;
   keyMiddleware?: EncryptKeyMiddleware;
+  splitPlan?: SplitStep[],
   streamMiddleware?: EncryptStreamMiddleware;
 };
 

lib/tdf3/src/client/index.ts

@@ -40,6 +40,7 @@ import {
   DecryptStreamMiddleware,
   EncryptKeyMiddleware,
   EncryptStreamMiddleware,
+  SplitStep,
 } from './builders.js';
 import { DecoratedReadableStream } from './DecoratedReadableStream.js';
 
@@ -221,7 +222,7 @@ export class Client {
    */
   readonly allowedKases: string[];
 
-  readonly kasPublicKey: Promise<KasPublicKeyInfo>;
+  readonly kasKeys: Record<string, Promise<KasPublicKeyInfo>>;
 
   readonly easEndpoint?: string;
 
@@ -328,13 +329,17 @@ export class Client {
       dpopKeys: clientConfig.dpopKeys,
     });
     if (clientConfig.kasPublicKey) {
-      this.kasPublicKey = Promise.resolve({
-        url: this.kasEndpoint,
-        algorithm: 'rsa:2048',
-        publicKey: clientConfig.kasPublicKey,
-      });
-    } else {
-      this.kasPublicKey = fetchKasPublicKey(this.kasEndpoint);
+      this.kasKeys[this.kasEndpoint] = Promise.resolve({
+          url: this.kasEndpoint,
+          algorithm: 'rsa:2048',
+          publicKey: clientConfig.kasPublicKey,
+        });
+    }
+    for (const kasEndpoint of this.allowedKases) {
+      if (kasEndpoint in this.kasKeys) {
+        continue;
+      }
+      this.kasKeys[kasEndpoint] = fetchKasPublicKey(this.kasEndpoint);
     }
   }
 
@@ -365,9 +370,11 @@ export class Client {
     eo,
     keyMiddleware = defaultKeyMiddleware,
     streamMiddleware = async (stream: DecoratedReadableStream) => stream,
+    splitPlan,
   }: EncryptParams): Promise<DecoratedReadableStream> {
     const dpopKeys = await this.dpopKeys;
     const kasPublicKey = await this.kasPublicKey;
+
     const policyObject = asPolicy(scope);
     validatePolicyObject(policyObject);
 
@@ -383,14 +390,20 @@ export class Client {
       eo.attributes.forEach((attr) => s.addJwtAttribute(attr));
       attributeSet = s;
     }
-    encryptionInformation.keyAccess.push(
-      await buildKeyAccess({
-        attributeSet,
-        type: offline ? 'wrapped' : 'remote',
-        url: kasPublicKey.url,
-        kid: kasPublicKey.kid,
-        publicKey: kasPublicKey.publicKey,
-        metadata,
+
+    const splits: SplitStep[] = splitPlan || [{kas: this.kasEndpoint}]
+    encryptionInformation.keyAccess = await Promise.all(
+      splits.map(async ({kas, split}) => {
+        const kasPublicKey = await this.kasKeys[kas];
+        return buildKeyAccess({
+          attributeSet,
+          type: offline ? 'wrapped' : 'remote',
+          url: kasPublicKey.url,
+          kid: kasPublicKey.kid,
+          publicKey: kasPublicKey.publicKey,
+          metadata,
+          split,
+        });
       })
     );
     const { keyForEncryption, keyForManifest } = await (keyMiddleware as EncryptKeyMiddleware)();

lib/tdf3/src/models/encryption-information.ts

@@ -24,8 +24,10 @@ export type Segment = {
   readonly encryptedSegmentSize?: number;
 };
 
+export type SplitType = 'split' | 'flat';
+
 export type EncryptionInformation = {
-  readonly type: string;
+  readonly type: SplitType;
   readonly keyAccess: KeyAccessObject[];
   readonly integrityInformation: {
     readonly rootSignature: {
@@ -75,19 +77,21 @@ export class SplitKey {
   }
 
   async getKeyAccessObjects(policy: Policy, keyInfo: KeyInfo): Promise<KeyAccessObject[]> {
+    const splitIds = [...new Set(this.keyAccess.map(({ split }) => split))].sort();
     const unwrappedKeySplitBuffers = await keySplit(
       new Uint8Array(keyInfo.unwrappedKeyBinary.asByteArray()),
-      this.keyAccess.length,
+      splitIds.length,
       this.cryptoService
     );
+    const splitsByName = Object.fromEntries(splitIds.map((split, index) => [split, unwrappedKeySplitBuffers[index]]));
 
     const keyAccessObjects = [];
-    for (let i = 0; i < this.keyAccess.length; i++) {
+    for (const item of this.keyAccess) {
       // use the key split to encrypt metadata for each key access object
-      const unwrappedKeySplitBuffer = unwrappedKeySplitBuffers[i];
+      const unwrappedKeySplitBuffer = splitsByName[item.split];
       const unwrappedKeySplitBinary = Binary.fromArrayBuffer(unwrappedKeySplitBuffer.buffer);
 
-      const metadata = this.keyAccess[i].metadata || '';
+      const metadata = item.metadata || '';
       const metadataStr = (
         typeof metadata === 'object'
           ? JSON.stringify(metadata)
@@ -112,7 +116,7 @@ export class SplitKey {
       };
 
       const encryptedMetadataStr = JSON.stringify(encryptedMetadataOb);
-      const keyAccessObject = await this.keyAccess[i].write(
+      const keyAccessObject = await item.write(
         policy,
         unwrappedKeySplitBuffer,
         encryptedMetadataStr
@@ -139,7 +143,7 @@ export class SplitKey {
     const policyForManifest = base64.encode(JSON.stringify(policy));
 
     return {
-      type: 'split',
+      type: 'flat',
       keyAccess: keyAccessObjects,
       method: {
         algorithm,

lib/tdf3/src/models/key-access.ts

@@ -17,7 +17,8 @@ export class Wrapped {
     public readonly url: string,
     public readonly kid: string | undefined,
     public readonly publicKey: string,
-    public readonly metadata: unknown
+    public readonly metadata: unknown,
+    public readonly split: string
   ) {}
 
   async write(
@@ -48,6 +49,9 @@ export class Wrapped {
     if (this.kid) {
       this.keyAccessObject.kid = this.kid;
     }
+    if (this.split?.length) {
+      this.keyAccessObject.split = this.split;
+    }
 
     return this.keyAccessObject;
   }
@@ -63,7 +67,8 @@ export class Remote {
     public readonly url: string,
     public readonly kid: string | undefined,
     public readonly publicKey: string,
-    public readonly metadata: unknown
+    public readonly metadata: unknown,
+    public readonly split: string
   ) {}
 
   async write(
@@ -103,6 +108,7 @@ export class Remote {
 export type KeyAccess = Remote | Wrapped;
 
 export type KeyAccessObject = {
+  split?: string;
   type: KeyAccessType;
   url: string;
   kid?: string;

lib/tdf3/src/tdf.ts

@@ -18,6 +18,8 @@ import {
   UpsertResponse,
   Wrapped as KeyAccessWrapped,
   KeyAccess,
+  KeyAccessObject,
+  SplitType,
 } from './models/index.js';
 import { base64 } from '../../src/encodings/index.js';
 import {
@@ -69,7 +71,7 @@ export type EncryptionOptions = {
   /**
    * Defaults to `split`, the currently only implmented key wrap algorithm.
    */
-  type?: string;
+  type?: SplitType;
   // Defaults to AES-256-GCM for the encryption.
   cipher?: string;
 };
@@ -91,6 +93,7 @@ export type BuildKeyAccess = {
   publicKey: string;
   attributeUrl?: string;
   metadata?: Metadata;
+  split?: string;
 };
 
 type Segment = {
@@ -339,6 +342,7 @@ export async function buildKeyAccess({
   kid,
   attributeUrl,
   metadata,
+  split = '',
 }: BuildKeyAccess): Promise<KeyAccess> {
   /** Internal function to keep it DRY */
   function createKeyAccess(
@@ -350,9 +354,9 @@ export async function buildKeyAccess({
   ) {
     switch (type) {
       case 'wrapped':
-        return new KeyAccessWrapped(kasUrl, kasKeyIdentifier, pubKey, metadata);
+        return new KeyAccessWrapped(kasUrl, kasKeyIdentifier, pubKey, metadata, split);
       case 'remote':
-        return new KeyAccessRemote(kasUrl, kasKeyIdentifier, pubKey, metadata);
+        return new KeyAccessRemote(kasUrl, kasKeyIdentifier, pubKey, metadata, split);
       default:
         throw new KeyAccessError(`buildKeyAccess: Key access type ${type} is unknown`);
     }
@@ -799,6 +803,29 @@ async function loadTDFStream(
   return { manifest, zipReader, centralDirectory };
 }
 
+export function splitLookupTableFactory(keyAccess: KeyAccessObject[], allowedKases: string[]): Record<string, Record<string, KeyAccessObject>> {
+  const splitIds = new Set(keyAccess.map(({split}) => split || ''));
+  const accessibleSplits = new Set(keyAccess.filter(({url}) => allowedKases.includes(url)).map(({split}) => split));
+  if (splitIds.size > accessibleSplits.size) {
+    const disallowedKases = new Set(keyAccess.filter(({url}) => !allowedKases.includes(url)).map(({url}) => url));
+    throw new KasDecryptError(
+      `Unreconstructable key - disallowed KASes include: [${JSON.stringify(disallowedKases)}] from splitIds [${JSON.stringify(splitIds)}]`
+    );
+  }
+  const splitPotentials: Record<string, Record<string, KeyAccessObject>> = Object.fromEntries([...splitIds].map(s => [s, {}]));
+  for (const kao of keyAccess) {
+    const disjunction  = splitPotentials[kao.split || ''];
+    if (kao.url in disjunction) {
+      throw new KasDecryptError(`TODO: Fallback to no split ids. Repetition found for [${kao.url}] on split [${kao.split}]`);
+    }
+    if (allowedKases.includes(kao.url)) {
+      disjunction[kao.url] = kao;
+    }
+  }
+  return splitPotentials;
+
+}
+
 async function unwrapKey({
   manifest,
   allowedKases,
@@ -815,17 +842,24 @@ async function unwrapKey({
   cryptoService: CryptoService;
 }) {
   if (authProvider === undefined) {
-    throw new Error('Upsert can be done without auth provider');
+    throw new KasDecryptError('Upsert can be done without auth provider');
   }
   const { keyAccess } = manifest.encryptionInformation;
+  const splitPotentials = splitLookupTableFactory(keyAccess, allowedKases);
+
   let responseMetadata;
   const isAppIdProvider = authProvider && isAppIdProviderCheck(authProvider);
   // Get key access information to know the KAS URLS
   const rewrappedKeys = await Promise.all(
-    keyAccess.map(async (keySplitInfo) => {
-      if (!allowedKases.includes(keySplitInfo.url)) {
-        throw new KasUpsertError(`Unexpected KAS url: [${keySplitInfo.url}]`);
+    Object.entries(splitPotentials).map(async ([splitId, potentials]) => {
+      if (!potentials || !Object.keys(potentials).length) {
+        throw new KasDecryptError(
+          `Unreconstructable key - no valid KAS found for split ${JSON.stringify(splitId)}`
+          );
       }
+      // TODO: If we have multiple ways of getting a value, try the 'best' way
+      // or maybe retry across all potential ways
+      const [keySplitInfo] = Object.values(potentials);
       const url = `${keySplitInfo.url}/${isAppIdProvider ? '' : 'v2/'}rewrap`;
 
       const ephemeralEncryptionKeys = await cryptoService.cryptoToPemPair(

lib/tests/mocha/unit/tdf.spec.ts

@@ -76,6 +76,14 @@ describe('TDF', () => {
   });
 });
 
+describe('splitLookupTableFactory', () => {
+  it('throws when empty', () => {
+
+  });
+});
+
+describe('', () => {});
+
 describe('fetchKasPublicKey', async () => {
   it('missing kas names throw', async () => {
     try {