From 889dbc59cf5595280c776b0848abda20ea370d29 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Bj=C3=A4reholt?= Date: Thu, 17 Sep 2026 17:12:46 +0200 Subject: [PATCH] fix(research): re-anchor the export sanitizer after aw-server-rust#677 streaming exports Both Research Edition legs fail on master since the v0.14.0 submodule bump (aac77dc) at "Patch research edition export hostname sanitizer": export.rs: expected exactly one export-sanitizer insertion point, found 0 aw-server-rust#677 (perf(export): serialize HTTP exports with bounded event buffering) rewrote the export path: `/api/0/export` and `/api/0/buckets//export` are now one-line calls to `BucketsExportRocket::new`, which spools the JSON through `aw_datastore::export_to_file`. No whole `BucketsExport` value exists at the endpoints any more, so the old needles in export.rs and bucket.rs are gone. The sanitizer needs the whole export (fail closed on unfiltered events, identity rewriting with collision detection across buckets), so it is now spliced into `BucketsExportRocket::new` instead: re-read the spooled JSON, sanitize, spool again. That is one insertion point for both endpoints, and it relaxes #677's memory bound only for research builds, whose exports are category-only and small. Fails closed as before if the call site drifts. Verified: patcher tests pass; patch applied to aw-server-rust b0fab73 and `cargo check -p aw-server` compiles the inserted code. --- scripts/patch_research_edition_export.py | 90 +++++++++------ .../test_patch_research_edition_export.py | 104 +++++++++--------- 2 files changed, 107 insertions(+), 87 deletions(-) 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():