Improve chat encryption, delivery, and session handling

Adds robust session serialization/deserialization for ratchet state, session cleanup for corrupted data, and improved error handling in chat encryption hooks. Enhances chat delivery to support both local and remote recipients, ensures conversation creation, and adds timeouts to key fetches. Also introduces utilities for deleting and clearing encrypted session data from storage.
2026-01-27 20:36:03 -08:00
parent 8c3dffaa4c
commit e92c747a5b
7 changed files with 328 additions and 29 deletions
@@ -12,14 +12,26 @@ export async function GET(request: NextRequest) {

    const handle = searchParams.get('handle');

-    // Helper to fetch and check
+    // Helper to fetch and check with timeout
    const tryFetch = async (url: string) => {
        console.log(`[Proxy] Fetching keys from: ${url}`);
-        const res = await fetch(url, { headers: { 'Accept': 'application/json' } });
-        if (res.ok) return await res.json();
-        const text = await res.text();
-        console.warn(`[Proxy] Fetch failed for ${url} (${res.status}): ${text}`);
-        return null;
+        try {
+            const controller = new AbortController();
+            const id = setTimeout(() => controller.abort(), 5000); // 5s timeout
+            const res = await fetch(url, {
+                headers: { 'Accept': 'application/json' },
+                signal: controller.signal
+            });
+            clearTimeout(id);
+
+            if (res.ok) return await res.json();
+            const text = await res.text();
+            console.warn(`[Proxy] Fetch failed for ${url} (${res.status}): ${text}`);
+            return null;
+        } catch (err: any) {
+            console.warn(`[Proxy] Fetch error for ${url}: ${err.name} - ${err.message}`);
+            return null;
+        }
    };

    // 1. Try Primary DID
@@ -1,12 +1,14 @@

 import { NextRequest, NextResponse } from 'next/server';
 import { db } from '@/db';
-import { chatInbox } from '@/db/schema';
+import { chatInbox, chatConversations, users } from '@/db/schema';
 import { requireSignedAction } from '@/lib/auth/verify-signature';
+import { eq, and } from 'drizzle-orm';
+import { signedFetch } from '@/lib/api/signed-fetch';

 /**
 * POST /api/chat/send
- * Deliver an encrypted envelope to a local user's inbox.
+ * Deliver an encrypted envelope to a local or remote user's inbox.
 */
 export async function POST(request: NextRequest) {
    try {
@@ -20,22 +22,130 @@ export async function POST(request: NextRequest) {
            return NextResponse.json({ error: 'Invalid action type' }, { status: 400 });
        }

-        const { recipientDid, recipientDeviceId, ciphertext } = data;
+        const { recipientDid, recipientDeviceId, ciphertext, nodeDomain, recipientHandle } = data;

        if (!recipientDid || !ciphertext) {
            return NextResponse.json({ error: 'Missing required delivery fields' }, { status: 400 });
        }

-        // 2. Insert into Inbox
-        await db.insert(chatInbox).values({
-            senderDid: user.did,
-            recipientDid,
-            recipientDeviceId: recipientDeviceId || null, // Null = broadcast/all? Or specific? V2.1 says per-device.
-            envelope: JSON.stringify(body),
-            expiresAt: new Date(Date.now() + 30 * 24 * 60 * 60 * 1000) // 30 days TTL
+        // 2. Check if recipient is local or remote
+        const recipientUser = await db.query.users.findFirst({
+            where: eq(users.did, recipientDid)
        });

-        return NextResponse.json({ success: true });
+        if (recipientUser) {
+            // LOCAL RECIPIENT - Store in local inbox
+            
+            // Ensure conversation exists for recipient
+            let conversation = await db.query.chatConversations.findFirst({
+                where: and(
+                    eq(chatConversations.participant1Id, recipientUser.id),
+                    eq(chatConversations.participant2Handle, user.handle)
+                )
+            });
+
+            if (!conversation) {
+                const [newConv] = await db.insert(chatConversations).values({
+                    participant1Id: recipientUser.id,
+                    participant2Handle: user.handle,
+                    lastMessageAt: new Date(),
+                    lastMessagePreview: '[Encrypted Message]'
+                }).returning();
+                conversation = newConv;
+            }
+
+            // Ensure conversation exists for sender
+            let senderConversation = await db.query.chatConversations.findFirst({
+                where: and(
+                    eq(chatConversations.participant1Id, user.id),
+                    eq(chatConversations.participant2Handle, recipientUser.handle)
+                )
+            });
+
+            if (!senderConversation) {
+                await db.insert(chatConversations).values({
+                    participant1Id: user.id,
+                    participant2Handle: recipientUser.handle,
+                    lastMessageAt: new Date(),
+                    lastMessagePreview: '[Encrypted Message]'
+                });
+            }
+
+            // Insert into local inbox
+            await db.insert(chatInbox).values({
+                senderDid: user.did,
+                recipientDid,
+                recipientDeviceId: recipientDeviceId || null,
+                envelope: JSON.stringify(body),
+                expiresAt: new Date(Date.now() + 30 * 24 * 60 * 60 * 1000)
+            });
+
+            return NextResponse.json({ success: true, local: true });
+
+        } else {
+            // REMOTE RECIPIENT - Forward to their node
+            
+            let targetNodeDomain = nodeDomain;
+            
+            if (!targetNodeDomain) {
+                // Try to extract from DID
+                if (recipientDid.startsWith('did:web:')) {
+                    targetNodeDomain = recipientDid.replace('did:web:', '');
+                } else {
+                    return NextResponse.json({ 
+                        error: 'Remote delivery requires node domain' 
+                    }, { status: 400 });
+                }
+            }
+
+            // Forward the signed envelope to the remote node
+            const remoteUrl = `https://${targetNodeDomain}/api/swarm/chat/inbox`;
+            
+            try {
+                const response = await fetch(remoteUrl, {
+                    method: 'POST',
+                    headers: { 'Content-Type': 'application/json' },
+                    body: JSON.stringify(body)
+                });
+
+                if (!response.ok) {
+                    const errorData = await response.json().catch(() => ({}));
+                    console.error('[Chat Send] Remote delivery failed:', errorData);
+                    return NextResponse.json({ 
+                        error: 'Remote delivery failed',
+                        details: errorData 
+                    }, { status: response.status });
+                }
+
+                // Create local conversation for sender if we have recipient handle
+                if (recipientHandle) {
+                    const existingConv = await db.query.chatConversations.findFirst({
+                        where: and(
+                            eq(chatConversations.participant1Id, user.id),
+                            eq(chatConversations.participant2Handle, recipientHandle)
+                        )
+                    });
+
+                    if (!existingConv) {
+                        await db.insert(chatConversations).values({
+                            participant1Id: user.id,
+                            participant2Handle: recipientHandle,
+                            lastMessageAt: new Date(),
+                            lastMessagePreview: '[Encrypted Message]'
+                        });
+                    }
+                }
+                
+                return NextResponse.json({ success: true, remote: true });
+
+            } catch (error: any) {
+                console.error('[Chat Send] Remote delivery error:', error);
+                return NextResponse.json({ 
+                    error: 'Failed to reach remote node',
+                    details: error.message 
+                }, { status: 500 });
+            }
+        }

    } catch (error: any) {
        console.error('Delivery failed:', error);
@@ -302,25 +302,32 @@ export default function ChatPage() {
            const data = await res.json();
            if (!data.user?.did) {
                alert('User not found or V2 not enabled.');
+                setSending(false);
                return;
            }

+            console.log('[Chat UI] Starting chat with:', data.user);
+
            // Send "Hello" to init session
            await sendMessage(data.user.did, '👋', data.user.nodeDomain, data.user.handle);

+            console.log('[Chat UI] Message sent, reloading conversations');
+
            setShowNewChat(false);
            setNewChatHandle('');
-            loadConversations(false);
+            await loadConversations(false);
            // Select the new conversation (we might need to find it)
            // For now just reload list.
        } catch (e: any) {
-            console.error('Start chat failed:', e);
+            console.error('[Chat UI] Start chat failed:', e);
            if (e.message.includes('Recipient keys not found')) {
                alert('This user has not set up secure chat yet. They need to log in to enable end-to-end encryption.');
            } else {
                alert('Failed to start chat: ' + e.message);
            }
-        } finally { setSending(false); }
+        } finally { 
+            setSending(false); 
+        }
    };

    const handleDeleteConversation = async (deleteFor: 'self' | 'both') => {
@@ -98,6 +98,17 @@ async function getItem(key: string): Promise<string | undefined> {
    });
 }

+async function deleteItem(key: string): Promise<void> {
+    if (!dbInstance) throw new Error('Database locked');
+    return new Promise((resolve, reject) => {
+        const tx = dbInstance!.transaction(STORE_NAME, 'readwrite');
+        const store = tx.objectStore(STORE_NAME);
+        const req = store.delete(key);
+        req.onsuccess = () => resolve();
+        req.onerror = () => reject(req.error);
+    });
+}
+
 // ----------------------------------------------------------------------------
 // 3. Encrypted Read/Write
 // ----------------------------------------------------------------------------
@@ -135,6 +146,41 @@ export async function loadEncrypted<T>(key: string): Promise<T | null> {
    }
 }

+/**
+ * Deletes an encrypted item from storage.
+ */
+export async function deleteEncrypted(key: string): Promise<void> {
+    await deleteItem(key);
+}
+
+/**
+ * Clears all session data (useful for recovery from corruption).
+ */
+export async function clearAllSessions(): Promise<void> {
+    if (!dbInstance) throw new Error('Database locked');
+    
+    return new Promise((resolve, reject) => {
+        const tx = dbInstance!.transaction(STORE_NAME, 'readwrite');
+        const store = tx.objectStore(STORE_NAME);
+        const req = store.openCursor();
+        
+        req.onsuccess = (event) => {
+            const cursor = (event.target as IDBRequest).result;
+            if (cursor) {
+                // Only delete session keys, not device keys
+                if (cursor.key.toString().startsWith('session:')) {
+                    cursor.delete();
+                }
+                cursor.continue();
+            } else {
+                resolve();
+            }
+        };
+        
+        req.onerror = () => reject(req.error);
+    });
+}
+
 // ----------------------------------------------------------------------------
 // 4. Specific Key Managers
 // ----------------------------------------------------------------------------
@@ -73,7 +73,7 @@ export async function importX25519PrivateKey(base64: string): Promise<CryptoKey>
        'pkcs8',
        binary,
        { name: 'X25519' },
-        false,
+        true, // Must be extractable for serialization
        ['deriveKey', 'deriveBits']
    );
 }
