✨ 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
opentdf · Jul 31, 2024 · 445e20c · 445e20c
1 parent 41cb5f4
commit 445e20c
Show file tree

Hide file tree

Showing 8 changed files with 110 additions and 48 deletions.
diff --git a/lib/src/access.ts b/lib/src/access.ts
@@ -1,4 +1,5 @@
 import { type AuthProvider } from './auth/auth.js';
+import { pemToCryptoPublicKey } from './utils.js';
 
 export class RewrapRequest {
   signedRequestToken = '';
@@ -49,13 +50,13 @@ export async function fetchWrappedKey(
   return response.json();
 }
 
-export async function fetchECKasPubKey(kasEndpoint: string): Promise<string> {
+export async function fetchECKasPubKey(kasEndpoint: string): Promise<CryptoKey> {
   const kasPubKeyResponse = await fetch(`${kasEndpoint}/kas_public_key?algorithm=ec:secp256r1`);
   if (!kasPubKeyResponse.ok) {
     throw new Error(
       `Unable to validate KAS [${kasEndpoint}]. Received [${kasPubKeyResponse.status}:${kasPubKeyResponse.statusText}]`
     );
   }
-  return kasPubKeyResponse.json();
+  const pem = await kasPubKeyResponse.json();
+  return pemToCryptoPublicKey(pem);
 }
-
diff --git a/lib/src/index.ts b/lib/src/index.ts
@@ -10,21 +10,8 @@ import {
 } from './nanotdf/index.js';
 import { keyAgreement } from './nanotdf-crypto/index.js';
 import { TypedArray, createAttribute, Policy } from './tdf/index.js';
-import { AuthProvider } from './auth/auth.js';
 import { fetchECKasPubKey } from './access.js';
 import { ClientConfig } from './nanotdf/Client.js';
-import { pemToCryptoPublicKey } from './utils.js';
-
-async function fetchKasPubKey(kasUrl: string): Promise<CryptoKey> {
-  const kasPubKeyResponse = await fetch(`${kasUrl}/kas_public_key?algorithm=ec:secp256r1`);
-  if (!kasPubKeyResponse.ok) {
-    throw new Error(
-      `Unable to validate KAS [${kasUrl}]. Received [${kasPubKeyResponse.status}:${kasPubKeyResponse.statusText}]`
-    );
-  }
-  const pem = await kasPubKeyResponse.json();
-  return pemToCryptoPublicKey(pem);
-}
 
 /**
  * NanoTDF SDK Client

diff --git a/lib/tdf3/index.ts b/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 {

diff --git a/lib/tdf3/src/client/builders.ts b/lib/tdf3/src/client/builders.ts
@@ -26,6 +26,11 @@ export type EncryptStreamMiddleware = (
   stream: DecoratedReadableStream
 ) => Promise<DecoratedReadableStream>;
 
+export type SplitStep = {
+  kas: string;
+  sid?: 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;
 };
 

diff --git a/lib/tdf3/src/client/index.ts b/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({
+      this.kasKeys[this.kasEndpoint] = Promise.resolve({
         url: this.kasEndpoint,
         algorithm: 'rsa:2048',
         publicKey: clientConfig.kasPublicKey,
       });
-    } else {
-      this.kasPublicKey = fetchKasPublicKey(this.kasEndpoint);
+    }
+    for (const kasEndpoint of this.allowedKases) {
+      if (kasEndpoint in this.kasKeys) {
+        continue;
+      }
+      this.kasKeys[kasEndpoint] = fetchKasPublicKey(this.kasEndpoint);
     }
   }
 
@@ -365,9 +370,10 @@ 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 +389,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, sid }) => {
+        const kasPublicKey = await this.kasKeys[kas];
+        return buildKeyAccess({
+          attributeSet,
+          type: offline ? 'wrapped' : 'remote',
+          url: kasPublicKey.url,
+          kid: kasPublicKey.kid,
+          publicKey: kasPublicKey.publicKey,
+          metadata,
+          sid,
+        });
       })
     );
     const { keyForEncryption, keyForManifest } = await (keyMiddleware as EncryptKeyMiddleware)();

diff --git a/lib/tdf3/src/models/encryption-information.ts b/lib/tdf3/src/models/encryption-information.ts
@@ -24,8 +24,10 @@ export type Segment = {
   readonly encryptedSegmentSize?: number;
 };
 
+export type SplitType = 'split';
+
 export type EncryptionInformation = {
-  readonly type: string;
+  readonly type: SplitType;
   readonly keyAccess: KeyAccessObject[];
   readonly integrityInformation: {
     readonly rootSignature: {
@@ -75,19 +77,23 @@ export class SplitKey {
   }
 
   async getKeyAccessObjects(policy: Policy, keyInfo: KeyInfo): Promise<KeyAccessObject[]> {
+    const splitIds = [...new Set(this.keyAccess.map(({ sid }) => sid))].sort();
     const unwrappedKeySplitBuffers = await keySplit(
       new Uint8Array(keyInfo.unwrappedKeyBinary.asByteArray()),
-      this.keyAccess.length,
+      splitIds.length,
       this.cryptoService
     );
+    const splitsByName = Object.fromEntries(
+      splitIds.map((sid, index) => [sid, 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.sid];
       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 +118,7 @@ export class SplitKey {
       };
 
       const encryptedMetadataStr = JSON.stringify(encryptedMetadataOb);
-      const keyAccessObject = await this.keyAccess[i].write(
+      const keyAccessObject = await item.write(
         policy,
         unwrappedKeySplitBuffer,
         encryptedMetadataStr

diff --git a/lib/tdf3/src/models/key-access.ts b/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 sid: string
   ) {}
 
   async write(
@@ -51,6 +52,9 @@ export class Wrapped {
     if (this.kid) {
       this.keyAccessObject.kid = this.kid;
     }
+    if (this.sid?.length) {
+      this.keyAccessObject.sid = this.sid;
+    }
 
     return this.keyAccessObject;
   }
@@ -66,7 +70,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 sid: string
   ) {}
 
   async write(
@@ -109,6 +114,7 @@ export class Remote {
 export type KeyAccess = Remote | Wrapped;
 
 export type KeyAccessObject = {
+  sid?: string;
   type: KeyAccessType;
   url: string;
   kid?: string;

diff --git a/lib/tdf3/src/tdf.ts b/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 {
@@ -70,7 +72,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;
 };
@@ -92,6 +94,7 @@ export type BuildKeyAccess = {
   publicKey: string;
   attributeUrl?: string;
   metadata?: Metadata;
+  sid?: string;
 };
 
 type Segment = {
@@ -340,6 +343,7 @@ export async function buildKeyAccess({
   kid,
   attributeUrl,
   metadata,
+  sid = '',
 }: BuildKeyAccess): Promise<KeyAccess> {
   /** Internal function to keep it DRY */
   function createKeyAccess(
@@ -351,9 +355,9 @@ export async function buildKeyAccess({
   ) {
     switch (type) {
       case 'wrapped':
-        return new KeyAccessWrapped(kasUrl, kasKeyIdentifier, pubKey, metadata);
+        return new KeyAccessWrapped(kasUrl, kasKeyIdentifier, pubKey, metadata, sid);
       case 'remote':
-        return new KeyAccessRemote(kasUrl, kasKeyIdentifier, pubKey, metadata);
+        return new KeyAccessRemote(kasUrl, kasKeyIdentifier, pubKey, metadata, sid);
       default:
         throw new KeyAccessError(`buildKeyAccess: Key access type ${type} is unknown`);
     }
@@ -800,6 +804,41 @@ 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(({ sid }) => sid || ''));
+  const accessibleSplits = new Set(
+    keyAccess.filter(({ url }) => allowedKases.includes(url)).map(({ sid }) => sid)
+  );
+  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.sid || ''];
+    if (kao.url in disjunction) {
+      throw new KasDecryptError(
+        `TODO: Fallback to no split ids. Repetition found for [${kao.url}] on split [${kao.sid}]`
+      );
+    }
+    if (allowedKases.includes(kao.url)) {
+      disjunction[kao.url] = kao;
+    }
+  }
+  return splitPotentials;
+}
+
 async function unwrapKey({
   manifest,
   allowedKases,
@@ -816,22 +855,25 @@ 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)) {
+    Object.entries(splitPotentials).map(async ([splitId, potentials]) => {
+      if (!potentials || !Object.keys(potentials).length) {
         throw new UnsafeUrlError(
-          `cannot decrypt TDF: [${keySplitInfo.url}] not on allowlist ${JSON.stringify(
-            allowedKases
-          )}`,
-          keySplitInfo.url
-        );
+          `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(