/build/source/nativelink-util/src/digest_hasher.rs

Source
// Copyright 2024 The NativeLink Authors. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// 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 std::sync::OnceLock;

use blake3::Hasher as Blake3Hasher;
use bytes::BytesMut;
use futures::Future;
use nativelink_config::stores::ConfigDigestHashFunction;
use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err};
use nativelink_metric::{
    MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent,
};
use nativelink_proto::build::bazel::remote::execution::v2::digest_function::Value as ProtoDigestFunction;
use opentelemetry::context::Context;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeekExt};

use crate::common::DigestInfo;
use crate::{fs, spawn_blocking};

static DEFAULT_DIGEST_HASHER_FUNC: OnceLock<DigestHasherFunc> = OnceLock::new();

/// Utility function to make a context with a specific hasher function set.
pub fn make_ctx_for_hash_func<H>(hasher: H) -> Result<Context, Error>
where
    H: TryInto<DigestHasherFunc>,
    H::Error: Into<Error>,
{
    let digest_hasher_func = hasher
        .try_into()
        .err_tip(|| "Could not convert into DigestHasherFunc")?0;

    let new_ctx = Context::current_with_value(digest_hasher_func);

    Ok(new_ctx)
}

/// Get the default hasher.
pub fn default_digest_hasher_func() -> DigestHasherFunc {
    *DEFAULT_DIGEST_HASHER_FUNC.get_or_init(|| DigestHasherFunc::Sha256)
}

/// Get the hasher requested by the client from the active context (set via
/// [`make_ctx_for_hash_func`]), falling back to the default hasher.
pub fn digest_hasher_func_from_context() -> DigestHasherFunc {
    Context::current()
        .get::<DigestHasherFunc>()
        .map_or_else(default_digest_hasher_func, |v| *v)
}

/// Sets the default hasher to use if no hasher was requested by the client.
pub fn set_default_digest_hasher_func(hasher: DigestHasherFunc) -> Result<(), Error> {
    DEFAULT_DIGEST_HASHER_FUNC
        .set(hasher)
        .map_err(|_| make_err!(Code::Internal, "default_digest_hasher_func already set"))
}

/// Supported digest hash functions.
#[derive(Copy, Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash, Serialize, Deserialize)]
pub enum DigestHasherFunc {
    Sha256,
    Blake3,
}

impl MetricsComponent for DigestHasherFunc {
    fn publish(
        &self,
        kind: MetricKind,
        field_metadata: MetricFieldData,
    ) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
        format!("{self:?}").publish(kind, field_metadata)
    }
}

impl DigestHasherFunc {
    pub fn hasher(&self) -> DigestHasherImpl {
        self.into()
    }

    #[must_use]
    pub const fn proto_digest_func(&self) -> ProtoDigestFunction {
        match self {
            Self::Sha256 => ProtoDigestFunction::Sha256,
            Self::Blake3 => ProtoDigestFunction::Blake3,
        }
    }
}

impl From<ConfigDigestHashFunction> for DigestHasherFunc {
    fn from(value: ConfigDigestHashFunction) -> Self {
        match value {
            ConfigDigestHashFunction::Sha256 => Self::Sha256,
            ConfigDigestHashFunction::Blake3 => Self::Blake3,
        }
    }
}

impl TryFrom<ProtoDigestFunction> for DigestHasherFunc {
    type Error = Error;

    fn try_from(value: ProtoDigestFunction) -> Result<Self, Self::Error> {
        match value {
            ProtoDigestFunction::Sha256 => Ok(Self::Sha256),
            ProtoDigestFunction::Blake3 => Ok(Self::Blake3),
            v => Err(make_input_err!(
                "Unknown or unsupported digest function for proto conversion {v:?}"
            )),
        }
    }
}

impl TryFrom<&str> for DigestHasherFunc {
    type Error = Error;

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        match value.to_uppercase().as_str() {
            "SHA256" => Ok(Self::Sha256),
            "BLAKE3" => Ok(Self::Blake3),
            v => Err(make_input_err!(
                "Unknown or unsupported digest function for string conversion: {v:?}"
            )),
        }
    }
}

impl core::fmt::Display for DigestHasherFunc {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::Sha256 => write!(f, "SHA256"),
            Self::Blake3 => write!(f, "BLAKE3"),
        }
    }
}

impl TryFrom<i32> for DigestHasherFunc {
    type Error = Error;

