(feat): more crypto migration

packages/chat/src/context.tsx

@@ -12,14 +12,39 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { useRouterState } from "@tanstack/react-router";
 import type { InfiniteData } from "@tanstack/react-query";
 import type { LiveMessage, RawMessage, RawMessages } from "./values";
+import {
+  deriveChatId,
+  deriveChatSecretId,
+  decryptChatText,
+  encryptChatText,
+  kemPublicKeyFromPublicKeyBundle,
+  ownKemPublicKeyFromKeyring,
+  randomChatSecret,
+  unwrapChatSecret,
+  wrapChatSecret,
+} from "@tensamin/crypto/chatSecret";
 import { useStorage } from "@tensamin/storage/context";
 import { useMTP } from "@tensamin/mtp";
 import { log } from "@tensamin/shared/log";
 import { useSession } from "@tensamin/storage/session";
+import { useUser } from "@tensamin/user/context";
 
 export const context = createContext<contextType | undefined>(undefined);
 
 const queryClient = new QueryClient();
+const CHAT_SECRET_VERSION = 1;
+
+function bytesFromProtocol(value: unknown): Uint8Array {
+  if (value instanceof Uint8Array) return value;
+  if (Array.isArray(value)) return new Uint8Array(value);
+  if (typeof value === "string") {
+    const bin = atob(value);
+    const out = new Uint8Array(bin.length);
+    for (let i = 0; i < bin.length; i++) out[i] = bin.charCodeAt(i);
+    return out;
+  }
+  throw new Error("expected protocol bytes");
+}
 
 function updateMessageStateBySendTime<
   T extends { SendTime: number; MessageState: RawMessage["MessageState"] },
@@ -50,14 +75,21 @@ function updateMessageStateBySendTime<
 
 export default function Provider({ children }: { children: ReactNode }) {
   const { load } = useStorage();
-  const { send, subscribePush, subscribeEncrypted, decryptEncryptedRecord } =
-    useMTP();
+  const { send, subscribePush } = useMTP();
+  const { get: getUser } = useUser();
   const { moveUserIdToTop } = useSession();
 
-  const [error] = useState("");
-  const [errorDescription] = useState("");
+  const [error, setError] = useState("");
+  const [errorDescription, setErrorDescription] = useState("");
 
   const [liveMessagesState, setLiveMessagesState] = useState<LiveMessage[]>([]);
+  const [currentChatSecretState, setCurrentChatSecretState] = useState<{
+    userId: number;
+    value: Uint8Array | null;
+  }>({
+    userId: 0,
+    value: null,
+  });
 
   const inputBoxRef = useRef<HTMLDivElement>(null);
 
@@ -71,40 +103,132 @@ export default function Provider({ children }: { children: ReactNode }) {
     return Number(rawId ?? 0);
   }, [locationSearch]);
 
