support encryption in core pkg

packages/core/jsr.json

@@ -1,7 +1,7 @@
 {
     "name": "@sinkr/core",
     "license": "MIT",
-    "version": "0.5.1",
+    "version": "0.6.0",
     "exports": {
         ".": "./src/index.ts",
         "./client": "./src/index.browser.ts",

packages/core/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@sinkr/core",
-  "version": "0.5.1",
+  "version": "0.6.0",
   "type": "module",
   "main": "src/index.ts",
   "exports": {
@@ -26,7 +26,8 @@
   },
   "prettier": "@sinkr/prettier-config",
   "dependencies": {
-    "@sinkr/validators": "npm:@jsr/[email protected]",
+    "@oslojs/encoding": "^1.1.0",
+    "@sinkr/validators": "npm:@jsr/[email protected]",
     "emittery": "^1.0.3",
     "undici": "^7.1.0",
     "zod": "^3.24.1"

packages/core/src/browser.ts

@@ -1,4 +1,5 @@
 import type { z } from "zod";
+import { decodeBase64url } from "@oslojs/encoding";
 import Emittery from "emittery";
 
 import type {
@@ -14,7 +15,7 @@ import type {
 import { ClientReceiveSchema } from "@sinkr/validators";
 
 import type { RealEventMap } from "./event-fallback";
-import type { UserInfo } from "./types";
+import type { EncryptionInput, UserInfo } from "./types";
 import {
   connectSymbol,
   countEventSymbol,
@@ -42,7 +43,10 @@ interface DefaultEvents {
 }
 
 type GenericMessageEvent<T> = Prettify<
-  Omit<z.infer<typeof ClientReceiveMessageSchema>, "message"> & { message: T }
+  Omit<z.infer<typeof ClientReceiveMessageSchema>, "message"> & {
+    message: T;
+    index?: number;
+  }
 >;
 
 type MappedEvents = {
@@ -65,15 +69,124 @@ function proxyRemoveEmit<T extends Emittery<any>>(emitter: T) {
   });
 }
 
+const SUPPORTED_HASHES = ["256", "384", "512"];
+
+async function importRSAJWK(key: JsonWebKey) {
+  if (key.kty !== "RSA" || !key.alg) {
+    throw new Error("Invalid key type");
+  }
+  const [_, algSubtype, hash] = key.alg.split("-");
+  if (algSubtype !== "OAEP" || !hash) {
+    throw new Error("Invalid key type");
+  }
+  if (!SUPPORTED_HASHES.includes(hash)) {
+    throw new Error("Unsupported hash");
+  }
+  if (!key.key_ops?.includes("encrypt")) {
+    throw new Error("Key does not support encryption");
+  }
+  return await crypto.subtle.importKey(
+    "jwk",
+    key,
+    {
+      name: "RSA-OAEP",
+      hash: `SHA-${hash}`,
+    },
+    true,
+    key.key_ops as KeyUsage[],
+  );
+}
+
+async function importAESGCMJWK(key: JsonWebKey) {
+  if (key.kty !== "oct") {
+    throw new Error("Invalid key type");
+  }
+  if (!key.key_ops?.includes("encrypt")) {
+    throw new Error("Key does not support encryption");
+  }
+  if (key.alg !== "A256GCM") {
+    throw new Error("Unsupported algorithm");
+  }
+  return await crypto.subtle.importKey(
+    "jwk",
+    key,
+    {
+      name: "AES-GCM",
+    },
+    true,
+    key.key_ops as KeyUsage[],
+  );
+}
+
+async function importUnknownJWK(key: JsonWebKey | CryptoKey) {
+  if (key instanceof CryptoKey) {
+    if (!key.usages.includes("encrypt")) {
+      throw new Error("Key does not support encryption");
+    }
+    switch (key.algorithm.name) {
+      case "RSA-OAEP":
+      case "AES-GCM":
+        return key;
+    }
+    throw new Error("Unsupported key type");
+  }
+  if (!key.kty) {
+    throw new Error("Invalid key type");
+  }
+  if (key.kty === "RSA") {
+    return await importRSAJWK(key);
+  }
+  if (key.kty === "oct") {
+    return await importAESGCMJWK(key);
+  }
+  throw new Error("Unsupported key type");
+}
+
+async function decrypt(ciphertext: string, key: CryptoKey) {
+  if (key.algorithm.name === "RSA-OAEP") {
+    const decoded = decodeBase64url(ciphertext);
+    return await crypto.subtle.decrypt(
+      {
+        name: "RSA-OAEP",
+      },
+      key,
+      decoded,
+    );
+  } else {
+    const [iv, msg] = ciphertext.split(".").map(decodeBase64url);
+    if (!iv || !msg) {
+      throw new Error("Invalid ciphertext");
+    }
+    return await crypto.subtle.decrypt(
+      {
+        name: "AES-GCM",
+        iv,
+      },
+      key,
+      msg,
+    );
+  }
+}
+
 class BrowserSinker extends Emittery<EventMapWithDefaults> {
   private ws: WebSocket | null = null;
 
   private channelCache = new Map<string, WeakRef<Channel | PresenceChannel>>();
 
+  private keyMap = new Map<string, CryptoKey>();
+
   constructor(private url: string) {
     super();
   }
 
+  /**
+   * Imports a decryption key for the client to decrypt messages.
+   */
+  async addDecryptionKey(key: EncryptionInput) {
+    const imported = await importUnknownJWK(key.key);
+    this.keyMap.set(key.keyId, imported);
+  }
+
   /**
    * Get a sinkr channel by name. If a presence channel is specified, alternative types and messages will be available.
    * @param channel The channel to listen to. Events will only fire if the current client is subscribed to the channel.
@@ -151,17 +264,52 @@ class BrowserSinker extends Emittery<EventMapWithDefaults> {
             break;
         }
       } else {
-        void this.emit(
-          data.data.event,
-          data.data as unknown as GenericMessageEvent<unknown>,
-        );
+        void this.emitMessage(data.data);
       }
     });
     this.ws.addEventListener("close", () => {
       void this.emit(disconnectSymbol);
       this.ws = null;
     });
   }
+
+  private async emitMessage<T>(
+    input: z.infer<typeof ClientReceiveMessageSchema>,
+  ) {
+    if (
+      input.message.type === "encrypted" ||
+      input.message.type === "encrypted-chunk"
+    ) {
+      const key = this.keyMap.get(input.message.keyId);
+      if (!key) {
+        return;
+      }
+      try {
+        const msg = await decrypt(input.message.ciphertext, key);
+        const msgDecoded = new TextDecoder().decode(msg);
+        const msgTransform = {
+          message: JSON.parse(msgDecoded) as T,
+          index: "index" in input.message ? input.message.index : undefined,
+        };
+        await this.emit(input.event, {
+          ...input,
+          message: msgTransform,
+        });
+      } catch (e) {
+        console.error(e);
+        return;
+      }
+    } else {
+      const msgTransform = {
+        message: input.message.message as T,
+        index: "index" in input.message ? input.message.index : undefined,
+      };
+      await this.emit(input.event, {
+        ...input,
+        message: msgTransform,
+      });
+    }
+  }
 }
 
 interface DefaultChannelEventMap {