Add proving machine tests (#749)

* Add actor mock helper and tests

* format tests

* fix tests

* wip fix tests

* address cr feedback

app/tests/utils/proving/actorMock.ts

@@ -0,0 +1,22 @@
+// SPDX-License-Identifier: BUSL-1.1; Copyright (c) 2025 Social Connect Labs, Inc.; Licensed under BUSL-1.1 (see LICENSE); Apache-2.0 from 2029-06-11
+
+import { jest } from '@jest/globals';
+
+// Minimal actor stub used to observe send calls and emit state transitions
+
+export const actorMock = {
+  start: jest.fn(),
+  stop: jest.fn(),
+  send: jest.fn(),
+  subscribe: jest.fn((cb: (state: any) => void) => {
+    (actorMock as any)._callback = cb;
+    return { unsubscribe: jest.fn() };
+  }),
+};
+
+export function emitState(stateValue: string) {
+  const cb = (actorMock as any)._callback;
+  if (cb) {
+    cb({ value: stateValue, matches: (v: string) => v === stateValue });
+  }
+}

app/tests/utils/proving/provingMachine.generatePayload.test.ts

@@ -0,0 +1,224 @@
+// SPDX-License-Identifier: BUSL-1.1; Copyright (c) 2025 Social Connect Labs, Inc.; Licensed under BUSL-1.1 (see LICENSE); Apache-2.0 from 2029-06-11
+
+import { jest } from '@jest/globals';
+
+import { useProtocolStore } from '../../../src/stores/protocolStore';
+import { useSelfAppStore } from '../../../src/stores/selfAppStore';
+import { useProvingStore } from '../../../src/utils/proving/provingMachine';
+
+jest.mock('xstate', () => {
+  const actual = jest.requireActual('xstate') as any;
+  const { actorMock } = require('./actorMock');
+  return { ...actual, createActor: jest.fn(() => actorMock) };
+});
+
+jest.mock('../../../src/utils/analytics', () => () => ({
+  trackEvent: jest.fn(),
+}));
+
+jest.mock('@selfxyz/common', () => {
+  const actual = jest.requireActual('@selfxyz/common') as any;
+  return {
+    ...actual,
+    getSolidityPackedUserContextData: jest.fn(() => '0x1234'),
+  };
+});
+
+jest.mock('../../../src/utils/proving/provingInputs', () => ({
+  generateTEEInputsRegister: jest.fn(() => ({
+    inputs: { r: 1 },
+    circuitName: 'reg',
+    endpointType: 'celo',
+    endpoint: 'https://reg',
+  })),
+  generateTEEInputsDSC: jest.fn(() => ({
+    inputs: { d: 1 },
+    circuitName: 'dsc',
+    endpointType: 'celo',
+    endpoint: 'https://dsc',
+  })),
+  generateTEEInputsDisclose: jest.fn(() => ({
+    inputs: { s: 1 },
+    circuitName: 'vc_and_disclose',
+    endpointType: 'https',
+    endpoint: 'https://dis',
+  })),
+}));
+
+jest.mock('../../../src/utils/proving/provingUtils', () => {
+  const actual = jest.requireActual(
+    '../../../src/utils/proving/provingUtils',
+  ) as any;
+  return {
+    ...actual,
+    getPayload: jest.fn(() => ({ mocked: true })),
+    encryptAES256GCM: jest.fn(() => ({
+      nonce: [0],
+      cipher_text: [1],
+      auth_tag: [2],
+    })),
+  };
+});
+
+const {
+  getPayload,
+  encryptAES256GCM,
+} = require('../../../src/utils/proving/provingUtils');
+const {
+  generateTEEInputsRegister,
+  generateTEEInputsDSC,
+  generateTEEInputsDisclose,
+} = require('../../../src/utils/proving/provingInputs');
+
+function setupDefaultStores() {
+  useProvingStore.setState({
+    circuitType: 'register',
+    passportData: { documentCategory: 'passport', mock: false },
+    secret: 'sec',
+    uuid: '123',
+    sharedKey: Buffer.alloc(32, 1),
+    env: 'prod',
+  });
+
+  useSelfAppStore.setState({
+    selfApp: createMockSelfApp(),
+  });
+
+  useProtocolStore.setState({
+    passport: createMockProtocolState(),
+    id_card: createMockProtocolState(),
+  });
+}
+
+function createMockSelfApp() {
+  return {
+    chainID: 42220 as const,
+    userId: 'u',
+    userDefinedData: '0x0',
+    endpointType: 'https' as const,
+    endpoint: 'https://e',
+    scope: 's',
+    sessionId: '',
+    appName: '',
+    logoBase64: '',
+    header: '',
+    userIdType: 'uuid' as const,
+    devMode: false,
+    disclosures: {},
+    version: 1,
+  };
+}
+
+function createMockProtocolState() {
+  return {
+    dsc_tree: 'tree',
+    csca_tree: [['a']],
+    commitment_tree: null,
+    deployed_circuits: null,
+    circuits_dns_mapping: null,
+    alternative_csca: {},
+    fetch_deployed_circuits: jest.fn(() => Promise.resolve()),
+    fetch_circuits_dns_mapping: jest.fn(() => Promise.resolve()),
+    fetch_csca_tree: jest.fn(() => Promise.resolve()),
+    fetch_dsc_tree: jest.fn(() => Promise.resolve()),
+    fetch_identity_tree: jest.fn(() => Promise.resolve()),
+    fetch_alternative_csca: jest.fn(() => Promise.resolve()),
+    fetch_all: jest.fn(() => Promise.resolve()),
+  };
+}
+
+describe('_generatePayload', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    setupDefaultStores();
+  });
+
+  it('register circuit', async () => {
+    useProvingStore.setState({ circuitType: 'register' });
+    const payload = await useProvingStore.getState()._generatePayload();
+    expect(generateTEEInputsRegister).toHaveBeenCalled();
+    expect(getPayload).toHaveBeenCalled();
+    expect(encryptAES256GCM).toHaveBeenCalled();
+    expect(useProvingStore.getState().endpointType).toBe('celo');
+    expect(payload.params).toEqual({
+      uuid: '123',
+      nonce: [0],
+      cipher_text: [1],
+      auth_tag: [2],
+    });
+  });
+
+  it('dsc circuit', async () => {
+    useProvingStore.setState({ circuitType: 'dsc' });
+    const payload = await useProvingStore.getState()._generatePayload();
+    expect(generateTEEInputsDSC).toHaveBeenCalled();
+    expect(generateTEEInputsDSC).toHaveBeenCalledWith(
+      expect.objectContaining({
+        documentCategory: 'passport',
+        mock: false,
+      }),
+      [['a']], // csca_tree
+      'prod', // env
+    );
+    expect(useProvingStore.getState().endpointType).toBe('celo');
+    expect(payload.params.uuid).toBe('123');
+    expect(payload.params).toEqual({
+      uuid: '123',
+      nonce: [0],
+      cipher_text: [1],
+      auth_tag: [2],
+    });
+  });
+
+  it('disclose circuit', async () => {
+    useProvingStore.setState({ circuitType: 'disclose' });
+    const payload = await useProvingStore.getState()._generatePayload();
+    expect(generateTEEInputsDisclose).toHaveBeenCalled();
+    expect(useProvingStore.getState().endpointType).toBe('https');
+    expect(payload.params.uuid).toBe('123');
+  });
+
+  it('handles missing passport data', async () => {
+    useProvingStore.setState({ passportData: null });
+    await expect(
+      useProvingStore.getState()._generatePayload(),
+    ).rejects.toThrow();
+  });
+
+  it('handles input generation failure', async () => {
+    // Reset all mocks first
+    jest.clearAllMocks();
+    setupDefaultStores();
+
+    generateTEEInputsRegister.mockImplementation(() => {
+      throw new Error('Input generation failed');
+    });
+    await expect(useProvingStore.getState()._generatePayload()).rejects.toThrow(
+      'Input generation failed',
+    );
+
+    // Restore the mock to its original implementation
+    generateTEEInputsRegister.mockRestore();
+  });
+
+  it('handles encryption failure', async () => {
+    // Reset all mocks first
+    jest.clearAllMocks();
+    setupDefaultStores();
+
+    // Restore all original implementations first
+    generateTEEInputsRegister.mockImplementation(() => ({
+      inputs: { r: 1 },
+      circuitName: 'reg',
+      endpointType: 'celo',
+      endpoint: 'https://reg',
+    }));
+
+    encryptAES256GCM.mockImplementation(() => {
+      throw new Error('Encryption failed');
+    });
+    await expect(useProvingStore.getState()._generatePayload()).rejects.toThrow(
+      'Encryption failed',
+    );
+  });
+});