+  const currentChatSecret = useMemo(() => {
+    if (currentChatSecretState.userId !== userIdValue) {
+      return null;
+    }
+
+    return currentChatSecretState.value;
+  }, [currentChatSecretState, userIdValue]);
+
+  useEffect(() => {
+    if (!userIdValue) return;
+
+    let active = true;
+
+    void (async () => {
+      try {
+        setError("");
+        setErrorDescription("");
+
+        const ownUserId = Number(await load("user_id"));
+        const keyring = await load("mtp_keyring");
+        const chatId = deriveChatId(ownUserId, userIdValue);
+        const secretId = deriveChatSecretId(chatId);
+
+        const existing = await send("GetChatSecret", {
+          UserId: String(ownUserId),
+          ChatId: chatId,
+          SecretId: secretId,
+        });
+
+        if (!existing.type.startsWith("Error")) {
+          const data = existing.data as Record<string, unknown>;
+          const secret = await unwrapChatSecret({
+            encryptedSecret: bytesFromProtocol(data.EncryptedSecret),
+            kemCiphertext: bytesFromProtocol(data.KemCiphertext),
+            keyring,
+            chatId: String(data.ChatId),
+            secretId: String(data.SecretId),
+            version: Number(data.VersionNumber),
+            wrappingScheme: String(data.WrappingScheme),
+          });
+
+          if (active) {
+            setCurrentChatSecretState({ userId: userIdValue, value: secret });
+          }
+          return;
+        }
+
+        const rawSecret = randomChatSecret();
+        const ownWrapped = await wrapChatSecret({
+          chatSecret: rawSecret,
+          recipientKemPublicKey: ownKemPublicKeyFromKeyring(keyring),
+          chatId,
+          secretId,
+          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(),
+        });
+
+        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(),
+        });
+
+        if (active) {
+          setCurrentChatSecretState({ userId: userIdValue, value: rawSecret });
+        }
+      } catch (err) {
+        log(1, "chat", "red", "Failed to initialize chat secret", err);
+        if (active) {
+          setError(err instanceof Error ? err.name : "Unknown Error");
+          setErrorDescription(err instanceof Error ? err.message : String(err));
+          setCurrentChatSecretState({ userId: userIdValue, value: null });
+        }
+      }
+    })();
+
+    return () => {
+      active = false;
+    };
+  }, [getUser, load, send, userIdValue]);
+
   const getMessages = useCallback(
     async (amount: number, offset: number) => {
-      const response = await send("EncryptedMessagesGet", {
-        Limit: amount,
-        SenderUserId: String(userIdValue),
-      });
-
-      if (response.type.startsWith("error")) {
-        throw new Error(response.type);
+      if (!currentChatSecret) {
+        return [];
       }
 
-      const encryptedMessages = (response.data.Messages ?? []) as Record<
-        string,
-        unknown
-      >[];
-      const decrypted = await Promise.all(
-        encryptedMessages.slice(offset, offset + amount).map(async (record) => {
-          const inner = await decryptEncryptedRecord(record);
-          const data = inner.data as Record<string, unknown>;
-          return {
-            NotEncrypted: false,
-            SendTime: Number(
-              data.SendTime ?? record.CreatedAt ?? record.createdAt,
-            ),
-            Content: String(data.Content ?? ""),
-            SentBySelf:
-              String(record.SenderUserId ?? record.senderUserId ?? "") ===
-              String(await load("user_id")),
-            MessageState: "received" as RawMessage["MessageState"],
-          };
-        }),
-      );
+      const messages = await send("MessagesGet", {
+        Amount: amount,
+        Offset: offset,
+        UserId: userIdValue,
+      });
 
-      const sorted = [...decrypted].sort((a, b) => a.SendTime - b.SendTime);
+      if (messages.type.startsWith("error")) {
+        throw new Error(messages.type);
+      }
+
+      const rawMessages = messages.data.Messages;
+      const sorted = [...rawMessages].sort((a, b) => a.SendTime - b.SendTime);
 
       if (sorted.length > 0) {
         const fetchedSendTimes = new Set(sorted.map((item) => item.SendTime));
@@ -118,9 +242,20 @@ export default function Provider({ children }: { children: ReactNode }) {
         });
       }
 
-      return sorted;
+      return await Promise.all(
+        sorted.map(async (message) => {
+          try {
+            return {
+              ...message,
+              Content: await decryptChatText(currentChatSecret, message.Content),
+            };
+          } catch {
+            return message;
+          }
+        }),
+      );
     },
-    [send, userIdValue, decryptEncryptedRecord, load],
+    [currentChatSecret, send, userIdValue],
   );
 
   const addLiveMessage = useCallback(
@@ -161,33 +296,75 @@ export default function Provider({ children }: { children: ReactNode }) {
     setLiveMessagesState([]);
   }, []);
 
