Updated stuff

src/app/ws.ts

@@ -17,6 +17,7 @@ export default class EnclaveWebSocket {
   onOpenQueue: Array<() => void>;
 
   private sendCounter = 0n;
+  private recvCounter = 0n;
   private sharedSecret: Uint8Array | null = null;
 
   private readonly handshakeReady: Promise<void>;
@@ -43,6 +44,9 @@ export default class EnclaveWebSocket {
       this.onOpenQueue.forEach((fun) => fun());
     };
 
+    // First binary message is the server's raw x25519 pubkey — this
+    // one-time listener handles the handshake, then hands off to
+    // the normal encrypted read loop.
     this.websocket.onmessage = (msg) => this.handleHandshakeMessage(msg);
   }
 
@@ -54,9 +58,13 @@ export default class EnclaveWebSocket {
       serverPubkey,
     );
 
+    // Respond with our own x25519 pubkey, in plaintext — this is the
+    // one message on either side that can't be encrypted yet, since
+    // the shared secret doesn't exist until both pubkeys are known.
     const myPublicKey = x25519.getPublicKey(this.myX25519PrivateKey);
     this.websocket.send(myPublicKey);
 
+    // From here on, every message is encrypted.
     this.websocket.onmessage = null;
     this.resolveHandshake();
   }
@@ -64,20 +72,27 @@ export default class EnclaveWebSocket {
   private nextSendNonce(): Uint8Array {
     const nonce = new Uint8Array(12);
     const view = new DataView(nonce.buffer);
-    view.setBigUint64(0, this.sendCounter, false); // Big-endian 64-bit counter
-    nonce[11] |= 0b1000_0000; // Client-to-server direction flag bit
+    view.setBigUint64(0, this.sendCounter, false); // big-endian, first 8 bytes
+    nonce[11] |= 0b1000_0000; // distinguishes client-send from server-send direction
     this.sendCounter += 1n;
     return nonce;
   }
 
-  public encrypt(plaintext: Uint8Array): Uint8Array {
+  private nextRecvNonce(): Uint8Array {
+    const nonce = new Uint8Array(12);
+    const view = new DataView(nonce.buffer);
+    view.setBigUint64(0, this.recvCounter, false);
+    this.recvCounter += 1n;
+    return nonce;
+  }
+
+  private encrypt(plaintext: Uint8Array): Uint8Array {
     if (!this.sharedSecret) throw new Error("Handshake not complete");
 
     const nonce = this.nextSendNonce();
     const cipher = chacha20poly1305(this.sharedSecret, nonce);
     const ciphertext = cipher.encrypt(plaintext);
 
-    // Prepend nonce (12 bytes) + ciphertext
     const out = new Uint8Array(nonce.length + ciphertext.length);
     out.set(nonce, 0);
     out.set(ciphertext, nonce.length);
@@ -86,12 +101,11 @@ export default class EnclaveWebSocket {
 
   public decrypt(data: Uint8Array): Uint8Array {
     if (!this.sharedSecret) throw new Error("Handshake not complete");
-    if (data.length < 12) {
+    if (data.length < 12)
       throw new Error("Message too short to contain a nonce");
-    }
 
-    const nonce = new Uint8Array(data.subarray(0, 12));
-    const ciphertext = new Uint8Array(data.subarray(12));
+    const nonce = data.slice(0, 12);
+    const ciphertext = data.slice(12);
 
     const cipher = chacha20poly1305(this.sharedSecret, nonce);
     return cipher.decrypt(ciphertext);
@@ -122,15 +136,8 @@ export default class EnclaveWebSocket {
     await this.handshakeReady;
 
     return new Promise((ok) => {
-      this.websocket.onmessage = async (msg) => {
-        let buffer: ArrayBuffer;
-        if (msg.data instanceof Blob) {
-          buffer = await msg.data.arrayBuffer();
-        } else {
-          buffer = msg.data as ArrayBuffer;
-        }
-
-        const encrypted = new Uint8Array(buffer);
+      this.websocket.onmessage = (msg) => {
+        const encrypted = new Uint8Array(msg.data as ArrayBuffer);
         const plaintext = this.decrypt(encrypted);
         const data = JSON.parse(new TextDecoder().decode(plaintext));
         ok(data);