<commit_before>use crate::cli::Command;
use mprovision as mp;
use std::env;
use std::io::{self, Write};
use std::path::Path;
use std::path::PathBuf;
use std::time::{Duration, SystemTime};

mod cli;

fn main() {
    if let Err(e) = cli::parse(env::args()).and_then(run) {
        e.exit();
    }
}

fn run(command: cli::Command) -> Result<(), cli::Error> {
    match command {
        Command::List(cli::ListParams {
            text,
            expire_in_days,
            directory,
        }) => list(&text, expire_in_days, directory),
        Command::ShowUuid(cli::ShowUuidParams { uuid, directory }) => show_uuid(&uuid, directory),
        Command::ShowFile(cli::ShowFileParams { file }) => show_file(&file),
        Command::Remove(cli::RemoveParams { ids, directory }) => remove(&ids, directory),
        Command::Clean(cli::CleanParams { directory }) => clean(directory),
    }
}

fn list(
    text: &Option<String>,
    expires_in_days: Option<u64>,
    directory: Option<PathBuf>,
) -> Result<(), cli::Error> {
    let dir = mp::with_directory(directory)?;
    let entries = mp::entries(&dir).map(|entries| entries.collect::<Vec<_>>())?;
    let date =
        expires_in_days.map(|days| SystemTime::now() + Duration::from_secs(days * 24 * 60 * 60));
    let filter_string = text.as_ref();
    let mut profiles = mp::filter(entries, |profile| match (date, filter_string) {
        (Some(date), Some(string)) => {
            profile.info.expiration_date <= date && profile.info.contains(string)
        }
        (Some(date), _) => profile.info.expiration_date <= date,
        (_, Some(string)) => profile.info.contains(string),
        (_, _) => true,
    });
    profiles.sort_by(|a, b| a.info.creation_date.cmp(&b.info.creation_date));
    if profiles.is_empty() {
        Ok(())
    } else {
        let stdout = io::stdout();
        let mut stdout = stdout.lock();
        for profile in &profiles {
            writeln!(&mut stdout, "{}\n", profile.info.description())?;
        }
        Ok(())
    }
}

fn show_uuid(uuid: &str, directory: Option<PathBuf>) -> Result<(), cli::Error> {
    let dir = mp::with_directory(directory)?;
    let profile = mp::find_by_uuid(&dir, &uuid)?;
    let xml = mp::show(&profile.path)?;
    writeln!(io::stdout(), "{}", xml)?;
    Ok(())
}

fn show_file(path: &Path) -> Result<(), cli::Error> {
    let xml = mp::show(&path)?;
    writeln!(io::stdout(), "{}", xml)?;
    Ok(())
}

fn remove(ids: &[String], directory: Option<PathBuf>) -> Result<(), cli::Error> {
    let dir = mp::with_directory(directory)?;
    let profiles = mp::find_by_ids(&dir, &ids)?;
    for profile in profiles {
        if mp::remove(&profile.path).is_ok() {
            println!("\nRemoved: {}", profile.info.description());
        } else {
            println!("\nError while removing '{}'", profile.info.uuid);
        }
    }
    Ok(())
}

fn clean(directory: Option<PathBuf>) -> Result<(), cli::Error> {
    fn concat(mut s1: String, s2: &str) -> String {
        s1.push_str(&s2);
        s1
    }

    let date = SystemTime::now();
    mp::with_directory(directory)
        .and_then(|dir| mp::entries(&dir).map(|entries| entries.collect::<Vec<_>>()))
        .map_err(|err| err.into())
        .map(|entries| mp::filter(entries, |profile| profile.info.expiration_date <= date))
        .and_then(|profiles| {
            if profiles.is_empty() {
                println!("All provisioning profiles are valid");
                Ok(())
            } else {
                profiles
                    .iter()
                    .map(|profile| {
                        std::fs::remove_file(&profile.path)
                            .map(|_| format!("'{}' was removed\n", profile.info.uuid))
                            .map_err(|err| format!("'{}' {}\n", profile.info.uuid, err))
                    })
                    .fold(Ok(String::new()), |acc, s| match (acc, s) {
                        (Ok(acc), Ok(s)) => Ok(concat(acc, &s)),
                        (Ok(acc), Err(s)) | (Err(acc), Ok(s)) | (Err(acc), Err(s)) => {
                            Err(concat(acc, &s))
                        }
                    })
                    .map(|s| println!("{}", s))
                    .map_err(cli::Error::Custom)
            }
        })
}
<commit_msg>Refactor out chaining combinators for the clean command<commit_after>use crate::cli::Command;
use mprovision as mp;
use std::env;
use std::io::{self, Write};
use std::path::Path;
use std::path::PathBuf;
use std::time::{Duration, SystemTime};

mod cli;

fn main() {
    if let Err(e) = cli::parse(env::args()).and_then(run) {
        e.exit();
    }
}

fn run(command: cli::Command) -> Result<(), cli::Error> {
    match command {
        Command::List(cli::ListParams {
            text,
            expire_in_days,
            directory,
        }) => list(&text, expire_in_days, directory),
        Command::ShowUuid(cli::ShowUuidParams { uuid, directory }) => show_uuid(&uuid, directory),
        Command::ShowFile(cli::ShowFileParams { file }) => show_file(&file),
        Command::Remove(cli::RemoveParams { ids, directory }) => remove(&ids, directory),
        Command::Clean(cli::CleanParams { directory }) => clean(directory),
    }
}

fn list(
    text: &Option<String>,
    expires_in_days: Option<u64>,
    directory: Option<PathBuf>,
) -> Result<(), cli::Error> {
    let dir = mp::with_directory(directory)?;
    let entries = mp::entries(&dir).map(|entries| entries.collect::<Vec<_>>())?;
    let date =
        expires_in_days.map(|days| SystemTime::now() + Duration::from_secs(days * 24 * 60 * 60));
    let filter_string = text.as_ref();
    let mut profiles = mp::filter(entries, |profile| match (date, filter_string) {
        (Some(date), Some(string)) => {
            profile.info.expiration_date <= date && profile.info.contains(string)
        }
        (Some(date), _) => profile.info.expiration_date <= date,
        (_, Some(string)) => profile.info.contains(string),
        (_, _) => true,
    });
    profiles.sort_by(|a, b| a.info.creation_date.cmp(&b.info.creation_date));
    if profiles.is_empty() {
        Ok(())
    } else {
        let stdout = io::stdout();
        let mut stdout = stdout.lock();
        for profile in &profiles {
            writeln!(&mut stdout, "{}\n", profile.info.description())?;
        }
        Ok(())
    }
}

fn show_uuid(uuid: &str, directory: Option<PathBuf>) -> Result<(), cli::Error> {
    let dir = mp::with_directory(directory)?;
    let profile = mp::find_by_uuid(&dir, &uuid)?;
    let xml = mp::show(&profile.path)?;
    writeln!(io::stdout(), "{}", xml)?;
    Ok(())
}

fn show_file(path: &Path) -> Result<(), cli::Error> {
    let xml = mp::show(&path)?;
    writeln!(io::stdout(), "{}", xml)?;
    Ok(())
}

fn remove(ids: &[String], directory: Option<PathBuf>) -> Result<(), cli::Error> {
    let dir = mp::with_directory(directory)?;
    let profiles = mp::find_by_ids(&dir, &ids)?;
    for profile in profiles {
        if mp::remove(&profile.path).is_ok() {
            println!("\nRemoved: {}", profile.info.description());
        } else {
            println!("\nError while removing '{}'", profile.info.uuid);
        }
    }
    Ok(())
}

fn clean(directory: Option<PathBuf>) -> Result<(), cli::Error> {
    let date = SystemTime::now();
    let dir = mp::with_directory(directory)?;
    let entries = mp::entries(&dir).map(|entries| entries.collect::<Vec<_>>())?;
    let profiles = mp::filter(entries, |profile| profile.info.expiration_date <= date);
    if profiles.is_empty() {
        println!("All provisioning profiles are valid");
    } else {
        for profile in profiles {
            if mp::remove(&profile.path).is_ok() {
                println!("\nRemoved: {}", profile.info.description());
            } else {
                println!("\nError while removing '{}'", profile.info.uuid);
            }
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>#![feature(rustc_private)]

extern crate env_logger;
extern crate getopts;
#[macro_use]
extern crate log;
extern crate log_settings;
extern crate miri;
extern crate rustc;
extern crate rustc_metadata;
extern crate rustc_driver;
extern crate rustc_errors;
extern crate rustc_codegen_utils;
extern crate rustc_interface;
extern crate syntax;

use std::str::FromStr;
use std::env;

use rustc_interface::interface;
use rustc::hir::def_id::LOCAL_CRATE;

struct MiriCompilerCalls {
    miri_config: miri::MiriConfig,
}

impl rustc_driver::Callbacks for MiriCompilerCalls {
    fn after_parsing(&mut self, compiler: &interface::Compiler) -> bool {
        let attr = (
            syntax::symbol::Symbol::intern("miri"),
            syntax::feature_gate::AttributeType::Whitelisted,
        );
        compiler.session().plugin_attributes.borrow_mut().push(attr);

        // Continue execution
        true
    }

    fn after_analysis(&mut self, compiler: &interface::Compiler) -> bool {
        init_late_loggers();
        compiler.session().abort_if_errors();

        compiler.global_ctxt().unwrap().peek_mut().enter(|tcx| {
            let (entry_def_id, _) = tcx.entry_fn(LOCAL_CRATE).expect("no main function found!");
            let mut config = self.miri_config.clone();

            // Add filename to `miri` arguments.
            config.args.insert(0, compiler.input().filestem().to_string());

            miri::eval_main(tcx, entry_def_id, config);
        });

        compiler.session().abort_if_errors();

        // Don't continue execution
        false
    }
}

fn init_early_loggers() {
    // Note that our `extern crate log` is *not* the same as rustc's; as a result, we have to
    // initialize them both, and we always initialize `miri`'s first.
    let env = env_logger::Env::new().filter("MIRI_LOG").write_style("MIRI_LOG_STYLE");
    env_logger::init_from_env(env);
    // We only initialize `rustc` if the env var is set (so the user asked for it).
    // If it is not set, we avoid initializing now so that we can initialize
    // later with our custom settings, and *not* log anything for what happens before
    // `miri` gets started.
    if env::var("RUSTC_LOG").is_ok() {
        rustc_driver::init_rustc_env_logger();
    }
}

fn init_late_loggers() {
    // We initialize loggers right before we start evaluation. We overwrite the `RUSTC_LOG`
    // env var if it is not set, control it based on `MIRI_LOG`.
    if let Ok(var) = env::var("MIRI_LOG") {
        if env::var("RUSTC_LOG").is_err() {
            // We try to be a bit clever here: if `MIRI_LOG` is just a single level
            // used for everything, we only apply it to the parts of rustc that are
            // CTFE-related. Otherwise, we use it verbatim for `RUSTC_LOG`.
            // This way, if you set `MIRI_LOG=trace`, you get only the right parts of
            // rustc traced, but you can also do `MIRI_LOG=miri=trace,rustc_mir::interpret=debug`.
            if log::Level::from_str(&var).is_ok() {
                env::set_var("RUSTC_LOG",
                    &format!("rustc::mir::interpret={0},rustc_mir::interpret={0}", var));
            } else {
                env::set_var("RUSTC_LOG", &var);
            }
            rustc_driver::init_rustc_env_logger();
        }
    }

    // If `MIRI_BACKTRACE` is set and `RUST_CTFE_BACKTRACE` is not, set `RUST_CTFE_BACKTRACE`.
    // Do this late, so we really only apply this to miri's errors.
    if let Ok(var) = env::var("MIRI_BACKTRACE") {
        if env::var("RUST_CTFE_BACKTRACE") == Err(env::VarError::NotPresent) {
            env::set_var("RUST_CTFE_BACKTRACE", &var);
        }
    }
}

/// Returns the "default sysroot" that Miri will use if no `--sysroot` flag is set.
/// Should be a compile-time constant.
fn compile_time_sysroot() -> Option<String> {
    if option_env!("RUSTC_STAGE").is_some() {
        // This is being built as part of rustc, and gets shipped with rustup.
        // We can rely on the sysroot computation in librustc.
        return None;
    }
    // For builds outside rustc, we need to ensure that we got a sysroot
    // that gets used as a default.  The sysroot computation in librustc would
    // end up somewhere in the build dir.
    // Taken from PR <https://github.com/Manishearth/rust-clippy/pull/911>.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    Some(match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => {
            option_env!("RUST_SYSROOT")
                .expect("To build Miri without rustup, set the `RUST_SYSROOT` env var at build time")
                .to_owned()
        }
    })
}