-  useEffect(() => {
-    return subscribeEncrypted("MessageSend", (data) => {
-      const rawData = data as unknown as {
-        Content?: unknown;
-        ReceiverId?: unknown;
-        SendTime?: unknown;
-      };
-      const sendTime = Number(rawData.SendTime);
-
-      if (!Number.isFinite(sendTime) || typeof rawData.Content !== "string") {
-        log(3, "chat", "yellow", "Ignoring invalid encrypted live message");
-        return;
-      }
-
-      addLiveMessage({
-        NotEncrypted: false,
-        SendTime: sendTime,
-        Content: rawData.Content,
-        SentBySelf: false,
-        MessageState: "received",
-      });
-    });
-  }, [addLiveMessage, subscribeEncrypted]);
-
   // Get live updates for message states
   useEffect(() => {
     return subscribePush((message) => {
+      if (message.type === "ChatSecretForward") {
+        const data = message.data as Record<string, unknown>;
+        void (async () => {
+          const ownUserId = Number(await load("user_id"));
+          if (String(data.RecipientUserId) !== String(ownUserId)) return;
+
+          const keyring = await load("mtp_keyring");
+          const chatId = String(data.ChatId);
+          const secretId = String(data.SecretId);
+          const secret = await unwrapChatSecret({
+            encryptedSecret: bytesFromProtocol(data.EncryptedSecret),
+            kemCiphertext: bytesFromProtocol(data.KemCiphertext),
+            keyring,
+            chatId,
+            secretId,
+            version: Number(data.VersionNumber),
+            wrappingScheme: String(data.WrappingScheme),
+          });
+          const ownWrapped = await wrapChatSecret({
+            chatSecret: secret,
+            recipientKemPublicKey: ownKemPublicKeyFromKeyring(keyring),
+            chatId,
+            secretId,
+            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(),
+          });
+
+          if (deriveChatId(ownUserId, userIdValue) === chatId) {
+            setCurrentChatSecretState({ userId: userIdValue, value: secret });
+          }
+        })().catch((err) => {
+          log(1, "chat", "red", "Failed to accept chat secret", err);
+        });
+        return;
+      }
+
+      if (message.type === "MessageLive") {
+        const data = message.data as {
+          Message?: RawMessage;
+          SenderId?: number;
+        };
+        if (!data.Message || !currentChatSecret) return;
+
+        void decryptChatText(currentChatSecret, data.Message.Content)
+          .catch(() => data.Message?.Content ?? "")
+          .then((content) => {
+            if (!data.Message) return;
+            addLiveMessage({
+              ...data.Message,
+              Content: content,
+              SentBySelf: false,
+            });
+          });
+        return;
+      }
+
       if (message.type !== "MessageState") {
         return;
       }
@@ -243,6 +420,7 @@ export default function Provider({ children }: { children: ReactNode }) {
       const queryKey = [
         "chat-messages",
         String(userIdValue),
+        currentChatSecret !== null,
       ] as const;
       queryClient.setQueryData<InfiniteData<RawMessages>>(
         queryKey,
@@ -278,7 +456,7 @@ export default function Provider({ children }: { children: ReactNode }) {
         },
       );
     });
-  }, [subscribePush, userIdValue]);
+  }, [addLiveMessage, currentChatSecret, load, send, subscribePush, userIdValue]);
 
   return (
     <QueryClientProvider client={queryClient}>
@@ -288,7 +466,7 @@ export default function Provider({ children }: { children: ReactNode }) {
           liveMessages: () => liveMessagesState,
           addLiveMessage,
           clearLiveMessages,
-          sharedSecret: "",
+          chatSecret: currentChatSecret,
           userId: userIdValue,
           inputBoxRef,
           error,
@@ -308,7 +486,7 @@ type contextType = {
     setFailed: (failed: boolean) => void;
   };
   clearLiveMessages: () => void;
-  sharedSecret: string;
+  chatSecret: Uint8Array | null;
   userId: number;
   inputBoxRef: React.RefObject<HTMLDivElement | null>;
   error: string;