fix(research): re-anchor the export sanitizer after aw-server-rust#677 streaming exports

scripts/patch_research_edition_export.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Patch aw-server-rust export endpoints for the Research Edition build.
+"""Patch aw-server-rust's export path for the Research Edition build.
 
 Run as part of the CI build for research edition:
 
@@ -8,6 +8,17 @@
 Standard builds never run this script, so `/api/0/export` stays byte-for-byte
 unchanged outside Research Edition. The patch is fail-closed: missing markers
 abort the build rather than shipping an unsanitized artifact.
+
+Where the sanitizer hooks in: since aw-server-rust#677 the export is streamed
+by `aw_datastore::export_to_file` with bounded event buffering, and both
+`/api/0/export` and `/api/0/buckets/<id>/export` are one-line calls to
+`BucketsExportRocket::new` in `endpoints/util.rs`. No whole `BucketsExport`
+value exists at the endpoints any more, so the sanitizer — which needs the
+whole export to fail closed on unfiltered events and to detect identity
+collisions across buckets — is spliced into `BucketsExportRocket::new`: the
+spooled JSON is re-read, sanitized, and spooled again. Research exports are
+category-only, so this relaxes the streaming memory bound only where the
+data is already small.
 """
 from __future__ import annotations
 
@@ -17,38 +28,47 @@
 
 MARKER = "RESEARCH_EDITION_EXPORT_SANITIZE"
 
-EXPORT_INSERT_NEEDLE = """        export.buckets.insert(bid, bucket);
-    }
-
-    Ok(export.into())
-"""
-
-EXPORT_INSERT_REPLACEMENT = f"""        export.buckets.insert(bid, bucket);
-    }}
-
-    // {MARKER}
-    let export = match super::export_sanitize::sanitize_buckets_export(export) {{
-        Ok(export) => export,
-        Err(err) => {{
-            return Err(HttpErrorJson::new(rocket::http::Status::Conflict, err))
-        }}
-    }};
-    Ok(export.into())
+# The two lines in BucketsExportRocket::new that spool the export and rewind
+# it. The sanitizer is inserted right after them and shadows `file`.
+EXPORT_INSERT_NEEDLE = """        let (mut file, name) = datastore.export_to_file(bucket_id, file)?;
+        file.seek(SeekFrom::Start(0)).map_err(io_error)?;
 """
 
-BUCKET_INSERT_NEEDLE = """    export.buckets.insert(bucket_id.into(), bucket);
-
-    Ok(export.into())
-"""
-
-BUCKET_INSERT_REPLACEMENT = f"""    export.buckets.insert(bucket_id.into(), bucket);
-
-    // {MARKER}
-    let export = match super::export_sanitize::sanitize_buckets_export(export) {{
-        Ok(export) => export,
-        Err(err) => return Err(HttpErrorJson::new(Status::Conflict, err)),
-    }};
-    Ok(export.into())
+EXPORT_INSERT_REPLACEMENT = f"""        let (mut file, name) = datastore.export_to_file(bucket_id, file)?;
+        file.seek(SeekFrom::Start(0)).map_err(io_error)?;
+        // {MARKER}
+        // The datastore streams the export with bounded event buffering, so no
+        // whole `BucketsExport` exists here. The Research Edition sanitizer
+        // needs one (fail closed on unfiltered events, identity rewriting with
+        // collision detection across buckets): re-read the spooled JSON,
+        // sanitize, and spool again. Research exports are category-only, so
+        // this relaxes the memory bound only where the data is already small.
+        let file = {{
+            let export: aw_models::BucketsExport =
+                serde_json::from_reader(std::io::BufReader::new(&file)).map_err(|err| {{
+                    error!("Failed to parse export for sanitizing: {{err}}");
+                    HttpErrorJson::new(
+                        Status::InternalServerError,
+                        "Failed to prepare export file".into(),
+                    )
+                }})?;
+            let export = super::export_sanitize::sanitize_buckets_export(export)
+                .map_err(|err| HttpErrorJson::new(Status::Conflict, err))?;
+            let mut sanitized = tempfile::tempfile().map_err(io_error)?;
+            {{
+                let mut writer = std::io::BufWriter::new(&mut sanitized);
+                serde_json::to_writer(&mut writer, &export).map_err(|err| {{
+                    error!("Failed to write sanitized export: {{err}}");
+                    HttpErrorJson::new(
+                        Status::InternalServerError,
+                        "Failed to prepare export file".into(),
+                    )
+                }})?;
+                std::io::Write::flush(&mut writer).map_err(io_error)?;
+            }}
+            sanitized.seek(SeekFrom::Start(0)).map_err(io_error)?;
+            sanitized
+        }};
 """
 
 MOD_NEEDLE = "mod export;\n"
