<commit_before>extern crate bazel_protos;
extern crate boxfuture;
extern crate clap;
extern crate fs;
extern crate itertools;
extern crate futures;
extern crate protobuf;

use boxfuture::{Boxable, BoxFuture};
use clap::{App, Arg, SubCommand};
use fs::hash::Fingerprint;
use fs::store::Store;
use fs::VFS;
use futures::future::{Future, join_all};
use itertools::Itertools;
use std::error::Error;
use std::ffi::OsString;
use std::fs::File;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::process::exit;
use std::sync::Arc;

#[derive(Debug)]
enum ExitCode {
  UnknownError = 1,
  NotFound = 2,
}

#[derive(Debug)]
struct ExitError(pub String, pub ExitCode);

impl From<String> for ExitError {
  fn from(s: String) -> Self {
    ExitError(s, ExitCode::UnknownError)
  }
}

fn main() {
  match execute(
    App::new("fs_util")
      .subcommand(
        SubCommand::with_name("file")
          .subcommand(
            SubCommand::with_name("cat")
              .about("Output the contents of a file by fingerprint.")
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              )),
          )
          .subcommand(
            SubCommand::with_name("save")
              .about(
                "Ingest a file by path, which allows it to be used in Directories/Snapshots. \
Outputs a fingerprint of its contents and its size in bytes, separated by a space.",
              )
              .arg(Arg::with_name("path").required(true).takes_value(true)),
          ),
      )
      .subcommand(
        SubCommand::with_name("directory")
          .subcommand(
            SubCommand::with_name("materialize")
              .about(
                "Materialize a directory by fingerprint to the filesystem. \
Destination must not exist before this command is run.",
              )
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              ))
              .arg(Arg::with_name("destination").required(true).takes_value(
                true,
              )),
          )
          .subcommand(
            SubCommand::with_name("save")
              .about(
                "Ingest a directory recursively. Saves all files found therein and saves Directory \
protos for each directory found. Outputs a fingerprint of the canonical top-level Directory proto \
and the size of the serialized proto in bytes, separated by a space.",
              )
              .arg(
                Arg::with_name("globs")
                  .required(true)
                  .takes_value(true)
                  .multiple(true)
                  .help(
                    "globs matching the files and directories which should be included in the \
directory, relative to the root.",
                  ),
              )
                .arg(Arg::with_name("root").long("root").required(true).takes_value(true).help(
                  "Root under which the globs live. The Directory proto produced will be relative \
to this directory.",
                )),
          )
          .subcommand(
            SubCommand::with_name("cat-proto")
              .about(
                "Output the bytes of a serialized Directory proto addressed by fingerprint.",
              )
              .arg(
                Arg::with_name("output-format")
                  .long("output-format")
                  .takes_value(true)
                  .default_value("binary")
                  .possible_values(&["binary", "text"]),
              )
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              )),
          ),
      )
      .subcommand(
        SubCommand::with_name("cat")
          .about(
            "Output the contents of a file or Directory proto addressed by fingerprint.",
          )
          .arg(Arg::with_name("fingerprint").required(true).takes_value(
            true,
          )),
      )
      .arg(
        Arg::with_name("local_store_path")
          .takes_value(true)
          .long("local_store_path")
          .required(true),
      )
      .get_matches(),
  ) {
    Ok(_) => {}
    Err(err) => {
      eprintln!("{}", err.0);
      exit(err.1 as i32)
    }
  };
}

fn execute(top_match: clap::ArgMatches) -> Result<(), ExitError> {
  let store_dir = top_match.value_of("local_store_path").unwrap();
  let store = Arc::new(Store::new(store_dir).map_err(|e| {
    format!(
      "Failed to open/create store for directory {}: {}",
      store_dir,
      e
    )
  })?);

  match top_match.subcommand() {
    ("file", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("cat", Some(args)) => {
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          match store.load_file_bytes(&fingerprint)? {
            Some(bytes) => {
              io::stdout().write(&bytes).unwrap();
              Ok(())
            }
            None => Err(ExitError(
              format!("File with fingerprint {} not found", fingerprint),
              ExitCode::NotFound,
            )),
          }
        }
        ("save", Some(args)) => {
          let path = PathBuf::from(args.value_of("path").unwrap());
          // Canonicalize path to guarantee that a relative path has a parent.
          let posix_fs = make_posix_fs(path
            .canonicalize()
            .map_err(|e| {
              format!("Error canonicalizing path {:?}: {}", path, e.description())
            })?
            .parent()
            .ok_or_else(|| {
              format!("File being saved must have parent but {:?} did not", path)
            })?);
          let file = posix_fs
            .stat(PathBuf::from(path.file_name().unwrap()))
            .unwrap();
          match file {
            fs::Stat::File(f) => {
              let (fingerprint, size_bytes) =
                save_file(store, Arc::new(posix_fs), f).wait().unwrap();
              Ok(println!("{} {}", fingerprint, size_bytes))
            }
            o => Err(
              format!(
                "Tried to save file {:?} but it was not a file, was a {:?}",
                path,
                o
              ).into(),
            ),
          }

        }
        (_, _) => unimplemented!(),
      }
    }
    ("directory", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("materialize", Some(args)) => {
          let destination = Path::new(args.value_of("destination").unwrap());
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          Ok(materialize_directory(store, destination, &fingerprint)?)
        }
        ("save", Some(args)) => {
          let posix_fs = Arc::new(make_posix_fs(args.value_of("root").unwrap()));
          let (fingerprint, size_bytes) = posix_fs
            .expand(fs::PathGlobs::create(
              &args
                .values_of("globs")
                .unwrap()
                .map(|s| s.to_string())
                .collect::<Vec<String>>(),
              &[],
            )?)
            .map_err(|e| format!("Error expanding globs: {}", e.description()))
            .and_then(move |paths| save_directory(store, posix_fs, paths))
            .wait()?;
          Ok(println!("{} {}", fingerprint, size_bytes))
        }
        ("cat-proto", Some(args)) => {
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          let proto_bytes = match args.value_of("output-format").unwrap() {
            "binary" => store.load_directory_proto_bytes(&fingerprint),
            "text" => {
              store.load_directory_proto(&fingerprint).map(|maybe_p| {
                maybe_p.map(|p| format!("{:?}\n", p).as_bytes().to_vec())
              })
            }
            format => Err(format!(
              "Unexpected value of --output-format arg: {}",
              format
            )),
          }?;
          match proto_bytes {
            Some(bytes) => {
              io::stdout().write(&bytes).unwrap();
              Ok(())
            }
            None => Err(ExitError(
              format!(
                "Directory with fingerprint {} not found",
                fingerprint
              ),
              ExitCode::NotFound,
            )),
          }
        }
        (_, _) => unimplemented!(),
      }
    }
    ("cat", Some(args)) => {
      let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
      let v = match store.load_file_bytes(&fingerprint)? {
        None => store.load_directory_proto_bytes(&fingerprint)?,
        some => some,
      };
      match v {
        Some(bytes) => {
          io::stdout().write(&bytes).unwrap();
          Ok(())
        }
        None => Err(ExitError(
          format!("Fingerprint {} not found", fingerprint),
          ExitCode::NotFound,
        )),
      }
    }

    (_, _) => unimplemented!(),
  }
}

fn make_posix_fs<P: AsRef<Path>>(root: P) -> fs::PosixFS {
  fs::PosixFS::new(&root, vec![]).unwrap()
}

fn save_file(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  file: fs::File,
) -> BoxFuture<(Fingerprint, usize), String> {
  posix_fs
    .read_file(&file)
    .map_err(move |err| {
      format!("Error reading file {:?}: {}", file, err.description())
    })
    .and_then(move |content| {
      Ok((
        store.store_file_bytes(&content.content)?,
        content.content.len(),
      ))
    })
    .to_boxed()
}

fn save_directory(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  mut path_stats: Vec<fs::PathStat>,
) -> BoxFuture<(Fingerprint, usize), String> {
  let mut file_futures: Vec<BoxFuture<bazel_protos::remote_execution::FileNode, String>> =
    Vec::new();
  let mut dir_futures: Vec<BoxFuture<bazel_protos::remote_execution::DirectoryNode, String>> =
    Vec::new();

  path_stats.sort_by(|a, b| a.path().cmp(b.path()));

  for (first_component, group) in
    &path_stats.into_iter().group_by(|s| {
      s.path().components().next().unwrap().as_os_str().to_owned()
    })
  {
    let mut path_group: Vec<fs::PathStat> = group.collect();
    if path_group.len() == 1 && path_group.get(0).unwrap().path().components().count() == 1 {
      // Exactly one entry with exactly one component indicates either a file in this directory, or
      // an empty directory.
      // If the child is a non-empty directory, or a file therein, there must be multiple PathStats
      // with that prefix component, and we will handle that in the recursive save_directory call.

      match path_group.pop().unwrap() {
        fs::PathStat::File { ref stat, .. } => {
          let is_executable = stat.is_executable;
          file_futures.push(
            save_file(store.clone(), posix_fs.clone(), stat.clone())
              .and_then(move |(fingerprint, size_bytes)| {
                let mut file_node = bazel_protos::remote_execution::FileNode::new();
                file_node.set_name(osstring_as_utf8(first_component)?);
                file_node.set_digest(fingerprint_to_digest(&fingerprint, size_bytes as i64));
                file_node.set_is_executable(is_executable);
                Ok(file_node)
              })
              .to_boxed(),
          );
        }
        fs::PathStat::Dir { .. } => {
          // Because there are no children of this Dir, it must be empty.
          let (fingerprint, size_bytes) = store
            .record_directory(&bazel_protos::remote_execution::Directory::new())
            .unwrap();
          let mut directory_node = bazel_protos::remote_execution::DirectoryNode::new();
          directory_node.set_name(osstring_as_utf8(first_component).unwrap());
          directory_node.set_digest(fingerprint_to_digest(&fingerprint, size_bytes as i64));
          dir_futures.push(futures::future::ok(directory_node).to_boxed());
        }
      }
    } else {
      dir_futures.push(
        save_directory(
          store.clone(),
          posix_fs.clone(),
          paths_of_child_dir(path_group),
        ).and_then(move |(fingerprint, size_bytes)| {
          let mut dir_node = bazel_protos::remote_execution::DirectoryNode::new();
          dir_node.set_name(osstring_as_utf8(first_component)?);
          dir_node.set_digest(fingerprint_to_digest(&fingerprint, size_bytes as i64));
          Ok(dir_node)
        })
          .to_boxed(),
      );
    }
  }
  join_all(dir_futures)
    .join(join_all(file_futures))
    .and_then(move |(dirs, files)| {
      let mut directory = bazel_protos::remote_execution::Directory::new();
      directory.set_directories(protobuf::RepeatedField::from_vec(dirs));
      directory.set_files(protobuf::RepeatedField::from_vec(files));
      store.record_directory(&directory)
    })
    .to_boxed()
}

