simplify

Author: sylvinus

.github/workflows/messages.yml

@@ -39,8 +39,8 @@ jobs:
           chown -R 1001:1001 data
       - name: Collect static files
         run: make collectstatic
-      - name: Create msg-imports bucket
-        run: make import-bucket
+      - name: Create object storage buckets
+        run: make create-buckets
       - name: Run backend tests
         run: make test-back-parallel
 

Makefile

@@ -100,7 +100,7 @@ bootstrap: ## Prepare the project for local development
 update:  ## Update the project with latest changes
 	@$(MAKE) data/media
 	@$(MAKE) data/static
-	@$(MAKE) import-bucket
+	@$(MAKE) create-buckets
 	@$(MAKE) create-env-files
 	@$(MAKE) build
 	@$(MAKE) collectstatic
@@ -161,11 +161,11 @@ restart-minimal: \
 	start-minimal
 .PHONY: restart-minimal
 
-import-bucket: ## create the message imports & blobs buckets in objectstorage
+create-buckets: ## create the message imports & blobs buckets in objectstorage
 	@$(COMPOSE) up -d objectstorage --wait
 	@$(MANAGE_DB) create_bucket --storage message-imports --expire-days 1
 	@$(MANAGE_DB) create_bucket --storage message-blobs
-.PHONY: import-bucket
+.PHONY: create-buckets
 
 shell-objectstorage: ## open a shell in the objectstorage container
 	@$(COMPOSE) run --rm --build objectstorage bash

src/backend/core/admin.py

@@ -1005,10 +1005,19 @@ class BlobAdmin(admin.ModelAdmin):
         "size",
         "size_compressed",
         "compression",
+        "storage_location",
+        "encryption_key_id",
         "created_at",
     )
     search_fields = ("mailbox__local_part", "mailbox__domain__name", "content_type")
-    list_filter = ("content_type", "compression", "created_at", "updated_at")
+    list_filter = (
+        "content_type",
+        "compression",
+        "storage_location",
+        "encryption_key_id",
+        "created_at",
+        "updated_at",
+    )
     autocomplete_fields = ("mailbox",)
 
     def get_queryset(self, request):

src/backend/core/api/viewsets/config.py

@@ -190,7 +190,8 @@ def get(self, request):
         ]
         dict_settings = {}
         for setting in array_settings:
-            dict_settings[setting] = getattr(settings, setting, None)
+            if hasattr(settings, setting):
+                dict_settings[setting] = getattr(settings, setting)
 
         # AI Features
         dict_settings["AI_ENABLED"] = is_ai_enabled()

src/backend/core/management/commands/verify_tiered_storage.py

@@ -17,7 +17,7 @@
 
 from core.enums import BlobStorageLocationChoices, CompressionTypeChoices
 from core.models import Blob
-from core.services.tiered_storage import TieredStorageService
+from core.services.tiered_storage import TieredStorageService, sha256_advisory_lock
 
 
 class Command(BaseCommand):
@@ -31,11 +31,6 @@ def add_arguments(self, parser):
             help="Verification mode: db-to-storage (check DB records have storage backing), "
             "storage-to-db (find orphans in storage), or full (both)",
         )
-        parser.add_argument(
-            "--fix",
-            action="store_true",
-            help="Fix issues: delete orphans from storage (missing blobs are only reported)",
-        )
         parser.add_argument(
             "--verify-hashes",
             action="store_true",
@@ -61,7 +56,6 @@ def add_arguments(self, parser):
 
     def handle(self, *args, **options):
         self.service = TieredStorageService()
-        self.fix = options["fix"]
         self.verify_hashes = options["verify_hashes"]
         self.limit = options["limit"]
         self.dry_run = options["dry_run"]
@@ -184,16 +178,6 @@ def verify_storage_to_db(self):
             if not refs_exist:
                 orphan_count += 1
                 self.stdout.write(self.style.WARNING(f"ORPHAN: {obj_name}"))
-                if self.fix:
-                    try:
-                        self.service.storage.delete(obj_name)
-                        self.stdout.write(
-                            self.style.SUCCESS(f"  -> Deleted orphan {obj_name}")
-                        )
-                    except Exception as e:  # pylint: disable=broad-except
-                        self.stdout.write(
-                            self.style.ERROR(f"  -> Failed to delete: {e}")
-                        )
 
             # Optionally verify hash
             if self.verify_hashes and refs_exist:
@@ -400,183 +384,145 @@ def re_encrypt_blobs(self):
             )
 
     def _re_encrypt_single_blob(self, blob: Blob, target_key_id: int) -> str:
-        """
-        Re-encrypt a single blob with the target key.
-
-        Uses select_for_update to prevent concurrent modifications.
-        For object storage blobs, keeps old content in memory so S3 can
-        be restored if the DB update fails.
+        """Re-encrypt a single blob (and its dedup cohort) with the target key.
 
-        Args:
-            blob: The blob to re-encrypt
-            target_key_id: The encryption key ID to use for re-encryption
+        Holds the per-sha256 advisory lock for the duration so concurrent
+        offload, cleanup, or another re-encrypt of the same content cannot
+        interleave.
 
-        Returns:
-            "success", "skipped", or "error"
+        Returns "success", "skipped", or "error".
         """
-        old_key_id = blob.encryption_key_id
-
         if self.dry_run:
-            location = (
-                "POSTGRES"
-                if blob.storage_location == BlobStorageLocationChoices.POSTGRES
-                else "OBJECT_STORAGE"
-            )
+            location = blob.get_storage_location_display()
             self.stdout.write(
                 f"  Would re-encrypt blob {blob.id} ({location}): "
-                f"key_id {old_key_id} -> {target_key_id}"
+                f"key_id {blob.encryption_key_id} -> {target_key_id}"
             )
             return "success"
 
+        sha256 = bytes(blob.sha256)
+
         if blob.storage_location == BlobStorageLocationChoices.POSTGRES:
-            with transaction.atomic():
-                blob = Blob.objects.select_for_update().get(id=blob.id)
-                if blob.encryption_key_id == target_key_id:
-                    return "skipped"
+            return self._re_encrypt_postgres_blob(blob, sha256, target_key_id)
+        return self._re_encrypt_object_storage_blob(blob, sha256, target_key_id)
+
+    def _re_encrypt_postgres_blob(
+        self, blob: Blob, sha256: bytes, target_key_id: int
+    ) -> str:
+        """Re-encrypt a single PostgreSQL-backed blob row in place."""
+        with transaction.atomic(), sha256_advisory_lock(sha256):
+            blob = Blob.objects.select_for_update().get(id=blob.id)
+            if blob.encryption_key_id == target_key_id:
+                return "skipped"
+            if blob.raw_content is None:
+                self.stdout.write(
+                    self.style.WARNING(
+                        f"  SKIP blob {blob.id}: no content in PostgreSQL"
+                    )
+                )
+                return "skipped"
 
-                old_key_id = blob.encryption_key_id
+            old_key_id = blob.encryption_key_id
+            decrypted = self.service.decrypt(bytes(blob.raw_content), old_key_id)
+            encrypted, new_key_id = self.service.encrypt(decrypted)
 
-                if blob.raw_content is None:
-                    self.stdout.write(
-                        self.style.WARNING(
-                            f"  SKIP blob {blob.id}: no content in PostgreSQL"
-                        )
-                    )
-                    return "skipped"
+            blob.raw_content = encrypted
+            blob.encryption_key_id = new_key_id
+            blob.save(
+                update_fields=["raw_content", "encryption_key_id", "size_compressed"]
+            )
 
-                decrypted = self.service.decrypt(bytes(blob.raw_content), old_key_id)
-                encrypted, new_key_id = self.service.encrypt(decrypted)
+        self.stdout.write(
+            f"  Re-encrypted blob {blob.id} (POSTGRES): "
+            f"key_id {old_key_id} -> {new_key_id}"
+        )
+        return "success"
 
-                blob.raw_content = encrypted
-                blob.encryption_key_id = new_key_id
-                blob.save(
-                    update_fields=[
-                        "raw_content",
-                        "encryption_key_id",
-                        "size_compressed",
-                    ]
-                )
+    def _re_encrypt_object_storage_blob(
+        self, blob: Blob, sha256: bytes, target_key_id: int
+    ) -> str:
+        """Re-encrypt the storage object backing an OBJECT_STORAGE cohort.
 