@@ -82,20 +102,18 @@ def patch_tree(repo_root: pathlib.Path) -> None:
     endpoints = (
         repo_root / "aw-server-rust" / "aw-server" / "src" / "endpoints"
     )
-    export_rs = endpoints / "export.rs"
-    bucket_rs = endpoints / "bucket.rs"
+    util_rs = endpoints / "util.rs"
     mod_rs = endpoints / "mod.rs"
     dest = endpoints / "export_sanitize.rs"
 
-    for required in (export_rs, bucket_rs, mod_rs):
+    for required in (util_rs, mod_rs):
         if not required.is_file():
             raise FileNotFoundError(f"expected Rust export source at {required}")
 
     shutil.copyfile(source, dest)
 
     _replace_once(mod_rs, MOD_NEEDLE, MOD_REPLACEMENT, "mod export_sanitize;")
-    _replace_once(export_rs, EXPORT_INSERT_NEEDLE, EXPORT_INSERT_REPLACEMENT, MARKER)
-    _replace_once(bucket_rs, BUCKET_INSERT_NEEDLE, BUCKET_INSERT_REPLACEMENT, MARKER)
+    _replace_once(util_rs, EXPORT_INSERT_NEEDLE, EXPORT_INSERT_REPLACEMENT, MARKER)
 
 
 def main() -> None:

scripts/tests/test_patch_research_edition_export.py

@@ -18,30 +18,31 @@
 CONFIG_SPEC.loader.exec_module(config_patcher)
 
 
-def _write_tree(tmp_path: Path, export: str, bucket: str, mod: str) -> Path:
+def _write_tree(tmp_path: Path, util: str, mod: str) -> Path:
     endpoints = tmp_path / "aw-server-rust" / "aw-server" / "src" / "endpoints"
     endpoints.mkdir(parents=True)
-    (endpoints / "export.rs").write_text(export, encoding="utf-8")
-    (endpoints / "bucket.rs").write_text(bucket, encoding="utf-8")
+    (endpoints / "util.rs").write_text(util, encoding="utf-8")
     (endpoints / "mod.rs").write_text(mod, encoding="utf-8")
     return tmp_path
 
 
-EXPORT_SRC = """use std::collections::HashMap;
-
-pub fn buckets_export() {
-    for (bid, mut bucket) in buckets.drain() {
-        export.buckets.insert(bid, bucket);
+# Shape of BucketsExportRocket::new after aw-server-rust#677: the export is
+# spooled to a tempfile by the datastore and rewound. Both export endpoints
+# go through this one function, so it is the single insertion point.
+UTIL_SRC = """impl BucketsExportRocket {
+    pub fn new(
+        datastore: &aw_datastore::Datastore,
+        bucket_id: Option<&str>,
+    ) -> Result<Self, HttpErrorJson> {
+        let file = tempfile::tempfile().map_err(io_error)?;
+        let (mut file, name) = datastore.export_to_file(bucket_id, file)?;
+        file.seek(SeekFrom::Start(0)).map_err(io_error)?;
+        let filename = match name {
+            Some(id) => format!("attachment; filename=aw-bucket-export_{id}.json"),
+            None => "attachment; filename=aw-buckets-export.json".into(),
+        };
+        Ok(Self { file, filename })
     }
-
-    Ok(export.into())
-}
-"""
-
-BUCKET_SRC = """pub fn bucket_export() {
-    export.buckets.insert(bucket_id.into(), bucket);
-
-    Ok(export.into())
 }
 """
 
@@ -51,42 +52,40 @@ def _write_tree(tmp_path: Path, export: str, bucket: str, mod: str) -> Path:
 """
 
 
-def test_patch_inserts_module_and_both_call_sites(tmp_path: Path):
-    root = _write_tree(tmp_path, EXPORT_SRC, BUCKET_SRC, MOD_SRC)
+def _util(root: Path) -> str:
+    return (root / "aw-server-rust/aw-server/src/endpoints/util.rs").read_text(
+        encoding="utf-8"
+    )
+
+
+def test_patch_inserts_module_and_call_site(tmp_path: Path):
+    root = _write_tree(tmp_path, UTIL_SRC, MOD_SRC)
 
     patcher.patch_tree(root)
 
