about summary refs log tree commit diff
path: root/tvix/castore/src/directoryservice/closure_validator.rs
blob: 6f263157a9fb290b0608847ae2b84eedf324690b (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
use std::collections::{HashMap, HashSet};

use bstr::ByteSlice;

use tracing::{instrument, trace};

use crate::{
    proto::{self, Directory},
    B3Digest, Error,
};

/// This can be used to validate a Directory closure (DAG of connected
/// Directories), and their insertion order.
///
/// Directories need to be inserted (via `add`), in an order from the leaves to
/// the root. During insertion, We validate as much as we can at that time:
///
///  - individual validation of Directory messages
///  - validation of insertion order (no upload of not-yet-known Directories)
///  - validation of size fields of referred Directories
///
/// Internally it keeps all received Directories (and their sizes) in a HashMap,
/// keyed by digest.
///
/// Once all Directories have been inserted, a finalize function can be
/// called to get a (deduplicated and) validated list of directories, in
/// insertion order.
/// During finalize, a check for graph connectivity is performed too, to ensure
/// there's no disconnected components, and only one root.
#[derive(Default)]
pub struct ClosureValidator {
    directories_and_sizes: HashMap<B3Digest, (Directory, u64)>,

    /// Keeps track of the last-inserted directory digest. Used to start the
    /// connectivity check.
    last_directory_digest: Option<B3Digest>,
}

impl ClosureValidator {
    /// Insert a new Directory into the closure.
    /// Perform individual Directory validation, validation of insertion order
    /// and size fields.
    #[instrument(level = "trace", skip_all, fields(directory.digest=%directory.digest()), err)]
    pub fn add(&mut self, directory: proto::Directory) -> Result<(), Error> {
        let digest = directory.digest();

        if self.directories_and_sizes.contains_key(&digest) {
            trace!("already seen, skipping");
            return Ok(());
        }

        // Do some general validation
        directory
            .validate()
            .map_err(|e| Error::InvalidRequest(e.to_string()))?;

        // Ensure the directory only refers to directories which we already accepted.
        for dir in &directory.directories {
            let dir_dgst = B3Digest::try_from(dir.digest.to_owned()).unwrap(); // validated

            // Ensure the digest has already been seen
            let (_, recorded_dir_size) =
                self.directories_and_sizes.get(&dir_dgst).ok_or_else(|| {
                    Error::InvalidRequest(format!(
                        "'{}' refers to unseen child dir: {}",
                        dir.name.as_bstr(),
                        dir_dgst
                    ))
                })?;

            // Ensure the size specified in the child node matches our records.
            if dir.size != *recorded_dir_size {
                return Err(Error::InvalidRequest(format!(
                    "'{}' has wrong size, specified {}, recorded {}",
                    dir.name.as_bstr(),
                    dir.size,
                    recorded_dir_size
                )));
            }
        }

        trace!("inserting");
        let directory_size = directory.size();
        self.directories_and_sizes
            .insert(digest.clone(), (directory, directory_size));
        self.last_directory_digest = Some(digest);

        Ok(())
    }

    /// Ensure that all inserted Directories are connected, then return a
    /// (deduplicated) and validated list of directories, in from-leaves-to-root
    /// order.
    /// In case no elements have been inserted, returns an empty list.
    #[instrument(level = "trace", skip_all, err)]
    pub(crate) fn finalize(mut self) -> Result<Vec<Directory>, Error> {
        if self.last_directory_digest.is_none() {
            return Ok(vec![]);
        }
        let root_digest = self.last_directory_digest.unwrap();

        // recursively walk all directories reachable from there,
        // ensure we visited all nodes that were uploaded.
        // If not, we might have received multiple disconnected graphs.

        // The list of directories we know we need to visit.
        // Once we're finished working, and (there's no errors), this in reversed order will
        // be a valid insertion order, and directories_and_sizes will be empty.
        let mut dirs_to_visit = Vec::with_capacity(self.directories_and_sizes.len());
        dirs_to_visit.push(
            self.directories_and_sizes
                .remove(&root_digest)
                .expect("root digest not found")
                .0,
        );
        // The set of digests seen while going visiting all directories.
        let mut seen_dir_digests = HashSet::new();

        // This loop moves gradually to the end of `dirs_to_visit`, while it's being
        // extended.
        // The loop stops once it reaches the end.
        let mut i = 0;
        while let Some(directory) = dirs_to_visit.get(i).map(|d| d.to_owned()) {
            // lookup all child directories and put them in the back of dirs_to_visit,
            // if they're not already there.
            for child_dir in &directory.directories {
                let child_digest = B3Digest::try_from(child_dir.digest.to_owned()).unwrap(); // validated

                // In case the digest is neither already visited nor already in dirs_to_visit,
                // add it to the end of it.
                // We don't need to check for the hash we're currently
                // visiting, as we can't self-reference.
                if !seen_dir_digests.contains(&child_digest) {
                    dirs_to_visit.push(
                        self.directories_and_sizes
                            .remove(&child_digest)
                            .expect("child digest not found")
                            .0,
                    );
                    seen_dir_digests.insert(child_digest);
                }
            }

            i += 1;
        }

        // check directories_and_sizes is empty.
        if !self.directories_and_sizes.is_empty() {
            if cfg!(debug_assertions) {
                return Err(Error::InvalidRequest(format!(
                    "found {} disconnected nodes: {:?}",
                    self.directories_and_sizes.len(),
                    self.directories_and_sizes
                )));
            } else {
                return Err(Error::InvalidRequest(format!(
                    "found {} disconnected nodes",
                    self.directories_and_sizes.len()
                )));
            }
        }

        // Reverse to have correct insertion order.
        dirs_to_visit.reverse();

        Ok(dirs_to_visit)
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        fixtures::{DIRECTORY_A, DIRECTORY_B, DIRECTORY_C},
        proto::{self, Directory},
    };
    use lazy_static::lazy_static;
    use rstest::rstest;

    lazy_static! {
        pub static ref BROKEN_DIRECTORY : Directory = Directory {
            symlinks: vec![proto::SymlinkNode {
                name: "".into(), // invalid name!
                target: "doesntmatter".into(),
            }],
            ..Default::default()
        };

        pub static ref BROKEN_PARENT_DIRECTORY: Directory = Directory {
            directories: vec![proto::DirectoryNode {
                name: "foo".into(),
                digest: DIRECTORY_A.digest().into(),
                size: DIRECTORY_A.size() + 42, // wrong!
            }],
            ..Default::default()
        };
    }

    use super::ClosureValidator;

    #[rstest]
    /// Uploading an empty directory should succeed.
    #[case::empty_directory(&[&*DIRECTORY_A], false, Some(vec![&*DIRECTORY_A]))]
    /// Uploading A, then B (referring to A) should succeed.
    #[case::simple_closure(&[&*DIRECTORY_A, &*DIRECTORY_B], false, Some(vec![&*DIRECTORY_A, &*DIRECTORY_B]))]
    /// Uploading A, then A, then C (referring to A twice) should succeed.
    /// We pretend to be a dumb client not deduping directories.
    #[case::same_child(&[&*DIRECTORY_A, &*DIRECTORY_A, &*DIRECTORY_C], false, Some(vec![&*DIRECTORY_A, &*DIRECTORY_C]))]
    /// Uploading A, then C (referring to A twice) should succeed.
    #[case::same_child_dedup(&[&*DIRECTORY_A, &*DIRECTORY_C], false, Some(vec![&*DIRECTORY_A, &*DIRECTORY_C]))]
    /// Uploading A, then C (referring to A twice), then B (itself referring to A) should fail during close,
    /// as B itself would be left unconnected.
    #[case::unconnected_node(&[&*DIRECTORY_A, &*DIRECTORY_C, &*DIRECTORY_B], false, None)]
    /// Uploading B (referring to A) should fail immediately, because A was never uploaded.
    #[case::dangling_pointer(&[&*DIRECTORY_B], true, None)]
    /// Uploading a directory failing validation should fail immediately.
    #[case::failing_validation(&[&*BROKEN_DIRECTORY], true, None)]
    /// Uploading a directory which refers to another Directory with a wrong size should fail.
    #[case::wrong_size_in_parent(&[&*DIRECTORY_A, &*BROKEN_PARENT_DIRECTORY], true, None)]
    fn test_uploads(
        #[case] directories_to_upload: &[&Directory],
        #[case] exp_fail_upload_last: bool,
        #[case] exp_finalize: Option<Vec<&Directory>>, // Some(_) if finalize successful, None if not.
    ) {
        let mut dcv = ClosureValidator::default();
        let len_directories_to_upload = directories_to_upload.len();

        for (i, d) in directories_to_upload.iter().enumerate() {
            let resp = dcv.add((*d).clone());
            if i == len_directories_to_upload - 1 && exp_fail_upload_last {
                assert!(resp.is_err(), "expect last put to fail");

                // We don't really care anymore what finalize() would return, as
                // the add() failed.
                return;
            } else {
                assert!(resp.is_ok(), "expect put to succeed");
            }
        }

        // everything was uploaded successfully. Test drain().
        let resp = dcv.finalize();

        match exp_finalize {
            Some(exp_drain) => {
                assert_eq!(
                    Vec::from_iter(exp_drain.iter().map(|e| (*e).to_owned())),
                    resp.expect("drain should succeed")
                );
            }
            None => {
                resp.expect_err("drain should fail");
            }
        }
    }
}