better E2E encryption, nicer UI, bug fixes, more

server/websocketHandler.ts

@@ -45,6 +45,13 @@ async function joinRoom(roomId: string, socket: WebSocket) {
         // for some reason, when you filter the array when the length is 1 it stays at 1, but we *know* that if its 1
         // then when this client disconnects, the room should be deleted since the room is empty
         if (room.length === 1) {
+            // give a 5 second grace period before deleting the room
+            setTimeout(() => {
+                if (rooms.get(roomId)?.length === 1) {
+                    console.log("Room is empty, deleting");
+                    deleteRoom(roomId);
+                }
+            }, 5000)
             deleteRoom(roomId);
             return;
         }
@@ -54,19 +61,9 @@ async function joinRoom(roomId: string, socket: WebSocket) {
 
     // TODO: consider letting rooms get larger than 2 clients
     if (room.length == 2) {
-        // A room key used to wrap the clients public keys during key exchange
-        let roomKey = await crypto.subtle.generateKey(
-            {
-                name: "AES-KW",
-                length: 256,
-            },
-            true,
-            ["wrapKey", "unwrapKey"],
-        )
-        let jsonWebKey = await crypto.subtle.exportKey("jwk", roomKey);
         room.forEach(async client => {
             // announce the room is ready, and tell each peer if they are the initiator
-            client.send(JSON.stringify({ type: SocketMessageType.ROOM_READY, data: { isInitiator: client !== socket, roomKey: { key: jsonWebKey } } }));
+            client.send(JSON.stringify({ type: SocketMessageType.ROOM_READY, data: { isInitiator: client !== socket } }));
         });
     }
 

src/components/RTCMessage.svelte

@@ -1,7 +1,7 @@
 <script lang="ts">
     import { writable, type Writable } from "svelte/store";
-    import { room, connectionState } from "../stores/roomStore";
+    import { room } from "../stores/roomStore";
-    import { connected } from "../stores/websocketStore";
+    import { webSocketConnected } from "../stores/websocketStore";
     import {
         isRTCConnected,
         dataChannelReady,
@@ -12,6 +12,7 @@
     import { WebRTCPacketType } from "../types/webrtc";
     import { ConnectionState } from "../types/websocket";
     import { MessageType } from "../types/message";
+    import { fade } from "svelte/transition";
 
     let inputMessage: Writable<string> = writable("");
     let inputFile = writable(null);
@@ -34,7 +35,6 @@
         // }
 
         if ($inputMessage) {
-            // $messages = [...$messages, `You: ${$inputMessage}`];
             $messages = [
                 ...$messages,
                 {
@@ -56,13 +56,37 @@
         }
     }
 
+    let canCloseLoadingOverlay = writable(false);
+    keyExchangeDone.subscribe((value) => {
+        console.log("Key exchange done:", value, $keyExchangeDone);
+        if (value) {
+            // provide a grace period for the user to see that the connection is established
+            setTimeout(() => {
+                canCloseLoadingOverlay.set(true);
+            }, 1000);
+        }
+    });
+
     function pickFile() {
         inputFileElement.click();
     }
 </script>
 
+<p>{$room?.id} - {$room?.connectionState} - {$webSocketConnected}</p>
+
 <!-- If we are in a room, connected to the websocket server, and the have been informed that we are connected to the room -->
-{#if $room !== null && $connected === true && $connectionState === ConnectionState.CONNECTED}
+{#if $room !== null && $webSocketConnected === true && $room.connectionState === ConnectionState.CONNECTED}
+    <div
+        class="flex flex-col sm:max-w-4/5 lg:max-w-3/5 min-h-[calc(5/12_*_100vh)]"
+    >
+        <div
+            class="flex-grow flex flex-col overflow-y-auto mb-4 p-2 bg-gray-800 rounded break-all relative"
+        >
+            {#if !$isRTCConnected || !$dataChannelReady || !$keyExchangeDone || !$canCloseLoadingOverlay}
+                <div
+                    transition:fade={{ duration: 300 }}
+                    class="absolute top-0 left-0 bottom-0 right-0 flex justify-center items-center flex-col bg-black/55 backdrop-blur-md"
+                >
                     {#if !$isRTCConnected}
                         <p>Waiting for peer to connect...</p>
                     {:else if !$dataChannelReady}
@@ -70,25 +94,70 @@
                     {:else if !$keyExchangeDone}
                         <p>Establishing a secure connection with the peer...</p>
                     {:else}
-        <div
+                        <p>
-            class="flex-grow overflow-y-auto mb-4 p-2 bg-gray-800 rounded break-all"
+                            Successfully established a secure connection to
+                            peer!
+                        </p>
+                    {/if}
+                    <div class="mt-2">
+                        {#if !$keyExchangeDone}
+                            <!-- loading spinner -->
+                            <svg
+                                class="animate-spin -ml-1 mr-3 h-5 w-5 text-white"
+                                xmlns="http://www.w3.org/2000/svg"
+                                fill="none"
+                                viewBox="0 0 24 24"
                             >
+                                <circle
+                                    class="opacity-25"
+                                    cx="12"
+                                    cy="12"
+                                    r="10"
+                                    stroke="currentColor"
+                                    stroke-width="4"
+                                />
+                                <path
+                                    class="opacity-75"
+                                    fill="currentColor"
+                                    d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"
+                                />
+                            </svg>
+                        {:else}
+                            <svg
+                                xmlns="http://www.w3.org/2000/svg"
+                                width="16"
+                                height="16"
+                                viewBox="0 0 24 24"
+                            >
+                                <path
+                                    fill="none"
+                                    stroke="currentColor"
+                                    stroke-linecap="round"
+                                    stroke-linejoin="round"
+                                    stroke-width="2"
+                                    d="M5 13l4 4L19 7"
+                                />
+                            </svg>
+                        {/if}
+                    </div>
+                </div>
+            {/if}
             {#each $messages as msg}
-                <div>
+                <div class="flex flex-row gap-2">
-                    <div class="w-fit h-max">
+                    <p class="break-keep">
                         {#if msg.initiator}
                             You:
                         {:else}
                             Peer:
                         {/if}
-                    </div>
+                    </p>
-                    <span>
+                    <p>
                         {#if msg.type === MessageType.TEXT}
                             {msg.data}
                         {:else}
                             Unknown message type: {msg.type}
                         {/if}
-                    </span>
+                    </p>
                 </div>
             {/each}
         </div>
@@ -98,19 +167,25 @@
             bind:this={inputFileElement}
             class="absolute opacity-0 -top-[9999px] -left-[9999px]"
         />
-        <div class="flex gap-2">
+        <div class="flex gap-2 w-full flex-row">
             <input
                 type="text"
                 bind:value={$inputMessage}
                 on:keyup={(e) => e.key === "Enter" && sendMessage()}
+                disabled={!$isRTCConnected ||
+                    !$dataChannelReady ||
+                    !$keyExchangeDone}
                 placeholder="Type your message..."
-                class="flex-grow p-2 rounded bg-gray-700 border border-gray-600 text-gray-100 placeholder-gray-400 focus:outline-none focus:ring-2 focus:ring-blue-500"
+                class="flex-grow p-2 rounded bg-gray-700 border border-gray-600 text-gray-100 placeholder-gray-400
+            focus:outline-none focus:ring-2 focus:ring-blue-500 disabled:opacity-50 disabled:cursor-not-allowed"
             />
             <button
                 on:click={pickFile}
-                disabled={!dataChannelReady}
+                disabled={!$isRTCConnected ||
+                    !$dataChannelReady ||
+                    !$keyExchangeDone}
                 aria-label="Pick file"
-                class="px-4 py-2 bg-blue-600 hover:bg-blue-700 text-white rounded disabled:opacity-50 disabled:cursor-not-allowed"
+                class="px-4 py-2 bg-blue-600 not-disabled:hover:bg-blue-700 text-white rounded disabled:opacity-50 disabled:cursor-not-allowed"
             >
                 <svg
                     xmlns="http://www.w3.org/2000/svg"
@@ -129,11 +204,13 @@
             </button>
             <button
                 on:click={sendMessage}
-                disabled={!dataChannelReady}
+                disabled={!$isRTCConnected ||
-                class="px-4 py-2 bg-blue-600 hover:bg-blue-700 text-white rounded disabled:opacity-50 disabled:cursor-not-allowed"
+                    !$dataChannelReady ||
+                    !$keyExchangeDone}
+                class="px-4 py-2 bg-blue-600 not-disabled:hover:bg-blue-700 text-white rounded disabled:opacity-50 disabled:cursor-not-allowed"
             >
                 Send
             </button>
         </div>
-    {/if}
+    </div>
 {/if}

src/lib/webrtc.ts

@@ -1,7 +1,7 @@
 import { get } from 'svelte/store';
-import { ws } from '../stores/websocketStore';
+import { WebSocketMessageType, ws } from '../stores/websocketStore';
-import { roomKey } from '../utils/webrtcUtil';
 import { WebRTCPacketType, type KeyStore, type WebRTCPeerCallbacks } from '../types/webrtc';
+import { clientKeyConfig } from '../shared/keyConfig';
 
 export class WebRTCPeer {
     private peer: RTCPeerConnection | null = null;
@@ -28,13 +28,13 @@ export class WebRTCPeer {
     }
 
     private sendIceCandidate(candidate: RTCIceCandidate) {
-        get(ws).send(JSON.stringify({
+        get(ws).send({
-            type: 'ice-candidate',
+            type: WebSocketMessageType.WEBRTC_ICE_CANDIDATE,
             data: {
                 roomId: this.roomId,
                 candidate: candidate,
             },
-        }))
+        })
     }
 
     public async initialize() {
@@ -98,7 +98,6 @@ export class WebRTCPeer {
 
         channel.onmessage = async (event: MessageEvent<ArrayBuffer>) => {
             console.log('data channel message:', event.data);
-
             // event is binary data, we need to parse it, convert it into a WebRTCMessage, and then decrypt it if
             // necessary
             let data = new Uint8Array(event.data);
@@ -116,32 +115,17 @@ export class WebRTCPeer {
 
                 console.log("Received key exchange", data.buffer);
 
-                // let textDecoder = new TextDecoder();
+                const textDecoder = new TextDecoder();
-                // let dataString = textDecoder.decode(data.buffer);
+                const jsonKey = JSON.parse(textDecoder.decode(data));
 
-                // console.log("Received key exchange", dataString);
+                console.log("Received key exchange", jsonKey);
 
-                // let json = JSON.parse(dataString);
+                this.keys.peersPublicKey = await window.crypto.subtle.importKey(
-
-                let unwrappingKey = get(roomKey);
-                if (!unwrappingKey.key) throw new Error("Room key not set");
-
-                this.keys.peersPublicKey = await window.crypto.subtle.unwrapKey(
                     "jwk",
-                    data,
+                    jsonKey,
-                    unwrappingKey.key,
+                    clientKeyConfig,
-                    {
-                        name: "AES-KW",
-                        length: 256,
-                    },
-                    {
-                        name: "RSA-OAEP",
-                        modulusLength: 4096,
-                        publicExponent: new Uint8Array([1, 0, 1]),
-                        hash: "SHA-256",
-                    },
                     true,
-                    ["encrypt"],
+                    ["wrapKey"],
                 );
 
                 // if our keys are not generated, start the reponding side of the key exchange
@@ -155,7 +139,32 @@ export class WebRTCPeer {
             }
 
             if (encrypted) {
-                data = new Uint8Array(await this.decrypt(data.buffer));
+                if (!this.keys.localKeys) {
+                    throw new Error("Local keypair not generated");
+                }
+
+                // start at 0 since the header is already sliced off
+                let keyLength = data[0] << 8 | data[1];
+
+                let aeskey = await window.crypto.subtle.unwrapKey(
+                    "raw",
+                    data.subarray(2, 2 + keyLength),
+                    this.keys.localKeys.privateKey,
+                    clientKeyConfig,
+                    {
+                        name: "AES-GCM",
+                        length: 256,
+                    },
+                    true,
+                    ["encrypt", "decrypt"],
+                )
+
+                let iv = data.subarray(2 + keyLength, 2 + keyLength + 16);
+                let encryptedData = data.subarray(2 + keyLength + 16);
+
+                console.log("Decrypting message", encryptedData);
+
+                data = new Uint8Array(await this.decrypt(encryptedData, aeskey, iv));
             }
 
             let message = {
@@ -187,13 +196,13 @@ export class WebRTCPeer {
 
             await this.peer.setLocalDescription(offer)
 
-            get(ws).send(JSON.stringify({
+            get(ws).send({
-                type: 'offer',
+                type: WebSocketMessageType.WEBRTC_OFFER,
                 data: {
                     roomId: this.roomId,
                     sdp: offer,
                 },
-            }));
+            });
         } catch (error) {
             console.info('Error creating offer:', error);
             // should trigger re-negotiation
@@ -221,13 +230,13 @@ export class WebRTCPeer {
 
             console.log("Sending answer", answer);
 
-            get(ws).send(JSON.stringify({
+            get(ws).send({
-                type: 'answer',
+                type: WebSocketMessageType.WERTC_ANSWER,
                 data: {
                     roomId: this.roomId,
                     sdp: answer,
                 },
-            }));
+            });
 
         } catch (error) {
             console.error('Error creating answer:', error);
@@ -248,20 +257,14 @@ export class WebRTCPeer {
 
     private async generateKeyPair() {
         console.log("Generating key pair");
+        // this key pair is used for wrapping the unique AES-GCM key for each message
         const keyPair = await window.crypto.subtle.generateKey(
-            {
+            clientKeyConfig,
-                name: "RSA-OAEP",
-                modulusLength: 4096,
-                publicExponent: new Uint8Array([1, 0, 1]),
-                hash: "SHA-256",
-            },
             true,
-            ["encrypt", "decrypt"],
+            ["wrapKey", "unwrapKey"],
         );
 
-        if (keyPair instanceof CryptoKey) {
+        console.log("generated key pair", keyPair);
-            throw new Error("Key pair not generated");
-        }
 
         this.keys.localKeys = keyPair;
     }
@@ -271,50 +274,40 @@ export class WebRTCPeer {
         await this.generateKeyPair();
         if (!this.keys.localKeys) throw new Error("Key pair not generated");
 
-        let wrappingKey = get(roomKey);
+        console.log("exporting key", this.keys.localKeys.publicKey);
-        if (!wrappingKey.key) throw new Error("Room key not set");
 
+        const exported = await window.crypto.subtle.exportKey("jwk", this.keys.localKeys.publicKey);
 
-        console.log("wrapping key", this.keys.localKeys.publicKey, wrappingKey.key);
+        // convert exported key to a string then pack that sting into an array buffer
-        const exported = await window.crypto.subtle.wrapKey(
+        const exportedKeyBuffer = new TextEncoder().encode(JSON.stringify(exported));
-            "jwk",
-            this.keys.localKeys.publicKey,
-            wrappingKey.key,
-            {
-                name: "AES-KW",
-                length: 256,
-            },
-        );
-
-        console.log("wrapping key exported", exported);
-
-        const exportedKeyBuffer = exported;
 
         console.log("exported key buffer", exportedKeyBuffer);
 
-        this.send(exportedKeyBuffer, WebRTCPacketType.KEY_EXCHANGE);
+        this.send(exportedKeyBuffer.buffer, WebRTCPacketType.KEY_EXCHANGE);
     }
 
-    private async encrypt(data: ArrayBuffer): Promise<ArrayBuffer> {
+    private async encrypt(data: Uint8Array<ArrayBuffer>, key: CryptoKey, iv: Uint8Array<ArrayBuffer>): Promise<ArrayBuffer> {
-        if (!this.keys.peersPublicKey) throw new Error("Peer's public key not set");
-
         return await window.crypto.subtle.encrypt(
             {
-                name: "RSA-OAEP",
+                name: "AES-GCM",
+                length: 256,
+                iv,
+                tagLength: 128,
             },
-            this.keys.peersPublicKey,
+            key,
             data,
         );
     }
 
-    private async decrypt(data: ArrayBuffer): Promise<ArrayBuffer> {
+    private async decrypt(data: Uint8Array<ArrayBuffer>, key: CryptoKey, iv: Uint8Array<ArrayBuffer>): Promise<ArrayBuffer> {
-        if (!this.keys.localKeys) throw new Error("Local keypair not generated");
-
         return await window.crypto.subtle.decrypt(
             {
-                name: "RSA-OAEP",
+                name: "AES-GCM",
+                length: 256,
+                iv,
+                tagLength: 128,
             },
-            this.keys.localKeys.privateKey,
+            key,
             data,
         );
     }
@@ -331,15 +324,36 @@ export class WebRTCPeer {
         if (this.keys.peersPublicKey && type != WebRTCPacketType.KEY_EXCHANGE) {
             console.log("Sending encrypted message", data);
 
-            let encryptedData = await this.encrypt(data);
+            let iv = window.crypto.getRandomValues(new Uint8Array(16));
+            let key = await window.crypto.subtle.generateKey(
+                {
+                    name: "AES-GCM",
+                    length: 256,
+                },
+                true,
+                ["encrypt", "decrypt"],
+            )
 
-            console.log("Encrypted data", encryptedData);
+            let encryptedData = await this.encrypt(new Uint8Array(data), key, iv);
+
+            let exportedKey = await window.crypto.subtle.wrapKey(
+                "raw",
+                key,
+                this.keys.peersPublicKey,
+                clientKeyConfig,
+            )
 
             header |= 1 << 7;
 
-            let buf = new Uint8Array(encryptedData.byteLength + 1);
+            let buf = new Uint8Array(encryptedData.byteLength + 3 + exportedKey.byteLength + iv.byteLength);
             buf[0] = header;
-            buf.subarray(1).set(new Uint8Array(encryptedData));
+            buf[1] = (exportedKey.byteLength >> 8) & 0xFF;
+            buf[2] = exportedKey.byteLength & 0xFF;
+            buf.subarray(3).set(new Uint8Array(exportedKey));
+            buf.subarray(3 + exportedKey.byteLength).set(new Uint8Array(iv));
+            buf.subarray(3 + exportedKey.byteLength + iv.byteLength).set(new Uint8Array(encryptedData));
+
+            console.log("Sending encrypted message", buf);
 
             this.dataChannel.send(buf.buffer);
         } else {

src/routes/+page.svelte

@@ -1,6 +1,10 @@
 <script lang="ts">
-    import { ws, connected } from "../stores/websocketStore";
+    import {
-    import { room, connectionState } from "../stores/roomStore";
+        ws,
+        webSocketConnected,
+        WebSocketMessageType,
+    } from "../stores/websocketStore";
+    import { room } from "../stores/roomStore";
     import { browser } from "$app/environment";
     import { peer, handleMessage } from "../utils/webrtcUtil";
     import { onDestroy, onMount } from "svelte";
@@ -8,13 +12,19 @@
     import { ConnectionState } from "../types/websocket";
 
     onMount(async () => {
-        $connectionState = ConnectionState.CONNECTING;
+        room.update((room) => ({
+            ...room,
+            connectionState: ConnectionState.CONNECTING,
+        }));
         $ws.addEventListener("message", handleMessage);
     });
 
     onDestroy(() => {
         if ($ws) {
-            $connectionState = ConnectionState.DISCONNECTED;
+            room.update((room) => ({
+                ...room,
+                connectionState: ConnectionState.DISCONNECTED,
+            }));
             $ws.removeEventListener("message", handleMessage);
         }
         if ($peer) {
@@ -26,20 +36,20 @@
 <div class="p-4">
     <h1>Welcome to Wormhole!</h1>
 
-    {#if $connected}
+    {#if $webSocketConnected}
         <button
             on:click={() => {
-                $ws.send(JSON.stringify({ type: "create" })); // send a message when the button is clicked
+                $ws.send({ type: WebSocketMessageType.CREATE_ROOM }); // send a message when the button is clicked
             }}>Create Room</button
         >
     {:else}
         <p>Connecting to server...</p>
     {/if}
 
-    {#if $room && browser}
+    {#if $room.id && browser}
         <p>Room created!</p>
         <p>Share this link with your friend:</p>
-        <a href={`${location.origin}/${$room}`}>{location.origin}/{$room}</a>
+        <a href={`${location.origin}/${$room}`}>{location.origin}/{$room.id}</a>
     {/if}
 
     <RtcMessage />

src/routes/[roomId]/+page.svelte

@@ -1,31 +1,47 @@
 <script lang="ts">
     import { page } from "$app/state";
     import { onDestroy, onMount } from "svelte";
-    import { room, connectionState } from "../../stores/roomStore";
+    import { room } from "../../stores/roomStore";
     import { error, handleMessage, peer } from "../../utils/webrtcUtil";
-    import { ws, connected } from "../../stores/websocketStore";
+    import {
+        ws,
+        webSocketConnected,
+        WebSocketMessageType,
+    } from "../../stores/websocketStore";
     import RtcMessage from "../../components/RTCMessage.svelte";
     import { ConnectionState } from "../../types/websocket";
 
-    if (!page.params.roomId) {
+    const roomId = page.params.roomId;
+    if (roomId === undefined) {
         throw new Error("Room ID not provided");
     }
 
     // subscribe to the websocket store
-    room.set(page.params.roomId);
+    room.update((room) => ({ ...room, id: roomId }));
 
     onMount(async () => {
         $ws.addEventListener("message", handleMessage);
 
-        $ws.onopen = () => {
+        webSocketConnected.subscribe((value) => {
-            $connectionState = ConnectionState.CONNECTING;
+            if (value) {
-            $ws.send(JSON.stringify({ type: "join", roomId: $room }));
+                $ws.send({ type: WebSocketMessageType.JOIN_ROOM, roomId });
-        };
+            }
+        });
+        // $ws.onopen = () => {
+        //     room.update((room) => ({
+        //         ...room,
+        //         connectionState: ConnectionState.CONNECTING,
+        //     }));
+        //     $ws.send({ type: WebSocketMessageType.JOIN_ROOM, roomId });
+        // };
     });
 
     onDestroy(() => {
         if ($ws) {
-            $connectionState = ConnectionState.DISCONNECTED;
+            room.update((room) => ({
+                ...room,
+                connectionState: ConnectionState.DISCONNECTED,
+            }));
             $ws.close();
         }
         if ($peer) {
@@ -37,7 +53,7 @@
 <div class="p-4">
     {#if $error}
         <p>Whoops! That room doesn't exist.</p>
-    {:else if !$connected || $connectionState === ConnectionState.CONNECTING}
+    {:else if !$webSocketConnected || $room.connectionState === ConnectionState.CONNECTING}
         <p>Connecting to server...</p>
     {:else}
         <RtcMessage />

src/shared/keyConfig.ts

@@ -0,0 +1,6 @@
+export const clientKeyConfig = {
+    name: "RSA-OAEP",
+    modulusLength: 4096,
+    publicExponent: new Uint8Array([1, 0, 1]),
+    hash: "SHA-256",
+};

src/stores/roomStore.ts

@@ -1,5 +1,13 @@
 import { writable, type Writable } from 'svelte/store';
 import { ConnectionState } from '../types/websocket';
 
-export let room: Writable<string | null> = writable(null);
+export interface Room {
-export let connectionState: Writable<ConnectionState> = writable(ConnectionState.DISCONNECTED);
+    id: string | null;
+    connectionState: ConnectionState;
+}
+
+export const room: Writable<Room> = writable({
+    id: null,
+    connectionState: ConnectionState.DISCONNECTED,
+    key: null,
+});

src/stores/websocketStore.ts

@@ -1,10 +1,125 @@
 import { writable } from 'svelte/store';
 import { browser } from '$app/environment';
 
-let socket: WebSocket | null = null;
+export enum WebSocketMessageType {
-export const connected = writable(false);
+    // room messages
+    CREATE_ROOM = "create",
+    JOIN_ROOM = "join",
 
-function createSocket(): WebSocket {
+    // response messages
+    ROOM_CREATED = "created",
+    ROOM_JOINED = "joined",
+    ROOM_READY = "ready",
+
+    // webrtc messages
+    WEBRTC_OFFER = "offer",
+    WERTC_ANSWER = "answer",
+    WEBRTC_ICE_CANDIDATE = "ice-candidate",
+
+    ERROR = "error",
+}
+
+export type WebSocketMessage =
+    | CreateRoomMessage
+    | JoinRoomMessage
+    | RoomCreatedMessage
+    | RoomJoinedMessage
+    | RoomReadyMessage
+    | OfferMessage
+    | AnswerMessage
+    | IceCandidateMessage
+    | ErrorMessage;
+
+interface ErrorMessage {
+    type: WebSocketMessageType.ERROR;
+    data: string;
+}
+
+interface CreateRoomMessage {
+    type: WebSocketMessageType.CREATE_ROOM;
+}
+
+interface JoinRoomMessage {
+    type: WebSocketMessageType.JOIN_ROOM;
+    roomId: string;
+}
+
+interface RoomCreatedMessage {
+    type: WebSocketMessageType.ROOM_CREATED;
+    data: string;
+}
+
+interface RoomJoinedMessage {
+    type: WebSocketMessageType.ROOM_JOINED;
+    roomId: string;
+}
+
+interface RoomReadyMessage {
+    type: WebSocketMessageType.ROOM_READY;
+    data: {
+        isInitiator: boolean;
+        roomKey: {
+            key: JsonWebKey;
+        };
+    };
+}
+
+interface OfferMessage {
+    type: WebSocketMessageType.WEBRTC_OFFER;
+    data: {
+        roomId: string;
+        sdp: RTCSessionDescriptionInit;
+    };
+}
+
+interface AnswerMessage {
+    type: WebSocketMessageType.WERTC_ANSWER;
+    data: {
+        roomId: string;
+        sdp: RTCSessionDescriptionInit;
+    };
+}
+
+interface IceCandidateMessage {
+    type: WebSocketMessageType.WEBRTC_ICE_CANDIDATE;
+    data: {
+        roomId: string;
+        candidate: RTCIceCandidateInit;
+    };
+}
+
+export class Socket {
+    private ws: WebSocket;
+
+    public addEventListener: typeof WebSocket.prototype.addEventListener;
+    public removeEventListener: typeof WebSocket.prototype.removeEventListener;
+    public dispatchEvent: typeof WebSocket.prototype.dispatchEvent;
+    public close: typeof WebSocket.prototype.close;
+
+    constructor(public url: string, public protocols?: string | string[] | undefined) {
+        this.ws = new WebSocket(url, protocols);
+
+        this.ws.addEventListener("open", () => {
+            console.log("WebSocket opened");
+        });
+
+        this.addEventListener = this.ws.addEventListener.bind(this.ws);
+        this.removeEventListener = this.ws.removeEventListener.bind(this.ws);
+        this.dispatchEvent = this.ws.dispatchEvent.bind(this.ws);
+        this.close = this.ws.close.bind(this.ws);
+    }
+
+    public send(message: WebSocketMessage) {
+        console.log("Sending message:", message);
+
+        this.ws.send(JSON.stringify(message));
+    }
+}
+
+let socket: Socket | null = null;
+export const webSocketConnected = writable(false);
+
+function createSocket(): Socket {
     if (!browser) {
         return null;
     }
@@ -14,16 +129,20 @@ function createSocket(): WebSocket {
     }
 
     const protocol = location.protocol === 'https:' ? 'wss:' : 'ws:';
-    socket = new WebSocket(`${protocol}//${location.host}/`);
+    socket = new Socket(`${protocol}//${location.host}/`);
 
     socket.addEventListener('open', () => {
-        connected.set(true);
+        webSocketConnected.set(true);
         console.log('Connected to websocket server');
     });
 
     socket.addEventListener('close', () => {
-        connected.set(false);
+        webSocketConnected.set(false);
-        console.log('Disconnected from websocket server');
+        socket = null;
+        console.log('Disconnected from websocket server, reconnecting...');
+        setTimeout(() => {
+            ws.set(createSocket());
+        }, 1000);
     });
 
     return socket;

src/types/websocket.ts

@@ -52,9 +52,6 @@ export interface SocketMessageRoomReady extends SocketMessageBase {
     data: {
         roomId: string;
         isInitiator: boolean;
-        roomKey: {
-            key: JsonWebKey;
-        };
     };
 }
 

src/utils/webrtcUtil.ts

@@ -1,17 +1,17 @@
 import { writable, get, type Writable } from "svelte/store";
 import { WebRTCPeer } from "$lib/webrtc";
 import { WebRTCPacketType } from "../types/webrtc";
-import { room, connectionState } from "../stores/roomStore";
+import { room } from "../stores/roomStore";
 import { ConnectionState } from "../types/websocket";
 import { messages } from "../stores/messageStore";
 import { MessageType, type Message } from "../types/message";
+import { WebSocketMessageType, type WebSocketMessage } from "../stores/websocketStore";
 
-export const error = writable(null);
+export const error: Writable<string | null> = writable(null);
 export let peer: Writable<WebRTCPeer | null> = writable(null);
 export let isRTCConnected: Writable<boolean> = writable(false);
 export let dataChannelReady: Writable<boolean> = writable(false);
 export let keyExchangeDone: Writable<boolean> = writable(false);
-export let roomKey: Writable<{ key: CryptoKey | null }> = writable({ key: null });
 
 const callbacks = {
     onConnected: () => {
@@ -64,53 +64,36 @@ const callbacks = {
     },
 };
 
+
+
 export async function handleMessage(event: MessageEvent) {
-    console.log("Message received:", event.data);
+    console.log("Message received:", event.data, typeof event.data);
-    const message = JSON.parse(event.data);
+    const message: WebSocketMessage = JSON.parse(event.data);
 
     switch (message.type) {
-        case "created":
+        case WebSocketMessageType.ROOM_CREATED:
-            connectionState.set(ConnectionState.CONNECTED);
             console.log("Room created:", message.data);
-            room.set(message.data);
+            room.update((room) => ({ ...room, id: message.data, connectionState: ConnectionState.CONNECTED }));
             return;
-        case "join":
+        case WebSocketMessageType.JOIN_ROOM:
-            console.log("new client joined room", message.data);
+            console.log("new client joined room");
             return;
-        case "joined":
+        case WebSocketMessageType.ROOM_JOINED:
-            connectionState.set(ConnectionState.CONNECTED);
+            room.update((room) => ({ ...room, connectionState: ConnectionState.CONNECTED }));
-            console.log("Joined room:", message.data);
+            console.log("Joined room");
             return;
-        case "error":
+        case WebSocketMessageType.ERROR:
             console.error("Error:", message.data);
             error.set(message.data);
             return;
-        case "ready":
+        case WebSocketMessageType.ROOM_READY:
-            const roomId = get(room);
+            let roomId = get(room).id;
 
-            if (!roomId) {
+            if (roomId === null) {
                 console.error("Room not set");
                 return;
             }
 
-            try {
-                // let iv = new ArrayBuffer(message.data.roomKey.iv)
-
-                let importedRoomKey = await window.crypto.subtle.importKey(
-                    "jwk",
-                    message.data.roomKey.key,
-                    {
-                        name: "AES-KW",
-                        length: 256,
-                    },
-                    true,
-                    ["wrapKey", "unwrapKey"],
-                )
-                roomKey.set({ key: importedRoomKey });
-            } catch (e) {
-                console.error("Error importing room key:", e);
-                return;
-            }
             peer.set(new WebRTCPeer(
                 roomId,
                 message.data.isInitiator,
@@ -129,26 +112,26 @@ export async function handleMessage(event: MessageEvent) {
     }
 
     switch (message.type) {
-        case "offer":
+        case WebSocketMessageType.WEBRTC_OFFER:
             console.log("Received offer");
             await get(peer)?.setRemoteDescription(
                 new RTCSessionDescription(message.data.sdp),
             );
             await get(peer)?.createAnswer();
             return;
-        case "answer":
+        case WebSocketMessageType.WERTC_ANSWER:
             console.log("Received answer");
             await get(peer)?.setRemoteDescription(
                 new RTCSessionDescription(message.data.sdp),
             );
             return;
-        case "ice-candidate":
+        case WebSocketMessageType.WEBRTC_ICE_CANDIDATE:
             console.log("Received ICE candidate");
             await get(peer)?.addIceCandidate(message.data.candidate);
             return;
         default:
             console.warn(
-                `Unknown message type: ${message.type} from ${get(room)}`,
+                `Unknown message type: ${message.type} from ${get(room).id}`,
             );
     }
 }