From ebdc639c882a6c076fc27b1e4f30b798d22883cc Mon Sep 17 00:00:00 2001 From: Evi Nova <66773372+Tranquil-Flow@users.noreply.github.com> Date: Tue, 3 Feb 2026 08:00:19 +1000 Subject: [PATCH] fix: resolve issue from losing connection mid proof verifcation (#1627) * fix: resolve issue from losing connection mid proof verifcation * fix: properly clean up socket resource to prevent memory leak * fix: proper memory leak fix + have both proactive/reactive reconnection * fix: reset scroll state on QR scans and improve WebSocket reconnection handling * chore: yarn fmt --- app/src/screens/verification/ProveScreen.tsx | 14 ++ .../src/proving/provingMachine.ts | 145 +++++++++++++++++- .../src/stores/selfAppStore.tsx | 33 ++-- 3 files changed, 166 insertions(+), 26 deletions(-) diff --git a/app/src/screens/verification/ProveScreen.tsx b/app/src/screens/verification/ProveScreen.tsx index 641d311d2..8f0183b4e 100644 --- a/app/src/screens/verification/ProveScreen.tsx +++ b/app/src/screens/verification/ProveScreen.tsx @@ -172,6 +172,20 @@ const ProveScreen: React.FC = () => { if (selectedAppRef.current?.sessionId !== selectedApp.sessionId) { hasInitializedScrollStateRef.current = false; setHasScrolledToBottom(false); + + // After state reset, check if content is short using current measurements. + // Use setTimeout(0) to ensure we read values AFTER React processes the reset, + // without adding measurements to dependencies (which causes race conditions). + setTimeout(() => { + const hasMeasurements = + scrollViewContentHeight > 0 && scrollViewHeight > 0; + const isShort = scrollViewContentHeight <= scrollViewHeight + 50; + + if (hasMeasurements && isShort) { + setHasScrolledToBottom(true); + hasInitializedScrollStateRef.current = true; + } + }, 0); } setDefaultDocumentTypeIfNeeded(); diff --git a/packages/mobile-sdk-alpha/src/proving/provingMachine.ts b/packages/mobile-sdk-alpha/src/proving/provingMachine.ts index 3b6b1b647..bcf7ec396 100644 --- a/packages/mobile-sdk-alpha/src/proving/provingMachine.ts +++ b/packages/mobile-sdk-alpha/src/proving/provingMachine.ts @@ -211,6 +211,7 @@ export interface ProvingState { sharedKey: Buffer | null; wsConnection: WebSocket | null; wsHandlers: WsHandlers | null; + wsReconnectAttempts: number; socketConnection: Socket | null; uuid: string | null; userConfirmed: boolean; @@ -251,6 +252,7 @@ export interface ProvingState { _handleWsOpen: (selfClient: SelfClient) => void; _handleWsError: (error: Event, selfClient: SelfClient) => void; _handleWsClose: (event: CloseEvent, selfClient: SelfClient) => void; + _reconnectTeeWebSocket: (selfClient: SelfClient) => Promise; _handlePassportNotSupported: (selfClient: SelfClient) => void; _handleAccountRecoveryChoice: (selfClient: SelfClient) => void; @@ -498,6 +500,7 @@ export const useProvingStore = create((set, get) => { sharedKey: null, wsConnection: null, wsHandlers: null, + wsReconnectAttempts: 0, socketConnection: null, uuid: null, userConfirmed: false, @@ -823,6 +826,8 @@ export const useProvingStore = create((set, get) => { reason: event.reason, }); const currentState = get().currentState; + + // Handle unexpected close during active proving states if ( currentState === 'init_tee_connexion' || currentState === 'proving' || @@ -836,11 +841,105 @@ export const useProvingStore = create((set, get) => { selfClient, ); } + + // In ready_to_prove state, attempt automatic reconnection to handle network interruptions. + // Users may lose connectivity briefly; reconnecting transparently improves UX. + if (currentState === 'ready_to_prove') { + const MAX_RECONNECT_ATTEMPTS = 3; + const attempts = get().wsReconnectAttempts; + + if (attempts < MAX_RECONNECT_ATTEMPTS) { + selfClient.logProofEvent('info', 'TEE WebSocket reconnection attempt', context, { + attempt: attempts + 1, + max_attempts: MAX_RECONNECT_ATTEMPTS, + }); + set({ wsConnection: null, wsReconnectAttempts: attempts + 1 }); + + const backoffMs = Math.min(1000 * Math.pow(2, attempts), 10000); + setTimeout(() => { + if (get().currentState === 'ready_to_prove') { + get()._reconnectTeeWebSocket(selfClient); + } + }, backoffMs); + return; + } + + selfClient.logProofEvent('error', 'TEE WebSocket reconnection exhausted', context, { + failure: 'PROOF_FAILED_CONNECTION', + attempts: MAX_RECONNECT_ATTEMPTS, + }); + get()._handleWebSocketMessage( + new MessageEvent('error', { + data: JSON.stringify({ error: 'WebSocket reconnection failed' }), + }), + selfClient, + ); + } + if (get().wsConnection) { set({ wsConnection: null }); } }, + /** + * Re-establishes the TEE WebSocket connection using stored circuit parameters. + * Called automatically when connection is lost in ready_to_prove state. + */ + _reconnectTeeWebSocket: async (selfClient: SelfClient): Promise => { + const context = createProofContext(selfClient, '_reconnectTeeWebSocket'); + const { passportData, circuitType } = get(); + + if (!passportData || !circuitType) { + selfClient.logProofEvent('error', 'Reconnect failed: missing prerequisites', context); + return false; + } + + const typedCircuitType = circuitType as 'disclose' | 'register' | 'dsc'; + const circuitName = + typedCircuitType === 'disclose' + ? passportData.documentCategory === 'aadhaar' + ? 'disclose_aadhaar' + : 'disclose' + : getCircuitNameFromPassportData(passportData, typedCircuitType as 'register' | 'dsc'); + + const wsRpcUrl = resolveWebSocketUrl(selfClient, typedCircuitType, passportData as PassportData, circuitName); + if (!wsRpcUrl) { + selfClient.logProofEvent('error', 'Reconnect failed: no WebSocket URL', context); + return false; + } + + selfClient.logProofEvent('info', 'TEE WebSocket reconnection started', context); + + return new Promise(resolve => { + const ws = new WebSocket(wsRpcUrl); + const RECONNECT_TIMEOUT_MS = 15000; + + const wsHandlers: WsHandlers = { + message: (event: MessageEvent) => get()._handleWebSocketMessage(event, selfClient), + open: () => { + selfClient.logProofEvent('info', 'TEE WebSocket reconnected', context); + set({ wsReconnectAttempts: 0 }); + resolve(true); + }, + error: (error: Event) => get()._handleWsError(error, selfClient), + close: (event: CloseEvent) => get()._handleWsClose(event, selfClient), + }; + + set({ wsConnection: ws, wsHandlers }); + ws.addEventListener('message', wsHandlers.message); + ws.addEventListener('open', wsHandlers.open); + ws.addEventListener('error', wsHandlers.error); + ws.addEventListener('close', wsHandlers.close); + + setTimeout(() => { + if (ws.readyState !== WebSocket.OPEN) { + selfClient.logProofEvent('warn', 'TEE WebSocket reconnection timeout', context); + resolve(false); + } + }, RECONNECT_TIMEOUT_MS); + }); + }, + init: async ( selfClient: SelfClient, circuitType: 'dsc' | 'disclose' | 'register', @@ -1293,7 +1392,7 @@ export const useProvingStore = create((set, get) => { close: (event: CloseEvent) => get()._handleWsClose(event, selfClient), }; - set({ wsConnection: ws, wsHandlers }); + set({ wsConnection: ws, wsHandlers, wsReconnectAttempts: 0 }); ws.addEventListener('message', wsHandlers.message); ws.addEventListener('open', wsHandlers.open); @@ -1318,7 +1417,8 @@ export const useProvingStore = create((set, get) => { startProving: async (selfClient: SelfClient) => { _checkActorInitialized(actor); const startTime = Date.now(); - const { wsConnection, sharedKey, passportData, secret, uuid } = get(); + let { wsConnection } = get(); + const { sharedKey, passportData, secret, uuid } = get(); const context = createProofContext(selfClient, 'startProving', { sessionId: uuid || get().uuid || 'unknown-session', }); @@ -1330,17 +1430,45 @@ export const useProvingStore = create((set, get) => { console.error('Cannot start proving: Not in ready_to_prove state.'); return; } - if (!wsConnection || !sharedKey || !passportData || !secret || !uuid) { + + // Check non-connection prerequisites first + if (!sharedKey || !passportData || !secret || !uuid) { selfClient.logProofEvent('error', 'Missing proving prerequisites', context, { failure: 'PROOF_FAILED_CONNECTION', }); - console.error('Cannot start proving: Missing wsConnection, sharedKey, passportData, secret, or uuid.'); + console.error('Cannot start proving: Missing sharedKey, passportData, secret, or uuid.'); actor!.send({ type: 'PROVE_ERROR' }); return; } + // Attempt reconnection if WebSocket is missing or not open + if (!wsConnection || wsConnection.readyState !== WebSocket.OPEN) { + selfClient.logProofEvent('warn', 'WebSocket not ready, attempting reconnection', context, { + wsConnectionExists: !!wsConnection, + readyState: wsConnection?.readyState, + }); + + const reconnected = await get()._reconnectTeeWebSocket(selfClient); + if (!reconnected) { + selfClient.logProofEvent('error', 'WebSocket reconnection failed', context, { + failure: 'PROOF_FAILED_CONNECTION', + }); + actor!.send({ type: 'PROVE_ERROR' }); + return; + } + + // Get the new connection after reconnection + wsConnection = get().wsConnection; + if (!wsConnection || wsConnection.readyState !== WebSocket.OPEN) { + selfClient.logProofEvent('error', 'Reconnected WebSocket not ready', context, { + failure: 'PROOF_FAILED_CONNECTION', + }); + actor!.send({ type: 'PROVE_ERROR' }); + return; + } + } + try { - // Emit event for FCM token registration selfClient.emit(SdkEvents.PROVING_BEGIN_GENERATION, { uuid, isMock: passportData?.mock ?? false, @@ -1350,7 +1478,12 @@ export const useProvingStore = create((set, get) => { selfClient.trackEvent(ProofEvents.PAYLOAD_GEN_STARTED); selfClient.logProofEvent('info', 'Payload generation started', context); const submitBody = await get()._generatePayload(selfClient); - wsConnection.send(JSON.stringify(submitBody)); + + const activeWsConnection = get().wsConnection; + if (!activeWsConnection) { + throw new Error('WebSocket connection lost during payload generation'); + } + activeWsConnection.send(JSON.stringify(submitBody)); selfClient.logProofEvent('info', 'Payload sent over WebSocket', context); selfClient.trackEvent(ProofEvents.PAYLOAD_SENT); selfClient.trackEvent(ProofEvents.PROVING_PROCESS_STARTED); diff --git a/packages/mobile-sdk-alpha/src/stores/selfAppStore.tsx b/packages/mobile-sdk-alpha/src/stores/selfAppStore.tsx index 5af1f350e..613c53f67 100644 --- a/packages/mobile-sdk-alpha/src/stores/selfAppStore.tsx +++ b/packages/mobile-sdk-alpha/src/stores/selfAppStore.tsx @@ -13,8 +13,6 @@ import { WS_DB_RELAYER } from '@selfxyz/common'; * Zustand state backing the in-app handoff between the SDK and the hosted Self * application. The store tracks the active websocket session, latest * {@link SelfApp} payload, and helper callbacks used by the proving machine. - * Consumers should treat the state as ephemeral and expect it to reset whenever - * the socket disconnects. */ export interface SelfAppState { selfApp: SelfApp | null; @@ -80,24 +78,17 @@ export const useSelfAppStore = create((set, get) => ({ socket.on('connect', () => {}); - // Listen for the event only once per connection attempt socket.once('self_app', (data: unknown) => { try { const appData: SelfApp = typeof data === 'string' ? JSON.parse(data) : (data as SelfApp); - // Basic validation if (!appData || typeof appData !== 'object' || !appData.sessionId) { - console.error('[SelfAppStore] Invalid app data received:', appData); - // Optionally clear the app data or handle the error appropriately + console.error('[SelfAppStore] Invalid app data received'); set({ selfApp: null }); return; } if (appData.sessionId !== get().sessionId) { - console.warn( - `[SelfAppStore] Received SelfApp for session ${ - appData.sessionId - }, but current session is ${get().sessionId}. Ignoring.`, - ); + console.warn('[SelfAppStore] Session mismatch, ignoring payload'); return; } @@ -109,20 +100,22 @@ export const useSelfAppStore = create((set, get) => ({ }); socket.on('connect_error', error => { - console.error('[SelfAppStore] Mobile WS connection error:', error); - // Clean up on connection error - get().cleanSelfApp(); + // Socket.io handles reconnection automatically with exponential backoff. + // State is preserved to allow seamless recovery when network returns. + console.error('[SelfAppStore] Connection error:', error.message); }); socket.on('error', error => { - console.error('[SelfAppStore] Mobile WS error:', error); - // Consider if cleanup is needed here as well + console.error('[SelfAppStore] Socket error:', error); }); - socket.on('disconnect', (_reason: string) => { - // Prevent cleaning up if disconnect was initiated by cleanSelfApp - if (get().socket === socket) { - set({ socket: null, sessionId: null, selfApp: null }); + socket.on('disconnect', (reason: string) => { + if (get().socket !== socket) return; + + // Only clear state on intentional disconnects. For transient network issues + // (transport close, ping timeout), socket.io reconnects automatically. + if (reason === 'io server disconnect' || reason === 'io client disconnect') { + set({ socket: null, sessionId: null }); } }); } catch (error) {