fn main() {
    init_early_loggers();

    // Parse our arguments and split them across `rustc` and `miri`.
    let mut validate = true;
    let mut seed: Option<u64> = None;
    let mut rustc_args = vec![];
    let mut miri_args = vec![];
    let mut after_dashdash = false;
    for arg in std::env::args() {
        if rustc_args.is_empty() {
            // Very first arg: for `rustc`.
            rustc_args.push(arg);
        }
        else if after_dashdash {
            // Everything that comes after are `miri` args.
            miri_args.push(arg);
        } else {
            match arg.as_str() {
                "-Zmiri-disable-validation" => {
                    validate = false;
                },
                "--" => {
                    after_dashdash = true;
                }
                arg if arg.starts_with("-Zmiri-seed=") => {
                    if seed.is_some() {
                        panic!("Cannot specify -Zmiri-seed multiple times!");
                    }
                    let seed_raw = hex::decode(arg.trim_start_matches("-Zmiri-seed=")).unwrap();
                    if seed_raw.len() > 8 {
                        panic!(format!("-Zmiri-seed must be at most 8 bytes, was {}", seed_raw.len()));
                    }

                    let mut bytes = [0; 8];
                    bytes[..seed_raw.len()].copy_from_slice(&seed_raw);
                    seed = Some(u64::from_be_bytes(bytes));

                },
                _ => {
                    rustc_args.push(arg);
                }
            }
        }
    }

    // Determine sysroot if needed.  Make sure we always call `compile_time_sysroot`
    // as that also does some sanity-checks of the environment we were built in.
    // FIXME: Ideally we'd turn a bad build env into a compile-time error, but
    // CTFE does not seem powerful enough for that yet.
    if let Some(sysroot) = compile_time_sysroot() {
        let sysroot_flag = "--sysroot".to_string();
        if !rustc_args.contains(&sysroot_flag) {
            // We need to overwrite the default that librustc would compute.
            rustc_args.push(sysroot_flag);
            rustc_args.push(sysroot);
        }
    }

    // Finally, add the default flags all the way in the beginning, but after the binary name.
    rustc_args.splice(1..1, miri::miri_default_args().iter().map(ToString::to_string));

    debug!("rustc arguments: {:?}", rustc_args);
    debug!("miri arguments: {:?}", miri_args);
    let miri_config = miri::MiriConfig { validate, args: miri_args, seed };
    let result = rustc_driver::report_ices_to_stderr_if_any(move || {
        rustc_driver::run_compiler(&rustc_args, &mut MiriCompilerCalls { miri_config }, None, None)
    }).and_then(|result| result);
    std::process::exit(result.is_err() as i32);
}
<commit_msg>Give a useful error message if user gives invalid random seed<commit_after>#![feature(rustc_private)]

extern crate env_logger;
extern crate getopts;
#[macro_use]
extern crate log;
extern crate log_settings;
extern crate miri;
extern crate rustc;
extern crate rustc_metadata;
extern crate rustc_driver;
extern crate rustc_errors;
extern crate rustc_codegen_utils;
extern crate rustc_interface;
extern crate syntax;

use std::str::FromStr;
use std::env;

use hex::FromHexError;

use rustc_interface::interface;
use rustc::hir::def_id::LOCAL_CRATE;

struct MiriCompilerCalls {
    miri_config: miri::MiriConfig,
}

impl rustc_driver::Callbacks for MiriCompilerCalls {
    fn after_parsing(&mut self, compiler: &interface::Compiler) -> bool {
        let attr = (
            syntax::symbol::Symbol::intern("miri"),
            syntax::feature_gate::AttributeType::Whitelisted,
        );
        compiler.session().plugin_attributes.borrow_mut().push(attr);

        // Continue execution
        true
    }

    fn after_analysis(&mut self, compiler: &interface::Compiler) -> bool {
        init_late_loggers();
        compiler.session().abort_if_errors();

        compiler.global_ctxt().unwrap().peek_mut().enter(|tcx| {
            let (entry_def_id, _) = tcx.entry_fn(LOCAL_CRATE).expect("no main function found!");
            let mut config = self.miri_config.clone();

            // Add filename to `miri` arguments.
            config.args.insert(0, compiler.input().filestem().to_string());

            miri::eval_main(tcx, entry_def_id, config);
        });

        compiler.session().abort_if_errors();

        // Don't continue execution
        false
    }
}

fn init_early_loggers() {
    // Note that our `extern crate log` is *not* the same as rustc's; as a result, we have to
    // initialize them both, and we always initialize `miri`'s first.
    let env = env_logger::Env::new().filter("MIRI_LOG").write_style("MIRI_LOG_STYLE");
    env_logger::init_from_env(env);
    // We only initialize `rustc` if the env var is set (so the user asked for it).
    // If it is not set, we avoid initializing now so that we can initialize
    // later with our custom settings, and *not* log anything for what happens before
    // `miri` gets started.
    if env::var("RUSTC_LOG").is_ok() {
        rustc_driver::init_rustc_env_logger();
    }
}

fn init_late_loggers() {
    // We initialize loggers right before we start evaluation. We overwrite the `RUSTC_LOG`
    // env var if it is not set, control it based on `MIRI_LOG`.
    if let Ok(var) = env::var("MIRI_LOG") {
        if env::var("RUSTC_LOG").is_err() {
            // We try to be a bit clever here: if `MIRI_LOG` is just a single level
            // used for everything, we only apply it to the parts of rustc that are
            // CTFE-related. Otherwise, we use it verbatim for `RUSTC_LOG`.
            // This way, if you set `MIRI_LOG=trace`, you get only the right parts of
            // rustc traced, but you can also do `MIRI_LOG=miri=trace,rustc_mir::interpret=debug`.
            if log::Level::from_str(&var).is_ok() {
                env::set_var("RUSTC_LOG",
                    &format!("rustc::mir::interpret={0},rustc_mir::interpret={0}", var));
            } else {
                env::set_var("RUSTC_LOG", &var);
            }
            rustc_driver::init_rustc_env_logger();
        }
    }

    // If `MIRI_BACKTRACE` is set and `RUST_CTFE_BACKTRACE` is not, set `RUST_CTFE_BACKTRACE`.
    // Do this late, so we really only apply this to miri's errors.
    if let Ok(var) = env::var("MIRI_BACKTRACE") {
        if env::var("RUST_CTFE_BACKTRACE") == Err(env::VarError::NotPresent) {
            env::set_var("RUST_CTFE_BACKTRACE", &var);
        }
    }
}

/// Returns the "default sysroot" that Miri will use if no `--sysroot` flag is set.
/// Should be a compile-time constant.
fn compile_time_sysroot() -> Option<String> {
    if option_env!("RUSTC_STAGE").is_some() {
        // This is being built as part of rustc, and gets shipped with rustup.
        // We can rely on the sysroot computation in librustc.
        return None;
    }
    // For builds outside rustc, we need to ensure that we got a sysroot
    // that gets used as a default.  The sysroot computation in librustc would
    // end up somewhere in the build dir.
    // Taken from PR <https://github.com/Manishearth/rust-clippy/pull/911>.
    let home = option_env!("RUSTUP_HOME").or(option_env!("MULTIRUST_HOME"));
    let toolchain = option_env!("RUSTUP_TOOLCHAIN").or(option_env!("MULTIRUST_TOOLCHAIN"));
    Some(match (home, toolchain) {
        (Some(home), Some(toolchain)) => format!("{}/toolchains/{}", home, toolchain),
        _ => {
            option_env!("RUST_SYSROOT")
                .expect("To build Miri without rustup, set the `RUST_SYSROOT` env var at build time")
                .to_owned()
        }
    })
}

fn main() {
    init_early_loggers();

    // Parse our arguments and split them across `rustc` and `miri`.
    let mut validate = true;
    let mut seed: Option<u64> = None;
    let mut rustc_args = vec![];
    let mut miri_args = vec![];
    let mut after_dashdash = false;
    for arg in std::env::args() {
        if rustc_args.is_empty() {
            // Very first arg: for `rustc`.
            rustc_args.push(arg);
        }
        else if after_dashdash {
            // Everything that comes after are `miri` args.
            miri_args.push(arg);
        } else {
            match arg.as_str() {
                "-Zmiri-disable-validation" => {
                    validate = false;
                },
                "--" => {
                    after_dashdash = true;
                }
                arg if arg.starts_with("-Zmiri-seed=") => {
                    if seed.is_some() {
                        panic!("Cannot specify -Zmiri-seed multiple times!");
                    }
                    let seed_raw = hex::decode(arg.trim_start_matches("-Zmiri-seed="))
                        .unwrap_or_else(|err| match err {
                            FromHexError::InvalidHexCharacter { .. } => panic!(
                                "-Zmiri-seed should only contain valid hex digits [0-9a-fA-F]"
                            ),
                            FromHexError::OddLength => panic!("-Zmiri-seed should have an even number of digits"),
                            err => panic!("Unknown error decoding -Zmiri-seed as hex: {:?}", err),
                        });
                    if seed_raw.len() > 8 {
                        panic!(format!("-Zmiri-seed must be at most 8 bytes, was {}", seed_raw.len()));
                    }

                    let mut bytes = [0; 8];
                    bytes[..seed_raw.len()].copy_from_slice(&seed_raw);
                    seed = Some(u64::from_be_bytes(bytes));

                },
                _ => {
                    rustc_args.push(arg);
                }
            }
        }
    }

    // Determine sysroot if needed.  Make sure we always call `compile_time_sysroot`
    // as that also does some sanity-checks of the environment we were built in.
    // FIXME: Ideally we'd turn a bad build env into a compile-time error, but
    // CTFE does not seem powerful enough for that yet.
    if let Some(sysroot) = compile_time_sysroot() {
        let sysroot_flag = "--sysroot".to_string();
        if !rustc_args.contains(&sysroot_flag) {
            // We need to overwrite the default that librustc would compute.
            rustc_args.push(sysroot_flag);
            rustc_args.push(sysroot);
        }
    }

    // Finally, add the default flags all the way in the beginning, but after the binary name.
    rustc_args.splice(1..1, miri::miri_default_args().iter().map(ToString::to_string));

    debug!("rustc arguments: {:?}", rustc_args);
    debug!("miri arguments: {:?}", miri_args);
    let miri_config = miri::MiriConfig { validate, args: miri_args, seed };
    let result = rustc_driver::report_ices_to_stderr_if_any(move || {
        rustc_driver::run_compiler(&rustc_args, &mut MiriCompilerCalls { miri_config }, None, None)
    }).and_then(|result| result);
    std::process::exit(result.is_err() as i32);
}
<|endoftext|>


<commit_before>pub mod bookitem;
pub mod bookconfig;

pub mod bookconfig_test;

