fix(nano): Store kid (#334)

* fix(nano): Fixes `url` property when with kid

- Adds a several tests for ResourceLocator
- Since the unit tests I've added use the hex encoding for ease of reading, this change also updates the hex encoding functions and greatly increases test coverage for them

* fix(nano): actually store kid

* Update utils.ts

* Update access.ts

* 🤖 🎨 Autoformat

* Update access.ts

* Update server.ts

* fixes

* 🤖 🎨 Autoformat

* fix for v1 kas public key fallback

---------

Co-authored-by: dmihalcik-virtru <[email protected]>
Co-authored-by: Tyler Biscoe <[email protected]>

Author: dmihalcik-virtru

lib/src/access.ts

@@ -50,15 +50,61 @@ export async function fetchWrappedKey(
   return response.json();
 }
 
-export async function fetchECKasPubKey(kasEndpoint: string): Promise<CryptoKey> {
-  const kasPubKeyResponse = await fetch(`${kasEndpoint}/kas_public_key?algorithm=ec:secp256r1`);
+export type KasPublicKeyAlgorithm = 'ec:secp256r1' | 'rsa:2048';
+
+export type KasPublicKeyInfo = {
+  url: string;
+  algorithm: KasPublicKeyAlgorithm;
+  kid?: string;
+  publicKey: string;
+  key: Promise<CryptoKey>;
+};
+
+/**
+ * If we have KAS url but not public key we can fetch it from KAS, fetching
+ * the value from `${kas}/kas_public_key`.
+ */
+
+export async function fetchECKasPubKey(kasEndpoint: string): Promise<KasPublicKeyInfo> {
+  validateSecureUrl(kasEndpoint);
+  const pkUrlV2 = `${kasEndpoint}/v2/kas_public_key?algorithm=ec:secp256r1&v=2`;
+  const kasPubKeyResponse = await fetch(pkUrlV2);
   if (!kasPubKeyResponse.ok) {
-    throw new Error(
-      `Unable to validate KAS [${kasEndpoint}]. Received [${kasPubKeyResponse.status}:${kasPubKeyResponse.statusText}]`
-    );
+    if (kasPubKeyResponse.status != 404) {
+      throw new Error(
+        `unable to load KAS public key from [${pkUrlV2}]. Received [${kasPubKeyResponse.status}:${kasPubKeyResponse.statusText}]`
+      );
+    }
+    console.log('falling back to v1 key');
+    // most likely a server that does not implement v2 endpoint, so no key identifier
+    const pkUrlV1 = `${kasEndpoint}/kas_public_key?algorithm=ec:secp256r1`;
+    const r2 = await fetch(pkUrlV1);
+    if (!r2.ok) {
+      throw new Error(
+        `unable to load KAS public key from [${pkUrlV1}]. Received [${r2.status}:${r2.statusText}]`
+      );
+    }
+    const pem = await r2.json();
+    console.log('pem returned', pem);
+    return {
+      key: pemToCryptoPublicKey(pem),
+      publicKey: pem,
+      url: kasEndpoint,
+      algorithm: 'ec:secp256r1',
+    };
+  }
+  const jsonContent = await kasPubKeyResponse.json();
+  const { publicKey, kid }: KasPublicKeyInfo = jsonContent;
+  if (!publicKey) {
+    throw new Error(`Invalid response from public key endpoint [${JSON.stringify(jsonContent)}]`);
   }
-  const pem = await kasPubKeyResponse.json();
-  return pemToCryptoPublicKey(pem);
+  return {
+    key: pemToCryptoPublicKey(publicKey),
+    publicKey,
+    url: kasEndpoint,
+    algorithm: 'ec:secp256r1',
+    ...(kid && { kid }),
+  };
 }
 
 const origin = (u: string): string => {

lib/src/encodings/hex.ts

@@ -1,19 +1,51 @@
+import { InvalidCharacterError } from './base64.js';
+
 export function encode(str: string): string {
   let hex = '';
   for (let i = 0; i < str.length; i++) {
-    hex += `${str.charCodeAt(i).toString(16)}`;
+    const s = str.charCodeAt(i).toString(16);
+    if (s.length < 2) {
+      hex += '0' + s;
+    } else if (s.length > 2) {
+      throw new InvalidCharacterError(`invalid input at char ${i} == [${hex.substring(i, i + 1)}]`);
+    } else {
+      hex += `${s}`;
+    }
   }
   return hex;
 }
 
 export function decode(hex: string): string {
+  if (hex.length & 1) {
+    throw new InvalidCharacterError('invalid input.');
+  }
   let str = '';
   for (let i = 0; i < hex.length; i += 2) {
-    str += String.fromCharCode(parseInt(hex.substr(i, 2), 16));
+    const b = parseInt(hex.substring(i, i + 2), 16);
+    if (isNaN(b)) {
+      throw new InvalidCharacterError(`invalid input at char ${i} == [${hex.substring(i, i + 2)}]`);
+    }
+    str += String.fromCharCode(b);
   }
   return str;
 }
 
+export function decodeArrayBuffer(hex: string): ArrayBuffer | never {
+  const binLength = hex.length >> 1; // 1 byte per 2 characters
+  if (hex.length & 1) {
+    throw new InvalidCharacterError('invalid input.');
+  }
+  const bytes = new Uint8Array(binLength);
+  for (let i = 0; i < hex.length; i += 2) {
+    const b = parseInt(hex.substring(i, i + 2), 16);
+    if (isNaN(b)) {
+      throw new InvalidCharacterError(`invalid input at char ${i} == [${hex.substring(i, i + 2)}]`);
+    }
+    bytes[i >> 1] = b;
+  }
+  return bytes.buffer;
+}
+
 export function encodeArrayBuffer(arrayBuffer: ArrayBuffer): string | never {
   if (typeof arrayBuffer !== 'object') {
     throw new TypeError('Expected input to be an ArrayBuffer Object');

lib/src/index.ts

@@ -155,14 +155,7 @@ export class NanoTDFClient extends Client {
     payloadIV[10] = lengthAsUint24[1];
     payloadIV[11] = lengthAsUint24[0];
 
-    return encrypt(
-      policyObjectAsStr,
-      this.kasPubKey,
-      this.kasUrl,
-      ephemeralKeyPair,
-      payloadIV,
-      data
-    );
+    return encrypt(policyObjectAsStr, this.kasPubKey, ephemeralKeyPair, payloadIV, data);
   }
 }
 
@@ -274,14 +267,13 @@ export class NanoTDFDatasetClient extends Client {
       // Generate a symmetric key.
       this.symmetricKey = await keyAgreement(
         ephemeralKeyPair.privateKey,
-        this.kasPubKey,
+        await this.kasPubKey.key,
         await getHkdfSalt(DefaultParams.magicNumberVersion)
       );
 
       const nanoTDFBuffer = await encrypt(
         policyObjectAsStr,
         this.kasPubKey,
-        this.kasUrl,
         ephemeralKeyPair,
         ivVector,
         data

lib/src/nanotdf/Client.ts

@@ -3,7 +3,7 @@ import * as base64 from '../encodings/base64.js';
 import { generateKeyPair, keyAgreement } from '../nanotdf-crypto/index.js';
 import getHkdfSalt from './helpers/getHkdfSalt.js';
 import DefaultParams from './models/DefaultParams.js';
-import { fetchWrappedKey, OriginAllowList } from '../access.js';
+import { fetchWrappedKey, KasPublicKeyInfo, OriginAllowList } from '../access.js';
 import { AuthProvider, isAuthProvider, reqSignature } from '../auth/providers.js';
 import { UnsafeUrlError } from '../errors.js';
 import { cryptoPublicToPem, pemToCryptoPublicKey, validateSecureUrl } from '../utils.js';
@@ -106,7 +106,7 @@ export default class Client {
     This is needed as the flow is very specific. Errors should be thrown if the necessary step is not completed.
   */
   protected kasUrl: string;
-  kasPubKey?: CryptoKey;
+  kasPubKey?: KasPublicKeyInfo;
   readonly authProvider: AuthProvider;
   readonly dpopEnabled: boolean;
   dissems: string[] = [];
@@ -341,6 +341,7 @@ export default class Client {
 
       return unwrappedKey;
     } catch (cause) {
+      console.error('rewrap fail', cause);
       throw new Error('Could not rewrap key with entity object.', { cause });
     }
   }

lib/src/nanotdf/encrypt.ts

@@ -15,21 +15,20 @@ import {
   digest,
   exportCryptoKey,
 } from '../nanotdf-crypto/index.js';
+import { KasPublicKeyInfo } from '../access.js';
 
 /**
  * Encrypt the plain data into nanotdf buffer
  *
  * @param policy Policy that will added to the nanotdf
- * @param kasPub
- * @param kasUrl KAS url as string or ResourceLocator
+ * @param kasInfo KAS url and public key data
  * @param ephemeralKeyPair SDK ephemeral key pair to generate symmetric key
  * @param iv
  * @param data The data to be encrypted
  */
 export default async function encrypt(
   policy: string,
-  kasPub: CryptoKey,
-  kasUrl: string | ResourceLocator,
+  kasInfo: KasPublicKeyInfo,
   ephemeralKeyPair: CryptoKeyPair,
   iv: Uint8Array,
   data: string | TypedArray | ArrayBuffer
@@ -40,17 +39,17 @@ export default async function encrypt(
   }
   const symmetricKey = await keyAgreement(
     ephemeralKeyPair.privateKey,
-    kasPub,
+    await kasInfo.key,
     // Get the hkdf salt params
     await getHkdfSalt(DefaultParams.magicNumberVersion)
   );
 
   // Construct the kas locator
   let kasResourceLocator;
-  if (kasUrl instanceof ResourceLocator) {
-    kasResourceLocator = kasUrl;
+  if (kasInfo.kid) {
+    kasResourceLocator = ResourceLocator.parse(kasInfo.url, kasInfo.kid);
   } else {
-    kasResourceLocator = ResourceLocator.parse(kasUrl);
+    kasResourceLocator = ResourceLocator.parse(kasInfo.url);
   }
 
   // Auth tag length for policy and payload

lib/src/nanotdf/models/ResourceLocator.ts

@@ -47,7 +47,7 @@ export default class ResourceLocator {
         protocolIdentifierByte[0] = 0x01;
         break;
       default:
-        throw new Error('Resource locator protocol is not supported.');
+        throw new Error('resource locator protocol unsupported');
     }
 
     // Set identifier padded length and protocol identifier byte
@@ -149,7 +149,7 @@ export default class ResourceLocator {
   }
 
   get url(): string | never {
-    switch (this.protocol) {
+    switch (this.protocol & 0xf) {
       case ProtocolEnum.Http:
         return 'http://' + this.body;
       case ProtocolEnum.Https:

lib/src/tdf/AttributeObject.ts

@@ -1,4 +1,4 @@
-import { cryptoPublicToPem } from '../utils.js';
+import { type KasPublicKeyInfo } from '../access.js';
 
 export interface AttributeObject {
   readonly attribute: string;
@@ -13,14 +13,14 @@ export interface AttributeObject {
 
 export async function createAttribute(
   attribute: string,
-  pubKey: CryptoKey,
+  pubKey: KasPublicKeyInfo,
   kasUrl: string
 ): Promise<AttributeObject> {
   return {
     attribute,
     isDefault: false,
     displayName: '',
-    pubKey: await cryptoPublicToPem(pubKey),
+    pubKey: pubKey.publicKey,
     kasUrl,
     schemaVersion: '1.1.0',
   };

lib/src/utils.ts

@@ -1,4 +1,6 @@
 import { type AxiosResponseHeaders, type RawAxiosResponseHeaders } from 'axios';
+import { exportSPKI, importX509 } from 'jose';
+
 import { base64 } from './encodings/index.js';
 import { pemCertToCrypto, pemPublicToCrypto } from './nanotdf-crypto/index.js';
 
@@ -126,5 +128,18 @@ export async function pemToCryptoPublicKey(pem: string): Promise<CryptoKey> {
   } else if (/-----BEGIN CERTIFICATE-----/.test(pem)) {
     return pemCertToCrypto(pem);
   }
-  throw new Error('unsupported pem type');
+  throw new Error(`unsupported pem type [${pem}]`);
+}
+
+export async function extractPemFromKeyString(keyString: string): Promise<string> {
+  let pem: string = keyString;
+
+  // Skip the public key extraction if we find that the KAS url provides a
+  // PEM-encoded key instead of certificate
+  if (keyString.includes('CERTIFICATE')) {
+    const cert = await importX509(keyString, 'RS256', { extractable: true });
+    pem = await exportSPKI(cert);
+  }
+
+  return pem;
 }

lib/tdf3/src/client/index.ts

@@ -14,7 +14,6 @@ import {
   buildKeyAccess,
   EncryptConfiguration,
   fetchKasPublicKey,
-  KasPublicKeyInfo,
   unwrapHtml,
   validatePolicyObject,
   readStream,
@@ -31,7 +30,7 @@ import {
   withHeaders,
 } from '../../../src/auth/auth.js';
 import EAS from '../../../src/auth/Eas.js';
-import { cryptoPublicToPem, validateSecureUrl } from '../../../src/utils.js';
+import { cryptoPublicToPem, pemToCryptoPublicKey, validateSecureUrl } from '../../../src/utils.js';
 
 import {
   EncryptParams,
@@ -50,7 +49,7 @@ import {
   type DecryptSource,
   EncryptParamsBuilder,
 } from './builders.js';
-import { OriginAllowList } from '../../../src/access.js';
+import { KasPublicKeyInfo, OriginAllowList } from '../../../src/access.js';
 import { TdfError } from '../../../src/errors.js';
 import { EntityObject } from '../../../src/tdf/EntityObject.js';
 import { Binary } from '../binary.js';
@@ -337,6 +336,7 @@ export class Client {
       this.kasKeys[this.kasEndpoint] = Promise.resolve({
         url: this.kasEndpoint,
         algorithm: 'rsa:2048',
+        key: pemToCryptoPublicKey(clientConfig.kasPublicKey),
         publicKey: clientConfig.kasPublicKey,
       });
     }