fn paths_of_child_dir(paths: Vec<fs::PathStat>) -> Vec<fs::PathStat> {
  paths
    .into_iter()
    .filter_map(|s| {
      if s.path().components().count() == 1 {
        return None;
      }
      Some(match s {
        fs::PathStat::File { path, stat } => {
          fs::PathStat::File {
            path: path.iter().skip(1).collect(),
            stat: stat,
          }
        }
        fs::PathStat::Dir { path, stat } => {
          fs::PathStat::Dir {
            path: path.iter().skip(1).collect(),
            stat: stat,
          }
        }
      })
    })
    .collect()
}

fn materialize_directory(
  store: Arc<Store>,
  destination: &Path,
  fingerprint: &Fingerprint,
) -> Result<(), ExitError> {
  let directory = store.load_directory_proto(&fingerprint)?.ok_or_else(|| {
    ExitError(
      format!("Directory with fingerprint {} not found", fingerprint),
      ExitCode::NotFound,
    )
  })?;
  make_clean_dir(&destination).map_err(|e| {
    format!(
      "Error making directory {:?}: {}",
      destination,
      e.description()
    )
  })?;
  for file_node in directory.get_files() {
    let fingerprint = &Fingerprint::from_hex_string(&file_node.get_digest().get_hash())?;
    match store.load_file_bytes(fingerprint)? {
      Some(bytes) => {
        let path = destination.join(file_node.get_name());
        File::create(&path)
          .and_then(|mut f| f.write_all(&bytes))
          .map_err(|e| {
            format!("Error writing file {:?}: {}", path, e.description())
          })?;
      }
      None => {
        return Err(ExitError(
          format!(
            "File with fingerprint {} not found",
            file_node.get_digest().get_hash()
          ),
          ExitCode::NotFound,
        ))
      }
    }
  }
  for directory_node in directory.get_directories() {
    materialize_directory(
      store.clone(),
      &destination.join(directory_node.get_name()),
      &Fingerprint::from_hex_string(directory_node.get_digest().get_hash())?,
    )?;
  }
  Ok(())
}

pub fn fingerprint_to_digest(
  fingerprint: &Fingerprint,
  size_bytes: i64,
) -> bazel_protos::remote_execution::Digest {
  let mut digest = bazel_protos::remote_execution::Digest::new();
  digest.set_hash(fingerprint.to_hex());
  digest.set_size_bytes(size_bytes);
  digest
}

fn osstring_as_utf8(path: OsString) -> Result<String, String> {
  path.into_string().map_err(|p| {
    format!("{:?}'s file_name is not representable in UTF8", p)
  })
}

fn make_clean_dir(path: &Path) -> Result<(), io::Error> {
  let parent = path.parent().ok_or_else(|| {
    io::Error::new(io::ErrorKind::NotFound, format!("{:?} had no parent", path))
  })?;
  std::fs::create_dir_all(parent)?;
  std::fs::create_dir(path)
}
<commit_msg>Rename local_store_path arg to local-store-path (#5092)<commit_after>extern crate bazel_protos;
extern crate boxfuture;
extern crate clap;
extern crate fs;
extern crate itertools;
extern crate futures;
extern crate protobuf;

use boxfuture::{Boxable, BoxFuture};
use clap::{App, Arg, SubCommand};
use fs::hash::Fingerprint;
use fs::store::Store;
use fs::VFS;
use futures::future::{Future, join_all};
use itertools::Itertools;
use std::error::Error;
use std::ffi::OsString;
use std::fs::File;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::process::exit;
use std::sync::Arc;

#[derive(Debug)]
enum ExitCode {
  UnknownError = 1,
  NotFound = 2,
}

#[derive(Debug)]
struct ExitError(pub String, pub ExitCode);

impl From<String> for ExitError {
  fn from(s: String) -> Self {
    ExitError(s, ExitCode::UnknownError)
  }
}

fn main() {
  match execute(
    App::new("fs_util")
      .subcommand(
        SubCommand::with_name("file")
          .subcommand(
            SubCommand::with_name("cat")
              .about("Output the contents of a file by fingerprint.")
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              )),
          )
          .subcommand(
            SubCommand::with_name("save")
              .about(
                "Ingest a file by path, which allows it to be used in Directories/Snapshots. \
Outputs a fingerprint of its contents and its size in bytes, separated by a space.",
              )
              .arg(Arg::with_name("path").required(true).takes_value(true)),
          ),
      )
      .subcommand(
        SubCommand::with_name("directory")
          .subcommand(
            SubCommand::with_name("materialize")
              .about(
                "Materialize a directory by fingerprint to the filesystem. \
Destination must not exist before this command is run.",
              )
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              ))
              .arg(Arg::with_name("destination").required(true).takes_value(
                true,
              )),
          )
          .subcommand(
            SubCommand::with_name("save")
              .about(
                "Ingest a directory recursively. Saves all files found therein and saves Directory \
protos for each directory found. Outputs a fingerprint of the canonical top-level Directory proto \
and the size of the serialized proto in bytes, separated by a space.",
              )
              .arg(
                Arg::with_name("globs")
                  .required(true)
                  .takes_value(true)
                  .multiple(true)
                  .help(
                    "globs matching the files and directories which should be included in the \
directory, relative to the root.",
                  ),
              )
                .arg(Arg::with_name("root").long("root").required(true).takes_value(true).help(
                  "Root under which the globs live. The Directory proto produced will be relative \
to this directory.",
                )),
          )
          .subcommand(
            SubCommand::with_name("cat-proto")
              .about(
                "Output the bytes of a serialized Directory proto addressed by fingerprint.",
              )
              .arg(
                Arg::with_name("output-format")
                  .long("output-format")
                  .takes_value(true)
                  .default_value("binary")
                  .possible_values(&["binary", "text"]),
              )
              .arg(Arg::with_name("fingerprint").required(true).takes_value(
                true,
              )),
          ),
      )
      .subcommand(
        SubCommand::with_name("cat")
          .about(
            "Output the contents of a file or Directory proto addressed by fingerprint.",
          )
          .arg(Arg::with_name("fingerprint").required(true).takes_value(
            true,
          )),
      )
      .arg(
        Arg::with_name("local-store-path")
          .takes_value(true)
          .long("local-store-path")
          .required(true),
      )
      .get_matches(),
  ) {
    Ok(_) => {}
    Err(err) => {
      eprintln!("{}", err.0);
      exit(err.1 as i32)
    }
  };
}

fn execute(top_match: clap::ArgMatches) -> Result<(), ExitError> {
  let store_dir = top_match.value_of("local-store-path").unwrap();
  let store = Arc::new(Store::new(store_dir).map_err(|e| {
    format!(
      "Failed to open/create store for directory {}: {}",
      store_dir,
      e
    )
  })?);

  match top_match.subcommand() {
    ("file", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("cat", Some(args)) => {
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          match store.load_file_bytes(&fingerprint)? {
            Some(bytes) => {
              io::stdout().write(&bytes).unwrap();
              Ok(())
            }
            None => Err(ExitError(
              format!("File with fingerprint {} not found", fingerprint),
              ExitCode::NotFound,
            )),
          }
        }
        ("save", Some(args)) => {
          let path = PathBuf::from(args.value_of("path").unwrap());
          // Canonicalize path to guarantee that a relative path has a parent.
          let posix_fs = make_posix_fs(path
            .canonicalize()
            .map_err(|e| {
              format!("Error canonicalizing path {:?}: {}", path, e.description())
            })?
            .parent()
            .ok_or_else(|| {
              format!("File being saved must have parent but {:?} did not", path)
            })?);
          let file = posix_fs
            .stat(PathBuf::from(path.file_name().unwrap()))
            .unwrap();
          match file {
            fs::Stat::File(f) => {
              let (fingerprint, size_bytes) =
                save_file(store, Arc::new(posix_fs), f).wait().unwrap();
              Ok(println!("{} {}", fingerprint, size_bytes))
            }
            o => Err(
              format!(
                "Tried to save file {:?} but it was not a file, was a {:?}",
                path,
                o
              ).into(),
            ),
          }

        }
        (_, _) => unimplemented!(),
      }
    }
    ("directory", Some(sub_match)) => {
      match sub_match.subcommand() {
        ("materialize", Some(args)) => {
          let destination = Path::new(args.value_of("destination").unwrap());
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          Ok(materialize_directory(store, destination, &fingerprint)?)
        }
        ("save", Some(args)) => {
          let posix_fs = Arc::new(make_posix_fs(args.value_of("root").unwrap()));
          let (fingerprint, size_bytes) = posix_fs
            .expand(fs::PathGlobs::create(
              &args
                .values_of("globs")
                .unwrap()
                .map(|s| s.to_string())
                .collect::<Vec<String>>(),
              &[],
            )?)
            .map_err(|e| format!("Error expanding globs: {}", e.description()))
            .and_then(move |paths| save_directory(store, posix_fs, paths))
            .wait()?;
          Ok(println!("{} {}", fingerprint, size_bytes))
        }
        ("cat-proto", Some(args)) => {
          let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
          let proto_bytes = match args.value_of("output-format").unwrap() {
            "binary" => store.load_directory_proto_bytes(&fingerprint),
            "text" => {
              store.load_directory_proto(&fingerprint).map(|maybe_p| {
                maybe_p.map(|p| format!("{:?}\n", p).as_bytes().to_vec())
              })
            }
            format => Err(format!(
              "Unexpected value of --output-format arg: {}",
              format
            )),
          }?;
          match proto_bytes {
            Some(bytes) => {
              io::stdout().write(&bytes).unwrap();
              Ok(())
            }
            None => Err(ExitError(
              format!(
                "Directory with fingerprint {} not found",
                fingerprint
              ),
              ExitCode::NotFound,
            )),
          }
        }
        (_, _) => unimplemented!(),
      }
    }
    ("cat", Some(args)) => {
      let fingerprint = Fingerprint::from_hex_string(args.value_of("fingerprint").unwrap())?;
      let v = match store.load_file_bytes(&fingerprint)? {
        None => store.load_directory_proto_bytes(&fingerprint)?,
        some => some,
      };
      match v {
        Some(bytes) => {
          io::stdout().write(&bytes).unwrap();
          Ok(())
        }
        None => Err(ExitError(
          format!("Fingerprint {} not found", fingerprint),
          ExitCode::NotFound,
        )),
      }
    }

    (_, _) => unimplemented!(),
  }
}

