(wip): fix chat crypto

packages/chat/src/context.tsx

@@ -35,6 +35,10 @@ const CHAT_SECRET_VERSION = 1;
 
 function bytesFromProtocol(value: unknown): Uint8Array {
   if (value instanceof Uint8Array) return value;
+  if (value instanceof ArrayBuffer) return new Uint8Array(value);
+  if (ArrayBuffer.isView(value)) {
+    return new Uint8Array(value.buffer, value.byteOffset, value.byteLength);
+  }
   if (Array.isArray(value)) return new Uint8Array(value);
   if (typeof value === "string") {
     const bin = atob(value);
@@ -42,6 +46,16 @@ function bytesFromProtocol(value: unknown): Uint8Array {
     for (let i = 0; i < bin.length; i++) out[i] = bin.charCodeAt(i);
     return out;
   }
+  if (typeof value === "object" && value !== null) {
+    const entries = Object.entries(value);
+    if (
+      entries.every(
+        ([key, item]) => /^\d+$/.test(key) && typeof item === "number",
+      )
+    ) {
+      return new Uint8Array(entries.map(([, item]) => item as number));
+    }
+  }
   throw new Error("expected protocol bytes");
 }
 
@@ -72,6 +86,16 @@ function updateMessageStateBySendTime<
   };
 }
 