pub use self::bookitem::{BookItem, BookItems};
pub use self::bookconfig::BookConfig;

use std::path::{Path, PathBuf};
use std::fs::{self, File};
use std::error::Error;
use std::io;
use std::io::Write;
use std::io::ErrorKind;
use std::process::Command;

use {theme, parse, utils};
use renderer::{Renderer, HtmlHandlebars};

pub struct MDBook {
    root: PathBuf,
    dest: PathBuf,
    src: PathBuf,
    theme_path: PathBuf,

    pub title: String,
    pub author: String,
    pub description: String,

    pub content: Vec<BookItem>,
    renderer: Box<Renderer>,

    livereload: Option<String>,

    /// Should `mdbook build` create files referenced from SUMMARY.md if they
    /// don't exist
    pub create_missing: bool,
}

impl MDBook {
    /// Create a new `MDBook` struct with root directory `root`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use std::path::Path;
    /// # fn main() {
    /// let book = MDBook::new(Path::new("root_dir"));
    /// # }
    /// ```
    ///
    /// In this example, `root_dir` will be the root directory of our book and is specified in function
    /// of the current working directory by using a relative path instead of an absolute path.
    ///
    /// Default directory paths:
    ///
    /// - source: `root/src`
    /// - output: `root/book`
    /// - theme: `root/theme`
    ///
    /// They can both be changed by using [`set_src()`](#method.set_src) and [`set_dest()`](#method.set_dest)

    pub fn new(root: &Path) -> MDBook {

        if !root.exists() || !root.is_dir() {
            warn!("{:?} No directory with that name", root);
        }

        MDBook {
            root: root.to_owned(),
            dest: root.join("book"),
            src: root.join("src"),
            theme_path: root.join("theme"),

            title: String::new(),
            author: String::new(),
            description: String::new(),

            content: vec![],
            renderer: Box::new(HtmlHandlebars::new()),

            livereload: None,
            create_missing: true,
        }
    }

    /// Returns a flat depth-first iterator over the elements of the book, it returns an [BookItem enum](bookitem.html):
    /// `(section: String, bookitem: &BookItem)`
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use mdbook::BookItem;
    /// # use std::path::Path;
    /// # fn main() {
    /// # let mut book = MDBook::new(Path::new("mybook"));
    /// for item in book.iter() {
    ///     match item {
    ///         &BookItem::Chapter(ref section, ref chapter) => {},
    ///         &BookItem::Affix(ref chapter) => {},
    ///         &BookItem::Spacer => {},
    ///     }
    /// }
    ///
    /// // would print something like this:
    /// // 1. Chapter 1
    /// // 1.1 Sub Chapter
    /// // 1.2 Sub Chapter
    /// // 2. Chapter 2
    /// //
    /// // etc.
    /// # }
    /// ```

    pub fn iter(&self) -> BookItems {
        BookItems {
            items: &self.content[..],
            current_index: 0,
            stack: Vec::new(),
        }
    }

    /// `init()` creates some boilerplate files and directories to get you started with your book.
    ///
    /// ```text
    /// book-test/
    /// ├── book
    /// └── src
    ///     ├── chapter_1.md
    ///     └── SUMMARY.md
    /// ```
    ///
    /// It uses the paths given as source and output directories and adds a `SUMMARY.md` and a
    /// `chapter_1.md` to the source directory.

    pub fn init(&mut self) -> Result<(), Box<Error>> {

        debug!("[fn]: init");

        if !self.root.exists() {
            fs::create_dir_all(&self.root).unwrap();
            info!("{:?} created", &self.root);
        }

        {

            if !self.dest.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.dest);
                try!(fs::create_dir_all(&self.dest));
            }

            if !self.src.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.src);
                try!(fs::create_dir_all(&self.src));
            }

            let summary = self.src.join("SUMMARY.md");

            if !summary.exists() {

                // Summary does not exist, create it

                debug!("[*]: {:?} does not exist, trying to create SUMMARY.md", self.src.join("SUMMARY.md"));
                let mut f = try!(File::create(&self.src.join("SUMMARY.md")));

                debug!("[*]: Writing to SUMMARY.md");

                try!(writeln!(f, "# Summary"));
                try!(writeln!(f, ""));
                try!(writeln!(f, "- [Chapter 1](./chapter_1.md)"));
            }
        }

        // parse SUMMARY.md, and create the missing item related file
        try!(self.parse_summary());

        debug!("[*]: constructing paths for missing files");
        for item in self.iter() {
            debug!("[*]: item: {:?}", item);
            let ch = match *item {
                BookItem::Spacer => continue,
                BookItem::Chapter(_, ref ch) |
                BookItem::Affix(ref ch) => ch,
            };
            if ch.path != PathBuf::new() {
                let path = self.src.join(&ch.path);

                if !path.exists() {
                    if !self.create_missing {
                        return Err(format!(
                            "'{}' referenced from SUMMARY.md does not exist.",
                            path.to_string_lossy()).into());
                    }
                    debug!("[*]: {:?} does not exist, trying to create file", path);
                    try!(::std::fs::create_dir_all(path.parent().unwrap()));
                    let mut f = try!(File::create(path));

                    // debug!("[*]: Writing to {:?}", path);
                    try!(writeln!(f, "# {}", ch.name));
                }
            }
        }

        debug!("[*]: init done");
        Ok(())
    }

    pub fn create_gitignore(&self) {
        let gitignore = self.get_gitignore();

        if !gitignore.exists() {
            // Gitignore does not exist, create it

            // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`. If it
            // is not, `strip_prefix` will return an Error.
            if !self.get_dest().starts_with(&self.root) {
                return;
            }

            let relative = self.get_dest()
                               .strip_prefix(&self.root)
                               .expect("Destination is not relative to root.");
            let relative = relative.to_str()
                                   .expect("Path could not be yielded into a string slice.");

            debug!("[*]: {:?} does not exist, trying to create .gitignore", gitignore);

            let mut f = File::create(&gitignore).expect("Could not create file.");

            debug!("[*]: Writing to .gitignore");

            writeln!(f, "{}", relative).expect("Could not write to file.");
        }
    }

    /// The `build()` method is the one where everything happens. First it parses `SUMMARY.md` to
    /// construct the book's structure in the form of a `Vec<BookItem>` and then calls `render()`
    /// method of the current renderer.
    ///
    /// It is the renderer who generates all the output files.
    pub fn build(&mut self) -> Result<(), Box<Error>> {
        debug!("[fn]: build");

        try!(self.init());

        // Clean output directory
        try!(utils::fs::remove_dir_content(&self.dest));

        try!(self.renderer.render(&self));

        Ok(())
    }

    pub fn get_gitignore(&self) -> PathBuf {
        self.root.join(".gitignore")
    }

    pub fn copy_theme(&self) -> Result<(), Box<Error>> {
        debug!("[fn]: copy_theme");

        let theme_dir = self.src.join("theme");

        if !theme_dir.exists() {
            debug!("[*]: {:?} does not exist, trying to create directory", theme_dir);
            try!(fs::create_dir(&theme_dir));
        }

        // index.hbs
        let mut index = try!(File::create(&theme_dir.join("index.hbs")));
        try!(index.write_all(theme::INDEX));

        // book.css
        let mut css = try!(File::create(&theme_dir.join("book.css")));
        try!(css.write_all(theme::CSS));

        // favicon.png
        let mut favicon = try!(File::create(&theme_dir.join("favicon.png")));
        try!(favicon.write_all(theme::FAVICON));

        // book.js
        let mut js = try!(File::create(&theme_dir.join("book.js")));
        try!(js.write_all(theme::JS));

        // highlight.css
        let mut highlight_css = try!(File::create(&theme_dir.join("highlight.css")));
        try!(highlight_css.write_all(theme::HIGHLIGHT_CSS));

        // highlight.js
        let mut highlight_js = try!(File::create(&theme_dir.join("highlight.js")));
        try!(highlight_js.write_all(theme::HIGHLIGHT_JS));

        Ok(())
    }

    pub fn write_file<P: AsRef<Path>>(&self, filename: P, content: &[u8]) -> Result<(), Box<Error>> {
        let path = self.get_dest().join(filename);
        try!(utils::fs::create_file(&path).and_then(|mut file| {
            file.write_all(content)
        }).map_err(|e| {
            io::Error::new(io::ErrorKind::Other, format!("Could not create {}: {}", path.display(), e))
        }));
        Ok(())
    }

    /// Parses the `book.json` file (if it exists) to extract the configuration parameters.
    /// The `book.json` file should be in the root directory of the book.
    /// The root directory is the one specified when creating a new `MDBook`

    pub fn read_config(mut self) -> Self {

        let config = BookConfig::new(&self.root)
                         .read_config(&self.root)
                         .to_owned();

        self.title = config.title;
        self.description = config.description;
        self.author = config.author;

        self.dest = config.dest;
        self.src = config.src;
        self.theme_path = config.theme_path;

        self
    }

    /// You can change the default renderer to another one by using this method. The only requirement
    /// is for your renderer to implement the [Renderer trait](../../renderer/renderer/trait.Renderer.html)
    ///
    /// ```no_run
    /// extern crate mdbook;
    /// use mdbook::MDBook;
    /// use mdbook::renderer::HtmlHandlebars;
    /// # use std::path::Path;
    ///
    /// fn main() {
    ///     let mut book = MDBook::new(Path::new("mybook"))
    ///                         .set_renderer(Box::new(HtmlHandlebars::new()));
    ///
    ///     // In this example we replace the default renderer by the default renderer...
    ///     // Don't forget to put your renderer in a Box
    /// }
    /// ```
    ///
    /// **note:** Don't forget to put your renderer in a `Box` before passing it to `set_renderer()`

    pub fn set_renderer(mut self, renderer: Box<Renderer>) -> Self {
        self.renderer = renderer;
        self
    }

    pub fn test(&mut self) -> Result<(), Box<Error>> {
        // read in the chapters
        try!(self.parse_summary());
        for item in self.iter() {

            if let BookItem::Chapter(_, ref ch) = *item {
                if ch.path != PathBuf::new() {

                    let path = self.get_src().join(&ch.path);

                    println!("[*]: Testing file: {:?}", path);

                    let output_result = Command::new("rustdoc")
                                            .arg(&path)
                                            .arg("--test")
                                            .output();
                    let output = try!(output_result);

                    if !output.status.success() {
                        return Err(Box::new(io::Error::new(ErrorKind::Other, format!(
                                        "{}\n{}",
                                        String::from_utf8_lossy(&output.stdout),
                                        String::from_utf8_lossy(&output.stderr)))) as Box<Error>);
                    }
                }
            }
        }
        Ok(())
    }

    pub fn get_root(&self) -> &Path {
        &self.root
    }

    pub fn set_dest(mut self, dest: &Path) -> Self {

        // Handle absolute and relative paths
        if dest.is_absolute() {
            self.dest = dest.to_owned();
        } else {
            let dest = self.root.join(dest).to_owned();
            self.dest = dest;
        }

        self
    }

    pub fn get_dest(&self) -> &Path {
        &self.dest
    }

    pub fn set_src(mut self, src: &Path) -> Self {

        // Handle absolute and relative paths
        if src.is_absolute() {
            self.src = src.to_owned();
        } else {
            let src = self.root.join(src).to_owned();
            self.src = src;
        }

        self
    }

    pub fn get_src(&self) -> &Path {
        &self.src
    }

    pub fn set_title(mut self, title: &str) -> Self {
        self.title = title.to_owned();
        self
    }

    pub fn get_title(&self) -> &str {
        &self.title
    }

    pub fn set_author(mut self, author: &str) -> Self {
        self.author = author.to_owned();
        self
    }

    pub fn get_author(&self) -> &str {
        &self.author
    }

    pub fn set_description(mut self, description: &str) -> Self {
        self.description = description.to_owned();
        self
    }

    pub fn get_description(&self) -> &str {
        &self.description
    }

    pub fn set_livereload(&mut self, livereload: String) -> &mut Self {
        self.livereload = Some(livereload);
        self
    }

    pub fn unset_livereload(&mut self) -> &Self {
        self.livereload = None;
        self
    }

    pub fn get_livereload(&self) -> Option<&String> {
        self.livereload.as_ref()
    }

    pub fn set_theme_path(mut self, theme_path: &Path) -> Self {
        self.theme_path = if theme_path.is_absolute() {
            theme_path.to_owned()
        } else {
            self.root.join(theme_path).to_owned()
        };
        self
    }

    pub fn get_theme_path(&self) -> &Path {
        &self.theme_path
    }

    // Construct book
    fn parse_summary(&mut self) -> Result<(), Box<Error>> {
        // When append becomes stable, use self.content.append() ...
        self.content = try!(parse::construct_bookitems(&self.src.join("SUMMARY.md")));
        Ok(())
    }
}
<commit_msg>Rewrite an emptiness check.<commit_after>pub mod bookitem;
pub mod bookconfig;

