blackout_target/

static_tree.rs

// Copyright 2019 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Generate a random directory tree that can be commited to disk all at once. The general formula
//! goes like this -
//!
//! ```
//! use {rand::{Rng, SeedableRng, rngs::StdRng}, static_tree::{EntryDistribution, DirectoryEntry}};
//! let mut rng = StdRng::seed_from_u64(seed);
//! let dist = EntryDistribution::new(depth);
//! let tree: DirectoryEntry = rng.sample(&dist);
//! tree.write_tree_at(root).expect("failed to write tree");
//! ```

use anyhow::Error;
use fidl_fuchsia_io as fio;
use rand::Rng;
use rand::distr::{Bernoulli, Distribution, StandardUniform};

/// A random distribution specialized to generation of random directory trees. This distribution
/// decreases the likelyhood of a directory being generated linearly relative to the depth of the
/// subset of the tree being generated, until it's 0% at the maximum depth.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct EntryDistribution {
    depth: u32,
    max_depth: u32,
}

impl EntryDistribution {
    /// Create a new `EntryDistribution` with a maximum depth of `max_depth`. This distribution is
    /// used for generating [`DirectoryEntry`]s and [`Entry`]s.
    pub fn new(max_depth: u32) -> EntryDistribution {
        EntryDistribution { depth: 1, max_depth }
    }

    /// get the entry distribution for the next level down on the directory tree.
    fn next_level(&self) -> EntryDistribution {
        debug_assert!(
            self.depth <= self.max_depth,
            "directory tree exceeded max depth ({} vs max of {}). programming error.",
            self.depth,
            self.max_depth
        );
        EntryDistribution { depth: self.depth + 1, max_depth: self.max_depth }
    }

    /// creates a bernoulli distribution where the likelyhood is progressively decreased at further
    /// depths until it's 0% at max_depth.
    fn directory_distribution(&self) -> Bernoulli {
        Bernoulli::from_ratio(self.max_depth - self.depth, self.max_depth).unwrap()
    }
}

/// A file entry in the generated tree. Contains a randomly generated u64 for a name and randomly
/// generated byte contents. The file size is random, in a range of 1 byte to 2^16 bytes (~65kb).
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct FileEntry {
    name: u64,
    contents: Vec<u8>,
}

impl Distribution<FileEntry> for StandardUniform {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> FileEntry {
        let size = usize::from(rng.random::<u16>());
        // unfortunately, we can't use sample_iter to generate the content here. the trait
        // definition for distribution requires the Rng type parameter to have ?Sized (or "maybe
        // sized"), but the sample_iter function requires the provided Rng be Sized, either through
        // the explicit Self: Sized on the Rng sample_iter function, or the implicit Sized present
        // on type parameters for the equivalent Distribution function.
        let mut contents = vec![0; size];
        rng.fill(contents.as_mut_slice());
        FileEntry { name: rng.random(), contents }
    }
}

impl FileEntry {
    async fn write_file_at(self, root: &fio::DirectoryProxy) -> Result<(), Error> {
        let file = fuchsia_fs::directory::open_file(
            root,
            &self.name.to_string(),
            fio::Flags::FLAG_MAYBE_CREATE | fio::PERM_WRITABLE,
        )
        .await?;
        fuchsia_fs::file::write(&file, &self.contents).await?;
        Ok(())
    }
}

/// A directory entry in the generated tree. Contains a randomly generated u64 for a name and a
/// vector of randomly generated directory entries, with a number of entries somewhere in the range
/// of [0, 6). The sample function generates the entire subtree depth-first before returning, and is
/// mutually recursive with the [`Entry`] sample function.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct DirectoryEntry {
    name: u64,
    entries: Vec<Entry>,
}

impl Distribution<DirectoryEntry> for EntryDistribution {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> DirectoryEntry {
        // each directory has a random number of entries in the range [0, 6)
        let num_entries = rng.random_range(0..6);
        let mut entries = vec![];
        let entry_dist = self.next_level();
        for _ in 0..num_entries {
            entries.push(rng.sample(&entry_dist));
        }
        DirectoryEntry { name: rng.random(), entries }
    }
}

impl DirectoryEntry {
    /// get the path the directory entry will be written to relative to a given root.
    pub fn get_name(&self) -> String {
        self.name.to_string()
    }

    /// Take the entire randomly generated tree and commit it to disk.
    pub fn write_tree_at<'a>(
        self,
        root: &'a fio::DirectoryProxy,
    ) -> futures::future::BoxFuture<'a, Result<(), Error>> {
        Box::pin(async {
            let this = fuchsia_fs::directory::create_directory(
                root,
                &self.get_name(),
                fio::PERM_READABLE | fio::PERM_WRITABLE,
            )
            .await?;
            for entry in self.entries {
                match entry {
                    Entry::File(file_entry) => file_entry.write_file_at(&this).await?,
                    Entry::Directory(dir_entry) => dir_entry.write_tree_at(&this).await?,
                }
            }
            Ok(())
        })
    }
}

/// An entry of either File or Directory. The chance of it being a directory is relative to the depth
/// recorded in EntryDistribution. The associated entry is also then generated.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Entry {
    /// This entry is a file.
    File(FileEntry),
    /// This entry is a directory
    Directory(DirectoryEntry),
}