fn make_posix_fs<P: AsRef<Path>>(root: P) -> fs::PosixFS {
  fs::PosixFS::new(&root, vec![]).unwrap()
}

fn save_file(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  file: fs::File,
) -> BoxFuture<(Fingerprint, usize), String> {
  posix_fs
    .read_file(&file)
    .map_err(move |err| {
      format!("Error reading file {:?}: {}", file, err.description())
    })
    .and_then(move |content| {
      Ok((
        store.store_file_bytes(&content.content)?,
        content.content.len(),
      ))
    })
    .to_boxed()
}

fn save_directory(
  store: Arc<Store>,
  posix_fs: Arc<fs::PosixFS>,
  mut path_stats: Vec<fs::PathStat>,
) -> BoxFuture<(Fingerprint, usize), String> {
  let mut file_futures: Vec<BoxFuture<bazel_protos::remote_execution::FileNode, String>> =
    Vec::new();
  let mut dir_futures: Vec<BoxFuture<bazel_protos::remote_execution::DirectoryNode, String>> =
    Vec::new();

  path_stats.sort_by(|a, b| a.path().cmp(b.path()));

  for (first_component, group) in
    &path_stats.into_iter().group_by(|s| {
      s.path().components().next().unwrap().as_os_str().to_owned()
    })
  {
    let mut path_group: Vec<fs::PathStat> = group.collect();
    if path_group.len() == 1 && path_group.get(0).unwrap().path().components().count() == 1 {
      // Exactly one entry with exactly one component indicates either a file in this directory, or
      // an empty directory.
      // If the child is a non-empty directory, or a file therein, there must be multiple PathStats
      // with that prefix component, and we will handle that in the recursive save_directory call.

      match path_group.pop().unwrap() {
        fs::PathStat::File { ref stat, .. } => {
          let is_executable = stat.is_executable;
          file_futures.push(
            save_file(store.clone(), posix_fs.clone(), stat.clone())
              .and_then(move |(fingerprint, size_bytes)| {
                let mut file_node = bazel_protos::remote_execution::FileNode::new();
                file_node.set_name(osstring_as_utf8(first_component)?);
                file_node.set_digest(fingerprint_to_digest(&fingerprint, size_bytes as i64));
                file_node.set_is_executable(is_executable);
                Ok(file_node)
              })
              .to_boxed(),
          );
        }
        fs::PathStat::Dir { .. } => {
          // Because there are no children of this Dir, it must be empty.
          let (fingerprint, size_bytes) = store
            .record_directory(&bazel_protos::remote_execution::Directory::new())
            .unwrap();
          let mut directory_node = bazel_protos::remote_execution::DirectoryNode::new();
          directory_node.set_name(osstring_as_utf8(first_component).unwrap());
          directory_node.set_digest(fingerprint_to_digest(&fingerprint, size_bytes as i64));
          dir_futures.push(futures::future::ok(directory_node).to_boxed());
        }
      }
    } else {
      dir_futures.push(
        save_directory(
          store.clone(),
          posix_fs.clone(),
          paths_of_child_dir(path_group),
        ).and_then(move |(fingerprint, size_bytes)| {
          let mut dir_node = bazel_protos::remote_execution::DirectoryNode::new();
          dir_node.set_name(osstring_as_utf8(first_component)?);
          dir_node.set_digest(fingerprint_to_digest(&fingerprint, size_bytes as i64));
          Ok(dir_node)
        })
          .to_boxed(),
      );
    }
  }
  join_all(dir_futures)
    .join(join_all(file_futures))
    .and_then(move |(dirs, files)| {
      let mut directory = bazel_protos::remote_execution::Directory::new();
      directory.set_directories(protobuf::RepeatedField::from_vec(dirs));
      directory.set_files(protobuf::RepeatedField::from_vec(files));
      store.record_directory(&directory)
    })
    .to_boxed()
}

fn paths_of_child_dir(paths: Vec<fs::PathStat>) -> Vec<fs::PathStat> {
  paths
    .into_iter()
    .filter_map(|s| {
      if s.path().components().count() == 1 {
        return None;
      }
      Some(match s {
        fs::PathStat::File { path, stat } => {
          fs::PathStat::File {
            path: path.iter().skip(1).collect(),
            stat: stat,
          }
        }
        fs::PathStat::Dir { path, stat } => {
          fs::PathStat::Dir {
            path: path.iter().skip(1).collect(),
            stat: stat,
          }
        }
      })
    })
    .collect()
}

fn materialize_directory(
  store: Arc<Store>,
  destination: &Path,
  fingerprint: &Fingerprint,
) -> Result<(), ExitError> {
  let directory = store.load_directory_proto(&fingerprint)?.ok_or_else(|| {
    ExitError(
      format!("Directory with fingerprint {} not found", fingerprint),
      ExitCode::NotFound,
    )
  })?;
  make_clean_dir(&destination).map_err(|e| {
    format!(
      "Error making directory {:?}: {}",
      destination,
      e.description()
    )
  })?;
  for file_node in directory.get_files() {
    let fingerprint = &Fingerprint::from_hex_string(&file_node.get_digest().get_hash())?;
    match store.load_file_bytes(fingerprint)? {
      Some(bytes) => {
        let path = destination.join(file_node.get_name());
        File::create(&path)
          .and_then(|mut f| f.write_all(&bytes))
          .map_err(|e| {
            format!("Error writing file {:?}: {}", path, e.description())
          })?;
      }
      None => {
        return Err(ExitError(
          format!(
            "File with fingerprint {} not found",
            file_node.get_digest().get_hash()
          ),
          ExitCode::NotFound,
        ))
      }
    }
  }
  for directory_node in directory.get_directories() {
    materialize_directory(
      store.clone(),
      &destination.join(directory_node.get_name()),
      &Fingerprint::from_hex_string(directory_node.get_digest().get_hash())?,
    )?;
  }
  Ok(())
}

pub fn fingerprint_to_digest(
  fingerprint: &Fingerprint,
  size_bytes: i64,
) -> bazel_protos::remote_execution::Digest {
  let mut digest = bazel_protos::remote_execution::Digest::new();
  digest.set_hash(fingerprint.to_hex());
  digest.set_size_bytes(size_bytes);
  digest
}

fn osstring_as_utf8(path: OsString) -> Result<String, String> {
  path.into_string().map_err(|p| {
    format!("{:?}'s file_name is not representable in UTF8", p)
  })
}

fn make_clean_dir(path: &Path) -> Result<(), io::Error> {
  let parent = path.parent().ok_or_else(|| {
    io::Error::new(io::ErrorKind::NotFound, format!("{:?} had no parent", path))
  })?;
  std::fs::create_dir_all(parent)?;
  std::fs::create_dir(path)
}
<|endoftext|>


<commit_before>extern crate react;

use react::*;

#[test]
fn input_cells_have_a_value() {
    let mut reactor = Reactor::new();
    let input = reactor.create_input(10);
    assert_eq!(reactor.value(CellID::Input(input)), Some(10));
}

#[test]
#[ignore]
fn an_input_cells_value_can_be_set() {
    let mut reactor = Reactor::new();
    let input = reactor.create_input(4);
    assert!(reactor.set_value(input, 20));
    assert_eq!(reactor.value(CellID::Input(input)), Some(20));
}

#[test]
#[ignore]
fn error_setting_a_nonexistent_input_cell() {
    let mut dummy_reactor = Reactor::new();
    let input = dummy_reactor.create_input(1);
    assert!(!Reactor::new().set_value(input, 0));
}

#[test]
#[ignore]
fn compute_cells_calculate_initial_value() {
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let output = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    assert_eq!(reactor.value(CellID::Compute(output)), Some(2));
}

#[test]
#[ignore]
fn compute_cells_take_inputs_in_the_right_order() {
    let mut reactor = Reactor::new();
    let one = reactor.create_input(1);
    let two = reactor.create_input(2);
    let output = reactor
        .create_compute(&[CellID::Input(one), CellID::Input(two)], |v| {
            v[0] + v[1] * 10
        })
        .unwrap();
    assert_eq!(reactor.value(CellID::Compute(output)), Some(21));
}

#[test]
#[ignore]
fn error_creating_compute_cell_if_input_doesnt_exist() {
    let mut dummy_reactor = Reactor::new();
    let input = dummy_reactor.create_input(1);
    assert_eq!(
        Reactor::new().create_compute(&[CellID::Input(input)], |_| 0),
        Err(CellID::Input(input))
    );
}

#[test]
#[ignore]
fn do_not_break_cell_if_creating_compute_cell_with_valid_and_invalid_input() {
    let mut dummy_reactor = Reactor::new();
    let _ = dummy_reactor.create_input(1);
    let dummy_cell = dummy_reactor.create_input(2);
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    assert_eq!(
        reactor.create_compute(&[CellID::Input(input), CellID::Input(dummy_cell)], |_| 0),
        Err(CellID::Input(dummy_cell))
    );
    assert!(reactor.set_value(input, 5));
    assert_eq!(reactor.value(CellID::Input(input)), Some(5));
}

#[test]
#[ignore]
fn compute_cells_update_value_when_dependencies_are_changed() {
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let output = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    assert_eq!(reactor.value(CellID::Compute(output)), Some(2));
    assert!(reactor.set_value(input, 3));
    assert_eq!(reactor.value(CellID::Compute(output)), Some(4));
}

#[test]
#[ignore]
fn compute_cells_can_depend_on_other_compute_cells() {
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let times_two = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] * 2)
        .unwrap();
    let times_thirty = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] * 30)
        .unwrap();
    let output = reactor
        .create_compute(
            &[CellID::Compute(times_two), CellID::Compute(times_thirty)],
            |v| v[0] + v[1],
        )
        .unwrap();
    assert_eq!(reactor.value(CellID::Compute(output)), Some(32));
    assert!(reactor.set_value(input, 3));
    assert_eq!(reactor.value(CellID::Compute(output)), Some(96));
}

/// A CallbackRecorder helps tests whether callbacks get called correctly.
/// You'll see it used in tests that deal with callbacks.
/// The names should be descriptive enough so that the tests make sense,
/// so it's not necessary to fully understand the implementation,
/// though you are welcome to.
struct CallbackRecorder {
    // Note that this `Cell` is https://doc.rust-lang.org/std/cell/
    // a mechanism to allow internal mutability,
    // distinct from the cells (input cells, compute cells) in the reactor
    value: std::cell::Cell<Option<isize>>,
}

