diff --git a/nativelink-store/src/verify_store.rs b/nativelink-store/src/verify_store.rs index 449758300..c5178c397 100644 --- a/nativelink-store/src/verify_store.rs +++ b/nativelink-store/src/verify_store.rs @@ -72,8 +72,7 @@ impl VerifyStore { .err_tip(|| "Failed to read chunk in check_update in verify store")?; sum_size += chunk.len() as u64; - let mut done = chunk.is_empty(); // Is EOF. - + // Ensure if a user sends us too much data we fail quickly. if let Some(expected_size) = maybe_expected_digest_size { match sum_size.cmp(&expected_size) { std::cmp::Ordering::Greater => { @@ -85,24 +84,25 @@ impl VerifyStore { )); } std::cmp::Ordering::Equal => { - let eof_chunk = rx - .recv() - .await - .err_tip(|| "Failed to read eof_chunk in verify store")?; - if !eof_chunk.is_empty() { - self.size_verification_failures.inc(); - return Err(make_input_err!( - "Expected EOF chunk when exact size was hit on insert in verify store - {}", - expected_size, - )); + // Ensure our next chunk is the EOF chunk. + // If this was an error it'll be caught on the .recv() + // on next cycle. + if let Ok(eof_chunk) = rx.peek().await { + if !eof_chunk.is_empty() { + self.size_verification_failures.inc(); + return Err(make_input_err!( + "Expected EOF chunk when exact size was hit on insert in verify store - {}", + expected_size, + )); + } } - done = true; } std::cmp::Ordering::Less => {} } } - if done { + // If is EOF. + if chunk.is_empty() { if let Some(expected_size) = maybe_expected_digest_size { if sum_size != expected_size { self.size_verification_failures.inc(); @@ -124,9 +124,6 @@ impl VerifyStore { )); } } - } - - if chunk.is_empty() { tx.send_eof().err_tip(|| "In verify_store::check_update")?; break; } @@ -141,13 +138,6 @@ impl VerifyStore { write_future .await .err_tip(|| "Failed to write chunk to inner store in verify store")?; - - // If are done, but not an empty `chunk`, it means we are at the exact - // size match and already received the EOF chunk above. - if done { - tx.send_eof().err_tip(|| "In verify_store::check_update")?; - break; - } } Ok(()) } diff --git a/nativelink-store/tests/verify_store_test.rs b/nativelink-store/tests/verify_store_test.rs index 21bc79313..d3c24b030 100644 --- a/nativelink-store/tests/verify_store_test.rs +++ b/nativelink-store/tests/verify_store_test.rs @@ -351,3 +351,31 @@ async fn verify_fails_immediately_on_too_much_data_sent_update() -> Result<(), E ); Ok(()) } + +#[nativelink_test] +async fn verify_size_and_hash_suceeds_on_small_data() -> Result<(), Error> { + let inner_store = MemoryStore::new(&nativelink_config::stores::MemoryStore::default()); + let store = VerifyStore::new( + &nativelink_config::stores::VerifyStore { + backend: nativelink_config::stores::StoreConfig::memory( + nativelink_config::stores::MemoryStore::default(), + ), + verify_size: true, + verify_hash: true, + }, + Store::new(inner_store.clone()), + ); + + /// This value is sha256("123"). + const HASH: &str = "a665a45920422f9d417e4867efdc4fb8a04a1f3fff1fa07e998e86f7f7a27ae3"; + const VALUE: &str = "123"; + let digest = DigestInfo::try_new(HASH, 3).unwrap(); + let result = store.update_oneshot(digest, VALUE.into()).await; + assert_eq!(result, Ok(()), "Expected success, got: {:?}", result); + assert_eq!( + inner_store.has(digest).await, + Ok(Some(VALUE.len())), + "Expected data to exist in store after update" + ); + Ok(()) +}