store(dedup): reject torn content-chunk reads instead of serving them as success

nativelink-store/src/dedup_store.rs

@@ -352,6 +352,55 @@ impl StoreDriver for DedupStore {
                     .await
                     .err_tip(|| "Failed to get_part in content_store in dedup_store")?;
 
+                // Content chunks are content-addressed and whole: a chunk must
+                // contain exactly the number of bytes its index entry declares.
+                // If the content store returns a chunk of any other size (a
+                // torn or partial read, or an index/content size disagreement)
+                // the streaming loop below would either panic on the
+                // `bytes_to_skip` assert or silently under-deliver and send EOF
+                // as success. Reject the mismatch here so it surfaces as an
+                // error instead.
+                let declared = index_entry.size_bytes();
+                let actual = data.len() as u64;
+                if actual != declared {
+                    // Every blob that references this chunk is unreadable while
+                    // it stays, and nothing else would notice: `has()` sums the
+                    // sizes the content store reports, and `update()` skips any
+                    // chunk the content store already has, so a re-upload would
+                    // keep the bad copy. Drop it, as `VerifyStore` drops a copy
+                    // that fails verification, so `has()` reports those blobs
+                    // missing and the next upload writes the chunk again. This
+                    // is best effort: a content store that does not support
+                    // `remove` keeps the chunk.
+                    match self
+                        .content_store
+                        .as_store_driver_pin()
+                        .remove(index_entry.into())
+                        .await
+                    {
+                        Ok(removed) => warn!(
+                            %index_entry,
+                            actual,
+                            declared,
+                            removed,
+                            "Dropped a dedup content chunk whose size does not match its digest"
+                        ),
+                        Err(err) => warn!(
+                            %index_entry,
+                            ?err,
+                            "Could not drop a dedup content chunk whose size does not match its digest"
+                        ),
+                    }
+                    // NotFound rather than the DataLoss `VerifyStore` returns: a
+                    // `FastSlowStore` above falls back to its slow tier only on
+                    // NotFound, and the slow tier may hold a good copy.
+                    return Err(make_err!(
+                        Code::NotFound,
+                        "DedupStore content chunk {} is {actual} bytes but its index entry declares {declared}; refusing to serve a torn read as success",
+                        index_entry,
+                    ));
+                }
+
                 Result::<_, Error>::Ok(data)
             })
             .buffered(self.max_concurrent_fetch_per_get);

nativelink-store/tests/dedup_store_test.rs

@@ -12,11 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::sync::Arc;
+
 use nativelink_config::stores::{DedupSpec, MemorySpec, StoreSpec};
 use nativelink_error::{Code, Error, ResultExt};
 use nativelink_macro::nativelink_test;
 use nativelink_store::cas_utils::ZERO_BYTE_DIGESTS;
-use nativelink_store::dedup_store::DedupStore;
+use nativelink_store::compression_store::WincodeConfig;
+use nativelink_store::dedup_store::{DedupIndex, DedupStore};
 use nativelink_store::memory_store::MemoryStore;
 use nativelink_util::common::DigestInfo;
 use nativelink_util::store_trait::{Store, StoreLike};
@@ -445,3 +448,177 @@ async fn empty_cas_upload_does_not_publish_dedup_index() -> Result<(), Error> {
     assert_eq!(content.len_for_test(), 0);
     Ok(())
 }