pub mod bookconfig_test;

pub use self::bookitem::{BookItem, BookItems};
pub use self::bookconfig::BookConfig;

use std::path::{Path, PathBuf};
use std::fs::{self, File};
use std::error::Error;
use std::io;
use std::io::Write;
use std::io::ErrorKind;
use std::process::Command;

use {theme, parse, utils};
use renderer::{Renderer, HtmlHandlebars};

pub struct MDBook {
    root: PathBuf,
    dest: PathBuf,
    src: PathBuf,
    theme_path: PathBuf,

    pub title: String,
    pub author: String,
    pub description: String,

    pub content: Vec<BookItem>,
    renderer: Box<Renderer>,

    livereload: Option<String>,

    /// Should `mdbook build` create files referenced from SUMMARY.md if they
    /// don't exist
    pub create_missing: bool,
}

impl MDBook {
    /// Create a new `MDBook` struct with root directory `root`
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use std::path::Path;
    /// # fn main() {
    /// let book = MDBook::new(Path::new("root_dir"));
    /// # }
    /// ```
    ///
    /// In this example, `root_dir` will be the root directory of our book and is specified in function
    /// of the current working directory by using a relative path instead of an absolute path.
    ///
    /// Default directory paths:
    ///
    /// - source: `root/src`
    /// - output: `root/book`
    /// - theme: `root/theme`
    ///
    /// They can both be changed by using [`set_src()`](#method.set_src) and [`set_dest()`](#method.set_dest)

    pub fn new(root: &Path) -> MDBook {

        if !root.exists() || !root.is_dir() {
            warn!("{:?} No directory with that name", root);
        }

        MDBook {
            root: root.to_owned(),
            dest: root.join("book"),
            src: root.join("src"),
            theme_path: root.join("theme"),

            title: String::new(),
            author: String::new(),
            description: String::new(),

            content: vec![],
            renderer: Box::new(HtmlHandlebars::new()),

            livereload: None,
            create_missing: true,
        }
    }

    /// Returns a flat depth-first iterator over the elements of the book, it returns an [BookItem enum](bookitem.html):
    /// `(section: String, bookitem: &BookItem)`
    ///
    /// ```no_run
    /// # extern crate mdbook;
    /// # use mdbook::MDBook;
    /// # use mdbook::BookItem;
    /// # use std::path::Path;
    /// # fn main() {
    /// # let mut book = MDBook::new(Path::new("mybook"));
    /// for item in book.iter() {
    ///     match item {
    ///         &BookItem::Chapter(ref section, ref chapter) => {},
    ///         &BookItem::Affix(ref chapter) => {},
    ///         &BookItem::Spacer => {},
    ///     }
    /// }
    ///
    /// // would print something like this:
    /// // 1. Chapter 1
    /// // 1.1 Sub Chapter
    /// // 1.2 Sub Chapter
    /// // 2. Chapter 2
    /// //
    /// // etc.
    /// # }
    /// ```

    pub fn iter(&self) -> BookItems {
        BookItems {
            items: &self.content[..],
            current_index: 0,
            stack: Vec::new(),
        }
    }

    /// `init()` creates some boilerplate files and directories to get you started with your book.
    ///
    /// ```text
    /// book-test/
    /// ├── book
    /// └── src
    ///     ├── chapter_1.md
    ///     └── SUMMARY.md
    /// ```
    ///
    /// It uses the paths given as source and output directories and adds a `SUMMARY.md` and a
    /// `chapter_1.md` to the source directory.

    pub fn init(&mut self) -> Result<(), Box<Error>> {

        debug!("[fn]: init");

        if !self.root.exists() {
            fs::create_dir_all(&self.root).unwrap();
            info!("{:?} created", &self.root);
        }

        {

            if !self.dest.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.dest);
                try!(fs::create_dir_all(&self.dest));
            }

            if !self.src.exists() {
                debug!("[*]: {:?} does not exist, trying to create directory", self.src);
                try!(fs::create_dir_all(&self.src));
            }

            let summary = self.src.join("SUMMARY.md");

            if !summary.exists() {

                // Summary does not exist, create it

                debug!("[*]: {:?} does not exist, trying to create SUMMARY.md", self.src.join("SUMMARY.md"));
                let mut f = try!(File::create(&self.src.join("SUMMARY.md")));

                debug!("[*]: Writing to SUMMARY.md");

                try!(writeln!(f, "# Summary"));
                try!(writeln!(f, ""));
                try!(writeln!(f, "- [Chapter 1](./chapter_1.md)"));
            }
        }

        // parse SUMMARY.md, and create the missing item related file
        try!(self.parse_summary());

        debug!("[*]: constructing paths for missing files");
        for item in self.iter() {
            debug!("[*]: item: {:?}", item);
            let ch = match *item {
                BookItem::Spacer => continue,
                BookItem::Chapter(_, ref ch) |
                BookItem::Affix(ref ch) => ch,
            };
            if ch.path.as_os_str().is_empty() {
                let path = self.src.join(&ch.path);

                if !path.exists() {
                    if !self.create_missing {
                        return Err(format!(
                            "'{}' referenced from SUMMARY.md does not exist.",
                            path.to_string_lossy()).into());
                    }
                    debug!("[*]: {:?} does not exist, trying to create file", path);
                    try!(::std::fs::create_dir_all(path.parent().unwrap()));
                    let mut f = try!(File::create(path));

                    // debug!("[*]: Writing to {:?}", path);
                    try!(writeln!(f, "# {}", ch.name));
                }
            }
        }

        debug!("[*]: init done");
        Ok(())
    }

    pub fn create_gitignore(&self) {
        let gitignore = self.get_gitignore();

        if !gitignore.exists() {
            // Gitignore does not exist, create it

            // Because of `src/book/mdbook.rs#L37-L39`, `dest` will always start with `root`. If it
            // is not, `strip_prefix` will return an Error.
            if !self.get_dest().starts_with(&self.root) {
                return;
            }

            let relative = self.get_dest()
                               .strip_prefix(&self.root)
                               .expect("Destination is not relative to root.");
            let relative = relative.to_str()
                                   .expect("Path could not be yielded into a string slice.");

            debug!("[*]: {:?} does not exist, trying to create .gitignore", gitignore);

            let mut f = File::create(&gitignore).expect("Could not create file.");

            debug!("[*]: Writing to .gitignore");

            writeln!(f, "{}", relative).expect("Could not write to file.");
        }
    }

    /// The `build()` method is the one where everything happens. First it parses `SUMMARY.md` to
    /// construct the book's structure in the form of a `Vec<BookItem>` and then calls `render()`
    /// method of the current renderer.
    ///
    /// It is the renderer who generates all the output files.
    pub fn build(&mut self) -> Result<(), Box<Error>> {
        debug!("[fn]: build");

        try!(self.init());

        // Clean output directory
        try!(utils::fs::remove_dir_content(&self.dest));

        try!(self.renderer.render(&self));

        Ok(())
    }

    pub fn get_gitignore(&self) -> PathBuf {
        self.root.join(".gitignore")
    }

    pub fn copy_theme(&self) -> Result<(), Box<Error>> {
        debug!("[fn]: copy_theme");

        let theme_dir = self.src.join("theme");

        if !theme_dir.exists() {
            debug!("[*]: {:?} does not exist, trying to create directory", theme_dir);
            try!(fs::create_dir(&theme_dir));
        }

        // index.hbs
        let mut index = try!(File::create(&theme_dir.join("index.hbs")));
        try!(index.write_all(theme::INDEX));

        // book.css
        let mut css = try!(File::create(&theme_dir.join("book.css")));
        try!(css.write_all(theme::CSS));

        // favicon.png
        let mut favicon = try!(File::create(&theme_dir.join("favicon.png")));
        try!(favicon.write_all(theme::FAVICON));

        // book.js
        let mut js = try!(File::create(&theme_dir.join("book.js")));
        try!(js.write_all(theme::JS));

        // highlight.css
        let mut highlight_css = try!(File::create(&theme_dir.join("highlight.css")));
        try!(highlight_css.write_all(theme::HIGHLIGHT_CSS));

        // highlight.js
        let mut highlight_js = try!(File::create(&theme_dir.join("highlight.js")));
        try!(highlight_js.write_all(theme::HIGHLIGHT_JS));

        Ok(())
    }

    pub fn write_file<P: AsRef<Path>>(&self, filename: P, content: &[u8]) -> Result<(), Box<Error>> {
        let path = self.get_dest().join(filename);
        try!(utils::fs::create_file(&path).and_then(|mut file| {
            file.write_all(content)
        }).map_err(|e| {
            io::Error::new(io::ErrorKind::Other, format!("Could not create {}: {}", path.display(), e))
        }));
        Ok(())
    }

    /// Parses the `book.json` file (if it exists) to extract the configuration parameters.
    /// The `book.json` file should be in the root directory of the book.
    /// The root directory is the one specified when creating a new `MDBook`

    pub fn read_config(mut self) -> Self {

        let config = BookConfig::new(&self.root)
                         .read_config(&self.root)
                         .to_owned();

        self.title = config.title;
        self.description = config.description;
        self.author = config.author;

        self.dest = config.dest;
        self.src = config.src;
        self.theme_path = config.theme_path;

        self
    }

    /// You can change the default renderer to another one by using this method. The only requirement
    /// is for your renderer to implement the [Renderer trait](../../renderer/renderer/trait.Renderer.html)
    ///
    /// ```no_run
    /// extern crate mdbook;
    /// use mdbook::MDBook;
    /// use mdbook::renderer::HtmlHandlebars;
    /// # use std::path::Path;
    ///
    /// fn main() {
    ///     let mut book = MDBook::new(Path::new("mybook"))
    ///                         .set_renderer(Box::new(HtmlHandlebars::new()));
    ///
    ///     // In this example we replace the default renderer by the default renderer...
    ///     // Don't forget to put your renderer in a Box
    /// }
    /// ```
    ///
    /// **note:** Don't forget to put your renderer in a `Box` before passing it to `set_renderer()`

    pub fn set_renderer(mut self, renderer: Box<Renderer>) -> Self {
        self.renderer = renderer;
        self
    }

    pub fn test(&mut self) -> Result<(), Box<Error>> {
        // read in the chapters
        try!(self.parse_summary());
        for item in self.iter() {

            if let BookItem::Chapter(_, ref ch) = *item {
                if ch.path != PathBuf::new() {

                    let path = self.get_src().join(&ch.path);

                    println!("[*]: Testing file: {:?}", path);

                    let output_result = Command::new("rustdoc")
                                            .arg(&path)
                                            .arg("--test")
                                            .output();
                    let output = try!(output_result);

                    if !output.status.success() {
                        return Err(Box::new(io::Error::new(ErrorKind::Other, format!(
                                        "{}\n{}",
                                        String::from_utf8_lossy(&output.stdout),
                                        String::from_utf8_lossy(&output.stderr)))) as Box<Error>);
                    }
                }
            }
        }
        Ok(())
    }

    pub fn get_root(&self) -> &Path {
        &self.root
    }

    pub fn set_dest(mut self, dest: &Path) -> Self {

        // Handle absolute and relative paths
        if dest.is_absolute() {
            self.dest = dest.to_owned();
        } else {
            let dest = self.root.join(dest).to_owned();
            self.dest = dest;
        }

        self
    }

    pub fn get_dest(&self) -> &Path {
        &self.dest
    }

    pub fn set_src(mut self, src: &Path) -> Self {

        // Handle absolute and relative paths
        if src.is_absolute() {
            self.src = src.to_owned();
        } else {
            let src = self.root.join(src).to_owned();
            self.src = src;
        }

        self
    }

    pub fn get_src(&self) -> &Path {
        &self.src
    }

    pub fn set_title(mut self, title: &str) -> Self {
        self.title = title.to_owned();
        self
    }

    pub fn get_title(&self) -> &str {
        &self.title
    }

    pub fn set_author(mut self, author: &str) -> Self {
        self.author = author.to_owned();
        self
    }

    pub fn get_author(&self) -> &str {
        &self.author
    }

    pub fn set_description(mut self, description: &str) -> Self {
        self.description = description.to_owned();
        self
    }

    pub fn get_description(&self) -> &str {
        &self.description
    }

    pub fn set_livereload(&mut self, livereload: String) -> &mut Self {
        self.livereload = Some(livereload);
        self
    }

    pub fn unset_livereload(&mut self) -> &Self {
        self.livereload = None;
        self
    }

    pub fn get_livereload(&self) -> Option<&String> {
        self.livereload.as_ref()
    }

    pub fn set_theme_path(mut self, theme_path: &Path) -> Self {
        self.theme_path = if theme_path.is_absolute() {
            theme_path.to_owned()
        } else {
            self.root.join(theme_path).to_owned()
        };
        self
    }

    pub fn get_theme_path(&self) -> &Path {
        &self.theme_path
    }

    // Construct book
    fn parse_summary(&mut self) -> Result<(), Box<Error>> {
        // When append becomes stable, use self.content.append() ...
        self.content = try!(parse::construct_bookitems(&self.src.join("SUMMARY.md")));
        Ok(())
    }
}
<|endoftext|>


