393 files

custom_components/maintenance_supporter/helpers/doc_archive.py

@@ -111,11 +111,14 @@ def _object_part_ids(hass: HomeAssistant, object_id: str) -> set[str]:
     return set()
 
 
-def build_documents_archive(hass: HomeAssistant, entry_ids: set[str] | None = None) -> bytes:
-    """Build a documents ZIP for the selected objects (None = all).
+class ArchiveTooLarge(ValueError):
+    """The selection's files exceed what an import accepts (MAX_ARCHIVE_BYTES)."""
 
-    Runs on the event loop for the metadata gather (reads config entries) but
-    the heavy blob reads happen here synchronously — call via the executor.
+
+def _gather_archive(hass: HomeAssistant, entry_ids: set[str] | None) -> tuple[Any, dict[str, Any], list[str]]:
+    """The manifest and the blob digests of the selected objects (None = all).
+
+    Metadata only — reads config entries and the document store.
     """
     from ..const import CONF_OBJECT
     from ..export import object_entries
@@ -141,25 +144,89 @@ def build_documents_archive(hass: HomeAssistant, entry_ids: set[str] | None = No
         if docs:
             manifest_objects.append({"object_id": object_id, "object_name": obj.get("name", ""), "documents": docs})
 
-    manifest = {"version": ARCHIVE_VERSION, "objects": manifest_objects}
+    return store, {"version": ARCHIVE_VERSION, "objects": manifest_objects}, sorted(blob_hashes)
 
-    buf = io.BytesIO()
-    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
+
+def _write_archive(store: Any, manifest: dict[str, Any], blob_hashes: list[str], target: Any) -> None:
+    """Write the ZIP to ``target`` (a path or binary file object) — blocking.
+
+    Each blob is streamed from disk into the archive (``ZipFile.write``), so
+    memory stays flat however large the documents are. Refuses up front a
+    selection whose files exceed MAX_ARCHIVE_BYTES: the import refuses such
+    an archive anyway, and the export used to assemble the whole ZIP in
+    memory with no ceiling (bug audit 2026-09-26).
+    """
+    blobs: list[tuple[str, Any]] = []
+    total = 0
+    for h in blob_hashes:
+        if store is None:
+            break
+        try:
+            path = store.blob_path(h)
+        except ValueError:
+            continue
+        if not path.is_file():
+            _LOGGER.warning("Documents archive: blob %s missing on disk, skipped", h[:12])
+            continue
+        total += path.stat().st_size
+        if total > MAX_ARCHIVE_BYTES:
+            raise ArchiveTooLarge(
+                f"The selected documents exceed {MAX_ARCHIVE_BYTES // (1024 * 1024)} MB, "
+                "more than an archive import accepts — export fewer objects at a time."
+            )
+        blobs.append((h, path))
+    with zipfile.ZipFile(target, "w", zipfile.ZIP_DEFLATED) as zf:
         zf.writestr(MANIFEST_NAME, json.dumps(manifest, ensure_ascii=False, indent=2))
-        for h in sorted(blob_hashes):
-            if store is None:
-                break
-            try:
-                path = store.blob_path(h)
-            except ValueError:
-                continue
-            if path.is_file():
-                zf.writestr(f"{BLOB_DIR}{h}", path.read_bytes())
-            else:
-                _LOGGER.warning("Documents archive: blob %s missing on disk, skipped", h[:12])
+        for h, path in blobs:
+            zf.write(path, arcname=f"{BLOB_DIR}{h}")
+
+
+def build_documents_archive(hass: HomeAssistant, entry_ids: set[str] | None = None) -> bytes:
+    """Build a documents ZIP in memory for the selected objects (None = all).
+
+    Blocking — call via the executor. The HTTP export streams through
+    :func:`async_build_documents_archive_file` instead; this in-memory form
+    stays for tooling and tests.
+    """
+    store, manifest, blob_hashes = _gather_archive(hass, entry_ids)
+    buf = io.BytesIO()
+    _write_archive(store, manifest, blob_hashes, buf)
     return buf.getvalue()
 
 
+async def async_build_documents_archive_file(hass: HomeAssistant, entry_ids: set[str] | None = None) -> str:
+    """Write the documents ZIP to a temporary file and return its path.
+
+    The caller streams it out and deletes it. Raises :class:`ArchiveTooLarge`
+    (nothing left behind) when the selection exceeds the import ceiling.
+    """
+    import os
+    import tempfile
+
+    store, manifest, blob_hashes = _gather_archive(hass, entry_ids)
+
+    def _write() -> str:
+        fd, path = tempfile.mkstemp(prefix="maintenance-documents-", suffix=".zip")
+        os.close(fd)
+        try:
+            _write_archive(store, manifest, blob_hashes, path)
+        except BaseException:
+            os.unlink(path)
+            raise
+        return path
+
+    return await hass.async_add_executor_job(_write)
+
+
+def _doc_key(doc: dict[str, Any]) -> tuple[str | None, str | None]:
+    """A document's identity for the idempotent re-import: kind + content
+    hash (file) or URL (link). Only strings count — a list-valued URL in a
+    crafted manifest was unhashable and crashed the set lookup."""
+    kind = doc.get("kind")
+    ref = doc.get("hash") or doc.get("url")
+    return (kind if isinstance(kind, str) else None, ref if isinstance(ref, str) else None)
+
+
 async def import_documents_archive(hass: HomeAssistant, data: bytes) -> dict[str, Any]:
     """Restore a documents ZIP: write blobs back, re-attach metadata.
 
@@ -203,67 +270,78 @@ async def import_documents_archive(hass: HomeAssistant, data: bytes) -> dict[str
     except (zipfile.BadZipFile, ValueError, json.JSONDecodeError, KeyError) as err:
         return {"error": f"invalid archive: {err}"}
 
+    # The manifest is untrusted JSON: a list at the top, an "objects" that is
+    # no list or a "documents" that is a number raised AttributeError /
+    # TypeError past the handler above — a bare 500 instead of the clean
+    # error (bug audit 2026-09-27). Wrong shapes are refused or skipped.
+    if not isinstance(manifest, dict):
+        return {"error": "invalid archive: the manifest is not an object"}
+    raw_objects = manifest.get("objects")
+    manifest_objects: list[tuple[dict[str, Any], list[Any]]] = [
+        (obj, docs)
+        for obj in (raw_objects if isinstance(raw_objects, list) else [])
+        if isinstance(obj, dict) and isinstance(docs := obj.get("documents", []), list)
+    ]
+
     # 1) Write back only blobs a manifest document actually references — an
     # archive carrying extra blobs must not litter /config with orphans that
     # ride every HA backup and are never refcounted (disk-fill hardening).
     referenced: set[str] = set()
-    for obj in manifest.get("objects", []):
-        if not isinstance(obj, dict):
-            continue
-        for m in obj.get("documents", []):
+    for _obj, docs in manifest_objects:
+        for m in docs:
             if isinstance(m, dict) and isinstance(m.get("hash"), str):
                 referenced.add(m["hash"])
 
     import hashlib
 
     written = 0
-    for digest, content in blobs.items():
-        if digest not in referenced:
-            _LOGGER.info("Documents archive: blob %s referenced by no document, skipped", digest[:12])
-            continue
-        if hashlib.sha256(content).hexdigest() != digest:
-            _LOGGER.warning("Documents archive: blob %s failed hash check, skipped", digest[:12])
-            continue
-        _, wrote_new = await hass.async_add_executor_job(store._store_blob_sync, content)
-        if wrote_new:
-            written += 1
-        store.notify_blob_added(digest)
-
-    # 2) Re-attach metadata to the matching object (id, then name).
-    ids, by_name = _object_name_map(hass)
     docs_created = 0
     objects_matched = 0
     # old → new document ids across every restored object, so history photos
     # and part doc links that pointed at the archived ids follow (the JSON
     # importer does the same for its own restore).
     doc_id_map: dict[str, str] = {}
-    for obj in manifest.get("objects", []):
-        if not isinstance(obj, dict):
-            continue
-        target = ids.get(str(obj.get("object_id") or "")) or by_name.get(str(obj.get("object_name") or ""))
-        if target is None:
-            _LOGGER.info("Documents archive: no object matches %r, its docs skipped", obj.get("object_name"))
-            continue
-        objects_matched += 1
-        # Skip docs already present on the target (idempotent re-import).
-        existing = store.for_object(target)
-        existing_keys = {(d.get("kind"), d.get("hash") or d.get("url")) for d in existing}
-        fresh = [
-            m
-            for m in obj.get("documents", [])
-            if isinstance(m, dict) and (m.get("kind"), m.get("hash") or m.get("url")) not in existing_keys
-        ]
-        if fresh:
-            # Keep task links that still resolve on the target (a same-instance
-            # restore) via an identity map over the object's current task ids;
-            # a cross-instance restore has fresh task ids, so those links drop
-            # here and are re-established by the JSON import's remap instead.
-            valid_task_ids = _object_task_ids(hass, target)
-            identity = {tid: tid for tid in valid_task_ids}
-            part_identity = {pid: pid for pid in _object_part_ids(hass, target)}
-            docs_created += await store.async_import_documents(
-                target, fresh, task_id_map=identity, part_id_map=part_identity, id_map=doc_id_map
-            )
+    # Under the store's blob lock from the first blob write to the last
+    # refcount bump: a document delete running meanwhile could remove a blob
+    # file the restore had just written (or found) before the restored
+    # document registered it — a document pointing at nothing (bug audit
+    # 2026-09-27; the upload and delete paths took the lock, this one not).
+    async with store.blob_lock:
+        for digest, content in blobs.items():
+            if digest not in referenced:
+                _LOGGER.info("Documents archive: blob %s referenced by no document, skipped", digest[:12])
+                continue
+            if hashlib.sha256(content).hexdigest() != digest:
+                _LOGGER.warning("Documents archive: blob %s failed hash check, skipped", digest[:12])
+                continue
+            _, wrote_new = await hass.async_add_executor_job(store._store_blob_sync, content)
+            if wrote_new:
+                written += 1
+            store.notify_blob_added(digest)
+
+        # 2) Re-attach metadata to the matching object (id, then name).
+        ids, by_name = _object_name_map(hass)
+        for obj, docs in manifest_objects:
+            target = ids.get(str(obj.get("object_id") or "")) or by_name.get(str(obj.get("object_name") or ""))
+            if target is None:
+                _LOGGER.info("Documents archive: no object matches %r, its docs skipped", obj.get("object_name"))
+                continue
+            objects_matched += 1
+            # Skip docs already present on the target (idempotent re-import).
+            existing = store.for_object(target)
+            existing_keys = {_doc_key(d) for d in existing}
+            fresh = [m for m in docs if isinstance(m, dict) and _doc_key(m) not in existing_keys]
+            if fresh:
+                # Keep task links that still resolve on the target (a same-instance
+                # restore) via an identity map over the object's current task ids;
+                # a cross-instance restore has fresh task ids, so those links drop
+                # here and are re-established by the JSON import's remap instead.
+                valid_task_ids = _object_task_ids(hass, target)
+                identity = {tid: tid for tid in valid_task_ids}
+                part_identity = {pid: pid for pid in _object_part_ids(hass, target)}
+                docs_created += await store.async_import_documents(
+                    target, fresh, task_id_map=identity, part_id_map=part_identity, id_map=doc_id_map
+                )
     if doc_id_map:
         await async_rewrite_doc_refs(hass, lambda old: doc_id_map.get(old, old))