packages/crypto/lib/worker/workerPool.ts

   1 import type { Remote } from 'comlink';
   2 import { releaseProxy, transferHandlers, wrap } from 'comlink';
   3
   4 import { getIsNetworkError } from '@proton/shared/lib/api/helpers/apiErrorHelper';
   5 import { captureMessage } from '@proton/shared/lib/helpers/sentry';
   6
   7 import type { Api as CryptoApi, ApiInterface as CryptoApiInterface } from './api';
   8 import { mainThreadTransferHandlers } from './transferHandlers';
   9
  10 export interface WorkerInitOptions {
  11     v6Canary?: boolean;
  12 }
  13
  14 export interface WorkerPoolInitOptions {
  15     poolSize?: number;
  16     openpgpConfigOptions?: WorkerInitOptions;
  17 }
  18
  19 export interface WorkerPoolInterface extends CryptoApiInterface {
  20     /**
  21      * Setup worker pool (singleton instance):
  22      * create and start workers, and initializes internal Crypto API (incl. pmcrypto and OpenPGP.js)
  23      * @param options.poolSize - number of workers to start; defaults to `Navigator.hardwareConcurrency()` if available, otherwise to 1.
  24      */
  25     init(options?: WorkerPoolInitOptions): Promise<void>;
  26
  27     /**
  28      * Close all workers, after clearing their internal key store.
  29      * After the pool has been destroyed, it is possible to `init()` it again.
  30      */
  31     destroy(): Promise<void>;
  32 }
  33
  34 const errorReporter = (err: Error) => {
  35     if (getIsNetworkError(err)) {
  36         captureMessage('Network error in crypto worker', {
  37             level: 'info',
  38             extra: { message: err.message },
  39         });
  40     }
  41
  42     throw err;
  43 };
  44
  45 const reportKeyCompatibilityErrorIfPresent = (compatibilityError: Error | null) => {
  46     compatibilityError &&
  47         captureMessage('Key compatibility error', {
  48             level: 'info',
  49             extra: { message: compatibilityError.message, stack: compatibilityError.stack },
  50         });
  51 };
  52
  53 // Singleton worker pool.
  54 export const CryptoWorkerPool: WorkerPoolInterface = (() => {
  55     let workerPool: Remote<CryptoApi>[] | null = null;
  56     let i = -1;
  57
  58     const initWorker = async (openpgpConfigOptions: WorkerInitOptions) => {
  59         // Webpack static analyser is not especially powerful at detecting web workers that require bundling,
  60         // see: https://github.com/webpack/webpack.js.org/issues/4898#issuecomment-823073304.
  61         // Harcoding the path here is the easiet way to get the worker to be bundled properly.
  62         const RemoteApi = wrap<typeof CryptoApi>(
  63             openpgpConfigOptions.v6Canary
  64                 ? new Worker(
  65                       new URL(
  66                           /* webpackChunkName: "crypto-worker-v6-canary" */
  67                           './worker_v6_canary.ts',
  68                           import.meta.url
  69                       )
  70                   )
  71                 : new Worker(
  72                       new URL(
  73                           /* webpackChunkName: "crypto-worker" */
  74                           './worker.ts',
  75                           import.meta.url
  76                       )
  77                   )
  78         );
  79
  80         const worker = await new RemoteApi();
  81         return worker;
  82     };
  83
  84     const destroyWorker = async (worker: Remote<CryptoApi>) => {
  85         await worker?.clearKeyStore();
  86         worker?.[releaseProxy]();
  87     };
  88
  89     /**
  90      * Get worker from the pool pool. By default, the workers are picked in a round-robin fashion, to balance the load.
  91      * However, this might not be desirable for operations like e.g. argon2, which is resource intensive and caches them
  92      * (wasm module & allocated memory) across calls.
  93      * @param [fixed] - whether to always return the same worker
  94      */
  95     const getWorker = (fixed = false): Remote<CryptoApi> => {
  96         if (workerPool == null) {
  97             throw new Error('Uninitialised worker pool');
  98         }
  99         if (fixed) {
 100             return workerPool[0];
 101         }
 102         i = (i + 1) % workerPool.length;
 103         return workerPool[i];
 104     };
 105
 106     // The return type is technically `Remote<CryptoApi>[]` but that removes some type inference capabilities that are
 107     // useful to type-check the internal worker pool functions.
 108     const getAllWorkers = (): CryptoApi[] => {
 109         if (workerPool == null) {
 110             throw new Error('Uninitialised worker pool');
 111         }
 112         return workerPool as any as CryptoApi[];
 113     };
 114
 115     return {
 116         init: async ({ poolSize = navigator.hardwareConcurrency || 1, openpgpConfigOptions = {} } = {}) => {
 117             if (workerPool !== null) {
 118                 throw new Error('worker pool already initialised');
 119             }
 120             // We load one worker early to ensure the browser serves the cached resources to the rest of the pool
 121             workerPool = [await initWorker(openpgpConfigOptions)];
 122             if (poolSize > 1) {
 123                 workerPool = workerPool.concat(
 124                     await Promise.all(new Array(poolSize - 1).fill(null).map(() => initWorker(openpgpConfigOptions)))
 125                 );
 126             }
 127             mainThreadTransferHandlers.forEach(({ name, handler }) => transferHandlers.set(name, handler));
 128         },
 129         destroy: async () => {
 130             workerPool && (await Promise.all(workerPool.map(destroyWorker)));
 131             workerPool = null;
 132         },
 133         // @ts-ignore marked as non-callable, unclear why, might be due to a limitation of type Remote
 134         encryptMessage: (opts) => getWorker().encryptMessage(opts).catch(errorReporter),
 135         decryptMessage: (opts) => getWorker().decryptMessage(opts).catch(errorReporter),
 136         // @ts-ignore marked as non-callable, unclear why, might be due to a limitation of type Remote
 137         signMessage: (opts) => getWorker().signMessage(opts).catch(errorReporter),
 138         // @ts-ignore marked as non-callable, unclear why, might be due to a limitation of type Remote
 139         verifyMessage: (opts) => getWorker().verifyMessage(opts),
 140         verifyCleartextMessage: (opts) => getWorker().verifyCleartextMessage(opts).catch(errorReporter),
 141         processMIME: (opts) => getWorker().processMIME(opts).catch(errorReporter),
 142         computeHash: (opts) => getWorker().computeHash(opts).catch(errorReporter),
 143         computeHashStream: (opts) => getWorker().computeHashStream(opts).catch(errorReporter),
 144         computeArgon2: (opts) => getWorker(true).computeArgon2(opts).catch(errorReporter),
 145
 146         generateSessionKey: (opts) => getWorker().generateSessionKey(opts).catch(errorReporter),
 147         generateSessionKeyForAlgorithm: (opts) => getWorker().generateSessionKeyForAlgorithm(opts).catch(errorReporter),
 148         encryptSessionKey: (opts) => getWorker().encryptSessionKey(opts).catch(errorReporter),
 149         decryptSessionKey: (opts) => getWorker().decryptSessionKey(opts).catch(errorReporter),
 150         importPrivateKey: async (opts) => {
 151             const [first, ...rest] = getAllWorkers();
 152             const result = await first.importPrivateKey(opts).catch(errorReporter);
 153             reportKeyCompatibilityErrorIfPresent(result._getCompatibilityError());
 154             await Promise.all(rest.map((worker) => worker.importPrivateKey(opts, result._idx)));
 155             return result;
 156         },
 157         importPublicKey: async (opts) => {
 158             const [first, ...rest] = getAllWorkers();
 159             const result = await first.importPublicKey(opts).catch(errorReporter);
 160             reportKeyCompatibilityErrorIfPresent(result._getCompatibilityError());
 161             await Promise.all(rest.map((worker) => worker.importPublicKey(opts, result._idx)));
 162             return result;
 163         },
 164         generateKey: async (opts) => {
 165             const [first, ...rest] = getAllWorkers();
 166             const keyReference = await first.generateKey(opts).catch(errorReporter);
 167             const key = await first.exportPrivateKey({ privateKey: keyReference, passphrase: null, format: 'binary' });
 168             await Promise.all(
 169                 rest.map((worker) => worker.importPrivateKey({ binaryKey: key, passphrase: null }, keyReference._idx))
 170             );
 171             return keyReference;
 172         },
 173         reformatKey: async (opts) => {
 174             const [first, ...rest] = getAllWorkers();
 175             const keyReference = await first.reformatKey(opts).catch(errorReporter);
 176             const key = await first.exportPrivateKey({ privateKey: keyReference, passphrase: null, format: 'binary' });
 177             await Promise.all(
 178                 rest.map((worker) => worker.importPrivateKey({ binaryKey: key, passphrase: null }, keyReference._idx))
 179             );
 180             return keyReference;
 181         },
 182         generateE2EEForwardingMaterial: (opts) => getWorker().generateE2EEForwardingMaterial(opts).catch(errorReporter),
 183         doesKeySupportE2EEForwarding: async (opts) =>
 184             getWorker().doesKeySupportE2EEForwarding(opts).catch(errorReporter),
 185         isE2EEForwardingKey: async (opts) => getWorker().isE2EEForwardingKey(opts).catch(errorReporter),
 186
 187         replaceUserIDs: async (opts) => {
 188             await Promise.all(getAllWorkers().map((worker) => worker.replaceUserIDs(opts)));
 189         },
 190         cloneKeyAndChangeUserIDs: async (opts) => {
 191             const [first, ...rest] = getAllWorkers();
 192             const keyReference = await first.cloneKeyAndChangeUserIDs(opts).catch(errorReporter);
 193             const key = await first.exportPrivateKey({ privateKey: keyReference, passphrase: null, format: 'binary' });
 194             await Promise.all(
 195                 rest.map((worker) => worker.importPrivateKey({ binaryKey: key, passphrase: null }, keyReference._idx))
 196             );
 197             return keyReference;
 198         },
 199         exportPublicKey: (opts) => getWorker().exportPublicKey(opts).catch(errorReporter),
 200         exportPrivateKey: (opts) => getWorker().exportPrivateKey(opts).catch(errorReporter),
 201         clearKeyStore: async () => {
 202             await Promise.all(getAllWorkers().map((worker) => worker.clearKeyStore()));
 203         },
 204         clearKey: async (opts) => {
 205             await Promise.all(getAllWorkers().map((worker) => worker.clearKey(opts)));
 206         },
 207
 208         isExpiredKey: (opts) => getWorker().isExpiredKey(opts).catch(errorReporter),
 209         isRevokedKey: (opts) => getWorker().isRevokedKey(opts).catch(errorReporter),
 210         canKeyEncrypt: (opts) => getWorker().canKeyEncrypt(opts).catch(errorReporter),
 211         getSHA256Fingerprints: (opts) => getWorker().getSHA256Fingerprints(opts),
 212         getMessageInfo: (opts) => getWorker().getMessageInfo(opts).catch(errorReporter),
 213         getKeyInfo: (opts) => getWorker().getKeyInfo(opts).catch(errorReporter),
 214         getSignatureInfo: (opts) => getWorker().getSignatureInfo(opts).catch(errorReporter),
 215         getArmoredKeys: (opts) => getWorker().getArmoredKeys(opts),
 216         getArmoredSignature: (opts) => getWorker().getArmoredSignature(opts),
 217         getArmoredMessage: (opts) => getWorker().getArmoredMessage(opts),
 218     } as WorkerPoolInterface; // casting needed to 'reuse' CryptoApi's parametric types declarations and preserve dynamic inference of
 219     // the output types based on the input ones.
 220 })();