<commit_before>//! Parallelization contracts, describing requirements for data movement along dataflow edges.
//!
//! Pacts describe how data should be exchanged between workers, and implement a method which
//! creates a pair of `Push` and `Pull` implementors from an `A: AsWorker`. These two endpoints
//! respectively distribute and collect data among workers according to the pact.
//!
//! The only requirement of a pact is that it not alter the number of `D` records at each time `T`.
//! The progress tracking logic assumes that this number is independent of the pact used.

use std::marker::PhantomData;

use communication::{Push, Pull, Data};
use communication::allocator::thread::{ThreadPusher, ThreadPuller};

use worker::AsWorker;
use dataflow::channels::pushers::Exchange as ExchangePusher;
use super::{Bundle, Message};

use logging::TimelyLogger as Logger;

/// A `ParallelizationContract` allocates paired `Push` and `Pull` implementors.
pub trait ParallelizationContract<T: 'static, D: 'static> {
    /// Type implementing `Push` produced by this pact.
    type Pusher: Push<Bundle<T, D>>+'static;
    /// Type implementing `Pull` produced by this pact.
    type Puller: Pull<Bundle<T, D>>+'static;
    /// Allocates a matched pair of push and pull endpoints implementing the pact.
    fn connect<A: AsWorker>(self, allocator: &mut A, identifier: usize, address: &[usize], logging: Option<Logger>) -> (Self::Pusher, Self::Puller);
}

/// A direct connection
pub struct Pipeline;
impl<T: 'static, D: 'static> ParallelizationContract<T, D> for Pipeline {
    type Pusher = LogPusher<T, D, ThreadPusher<Bundle<T, D>>>;
    type Puller = LogPuller<T, D, ThreadPuller<Bundle<T, D>>>;
    fn connect<A: AsWorker>(self, allocator: &mut A, identifier: usize, address: &[usize], logging: Option<Logger>) -> (Self::Pusher, Self::Puller) {
        let (pusher, puller) = allocator.pipeline::<Message<T, D>>(identifier, address);
        // // ignore `&mut A` and use thread allocator
        // let (pusher, puller) = Thread::new::<Bundle<T, D>>();
        (LogPusher::new(pusher, allocator.index(), allocator.index(), identifier, logging.clone()),
         LogPuller::new(puller, allocator.index(), identifier, logging.clone()))
    }
}

/// An exchange between multiple observers by data
pub struct Exchange<D, F: Fn(&D)->u64+'static> { hash_func: F, phantom: PhantomData<D>, }
impl<D, F: Fn(&D)->u64> Exchange<D, F> {
    /// Allocates a new `Exchange` pact from a distribution function.
    pub fn new(func: F) -> Exchange<D, F> {
        Exchange {
            hash_func:  func,
            phantom:    PhantomData,
        }
    }
}

// Exchange uses a `Box<Pushable>` because it cannot know what type of pushable will return from the allocator.
impl<T: Eq+Data+Clone, D: Data+Clone, F: Fn(&D)->u64+'static> ParallelizationContract<T, D> for Exchange<D, F> {
    // TODO: The closure in the type prevents us from naming it.
    //       Could specialize `ExchangePusher` to a time-free version.
    type Pusher = Box<Push<Bundle<T, D>>>;
    type Puller = Box<Pull<Bundle<T, D>>>;
    fn connect<A: AsWorker>(self, allocator: &mut A, identifier: usize, address: &[usize], logging: Option<Logger>) -> (Self::Pusher, Self::Puller) {
        let (senders, receiver) = allocator.allocate::<Message<T, D>>(identifier, address);
        let senders = senders.into_iter().enumerate().map(|(i,x)| LogPusher::new(x, allocator.index(), i, identifier, logging.clone())).collect::<Vec<_>>();
        (Box::new(ExchangePusher::new(senders, move |_, d| (self.hash_func)(d))), Box::new(LogPuller::new(receiver, allocator.index(), identifier, logging.clone())))
    }
}

// /// An exchange between multiple observers by time and data
// pub struct TimeExchange<D, T, F: Fn(&T, &D)->u64+'static> { hash_func: F, phantom: PhantomData<(T, D)>, }
// impl<D, T, F: Fn(&T, &D)->u64> TimeExchange<D, T, F> {
//     /// Allocates a new `TimeExchange` pact from a distribution function.
//     pub fn new(func: F) -> TimeExchange<D, T, F> {
//         TimeExchange {
//             hash_func:  func,
//             phantom:    PhantomData,
//         }
//     }
// }

// impl<T: Eq+Data+Clone, D: Data+Clone, F: Fn(&T, &D)->u64+'static> ParallelizationContract<T, D> for TimeExchange<D, T, F> {
//     type Pusher = ExchangePusher<T, D, Pusher<T, D, Box<Push<CommMessage<Message<T, D>>>>>, F>;
//     type Puller = Puller<T, D, Box<Pull<CommMessage<Message<T, D>>>>>;
//     fn connect<A: AsWorker>(self, allocator: &mut A, identifier: usize, logging: Logger) -> (Self::Pusher, Self::Puller) {
//         let (senders, receiver, channel_id) = allocator.allocate::<Message<T, D>>();
//         let senders = senders.into_iter().enumerate().map(|(i,x)| Pusher::new(x, allocator.index(), i, identifier, channel_id, logging.clone())).collect::<Vec<_>>();
//         (ExchangePusher::new(senders, self.hash_func), Puller::new(receiver, allocator.index(), identifier, channel_id, logging.clone()))
//     }
// }

/// Wraps a `Message<T,D>` pusher to provide a `Push<(T, Content<D>)>`.
pub struct LogPusher<T, D, P: Push<Bundle<T, D>>> {
    pusher: P,
    channel: usize,
    counter: usize,
    source: usize,
    target: usize,
    phantom: ::std::marker::PhantomData<(T, D)>,
    logging: Option<Logger>,
}
impl<T, D, P: Push<Bundle<T, D>>> LogPusher<T, D, P> {
    /// Allocates a new pusher.
    pub fn new(pusher: P, source: usize, target: usize, channel: usize, logging: Option<Logger>) -> Self {
        LogPusher {
            pusher,
            channel,
            counter: 0,
            source,
            target,
            phantom: ::std::marker::PhantomData,
            logging,
        }
    }
}

impl<T, D, P: Push<Bundle<T, D>>> Push<Bundle<T, D>> for LogPusher<T, D, P> {
    #[inline(always)]
    fn push(&mut self, pair: &mut Option<Bundle<T, D>>) {
        if let Some(bundle) = pair {
            self.counter += 1;
            self.logging.as_ref().map(|l| l.log(::logging::MessagesEvent {
                is_send: true,
                channel: self.channel,
                source: self.source,
                target: self.target,
                seq_no: self.counter-1,
                length: bundle.data.len(),
            }));
        }
        self.pusher.push(pair);
    }
}

/// Wraps a `Message<T,D>` puller to provide a `Pull<(T, Content<D>)>`.
pub struct LogPuller<T, D, P: Pull<Bundle<T, D>>> {
    puller: P,
    channel: usize,
    index: usize,
    phantom: ::std::marker::PhantomData<(T, D)>,
    logging: Option<Logger>,
}
impl<T, D, P: Pull<Bundle<T, D>>> LogPuller<T, D, P> {
    /// Allocates a new `Puller`.
    pub fn new(puller: P, index: usize, channel: usize, logging: Option<Logger>) -> Self {
        LogPuller {
            puller,
            channel,
            index,
            phantom: ::std::marker::PhantomData,
            logging,
        }
    }
}

impl<T, D, P: Pull<Bundle<T, D>>> Pull<Bundle<T, D>> for LogPuller<T, D, P> {
    #[inline(always)]
    fn pull(&mut self) -> &mut Option<Bundle<T,D>> {
        let result = self.puller.pull();
        if let Some(bundle) = result {
            let channel = self.channel;
            let target = self.index;
            self.logging.as_ref().map(|l| l.log(::logging::MessagesEvent {
                is_send: false,
                channel,
                source: bundle.from,
                target,
                seq_no: bundle.seq,
                length: bundle.data.len(),
            }));
        }
        result
    }
}
<commit_msg>imprint sequence number<commit_after>//! Parallelization contracts, describing requirements for data movement along dataflow edges.
//!
//! Pacts describe how data should be exchanged between workers, and implement a method which
//! creates a pair of `Push` and `Pull` implementors from an `A: AsWorker`. These two endpoints
//! respectively distribute and collect data among workers according to the pact.
//!
//! The only requirement of a pact is that it not alter the number of `D` records at each time `T`.
//! The progress tracking logic assumes that this number is independent of the pact used.

use std::marker::PhantomData;

use communication::{Push, Pull, Data};
use communication::allocator::thread::{ThreadPusher, ThreadPuller};

use worker::AsWorker;
use dataflow::channels::pushers::Exchange as ExchangePusher;
use super::{Bundle, Message};

use logging::TimelyLogger as Logger;

