Skip to content

Revert "Reject short and torn uploads in ByteStream and FilesystemStore, and report QueryWriteStatus honestly (#2881)" #2953

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
MarcusSorealheis merged 1 commit into from
Oct 16, 2026
Merged

Revert "Reject short and torn uploads in ByteStream and FilesystemStore, and report QueryWriteStatus honestly (#2881)" #2953

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 2 additions & 29 deletions nativelink-service/src/bytestream_server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")?;
Expand Down Expand Up @@ -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;
}
}
Expand Down Expand Up @@ -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,
}))
}
}
Expand Down
122 changes: 3 additions & 119 deletions nativelink-service/tests/bytestream_server_test.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -357,61 +357,6 @@ pub async fn chunked_stream_receives_all_data() -> Result<(), Box<dyn core::erro
Ok(())
}

// Regression: a Write that signals finish_write before sending the full
// declared blob size MUST be rejected, even for a store with no size/hash
// verification (e.g. a plain memory/filesystem fast tier). Otherwise the CAS
// admits a truncated artifact that later reads serve under a digest it does not
// hash to -- silent build corruption.
#[nativelink_test]
pub async fn short_finish_write_is_rejected_and_not_committed()
-> Result<(), Box<dyn core::error::Error>> {
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<dyn core::error::Error>> {
Expand Down Expand Up @@ -1903,62 +1848,6 @@ pub async fn test_query_write_status_smoke_test() -> Result<(), Box<dyn core::er
Ok(())
}

// Regression: a client that declares a digest of N bytes but sends fewer and
// then sets finish_write must be rejected, and must not leave a readable,
// hash-mismatched partial blob behind. Previously the server accepted the short
// stream, reported committed_size == N, stored the partial bytes, and served
// them on a subsequent Read with OK (found by the rechaos chaos monkey).
#[nativelink_test]
pub async fn short_finish_write_is_rejected() -> Result<(), Box<dyn core::error::Error>> {
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<dyn core::error::Error>> {
const MAX_MESSAGE_SIZE: usize = 1024 * 1024; // 1MB.
Expand Down Expand Up @@ -1987,12 +1876,9 @@ pub async fn max_decoding_message_size_test() -> Result<(), Box<dyn core::error:

{
// Test to ensure if we send exactly our max message size, it will succeed.
let data_len = MAX_MESSAGE_SIZE - WRITE_REQUEST_MSG_WRAPPER_SIZE;
let data = Bytes::from(vec![0u8; data_len]);
let data = Bytes::from(vec![0u8; MAX_MESSAGE_SIZE - WRITE_REQUEST_MSG_WRAPPER_SIZE]);
let write_request = WriteRequest {
// The declared blob size must equal the bytes actually sent, otherwise
// the server (correctly) rejects it as a short finish_write.
resource_name: make_resource_name(data_len),
resource_name: make_resource_name(MAX_MESSAGE_SIZE),
write_offset: 0,
finish_write: true,
data,
Expand All @@ -2013,9 +1899,7 @@ pub async fn max_decoding_message_size_test() -> Result<(), Box<dyn core::error:
MAX_MESSAGE_SIZE - WRITE_REQUEST_MSG_WRAPPER_SIZE + 1
]);
let write_request = WriteRequest {
resource_name: make_resource_name(
MAX_MESSAGE_SIZE - WRITE_REQUEST_MSG_WRAPPER_SIZE + 1,
),
resource_name: make_resource_name(MAX_MESSAGE_SIZE),
write_offset: 0,
finish_write: true,
data,
Expand Down
20 changes: 2 additions & 18 deletions nativelink-store/src/filesystem_store.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1842,7 +1842,6 @@ impl<Fe: FileEntry> FilesystemStore<Fe> {
mut temp_file: FileSlot,
final_key: StoreKey<'static>,
mut reader: DropCloserReadHalf,
upload_size: UploadSizeInfo,
) -> Result<u64, Error> {
let mut data_size = 0;
loop {
Expand All @@ -1861,21 +1860,6 @@ impl<Fe: FileEntry> FilesystemStore<Fe> {
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")
Expand Down Expand Up @@ -2138,7 +2122,7 @@ impl<Fe: FileEntry> StoreDriver for FilesystemStore<Fe> {
self: Pin<&Self>,
key: StoreKey<'_>,
mut reader: DropCloserReadHalf,
upload_size: UploadSizeInfo,
_upload_size: UploadSizeInfo,
) -> Result<u64, Error> {
if is_zero_digest(key.borrow()) {
// don't need to add, because zero length files are just assumed to exist.
Expand All @@ -2157,7 +2141,7 @@ impl<Fe: FileEntry> StoreDriver for FilesystemStore<Fe> {

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!(
Expand Down
67 changes: 0 additions & 67 deletions nativelink-store/tests/filesystem_store_test.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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::<FileEntryImpl>::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(())
}
Loading