1.1.3

packages/shared/src/components/channel/Messages.tsx

@@ -1,4 +1,4 @@
-import { useMutation, usePaginatedQuery, useQuery } from 'convex/react';
+import { useAction, useMutation, usePaginatedQuery, useQuery } from 'convex/react';
 import { useCallback, useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react';
 import { api } from '../../../../../convex/_generated/api';
 import { usePlatform } from '../../platform';
@@ -57,12 +57,26 @@ const TAG_LENGTH = 32;
 const decryptionCache = new Map<string, string>();
 const MAX_CACHE = 2000;
 
-function namespacedKey(userId: string | null, id: string): string {
-	return `${userId ?? 'anon'}:${id}`;
+// Cache keys are `user:id[:source]`. Passing `source` (the current
+// ciphertext) scopes cache hits to a specific version of the
+// message, so when messages.edit swaps the ciphertext a subsequent
+// fetch naturally misses and re-decrypts. The orphaned entry under
+// the old source ages out via MAX_CACHE LRU.
+function namespacedKey(
+	userId: string | null,
+	id: string,
+	source?: string,
+): string {
+	return source ? `${userId ?? 'anon'}:${id}:${source}` : `${userId ?? 'anon'}:${id}`;
 }
 
-function cacheSet(userId: string | null, id: string, content: string) {
-	const key = namespacedKey(userId, id);
+function cacheSet(
+	userId: string | null,
+	id: string,
+	content: string,
+	source?: string,
+) {
+	const key = namespacedKey(userId, id, source);
 	if (decryptionCache.size >= MAX_CACHE) {
 		const firstKey = decryptionCache.keys().next().value;
 		if (firstKey !== undefined) decryptionCache.delete(firstKey);
@@ -70,8 +84,12 @@ function cacheSet(userId: string | null, id: string, content: string) {
 	decryptionCache.set(key, content);
 }
 
-function cacheGet(userId: string | null, id: string): string | undefined {
-	return decryptionCache.get(namespacedKey(userId, id));
+function cacheGet(
+	userId: string | null,
+	id: string,
+	source?: string,
+): string | undefined {
+	return decryptionCache.get(namespacedKey(userId, id, source));
 }
 
 // Exposed for the logout hook to flush plaintext from memory proactively,
@@ -320,6 +338,15 @@ export function Messages({ channelId, onReply }: MessagesProps) {
 
 	const [decryptedMap, setDecryptedMap] = useState<Map<string, string>>(new Map());
 
+	// Fingerprint of each decrypted entry's source ciphertext.
+	// `messages.edit` changes the ciphertext while keeping the id,
+	// so the decrypt effects below gate on "has id AND source
+	// matches" — without this, the stale plaintext would stay pinned
+	// until the next reload. Populated from cache hits and successful
+	// decrypts. Refs (not state) so staleness checks see the latest
+	// value without a re-render in between.
+	const plaintextSourceRef = useRef<Map<string, string>>(new Map());
+
 	// Reply previews — separate map keyed by the *child* message id so
 	// the per-row render can grab the parent's plaintext without
 	// re-decrypting on every paint. Filled by the effect below.
@@ -347,13 +374,19 @@ export function Messages({ channelId, onReply }: MessagesProps) {
 		if (!pagedMessages || pagedMessages.length === 0) return;
 		let changed = false;
 		let next: Map<string, string> | null = null;
+		const sourceRef = plaintextSourceRef.current;
 		for (const msg of pagedMessages as any[]) {
 			const id = msg.id as string;
-			if (decryptedMap.has(id)) continue;
-			const cached = cacheGet(userId, id);
+			const source = msg.ciphertext as string;
+			// Already decrypted against this exact ciphertext — nothing
+			// to hydrate. If the source has changed (edit), fall through
+			// so we can try the cache keyed by the new source.
+			if (decryptedMap.has(id) && sourceRef.get(id) === source) continue;
+			const cached = cacheGet(userId, id, source);
 			if (cached !== undefined) {
 				if (!next) next = new Map(decryptedMap);
 				next.set(id, cached);
+				sourceRef.set(id, source);
 				changed = true;
 			}
 		}
@@ -450,12 +483,20 @@ export function Messages({ channelId, onReply }: MessagesProps) {
 						nonce: string;
 						tag: string;
 						key: string;
+						ciphertext: string;
 				  };
 			const jobs: Job[] = [];
+			const sourceRef = plaintextSourceRef.current;
 			for (const msg of pagedMessages as any[]) {
 				const id = msg.id as string;
-				if (decryptedMap.has(id)) continue;
-				if (cacheGet(userId, id) !== undefined) continue;
+				const source = msg.ciphertext as string;
+				// Gate on (has decrypted state) AND (source matches).
+				// The source check is what lets edits re-decrypt: after
+				// `messages.edit` swaps the ciphertext, the entry in
+				// `plaintextSourceRef` still points at the old source,
+				// so this branch falls through to queue a fresh decrypt.
+				if (decryptedMap.has(id) && sourceRef.get(id) === source) continue;
+				if (cacheGet(userId, id, source) !== undefined) continue;
 				if (!msg.ciphertext || msg.ciphertext.length < TAG_LENGTH) {
 					jobs.push({ kind: 'sentinel', id, value: '[Invalid Encrypted Message]' });
 					continue;
@@ -474,6 +515,7 @@ export function Messages({ channelId, onReply }: MessagesProps) {
 					tag: msg.ciphertext.slice(-TAG_LENGTH),
 					nonce: msg.nonce,
 					key: keyForVersion,
+					ciphertext: msg.ciphertext,
 				});
 			}
 			if (jobs.length === 0) return;
@@ -484,7 +526,7 @@ export function Messages({ channelId, onReply }: MessagesProps) {
 			const results = await Promise.all(
 				jobs.map(async (j) => {
 					if (j.kind === 'sentinel') {
-						return { id: j.id, value: j.value, cache: false };
+						return { id: j.id, value: j.value, cache: false, source: null as string | null };
 					}
 					try {
 						const plaintext = await crypto.decryptData(
@@ -493,19 +535,22 @@ export function Messages({ channelId, onReply }: MessagesProps) {
 							j.nonce,
 							j.tag,
 						);
-						return { id: j.id, value: plaintext, cache: true };
+						return { id: j.id, value: plaintext, cache: true, source: j.ciphertext };
 					} catch {
-						return { id: j.id, value: '[Unable to decrypt]', cache: false };
+						return { id: j.id, value: '[Unable to decrypt]', cache: false, source: null };
 					}
 				}),
 			);
 			if (cancelled) return;
-			const next = new Map(decryptedMap);
 			for (const r of results) {
-				if (r.cache) cacheSet(userId, r.id, r.value);
-				next.set(r.id, r.value);
+				if (r.cache && r.source) cacheSet(userId, r.id, r.value, r.source);
+				if (r.source) plaintextSourceRef.current.set(r.id, r.source);
 			}
-			setDecryptedMap(next);
+			setDecryptedMap((prev) => {
+				const next = new Map(prev);
+				for (const r of results) next.set(r.id, r.value);
+				return next;
+			});
 		})();
 		return () => {
 			cancelled = true;
@@ -708,6 +753,55 @@ export function Messages({ channelId, onReply }: MessagesProps) {
 			});
 	}, [pagedMessages, decryptedMap, replyPreviewMap, channelId]);
 
+	// ── Edit flow ───────────────────────────────────────────────
+	// Mirrors the send path: we re-encrypt the new plaintext under
+	// the message's *original* keyVersion key (so the server's
+	// unchanged `keyVersion` still decrypts), re-sign the ciphertext,
+	// and produce an `edit:...` auth signature for the action guard.
+	// Bails silently if we no longer have the key for that version —
+	// shouldn't happen in practice (rotations keep old entries) but
+	// we'd rather cancel than corrupt a message.
+	const editMessageAction = useAction(api.messageActions.edit);
+	const handleEditMessage = useCallback(
+		async (messageId: string, newText: string) => {
+			if (!pagedMessages || !userId) throw new Error('Not ready.');
+			const signingKey =
+				typeof sessionStorage !== 'undefined'
+					? sessionStorage.getItem('signingKey')
+					: null;
+			if (!signingKey) throw new Error('No signing key in session.');
+			const raw = (pagedMessages as any[]).find((m: any) => m.id === messageId);
+			if (!raw) throw new Error('Message not in current page.');
+			const msgKeyVersion = Number(raw.key_version ?? 1);
+			const key = channelKeysByVersion.get(msgKeyVersion);
+			if (!key) throw new Error('Missing channel key for this message.');
+			// Messages with attachments wrap the caption as { text } — preserve
+			// that envelope so the attachment list doesn't get dropped. Pure
+			// text messages stay plain strings for forward-compat with older
+			// decryption paths that don't `JSON.parse`.
+			const hadAttachments = Array.isArray(raw?.attachments) && raw.attachments.length > 0;
+			const payload = hadAttachments ? JSON.stringify({ text: newText }) : newText;
+			const { content, iv, tag } = await crypto.encryptData(payload, key);
+			const ciphertext = content + tag;
+			const signature = await crypto.signMessage(signingKey, ciphertext);
+			const authTimestamp = Date.now();
+			const authSignature = await crypto.signMessage(
+				signingKey,
+				`edit:${messageId}:${userId}:${authTimestamp}`,
+			);
+			await editMessageAction({
+				id: messageId as any,
+				userId: userId as any,
+				ciphertext,
+				nonce: iv,
+				signature,
+				authTimestamp,
+				authSignature,
+			});
+		},
+		[pagedMessages, userId, channelKeysByVersion, crypto, editMessageAction],
+	);
+
 	// Pinned-to-bottom tracking. Matches the new UI's approach: a single
 	// boolean in a ref, updated on every onScroll event via isNearBottom.
 	// MutationObserver + ResizeObserver below drive all auto-scroll
@@ -1172,6 +1266,7 @@ export function Messages({ channelId, onReply }: MessagesProps) {
 									messages={item.group as any}
 									channelId={channelId}
 									onReply={onReply}
+									onEditMessage={handleEditMessage}
 								/>,
 							);
 						} else {