snix_castore/directoryservice/
object_store.rs

1use std::collections::HashMap;
2use std::collections::hash_map;
3use std::sync::Arc;
4
5use data_encoding::HEXLOWER;
6use futures::SinkExt;
7use futures::StreamExt;
8use futures::TryStreamExt;
9use futures::stream::BoxStream;
10use object_store::ObjectStoreExt;
11use object_store::{ObjectStore, path::Path};
12use prost::Message;
13use tokio::io::AsyncWriteExt;
14use tokio_util::codec::LengthDelimitedCodec;
15use tonic::async_trait;
16use tracing::{Level, instrument, trace, warn};
17use url::Url;
18
19use super::{Directory, DirectoryPutter, DirectoryService};
20use crate::composition::{CompositionContext, ServiceBuilder};
21use crate::directoryservice::directory_graph::DirectoryGraphBuilder;
22use crate::directoryservice::order_validator::{self, LeavesToRoot, OrderValidator, RootToLeaves};
23use crate::{B3Digest, Node, proto};
24
25/// Stores directory closures in an object store.
26/// Notably, this makes use of the option to disallow accessing child directories except when
27/// fetching them recursively via the top-level directory, since all batched writes
28/// (using `put_multiple_start`) are stored in a single object.
29/// Directories are stored in a length-delimited format with a 1MiB limit. The length field is a
30/// u32 and the directories are stored in root-to-leaves topological order, the same way they will
31/// be returned to the client in get_recursive.
32#[derive(Clone)]
33pub struct ObjectStoreDirectoryService {
34    instance_name: String,
35    object_store: Arc<dyn ObjectStore>,
36    base_path: Path,
37}
38
39#[instrument(level=Level::TRACE, skip_all,fields(base_path=%base_path,blob.digest=%digest),ret(Display))]
40fn derive_dirs_path(base_path: &Path, digest: &B3Digest) -> Path {
41    base_path
42        .clone()
43        .join("dirs")
44        .join("b3")
45        .join(HEXLOWER.encode(&digest.as_slice()[..2]))
46        .join(HEXLOWER.encode(digest.as_slice()))
47}
48
49/// Helper function, parsing protobuf-encoded Directories into [crate::Directory].
50fn parse_proto_directory(encoded_directory: &[u8]) -> Result<crate::Directory, Error> {
51    let directory_proto = proto::Directory::decode(encoded_directory)?;
52
53    Ok(Directory::try_from(directory_proto)?)
54}
55
56#[allow(clippy::identity_op)]
57const MAX_FRAME_LENGTH: usize = 1 * 1024 * 1024 * 1000; // 1 MiB
58//
59impl ObjectStoreDirectoryService {
60    /// Constructs a new [ObjectStoreDirectoryService] from a [Url] supported by
61    /// [object_store].
62    /// Any path suffix becomes the base path of the object store.
63    /// additional options, the same as in [object_store::parse_url_opts] can
64    /// be passed.
65    pub fn parse_url_opts<I, K, V>(url: &Url, options: I) -> Result<Self, object_store::Error>
66    where
67        I: IntoIterator<Item = (K, V)>,
68        K: AsRef<str>,
69        V: Into<String>,
70    {
71        let (object_store, path) = object_store::parse_url_opts(url, options)?;
72
73        Ok(Self {
74            instance_name: "root".into(),
75            object_store: Arc::new(object_store),
76            base_path: path,
77        })
78    }
79
80    /// Like [Self::parse_url_opts], except without the options.
81    pub fn parse_url(url: &Url) -> Result<Self, object_store::Error> {
82        Self::parse_url_opts(url, Vec::<(String, String)>::new())
83    }
84
85    pub fn new(instance_name: String, object_store: Arc<dyn ObjectStore>, base_path: Path) -> Self {
86        Self {
87            instance_name,
88            object_store,
89            base_path,
90        }
91    }
92}
93
94#[async_trait]
95impl DirectoryService for ObjectStoreDirectoryService {
96    /// This is the same steps as for get_recursive anyways, so we just call get_recursive and
97    /// return the first element of the stream and drop the request.
98    #[instrument(level = "trace", skip_all, fields(directory.digest = %digest, instance_name = %self.instance_name))]
99    async fn get(&self, digest: &B3Digest) -> Result<Option<Directory>, super::Error> {
100        self.get_recursive(digest).take(1).next().await.transpose()
101    }
102
103    #[instrument(level = "trace", skip_all, fields(directory.digest = %directory.digest(), instance_name = %self.instance_name))]
104    async fn put(&self, directory: Directory) -> Result<B3Digest, super::Error> {
105        // Ensure the directory doesn't contain other directory children
106        if directory
107            .nodes()
108            .any(|(_, e)| matches!(e, Node::Directory { .. }))
109        {
110            Err(Error::PutForDirectoryWithChildren)?
111        }
112
113        let mut handle = self.put_multiple_start();
114        handle.put(directory).await?;
115        handle.close().await
116    }
117
118    #[instrument(level = "trace", skip_all, fields(directory.digest = %root_directory_digest, instance_name = %self.instance_name))]
119    fn get_recursive(
120        &self,
121        root_directory_digest: &B3Digest,
122    ) -> BoxStream<'_, Result<Directory, super::Error>> {
123        // Check that we are not passing on bogus from the object store to the client, and that the
124        // trust chain from the root digest to the leaves is intact.
125        let dir_path = derive_dirs_path(&self.base_path, root_directory_digest);
126        let object_store = &self.object_store;
127        let root_directory_digest = *root_directory_digest;
128
129        async_stream::try_stream! {
130                let bytes_stream = match object_store.get(&dir_path).await {
131                    Ok(v) => v.into_stream(),
132                    Err(object_store::Error::NotFound { .. }) => {
133                        return;
134                    }
135                    Err(e) => Err(Error::ObjectStore(e))?,
136                };
137
138                // get a reader of the response body.
139                let r = tokio_util::io::StreamReader::new(bytes_stream);
140                let decompressed_stream = async_compression::tokio::bufread::ZstdDecoder::new(r);
141
142                // the subdirectories are stored in a length delimited format
143                let mut encoded_directories = LengthDelimitedCodec::builder()
144                    .max_frame_length(MAX_FRAME_LENGTH)
145                    .length_field_type::<u32>()
146                    .new_read(decompressed_stream)
147                    .err_into::<Error>();
148
149                let mut order_validator = RootToLeaves::new_with_root_digest(root_directory_digest);
150                while let Some(encoded_directory) = encoded_directories.try_next().await? {
151                    // hash the encoded proto message, only proceed if we would accept a directory with this digest.
152                    let digest = B3Digest::from(blake3::hash(&encoded_directory));
153                    if !order_validator.would_accept(&digest) {
154                        Err(Error::UnexpectedDigest(digest))?;
155                    }
156
157                    // only then proceed with parsing
158                    let directory = parse_proto_directory(&encoded_directory)?;
159
160                    // The directory can still be rejected for other reasons.
161                    order_validator.try_accept(&directory).map_err(Error::DirectoryOrdering)?;
162
163                    yield directory;
164                }
165
166                order_validator.finalize().map_err(Error::DirectoryOrdering)?;
167        }
168        .boxed()
169    }
170
171    #[instrument(skip_all)]
172    fn put_multiple_start(&self) -> Box<dyn DirectoryPutter + '_>
173    where
174        Self: Clone,
175    {
176        Box::new(ObjectStoreDirectoryPutter::new(
177            self.object_store.clone(),
178            &self.base_path,
179        ))
180    }
181}
182
183#[derive(thiserror::Error, Debug)]
184enum Error {
185    #[error("wrong arguments: {0}")]
186    WrongConfig(&'static str),
187    #[error("put() may only be used for directories without children")]
188    PutForDirectoryWithChildren,
189
190    #[error("Directory Graph ordering error")]
191    DirectoryOrdering(#[from] order_validator::OrderingError),
192    #[error("next directory in batch has unexpected digest {0}")]
193    UnexpectedDigest(B3Digest),
194    #[error("failed to decode protobuf: {0}")]
195    ProtobufDecode(#[from] prost::DecodeError),
196    #[error("failed to validate directory: {0}")]
197    DirectoryValidation(#[from] crate::DirectoryError),
198
199    #[error("DirectoryPutter already closed")]
200    DirectoryPutterAlreadyClosed,
201
202    #[error("ObjectStore error: {0}")]
203    ObjectStore(#[from] object_store::Error),
204
205    #[error("io error: {0}")]
206    IO(#[from] std::io::Error),
207}
208impl From<Error> for super::Error {
209    fn from(value: Error) -> Self {
210        Self(Box::new(value))
211    }
212}
213
214#[derive(serde::Deserialize)]
215#[serde(deny_unknown_fields)]
216pub struct ObjectStoreDirectoryServiceConfig {
217    object_store_url: String,
218    #[serde(default)]
219    object_store_options: HashMap<String, String>,
220}
221
222impl TryFrom<url::Url> for ObjectStoreDirectoryServiceConfig {
223    type Error = Box<dyn std::error::Error + Send + Sync>;
224    fn try_from(url: url::Url) -> Result<Self, Self::Error> {
225        // We need to convert the URL to string, strip the prefix there, and then
226        // parse it back as url, as Url::set_scheme() rejects some of the transitions we want to do.
227        let trimmed_url = {
228            let s = url.to_string();
229            let mut url = Url::parse(s.strip_prefix("objectstore+").ok_or(Error::WrongConfig(
230                "Missing objectstore+ part in URI scheme",
231            ))?)?;
232            // trim the query pairs, they might contain credentials or local settings we don't want to send as-is.
233            url.set_query(None);
234            url
235        };
236        Ok(ObjectStoreDirectoryServiceConfig {
237            object_store_url: trimmed_url.into(),
238            object_store_options: url
239                .query_pairs()
240                .into_iter()
241                .map(|(k, v)| (k.to_string(), v.to_string()))
242                .collect(),
243        })
244    }
245}
246
247#[async_trait]
248impl ServiceBuilder for ObjectStoreDirectoryServiceConfig {
249    type Output = dyn DirectoryService;
250    async fn build<'a>(
251        &'a self,
252        instance_name: &str,
253        _context: &CompositionContext,
254    ) -> Result<Arc<Self::Output>, Box<dyn std::error::Error + Send + Sync>> {
255        let opts = {
256            let mut opts: HashMap<&str, _> = self
257                .object_store_options
258                .iter()
259                .map(|(k, v)| (k.as_str(), v.as_str()))
260                .collect();
261
262            if let hash_map::Entry::Vacant(e) =
263                opts.entry(object_store::ClientConfigKey::UserAgent.as_ref())
264            {
265                e.insert(crate::USER_AGENT);
266            }
267
268            opts
269        };
270
271        let (object_store, path) =
272            object_store::parse_url_opts(&self.object_store_url.parse()?, opts)?;
273        Ok(Arc::new(ObjectStoreDirectoryService::new(
274            instance_name.to_string(),
275            Arc::new(object_store),
276            path,
277        )))
278    }
279}
280
281struct ObjectStoreDirectoryPutter<'a> {
282    object_store: Arc<dyn ObjectStore>,
283    base_path: &'a Path,
284
285    builder: Option<DirectoryGraphBuilder<LeavesToRoot>>,
286}
287
288impl<'a> ObjectStoreDirectoryPutter<'a> {
289    fn new(object_store: Arc<dyn ObjectStore>, base_path: &'a Path) -> Self {
290        Self {
291            object_store,
292            base_path,
293            builder: Some(DirectoryGraphBuilder::<LeavesToRoot>::new()),
294        }
295    }
296}
297
298#[async_trait]
299impl DirectoryPutter for ObjectStoreDirectoryPutter<'_> {
300    #[instrument(level = "trace", skip_all, fields(directory.digest=%directory.digest()), err)]
301    async fn put(&mut self, directory: Directory) -> Result<(), super::Error> {
302        let builder = self
303            .builder
304            .as_mut()
305            .ok_or_else(|| Error::DirectoryPutterAlreadyClosed)?;
306
307        builder
308            .try_insert(directory)
309            .map_err(Error::DirectoryOrdering)?;
310
311        Ok(())
312    }
313
314    #[instrument(level = "trace", skip_all, ret, err)]
315    async fn close(&mut self) -> Result<B3Digest, super::Error> {
316        let builder = self
317            .builder
318            .take()
319            .ok_or_else(|| Error::DirectoryPutterAlreadyClosed)?;
320
321        // Retrieve the validated directories.
322        let directory_graph = builder.build().map_err(Error::DirectoryOrdering)?;
323        let root_digest = directory_graph.root().digest();
324
325        let dir_path = derive_dirs_path(self.base_path, &root_digest);
326
327        match self.object_store.head(&dir_path).await {
328            // directory tree already exists, nothing to do
329            Ok(_) => {
330                trace!("directory tree already exists");
331            }
332
333            // directory tree does not yet exist, compress and upload.
334            Err(object_store::Error::NotFound { .. }) => {
335                trace!("uploading directory tree");
336
337                let object_store_writer =
338                    object_store::buffered::BufWriter::new(self.object_store.clone(), dir_path);
339                let compressed_writer =
340                    async_compression::tokio::write::ZstdEncoder::new(object_store_writer);
341                let mut directories_sink = LengthDelimitedCodec::builder()
342                    .max_frame_length(MAX_FRAME_LENGTH)
343                    .length_field_type::<u32>()
344                    .new_write(compressed_writer);
345
346                // Drain the graph in *Root-To-Leaves*, order, as that's how we write it to storage.
347                for directory in directory_graph.drain_root_to_leaves() {
348                    directories_sink
349                        .send(proto::Directory::from(directory).encode_to_vec().into())
350                        .await
351                        .map_err(Error::IO)?;
352                }
353
354                let mut compressed_writer = directories_sink.into_inner();
355                compressed_writer.shutdown().await.map_err(Error::IO)?;
356            }
357            // other error
358            Err(err) => Err(Error::ObjectStore(err))?,
359        }
360
361        Ok(root_digest)
362    }
363}
snix_castore/directoryservice/object_store.rs

snix_castore/directoryservice/
object_store.rs