impl CallbackRecorder {
    fn new() -> Self {
        CallbackRecorder {
            value: std::cell::Cell::new(None),
        }
    }

    fn expect_to_have_been_called_with(&self, v: isize) {
        assert_ne!(
            self.value.get(),
            None,
            "Callback was not called, but should have been"
        );
        assert_eq!(
            self.value.replace(None),
            Some(v),
            "Callback was called with incorrect value"
        );
    }

    fn expect_not_to_have_been_called(&self) {
        assert_eq!(
            self.value.get(),
            None,
            "Callback was called, but should not have been"
        );
    }

    fn callback_called(&self, v: isize) {
        assert_eq!(
            self.value.replace(Some(v)),
            None,
            "Callback was called too many times; can't be called with {}",
            v
        );
    }
}

#[test]
#[ignore]
fn compute_cells_fire_callbacks() {
    let cb = CallbackRecorder::new();
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let output = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    assert!(
        reactor
            .add_callback(output, |v| cb.callback_called(v))
            .is_some()
    );
    assert!(reactor.set_value(input, 3));
    cb.expect_to_have_been_called_with(4);
}

#[test]
#[ignore]
fn error_adding_callback_to_nonexistent_cell() {
    let mut dummy_reactor = Reactor::new();
    let input = dummy_reactor.create_input(1);
    let output = dummy_reactor
        .create_compute(&[CellID::Input(input)], |_| 0)
        .unwrap();
    assert_eq!(
        Reactor::new().add_callback(output, |_: usize| println!("hi")),
        None
    );
}

#[test]
#[ignore]
fn callbacks_only_fire_on_change() {
    let cb = CallbackRecorder::new();
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let output = reactor
        .create_compute(
            &[CellID::Input(input)],
            |v| if v[0] < 3 { 111 } else { 222 },
        )
        .unwrap();
    assert!(
        reactor
            .add_callback(output, |v| cb.callback_called(v))
            .is_some()
    );

    assert!(reactor.set_value(input, 2));
    cb.expect_not_to_have_been_called();
    assert!(reactor.set_value(input, 4));
    cb.expect_to_have_been_called_with(222);
}

#[test]
#[ignore]
fn callbacks_can_be_added_and_removed() {
    let cb1 = CallbackRecorder::new();
    let cb2 = CallbackRecorder::new();
    let cb3 = CallbackRecorder::new();

    let mut reactor = Reactor::new();
    let input = reactor.create_input(11);
    let output = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();

    let callback = reactor
        .add_callback(output, |v| cb1.callback_called(v))
        .unwrap();
    assert!(
        reactor
            .add_callback(output, |v| cb2.callback_called(v))
            .is_some()
    );

    assert!(reactor.set_value(input, 31));
    cb1.expect_to_have_been_called_with(32);
    cb2.expect_to_have_been_called_with(32);

    assert!(reactor.remove_callback(output, callback).is_ok());
    assert!(
        reactor
            .add_callback(output, |v| cb3.callback_called(v))
            .is_some()
    );

    assert!(reactor.set_value(input, 41));
    cb1.expect_not_to_have_been_called();
    cb2.expect_to_have_been_called_with(42);
    cb3.expect_to_have_been_called_with(42);
}

#[test]
#[ignore]
fn removing_a_callback_multiple_times_doesnt_interfere_with_other_callbacks() {
    let cb1 = CallbackRecorder::new();
    let cb2 = CallbackRecorder::new();

    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let output = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    let callback = reactor
        .add_callback(output, |v| cb1.callback_called(v))
        .unwrap();
    assert!(
        reactor
            .add_callback(output, |v| cb2.callback_called(v))
            .is_some()
    );
    // We want the first remove to be Ok, but the others should be errors.
    assert!(reactor.remove_callback(output, callback).is_ok());
    for _ in 1..5 {
        assert_eq!(
            reactor.remove_callback(output, callback),
            Err(RemoveCallbackError::NonexistentCallback)
        );
    }

    assert!(reactor.set_value(input, 2));
    cb1.expect_not_to_have_been_called();
    cb2.expect_to_have_been_called_with(3);
}

#[test]
#[ignore]
fn callbacks_should_only_be_called_once_even_if_multiple_dependencies_change() {
    let cb = CallbackRecorder::new();
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let plus_one = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    let minus_one1 = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] - 1)
        .unwrap();
    let minus_one2 = reactor
        .create_compute(&[CellID::Compute(minus_one1)], |v| v[0] - 1)
        .unwrap();
    let output = reactor
        .create_compute(
            &[CellID::Compute(plus_one), CellID::Compute(minus_one2)],
            |v| v[0] * v[1],
        )
        .unwrap();
    assert!(
        reactor
            .add_callback(output, |v| cb.callback_called(v))
            .is_some()
    );
    assert!(reactor.set_value(input, 4));
    cb.expect_to_have_been_called_with(10);
}

#[test]
#[ignore]
fn callbacks_should_not_be_called_if_dependencies_change_but_output_value_doesnt_change() {
    let cb = CallbackRecorder::new();
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let plus_one = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    let minus_one = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] - 1)
        .unwrap();
    let always_two = reactor
        .create_compute(
            &[CellID::Compute(plus_one), CellID::Compute(minus_one)],
            |v| v[0] - v[1],
        )
        .unwrap();
    assert!(
        reactor
            .add_callback(always_two, |v| cb.callback_called(v))
            .is_some()
    );
    for i in 2..5 {
        assert!(reactor.set_value(input, i));
        cb.expect_not_to_have_been_called();
    }
}

#[test]
#[ignore]
fn test_adder_with_boolean_values() {
    // This is a digital logic circuit called an adder:
    // https://en.wikipedia.org/wiki/Adder_(electronics)
    let mut reactor = Reactor::new();
    let a = reactor.create_input(false);
    let b = reactor.create_input(false);
    let carry_in = reactor.create_input(false);

    let a_xor_b = reactor
        .create_compute(&[CellID::Input(a), CellID::Input(b)], |v| v[0] ^ v[1])
        .unwrap();
    let sum = reactor
        .create_compute(&[CellID::Compute(a_xor_b), CellID::Input(carry_in)], |v| {
            v[0] ^ v[1]
        })
        .unwrap();

    let a_xor_b_and_cin = reactor
        .create_compute(&[CellID::Compute(a_xor_b), CellID::Input(carry_in)], |v| {
            v[0] && v[1]
        })
        .unwrap();
    let a_and_b = reactor
        .create_compute(&[CellID::Input(a), CellID::Input(b)], |v| v[0] && v[1])
        .unwrap();
    let carry_out = reactor
        .create_compute(
            &[CellID::Compute(a_xor_b_and_cin), CellID::Compute(a_and_b)],
            |v| v[0] || v[1],
        )
        .unwrap();

    let tests = &[
        (false, false, false, false, false),
        (false, false, true, false, true),
        (false, true, false, false, true),
        (false, true, true, true, false),
        (true, false, false, false, true),
        (true, false, true, true, false),
        (true, true, false, true, false),
        (true, true, true, true, true),
    ];

    for &(aval, bval, cinval, expected_cout, expected_sum) in tests {
        assert!(reactor.set_value(a, aval));
        assert!(reactor.set_value(b, bval));
        assert!(reactor.set_value(carry_in, cinval));

        assert_eq!(reactor.value(CellID::Compute(sum)), Some(expected_sum));
        assert_eq!(
            reactor.value(CellID::Compute(carry_out)),
            Some(expected_cout)
        );
    }
}
<commit_msg>react: fix inappropriate use of [iu]size<commit_after>extern crate react;

use react::*;

#[test]
fn input_cells_have_a_value() {
    let mut reactor = Reactor::new();
    let input = reactor.create_input(10);
    assert_eq!(reactor.value(CellID::Input(input)), Some(10));
}

#[test]
#[ignore]
fn an_input_cells_value_can_be_set() {
    let mut reactor = Reactor::new();
    let input = reactor.create_input(4);
    assert!(reactor.set_value(input, 20));
    assert_eq!(reactor.value(CellID::Input(input)), Some(20));
}

#[test]
#[ignore]
fn error_setting_a_nonexistent_input_cell() {
    let mut dummy_reactor = Reactor::new();
    let input = dummy_reactor.create_input(1);
    assert!(!Reactor::new().set_value(input, 0));
}

#[test]
#[ignore]
fn compute_cells_calculate_initial_value() {
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let output = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    assert_eq!(reactor.value(CellID::Compute(output)), Some(2));
}

#[test]
#[ignore]
fn compute_cells_take_inputs_in_the_right_order() {
    let mut reactor = Reactor::new();
    let one = reactor.create_input(1);
    let two = reactor.create_input(2);
    let output = reactor
        .create_compute(&[CellID::Input(one), CellID::Input(two)], |v| {
            v[0] + v[1] * 10
        })
        .unwrap();
    assert_eq!(reactor.value(CellID::Compute(output)), Some(21));
}

#[test]
#[ignore]
fn error_creating_compute_cell_if_input_doesnt_exist() {
    let mut dummy_reactor = Reactor::new();
    let input = dummy_reactor.create_input(1);
    assert_eq!(
        Reactor::new().create_compute(&[CellID::Input(input)], |_| 0),
        Err(CellID::Input(input))
    );
}

#[test]
#[ignore]
fn do_not_break_cell_if_creating_compute_cell_with_valid_and_invalid_input() {
    let mut dummy_reactor = Reactor::new();
    let _ = dummy_reactor.create_input(1);
    let dummy_cell = dummy_reactor.create_input(2);
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    assert_eq!(
        reactor.create_compute(&[CellID::Input(input), CellID::Input(dummy_cell)], |_| 0),
        Err(CellID::Input(dummy_cell))
    );
    assert!(reactor.set_value(input, 5));
    assert_eq!(reactor.value(CellID::Input(input)), Some(5));
}

#[test]
#[ignore]
fn compute_cells_update_value_when_dependencies_are_changed() {
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let output = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    assert_eq!(reactor.value(CellID::Compute(output)), Some(2));
    assert!(reactor.set_value(input, 3));
    assert_eq!(reactor.value(CellID::Compute(output)), Some(4));
}

