Fix crypto and message duplication

frontend/src/auth/crypto.ts

@@ -8,30 +8,33 @@ import type { BackupBlob, UploadPublicKeyRequest } from "../core/types";
 let currentPublicKey: Uint8Array | null = null;
 let currentPrivateKey: Uint8Array | null = null;
 
-async function fetchPublicKey(): Promise<Uint8Array | null> {
+async function fetchPublicKey(token?: string): Promise<Uint8Array | null> {
-	const res = await fetch(`${API_BASE_URL}/crypto/public-key`, { method: "GET", headers: getAuthHeaders(true) });
+	const headers = token ? { 'Authorization': `Bearer ${token}` } : getAuthHeaders(true);
+	const res = await fetch(`${API_BASE_URL}/crypto/public-key`, { method: "GET", headers });
 	if (!res.ok) return null;
 	const data = await res.json();
 	if (!data?.publicKey) return null;
 	return ub64(data.publicKey);
 }
 
-async function uploadPublicKey(publicKey: Uint8Array): Promise<void> {
+async function uploadPublicKey(publicKey: Uint8Array, token?: string): Promise<void> {
 	const payload: UploadPublicKeyRequest = { 
 		publicKey: b64(publicKey) 
 	}
 
+	const headers = token ? { 'Authorization': `Bearer ${token}` } : getAuthHeaders(true);
 	await fetch(`${API_BASE_URL}/crypto/public-key`, {
 		method: "POST",
-		headers: getAuthHeaders(true),
+		headers,
 		body: JSON.stringify(payload)
 	});
 }
 
-async function fetchBackupBlob(): Promise<string | null> {
+async function fetchBackupBlob(token?: string): Promise<string | null> {
+	const headers = token ? { 'Authorization': `Bearer ${token}` } : getAuthHeaders(true);
 	const res = await fetch(`${API_BASE_URL}/crypto/backup`, { 
 		method: "GET",
-		headers: getAuthHeaders(true) 
+		headers 
 	});
 	if (res.ok) {
 		const response: BackupBlob = await res.json();
@@ -41,12 +44,13 @@ async function fetchBackupBlob(): Promise<string | null> {
 	}
 }
 
-async function uploadBackupBlob(blobJson: string): Promise<void> {
+async function uploadBackupBlob(blobJson: string, token?: string): Promise<void> {
 	const payload: BackupBlob = { blob: blobJson }
 
+	const headers = token ? { 'Authorization': `Bearer ${token}` } : getAuthHeaders(true);
 	await fetch(`${API_BASE_URL}/crypto/backup`, {
 		method: "POST",
-		headers: getAuthHeaders(true),
+		headers,
 		body: JSON.stringify(payload)
 	});
 }
@@ -61,16 +65,16 @@ export function getCurrentKeys(): UserKeyPairMemory | null {
 	return null;
 }
 
-export async function ensureKeysOnLogin(password: string): Promise<UserKeyPairMemory> {
+export async function ensureKeysOnLogin(password: string, token?: string): Promise<UserKeyPairMemory> {
 	// Try to restore from backup
-	const blobJson = await fetchBackupBlob();
+	const blobJson = await fetchBackupBlob(token);
 	if (blobJson) {
 		const blob = decodeBlob(blobJson);
 		const bundle = await decryptBackupWithPassword(password, blob);
 		currentPrivateKey = bundle.privateKey;
 		// Ensure public key exists on server; if not, derive from private (not possible via libsafely), so keep previous
 		// In our simple scheme, we rely on server having the public key or we reupload generated one on first setup
-		const serverPub = await fetchPublicKey();
+		const serverPub = await fetchPublicKey(token);
 		if (serverPub) {
 			currentPublicKey = serverPub;
 		} else {
@@ -78,9 +82,9 @@ export async function ensureKeysOnLogin(password: string): Promise<UserKeyPairMe
 			const pair = generateX25519KeyPair();
 			currentPublicKey = pair.publicKey;
 			currentPrivateKey = pair.privateKey;
-			await uploadPublicKey(currentPublicKey);
+			await uploadPublicKey(currentPublicKey, token);
 			const newBlob = await encryptBackupWithPassword(password, { version: 1, privateKey: currentPrivateKey });
-			await uploadBackupBlob(encodeBlob(newBlob));
+			await uploadBackupBlob(encodeBlob(newBlob), token);
 		}
 		return { publicKey: currentPublicKey!, privateKey: currentPrivateKey! };
 	}
@@ -89,9 +93,9 @@ export async function ensureKeysOnLogin(password: string): Promise<UserKeyPairMe
 	const pair = generateX25519KeyPair();
 	currentPublicKey = pair.publicKey;
 	currentPrivateKey = pair.privateKey;
-	await uploadPublicKey(currentPublicKey);
+	await uploadPublicKey(currentPublicKey, token);
 	const encBlob = await encryptBackupWithPassword(password, { version: 1, privateKey: currentPrivateKey });
-	await uploadBackupBlob(encodeBlob(encBlob));
+	await uploadBackupBlob(encodeBlob(encBlob), token);
 	return { publicKey: currentPublicKey, privateKey: currentPrivateKey };
 }
 

frontend/src/ui/hooks/useChat.ts

@@ -1,4 +1,4 @@
-import { useEffect, useCallback } from "react";
+import { useEffect, useCallback, useRef } from "react";
 import { useAppState } from "../state";
 import { request, websocket } from "../../websocket";
 import { API_BASE_URL } from "../../core/config";
@@ -19,9 +19,11 @@ export function useChat() {
         user 
     } = useAppState();
 
+    const messagesLoadedRef = useRef(false);
+
     // Load messages for the current chat
     const loadMessages = useCallback(async () => {
-        if (!user.authToken) return;
+        if (!user.authToken || messagesLoadedRef.current) return;
 
         try {
             const response = await fetch(`${API_BASE_URL}/get_messages`, {
@@ -38,6 +40,7 @@ export function useChat() {
                     });
                 }
             }
+            messagesLoadedRef.current = true;
         } catch (error) {
             console.error("Error loading messages:", error);
         }
@@ -99,12 +102,20 @@ export function useChat() {
         return () => {
             websocket.removeEventListener("message", handleWebSocketMessage);
         };
-    }, [addMessage, user.currentUser]);
+    }, [addMessage, updateMessage, removeMessage, user.currentUser]);
 
-    // Load messages when component mounts or chat changes
+    // Load messages only once when component mounts and user is authenticated
     useEffect(() => {
-        loadMessages();
+        if (user.authToken && !messagesLoadedRef.current) {
-    }, [loadMessages]);
+            loadMessages();
+        }
+    }, [user.authToken, loadMessages]);
+
+    // Reset messages loaded flag and clear messages when chat changes
+    useEffect(() => {
+        messagesLoadedRef.current = false;
+        clearMessages(); // Clear messages when switching chats
+    }, [chat.currentChat, clearMessages]);
 
     return {
         messages: chat.messages,

frontend/src/ui/screen/LoginScreen.tsx

@@ -55,13 +55,16 @@ export default function LoginScreen() {
                             
                             if (response.ok) {
                                 const data: LoginResponse = await response.json();
-                                // Store the JWT token
+                                // Store the JWT token first
-                                setUser(data.token, data.user)
+                                setUser(data.token, data.user);
+                                
+                                // Setup keys with the token we just received
                                 try {
-                                    await ensureKeysOnLogin(password);
+                                    await ensureKeysOnLogin(password, data.token);
                                 } catch (e) {
                                     console.error("Key setup failed:", e);
                                 }
+                                
                                 setCurrentPage("chat");
                                 // initializeProfile(); // Initialize profile after login
                             } else {

frontend/src/ui/state.ts

@@ -49,12 +49,20 @@ export const useAppState = create<AppState>((set, get) => ({
         dmUsers: [],
         activeDm: null
     },
-    addMessage: (message: Message) => set((state) => ({
+    addMessage: (message: Message) => set((state) => {
-        chat: {
+        // Check if message already exists to prevent duplicates
-            ...state.chat,
+        const messageExists = state.chat.messages.some(msg => msg.id === message.id);
-            messages: [...state.chat.messages, message]
+        if (messageExists) {
+            return state; // Return unchanged state if message already exists
         }
-    })),
+        
+        return {
+            chat: {
+                ...state.chat,
+                messages: [...state.chat.messages, message]
+            }
+        };
+    }),
     updateMessage: (messageId: number, updatedMessage: Partial<Message>) => set((state) => ({
         chat: {
             ...state.chat,