@getanthill/datastore/src/services/fhe/index.ts

Event-Sourced Datastore

/**
 * @credits https://github.com/s0l0ist/node-seal
 *
 * @see https://medium.com/@s0l0ist/homomorphic-encryption-for-web-apps-a3fa52e9f03d
 * @see https://medium.com/@s0l0ist/homomorphic-encryption-for-web-apps-b615fb64d2a2
 */
import type { SEALLibrary } from 'node-seal/implementation/seal';
import type { CipherText } from 'node-seal/implementation/cipher-text';
import type { PlainText } from 'node-seal/implementation/plain-text';

import vm from 'vm';

// @ts-ignore
import { Context } from 'node-seal/implementation/context';
import { SecretKey } from 'node-seal/implementation/secret-key';
import { PublicKey } from 'node-seal/implementation/public-key';
import { KeyGenerator } from 'node-seal/implementation/key-generator';
import { RelinKeys } from 'node-seal/implementation/relin-keys';
import { GaloisKeys } from 'node-seal/implementation/galois-keys';
import { Evaluator } from 'node-seal/implementation/evaluator';

import { mapValuesDeep } from '../../utils';
import { overloadAsString } from './operator-overload';
import { CKKSEncoder } from 'node-seal/implementation/ckks-encoder';
import { BatchEncoder } from 'node-seal/implementation/batch-encoder';

let sealPromise: Promise<SEALLibrary> | null = null;

interface Deletable {
  delete: Function;
}

export interface FullyHomomorphicEncryptionClientConfig {
  scheme?: 'bfv' | 'bgv' | 'ckks';
  polyModulusDegree?: number;
  bitSize?: number;
  bitSizes?: Array<number>;
}

export default class FullyHomomorphicEncryptionClient {
  private _config: FullyHomomorphicEncryptionClientConfig;
  private _seal: SEALLibrary | null = null;
  private _context: Context | null = null;
  private _keys: {
    secret: SecretKey;
    public: PublicKey;
    relin: RelinKeys;
    galois: GaloisKeys;
  } | null = null;
  private _evaluator: Evaluator | null = null;
  /**
   * @see https://github.com/s0l0ist/node-seal?tab=readme-ov-file#caveats
   * Must have a `delete()` method
   */
  private _instances: Array<Deletable> = [];

  /**
   * Sample cipher instance to test the constructor equality
   * @see computeFromAny()
   */
  private _cipherSample: CipherText | null = null;
  private _encoder: CKKSEncoder | BatchEncoder | null = null;

  constructor(config: FullyHomomorphicEncryptionClientConfig) {
    this._config = config;
  }

  /**
   * @see https://github.com/s0l0ist/node-seal
   */
  async connect() {
    if (this._seal !== null) {
      return this;
    }

    if (sealPromise === null) {
      const SEAL = require('node-seal');

      sealPromise = SEAL();
    }

    this._seal = await sealPromise;

    this.init();
  }

  async init() {
    await this.connect();

    // Create a new EncryptionParameters
    const schemeType = this.seal.SchemeType[this.scheme];
    const securityLevel = this.seal.SecurityLevel.tc128;
    const polyModulusDegree = this.polyModulusDegree;
    const bitSizes = this.bitSizes;
    const bitSize = this.bitSize;

    const encParms = this.seal.EncryptionParameters(schemeType);

    // Assign Poly Modulus Degree
    encParms.setPolyModulusDegree(polyModulusDegree);

    // Create a suitable set of CoeffModulus primes
    encParms.setCoeffModulus(
      this.seal.CoeffModulus.Create(
        polyModulusDegree,
        Int32Array.from(bitSizes),
      ),
    );

    // Assign a PlainModulus (only for bfv/bgv scheme type)
    if (this.scheme !== 'ckks') {
      encParms.setPlainModulus(
        this.seal.PlainModulus.Batching(polyModulusDegree, bitSize),
      );
    }

    // Create a new Context
    this._context = this.seal.Context(encParms, true, securityLevel);

    // Helper to check if the Context was created successfully
    /* istanbul ignore next */
    if (!this._context.parametersSet()) {
      throw new Error(
        'Could not set the parameters in the given context. Please try different encryption parameters.',
      );
    }

    this.createKeys();

    this._evaluator = this.seal.Evaluator(this.context!);
    this.addInstance(this._evaluator);

    this._cipherSample = this.seal.CipherText();
    this.addInstance(this._cipherSample);

    if (this.scheme === 'ckks') {
      this._encoder = this.seal.CKKSEncoder(this._context!);
    } else {
      this._encoder = this.seal.BatchEncoder(this._context!);
    }
    this.addInstance(this._encoder);

    return this;
  }

  async disconnect() {
    if (this._seal !== null) {
      this.flushInstances();
      this._seal = null;
    }

    return this;
  }

