(feat): more crypto migration

packages/chat/src/components/input.tsx

@@ -15,6 +15,7 @@ import { useChat } from "../context";
 import { useMTP } from "@tensamin/mtp";
 import { log, toast } from "@tensamin/shared/log";
 import { cn, useIsMobile } from "@tensamin/ui";
+import { encryptChatText } from "@tensamin/crypto/chatSecret";
 
 import { useSession } from "@tensamin/storage/session";
 import GifPicker from "./gifPicker";
@@ -28,8 +29,8 @@ export default function InputComponent({
 }) {
   const [invertEnterBehavior, setInvertEnterBehavior] = React.useState(false);
 
-  const { sendEncrypted } = useMTP();
-  const { addLiveMessage, userId, inputBoxRef } = useChat();
+  const { send } = useMTP();
+  const { addLiveMessage, chatSecret, userId, inputBoxRef } = useChat();
   const { load, save } = useStorage();
   const { moveUserIdToTop } = useSession();
   const gifPopoverRef = React.useRef<HTMLDivElement>(null);
@@ -64,6 +65,11 @@ export default function InputComponent({
       return;
     }
 
+    if (!chatSecret) {
+      toast("error", "Still getting chat secret...");
+      return;
+    }
+
     log(3, "chat", "purple", "Message send init, adding live message ...");
 
     const reference = addLiveMessage({
@@ -74,38 +80,33 @@ export default function InputComponent({
       MessageState: "awaiting",
     });
 
-    log(3, "chat", "purple", "Live message added, sending encrypted frame...");
+    log(3, "chat", "purple", "Live message added, encrypting...");
 
-    void load("user_id")
-      .then((ownUserId) =>
-        sendEncrypted(
-          "MessageSend",
-          {
-            Content: currentValue,
-            ReceiverId: userId,
-            SendTime: time,
-          },
-          {
-            senderUserId: String(ownUserId),
-            recipientUserId: String(userId),
-          },
-        ),
-      )
-      .catch((e) => {
-        log(0, "Chat", "red", "Failed to send encrypted message", e, {
-          ReceiverId: userId,
-          SendTime: time,
-        });
+    const encryptedContent = await encryptChatText(chatSecret, currentValue).catch(
+      (err) => {
+        toast("error", "Failed to encrypt message", String(err));
         reference.setFailed(true);
-        toast(
-          "error",
-          e instanceof Error && e.message.includes("public key")
-            ? "Recipient has no encryption public key available"
-            : "Failed to send encrypted message",
-        );
-      });
+      },
+    );
 
-    log(3, "chat", "purple", "Encrypted message send queued");
+    if (!encryptedContent) return;
+
+    log(3, "chat", "purple", "Content encrypted, sending message...");
+
+    send("MessageSend", {
+      Content: encryptedContent,
+      ReceiverId: userId,
+      SendTime: time,
+    }).catch((e) => {
+      log(0, "Chat", "red", "Failed to send message", e, {
+        ReceiverId: userId,
+        SendTime: time,
+      });
+      reference.setFailed(true);
+      toast("error", "Failed to send message");
+    });
+
+    log(3, "chat", "purple", "Message sent");
 
     moveUserIdToTop(userId);
 

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;

packages/chat/src/screen.tsx

@@ -79,6 +79,7 @@ export default function Screen() {
     liveMessages,
     clearLiveMessages,
     userId,
+    chatSecret,
     error,
     errorDescription,
   } = useChat();
@@ -105,12 +106,13 @@ export default function Screen() {
   const [value, setValue] = React.useState("");
 
   const hasValidChatUser = Number.isSafeInteger(userId) && userId > 0;
+  const hasChatSecret = chatSecret !== null;
 
   const messagesQuery = useInfiniteQuery({
-    queryKey: ["chat-messages", String(userId)],
+    queryKey: ["chat-messages", String(userId), hasChatSecret],
     initialPageParam: 0,
     queryFn: ({ pageParam }) => getMessages(PAGE_SIZE, Number(pageParam)),
-    enabled: hasValidChatUser,
+    enabled: hasValidChatUser && hasChatSecret,
     getNextPageParam: (lastPage, allPages) => {
       if (lastPage.length < PAGE_SIZE) {
         return undefined;