-    export = (root / "aw-server-rust/aw-server/src/endpoints/export.rs").read_text(
-        encoding="utf-8"
-    )
-    bucket = (root / "aw-server-rust/aw-server/src/endpoints/bucket.rs").read_text(
-        encoding="utf-8"
-    )
+    util = _util(root)
     mod = (root / "aw-server-rust/aw-server/src/endpoints/mod.rs").read_text(
         encoding="utf-8"
     )
     copied = root / "aw-server-rust/aw-server/src/endpoints/export_sanitize.rs"
 
     assert copied.is_file()
     assert "mod export_sanitize;" in mod
-    assert patcher.MARKER in export
-    assert patcher.MARKER in bucket
-    assert "sanitize_buckets_export" in export
-    assert "sanitize_buckets_export" in bucket
-    assert "Status::Conflict" in export
-    assert "Status::Conflict" in bucket
+    assert util.count(patcher.MARKER) == 1
+    assert "sanitize_buckets_export" in util
+    assert "Status::Conflict" in util
+    # The sanitized spool replaces `file` before the filename is chosen and
+    # the struct is built, so the response body is the sanitized JSON.
+    assert util.index(patcher.MARKER) < util.index("let filename = match name")
+    assert "sanitized.seek(SeekFrom::Start(0))" in util
 
 
 def test_patch_is_idempotent(tmp_path: Path):
-    root = _write_tree(tmp_path, EXPORT_SRC, BUCKET_SRC, MOD_SRC)
+    root = _write_tree(tmp_path, UTIL_SRC, MOD_SRC)
     patcher.patch_tree(root)
-    first = (root / "aw-server-rust/aw-server/src/endpoints/export.rs").read_text(
-        encoding="utf-8"
-    )
+    first = _util(root)
     patcher.patch_tree(root)
-    second = (root / "aw-server-rust/aw-server/src/endpoints/export.rs").read_text(
-        encoding="utf-8"
-    )
+    second = _util(root)
     assert first == second
     mod = (root / "aw-server-rust/aw-server/src/endpoints/mod.rs").read_text(
         encoding="utf-8"
@@ -95,27 +94,30 @@ def test_patch_is_idempotent(tmp_path: Path):
 
 
 def test_patch_fails_closed_without_export_marker(tmp_path: Path):
-    root = _write_tree(tmp_path, "fn buckets_export() {}\n", BUCKET_SRC, MOD_SRC)
+    root = _write_tree(tmp_path, "impl BucketsExportRocket {}\n", MOD_SRC)
     with pytest.raises(ValueError, match="insertion point"):
         patcher.patch_tree(root)
 
 
+def test_patch_fails_closed_on_pre_677_endpoints(tmp_path: Path):
+    # A tree where the endpoints still build `export.buckets` themselves has
+    # no spooling call site; the patch must refuse rather than ship unsanitized.
+    legacy_util = "impl BucketsExportRocket {\n    pub fn new() {}\n}\n"
+    root = _write_tree(tmp_path, legacy_util, MOD_SRC)
+    with pytest.raises(ValueError, match="found 0"):
+        patcher.patch_tree(root)
+
+
 def test_live_tree_is_patchable_or_already_patched():
     root = Path(__file__).resolve().parents[2]
-    export = root / "aw-server-rust/aw-server/src/endpoints/export.rs"
-    bucket = root / "aw-server-rust/aw-server/src/endpoints/bucket.rs"
-    if not export.is_file() or not bucket.is_file():
+    util = root / "aw-server-rust/aw-server/src/endpoints/util.rs"
+    if not util.is_file():
         pytest.skip("aw-server-rust not checked out")
-    export_text = export.read_text(encoding="utf-8")
-    bucket_text = bucket.read_text(encoding="utf-8")
+    util_text = util.read_text(encoding="utf-8")
     assert (
-        patcher.MARKER in export_text
-        or export_text.count(patcher.EXPORT_INSERT_NEEDLE) == 1
-    )
-    assert (
-        patcher.MARKER in bucket_text
-        or bucket_text.count(patcher.BUCKET_INSERT_NEEDLE) == 1
-    )
+        patcher.MARKER in util_text
+        or util_text.count(patcher.EXPORT_INSERT_NEEDLE) == 1
+    ), "BucketsExportRocket::new no longer matches the research export patch; update EXPORT_INSERT_NEEDLE"
 
 
 def test_sanitizer_allowlist_covers_config_categories():