wip encrypt with attachments

src/backend/core/api/serializers.py

@@ -976,6 +976,14 @@ class EncryptDocumentSerializer(serializers.Serializer):
 
     content = serializers.CharField(required=True)
     encryptedSymmetricKeyPerUser = serializers.DictField(child=serializers.CharField(), required=True)
+    attachmentKeyMapping = serializers.DictField(
+        child=serializers.CharField(),
+        required=False,
+        default=dict,
+        help_text="Mapping of original attachment key to new encrypted attachment key. "
+        "During encryption, existing attachments are uploaded encrypted under new keys. "
+        "This mapping tells the backend to copy each new key over the original and clean up.",
+    )
 
 
 class RemoveEncryptionSerializer(serializers.Serializer):

src/backend/core/api/viewsets.py

@@ -1480,19 +1480,10 @@ class DocumentViewSet(
         serializer = serializers.FileUploadSerializer(data=request.data)
         serializer.is_valid(raise_exception=True)
 
+        # Normally encrypted attachments would be only allowed on encrypted documents and vice-versa
+        # but since during encryption/decryption we upload all attachments before the switch, we cannot enforce this rule
         is_file_encrypted = serializer.validated_data.get("is_encrypted", False)
 
-        # Encrypted attachments are only allowed on encrypted documents
-        if is_file_encrypted and not document.is_encrypted:
-            raise drf.exceptions.ValidationError({
-                "is_encrypted":
-                "Cannot upload encrypted attachments to a non-encrypted document."
-            })
-
-        # Generate a generic yet unique filename to store the image in object storage
-        file_id = uuid.uuid4()
-        ext = serializer.validated_data["expected_extension"]
-
         # For encrypted files, set status to READY immediately since the server
         # cannot inspect ciphertext for malware scanning.
         initial_status = (
@@ -1513,6 +1504,10 @@ class DocumentViewSet(
         if is_file_encrypted:
             extra_args["Metadata"]["is_encrypted"] = "true"
 
+        # Generate a generic yet unique filename to store the image in object storage
+        file_id = uuid.uuid4()
+        ext = serializer.validated_data["expected_extension"]
+
         file_unsafe = ""
         if serializer.validated_data["is_unsafe"]:
             extra_args["Metadata"]["is_unsafe"] = "true"
@@ -2058,6 +2053,7 @@ class DocumentViewSet(
 
         content = serializer.validated_data["content"]
         encryptedSymmetricKeyPerUser = serializer.validated_data["encryptedSymmetricKeyPerUser"]
+        attachment_key_mapping = serializer.validated_data.get("attachmentKeyMapping", {})
 
         # Prevent encryption if there are pending invitations
         if document.invitations.exists():
@@ -2092,11 +2088,34 @@ class DocumentViewSet(
                 f'Only users with access should have encrypted symmetric keys.'
             })
 
+        # Remove old unencrypted attachment keys from the allowed list.
+        # The frontend uploaded encrypted copies under new keys and updated the
+        # Yjs content to reference them.
+        if attachment_key_mapping:
+            old_keys = set(attachment_key_mapping.keys())
+            document.attachments = [
+                k for k in (document.attachments or []) if k not in old_keys
+            ]
+
         # Update the document content and encryption status
         document.content = content  # This will be cached and saved to object storage
         document.is_encrypted = True
         document.save()
 
+        # Clean up old S3 objects only after the DB transaction has committed,
+        # so a deletion failure can never affect the encrypt operation.
+        if attachment_key_mapping:
+            def _cleanup_old_attachments():
+                s3_client = default_storage.connection.meta.client
+                bucket_name = default_storage.bucket_name
+                for old_key in attachment_key_mapping:
+                    try:
+                        s3_client.delete_object(Bucket=bucket_name, Key=old_key)
+                    except ClientError:
+                        logger.warning("Failed to delete old attachment %s", old_key)
+
+            transaction.on_commit(_cleanup_old_attachments)
+
         # Store the encrypted symmetric keys in DocumentAccess for each user
         for user_id, encrypted_key in encryptedSymmetricKeyPerUser.items():
             try:

src/frontend/apps/impress/src/features/docs/doc-editor/components/BlockNoteEditor.tsx

@@ -154,7 +154,7 @@ export const BlockNoteEditor = ({
       };
       const mime = mimeMap[ext] || 'application/octet-stream';
 
-      const blob = new Blob([decryptedBytes], { type: mime });
+      const blob = new Blob([decryptedBytes as BlobPart], { type: mime });
       const blobUrl = URL.createObjectURL(blob);
       blobUrlCacheRef.current.set(url, blobUrl);
 

src/frontend/apps/impress/src/features/docs/doc-editor/components/custom-blocks/PdfBlock.tsx

@@ -17,7 +17,7 @@ import { useCallback, useEffect, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { createGlobalStyle, css } from 'styled-components';
 
-import { Box, Button, Icon, Loading } from '@/components';
+import { Box, Icon, Loading } from '@/components';
 
 import { ANALYZE_URL } from '../../conf';
 import { DocsBlockNoteEditor } from '../../types';

src/frontend/apps/impress/src/features/docs/doc-management/api/useEncryptDoc.tsx

@@ -11,6 +11,7 @@ interface EncryptDocProps {
   docId: string;
   content: Uint8Array<ArrayBufferLike>;
   encryptedSymmetricKeyPerUser: Record<string, ArrayBuffer>;
+  attachmentKeyMapping?: Record<string, string>;
 }
 
 export const encryptDoc = async ({
@@ -33,6 +34,7 @@ export const encryptDoc = async ({
       ...params,
       content: toBase64(params.content),
       encryptedSymmetricKeyPerUser: base64EncryptedSymmetricKeyPerUser,
+      attachmentKeyMapping: params.attachmentKeyMapping || {},
     }),
   });
 

src/frontend/apps/impress/src/features/docs/doc-management/components/ModalEncryptDoc.tsx

@@ -1,5 +1,6 @@
 import {
   Button,
+  Loader,
   Modal,
   ModalSize,
   VariantType,
@@ -8,6 +9,9 @@ import {
 import { useTranslation } from 'react-i18next';
 import * as Y from 'yjs';
 
+import { backendUrl } from '@/api';
+import { useState } from 'react';
+
 import { Box, ButtonCloseModal, Text, TextErrors } from '@/components';
 import { useUserUpdate } from '@/core/api/useUserUpdate';
 import {
@@ -17,6 +21,7 @@ import {
   getEncryptionDB,
   prepareEncryptedSymmetricKeysForUsers,
 } from '@/docs/doc-collaboration';
+import { createDocAttachment } from '@/docs/doc-editor/api';
 import { toBase64 } from '@/docs/doc-editor';
 import { useAuth } from '@/features/auth';
 import {
@@ -27,6 +32,158 @@ import {
   useProviderStore,
 } from '@/features/docs/doc-management';
 import { useKeyboardAction } from '@/hooks';
+import { Spinner } from '@gouvfr-lasuite/ui-kit';
+
+function traverseYDoc(
+  node: Y.XmlElement | Y.XmlFragment,
+  callback: (el: Y.XmlElement) => void,
+) {
+  if (node instanceof Y.XmlElement) {
+    callback(node);
+  }
+
+  node.toArray().forEach((child) => {
+    if (child instanceof Y.XmlElement || child instanceof Y.XmlFragment) {
+      traverseYDoc(child, callback);
+    }
+  });
+}
+
+const UUID =
+  '[\\da-fA-F]{8}-[\\da-fA-F]{4}-[\\da-fA-F]{4}-[\\da-fA-F]{4}-[\\da-fA-F]{12}';
+const ATTACHMENT_KEY_REGEX = new RegExp(
+  `^/media/(${UUID}/attachments/${UUID}(?:-unsafe)?\\.[a-zA-Z0-9]{1,10})$`,
+);
+
+type ExtractedKeyMetadata = {
+  mediaUrl: string;
+  name?: string;
+  nodes: Y.XmlElement[];
+};
+
+// extract unique attachment keys from the Yjs document
+const extractAttachmentKeysAndMetadata = (
+  yDoc: Y.Doc,
+): Map<string, ExtractedKeyMetadata> => {
+  const fragment = yDoc.getXmlFragment('document-store');
+
+  // for each key keep track of nodes
+  const keysAndMetadata = new Map<string, ExtractedKeyMetadata>();
+
+  yDoc.transact(() => {
+    traverseYDoc(fragment, (node) => {
+      const urlAttributeValue = node.getAttribute('url');
+
+      if (urlAttributeValue) {
+        // url should always be valid
+        const url = new URL(urlAttributeValue);
+
+        // applying the test only on the pathname since hostname can vary
+        const match = ATTACHMENT_KEY_REGEX.exec(url.pathname);
+
+        if (match) {
+          const key = match[1];
+          const keyMetadata = keysAndMetadata.get(key);
+
+          if (keyMetadata) {
+            keyMetadata.nodes.push(node);
+          } else {
+            // avoid any unexpected parts
+            url.search = '';
+            url.hash = '';
+
+            keysAndMetadata.set(key, {
+              mediaUrl: url.toString(),
+              name: node.getAttribute('name'),
+              nodes: [node],
+            });
+          }
+        }
+      }
+    });
+  });
+
+  return keysAndMetadata;
+};
+
+/**
+ * encrypt existing unencrypted attachments and return:
+ * - a modified Yjs state with URLs pointing to new encrypted files
+ * - a mapping of old S3 keys to new ones (for backend cleanup)
+ *
+ * originals are never modified so if the process fails midway the document
+ * still works with its original unencrypted attachments.
+ */
+const encryptRemoteAttachments = async (
+  yDoc: Y.Doc,
+  docId: string,
+  symmetricKey: CryptoKey,
+): Promise<Record<string, string>> => {
+  const attachmentKeysAndMetadata = extractAttachmentKeysAndMetadata(yDoc);
+
+  // if no attachment it's straightforward
+  if (attachmentKeysAndMetadata.size === 0) {
+    return {};
+  }
+
+  // otherwise upload encrypted copies as new attachments and collect the mapping
+  const attachmentKeyMapping: Record<string, string> = {};
+
+  for (const [oldAttachmentKey, oldAttachmentMetadata] of Array.from(
+    attachmentKeysAndMetadata.entries(),
+  )) {
+    const response = await fetch(oldAttachmentMetadata.mediaUrl, {
+      credentials: 'include',
+    });
+    if (!response.ok) {
+      throw new Error('attachment cannot be fetch');
+    }
+
+    const fileBytes = new Uint8Array(await response.arrayBuffer());
+    const encryptedBytes = await encryptContent(fileBytes, symmetricKey);
+
+    const fileName = oldAttachmentMetadata.name ?? 'file'; // since encrypted we could not reuse the file name that can be stored as clear text
+    const encryptedFile = new File([encryptedBytes], fileName, {
+      type: 'application/octet-stream',
+    });
+
+    const body = new FormData();
+    body.append('file', encryptedFile);
+    body.append('is_encrypted', 'true');
+
+    const result = await createDocAttachment({ docId, body });
+
+    // result.file is like "/api/v1.0/documents/{id}/media-check/?key={newKey}"
+    const newKey = new URL(
+      result.file,
+      window.location.origin,
+    ).searchParams.get('key');
+
+    if (!newKey) {
+      throw new Error('file key must be provided once uploaded');
+    }
+
+    attachmentKeyMapping[oldAttachmentKey] = newKey;
+  }
+
+  // once uploaded, we can update all nodes referencing attachments with their new key
+  yDoc.transact(() => {
+    for (const [oldAttachmentKey, oldAttachmentMetadata] of Array.from(
+      attachmentKeysAndMetadata.entries(),
+    )) {
+      const newMediaUrl = oldAttachmentMetadata.mediaUrl.replace(
+        oldAttachmentKey,
+        attachmentKeyMapping[oldAttachmentKey],
+      );
+
+      for (const node of oldAttachmentMetadata.nodes) {
+        node.setAttribute('url', newMediaUrl);
+      }
+    }
+  });
+
+  return attachmentKeyMapping;
+};
 
 interface ModalEncryptDocProps {
   doc: Doc;
@@ -51,129 +208,156 @@ export const ModalEncryptDoc = ({
   const { user } = useAuth();
   const { mutateAsync: updateUser } = useUserUpdate();
 
+  const [isPending, setIsPending] = useState(false);
+
   const {
-    mutate: encryptDoc,
+    mutateAsync: encryptDoc,
     isError,
     error,
   } = useEncryptDoc({
     listInvalidQueries: [KEY_DOC, KEY_LIST_DOC],
-    options: {
-      onSuccess: () => {
-        onSuccess && onSuccess(doc);
-        onClose();
-
-        toast(t('The document has been encrypted.'), VariantType.SUCCESS, {
-          duration: 4000,
-        });
-      },
-    },
   });
 
   const keyboardAction = useKeyboardAction();
 
   const handleClose = () => {
+    if (isPending) {
+      return;
+    }
     onClose();
   };
 
   const handleEncrypt = async () => {
-    if (!provider || !user) {
+    if (!provider || !user || isPending) {
       return;
     }
 
-    let currentUserPublicKeyFromThisOnboardingSession: ArrayBuffer | null =
-      null;
+    setIsPending(true);
 
-    // Perform the onboarding if that's the first time using encryption on this device
-    if (!encryptionSettings) {
-      // TODO: trigger the onboarding, either by creating or retrieving a key from another device
-      // TODO: probably the logic should be at a device key level, not user one?
+    try {
+      let currentUserPublicKeyFromThisOnboardingSession: ArrayBuffer | null =
+        null;
 
-      const userKeyPair = await generateUserKeyPair();
+      // Perform the onboarding if that's the first time using encryption on this device
+      if (!encryptionSettings) {
+        // TODO: trigger the onboarding, either by creating or retrieving a key from another device
+        // TODO: probably the logic should be at a device key level, not user one?
 
-      const encryptionDatabase = await getEncryptionDB();
+        const userKeyPair = await generateUserKeyPair();
 
-      // TODO: it should use transaction
-      // encryptionDatabase.transaction
-      await encryptionDatabase.put(
-        'privateKey',
-        userKeyPair.privateKey,
-        `user:${user.id}`,
-      );
-      await encryptionDatabase.put(
-        'publicKey',
-        userKeyPair.publicKey,
-        `user:${user.id}`,
+        const encryptionDatabase = await getEncryptionDB();
+
+        // TODO: it should use transaction
+        // encryptionDatabase.transaction
+        await encryptionDatabase.put(
+          'privateKey',
+          userKeyPair.privateKey,
+          `user:${user.id}`,
+        );
+        await encryptionDatabase.put(
+          'publicKey',
+          userKeyPair.publicKey,
+          `user:${user.id}`,
+        );
+
+        const rawPublicKey = await crypto.subtle.exportKey(
+          'spki',
+          userKeyPair.publicKey,
+        );
+
+        // TODO: it should throw if the backend has already a public key (so the user can with concious forget the old one (but here he did the onboarding already so... it was probably a new device))
+        await updateUser({
+          id: user.id,
+          encryption_public_key: toBase64(new Uint8Array(rawPublicKey)),
+        });
+
+        currentUserPublicKeyFromThisOnboardingSession = rawPublicKey;
+
+        // TODO: should check encryptionSettings will update, otherwise hard refresh is needed
+        window.location.reload();
+
+        return;
+      }
+
+      const documentSymmetricKey = await generateSymmetricKey();
+
+      // Their public key are base64 encoded, decoding the whole
+      const usersPublicKeys: Record<string, ArrayBuffer> = {};
+
+      if (doc.accesses_public_keys_per_user) {
+        // TODO:
+        // TODO: should throw if missing public keys according to current accesses
+        // TODO:
+
+        for (const [userId, publicKey] of Object.entries(
+          doc.accesses_public_keys_per_user,
+        )) {
+          usersPublicKeys[userId] = Buffer.from(publicKey, 'base64').buffer;
+        }
+
+        // if the onboarding has been done directly in this encryption flow, the backend has not yet told the frontend
+        // about the current user key, so just patching the mapping with this new public key
+        if (currentUserPublicKeyFromThisOnboardingSession) {
+          usersPublicKeys[user.id] =
+            currentUserPublicKeyFromThisOnboardingSession;
+        }
+      } else {
+        // if it has been not provided it's weird because it should only happen for people not authenticated
+        throw new Error(`"accesses_public_keys_per_user" should be provided`);
+      }
+
+      // Prepare encrypted symmetric keys for all users with access
+      const encryptedSymmetricKeyPerUser =
+        await prepareEncryptedSymmetricKeysForUsers(
+          documentSymmetricKey,
+          usersPublicKeys,
+        );
+
+      // clone the Yjs document since performing changes during encryption that require backend confirmation
+      // once successfully done it can be used locally
+      const ongoingDoc = new Y.Doc();
+      Y.applyUpdate(ongoingDoc, Y.encodeStateAsUpdate(provider.document));
+
+      // encrypt existing attachments
+      const attachmentKeyMapping = await encryptRemoteAttachments(
+        ongoingDoc,
+        doc.id,
+        documentSymmetricKey,
       );
 
-      const rawPublicKey = await crypto.subtle.exportKey(
-        'spki',
-        userKeyPair.publicKey,
+      const ongoingDocState = Y.encodeStateAsUpdate(ongoingDoc);
+
+      const encryptedContent = await encryptContent(
+        new Uint8Array(ongoingDocState),
+        documentSymmetricKey,
       );
 
-      // TODO: it should throw if the backend has already a public key (so the user can with concious forget the old one (but here he did the onboarding already so... it was probably a new device))
-      await updateUser({
-        id: user.id,
-        encryption_public_key: toBase64(new Uint8Array(rawPublicKey)),
+      // TODO:
+      // TODO: if none it should at least make it for the current user
+      // TODO: so it makes sense `accesses_public_keys_per_user` is always passed?
+      // TODO:
+
+      await encryptDoc({
+        docId: doc.id,
+        content: encryptedContent,
+        encryptedSymmetricKeyPerUser,
+        attachmentKeyMapping,
       });
 
-      currentUserPublicKeyFromThisOnboardingSession = rawPublicKey;
+      // since the encrypted state has been committed with success with can apply it locally with adjusted encrypted attachments
+      Y.applyUpdate(provider!.document, Y.encodeStateAsUpdate(ongoingDoc));
 
-      // TODO: should check encryptionSettings will update, otherwise hard refresh is needed
-      window.location.reload();
+      onSuccess?.(doc);
+      onClose();
 
-      return;
+      toast(t('The document has been encrypted.'), VariantType.SUCCESS, {
+        duration: 4000,
+      });
+
+      ongoingDoc.destroy();
+    } finally {
+      setIsPending(false);
     }
-
-    const documentSymmetricKey = await generateSymmetricKey();
-
-    const state = Y.encodeStateAsUpdate(provider.document);
-    const encryptedContent = await encryptContent(
-      new Uint8Array(state),
-      documentSymmetricKey,
-    );
-
-    // Their public key are base64 encoded, decoding the whole
-    const usersPublicKeys: Record<string, ArrayBuffer> = {};
-
-    if (doc.accesses_public_keys_per_user) {
-      // TODO:
-      // TODO: should throw if missing public keys according to current accesses
-      // TODO:
-
-      for (const [userId, publicKey] of Object.entries(
-        doc.accesses_public_keys_per_user,
-      )) {
-        usersPublicKeys[userId] = Buffer.from(publicKey, 'base64').buffer;
-      }
-
-      // if the onboarding has been done directly in this encryption flow, the backend has not yet told the frontend
-      // about the current user key, so just patching the mapping with this new public key
-      if (currentUserPublicKeyFromThisOnboardingSession) {
-        usersPublicKeys[user.id] =
-          currentUserPublicKeyFromThisOnboardingSession;
-      }
-    } else {
-      // if it has been not provided it's weird because it should only happen for people not authenticated
-      throw new Error(`"accesses_public_keys_per_user" should be provided`);
-    }
-
-    // Prepare encrypted symmetric keys for all users with access
-    const encryptedSymmetricKeyPerUser =
-      await prepareEncryptedSymmetricKeysForUsers(
-        documentSymmetricKey,
-        usersPublicKeys,
-      );
-
-    // TODO:
-    // TODO: if none it should at least make it for the current user
-    // TODO: so it makes sense `accesses_public_keys_per_user` is always passed?
-    // TODO:
-
-    encryptDoc({
-      docId: doc.id,
-      content: encryptedContent,
-      encryptedSymmetricKeyPerUser,
-    });
   };
 
   const handleCloseKeyDown = keyboardAction(handleClose);
@@ -182,7 +366,7 @@ export const ModalEncryptDoc = ({
   return (
     <Modal
       isOpen
-      closeOnClickOutside
+      closeOnClickOutside={!isPending}
       hideCloseButton
       onClose={handleClose}
       aria-describedby="modal-encrypt-doc-title"
@@ -193,6 +377,7 @@ export const ModalEncryptDoc = ({
             fullWidth
             onClick={handleClose}
             onKeyDown={handleCloseKeyDown}
+            disabled={isPending}
           >
             {t('Cancel')}
           </Button>
@@ -201,6 +386,14 @@ export const ModalEncryptDoc = ({
             fullWidth
             onClick={handleEncrypt}
             onKeyDown={handleEncryptKeyDown}
+            disabled={isPending}
+            icon={
+              isPending ? (
+                <div>
+                  <Spinner size="sm" />
+                </div>
+              ) : undefined
+            }
           >
             {t('Confirm')}
           </Button>
@@ -227,6 +420,7 @@ export const ModalEncryptDoc = ({
             aria-label={t('Close the encrypt modal')}
             onClick={handleClose}
             onKeyDown={handleCloseKeyDown}
+            disabled={isPending}
           />
         </Box>
       }