/build/source/nativelink-store/src/verify_store.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 core::pin::Pin;
use std::sync::Arc;

use async_trait::async_trait;
use nativelink_config::stores::VerifySpec;
use nativelink_error::{Error, ResultExt, make_input_err};
use nativelink_metric::MetricsComponent;
use nativelink_util::buf_channel::{
    DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair,
};
use nativelink_util::common::PackedHash;
use nativelink_util::digest_hasher::{
    DigestHasher, DigestHasherFunc, digest_hasher_func_from_context,
};
use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator};
use nativelink_util::metrics_utils::CounterWithTime;
use nativelink_util::store_trait::{
    RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
};

#[derive(Debug, MetricsComponent)]
pub struct VerifyStore {
    #[metric(group = "inner_store")]
    inner_store: Store,
    #[metric(help = "If the verification store is verifying the size of the data")]
    verify_size: bool,
    #[metric(help = "If the verification store is verifying the hash of the data")]
    verify_hash: bool,

    // Metrics.
    #[metric(help = "Number of failures the verification store had due to size mismatches")]
    size_verification_failures: CounterWithTime,
    #[metric(help = "Number of failures the verification store had due to hash mismatches")]
    hash_verification_failures: CounterWithTime,
}

impl VerifyStore {
    pub fn new(spec: &VerifySpec, inner_store: Store) -> Arc<Self> {
        Arc::new(Self {
            inner_store,
            verify_size: spec.verify_size,
            verify_hash: spec.verify_hash,
            size_verification_failures: CounterWithTime::default(),
            hash_verification_failures: CounterWithTime::default(),
        })
    }

    async fn inner_check_update<D: DigestHasher>(
        &self,
        mut tx: DropCloserWriteHalf,
        mut rx: DropCloserReadHalf,
        maybe_expected_digest_size: Option<u64>,
        original_hash: &PackedHash,
        digest_function: Option<DigestHasherFunc>,
        mut maybe_hasher: Option<&mut D>,
    ) -> Result<u64, Error> {
        let mut sum_size: u64 = 0;
        loop {
            let chunk = rx
                .recv()
                .await
                .err_tip(|| "Failed to read chunk in check_update in verify store")?0;
            sum_size += chunk.len() as u64;

            // Ensure if a user sends us too much data we fail quickly.
            if let Some(expected_size18) = maybe_expected_digest_size {
                match sum_size.cmp(&expected_size) {
                    core::cmp::Ordering::Greater => {
                        self.size_verification_failures.inc();
                        return Err(make_input_err!(
                            "Expected size {} but already received {} on insert",
                            expected_size,
                            sum_size
                        ));
                    }
                    core::cmp::Ordering::Equal => {
                        // Ensure our next chunk is the EOF chunk.
                        // If this was an error it'll be caught on the .recv()
                        // on next cycle.
                        if let Ok(eof_chunk) = rx.peek().await
                            && !eof_chunk.is_empty()
                        {
                            self.size_verification_failures.inc();
                            return Err(make_input_err!(
                                "Expected EOF chunk when exact size was hit on insert in verify store - {}",
                                expected_size,
                            ));
                        }
                    }
                    core::cmp::Ordering::Less => {}
                }
            }

            // If is EOF.
            if chunk.is_empty() {
                if let Some(expected_size7) = maybe_expected_digest_size
                    && sum_size != expected_size
                {
                    self.size_verification_failures.inc();
                    return Err(make_input_err!(
                        "Expected size {} but got size {} on insert",
                        expected_size,
                        sum_size
                    ));
                }
                if let Some(hasher8) = maybe_hasher.as_mut() {
                    let digest = hasher.finalize_digest();
                    let hash_result = digest.packed_hash();
                    if original_hash != hash_result {
                        self.hash_verification_failures.inc();
                        let Some(digest_function) = digest_function else {
                            return Err(make_input_err!(
                                "Hash verification failed without a digest function"
                            ));
                        };
                        return Err(make_input_err!(
                            "Hash verification using {digest_function} failed: client declared {digest_function}:{original_hash}, but the server computed {digest_function}:{hash_result}",
                        ));
                    }
                }
                tx.send_eof().err_tip(|| "In verify_store::check_update")?0;
                break;
            }

            // This will allows us to hash while sending data to another thread.
            let write_future = tx.send(chunk.clone());

            if let Some(hasher9) = maybe_hasher.as_mut() {
                hasher.update(chunk.as_ref());
            }6

            write_future
                .await
                .err_tip(|| "Failed to write chunk to inner store in verify store")?0;
        }
        Ok(sum_size)
    }
}

