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
This commit is contained in:
Evi Nova
2026-02-03 08:00:19 +10:00
committed by GitHub
parent a7b790d41c
commit ebdc639c88
3 changed files with 166 additions and 26 deletions

View File

@@ -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();

View File

@@ -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<boolean>;
_handlePassportNotSupported: (selfClient: SelfClient) => void;
_handleAccountRecoveryChoice: (selfClient: SelfClient) => void;
@@ -498,6 +500,7 @@ export const useProvingStore = create<ProvingState>((set, get) => {
sharedKey: null,
wsConnection: null,
wsHandlers: null,
wsReconnectAttempts: 0,
socketConnection: null,
uuid: null,
userConfirmed: false,
@@ -823,6 +826,8 @@ export const useProvingStore = create<ProvingState>((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<ProvingState>((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<boolean> => {
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<ProvingState>((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<ProvingState>((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<ProvingState>((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<ProvingState>((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);

View File

@@ -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<SelfAppState>((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<SelfAppState>((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) {