+function assertProtocolSuccess(type: string, response: { type: string }) {
+  if (response.type.startsWith("Error")) {
+    throw new Error(`${type} failed: ${response.type}`);
+  }
+}
+
+function protocolBytes(bytes: Uint8Array): Uint8Array<ArrayBuffer> {
+  return new Uint8Array(bytes);
+}
+
 export default function Provider({ children }: { children: ReactNode }) {
   const { load } = useStorage();
   const { send, subscribePush } = useMTP();
@@ -125,6 +149,34 @@ export default function Provider({ children }: { children: ReactNode }) {
         const chatId = deriveChatId(ownUserId, userIdValue);
         const secretId = deriveChatSecretId(chatId);
 
+        async function forwardChatSecret(chatSecret: Uint8Array) {
+          const peerUser = await getUser(userIdValue);
+          const peerWrapped = await wrapChatSecret({
+            chatSecret,
+            recipientKemPublicKey: kemPublicKeyFromPublicKeyBundle(
+              peerUser.PublicKey,
+            ),
+            chatId,
+            secretId,
+            version: CHAT_SECRET_VERSION,
+          });
+
+          assertProtocolSuccess(
+            "ChatSecretForward",
+            await send("ChatSecretForward", {
+              ChatId: chatId,
+              SenderUserId: String(ownUserId),
+              RecipientUserId: String(userIdValue),
+              SecretId: secretId,
+              VersionNumber: CHAT_SECRET_VERSION,
+              EncryptedSecret: protocolBytes(peerWrapped.encryptedSecret),
+              KemCiphertext: protocolBytes(peerWrapped.kemCiphertext),
+              WrappingScheme: peerWrapped.wrappingScheme,
+              CreatedAt: Date.now(),
+            }),
+          );
+        }
+
         const existing = await send("GetChatSecret", {
           UserId: String(ownUserId),
           ChatId: chatId,
@@ -146,6 +198,8 @@ export default function Provider({ children }: { children: ReactNode }) {
           if (active) {
             setCurrentChatSecretState({ userId: userIdValue, value: secret });
           }
+
+          await forwardChatSecret(secret);
           return;
         }
 
@@ -158,39 +212,21 @@ export default function Provider({ children }: { children: ReactNode }) {
           version: CHAT_SECRET_VERSION,
         });
 
-        await send("SetChatSecret", {
-          UserId: String(ownUserId),
-          ChatId: chatId,
-          SecretId: secretId,
-          VersionNumber: CHAT_SECRET_VERSION,
-          EncryptedSecret: Array.from(ownWrapped.encryptedSecret),
-          KemCiphertext: Array.from(ownWrapped.kemCiphertext),
-          WrappingScheme: ownWrapped.wrappingScheme,
-          CreatedAt: Date.now(),
-        });
+        assertProtocolSuccess(
+          "SetChatSecret",
+          await send("SetChatSecret", {
+            UserId: String(ownUserId),
+            ChatId: chatId,
+            SecretId: secretId,
+            VersionNumber: CHAT_SECRET_VERSION,
+            EncryptedSecret: protocolBytes(ownWrapped.encryptedSecret),
+            KemCiphertext: protocolBytes(ownWrapped.kemCiphertext),
+            WrappingScheme: ownWrapped.wrappingScheme,
+            CreatedAt: Date.now(),
+          }),
+        );
 
-        const peerUser = await getUser(userIdValue);
-        const peerWrapped = await wrapChatSecret({
-          chatSecret: rawSecret,
-          recipientKemPublicKey: kemPublicKeyFromPublicKeyBundle(
-            peerUser.PublicKey,
-          ),
-          chatId,
-          secretId,
-          version: CHAT_SECRET_VERSION,
-        });
-
-        await send("ChatSecretForward", {
-          ChatId: chatId,
-          SenderUserId: String(ownUserId),
-          RecipientUserId: String(userIdValue),
-          SecretId: secretId,
-          VersionNumber: CHAT_SECRET_VERSION,
-          EncryptedSecret: Array.from(peerWrapped.encryptedSecret),
-          KemCiphertext: Array.from(peerWrapped.kemCiphertext),
-          WrappingScheme: peerWrapped.wrappingScheme,
-          CreatedAt: Date.now(),
-        });
+        await forwardChatSecret(rawSecret);
 
         if (active) {
           setCurrentChatSecretState({ userId: userIdValue, value: rawSecret });
@@ -317,6 +353,68 @@ export default function Provider({ children }: { children: ReactNode }) {
           const keyring = await load("mtp_keyring");
           const chatId = String(data.ChatId);
           const secretId = String(data.SecretId);
+          const senderUserId = Number(data.SenderUserId);
+          if (!Number.isSafeInteger(senderUserId) || senderUserId <= 0) {
+            throw new Error("Invalid chat secret sender id");
+          }
+
+          const existing = await send("GetChatSecret", {
+            UserId: String(ownUserId),
+            ChatId: chatId,
+            SecretId: secretId,
+          });
+
+          if (!existing.type.startsWith("Error") && senderUserId > ownUserId) {
+            const existingData = existing.data as Record<string, unknown>;
+            const existingSecret = await unwrapChatSecret({
+              encryptedSecret: bytesFromProtocol(existingData.EncryptedSecret),
+              kemCiphertext: bytesFromProtocol(existingData.KemCiphertext),
+              keyring,
+              chatId: String(existingData.ChatId),
+              secretId: String(existingData.SecretId),
+              version: Number(existingData.VersionNumber),
+              wrappingScheme: String(existingData.WrappingScheme),
+            });
+            const senderUser = await getUser(senderUserId);
+            const senderWrapped = await wrapChatSecret({
+              chatSecret: existingSecret,
+              recipientKemPublicKey: kemPublicKeyFromPublicKeyBundle(
+                senderUser.PublicKey,
+              ),
+              chatId,
+              secretId,
+              version: Number(existingData.VersionNumber),
+            });
+
+            assertProtocolSuccess(
+              "ChatSecretForward",
+              await send("ChatSecretForward", {
+                ChatId: chatId,
+                SenderUserId: String(ownUserId),
+                RecipientUserId: String(senderUserId),
+                SecretId: String(existingData.SecretId),
+                VersionNumber: Number(existingData.VersionNumber),
+                EncryptedSecret: protocolBytes(senderWrapped.encryptedSecret),
+                KemCiphertext: protocolBytes(senderWrapped.kemCiphertext),
+                WrappingScheme: senderWrapped.wrappingScheme,
+                CreatedAt: Date.now(),
+              }),
+            );
+
+            log(
+              2,
+              "chat",
+              "purple",
+              "Ignoring forwarded chat secret because local user owns the deterministic secret",
+              {
+                chatId,
+                ownUserId,
+                senderUserId,
+              },
+            );
+            return;
+          }
+
           const secret = await unwrapChatSecret({
             encryptedSecret: bytesFromProtocol(data.EncryptedSecret),
             kemCiphertext: bytesFromProtocol(data.KemCiphertext),
@@ -334,19 +432,25 @@ export default function Provider({ children }: { children: ReactNode }) {
             version: Number(data.VersionNumber),
           });
 
-          await send("SetChatSecret", {
-            UserId: String(ownUserId),
-            ChatId: chatId,
-            SecretId: secretId,
-            VersionNumber: Number(data.VersionNumber),
-            EncryptedSecret: Array.from(ownWrapped.encryptedSecret),
-            KemCiphertext: Array.from(ownWrapped.kemCiphertext),
-            WrappingScheme: ownWrapped.wrappingScheme,
-            CreatedAt: Date.now(),
-          });
+          assertProtocolSuccess(
+            "SetChatSecret",
+            await send("SetChatSecret", {
+              UserId: String(ownUserId),
+              ChatId: chatId,
+              SecretId: secretId,
+              VersionNumber: Number(data.VersionNumber),
+              EncryptedSecret: protocolBytes(ownWrapped.encryptedSecret),
+              KemCiphertext: protocolBytes(ownWrapped.kemCiphertext),
+              WrappingScheme: ownWrapped.wrappingScheme,
+              CreatedAt: Date.now(),
+            }),
+          );
 
           if (deriveChatId(ownUserId, userIdValue) === chatId) {
             setCurrentChatSecretState({ userId: userIdValue, value: secret });
+            void queryClient.invalidateQueries({
+              queryKey: ["chat-messages", String(userIdValue)],
+            });
           }
         })().catch((err) => {
           log(1, "chat", "red", "Failed to accept chat secret", err);
@@ -474,6 +578,7 @@ export default function Provider({ children }: { children: ReactNode }) {
   }, [
     addLiveMessage,
     currentChatSecret,
+    getUser,
     load,
     send,
     subscribePush,

packages/crypto/src/chatSecret.ts

@@ -1,6 +1,7 @@
-import { crypto } from "mtp";
+import { base64ToBytes, bytesToBase64, crypto } from "mtp";
 
 const textEncoder = new TextEncoder();
+const textDecoder = new TextDecoder();
 
 export const CHAT_SECRET_WRAPPING_SCHEME =
   "mtp-chat-secret-kem-chacha20poly1305-hkdf-sha256-v1";
@@ -73,11 +74,16 @@ export async function unwrapChatSecret(args: {
   wrappingScheme: string;
 }): Promise<Uint8Array> {
   if (args.wrappingScheme !== CHAT_SECRET_WRAPPING_SCHEME) {
-    throw new Error(`Unsupported chat secret wrapping scheme: ${args.wrappingScheme}`);
+    throw new Error(
+      `Unsupported chat secret wrapping scheme: ${args.wrappingScheme}`,
+    );
   }
 
   const ownKeys = crypto.keyringToKeys(args.keyring);
-  const sharedSecret = crypto.decapsulate(ownKeys.kemSecretKey, args.kemCiphertext);
+  const sharedSecret = crypto.decapsulate(
+    ownKeys.kemSecretKey,
+    args.kemCiphertext,
+  );
 
   try {
     const wrappingKey = deriveWrappingKey({
@@ -103,7 +109,9 @@ export async function encryptChatText(
 ): Promise<string> {
   const key = deriveMessageKey(chatSecret);
   try {
-    return await crypto.encryptText(key, plaintext);
+    return bytesToBase64(
+      await crypto.encrypt(key, textEncoder.encode(plaintext)),
+    );
   } finally {
     key.fill(0);
   }
@@ -115,7 +123,9 @@ export async function decryptChatText(
 ): Promise<string> {
   const key = deriveMessageKey(chatSecret);
   try {
-    return await crypto.decryptText(key, ciphertext);
+    return textDecoder.decode(
+      await crypto.decrypt(key, base64ToBytes(ciphertext)),
+    );
   } finally {
     key.fill(0);
   }

packages/mtp/src/context.tsx

@@ -111,7 +111,9 @@ function validateResponse<T extends keyof Schemas & string>(
     return message as ProtocolMessage<T>;
   }
 
-  const schema = schemas[type]?.response;
+  const schema =
+    schemas[message.type as keyof Schemas & string]?.response ??
+    schemas[type]?.response;
   if (!schema) {
     return message as ProtocolMessage<T>;
   }
@@ -574,15 +576,7 @@ export function Provider(props: {
         subscribePush,
       });
     }
-  }, [
-    connected,
-    identified,
-    mtpUrl,
-    send,
-    subscribe,
-    subscribePush,
-    mtpRef,
-  ]);
+  }, [connected, identified, mtpUrl, send, subscribe, subscribePush, mtpRef]);
 
   const sendQueued: BoundSendFn = useMemo(
     () => async (type, data, options) => {

packages/shared/src/data.ts

@@ -18,6 +18,20 @@ const bytesLike = z.union([
   z.base64(),
 ]);
 
+const protocolBytes = z.instanceof(Uint8Array);
+
+const chatSecretResponse = z.object({
+  UserId: z.string(),
+  ChatId: z.string(),
+  SecretId: z.string(),
+  VersionNumber: z.number(),
+  EncryptedSecret: bytesLike,
+  KemCiphertext: bytesLike,
+  WrappingScheme: z.string(),
+  CreatedAt: z.number(),
+  UpdatedAt: z.number(),
+});
+
 export const Message = z.object({
   NotEncrypted: z.boolean().optional(),
   SentBySelf: z.boolean().optional(),
@@ -257,8 +271,8 @@ export const mtp = {
       ChatId: z.string(),
       SecretId: z.string(),
       VersionNumber: z.number(),
-      EncryptedSecret: bytesLike,
-      KemCiphertext: bytesLike,
+      EncryptedSecret: protocolBytes,
+      KemCiphertext: protocolBytes,
       WrappingScheme: z.string(),
       CreatedAt: z.number(),
     }),
@@ -270,21 +284,11 @@ export const mtp = {
       ChatId: z.string(),
       SecretId: z.string().optional(),
     }),
-    response: z.object({}),
+    response: chatSecretResponse,
   },
   ChatSecretResponse: {
     request: z.object({}).optional(),
-    response: z.object({
-      UserId: z.string(),
-      ChatId: z.string(),
-      SecretId: z.string(),
-      VersionNumber: z.number(),
-      EncryptedSecret: bytesLike,
-      KemCiphertext: bytesLike,
-      WrappingScheme: z.string(),
-      CreatedAt: z.number(),
-      UpdatedAt: z.number(),
-    }),
+    response: chatSecretResponse,
   },
   ChatSecretForward: {
     request: z.object({
@@ -293,8 +297,8 @@ export const mtp = {
       RecipientUserId: z.string(),
       SecretId: z.string(),
       VersionNumber: z.number(),
-      EncryptedSecret: bytesLike,
-      KemCiphertext: bytesLike,
+      EncryptedSecret: protocolBytes,
+      KemCiphertext: protocolBytes,
       WrappingScheme: z.string(),
       CreatedAt: z.number(),
     }),