/// A `ParallelizationContract` allocates paired `Push` and `Pull` implementors.
pub trait ParallelizationContract<T: 'static, D: 'static> {
    /// Type implementing `Push` produced by this pact.
    type Pusher: Push<Bundle<T, D>>+'static;
    /// Type implementing `Pull` produced by this pact.
    type Puller: Pull<Bundle<T, D>>+'static;
    /// Allocates a matched pair of push and pull endpoints implementing the pact.
    fn connect<A: AsWorker>(self, allocator: &mut A, identifier: usize, address: &[usize], logging: Option<Logger>) -> (Self::Pusher, Self::Puller);
}

/// A direct connection
pub struct Pipeline;
impl<T: 'static, D: 'static> ParallelizationContract<T, D> for Pipeline {
    type Pusher = LogPusher<T, D, ThreadPusher<Bundle<T, D>>>;
    type Puller = LogPuller<T, D, ThreadPuller<Bundle<T, D>>>;
    fn connect<A: AsWorker>(self, allocator: &mut A, identifier: usize, address: &[usize], logging: Option<Logger>) -> (Self::Pusher, Self::Puller) {
        let (pusher, puller) = allocator.pipeline::<Message<T, D>>(identifier, address);
        // // ignore `&mut A` and use thread allocator
        // let (pusher, puller) = Thread::new::<Bundle<T, D>>();
        (LogPusher::new(pusher, allocator.index(), allocator.index(), identifier, logging.clone()),
         LogPuller::new(puller, allocator.index(), identifier, logging.clone()))
    }
}

/// An exchange between multiple observers by data
pub struct Exchange<D, F: Fn(&D)->u64+'static> { hash_func: F, phantom: PhantomData<D>, }
impl<D, F: Fn(&D)->u64> Exchange<D, F> {
    /// Allocates a new `Exchange` pact from a distribution function.
    pub fn new(func: F) -> Exchange<D, F> {
        Exchange {
            hash_func:  func,
            phantom:    PhantomData,
        }
    }
}

// Exchange uses a `Box<Pushable>` because it cannot know what type of pushable will return from the allocator.
impl<T: Eq+Data+Clone, D: Data+Clone, F: Fn(&D)->u64+'static> ParallelizationContract<T, D> for Exchange<D, F> {
    // TODO: The closure in the type prevents us from naming it.
    //       Could specialize `ExchangePusher` to a time-free version.
    type Pusher = Box<Push<Bundle<T, D>>>;
    type Puller = Box<Pull<Bundle<T, D>>>;
    fn connect<A: AsWorker>(self, allocator: &mut A, identifier: usize, address: &[usize], logging: Option<Logger>) -> (Self::Pusher, Self::Puller) {
        let (senders, receiver) = allocator.allocate::<Message<T, D>>(identifier, address);
        let senders = senders.into_iter().enumerate().map(|(i,x)| LogPusher::new(x, allocator.index(), i, identifier, logging.clone())).collect::<Vec<_>>();
        (Box::new(ExchangePusher::new(senders, move |_, d| (self.hash_func)(d))), Box::new(LogPuller::new(receiver, allocator.index(), identifier, logging.clone())))
    }
}

// /// An exchange between multiple observers by time and data
// pub struct TimeExchange<D, T, F: Fn(&T, &D)->u64+'static> { hash_func: F, phantom: PhantomData<(T, D)>, }
// impl<D, T, F: Fn(&T, &D)->u64> TimeExchange<D, T, F> {
//     /// Allocates a new `TimeExchange` pact from a distribution function.
//     pub fn new(func: F) -> TimeExchange<D, T, F> {
//         TimeExchange {
//             hash_func:  func,
//             phantom:    PhantomData,
//         }
//     }
// }

// impl<T: Eq+Data+Clone, D: Data+Clone, F: Fn(&T, &D)->u64+'static> ParallelizationContract<T, D> for TimeExchange<D, T, F> {
//     type Pusher = ExchangePusher<T, D, Pusher<T, D, Box<Push<CommMessage<Message<T, D>>>>>, F>;
//     type Puller = Puller<T, D, Box<Pull<CommMessage<Message<T, D>>>>>;
//     fn connect<A: AsWorker>(self, allocator: &mut A, identifier: usize, logging: Logger) -> (Self::Pusher, Self::Puller) {
//         let (senders, receiver, channel_id) = allocator.allocate::<Message<T, D>>();
//         let senders = senders.into_iter().enumerate().map(|(i,x)| Pusher::new(x, allocator.index(), i, identifier, channel_id, logging.clone())).collect::<Vec<_>>();
//         (ExchangePusher::new(senders, self.hash_func), Puller::new(receiver, allocator.index(), identifier, channel_id, logging.clone()))
//     }
// }

/// Wraps a `Message<T,D>` pusher to provide a `Push<(T, Content<D>)>`.
pub struct LogPusher<T, D, P: Push<Bundle<T, D>>> {
    pusher: P,
    channel: usize,
    counter: usize,
    source: usize,
    target: usize,
    phantom: ::std::marker::PhantomData<(T, D)>,
    logging: Option<Logger>,
}
impl<T, D, P: Push<Bundle<T, D>>> LogPusher<T, D, P> {
    /// Allocates a new pusher.
    pub fn new(pusher: P, source: usize, target: usize, channel: usize, logging: Option<Logger>) -> Self {
        LogPusher {
            pusher,
            channel,
            counter: 0,
            source,
            target,
            phantom: ::std::marker::PhantomData,
            logging,
        }
    }
}

impl<T, D, P: Push<Bundle<T, D>>> Push<Bundle<T, D>> for LogPusher<T, D, P> {
    #[inline(always)]
    fn push(&mut self, pair: &mut Option<Bundle<T, D>>) {
        if let Some(bundle) = pair {
            self.counter += 1;
            // Stamp the sequence number and source.
            // FIXME: Awkward moment/logic.
            if let Some(message) = bundle.if_mut() {
                message.seq = self.counter-1;
                message.from = self.source;
            }

            self.logging.as_ref().map(|l| l.log(::logging::MessagesEvent {
                is_send: true,
                channel: self.channel,
                source: self.source,
                target: self.target,
                seq_no: self.counter-1,
                length: bundle.data.len(),
            }));
        }
        self.pusher.push(pair);
    }
}

/// Wraps a `Message<T,D>` puller to provide a `Pull<(T, Content<D>)>`.
pub struct LogPuller<T, D, P: Pull<Bundle<T, D>>> {
    puller: P,
    channel: usize,
    index: usize,
    phantom: ::std::marker::PhantomData<(T, D)>,
    logging: Option<Logger>,
}
impl<T, D, P: Pull<Bundle<T, D>>> LogPuller<T, D, P> {
    /// Allocates a new `Puller`.
    pub fn new(puller: P, index: usize, channel: usize, logging: Option<Logger>) -> Self {
        LogPuller {
            puller,
            channel,
            index,
            phantom: ::std::marker::PhantomData,
            logging,
        }
    }
}

impl<T, D, P: Pull<Bundle<T, D>>> Pull<Bundle<T, D>> for LogPuller<T, D, P> {
    #[inline(always)]
    fn pull(&mut self) -> &mut Option<Bundle<T,D>> {
        let result = self.puller.pull();
        if let Some(bundle) = result {
            let channel = self.channel;
            let target = self.index;
            self.logging.as_ref().map(|l| l.log(::logging::MessagesEvent {
                is_send: false,
                channel,
                source: bundle.from,
                target,
                seq_no: bundle.seq,
                length: bundle.data.len(),
            }));
        }
        result
    }
}
<|endoftext|>


<commit_before>use std::io;
use byteorder::{ByteOrder, LittleEndian};
use key;

pub struct ChaCha20 {
    state: [u32; 16],
}

pub fn gen_nonce() -> io::Result<[u8; 12]> {
    Ok(key::gen()?)
}

impl ChaCha20 {
    /// Be sure never to encrypt with a given (key, nonce) pair more than once
    pub fn new(key: &[u8; 32], nonce: &[u8; 12]) -> Self {
        let mut chacha20 = Self { state: [0; 16] };
        Self::setup_state(&mut chacha20.state, &key, &nonce);
        chacha20
    }

    pub fn get_block(&self, counter: u32) -> [u8; 64] {
        Self::serialize_block(self.block(counter))
    }

    fn block(&self, counter: u32) -> [u32; 16] {
        let mut state = [0; 16];
        state.copy_from_slice(&self.state);
        state[12] = counter;
        for _ in 0..10 {
            Self::inner_block(&mut state);
        }
        for (x, &y) in state.iter_mut().zip(&self.state) {
            *x = x.wrapping_add(y);
        }
        state[12] = state[12].wrapping_add(counter);
        state
    }

    fn setup_state(state: &mut [u32; 16], key: &[u8; 32], nonce: &[u8; 12]) {
        state[0] = 0x61707865;
        state[1] = 0x3320646e;
        state[2] = 0x79622d32;
        state[3] = 0x6b206574;
        state[4..12].copy_from_slice(&Self::convert_key(&key));
        state[13..].copy_from_slice(&Self::convert_nonce(&nonce));
    }

    fn inner_block(state: &mut [u32; 16]) {
        Self::quarter_round(state, 0, 4, 8, 12);
        Self::quarter_round(state, 1, 5, 9, 13);
        Self::quarter_round(state, 2, 6, 10, 14);
        Self::quarter_round(state, 3, 7, 11, 15);
        Self::quarter_round(state, 0, 5, 10, 15);
        Self::quarter_round(state, 1, 6, 11, 12);
        Self::quarter_round(state, 2, 7, 8, 13);
        Self::quarter_round(state, 3, 4, 9, 14);
    }

    fn quarter_round(state: &mut [u32; 16], i: usize, j: usize, k: usize, l: usize) {
        let mut a = state[i];
        let mut b = state[j];
        let mut c = state[k];
        let mut d = state[l];
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(16);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(12);
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(8);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(7);
        state[i] = a;
        state[j] = b;
        state[k] = c;
        state[l] = d;
    }

    fn convert_key(key: &[u8; 32]) -> [u32; 8] {
        let mut ret = [0; 8];
        for (byte, chunk) in ret.iter_mut().zip(key.chunks(4)) {
            *byte = LittleEndian::read_u32(chunk);
        }
        ret
    }

    fn convert_nonce(nonce: &[u8; 12]) -> [u32; 3] {
        let mut ret = [0; 3];
        for (byte, chunk) in ret.iter_mut().zip(nonce.chunks(4)) {
            *byte = LittleEndian::read_u32(chunk);
        }
        ret
    }

    fn serialize_block(block: [u32; 16]) -> [u8; 64] {
        let mut ret = [0; 64];
        for (chunk, &byte) in ret.chunks_mut(4).zip(&block) {
            LittleEndian::write_u32(chunk, byte);
        }
        ret
    }
}

#[cfg(test)]
mod tests {
    use chacha20::*;

