Skip to main content

max / makenotwork

server: cap buffered scan download to recorded size + slack New StorageBackend::download_object_buf_capped (default method streaming via download_stream, backed by a testable read_bytestream_capped) bounds the in-memory scan branch at (file_size_bytes + SCAN_SPOOL_SLACK_BYTES).min(SCAN_MAX_MEMORY_BYTES), mirroring the spool path so a mis-recorded size can't pull an unbounded body into RAM. 3 cap tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Author: Max Johnson <me@maxj.phd> · 2026-07-05 15:50 UTC
Commit: f58201dbdf20b15a70a80cf1c01f3d1536369523
Parent: 6349845
2 files changed, +75 insertions, -2 deletions
M server/src/scanning/worker.rs +9 -2
@@ -330,9 +330,16 @@ async fn run_pipeline_and_decide(
330 330 );
331 331 super::too_large_to_scan(job.file_size_bytes as u64)
332 332 } else if (job.file_size_bytes as usize) < constants::SCAN_MAX_MEMORY_BYTES {
333 - // `download_object_buf` returns the aggregated body as `Bytes` (no
333 + // `download_object_buf_capped` returns the aggregated body as `Bytes` (no
334 334 // `to_vec` copy); `scan` takes it directly (Run #2 Performance SERIOUS).
335 - let data = backend.download_object_buf(&job.s3_key).await?;
335 + // Bound the aggregation by the recorded size + slack, never above the
336 + // in-memory threshold: `file_size_bytes` is asserted at upload and could
337 + // under-report the real object, so cap the buffered read like the spool
338 + // path does rather than pulling an unbounded body into RAM (Run 22 Perf).
339 + let cap = (job.file_size_bytes as u64)
340 + .saturating_add(constants::SCAN_SPOOL_SLACK_BYTES)
341 + .min(constants::SCAN_MAX_MEMORY_BYTES as u64);
342 + let data = backend.download_object_buf_capped(&job.s3_key, cap).await?;
336 343 let _permit = ctx.scan_semaphore.acquire().await?;
337 344 Arc::clone(&ctx.pipeline).scan(data, file_type).await
338 345 } else {
M server/src/storage.rs +66
@@ -353,6 +353,36 @@ impl S3Bucket {
353 353 }
354 354 }
355 355
356 + /// Aggregate a `ByteStream` into memory, aborting once more than `max_bytes`
357 + /// have been read. Backs [`StorageBackend::download_object_buf_capped`]; factored
358 + /// out as a free function so the cap logic is unit-testable without a full
359 + /// backend. `label` is only used in the error message (the object key).
360 + pub(crate) async fn read_bytestream_capped(
361 + mut stream: s3_storage::ByteStream,
362 + label: &str,
363 + max_bytes: u64,
364 + ) -> Result<bytes::Bytes> {
365 + let mut buf = bytes::BytesMut::new();
366 + let mut read: u64 = 0;
367 + loop {
368 + match stream.try_next().await {
369 + Ok(Some(chunk)) => {
370 + read += chunk.len() as u64;
371 + if read > max_bytes {
372 + return Err(AppError::Storage(format!(
373 + "object {label} exceeds scan in-memory cap ({read} > {max_bytes} bytes); \
374 + recorded size under-reported the real object"
375 + )));
376 + }
377 + buf.extend_from_slice(&chunk);
378 + }
379 + Ok(None) => break,
380 + Err(e) => return Err(AppError::Storage(format!("read object from S3: {e}"))),
381 + }
382 + }
383 + Ok(buf.freeze())
384 + }
385 +
356 386 /// Abstract storage backend — implemented by `S3Client` (production) and
357 387 /// `InMemoryStorage` (tests). Routes access storage through this trait.
358 388 #[async_trait::async_trait]
@@ -374,6 +404,17 @@ pub trait StorageBackend: Send + Sync {
374 404 /// drive the stream to disk (scanner spool) or to a layer that consumes
375 405 /// chunks directly.
376 406 async fn download_stream(&self, s3_key: &str) -> Result<s3_storage::ByteStream>;
407 + /// Download into memory like [`download_object_buf`], but abort if the body
408 + /// exceeds `max_bytes`. The scanner routes files it *recorded* as small to an
409 + /// in-memory branch, but `file_size_bytes` is asserted at upload time and can
410 + /// under-report the real object; this bounds the aggregation so a mis-recorded
411 + /// or abusive object can't pull an unbounded body into RAM — the independent
412 + /// ceiling the spool path already enforces (Run 22 Perf). Streams via
413 + /// `download_stream`, so no backend can hand back the whole body up front.
414 + async fn download_object_buf_capped(&self, s3_key: &str, max_bytes: u64) -> Result<bytes::Bytes> {
415 + let stream = self.download_stream(s3_key).await?;
416 + read_bytestream_capped(stream, s3_key, max_bytes).await
417 + }
377 418 /// Read up to the first `len` bytes of an object. Production overrides this
378 419 /// with a ranged `GetObject` so a content sniff transfers only the header,
379 420 /// not the whole object. The default streams and stops early — correct, but
@@ -987,6 +1028,31 @@ impl StorageBackend for S3Client {
987 1028 mod tests {
988 1029 use super::*;
989 1030
1031 + #[tokio::test]
1032 + async fn capped_read_aborts_when_body_exceeds_cap() {
1033 + // A recorded-small object whose real body is larger must not aggregate
1034 + // past the cap (Run 22 Perf — the buffered scan-download ceiling).
1035 + let stream = s3_storage::ByteStream::from(vec![0u8; 100]);
1036 + let err = read_bytestream_capped(stream, "k", 50).await.unwrap_err();
1037 + assert!(matches!(err, AppError::Storage(_)), "expected Storage error, got {err:?}");
1038 + }
1039 +
1040 + #[tokio::test]
1041 + async fn capped_read_allows_body_within_cap() {
1042 + let stream = s3_storage::ByteStream::from(vec![7u8; 40]);
1043 + let out = read_bytestream_capped(stream, "k", 50).await.unwrap();
1044 + assert_eq!(out.len(), 40);
1045 + assert!(out.iter().all(|&b| b == 7));
1046 + }
1047 +
1048 + #[tokio::test]
1049 + async fn capped_read_allows_body_exactly_at_cap() {
1050 + // The abort condition is strictly `>`, so a body equal to the cap passes.
1051 + let stream = s3_storage::ByteStream::from(vec![1u8; 50]);
1052 + let out = read_bytestream_capped(stream, "k", 50).await.unwrap();
1053 + assert_eq!(out.len(), 50);
1054 + }
1055 +
990 1056 #[test]
991 1057 fn extract_key_cdn_form() {
992 1058 let key = extract_s3_key_from_url(