#[test]
#[ignore]
fn compute_cells_can_depend_on_other_compute_cells() {
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let times_two = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] * 2)
        .unwrap();
    let times_thirty = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] * 30)
        .unwrap();
    let output = reactor
        .create_compute(
            &[CellID::Compute(times_two), CellID::Compute(times_thirty)],
            |v| v[0] + v[1],
        )
        .unwrap();
    assert_eq!(reactor.value(CellID::Compute(output)), Some(32));
    assert!(reactor.set_value(input, 3));
    assert_eq!(reactor.value(CellID::Compute(output)), Some(96));
}

/// A CallbackRecorder helps tests whether callbacks get called correctly.
/// You'll see it used in tests that deal with callbacks.
/// The names should be descriptive enough so that the tests make sense,
/// so it's not necessary to fully understand the implementation,
/// though you are welcome to.
struct CallbackRecorder {
    // Note that this `Cell` is https://doc.rust-lang.org/std/cell/
    // a mechanism to allow internal mutability,
    // distinct from the cells (input cells, compute cells) in the reactor
    value: std::cell::Cell<Option<i32>>,
}

impl CallbackRecorder {
    fn new() -> Self {
        CallbackRecorder {
            value: std::cell::Cell::new(None),
        }
    }

    fn expect_to_have_been_called_with(&self, v: i32) {
        assert_ne!(
            self.value.get(),
            None,
            "Callback was not called, but should have been"
        );
        assert_eq!(
            self.value.replace(None),
            Some(v),
            "Callback was called with incorrect value"
        );
    }

    fn expect_not_to_have_been_called(&self) {
        assert_eq!(
            self.value.get(),
            None,
            "Callback was called, but should not have been"
        );
    }

    fn callback_called(&self, v: i32) {
        assert_eq!(
            self.value.replace(Some(v)),
            None,
            "Callback was called too many times; can't be called with {}",
            v
        );
    }
}

#[test]
#[ignore]
fn compute_cells_fire_callbacks() {
    let cb = CallbackRecorder::new();
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let output = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    assert!(
        reactor
            .add_callback(output, |v| cb.callback_called(v))
            .is_some()
    );
    assert!(reactor.set_value(input, 3));
    cb.expect_to_have_been_called_with(4);
}

#[test]
#[ignore]
fn error_adding_callback_to_nonexistent_cell() {
    let mut dummy_reactor = Reactor::new();
    let input = dummy_reactor.create_input(1);
    let output = dummy_reactor
        .create_compute(&[CellID::Input(input)], |_| 0)
        .unwrap();
    assert_eq!(
        Reactor::new().add_callback(output, |_: u32| println!("hi")),
        None
    );
}

#[test]
#[ignore]
fn callbacks_only_fire_on_change() {
    let cb = CallbackRecorder::new();
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let output = reactor
        .create_compute(
            &[CellID::Input(input)],
            |v| if v[0] < 3 { 111 } else { 222 },
        )
        .unwrap();
    assert!(
        reactor
            .add_callback(output, |v| cb.callback_called(v))
            .is_some()
    );

    assert!(reactor.set_value(input, 2));
    cb.expect_not_to_have_been_called();
    assert!(reactor.set_value(input, 4));
    cb.expect_to_have_been_called_with(222);
}

#[test]
#[ignore]
fn callbacks_can_be_added_and_removed() {
    let cb1 = CallbackRecorder::new();
    let cb2 = CallbackRecorder::new();
    let cb3 = CallbackRecorder::new();

    let mut reactor = Reactor::new();
    let input = reactor.create_input(11);
    let output = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();

    let callback = reactor
        .add_callback(output, |v| cb1.callback_called(v))
        .unwrap();
    assert!(
        reactor
            .add_callback(output, |v| cb2.callback_called(v))
            .is_some()
    );

    assert!(reactor.set_value(input, 31));
    cb1.expect_to_have_been_called_with(32);
    cb2.expect_to_have_been_called_with(32);

    assert!(reactor.remove_callback(output, callback).is_ok());
    assert!(
        reactor
            .add_callback(output, |v| cb3.callback_called(v))
            .is_some()
    );

    assert!(reactor.set_value(input, 41));
    cb1.expect_not_to_have_been_called();
    cb2.expect_to_have_been_called_with(42);
    cb3.expect_to_have_been_called_with(42);
}

#[test]
#[ignore]
fn removing_a_callback_multiple_times_doesnt_interfere_with_other_callbacks() {
    let cb1 = CallbackRecorder::new();
    let cb2 = CallbackRecorder::new();

    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let output = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    let callback = reactor
        .add_callback(output, |v| cb1.callback_called(v))
        .unwrap();
    assert!(
        reactor
            .add_callback(output, |v| cb2.callback_called(v))
            .is_some()
    );
    // We want the first remove to be Ok, but the others should be errors.
    assert!(reactor.remove_callback(output, callback).is_ok());
    for _ in 1..5 {
        assert_eq!(
            reactor.remove_callback(output, callback),
            Err(RemoveCallbackError::NonexistentCallback)
        );
    }

    assert!(reactor.set_value(input, 2));
    cb1.expect_not_to_have_been_called();
    cb2.expect_to_have_been_called_with(3);
}

#[test]
#[ignore]
fn callbacks_should_only_be_called_once_even_if_multiple_dependencies_change() {
    let cb = CallbackRecorder::new();
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let plus_one = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    let minus_one1 = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] - 1)
        .unwrap();
    let minus_one2 = reactor
        .create_compute(&[CellID::Compute(minus_one1)], |v| v[0] - 1)
        .unwrap();
    let output = reactor
        .create_compute(
            &[CellID::Compute(plus_one), CellID::Compute(minus_one2)],
            |v| v[0] * v[1],
        )
        .unwrap();
    assert!(
        reactor
            .add_callback(output, |v| cb.callback_called(v))
            .is_some()
    );
    assert!(reactor.set_value(input, 4));
    cb.expect_to_have_been_called_with(10);
}

#[test]
#[ignore]
fn callbacks_should_not_be_called_if_dependencies_change_but_output_value_doesnt_change() {
    let cb = CallbackRecorder::new();
    let mut reactor = Reactor::new();
    let input = reactor.create_input(1);
    let plus_one = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] + 1)
        .unwrap();
    let minus_one = reactor
        .create_compute(&[CellID::Input(input)], |v| v[0] - 1)
        .unwrap();
    let always_two = reactor
        .create_compute(
            &[CellID::Compute(plus_one), CellID::Compute(minus_one)],
            |v| v[0] - v[1],
        )
        .unwrap();
    assert!(
        reactor
            .add_callback(always_two, |v| cb.callback_called(v))
            .is_some()
    );
    for i in 2..5 {
        assert!(reactor.set_value(input, i));
        cb.expect_not_to_have_been_called();
    }
}