impl Distribution<Entry> for EntryDistribution {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Entry {
        if rng.sample(self.directory_distribution()) {
            Entry::Directory(rng.sample(self))
        } else {
            Entry::File(rng.sample(StandardUniform))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{DirectoryEntry, Entry, EntryDistribution, FileEntry};
    use fs_management::Minfs;
    use fuchsia_async as fasync;
    use ramdevice_client::RamdiskClient;
    use rand::{Rng as _, RngCore};

    struct StepRng {
        state: u64,
        increment: u64,
    }

    impl StepRng {
        pub fn new(initial: u64, increment: u64) -> Self {
            Self { state: initial, increment: increment }
        }
    }

    impl RngCore for StepRng {
        fn next_u32(&mut self) -> u32 {
            self.next_u64() as u32
        }

        fn next_u64(&mut self) -> u64 {
            let r = self.state;
            self.state = self.state.wrapping_add(self.increment);
            r
        }

        fn fill_bytes(&mut self, dst: &mut [u8]) {
            rand_core::impls::fill_bytes_via_next(self, dst)
        }
    }

    // this fixture will have to get updated any time the generation logic changes. depending on the
    // nature of the change, the comparison logic may have to be changed as well.
    fn get_fixture(initial: u64, increment: u64, depth: u32) -> DirectoryEntry {
        // confirm that the rng and depth used to generate this tree are the same.
        assert_eq!(initial, 0xFFFF0000, "initial value changed - update fixture");
        assert_eq!(increment, 1, "increment value changed - update fixture");
        assert_eq!(depth, 2, "depth value changed - update fixture");
        DirectoryEntry {
            name: 4294901785,
            entries: vec![
                Entry::File(FileEntry { name: 4294901764, contents: vec![3, 0] }),
                Entry::File(FileEntry { name: 4294901768, contents: vec![7, 0, 255, 255, 0, 0] }),
                Entry::File(FileEntry {
                    name: 4294901773,
                    contents: vec![11, 0, 255, 255, 0, 0, 0, 0, 12, 0],
                }),
                Entry::File(FileEntry {
                    name: 4294901778,
                    contents: vec![16, 0, 255, 255, 0, 0, 0, 0, 17, 0, 255, 255, 0, 0, 0],
                }),
                Entry::File(FileEntry {
                    name: 4294901784,
                    contents: vec![
                        21, 0, 255, 255, 0, 0, 0, 0, 22, 0, 255, 255, 0, 0, 0, 0, 23, 0, 255, 255,
                    ],
                }),
            ],
        }
    }

    #[test]
    fn gen_tree() {
        let initial = 0xFFFF0000;
        let increment = 1;
        let depth = 2;
        // make sure we are generating a tree at all.
        let mut rng = StepRng::new(initial, increment);
        let dist = EntryDistribution::new(depth);
        let tree: DirectoryEntry = rng.sample(dist);
        // this skips the contents for the files, because they are huge, so we do our own manual
        // comparison.
        let fixture = get_fixture(initial, increment, depth);
        assert_eq!(fixture, tree);
        for i in 0..fixture.entries.len() {
            match &fixture.entries[i] {
                Entry::File(fixture_file_entry) => match &tree.entries[i] {
                    Entry::File(tree_file_entry) => {
                        assert_eq!(fixture_file_entry.name, tree_file_entry.name)
                    }
                    _ => panic!("expected a file in generated tree"),
                },
                _ => panic!("expected a file in fixture tree"),
            }
        }
    }

    #[test]
    fn same_rng_same_tree() {
        // this confirms an assumption about the properties of tree generation that we rely on for
        // consistency checking.
        let dist = EntryDistribution::new(5);

        let tree1: DirectoryEntry = StepRng::new(1337, 1).sample(dist);
        let tree2: DirectoryEntry = StepRng::new(1337, 1).sample(dist);

        assert_eq!(tree1, tree2);
    }

    #[fasync::run_singlethreaded(test)]
    async fn write_tree() {
        let root = "/test-root";
        let initial = 0xFFFF0000;
        let increment = 1;
        let depth = 2;

        let mut rng = StepRng::new(initial, increment);
        let dist = EntryDistribution::new(depth);
        let tree: DirectoryEntry = rng.sample(dist);

        let ramdisk = RamdiskClient::create(512, 1 << 16).await.expect("failed to make ramdisk");
        let controller = ramdisk.open_controller().expect("invalid controller");
        let mut minfs = Minfs::new(controller);

        minfs.format().await.expect("failed to format minfs");
        let mut minfs = minfs.serve().await.expect("failed to mount minfs");
        minfs.bind_to_path(root).expect("failed to bind path");

        tree.write_tree_at(minfs.root()).await.expect("failed to write tree");

        let fixture = get_fixture(initial, increment, depth);
        let path = std::path::PathBuf::from(format!("{}/{}", root, fixture.name));
        assert!(path.is_dir(), "{}", path.display());
        for (i, entry) in std::fs::read_dir(&path).expect("failed to read directory").enumerate() {
            let entry = entry.expect("failed to read entry");
            let file_type = entry.file_type().expect("failed to get file type");
            assert!(file_type.is_file());
            let file_name =
                entry.file_name().into_string().expect("failed to convert file name to string");
            let expected_name = match &fixture.entries[i] {
                Entry::File(f) => f.name,
                Entry::Directory(_) => panic!("expected a file in the fixture tree"),
            };
            assert_eq!(file_name, expected_name.to_string());
        }

        minfs.shutdown().await.expect("failed to unmount minfs");
        ramdisk.destroy().await.expect("failed to destroy ramdisk");
    }
}