#[async_trait]
impl StoreDriver for VerifyStore {
    async fn post_init(self: Arc<Self>) -> Result<(), Error> {
        self.inner_store.clone().into_inner().post_init().await?;
        Ok(())
    }

    async fn has_with_results(
        self: Pin<&Self>,
        digests: &[StoreKey<'_>],
        results: &mut [Option<u64>],
    ) -> Result<(), Error> {
        self.inner_store.has_with_results(digests, results).await
    }

    async fn update(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        reader: DropCloserReadHalf,
        size_info: UploadSizeInfo,
    ) -> Result<u64, Error> {
        let StoreKey::Digest(digest) = key else {
            return Err(make_input_err!(
                "Only digests are supported in VerifyStore. Got {key:?}"
            ));
        };
        let digest_size = digest.size_bytes();
        if let UploadSizeInfo::ExactSize(expected_size) = size_info
            && self.verify_size
            && expected_size != digest_size
        {
            self.size_verification_failures.inc();
            return Err(make_input_err!(
                "Expected size to match. Got {} but digest says {} on update",
                expected_size,
                digest_size
            ));
        }

        let digest_function = if self.verify_hash {
            Some(digest_hasher_func_from_context())
        } else {
            None
        };
        let mut hasher = digest_function.map(|digest_function| digest_function.hasher());

        let maybe_digest_size = if self.verify_size {
            Some(digest_size)
        } else {
            None
        };
        let (tx, rx) = make_buf_channel_pair();

        let update_fut = self.inner_store.update(digest, rx, size_info);
        let check_fut = self.inner_check_update(
            tx,
            reader,
            maybe_digest_size,
            digest.packed_hash(),
            digest_function,
            hasher.as_mut(),
        );

        let (update_res, check_res) = tokio::join!(update_fut, check_fut);

        match (update_res, check_res) {
            // Prioritize the check future's error, as it's more specific.
            (_, Err(e)) | (Err(e), Ok(_)) => Err(e),
            (Ok(size), Ok(_)) => Ok(size),
        }
    }

    async fn get_part(
        self: Pin<&Self>,
        key: StoreKey<'_>,
        writer: &mut DropCloserWriteHalf,
        offset: u64,
        length: Option<u64>,
    ) -> Result<(), Error> {
        self.inner_store.get_part(key, writer, offset, length).await
    }

    fn inner_store(&self, _digest: Option<StoreKey>) -> &'_ dyn StoreDriver {
        self
    }

    fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) {
        self
    }

    fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
        self
    }

    fn register_remove_callback(self: Arc<Self>, callback: RemoveCallback) -> Result<(), Error> {
        self.inner_store.register_remove_callback(callback)
    }
}

default_health_status_indicator!(VerifyStore);

Coverage Report