#[test]
#[ignore]
fn test_adder_with_boolean_values() {
    // This is a digital logic circuit called an adder:
    // https://en.wikipedia.org/wiki/Adder_(electronics)
    let mut reactor = Reactor::new();
    let a = reactor.create_input(false);
    let b = reactor.create_input(false);
    let carry_in = reactor.create_input(false);

    let a_xor_b = reactor
        .create_compute(&[CellID::Input(a), CellID::Input(b)], |v| v[0] ^ v[1])
        .unwrap();
    let sum = reactor
        .create_compute(&[CellID::Compute(a_xor_b), CellID::Input(carry_in)], |v| {
            v[0] ^ v[1]
        })
        .unwrap();

    let a_xor_b_and_cin = reactor
        .create_compute(&[CellID::Compute(a_xor_b), CellID::Input(carry_in)], |v| {
            v[0] && v[1]
        })
        .unwrap();
    let a_and_b = reactor
        .create_compute(&[CellID::Input(a), CellID::Input(b)], |v| v[0] && v[1])
        .unwrap();
    let carry_out = reactor
        .create_compute(
            &[CellID::Compute(a_xor_b_and_cin), CellID::Compute(a_and_b)],
            |v| v[0] || v[1],
        )
        .unwrap();

    let tests = &[
        (false, false, false, false, false),
        (false, false, true, false, true),
        (false, true, false, false, true),
        (false, true, true, true, false),
        (true, false, false, false, true),
        (true, false, true, true, false),
        (true, true, false, true, false),
        (true, true, true, true, true),
    ];

    for &(aval, bval, cinval, expected_cout, expected_sum) in tests {
        assert!(reactor.set_value(a, aval));
        assert!(reactor.set_value(b, bval));
        assert!(reactor.set_value(carry_in, cinval));

        assert_eq!(reactor.value(CellID::Compute(sum)), Some(expected_sum));
        assert_eq!(
            reactor.value(CellID::Compute(carry_out)),
            Some(expected_cout)
        );
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::net::SocketAddr;

use anyhow::Result;
use crossbeam_channel::Receiver;
use laminar::SocketEvent;
use log::{error, info};

use zombenum_shared::components::{GameEntityId, PlayerId};
use zombenum_shared::messages::client::*;

/// This is a helper struct for managing events from multiple users.
/// All events are sorted by their variant and combined with the PlayerId of the user that
/// triggered that event.
///
/// This should allow us to efficiently iterate and look up events by type in our systems.
#[derive(Default)]
pub struct PlayerEvents {
    pub join_game: Vec<PlayerId>,
    pub leave_game: Vec<PlayerId>,
    pub player_move: Vec<(PlayerId, PlayerMoveInputEvent)>,
    pub player_shoot: Vec<(PlayerId, PlayerShootInputEvent)>,
    pub player_use: Vec<(PlayerId, GameEntityId)>,
}

/// Get all received events and sort them in a PlayerEvents struct.
///
/// This function also handles meta events, such as socket connects, disconnects and timeouts.
pub fn get_user_events(
    receiver: &Receiver<SocketEvent>,
    connected_users: &mut HashMap<SocketAddr, PlayerId>,
) -> Result<PlayerEvents> {
    let mut events = PlayerEvents::default();
    while let Ok(event) = receiver.try_recv() {
        match event {
            SocketEvent::Packet(packet) => {
                // Get some meta information and do some preprocessing
                let client_addr = packet.addr();
                let client_ip = client_addr.ip();
                let payload = packet.payload();
                let message: ClientMessage = bincode::deserialize(payload)?;

                info!("Received {:#?} from {:?}", message, client_ip);
                match message {
                    ClientMessage::JoinGame(player_id) => {
                        // A user wants to join the game. Register the user.
                        connected_users.insert(packet.addr(), player_id.clone());
                        events.join_game.push(player_id)
                    }
                    ClientMessage::LeaveGame => {
                        // A user wants to leave the game.
                        // Display an error if we don't know this user.
                        if let Some(player_id) = connected_users.get(&packet.addr()) {
                            events.leave_game.push(player_id.clone());
                        } else {
                            connected_users.remove(&packet.addr());
                            error!(
                                "Unknown user with addr {:?} just disconnected",
                                packet.addr()
                            );
                        }
                    }

                    ClientMessage::PlayerInputEvent(player_input_event) => {
                        // Get the user id from the socket address.
                        // Show an error and ignore the packet, if we cannot find an associated
                        // user.
                        let player_id = if let Some(player_id) = connected_users.get(&packet.addr())
                        {
                            player_id.clone()
                        } else {
                            error!("Unknown user sent event with addr {:?}", packet.addr());
                            continue;
                        };

                        get_user_input_event(&mut events, player_input_event.event_type, player_id);
                    }
                }
            }

            SocketEvent::Connect(address) => {
                info!("Client connected: {}", address);
            }
            SocketEvent::Disconnect(address) => {
                connected_users.remove(&address);
                info!("Client disconnected: {}", address);
            }
            SocketEvent::Timeout(address) => {
                connected_users.remove(&address);
                info!("Client timed out: {}", address);
            }
        }
    }

    Ok(events)
}

/// Add a given PlayerInputEvent to the list of PlayerEvents
fn get_user_input_event(
    events: &mut PlayerEvents,
    event: PlayerInputEventType,
    player_id: PlayerId,
) {
    match event {
        PlayerInputEventType::Move(event) => {
            events.player_move.push((player_id, event));
        }
        PlayerInputEventType::Shoot(event) => {
            events.player_shoot.push((player_id, event));
        }
        PlayerInputEventType::Use(entity) => {
            events.player_use.push((player_id, entity));
        }
        _ => error!("Unhandled event {:?}", event),
    }
}
<commit_msg>Better server-side event handling<commit_after>use std::collections::HashMap;
use std::net::SocketAddr;

use anyhow::Result;
use crossbeam_channel::Receiver;
use laminar::SocketEvent;
use log::{error, info};

use zombenum_shared::components::{GameEntityId, PlayerId};
use zombenum_shared::messages::client::*;

/// This is a helper struct for managing events from multiple users.
/// All events are sorted by their variant and combined with the PlayerId of the user that
/// triggered that event.
///
/// This should allow us to efficiently iterate and look up events by type in our systems.
#[derive(Default)]
pub struct PlayerEvents {
    pub join_game: Vec<PlayerId>,
    pub leave_game: Vec<PlayerId>,
    pub player_move: Vec<(PlayerId, PlayerMoveInputEvent)>,
    pub player_shoot: Vec<(PlayerId, PlayerShootInputEvent)>,
    pub player_use: Vec<(PlayerId, GameEntityId)>,
}

/// Get all received events and sort them in a PlayerEvents struct.
///
/// This function also handles meta events, such as socket connects, disconnects and timeouts.
pub fn get_user_events(
    receiver: &Receiver<SocketEvent>,
    connected_users: &mut HashMap<SocketAddr, PlayerId>,
) -> Result<PlayerEvents> {
    let mut events = PlayerEvents::default();
    while let Ok(event) = receiver.try_recv() {
        match event {
            SocketEvent::Packet(packet) => {
                // Get some meta information and do some preprocessing
                let client_addr = packet.addr();
                let client_ip = client_addr.ip();
                let payload = packet.payload();
                let message: ClientMessage = bincode::deserialize(payload)?;

                info!("Received {:#?} from {:?}", message, client_ip);
                match message {
                    ClientMessage::JoinGame(player_id) => {
                        // A user wants to join the game. Register the user.
                        connected_users.insert(packet.addr(), player_id.clone());
                        events.join_game.push(player_id)
                    }
                    ClientMessage::LeaveGame => {
                        // A user wants to leave the game.
                        // Display an error if we don't know this user.
                        if let Some(player_id) = connected_users.get(&packet.addr()) {
                            events.leave_game.push(player_id.clone());
                        } else {
                            connected_users.remove(&packet.addr());
                            error!(
                                "Unknown user with addr {:?} just disconnected",
                                packet.addr()
                            );
                        }
                    }

                    ClientMessage::PlayerInputEvent(player_input_event) => {
                        // Get the user id from the socket address.
                        // Show an error and ignore the packet, if we cannot find an associated
                        // user.
                        let player_id = if let Some(player_id) = connected_users.get(&packet.addr())
                        {
                            player_id.clone()
                        } else {
                            error!("Unknown user sent event with addr {:?}", packet.addr());
                            continue;
                        };

                        get_user_input_event(&mut events, player_input_event, player_id);
                    }
                }
            }

            SocketEvent::Connect(address) => {
                info!("Client connected: {}", address);
            }
            SocketEvent::Disconnect(address) => {
                connected_users.remove(&address);
                info!("Client disconnected: {}", address);
            }
            SocketEvent::Timeout(address) => {
                connected_users.remove(&address);
                info!("Client timed out: {}", address);
            }
        }
    }

    Ok(events)
}

/// Add a given PlayerInputEvent to the list of PlayerEvents
fn get_user_input_event(events: &mut PlayerEvents, event: PlayerInputEvent, player_id: PlayerId) {
    match event.event_type {
        PlayerInputEventType::Move(event) => {
            events.player_move.push((player_id, event));
        }
        PlayerInputEventType::Shoot(event) => {
            events.player_shoot.push((player_id, event));
        }
        PlayerInputEventType::Use(entity) => {
            events.player_use.push((player_id, entity));
        }
        _ => error!("Unhandled event {:?}", event),
    }
}
<|endoftext|>


<commit_before>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use bytes::BytesMut;
use citaprotocol::pubsub_message_to_network_message;
use config;
use config::NetConfig;
use libproto::{Message, OperateType};
use native_tls::{self, TlsConnector};
use notify::DebouncedEvent;
use std::collections::HashMap;
use std::convert::TryInto;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::net::SocketAddr;
use std::net::{Shutdown, TcpStream};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use tokio;
use tokio::prelude::*;
use tokio::util::FutureExt;
use tokio_tls;

const TIMEOUT: u64 = 15;
const ROOT_CERT_FILE: &str = "rootCA.crt";

pub enum RealStream {
    CryptStream(tokio_tls::TlsStream<tokio::net::TcpStream>),
    NormalStream(TcpStream),
}

impl RealStream {
    pub fn shutdown(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.shutdown().map(|_| ()),
            RealStream::NormalStream(ref mut tcp) => tcp.shutdown(Shutdown::Both),
        }
    }
}

impl io::Write for RealStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.write(buf),
            RealStream::NormalStream(ref mut tcp) => tcp.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.flush(),
            RealStream::NormalStream(ref mut tcp) => tcp.flush(),
        }
    }
}

pub enum Task {
    Broadcast((String, Message)),
    Update(NetConfig),
    NewTCP((u32, SocketAddr, RealStream, String)),
}

/// Manager unconnected address
struct Manager {
    need_connect: Vec<(u32, SocketAddr, String)>,
    connect_receiver: Receiver<(u32, SocketAddr, String)>,
    task_sender: Sender<Task>,
    enable_tls: bool,
}

fn generate_tls_connector(path: &str) -> Option<TlsConnector> {
    let mut file = File::open(path).expect("Not has cert file");
    let mut pem = vec![];
    file.read_to_end(&mut pem).expect("Cert File read error");
    let root_ca =
        native_tls::Certificate::from_pem(&pem).expect("Cert file content not right, pem ?");

    native_tls::TlsConnector::builder()
        .min_protocol_version(Some(native_tls::Protocol::Tlsv11))
        .max_protocol_version(Some(native_tls::Protocol::Tlsv12))
        .add_root_certificate(root_ca)
        .build()
        .ok()
}

impl Manager {
    fn new(
        task_sender: Sender<Task>,
        connect_receiver: Receiver<(u32, SocketAddr, String)>,
        enable_tls: bool,
    ) -> Self {
        Manager {
            need_connect: Vec::new(),
            connect_receiver,
            task_sender,
            enable_tls,
        }
    }

    fn run(&mut self) {
        let (tls_connector, mut rt) = if self.enable_tls {
            let tls_connector = generate_tls_connector(ROOT_CERT_FILE);
            let rt = tokio::runtime::current_thread::Runtime::new().ok();
            if tls_connector.is_none() || rt.is_none() {
                panic!("TLS connector not generated");
            }
            (tls_connector, rt)
        } else {
            (None, None)
        };

        loop {
            while let Ok(message) = self.connect_receiver.try_recv() {
                self.need_connect.push(message);
            }
            let mut new_need_connect = Vec::new();
            while let Some((id, addr, common_name)) = self.need_connect.pop() {
                match tls_connector.clone() {
                    Some(tls_connect) => {
                        let common_name_clone = common_name.clone();
                        let task = tokio::net::TcpStream::connect(&addr)
                            .and_then(move |socket| {
                                tokio_tls::TlsConnector::from(tls_connect)
                                    .connect(&common_name_clone, socket)
                                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                            })
                            .timeout(Duration::from_secs(TIMEOUT));
                        match rt.as_mut().unwrap().block_on(task) {
                            Ok(tls) => {
                                self.task_sender
                                    .send(Task::NewTCP((
                                        id,
                                        addr,
                                        RealStream::CryptStream(tls),
                                        common_name.clone(),
                                    )))
                                    .unwrap();
                            }
                            Err(e) => {
                                warn!("TLS connect {} failed, error: {}", addr, e);
                                new_need_connect.push((id, addr, common_name));
                            }
                        };
                    }
                    None => match TcpStream::connect_timeout(&addr, Duration::from_secs(TIMEOUT)) {
                        Ok(tcp) => {
                            self.task_sender
                                .send(Task::NewTCP((
                                    id,
                                    addr,
                                    RealStream::NormalStream(tcp),
                                    common_name.clone(),
                                )))
                                .unwrap();
                        }
                        Err(e) => {
                            warn!(
                                "Node{}, {} unable to establish connection, error: {}",
                                id, addr, e
                            );
                            new_need_connect.push((id, addr, common_name));
                        }
                    },
                }
            }
            self.need_connect = new_need_connect;

            if !self.need_connect.is_empty() {
                debug!(
                    "Complete a round of attempts to connect, \
                     left {} address for the next round of processing",
                    self.need_connect.len()
                );
            }

            thread::sleep(Duration::from_secs(TIMEOUT));
        }
    }
}

/// Manage p2p networks
pub struct Connections {
    id_card: u32,
    /// list of peer: id, addr, tcp_connect
    peers: HashMap<(u32, SocketAddr, String), RealStream>,
    pub is_pause: Arc<AtomicBool>,
    pub connect_number: Arc<AtomicUsize>,
    task_receiver: Receiver<Task>,
    connect_sender: Sender<(u32, SocketAddr, String)>,
}

