feat: Add support for multiple digest functions with backward compatibility

Author: asr2003

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,28 @@ 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");
+
+        if !is_supported_digest_function(digest_function) {
+            return Err(make_input_err!(
+                "Unsupported digest_function: {} in resource_name '{}'",
+                digest_function,
+                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
         })
@@ -516,6 +532,18 @@ impl StoreDriver for GrpcStore {
         keys: &[StoreKey<'_>],
         results: &mut [Option<u64>],
     ) -> Result<(), Error> {
+        let digest_function = ActiveOriginContext::get_value(&ACTIVE_HASHER_FUNC)
+            .err_tip(|| "In GrpcStore::has_with_results")?
+            .map_or_else(default_digest_hasher_func, |v| *v)
+            .to_string();
+
+        if !is_supported_digest_function(&digest_function) {
+            return Err(make_input_err!(
+                "Unsupported digest_function: {}",
+                digest_function
+            ));
+        }
+
         if matches!(self.store_type, nativelink_config::stores::StoreType::Ac) {
             keys.iter()
                 .zip(results.iter_mut())

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 std::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 std::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(())
+}