snix_store/nar/
listing.rs

1use futures::TryStreamExt;
2use nix_compat::nar::listing::{Listing, ListingEntry, ListingVersion};
3use snix_castore::{
4    B3Digest, Directory, Node,
5    directoryservice::{DirectoryService, OrderingError, RootToLeavesValidator},
6};
7use std::{
8    collections::{BTreeMap, HashMap},
9    sync::atomic::{AtomicU64, Ordering},
10};
11
12/// A writer that only counts the number of bytes written so far,
13/// updating a &AtomicU64.
14/// We can't use `count_write::CountWrite`.
15/// As we hand out a &mut Write to the NAR writer, we can't have other read-only
16/// references until we drop it, yet we need to be able to access the current
17/// offsets to put in the listing scructure we build up at the same time.
18struct CountingWriter<'a> {
19    bytes_written: &'a AtomicU64,
20}
21
22impl std::io::Write for CountingWriter<'_> {
23    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
24        self.bytes_written
25            .fetch_add(buf.len() as u64, Ordering::Relaxed);
26        Ok(buf.len())
27    }
28
29    fn flush(&mut self) -> std::io::Result<()> {
30        Ok(())
31    }
32}
33
34/// Accepts a [Node] pointing to the root of a (store) path,
35/// and uses the passed [DirectoryService] to expand the entire structure.
36/// Then assembles a [Listing].
37pub async fn produce_listing<DS>(root_node: &Node, directory_service: &DS) -> Result<Listing, Error>
38where
39    DS: DirectoryService,
40{
41    let mut directories: HashMap<B3Digest, Directory> = HashMap::new();
42
43    if let Node::Directory { digest, .. } = root_node {
44        let mut directories_stream = directory_service.get_recursive(digest);
45        let mut validator = RootToLeavesValidator::new_with_root_digest(digest.to_owned());
46        while let Some(directory) = directories_stream.try_next().await? {
47            validator.try_accept(&directory)?;
48            directories.insert(directory.digest(), directory);
49        }
50    };
51
52    let bytes_written = AtomicU64::new(0);
53    let mut writer = CountingWriter {
54        bytes_written: &bytes_written,
55    };
56    let nar_node = nix_compat::nar::writer::open(&mut writer).expect("writing");
57
58    // bytes_written.fetch_min(20, Ordering::Relaxed);
59
60    Ok(Listing::V1 {
61        root: produce_listing_inner(root_node, &bytes_written, nar_node, &|digest| {
62            // The lookup is infallible, as we validate the closures, or we never call the function at all.
63            directories
64                .get(digest)
65                .expect("Snix bug: lookup of unknown directory")
66        })?,
67        version: ListingVersion,
68    })
69}
70
71fn produce_listing_inner<'d, F>(
72    node: &Node,
73    bytes_written: &AtomicU64,
74    writer_node: nix_compat::nar::writer::Node<'_, impl std::io::Write>,
75    get_directory: &'d F,
76) -> Result<ListingEntry, Error>
77where
78    F: Fn(&B3Digest) -> &'d Directory + 'd,
79{
80    Ok(match node {
81        Node::Directory { digest, .. } => {
82            let mut writer_entries = writer_node.directory().expect("nar writing");
83
84            let nodes_it = get_directory(digest).nodes();
85
86            let mut child_entries: BTreeMap<String, ListingEntry> = BTreeMap::new();
87
88            for (child_name, child_node) in nodes_it {
89                let writer_child_node = writer_entries
90                    .entry(child_name.as_ref())
91                    .expect("nar writing");
92                child_entries.insert(
93                    str::from_utf8(child_name.as_ref())
94                        .or(Err(Error::PathComponentIsNoString))?
95                        .to_owned(),
96                    produce_listing_inner(
97                        child_node,
98                        bytes_written,
99                        writer_child_node,
100                        get_directory,
101                    )?,
102                );
103            }
104
105            writer_entries.close().expect("nar writing");
106
107            ListingEntry::Directory {
108                entries: child_entries,
109            }
110        }
111        Node::File {
112            size, executable, ..
113        } => {
114            let (w, hdl) = writer_node
115                .file_manual_write(*executable, *size)
116                .expect("nar writing");
117
118            // This is the offset we want to store, before we add to the counter.
119            let nar_offset = bytes_written.fetch_add(*size, Ordering::Relaxed);
120
121            hdl.close(w).expect("nar writing");
122
123            ListingEntry::Regular {
124                size: *size,
125                executable: *executable,
126                nar_offset,
127            }
128        }
129        Node::Symlink { target } => {
130            writer_node.symlink(target.as_ref()).expect("nar writing");
131
132            ListingEntry::Symlink {
133                target: str::from_utf8(target.as_ref())
134                    .or(Err(Error::SymlinkTargetIsNoString))?
135                    .to_owned(),
136            }
137        }
138    })
139}
140#[derive(Debug, thiserror::Error)]
141pub enum Error {
142    #[error("symlink target is not a string")]
143    SymlinkTargetIsNoString,
144    #[error("path component is not a string")]
145    PathComponentIsNoString,
146    #[error("from directoryservice: {0}")]
147    DirectoryService(#[from] snix_castore::directoryservice::Error),
148    #[error("ordering error from directoryservice: {0}")]
149    OrderingError(#[from] OrderingError),
150}
151
152#[cfg(test)]
153mod tests {
154    use super::produce_listing;
155    use crate::fixtures::{
156        CASTORE_NODE_COMPLICATED, CASTORE_NODE_HELLOWORLD, CASTORE_NODE_SYMLINK,
157    };
158    use rstest::rstest;
159    use snix_castore::Directory;
160    use snix_castore::fixtures::{DIRECTORY_COMPLICATED, DIRECTORY_WITH_KEEP};
161    use snix_castore::{
162        Node, directoryservice::DirectoryService, utils::gen_test_directory_service,
163    };
164
165    #[tokio::test]
166    #[rstest]
167    #[case::symlink(&CASTORE_NODE_SYMLINK, vec![],
168        r#"{"root":{"target":"/nix/store/somewhereelse","type":"symlink"},"version":1}"#)]
169    #[case::blob(&CASTORE_NODE_HELLOWORLD, vec![],
170        r#"{"root":{"narOffset":96,"size":12,"type":"regular"},"version":1}"#)]
171    #[case::complicated(&CASTORE_NODE_COMPLICATED, vec![&*DIRECTORY_COMPLICATED, &*DIRECTORY_WITH_KEEP],
172        r#"{"root":{"entries":{".keep":{"narOffset":232,"size":0,"type":"regular"},"aa":{"target":"/nix/store/somewhereelse","type":"symlink"},"keep":{"entries":{".keep":{"narOffset":760,"size":0,"type":"regular"}},"type":"directory"}},"type":"directory"},"version":1}"#)]
173    async fn produce_listing_test(
174        #[case] root_node: &Node,
175        #[case] directories: Vec<&Directory>,
176        #[case] exp_json: &str,
177    ) {
178        let svc = gen_test_directory_service();
179        for directory in directories {
180            svc.put(directory.to_owned()).await.expect("must insert");
181        }
182
183        let listing = produce_listing(root_node, &svc)
184            .await
185            .expect("must succeed");
186        let actual_json = serde_json::to_string(&listing).expect("must serialize");
187
188        assert_eq!(exp_json, actual_json);
189    }
190}
snix_store/nar/listing.rs

snix_store/nar/
listing.rs
