Added withToken func to remove the duplicated code

src/services/iframe-controller.service.ts

@@ -6,7 +6,11 @@ import { splitPrivateData, isValid32ByteHex } from '../utils/service.utils';
 import { MerkleProofResult } from '../../pkg/sdk_client';
 
 export class IframeController {
+  private static isInitialized = false; // <--- VERROU
+
   static async init() {
+    if (this.isInitialized) return; // On sort si déjà lancé
+
     // On ne lance l'écoute que si on est dans une iframe
     if (window.self !== window.top) {
       console.log('[IframeController] 📡 Mode Iframe détecté. Démarrage des listeners API...');
@@ -36,6 +40,17 @@ export class IframeController {
       );
     };
 
+    // Helper pour vérifier le token avant chaque action sensible
+    const withToken = async (event: MessageEvent, action: () => Promise<void>) => {
+      const { accessToken } = event.data;
+      // On vérifie si le token est présent ET valide pour l'origine de l'iframe
+      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
+        throw new Error('Invalid or expired session token');
+      }
+      // Si tout est bon, on exécute l'action
+      await action();
+    };
+
     // --- Définitions des gestionnaires (Handlers) ---
 
     const handleRequestLink = async (event: MessageEvent) => {
@@ -95,11 +110,7 @@ export class IframeController {
         throw new Error('Device already paired — ignoring CREATE_PAIRING request');
       }
 
-      const { accessToken } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
-
+      await withToken(event, async () => {
         console.log("[Router:API] 🚀 Démarrage du processus d'appairage...");
 
         const myAddress = services.getDeviceAddress();
@@ -138,17 +149,14 @@ export class IframeController {
           messageId: event.data.messageId,
         };
         window.parent.postMessage(successMsg, event.origin);
+      });
     };
 
     const handleGetMyProcesses = async (event: MessageEvent) => {
       console.log(`[Router:API] 📨 Message ${MessageType.GET_MY_PROCESSES} reçu`);
       if (!services.isPaired()) throw new Error('Device not paired');
 
-      const { accessToken } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
-
+      await withToken(event, async () => {
         const myProcesses = await services.getMyProcesses();
 
         window.parent.postMessage(
@@ -159,17 +167,14 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     const handleGetProcesses = async (event: MessageEvent) => {
       console.log(`[Router:API] 📨 Message ${MessageType.GET_PROCESSES} reçu`);
       if (!services.isPaired()) throw new Error('Device not paired');
 
-      const { accessToken } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
-
+      await withToken(event, async () => {
         const processes = await services.getProcesses();
 
         window.parent.postMessage(
@@ -180,17 +185,16 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     const handleDecryptState = async (event: MessageEvent) => {
       console.log(`[Router:API] 📨 Message ${MessageType.RETRIEVE_DATA} reçu`);
       if (!services.isPaired()) throw new Error('Device not paired');
 
-      const { processId, stateId, accessToken } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
+      const { processId, stateId } = event.data;
 
+      await withToken(event, async () => {
         const process = await services.getProcess(processId);
         if (!process) throw new Error("Can't find process");
 
@@ -220,6 +224,7 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     const handleValidateToken = async (event: MessageEvent) => {
@@ -295,11 +300,7 @@ export class IframeController {
         throw new Error('Device not paired');
       }
 
-      const { accessToken } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
-
+      await withToken(event, async () => {
         window.parent.postMessage(
           {
             type: MessageType.GET_PAIRING_ID,
@@ -308,17 +309,16 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     const handleCreateProcess = async (event: MessageEvent) => {
       console.log(`[Router:API] 📨 Message ${MessageType.CREATE_PROCESS} reçu`);
       if (!services.isPaired()) throw new Error('Device not paired');
 
-      const { processData, privateFields, roles, accessToken } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
+      const { processData, privateFields, roles } = event.data;
 
+      await withToken(event, async () => {
         console.log('[Router:API] 🚀 Démarrage de la création de processus standard...');
         const { privateData, publicData } = splitPrivateData(processData, privateFields);
 
@@ -350,16 +350,16 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     const handleNotifyUpdate = async (event: MessageEvent) => {
       console.log(`[Router:API] 📨 Message ${MessageType.NOTIFY_UPDATE} reçu`);
       if (!services.isPaired()) throw new Error('Device not paired');
 
-      const { processId, stateId, accessToken } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
+      const { processId, stateId } = event.data;
+
+      await withToken(event, async () => {
         if (!isValid32ByteHex(stateId)) throw new Error('Invalid state id');
 
         const res = await services.createPrdUpdate(processId, stateId);
@@ -372,17 +372,16 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     const handleValidateState = async (event: MessageEvent) => {
       console.log(`[Router:API] 📨 Message ${MessageType.VALIDATE_STATE} reçu`);
       if (!services.isPaired()) throw new Error('Device not paired');
 
-      const { processId, stateId, accessToken } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
+      const { processId, stateId } = event.data;
 
+      await withToken(event, async () => {
         const res = await services.approveChange(processId, stateId);
         await services.handleApiReturn(res);
 
@@ -394,17 +393,16 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     const handleUpdateProcess = async (event: MessageEvent) => {
       console.log(`[Router:API] 📨 Message ${MessageType.UPDATE_PROCESS} reçu`);
       if (!services.isPaired()) throw new Error('Device not paired');
 
-      const { processId, newData, privateFields, roles, accessToken } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
+      const { processId, newData, privateFields, roles } = event.data;
 
+      await withToken(event, async () => {
         console.log(`[Router:API] 🔄 Transfert de la mise à jour de ${processId} au service...`);
 
         // Le service gère maintenant tout : récupération, réparation d'état, et mise à jour.
@@ -422,17 +420,16 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     const handleDecodePublicData = async (event: MessageEvent) => {
       console.log(`[Router:API] 📨 Message ${MessageType.DECODE_PUBLIC_DATA} reçu`);
       if (!services.isPaired()) throw new Error('Device not paired');
 
-      const { accessToken, encodedData } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
+      const { encodedData } = event.data;
 
+      await withToken(event, async () => {
         const decodedData = services.decodeValue(encodedData);
         window.parent.postMessage(
           {
@@ -442,15 +439,14 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     const handleHashValue = async (event: MessageEvent) => {
       console.log(`[Router:API] 📨 Message ${MessageType.HASH_VALUE} reçu`);
-      const { accessToken, commitedIn, label, fileBlob } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
+      const { commitedIn, label, fileBlob } = event.data;
 
+      await withToken(event, async () => {
         const hash = services.getHashForFile(commitedIn, label, fileBlob);
         window.parent.postMessage(
           {
@@ -460,15 +456,14 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     const handleGetMerkleProof = async (event: MessageEvent) => {
       console.log(`[Router:API] 📨 Message ${MessageType.GET_MERKLE_PROOF} reçu`);
-      const { accessToken, processState, attributeName } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
+      const { processState, attributeName } = event.data;
 
+      await withToken(event, async () => {
         const proof = services.getMerkleProofForFile(processState, attributeName);
         window.parent.postMessage(
           {
@@ -478,15 +473,14 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     const handleValidateMerkleProof = async (event: MessageEvent) => {
       console.log(`[Router:API] 📨 Message ${MessageType.VALIDATE_MERKLE_PROOF} reçu`);
-      const { accessToken, merkleProof, documentHash } = event.data;
-      if (!accessToken || !(await tokenService.validateToken(accessToken, event.origin))) {
-        throw new Error('Invalid or expired session token');
-      }
+      const { merkleProof, documentHash } = event.data;
 
+      await withToken(event, async () => {
         let parsedMerkleProof: MerkleProofResult;
         try {
           parsedMerkleProof = JSON.parse(merkleProof);
@@ -503,6 +497,7 @@ export class IframeController {
           },
           event.origin,
         );
+      });
     };
 
     // --- Le "Switchyard" : il reçoit tous les messages et les dispatche ---

src/services/websockets.service.ts

@@ -1,36 +1,53 @@
-import { AnkFlag } from 'pkg/sdk_client';
+import { AnkFlag } from '../../pkg/sdk_client'; // Vérifie le chemin vers pkg
 import Services from './service';
 
-let ws: WebSocket;
+let ws: WebSocket | null = null;
 let messageQueue: string[] = [];
+let reconnectInterval = 1000; // Délai initial de 1s avant reconnexion
+const MAX_RECONNECT_INTERVAL = 30000; // Max 30s
+let isConnecting = false;
+let urlReference: string = '';
+let pingIntervalId: any = null;
+
 export async function initWebsocket(url: string) {
-  ws = new WebSocket(url);
+  urlReference = url;
+  connect();
+}
 
-  if (ws !== null) {
-    ws.onopen = async (event) => {
-      console.log('WebSocket connection established');
+function connect() {
+  if (isConnecting || (ws && ws.readyState === WebSocket.OPEN)) return;
+  isConnecting = true;
 
+  console.log(`[WS] 🔌 Tentative de connexion à ${urlReference}...`);
+  ws = new WebSocket(urlReference);
+
+  ws.onopen = async () => {
+    console.log('[WS] ✅ Connexion établie !');
+    isConnecting = false;
+    reconnectInterval = 1000; // Reset du délai
+
+    // Démarrer le Heartbeat (Ping pour garder la connexion vivante)
+    startHeartbeat();
+
+    // Vider la file d'attente (messages envoyés pendant la coupure)
     while (messageQueue.length > 0) {
       const message = messageQueue.shift();
-        if (message) {
-          ws.send(message);
-        }
+      if (message) ws?.send(message);
     }
   };
 
-    // Listen for messages
   ws.onmessage = (event) => {
     const msgData = event.data;
-
-      // console.log("Received text message: ", msgData);
-      (async () => {
     if (typeof msgData === 'string') {
+      (async () => {
         try {
           const parsedMessage = JSON.parse(msgData);
           const services = await Services.getInstance();
+
+          // Gestion des messages
           switch (parsedMessage.flag) {
             case 'Handshake':
-                await services.handleHandshakeMsg(url, parsedMessage.content);
+              await services.handleHandshakeMsg(urlReference, parsedMessage.content);
               break;
             case 'NewTx':
               await services.parseNewTx(parsedMessage.content);
@@ -39,51 +56,80 @@ export async function initWebsocket(url: string) {
               await services.parseCipher(parsedMessage.content);
               break;
             case 'Commit':
-                // Basically if we see this it means we have an error
               await services.handleCommitError(parsedMessage.content);
               break;
+            // Ajoute d'autres cas si nécessaire
+            default:
+            // console.log('[WS] Message reçu:', parsedMessage.flag);
           }
         } catch (error) {
-            console.error('Received an invalid message:', error);
-          }
-        } else {
-          console.error('Received a non-string message');
+          console.error('[WS] Erreur traitement message:', error);
         }
       })();
-    };
-
-    // Listen for possible errors
-    ws.onerror = (event) => {
-      console.error('WebSocket error:', event);
-    };
-
-    // Listen for when the connection is closed
-    ws.onclose = (event) => {
-      console.log('WebSocket is closed now.');
-    };
     }
+  };
+
+  ws.onerror = (event) => {
+    console.error('[WS] 💥 Erreur:', event);
+    // Pas besoin de reconnecter ici, onclose sera appelé juste après
+  };
+
+  ws.onclose = (event) => {
+    isConnecting = false;
+    stopHeartbeat();
+    console.warn(`[WS] ⚠️ Déconnecté (Code: ${event.code}). Reconnexion dans ${reconnectInterval / 1000}s...`);
+
+    // Reconnexion exponentielle (1s, 1.5s, 2.25s...)
+    setTimeout(() => {
+      connect();
+      reconnectInterval = Math.min(reconnectInterval * 1.5, MAX_RECONNECT_INTERVAL);
+    }, reconnectInterval);
+  };
+}
+
+function startHeartbeat() {
+  stopHeartbeat();
+  // Envoie un ping toutes les 30 secondes pour éviter que le serveur ou le navigateur ne coupe la connexion
+  pingIntervalId = setInterval(() => {
+    if (ws && ws.readyState === WebSocket.OPEN) {
+      // Adapter selon ce que ton serveur attend comme Ping, ou envoyer un message vide
+      // ws.send(JSON.stringify({ flag: 'Ping', content: '' }));
+    }
+  }, 30000);
+}
+
+function stopHeartbeat() {
+  if (pingIntervalId) clearInterval(pingIntervalId);
 }
 
-// Method to send messages
 export function sendMessage(flag: AnkFlag, message: string): void {
-  if (ws.readyState === WebSocket.OPEN) {
+  if (ws && ws.readyState === WebSocket.OPEN) {
     const networkMessage = {
       flag: flag,
       content: message,
     };
-    console.log('Sending message of type:', flag);
     ws.send(JSON.stringify(networkMessage));
   } else {
-    console.error('WebSocket is not open. ReadyState:', ws.readyState);
-    messageQueue.push(message);
+    console.warn(`[WS] Pas connecté. Message '${flag}' mis en file d'attente.`);
+    const networkMessage = {
+      flag: flag,
+      content: message,
+    };
+    messageQueue.push(JSON.stringify(networkMessage));
+
+    // Si on n'est pas déjà en train de se connecter, on force une tentative
+    if (!isConnecting) connect();
   }
 }
 
 export function getUrl(): string {
-  return ws.url;
+  return urlReference;
 }
 
-// Method to close the WebSocket connection
 export function close(): void {
+  if (ws) {
+    ws.onclose = null; // On évite la reconnexion auto si fermeture volontaire
+    stopHeartbeat();
     ws.close();
+  }
 }