feat: add support for multiple digest functions with validation tests

Author: asr2003

nativelink-store/BUILD.bazel

@@ -101,6 +101,7 @@ rust_test_suite(
         "tests/filesystem_store_test.rs",
         "tests/gcs_client_test.rs",
         "tests/gcs_store_test.rs",
+        "tests/grpc_store_test.rs",
         "tests/memory_store_test.rs",
         "tests/redis_store_test.rs",
         "tests/ref_store_test.rs",

nativelink-store/src/grpc_store.rs

@@ -44,7 +44,7 @@ use nativelink_util::health_utils::HealthStatusIndicator;
 use nativelink_util::proto_stream_utils::{
     FirstStream, WriteRequestStreamWrapper, WriteState, WriteStateWrapper,
 };
-use nativelink_util::resource_info::ResourceInfo;
+use nativelink_util::resource_info::{ResourceInfo, is_supported_digest_function};
 use nativelink_util::retry::{Retrier, RetryResult};
 use nativelink_util::store_trait::{StoreDriver, StoreKey, UploadSizeInfo};
 use nativelink_util::{default_health_status_indicator, tls_utils};
@@ -278,12 +278,22 @@ impl GrpcStore {
     where
         R: IntoRequest<ReadRequest>,
     {
+        const IS_UPLOAD_FALSE: bool = false;
+
+        let request = self.get_read_request(grpc_request.into_request().into_inner())?;
+        let resource_name = &request.resource_name;
+        let resource_info = ResourceInfo::new(resource_name, IS_UPLOAD_FALSE)
+            .err_tip(|| "Failed to parse resource_name in GrpcStore::read")?;
+
+        let digest_function = resource_info.digest_function.as_deref().unwrap_or("sha256");
+
+        Self::validate_digest_function(digest_function, Some(resource_name))?;
+
         error_if!(
             matches!(self.store_type, nativelink_config::stores::StoreType::Ac),
             "CAS operation on AC store"
         );
 
-        let request = self.get_read_request(grpc_request.into_request().into_inner())?;
         self.perform_request(request, |request| async move {
             self.read_internal(request).await
         })
@@ -505,6 +515,23 @@ impl GrpcStore {
             .await
             .map(|_| ())
     }
+
+    pub fn validate_digest_function(
+        digest_function: &str,
+        resource_name: Option<&str>,
+    ) -> Result<(), Error> {
+        if !is_supported_digest_function(digest_function) {
+            return Err(make_input_err!(
+                "Unsupported digest_function: {}{}",
+                digest_function,
+                match resource_name {
+                    Some(name) => format!(" in resource_name '{name}'"),
+                    None => String::new(),
+                }
+            ));
+        }
+        Ok(())
+    }
 }
 
 #[async_trait]
@@ -516,6 +543,12 @@ impl StoreDriver for GrpcStore {
         keys: &[StoreKey<'_>],
         results: &mut [Option<u64>],
     ) -> Result<(), Error> {
+        let digest_function = Context::current()
+            .get::<DigestHasherFunc>()
+            .map_or_else(default_digest_hasher_func, |v| *v)
+            .to_string();
+        GrpcStore::validate_digest_function(&digest_function, None)?;
+
         if matches!(self.store_type, nativelink_config::stores::StoreType::Ac) {
             keys.iter()
                 .zip(results.iter_mut())

nativelink-store/tests/grpc_store_test.rs

@@ -0,0 +1,57 @@
+// Copyright 2025 The NativeLink Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use nativelink_store::grpc_store::GrpcStore;
+use nativelink_util::resource_info::{ResourceInfo, is_supported_digest_function};
+use opentelemetry::context::Context;
+
+#[test]
+fn test_is_supported_digest_function() {
+    assert!(is_supported_digest_function("sha256"));
+    assert!(is_supported_digest_function("sha512"));
+    assert!(!is_supported_digest_function("crc32"));
+}
+
+#[test]
+fn test_read_rejects_invalid_digest_function() -> Result<(), Box<dyn core::error::Error>> {
+    const RESOURCE_NAME: &str = "instance_name/blobs/sha256/0123456789abcdef0123456789abcdef/123";
+    let mut resource_info = ResourceInfo::new(RESOURCE_NAME, false)?;
+    resource_info.digest_function = Some("sha3".into());
+    let digest_func = resource_info.digest_function.clone().unwrap();
+
+    let result = GrpcStore::validate_digest_function(&digest_func, Some(RESOURCE_NAME));
+    assert!(result.is_err(), "Expected error on invalid digest_function");
+    let msg = result.unwrap_err().to_string();
+    assert!(
+        msg.contains("Unsupported digest_function"),
+        "Unexpected error: {}",
+        msg
+    );
+    Ok(())
+}
+
+#[test]
+fn test_has_with_results_rejects_invalid_digest_function_in_context() {
+    let ctx = Context::current().with_value("sha3_256".to_string());
+    let _guard = ctx.attach();
+
+    let result = GrpcStore::validate_digest_function("sha3_256", None);
+    assert!(result.is_err(), "Expected error from context digest check");
+    let msg = result.unwrap_err().to_string();
+    assert!(
+        msg.contains("Unsupported digest_function"),
+        "Unexpected error: {}",
+        msg
+    );
+}

nativelink-util/src/resource_info.rs

@@ -206,6 +206,10 @@ enum State {
     OptionalMetadata,
 }
 
+pub fn is_supported_digest_function(digest_function: &str) -> bool {
+    DIGEST_FUNCTIONS.contains(&digest_function.to_lowercase().as_str())
+}
+
 // Iterate backwards looking for "(compressed-)blobs", once found, move forward
 // populating the output struct. This recursive function utilises the stack to
 // temporarily hold the reference to the previous item reducing the need for

nativelink-util/tests/resource_info_test.rs

@@ -15,7 +15,7 @@
 use std::borrow::Cow;
 
 use nativelink_macro::nativelink_test;
-use nativelink_util::resource_info::ResourceInfo;
+use nativelink_util::resource_info::{ResourceInfo, is_supported_digest_function};
 use pretty_assertions::assert_eq;
 
 #[nativelink_test]
@@ -707,3 +707,31 @@ async fn write_invalid_size_test() -> Result<(), Box<dyn core::error::Error>> {
     assert!(ResourceInfo::new(RESOURCE_NAME, true).is_err());
     Ok(())
 }
+
+#[nativelink_test]
+async fn test_supported_digest_functions() -> Result<(), Box<dyn core::error::Error>> {
+    assert_eq!(is_supported_digest_function("sha256"), true);
+    assert_eq!(is_supported_digest_function("sha1"), true);
+    assert_eq!(is_supported_digest_function("md5"), true);
+    assert_eq!(is_supported_digest_function("vso"), true);
+    assert_eq!(is_supported_digest_function("sha384"), true);
+    assert_eq!(is_supported_digest_function("sha512"), true);
+    assert_eq!(is_supported_digest_function("murmur3"), true);
+    assert_eq!(is_supported_digest_function("sha256tree"), true);
+    assert_eq!(is_supported_digest_function("blake3"), true);
+
+    Ok(())
+}
+
+#[nativelink_test]
+async fn test_unsupported_digest_functions() -> Result<(), Box<dyn core::error::Error>> {
+    assert_eq!(is_supported_digest_function("sha3"), false);
+    assert_eq!(
+        is_supported_digest_function("invalid_digest_function"),
+        false
+    );
+    assert_eq!(is_supported_digest_function("boo"), false);
+    assert_eq!(is_supported_digest_function("random_hash"), false);
+
+    Ok(())
+}