    fn try_from(value: i32) -> Result<Self, Self::Error> {
        // Zero means not-set.
        if value == 0 {
            return Ok(default_digest_hasher_func());
        }
        match ProtoDigestFunction::try_from(value) {
            Ok(ProtoDigestFunction::Sha256) => Ok(Self::Sha256),
            Ok(ProtoDigestFunction::Blake3) => Ok(Self::Blake3),
            value => Err(make_input_err!(
                "Unknown or unsupported digest function for int conversion: {:?}",
                value.map(|v| v.as_str_name())
            )),
        }
    }
}

impl From<&DigestHasherFunc> for DigestHasherImpl {
    fn from(value: &DigestHasherFunc) -> Self {
        let hash_func_impl = match value {
            DigestHasherFunc::Sha256 => DigestHasherFuncImpl::Sha256(Sha256::new()),
            DigestHasherFunc::Blake3 => DigestHasherFuncImpl::Blake3(Box::default()),
        };
        Self {
            hashed_size: 0,
            hash_func_impl,
        }
    }
}

/// Wrapper to compute a hash of arbitrary data.
pub trait DigestHasher {
    /// Update the hasher with some additional data.
    fn update(&mut self, input: &[u8]);

    /// Finalize the hash function and collect the results into a digest.
    fn finalize_digest(&mut self) -> DigestInfo;

    /// Specialized version of the hashing function that is optimized for
    /// handling files. These optimizations take into account things like,
    /// the file size and the hasher algorithm to decide how to best process
    /// the file and feed it into the hasher.
    fn digest_for_file(
        self,
        file_path: impl AsRef<std::path::Path>,
        file: fs::FileSlot,
        size_hint: Option<u64>,
    ) -> impl Future<Output = Result<(DigestInfo, fs::FileSlot), Error>>;

    /// Utility function to compute a hash from a generic reader.
    fn compute_from_reader<R: AsyncRead + Unpin + Send>(
        &mut self,
        mut reader: R,
    ) -> impl Future<Output = Result<DigestInfo, Error>> {
        async move {
            let mut chunk = BytesMut::with_capacity(fs::DEFAULT_READ_BUFF_SIZE);
            loop {
                reader
                    .read_buf(&mut chunk)
                    .await
                    .err_tip(|| "Could not read chunk during compute_from_reader")?0;
                if chunk.is_empty() {
                    break; // EOF.
                }
                DigestHasher::update(self, &chunk);
                chunk.clear();
            }
            Ok(DigestHasher::finalize_digest(self))
        }
    }
}

#[expect(
    variant_size_differences,
    reason = "some variants are already boxed; this is acceptable"
)]
#[derive(Debug)]
pub enum DigestHasherFuncImpl {
    Sha256(Sha256),
    Blake3(Box<Blake3Hasher>), // Box because Blake3Hasher is 1.3kb in size.
}

/// The individual implementation of the hash function.
#[derive(Debug)]
pub struct DigestHasherImpl {
    hashed_size: u64,
    hash_func_impl: DigestHasherFuncImpl,
}

impl DigestHasherImpl {
    #[inline]
    async fn hash_file(
        &mut self,
        mut file: fs::FileSlot,
    ) -> Result<(DigestInfo, fs::FileSlot), Error> {
        let digest = self
            .compute_from_reader(&mut file)
            .await
            .err_tip(|| "In digest_for_file")?0;
        Ok((digest, file))
    }
}

impl DigestHasher for DigestHasherImpl {
    #[inline]
    fn update(&mut self, input: &[u8]) {
        self.hashed_size += input.len() as u64;
        match &mut self.hash_func_impl {
            DigestHasherFuncImpl::Sha256(h) => sha2::digest::Update::update(h, input),
            DigestHasherFuncImpl::Blake3(h) => {
                Blake3Hasher::update(h, input);
            }
        }
    }

    #[inline]
    fn finalize_digest(&mut self) -> DigestInfo {
        let hash = match &mut self.hash_func_impl {
            DigestHasherFuncImpl::Sha256(h) => h.finalize_reset().into(),
            DigestHasherFuncImpl::Blake3(h) => h.finalize().into(),
        };
        DigestInfo::new(hash, self.hashed_size)
    }