app/tests/utils/proving/provingMachine.init.test.ts

@@ -0,0 +1,167 @@
+// SPDX-License-Identifier: BUSL-1.1; Copyright (c) 2025 Social Connect Labs, Inc.; Licensed under BUSL-1.1 (see LICENSE); Apache-2.0 from 2029-06-11
+
+import { jest } from '@jest/globals';
+
+import { useProvingStore } from '../../../src/utils/proving/provingMachine';
+import { emitState } from './actorMock';
+
+jest.mock('xstate', () => {
+  const actual = jest.requireActual('xstate') as any;
+  const { actorMock } = require('./actorMock');
+  return { ...actual, createActor: jest.fn(() => actorMock) };
+});
+
+jest.mock('../../../src/providers/passportDataProvider', () => ({
+  loadSelectedDocument: jest.fn(),
+}));
+
+jest.mock('../../../src/providers/authProvider', () => ({
+  unsafe_getPrivateKey: jest.fn(),
+}));
+
+jest.mock('../../../src/utils/analytics', () => () => ({
+  trackEvent: jest.fn(),
+}));
+
+// Mock uuid v4 function
+jest.mock('uuid', () => ({
+  v4: jest.fn(() => 'uuid'),
+}));
+
+const {
+  loadSelectedDocument,
+} = require('../../../src/providers/passportDataProvider');
+const { unsafe_getPrivateKey } = require('../../../src/providers/authProvider');
+const { actorMock } = require('./actorMock');
+
+describe('provingMachine init', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    useProvingStore.setState({});
+
+    // Mock WebSocket
+    const mockWebSocket = {
+      send: jest.fn(),
+      close: jest.fn(),
+      addEventListener: jest.fn(),
+      removeEventListener: jest.fn(),
+    } as any;
+    useProvingStore.setState({ wsConnection: mockWebSocket });
+  });
+
+  it('handles missing document', async () => {
+    loadSelectedDocument.mockResolvedValue(null);
+    await useProvingStore.getState().init('register');
+    expect(actorMock.send).toHaveBeenCalledWith({
+      type: 'PASSPORT_DATA_NOT_FOUND',
+    });
+    emitState('passport_data_not_found');
+    expect(useProvingStore.getState().currentState).toBe(
+      'passport_data_not_found',
+    );
+  });
+
+  it('initializes state with document and secret', async () => {
+    loadSelectedDocument.mockResolvedValue({
+      data: { documentCategory: 'passport', mock: false },
+    });
+    unsafe_getPrivateKey.mockResolvedValue('mysecret');
+    await useProvingStore.getState().init('register');
+    expect(useProvingStore.getState().passportData).toEqual({
+      documentCategory: 'passport',
+      mock: false,
+    });
+    expect(useProvingStore.getState().secret).toBe('mysecret');
+    expect(useProvingStore.getState().env).toBe('prod');
+    expect(useProvingStore.getState().circuitType).toBe('register');
+  });
+
+  it('handles document loading error', async () => {
+    loadSelectedDocument.mockRejectedValue(new Error('Network error'));
+    await expect(useProvingStore.getState().init('register')).rejects.toThrow(
+      'Network error',
+    );
+  });
+
+  it('handles invalid document structure', async () => {
+    loadSelectedDocument.mockResolvedValue({ data: null });
+    await expect(useProvingStore.getState().init('register')).rejects.toThrow();
+  });
+
+  it('initializes state with document and secret for different circuit types', async () => {
+    const circuitTypes = ['register', 'dsc', 'disclose'] as const;
+
+    for (const circuitType of circuitTypes) {
+      loadSelectedDocument.mockResolvedValue({
+        data: { documentCategory: 'passport', mock: false },
+      });
+      unsafe_getPrivateKey.mockResolvedValue('mysecret');
+
+      await useProvingStore.getState().init(circuitType);
+
+      expect(useProvingStore.getState().circuitType).toBe(circuitType);
+      // Add more assertions for circuit-specific initialization
+    }
+  });
+
+  it('_handleWsClose handles different close codes', () => {
+    const closeCodes = [1000, 1001, 1006, 1011];
+    closeCodes.forEach(code => {
+      jest.clearAllMocks();
+      useProvingStore.setState({ currentState: 'proving' });
+      const event: any = { code, reason: `Close code ${code}`, type: 'close' };
+      useProvingStore.getState()._handleWsClose(event);
+      expect(actorMock.send).toHaveBeenCalledWith({ type: 'PROVE_ERROR' });
+    });
+  });
+
+  it('_handleWsClose ignores close during non-proving states', () => {
+    (['idle', 'completed', 'error'] as const).forEach(state => {
+      jest.clearAllMocks();
+      useProvingStore.setState({ currentState: state });
+      const event: any = { code: 1000, reason: '', type: 'close' };
+      useProvingStore.getState()._handleWsClose(event);
+      expect(actorMock.send).not.toHaveBeenCalled();
+    });
+  });
+
+  it('_handleWsOpen sends hello', () => {
+    // Verify initial state
+    expect(useProvingStore.getState().uuid).toBeNull();
+
+    useProvingStore.getState()._handleWsOpen();
+    const ws = useProvingStore.getState().wsConnection as any;
+    expect(ws.send).toHaveBeenCalled();
+    const sent = JSON.parse(ws.send.mock.calls[0][0]);
+
+    // Verify complete message structure
+    expect(sent).toEqual({
+      jsonrpc: '2.0',
+      method: 'openpassport_hello',
+      id: 1,
+      params: {
+        user_pubkey: expect.any(Array),
+        uuid: 'uuid',
+      },
+    });
+    expect(sent.params.uuid).toBe('uuid');
+    expect(useProvingStore.getState().uuid).toBe('uuid');
+  });
+
+  it('_handleWebSocketMessage handles malformed JSON', async () => {
+    const message = new MessageEvent('message', {
+      data: 'invalid json{',
+    });
+    await useProvingStore.getState()._handleWebSocketMessage(message);
+    expect(actorMock.send).toHaveBeenCalledWith({ type: 'PROVE_ERROR' });
+  });
+
+  it('_handleWebSocketMessage handles missing attestation field', async () => {
+    const message = new MessageEvent('message', {
+      data: JSON.stringify({ result: {} }),
+    });
+    await useProvingStore.getState()._handleWebSocketMessage(message);
+    // This should just log a warning, not send PROVE_ERROR
+    expect(actorMock.send).not.toHaveBeenCalled();
+  });
+});

