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