feat: client recovery sliver (#2861)

* client recovery sliver

* hook sliver recovery into retrieve slivers call

* add tests

* clean up before review

* address review comments

* more comments

Author: halfprice

crates/walrus-core/src/encoding/slivers.rs

@@ -269,8 +269,8 @@ impl<T: EncodingAxis> SliverData<T> {
     /// Returns the recovered [`SliverData`] if decoding succeeds or `None` if decoding fails.
     ///
     /// The recovery symbols are only the raw symbols without the Merkle proofs.
-    pub fn recover_sliver_from_decoding_symbols<I>(
-        recovery_symbols: I,
+    fn recover_sliver_from_decoding_symbols<I>(
+        decoding_symbols: I,
         target_index: SliverIndex,
         symbol_size: NonZeroU16,
         config: EncodingConfigEnum,
@@ -279,21 +279,52 @@ impl<T: EncodingAxis> SliverData<T> {
         I: IntoIterator,
         I::IntoIter: Iterator<Item = DecodingSymbol<T>> + ExactSizeIterator,
     {
-        let recovery_symbols = recovery_symbols.into_iter();
+        let decoding_symbols = decoding_symbols.into_iter();
 
         // Note: The following code may have to be changed if we add encodings that require a
         // variable number of symbols to recover a sliver.
         let RequiredCount::Exact(n_symbols_required) = config.n_symbols_for_recovery::<T>();
-        if recovery_symbols.len() < n_symbols_required {
+        if decoding_symbols.len() < n_symbols_required {
             // We don't even have to attempt decoding if we don't have enough recovery symbols.
             return Err(DecodeError::DecodingUnsuccessful);
         }
 
         config
-            .decode_from_decoding_symbols(symbol_size, recovery_symbols)
+            .decode_from_decoding_symbols(symbol_size, decoding_symbols)
             .map(|data| SliverData::new(data, symbol_size, target_index))
     }
 
+    /// Attempts to recover a sliver from the provided decoding symbols.
+    ///
+    /// This is a user facing API that is used to reconstruct a sliver from decoding symbols
+    /// without per symbol Merkle proofs.
+    ///
+    /// The reconstructed sliver is verified against the metadata to ensure it is consistent.
+    pub fn try_recover_sliver_from_decoding_symbols<I>(
+        decoding_symbols: I,
+        target_index: SliverIndex,
+        metadata: &BlobMetadata,
+        encoding_config: &EncodingConfig,
+    ) -> Result<SliverData<T>, SliverRecoveryOrVerificationError>
+    where
+        I: IntoIterator,
+        I::IntoIter: Iterator<Item = DecodingSymbol<T>> + ExactSizeIterator,
+    {
+        let symbol_size = metadata.symbol_size(encoding_config)?;
+        let config_enum = encoding_config.get_for_type(metadata.encoding_type());
+
+        let sliver = Self::recover_sliver_from_decoding_symbols(
+            decoding_symbols,
+            target_index,
+            symbol_size,
+            config_enum,
+        )
+        .map_err(|_| SliverRecoveryError::DecodingFailure)?;
+
+        sliver.verify(encoding_config, metadata)?;
+        Ok(sliver)
+    }
+
     /// Attempts to recover a sliver from the provided recovery symbols.
     ///
     /// The function *does not* verify the recovery symbols. It is the caller's responsibility to

crates/walrus-e2e-tests/tests/test_client.rs

@@ -39,9 +39,11 @@ use walrus_core::{
     SliverPairIndex,
     encoding::{
         BLOB_TYPE_ATTRIBUTE_KEY,
+        EncodingAxis,
         EncodingFactory as _,
         Primary,
         QUILT_TYPE_VALUE,
+        Secondary,
         encoded_blob_length_for_n_shards,
         quilt_encoding::{QuiltApi, QuiltStoreBlob, QuiltVersionV1},
     },
@@ -56,6 +58,7 @@ use walrus_sdk::{
         StoreArgs,
         StoreBlobsApi as _,
         WalrusNodeClient,
+        byte_range_read_client::{ByteRangeReadClient, ByteRangeReadClientConfig},
         client_types::WalrusStoreBlob,
         quilt_client::QuiltClientConfig,
         responses::{BlobStoreResult, QuiltStoreResult},
@@ -79,6 +82,7 @@ use walrus_sdk::{
 };
 use walrus_service::test_utils::{
     StorageNodeHandleTrait,
+    TestCluster,
     TestNodesConfig,
     UnusedSocketAddress,
     test_cluster::{self, FROST_PER_NODE_WEIGHT},
@@ -1147,9 +1151,10 @@ async fn test_store_quilt(blobs_to_create: u32) -> TestResult {
         .collect::<Vec<_>>();
 
     // Store the quilt.
-    let quilt_client = client
-        .quilt_client()
-        .with_config(QuiltClientConfig::new(6, Duration::from_mins(1)));
+    let quilt_client =
+        client
+            .quilt_client()
+            .with_config(QuiltClientConfig::new(6, Duration::from_mins(1), 100));
     let quilt = quilt_client
         .construct_quilt::<QuiltVersionV1>(&quilt_store_blobs, encoding_type)
         .await?;
@@ -2856,22 +2861,34 @@ async fn test_store_with_upload_relay_with_tip() {
         .expect("shutdown upload relay");
 }
 
-/// Tests client byte range read functionality.
-#[ignore = "ignore E2E tests by default"]
-#[walrus_simtest]
-async fn test_byte_range_read_client() -> TestResult {
+fn fail_random_node(walrus_cluster: &TestCluster) {
+    let node_to_stop = thread_rng().gen_range(0..5);
+    tracing::info!("failing node {node_to_stop}");
+    walrus_cluster.nodes[node_to_stop].cancel();
+}
+
+// Tests client byte range read functionality, with different node failure scenarios.
+async_param_test! {
+    #[ignore = "ignore E2E tests by default"]
+    #[walrus_simtest]
+    test_byte_range_read_client -> TestResult : [
+        with_node_failure: (true),
+        without_node_failure: (false),
+    ]
+}
+async fn test_byte_range_read_client(with_node_failure: bool) -> TestResult {
     walrus_test_utils::init_tracing();
 
     let test_nodes_config = TestNodesConfig::builder()
         .with_node_weights(&[7, 7, 7, 7, 7])
         .build();
     let test_cluster_builder =
         test_cluster::E2eTestSetupBuilder::new().with_test_nodes_config(test_nodes_config);
-    let (_sui_cluster_handle, _cluster, client, _, _) = test_cluster_builder.build().await?;
+    let (_sui_cluster_handle, walrus_cluster, client, _, _) = test_cluster_builder.build().await?;
     let client = client.as_ref();
 
     // Generate a non zero blob size.
-    let blob_size: u64 = thread_rng().gen_range(1..=1000000);
+    let blob_size: u64 = thread_rng().gen_range(100..=1000000);
     tracing::info!("blob size: {blob_size}");
     let blobs = walrus_test_utils::random_data_list(
         usize::try_from(blob_size).expect("blob size should be valid"),
@@ -2890,6 +2907,10 @@ async fn test_byte_range_read_client() -> TestResult {
         .blob_id()
         .expect("blob ID should be present");
 
+    if with_node_failure {
+        fail_random_node(&walrus_cluster);
+    }
+
     // Test reading byte range from the blob. We do 20 randomly generated byte range reads and
     // verify that the data is correct.
     for _ in 0..20 {
@@ -2995,15 +3016,30 @@ async fn test_byte_range_read_size_too_large() -> TestResult {
     Ok(())
 }
 
-/// Tests that streaming a blob returns the correct data.
-#[ignore = "ignore E2E tests by default"]
-#[walrus_simtest]
-async fn test_streaming_blob() -> TestResult {
+// Tests that streaming a blob returns the correct data.
+async_param_test! {
+    #[ignore = "ignore E2E tests by default"]
+    #[walrus_simtest]
+    test_streaming_blob -> TestResult : [
+        with_node_failure: (true),
+        without_node_failure: (false),
+    ]
+}
+async fn test_streaming_blob(with_node_failure: bool) -> TestResult {
     walrus_test_utils::init_tracing();
 
+    // Create a test cluster with 5 nodes with the same weight. This makes sure that when losing
+    // any node, it still has a quorum.
+    let test_nodes_config = TestNodesConfig::builder()
+        .with_node_weights(&[7, 7, 7, 7, 7])
+        .build();
+
     // Setup test cluster
-    let (_sui_cluster_handle, _cluster, client, _, _) =
-        test_cluster::E2eTestSetupBuilder::new().build().await?;
+    let (_sui_cluster_handle, walrus_cluster, client, _, _) =
+        test_cluster::E2eTestSetupBuilder::new()
+            .with_test_nodes_config(test_nodes_config)
+            .build()
+            .await?;
 
     // Generate and store a test blob (~100KB to span multiple slivers)
     let blob_size = 100_000;
@@ -3023,6 +3059,10 @@ async fn test_streaming_blob() -> TestResult {
         .blob_id()
         .expect("blob ID should be present");
 
+    if with_node_failure {
+        fail_random_node(&walrus_cluster);
+    }
+
     // Create a read-only client for streaming (SuiReadClient implements Clone)
     let sui_read_client = client.inner.sui_client().read_client().clone();
     let config = client.inner.config().clone();
@@ -3417,5 +3457,178 @@ async fn test_aggregator_get_blob_by_blob_id() -> TestResult {
     assert_eq!(content_type.to_str().unwrap(), "application/json");
 
     aggregator.shutdown().await?;
+
+    Ok(())
+}
+
+// Tests client side recovering slivers and reconstructing the blob.
+// This test uploads a blob, recovers all slivers, and verifies the reconstructed blob
+// matches the original.
+async fn run_test_client_recover_slivers<E: EncodingAxis>(
+    client: &WalrusNodeClient<SuiContractClient>,
+) -> TestResult {
+    // Generate a random blob.
+    let blob_size: u64 = thread_rng().gen_range(10000..=100000);
+    tracing::info!("blob size: {blob_size}");
+    let blobs = walrus_test_utils::random_data_list(
+        usize::try_from(blob_size).expect("blob size should be valid"),
+        1,
+    );
+    let original_blob = &blobs[0];
+
+    // Store the blob on Walrus.
+    let blob_read_result = client
+        .reserve_and_store_blobs_retry_committees(
+            blobs.clone(),
+            vec![],
+            &StoreArgs::default_with_epochs(5).with_persistence(BlobPersistence::Permanent),
+        )
+        .await?;
+
+    let blob_id = blob_read_result[0]
+        .blob_id()
+        .expect("blob ID should be present");
+
+    tracing::info!("stored blob {blob_id}");
+    let current_epoch = client.sui_client().current_epoch().await?;
+    let metadata = client.retrieve_metadata(current_epoch, &blob_id).await?;
+
+    // Get encoding config and calculate the number of systematic slivers for the encoding axis.
+    let encoding_config = client
+        .encoding_config()
+        .get_for_type(metadata.metadata().encoding_type());
+    let n_systematic = encoding_config.n_systematic_slivers::<E>();
+
+    // Create a list of all systematic sliver indices.
+    let all_primary_sliver_indices: Vec<_> = (0..n_systematic.get())
+        .map(walrus_core::SliverIndex::new)
+        .collect();
+
+    tracing::info!("recovering {} slivers", n_systematic);
+
+    // Recover all primary slivers using the recover_slivers method.
+    let recovered_slivers = client
+        .recover_slivers::<E>(&metadata, &all_primary_sliver_indices, current_epoch)
+        .await?;
+
+    tracing::info!("recovered {} slivers", recovered_slivers.len());
+
+    // Verify we got all the slivers.
+    assert_eq!(
+        recovered_slivers.len(),
+        usize::from(n_systematic.get()),
+        "should recover all systematic slivers"
+    );
+
+    // Sort slivers by index to ensure they're in order.
+    let mut sorted_slivers = recovered_slivers;
+    sorted_slivers.sort_by_key(|s| s.index.get());
+
+    // Reconstruct the blob from the recovered slivers.
+    let mut reconstructed_blob = Vec::new();
+    if E::IS_PRIMARY {
+        for sliver in &sorted_slivers {
+            reconstructed_blob.extend_from_slice(sliver.symbols.data());
+        }
+    } else {
+        // For secondary slivers, slivers are in columns. So the reconstructed blob needs to
+        // first take the first bytes from each sliver to form a row, and then concatenate the rows
+        // to form the reconstructed blob.
+        let row_length = sorted_slivers[0].len() / sorted_slivers[0].symbols.symbol_usize();
+        for row in 0..row_length {
+            for sliver in sorted_slivers.iter() {
+                reconstructed_blob.extend_from_slice(
+                    &sliver.symbols.data()[row * sliver.symbols.symbol_usize()
+                        ..(row + 1) * sliver.symbols.symbol_usize()],
+                );
+            }
+        }
+    }
+
+    // Truncate to the original blob size (last sliver may have padding).
+    reconstructed_blob.truncate(original_blob.len());
+
+    // Verify the reconstructed blob matches the original.
+    assert_eq!(
+        reconstructed_blob, *original_blob,
+        "reconstructed blob should match original"
+    );
+
+    Ok(())
+}
+
+/// Tests client side recovering sliver functionality.
+#[ignore = "ignore E2E tests by default"]
+#[walrus_simtest]
+async fn test_client_recover_slivers() -> TestResult {
+    walrus_test_utils::init_tracing();
+
+    let test_nodes_config = TestNodesConfig::builder()
+        .with_node_weights(&[7, 7, 7, 7, 7])
+        .build();
+    let test_cluster_builder =
+        test_cluster::E2eTestSetupBuilder::new().with_test_nodes_config(test_nodes_config);
+    let (_sui_cluster_handle, _cluster, client, _, _) = test_cluster_builder.build().await?;
+    let client = client.as_ref();
+
+    for _ in 0..10 {
+        run_test_client_recover_slivers::<Primary>(client).await?;
+        run_test_client_recover_slivers::<Secondary>(client).await?;
+    }
+
+    Ok(())
+}
+
+/// Test client side recovery with short timeout returns proper error.
+#[ignore = "ignore E2E tests by default"]
+#[walrus_simtest]
+async fn test_client_recover_slivers_timeout() -> TestResult {
+    walrus_test_utils::init_tracing();
+
+    let test_nodes_config = TestNodesConfig::builder()
+        .with_node_weights(&[7, 7, 7, 7, 7])
+        .build();
+    let test_cluster_builder =
+        test_cluster::E2eTestSetupBuilder::new().with_test_nodes_config(test_nodes_config);
+    let (_sui_cluster_handle, _cluster, client, _, _) = test_cluster_builder.build().await?;
+    let client = client.as_ref();
+
+    // Generate a random blob.
+    let blob_size: u64 = 10000;
+    let blobs = walrus_test_utils::random_data_list(
+        usize::try_from(blob_size).expect("blob size should be valid"),
+        1,
+    );
+
+    // Store the blob on Walrus.
+    let blob_read_result = client
+        .reserve_and_store_blobs_retry_committees(
+            blobs.clone(),
+            vec![],
+            &StoreArgs::default_with_epochs(5).with_persistence(BlobPersistence::Permanent),
+        )
+        .await?;
+
+    let blob_id = blob_read_result[0]
+        .blob_id()
+        .expect("blob ID should be present");
+
+    let byte_range_client = ByteRangeReadClient::new(
+        client,
+        ByteRangeReadClientConfig {
+            timeout: Duration::from_secs(0),
+            ..Default::default()
+        },
+    );
+
+    // Given that the timeout is 0, the read byte range should fail with error.
+    let result = byte_range_client.read_byte_range(&blob_id, 0, 100).await;
+    assert!(
+        result.is_err()
+            && result
+                .unwrap_err()
+                .to_string()
+                .contains("failed to retrieve some slivers")
+    );
     Ok(())
 }