  async clone() {
    const client = new FullyHomomorphicEncryptionClient(this._config);

    await client.connect();

    client._context = this.context;
    client._keys = this.keys;

    return client;
  }

  get seal() {
    if (this._seal === null) {
      throw new Error('Node seal must be initialized first with `.connect()`');
    }

    return this._seal;
  }

  get context() {
    return this._context;
  }

  get keys() {
    return this._keys;
  }

  get scheme() {
    return this._config.scheme ?? 'bgv';
  }

  get bitSize() {
    return this._config.bitSize ?? 20;
  }

  get bitSizes() {
    return this._config.bitSizes ?? [36, 36, 37];
  }
  get polyModulusDegree() {
    return this._config.polyModulusDegree ?? 4096;
  }

  get evaluator() {
    return this._evaluator!;
  }

  addInstance(instance: Deletable) {
    this._instances.push(instance);
  }

  flushInstances() {
    for (const instance of this._instances) {
      instance.delete();
    }

    this._instances = [];
  }

  createKeys(_keyGenerator?: KeyGenerator) {
    // Create a new KeyGenerator (use uploaded keys if applicable)
    const keyGenerator =
      _keyGenerator ?? this.seal.KeyGenerator(this._context!);

    // Get the SecretKey from the keyGenerator
    const Secret_key_A_ = keyGenerator.secretKey();

    // Get the PublicKey from the keyGenerator
    const Public_key_A_ = keyGenerator.createPublicKey();

    // Create a new RelinKey
    const Relin_key_A_ = keyGenerator.createRelinKeys();

    // Create a new GaloisKey
    const Galois_key_A_ = keyGenerator.createGaloisKeys();

    this._keys = {
      secret: Secret_key_A_,
      public: Public_key_A_,
      relin: Relin_key_A_,
      galois: Galois_key_A_,
    };
  }

  encode(arr: Array<number>) {
    if (this.scheme === 'ckks') {
      const float64 = Float64Array.from(arr);

      return (this._encoder! as CKKSEncoder).encode(
        float64,
        Math.pow(2, this.bitSize),
      );
    }

    const int32 = Int32Array.from(arr);

    return (this._encoder as BatchEncoder).encode(int32);
  }

  createCypher(value: Array<number>): CipherText {
    const plainText = this.encode(value);

    const encryptor = this.seal.Encryptor(this._context!, this.keys!.public);

    return encryptor.encrypt(plainText!)!;
  }

  decode(plainText: PlainText): Array<number> {
    if (this.scheme === 'ckks') {
      const float64 = this._encoder!.decode(plainText);

      return Array.from(float64);
    }

    const int32 = this._encoder!.decode(plainText!);

    return Array.from(int32);
  }

  fromCypher(cypher: CipherText): Array<number> {
    const decryptor = this.seal.Decryptor(this._context!, this._keys!.secret);

    const plainText = decryptor.decrypt(cypher);

    return this.decode(plainText!);
  }

  async computeFromAnyToCiphers(
    handler: Function,
    ...args: Array<any>
  ): Promise<any> {
    return new Promise((resolve, reject) => {
      /* istanbul ignore next */
      const rawScript = overloadAsString(handler, args, this);

      const script = new vm.Script(rawScript);
      const context = {
        client: this,
        handler,
        args,
        resolve,
        reject,
        console,
      };
      script.runInNewContext(context, {
        timeout: 120_000,
        breakOnSigint: true,
      });
    });
  }

  async computeFromAny(handler: Function, ...args: Array<any>): Promise<any> {
    return this.computeFromAnyToCiphers(handler, ...args).then((res) => {
      if (res.constructor === this._cipherSample?.constructor) {
        return res.save();
      }

      const result = mapValuesDeep(res, (val) => {
        if (!!val && typeof val.save === 'function') {
          return val.save();
        }

        return val;
      });

      return result;
    });
  }

  async compute(handler: Function, ..._args: Array<string>) {
    const client = await this.clone();

    const value = await client.computeFromAny(handler, ..._args);

    /**
     * @see https://github.com/s0l0ist/node-seal?tab=readme-ov-file#caveats
     * and specifically Garbage Collection
     */
    client.flushInstances();

    return value;
  }

  compile(script: string) {
    const body = script.split('\n').slice(1, -1).join('\n');
    const args = script
      .split('\n', 1)[0]
      .replace(/^(async )?function( [a-z]{1,30})?\(/, '')
      .split(')', 1)[0]!
      .split(',')
      .map((t) => t.trim());

    const AsyncFunction = Object.getPrototypeOf(
      /* istanbul ignore next */
      async function () {},
    ).constructor;

    var retFn = AsyncFunction(...args, body);

    return retFn;
  }

  // FHE functions
  relinearize(val: CipherText) {
    return this.evaluator.relinearize(val, this.keys!.relin);
  }
}