    const KEY: [u8; 32] = [
        0x00,
        0x01,
        0x02,
        0x03,
        0x04,
        0x05,
        0x06,
        0x07,
        0x08,
        0x09,
        0x0a,
        0x0b,
        0x0c,
        0x0d,
        0x0e,
        0x0f,
        0x10,
        0x11,
        0x12,
        0x13,
        0x14,
        0x15,
        0x16,
        0x17,
        0x18,
        0x19,
        0x1a,
        0x1b,
        0x1c,
        0x1d,
        0x1e,
        0x1f,
    ];
    const NONCE: [u8; 12] = [
        0x00,
        0x00,
        0x00,
        0x09,
        0x00,
        0x00,
        0x00,
        0x4a,
        0x00,
        0x00,
        0x00,
        0x00,
    ];
    const BLOCK_ONE: [u8; 64] = [
        0x10,
        0xf1,
        0xe7,
        0xe4,
        0xd1,
        0x3b,
        0x59,
        0x15,
        0x50,
        0x0f,
        0xdd,
        0x1f,
        0xa3,
        0x20,
        0x71,
        0xc4,
        0xc7,
        0xd1,
        0xf4,
        0xc7,
        0x33,
        0xc0,
        0x68,
        0x03,
        0x04,
        0x22,
        0xaa,
        0x9a,
        0xc3,
        0xd4,
        0x6c,
        0x4e,
        0xd2,
        0x82,
        0x64,
        0x46,
        0x07,
        0x9f,
        0xaa,
        0x09,
        0x14,
        0xc2,
        0xd7,
        0x05,
        0xd9,
        0x8b,
        0x02,
        0xa2,
        0xb5,
        0x12,
        0x9c,
        0xd1,
        0xde,
        0x16,
        0x4e,
        0xb9,
        0xcb,
        0xd0,
        0x83,
        0xe8,
        0xa2,
        0x50,
        0x3c,
        0x4e,
    ];
    const SETUP_STATE: [u32; 16] = [
        0x61707865,
        0x3320646e,
        0x79622d32,
        0x6b206574,
        0x03020100,
        0x07060504,
        0x0b0a0908,
        0x0f0e0d0c,
        0x13121110,
        0x17161514,
        0x1b1a1918,
        0x1f1e1d1c,
        0x00000000,
        0x09000000,
        0x4a000000,
        0x00000000,
    ];
    const FINAL_STATE: [u32; 16] = [
        0xe4e7f110,
        0x15593bd1,
        0x1fdd0f50,
        0xc47120a3,
        0xc7f4d1c7,
        0x0368c033,
        0x9aaa2204,
        0x4e6cd4c3,
        0x466482d2,
        0x09aa9f07,
        0x05d7c214,
        0xa2028bd9,
        0xd19c12b5,
        0xb94e16de,
        0xe883d0cb,
        0x4e3c50a2,
    ];

    #[test]
    fn test_new() {
        let chacha20 = ChaCha20::new(&KEY, &NONCE);
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &NONCE);
        assert_eq!(state, chacha20.state);
    }

    fn check_serialized_block(block: &[u8]) {
        assert_eq!(BLOCK_ONE.len(), block.len());
        for (lhs, rhs) in BLOCK_ONE.iter().zip(block) {
            assert_eq!(lhs, rhs);
        }
    }

    #[test]
    fn test_get_block() {
        let mut chacha20 = ChaCha20 { state: [0; 16] };
        ChaCha20::setup_state(&mut chacha20.state, &KEY, &NONCE);
        check_serialized_block(&chacha20.get_block(1));
    }

    #[test]
    fn test_block() {
        let mut chacha20 = ChaCha20 { state: [0; 16] };
        ChaCha20::setup_state(&mut chacha20.state, &KEY, &NONCE);
        assert_eq!(FINAL_STATE, chacha20.block(1));
    }

    #[test]
    fn test_setup_state() {
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &NONCE);
        assert_eq!(SETUP_STATE, state);
    }

    #[test]
    fn test_inner_block() {
        let mut state = SETUP_STATE;
        state[12] = 1;
        for _ in 0..10 {
            ChaCha20::inner_block(&mut state);
        }
        assert_eq!(
            [
                0x837778ab,
                0xe238d763,
                0xa67ae21e,
                0x5950bb2f,
                0xc4f2d0c7,
                0xfc62bb2f,
                0x8fa018fc,
                0x3f5ec7b7,
                0x335271c2,
                0xf29489f3,
                0xeabda8fc,
                0x82e46ebd,
                0xd19c12b4,
                0xb04e16de,
                0x9e83d0cb,
                0x4e3c50a2,
            ],
            state
        );
    }

    #[test]
    fn test_quarter_round() {
        let mut state = [0; 16];
        state[0] = 0x11111111;
        state[1] = 0x01020304;
        state[2] = 0x9b8d6f43;
        state[3] = 0x01234567;
        ChaCha20::quarter_round(&mut state, 0, 1, 2, 3);
        assert_eq!([0xea2a92f4, 0xcb1cf8ce, 0x4581472e, 0x5881c4bb], state[..4]);

        state = [
            0x879531e0,
            0xc5ecf37d,
            0x516461b1,
            0xc9a62f8a,
            0x44c20ef3,
            0x3390af7f,
            0xd9fc690b,
            0x2a5f714c,
            0x53372767,
            0xb00a5631,
            0x974c541a,
            0x359e9963,
            0x5c971061,
            0x3d631689,
            0x2098d9d6,
            0x91dbd320,
        ];
        ChaCha20::quarter_round(&mut state, 2, 7, 8, 13);
        assert_eq!(
            [
                0x879531e0,
                0xc5ecf37d,
                0xbdb886dc,
                0xc9a62f8a,
                0x44c20ef3,
                0x3390af7f,
                0xd9fc690b,
                0xcfacafd2,
                0xe46bea80,
                0xb00a5631,
                0x974c541a,
                0x359e9963,
                0x5c971061,
                0xccc07c79,
                0x2098d9d6,
                0x91dbd320,
            ],
            state
        );
    }

    #[test]
    fn test_convert_key() {
        assert_eq!(
            [
                0x03020100,
                0x07060504,
                0x0b0a0908,
                0x0f0e0d0c,
                0x13121110,
                0x17161514,
                0x1b1a1918,
                0x1f1e1d1c,
            ],
            ChaCha20::convert_key(&KEY)
        );
    }

    #[test]
    fn test_convert_nonce() {
        assert_eq!(
            [0x09000000, 0x4a000000, 0x00000000],
            ChaCha20::convert_nonce(&NONCE)
        );
    }

    #[test]
    fn test_serialize_block() {
        check_serialized_block(&ChaCha20::serialize_block(FINAL_STATE));
    }
}
<commit_msg>chacha20: refactor little-endian and block functions<commit_after>use std::io;
use byteorder::{ByteOrder, LittleEndian};
use key;

pub struct ChaCha20 {
    state: [u32; 16],
}

pub fn gen_nonce() -> io::Result<[u8; 12]> {
    Ok(key::gen()?)
}

impl ChaCha20 {
    /// Be sure never to encrypt with a given (key, nonce) pair more than once
    pub fn new(key: &[u8; 32], nonce: &[u8; 12]) -> Self {
        let mut chacha20 = Self { state: [0; 16] };
        Self::setup_state(&mut chacha20.state, &key, &nonce);
        chacha20
    }

    pub fn get_block(&self, counter: u32) -> [u8; 64] {
        let mut state = [0; 16];
        self.transform_state(&mut state, counter);
        Self::serialize_block(state)
    }

    fn setup_state(state: &mut [u32; 16], key: &[u8; 32], nonce: &[u8; 12]) {
        state[0] = 0x61707865;
        state[1] = 0x3320646e;
        state[2] = 0x79622d32;
        state[3] = 0x6b206574;
        Self::to_le(&mut state[4..12], key);
        Self::to_le(&mut state[13..], nonce);
    }

    fn transform_state(&self, state: &mut [u32; 16], counter: u32) {
        state.copy_from_slice(&self.state);
        state[12] = counter;
        for _ in 0..10 {
            Self::inner_block(state);
        }
        for (x, &y) in state.iter_mut().zip(&self.state) {
            *x = x.wrapping_add(y);
        }
        state[12] = state[12].wrapping_add(counter);
    }

    fn inner_block(state: &mut [u32; 16]) {
        Self::quarter_round(state, 0, 4, 8, 12);
        Self::quarter_round(state, 1, 5, 9, 13);
        Self::quarter_round(state, 2, 6, 10, 14);
        Self::quarter_round(state, 3, 7, 11, 15);
        Self::quarter_round(state, 0, 5, 10, 15);
        Self::quarter_round(state, 1, 6, 11, 12);
        Self::quarter_round(state, 2, 7, 8, 13);
        Self::quarter_round(state, 3, 4, 9, 14);
    }

    fn quarter_round(state: &mut [u32; 16], i: usize, j: usize, k: usize, l: usize) {
        let mut a = state[i];
        let mut b = state[j];
        let mut c = state[k];
        let mut d = state[l];
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(16);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(12);
        a = a.wrapping_add(b);
        d ^= a;
        d = d.rotate_left(8);
        c = c.wrapping_add(d);
        b ^= c;
        b = b.rotate_left(7);
        state[i] = a;
        state[j] = b;
        state[k] = c;
        state[l] = d;
    }

    fn to_le(dest: &mut [u32], src: &[u8]) {
        for (byte, chunk) in dest.iter_mut().zip(src.chunks(4)) {
            *byte = LittleEndian::read_u32(chunk);
        }
    }

    fn serialize_block(block: [u32; 16]) -> [u8; 64] {
        let mut ret = [0; 64];
        for (chunk, &byte) in ret.chunks_mut(4).zip(&block) {
            LittleEndian::write_u32(chunk, byte);
        }
        ret
    }
}

#[cfg(test)]
mod tests {
    use chacha20::*;

    const KEY: [u8; 32] = [
        0x00,
        0x01,
        0x02,
        0x03,
        0x04,
        0x05,
        0x06,
        0x07,
        0x08,
        0x09,
        0x0a,
        0x0b,
        0x0c,
        0x0d,
        0x0e,
        0x0f,
        0x10,
        0x11,
        0x12,
        0x13,
        0x14,
        0x15,
        0x16,
        0x17,
        0x18,
        0x19,
        0x1a,
        0x1b,
        0x1c,
        0x1d,
        0x1e,
        0x1f,
    ];
    const NONCE: [u8; 12] = [
        0x00,
        0x00,
        0x00,
        0x09,
        0x00,
        0x00,
        0x00,
        0x4a,
        0x00,
        0x00,
        0x00,
        0x00,
    ];
    const BLOCK_ONE: [u8; 64] = [
        0x10,
        0xf1,
        0xe7,
        0xe4,
        0xd1,
        0x3b,
        0x59,
        0x15,
        0x50,
        0x0f,
        0xdd,
        0x1f,
        0xa3,
        0x20,
        0x71,
        0xc4,
        0xc7,
        0xd1,
        0xf4,
        0xc7,
        0x33,
        0xc0,
        0x68,
        0x03,
        0x04,
        0x22,
        0xaa,
        0x9a,
        0xc3,
        0xd4,
        0x6c,
        0x4e,
        0xd2,
        0x82,
        0x64,
        0x46,
        0x07,
        0x9f,
        0xaa,
        0x09,
        0x14,
        0xc2,
        0xd7,
        0x05,
        0xd9,
        0x8b,
        0x02,
        0xa2,
        0xb5,
        0x12,
        0x9c,
        0xd1,
        0xde,
        0x16,
        0x4e,
        0xb9,
        0xcb,
        0xd0,
        0x83,
        0xe8,
        0xa2,
        0x50,
        0x3c,
        0x4e,
    ];
    const SETUP_STATE: [u32; 16] = [
        0x61707865,
        0x3320646e,
        0x79622d32,
        0x6b206574,
        0x03020100,
        0x07060504,
        0x0b0a0908,
        0x0f0e0d0c,
        0x13121110,
        0x17161514,
        0x1b1a1918,
        0x1f1e1d1c,
        0x00000000,
        0x09000000,
        0x4a000000,
        0x00000000,
    ];
    const FINAL_STATE: [u32; 16] = [
        0xe4e7f110,
        0x15593bd1,
        0x1fdd0f50,
        0xc47120a3,
        0xc7f4d1c7,
        0x0368c033,
        0x9aaa2204,
        0x4e6cd4c3,
        0x466482d2,
        0x09aa9f07,
        0x05d7c214,
        0xa2028bd9,
        0xd19c12b5,
        0xb94e16de,
        0xe883d0cb,
        0x4e3c50a2,
    ];

