Merge pull request #41 from dallen4/msg-mutex-resilience

Optimization: Message Mutex Lock & Dynamic Replay

.github/workflows/web_ci_workflow.yml

@@ -1,7 +1,8 @@
 name: Playwright Tests
+
 on:
-  push:
-    branches: [main, alpha]
+  deployment_status:
+
 jobs:
   test:
     timeout-minutes: 20
@@ -19,12 +20,11 @@ jobs:
         shell: bash
         env:
           STAGE: ${{ github.ref_name }}
-        run: |
-          if [[ $STAGE == "alpha" ]]; then BASE_URI="https://alpha.deadrop.io"; else BASE_URI="https://deadrop.io"; fi
-          CI=true TEST_URI=$BASE_URI yarn web playwright test --trace on
+          BASE_URL: ${{ github.event.deployment_status.target_url }}
+        run: CI=true TEST_URI=$BASE_URL yarn web playwright test --trace on
       - uses: actions/upload-artifact@v3
         if: always()
         with:
           name: playwright-report
-          path: test-results/
+          path: web/test-results/
           retention-days: 30

shared/lib/constants.ts

@@ -54,3 +54,13 @@ export enum MessageType {
     Verify = 'verify',
     ConfirmVerification = 'confirm',
 }
+
+export const DropMessageOrderMap = new Map([
+    [MessageType.Handshake, MessageType.Handshake],
+    [MessageType.Payload, MessageType.Verify],
+]);
+
+export const GrabMessageOrderMap = new Map([
+    [MessageType.Handshake, MessageType.Payload],
+    [MessageType.Verify, MessageType.ConfirmVerification],
+]);

shared/lib/peer.ts

@@ -10,12 +10,14 @@ const onUnload = (e: BeforeUnloadEvent) => {
     return 'Are you sure you want to leave?';
 };
 