    async fn digest_for_file(
        mut self,
        file_path: impl AsRef<std::path::Path>,
        mut file: fs::FileSlot,
        size_hint: Option<u64>,
    ) -> Result<(DigestInfo, fs::FileSlot), Error> {
        let file_position = file
            .stream_position()
            .await
            .err_tip(|| "Couldn't get stream position in digest_for_file")?0;
        if file_position != 0 {
            return self.hash_file(file).await;
        }
        // If we are a small file, it's faster to just do it the "slow" way.
        // Great read: https://github.com/david-slatinek/c-read-vs.-mmap
        if let Some(size_hint) = size_hint
            && size_hint <= fs::DEFAULT_READ_BUFF_SIZE as u64
        {
            return self.hash_file(file).await;
        }
        let file_path = file_path.as_ref().to_path_buf();
        match self.hash_func_impl {
            DigestHasherFuncImpl::Sha256(_) => self.hash_file(file).await,
            DigestHasherFuncImpl::Blake3(mut hasher) => {
                spawn_blocking!("digest_for_file", move || {
                    hasher.update_mmap(file_path).map_err(|e| {
                        Error::from_std_err(Code::Internal, &e)
                            .append("Error in blake3's update_mmap")
                    })?;
                    Result::<_, Error>::Ok((
                        DigestInfo::new(hasher.finalize().into(), hasher.count()),
                        file,
                    ))
                })
                .await
                .err_tip(|| "Could not spawn blocking task in digest_for_file")?
            }
        }
    }
}

Coverage Report

Created: 2026-07-21 15:28