    #[test]
    fn test_new() {
        let chacha20 = ChaCha20::new(&KEY, &NONCE);
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &NONCE);
        assert_eq!(state, chacha20.state);
    }

    fn check_serialized_block(block: &[u8]) {
        assert_eq!(BLOCK_ONE.len(), block.len());
        for (lhs, rhs) in BLOCK_ONE.iter().zip(block) {
            assert_eq!(lhs, rhs);
        }
    }

    #[test]
    fn test_get_block() {
        let mut chacha20 = ChaCha20 { state: [0; 16] };
        ChaCha20::setup_state(&mut chacha20.state, &KEY, &NONCE);
        check_serialized_block(&chacha20.get_block(1));
    }

    #[test]
    fn test_setup_state() {
        let mut state = [0; 16];
        ChaCha20::setup_state(&mut state, &KEY, &NONCE);
        assert_eq!(SETUP_STATE, state);
    }

    #[test]
    fn test_transform_state() {
        let chacha20 = ChaCha20 { state: SETUP_STATE };
        let mut state = [0; 16];
        chacha20.transform_state(&mut state, 1);
        assert_eq!(FINAL_STATE, state);
    }

    #[test]
    fn test_inner_block() {
        let mut state = SETUP_STATE;
        state[12] = 1;
        for _ in 0..10 {
            ChaCha20::inner_block(&mut state);
        }
        assert_eq!(
            [
                0x837778ab,
                0xe238d763,
                0xa67ae21e,
                0x5950bb2f,
                0xc4f2d0c7,
                0xfc62bb2f,
                0x8fa018fc,
                0x3f5ec7b7,
                0x335271c2,
                0xf29489f3,
                0xeabda8fc,
                0x82e46ebd,
                0xd19c12b4,
                0xb04e16de,
                0x9e83d0cb,
                0x4e3c50a2,
            ],
            state
        );
    }

    #[test]
    fn test_quarter_round() {
        let mut state = [0; 16];
        state[0] = 0x11111111;
        state[1] = 0x01020304;
        state[2] = 0x9b8d6f43;
        state[3] = 0x01234567;
        ChaCha20::quarter_round(&mut state, 0, 1, 2, 3);
        assert_eq!([0xea2a92f4, 0xcb1cf8ce, 0x4581472e, 0x5881c4bb], state[..4]);

        state = [
            0x879531e0,
            0xc5ecf37d,
            0x516461b1,
            0xc9a62f8a,
            0x44c20ef3,
            0x3390af7f,
            0xd9fc690b,
            0x2a5f714c,
            0x53372767,
            0xb00a5631,
            0x974c541a,
            0x359e9963,
            0x5c971061,
            0x3d631689,
            0x2098d9d6,
            0x91dbd320,
        ];
        ChaCha20::quarter_round(&mut state, 2, 7, 8, 13);
        assert_eq!(
            [
                0x879531e0,
                0xc5ecf37d,
                0xbdb886dc,
                0xc9a62f8a,
                0x44c20ef3,
                0x3390af7f,
                0xd9fc690b,
                0xcfacafd2,
                0xe46bea80,
                0xb00a5631,
                0x974c541a,
                0x359e9963,
                0x5c971061,
                0xccc07c79,
                0x2098d9d6,
                0x91dbd320,
            ],
            state
        );
    }

    #[test]
    fn test_to_le() {
        let mut key = [0; 8];
        ChaCha20::to_le(&mut key, &KEY);
        assert_eq!(
            [
                0x03020100,
                0x07060504,
                0x0b0a0908,
                0x0f0e0d0c,
                0x13121110,
                0x17161514,
                0x1b1a1918,
                0x1f1e1d1c,
            ],
            key
        );

        let mut nonce = [0; 3];
        ChaCha20::to_le(&mut nonce, &NONCE);
        assert_eq!([0x09000000, 0x4a000000, 0x00000000], nonce);
    }

    #[test]
    fn test_serialize_block() {
        check_serialized_block(&ChaCha20::serialize_block(FINAL_STATE));
    }
}
<|endoftext|>


<commit_before>use super::{InvalidDatabaseError, Reader, Decoder, IoError};
use std::io::net::ip::IpAddr;
use std::from_str::FromStr;
use serialize::Decodable;

#[test]
fn test_decoder() {

    #[deriving(Decodable, Show, Eq)]
    struct MapXType {
        arrayX: ~[uint],
        utf8_stringX: ~str
    };

    #[deriving(Decodable, Show, Eq)]
    struct MapType {
        mapX: MapXType
    };

    #[deriving(Decodable, Show)]
    struct TestType {
        array:       ~[uint],
        boolean:     bool,
        bytes:       ~[u8],
        double:      f64,
        float:       f32,
        int32:       i32,
        map:         MapType,
        uint16:      u16,
        uint32:      u32,
        uint64:      u64,
        uint128:     ~[u8],
        utf8_string: ~str
    }

    let r = Reader::open("test-data/test-data/MaxMind-DB-test-decoder.mmdb").unwrap();
    let ip: IpAddr = FromStr::from_str("::1.1.1.0").unwrap();
    let raw_data = r.lookup(ip);

    let mut decoder = Decoder::new(raw_data.unwrap());
    let result: TestType = match Decodable::decode(&mut decoder) {
        Ok(v) => v,
        Err(e) => fail!("Decoding error: {}", e)
    };

    assert_eq!(result.array, ~[ 1u, 2u, 3u ]);
    assert_eq!(result.boolean, true);
    assert_eq!(result.bytes, ~[0u8, 0u8, 0u8, 42u8])
    assert_eq!(result.double, 42.123456);
    assert_eq!(result.float, 1.1);
    assert_eq!(result.int32, -268435456);

    assert_eq!(result.map, MapType{ mapX: MapXType{ arrayX: ~[7,8,9], utf8_stringX: "hello".to_owned()}});

    assert_eq!(result.uint16, 100);
    assert_eq!(result.uint32, 268435456);
    assert_eq!(result.uint64, 1152921504606846976);
    assert_eq!(result.uint128, ~[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]);

    assert_eq!(result.utf8_string,  "unicode! ☯ - ♫".to_owned());
}

#[test]
fn test_broken_database() {
    let r = Reader::open("test-data/test-data/GeoIP2-City-Test-Broken-Double-Format.mmdb").unwrap();
    let ip: IpAddr = FromStr::from_str("2001:220::").unwrap();
    let result = r.lookup(ip);
    assert_eq!(result, Err(InvalidDatabaseError("double of size 2".to_owned())));
}

#[test]
fn test_missing_database() {
    let r = Reader::open("file-does-not-exist.mmdb");
    match r {
        Ok(_) => fail!("Received Reader when opening non-existent file"),
        Err(IoError(_)) => assert!(true),
        Err(_) => assert!(false)
    }
}

#[test]
fn test_non_database() {
    let r = Reader::open("README.md");
    match r {
        Ok(_) => fail!("Received Reader when opening a non-MMDB file"),
        Err(e) => assert_eq!(e, InvalidDatabaseError("Could not find MaxMind DB metadata in file.".to_owned()))

    }
}

<commit_msg>Added metadata check<commit_after>use super::{InvalidDatabaseError, Reader, Decoder, IoError};
use std::io::net::ip::IpAddr;
use std::from_str::FromStr;
use serialize::Decodable;

#[test]
fn test_decoder() {

    #[deriving(Decodable, Show, Eq)]
    struct MapXType {
        arrayX: ~[uint],
        utf8_stringX: ~str
    };

    #[deriving(Decodable, Show, Eq)]
    struct MapType {
        mapX: MapXType
    };

    #[deriving(Decodable, Show)]
    struct TestType {
        array:       ~[uint],
        boolean:     bool,
        bytes:       ~[u8],
        double:      f64,
        float:       f32,
        int32:       i32,
        map:         MapType,
        uint16:      u16,
        uint32:      u32,
        uint64:      u64,
        uint128:     ~[u8],
        utf8_string: ~str
    }

    let r = Reader::open("test-data/test-data/MaxMind-DB-test-decoder.mmdb").unwrap();
    let ip: IpAddr = FromStr::from_str("::1.1.1.0").unwrap();
    let raw_data = r.lookup(ip);

    let mut decoder = Decoder::new(raw_data.unwrap());
    let result: TestType = match Decodable::decode(&mut decoder) {
        Ok(v) => v,
        Err(e) => fail!("Decoding error: {}", e)
    };

    assert_eq!(result.array, ~[ 1u, 2u, 3u ]);
    assert_eq!(result.boolean, true);
    assert_eq!(result.bytes, ~[0u8, 0u8, 0u8, 42u8])
    assert_eq!(result.double, 42.123456);
    assert_eq!(result.float, 1.1);
    assert_eq!(result.int32, -268435456);

    assert_eq!(result.map, MapType{ mapX: MapXType{ arrayX: ~[7,8,9], utf8_stringX: "hello".to_owned()}});

    assert_eq!(result.uint16, 100);
    assert_eq!(result.uint32, 268435456);
    assert_eq!(result.uint64, 1152921504606846976);
    assert_eq!(result.uint128, ~[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]);

    assert_eq!(result.utf8_string,  "unicode! ☯ - ♫".to_owned());
}

#[test]
fn test_broken_database() {
    let r = Reader::open("test-data/test-data/GeoIP2-City-Test-Broken-Double-Format.mmdb").unwrap();
    let ip: IpAddr = FromStr::from_str("2001:220::").unwrap();
    let result = r.lookup(ip);
    assert_eq!(result, Err(InvalidDatabaseError("double of size 2".to_owned())));
}

#[test]
fn test_missing_database() {
    let r = Reader::open("file-does-not-exist.mmdb");
    match r {
        Ok(_) => fail!("Received Reader when opening non-existent file"),
        Err(IoError(_)) => assert!(true),
        Err(_) => assert!(false)
    }
}

#[test]
fn test_non_database() {
    let r = Reader::open("README.md");
    match r {
        Ok(_) => fail!("Received Reader when opening a non-MMDB file"),
        Err(e) => assert_eq!(e, InvalidDatabaseError("Could not find MaxMind DB metadata in file.".to_owned()))

    }
}

#[test]
fn test_reader() {
    let sizes = [24u, 28, 32];
    for record_size in sizes.iter() {
        let versions = [4u, 6];
        for ip_version in versions.iter() {
            let filename = format!("test-data/test-data/MaxMind-DB-test-ipv{}-{}.mmdb",
             ip_version, record_size);
            let reader = Reader::open(filename).unwrap();

            check_metadata(reader, *ip_version, *record_size);

            // if ip_version == 4 {
            //     check_ip_v4(reader)
            // } else {
            //     check_ip_v6(reader)
            // }
        }
    }
}

fn check_metadata(reader: Reader, ip_version: uint, record_size: uint) {
    let metadata = reader.metadata;

    assert_eq!(metadata.binary_format_major_version,  2u16)

    assert_eq!(metadata.binary_format_minor_version,  0u16)
    assert!(metadata.build_epoch >= 1397457605)
    assert_eq!(metadata.database_type, "Test".to_owned())

    assert_eq!(*metadata.description.find(&"en".to_owned()).unwrap(),
                   "Test Database".to_owned());
    assert_eq!(*metadata.description.find(&"zh".to_owned()).unwrap(),
                   "Test Database Chinese".to_owned());

    assert_eq!(metadata.ip_version,  ip_version as u16)
    assert_eq!(metadata.languages, ~["en".to_owned(), "zh".to_owned()])

    if ip_version == 4 {
        assert_eq!(metadata.node_count,  37)
    } else {
        assert_eq!(metadata.node_count,  160)
    }

    assert_eq!(metadata.record_size,  record_size as u16)
}
<|endoftext|>