feat: Add support for multiple digest functions with backward compati…

…bility
TraceMachina · Dec 18, 2024 · 8b954aa · 8b954aa
1 parent f672af7
commit 8b954aa
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 2 deletions.
diff --git 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
@@ -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