Line	Count	Source
1		// Copyright 2024 The NativeLink Authors. All rights reserved.
2		//
3		// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
4		// you may not use this file except in compliance with the License.
5		// You may obtain a copy of the License at
6		//
7		// See LICENSE file for details
8		//
9		// Unless required by applicable law or agreed to in writing, software
10		// distributed under the License is distributed on an "AS IS" BASIS,
11		// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12		// See the License for the specific language governing permissions and
13		// limitations under the License.
14
15		use std::sync::OnceLock;
16
17		use blake3::Hasher as Blake3Hasher;
18		use bytes::BytesMut;
19		use futures::Future;
20		use nativelink_config::stores::ConfigDigestHashFunction;
21		use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err};
22		use nativelink_metric::{
23		MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent,
24		};
25		use nativelink_proto::build::bazel::remote::execution::v2::digest_function::Value as ProtoDigestFunction;
26		use opentelemetry::context::Context;
27		use serde::{Deserialize, Serialize};
28		use sha2::{Digest, Sha256};
29		use tokio::io::{AsyncRead, AsyncReadExt, AsyncSeekExt};
30
31		use crate::common::DigestInfo;
32		use crate::{fs, spawn_blocking};
33
34		static DEFAULT_DIGEST_HASHER_FUNC: OnceLock<DigestHasherFunc> = OnceLock::new();
35
36		/// Utility function to make a context with a specific hasher function set.
37	82	pub fn make_ctx_for_hash_func<H>(hasher: H) -> Result<Context, Error>
38	82	where
39	82	H: TryInto<DigestHasherFunc>,
40	82	H::Error: Into<Error>,
41		{
42	82	let digest_hasher_func = hasher
43	82	.try_into()
44	82	.err_tip(\|\| "Could not convert into DigestHasherFunc") ?0 ;
45
46	82	let new_ctx = Context::current_with_value(digest_hasher_func);
47
48	82	Ok(new_ctx)
49	82	}
50
51		/// Get the default hasher.
52	282	pub fn default_digest_hasher_func() -> DigestHasherFunc {
53	282	*DEFAULT_DIGEST_HASHER_FUNC.get_or_init(\|\| DigestHasherFunc::Sha256)
54	282	}
55
56		/// Get the hasher requested by the client from the active context (set via
57		/// [`make_ctx_for_hash_func`]), falling back to the default hasher.
58	5	pub fn digest_hasher_func_from_context() -> DigestHasherFunc {
59	5	Context::current()
60	5	.get::<DigestHasherFunc>()
61	5	.map_or_else(default_digest_hasher_func, \|v\| *v)
62	5	}
63
64		/// Sets the default hasher to use if no hasher was requested by the client.
65	0	pub fn set_default_digest_hasher_func(hasher: DigestHasherFunc) -> Result<(), Error> {
66	0	DEFAULT_DIGEST_HASHER_FUNC
67	0	.set(hasher)
68	0	.map_err(\|_\| make_err!(Code::Internal, "default_digest_hasher_func already set"))
69	0	}
70
71		/// Supported digest hash functions.
72		#[derive(Copy, Clone, Debug, Ord, PartialOrd, Eq, PartialEq, Hash, Serialize, Deserialize)]
73		pub enum DigestHasherFunc {
74		Sha256,
75		Blake3,
76		}
77
78		impl MetricsComponent for DigestHasherFunc {
79	0	fn publish(
80	0	&self,
81	0	kind: MetricKind,
82	0	field_metadata: MetricFieldData,
83	0	) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
84	0	format!("{self:?}").publish(kind, field_metadata)
85	0	}
86		}
87
88		impl DigestHasherFunc {
89	5.24k	pub fn hasher(&self) -> DigestHasherImpl {
90	5.24k	self.into()
91	5.24k	}
92
93		#[must_use]
94	265	pub const fn proto_digest_func(&self) -> ProtoDigestFunction {
95	265	match self {
96	263	Self::Sha256 => ProtoDigestFunction::Sha256,
97	2	Self::Blake3 => ProtoDigestFunction::Blake3,
98		}
99	265	}
100		}
101
102		impl From<ConfigDigestHashFunction> for DigestHasherFunc {
103	0	fn from(value: ConfigDigestHashFunction) -> Self {
104	0	match value {
105	0	ConfigDigestHashFunction::Sha256 => Self::Sha256,
106	0	ConfigDigestHashFunction::Blake3 => Self::Blake3,
107		}
108	0	}
109		}
110
111		impl TryFrom<ProtoDigestFunction> for DigestHasherFunc {
112		type Error = Error;
113
114	8	fn try_from(value: ProtoDigestFunction) -> Result<Self, Self::Error> {
115	8	match value {
116	6	ProtoDigestFunction::Sha256 => Ok(Self::Sha256),
117	0	ProtoDigestFunction::Blake3 => Ok(Self::Blake3),
118	2	v => Err(make_input_err!(
119	2	"Unknown or unsupported digest function for proto conversion {v:?}"
120	2	)),
121		}
122	8	}
123		}
124
125		impl TryFrom<&str> for DigestHasherFunc {
126		type Error = Error;
127
128	0	fn try_from(value: &str) -> Result<Self, Self::Error> {
129	0	match value.to_uppercase().as_str() {
130	0	"SHA256" => Ok(Self::Sha256),
131	0	"BLAKE3" => Ok(Self::Blake3),
132	0	v => Err(make_input_err!(
133	0	"Unknown or unsupported digest function for string conversion: {v:?}"
134	0	)),
135		}
136	0	}
137		}
138
139		impl core::fmt::Display for DigestHasherFunc {
140	29	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
141	29	match self {
142	27	Self::Sha256 => write!(f, "SHA256"),
143	2	Self::Blake3 => write!(f, "BLAKE3"),
144		}
145	29	}
146		}
147
148		impl TryFrom<i32> for DigestHasherFunc {
149		type Error = Error;
150
151	80	fn try_from(value: i32) -> Result<Self, Self::Error> {
152		// Zero means not-set.
153	80	if value == 0 {
154	29	return Ok(default_digest_hasher_func());
155	51	}
156	51	match ProtoDigestFunction::try_from(value) {
157	48	Ok(ProtoDigestFunction::Sha256) => Ok(Self::Sha256),
158	3	Ok(ProtoDigestFunction::Blake3) => Ok(Self::Blake3),
159	0	value => Err(make_input_err!(
160		"Unknown or unsupported digest function for int conversion: {:?}",
161	0	value.map(\|v\| v.as_str_name())
162		)),
163		}
164	80	}
165		}
166
167		impl From<&DigestHasherFunc> for DigestHasherImpl {
168	5.24k	fn from(value: &DigestHasherFunc) -> Self {
169	5.24k	let hash_func_impl = match value {
170	199	DigestHasherFunc::Sha256 => DigestHasherFuncImpl::Sha256(Sha256::new()),
171	5.04k	DigestHasherFunc::Blake3 => DigestHasherFuncImpl::Blake3(Box::default()),
172		};
173	5.24k	Self {
174	5.24k	hashed_size: 0,
175	5.24k	hash_func_impl,
176	5.24k	}
177	5.24k	}
178		}
179
180		/// Wrapper to compute a hash of arbitrary data.
181		pub trait DigestHasher {
182		/// Update the hasher with some additional data.
183		fn update(&mut self, input: &[u8]);
184
185		/// Finalize the hash function and collect the results into a digest.
186		fn finalize_digest(&mut self) -> DigestInfo;
187
188		/// Specialized version of the hashing function that is optimized for
189		/// handling files. These optimizations take into account things like,
190		/// the file size and the hasher algorithm to decide how to best process
191		/// the file and feed it into the hasher.
192		fn digest_for_file(
193		self,
194		file_path: impl AsRef<std::path::Path>,
195		file: fs::FileSlot,
196		size_hint: Option<u64>,
197		) -> impl Future<Output = Result<(DigestInfo, fs::FileSlot), Error>>;
198
199		/// Utility function to compute a hash from a generic reader.
200	13	fn compute_from_reader<R: AsyncRead + Unpin + Send>(
201	13	&mut self,
202	13	mut reader: R,
203	13	) -> impl Future<Output = Result<DigestInfo, Error>> {
204	13	async move {
205	13	let mut chunk = BytesMut::with_capacity(fs::DEFAULT_READ_BUFF_SIZE);
206		loop {
207	23	reader
208	23	.read_buf(&mut chunk)
209	23	.await
210	23	.err_tip(\|\| "Could not read chunk during compute_from_reader") ?0 ;
211	23	if chunk.is_empty() {
212	13	break; // EOF.
213	10	}
214	10	DigestHasher::update(self, &chunk);
215	10	chunk.clear();
216		}
217	13	Ok(DigestHasher::finalize_digest(self))
218	13	}
219	13	}
220		}
221
222		#[expect(
223		variant_size_differences,
224		reason = "some variants are already boxed; this is acceptable"
225		)]
226		#[derive(Debug)]
227		pub enum DigestHasherFuncImpl {
228		Sha256(Sha256),
229		Blake3(Box<Blake3Hasher>), // Box because Blake3Hasher is 1.3kb in size.
230		}
231
232		/// The individual implementation of the hash function.
233		#[derive(Debug)]
234		pub struct DigestHasherImpl {
235		hashed_size: u64,
236		hash_func_impl: DigestHasherFuncImpl,
237		}
238
239		impl DigestHasherImpl {
240		#[inline]
241	0	async fn hash_file(
242	0	&mut self,
243	0	mut file: fs::FileSlot,
244	13	) -> Result<(DigestInfo, fs::FileSlot), Error> {
245	13	let digest = self
246	13	.compute_from_reader(&mut file)
247	13	.await
248	13	.err_tip(\|\| "In digest_for_file") ?0 ;
249	13	Ok((digest, file))
250	13	}
251		}
252
253		impl DigestHasher for DigestHasherImpl {
254		#[inline]
255	5.25k	fn update(&mut self, input: &[u8]) {
256	5.25k	self.hashed_size += input.len() as u64;
257	5.25k	match &mut self.hash_func_impl {
258	201	DigestHasherFuncImpl::Sha256(h) => sha2::digest::Update::update(h, input),
259	5.04k	DigestHasherFuncImpl::Blake3(h) => {
260	5.04k	Blake3Hasher::update(h, input);
261	5.04k	}
262		}
263	5.25k	}
264
265		#[inline]
266	5.24k	fn finalize_digest(&mut self) -> DigestInfo {
267	5.24k	let hash = match &mut self.hash_func_impl {
268	198	DigestHasherFuncImpl::Sha256(h) => h.finalize_reset().into(),
269	5.04k	DigestHasherFuncImpl::Blake3(h) => h.finalize().into(),
270		};
271	5.24k	DigestInfo::new(hash, self.hashed_size)
272	5.24k	}
273
274	13	async fn digest_for_file(
275	13	mut self,
276	13	file_path: impl AsRef<std::path::Path>,
277	13	mut file: fs::FileSlot,
278	13	size_hint: Option<u64>,
279	13	) -> Result<(DigestInfo, fs::FileSlot), Error> {
280	13	let file_position = file
281	13	.stream_position()
282	13	.await
283	13	.err_tip(\|\| "Couldn't get stream position in digest_for_file") ?0 ;
284	13	if file_position != 0 {
285	0	return self.hash_file(file).await;
286	13	}
287		// If we are a small file, it's faster to just do it the "slow" way.
288		// Great read: https://github.com/david-slatinek/c-read-vs.-mmap
289	13	if let Some(size_hint) = size_hint
290	13	&& size_hint <= fs::DEFAULT_READ_BUFF_SIZE as u64
291		{
292	13	return self.hash_file(file).await;
293	0	}
294	0	let file_path = file_path.as_ref().to_path_buf();
295	0	match self.hash_func_impl {
296	0	DigestHasherFuncImpl::Sha256(_) => self.hash_file(file).await,
297	0	DigestHasherFuncImpl::Blake3(mut hasher) => {
298	0	spawn_blocking!("digest_for_file", move \|\| {
299	0	hasher.update_mmap(file_path).map_err(\|e\| {
300	0	Error::from_std_err(Code::Internal, &e)
301	0	.append("Error in blake3's update_mmap")
302	0	})?;
303	0	Result::<_, Error>::Ok((
304	0	DigestInfo::new(hasher.finalize().into(), hasher.count()),
305	0	file,
306	0	))
307	0	})
308	0	.await
309	0	.err_tip(\|\| "Could not spawn blocking task in digest_for_file")?
310		}
311		}
312	13	}
313		}