Fix: Websocket stale closure & O(n²) packet decode issue

packages/nitro-react/src/components/CloseButton.tsx

@@ -2,7 +2,7 @@ import { forwardRef, type HTMLAttributes } from 'react';
 
 import { cn, cva, type VariantProps } from '#base/utils';
 
-const closebuttonVariantsConfig = {
+const closeButtonVariantsConfig = {
     variant: {
         // default
         '0': '[border-image-source:var(--closebutton-0-default-src)] [border-image-slice:0_0_0_0_fill] [border-image-width:0px_0px_0px_0px] [border-image-repeat:stretch] hover:[border-image-source:var(--closebutton-0-hovering-src)] hover:[border-image-slice:0_0_0_0_fill] hover:[border-image-width:0px_0px_0px_0px] hover:[border-image-repeat:stretch] active:[border-image-source:var(--closebutton-0-pressed-src)] active:[border-image-slice:0_0_0_0_fill] active:[border-image-width:0px_0px_0px_0px] active:[border-image-repeat:stretch]',
@@ -17,30 +17,30 @@ const closebuttonVariantsConfig = {
     },
 } as const;
 
-const closebuttonVariants = cva(
+const closeButtonVariants = cva(
     '',
     {
-        variants: closebuttonVariantsConfig,
+        variants: closeButtonVariantsConfig,
         defaultVariants: {
             variant: '0',
         },
     }
 );
 
-type ClosebuttonVariantProps = VariantProps<typeof closebuttonVariantsConfig>;
+type CloseButtonVariantProps = VariantProps<typeof closeButtonVariantsConfig>;
 
-interface ClosebuttonProps extends HTMLAttributes<HTMLDivElement>, ClosebuttonVariantProps {
+interface CloseButtonProps extends HTMLAttributes<HTMLDivElement>, CloseButtonVariantProps {
     className?: string;
 }
 
-export const Closebutton = forwardRef<HTMLDivElement, ClosebuttonProps>(
+export const CloseButton = forwardRef<HTMLDivElement, CloseButtonProps>(
     ({ className, variant, ...props }, ref) => (
         <div
             ref={ref}
-            className={cn(closebuttonVariants({ variant }), className)}
+            className={cn(closeButtonVariants({ variant }), className)}
             {...props}
         />
     )
 );
 
-Closebutton.displayName = 'Closebutton';
+CloseButton.displayName = 'CloseButton';

packages/nitro-react/src/context/communication/WebSocketContext.tsx

@@ -3,6 +3,7 @@ import { createContext } from 'react';
 
 type WebSocketStore = {
     isAuthenticated: boolean,
+    isDisconnected: boolean,
     connect: () => void,
     send: <T extends object>(...packets: IOutgoingPacket<T>[]) => void,
     subscribe: <T extends object>(event: IncomingPacketConstructor<T>, handler: (data: T) => void) => void;

packages/nitro-react/src/context/communication/WebSocketContextProvider.tsx

@@ -15,9 +15,11 @@ type ProviderProps = {
     children: ReactNode;
 }
 
+type ConnectionPhase = 'idle' | 'connecting' | 'authenticating' | 'awaitingHandlers' | 'ready' | 'closed';
+
 export const WebSocketContextProvider = ({ children }: ProviderProps) => {
-    const [isAuthenticated, setIsAuthenticated] = useState<boolean>(false);
-    const [isReady, setIsReady] = useState<boolean>(false);
+    const phase = useRef<ConnectionPhase>('idle');
+    const [renderedPhase, setRenderedPhase] = useState<ConnectionPhase>('idle');
     const { incomingByHeader, incomingCtors, incomingHeaderByCtor, registerManyIncoming } = useCommunicationIncoming();
     const { outgoingHeaderByComposerName, registerManyOutgoing } = useCommunicationOutgoing();
     //const socketUrl = useConfigurationStore(x => x.config['socket.url'] as string) ?? undefined;
@@ -27,19 +29,29 @@ export const WebSocketContextProvider = ({ children }: ProviderProps) => {
     const listeners = useRef<Map<IncomingPacketConstructor<object>, Array<(data: object) => void>>>(new Map());
     const pendingClientMessages = useRef<IOutgoingPacket<object>[]>([]);
     const pendingServerMessages = useRef<IMessageDataWrapper[]>([]);
+    const hasConnected = useRef<boolean>(false);
 
     const connect = () => {
         try {
             const params = new URLSearchParams(window.location.search);
             const socketUrl = params.get('socketUrl') ?? '';
 
             if (!socketUrl || !socketUrl.length || ws.current) return;
+            if (hasConnected.current) return;
+
+            hasConnected.current = true;
+
+            const socket = new WebSocket(socketUrl);
+
+            ws.current = socket;
+
+            socket.binaryType = 'arraybuffer';
 
-            ws.current = new WebSocket(socketUrl);
+            setPhase('connecting');
 
-            ws.current.binaryType = 'arraybuffer';
+            socket.onopen = () => {
+                setPhase('authenticating');
 
-            ws.current.onopen = (event: Event) => {
                 send(new ClientHelloComposer({
                     production: production,
                     platform: 'WEB',
@@ -56,17 +68,25 @@ export const WebSocketContextProvider = ({ children }: ProviderProps) => {
                 }));
             };
 
-            ws.current.onerror = (event: Event) => {
+            socket.onerror = (event: Event) => {
                 NitroLogger.error('WebSocket error:', event);
             };
 
-            ws.current.onclose = (event: CloseEvent) => {
+            socket.onclose = (event: CloseEvent) => {
                 NitroLogger.warn('WebSocket closed:', event.code, event.reason);
-                setIsAuthenticated(false);
-                setIsReady(false);
+
+                if (ws.current !== socket) return;
+
+                ws.current = undefined;
+                wsBuffer.current = new ArrayBuffer(0);
+
+                pendingClientMessages.current = [];
+                pendingServerMessages.current = [];
+
+                setPhase('closed');
             };
 
-            ws.current.onmessage = (event: MessageEvent<ArrayBuffer>) => {
+            socket.onmessage = (event: MessageEvent<ArrayBuffer>) => {
                 const array = new Uint8Array(wsBuffer.current.byteLength + event.data.byteLength);
 
                 array.set(new Uint8Array(wsBuffer.current), 0);
@@ -85,15 +105,7 @@ export const WebSocketContextProvider = ({ children }: ProviderProps) => {
         try {
             if (wsBuffer.current.byteLength === 0) return;
 
-            const wrappers = decodeWrappers();
-
-            if (isAuthenticated && !isReady) {
-                pendingServerMessages.current.push(...wrappers);
-
-                return;
-            }
-
-            processWrappers(...wrappers);
+            dispatchWrappers(decodeWrappers());
         } catch (err) {
             NitroLogger.error(err);
         }
@@ -155,53 +167,69 @@ export const WebSocketContextProvider = ({ children }: ProviderProps) => {
 
         const reader = new BinaryReader(wsBuffer.current);
 
-        while (wsBuffer.current.byteLength) {
-            if (wsBuffer.current.byteLength < 4) break;
+        let consumed = 0;
 
-            const length = reader.readInt();
+        try {
+            while (wsBuffer.current.byteLength - consumed >= 4) {
+                const length = reader.readInt();
 
-            if (length > (wsBuffer.current.byteLength - 4)) break;
+                if (length < 2) {
+                    NitroLogger.error(`WebSocket: Malformed packet length: ${length}`);
+                    ws.current?.close();
+                    break;
+                }
+
+                if (length > reader.remaining()) break;
 
-            const extracted = reader.readBytes(length);
+                const extracted = reader.readBytes(length);
 
-            wrappers.push(new EvaWireDataWrapper(extracted.readShort(), extracted));
+                wrappers.push(new EvaWireDataWrapper(extracted.readShort(), extracted));
 
-            wsBuffer.current = wsBuffer.current.slice(length + 4);
+                consumed += length + 4;
+            }
+        } catch (err) {
+            NitroLogger.error(err);
         }
 
+        if (consumed) wsBuffer.current = wsBuffer.current.slice(consumed);
+
         return wrappers;
     }
 
-    const processWrappers = (...wrappers: IMessageDataWrapper[]) => {
+    const processWrapper = (wrapper: IMessageDataWrapper) => {
         try {
-            if (!wrappers || !wrappers.length) return;
-
-            for (const wrapper of wrappers) {
-                try {
-                    const ctor = incomingByHeader.current.get(wrapper.header);
+            const ctor = incomingByHeader.current.get(wrapper.header);
 
-                    if (!ctor) continue;
+            if (!ctor) return;
 
-                    const handlers = listeners.current.get(ctor);
+            const handlers = listeners.current.get(ctor);
 
-                    if (!handlers?.length) continue;
+            if (!handlers?.length) return;
 
-                    const parsed = new ctor().parse(wrapper);
+            const parsed = new ctor().parse(wrapper);
 
-                    for (const handle of handlers) handle(parsed);
-                } catch (err) {
-                    NitroLogger.error(err);
-                }
-            }
+            for (const handle of handlers) handle(parsed);
         } catch (err) {
             NitroLogger.error(err);
         }
     }
 
+    const dispatchWrappers = (wrappers: IMessageDataWrapper[]) => {
+        for (let index = 0; index < wrappers.length; index++) {
+            if (phase.current === 'awaitingHandlers') {
+                pendingServerMessages.current.push(...wrappers.slice(index));
+
+                return;
+            }
+
+            processWrapper(wrappers[index]);
+        }
+    }
+
     const send = <T extends object,>(...packets: IOutgoingPacket<T>[]) => {
         if (!packets?.length) return;
 
-        if (isAuthenticated && !isReady) {
+        if (phase.current === 'awaitingHandlers') {
             pendingClientMessages.current.push(...packets);
 
             return;
@@ -240,6 +268,12 @@ export const WebSocketContextProvider = ({ children }: ProviderProps) => {
         }
     }
 
+    const setPhase = (next: ConnectionPhase) => {
+        phase.current = next;
+
+        setRenderedPhase(next);
+    }
+
     const subscribe = <T extends object>(
         event: IncomingPacketConstructor<T>,
         handler: (data: T) => void
@@ -270,35 +304,40 @@ export const WebSocketContextProvider = ({ children }: ProviderProps) => {
     };
 
     const setReady = () => {
-        if (isReady) return;
+        if (phase.current !== 'awaitingHandlers') return;
 
         const pendingClient = pendingClientMessages.current;
         const pendingServer = pendingServerMessages.current;
 
         pendingServerMessages.current = [];
         pendingClientMessages.current = [];
 
-        setIsReady(true);
+        setPhase('ready');
 
-        processWrappers(...pendingServer);
+        dispatchWrappers(pendingServer);
         sendRaw(...pendingClient);
     }
 
     useEffect(() => {
         registerManyIncoming(GetIncomingPackets());
         registerManyOutgoing(GetOutgoingPackets());
+
+        return subscribe(AuthenticationOKMessage, () => setPhase('awaitingHandlers'));
     }, []);
 
-    useEffect(() => {
-        if (isAuthenticated) return;
+    useEffect(() => () => {
+        const socket = ws.current;
+
+        ws.current = undefined;
+
+        socket?.close(1000, 'Client shutting down');
+    }, []);
 
-        return subscribe(AuthenticationOKMessage, data => {
-            setIsAuthenticated(true);
-        });
-    }, [isAuthenticated, subscribe]);
+    const isAuthenticated = renderedPhase === 'awaitingHandlers' || renderedPhase === 'ready';
+    const isDisconnected = renderedPhase === 'closed';
 
     return (
-        <WebSocketContext.Provider value={{ isAuthenticated, connect, send, subscribe, setReady }}>
+        <WebSocketContext.Provider value={{ isAuthenticated, isDisconnected, connect, send, subscribe, setReady }}>
             {children}
         </WebSocketContext.Provider>
     );

packages/nitro-react/src/views/room-widgets/object-infostand/InfostandFurniView.tsx

@@ -1,7 +1,7 @@
 import { FurniturePickupMode, IRoomFurnitureData } from "@nitrodevco/nitro-api";
 
 import { Border, Button } from "#base/components";
-import { CloseButton } from "#base/components/Closebutton";
+import { CloseButton } from "#base/components/CloseButton";
 import { FurnitureImage } from "#base/components/FurnitureImage";
 import { useLocalizationStore } from "#base/stores";