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

Event-Sourced Datastore

import type App from '../../App';

import setup from '../../../test/setup';

import FHEClient from '.';

describe('services/fhe', () => {
  let client: FHEClient;
  let app: App;

  beforeAll(async () => {
    app = await setup.build();

    client = new FHEClient({});

    await client.connect();
  });

  afterAll(async () => {
    client.disconnect();
    jest.resetAllMocks();
    await setup.teardownDb(app.services.mongodb);
  });

  describe('#connect / #init / #end', () => {
    it('binds the seal implementation on connect', async () => {
      const client = new FHEClient({});
      let error;
      try {
        client.seal;
      } catch (err) {
        error = err;
      }

      expect(error.message).toEqual(
        'Node seal must be initialized first with `.connect()`',
      );
    });

    it('binds the seal implementation on connect', async () => {
      const client = new FHEClient({});
      // @ts-expect-error For tests only
      expect(client._seal).toEqual(null);

      await client.connect();
      // @ts-expect-error For tests only
      expect(client._seal).not.toEqual(null);
    });

    it('binds the seal implementation on connect only once', async () => {
      const client = new FHEClient({});
      // @ts-expect-error For tests only
      expect(client._seal).toEqual(null);

      await client.connect();

      await client.connect();
      // @ts-expect-error For tests only
      expect(client._seal).not.toEqual(null);
    });

    it('noops on disconnect of seal on not connected instance', async () => {
      const client = new FHEClient({});
      // @ts-expect-error For tests only
      expect(client._seal).toEqual(null);

      await client.disconnect();
      // @ts-expect-error For tests only
      expect(client._seal).toEqual(null);
    });

    it('removes the seal library from the instance', async () => {
      const client = new FHEClient({});
      // @ts-expect-error For tests only
      expect(client._seal).toEqual(null);

      await client.connect();
      // @ts-expect-error For tests only
      expect(client._seal).not.toEqual(null);

      await client.disconnect();
      // @ts-expect-error For tests only
      expect(client._seal).toEqual(null);
    });
  });

  describe('#createCypher', () => {
    it('cyphers are not comparable each others', async () => {
      // Given
      const a = client.createCypher([1]);
      const b = client.createCypher([1]);

      // When
      //

      // Then
      expect(a.save()).not.toEqual(b.save());
    });

    it('cyphers can be safely loaded after save', async () => {
      // Given
      const a = client.createCypher([1]);

      // When
      const savedResult = a.save();
      const b = client.seal.CipherText();
      b.load(client.context!, savedResult);

      // Then
      expect(a.save()).toEqual(b.save());
    });
  });

  describe('#computeFromAnyToCiphers', () => {
    it('results of the exact same computation are the exactly the same', async () => {
      // Given
      const a = client.createCypher([1]);
      const b = client.createCypher([1]);

      // When
      const c1 = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a - b;
        },
        a,
        b,
      );
      const c2 = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a - b;
        },
        a,
        b,
      );

      // Then
      expect(c1.save()).toEqual(c2.save());
    });

    it('results of the exact same computation but with different cyphers are different', async () => {
      // Given
      const a = client.createCypher([1]);
      const b = client.createCypher([1]);
      const c = client.createCypher([1]);

      // When
      const c1 = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a - b;
        },
        a,
        b,
      );
      const c2 = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a - b;
        },
        a,
        c,
      );

      // Then
      expect(c1.save()).not.toEqual(c2.save());
    });

    it('computes the sum of 2 integers', async () => {
      // Given
      const a = client.createCypher([1]);
      const b = client.createCypher([2]);

      // When
      const c = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a + b;
        },
        a,
        b,
      );

      // Then
      expect(client.fromCypher(c).slice(0, 1)).toEqual([3]);
    });

    it('computes the sum of 2 array of integers', async () => {
      // Given
      const a = client.createCypher([1, 2]);
      const b = client.createCypher([3, 4]);

      // When
      const c = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a + b;
        },
        a,
        b,
      );

      // Then
      expect(client.fromCypher(c).slice(0, 2)).toEqual([4, 6]);
    });

    it('computes the sum of 2 array of integers', async () => {
      // Given
      const a = client.createCypher([1, 2, 3]);
      const b = client.createCypher([2, 3, 4]);

      // When
      const c = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a + b;
        },
        a,
        b,
      );

      // Then
      expect(client.fromCypher(c).slice(0, 3)).toEqual([3, 5, 7]);
    });

    it('computes the multiplication of 2 integers', async () => {
      // Given
      const a = client.createCypher([2]);
      const b = client.createCypher([3]);

      // When
      const c = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a * b;
        },
        a,
        b,
      );

      // Then
      expect(client.fromCypher(c).slice(0, 1)).toEqual([6]);
    });

    it('computes the sum of an integer with the number 1', async () => {
      // Given
      const a = client.createCypher([1]);

      // When
      const c = await client.computeFromAnyToCiphers(async function handler(a) {
        return a + 1;
      }, a);

      // Then
      expect(client.fromCypher(c).slice(0, 1)).toEqual([2]);
    });

    it('computes the sum of an array of integers with the number 1 scalar', async () => {
      // Given
      const a = client.createCypher([1, 2, 3]);

      // When
      const c = await client.computeFromAnyToCiphers(async function handler(a) {
        return a + 1;
      }, a);

      // Then
      expect(client.fromCypher(c).slice(0, 3)).toEqual([2, 3, 4]);
    });

    it('computes an operation with priorities', async () => {
      // Given
      const a = client.createCypher([1]);

      // When
      const c = await client.computeFromAnyToCiphers(async function handler(a) {
        const b = 3 - a;

        return (-a + 3) * b - 2;
      }, a);

      // Then
      expect(client.fromCypher(c).slice(0, 1)).toEqual([2]);
    });

    it('computes an operation with self negation', async () => {
      // Given
      const a = client.createCypher([1]);

      // When
      const c = await client.computeFromAnyToCiphers(async function handler(a) {
        return -a;
      }, a);

      // Then
      expect(client.fromCypher(c).slice(0, 1)).toEqual([-1]);
    });

    it('computes an operation with assignment', async () => {
      // Given
      const a = client.createCypher([1]);

      // When
      const c = await client.computeFromAnyToCiphers(async function handler(a) {
        const b = a + 2;
        return a - b;
      }, a);

      // Then
      expect(client.fromCypher(c).slice(0, 1)).toEqual([-2]);
    });

    it('computes an operation with function expression', async () => {
      // Given
      const a = client.createCypher([1]);

      // When
      const c = await client.computeFromAnyToCiphers(async function handler(a) {
        function b(c) {
          return c + 2;
        }
        return a - b(a);
      }, a);

      // Then
      expect(client.fromCypher(c).slice(0, 1)).toEqual([-2]);
    });

    it.skip('computes a Math.max operation', async () => {
      // Given
      const a = client.createCypher([1]);
      const b = client.createCypher([2]);

      // When
      const c = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return Math.max.call(this, a, b);
        },
        a,
        b,
      );

      // Then
      expect(client.fromCypher(c).slice(0, 1)).toEqual([-2]);
    });

    it('computes the multiplication of 2 floats', async () => {
      // Given
      const client = new FHEClient({
        scheme: 'ckks',
      });

      await client.connect();

      const a = client.createCypher([2.5]);
      const b = client.createCypher([3.3]);

      // When
      const c = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a * b;
        },
        a,
        b,
      );

      // Then
      const target = [2.5 * 3.3];
      expect(
        client
          .fromCypher(c)
          .slice(0, 1)
          .map((v, i) => Math.abs(v - target[i]) / target[i])
          .reduce((s, c) => s + c, 0),
      ).toBeLessThanOrEqual(0.005);
    });

    it('computes the multiplication of 2 high floats', async () => {
      // Given
      const client = new FHEClient({
        scheme: 'ckks',
      });

      await client.connect();

      const a = client.createCypher([123456.789]);
      const b = client.createCypher([987654.321]);

      // When
      const c = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a * b;
        },
        a,
        b,
      );

      // Then
      const target = [123456.789 * 987654.321];
      expect(
        client
          .fromCypher(c)
          .slice(0, 1)
          .map((v, i) => Math.abs(v - target[i]) / target[i])
          .reduce((s, c) => s + c, 0),
      ).toBeLessThanOrEqual(0.005);
    });

    it('computes the addition of 2 integers from different clients', async () => {
      // Given
      const alice = new FHEClient({});
      await alice.connect();

      const bernard = new FHEClient({});
      await bernard.connect();

      // Create new public from the one given by Alice:
      const UploadedPublicKey = bernard.seal.PublicKey();
      UploadedPublicKey.load(bernard.context!, alice.keys!.public.save());
      bernard.keys!.public = UploadedPublicKey;

      const a = alice.createCypher([2]);

      // The cypher is then created for Alice or Bernard computation
      const b = bernard.createCypher([3]);

      // When
      const aliceComputationResult = await alice.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a + b;
        },
        a,
        b,
      );
      const bernardComputationResult = await alice.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a + b;
        },
        a,
        b,
      );

      // Then
      // But only Alice can read the result because of the secret key
      expect(alice.fromCypher(aliceComputationResult).slice(0, 1)).toEqual([
        2 + 3,
      ]);
      expect(
        bernard.fromCypher(aliceComputationResult).slice(0, 1),
      ).not.toEqual([2 + 3]);
      expect(alice.fromCypher(bernardComputationResult).slice(0, 1)).toEqual([
        2 + 3,
      ]);
    });

    it('computes the addition of 2 integers given by Alice to Bernard', async () => {
      // Given
      const alice = new FHEClient({});
      await alice.connect();

      const a = alice.createCypher([2]);
      const b = alice.createCypher([3]);
      const alicePublicKey = alice.keys!.public.save();

      const bernard = new FHEClient({});
      await bernard.connect();

      // Create new public from the one given by Alice:
      const UploadedPublicKey = bernard.seal.PublicKey();
      UploadedPublicKey.load(bernard.context!, alicePublicKey);
      bernard.keys!.public = UploadedPublicKey;

      // When
      const c = await bernard.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a + b;
        },
        a,
        b,
      );

      // Then
      const savedResult = c.save();
      const cipherA = alice.seal.CipherText();
      cipherA.load(alice.context!, savedResult);

      expect(alice.fromCypher(cipherA).slice(0, 1)).toEqual([2 + 3]);
    });
  });

  describe('#computeFromAny', () => {
    it('computes a sum from raw values', async () => {
      // Given
      const state = {
        data: client.createCypher([1]).save(),
      };
      const event = {
        data: client.createCypher([2]).save(),
      };

      // When
      const res = await client.computeFromAny(
        async (state, event) => {
          const data = event.data + state.data + 2;

          return [
            {
              ...state,
              data,
            },
          ];
        },
        state,
        event,
      );

      // Then
      const cipher = client.seal.CipherText();
      cipher.load(client.context!, res[0].data);
      expect(client.fromCypher(cipher).slice(0, 1)).toEqual([5]);
    });

    it('computes a sum from raw values', async () => {
      // Given
      const state = null;
      const event = {
        data: client.createCypher([2]).save(),
      };

      // When
      const res = await client.computeFromAny(
        async (state, event) => {
          const data = event.data + (state?.data ?? 1) + 2;

          return [
            {
              ...state,
              data,
            },
          ];
        },
        state,
        event,
      );

      // Then
      const cipher = client.seal.CipherText();
      cipher.load(client.context!, res[0].data);
      expect(client.fromCypher(cipher).slice(0, 1)).toEqual([5]);
    });

    it('computes the addition assignation', async () => {
      // Given
      const state = {
        data: client.createCypher([1]).save(),
      };
      const event = {
        data: client.createCypher([2]).save(),
      };

      // When
      const res = await client.computeFromAny(
        async (state, event) => {
          state.data += 2;

          return [state];
        },
        state,
        event,
      );

      // Then
      const cipher = client.seal.CipherText();
      cipher.load(client.context!, res[0].data);
      expect(client.fromCypher(cipher).slice(0, 1)).toEqual([3]);
    });

    it('computes the subtraction assignation', async () => {
      // Given
      const state = {
        data: client.createCypher([1]).save(),
      };
      const event = {
        data: client.createCypher([2]).save(),
      };

      // When
      const res = await client.computeFromAny(
        async (state, event) => {
          state.data -= 2;

          return [state];
        },
        state,
        event,
      );

      // Then
      const cipher = client.seal.CipherText();
      cipher.load(client.context!, res[0].data);
      expect(client.fromCypher(cipher).slice(0, 1)).toEqual([-1]);
    });

    it('computes the multiplication assignation', async () => {
      // Given
      const state = {
        data: client.createCypher([2]).save(),
      };
      const event = {
        data: client.createCypher([3]).save(),
      };

      // When
      const res = await client.computeFromAny(
        async (state, event) => {
          state.data *= event.data;

          return [state];
        },
        state,
        event,
      );

      // Then
      const cipher = client.seal.CipherText();
      cipher.load(client.context!, res[0].data);
      expect(client.fromCypher(cipher).slice(0, 1)).toEqual([6]);
    });

    it('computes a subtraction from raw values', async () => {
      // Given
      const state = {
        data: client.createCypher([2]).save(),
      };
      const event = {
        data: client.createCypher([1]).save(),
      };

      // When
      const res = await client.computeFromAny(
        async (state, event) => {
          const data = state.data - event.data;

          return [
            {
              ...state,
              data,
            },
          ];
        },
        state,
        event,
      );

      // Then
      const cipher = client.seal.CipherText();
      cipher.load(client.context!, res[0].data);
      expect(client.fromCypher(cipher).slice(0, 1)).toEqual([1]);
    });

    it('computes the negation of a raw value', async () => {
      // Given
      const state = {
        data: client.createCypher([2]).save(),
      };
      const event = {
        data: client.createCypher([1]).save(),
      };

      // When
      const res = await client.computeFromAny(
        async (state, event) => {
          const data = -event.data;

          return [
            {
              ...state,
              data,
            },
          ];
        },
        state,
        event,
      );

      // Then
      const cipher = client.seal.CipherText();
      cipher.load(client.context!, res[0].data);
      expect(client.fromCypher(cipher).slice(0, 1)).toEqual([-1]);
    });

    it('computes the multiplication of raw values', async () => {
      // Given
      const state = {
        data: client.createCypher([2]).save(),
      };
      const event = {
        data: client.createCypher([3]).save(),
      };

      // When
      const res = await client.computeFromAny(
        async (state, event) => {
          const data = state.data * event.data;

          return [
            {
              ...state,
              data,
            },
          ];
        },
        state,
        event,
      );

      // Then
      const cipher = client.seal.CipherText();
      cipher.load(client.context!, res[0].data);
      expect(client.fromCypher(cipher).slice(0, 1)).toEqual([6]);
    });
  });

  describe('#vm', () => {
    it('computes the sum of 2 integers in a vm', async () => {
      // Given
      const a = client.createCypher([1]).save();
      const b = client.createCypher([2]).save();

      // When
      const c = await client.compute(
        async function handler(a, b) {
          return a + b;
        },
        a,
        b,
      );

      // Then
      const cipher = client.seal.CipherText();
      cipher.load(client.context!, c);
      expect(client.fromCypher(cipher).slice(0, 1)).toEqual([3]);
    });

    it('computes the sum of 2 integers in a vm with a sync function', async () => {
      // Given
      const a = client.createCypher([1]).save();
      const b = client.createCypher([2]).save();

      // When
      const c = await client.compute(
        function handler(a, b) {
          return a + b;
        },
        a,
        b,
      );

      // Then
      const cipher = client.seal.CipherText();
      cipher.load(client.context!, c);
      expect(client.fromCypher(cipher).slice(0, 1)).toEqual([3]);
    });

    it('computes the sum of 2 integers with a predefined script', async () => {
      // Given
      const a = client.createCypher([1]).save();
      const b = client.createCypher([2]).save();
      const fn = client.compile(
        function handler(a, b) {
          return a + b;
        }.toString(),
      );

      // When
      const c = await client.compute(fn, a, b);

      // Then
      const cipher = client.seal.CipherText();
      cipher.load(client.context!, c);
      expect(client.fromCypher(cipher).slice(0, 1)).toEqual([3]);
    });
  });

  describe('bgv', () => {
    it('sum 2 integers together', async () => {
      // Given
      const client = new FHEClient({
        scheme: 'bgv',
      });
      await client.connect();

      const a = client.createCypher([1, 2]);
      const b = client.createCypher([3, 4]);

      // When
      const c = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a + b;
        },
        a,
        b,
      );

      // Then
      expect(client.fromCypher(c).slice(0, 2)).toEqual([4, 6]);
    });
  });

  describe('bfv', () => {
    it('sum 2 integers together', async () => {
      // Given
      const client = new FHEClient({
        scheme: 'bfv',
      });
      await client.connect();

      const a = client.createCypher([1, 2]);
      const b = client.createCypher([3, 4]);

      // When
      const c = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a + b;
        },
        a,
        b,
      );

      // Then
      expect(client.fromCypher(c).slice(0, 2)).toEqual([4, 6]);
    });
  });

  describe('ckks', () => {
    it('sum 2 integers together', async () => {
      // Given
      const client = new FHEClient({
        scheme: 'ckks',
      });
      await client.connect();

      const a = client.createCypher([1.1, 2.1]);
      const b = client.createCypher([3.3, 4.4]);

      // When
      const c = await client.computeFromAnyToCiphers(
        async function handler(a, b) {
          return a + b;
        },
        a,
        b,
      );

      // Then
      const target = [4.4, 6.5];
      expect(
        client
          .fromCypher(c)
          .slice(0, 2)
          .map((v, i) => Math.abs(v - target[i]) / target[i])
          .reduce((s, c) => s + c, 0) / 2,
      ).toBeLessThanOrEqual(0.005);
    });
  });
});