snix_store/nar/
listing.rs

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

snix_store/nar/
listing.rs