@@ -234,6 +234,9 @@ export function arrayBufferToBase64(buffer: ArrayBuffer): string {
 }

 export function base64ToArrayBuffer(base64: string): ArrayBuffer {
+    if (!base64 || typeof base64 !== 'string') {
+        throw new Error('Invalid base64 input: expected non-empty string');
+    }
    // Handle URL safe base64 if needed, but we assume standard
    const binary = atob(base64.replace(/-/g, '+').replace(/_/g, '/'));
    const bytes = new Uint8Array(binary.length);
@@ -13,6 +13,7 @@ import {
    encrypt as aeadEncrypt,
    decrypt as aeadDecrypt,
    importX25519PublicKey,
+    importX25519PrivateKey,
    exportKey,
    generateX25519KeyPair,
    base64ToArrayBuffer,
@@ -328,3 +329,62 @@ export async function ratchetDecrypt(

    return { plaintext, newState: state };
 }
+
+// ----------------------------------------------------------------------------
+// 6. Serialization Helpers (CRITICAL: CryptoKeys and Buffers don't JSON stringify)
+// ----------------------------------------------------------------------------
+
+export interface SerializedRatchetState {
+    dhPair: { pub: string, priv: string };
+    remoteDhPub: string;
+    rootKey: string;
+    chainKeySend: string;
+    chainKeyRecv: string;
+    ns: number;
+    nr: number;
+    pn: number;
+}
+
+export async function serializeRatchetState(state: RatchetState): Promise<SerializedRatchetState> {
+    return {
+        dhPair: {
+            pub: await exportKey(state.dhPair.publicKey),
+            priv: await exportKey(state.dhPair.privateKey)
+        },
+        remoteDhPub: await exportKey(state.remoteDhPub),
+        rootKey: arrayBufferToBase64(state.rootKey),
+        chainKeySend: arrayBufferToBase64(state.chainKeySend),
+        chainKeyRecv: arrayBufferToBase64(state.chainKeyRecv),
+        ns: state.ns,
+        nr: state.nr,
+        pn: state.pn
+    };
+}
+
+export async function deserializeRatchetState(data: SerializedRatchetState): Promise<RatchetState> {
+    // Validate integrity - check all required fields exist (but allow empty strings for buffers that can be empty)
+    if (!data || 
+        !data.rootKey || 
+        !data.dhPair || 
+        !data.dhPair.pub || 
+        !data.dhPair.priv ||
+        !data.remoteDhPub ||
+        data.chainKeySend === undefined || data.chainKeySend === null ||
+        data.chainKeyRecv === undefined || data.chainKeyRecv === null) {
+        throw new Error('Invalid serialized state: missing required fields');
+    }
+
+    return {
+        dhPair: {
+            publicKey: await importX25519PublicKey(data.dhPair.pub),
+            privateKey: await importX25519PrivateKey(data.dhPair.priv)
+        },
+        remoteDhPub: await importX25519PublicKey(data.remoteDhPub),
+        rootKey: base64ToArrayBuffer(data.rootKey),
+        chainKeySend: data.chainKeySend ? base64ToArrayBuffer(data.chainKeySend) : new Uint8Array(0).buffer,
+        chainKeyRecv: data.chainKeyRecv ? base64ToArrayBuffer(data.chainKeyRecv) : new Uint8Array(0).buffer,
+        ns: data.ns,
+        nr: data.nr,
+        pn: data.pn
+    };
+}
@@ -8,7 +8,9 @@ import {
  storeDeviceKeys,
  isStorageUnlocked,
  storeEncrypted,
-  loadEncrypted
+  loadEncrypted,
+  deleteEncrypted,
+  clearAllSessions
 } from '@/lib/crypto/chat-storage';
 import {
  generateX25519KeyPair,
@@ -24,7 +26,9 @@ import {
  ratchetDecrypt,
  x3dhSender,
  x3dhReceiver,
-  RatchetState
+  RatchetState,
+  serializeRatchetState,
+  deserializeRatchetState
 } from '@/lib/crypto/ratchet';
 import { useUserIdentity } from './useUserIdentity';
 import { v4 as uuidv4 } from 'uuid';
@@ -54,6 +58,18 @@ export function useChatEncryption() {
      setIsLocked(!unlocked);

      if (unlocked) {
+        // One-time cleanup of corrupted sessions (migration fix)
+        if (!sessionStorage.getItem('synapsis_sessions_cleaned_v2')) {
+          try {
+            console.log('[Chat] Running one-time session cleanup...');
+            await clearAllSessions();
+            sessionStorage.setItem('synapsis_sessions_cleaned_v2', 'true');
+            console.log('[Chat] Session cleanup complete');
+          } catch (e) {
+            console.error('[Chat] Session cleanup failed:', e);
+          }
+        }
+
        try {
          const keys = await loadDeviceKeys();
          if (keys) {
@@ -281,13 +297,21 @@ export function useChatEncryption() {
    let bundles: any[];
    try {
      console.log(`[Chat] Fetching keys via proxy: ${proxyUrl}`);
-      const res = await fetch(proxyUrl);
+      const controller = new AbortController();
+      const id = setTimeout(() => controller.abort(), 8000); // 8s timeout (proxy has 5s)
+
+      const res = await fetch(proxyUrl, {
+        signal: controller.signal
+      });
+      clearTimeout(id);
+
      if (!res.ok) {
        const data = await res.json();
        console.error(`[Chat] Bundle fetch failed (${res.status}):`, data);
        throw new Error(`Recipient keys not found (Status: ${res.status})`);
      }
      bundles = await res.json();
+      console.log(`[Chat] Fetched ${bundles.length} device bundle(s):`, bundles);
    } catch (err: any) {
      console.error(`[Chat] Network error fetching bundles:`, err);
      throw new Error(`Failed to resolve recipient keys: ${err.message}`);
@@ -300,6 +324,8 @@ export function useChatEncryption() {
    const localKeys = await loadDeviceKeys();
    if (!localKeys) throw new Error('Keys lost');

+    console.log(`[Chat] Processing ${bundles.length} recipient device(s)...`);
+
    // 2. Loop through all devices
    for (const bundle of bundles) {
      // IMPORTANT: Use the DID from the bundle! 
@@ -311,12 +337,24 @@ export function useChatEncryption() {

      let state = sessionsRef.current.get(sessionKey);
      if (!state) {
-        state = await loadEncrypted<RatchetState>(sessionKey) || undefined;
+        const stored = await loadEncrypted<any>(sessionKey);
+        if (stored) {
+          try {
+            // If stored is old/broken (missing rootKey string), this throws.
+            state = await deserializeRatchetState(stored);
+          } catch (e) {
+            console.warn('[Chat] Found corrupted session state, resetting:', e);
+            // Delete the corrupted session from storage
+            await deleteEncrypted(sessionKey);
+            state = undefined;
+          }
+        }
      }

      let headerData: any = null;

      if (!state) {
+        console.log(`[Chat] Initializing new session for device ${bundle.deviceId}...`);
        // X3DH Init
        const remoteIdentityKey = await importX25519PublicKey(bundle.identityKey);
        const remoteSignedPreKey = await importX25519PublicKey(bundle.signedPreKey.key);
@@ -339,10 +377,13 @@ export function useChatEncryption() {
        };
      }

+      console.log(`[Chat] Encrypting message for device ${bundle.deviceId}...`);
      const { ciphertext, newState } = await ratchetEncrypt(state, content);

      sessionsRef.current.set(sessionKey, newState);
-      await storeEncrypted(sessionKey, newState);
+      // Serialize before storing
+      const serialized = await serializeRatchetState(newState);
+      await storeEncrypted(sessionKey, serialized);

      // Payload
      const payload = {
@@ -357,18 +398,29 @@ export function useChatEncryption() {
      const fullData = {
        recipientDid,
        recipientDeviceId: bundle.deviceId,
-        ciphertext: JSON.stringify(payload)
+        ciphertext: JSON.stringify(payload),
+        nodeDomain: nodeDomain || undefined, // Include for remote routing
+        recipientHandle: recipientHandle || undefined // Include for conversation creation
      };

      const action = await signUserAction('chat.deliver', fullData);

-      await fetch('/api/chat/send', {
+      const sendRes = await fetch('/api/chat/send', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(action)
      });
+
+      if (!sendRes.ok) {
+        const errorData = await sendRes.json().catch(() => ({ error: 'Unknown error' }));
+        console.error('[Chat] Send failed:', errorData);
+        throw new Error(`Failed to send message: ${errorData.error || sendRes.statusText}`);
+      }
+
+      console.log(`[Chat] Message sent successfully to device ${bundle.deviceId}`);
    }

+    console.log('[Chat] All messages sent successfully');
  }, [isReady, identity, signUserAction]);

  /**
@@ -397,7 +449,16 @@ export function useChatEncryption() {
      const sessionKey = `session:${senderDid}:${senderDeviceId}`;
      let state = sessionsRef.current.get(sessionKey);
      if (!state) {
-        state = await loadEncrypted<RatchetState>(sessionKey) || undefined;
+        const stored = await loadEncrypted<any>(sessionKey);
+        if (stored) {
+          try {
+            state = await deserializeRatchetState(stored);
+          } catch (e) {
+            console.warn('[Chat] Corrupted session for decryption, treating as new/lost:', e);
+            // Delete the corrupted session from storage
+            await deleteEncrypted(sessionKey);
+          }
+        }
      }

      // 3. X3DH Receiver Init if needed