tools/blackout/blackout-target/src/static_tree.rs - fuchsia - Git at Google

 // 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,
     fidl_fuchsia_io as fio,
     rand::{distributions, Rng},
 };

 /// 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) -> distributions::Bernoulli {
         distributions::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 distributions::Distribution<FileEntry> for distributions::Standard {
     fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> FileEntry {
         let size = rng.gen_range(1..1 << 16);
         // 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.gen(), 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::OpenFlags::CREATE | fio::OpenFlags::RIGHT_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 distributions::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.gen_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.gen(), 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::OpenFlags::RIGHT_READABLE | fio::OpenFlags::RIGHT_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 distributions::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(distributions::Standard))
         }
     }
 }

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

     // 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, 0xdeadbeef, "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: 3735928580,
             entries: vec![
                 Entry::File(FileEntry { name: 3735928563, contents: vec![242] }),
                 Entry::File(FileEntry { name: 3735928567, contents: vec![246] }),
                 Entry::File(FileEntry { name: 3735928571, contents: vec![250] }),
                 Entry::File(FileEntry { name: 3735928575, contents: vec![254] }),
                 Entry::File(FileEntry { name: 3735928579, contents: vec![2] }),
             ],
         }
     }

     #[test]
     fn gen_tree() {
         let initial = 0xdeadbeef;
         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 = 0xdeadbeef;
         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 mut ramdisk =
             RamdiskClient::create(512, 1 << 16).await.expect("failed to make ramdisk");
         let controller = ramdisk.take_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");
     }
 }
	// 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,
	fidl_fuchsia_io as fio,
	rand::{distributions, Rng},
	};

	/// 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) -> distributions::Bernoulli {
	distributions::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 distributions::Distribution<FileEntry> for distributions::Standard {
	fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> FileEntry {
	let size = rng.gen_range(1..1 << 16);
	// 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.gen(), 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::OpenFlags::CREATE \| fio::OpenFlags::RIGHT_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 distributions::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.gen_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.gen(), 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::OpenFlags::RIGHT_READABLE \| fio::OpenFlags::RIGHT_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 distributions::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(distributions::Standard))
	}
	}
	}

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

	// 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, 0xdeadbeef, "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: 3735928580,
	entries: vec![
	Entry::File(FileEntry { name: 3735928563, contents: vec![242] }),
	Entry::File(FileEntry { name: 3735928567, contents: vec![246] }),
	Entry::File(FileEntry { name: 3735928571, contents: vec![250] }),
	Entry::File(FileEntry { name: 3735928575, contents: vec![254] }),
	Entry::File(FileEntry { name: 3735928579, contents: vec![2] }),
	],
	}
	}

	#[test]
	fn gen_tree() {
	let initial = 0xdeadbeef;
	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 = 0xdeadbeef;
	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 mut ramdisk =
	RamdiskClient::create(512, 1 << 16).await.expect("failed to make ramdisk");
	let controller = ramdisk.take_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");
	}
	}