diff --git a/nativelink-service/src/bytestream_server.rs b/nativelink-service/src/bytestream_server.rs index c18cf49dff..d42c2f6aca 100644 --- a/nativelink-service/src/bytestream_server.rs +++ b/nativelink-service/src/bytestream_server.rs @@ -944,17 +944,6 @@ 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")?; @@ -1069,14 +1058,6 @@ 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; } } @@ -1392,17 +1373,9 @@ 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: committed_size.try_into().unwrap_or(i64::MAX), - complete, + committed_size: item_size.try_into().unwrap_or(i64::MAX), + complete: true, })) } } diff --git a/nativelink-service/tests/bytestream_server_test.rs b/nativelink-service/tests/bytestream_server_test.rs index 972239cdee..20da09e6bd 100644 --- a/nativelink-service/tests/bytestream_server_test.rs +++ b/nativelink-service/tests/bytestream_server_test.rs @@ -357,61 +357,6 @@ 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> { @@ -1903,62 +1848,6 @@ 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. @@ -1987,12 +1876,9 @@ 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 { @@ -1861,21 +1860,6 @@ 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") @@ -2138,7 +2122,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. @@ -2157,7 +2141,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, upload_size) + self.update_file(entry, temp_file, key.into_owned(), reader) .await .err_tip(|| { format!( diff --git a/nativelink-store/tests/filesystem_store_test.rs b/nativelink-store/tests/filesystem_store_test.rs index 3381af555d..27f8bf1833 100644 --- a/nativelink-store/tests/filesystem_store_test.rs +++ b/nativelink-store/tests/filesystem_store_test.rs @@ -3118,70 +3118,3 @@ 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(()) -}