From 0259ef3e29e013c47bbc206b48a7ec6af2c0dab3 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 15:36:10 -0400 Subject: [PATCH 1/2] fix(service): reject ByteStream finish_write before the declared blob size A ByteStream Write that sets finish_write=true before sending the full size declared in its resource name was accepted as a complete upload: the handler returned OK with committed_size equal to the full declared size, the blob then showed as present in FindMissingBlobs, and Read served the truncated bytes under a digest they do not hash to -- silent CAS content corruption. Both the streaming path (process_client_stream) and the oneshot path (inner_write_oneshot) only guarded the too-many-bytes case and committed on finish_write without checking that the received byte count equalled the declared size. A VerifyStore-fronted write path catches this on EOF, but an unverified store tier (filesystem/memory fast tier) commits the partial. Reject finish_write with InvalidArgument when bytes_received != expected_size in both paths. The compressed path already verifies decoded size and digest. Also fixes a latent same-class defect in max_decoding_message_size_test, whose success case declared a size that did not match the bytes it sent and only passed because of this bug. Co-Authored-By: Claude Opus 4.8 --- nativelink-service/src/bytestream_server.rs | 19 ++++++ .../tests/bytestream_server_test.rs | 61 ++++++++++++++++++- 2 files changed, 78 insertions(+), 2 deletions(-) diff --git a/nativelink-service/src/bytestream_server.rs b/nativelink-service/src/bytestream_server.rs index d42c2f6acad..e8ae7bc0007 100644 --- a/nativelink-service/src/bytestream_server.rs +++ b/nativelink-service/src/bytestream_server.rs @@ -944,6 +944,17 @@ impl ByteStreamServer { return Err(make_input_err!("Received more bytes than expected")); } if write_request.finish_write { + // The client declared the blob size in the resource name. If it + // signals completion before sending that many bytes, the upload + // is incomplete and MUST be rejected: committing it would admit a + // truncated artifact into the CAS under a digest it does not hash + // to, which later reads would serve as if it were complete. + let bytes_written = tx.get_bytes_written(); + if bytes_written != expected_size { + return Err(make_input_err!( + "Client finished writing before sending all data. Expected {expected_size} bytes but received {bytes_written}" + )); + } // Gracefully close our stream. tx.send_eof() .err_tip(|| "Failed to send EOF in ByteStream::write")?; @@ -1058,6 +1069,14 @@ impl ByteStreamServer { } if write_request.finish_write { + // Reject an incomplete upload: a finish_write before the declared + // size would otherwise commit a truncated blob under a digest it + // does not hash to. See inner_write for the full rationale. + if bytes_received != expected_size { + return Err(make_input_err!( + "Client finished writing before sending all data. Expected {expected_size} bytes but received {bytes_received}" + )); + } break; } } diff --git a/nativelink-service/tests/bytestream_server_test.rs b/nativelink-service/tests/bytestream_server_test.rs index 20da09e6bde..d67bf44b142 100644 --- a/nativelink-service/tests/bytestream_server_test.rs +++ b/nativelink-service/tests/bytestream_server_test.rs @@ -357,6 +357,60 @@ pub async fn chunked_stream_receives_all_data() -> Result<(), Box Result<(), Box> { + let store_manager = make_store_manager().await?; + let bs_server = Arc::new( + make_bytestream_server(store_manager.as_ref(), None).expect("Failed to make server"), + ); + let store = store_manager.get_store("main_cas").unwrap(); + + let raw_data = b"12456789abcdefghijk"; + let declared_len = raw_data.len(); + + let (tx, join_handle) = make_stream_and_writer_spawn(bs_server, Some(CompressionEncoding::Gzip)); + + // Declare the full length in the resource name, but only send a short + // prefix and immediately set finish_write=true. + const SHORT_PREFIX: usize = 5; + let resource_name = format!( + "{}/uploads/{}/blobs/{}/{}", + INSTANCE_NAME, "4dcec57e-1389-4ab5-b188-4a59f22ceb4b", HASH1, declared_len + ); + let write_request = WriteRequest { + resource_name, + write_offset: 0, + finish_write: true, + data: raw_data[..SHORT_PREFIX].into(), + }; + tx.send(Frame::data(encode_stream_proto(&write_request)?)) + .await?; + + // The write MUST fail (not return OK with a full committed_size). + let server_result = join_handle.await.expect("Failed to join"); + assert!( + server_result.is_err(), + "Short finish_write must be rejected, got Ok: {server_result:?}" + ); + + // The blob MUST NOT be present under its declared digest. + assert!( + store + .has(DigestInfo::try_new(HASH1, declared_len)?) + .await? + .is_none(), + "Truncated blob must not be committed to the store", + ); + + Ok(()) +} + #[nativelink_test] pub async fn sha256_write_without_digest_segment_populates_cache() -> Result<(), Box> { @@ -1876,9 +1930,12 @@ pub async fn max_decoding_message_size_test() -> Result<(), Box Date: Fri, 2 Oct 2026 18:51:58 -0400 Subject: [PATCH 2/2] fix(cas): reject short/torn uploads and report honest write status Unifies the handler-level short finish_write rejection (N13, e5d1b0db) with two further CAS-integrity fixes for the same disclosure (TRA-1373): - filesystem_store: refuse to emplace a blob whose byte count disagrees with an ExactSize upload hint. A torn/aborted upload stream closes the writer half, which arrives at the store as a clean EOF; without this commit-point guard the partial bytes are stored under the full digest key regardless of any handler check (finish_write never arrives on the abort path). - bytestream_server: QueryWriteStatus clamps committed_size to the declared digest size (the filesystem store reports block-rounded on-disk footprint, e.g. 17 -> 4096) and only reports complete once the store holds at least that many bytes, instead of unconditional complete=true. Regression tests: - short_finish_write_is_rejected_and_not_committed (oneshot path, N13) - short_finish_write_is_rejected (streaming multi-message path) - aborted_mid_stream_upload_is_not_committed (EOF without finish_write; only the store-level guard can catch this) - max_decoding_message_size_test updated: it asserted the old short-upload-accepted behavior. Co-Authored-By: Claude Opus 4.8 --- nativelink-service/src/bytestream_server.rs | 12 +++- .../tests/bytestream_server_test.rs | 65 +++++++++++++++++- nativelink-store/src/filesystem_store.rs | 20 +++++- .../tests/filesystem_store_test.rs | 67 +++++++++++++++++++ 4 files changed, 157 insertions(+), 7 deletions(-) diff --git a/nativelink-service/src/bytestream_server.rs b/nativelink-service/src/bytestream_server.rs index e8ae7bc0007..c18cf49dff6 100644 --- a/nativelink-service/src/bytestream_server.rs +++ b/nativelink-service/src/bytestream_server.rs @@ -1392,9 +1392,17 @@ impl ByteStreamServer { complete: false, })); }; + // `item_size` is the store's reported size; for the filesystem store + // that is the block-rounded on-disk footprint, not the logical byte + // count. Clamp the reported committed_size to the declared digest size + // and only report `complete` once the store holds at least that many + // bytes, so QueryWriteStatus never over-reports nor falsely completes. + let expected_size = digest.size_bytes(); + let complete = item_size >= expected_size; + let committed_size = if complete { expected_size } else { item_size }; Ok(Response::new(QueryWriteStatusResponse { - committed_size: item_size.try_into().unwrap_or(i64::MAX), - complete: true, + committed_size: committed_size.try_into().unwrap_or(i64::MAX), + complete, })) } } diff --git a/nativelink-service/tests/bytestream_server_test.rs b/nativelink-service/tests/bytestream_server_test.rs index d67bf44b142..972239cdee7 100644 --- a/nativelink-service/tests/bytestream_server_test.rs +++ b/nativelink-service/tests/bytestream_server_test.rs @@ -365,6 +365,7 @@ pub async fn chunked_stream_receives_all_data() -> Result<(), Box Result<(), Box> { + const SHORT_PREFIX: usize = 5; let store_manager = make_store_manager().await?; let bs_server = Arc::new( make_bytestream_server(store_manager.as_ref(), None).expect("Failed to make server"), @@ -374,11 +375,11 @@ pub async fn short_finish_write_is_rejected_and_not_committed() let raw_data = b"12456789abcdefghijk"; let declared_len = raw_data.len(); - let (tx, join_handle) = make_stream_and_writer_spawn(bs_server, Some(CompressionEncoding::Gzip)); + let (tx, join_handle) = + make_stream_and_writer_spawn(bs_server, Some(CompressionEncoding::Gzip)); // Declare the full length in the resource name, but only send a short // prefix and immediately set finish_write=true. - const SHORT_PREFIX: usize = 5; let resource_name = format!( "{}/uploads/{}/blobs/{}/{}", INSTANCE_NAME, "4dcec57e-1389-4ab5-b188-4a59f22ceb4b", HASH1, declared_len @@ -1902,6 +1903,62 @@ pub async fn test_query_write_status_smoke_test() -> Result<(), Box Result<(), Box> { + let store_manager = make_store_manager().await?; + let bs_server = Arc::new( + make_bytestream_server(store_manager.as_ref(), None).expect("Failed to make server"), + ); + let store = store_manager.get_store("main_cas").unwrap(); + + // Declare a 19-byte blob, but only send 12 bytes across two messages and + // then finish_write (multi-message forces the streaming write path). + let raw_data = b"12456789abcdefghijk"; + let resource_name = format!( + "{}/uploads/{}/blobs/{}/{}", + INSTANCE_NAME, + "4dcec57e-1389-4ab5-b188-4a59f22ceb4b", + HASH1, + raw_data.len() + ); + let (tx, join_handle) = + make_stream_and_writer_spawn(bs_server, Some(CompressionEncoding::Gzip)); + + let mut write_request = WriteRequest { + resource_name, + write_offset: 0, + finish_write: false, + data: raw_data[..8].into(), + }; + tx.send(Frame::data(encode_stream_proto(&write_request)?)) + .await?; + // Short final chunk: total 12 < declared 19. + write_request.write_offset = 8; + write_request.data = raw_data[8..12].into(); + write_request.finish_write = true; + tx.send(Frame::data(encode_stream_proto(&write_request)?)) + .await?; + + let server_result = join_handle.await.expect("Failed to join"); + assert!( + server_result.is_err(), + "short finish_write must be rejected, got {server_result:?}" + ); + assert!( + store + .has(DigestInfo::try_new(HASH1, raw_data.len())?) + .await? + .is_none(), + "a rejected short upload must not leave a readable partial blob", + ); + Ok(()) +} + #[nativelink_test] pub async fn max_decoding_message_size_test() -> Result<(), Box> { const MAX_MESSAGE_SIZE: usize = 1024 * 1024; // 1MB. @@ -1956,7 +2013,9 @@ pub async fn max_decoding_message_size_test() -> Result<(), Box FilesystemStore { mut temp_file: FileSlot, final_key: StoreKey<'static>, mut reader: DropCloserReadHalf, + upload_size: UploadSizeInfo, ) -> Result { let mut data_size = 0; loop { @@ -1860,6 +1861,21 @@ impl FilesystemStore { data_size += data_len as u64; } + // Refuse to commit a blob whose byte count disagrees with its declared + // digest size. A torn or aborted upload stream closes the writer half, + // which arrives here as a normal EOF; without this check the partial + // bytes would be emplaced under the full digest's key and later read + // back OK with a mismatched hash (found by the rechaos chaos monkey). + if let UploadSizeInfo::ExactSize(expected_size) = upload_size + && data_size != expected_size + { + return Err(make_err!( + Code::InvalidArgument, + "Filesystem store received {data_size} bytes but the digest declares \ + {expected_size}; refusing to store a partial blob" + )); + } + let permit = if let Some(sem) = &self.write_semaphore { Some(sem.acquire().await.map_err(|err| { Error::from_std_err(Code::Internal, &err).append("Write semaphore closed") @@ -2122,7 +2138,7 @@ impl StoreDriver for FilesystemStore { self: Pin<&Self>, key: StoreKey<'_>, mut reader: DropCloserReadHalf, - _upload_size: UploadSizeInfo, + upload_size: UploadSizeInfo, ) -> Result { if is_zero_digest(key.borrow()) { // don't need to add, because zero length files are just assumed to exist. @@ -2141,7 +2157,7 @@ impl StoreDriver for FilesystemStore { let (entry, temp_file, temp_full_path) = self.make_temp_file(temp_key).await?; - self.update_file(entry, temp_file, key.into_owned(), reader) + self.update_file(entry, temp_file, key.into_owned(), reader, upload_size) .await .err_tip(|| { format!( diff --git a/nativelink-store/tests/filesystem_store_test.rs b/nativelink-store/tests/filesystem_store_test.rs index 27f8bf1833c..3381af555d2 100644 --- a/nativelink-store/tests/filesystem_store_test.rs +++ b/nativelink-store/tests/filesystem_store_test.rs @@ -3118,3 +3118,70 @@ async fn partial_migration_stays_readable_and_retries_next_startup() -> Result<( Ok(()) } + +// Regression (TRA-1373, aborted-mid-stream variant): a ByteStream Write that is +// torn mid-stream (client abort, injected UNAVAILABLE, etc.) closes the writer +// half, which arrives at the store as a normal EOF with fewer bytes than the +// digest declares -- finish_write is never sent, so a handler-level check can +// never fire. The store itself must refuse to emplace the partial blob at the +// commit point when the upload declared an exact size. +#[nativelink_test] +async fn aborted_mid_stream_upload_is_not_committed() -> Result<(), Error> { + const DECLARED_SIZE: usize = 19; + const PARTIAL: &[u8] = b"12456789abcd"; // 12 bytes < declared 19. + + let content_path = make_temp_path("content_path"); + let temp_path = make_temp_path("temp_path"); + + let store = Box::pin( + FilesystemStore::::new(&FilesystemSpec { + content_path, + temp_path, + read_buffer_size: 1, + ..Default::default() + }) + .await?, + ); + + let digest = DigestInfo::try_new(HASH1, DECLARED_SIZE)?; + + let (mut writer, reader) = make_buf_channel_pair(); + let store_ref = &store; + let update_fut = async move { + store_ref + .update( + digest, + reader, + UploadSizeInfo::ExactSize(DECLARED_SIZE as u64), + ) + .await + }; + // Simulate the torn stream: partial bytes, then the writer half closes + // cleanly (EOF) without ever having sent the declared byte count. + let writer_fut = async move { + writer.send(Bytes::from_static(PARTIAL)).await?; + writer.send_eof()?; + Ok::<_, Error>(()) + }; + + let (update_result, writer_result) = tokio::join!(update_fut, writer_fut); + writer_result?; + + assert!( + update_result.is_err(), + "A torn upload (EOF before the declared size) must not be committed, got {update_result:?}" + ); + assert_eq!( + update_result.unwrap_err().code, + Code::InvalidArgument, + "Expected the store-level size guard to reject the partial blob" + ); + + // The partial blob MUST NOT be present under its declared digest. + assert_eq!( + store.has(digest).await?, + None, + "Partial blob must not be readable under the full digest" + ); + Ok(()) +}