From 6b47cd75b60d443f2928a84a19422c605929a82e Mon Sep 17 00:00:00 2001
From: asr2003 <162500856+asr2003@users.noreply.github.com>
Date: Tue, 17 Dec 2024 11:36:37 +0530
Subject: [PATCH]  feat: Add support for multiple digest functions with
 backward compatibility

- Updated `ResourceInfo` to default `digest_function` to `sha256` if not provided.
- Ensured backward compatibility for existing functionality.
- Updated tests to reflect default `sha256` behavior when `digest_function` is not set
---
 nativelink-store/src/grpc_store.rs   | 32 ++++++++++++++++++++++++++--
 nativelink-util/src/resource_info.rs |  4 ++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs
index aaa1f0b0c..5a0dbd9ba 100644
--- a/nativelink-store/src/grpc_store.rs
+++ b/nativelink-store/src/grpc_store.rs
@@ -45,7 +45,7 @@ use nativelink_util::origin_context::ActiveOriginContext;
 use nativelink_util::proto_stream_utils::{
     FirstStream, WriteRequestStreamWrapper, WriteState, WriteStateWrapper,
 };
-use nativelink_util::resource_info::ResourceInfo;
+use nativelink_util::resource_info::{is_supported_digest_function, ResourceInfo};
 use nativelink_util::retry::{Retrier, RetryResult};
 use nativelink_util::store_trait::{StoreDriver, StoreKey, UploadSizeInfo};
 use nativelink_util::{default_health_status_indicator, tls_utils};
@@ -276,12 +276,28 @@ impl GrpcStore {
         &self,
         grpc_request: impl IntoRequest<ReadRequest>,
     ) -> Result<impl Stream<Item = Result<ReadResponse, Status>>, Error> {
+        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
         })
@@ -514,6 +530,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())
diff --git a/nativelink-util/src/resource_info.rs b/nativelink-util/src/resource_info.rs
index b3848d8b9..60fa95c2e 100644
--- a/nativelink-util/src/resource_info.rs
+++ b/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