diff --git a/nativelink-service/src/bytestream_server.rs b/nativelink-service/src/bytestream_server.rs index d42c2f6acad..c18cf49dff6 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; } } @@ -1373,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 20da09e6bde..972239cdee7 100644 --- a/nativelink-service/tests/bytestream_server_test.rs +++ b/nativelink-service/tests/bytestream_server_test.rs @@ -357,6 +357,61 @@ 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"), + ); + 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. + 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> { @@ -1848,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. @@ -1876,9 +1987,12 @@ pub async fn max_decoding_message_size_test() -> Result<(), Box 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(()) +}