first commit

backend/app/services/raw_mime_storage.py

@@ -0,0 +1,101 @@
+import asyncio
+import hashlib
+from dataclasses import dataclass
+
+import google_crc32c
+from google.api_core.exceptions import NotFound, PreconditionFailed
+from google.cloud import storage
+
+from app.config import get_settings
+
+
+@dataclass(frozen=True)
+class StoredMime:
+    bucket: str
+    name: str
+    generation: int
+    size: int
+    crc32c: str
+
+
+class RawMimeStorage:
+    def __init__(self, bucket: str, prefix: str = "raw", project: str | None = None):
+        self.bucket_name = bucket
+        self.prefix = prefix.strip("/")
+        self.client = storage.Client(project=project)
+
+    async def put(self, raw: bytes, delivery_key: str) -> StoredMime:
+        return await asyncio.to_thread(self._put, raw, delivery_key)
+
+    def _put(self, raw: bytes, delivery_key: str) -> StoredMime:
+        name = f"{self.prefix}/{delivery_key[:2]}/{delivery_key}.eml"
+        blob = self.client.bucket(self.bucket_name).blob(name)
+        checksum = google_crc32c.Checksum(raw).hexdigest().decode("ascii")
+        try:
+            blob.upload_from_string(
+                raw,
+                content_type="message/rfc822",
+                checksum="crc32c",
+                if_generation_match=0,
+            )
+            blob.reload()
+        except PreconditionFailed:
+            blob.reload()
+        return StoredMime(
+            bucket=self.bucket_name,
+            name=name,
+            generation=int(blob.generation),
+            size=len(raw),
+            crc32c=checksum,
+        )
+
+    async def get(self, reference: dict) -> bytes:
+        return await asyncio.to_thread(self._get, reference)
+
+    def _get(self, reference: dict) -> bytes:
+        generation = int(reference["generation"])
+        blob = self.client.bucket(reference["bucket"]).blob(
+            reference["name"],
+            generation=generation,
+        )
+        return blob.download_as_bytes(if_generation_match=generation)
+
+    async def delete(self, reference: dict) -> None:
+        await asyncio.to_thread(self._delete, reference)
+
+    def _delete(self, reference: dict) -> None:
+        generation = int(reference["generation"])
+        blob = self.client.bucket(reference["bucket"]).blob(
+            reference["name"],
+            generation=generation,
+        )
+        try:
+            blob.delete(if_generation_match=generation)
+        except NotFound:
+            # A prior delete may have removed GCS before Mongo acknowledged it.
+            pass
+
+
+def delivery_digest(raw: bytes, recipients: list[str]) -> str:
+    digest = hashlib.sha256()
+    digest.update(raw)
+    digest.update(b"\0")
+    digest.update("\0".join(sorted(set(recipients))).encode())
+    return digest.hexdigest()
+
+
+_storage: RawMimeStorage | None = None
+
+
+def get_raw_mime_storage() -> RawMimeStorage:
+    global _storage
+    if _storage is None:
+        settings = get_settings()
+        if not settings.gcs_raw_mime_bucket:
+            raise RuntimeError("GCS_RAW_MIME_BUCKET is required for inbound mail.")
+        _storage = RawMimeStorage(
+            settings.gcs_raw_mime_bucket,
+            settings.gcs_raw_mime_prefix,
+            settings.gcs_project,
+        )
+    return _storage