diff --git a/scripts/patch_research_edition_export.py b/scripts/patch_research_edition_export.py index 2f5cb35d8..8d685c5e1 100644 --- a/scripts/patch_research_edition_export.py +++ b/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//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: diff --git a/scripts/tests/test_patch_research_edition_export.py b/scripts/tests/test_patch_research_edition_export.py index 4d82002de..0f1d5bfb4 100644 --- a/scripts/tests/test_patch_research_edition_export.py +++ b/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 { + 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,17 +52,18 @@ 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" ) @@ -69,24 +71,21 @@ def test_patch_inserts_module_and_both_call_sites(tmp_path: Path): 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():