Created: 2026-08-24 08:15

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 core::pin::Pin;
16		use std::sync::Arc;
17
18		use async_trait::async_trait;
19		use nativelink_config::stores::VerifySpec;
20		use nativelink_error::{Error, ResultExt, make_input_err};
21		use nativelink_metric::MetricsComponent;
22		use nativelink_util::buf_channel::{
23		DropCloserReadHalf, DropCloserWriteHalf, make_buf_channel_pair,
24		};
25		use nativelink_util::common::PackedHash;
26		use nativelink_util::digest_hasher::{
27		DigestHasher, DigestHasherFunc, digest_hasher_func_from_context,
28		};
29		use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator};
30		use nativelink_util::metrics_utils::CounterWithTime;
31		use nativelink_util::store_trait::{
32		RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
33		};
34
35		#[derive(Debug, MetricsComponent)]
36		pub struct VerifyStore {
37		#[metric(group = "inner_store")]
38		inner_store: Store,
39		#[metric(help = "If the verification store is verifying the size of the data")]
40		verify_size: bool,
41		#[metric(help = "If the verification store is verifying the hash of the data")]
42		verify_hash: bool,
43
44		// Metrics.
45		#[metric(help = "Number of failures the verification store had due to size mismatches")]
46		size_verification_failures: CounterWithTime,
47		#[metric(help = "Number of failures the verification store had due to hash mismatches")]
48		hash_verification_failures: CounterWithTime,
49		}
50
51		impl VerifyStore {
52	13	pub fn new(spec: &VerifySpec, inner_store: Store) -> Arc<Self> {
53	13	Arc::new(Self {
54	13	inner_store,
55	13	verify_size: spec.verify_size,
56	13	verify_hash: spec.verify_hash,
57	13	size_verification_failures: CounterWithTime::default(),
58	13	hash_verification_failures: CounterWithTime::default(),
59	13	})
60	13	}
61
62	13	async fn inner_check_update<D: DigestHasher>(
63	13	&self,
64	13	mut tx: DropCloserWriteHalf,
65	13	mut rx: DropCloserReadHalf,
66	13	maybe_expected_digest_size: Option<u64>,
67	13	original_hash: &PackedHash,
68	13	digest_function: Option<DigestHasherFunc>,
69	13	mut maybe_hasher: Option<&mut D>,
70	13	) -> Result<u64, Error> {
71	13	let mut sum_size: u64 = 0;
72		loop {
73	28	let chunk = rx
74	28	.recv()
75	28	.await
76	28	.err_tip(\|\| "Failed to read chunk in check_update in verify store") ?0 ;
77	28	sum_size += chunk.len() as u64;
78
79		// Ensure if a user sends us too much data we fail quickly.
80	28	if let Some( expected_size18 ) = maybe_expected_digest_size {
81	18	match sum_size.cmp(&expected_size) {
82		core::cmp::Ordering::Greater => {
83	1	self.size_verification_failures.inc();
84	1	return Err(make_input_err!(
85	1	"Expected size {} but already received {} on insert",
86	1	expected_size,
87	1	sum_size
88	1	));
89		}
90		core::cmp::Ordering::Equal => {
91		// Ensure our next chunk is the EOF chunk.
92		// If this was an error it'll be caught on the .recv()
93		// on next cycle.
94	12	if let Ok(eof_chunk) = rx.peek().await
95	12	&& !eof_chunk.is_empty()
96		{
97	0	self.size_verification_failures.inc();
98	0	return Err(make_input_err!(
99	0	"Expected EOF chunk when exact size was hit on insert in verify store - {}",
100	0	expected_size,
101	0	));
102	12	}
103		}
104	5	core::cmp::Ordering::Less => {}
105		}
106	10	}
107
108		// If is EOF.
109	27	if chunk.is_empty() {
110	12	if let Some( expected_size7 ) = maybe_expected_digest_size
111	7	&& sum_size != expected_size
112		{
113	1	self.size_verification_failures.inc();
114	1	return Err(make_input_err!(
115	1	"Expected size {} but got size {} on insert",
116	1	expected_size,
117	1	sum_size
118	1	));
119	11	}
120	11	if let Some( hasher8 ) = maybe_hasher.as_mut() {
121	8	let digest = hasher.finalize_digest();
122	8	let hash_result = digest.packed_hash();
123	8	if original_hash != hash_result {
124	2	self.hash_verification_failures.inc();
125	2	let Some(digest_function) = digest_function else {
126	0	return Err(make_input_err!(
127	0	"Hash verification failed without a digest function"
128	0	));
129		};
130	2	return Err(make_input_err!(
131	2	"Hash verification using {digest_function} failed: client declared {digest_function}:{original_hash}, but the server computed {digest_function}:{hash_result}",
132	2	));
133	6	}
134	3	}
135	9	tx.send_eof().err_tip(\|\| "In verify_store::check_update") ?0 ;
136	9	break;
137	15	}
138
139		// This will allows us to hash while sending data to another thread.
140	15	let write_future = tx.send(chunk.clone());
141
142	15	if let Some( hasher9 ) = maybe_hasher.as_mut() {
143	9	hasher.update(chunk.as_ref());
144	9	}6
145
146	15	write_future
147	15	.await
148	15	.err_tip(\|\| "Failed to write chunk to inner store in verify store") ?0 ;
149		}
150	9	Ok(sum_size)
151	13	}
152		}
153
154		#[async_trait]
155		impl StoreDriver for VerifyStore {
156	0	async fn post_init(self: Arc<Self>) -> Result<(), Error> {
157		self.inner_store.clone().into_inner().post_init().await?;
158		Ok(())
159	0	}
160
161		async fn has_with_results(
162		self: Pin<&Self>,
163		digests: &[StoreKey<'_>],
164		results: &mut [Option<u64>],
165	0	) -> Result<(), Error> {
166		self.inner_store.has_with_results(digests, results).await
167	0	}
168
169		async fn update(
170		self: Pin<&Self>,
171		key: StoreKey<'_>,
172		reader: DropCloserReadHalf,
173		size_info: UploadSizeInfo,
174	13	) -> Result<u64, Error> {
175		let StoreKey::Digest(digest) = key else {
176		return Err(make_input_err!(
177		"Only digests are supported in VerifyStore. Got {key:?}"
178		));
179		};
180		let digest_size = digest.size_bytes();
181		if let UploadSizeInfo::ExactSize(expected_size) = size_info
182		&& self.verify_size
183		&& expected_size != digest_size
184		{
185		self.size_verification_failures.inc();
186		return Err(make_input_err!(
187		"Expected size to match. Got {} but digest says {} on update",
188		expected_size,
189		digest_size
190		));
191		}
192
193		let digest_function = if self.verify_hash {
194		Some(digest_hasher_func_from_context())
195		} else {
196		None
197		};
198	8	let mut hasher = digest_function.map(\|digest_function\| digest_function.hasher());
199
200		let maybe_digest_size = if self.verify_size {
201		Some(digest_size)
202		} else {
203		None
204		};
205		let (tx, rx) = make_buf_channel_pair();
206
207		let update_fut = self.inner_store.update(digest, rx, size_info);
208		let check_fut = self.inner_check_update(
209		tx,
210		reader,
211		maybe_digest_size,
212		digest.packed_hash(),
213		digest_function,
214		hasher.as_mut(),
215		);
216
217		let (update_res, check_res) = tokio::join!(update_fut, check_fut);
218
219		match (update_res, check_res) {
220		// Prioritize the check future's error, as it's more specific.
221		(_, Err(e)) \| (Err(e), Ok(_)) => Err(e),
222		(Ok(size), Ok(_)) => Ok(size),
223		}
224	13	}
225
226		async fn get_part(
227		self: Pin<&Self>,
228		key: StoreKey<'_>,
229		writer: &mut DropCloserWriteHalf,
230		offset: u64,
231		length: Option<u64>,
232	3	) -> Result<(), Error> {
233		self.inner_store.get_part(key, writer, offset, length).await
234	3	}
235
236	5	fn inner_store(&self, _digest: Option<StoreKey>) -> &'_ dyn StoreDriver {
237	5	self
238	5	}
239
240	5	fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) {
241	5	self
242	5	}
243
244	0	fn as_any_arc(self: Arc<Self>) -> Arc<dyn core::any::Any + Sync + Send + 'static> {
245	0	self
246	0	}
247
248	0	fn register_remove_callback(self: Arc<Self>, callback: RemoveCallback) -> Result<(), Error> {
249	0	self.inner_store.register_remove_callback(callback)
250	0	}
251		}
252
253		default_health_status_indicator!(VerifyStore);