From 8b954aa5b6ee59dcb636c66ad0dcdce45f3f0d0e Mon Sep 17 00:00:00 2001 From: asr2003 <162500856+asr2003@users.noreply.github.com> Date: Wed, 18 Dec 2024 03:16:00 +0000 Subject: [PATCH 1/2] feat: Add support for multiple digest functions with backward compatibility --- 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, ) -> Result>, 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], ) -> 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 From 25075981db060a39a36d5064f6e37fc13e80be6f Mon Sep 17 00:00:00 2001 From: asr2003 <162500856+asr2003@users.noreply.github.com> Date: Wed, 18 Dec 2024 21:02:21 +0530 Subject: [PATCH 2/2] Update resource_info_test.rs --- nativelink-util/tests/resource_info_test.rs | 30 ++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/nativelink-util/tests/resource_info_test.rs b/nativelink-util/tests/resource_info_test.rs index 2321d76ca..370a1a4a3 100644 --- a/nativelink-util/tests/resource_info_test.rs +++ b/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::{is_supported_digest_function, ResourceInfo}; use pretty_assertions::assert_eq; #[nativelink_test] @@ -707,3 +707,31 @@ async fn write_invalid_size_test() -> Result<(), Box> { assert!(ResourceInfo::new(RESOURCE_NAME, true).is_err()); Ok(()) } + +#[nativelink_test] +async fn test_supported_digest_functions() -> Result<(), Box> { + 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> { + 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(()) +}