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

snix_castore/directoryservice/
object_store.rs
