encryption, code cleanup, nice types, bug fixes, and more

README.md

@@ -1,38 +1,17 @@
-# sv
+# Wormhole
+(needs a different name I think because I dont want to confuse it with wormhole.app)
 
-Everything you need to build a Svelte project, powered by [`sv`](https://github.com/sveltejs/cli).
+A peer-to-peer encrypted file sharing app.
 
-## Creating a project
+## Features
+- E2E communication
+- P2P file sharing
+- P2P chat
 
-If you're seeing this, you've probably already done this step. Congrats!
+Your data is peer-to-peer encrypted and only accessible to the people you share it with, it never touches any servers.
 
-```sh
+## How to use
-# create a new project in the current directory
+1. clone the repo
-npx sv create
+2. run `bun install`
-
+3. run `bun run dev --host` (webrtc doesnt co-operate with localhost connections, so connect via 127.0.0.1)
-# create a new project in my-app
+4. open the browser at http://127.0.0.1:5173
-npx sv create my-app
-```
-
-## Developing
-
-Once you've created a project and installed dependencies with `npm install` (or `pnpm install` or `yarn`), start a development server:
-
-```sh
-npm run dev
-
-# or start the server and open the app in a new browser tab
-npm run dev -- --open
-```
-
-## Building
-
-To create a production version of your app:
-
-```sh
-npm run build
-```
-
-You can preview the production build with `npm run preview`.
-
-> To deploy your app, you may need to install an [adapter](https://svelte.dev/docs/kit/adapters) for your target environment.

server/websocketHandler.ts

@@ -1,44 +1,22 @@
 import { WebSocketServer } from "ws";
 import type { WebSocket } from "ws";
+import { SocketMessageType, type SocketMessage } from "../src/types/websocket";
 
 // TODO: remove stale rooms somehow
 const rooms = new Map<string, WebSocket[]>();
 
-enum MessageType {
+async function createRoom(socket: WebSocket): Promise<string> {
-    // requests
-    CREATE_ROOM = 'create',
-    JOIN_ROOM = 'join',
-
-    // responses
-    ROOM_CREATED = 'created',
-    ROOM_JOINED = 'joined',
-    ROOM_READY = 'ready',
-
-    // webrtc
-    ICE_CANDIDATE = 'ice-candidate',
-    OFFER = 'offer',
-    ANSWER = 'answer',
-
-    ERROR = 'error',
-}
-
-type Message = {
-    type: MessageType;
-    data: any;
-};
-
-function createRoom(socket: WebSocket): string {
     let roomId = Math.random().toString(36).substring(2, 10);
     rooms.set(roomId, []);
 
-    socket.send(JSON.stringify({ type: MessageType.ROOM_CREATED, data: roomId }));
+    socket.send(JSON.stringify({ type: SocketMessageType.ROOM_CREATED, data: roomId }));
 
-    joinRoom(roomId, socket);
+    await joinRoom(roomId, socket);
 
     return roomId;
 }
 
-function joinRoom(roomId: string, socket: WebSocket) {
+async function joinRoom(roomId: string, socket: WebSocket) {
     let room = rooms.get(roomId);
     console.log(room?.length);
 
@@ -48,13 +26,13 @@ function joinRoom(roomId: string, socket: WebSocket) {
     }
 
     if (room.length == 2) {
-        socket.send(JSON.stringify({ type: MessageType.ERROR, data: 'Room is full' }));
+        socket.send(JSON.stringify({ type: SocketMessageType.ERROR, data: 'Room is full' }));
         return;
     }
 
     // notify all clients in the room of the new client, except the client itself
     room.forEach(client => {
-        client.send(JSON.stringify({ type: MessageType.JOIN_ROOM, data: roomId }));
+        client.send(JSON.stringify({ type: SocketMessageType.JOIN_ROOM, data: roomId }));
     });
     room.push(socket);
 
@@ -76,11 +54,23 @@ function joinRoom(roomId: string, socket: WebSocket) {
 
     // TODO: consider letting rooms get larger than 2 clients
     if (room.length == 2) {
-        room.forEach(client => {
+        // 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: MessageType.ROOM_READY, data: { isInitiator: client !== socket } }));
+            client.send(JSON.stringify({ type: SocketMessageType.ROOM_READY, data: { isInitiator: client !== socket, roomKey: { key: jsonWebKey } } }));
         });
     }
+
+    console.log("Room created:", roomId, room.length);
 }
 
 function deleteRoom(roomId: string) {
@@ -90,57 +80,50 @@ function deleteRoom(roomId: string) {
 export function confgiureWebsocketServer(ws: WebSocketServer) {
     ws.on('connection', socket => {
         // Handle messages from the client
-        socket.on('message', event => {
+        socket.on('message', async event => {
-            let message;
+            let message: SocketMessage | undefined = undefined;
 
             if (event instanceof Buffer) { // Assuming JSON is sent as a string
                 try {
-                    const jsonObject = JSON.parse(Buffer.from(event).toString());
+                    message = JSON.parse(Buffer.from(event).toString());
-                    // TODO: validate the message
-                    message = jsonObject as Message;
                 } catch (e) {
                     console.error("Error parsing JSON:", e);
                 }
             }
 
-            if (!message) {
+            if (message === undefined) {
                 console.log("Received non-JSON message:", event);
                 // If the message is not JSON, send an error message
-                socket.send(JSON.stringify({ type: MessageType.ERROR, data: 'Invalid message' }));
+                socket.send(JSON.stringify({ type: SocketMessageType.ERROR, data: 'Invalid message' }));
                 return;
             }
 
-            let { type } = message;
+            switch (message.type) {
-
+                case SocketMessageType.CREATE_ROOM:
-            // coerce type to a MessageType enum
-            type = type as MessageType;
-
-            switch (type) {
-                case MessageType.CREATE_ROOM:
                     // else, create a new room
-                    createRoom(socket);
+                    await createRoom(socket);
                     break;
-                case MessageType.JOIN_ROOM:
+                case SocketMessageType.JOIN_ROOM:
                     // if join message has a roomId, join the room
-                    if (!message.data) {
+                    if (!message.roomId) {
-                        socket.send(JSON.stringify({ type: MessageType.ERROR, data: 'Invalid message' }));
+                        socket.send(JSON.stringify({ type: SocketMessageType.ERROR, data: 'Invalid message' }));
                         return;
                     }
 
                     // if the user tries to join a room that doesnt exist, send an error message
-                    if (rooms.get(message.data) == undefined) {
+                    if (rooms.get(message.roomId) == undefined) {
-                        socket.send(JSON.stringify({ type: MessageType.ERROR, data: 'Invalid roomId' }));
+                        socket.send(JSON.stringify({ type: SocketMessageType.ERROR, data: 'Invalid roomId' }));
                         return;
                     }
 
-                    joinRoom(message.data, socket);
+                    await joinRoom(message.roomId, socket);
 
                     // the client is now in the room and the peer knows about it
-                    socket.send(JSON.stringify({ type: MessageType.ROOM_JOINED, data: null }));
+                    socket.send(JSON.stringify({ type: SocketMessageType.ROOM_JOINED, roomId: message.roomId }));
                     break;
-                case MessageType.OFFER:
+                case SocketMessageType.OFFER:
-                case MessageType.ANSWER:
+                case SocketMessageType.ANSWER:
-                case MessageType.ICE_CANDIDATE:
+                case SocketMessageType.ICE_CANDIDATE:
                     // relay these messages to the other peers in the room
                     const room = rooms.get(message.data.roomId);
 
@@ -153,8 +136,8 @@ export function confgiureWebsocketServer(ws: WebSocketServer) {
                     }
                     break;
                 default:
-                    console.warn(`Unknown message type: ${type}`);
+                    console.warn(`Unknown message type: ${message.type}`);
-                    socket.send(JSON.stringify({ type: MessageType.ERROR, data: 'Unknown message type' }));
+                    socket.send(JSON.stringify({ type: SocketMessageType.ERROR, data: 'Unknown message type' }));
                     break;
             }
         });

src/components/RTCMessage.svelte

@@ -1,13 +1,17 @@
 <script lang="ts">
     import { writable, type Writable } from "svelte/store";
-    import { room } from "../stores/roomStore";
+    import { room, connectionState } from "../stores/roomStore";
     import { connected } from "../stores/websocketStore";
     import {
         isRTCConnected,
         dataChannelReady,
-        messages,
         peer,
+        keyExchangeDone,
     } from "../utils/webrtcUtil";
+    import { messages } from "../stores/messageStore";
+    import { WebRTCPacketType } from "../types/webrtc";
+    import { ConnectionState } from "../types/websocket";
+    import { MessageType } from "../types/message";
 
     let inputMessage: Writable<string> = writable("");
     let inputFile = writable(null);
@@ -23,15 +27,31 @@
             return;
         }
 
-        if ($inputFile != null && $inputFile[0] !== undefined) {
+        // if ($inputFile != null && $inputFile[0] !== undefined) {
-            $messages = [...$messages, `You: ${$inputFile[0].name}`];
+        //     $messages = [...$messages, `You: ${$inputFile[0].name}`];
-            $peer.send($inputFile[0]);
+        //     $peer.send($inputFile[0]);
-            $inputFile = null;
+        //     $inputFile = null;
-        }
+        // }
 
         if ($inputMessage) {
-            $messages = [...$messages, `You: ${$inputMessage}`];
+            // $messages = [...$messages, `You: ${$inputMessage}`];
-            $peer.send($inputMessage);
+            $messages = [
+                ...$messages,
+                {
+                    initiator: true,
+                    type: MessageType.TEXT,
+                    data: $inputMessage,
+                },
+            ];
+            $peer.send(
+                new TextEncoder().encode(
+                    JSON.stringify({
+                        type: MessageType.TEXT,
+                        data: $inputMessage,
+                    }),
+                ).buffer,
+                WebRTCPacketType.MESSAGE,
+            );
             $inputMessage = "";
         }
     }
@@ -41,15 +61,35 @@
     }
 </script>
 
-{#if $room !== null && $connected === true}
+<!-- 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 !$isRTCConnected}
         <p>Waiting for peer to connect...</p>
     {:else if !$dataChannelReady}
         <p>Establishing data channel...</p>
+    {:else if !$keyExchangeDone}
+        <p>Establishing a secure connection with the peer...</p>
     {:else}
-        <div class="flex-grow overflow-y-auto mb-4 p-2 bg-gray-800 rounded">
+        <div
+            class="flex-grow overflow-y-auto mb-4 p-2 bg-gray-800 rounded break-all"
+        >
             {#each $messages as msg}
-                <p>{msg}</p>
+                <div>
+                    <div class="w-fit h-max">
+                        {#if msg.initiator}
+                            You:
+                        {:else}
+                            Peer:
+                        {/if}
+                    </div>
+                    <span>
+                        {#if msg.type === MessageType.TEXT}
+                            {msg.data}
+                        {:else}
+                            Unknown message type: {msg.type}
+                        {/if}
+                    </span>
+                </div>
             {/each}
         </div>
         <input

src/lib/webrtc.ts

@@ -1,13 +1,7 @@
 import { get } from 'svelte/store';
 import { ws } from '../stores/websocketStore';
-
+import { roomKey } from '../utils/webrtcUtil';
-interface WebRTCPeerCallbacks {
+import { WebRTCPacketType, type KeyStore, type WebRTCPeerCallbacks } from '../types/webrtc';
-    onConnected: () => void;
-    onMessage: (message: string | ArrayBuffer) => void;
-    onDataChannelOpen: () => void;
-    onNegotiationNeeded: () => void;
-    onError: (error: any) => void;
-}
 
 export class WebRTCPeer {
     private peer: RTCPeerConnection | null = null;
@@ -15,10 +9,16 @@ export class WebRTCPeer {
     private isInitiator: boolean;
     private roomId: string;
     private callbacks: WebRTCPeerCallbacks;
+    private keys: KeyStore = {
+        localKeys: null,
+        peersPublicKey: null,
+    };
 
     private iceServers = [
-        { urls: 'stun:stun.l.google.com:19302' },
+        { urls: "stun:stun.l.google.com:19302" },
-        { urls: 'stun:stun1.l.google.com:19302' },
+        { urls: "stun:stun.l.google.com:5349" },
+        { urls: "stun:stun1.l.google.com:3478" },
+        { urls: "stun:stun1.l.google.com:5349" },
     ];
 
     constructor(roomId: string, isInitiator: boolean, callbacks: WebRTCPeerCallbacks) {
@@ -80,14 +80,90 @@ export class WebRTCPeer {
     }
 
     private setupDataChannelEvents(channel: RTCDataChannel) {
-        channel.onopen = () => {
+        channel.binaryType = "arraybuffer";
+
+        channel.onopen = async () => {
             console.log('data channel open');
             this.callbacks.onDataChannelOpen();
+
+            try {
+                if (this.isInitiator) {
+                    await this.startKeyExchange();
+                }
+            } catch (e) {
+                console.error("Error starting key exchange:", e);
+                this.callbacks.onError(e);
+            }
         };
 
-        channel.onmessage = (event) => {
+        channel.onmessage = async (event: MessageEvent<ArrayBuffer>) => {
             console.log('data channel message:', event.data);
-            this.callbacks.onMessage(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);
+            const encrypted = (data[0] >> 7) & 1;
+            const type = data[0] & 0b01111111;
+            data = data.slice(1);
+
+            console.log("parsed data", data, encrypted, type);
+
+            if (type == WebRTCPacketType.KEY_EXCHANGE) {
+                if (this.keys.peersPublicKey) {
+                    console.error("Key exchange already done");
+                    return;
+                }
+
+                console.log("Received key exchange", data.buffer);
+
+                // let textDecoder = new TextDecoder();
+                // let dataString = textDecoder.decode(data.buffer);
+
+                // console.log("Received key exchange", dataString);
+
+                // let json = JSON.parse(dataString);
+
+                let unwrappingKey = get(roomKey);
+                if (!unwrappingKey.key) throw new Error("Room key not set");
+
+                this.keys.peersPublicKey = await window.crypto.subtle.unwrapKey(
+                    "jwk",
+                    data,
+                    unwrappingKey.key,
+                    {
+                        name: "AES-KW",
+                        length: 256,
+                    },
+                    {
+                        name: "RSA-OAEP",
+                        modulusLength: 4096,
+                        publicExponent: new Uint8Array([1, 0, 1]),
+                        hash: "SHA-256",
+                    },
+                    true,
+                    ["encrypt"],
+                );
+
+                // if our keys are not generated, start the reponding side of the key exchange
+                if (!this.keys.localKeys) {
+                    await this.startKeyExchange();
+                }
+
+                // by this point, both peers should have exchanged their keys
+                this.callbacks.onKeyExchangeDone();
+                return;
+            }
+
+            if (encrypted) {
+                data = new Uint8Array(await this.decrypt(data.buffer));
+            }
+
+            let message = {
+                type: type as WebRTCPacketType,
+                data: data.buffer,
+            };
+
+            this.callbacks.onMessage(message);
         };
 
         channel.onclose = () => {
@@ -105,8 +181,11 @@ export class WebRTCPeer {
         if (!this.peer) throw new Error('Peer not initialized');
 
         try {
-            const offer = await this.peer.createOffer();
+            const offer = await this.peer.createOffer()
-            await this.peer.setLocalDescription(offer);
+
+            console.log("Sending offer", offer);
+
+            await this.peer.setLocalDescription(offer)
 
             get(ws).send(JSON.stringify({
                 type: 'offer',
@@ -115,10 +194,9 @@ export class WebRTCPeer {
                     sdp: offer,
                 },
             }));
-
         } catch (error) {
-            console.error('Error creating offer:', error);
+            console.info('Error creating offer:', error);
-            this.callbacks.onError(error);
+            // should trigger re-negotiation
         }
     }
 
@@ -141,6 +219,8 @@ export class WebRTCPeer {
             const answer = await this.peer.createAnswer();
             await this.peer.setLocalDescription(answer);
 
+            console.log("Sending answer", answer);
+
             get(ws).send(JSON.stringify({
                 type: 'answer',
                 data: {
@@ -166,9 +246,112 @@ export class WebRTCPeer {
         }
     }
 
-    public send(data: string | ArrayBuffer) {
+    private async generateKeyPair() {
+        console.log("Generating key pair");
+        const keyPair = await window.crypto.subtle.generateKey(
+            {
+                name: "RSA-OAEP",
+                modulusLength: 4096,
+                publicExponent: new Uint8Array([1, 0, 1]),
+                hash: "SHA-256",
+            },
+            true,
+            ["encrypt", "decrypt"],
+        );
+
+        if (keyPair instanceof CryptoKey) {
+            throw new Error("Key pair not generated");
+        }
+
+        this.keys.localKeys = keyPair;
+    }
+
+    private async startKeyExchange() {
+        console.log("Starting key exchange");
+        await this.generateKeyPair();
+        if (!this.keys.localKeys) throw new Error("Key pair not generated");
+
+        let wrappingKey = get(roomKey);
+        if (!wrappingKey.key) throw new Error("Room key not set");
+
+
+        console.log("wrapping key", this.keys.localKeys.publicKey, wrappingKey.key);
+        const exported = await window.crypto.subtle.wrapKey(
+            "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);
+    }
+
+    private async encrypt(data: 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",
+            },
+            this.keys.peersPublicKey,
+            data,
+        );
+    }
+
+    private async decrypt(data: ArrayBuffer): Promise<ArrayBuffer> {
+        if (!this.keys.localKeys) throw new Error("Local keypair not generated");
+
+        return await window.crypto.subtle.decrypt(
+            {
+                name: "RSA-OAEP",
+            },
+            this.keys.localKeys.privateKey,
+            data,
+        );
+    }
+
+    public async send(data: ArrayBuffer, type: WebRTCPacketType) {
+        console.log("Sending message of type", type, "with data", data);
+
         if (!this.dataChannel || this.dataChannel.readyState !== 'open') throw new Error('Data channel not initialized');
-        this.dataChannel.send(data);
+
+        console.log(this.keys)
+        let header = (type & 0x7F);
+
+        // the key exchange is done, encrypt the message
+        if (this.keys.peersPublicKey && type != WebRTCPacketType.KEY_EXCHANGE) {
+            console.log("Sending encrypted message", data);
+
+            let encryptedData = await this.encrypt(data);
+
+            console.log("Encrypted data", encryptedData);
+
+            header |= 1 << 7;
+
+            let buf = new Uint8Array(encryptedData.byteLength + 1);
+            buf[0] = header;
+            buf.subarray(1).set(new Uint8Array(encryptedData));
+
+            this.dataChannel.send(buf.buffer);
+        } else {
+            console.log("Sending unencrypted message", data);
+            // the key exchange is not done yet, send the message unencrypted
+
+            let buf = new Uint8Array(data.byteLength + 1);
+            buf[0] = header;
+            buf.subarray(1).set(new Uint8Array(data));
+
+            this.dataChannel.send(buf.buffer);
+        }
     }
 
     public close() {

src/routes/+page.svelte

@@ -1,17 +1,20 @@
 <script lang="ts">
     import { ws, connected } from "../stores/websocketStore";
-    import { room } from "../stores/roomStore";
+    import { room, connectionState } from "../stores/roomStore";
     import { browser } from "$app/environment";
     import { peer, handleMessage } from "../utils/webrtcUtil";
     import { onDestroy, onMount } from "svelte";
     import RtcMessage from "../components/RTCMessage.svelte";
+    import { ConnectionState } from "../types/websocket";
 
     onMount(async () => {
+        $connectionState = ConnectionState.CONNECTING;
         $ws.addEventListener("message", handleMessage);
     });
 
     onDestroy(() => {
         if ($ws) {
+            $connectionState = ConnectionState.DISCONNECTED;
             $ws.removeEventListener("message", handleMessage);
         }
         if ($peer) {

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

@@ -1,10 +1,11 @@
 <script lang="ts">
     import { page } from "$app/state";
     import { onDestroy, onMount } from "svelte";
-    import { room } from "../../stores/roomStore";
+    import { room, connectionState } from "../../stores/roomStore";
     import { error, handleMessage, peer } from "../../utils/webrtcUtil";
-    import { ws } from "../../stores/websocketStore";
+    import { ws, connected } from "../../stores/websocketStore";
     import RtcMessage from "../../components/RTCMessage.svelte";
+    import { ConnectionState } from "../../types/websocket";
 
     if (!page.params.roomId) {
         throw new Error("Room ID not provided");
@@ -17,12 +18,14 @@
         $ws.addEventListener("message", handleMessage);
 
         $ws.onopen = () => {
-            $ws.send(JSON.stringify({ type: "join", data: $room }));
+            $connectionState = ConnectionState.CONNECTING;
+            $ws.send(JSON.stringify({ type: "join", roomId: $room }));
         };
     });
 
     onDestroy(() => {
         if ($ws) {
+            $connectionState = ConnectionState.DISCONNECTED;
             $ws.close();
         }
         if ($peer) {
@@ -34,6 +37,8 @@
 <div class="p-4">
     {#if $error}
         <p>Whoops! That room doesn't exist.</p>
+    {:else if !$connected || $connectionState === ConnectionState.CONNECTING}
+        <p>Connecting to server...</p>
     {:else}
         <RtcMessage />
     {/if}

src/stores/messageStore.ts

@@ -0,0 +1,4 @@
+import { writable, type Writable } from "svelte/store";
+import type { Message } from "../types/message";
+
+export let messages: Writable<Message[]> = writable([]);

src/stores/roomStore.ts

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

src/types/message.ts

@@ -0,0 +1,63 @@
+export enum MessageType {
+    // chat packets
+    TEXT = 0,
+    // user offers to send a file
+    FILE_OFFER = 1,
+    // user downloads a file offered by the peer
+    FILE_REQUEST = 2,
+
+    // file packets
+    FILE = 3,
+
+    ERROR = 255
+}
+
+export type Message =
+    | TextMessage
+    | FileOfferMessage
+    | FileRequestMessage
+    | FileMessage
+    | ErrorMessage;
+
+interface BaseMessage {
+    initiator: boolean;
+}
+
+// ----- chat packets -----
+export interface TextMessage extends BaseMessage {
+    type: MessageType.TEXT;
+    data: string;
+}
+
+export interface FileOfferMessage extends BaseMessage {
+    type: MessageType.FILE_OFFER;
+    data: {
+        fileName: string;
+        fileSize: number;
+        // randomly generated to identify the file so that multiple files with the same name can be uploaded
+        id: string;
+    };
+}
+
+export interface FileRequestMessage extends BaseMessage {
+    type: MessageType.FILE_REQUEST;
+    data: {
+        id: string;
+    };
+}
+
+// ----- file packets -----
+export interface FileMessage extends BaseMessage {
+    type: MessageType.FILE;
+    data: {
+        id: string;
+        fileName: string;
+        fileSize: number;
+        data: ArrayBuffer;
+    };
+}
+
+export interface ErrorMessage extends BaseMessage {
+    type: MessageType.ERROR;
+    data: string;
+}

src/types/webrtc.ts

@@ -0,0 +1,27 @@
+export interface WebRTCPeerCallbacks {
+    onConnected: () => void;
+    onMessage: (message: { type: WebRTCPacketType, data: ArrayBuffer }) => void;
+    onDataChannelOpen: () => void;
+    onKeyExchangeDone: () => void;
+    onNegotiationNeeded: () => void;
+    onError: (error: any) => void;
+}
+
+// max 7 bits for the type
+export enum WebRTCPacketType {
+    // all bits set
+    KEY_EXCHANGE = 127,
+    MESSAGE = 0,
+}
+
+export interface WebRTCPacket {
+    encrypted: boolean; // 1 bit
+    type: WebRTCPacketType; // 7 bits
+
+    data: ArrayBuffer;
+}
+
+export interface KeyStore {
+    localKeys: CryptoKeyPair | null;
+    peersPublicKey: CryptoKey | null;
+}

src/types/websocket.ts

@@ -0,0 +1,99 @@
+export enum ConnectionState {
+    CONNECTING,
+    CONNECTED,
+    DISCONNECTED,
+}
+
+export enum SocketMessageType {
+    // requests
+    CREATE_ROOM = 'create',
+    JOIN_ROOM = 'join',
+
+    // responses
+    ROOM_CREATED = 'created',
+    ROOM_JOINED = 'joined',
+    ROOM_READY = 'ready',
+
+    // webrtc
+    ICE_CANDIDATE = 'ice-candidate',
+    OFFER = 'offer',
+    ANSWER = 'answer',
+
+    ERROR = 'error',
+}
+
+type SocketMessageBase = {
+    type: SocketMessageType;
+};
+
+export interface SocketMessageCreateRoom extends SocketMessageBase {
+    type: SocketMessageType.CREATE_ROOM;
+}
+
+export interface SocketMessageJoinRoom extends SocketMessageBase {
+    type: SocketMessageType.JOIN_ROOM;
+    roomId: string;
+}
+
+export interface SocketMessageRoomCreated extends SocketMessageBase {
+    type: SocketMessageType.ROOM_CREATED;
+    data: {
+        roomId: string;
+    };
+}
+
+export interface SocketMessageRoomJoined extends SocketMessageBase {
+    type: SocketMessageType.ROOM_JOINED;
+    roomId: string;
+}
+
+export interface SocketMessageRoomReady extends SocketMessageBase {
+    type: SocketMessageType.ROOM_READY;
+    data: {
+        roomId: string;
+        isInitiator: boolean;
+        roomKey: {
+            key: JsonWebKey;
+        };
+    };
+}
+
+export interface SocketMessageIceCandidate extends SocketMessageBase {
+    type: SocketMessageType.ICE_CANDIDATE;
+    data: {
+        roomId: string;
+        candidate: RTCIceCandidate;
+    };
+}
+
+export interface SocketMessageOffer extends SocketMessageBase {
+    type: SocketMessageType.OFFER;
+    data: {
+        roomId: string;
+        sdp: RTCSessionDescription;
+    };
+}
+
+export interface SocketMessageAnswer extends SocketMessageBase {
+    type: SocketMessageType.ANSWER;
+    data: {
+        roomId: string;
+        sdp: RTCSessionDescription;
+    };
+}
+
+export interface SocketMessageError extends SocketMessageBase {
+    type: SocketMessageType.ERROR;
+    data: string;
+}
+
+export type SocketMessage =
+    | SocketMessageCreateRoom
+    | SocketMessageJoinRoom
+    | SocketMessageRoomCreated
+    | SocketMessageRoomJoined
+    | SocketMessageRoomReady
+    | SocketMessageIceCandidate
+    | SocketMessageOffer
+    | SocketMessageAnswer
+    | SocketMessageError;

src/utils/webrtcUtil.ts

@@ -1,14 +1,17 @@
-import { onDestroy, onMount } from "svelte";
 import { writable, get, type Writable } from "svelte/store";
 import { WebRTCPeer } from "$lib/webrtc";
-import { connected, ws } from "../stores/websocketStore";
+import { WebRTCPacketType } from "../types/webrtc";
-import { room } from "../stores/roomStore";
+import { room, connectionState } from "../stores/roomStore";
+import { ConnectionState } from "../types/websocket";
+import { messages } from "../stores/messageStore";
+import { MessageType, type Message } from "../types/message";
 
 export const error = writable(null);
 export let peer: Writable<WebRTCPeer | null> = writable(null);
-export let messages: Writable<string[]> = writable([]);
 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: () => {
@@ -16,35 +19,48 @@ const callbacks = {
         isRTCConnected.set(true);
     },
     //! TODO: come up with a more complex room system. This is largely for testing purposes
-    onMessage: (message: string | ArrayBuffer) => {
+    onMessage: (message: { type: WebRTCPacketType, data: ArrayBuffer }) => {
-        console.log("Received message:", message);
+        // onMessage: (message: string | ArrayBuffer) => {
-        if (typeof message === 'object' && message instanceof Blob) {
+        console.log("WebRTC Received message:", message);
-            // download the file
+        // if (typeof message === 'object' && message instanceof Blob) {
-            const url = URL.createObjectURL(message);
+        //     // download the file
-            const a = document.createElement('a');
+        //     const url = URL.createObjectURL(message);
-            a.href = url;
+        //     const a = document.createElement('a');
-            a.download = message.name;
+        //     a.href = url;
-            document.body.appendChild(a);
+        //     a.download = message.name;
-            a.click();
+        //     document.body.appendChild(a);
-            setTimeout(() => {
+        //     a.click();
-                document.body.removeChild(a);
+        //     setTimeout(() => {
-                window.URL.revokeObjectURL(url);
+        //         document.body.removeChild(a);
-            }, 100);
+        //         window.URL.revokeObjectURL(url);
-        }
+        //     }, 100);
+        // }
 
-        messages.set([...get(messages), `Peer: ${message}`]);
+        console.log("Received message:", message);
+
+        // TODO: fixup
+        if (message.type === WebRTCPacketType.MESSAGE) {
+            let textDecoder = new TextDecoder();
+            let json: Message = JSON.parse(textDecoder.decode(message.data));
+            json.initiator = false;
+            messages.set([...get(messages), json]);
+        }
     },
     onDataChannelOpen: () => {
         console.log("Data channel open");
         dataChannelReady.set(true);
     },
+    onKeyExchangeDone: async () => {
+        console.log("Key exchange done");
+        keyExchangeDone.set(true);
+    },
     onNegotiationNeeded: async () => {
         console.log("Negotiation needed");
         await get(peer)?.createOffer();
     },
     onError: (error: any) => {
         console.error("Error:", error);
-        messages.set([...get(messages), `Error: ${error}`]);
+        messages.set([...get(messages), { initiator: false, type: MessageType.ERROR, data: error }]);
     },
 };
 
@@ -54,10 +70,15 @@ export async function handleMessage(event: MessageEvent) {
 
     switch (message.type) {
         case "created":
+            connectionState.set(ConnectionState.CONNECTED);
             console.log("Room created:", message.data);
             room.set(message.data);
             return;
+        case "join":
+            console.log("new client joined room", message.data);
+            return;
         case "joined":
+            connectionState.set(ConnectionState.CONNECTED);
             console.log("Joined room:", message.data);
             return;
         case "error":
@@ -72,6 +93,24 @@ export async function handleMessage(event: MessageEvent) {
                 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,