app/tests/utils/proving/provingMachine.websocket.test.ts

@@ -0,0 +1,102 @@
+// SPDX-License-Identifier: BUSL-1.1; Copyright (c) 2025 Social Connect Labs, Inc.; Licensed under BUSL-1.1 (see LICENSE); Apache-2.0 from 2029-06-11
+
+import { jest } from '@jest/globals';
+
+import { useProvingStore } from '../../../src/utils/proving/provingMachine';
+
+jest.mock('xstate', () => {
+  const actual = jest.requireActual('xstate') as any;
+  const { actorMock } = require('./actorMock');
+  return { ...actual, createActor: jest.fn(() => actorMock) };
+});
+
+jest.mock('../../../src/utils/analytics', () => () => ({
+  trackEvent: jest.fn(),
+}));
+jest.mock('uuid', () => ({ v4: jest.fn(() => 'uuid') }));
+
+jest.mock('../../../src/utils/proving/attest', () => ({
+  getPublicKey: jest.fn(() => '04' + 'a'.repeat(128)),
+  verifyAttestation: jest.fn(() => Promise.resolve(true)),
+}));
+jest.mock('../../../src/utils/proving/provingUtils', () => ({
+  ec: { keyFromPublic: jest.fn(() => ({ getPublic: jest.fn() })) },
+  clientKey: { derive: jest.fn(() => ({ toArray: () => Array(32).fill(1) })) },
+  clientPublicKeyHex: '00',
+}));
+
+jest.mock('../../../src/providers/passportDataProvider', () => ({
+  loadSelectedDocument: jest.fn(() =>
+    Promise.resolve({ data: { documentCategory: 'passport', mock: false } }),
+  ),
+}));
+
+jest.mock('../../../src/providers/authProvider', () => ({
+  unsafe_getPrivateKey: jest.fn(() => Promise.resolve('sec')),
+}));
+
+const { actorMock } = require('./actorMock');
+
+const { verifyAttestation } = require('../../../src/utils/proving/attest');
+
+describe('websocket handlers', () => {
+  beforeEach(async () => {
+    jest.clearAllMocks();
+    useProvingStore.setState({
+      wsConnection: {
+        send: jest.fn(),
+        addEventListener: jest.fn(),
+        removeEventListener: jest.fn(),
+        close: jest.fn(),
+      } as any,
+    });
+    await useProvingStore.getState().init('register');
+    useProvingStore.setState({
+      wsConnection: {
+        send: jest.fn(),
+        addEventListener: jest.fn(),
+        removeEventListener: jest.fn(),
+        close: jest.fn(),
+      } as any,
+    });
+  });
+
+  it('_handleWsOpen sends hello', () => {
+    useProvingStore.getState()._handleWsOpen();
+    const ws = useProvingStore.getState().wsConnection as any;
+    expect(ws.send).toHaveBeenCalled();
+    const sent = JSON.parse(ws.send.mock.calls[0][0]);
+    expect(sent.params.uuid).toBe('uuid');
+    expect(useProvingStore.getState().uuid).toBe('uuid');
+  });
+
+  it('_handleWebSocketMessage processes attestation', async () => {
+    const message = new MessageEvent('message', {
+      data: JSON.stringify({ result: { attestation: 'a' } }),
+    });
+    await useProvingStore.getState()._handleWebSocketMessage(message);
+    expect(verifyAttestation).toHaveBeenCalled();
+    expect(
+      actorMock.send.mock.calls.some(
+        (c: any) => c[0].type === 'CONNECT_SUCCESS',
+      ),
+    ).toBe(true);
+  });
+
+  it('_handleWebSocketMessage handles error', async () => {
+    const message = new MessageEvent('message', {
+      data: JSON.stringify({ error: 'oops' }),
+    });
+    await useProvingStore.getState()._handleWebSocketMessage(message);
+    const lastCall = actorMock.send.mock.calls.pop();
+    expect(lastCall[0]).toEqual({ type: 'PROVE_ERROR' });
+  });
+
+  it('_handleWsClose triggers failure during proving', () => {
+    useProvingStore.setState({ currentState: 'proving' });
+    const event: any = { code: 1000, reason: '', type: 'close' };
+    useProvingStore.getState()._handleWsClose(event);
+    const last = actorMock.send.mock.calls.pop();
+    expect(last[0]).toEqual({ type: 'PROVE_ERROR' });
+  });
+});

