modified approach

Author: MarcusSorealheis

nativelink-store/src/gcs_store.rs

@@ -31,6 +31,7 @@ use nativelink_util::retry::{Retrier, RetryResult};
 use nativelink_util::store_trait::{StoreDriver, StoreKey, UploadSizeInfo};
 use rand::Rng;
 use tokio::time::sleep;
+use tracing::warn;
 
 use crate::cas_utils::is_zero_digest;
 use crate::gcs_client::client::{GcsClient, GcsOperations};
@@ -222,6 +223,18 @@ where
         mut reader: DropCloserReadHalf,
         upload_size: UploadSizeInfo,
     ) -> Result<(), Error> {
+        // Check if this is a zero-byte file
+        if is_zero_digest(digest.borrow()) {
+            warn!(
+                "Attempted to upload zero-byte file to GCS. This is not supported. \
+                Consider using a size partitioning store to restrict zero-byte files."
+            );
+            // Consume any data in the reader (should be empty) and return success
+            // to maintain compatibility with existing behavior
+            drop(reader.consume(None).await?);
+            return Ok(());
+        }
+
         let object_path = self.make_object_path(&digest);
 
         reader.set_max_recent_data_size(

nativelink-store/tests/gcs_store_test.rs

@@ -346,43 +346,44 @@ async fn get_part_with_range() -> Result<(), Error> {
     Ok(())
 }
 
-#[nativelink_test]
-async fn get_part_zero_digest() -> Result<(), Error> {
-    // Create mock GCS operations
-    let mock_ops = Arc::new(MockGcsOperations::new());
-    let ops_as_trait: Arc<dyn GcsOperations> = mock_ops.clone();
-    let store = create_test_store(ops_as_trait).await?;
-
-    // Create a zero digest
-    let digest = DigestInfo::new(Sha256::new().finalize().into(), 0);
-    let store_key: StoreKey = to_store_key(digest);
-    let (mut tx, mut rx) = make_buf_channel_pair();
-
-    // Start get_part operation
-    let store_clone = store.clone();
-    let get_fut = nativelink_util::spawn!("get_part_full_task", async move {
-        store_clone.get_part(store_key, &mut tx, 0, None).await
-    });
-
-    // Receive the data - should be empty
-    let received_data = rx.consume(Some(100)).await?;
-    assert_eq!(
-        received_data.as_ref(),
-        b"",
-        "Received data should be empty for zero digest"
-    );
-    get_fut.await??;
-
-    // Verify no mock operations were called (zero digest is special-cased)
-    let call_counts = mock_ops.get_call_counts();
-    assert_eq!(
-        call_counts.read_calls.load(Ordering::Relaxed),
-        0,
-        "No read operations should have been called for zero digest"
-    );
-
-    Ok(())
-}
+// I'll delete this but it did not actually test and prevent the issue we were addressing.
+// #[nativelink_test]
+// async fn get_part_zero_digest() -> Result<(), Error> {
+//     // Create mock GCS operations
+//     let mock_ops = Arc::new(MockGcsOperations::new());
+//     let ops_as_trait: Arc<dyn GcsOperations> = mock_ops.clone();
+//     let store = create_test_store(ops_as_trait).await?;
+
+//     // Create a zero digest
+//     let digest = DigestInfo::new(Sha256::new().finalize().into(), 0);
+//     let store_key: StoreKey = to_store_key(digest);
+//     let (mut tx, mut rx) = make_buf_channel_pair();
+
+//     // Start get_part operation
+//     let store_clone = store.clone();
+//     let get_fut = nativelink_util::spawn!("get_part_full_task", async move {
+//         store_clone.get_part(store_key, &mut tx, 0, None).await
+//     });
+
+//     // Receive the data - should be empty
+//     let received_data = rx.consume(Some(100)).await?;
+//     assert_eq!(
+//         received_data.as_ref(),
+//         b"",
+//         "Received data should be empty for zero digest"
+//     );
+//     get_fut.await??;
+
+//     // Verify no mock operations were called (zero digest is special-cased)
+//     let call_counts = mock_ops.get_call_counts();
+//     assert_eq!(
+//         call_counts.read_calls.load(Ordering::Relaxed),
+//         0,
+//         "No read operations should have been called for zero digest"
+//     );
+
+//     Ok(())
+// }
 
 #[nativelink_test]
 async fn test_expired_object() -> Result<(), Error> {
@@ -609,3 +610,63 @@ fn create_object_path(key: &StoreKey) -> ObjectPath {
         &format!("{}{}", KEY_PREFIX, key.as_str()),
     )
 }
+
+#[nativelink_test]
+async fn test_zero_byte_upload_prevention() -> Result<(), Error> {
+    // Create mock GCS operations
+    let mock_ops = Arc::new(MockGcsOperations::new());
+    let ops_as_trait: Arc<dyn GcsOperations> = mock_ops.clone();
+    let store = create_test_store(ops_as_trait).await?;
+
+    // Create a zero-byte digest
+    let zero_digest = DigestInfo::new(Sha256::new().finalize().into(), 0);
+    let store_key: StoreKey = to_store_key(zero_digest);
+    let (mut tx, rx) = make_buf_channel_pair();
+
+    // Start update operation for zero-byte file
+    let store_clone = store.clone();
+    let update_fut = nativelink_util::spawn!("update_zero_byte_task", async move {
+        store_clone
+            .update(store_key, rx, UploadSizeInfo::ExactSize(0))
+            .await
+    });
+
+    // Send EOF immediately (zero bytes)
+    tx.send_eof()?;
+
+    // The update should complete but no upload should occur
+    let result = update_fut.await?;
+    assert!(result.is_ok(), "Update should complete without error");
+
+    // Verify that no write operations were called
+    let requests = mock_ops.get_requests().await;
+    let write_request_count = requests
+        .iter()
+        .filter(|req| {
+            matches!(
+                req,
+                MockRequest::Write { .. } | MockRequest::StartResumable { .. }
+            )
+        })
+        .count();
+
+    assert_eq!(
+        write_request_count, 0,
+        "No write requests should be made for zero-byte files"
+    );
+
+    // Verify call counts
+    let call_counts = mock_ops.get_call_counts();
+    assert_eq!(
+        call_counts.write_calls.load(Ordering::Relaxed),
+        0,
+        "No write operations should have been called"
+    );
+    assert_eq!(
+        call_counts.start_resumable_calls.load(Ordering::Relaxed),
+        0,
+        "No resumable upload operations should have been called"
+    );
+
+    Ok(())
+}