WIP - simplified encrypted chat interface

lib_crypto.ts

(async () => {
  const atob = window.atob;
  const btoa = window.btoa;

  const DB_NAME = 'crypto_keys';
  const DB_VERSION = 1;
  const SIGNING_KEYS_STORE = 'signingKeys';
  const ENCRYPTION_KEYS_STORE = 'encryptionKeys';

  function base64ToBuffer(base64: string): ArrayBuffer {
    const binaryString = atob(base64);
    const len = binaryString.length;
    const bytes = new Uint8Array(len);
    for (let i = 0; i < len; i++) bytes[i] = binaryString.charCodeAt(i);
    return bytes.buffer;
  }

  function bufferToBase64(buffer: ArrayBuffer): string {
    const bytes = new Uint8Array(buffer);
    let binary = '';
    bytes.forEach((byte) => binary += String.fromCharCode(byte));
    return btoa(binary);
  }

  async function cryptoKeyToBase64(key: CryptoKey, keyType: 'public' | 'private'): Promise<string> {
    const format = keyType === 'public' ? 'spki' : 'pkcs8';
    const exportedKey = await crypto.subtle.exportKey(format, key);
    return bufferToBase64(exportedKey);
  }

  async function base64ToCryptoKey(base64: string, keyType: 'public' | 'private', algorithm: RsaHashedImportParams, keyUsages: KeyUsage[]): Promise<CryptoKey> {
    const format = keyType === 'public' ? 'spki' : 'pkcs8';
    const keyBuffer = base64ToBuffer(base64);
    return crypto.subtle.importKey(format, keyBuffer, algorithm, true, keyUsages);
  }

  async function initDB(): Promise<IDBDatabase> {
    return new Promise((resolve, reject) => {
      const request = indexedDB.open(DB_NAME, DB_VERSION);
      request.onupgradeneeded = (event) => {
        const db = request.result;
        if (!db.objectStoreNames.contains(SIGNING_KEYS_STORE)) db.createObjectStore(SIGNING_KEYS_STORE, { keyPath: 'id' });
        if (!db.objectStoreNames.contains(ENCRYPTION_KEYS_STORE)) db.createObjectStore(ENCRYPTION_KEYS_STORE, { keyPath: 'id' });
      };
      request.onerror = () => reject(request.error);
      request.onsuccess = () => resolve(request.result);
    });
  }

  const dbPromise = initDB();

  async function storeKeyPair(storeName: string, keyPair: CryptoKeyPair, id: string): Promise<void> {
    const db = await dbPromise;

    // Perform key conversions outside and before the transaction.
    const publicKeyBase64 = await cryptoKeyToBase64(keyPair.publicKey, 'public');
    const privateKeyBase64 = await cryptoKeyToBase64(keyPair.privateKey, 'private');

    // Once conversions are done, start the transaction.
    await new Promise<void>((resolve, reject) => {
      const tx = db.transaction(storeName, 'readwrite');
      const store = tx.objectStore(storeName);
      const request = store.put({ id, publicKey: publicKeyBase64, privateKey: privateKeyBase64 });

      request.onsuccess = () => resolve();
      request.onerror = () => reject(request.error);
    });
  }

  async function generateKeyPair(type: 'encryption' | 'signing'): Promise<CryptoKeyPair> {
    return crypto.subtle.generateKey({
      hash: { name: 'SHA-512' },
      name: type === 'encryption' ? 'RSA-OAEP' : 'RSA-PSS',
      modulusLength: 4096,
      publicExponent: new Uint8Array([1, 0, 1]),
    }, true, type === 'encryption' ? ['encrypt', 'decrypt'] : ['sign', 'verify']) as Promise<CryptoKeyPair>;
  }

  async function getKeyPair(storeName: string, id: string): Promise<{ publicKey: string, privateKey: string }> {
    const db = await dbPromise;
    const tx = db.transaction(storeName, 'readonly');
    const store = tx.objectStore(storeName);
    return new Promise((resolve, reject) => {
      const request = store.get(id);
      request.onsuccess = async () => {
        if (request.result) {
          resolve(request.result);
        } else {
          const keyPair = await generateKeyPair(storeName === SIGNING_KEYS_STORE ? 'signing' : 'encryption');
          const publicKeyBase64 = await cryptoKeyToBase64(keyPair.publicKey, 'public');
          const privateKeyBase64 = await cryptoKeyToBase64(keyPair.privateKey, 'private');
          await storeKeyPair(storeName, keyPair, id);
          resolve({ publicKey: publicKeyBase64, privateKey: privateKeyBase64 });
        }
      };
      request.onerror = () => reject(request.error);
    });
  }

  async function exportPublicKey(context: string, type: 'encryption' | 'signing'): Promise<string> {
    const { publicKey } = await getKeyPair(type === 'signing' ? SIGNING_KEYS_STORE : ENCRYPTION_KEYS_STORE, context);
    return publicKey;
  }

  async function encryptMessage(message: string, publicKeyBase64: string): Promise<string> {
    const publicKey = await base64ToCryptoKey(publicKeyBase64, 'public', { name: 'RSA-OAEP', hash: { name: 'SHA-512' } } as RsaHashedImportParams, ['encrypt']);
    const encoder = new TextEncoder();
    const encodedMessage = encoder.encode(message);
    const ciphertextBuffer = await crypto.subtle.encrypt({ name: 'RSA-OAEP' }, publicKey, encodedMessage);
    return bufferToBase64(ciphertextBuffer);
  }

  async function decryptMessage(context: string, message: string): Promise<string> {
    const { privateKey: privateKeyBase64 } = await getKeyPair(ENCRYPTION_KEYS_STORE, context);
    const privateKey = await base64ToCryptoKey(privateKeyBase64, 'private', { name: 'RSA-OAEP', hash: { name: 'SHA-512' } } as RsaHashedImportParams, ['decrypt']);
    const ciphertextBuffer = base64ToBuffer(message);
    const decryptedMessage = await crypto.subtle.decrypt({ name: 'RSA-OAEP' }, privateKey, ciphertextBuffer);
    return new TextDecoder().decode(decryptedMessage);
  }

  async function signMessage(context: string, message: string): Promise<string> {
    const { privateKey: privateKeyBase64 } = await getKeyPair(SIGNING_KEYS_STORE, context);
    const privateKey = await base64ToCryptoKey(privateKeyBase64, 'private', { name: 'RSA-PSS', hash: { name: 'SHA-512' } } as RsaHashedImportParams, ['sign']);
    const encoder = new TextEncoder();
    const messageBuffer = encoder.encode(message);
    const signatureBuffer = await crypto.subtle.sign({ name: 'RSA-PSS', saltLength: 32 }, privateKey, messageBuffer);
    return bufferToBase64(signatureBuffer);
  }

  async function verifySignature(message: string, signature: string, publicKey: string): Promise<boolean> {
    const publicCryptoKey = await base64ToCryptoKey(publicKey, 'public', { name: 'RSA-PSS', hash: { name: 'SHA-512' } } as RsaHashedImportParams, ['verify']);
    const messageBuffer = new TextEncoder().encode(message);
    const signatureBuffer = base64ToBuffer(signature);
    return crypto.subtle.verify({ name: 'RSA-PSS', saltLength: 32 }, publicCryptoKey, signatureBuffer, messageBuffer);
  }

  async function encryptAndSign(context: string, message: string, publicKey: string): Promise<{ ciphertext: string; signature: string }> {
    const ciphertext = await encryptMessage(message, publicKey);
    const signature = await signMessage(context, ciphertext);
    return { ciphertext, signature };
  }

  async function verifyAndDecrypt(context: string, message: string, signature: string, publicKey: string): Promise<string | null> {
    const isValidSignature = await verifySignature(message, signature, publicKey);
    if (!isValidSignature) throw new Error('Signature verification failed');
    const decryptedMessage = await decryptMessage(context, message);
    return decryptedMessage;
  }
  
  const lib_crypto = Object.freeze({
    exportPublicKey,
    encryptMessage, decryptMessage,
    signMessage, verifySignature,
    encryptAndSign, verifyAndDecrypt,
  });

  Object.defineProperty(window, 'lib_crypto', {
    value: lib_crypto,
    writable: false,
    enumerable: true,
    configurable: false,
  });
})();