+
+/// Builds a `DedupStore` whose index for `VALID_HASH1` names one content chunk
+/// of `declared_size` bytes, and stores `stored` under that chunk's digest. A
+/// `stored` of any other length is a chunk whose bytes disagree with its index
+/// entry, as a torn write or a misbehaving content store would leave it.
+async fn store_with_one_chunk(
+    declared_size: usize,
+    stored: Vec<u8>,
+) -> Result<(Arc<DedupStore>, Arc<MemoryStore>, DigestInfo, DigestInfo), Error> {
+    let index_store = MemoryStore::new(&MemorySpec::default());
+    let content_store = MemoryStore::new(&MemorySpec::default());
+    let store = DedupStore::new(
+        &make_default_config(),
+        Store::new(index_store.clone()),
+        Store::new(content_store.clone()),
+    )?;
+
+    let chunk_digest = DigestInfo::try_new(VALID_HASH2, declared_size).unwrap();
+    content_store
+        .update_oneshot(chunk_digest, stored.into())
+        .await
+        .err_tip(|| "Failed to write content chunk")?;
+
+    let serialized = wincode::config::serialize(
+        &DedupIndex {
+            entries: vec![chunk_digest],
+        },
+        WincodeConfig::new(),
+    )
+    .unwrap();
+    let top_digest = DigestInfo::try_new(VALID_HASH1, declared_size).unwrap();
+    index_store
+        .update_oneshot(top_digest, serialized.into())
+        .await
+        .err_tip(|| "Failed to write index entry")?;
+    Ok((store, content_store, top_digest, chunk_digest))
+}
+
+// A content chunk must contain exactly the number of bytes its index entry
+// declares. A short chunk used to be served as a truncated read that ended in a
+// normal EOF. The read must fail with NotFound, and the bad chunk must be gone
+// afterwards so `has()` stops reporting the blob as present.
+#[nativelink_test]
+async fn get_part_short_chunk_is_error_test() -> Result<(), Error> {
+    let (store, content_store, top_digest, chunk_digest) =
+        store_with_one_chunk(100, vec![7u8; 10]).await?;
+
+    let result = store.get_part_unchunked(top_digest, 0, Some(100)).await;
+    let err = result.expect_err("a short chunk must not be served as a successful read");
+    assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}");
+
+    assert_eq!(
+        content_store.has(chunk_digest).await?,
+        None,
+        "the short chunk must be dropped from the content store"
+    );
+    assert_eq!(
+        store.has(top_digest).await?,
+        None,
+        "a blob whose chunk was dropped must no longer be reported as present"
+    );
+    Ok(())
+}
+
+// A read that starts inside a short chunk used to trip the `bytes_to_skip`
+// assert and panic. It must fail like any other read of that chunk.
+#[nativelink_test]
+async fn get_part_offset_inside_short_chunk_is_error_test() -> Result<(), Error> {
+    let (store, _content_store, top_digest, _chunk_digest) =
+        store_with_one_chunk(100, vec![7u8; 10]).await?;
+
+    let result = store.get_part_unchunked(top_digest, 50, Some(10)).await;
+    let err = result.expect_err("a read inside a short chunk must fail, not panic");
+    assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}");
+    Ok(())
+}
+
+// A chunk longer than its index entry declares is just as wrong as a short one.
+#[nativelink_test]
+async fn get_part_long_chunk_is_error_test() -> Result<(), Error> {
+    let (store, content_store, top_digest, chunk_digest) =
+        store_with_one_chunk(100, vec![7u8; 150]).await?;
+
+    let result = store.get_part_unchunked(top_digest, 0, None).await;
+    let err = result.expect_err("a long chunk must not be served as a successful read");
+    assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}");
+    assert_eq!(content_store.has(chunk_digest).await?, None);
+    Ok(())
+}
+
+// End to end: a blob with one corrupted chunk fails to read, is then reported
+// missing, and the client's re-upload makes it readable again. Before the bad
+// chunk was dropped, `update()` skipped it because the content store still had
+// a value under its digest, so the blob stayed unreadable forever.
+#[nativelink_test]
+async fn short_chunk_is_rewritten_by_the_next_upload_test() -> Result<(), Error> {
+    let index_store = MemoryStore::new(&MemorySpec::default());
+    let content_store = MemoryStore::new(&MemorySpec::default());
+    let store = DedupStore::new(
+        &make_default_config(),
+        Store::new(index_store.clone()),
+        Store::new(content_store.clone()),
+    )?;
+
+    let original_data = make_random_data(MEGABYTE_SZ);
+    let digest = DigestInfo::try_new(VALID_HASH1, MEGABYTE_SZ).unwrap();
+    store
+        .update_oneshot(digest, original_data.clone().into())
+        .await
+        .err_tip(|| "Failed to write data to dedup store")?;
+
+    // Truncate one chunk from the middle of the blob in place.
+    let index_bytes = index_store.get_part_unchunked(digest, 0, None).await?;
+    let index: DedupIndex =
+        wincode::config::deserialize(&index_bytes, WincodeConfig::new()).unwrap();
+    assert!(index.entries.len() > 2, "expected several chunks");
+    let victim = index.entries[index.entries.len() / 2];
+    let victim_bytes = content_store.get_part_unchunked(victim, 0, None).await?;
+    content_store
+        .update_oneshot(victim, victim_bytes.slice(..victim_bytes.len() / 2))
+        .await?;
+
+    let err = store
+        .get_part_unchunked(digest, 0, None)
+        .await
+        .expect_err("a blob with a truncated chunk must not read back");
+    assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}");
+    assert_eq!(
+        store.has(digest).await?,
+        None,
+        "the blob must now be missing"
+    );
+
+    store
+        .update_oneshot(digest, original_data.clone().into())
+        .await
+        .err_tip(|| "Failed to re-upload data to dedup store")?;
+    let rt_data = store
+        .get_part_unchunked(digest, 0, None)
+        .await
+        .err_tip(|| "The re-uploaded blob must read back")?;
+    assert_eq!(rt_data, original_data, "Expected round trip data to match");
+    Ok(())
+}
+
+// Guard against over-correction: a bounded request for MORE bytes than the blob
+// actually contains must still succeed and return the available bytes, matching
+// the clamp-to-available semantics of the underlying stores. The
+// size-consistency check must not turn a legitimate over-length request into an
+// error.
+#[nativelink_test]
+async fn get_part_over_length_request_still_ok_test() -> Result<(), Error> {
+    const DATA_SIZE: usize = 64 * 1024;
+
+    let store = DedupStore::new(
+        &make_default_config(),
+        Store::new(MemoryStore::new(&MemorySpec::default())),
+        Store::new(MemoryStore::new(&MemorySpec::default())),
+    )?;
+
+    let original_data = make_random_data(DATA_SIZE);
+    let digest = DigestInfo::try_new(VALID_HASH1, DATA_SIZE).unwrap();
+    store
+        .update_oneshot(digest, original_data.clone().into())
+        .await
+        .err_tip(|| "Failed to write data to dedup store")?;
+
+    let rt = store
+        .get_part_unchunked(digest, 0, Some((DATA_SIZE as u64) * 2))
+        .await
+        .err_tip(|| "Over-length bounded read should still succeed")?;
+    assert_eq!(rt, original_data, "Expected available bytes to be returned");
+    Ok(())
+}