impl Connections {
    pub fn new(config: &config::NetConfig) -> (Self, Sender<Task>) {
        let id_card = config.id_card.unwrap();
        let (task_sender, task_receiver) = channel();
        let (connect_sender, connect_receiver) = channel();

        let connect_task_sender = task_sender.clone();
        let enable_tls = config.enable_tls.unwrap_or(false);
        thread::spawn(move || {
            Manager::new(connect_task_sender, connect_receiver, enable_tls).run()
        });

        if let Some(peers) = config.peers.as_ref() {
            for peer in peers.iter() {
                let id_card: u32 = peer.id_card.unwrap();
                let addr = format!("{}:{}", peer.ip.clone().unwrap(), peer.port.unwrap())
                    .parse()
                    .unwrap();
                connect_sender
                    .send((id_card, addr, peer.common_name.clone().unwrap_or_default()))
                    .unwrap();
            }
        }

        (
            Connections {
                id_card,
                peers: HashMap::new(),
                is_pause: Arc::new(AtomicBool::new(false)),
                connect_number: Arc::new(AtomicUsize::new(0)),
                task_receiver,
                connect_sender,
            },
            task_sender,
        )
    }

    fn is_send(id_card: u32, origin: u32, operate: OperateType) -> bool {
        match operate {
            OperateType::Broadcast => true,
            OperateType::Single => id_card == origin,
            OperateType::Subtract => id_card != origin,
        }
    }

    pub fn run(&mut self) {
        loop {
            match self
                .task_receiver
                .recv_timeout(Duration::from_secs(TIMEOUT))
            {
                Ok(task) => match task {
                    Task::Broadcast((key, message)) => {
                        if !self.is_pause.load(Ordering::SeqCst) {
                            self.broadcast(key, message)
                        }
                    }
                    Task::NewTCP(tcp) => {
                        self.peers.insert((tcp.0, tcp.1, tcp.3), tcp.2);
                        self.connect_number.fetch_add(1, Ordering::Relaxed);
                    }
                    Task::Update(config) => self.update(config),
                },
                Err(_) => {
                    self.heart_beat();
                }
            }
        }
    }

    fn update(&mut self, config: config::NetConfig) {
        // Update configuration
        match config.peers {
            Some(peers) => {
                let config_peers = peers
                    .into_iter()
                    .map(|peer| {
                        let id_card: u32 = peer.id_card.unwrap();
                        let addr = format!("{}:{}", peer.ip.unwrap(), peer.port.unwrap())
                            .parse()
                            .unwrap();
                        (id_card, addr, peer.common_name.unwrap_or_default())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                let remove_peers = self
                    .peers
                    .keys()
                    .filter(|peer| !config_peers.contains(&peer))
                    .map(|ref peer| {
                        info!("Remove peer {}, {},{}", peer.0, peer.1, peer.2);
                        (peer.0, peer.1, peer.2.clone())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                config_peers
                    .into_iter()
                    .filter(|peer| {
                        self.peers
                            .keys()
                            .find(|&current_peer| current_peer == peer)
                            .is_none()
                    })
                    .for_each(|peer| {
                        info!("Add peer {}, {}", peer.0, peer.1);
                        let _ = self.connect_sender.send(peer);
                    });

                self.close(Some(remove_peers), false);
            }
            None => {
                info!("clear all peers after update!");
                self.close::<Vec<(u32, SocketAddr, String)>>(None, false);
            }
        }
    }

    fn broadcast(&mut self, key: String, mut msg: Message) {
        let origin = msg.get_origin();
        let operate = msg.get_operate();
        msg.set_origin(self.id_card);

        trace!("Broadcast msg {:?} from key {}", msg, key);
        let msg_bytes: Vec<u8> = msg.try_into().unwrap();

        let mut buf = BytesMut::with_capacity(4 + 4 + 1 + key.len() + msg_bytes.len());
        pubsub_message_to_network_message(&mut buf, Some((key, msg_bytes)));

        let mut peers = Vec::new();
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            if Connections::is_send(peer.0, origin, operate) {
                match stream.write_all(&buf) {
                    Ok(_) => {
                        peers.push(peer.0);
                    }
                    Err(e) => {
                        warn!("Node{} {} is shutdown, err: {}", peer.0, peer.1, e);
                        remove_peers.push(peer.clone());
                    }
                }
            }
        }
        self.close(Some(remove_peers), true);
        trace!(
            "{:?} broadcast msg to nodes {:?} {:?}",
            self.id_card,
            operate,
            peers
        );
    }

    fn close<T: ::std::iter::IntoIterator<Item = (u32, SocketAddr, String)>>(
        &mut self,
        peers: Option<T>,
        need_reconnect: bool,
    ) {
        match peers {
            Some(peers) => peers.into_iter().for_each(|peer| {
                if let Some(mut stream) = self.peers.remove(&peer) {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                    if need_reconnect {
                        self.connect_sender.send(peer).unwrap();
                    }
                    self.connect_number.fetch_sub(1, Ordering::Relaxed);
                }
            }),
            None => {
                self.peers.iter_mut().for_each(|(peer, stream)| {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                });
                if need_reconnect {
                    self.peers.iter().for_each(|(ref peer, _)| {
                        self.connect_sender
                            .send((peer.0, peer.1, peer.2.clone()))
                            .unwrap();
                    })
                }
                self.peers.clear();
                self.connect_number.store(0, Ordering::Relaxed);
            }
        }
    }

    fn heart_beat(&mut self) {
        let mut buf = BytesMut::with_capacity(4 + 4);
        pubsub_message_to_network_message(&mut buf, None);
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            match stream.write(&buf) {
                Ok(_) => {
                    let _ = stream.flush();
                }
                Err(e) => {
                    warn!(
                        "Node{} {} {} is shutdown, err: {}",
                        peer.0, peer.1, peer.2, e
                    );
                    remove_peers.push((peer.0, peer.1, peer.2.clone()));
                }
            };
        }
        self.close(Some(remove_peers), true);
    }
}

pub fn manage_connect(config_path: &str, rx: Receiver<DebouncedEvent>, task_send: Sender<Task>) {
    let config = String::from(config_path);

    thread::spawn(move || loop {
        match rx.recv() {
            Ok(event) => match event {
                DebouncedEvent::Create(path_buf) | DebouncedEvent::Write(path_buf) => {
                    if path_buf.is_file() {
                        let file_name = path_buf.file_name().unwrap().to_str().unwrap();
                        if file_name == config.as_str() {
                            info!("file {} changed, will auto reload!", file_name);
                            let config = NetConfig::new(config.as_str());
                            let _ = task_send.send(Task::Update(config));
                        }
                    }
                }
                _ => trace!("file notify event: {:?}", event),
            },
            Err(e) => warn!("watch error: {:?}", e),
        }
    });
}

#[cfg(test)]
mod test {
    use super::Connections;
    use libproto::OperateType;
    #[test]
    fn is_send_msg() {
        assert!(Connections::is_send(0, 0, OperateType::Broadcast));
        assert!(Connections::is_send(0, 1, OperateType::Broadcast));

        assert!(Connections::is_send(0, 0, OperateType::Single));
        assert!(!Connections::is_send(0, 1, OperateType::Single));

        assert!(!Connections::is_send(0, 0, OperateType::Subtract));
        assert!(Connections::is_send(0, 1, OperateType::Subtract));
    }
}
<commit_msg>support hostname<commit_after>// CITA
// Copyright 2016-2017 Cryptape Technologies LLC.

// This program is free software: you can redistribute it
// and/or modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation,
// either version 3 of the License, or (at your option) any
// later version.

// This program is distributed in the hope that it will be
// useful, but WITHOUT ANY WARRANTY; without even the implied
// warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
// PURPOSE. See the GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use bytes::BytesMut;
use citaprotocol::pubsub_message_to_network_message;
use config;
use config::NetConfig;
use libproto::{Message, OperateType};
use native_tls::{self, TlsConnector};
use notify::DebouncedEvent;
use std::collections::HashMap;
use std::convert::TryInto;
use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::net::{Shutdown, TcpStream};
use std::net::{SocketAddr, ToSocketAddrs};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::Arc;
use std::thread;
use std::time::Duration;
use tokio;
use tokio::prelude::*;
use tokio::util::FutureExt;
use tokio_tls;

const TIMEOUT: u64 = 15;
const ROOT_CERT_FILE: &str = "rootCA.crt";

pub enum RealStream {
    CryptStream(tokio_tls::TlsStream<tokio::net::TcpStream>),
    NormalStream(TcpStream),
}

impl RealStream {
    pub fn shutdown(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.shutdown().map(|_| ()),
            RealStream::NormalStream(ref mut tcp) => tcp.shutdown(Shutdown::Both),
        }
    }
}

impl io::Write for RealStream {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.write(buf),
            RealStream::NormalStream(ref mut tcp) => tcp.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self {
            RealStream::CryptStream(ref mut tls) => tls.flush(),
            RealStream::NormalStream(ref mut tcp) => tcp.flush(),
        }
    }
}

pub enum Task {
    Broadcast((String, Message)),
    Update(NetConfig),
    NewTCP((u32, SocketAddr, RealStream, String)),
}

/// Manager unconnected address
struct Manager {
    need_connect: Vec<(u32, SocketAddr, String)>,
    connect_receiver: Receiver<(u32, SocketAddr, String)>,
    task_sender: Sender<Task>,
    enable_tls: bool,
}

fn generate_tls_connector(path: &str) -> Option<TlsConnector> {
    let mut file = File::open(path).expect("Not has cert file");
    let mut pem = vec![];
    file.read_to_end(&mut pem).expect("Cert File read error");
    let root_ca =
        native_tls::Certificate::from_pem(&pem).expect("Cert file content not right, pem ?");

    native_tls::TlsConnector::builder()
        .min_protocol_version(Some(native_tls::Protocol::Tlsv11))
        .max_protocol_version(Some(native_tls::Protocol::Tlsv12))
        .add_root_certificate(root_ca)
        .build()
        .ok()
}

impl Manager {
    fn new(
        task_sender: Sender<Task>,
        connect_receiver: Receiver<(u32, SocketAddr, String)>,
        enable_tls: bool,
    ) -> Self {
        Manager {
            need_connect: Vec::new(),
            connect_receiver,
            task_sender,
            enable_tls,
        }
    }