-const removeOnUnloadListener = () => {
-    window.onbeforeunload = null;
-    window.removeEventListener('beforeunload', onUnload);
+export const removeOnUnloadListener = () => {
+    if (!isServer) {
+        window.onbeforeunload = null;
+        window.removeEventListener('beforeunload', onUnload);
+    }
 };
 
-function createPeer(id: string, url: string) {
+export function createPeer(id: string, url: string) {
     const server = new URL(url);
 
     const peer = new Peer(id, {
@@ -39,13 +41,9 @@ function createPeer(id: string, url: string) {
         }
     });
 
-    peer.on('disconnected', () => {
-        if (!isServer) removeOnUnloadListener();
-    });
+    peer.on('disconnected', removeOnUnloadListener);
 
-    peer.on('close', () => {
-        if (!isServer) removeOnUnloadListener();
-    });
+    peer.on('close', removeOnUnloadListener);
 
     return new Promise<Peer>((resolve) => {
         peer.on('open', (id: string) => {
@@ -55,5 +53,3 @@ function createPeer(id: string, url: string) {
         });
     });
 }
-
-export { createPeer };

shared/types/messages.ts

@@ -27,3 +27,5 @@ export interface ConfirmIntegrityMessage extends BaseMessage {
     type: MessageType.ConfirmVerification;
     verified: boolean;
 }
+
+export type MessageHandler = (msg: BaseMessage) => Promise<void>;

web/api/drops.ts

@@ -1,5 +1,5 @@
 import { generateIV } from '@shared/lib/util';
-import { getRedis } from 'lib/redis';
+import { getRedis } from 'api/redis';
 import { formatDropKey } from 'lib/util';
 import { nanoid } from 'nanoid';
 

web/api/limiter.ts

@@ -1,5 +1,5 @@
 import { hashRaw } from '@shared/lib/crypto/operations';
-import { getRedis } from '../lib/redis';
+import { getRedis } from './redis';
 
 const DAY_IN_SEC = 60 * 60 * 24;
 

web/api/middleware/cors.ts

@@ -3,10 +3,10 @@ import Cors from 'cors';
 export const cors = Cors({
     methods: ['POST', 'GET', 'DELETE'],
     origin: (origin, callback) => {
-        console.log(origin);
         if (
             !origin ||
             origin.endsWith('deadrop.io') ||
+            origin.endsWith('dallen4.vercel.app') ||
             origin.includes('vscode-webview:')
         )
             callback(null, true);
@@ -15,6 +15,9 @@ export const cors = Cors({
             origin.startsWith('http://localhost:')
         )
             callback(null, true);
-        else callback(new Error('Invalid origin'));
+        else {
+            console.log(`Invalid origin: ${origin}`);
+            callback(new Error('Invalid origin'));
+        }
     },
 });

web/hooks/use-drop.tsx

@@ -16,7 +16,12 @@ import type { DataConnection } from 'peerjs';
 import { useRef } from 'react';
 import { useMachine } from '@xstate/react/lib/useMachine';
 import { dropMachine, initDropContext } from '@shared/lib/machines/drop';
-import { DropEventType, DropState, MessageType } from '@shared/lib/constants';
+import {
+    DropEventType,
+    DropMessageOrderMap,
+    DropState,
+    MessageType,
+} from '@shared/lib/constants';
 import { generateGrabUrl } from 'lib/util';
 import { deleteReq, post } from 'lib/fetch';
 import { DROP_API_PATH } from 'config/paths';
@@ -32,16 +37,54 @@ import {
 import { encryptFile, hashFile } from 'lib/crypto';
 import { showNotification } from '@mantine/notifications';
 import { IconX } from '@tabler/icons';
+import { withMessageLock } from 'lib/messages';
 
 export const useDrop = () => {
     const logsRef = useRef<Array<string>>([]);
     const contextRef = useRef<DropContext>(initDropContext());
+    const timersRef = useRef(new Map<MessageType, NodeJS.Timeout>());
 
     const [{ value: state }, send] = useMachine(dropMachine);
 
     const pushLog = (message: string) => logsRef.current.push(message);
 
+    const clearTimer = (msgType: MessageType) => {
+        const timerId = timersRef.current.get(msgType);
+
+        if (timerId) {
+            clearTimeout(timerId);
+            timersRef.current.delete(msgType);
+        }
+    };
+
+    const sendMessage = async (msg: BaseMessage, retryCount: number = 0) => {
+        if (!contextRef.current.connection) return;
+
+        const expectedType = DropMessageOrderMap.get(msg.type)!;
+
+        clearTimer(expectedType);
+
+        if (retryCount >= 3) {
+            showNotification({
+                message:
+                    'Connection may be unstable, please try your drop again',
+                color: 'red',
+                icon: <IconX />,
+                autoClose: 4500,
+            });
+            console.error(`Attempt limit exceeded for type: ${msg.type}`);
+            return;
+        }
+
+        contextRef.current.connection.send(msg);
+
+        const timer = setTimeout(() => sendMessage(msg, retryCount + 1), 1000);
+        timersRef.current.set(expectedType, timer);
+    };
+
     const onMessage = async (msg: BaseMessage) => {
+        clearTimer(msg.type);
+
         if (msg.type === MessageType.Handshake) {
             const { input } = msg as HandshakeMessage;
 
@@ -77,7 +120,7 @@ export const useDrop = () => {
                 verified,
             };
 
-            contextRef.current.connection!.send(message);
+            sendMessage(message);
 
             pushLog('Integrity confirmation sent, completing drop...');
 
@@ -102,7 +145,8 @@ export const useDrop = () => {
 
         contextRef.current.connection = connection;
 
-        connection.on('data', onMessage);
+        const handlerWithLock = withMessageLock(onMessage, pushLog);
+        connection.on('data', handlerWithLock);
 
         send({ type: DropEventType.Connect, connection });
 
@@ -240,14 +284,16 @@ export const useDrop = () => {
                 : undefined,
         };
 
-        contextRef.current.connection!.send(message);
+        sendMessage(message);
 
         pushLog('Payload dropped, awaiting response...');
 
         send({ type: DropEventType.Drop });
     };
 
-    const cleanup = () => {
+    const cleanup = async () => {
+        const { removeOnUnloadListener } = await import('shared/lib/peer');
+
         contextRef.current.connection!.close();
         contextRef.current.peer!.disconnect();
         contextRef.current.peer!.destroy();
@@ -261,6 +307,8 @@ export const useDrop = () => {
             ),
         );
 
+        removeOnUnloadListener();
+
         contextRef.current = initDropContext();
     };
 

web/hooks/use-grab.ts → web/hooks/use-grab.tsx

@@ -5,7 +5,12 @@ import type {
     GrabContext,
     InitGrabEvent,
 } from '@shared/types/grab';
-import { GrabEventType, GrabState, MessageType } from '@shared/lib/constants';
+import {
+    GrabEventType,
+    GrabMessageOrderMap,
+    GrabState,
+    MessageType,
+} from '@shared/lib/constants';
 import { useRef } from 'react';
 import { get } from 'lib/fetch';
 import { useRouter } from 'next/router';
@@ -28,18 +33,58 @@ import {
     importKey,
 } from '@shared/lib/crypto/operations';
 import { decryptFile, hashFile } from 'lib/crypto';
+import { withMessageLock } from 'lib/messages';
+import { showNotification } from '@mantine/notifications';
+import { IconX } from '@tabler/icons';
 
 export const useGrab = () => {
     const router = useRouter();
 
     const logsRef = useRef<Array<string>>([]);
     const contextRef = useRef<GrabContext>(initGrabContext());
+    const timersRef = useRef(new Map<MessageType, NodeJS.Timeout>());
 
     const [{ value: state }, send] = useMachine(grabMachine);
 
     const pushLog = (message: string) => logsRef.current.push(message);
 
+    const clearTimer = (msgType: MessageType) => {
+        const timerId = timersRef.current.get(msgType);
+
+        if (timerId) {
+            clearTimeout(timerId);
+            timersRef.current.delete(msgType);
+        }
+    };
+
+    const sendMessage = async (msg: BaseMessage, retryCount: number = 0) => {
+        if (!contextRef.current.connection) return;
+
+        const expectedType = GrabMessageOrderMap.get(msg.type)!;
+
+        clearTimer(expectedType);
+
+        if (retryCount >= 3) {
+            showNotification({
+                message:
+                    'Connection may be unstable, please try your drop again',
+                color: 'red',
+                icon: <IconX />,
+                autoClose: 4500,
+            });
+            console.error(`Attempt limit exceeded for type: ${msg.type}`);
+            return;
+        }
+
+        contextRef.current.connection.send(msg);
+
+        const timer = setTimeout(() => sendMessage(msg, retryCount + 1), 1000);
+        timersRef.current.set(expectedType, timer);
+    };
+
     const onMessage = async (msg: BaseMessage) => {
+        clearTimer(msg.type);
+
         if (msg.type === MessageType.Handshake) {
             const { input } = msg as HandshakeMessage;
 
@@ -139,6 +184,19 @@ export const useGrab = () => {
             id: dropId,
         });
 
+        if (!resp) {
+            pushLog(`Drop instance ${dropId} not found, closing connection...`);
+            showNotification({
+                message: 'Drop not found, check your link',
+                color: 'red',
+                icon: <IconX />,
+                autoClose: 3000,
+            });
+
+            cleanup();
+            return;
+        }
+
         const { peerId: dropperId, nonce } = resp;
 
         contextRef.current.id = dropId;
@@ -167,7 +225,8 @@ export const useGrab = () => {
             send({ type: GrabEventType.Connect });
         });
 
-        connection.on('data', onMessage);
+        const handlerWithLock = withMessageLock(onMessage, pushLog);
+        connection.on('data', handlerWithLock);
 
         contextRef.current.connection = connection;
     };
@@ -187,12 +246,17 @@ export const useGrab = () => {
         connection!.send(message);
     };
 
-    const cleanup = () => {
-        if (contextRef.current.connection!.open)
+    const cleanup = async () => {
+        const { removeOnUnloadListener } = await import('shared/lib/peer');
+
+        if (contextRef.current.connection?.open)
             contextRef.current.connection!.close();
 
-        contextRef.current.peer!.disconnect();
-        contextRef.current.peer!.destroy();
+        contextRef.current.peer?.disconnect();
+        contextRef.current.peer?.destroy();
+
+        removeOnUnloadListener();
+
         contextRef.current = initGrabContext();
     };
 

web/lib/fetch.ts

@@ -18,6 +18,8 @@ export const get = async <Data>(
         },
     });
 
+    if (res.status === 404) return null;
+
     const data: Data = await res.json();
 
     return data;