+        Stages the new ciphertext at a temp key so a crash between DB
+        commit and storage promote leaves storage and DB consistent.
+        Concurrency is handled by the per-sha256 advisory lock; the
+        bulk-update flips every sibling to the new key in one round trip.
+        """
+        if not self.service.enabled:
             self.stdout.write(
-                f"  Re-encrypted blob {blob.id} (POSTGRES): "
-                f"key_id {old_key_id} -> {new_key_id}"
+                self.style.WARNING(
+                    f"  SKIP blob {blob.id}: object storage not configured"
+                )
             )
+            return "skipped"
 
-        else:
-            # Object storage blob
-            if not self.service.enabled:
-                self.stdout.write(
-                    self.style.WARNING(
-                        f"  SKIP blob {blob.id}: object storage not configured"
-                    )
-                )
-                return "skipped"
+        storage_key = self.service.compute_storage_key(sha256)
+        temp_key = None
 
-            storage_key = self.service.compute_storage_key(bytes(blob.sha256))
+        try:
+            with transaction.atomic(), sha256_advisory_lock(sha256):
+                cohort = Blob.objects.filter(
+                    sha256=sha256,
+                    storage_location=BlobStorageLocationChoices.OBJECT_STORAGE,
+                ).exclude(encryption_key_id=target_key_id)
+                old_blob = cohort.first()
+                if old_blob is None:
+                    return "skipped"
+                old_key_id = old_blob.encryption_key_id
 
-            # Read & re-encrypt outside the DB transaction.
-            try:
-                with self.service.storage.open(storage_key, "rb") as f:
-                    old_encrypted = f.read()
-            except FileNotFoundError:
-                self.stderr.write(
-                    self.style.ERROR(
-                        f"  ERROR blob {blob.id}: not found in storage at {storage_key}"
+                try:
+                    with self.service.storage.open(storage_key, "rb") as f:
+                        old_encrypted = f.read()
+                except FileNotFoundError:
+                    self.stderr.write(
+                        self.style.ERROR(
+                            f"  ERROR blob {blob.id}: not found in storage at "
+                            f"{storage_key}"
+                        )
                     )
-                )
-                return "error"
+                    return "error"
 
-            decrypted = self.service.decrypt(old_encrypted, blob.encryption_key_id)
-            encrypted, new_key_id = self.service.encrypt(decrypted)
+                decrypted = self.service.decrypt(old_encrypted, old_key_id)
+                encrypted, new_key_id = self.service.encrypt(decrypted)
 
-            # Stage the new ciphertext at a temp key so the canonical object
-            # is not touched until the DB transaction commits. If the DB
-            # update fails, we drop the temp object and storage stays
-            # consistent with the DB.
-            temp_key = f"{storage_key}.tmp.{new_key_id}"
-            self.service.storage.save(temp_key, ContentFile(encrypted))
+                temp_key = f"{storage_key}.tmp.{new_key_id}"
+                self.service.storage.save(temp_key, ContentFile(encrypted))
 
-            promote_done = {"value": False}
+                cohort.update(encryption_key_id=new_key_id)
 
-            def _promote_temp():
-                try:
-                    self.service.storage.save(storage_key, ContentFile(encrypted))
-                finally:
-                    try:
-                        self.service.storage.delete(temp_key)
-                    except Exception as cleanup_err:  # pylint: disable=broad-except
-                        self.stderr.write(
-                            self.style.ERROR(
-                                f"  WARN blob {blob.id}: failed to delete temp "
-                                f"{temp_key}: {cleanup_err}"
-                            )
-                        )
-                promote_done["value"] = True
+                staged_temp = temp_key
 
-            try:
-                with transaction.atomic():
-                    # The storage object is shared by every Blob row that has
-                    # the same sha256 and is in OBJECT_STORAGE (deduplication
-                    # in upload_blob). Lock the entire cohort so we can flip
-                    # all of them to new_key_id atomically — otherwise the
-                    # promote rewrites the object with new ciphertext while
-                    # un-rotated siblings still hold the old key_id.
-                    siblings = list(
-                        Blob.objects.select_for_update().filter(
-                            sha256=blob.sha256,
-                            storage_location=BlobStorageLocationChoices.OBJECT_STORAGE,
-                        )
-                    )
-                    if not siblings:
-                        # Row was deleted between iteration and lock.
-                        try:
-                            self.service.storage.delete(temp_key)
-                        except Exception as cleanup_err:  # pylint: disable=broad-except
-                            self.stderr.write(
-                                self.style.ERROR(
-                                    f"  WARN blob {blob.id}: failed to delete "
-                                    f"temp {temp_key}: {cleanup_err}"
-                                )
-                            )
-                        return "skipped"
-
-                    pending = [
-                        s for s in siblings if s.encryption_key_id != target_key_id
-                    ]
-                    if not pending:
-                        # Another worker rotated this cohort already.
+                def _promote_temp():
+                    try:
+                        self.service.storage.save(storage_key, ContentFile(encrypted))
+                    finally:
                         try:
-                            self.service.storage.delete(temp_key)
+                            self.service.storage.delete(staged_temp)
                         except Exception as cleanup_err:  # pylint: disable=broad-except
                             self.stderr.write(
                                 self.style.ERROR(
                                     f"  WARN blob {blob.id}: failed to delete "
-                                    f"temp {temp_key}: {cleanup_err}"
+                                    f"temp {staged_temp}: {cleanup_err}"
                                 )
                             )
-                        return "skipped"
 
-                    old_key_id = pending[0].encryption_key_id
-                    Blob.objects.filter(id__in=[s.id for s in pending]).update(
-                        encryption_key_id=new_key_id
-                    )
-                    transaction.on_commit(_promote_temp)
-            except Exception:
-                if not promote_done["value"]:
-                    try:
-                        self.service.storage.delete(temp_key)
-                    except Exception as cleanup_err:  # pylint: disable=broad-except
-                        self.stderr.write(
-                            self.style.ERROR(
-                                f"  WARN blob {blob.id}: failed to delete temp "
-                                f"{temp_key}: {cleanup_err}"
-                            )
+                transaction.on_commit(_promote_temp)
+                # Promotion now owns temp_key; don't double-delete on rollback.
+                temp_key = None
+        finally:
+            if temp_key is not None:
+                try:
+                    self.service.storage.delete(temp_key)
+                except Exception as cleanup_err:  # pylint: disable=broad-except
+                    self.stderr.write(
+                        self.style.ERROR(
+                            f"  WARN blob {blob.id}: failed to delete temp "
+                            f"{temp_key}: {cleanup_err}"
                         )
-                raise
-
-            self.stdout.write(
-                f"  Re-encrypted blob {blob.id} (OBJECT_STORAGE): "
-                f"key_id {old_key_id} -> {new_key_id}"
-            )
+                    )
 
+        self.stdout.write(
+            f"  Re-encrypted blob {blob.id} (OBJECT_STORAGE): "
+            f"key_id {old_key_id} -> {new_key_id}"
+        )
         return "success"