    fn run(&mut self) {
        let (tls_connector, mut rt) = if self.enable_tls {
            let tls_connector = generate_tls_connector(ROOT_CERT_FILE);
            let rt = tokio::runtime::current_thread::Runtime::new().ok();
            if tls_connector.is_none() || rt.is_none() {
                panic!("TLS connector not generated");
            }
            (tls_connector, rt)
        } else {
            (None, None)
        };

        loop {
            while let Ok(message) = self.connect_receiver.try_recv() {
                self.need_connect.push(message);
            }
            let mut new_need_connect = Vec::new();
            while let Some((id, addr, common_name)) = self.need_connect.pop() {
                match tls_connector.clone() {
                    Some(tls_connect) => {
                        let common_name_clone = common_name.clone();
                        let task = tokio::net::TcpStream::connect(&addr)
                            .and_then(move |socket| {
                                tokio_tls::TlsConnector::from(tls_connect)
                                    .connect(&common_name_clone, socket)
                                    .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
                            })
                            .timeout(Duration::from_secs(TIMEOUT));
                        match rt.as_mut().unwrap().block_on(task) {
                            Ok(tls) => {
                                self.task_sender
                                    .send(Task::NewTCP((
                                        id,
                                        addr,
                                        RealStream::CryptStream(tls),
                                        common_name.clone(),
                                    )))
                                    .unwrap();
                            }
                            Err(e) => {
                                warn!("TLS connect {} failed, error: {}", addr, e);
                                new_need_connect.push((id, addr, common_name));
                            }
                        };
                    }
                    None => match TcpStream::connect_timeout(&addr, Duration::from_secs(TIMEOUT)) {
                        Ok(tcp) => {
                            self.task_sender
                                .send(Task::NewTCP((
                                    id,
                                    addr,
                                    RealStream::NormalStream(tcp),
                                    common_name.clone(),
                                )))
                                .unwrap();
                        }
                        Err(e) => {
                            warn!(
                                "Node{}, {} unable to establish connection, error: {}",
                                id, addr, e
                            );
                            new_need_connect.push((id, addr, common_name));
                        }
                    },
                }
            }
            self.need_connect = new_need_connect;

            if !self.need_connect.is_empty() {
                debug!(
                    "Complete a round of attempts to connect, \
                     left {} address for the next round of processing",
                    self.need_connect.len()
                );
            }

            thread::sleep(Duration::from_secs(TIMEOUT));
        }
    }
}

/// Manage p2p networks
pub struct Connections {
    id_card: u32,
    /// list of peer: id, addr, tcp_connect
    peers: HashMap<(u32, SocketAddr, String), RealStream>,
    pub is_pause: Arc<AtomicBool>,
    pub connect_number: Arc<AtomicUsize>,
    task_receiver: Receiver<Task>,
    connect_sender: Sender<(u32, SocketAddr, String)>,
}

impl Connections {
    pub fn new(config: &config::NetConfig) -> (Self, Sender<Task>) {
        let id_card = config.id_card.unwrap();
        let (task_sender, task_receiver) = channel();
        let (connect_sender, connect_receiver) = channel();

        let connect_task_sender = task_sender.clone();
        let enable_tls = config.enable_tls.unwrap_or(false);
        thread::spawn(move || {
            Manager::new(connect_task_sender, connect_receiver, enable_tls).run()
        });

        if let Some(peers) = config.peers.as_ref() {
            for peer in peers.iter() {
                let id_card: u32 = peer.id_card.unwrap();
                let addr = format!("{}:{}", peer.ip.clone().unwrap(), peer.port.unwrap())
                    .to_socket_addrs()
                    .unwrap()
                    .next()
                    .unwrap();
                connect_sender
                    .send((id_card, addr, peer.common_name.clone().unwrap_or_default()))
                    .unwrap();
            }
        }

        (
            Connections {
                id_card,
                peers: HashMap::new(),
                is_pause: Arc::new(AtomicBool::new(false)),
                connect_number: Arc::new(AtomicUsize::new(0)),
                task_receiver,
                connect_sender,
            },
            task_sender,
        )
    }

    fn is_send(id_card: u32, origin: u32, operate: OperateType) -> bool {
        match operate {
            OperateType::Broadcast => true,
            OperateType::Single => id_card == origin,
            OperateType::Subtract => id_card != origin,
        }
    }

    pub fn run(&mut self) {
        loop {
            match self
                .task_receiver
                .recv_timeout(Duration::from_secs(TIMEOUT))
            {
                Ok(task) => match task {
                    Task::Broadcast((key, message)) => {
                        if !self.is_pause.load(Ordering::SeqCst) {
                            self.broadcast(key, message)
                        }
                    }
                    Task::NewTCP(tcp) => {
                        self.peers.insert((tcp.0, tcp.1, tcp.3), tcp.2);
                        self.connect_number.fetch_add(1, Ordering::Relaxed);
                    }
                    Task::Update(config) => self.update(config),
                },
                Err(_) => {
                    self.heart_beat();
                }
            }
        }
    }

    fn update(&mut self, config: config::NetConfig) {
        // Update configuration
        match config.peers {
            Some(peers) => {
                let config_peers = peers
                    .into_iter()
                    .map(|peer| {
                        let id_card: u32 = peer.id_card.unwrap();
                        let addr = format!("{}:{}", peer.ip.unwrap(), peer.port.unwrap())
                            .to_socket_addrs()
                            .unwrap()
                            .next()
                            .unwrap();
                        (id_card, addr, peer.common_name.unwrap_or_default())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                let remove_peers = self
                    .peers
                    .keys()
                    .filter(|peer| !config_peers.contains(&peer))
                    .map(|ref peer| {
                        info!("Remove peer {}, {},{}", peer.0, peer.1, peer.2);
                        (peer.0, peer.1, peer.2.clone())
                    })
                    .collect::<Vec<(u32, SocketAddr, String)>>();

                config_peers
                    .into_iter()
                    .filter(|peer| {
                        self.peers
                            .keys()
                            .find(|&current_peer| current_peer == peer)
                            .is_none()
                    })
                    .for_each(|peer| {
                        info!("Add peer {}, {}", peer.0, peer.1);
                        let _ = self.connect_sender.send(peer);
                    });

                self.close(Some(remove_peers), false);
            }
            None => {
                info!("clear all peers after update!");
                self.close::<Vec<(u32, SocketAddr, String)>>(None, false);
            }
        }
    }

    fn broadcast(&mut self, key: String, mut msg: Message) {
        let origin = msg.get_origin();
        let operate = msg.get_operate();
        msg.set_origin(self.id_card);

        trace!("Broadcast msg {:?} from key {}", msg, key);
        let msg_bytes: Vec<u8> = msg.try_into().unwrap();

        let mut buf = BytesMut::with_capacity(4 + 4 + 1 + key.len() + msg_bytes.len());
        pubsub_message_to_network_message(&mut buf, Some((key, msg_bytes)));

        let mut peers = Vec::new();
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            if Connections::is_send(peer.0, origin, operate) {
                match stream.write_all(&buf) {
                    Ok(_) => {
                        peers.push(peer.0);
                    }
                    Err(e) => {
                        warn!("Node{} {} is shutdown, err: {}", peer.0, peer.1, e);
                        remove_peers.push(peer.clone());
                    }
                }
            }
        }
        self.close(Some(remove_peers), true);
        trace!(
            "{:?} broadcast msg to nodes {:?} {:?}",
            self.id_card,
            operate,
            peers
        );
    }

    fn close<T: ::std::iter::IntoIterator<Item = (u32, SocketAddr, String)>>(
        &mut self,
        peers: Option<T>,
        need_reconnect: bool,
    ) {
        match peers {
            Some(peers) => peers.into_iter().for_each(|peer| {
                if let Some(mut stream) = self.peers.remove(&peer) {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                    if need_reconnect {
                        self.connect_sender.send(peer).unwrap();
                    }
                    self.connect_number.fetch_sub(1, Ordering::Relaxed);
                }
            }),
            None => {
                self.peers.iter_mut().for_each(|(peer, stream)| {
                    let _ = stream.shutdown().map_err(|err| {
                        warn!("Shutdown {} - {} failed: {}", peer.0, peer.1, err);
                    });
                });
                if need_reconnect {
                    self.peers.iter().for_each(|(ref peer, _)| {
                        self.connect_sender
                            .send((peer.0, peer.1, peer.2.clone()))
                            .unwrap();
                    })
                }
                self.peers.clear();
                self.connect_number.store(0, Ordering::Relaxed);
            }
        }
    }

    fn heart_beat(&mut self) {
        let mut buf = BytesMut::with_capacity(4 + 4);
        pubsub_message_to_network_message(&mut buf, None);
        let mut remove_peers = Vec::new();
        for (peer, stream) in self.peers.iter_mut() {
            match stream.write(&buf) {
                Ok(_) => {
                    let _ = stream.flush();
                }
                Err(e) => {
                    warn!(
                        "Node{} {} {} is shutdown, err: {}",
                        peer.0, peer.1, peer.2, e
                    );
                    remove_peers.push((peer.0, peer.1, peer.2.clone()));
                }
            };
        }
        self.close(Some(remove_peers), true);
    }
}

pub fn manage_connect(config_path: &str, rx: Receiver<DebouncedEvent>, task_send: Sender<Task>) {
    let config = String::from(config_path);

    thread::spawn(move || loop {
        match rx.recv() {
            Ok(event) => match event {
                DebouncedEvent::Create(path_buf) | DebouncedEvent::Write(path_buf) => {
                    if path_buf.is_file() {
                        let file_name = path_buf.file_name().unwrap().to_str().unwrap();
                        if file_name == config.as_str() {
                            info!("file {} changed, will auto reload!", file_name);
                            let config = NetConfig::new(config.as_str());
                            let _ = task_send.send(Task::Update(config));
                        }
                    }
                }
                _ => trace!("file notify event: {:?}", event),
            },
            Err(e) => warn!("watch error: {:?}", e),
        }
    });
}

#[cfg(test)]
mod test {
    use super::Connections;
    use libproto::OperateType;
    #[test]
    fn is_send_msg() {
        assert!(Connections::is_send(0, 0, OperateType::Broadcast));
        assert!(Connections::is_send(0, 1, OperateType::Broadcast));

        assert!(Connections::is_send(0, 0, OperateType::Single));
        assert!(!Connections::is_send(0, 1, OperateType::Single));

        assert!(!Connections::is_send(0, 0, OperateType::Subtract));
        assert!(Connections::is_send(0, 1, OperateType::Subtract));
    }
}
<|endoftext|>