app/tests/utils/proving/provingUtils.test.ts

@@ -0,0 +1,150 @@
+// SPDX-License-Identifier: BUSL-1.1; Copyright (c) 2025 Social Connect Labs, Inc.; Licensed under BUSL-1.1 (see LICENSE); Apache-2.0 from 2029-06-11
+
+import forge from 'node-forge';
+
+import {
+  encryptAES256GCM,
+  getPayload,
+  getWSDbRelayerUrl,
+} from '../../../src/utils/proving/provingUtils';
+
+describe('provingUtils', () => {
+  it('encryptAES256GCM encrypts and decrypts correctly', () => {
+    const key = forge.random.getBytesSync(32);
+    const plaintext = 'hello world';
+    const encrypted = encryptAES256GCM(plaintext, forge.util.createBuffer(key));
+
+    const decipher = forge.cipher.createDecipher(
+      'AES-GCM',
+      forge.util.createBuffer(key),
+    );
+    decipher.start({
+      iv: Buffer.from(encrypted.nonce).toString('binary'),
+      tagLength: 128,
+      tag: forge.util.createBuffer(
+        Buffer.from(encrypted.auth_tag).toString('binary'),
+      ),
+    });
+    decipher.update(
+      forge.util.createBuffer(
+        Buffer.from(encrypted.cipher_text).toString('binary'),
+      ),
+    );
+    const success = decipher.finish();
+    const decrypted = decipher.output.toString();
+
+    expect(success).toBe(true);
+    expect(decrypted).toBe(plaintext);
+  });
+
+  it('getPayload returns disclose payload', () => {
+    const inputs = { foo: 'bar' };
+    const payload = getPayload(
+      inputs,
+      'disclose',
+      'vc_and_disclose',
+      'https',
+      'https://example.com',
+      2,
+      '0xabc',
+    );
+    expect(payload).toEqual({
+      type: 'disclose',
+      endpointType: 'https',
+      endpoint: 'https://example.com',
+      onchain: false,
+      circuit: { name: 'vc_and_disclose', inputs: JSON.stringify(inputs) },
+      version: 2,
+      userDefinedData: '0xabc',
+    });
+  });
+
+  it('getPayload returns register payload', () => {
+    const payload = getPayload(
+      { a: 1 },
+      'register',
+      'register_circuit',
+      'celo',
+      'https://self.xyz',
+    );
+    expect(payload).toEqual({
+      type: 'register',
+      onchain: true,
+      endpointType: 'celo',
+      circuit: { name: 'register_circuit', inputs: JSON.stringify({ a: 1 }) },
+    });
+  });
+
+  it('getWSDbRelayerUrl handles endpoint types', () => {
+    expect(getWSDbRelayerUrl('celo')).toBe('wss://websocket.self.xyz');
+    expect(getWSDbRelayerUrl('https')).toBe('wss://websocket.self.xyz');
+    expect(getWSDbRelayerUrl('staging_celo')).toBe(
+      'wss://websocket.staging.self.xyz',
+    );
+  });
+
+  it('getPayload handles various inputs', () => {
+    // Test with null input - should work since JSON.stringify handles it
+    const payload1 = getPayload(
+      null,
+      'disclose',
+      'vc_and_disclose',
+      'https',
+      'https://example.com',
+    );
+    expect(payload1.circuit.inputs).toBe('null');
+
+    // Test with empty string circuit type - should work since it's just used as-is
+    const payload2 = getPayload(
+      {},
+      'disclose',
+      'vc_and_disclose',
+      'https',
+      'https://example.com',
+    );
+    expect(payload2.circuit.inputs).toBe('{}');
+
+    // Test with empty circuit name - should work since it's just used as-is
+    const payload3 = getPayload(
+      {},
+      'disclose',
+      '',
+      'https',
+      'https://example.com',
+    );
+    expect(payload3.circuit.name).toBe('');
+  });
+
+  it('getPayload handles special characters in inputs', () => {
+    const inputs = { message: 'Hello "World" & <script>alert("xss")</script>' };
+    const payload = getPayload(
+      inputs,
+      'disclose',
+      'vc_and_disclose',
+      'https',
+      'https://example.com',
+    );
+
+    // JSON.stringify will escape quotes, so we should expect the escaped version
+    expect(payload.circuit.inputs).toContain('Hello \\"World\\"');
+    expect(JSON.parse(payload.circuit.inputs)).toEqual(inputs);
+  });
+
+  it('encryptAES256GCM handles empty plaintext', () => {
+    const key = forge.random.getBytesSync(32);
+    const plaintext = '';
+    const encrypted = encryptAES256GCM(plaintext, forge.util.createBuffer(key));
+
+    expect(encrypted.cipher_text).toBeDefined();
+    expect(encrypted.nonce).toBeDefined();
+    expect(encrypted.auth_tag).toBeDefined();
+  });
+
+  it('encryptAES256GCM handles large plaintext', () => {
+    const key = forge.random.getBytesSync(32);
+    const plaintext = 'a'.repeat(10000);
+    const encrypted = encryptAES256GCM(plaintext, forge.util.createBuffer(key));
+
+    expect(encrypted.cipher_text.length).toBeGreaterThan(0);
+  });
+});